fix: fix in-place observed function generation

SciML · Oct 24, 2024 · 31d8273 · 31d8273
1 parent 9b439b7
commit 31d8273
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
@@ -576,8 +576,8 @@ function build_explicit_observed_function(sys, ts;
         iip_fn = build_function(ts,
             args...;
             postprocess_fbody = pre,
-            wrap_code = array_wrapper .∘ wrap_assignments(isscalar, obsexprs) .∘
-                        mtkparams_wrapper,
+            wrap_code = mtkparams_wrapper .∘ array_wrapper .∘
+                        wrap_assignments(isscalar, obsexprs),
             expression = Val{true})[2]
         if !expression
             iip_fn = eval_or_rgf(iip_fn; eval_expression, eval_module)

diff --git a/test/odesystem.jl b/test/odesystem.jl
@@ -1447,3 +1447,15 @@ end
     @parameters p
     @test_nowarn ODESystem(Equation[], t; parameter_dependencies = [p ~ 1.0], name = :a)
 end
+
+@testset "Inplace observed" begin
+    @variables x(t)
+    @parameters p[1:2] q
+    @mtkbuild sys = ODESystem(D(x) ~ sum(p) * x + q * t, t)
+    prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0), [p => ones(2), q => 2])
+    obsfn = ModelingToolkit.build_explicit_observed_function(
+        sys, [p..., q], return_inplace = true)[2]
+    buf = zeros(3)
+    obsfn(buf, prob.u0, prob.p, 0.0)
+    @test buf ≈ [1.0, 1.0, 2.0]
+end