feat: support inplace observed

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -201,6 +201,17 @@ function generate_custom_function(sys::AbstractSystem, exprs, dvs = unknowns(sys
     end
 end
 
+function wrap_assignments(isscalar, assignments; let_block = false)
+    function wrapper(expr)
+        Func(expr.args, [], Let(assignments, expr.body, let_block))
+    end
+    if isscalar
+        wrapper
+    else
+        wrapper, wrapper
+    end
+end
+
 function wrap_array_vars(sys::AbstractSystem, exprs; dvs = unknowns(sys))
     isscalar = !(exprs isa AbstractArray)
     array_vars = Dict{Any, AbstractArray{Int}}()

src/systems/diffeqs/odesystem.jl

@@ -382,6 +382,7 @@ function build_explicit_observed_function(sys, ts;
         checkbounds = true,
         drop_expr = drop_expr,
         ps = full_parameters(sys),
+        return_inplace = false,
         op = Operator,
         throw = true)
     if (isscalar = symbolic_type(ts) !== NotSymbolic())
@@ -479,16 +480,21 @@ function build_explicit_observed_function(sys, ts;
     if inputs === nothing
         args = [dvs, ps..., ivs...]
     else
-        ipts = DestructuredArgs(inputs, inbounds = !checkbounds)
+        ipts = DestructuredArgs(unwrap.(inputs), inbounds = !checkbounds)
         args = [dvs, ipts, ps..., ivs...]
     end
     pre = get_postprocess_fbody(sys)
-
-    ex = Func(args, [],
-             pre(Let(obsexprs,
-                 isscalar ? ts[1] : MakeArray(ts, output_type),
-                 false))) |> wrap_array_vars(sys, ts)[1] |> toexpr
-    expression ? ex : drop_expr(@RuntimeGeneratedFunction(ex))
+    res = build_function(isscalar ? ts[1] : ts,
+        args...;
+        postprocess_fbody = pre,
+        wrap_code = wrap_array_vars(
+            sys, isscalar ? ts[1] : ts) .∘ wrap_assignments(isscalar, obsexprs),
+        expression = Val{expression})
+    if isscalar || return_inplace
+        return res
+    else
+        return res[1]
+    end
 end
 
 function _eq_unordered(a, b)

test/odesystem.jl

@@ -1182,3 +1182,15 @@ end
 
     @test_nowarn ForwardDiff.derivative(P -> x_at_1(P), 1.0)
 end
+
+@testset "Inplace observed functions" begin
+    @parameters P
+    @variables x(t)
+    sys = structural_simplify(ODESystem([D(x) ~ P], t, [x], [P]; name = :sys))
+    obsfn = ModelingToolkit.build_explicit_observed_function(
+        sys, [x + 1, x + P, x + t], return_inplace = true)
+    ps = ModelingToolkit.MTKParameters(sys, [P => 2.0])
+    buffer = zeros(3)
+    @test_nowarn obsfn(buffer, [1.0], ps..., 3.0)
+    @test buffer ≈ [2.0, 3.0, 4.0]
+end