Add observed function for NonlinearFunctions

Author: YingboMa

src/systems/diffeqs/odesystem.jl

@@ -234,7 +234,8 @@ i.e. there are no cycles.
 function build_explicit_observed_function(
         sys, syms;
         expression=false,
-        output_type=Array)
+        output_type=Array,
+    )
 
     if (isscalar = !(syms isa Vector))
         syms = [syms]
@@ -252,13 +253,13 @@ function build_explicit_observed_function(
         output[i] = obs[idx].rhs
     end
 
+    dvs = DestructuredArgs(states(sys))
+    ps = DestructuredArgs(parameters(sys))
+    iv = independent_variable(sys)
+    args = iv === nothing ? [dvs, ps] : [dv, ps, iv]
+
     ex = Func(
-        [
-         DestructuredArgs(states(sys))
-         DestructuredArgs(parameters(sys))
-         independent_variable(sys)
-        ],
-        [],
+        args, [],
         Let(
             map(eq -> eq.lhs←eq.rhs, obs[1:maxidx]),
             isscalar ? output[1] : MakeArray(output, output_type)

src/systems/nonlinear/nonlinearsystem.jl

@@ -143,10 +143,19 @@ function DiffEqBase.NonlinearFunction{iip}(sys::NonlinearSystem, dvs = states(sy
         _jac = nothing
     end
 
+    observedfun = let sys = sys, dict = Dict()
+        function generated_observed(obsvar, u, p)
+            obs = get!(dict, value(obsvar)) do
+                build_explicit_observed_function(sys, obsvar)
+            end
+            obs(u, p)
+        end
+    end
+
     NonlinearFunction{iip}(f,
                      jac = _jac === nothing ? nothing : _jac,
                      jac_prototype = sparse ? similar(sys.jac[],Float64) : nothing,
-                     syms = Symbol.(states(sys)))
+                     syms = Symbol.(states(sys)), observed = observedfun)
 end
 
 """

test/reduction.jl

@@ -200,6 +200,8 @@ nlprob = NonlinearProblem(reducedsys, u0, pp)
 reducedsol = solve(nlprob, NewtonRaphson())
 residual = fill(100.0, length(states(reducedsys)))
 nlprob.f(residual, reducedsol.u, pp)
+@test hypot(nlprob.f.observed(u2, reducedsol.u, pp), nlprob.f.observed(u1, reducedsol.u, pp)) * pp ≈ reducedsol.u atol=1e-9
+
 @test all(x->abs(x) < 1e-5, residual)
 
 N = 5