init.

src/Catalyst.jl

@@ -114,20 +114,21 @@ export hc_steady_states
 
 ### Spatial Reaction Networks ###
 
-# spatial reactions
+# Spatial reactions.
 include("spatial_reaction_systems/spatial_reactions.jl")
 export TransportReaction, TransportReactions, @transport_reaction
 export isedgeparameter
 
-# lattice reaction systems
+# Lattice reaction systems
 include("spatial_reaction_systems/lattice_reaction_systems.jl")
 export LatticeReactionSystem
 export spatial_species, vertex_parameters, edge_parameters
 
-# variosu utility functions
+# Various utility functions
 include("spatial_reaction_systems/utility.jl")
 
-# spatial lattice ode systems.
+# Specific spatial problem types.
 include("spatial_reaction_systems/spatial_ODE_systems.jl")
+include("spatial_reaction_systems/lattice_jump_systems.jl")
 
 end # module

src/spatial_reaction_systems/lattice_jump_systems.jl

@@ -0,0 +1,81 @@
+### JumpProblem ###
+
+# Builds a spatial DiscreteProblem from a Lattice Reaction System.
+function DiffEqBase.DiscreteProblem(lrs::LatticeReactionSystem, u0_in, tspan, p_in = DiffEqBase.NullParameters(), args...; kwargs...)
+    is_transport_system(lrs) || error("Currently lattice Jump simulations only supported when all spatial reactions are transport reactions.")
+
+    # Converts potential symmaps to varmaps
+    # Vertex and edge parameters may be given in a tuple, or in a common vector, making parameter case complicated.
+    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)    
+
+    # Converts u0 and p to their internal forms.
+    # 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), lrs.num_verts)                                   
+    # 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(lrs.rs, u0, tspan, [vert_ps, edge_ps], args...; kwargs...)
+end
+
+# Builds a spatial JumpProblem from a DiscreteProblem containg a Lattice Reaction System.
+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.")
+
+    # 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])`).
+    hopping_constants = make_hopping_constants(dprob, lrs)
+    non_spat_dprob = DiscreteProblem(reshape(dprob.u0, lrs.num_species, lrs.num_verts), dprob.tspan, first.(dprob.p[1]))
+    majumps = make_majumps(non_spat_dprob, lrs.rs)
+
+    return JumpProblem(non_spat_dprob, aggregator, majumps; 
+                       hopping_constants, spatial_system = lrs.lattice, 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))
+    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(lrs.lattice.fadjlist[j])) 
+                         for i in 1:(lrs.num_species), j in 1:(lrs.num_verts)]
+
+    # For each edge, finds each position in `hopping_constants`.
+    for (e_idx, e) in enumerate(edges(lrs.lattice))
+        dst_idx = findfirst(isequal(e.dst), lrs.lattice.fadjlist[e.src])      
+        # For each species, sets that hopping rate.  
+        for s_idx in 1:(lrs.num_species)
+            hopping_constants[s_idx, e.src][dst_idx] = get_component_value(all_diff_rates[s_idx], e_idx)
+        end
+    end
+
+    return hopping_constants
+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 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

test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl

@@ -0,0 +1,157 @@
+### Preparations ###
+
+# Fetch packages.
+using JumpProcesses
+using Random, Statistics, SparseArrays, Test
+
+# Fetch test networks.
+include("../spatial_test_networks.jl")
+
+### Correctness Tests ###
+
+# Tests that there are no errors during runs for a variety of input forms.
+let
+    for grid in [small_2d_grid, short_path, small_directed_cycle]
+        for srs in [Vector{TransportReaction}(), SIR_srs_1, SIR_srs_2]
+            lrs = LatticeReactionSystem(SIR_system, srs, grid)
+            u0_1 = [:S => 999, :I => 1, :R => 0]
+            u0_2 = [:S => round.(Int64, 500.0 .+ 500.0 * rand_v_vals(lrs.lattice)), :I => 1, :R => 0, ]
+            u0_3 = [:S => 950, :I => round.(Int64, 50 * rand_v_vals(lrs.lattice)), :R => round.(Int64, 50 * rand_v_vals(lrs.lattice))]
+            u0_4 = [:S => round.(500.0 .+ 500.0 * rand_v_vals(lrs.lattice)), :I => round.(50 * rand_v_vals(lrs.lattice)), :R => round.(50 * rand_v_vals(lrs.lattice))]
+            u0_5 = make_u0_matrix(u0_3, vertices(lrs.lattice), map(s -> Symbol(s.f), species(lrs.rs)))
+            for u0 in [u0_1, u0_2, u0_3, u0_4, u0_5]
+                p1 = [:α => 0.1 / 1000, :β => 0.01]
+                p2 = [:α => 0.1 / 1000, :β => 0.02 * rand_v_vals(lrs.lattice)]
+                p3 = [
+                    :α => 0.1 / 2000 * rand_v_vals(lrs.lattice),
+                    :β => 0.02 * rand_v_vals(lrs.lattice),
+                ]
+                p4 = make_u0_matrix(p1, vertices(lrs.lattice), Symbol.(parameters(lrs.rs)))
+                for pV in [p1] #, p2, p3, p4] # Removed until spatial non-diffusion parameters are supported.
+                    pE_1 = map(sp -> sp => 0.01, ModelingToolkit.getname.(edge_parameters(lrs)))
+                    pE_2 = map(sp -> sp => 0.01, ModelingToolkit.getname.(edge_parameters(lrs)))
+                    pE_3 = map(sp -> sp => rand_e_vals(lrs.lattice, 0.01), ModelingToolkit.getname.(edge_parameters(lrs)))
+                    pE_4 = make_u0_matrix(pE_3, edges(lrs.lattice), ModelingToolkit.getname.(edge_parameters(lrs)))
+                    for pE in [pE_1, pE_2, pE_3, pE_4]
+                        dprob = DiscreteProblem(lrs, u0, (0.0, 100.0), (pV, pE))
+                        jprob = JumpProblem(lrs, dprob, NSM())
+                        @test SciMLBase.successful_retcode(solve(jprob, SSAStepper()))
+                    end
+                end
+            end
+        end
+    end
+end
+
+### Input Handling Tests ###
+
+# Tests that the correct hopping rates and initial conditions are generated.
+# In this base case, hopping rates should be on the form D_{s,i,j}.
+let
+    # Prepares the system.
+    lrs = LatticeReactionSystem(SIR_system, SIR_srs_2, small_2d_grid)
+
+    # Prepares various u0 input types.
+    u0_1 = [:I => 2.0, :S => 1.0, :R => 3.0]
+    u0_2 = [:I => fill(2., nv(small_2d_grid)), :S => 1.0, :R => 3.0]
+    u0_3 = [1.0, 2.0, 3.0]
+    u0_4 = [1.0, fill(2., nv(small_2d_grid)), 3.0]
+    u0_5 = permutedims(hcat(fill(1., nv(small_2d_grid)), fill(2., nv(small_2d_grid)), fill(3., nv(small_2d_grid))))
+
+    # Prepare various (compartment) parameter input types.
+    pV_1 = [:β => 0.2, :α => 0.1]
+    pV_2 = [:β => fill(0.2, nv(small_2d_grid)), :α => 1.0]
+    pV_3 = [0.1, 0.2]
+    pV_4 = [0.1, fill(0.2, nv(small_2d_grid))]
+    pV_5 = permutedims(hcat(fill(0.1, nv(small_2d_grid)), fill(0.2, nv(small_2d_grid))))
+
+    # Prepare various (diffusion) parameter input types.
+    pE_1 = [:dI => 0.02, :dS => 0.01, :dR => 0.03]
+    pE_2 = [:dI => 0.02, :dS => fill(0.01, ne(small_2d_grid)), :dR => 0.03]
+    pE_3 = [0.01, 0.02, 0.03]
+    pE_4 = [fill(0.01, ne(small_2d_grid)), 0.02, 0.03]
+    pE_5 = permutedims(hcat(fill(0.01, ne(small_2d_grid)), fill(0.02, ne(small_2d_grid)), fill(0.03, ne(small_2d_grid))))
+
+    # Checks hopping rates and u0 are correct.
+    true_u0 = [fill(1.0, 1, 25); fill(2.0, 1, 25); fill(3.0, 1, 25)]
+    true_hopping_rates = cumsum.([fill(dval, length(v)) for dval in [0.01,0.02,0.03], v in small_2d_grid.fadjlist])
+    true_maj_scaled_rates = [0.1, 0.2]
+    true_maj_reactant_stoch = [[1 => 1, 2 => 1], [2 => 1]]
+    true_maj_net_stoch = [[1 => -1, 2 => 1], [2 => -1, 3 => 1]]
+    for u0 in [u0_1, u0_2, u0_3, u0_4, u0_5]
+        # Provides parameters as a tupple.
+        for pV in [pV_1, pV_3], pE in [pE_1, pE_2, pE_3, pE_4, pE_5]
+            dprob = DiscreteProblem(lrs, u0, (0.0, 100.0), (pV,pE))
+            jprob = JumpProblem(lrs, dprob, NSM())
+            @test jprob.prob.u0 == true_u0
+            @test jprob.discrete_jump_aggregation.hop_rates.hop_const_cumulative_sums == true_hopping_rates
+            @test jprob.massaction_jump.scaled_rates == true_maj_scaled_rates
+            @test jprob.massaction_jump.reactant_stoch  == true_maj_reactant_stoch
+            @test jprob.massaction_jump.net_stoch == true_maj_net_stoch
+        end
+        # Provides parameters as a combined vector.
+        for pV in [pV_1], pE in [pE_1, pE_2]
+            dprob = DiscreteProblem(lrs, u0, (0.0, 100.0), [pE; pV])
+            jprob = JumpProblem(lrs, dprob, NSM())
+            @test jprob.prob.u0 == true_u0
+            @test jprob.discrete_jump_aggregation.hop_rates.hop_const_cumulative_sums == true_hopping_rates
+            @test jprob.massaction_jump.scaled_rates == true_maj_scaled_rates
+            @test jprob.massaction_jump.reactant_stoch  == true_maj_reactant_stoch
+            @test jprob.massaction_jump.net_stoch == true_maj_net_stoch
+        end
+    end
+end
+
+### ABC Model Test (from JumpProcesses) ###
+let 
+    # Preparations (stuff used in JumpProcesses examples ported over here, could be written directly into code).
+    Nsims = 100
+    reltol = 0.05
+    non_spatial_mean = [65.7395, 65.7395, 434.2605] #mean of 10,000 simulations
+    dim = 1
+    linear_size = 5
+    num_nodes = linear_size^dim
+    dims = Tuple(repeat([linear_size], dim))
+    domain_size = 1.0 #μ-meter
+    mesh_size = domain_size / linear_size
+    rates = [0.1 / mesh_size, 1.0]
+    diffusivity = 1.0
+    num_species = 3
+
+    # Make model.
+    rn = @reaction_network begin
+        (kB,kD), A + B <--> C
+    end
+    tr_1 = @transport_reaction D A
+    tr_2 = @transport_reaction D B
+    tr_3 = @transport_reaction D C
+    lattice = Graphs.grid(dims)
+    lrs = LatticeReactionSystem(rn, [tr_1, tr_2, tr_3], lattice)
+
+    # Set simulation parameters and create problems
+    u0 = [:A => [0,0,500,0,0], :B => [0,0,500,0,0], :C => 0]
+    tspan = (0.0, 10.0)
+    pV = [:kB => rates[1], :kD => rates[2]]
+    pE = [:D => diffusivity]
+    dprob = DiscreteProblem(lrs, u0, tspan, (pV, pE))
+    jump_problems = [JumpProblem(lrs, dprob, alg(); save_positions = (false, false)) for alg in [NSM, DirectCRDirect]] # NRM doesn't work. Might need Cartesian grid.
+
+    # Tests.
+    function get_mean_end_state(jump_prob, Nsims)
+        end_state = zeros(size(jump_prob.prob.u0))
+        for i in 1:Nsims
+            sol = solve(jump_prob, SSAStepper())
+            end_state .+= sol.u[end]
+        end
+        end_state / Nsims
+    end
+    for jprob in jump_problems
+        solution = solve(jprob, SSAStepper())
+        mean_end_state = get_mean_end_state(jprob, Nsims)
+        mean_end_state = reshape(mean_end_state, num_species, num_nodes)
+        diff = sum(mean_end_state, dims = 2) - non_spatial_mean
+        for (i, d) in enumerate(diff)
+            @test abs(d) < reltol * non_spatial_mean[i]
+        end
+    end
+end