Add observe fun for SteadyStateProblem

Author: YingboMa

src/systems/diffeqs/abstractodesystem.jl

@@ -145,6 +145,7 @@ function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
                                      eval_expression = true,
                                      sparse = false, simplify=false,
                                      eval_module = @__MODULE__,
+                                     steady_state = false,
                                      kwargs...) where {iip}
 
     f_gen = generate_function(sys, dvs, ps; expression=Val{eval_expression}, expression_module=eval_module, kwargs...)
@@ -178,12 +179,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)
+    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)
+                end
+                obs(u, p, t)
             end
-            obs(u, p, t)
         end
     end
 
@@ -228,9 +240,28 @@ 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)
@@ -271,6 +302,7 @@ function ODEFunctionExpr{iip}(sys::AbstractODESystem, dvs = states(sys),
                           jac_prototype = $jp_expr,
                           syms = $(Symbol.(states(sys))),
                           indepsym = $(QuoteNode(Symbol(independent_variable(sys)))),
+                          observed = $observedfun,
                          )
     end
     !linenumbers ? striplines(ex) : ex
@@ -379,7 +411,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,
@@ -393,13 +425,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,
@@ -417,7 +449,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

test/reduction.jl

@@ -143,6 +143,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