Merge pull request #1143 from isaacsas/pmapper-for-jumpsystems

switch MassActionJumps to use parameter vectors

Author: ChrisRackauckas

Project.toml

@@ -50,7 +50,7 @@ ConstructionBase = "1"
 DataStructures = "0.17, 0.18"
 DiffEqBase = "6.54.0"
 DiffEqCallbacks = "2.16"
-DiffEqJump = "6.7.5"
+DiffEqJump = "7.0"
 DiffRules = "0.1, 1.0"
 Distributions = "0.23, 0.24, 0.25"
 DocStringExtensions = "0.7, 0.8"

src/systems/jumps/jumpsystem.jl

@@ -185,24 +185,13 @@ function numericnstoich(mtrs::Vector{Pair{V,W}}, statetoid) where {V,W}
     sort!(ns)
 end
 
-# assemble a numeric MassActionJump from a MT MassActionJump representing one rx.
-function assemble_maj(maj::MassActionJump, statetoid, subber, invttype)
-    rval = subber(maj.scaled_rates)
-    rs   = numericrstoich(maj.reactant_stoch, statetoid)
-    ns   = numericnstoich(maj.net_stoch, statetoid)
-    maj  = MassActionJump(convert(invttype, value(rval)), rs, ns, scale_rates = false)
-    maj
+# assemble a numeric MassActionJump from a MT symbolics MassActionJumps
+function assemble_maj(majv::Vector{U}, statetoid, pmapper) where {U <: MassActionJump}
+    rs = [numericrstoich(maj.reactant_stoch, statetoid) for maj in majv]
+    ns = [numericnstoich(maj.net_stoch, statetoid) for maj in majv]
+    MassActionJump(rs, ns; param_mapper=pmapper, nocopy=true)
 end
 
-# For MassActionJumps that contain many reactions
-# function assemble_maj(maj::MassActionJump{U,V,W}, statetoid, subber,
-#                       invttype) where {U <: AbstractVector,V,W}
-#     rval = [convert(invttype,numericrate(sr, subber)) for sr in maj.scaled_rates]
-#     rs   = [numericrstoich(rs, statetoid) for rs in maj.reactant_stoch]
-#     ns   = [numericnstoich(ns, statetoid) for ns in maj.net_stoch]
-#     maj  = MassActionJump(rval, rs, ns, scale_rates = false)
-#     maj
-# end
 """
 ```julia
 function DiffEqBase.DiscreteProblem(sys::JumpSystem, u0map, tspan,
@@ -287,14 +276,13 @@ function DiffEqJump.JumpProblem(js::JumpSystem, prob, aggregator; kwargs...)
 
     # handling parameter substition and empty param vecs
     p = (prob.p isa DiffEqBase.NullParameters || prob.p === nothing) ? Num[] : prob.p
-    parammap  = map((x,y)->Pair(x,y), parameters(js), p)
-    subber    = substituter(parammap)
 
-    majs = MassActionJump[assemble_maj(j, statetoid, subber, invttype) for j in eqs.x[1]]
+    majpmapper = JumpSysMajParamMapper(js, p; jseqs=eqs, rateconsttype=invttype)
+    majs = isempty(eqs.x[1]) ? nothing : assemble_maj(eqs.x[1], statetoid, majpmapper)
     crjs = ConstantRateJump[assemble_crj(js, j, statetoid) for j in eqs.x[2]]
     vrjs = VariableRateJump[assemble_vrj(js, j, statetoid) for j in eqs.x[3]]
     ((prob isa DiscreteProblem) && !isempty(vrjs)) && error("Use continuous problems such as an ODEProblem or a SDEProblem with VariableRateJumps")
-    jset = JumpSet(Tuple(vrjs), Tuple(crjs), nothing, isempty(majs) ? nothing : majs)
+    jset = JumpSet(Tuple(vrjs), Tuple(crjs), nothing, majs)
 
     if needs_vartojumps_map(aggregator) || needs_depgraph(aggregator)
         jdeps = asgraph(js)
@@ -306,7 +294,8 @@ function DiffEqJump.JumpProblem(js::JumpSystem, prob, aggregator; kwargs...)
         vtoj = nothing; jtov = nothing; jtoj = nothing
     end
 
-    JumpProblem(prob, aggregator, jset; dep_graph=jtoj, vartojumps_map=vtoj, jumptovars_map=jtov, kwargs...)
+    JumpProblem(prob, aggregator, jset; dep_graph=jtoj, vartojumps_map=vtoj, jumptovars_map=jtov, 
+                scale_rates=false, nocopy=true, kwargs...)
 end
 
 
@@ -338,3 +327,49 @@ function modified_states!(mstates, jump::MassActionJump, sts)
         any(isequal(state), sts) && push!(mstates, state)
     end
 end
+
+
+
+###################### parameter mapper ###########################
+struct JumpSysMajParamMapper{U,V,W}
+    paramexprs::U     # the parameter expressions to use for each jump rate constant
+    sympars::V        # parameters(sys) from the underlying JumpSystem
+    subdict           # mapping from an element of parameters(sys) to its current numerical value
+end
+
+function JumpSysMajParamMapper(js::JumpSystem, p; jseqs=nothing, rateconsttype=Float64)
+    eqs        = (jseqs === nothing) ? equations(js) : jseqs
+    paramexprs = [maj.scaled_rates for maj in eqs.x[1]]
+    psyms      = parameters(js)
+    paramdict  = Dict(value(k) => value(v)  for (k, v) in zip(psyms,p))
+    JumpSysMajParamMapper{typeof(paramexprs),typeof(psyms),rateconsttype}(paramexprs, psyms, paramdict)
+end
+
+function updateparams!(ratemap::JumpSysMajParamMapper{U,V,W}, params) where {U <: AbstractArray, V <: AbstractArray, W}
+    for (i,p) in enumerate(params)
+        sympar = ratemap.sympars[i]
+        ratemap.subdict[sympar] = p
+    end
+    nothing
+end
+
+function updateparams!(::JumpSysMajParamMapper{U,V,W}, params::Nothing) where {U <: AbstractArray, V <: AbstractArray, W}
+    nothing
+end
+
+
+# create the initial parameter vector for use in a MassActionJump
+function (ratemap::JumpSysMajParamMapper{U,V,W})(params) where {U <: AbstractArray, V <: AbstractArray, W}
+    updateparams!(ratemap, params)
+    [convert(W,value(substitute(paramexpr, ratemap.subdict))) for paramexpr in ratemap.paramexprs]
+end
+
+# update a maj with parameter vectors
+function (ratemap::JumpSysMajParamMapper{U,V,W})(maj::MassActionJump, newparams; scale_rates, kwargs...) where {U <: AbstractArray, V <: AbstractArray, W}
+    updateparams!(ratemap, newparams)
+    for i in 1:get_num_majumps(maj)
+        maj.scaled_rates[i] = convert(W,value(substitute(ratemap.paramexprs[i], ratemap.subdict)))
+    end
+    scale_rates && DiffEqJump.scalerates!(maj.scaled_rates, maj.reactant_stoch)
+    nothing
+end

test/jumpsystem.jl

@@ -116,8 +116,6 @@ m2 = getmean(jprob,Nsims)
 maj1 = MassActionJump(2*β/2, [S => 1, I => 1], [S => -1, I => 1])
 maj2 = MassActionJump(γ, [I => 1], [I => -1, R => 1])
 js3   = JumpSystem([maj1,maj2], t, [S,I,R], [β,γ])
-statetoid = Dict(MT.value(state) => i for (i,state) in enumerate(states(js)))
-ptoid     = Dict(MT.value(par) => i for (i,par) in enumerate(parameters(js)))
 dprob = DiscreteProblem(js3, u₀map, tspan, parammap)
 jprob = JumpProblem(js3, dprob, Direct())
 m3 = getmean(jprob,Nsims)
@@ -136,8 +134,6 @@ m4 = getmean(jprobc,Nsims)
 maj1 = MassActionJump(2.0, [0 => 1], [S => 1])
 maj2 = MassActionJump(γ, [S => 1], [S => -1])
 js4   = JumpSystem([maj1,maj2], t, [S], [β,γ])
-statetoid = Dict(MT.value(state) => i for (i,state) in enumerate(states(js)))
-ptoid     = Dict(MT.value(par) => i for (i,par) in enumerate(parameters(js)))
 dprob = DiscreteProblem(js4, [S => 999], (0,1000.), [β => 100.,γ => .01])
 jprob = JumpProblem(js4, dprob, Direct())
 m4 = getmean(jprob,Nsims)
@@ -147,8 +143,6 @@ m4 = getmean(jprob,Nsims)
 maj1 = MassActionJump(2.0, [0 => 1], [S => 1])
 maj2 = MassActionJump(γ, [S => 2], [S => -1])
 js4   = JumpSystem([maj1,maj2], t, [S], [β,γ])
-statetoid = Dict(MT.value(state) => i for (i,state) in enumerate(states(js)))
-ptoid     = Dict(MT.value(par) => i for (i,par) in enumerate(parameters(js)))
 dprob = DiscreteProblem(js4, [S => 999], (0,1000.), [β => 100.,γ => .01])
 jprob = JumpProblem(js4, dprob, Direct())
 sol = solve(jprob, SSAStepper());
@@ -158,4 +152,26 @@ sol = solve(jprob, SSAStepper());
   sys1 = JumpSystem([maj1, maj2], t, [S], [β, γ], name = :sys1)
   sys2 = JumpSystem([maj1, maj2], t, [S], [β, γ], name = :sys1)
   @test_throws ArgumentError JumpSystem([sys1.γ ~ sys2.γ], t, [], [], systems = [sys1, sys2])
-end
+end
+
+# test if param mapper is setup correctly for callbacks
+@parameters k1 k2 k3
+@variables A(t) B(t)
+maj1 = MassActionJump(k1*k3, [0 => 1], [A => -1, B => 1])
+maj2 = MassActionJump(k2, [B => 1], [A => 1, B => -1])
+js5 = JumpSystem([maj1,maj2], t, [A,B], [k1,k2,k3])
+p  = [k1 => 2.0, k2 => 0.0, k3 => .5]
+u₀ = [A => 100, B => 0]
+tspan = (0.0,2000.0)
+dprob = DiscreteProblem(js5, u₀, tspan, p)
+jprob = JumpProblem(js5, dprob, Direct(), save_positions=(false,false))
+@test all(jprob.massaction_jump.scaled_rates .== [1.0,0.0])
+
+pcondit(u,t,integrator) = t==1000.0
+function paffect!(integrator)
+  integrator.p[1] = 0.0
+  integrator.p[2] = 1.0
+  reset_aggregated_jumps!(integrator)
+end
+sol = solve(jprob, SSAStepper(), tstops=[1000.0], callback=DiscreteCallback(pcondit,paffect!))
+@test sol[1,end] == 100

test/reactionsystem.jl

@@ -192,12 +192,13 @@ jumps[19] = VariableRateJump((u,p,t) -> p[19]*u[1]*t, integrator -> (integrator.
 jumps[20] = VariableRateJump((u,p,t) -> p[20]*t*u[1]*binomial(u[2],2)*u[3], integrator -> (integrator.u[2] -= 2; integrator.u[3] -= 1; integrator.u[4] += 2))
 
 statetoid = Dict(state => i for (i,state) in enumerate(states(js)))
-parammap = map((x,y)->Pair(x,y),parameters(js),pars)
-for i in midxs
-  maj = MT.assemble_maj(equations(js)[i], statetoid, ModelingToolkit.substituter(parammap),eltype(pars))
-  @test abs(jumps[i].scaled_rates - maj.scaled_rates) < 100*eps()
-  @test jumps[i].reactant_stoch == maj.reactant_stoch
-  @test jumps[i].net_stoch == maj.net_stoch
+jspmapper = ModelingToolkit.JumpSysMajParamMapper(js, pars)
+symmaj = MT.assemble_maj(equations(js).x[1], statetoid, jspmapper)
+maj    = MassActionJump(symmaj.param_mapper(pars), symmaj.reactant_stoch, symmaj.net_stoch, symmaj.param_mapper, scale_rates=false)
+for i in midxs  
+  @test abs(jumps[i].scaled_rates - maj.scaled_rates[i]) < 100*eps()
+  @test jumps[i].reactant_stoch == maj.reactant_stoch[i]
+  @test jumps[i].net_stoch == maj.net_stoch[i]
 end
 for i in cidxs
   crj = MT.assemble_crj(js, equations(js)[i], statetoid)