update JumpProblem for latest version

Author: TorkelE

src/Catalyst.jl

@@ -183,6 +183,9 @@ export make_edge_p_values, make_directed_edge_values
 include("spatial_reaction_systems/spatial_ODE_systems.jl")
 include("spatial_reaction_systems/lattice_jump_systems.jl")
 
+# General spatial modelling utility functions.
+include("spatial_reaction_systems/utility.jl")
+
 ### ReactionSystem Serialisation ###
 # Has to be at the end (because it uses records of all metadata declared by Catalyst).
 include("reactionsystem_serialisation/serialisation_support.jl")

src/spatial_reaction_systems/lattice_jump_systems.jl

@@ -7,28 +7,25 @@ function DiffEqBase.DiscreteProblem(lrs::LatticeReactionSystem, u0_in, tspan,
         error("Currently lattice Jump simulations only supported when all spatial reactions are transport reactions.")
     end
 
-    # Converts potential symmaps to varmaps
-    # Vertex and edge parameters may be given in a tuple, or in a common vector, making parameter case complicated.
+    # Converts potential symmaps to varmaps.
     u0_in = symmap_to_varmap(lrs, u0_in)
-    p_in = (p_in isa Tuple{<:Any, <:Any}) ?
-           (symmap_to_varmap(lrs, p_in[1]), symmap_to_varmap(lrs, p_in[2])) :
-           symmap_to_varmap(lrs, p_in)
+    p_in = symmap_to_varmap(lrs, p_in)    
 
     # Converts u0 and p to their internal forms.
+    # u0 is simply a vector with all the species' initial condition values across all vertices.
     # u0 is [spec 1 at vert 1, spec 2 at vert 1, ..., spec 1 at vert 2, ...].
-    u0 = lattice_process_u0(u0_in, species(lrs), num_verts(lrs))                                   
-    # Both vert_ps and edge_ps becomes vectors of vectors. Each have 1 element for each parameter. 
-    # These elements are length 1 vectors (if the parameter is uniform), 
-    # or length num_verts/nE, with unique values for each vertex/edge (for vert_ps/edge_ps, respectively).
-    vert_ps, edge_ps = lattice_process_p(p_in, vertex_parameters(lrs),
-        edge_parameters(lrs), lrs)
-
-    # Returns a DiscreteProblem.
-    # Previously, a Tuple was used for (vert_ps, edge_ps), but this was converted to a Vector internally.
-    return DiscreteProblem(u0, tspan, [vert_ps, edge_ps], args...; kwargs...)
+    u0 = lattice_process_u0(u0_in, species(lrs), lrs)                                
+    # vert_ps and `edge_ps` are vector maps, taking each parameter's Symbolics representation to its value(s).
+    # vert_ps values are vectors. Here, index (i) is a parameter's value in vertex i.
+    # edge_ps values are sparse matrices. Here, index (i,j) is a parameter's value in the edge from vertex i to vertex j.
+    # Uniform vertex/edge parameters store only a single value (a length 1 vector, or size 1x1 sparse matrix).
+    vert_ps, edge_ps = lattice_process_p(p_in, vertex_parameters(lrs), edge_parameters(lrs), lrs)  
+
+    # Returns a DiscreteProblem (which basically just stores the processed input).
+    return DiscreteProblem(u0, tspan, [vert_ps; edge_ps], args...; kwargs...)
 end
 
-# Builds a spatial JumpProblem from a DiscreteProblem containg a Lattice Reaction System.
+# Builds a spatial JumpProblem from a DiscreteProblem containing a `LatticeReactionSystem`.
 function JumpProcesses.JumpProblem(lrs::LatticeReactionSystem, dprob, aggregator, args...; name = nameof(reactionsystem(lrs)), 
                                    combinatoric_ratelaws = get_combinatoric_ratelaws(reactionsystem(lrs)), kwargs...)
     # Error checks.
@@ -37,50 +34,53 @@ function JumpProcesses.JumpProblem(lrs::LatticeReactionSystem, dprob, aggregator
     end
 
     # Computes hopping constants and mass action jumps (requires some internal juggling).
-    # 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
+    # Hence the parameters of this one does not 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, num_species(lrs), num_verts(lrs)), dprob.tspan, first.(dprob.p[1]))
 
+    # Creates and returns a spatial JumpProblem (masked lattices are not supported by these).
+    spatial_system = has_masked_lattice(lrs) ? get_lattice_graph(lrs) : lattice(lrs)
     return JumpProblem(non_spat_dprob, aggregator, sma_jumps; 
-                       hopping_constants, spatial_system = lattice(lrs), name, kwargs...)
+                       hopping_constants, spatial_system , name, kwargs...)
 end
 
 # Creates the hopping constants from a discrete problem and a lattice reaction system.
 function make_hopping_constants(dprob::DiscreteProblem, lrs::LatticeReactionSystem)
     # Creates the all_diff_rates vector, containing for each species, its transport rate across all edges.
     # If transport rate is uniform for one species, the vector have a single element, else one for each edge.
-    spatial_rates_dict = Dict(compute_all_transport_rates(dprob.p[1], dprob.p[2], lrs))
+    spatial_rates_dict = Dict(compute_all_transport_rates(Dict(dprob.p), lrs))
     all_diff_rates = [haskey(spatial_rates_dict, s) ? spatial_rates_dict[s] : [0.0]
                       for s in species(lrs)]
 
-    # Creates the hopping constant Matrix. It contains one element for each combination of species and vertex.
-    # Each element is a Vector, containing the outgoing hopping rates for that species, from that vertex, on that edge.
-    hopping_constants = [Vector{Float64}(undef, length(lattice(lrs).fadjlist[j])) 
-                         for i in 1:(num_species(lrs)), j in 1:(num_verts(lrs))]
-
-    # For each edge, finds each position in `hopping_constants`.
-    for (e_idx, e) in enumerate(edges(lattice(lrs)))
-        dst_idx = findfirst(isequal(e.dst), lattice(lrs).fadjlist[e.src])      
-        # For each species, sets that hopping rate.  
-        for s_idx in 1:(num_species(lrs))
-            hopping_constants[s_idx, e.src][dst_idx] = get_component_value(all_diff_rates[s_idx], e_idx)
-        end
+    # Creates an array (of the same size as the hopping constant array) containing all edges.
+    # First the array is a NxM matrix (number of species x number of vertices). Each element is a 
+    # vector containing all edges leading out from that vertex (sorted by destination index).
+    edge_array = [Pair{Int64,Int64}[] for _1 in 1:num_species(lrs), _2 in 1:num_verts(lrs)]
+    for e in edge_iterator(lrs), s_idx in 1:num_species(lrs)
+        push!(edge_array[s_idx, e[1]], e)
     end
-
+    foreach(e_vec -> sort!(e_vec; by = e -> e[2]), edge_array)
+
+    # Creates the hopping constants array. It has the same shape as the edge array, but each 
+    # element is that species transportation rate along that edge
+    hopping_constants = [
+        [Catalyst.get_edge_value(all_diff_rates[s_idx], e) for e in edge_array[s_idx, src_idx]]
+        for s_idx in 1:num_species(lrs), src_idx in 1:num_verts(lrs)
+    ]
     return hopping_constants
 end
 
 # 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.
+# Not sure if there is any form of performance improvement from that though. Likely 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(reactionsystem(lrs))]
+    is_spatials = [has_spatial_vertex_component(rx.rate, dprob.p) 
+        for rx in reactions(reactionsystem(lrs))]
 
     # Creates templates for the rates (uniform and spatial) and the stoichiometries.
     # We cannot fetch reactant_stoich and net_stoich from a (non-spatial) MassActionJump.
@@ -91,10 +91,10 @@ function make_spatial_majumps(dprob, lrs::LatticeReactionSystem)
     net_stoich = Vector{Vector{Pair{Int64, Int64}}}(undef, length(reactions(reactionsystem(lrs))))
 
     # Loops through reactions with non-spatial rates, computes their rates and stoichiometries.
-    cur_rx = 1;
+    cur_rx = 1
     for (is_spat, rx) in zip(is_spatials, reactions(reactionsystem(lrs)))
         is_spat && continue
-        u_rates[cur_rx] = compute_vertex_value(rx.rate, lrs; vert_ps = dprob.p[1])[1]
+        u_rates[cur_rx] = compute_vertex_value(rx.rate, lrs; ps = dprob.p)[1]
         substoich_map = Pair.(rx.substrates, rx.substoich)
         reactant_stoich[cur_rx] = int_map(substoich_map, reactionsystem(lrs))
         net_stoich[cur_rx] = int_map(rx.netstoich, reactionsystem(lrs))
@@ -103,8 +103,7 @@ function make_spatial_majumps(dprob, lrs::LatticeReactionSystem)
     # Loops through reactions with spatial rates, computes their rates and stoichiometries.
     for (is_spat, rx) in zip(is_spatials, reactions(reactionsystem(lrs)))    
         is_spat || continue
-        s_rates[cur_rx - length(u_rates), :] = compute_vertex_value(rx.rate, lrs;
-            vert_ps = dprob.p[1])
+        s_rates[cur_rx - length(u_rates), :] .= compute_vertex_value(rx.rate, lrs; ps = dprob.p)
         substoich_map = Pair.(rx.substrates, rx.substoich)
         reactant_stoich[cur_rx] = int_map(substoich_map, reactionsystem(lrs))
         net_stoich[cur_rx] = int_map(rx.netstoich, reactionsystem(lrs))
@@ -114,7 +113,7 @@ function make_spatial_majumps(dprob, lrs::LatticeReactionSystem)
     isempty(u_rates) && (u_rates = nothing)
     (count(is_spatials) == 0) && (s_rates = nothing)
 
-    return SpatialMassActionJump(u_rates, s_rates, reactant_stoich, net_stoich)
+    return SpatialMassActionJump(u_rates, s_rates, reactant_stoich, net_stoich, nothing)
 end
 
 ### Extra ###

src/spatial_reaction_systems/lattice_reaction_systems.jl

@@ -307,6 +307,16 @@ grid_dims(lrs::LatticeReactionSystem) = grid_dims(lattice(lrs))
 grid_dims(lattice::GridLattice{N,T}) where {N,T} = return N
 grid_dims(lattice::Graph) = error("Grid dimensions is only defined for LatticeReactionSystems with grid-based lattices (not graph-based).")
 
+"""
+    get_lattice_graph(lrs::LatticeReactionSystem)
+
+Returns lrs's lattice, but in as a graph. Currently does not work for Cartesian lattices.
+"""
+function get_lattice_graph(lrs::LatticeReactionSystem)
+    has_graph_lattice(lrs) && return lattice(lrs)
+    return Graphs.SimpleGraphFromIterator(Graphs.SimpleEdge(e[1], e[2]) 
+        for e in edge_iterator(lrs))
+end
 
 ### Catalyst-based Getters ###
 
@@ -393,7 +403,7 @@ end
 D_vals = make_edge_p_values(lrs, make_edge_p_value)
 ```
 """
-function make_edge_p_values(lrs::LatticeReactionSystem, make_edge_p_value::Function, )
+function make_edge_p_values(lrs::LatticeReactionSystem, make_edge_p_value::Function)
     if has_graph_lattice(lrs)
         error("The `make_edge_p_values` function is only meant for lattices with (Cartesian or masked) grid structures. It cannot be applied to graph lattices.")
     end

src/spatial_reaction_systems/spatial_ODE_systems.jl

@@ -186,7 +186,7 @@ function DiffEqBase.ODEProblem(lrs::LatticeReactionSystem, u0_in, tspan,
     u0 = lattice_process_u0(u0_in, species(lrs), lrs)                                       
     # vert_ps and `edge_ps` are vector maps, taking each parameter's Symbolics representation to its value(s).
     # vert_ps values are vectors. Here, index (i) is a parameter's value in vertex i.
-    # edge_ps becomes a sparse matrix. Here, index (i,j) is a parameter's value in the edge from vertex i to vertex j.
+    # edge_ps values are sparse matrices. Here, index (i,j) is a parameter's value in the edge from vertex i to vertex j.
     # Uniform vertex/edge parameters store only a single value (a length 1 vector, or size 1x1 sparse matrix).
     # In the `ODEProblem` vert_ps and edge_ps are merged (but for building the ODEFunction, they are separate).
     vert_ps, edge_ps = lattice_process_p(p_in, vertex_parameters(lrs), edge_parameters(lrs), lrs)  
@@ -321,6 +321,9 @@ end
 
 # Defines the forcing functor's effect on the (spatial) ODE system.
 function (f_func::LatticeTransportODEf)(du, u, p, t)
+    println(du)
+    println(u)
+    println(p)
     # Updates for non-spatial reactions.
     for vert_i in 1:(f_func.num_verts)
         # Gets the indices of all the species at vertex i.

src/spatial_reaction_systems/utility.jl

@@ -327,36 +327,25 @@ end
 # 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)
+function compute_vertex_value(exp, lrs::LatticeReactionSystem; u = [], ps = [])
     # Finds the symbols in the expression. Checks that all correspond to unknowns 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.")
+        error("An edge parameter was encountered in expressions: $exp. Here, only vertex-based components are expected.")
     end
-    # Creates a Function tha computes the expressions value for a parameter set.
+
+    # Creates a Function that 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)
+    value_dict = Dict(vcat(u, ps))
 
     # 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]...)]
+    if all(length(value_dict[sym]) == 1 for sym in relevant_syms)
+        return [exp_func([value_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:num_verts(lrs)]
+    return [exp_func([get_vertex_value(value_dict[sym], vert_idx) for sym in relevant_syms]...) 
+                                                                            for vert_idx in 1:num_verts(lrs)]
 end
 
 ### System Property Checks ###
@@ -373,26 +362,10 @@ function has_spatial_edge_component(exp, lrs::LatticeReactionSystem, edge_ps)
     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(reactionsystem(lrs), 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(reactionsystem(lrs), sym) for sym in exp_u]
-        any(length(u[u_idx]) != 1 for u_idx in u_idxs) && return true
-    end
-
-    return false
+# For a Symbolic expression, and a parameter set, checks if any relevant parameters have a 
+# spatial component. Filters out any parameters that are edge parameters.
+function has_spatial_vertex_component(exp, ps)
+    relevant_syms = Symbolics.get_variables(exp)
+    value_dict = Dict(filter(p -> p[2] isa Vector, ps))
+    return any(length(value_dict[sym]) > 1 for sym in  relevant_syms)
 end

test/runtests.jl

@@ -9,6 +9,7 @@ using SafeTestsets, Test
 
 ### Run Tests ###
 @time begin
+    @time @safetestset "Jump Lattice Systems Simulations" begin include("spatial_modelling/lattice_reaction_systems_jumps.jl") end
 
     # Tests the `ReactionSystem` structure and its properties.
     @time @safetestset "Reaction Structure" begin include("reactionsystem_core/reaction.jl") end
@@ -53,8 +54,9 @@ using SafeTestsets, Test
     # Tests spatial modelling and simulations.
     @time @safetestset "PDE Systems Simulations" begin include("spatial_modelling/simulate_PDEs.jl") end
     @time @safetestset "Lattice Reaction Systems" begin include("spatial_modelling/lattice_reaction_systems.jl") end
+    @time @safetestset "Spatial Lattice Variants" begin include("spatial_modelling/lattice_reaction_systems_lattice_types.jl") end
     @time @safetestset "ODE Lattice Systems Simulations" begin include("spatial_modelling/lattice_reaction_systems_ODEs.jl") end
-    @time @safetestset "Jump Lattice Systems Simulations" begin include("spatial_reaction_systems/lattice_reaction_systems_jumps.jl") end
+    @time @safetestset "Jump Lattice Systems Simulations" begin include("spatial_modelling/lattice_reaction_systems_jumps.jl") end
 
     # Tests network visualisation.
     @time @safetestset "Latexify" begin include("visualisation/latexify.jl") end