From bfa8f907761dc7bc377a428de0b3d2027e253b3f Mon Sep 17 00:00:00 2001
From: Torkel <torkel.loman@gmail.com>
Date: Wed, 29 May 2024 19:47:10 -0400
Subject: [PATCH 01/15] comment out SI doc (broken on v1.10.3, need 1.10.2 or
 new Julia release)

---
 docs/pages.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/pages.jl b/docs/pages.jl
index fd9cfa18ab..abfd24c0f1 100644
--- a/docs/pages.jl
+++ b/docs/pages.jl
@@ -51,7 +51,7 @@ pages = Any[
         # ODE parameter fitting using Turing.
         # SDE/Jump fitting.
         "inverse_problems/behaviour_optimisation.md",
-        "inverse_problems/structural_identifiability.md",
+        #"inverse_problems/structural_identifiability.md",
         # Practical identifiability.
         "inverse_problems/global_sensitivity_analysis.md",
         "Inverse problem examples" => Any[

From c86c500aeaa1aa60262df0cd80a2b3c57da6fa48 Mon Sep 17 00:00:00 2001
From: Torkel <torkel.loman@gmail.com>
Date: Thu, 30 May 2024 13:39:38 -0400
Subject: [PATCH 02/15] up

---
 docs/pages.jl                                             | 4 +---
 .../steady_state_functionality/bifurcation_diagrams.md    | 8 ++++----
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/docs/pages.jl b/docs/pages.jl
index abfd24c0f1..260c491979 100644
--- a/docs/pages.jl
+++ b/docs/pages.jl
@@ -26,13 +26,11 @@ pages = Any[
     ],
     "Model simulation" => Any[
         "model_simulation/simulation_introduction.md",
-        # Simulation introduction.
         "model_simulation/simulation_plotting.md",
         "model_simulation/simulation_structure_interfacing.md",
         "model_simulation/ensemble_simulations.md",
         # Stochastic simulation statistical analysis.
         "model_simulation/ode_simulation_performance.md",
-        # ODE Performance considerations/advice.
         # SDE Performance considerations/advice.
         # Jump Performance considerations/advice.
         # Finite state projection
@@ -51,7 +49,7 @@ pages = Any[
         # ODE parameter fitting using Turing.
         # SDE/Jump fitting.
         "inverse_problems/behaviour_optimisation.md",
-        #"inverse_problems/structural_identifiability.md",
+        #"inverse_problems/structural_identifiability.md", # Broken on Julia v1.10.3, requires v1.10.2 or 1.10.4.
         # Practical identifiability.
         "inverse_problems/global_sensitivity_analysis.md",
         "Inverse problem examples" => Any[
diff --git a/docs/src/steady_state_functionality/bifurcation_diagrams.md b/docs/src/steady_state_functionality/bifurcation_diagrams.md
index a76ff91393..cebe51dc1c 100644
--- a/docs/src/steady_state_functionality/bifurcation_diagrams.md
+++ b/docs/src/steady_state_functionality/bifurcation_diagrams.md
@@ -43,7 +43,7 @@ nothing # hide
 
 Finally, we compute our bifurcation diagram using:
 ```@example ex1
-bif_dia = bifurcationdiagram(bprob, PALC(), 2, opts_br; bothside = true)
+bif_dia = bifurcationdiagram(bprob, PALC(), 2, (args...) -> opts_br; bothside = true)
 nothing # hide
 ```
 Where `PALC()` designates that we wish to use the pseudo arclength continuation method to track our solution. The third argument (`2`) designates the maximum number of recursions when branches of branches are computed (branches appear as continuation encounters certain bifurcation points). For diagrams with highly branched structures (rare for many common small chemical reaction networks) this input is important. Finally, `bothside = true` designates that we wish to perform continuation on both sides of the initial point (which is typically the case). 
@@ -69,7 +69,7 @@ opt_newton = NewtonPar(tol = 1e-9, max_iterations = 1000)
 opts_br = ContinuationPar(p_min = p_span[1], p_max = p_span[2], 
                           dsmin = 0.001, dsmax = 0.01, max_steps = 1000,
                           newton_options = opt_newton)
-bif_dia = bifurcationdiagram(bprob, PALC(), 2, opts_br; bothside=true)
+bif_dia = bifurcationdiagram(bprob, PALC(), 2, (args...) -> opts_br; bothside = true)
 nothing # hide
 ```
 (however, in this case these additional settings have no significant effect on the result)
@@ -79,7 +79,7 @@ Let's consider the previous case, but instead compute the bifurcation diagram ov
 ```@example ex1
 p_span = (2.0, 15.0)
 opts_br = ContinuationPar(p_min = p_span[1], p_max = p_span[2], max_steps = 1000)
-bif_dia = bifurcationdiagram(bprob, PALC(), 2, opts_br= true)
+bif_dia = bifurcationdiagram(bprob, PALC(), 2, (args...) -> opts_br; bothside = true)
 plot(bif_dia; xguide = "k1", yguide = "X")
 ```
 Here, in the bistable region, we only see a single branch. The reason is that the continuation algorithm starts at our initial guess (here made at $k1 = 4.0$ for $(X,Y) = (5.0,2.0)$) and tracks the diagram from there. However, with the upper bound set at $k1=15.0$ the bifurcation diagram has a disjoint branch structure, preventing the full diagram from being computed by continuation alone. In this case it could be solved by increasing the bound from $k1=15.0$, however, this is not possible in all cases. In these cases, *deflation* can be used. This is described in the [BifurcationKit documentation](https://bifurcationkit.github.io/BifurcationKitDocs.jl/dev/tutorials/tutorials2/#Snaking-computed-with-deflation).
@@ -103,7 +103,7 @@ bprob = BifurcationProblem(kinase_model, u_guess, p_start, :d; plot_var = :Xp, u
 
 p_span = (0.1, 10.0)
 opts_br = ContinuationPar(p_min = p_span[1], p_max = p_span[2], max_steps = 1000)
-bif_dia = bifurcationdiagram(bprob, PALC(), 2, opts_br; bothside = true)
+bif_dia = bifurcationdiagram(bprob, PALC(), 2, (args...) -> opts_br; bothside = true)
 plot(bif_dia; xguide = "d", yguide = "Xp")
 ```
 This bifurcation diagram does not contain any interesting features (such as bifurcation points), and only shows how the steady state concentration of $Xp$ is reduced as $d$ increases. 

From c5398de85d2999198388dde3c8fae6d56aa02de4 Mon Sep 17 00:00:00 2001
From: Torkel <torkel.loman@gmail.com>
Date: Fri, 31 May 2024 10:47:15 -0400
Subject: [PATCH 03/15] add DiffEq doc dependency again

---
 docs/Project.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/Project.toml b/docs/Project.toml
index 246727a23e..ebb086fda6 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -5,6 +5,7 @@ CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 Catalyst = "479239e8-5488-4da2-87a7-35f2df7eef83"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DiffEqParamEstim = "1130ab10-4a5a-5621-a13d-e4788d82bd4c"
+DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DynamicalSystems = "61744808-ddfa-5f27-97ff-6e42cc95d634"

From cf83b4a4b20022ba77d1d2d5cd2b6794b9eded31 Mon Sep 17 00:00:00 2001
From: Torkel <torkel.loman@gmail.com>
Date: Sat, 2 Dec 2023 12:52:26 -0500
Subject: [PATCH 04/15] init.

---
 src/Catalyst.jl                               |   9 +-
 .../lattice_jump_systems.jl                   |  81 +++++++++
 .../lattice_reaction_systems_jumps.jl         | 157 ++++++++++++++++++
 3 files changed, 243 insertions(+), 4 deletions(-)
 create mode 100644 src/spatial_reaction_systems/lattice_jump_systems.jl
 create mode 100644 test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl

diff --git a/src/Catalyst.jl b/src/Catalyst.jl
index d80c115119..ad889bef22 100644
--- a/src/Catalyst.jl
+++ b/src/Catalyst.jl
@@ -166,20 +166,21 @@ export make_si_ode
 
 ### 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
diff --git a/src/spatial_reaction_systems/lattice_jump_systems.jl b/src/spatial_reaction_systems/lattice_jump_systems.jl
new file mode 100644
index 0000000000..2ecfe3fa56
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl b/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl
new file mode 100644
index 0000000000..f41f1d3090
--- /dev/null
+++ b/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
\ No newline at end of file

From f19582be1d5c8aa5394c4e2a6b1b20b23ee12b81 Mon Sep 17 00:00:00 2001
From: Torkel <torkel.loman@gmail.com>
Date: Sat, 2 Dec 2023 13:00:22 -0500
Subject: [PATCH 05/15] test up

---
 .../lattice_reaction_systems_jumps.jl             | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl b/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl
index f41f1d3090..8b7db9c105 100644
--- a/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl
+++ b/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl
@@ -102,17 +102,19 @@ let
     end
 end
 
-### ABC Model Test (from JumpProcesses) ###
+### Tests taken from JumpProcesses ###
+
+# ABC Model Test
 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
+    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
+    domain_size = 1.0 # μ-meter.
     mesh_size = domain_size / linear_size
     rates = [0.1 / mesh_size, 1.0]
     diffusivity = 1.0
@@ -128,15 +130,16 @@ let
     lattice = Graphs.grid(dims)
     lrs = LatticeReactionSystem(rn, [tr_1, tr_2, tr_3], lattice)
 
-    # Set simulation parameters and create problems
+    # 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.
+    # NRM could be added, but doesn't work. Might need Cartesian grid.
+    jump_problems = [JumpProblem(lrs, dprob, alg(); save_positions = (false, false)) for alg in [NSM, DirectCRDirect]] 
 
-    # Tests.
+    # Run tests.
     function get_mean_end_state(jump_prob, Nsims)
         end_state = zeros(size(jump_prob.prob.u0))
         for i in 1:Nsims

From fcff2e8e6c3c498995ef3b3647fb1162f77bab34 Mon Sep 17 00:00:00 2001
From: Torkel <torkel.loman@gmail.com>
Date: Sat, 30 Dec 2023 11:55:20 +0100
Subject: [PATCH 06/15] up

---
 .../lattice_reaction_systems_jumps.jl                        | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl b/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl
index 8b7db9c105..1eb1566c30 100644
--- a/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl
+++ b/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl
@@ -7,7 +7,8 @@ using Random, Statistics, SparseArrays, Test
 # Fetch test networks.
 include("../spatial_test_networks.jl")
 
-### Correctness Tests ###
+
+### General Tests ###
 
 # Tests that there are no errors during runs for a variety of input forms.
 let
@@ -43,6 +44,7 @@ let
     end
 end
 
+
 ### Input Handling Tests ###
 
 # Tests that the correct hopping rates and initial conditions are generated.
@@ -102,6 +104,7 @@ let
     end
 end
 
+
 ### Tests taken from JumpProcesses ###
 
 # ABC Model Test

From 50c20beb5ea32ed93c2823e38b8bd60f67053515 Mon Sep 17 00:00:00 2001
From: Torkel <torkel.loman@gmail.com>
Date: Sun, 31 Dec 2023 13:44:32 +0100
Subject: [PATCH 07/15] init

---
 .../lattice_jump_systems.jl                        | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/spatial_reaction_systems/lattice_jump_systems.jl b/src/spatial_reaction_systems/lattice_jump_systems.jl
index 2ecfe3fa56..462c926294 100644
--- a/src/spatial_reaction_systems/lattice_jump_systems.jl
+++ b/src/spatial_reaction_systems/lattice_jump_systems.jl
@@ -36,9 +36,9 @@ function JumpProcesses.JumpProblem(lrs::LatticeReactionSystem, dprob, aggregator
     # 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)
+    sma_jumps = make_spatial_majumps(non_spat_dprob, dprob, lrs)
 
-    return JumpProblem(non_spat_dprob, aggregator, majumps; 
+    return JumpProblem(non_spat_dprob, aggregator, sma_jumps; 
                        hopping_constants, spatial_system = lrs.lattice, name, kwargs...)
 end
 
@@ -66,6 +66,12 @@ 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)
+    
+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.
@@ -74,8 +80,8 @@ function make_majumps(non_spat_dprob, rs::ReactionSystem)
     eqs = equations(js)
     invttype = non_spat_dprob.tspan[1] === nothing ? Float64 : typeof(1 / non_spat_dprob.tspan[2])
 
-    # Assembles the mass action jumps.
+    # 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
\ No newline at end of file
+end

From 38c28691f7e1d2d13b62be11ba1355e9e7b8e4e7 Mon Sep 17 00:00:00 2001
From: Torkel <torkel.loman@gmail.com>
Date: Wed, 3 Jan 2024 14:36:39 +0100
Subject: [PATCH 08/15] Use SpatialMassActionJump

---
 src/Catalyst.jl                               |  6 +-
 .../lattice_jump_systems.jl                   | 87 +++++++++++++----
 src/spatial_reaction_systems/utility.jl       | 96 +++++++++++++++++++
 .../lattice_reaction_systems_jumps.jl         | 56 +++++++++++
 4 files changed, 222 insertions(+), 23 deletions(-)

diff --git a/src/Catalyst.jl b/src/Catalyst.jl
index ad889bef22..2aab60fb54 100644
--- a/src/Catalyst.jl
+++ b/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
diff --git a/src/spatial_reaction_systems/lattice_jump_systems.jl b/src/spatial_reaction_systems/lattice_jump_systems.jl
index 462c926294..6287936889 100644
--- a/src/spatial_reaction_systems/lattice_jump_systems.jl
+++ b/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
diff --git a/src/spatial_reaction_systems/utility.jl b/src/spatial_reaction_systems/utility.jl
index 03271562e8..691ba79a3c 100644
--- a/src/spatial_reaction_systems/utility.jl
+++ b/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
\ No newline at end of file
diff --git a/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl b/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl
index 1eb1566c30..6721e707cd 100644
--- a/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl
+++ b/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

From 89fcdf30d9c2c3916378c69a9adf1805655153b1 Mon Sep 17 00:00:00 2001
From: Torkel <torkel.loman@gmail.com>
Date: Thu, 25 Jan 2024 17:42:34 -0500
Subject: [PATCH 09/15] test up

---
 .../lattice_reaction_systems_jumps.jl                       | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl b/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl
index 6721e707cd..8fc019d8bf 100644
--- a/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl
+++ b/test/spatial_reaction_systems/lattice_reaction_systems_jumps.jl
@@ -87,9 +87,8 @@ let
             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
+            @test all(issetequal(ns1, ns2) for (ns1, ns2) in zip(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]
@@ -97,9 +96,8 @@ let
             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
+            @test all(issetequal(ns1, ns2) for (ns1, ns2) in zip(jprob.massaction_jump.net_stoch, true_maj_net_stoch))
         end
     end
 end

From 16938691527001b74a5608ea2119586631e01d58 Mon Sep 17 00:00:00 2001
From: Torkel <torkel.loman@gmail.com>
Date: Thu, 16 May 2024 08:41:50 -0400
Subject: [PATCH 10/15] update for new MTK Parameter structure

---
 .../lattice_jump_systems.jl                      | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/src/spatial_reaction_systems/lattice_jump_systems.jl b/src/spatial_reaction_systems/lattice_jump_systems.jl
index 6287936889..818da05d1a 100644
--- a/src/spatial_reaction_systems/lattice_jump_systems.jl
+++ b/src/spatial_reaction_systems/lattice_jump_systems.jl
@@ -2,7 +2,9 @@
 
 # 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.")
+    if !is_transport_system(lrs) 
+        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.
@@ -18,20 +20,20 @@ function DiffEqBase.DiscreteProblem(lrs::LatticeReactionSystem, u0_in, tspan, p_
     # 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...)
+    return DiscreteProblem(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.
-    # 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?") 
-    
+    if !isnothing(dprob.f.sys)
+        error("Unexpected `DiscreteProblem` passed into `JumpProblem`. Was a `LatticeReactionSystem` used as input to the initial `DiscreteProblem`?")
+    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).

From 09bd10a0d365df5c67c2786d9f95fef1b0b2371c Mon Sep 17 00:00:00 2001
From: Torkel <torkel.loman@gmail.com>
Date: Fri, 31 May 2024 14:28:27 -0400
Subject: [PATCH 11/15] rebase fix

---
 test/runtests.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/runtests.jl b/test/runtests.jl
index 8f255921d7..9196ca91c8 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -61,6 +61,7 @@ using SafeTestsets, Test
         @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 "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
     #end
 
     #if GROUP == "All" || GROUP == "Visualisation-Extensions"

From 67ab88808507ec639466c2f50a0a53eaf3948324 Mon Sep 17 00:00:00 2001
From: Torkel <torkel.loman@gmail.com>
Date: Fri, 31 May 2024 16:02:51 -0400
Subject: [PATCH 12/15] up

---
 docs/src/assets/Project.toml                         | 1 +
 docs/src/model_simulation/simulation_introduction.md | 7 +++++--
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/docs/src/assets/Project.toml b/docs/src/assets/Project.toml
index 246727a23e..ebb086fda6 100644
--- a/docs/src/assets/Project.toml
+++ b/docs/src/assets/Project.toml
@@ -5,6 +5,7 @@ CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 Catalyst = "479239e8-5488-4da2-87a7-35f2df7eef83"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DiffEqParamEstim = "1130ab10-4a5a-5621-a13d-e4788d82bd4c"
+DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DynamicalSystems = "61744808-ddfa-5f27-97ff-6e42cc95d634"
diff --git a/docs/src/model_simulation/simulation_introduction.md b/docs/src/model_simulation/simulation_introduction.md
index 6234dc4a0a..81894ac74e 100644
--- a/docs/src/model_simulation/simulation_introduction.md
+++ b/docs/src/model_simulation/simulation_introduction.md
@@ -210,6 +210,9 @@ Next, let us consider a simulation for another parameter set:
 ```@example simulation_intro_sde
 sprob = remake(sprob; u0 = [:X1 => 100.0, :X2 => 200.0], p = [:k1 => 200.0, :k2 => 500.0])
 sol = solve(sprob, STrapezoid())
+nothing # hide
+```
+```@example simulation_intro_sde
 sol = solve(sprob, STrapezoid(); seed = 12345) # hide
 plot(sol)
 ```
@@ -317,9 +320,9 @@ plot(sol)
 ### [Designating aggregators and simulation methods for jump simulations](@id simulation_intro_jumps_solver_designation)
 Jump simulations (just like ODEs and SDEs) are performed using solver methods. Unlike ODEs and SDEs, jump simulations are carried out by two different types of methods acting in tandem. First, an *aggregator* method is used to (after each reaction) determine the time to, and type of, the next reaction. Next, a simulation method is used to actually carry out the simulation.
 
-Several different aggregators are available (a full list is provided [here](https://docs.sciml.ai/JumpProcesses/stable/jump_types/#Jump-Aggregators-for-Exact-Simulation)). To designate a specific one, provide it as the third argument to the `JumpProblem`. E.g. to designate that Gillespie's direct method (`Direct`) should be used, use:
+Several different aggregators are available (a full list is provided [here](https://docs.sciml.ai/JumpProcesses/stable/jump_types/#Jump-Aggregators-for-Exact-Simulation)). To designate a specific one, provide it as the third argument to the `JumpProblem`. E.g. to designate that the sorting direct method (`SortingDirect`) should be used, use:
 ```@example simulation_intro_jumps
-jprob = JumpProblem(two_state_model, dprob, Direct())
+jprob = JumpProblem(two_state_model, dprob, SortingDirect())
 nothing # hide
 ```
 Especially for large systems, the choice of aggregator is relevant to simulation performance. A guide for aggregator selection is provided [here](@ref ref).

From abd6dcafddb1ddabdd7598f4229c24401e77d0f7 Mon Sep 17 00:00:00 2001
From: Torkel <torkel.loman@gmail.com>
Date: Fri, 31 May 2024 16:03:41 -0400
Subject: [PATCH 13/15] up

---
 docs/pages.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/docs/pages.jl b/docs/pages.jl
index 96e6587bf3..941e906c3a 100644
--- a/docs/pages.jl
+++ b/docs/pages.jl
@@ -18,7 +18,7 @@ pages = Any[
         "model_creation/network_analysis.md",
         "model_creation/chemistry_related_functionality.md",
         "Model creation examples" => Any[
-            "model_creation/examples/basic_CRN_library.md",
+            #"model_creation/examples/basic_CRN_library.md",
             "model_creation/examples/programmatic_generative_linear_pathway.md",
             "model_creation/examples/hodgkin_huxley_equation.md",
             "model_creation/examples/smoluchowski_coagulation_equation.md"
@@ -29,9 +29,9 @@ pages = Any[
         # Simulation introduction.
         "model_simulation/simulation_plotting.md",
         "model_simulation/simulation_structure_interfacing.md",
-        "model_simulation/ensemble_simulations.md",
+        #"model_simulation/ensemble_simulations.md",
         # Stochastic simulation statistical analysis.
-        "model_simulation/ode_simulation_performance.md",
+        #"model_simulation/ode_simulation_performance.md",
         # ODE Performance considerations/advice.
         # SDE Performance considerations/advice.
         # Jump Performance considerations/advice.
@@ -39,7 +39,7 @@ pages = Any[
     ],
     "Steady state analysis" => Any[
         "steady_state_functionality/homotopy_continuation.md",
-        "steady_state_functionality/nonlinear_solve.md",
+        #"steady_state_functionality/nonlinear_solve.md",
         "steady_state_functionality/steady_state_stability_computation.md",        
         "steady_state_functionality/bifurcation_diagrams.md",
         "steady_state_functionality/dynamical_systems.md"

From b7047cf99109e5dbe5b65d80017329e392e155c5 Mon Sep 17 00:00:00 2001
From: Torkel <torkel.loman@gmail.com>
Date: Fri, 31 May 2024 17:35:06 -0400
Subject: [PATCH 14/15] up

---
 .github/workflows/Documentation.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/Documentation.yml b/.github/workflows/Documentation.yml
index 7e9af3e1dd..fc1871b19a 100644
--- a/.github/workflows/Documentation.yml
+++ b/.github/workflows/Documentation.yml
@@ -14,7 +14,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@latest
         with:
-          version: '1'
+          version: '1.10.2'
       - name: Install dependencies
         run: julia --project=docs/ -e 'ENV["JULIA_PKG_SERVER"] = ""; using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
       - name: Build and deploy

From 013d58d364de49f4142a98e5a15890fbc1899473 Mon Sep 17 00:00:00 2001
From: Torkel Loman <torkel.loman@gmail.com>
Date: Sat, 1 Jun 2024 04:49:58 +0200
Subject: [PATCH 15/15] Update pages.jl

comment out graphviz stuff
---
 docs/pages.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/pages.jl b/docs/pages.jl
index 1a54135768..bd716f4327 100644
--- a/docs/pages.jl
+++ b/docs/pages.jl
@@ -14,7 +14,7 @@ pages = Any[
         # Events.
         #"model_creation/parametric_stoichiometry.md",# Distributed parameters, rates, and initial conditions.
         # Loading and writing models to files.
-        "model_creation/model_visualisation.md",
+        #"model_creation/model_visualisation.md",
         #"model_creation/network_analysis.md",
         "model_creation/chemistry_related_functionality.md",
         "Model creation examples" => Any[
@@ -67,4 +67,4 @@ pages = Any[
     # ],
     #"FAQs" => "faqs.md",
     #"API" => "api.md"
-]
\ No newline at end of file
+]