start adding JumpSystem features

Author: isaacsas

Project.toml

@@ -6,6 +6,7 @@ version = "3.0.2"
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+DiffEqJump = "c894b116-72e5-5b58-be3c-e6d8d4ac2b12"
 DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"

src/ModelingToolkit.jl

@@ -5,6 +5,7 @@ using StaticArrays, LinearAlgebra, SparseArrays
 using Latexify, Unitful, ArrayInterface
 using MacroTools
 using UnPack: @unpack
+using DiffEqJump: VariableRateJump, ConstantRateJump, MassActionJump
 
 using Base.Threads
 import MacroTools: splitdef, combinedef, postwalk, striplines
@@ -87,6 +88,8 @@ include("systems/diffeqs/index_reduction.jl")
 include("systems/diffeqs/modelingtoolkitize.jl")
 include("systems/diffeqs/validation.jl")
 
+include("systems/jumps/jumpsystem.jl")
+
 include("systems/nonlinear/nonlinearsystem.jl")
 
 include("systems/optimization/optimizationsystem.jl")

src/systems/jumps/jumpsystem.jl

@@ -0,0 +1,56 @@
+JumpType = Union{VariableRateJump, ConstantRateJump, MassActionJump}
+
+struct JumpSystem <: AbstractSystem
+    eqs::Vector{JumpType}
+    iv::Variable
+    states::Vector{Variable}
+    ps::Vector{Variable}
+    name::Symbol
+    systems::Vector{JumpSystem}
+end
+
+function JumpSystem(eqs, iv, states, ps; systems = JumpSystem[], 
+                                          name = gensym(:JumpSystem))
+    JumpSystem(eqs, iv, convert.(Variable, states), convert.(Variable, ps), name, systems)
+end
+
+function generate_rate_function(js, rate)
+    f = striplines(build_function(rate, states(js), parameters(js), independent_variable(js)))
+end
+
+function generate_affect_function(js, affect)
+    # waiting on integrator build_function form
+end
+
+function assemble_vrj(js, vrj)
+    rate   = generate_rate_function(js, vrj.rate)
+    affect = generate_affect_function(js, vrj.affect)
+    VariableRateJump(rate, affect)
+end
+
+function assemble_crj(js, crj)
+    rate   = generate_rate_function(js, vrj.rate)
+    affect = generate_affect_function(js, vrj.affect)
+    ConstantRateJump(rate, affect)
+end
+
+function assemble_maj(maj, states_to_idxs, ps_to_idxs; scale_rate=false)
+
+    # mass action scaled_rates need to be a Number, but 
+    # operations are numbers, so can't check the type directly
+    @assert !isa(maj.scaled_rates, Union{Operation,Variable})
+
+    rstype = fieldtypes(eltype(maj.reactant_stoch))[2]
+    rs     = Vector{Pair{valtype(states_to_idxs),rstype}}()
+    for (spec,stoich) in maj.reactant_stoch
+        push!(rs, states_to_idxs[spec] => stoich)
+    end
+
+    nstype = fieldtypes(eltype(maj.net_stoch))[2]
+    ns     = Vector{Pair{valtype(states_to_idxs),nstype}}()
+    for (spec,stoich) in maj.net_stoch
+        push!(ns, states_to_idxs[spec] => stoich)
+    end
+
+    MassActionJump(maj.scaled_rates, rs, ns, scale_rates=scale_rate)
+end