Use SpatialMassActionJump

Author: TorkelE

src/Catalyst.jl

@@ -6,9 +6,9 @@ module Catalyst
 using DocStringExtensions
 using SparseArrays, DiffEqBase, Reexport, Setfield
 using LaTeXStrings, Latexify, Requires
-using JumpProcesses: JumpProcesses,
-                     JumpProblem, MassActionJump, ConstantRateJump,
-                     VariableRateJump
+using JumpProcesses: JumpProcesses, JumpProblem, 
+                     MassActionJump, ConstantRateJump, VariableRateJump,
+                     SpatialMassActionJump
 
 # ModelingToolkit imports and convenience functions we use
 using ModelingToolkit

src/spatial_reaction_systems/lattice_jump_systems.jl

@@ -28,15 +28,19 @@ end
 function JumpProcesses.JumpProblem(lrs::LatticeReactionSystem, dprob, aggregator, args...; name = nameof(lrs.rs), 
                                    combinatoric_ratelaws = get_combinatoric_ratelaws(lrs.rs), kwargs...)
     # Error checks.
-    (dprob.p isa Vector{Vector{Vector{Float64}}}) || dprob.p isa Vector{Vector} || error("Parameters in input DiscreteProblem is of an unexpected type: $(typeof(dprob.p)). Was a LatticeReactionProblem passed into the DiscreteProblem when it was created?") # The second check (Vector{Vector} is needed becaus on the CI server somehow the Tuple{..., ...} is covnerted into a Vector[..., ...]). It does not happen when I run tests locally, so no ideal how to fix.
-    any(length.(dprob.p[1]) .> 1) && error("Spatial reaction rates are currently not supported in lattice jump simulations.")
+    # The second check (Vector{Vector} is needed because on the CI server somehow the Tuple{..., ...} is converted into a Vector[..., ...]).
+    # It does not happen when I run tests locally, so no ideal how to fix.
+    (dprob.p isa Vector{Vector{Vector{Float64}}}) || dprob.p isa Vector{Vector} || error("Parameters in input DiscreteProblem is of an unexpected type: $(typeof(dprob.p)). Was a LatticeReactionProblem passed into the DiscreteProblem when it was created?") 
 
     # Computes hopping constants and mass action jumps (requires some internal juggling).
-    # The non-spatial DiscreteProblem have a u0 matrix with entries for all combinations of species and vertexes.
     # Currently, JumpProcesses requires uniform vertex parameters (hence `p=first.(dprob.p[1])`).
+    # Currently, the resulting JumpProblem does not depend on parameters (no way to incorporate these).
+    # Hence the parameters of this one does nto actually matter. If at some point JumpProcess can 
+    # handle parameters this can be updated and improved.
+    # The non-spatial DiscreteProblem have a u0 matrix with entries for all combinations of species and vertexes.
     hopping_constants = make_hopping_constants(dprob, lrs)
+    sma_jumps = make_spatial_majumps(dprob, lrs)
     non_spat_dprob = DiscreteProblem(reshape(dprob.u0, lrs.num_species, lrs.num_verts), dprob.tspan, first.(dprob.p[1]))
-    sma_jumps = make_spatial_majumps(non_spat_dprob, dprob, lrs)
 
     return JumpProblem(non_spat_dprob, aggregator, sma_jumps; 
                        hopping_constants, spatial_system = lrs.lattice, name, kwargs...)
@@ -66,22 +70,65 @@ function make_hopping_constants(dprob::DiscreteProblem, lrs::LatticeReactionSyst
     return hopping_constants
 end
 
-# Creates the (spatial) mass action jumps from a (spatial) DiscreteProblem its non-spatial version, and a LatticeReactionSystem.
-function make_spatial_majumps(non_spat_dprob, dprob, rs::LatticeReactionSystem)
-    ma_jumps = make_majumps(non_spat_dprob, lrs.rs)
-    
+# Creates a SpatialMassActionJump struct from a (spatial) DiscreteProblem and a LatticeReactionSystem.
+# Could implementation a version which, if all reaction's rates are uniform, returns a MassActionJump.
+# Not sure if there is any form of performance improvement from that though. Possibly is not the case.
+function make_spatial_majumps(dprob, lrs::LatticeReactionSystem)
+    # Creates a vector, storing which reactions have spatial components.
+    is_spatials = [Catalyst.has_spatial_vertex_component(rx.rate, lrs; vert_ps = dprob.p[1]) for rx in reactions(lrs.rs)]
+
+    # Creates templates for the rates (uniform and spatial) and the stoichiometries.
+    # We cannot fetch reactant_stoich and net_stoich from a (non-spatial) MassActionJump.
+    # The reason is that we need to re-order the reactions so that uniform appears first, and spatial next.
+    u_rates = Vector{Float64}(undef, length(reactions(lrs.rs)) - count(is_spatials))
+    s_rates = Matrix{Float64}(undef, count(is_spatials), lrs.num_verts)
+    reactant_stoich = Vector{Vector{Pair{Int64, Int64}}}(undef, length(reactions(lrs.rs)))
+    net_stoich = Vector{Vector{Pair{Int64, Int64}}}(undef, length(reactions(lrs.rs)))
+
+    # Loops through reactions with non-spatial rates, computes their rates and stoichiometries.
+    cur_rx = 1;
+    for (is_spat, rx) in zip(is_spatials, reactions(lrs.rs))
+        is_spat && continue
+        u_rates[cur_rx] = compute_vertex_value(rx.rate, lrs; vert_ps = dprob.p[1])[1]
+        substoich_map = Pair.(rx.substrates, rx.substoich)
+        reactant_stoich[cur_rx] = int_map(substoich_map, lrs.rs)
+        net_stoich[cur_rx] = int_map(rx.netstoich, lrs.rs)
+        cur_rx += 1
+    end
+    # Loops through reactions with spatial rates, computes their rates and stoichiometries.
+    for (is_spat, rx) in zip(is_spatials, reactions(lrs.rs))    
+        is_spat || continue
+        s_rates[cur_rx-length(u_rates),:] = compute_vertex_value(rx.rate, lrs; vert_ps = dprob.p[1])
+        substoich_map = Pair.(rx.substrates, rx.substoich)
+        reactant_stoich[cur_rx] = int_map(substoich_map, lrs.rs)
+        net_stoich[cur_rx] = int_map(rx.netstoich, lrs.rs)
+        cur_rx += 1
+    end
+    # SpatialMassActionJump expects empty rate containers to be nothing.
+    isempty(u_rates) && (u_rates = nothing)
+    (count(is_spatials)==0) && (s_rates = nothing)
+
+    return SpatialMassActionJump(u_rates, s_rates, reactant_stoich, net_stoich)
 end
 
-# Creates the (non-spatial) mass action jumps from a (non-spatial) DiscreteProblem (and its Reaction System of origin).
-function make_majumps(non_spat_dprob, rs::ReactionSystem)
-    # Computes various required inputs for assembling the mass action jumps.
-    js = convert(JumpSystem, rs)
-    statetoid = Dict(ModelingToolkit.value(state) => i for (i, state) in enumerate(states(rs)))
-    eqs = equations(js)
-    invttype = non_spat_dprob.tspan[1] === nothing ? Float64 : typeof(1 / non_spat_dprob.tspan[2])
-
-    # Assembles the non-spatial mass action jumps.
-    p = (non_spat_dprob.p isa DiffEqBase.NullParameters || non_spat_dprob.p === nothing) ? Num[] : non_spat_dprob.p    
-    majpmapper = ModelingToolkit.JumpSysMajParamMapper(js, p; jseqs = eqs, rateconsttype = invttype)
-    return ModelingToolkit.assemble_maj(eqs.x[1], statetoid, majpmapper)
+### Extra ###
+
+# Temporary. Awaiting implementation in SII, or proper implementation withinCatalyst (with more general functionality).
+function int_map(map_in, sys) where {T,S}
+    return [ModelingToolkit.variable_index(sys, pair[1]) => pair[2] for pair in map_in]
 end
+
+# Currently unused. If we want to create certain types of MassActionJumps (instead of SpatialMassActionJumps) we can take this one back.
+# Creates the (non-spatial) mass action jumps from a (non-spatial) DiscreteProblem (and its Reaction System of origin).
+# function make_majumps(non_spat_dprob, rs::ReactionSystem)
+#     # Computes various required inputs for assembling the mass action jumps.
+#     js = convert(JumpSystem, rs)
+#     statetoid = Dict(ModelingToolkit.value(state) => i for (i, state) in enumerate(states(rs)))
+#     eqs = equations(js)
+#     invttype = non_spat_dprob.tspan[1] === nothing ? Float64 : typeof(1 / non_spat_dprob.tspan[2])
+# 
+#     # Assembles the non-spatial mass action jumps.
+#     p = (non_spat_dprob.p isa DiffEqBase.NullParameters || non_spat_dprob.p === nothing) ? Num[] : non_spat_dprob.p    
+#     majpmapper = ModelingToolkit.JumpSysMajParamMapper(js, p; jseqs = eqs, rateconsttype = invttype)
+#     return ModelingToolkit.assemble_maj(eqs.x[1], statetoid, majpmapper)
+# end

src/spatial_reaction_systems/utility.jl

@@ -295,3 +295,99 @@ end
 function matrix_expand_component_values(values::Vector{<:Vector}, n)
     reshape(expand_component_values(values, n), length(values), n)
 end
+
+# For an expression, computes its values using the provided state and parameter vectors.
+# The expression is assumed to be valid in edges (and can have edges parameter components).
+# If some component is non-uniform, output is a vector of length equal to the number of vertexes.
+# If all components are uniform, the output is a length one vector.
+function compute_edge_value(exp, lrs::LatticeReactionSystem, edge_ps)
+    # Finds the symbols in the expression. Checks that all correspond to edge parameters.
+    relevant_syms = Symbolics.get_variables(exp)
+    if !all(any(isequal(sym, p) for p in edge_parameters(lrs)) for sym in relevant_syms) 
+        error("An non-edge parameter was encountered in expressions: $exp. Here, only edge parameters are expected.")
+    end
+
+    # Creates a Function tha computes the expressions value for a parameter set.
+    exp_func = drop_expr(@RuntimeGeneratedFunction(build_function(exp, relevant_syms...)))
+    # Creates a dictionary with the value(s) for all edge parameters.
+    sym_val_dict = vals_to_dict(edge_parameters(lrs), edge_ps)
+
+    # If all values are uniform, compute value once. Else, do it at all edges.
+    if !has_spatial_edge_component(exp, lrs, edge_ps)
+        return [exp_func([sym_val_dict[sym][1] for sym in relevant_syms]...)]
+    end
+    return [exp_func([get_component_value(sym_val_dict[sym], idxE) for sym in relevant_syms]...) 
+                                                                            for idxE in 1:lrs.num_edges]
+end
+
+# For an expression, computes its values using the provided state and parameter vectors.
+# The expression is assumed to be valid in vertexes (and can have vertex parameter and state components).
+# If at least one component is non-uniform, output is a vector of length equal to the number of vertexes.
+# If all components are uniform, the output is a length one vector.
+function compute_vertex_value(exp, lrs::LatticeReactionSystem; u=nothing, vert_ps=nothing)
+    # Finds the symbols in the expression. Checks that all correspond to states or vertex parameters.
+    relevant_syms = Symbolics.get_variables(exp)
+    if any(any(isequal(sym) in edge_parameters(lrs)) for sym in relevant_syms) 
+        error("An edge parameter was encountered in expressions: $exp. Here, on vertex-based components are expected.")
+    end
+    # Creates a Function tha computes the expressions value for a parameter set.
+    exp_func = drop_expr(@RuntimeGeneratedFunction(build_function(exp, relevant_syms...)))
+    # Creates a dictionary with the value(s) for all edge parameters.
+    if !isnothing(u) && !isnothing(vert_ps)
+        all_syms = [species(lrs); vertex_parameters(lrs)]
+        all_vals = [u; vert_ps]
+    elseif !isnothing(u) && isnothing(vert_ps)
+        all_syms = species(lrs)
+        all_vals = u
+
+    elseif isnothing(u) && !isnothing(vert_ps)
+        all_syms = vertex_parameters(lrs)
+        all_vals = vert_ps
+    else
+        error("Either u or vertex_ps have to be provided to has_spatial_vertex_component.")
+    end
+    sym_val_dict = vals_to_dict(all_syms, all_vals)
+    
+    # If all values are uniform, compute value once. Else, do it at all edges.
+    if !has_spatial_vertex_component(exp, lrs; u, vert_ps)
+        return [exp_func([sym_val_dict[sym][1] for sym in relevant_syms]...)]
+    end
+    return [exp_func([get_component_value(sym_val_dict[sym], idxV) for sym in relevant_syms]...) 
+                                                                            for idxV in 1:lrs.num_verts]
+end
+
+### System Property Checks ###
+
+# For a Symbolic expression, a LatticeReactionSystem, and a parameter list of the internal format:
+# Checks if any edge parameter in the expression have a spatial component (that is, is not uniform).
+function has_spatial_edge_component(exp, lrs::LatticeReactionSystem, edge_ps)
+    # Finds the edge parameters in the expression. Computes their indexes.
+    exp_syms = Symbolics.get_variables(exp)
+    exp_edge_ps = filter(sym -> any(isequal(sym), edge_parameters(lrs)), exp_syms)
+    p_idxs = [findfirst(isequal(sym, edge_p) for edge_p in edge_parameters(lrs)) for sym in exp_syms]
+    # Checks if any of the corresponding value vectors have length != 1 (that is, is not uniform).
+    return any(length(edge_ps[p_idx]) != 1 for p_idx in p_idxs)
+end
+
+# For a Symbolic expression, a LatticeReactionSystem, and a parameter list of the internal format (vector of vectors):
+# Checks if any vertex parameter in the expression have a spatial component (that is, is not uniform).
+function has_spatial_vertex_component(exp, lrs::LatticeReactionSystem; u=nothing, vert_ps=nothing)
+    # Finds all the symbols in the expression.
+    exp_syms = Symbolics.get_variables(exp)
+
+    # If vertex parameter values where given, checks if any of these have non-uniform values.
+    if !isnothing(vert_ps)
+        exp_vert_ps = filter(sym -> any(isequal(sym), vertex_parameters(lrs)), exp_syms)
+        p_idxs = [ModelingToolkit.parameter_index(lrs.rs, sym) for sym in exp_vert_ps]
+        any(length(vert_ps[p_idx]) != 1 for p_idx in p_idxs) && return true
+    end
+
+    # If states values where given, checks if any of these have non-uniform values.
+    if !isnothing(u)
+        exp_u = filter(sym -> any(isequal(sym), species(lrs)), exp_syms)
+        u_idxs = [ModelingToolkit.variable_index(lrs.rs, sym) for sym in exp_u]
+        any(length(u[u_idx]) != 1 for u_idx in u_idxs) && return true
+    end
+
+    return false
+end

test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl

@@ -105,6 +105,62 @@ let
 end
 
 
+### SpatialMassActionJump Testing ###
+
+# Checks that the correct structure is produced.
+let 
+    # Network for reference:
+    # A, ∅ → X
+    # 1, 2X + Y → 3X
+    # B, X → Y
+    # 1, X → ∅
+    # srs = [@transport_reaction dX X]
+    # Create LatticeReactionSystem
+    lrs = LatticeReactionSystem(brusselator_system, brusselator_srs_1, small_3d_grid)
+
+    # Create JumpProblem
+    u0 = [:X => 1, :Y => rand(1:10, lrs.num_verts)]
+    tspan = (0.0, 100.0)
+    ps = [:A => 1.0, :B => 5.0 .+ rand(lrs.num_verts), :dX => rand(lrs.num_edges)]
+    dprob = DiscreteProblem(lrs, u0, tspan, ps)
+    jprob = JumpProblem(lrs, dprob, NSM())
+
+    # Checks internal structures.
+    jprob.massaction_jump.uniform_rates == [1.0, 0.5 ,10.] # 0.5 is due to combinatoric /2! in (2X + Y).
+    jprob.massaction_jump.spatial_rates[1,:] == ps[2][2]
+    # Test when new SII functions are ready, or we implement them in Catalyst.
+    # @test isequal(to_int(getfield.(reactions(lrs.rs), :netstoich)), jprob.massaction_jump.net_stoch)
+    # @test isequal(to_int(Pair.(getfield.(reactions(lrs.rs), :substrates),getfield.(reactions(lrs.rs), :substoich))), jprob.massaction_jump.net_stoch)
+
+    # Checks that problem can be simulated.
+    @test SciMLBase.successful_retcode(solve(jprob, SSAStepper()))
+end
+
+# Checks that simulations gives a correctly heterogeneous solution.
+let 
+    # Create model.
+    birth_death_network = @reaction_network begin
+        (p,d), 0 <--> X
+    end
+    srs = [(@transport_reaction D X)]
+    lrs = LatticeReactionSystem(birth_death_network, srs, very_small_2d_grid)
+    
+    # Create JumpProblem.
+    u0 = [:X => 1]
+    tspan = (0.0, 100.0)
+    ps = [:p => [0.1, 1.0, 10.0, 100.0], :d => 1.0, :D => 0.0]    
+    dprob = DiscreteProblem(lrs, u0, tspan, ps)
+    jprob = JumpProblem(lrs, dprob, NSM())
+
+    # Simulate model (a few repeats to ensure things don't succeed by change for uniform rates).
+    # Check that higher p gives higher mean.
+    for i = 1:5 
+        sol = solve(jprob, SSAStepper(); saveat = 1., seed = i*1234)
+        @test mean(getindex.(sol.u, 1)) < mean(getindex.(sol.u, 2)) < mean(getindex.(sol.u, 3)) < mean(getindex.(sol.u, 4))
+    end
+end
+
+
 ### Tests taken from JumpProcesses ###
 
 # ABC Model Test