Merge pull request #2946 from isaacsas/better_vrj_support

ChrisRackauckas · web-flow · commit e4e7d97e8d3f · 2024-08-11T07:14:26.000-04:00
support JumpProblems over ODEProblems
diff --git a/src/systems/dependency_graphs.jl b/src/systems/dependency_graphs.jl
@@ -36,8 +36,8 @@ equation_dependencies(jumpsys)
 equation_dependencies(jumpsys, variables = parameters(jumpsys))
 ```
 """
-function equation_dependencies(sys::AbstractSystem; variables = unknowns(sys))
-    eqs = equations(sys)
+function equation_dependencies(sys::AbstractSystem; variables = unknowns(sys),
+        eqs = equations(sys))
     deps = Set()
     depeqs_to_vars = Vector{Vector}(undef, length(eqs))
 
@@ -114,8 +114,9 @@ digr = asgraph(jumpsys)
 ```
 """
 function asgraph(sys::AbstractSystem; variables = unknowns(sys),
-        variablestoids = Dict(v => i for (i, v) in enumerate(variables)))
-    asgraph(equation_dependencies(sys, variables = variables), variablestoids)
+        variablestoids = Dict(v => i for (i, v) in enumerate(variables)),
+        eqs = equations(sys))
+    asgraph(equation_dependencies(sys; variables, eqs), variablestoids)
 end
 
 """
@@ -141,8 +142,7 @@ variable_dependencies(jumpsys)
 ```
 """
 function variable_dependencies(sys::AbstractSystem; variables = unknowns(sys),
-        variablestoids = nothing)
-    eqs = equations(sys)
+        variablestoids = nothing, eqs = equations(sys))
     vtois = isnothing(variablestoids) ? Dict(v => i for (i, v) in enumerate(variables)) :
             variablestoids
 
@@ -193,8 +193,8 @@ dg = asdigraph(digr, jumpsys)
 ```
 """
 function asdigraph(g::BipartiteGraph, sys::AbstractSystem; variables = unknowns(sys),
-        equationsfirst = true)
-    neqs = length(equations(sys))
+        equationsfirst = true, eqs = equations(sys))
+    neqs = length(eqs)
     nvars = length(variables)
     fadjlist = deepcopy(g.fadjlist)
     badjlist = deepcopy(g.badjlist)
diff --git a/src/systems/jumps/jumpsystem.jl b/src/systems/jumps/jumpsystem.jl
@@ -194,6 +194,10 @@ function JumpSystem(eqs, iv, unknowns, ps;
         metadata, gui_metadata, checks = checks)
 end
 
+has_massactionjumps(js::JumpSystem) = !isempty(equations(js).x[1])
+has_constantratejumps(js::JumpSystem) = !isempty(equations(js).x[2])
+has_variableratejumps(js::JumpSystem) = !isempty(equations(js).x[3])
+
 function generate_rate_function(js::JumpSystem, rate)
     consts = collect_constants(rate)
     if !isempty(consts) # The SymbolicUtils._build_function method of this case doesn't support postprocess_fbody
@@ -311,7 +315,7 @@ end
 ```julia
 DiffEqBase.DiscreteProblem(sys::JumpSystem, u0map, tspan,
                            parammap = DiffEqBase.NullParameters;
-                           use_union = false,
+                           use_union = true,
                            kwargs...)
 ```
 
@@ -331,7 +335,6 @@ dprob = DiscreteProblem(complete(js), u₀map, tspan, parammap)
 """
 function DiffEqBase.DiscreteProblem(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothing},
         parammap = DiffEqBase.NullParameters();
-        checkbounds = false,
         use_union = true,
         eval_expression = false,
         eval_module = @__MODULE__,
@@ -385,7 +388,7 @@ struct DiscreteProblemExpr{iip} end
 
 function DiscreteProblemExpr{iip}(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothing},
         parammap = DiffEqBase.NullParameters();
-        use_union = false,
+        use_union = true,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed `JumpSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `DiscreteProblemExpr`")
@@ -412,6 +415,60 @@ function DiscreteProblemExpr{iip}(sys::JumpSystem, u0map, tspan::Union{Tuple, No
     end
 end
 
+"""
+```julia
+DiffEqBase.ODEProblem(sys::JumpSystem, u0map, tspan,
+                           parammap = DiffEqBase.NullParameters;
+                           use_union = true,
+                           kwargs...)
+```
+
+Generates a blank ODEProblem for a pure jump JumpSystem to utilize as its `prob.prob`. This
+is used in the case where there are no ODEs and no SDEs associated with the system but there
+are jumps with an explicit time dependency (i.e. `VariableRateJump`s). If no jumps have an
+explicit time dependence, i.e. all are `ConstantRateJump`s or `MassActionJump`s then
+`DiscreteProblem` should be preferred for performance reasons.
+
+Continuing the example from the [`JumpSystem`](@ref) definition:
+
+```julia
+using DiffEqBase, JumpProcesses
+u₀map = [S => 999, I => 1, R => 0]
+parammap = [β => 0.1 / 1000, γ => 0.01]
+tspan = (0.0, 250.0)
+oprob = ODEProblem(complete(js), u₀map, tspan, parammap)
+```
+"""
+function DiffEqBase.ODEProblem(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothing},
+        parammap = DiffEqBase.NullParameters();
+        use_union = true,
+        eval_expression = false,
+        eval_module = @__MODULE__,
+        kwargs...)
+    if !iscomplete(sys)
+        error("A completed `JumpSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `DiscreteProblem`")
+    end
+    dvs = unknowns(sys)
+    ps = parameters(sys)
+
+    defs = defaults(sys)
+    defs = mergedefaults(defs, parammap, ps)
+    defs = mergedefaults(defs, u0map, dvs)
+
+    u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = false)
+    if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        p = MTKParameters(sys, parammap, u0map; eval_expression, eval_module)
+    else
+        p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = false, use_union)
+    end
+
+    observedfun = ObservedFunctionCache(sys; eval_expression, eval_module)
+
+    f = (du, u, p, t) -> (du .= 0; nothing)
+    df = ODEFunction(f; sys, observed = observedfun)
+    ODEProblem(df, u0, tspan, p; kwargs...)
+end
+
 """
 ```julia
 DiffEqBase.JumpProblem(js::JumpSystem, prob, aggregator; kwargs...)
@@ -449,10 +506,12 @@ function JumpProcesses.JumpProblem(js::JumpSystem, prob,
         error("Use continuous problems such as an ODEProblem or a SDEProblem with VariableRateJumps")
     jset = JumpSet(Tuple(vrjs), Tuple(crjs), nothing, majs)
 
+    # dep graphs are only for constant rate jumps
+    nonvrjs = ArrayPartition(eqs.x[1], eqs.x[2])
     if needs_vartojumps_map(aggregator) || needs_depgraph(aggregator) ||
        (aggregator isa JumpProcesses.NullAggregator)
-        jdeps = asgraph(js)
-        vdeps = variable_dependencies(js)
+        jdeps = asgraph(js; eqs = nonvrjs)
+        vdeps = variable_dependencies(js; eqs = nonvrjs)
         vtoj = jdeps.badjlist
         jtov = vdeps.badjlist
         jtoj = needs_depgraph(aggregator) ? eqeq_dependencies(jdeps, vdeps).fadjlist :
diff --git a/test/dep_graphs.jl b/test/dep_graphs.jl
@@ -1,5 +1,5 @@
 using Test
-using ModelingToolkit, Graphs, JumpProcesses
+using ModelingToolkit, Graphs, JumpProcesses, RecursiveArrayTools
 using ModelingToolkit: t_nounits as t, D_nounits as D
 import ModelingToolkit: value
 
@@ -16,11 +16,11 @@ j₅ = ConstantRateJump(k1 * I, [R ~ R + 1])
 j₆ = VariableRateJump(k1 * k2 / (1 + t) * S, [S ~ S - 1, R ~ R + 1])
 eqs = [j₁, j₂, j₃, j₄, j₅, j₆]
 @named js = JumpSystem(eqs, t, [S, I, R], [k1, k2])
-S = value(S);
-I = value(I);
-R = value(R);
-k1 = value(k1);
-k2 = value(k2);
+S = value(S)
+I = value(I)
+R = value(R)
+k1 = value(k1)
+k2 = value(k2)
 # eq to vars they depend on
 eq_sdeps = [Variable[], [S], [S, I], [S, R], [I], [S]]
 eq_sidepsf = [Int[], [1], [1, 2], [1, 3], [2], [1]]
@@ -72,6 +72,80 @@ end
 dg4 = varvar_dependencies(depsbg, deps2)
 @test dg == dg4
 
+# testing when ignoring VariableRateJumps
+let
+    @parameters k1 k2
+    @variables S(t) I(t) R(t)
+    j₁ = MassActionJump(k1, [0 => 1], [S => 1])
+    j₂ = MassActionJump(k1, [S => 1], [S => -1])
+    j₃ = MassActionJump(k2, [S => 1, I => 1], [S => -1, I => 1])
+    j₄ = MassActionJump(k2, [S => 2, R => 1], [R => -1])
+    j₅ = ConstantRateJump(k1 * I, [R ~ R + 1])
+    j₆ = VariableRateJump(k1 * k2 / (1 + t) * S, [S ~ S - 1, R ~ R + 1])
+    eqs = [j₁, j₂, j₃, j₄, j₅, j₆]
+    @named js = JumpSystem(eqs, t, [S, I, R], [k1, k2])
+    S = value(S)
+    I = value(I)
+    R = value(R)
+    k1 = value(k1)
+    k2 = value(k2)
+    # eq to vars they depend on
+    eq_sdeps = [Variable[], [S], [S, I], [S, R], [I]]
+    eq_sidepsf = [Int[], [1], [1, 2], [1, 3], [2]]
+    eq_sidepsb = [[2, 3, 4], [3, 5], [4]]
+
+    # filter out vrjs in making graphs
+    eqs = ArrayPartition(equations(js).x[1], equations(js).x[2])
+    deps = equation_dependencies(js; eqs)
+    @test length(deps) == length(eq_sdeps)
+    @test all(i -> isequal(Set(eq_sdeps[i]), Set(deps[i])), 1:length(eqs))
+    depsbg = asgraph(js; eqs)
+    @test depsbg.fadjlist == eq_sidepsf
+    @test depsbg.badjlist == eq_sidepsb
+
+    # eq to params they depend on
+    eq_pdeps = [[k1], [k1], [k2], [k2], [k1]]
+    eq_pidepsf = [[1], [1], [2], [2], [1]]
+    eq_pidepsb = [[1, 2, 5], [3, 4]]
+    deps = equation_dependencies(js; variables = parameters(js), eqs)
+    @test length(deps) == length(eq_pdeps)
+    @test all(i -> isequal(Set(eq_pdeps[i]), Set(deps[i])), 1:length(eqs))
+    depsbg2 = asgraph(js; variables = parameters(js), eqs)
+    @test depsbg2.fadjlist == eq_pidepsf
+    @test depsbg2.badjlist == eq_pidepsb
+
+    # var to eqs that modify them
+    s_eqdepsf = [[1, 2, 3], [3], [4, 5]]
+    s_eqdepsb = [[1], [1], [1, 2], [3], [3]]
+    ne = 6
+    bg = BipartiteGraph(ne, s_eqdepsf, s_eqdepsb)
+    deps2 = variable_dependencies(js; eqs)
+    @test isequal(bg, deps2)
+
+    # eq to eqs that depend on them
+    eq_eqdeps = [[2, 3, 4], [2, 3, 4], [2, 3, 4, 5], [4], [4], [2, 3, 4]]
+    dg = SimpleDiGraph(5)
+    for (eqidx, eqdeps) in enumerate(eq_eqdeps)
+        for eqdepidx in eqdeps
+            add_edge!(dg, eqidx, eqdepidx)
+        end
+    end
+    dg3 = eqeq_dependencies(depsbg, deps2)
+    @test dg == dg3
+
+    # var to vars that depend on them
+    var_vardeps = [[1, 2, 3], [1, 2, 3], [3]]
+    ne = 7
+    dg = SimpleDiGraph(3)
+    for (vidx, vdeps) in enumerate(var_vardeps)
+        for vdepidx in vdeps
+            add_edge!(dg, vidx, vdepidx)
+        end
+    end
+    dg4 = varvar_dependencies(depsbg, deps2)
+    @test dg == dg4
+end
+
 #####################################
 #       testing for ODE/SDEs
 #####################################
diff --git a/test/jumpsystem.jl b/test/jumpsystem.jl
