Merge pull request #926 from SciML/myb/obs

Add observed function for `NonlinearFunction`s

Author: YingboMa

src/systems/diffeqs/abstractodesystem.jl

@@ -158,6 +158,7 @@ function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
                                      eval_expression = true,
                                      sparse = false, simplify=false,
                                      eval_module = @__MODULE__,
+                                     steady_state = false,
                                      checkbounds=false,
                                      kwargs...) where {iip}
 
@@ -194,12 +195,23 @@ function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
 
     _M = (u0 === nothing || M == I) ? M : ArrayInterface.restructure(u0 .* u0',M)
 
-    observedfun = let sys = sys, dict = Dict()
-        function generated_observed(obsvar, u, p, t)
-            obs = get!(dict, value(obsvar)) do
-                build_explicit_observed_function(sys, obsvar; checkbounds=checkbounds)
+    observedfun = if steady_state
+        let sys = sys, dict = Dict()
+            function generated_observed(obsvar, u, p, t=Inf)
+                obs = get!(dict, value(obsvar)) do
+                    build_explicit_observed_function(sys, obsvar)
+                end
+                obs(u, p, t)
+            end
+        end
+    else
+        let sys = sys, dict = Dict()
+            function generated_observed(obsvar, u, p, t)
+                obs = get!(dict, value(obsvar)) do
+                    build_explicit_observed_function(sys, obsvar; checkbounds=checkbounds)
+                end
+                obs(u, p, t)
             end
-            obs(u, p, t)
         end
     end
 
@@ -301,9 +313,30 @@ function ODEFunctionExpr{iip}(sys::AbstractODESystem, dvs = states(sys),
                                      jac = false,
                                      linenumbers = false,
                                      sparse = false, simplify=false,
+                                     steady_state = false,
                                      kwargs...) where {iip}
 
     f_oop, f_iip = generate_function(sys, dvs, ps; expression=Val{true}, kwargs...)
+
+    dict = Dict()
+    #=
+    observedfun = if steady_state
+        :(function generated_observed(obsvar, u, p, t=Inf)
+              obs = get!($dict, value(obsvar)) do
+                  build_explicit_observed_function($sys, obsvar)
+              end
+              obs(u, p, t)
+          end)
+    else
+        :(function generated_observed(obsvar, u, p, t)
+              obs = get!($dict, value(obsvar)) do
+                  build_explicit_observed_function($sys, obsvar)
+              end
+              obs(u, p, t)
+          end)
+    end
+    =#
+
     fsym = gensym(:f)
     _f = :($fsym = ModelingToolkit.ODEFunctionClosure($f_oop, $f_iip))
     tgradsym = gensym(:tgrad)
@@ -579,7 +612,7 @@ end
 
 """
 ```julia
-function DiffEqBase.SteadyStateProblem(sys::AbstractODESystem,u0map,tspan,
+function DiffEqBase.SteadyStateProblem(sys::AbstractODESystem,u0map,
                                     parammap=DiffEqBase.NullParameters();
                                     version = nothing, tgrad=false,
                                     jac = false,
@@ -593,13 +626,13 @@ symbolically calculating numerical enhancements.
 function DiffEqBase.SteadyStateProblem{iip}(sys::AbstractODESystem,u0map,
                                             parammap=DiffEqBase.NullParameters();
                                             kwargs...) where iip
-    f, u0, p = process_DEProblem(ODEFunction{iip}, sys, u0map, parammap; kwargs...)
-    SteadyStateProblem(f,u0,p;kwargs...)
+    f, u0, p = process_DEProblem(ODEFunction{iip}, sys, u0map, parammap; steady_state = true, kwargs...)
+    SteadyStateProblem{iip}(f,u0,p;kwargs...)
 end
 
 """
 ```julia
-function DiffEqBase.SteadyStateProblemExpr(sys::AbstractODESystem,u0map,tspan,
+function DiffEqBase.SteadyStateProblemExpr(sys::AbstractODESystem,u0map,
                                     parammap=DiffEqBase.NullParameters();
                                     version = nothing, tgrad=false,
                                     jac = false,
@@ -617,7 +650,7 @@ struct SteadyStateProblemExpr{iip} end
 function SteadyStateProblemExpr{iip}(sys::AbstractODESystem,u0map,
                                     parammap=DiffEqBase.NullParameters();
                                     kwargs...) where iip
-    f, u0, p = process_DEProblem(ODEFunctionExpr{iip}, sys, u0map, parammap; kwargs...)
+    f, u0, p = process_DEProblem(ODEFunctionExpr{iip}, sys, u0map, parammap;steady_state = true, kwargs...)
     linenumbers = get(kwargs, :linenumbers, true)
     ex = quote
         f = $f

src/systems/diffeqs/odesystem.jl

@@ -252,13 +252,13 @@ function build_explicit_observed_function(
         output[i] = obs[idx].rhs
     end
 
+    dvs = DestructuredArgs(states(sys), inbounds=!checkbounds)
+    ps = DestructuredArgs(parameters(sys), inbounds=!checkbounds)
+    iv = independent_variable(sys)
+    args = iv === nothing ? [dvs, ps] : [dvs, ps, iv]
+
     ex = Func(
-        [
-         DestructuredArgs(states(sys), inbounds=!checkbounds)
-         DestructuredArgs(parameters(sys), inbounds=!checkbounds)
-         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

@@ -142,10 +142,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

@@ -147,6 +147,10 @@ u0 = [
 prob1 = ODEProblem(reduced_system, u0, (0.0, 100.0), pp)
 solve(prob1, Rodas5())
 
+prob2 = SteadyStateProblem(reduced_system, u0, pp)
+@test prob2.f.observed(lorenz2.u, prob2.u0, pp) === 1.0
+
+
 # issue #724 and #716
 let
     @parameters t
@@ -204,6 +208,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