Merge pull request #2797 from AayushSabharwal/as/inplace-observed

feat: support inplace observed

Author: ChrisRackauckas

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,31 @@ 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))
+    # Need to keep old method of building the function since it uses `output_type`,
+    # which can't be provided to `build_function`
+    oop_fn = Func(args, [],
+                 pre(Let(obsexprs,
+                     isscalar ? ts[1] : MakeArray(ts, output_type),
+                     false))) |> wrap_array_vars(sys, ts)[1] |> toexpr
+    oop_fn = expression ? oop_fn : drop_expr(@RuntimeGeneratedFunction(oop_fn))
+
+    if !isscalar
+        iip_fn = build_function(ts,
+            args...;
+            postprocess_fbody = pre,
+            wrap_code = wrap_array_vars(sys, ts) .∘ wrap_assignments(isscalar, obsexprs),
+            expression = Val{expression})[2]
+    end
+    if isscalar || !return_inplace
+        return oop_fn
+    else
+        return oop_fn, iip_fn
+    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)[2]
+    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