Merge pull request #1463 from isaacsas/observed_with_jumps

add observed to JumpProblems

Author: ChrisRackauckas

Project.toml

@@ -68,7 +68,7 @@ Reexport = "0.2, 1"
 Requires = "1.0"
 RuntimeGeneratedFunctions = "0.4.3, 0.5"
 SafeTestsets = "0.0.1"
-SciMLBase = "1.3"
+SciMLBase = "1.26.2"
 Setfield = "0.7, 0.8"
 SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0, 2"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"

src/systems/jumps/jumpsystem.jl

@@ -219,12 +219,24 @@ dprob = DiscreteProblem(js, u₀map, tspan, parammap)
 ```
 """
 function DiffEqBase.DiscreteProblem(sys::JumpSystem, u0map, tspan::Union{Tuple,Nothing},
-                                    parammap=DiffEqBase.NullParameters(); kwargs...)
+                                    parammap=DiffEqBase.NullParameters(); checkbounds=false, kwargs...)
     defs = defaults(sys)
     u0 = varmap_to_vars(u0map, states(sys); defaults=defs)
     p  = varmap_to_vars(parammap, parameters(sys); defaults=defs)
     f  = DiffEqBase.DISCRETE_INPLACE_DEFAULT
-    df = DiscreteFunction{true,true}(f, syms=Symbol.(states(sys)))
+    
+    # just taken from abstractodesystem.jl for ODEFunction def
+    obs = observed(sys)
+    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)
+                end
+                obs(u, p, t)
+            end
+        end
+
+    df = DiscreteFunction{true,true}(f, syms=Symbol.(states(sys)), observed=observedfun)
     DiscreteProblem(df, u0, tspan, p; kwargs...)
 end
 

test/jumpsystem.jl

@@ -71,6 +71,14 @@ function getmean(jprob,Nsims)
 end
 m = getmean(jprob,Nsims)
 
+@variables S2(t)
+obs = [S2 ~ 2*S]
+@named js2b = JumpSystem([j₁,j₃], t, [S,I,R], [β,γ], observed=obs)
+dprob = DiscreteProblem(js2b, u₀map, tspan, parammap)
+jprob = JumpProblem(js2b, dprob, Direct(), save_positions=(false,false))
+sol = solve(jprob, SSAStepper(), saveat=tspan[2]/10)
+@test all(2 .* sol[S] .== sol[S2])
+
 # test save_positions is working
 jprob = JumpProblem(js2, dprob, Direct(), save_positions=(false,false))
 sol = solve(jprob, SSAStepper(), saveat=1.0)