From aa9edad238b66a1f960795736f2f3cf6a4259bfe Mon Sep 17 00:00:00 2001
From: Vasily Ilin <vasilin97@gmail.com>
Date: Sat, 2 Sep 2023 22:28:14 -0400
Subject: [PATCH 01/12] Create `SpatialConstantRateJump`.

---
 src/spatial/spatial_constant_rate_jump.jl | 6 ++++++
 1 file changed, 6 insertions(+)
 create mode 100644 src/spatial/spatial_constant_rate_jump.jl

diff --git a/src/spatial/spatial_constant_rate_jump.jl b/src/spatial/spatial_constant_rate_jump.jl
new file mode 100644
index 00000000..77403475
--- /dev/null
+++ b/src/spatial/spatial_constant_rate_jump.jl
@@ -0,0 +1,6 @@
+struct SpatialConstantRateJump{F1, F2} <: AbstractJump
+    """Function `rate(u,p,t)` that returns the jump's current rate."""
+    rate::F1
+    """Function `affect(integrator)` that updates the state for one occurrence of the jump."""
+    affect!::F2
+end

From 49a6b495ea84796ba876a0b18977704a8b259071 Mon Sep 17 00:00:00 2001
From: Vasily Ilin <vasilin97@gmail.com>
Date: Sat, 2 Sep 2023 22:28:45 -0400
Subject: [PATCH 02/12] Hook up spatial constant rate jumps to `RxRates`.

---
 src/spatial/reaction_rates.jl | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/spatial/reaction_rates.jl b/src/spatial/reaction_rates.jl
index 53f3751d..38af6e16 100644
--- a/src/spatial/reaction_rates.jl
+++ b/src/spatial/reaction_rates.jl
@@ -3,7 +3,7 @@ A file with structs and functions for sampling reactions and updating reaction r
 """
 
 ### spatial rx rates ###
-struct RxRates{F, M}
+struct RxRates{F, M, C}
     "rx_rates[i,j] is rate of reaction i at site j"
     rates::Matrix{F}
 
@@ -12,18 +12,21 @@ struct RxRates{F, M}
 
     "AbstractMassActionJump"
     ma_jumps::M
+
+    cr_jumps::C
 end
 
 """
-    RxRates(num_sites::Int, ma_jumps::M) where {M}
+    RxRates(num_sites::Int, ma_jumps::M, cr_jumps::C) where {M, C}
 
 initializes RxRates with zero rates
 """
-function RxRates(num_sites::Int, ma_jumps::M) where {M}
+function RxRates(num_sites::Int, ma_jumps::M, cr_jumps::C) where {M, C}
     numrxjumps = get_num_majumps(ma_jumps)
     rates = zeros(Float64, numrxjumps, num_sites)
-    RxRates{Float64, M}(rates, vec(sum(rates, dims = 1)), ma_jumps)
+    RxRates{Float64, M}(rates, vec(sum(rates, dims = 1)), ma_jumps, cr_jumps)
 end
+RxRates(num_sites::Int, ma_jumps::M) where {M} = RxRates(num_sites, ma_jumps, ConstantRateJump[])
 
 num_rxs(rx_rates::RxRates) = get_num_majumps(rx_rates.ma_jumps)
 

From 964afaf3e7012ed34336f8252c43269f5d9cc23d Mon Sep 17 00:00:00 2001
From: Vasily Ilin <vasilin97@gmail.com>
Date: Sat, 2 Sep 2023 23:23:08 -0400
Subject: [PATCH 03/12] Add constant rate jump support to RxRates and utils.jl.

---
 src/spatial/reaction_rates.jl | 48 +++++++++++++++++++++++------------
 src/spatial/utils.jl          |  5 ++--
 2 files changed, 34 insertions(+), 19 deletions(-)

diff --git a/src/spatial/reaction_rates.jl b/src/spatial/reaction_rates.jl
index 38af6e16..149c3b92 100644
--- a/src/spatial/reaction_rates.jl
+++ b/src/spatial/reaction_rates.jl
@@ -1,5 +1,6 @@
 """
-A file with structs and functions for sampling reactions and updating reaction rates in spatial SSAs
+A file with structs and functions for sampling reactions and updating reaction rates in spatial SSAs.
+Massaction jumps go first in the indexing, then constant rate jumps.
 """
 
 ### spatial rx rates ###
@@ -13,6 +14,7 @@ struct RxRates{F, M, C}
     "AbstractMassActionJump"
     ma_jumps::M
 
+    "indexable collection of SpatialConstantRateJump"
     cr_jumps::C
 end
 
@@ -28,7 +30,7 @@ function RxRates(num_sites::Int, ma_jumps::M, cr_jumps::C) where {M, C}
 end
 RxRates(num_sites::Int, ma_jumps::M) where {M} = RxRates(num_sites, ma_jumps, ConstantRateJump[])
 
-num_rxs(rx_rates::RxRates) = get_num_majumps(rx_rates.ma_jumps)
+num_rxs(rx_rates::RxRates) = get_num_majumps(rx_rates.ma_jumps) + length(rx_rates.cr_jumps)
 
 """
     reset!(rx_rates::RxRates)
@@ -51,26 +53,22 @@ function total_site_rx_rate(rx_rates::RxRates, site)
 end
 
 """
-    update_rx_rates!(rx_rates, rxs, u, site)
+    update_rx_rates!(rx_rates, rxs, integrator, site)
 
 update rates of all reactions in rxs at site
 """
-function update_rx_rates!(rx_rates::RxRates{F, M}, rxs, integrator,
-                          site) where {F, M <: MassActionJump}
+function update_rx_rates!(rx_rates::RxRates, rxs, integrator,
+                          site)
     u = integrator.u
     ma_jumps = rx_rates.ma_jumps
     @inbounds for rx in rxs
-        set_rx_rate_at_site!(rx_rates, site, rx,
-                             evalrxrate((@view u[:, site]), rx, ma_jumps))
-    end
-end
-
-function update_rx_rates!(rx_rates::RxRates{F, M}, rxs, integrator,
-                          site) where {F, M <: SpatialMassActionJump}
-    u = integrator.u
-    ma_jumps = rx_rates.ma_jumps
-    @inbounds for rx in rxs
-        set_rx_rate_at_site!(rx_rates, site, rx, evalrxrate(u, rx, ma_jumps, site))
+        if is_massaction(rx_rates, rx)
+            rate = eval_massaction_rate(u, rx, ma_jumps, site)
+            set_rx_rate_at_site!(rx_rates, site, rx, rate)
+        else
+            cr_jump = rx_rates.cr_jumps[rx - get_num_majumps(ma_jumps)]
+            set_rx_rate_at_site!(rx_rates, site, rx, cr_jump.rate(u, integrator.p, integrator.t))
+        end
     end
 end
 
@@ -84,6 +82,16 @@ function sample_rx_at_site(rx_rates::RxRates, site, rng)
                   rand(rng) * total_site_rx_rate(rx_rates, site))
 end
 
+function execute_rx_at_site!(integrator, rx_rates::RxRates, rx, site)
+    if is_massaction(rx_rates, rx)
+        @inbounds executerx!((@view integrator.u[:, site]), rx,
+                             rx_rates.ma_jumps)
+    else
+        cr_jump = rx_rates.cr_jumps[rx - get_num_majumps(rx_rates.ma_jumps)]
+        cr_jump.affect!(integrator)
+    end
+end
+
 # helper functions
 function set_rx_rate_at_site!(rx_rates::RxRates, site, rx, rate)
     @inbounds old_rate = rx_rates.rates[rx, site]
@@ -96,3 +104,11 @@ function Base.show(io::IO, ::MIME"text/plain", rx_rates::RxRates)
     num_rxs, num_sites = size(rx_rates.rates)
     println(io, "RxRates with $num_rxs reactions and $num_sites sites")
 end
+
+"Return true if jump is a massaction jump."
+function is_massaction(rx_rates::RxRates, rx)
+    rx <= get_num_majumps(rx_rates.ma_jumps)
+end
+
+eval_massaction_rate(u, rx, ma_jumps::M, site) where {M <: SpatialMassActionJump} = evalrxrate(u, rx, ma_jumps, site)
+eval_massaction_rate(u, rx, ma_jumps::M, site) where {M <: MassActionJump} = evalrxrate((@view u[:, site]), rx, ma_jumps)
\ No newline at end of file
diff --git a/src/spatial/utils.jl b/src/spatial/utils.jl
index ddb9db41..d38f797b 100644
--- a/src/spatial/utils.jl
+++ b/src/spatial/utils.jl
@@ -9,7 +9,7 @@ struct SpatialJump{J}
     "source location"
     src::J
 
-    "index of jump as a hop or reaction"
+    "index of jump as a hop or reaction, hops first, then massaction reactions, then constant rate reactions"
     jidx::Int
 
     "destination location, equal to src for within-site reactions"
@@ -69,8 +69,7 @@ function update_state!(p, integrator)
         execute_hop!(integrator, jump.src, jump.dst, jump.jidx)
     else
         rx_index = reaction_id_from_jump(p, jump)
-        @inbounds executerx!((@view integrator.u[:, jump.src]), rx_index,
-                             p.rx_rates.ma_jumps)
+        execute_rx_at_site!(integrator, p.rx_rates, rx_index, jump.src)
     end
     # save jump that was just exectued
     p.prev_jump = jump

From 443dca3eb1533f00353dc2ddcaada34a9e3af6d9 Mon Sep 17 00:00:00 2001
From: Vasily Ilin <vasilin97@gmail.com>
Date: Sun, 3 Sep 2023 00:00:14 -0400
Subject: [PATCH 04/12] Add spatial constant rate jump to one unit test.

---
 src/JumpProcesses.jl                      |  3 ++-
 src/spatial/directcrdirect.jl             |  6 +++---
 src/spatial/nsm.jl                        |  6 +++---
 src/spatial/reaction_rates.jl             |  6 +++---
 src/spatial/spatial_constant_rate_jump.jl |  2 +-
 test/spatial/reaction_rates.jl            | 22 ++++++++++++++--------
 6 files changed, 26 insertions(+), 19 deletions(-)

diff --git a/src/JumpProcesses.jl b/src/JumpProcesses.jl
index d0476921..4f1c0150 100644
--- a/src/JumpProcesses.jl
+++ b/src/JumpProcesses.jl
@@ -56,6 +56,7 @@ include("aggregators/coevolve.jl")
 
 # spatial:
 include("spatial/spatial_massaction_jump.jl")
+include("spatial/spatial_constant_rate_jump.jl")
 include("spatial/topology.jl")
 include("spatial/hop_rates.jl")
 include("spatial/reaction_rates.jl")
@@ -98,7 +99,7 @@ export ExtendedJumpArray
 
 # spatial structs and functions
 export CartesianGrid, CartesianGridRej
-export SpatialMassActionJump
+export SpatialMassActionJump, SpatialConstantRateJump
 export outdegree, num_sites, neighbors
 export NSM, DirectCRDirect
 
diff --git a/src/spatial/directcrdirect.jl b/src/spatial/directcrdirect.jl
index a0b7ad3f..52f76640 100644
--- a/src/spatial/directcrdirect.jl
+++ b/src/spatial/directcrdirect.jl
@@ -16,8 +16,8 @@ mutable struct DirectCRDirectJumpAggregation{T, S, F1, F2, RNG, J, RX, HOP, DEPG
     rx_rates::RX
     hop_rates::HOP
     site_rates::Vector{T}
-    rates::F1 #rates for constant-rate jumps
-    affects!::F2 #affects! function determines the effect of constant-rate jumps
+    rates::F1 # legacy, not used
+    affects!::F2 # legacy, not used
     save_positions::Tuple{Bool, Bool}
     rng::RNG
     dep_gr::DEPGR #dep graph is same for each site
@@ -199,4 +199,4 @@ end
 
 number of constant rate jumps
 """
-num_constant_rate_jumps(aggregator::DirectCRDirectJumpAggregation) = 0
+num_constant_rate_jumps(aggregator::DirectCRDirectJumpAggregation) = length(aggregator.rx_rates.cr_jumps)
diff --git a/src/spatial/nsm.jl b/src/spatial/nsm.jl
index 92f3a2cf..a79c46fa 100644
--- a/src/spatial/nsm.jl
+++ b/src/spatial/nsm.jl
@@ -12,8 +12,8 @@ mutable struct NSMJumpAggregation{T, S, F1, F2, RNG, J, RX, HOP, DEPGR, VJMAP, J
     end_time::T
     rx_rates::RX
     hop_rates::HOP
-    rates::F1 #rates for constant-rate jumps
-    affects!::F2 #affects! function determines the effect of constant-rate jumps
+    rates::F1 # legacy, not used
+    affects!::F2 # legacy, not used
     save_positions::Tuple{Bool, Bool}
     rng::RNG
     dep_gr::DEPGR #dep graph is same for each site
@@ -187,4 +187,4 @@ end
 
 number of constant rate jumps
 """
-num_constant_rate_jumps(aggregator::NSMJumpAggregation) = 0
+num_constant_rate_jumps(aggregator::NSMJumpAggregation) = length(aggregator.rx_rates.cr_jumps)
diff --git a/src/spatial/reaction_rates.jl b/src/spatial/reaction_rates.jl
index 149c3b92..9a7af35d 100644
--- a/src/spatial/reaction_rates.jl
+++ b/src/spatial/reaction_rates.jl
@@ -24,9 +24,9 @@ end
 initializes RxRates with zero rates
 """
 function RxRates(num_sites::Int, ma_jumps::M, cr_jumps::C) where {M, C}
-    numrxjumps = get_num_majumps(ma_jumps)
+    numrxjumps = get_num_majumps(ma_jumps) + length(cr_jumps)
     rates = zeros(Float64, numrxjumps, num_sites)
-    RxRates{Float64, M}(rates, vec(sum(rates, dims = 1)), ma_jumps, cr_jumps)
+    RxRates{Float64, M, C}(rates, vec(sum(rates, dims = 1)), ma_jumps, cr_jumps)
 end
 RxRates(num_sites::Int, ma_jumps::M) where {M} = RxRates(num_sites, ma_jumps, ConstantRateJump[])
 
@@ -67,7 +67,7 @@ function update_rx_rates!(rx_rates::RxRates, rxs, integrator,
             set_rx_rate_at_site!(rx_rates, site, rx, rate)
         else
             cr_jump = rx_rates.cr_jumps[rx - get_num_majumps(ma_jumps)]
-            set_rx_rate_at_site!(rx_rates, site, rx, cr_jump.rate(u, integrator.p, integrator.t))
+            set_rx_rate_at_site!(rx_rates, site, rx, cr_jump.rate(u, integrator.p, integrator.t, site))
         end
     end
 end
diff --git a/src/spatial/spatial_constant_rate_jump.jl b/src/spatial/spatial_constant_rate_jump.jl
index 77403475..fdbb702b 100644
--- a/src/spatial/spatial_constant_rate_jump.jl
+++ b/src/spatial/spatial_constant_rate_jump.jl
@@ -1,5 +1,5 @@
 struct SpatialConstantRateJump{F1, F2} <: AbstractJump
-    """Function `rate(u,p,t)` that returns the jump's current rate."""
+    """Function `rate(u,p,t,site)` that returns the jump's current rate."""
     rate::F1
     """Function `affect(integrator)` that updates the state for one occurrence of the jump."""
     affect!::F2
diff --git a/test/spatial/reaction_rates.jl b/test/spatial/reaction_rates.jl
index e7e7fea4..373de410 100644
--- a/test/spatial/reaction_rates.jl
+++ b/test/spatial/reaction_rates.jl
@@ -9,8 +9,10 @@ const JP = JumpProcesses
 # sample_rx_at_site
 
 # Dummy integrator to test update_rx_rates!
-struct DummyIntegrator{S}
-    u::S
+struct DummyIntegrator{U,P,T}
+    u::U # state
+    p::P # parameters
+    t::T # time
 end
 
 io = IOBuffer()
@@ -22,22 +24,26 @@ num_species = 3
 reactstoch = [[1 => 1, 2 => 1], [3 => 1]]
 netstoch = [[1 => -1, 2 => -1, 3 => 1], [1 => 1, 2 => 1, 3 => -1]]
 rates = [0.1, 1.0]
-num_rxs = length(rates)
 ma_jumps = MassActionJump(rates, reactstoch, netstoch)
 spatial_ma_jumps = SpatialMassActionJump(rates, reactstoch, netstoch)
+rate_fn = (u, p, t, site) -> 1.0
+affect_fn!(integrator) = nothing # a dummy reaction, does nothing
+cr_jumps = [SpatialConstantRateJump(rate_fn, affect_fn!)]
+num_rxs = 3
 u = ones(Int, num_species, num_nodes)
-integrator = DummyIntegrator(u)
+integrator = DummyIntegrator(u, nothing, nothing)
 rng = StableRNG(12345)
 
 # Tests for RxRates
-rx_rates_list = [JP.RxRates(num_nodes, ma_jumps), JP.RxRates(num_nodes, spatial_ma_jumps)]
+rx_rates_list = [JP.RxRates(num_nodes, ma_jumps, cr_jumps), JP.RxRates(num_nodes, spatial_ma_jumps, cr_jumps)]
 for rx_rates in rx_rates_list
-    @test JP.num_rxs(rx_rates) == length(rates)
+    @test JP.num_rxs(rx_rates) == num_rxs
     show(io, "text/plain", rx_rates)
     for site in 1:num_nodes
         JP.update_rx_rates!(rx_rates, 1:num_rxs, integrator, site)
-        @test JP.total_site_rx_rate(rx_rates, site) == 1.1
-        rx_props = [JP.evalrxrate(u[:, site], rx, ma_jumps) for rx in 1:num_rxs]
+        @test JP.total_site_rx_rate(rx_rates, site) == 2.1
+        majump_props = [JP.evalrxrate(u[:, site], rx, ma_jumps) for rx in 1:2]
+        rx_props = [majump_props..., 1.0]
         rx_probs = rx_props / sum(rx_props)
         d = Dict{Int, Int}()
         for i in 1:num_samples

From eddb7bf4e8444d452c50e2e19a059ffd92caf604 Mon Sep 17 00:00:00 2001
From: Vasily Ilin <vasilin97@gmail.com>
Date: Sun, 3 Sep 2023 00:29:24 -0400
Subject: [PATCH 05/12] Add RxRates constructor from constant rate jumps only.

---
 src/spatial/reaction_rates.jl          | 3 ++-
 src/spatial/spatial_massaction_jump.jl | 3 +++
 test/spatial/reaction_rates.jl         | 7 ++++++-
 3 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/spatial/reaction_rates.jl b/src/spatial/reaction_rates.jl
index 9a7af35d..5cba241f 100644
--- a/src/spatial/reaction_rates.jl
+++ b/src/spatial/reaction_rates.jl
@@ -28,7 +28,8 @@ function RxRates(num_sites::Int, ma_jumps::M, cr_jumps::C) where {M, C}
     rates = zeros(Float64, numrxjumps, num_sites)
     RxRates{Float64, M, C}(rates, vec(sum(rates, dims = 1)), ma_jumps, cr_jumps)
 end
-RxRates(num_sites::Int, ma_jumps::M) where {M} = RxRates(num_sites, ma_jumps, ConstantRateJump[])
+RxRates(num_sites::Int, ma_jumps::M) where {M<:AbstractMassActionJump} = RxRates(num_sites, ma_jumps, ConstantRateJump[])
+RxRates(num_sites::Int, cr_jumps::C) where {C} = RxRates(num_sites, SpatialMassActionJump(nothing), cr_jumps)
 
 num_rxs(rx_rates::RxRates) = get_num_majumps(rx_rates.ma_jumps) + length(rx_rates.cr_jumps)
 
diff --git a/src/spatial/spatial_massaction_jump.jl b/src/spatial/spatial_massaction_jump.jl
index 34504f1b..559595a3 100644
--- a/src/spatial/spatial_massaction_jump.jl
+++ b/src/spatial/spatial_massaction_jump.jl
@@ -89,6 +89,9 @@ function SpatialMassActionJump(ma_jumps::MassActionJump{T, S, U, V}; scale_rates
                           scale_rates = scale_rates, useiszero = useiszero, nocopy = nocopy)
 end
 
+function SpatialMassActionJump(nothing)
+    SpatialMassActionJump([], [], [])
+end
 ##############################################
 
 function get_num_majumps(smaj::SpatialMassActionJump{Nothing, Nothing, S, U, V}) where
diff --git a/test/spatial/reaction_rates.jl b/test/spatial/reaction_rates.jl
index 373de410..719224f8 100644
--- a/test/spatial/reaction_rates.jl
+++ b/test/spatial/reaction_rates.jl
@@ -34,6 +34,11 @@ u = ones(Int, num_species, num_nodes)
 integrator = DummyIntegrator(u, nothing, nothing)
 rng = StableRNG(12345)
 
+# Test constructors
+@test JP.RxRates(num_nodes, ma_jumps).ma_jumps == ma_jumps
+@test JP.RxRates(num_nodes, spatial_ma_jumps).ma_jumps == spatial_ma_jumps
+@test JP.RxRates(num_nodes, cr_jumps).cr_jumps == cr_jumps
+
 # Tests for RxRates
 rx_rates_list = [JP.RxRates(num_nodes, ma_jumps, cr_jumps), JP.RxRates(num_nodes, spatial_ma_jumps, cr_jumps)]
 for rx_rates in rx_rates_list
@@ -58,4 +63,4 @@ for rx_rates in rx_rates_list
     for site in 1:num_nodes
         @test JP.total_site_rx_rate(rx_rates, site) == 0.0
     end
-end
+end
\ No newline at end of file

From d0115c5be151a7c220d17d1ba4bf6b7985d1b93e Mon Sep 17 00:00:00 2001
From: Vasily Ilin <vasilin97@gmail.com>
Date: Sun, 3 Sep 2023 00:45:33 -0400
Subject: [PATCH 06/12] Remove SpatialConstantRateJump.

Just use ConstantRateJump with `rate(u,p,t,site)`.
---
 src/JumpProcesses.jl                      | 2 +-
 src/spatial/reaction_rates.jl             | 2 +-
 src/spatial/spatial_constant_rate_jump.jl | 6 ------
 test/spatial/reaction_rates.jl            | 2 +-
 4 files changed, 3 insertions(+), 9 deletions(-)
 delete mode 100644 src/spatial/spatial_constant_rate_jump.jl

diff --git a/src/JumpProcesses.jl b/src/JumpProcesses.jl
index 4f1c0150..d9acd183 100644
--- a/src/JumpProcesses.jl
+++ b/src/JumpProcesses.jl
@@ -99,7 +99,7 @@ export ExtendedJumpArray
 
 # spatial structs and functions
 export CartesianGrid, CartesianGridRej
-export SpatialMassActionJump, SpatialConstantRateJump
+export SpatialMassActionJump
 export outdegree, num_sites, neighbors
 export NSM, DirectCRDirect
 
diff --git a/src/spatial/reaction_rates.jl b/src/spatial/reaction_rates.jl
index 5cba241f..070e73c4 100644
--- a/src/spatial/reaction_rates.jl
+++ b/src/spatial/reaction_rates.jl
@@ -14,7 +14,7 @@ struct RxRates{F, M, C}
     "AbstractMassActionJump"
     ma_jumps::M
 
-    "indexable collection of SpatialConstantRateJump"
+    "indexable collection of ConstantRateJump"
     cr_jumps::C
 end
 
diff --git a/src/spatial/spatial_constant_rate_jump.jl b/src/spatial/spatial_constant_rate_jump.jl
deleted file mode 100644
index fdbb702b..00000000
--- a/src/spatial/spatial_constant_rate_jump.jl
+++ /dev/null
@@ -1,6 +0,0 @@
-struct SpatialConstantRateJump{F1, F2} <: AbstractJump
-    """Function `rate(u,p,t,site)` that returns the jump's current rate."""
-    rate::F1
-    """Function `affect(integrator)` that updates the state for one occurrence of the jump."""
-    affect!::F2
-end
diff --git a/test/spatial/reaction_rates.jl b/test/spatial/reaction_rates.jl
index 719224f8..165b5ac6 100644
--- a/test/spatial/reaction_rates.jl
+++ b/test/spatial/reaction_rates.jl
@@ -28,7 +28,7 @@ ma_jumps = MassActionJump(rates, reactstoch, netstoch)
 spatial_ma_jumps = SpatialMassActionJump(rates, reactstoch, netstoch)
 rate_fn = (u, p, t, site) -> 1.0
 affect_fn!(integrator) = nothing # a dummy reaction, does nothing
-cr_jumps = [SpatialConstantRateJump(rate_fn, affect_fn!)]
+cr_jumps = [ConstantRateJump(rate_fn, affect_fn!)]
 num_rxs = 3
 u = ones(Int, num_species, num_nodes)
 integrator = DummyIntegrator(u, nothing, nothing)

From 3dc95965fcb9e2019c9a74cbfcb5ca66c88dad2e Mon Sep 17 00:00:00 2001
From: Vasily Ilin <vasilin97@gmail.com>
Date: Sun, 3 Sep 2023 00:46:08 -0400
Subject: [PATCH 07/12] Remove an `include`.

---
 src/JumpProcesses.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/JumpProcesses.jl b/src/JumpProcesses.jl
index d9acd183..d0476921 100644
--- a/src/JumpProcesses.jl
+++ b/src/JumpProcesses.jl
@@ -56,7 +56,6 @@ include("aggregators/coevolve.jl")
 
 # spatial:
 include("spatial/spatial_massaction_jump.jl")
-include("spatial/spatial_constant_rate_jump.jl")
 include("spatial/topology.jl")
 include("spatial/hop_rates.jl")
 include("spatial/reaction_rates.jl")

From 666d2cc98c52b5fe63b48a8dde54bd419733403c Mon Sep 17 00:00:00 2001
From: Vasily Ilin <vasilin97@gmail.com>
Date: Sun, 3 Sep 2023 01:05:33 -0400
Subject: [PATCH 08/12] Add spatial ConstantRateJump to ABC test.

---
 src/jumps.jl        |  1 +
 src/problem.jl      |  5 +++--
 test/spatial/ABC.jl | 18 ++++++++++++++++++
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/src/jumps.jl b/src/jumps.jl
index d80f3d92..0c70a73d 100644
--- a/src/jumps.jl
+++ b/src/jumps.jl
@@ -499,6 +499,7 @@ function JumpSet(vj, cj, rj, maj::MassActionJump{S, T, U, V}) where {S <: Number
 end
 
 JumpSet(jump::ConstantRateJump) = JumpSet((), (jump,), nothing, nothing)
+JumpSet(jumps::AbstractVector{ConstantRateJump}) = JumpSet((), jumps, nothing, nothing)
 JumpSet(jump::VariableRateJump) = JumpSet((jump,), (), nothing, nothing)
 JumpSet(jump::RegularJump) = JumpSet((), (), jump, nothing)
 JumpSet(jump::AbstractMassActionJump) = JumpSet((), (), nothing, jump)
diff --git a/src/problem.jl b/src/problem.jl
index 192c7b8a..109e2459 100644
--- a/src/problem.jl
+++ b/src/problem.jl
@@ -198,12 +198,13 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS
 
     ## Spatial jumps handling
     if spatial_system !== nothing && hopping_constants !== nothing
-        (num_crjs(jumps) == num_vrjs(jumps) == 0) ||
-            error("Spatial aggregators only support MassActionJumps currently.")
+        (num_vrjs(jumps) == 0) ||
+            error("Spatial aggregators currently only support MassActionJumps and ConstantRateJumps.")
 
         if is_spatial(aggregator)
             kwargs = merge((; hopping_constants, spatial_system), kwargs)
         else
+            # TODO(vilin97): Handle flattening of constant rate jumps.
             prob, maj = flatten(maj, prob, spatial_system, hopping_constants; kwargs...)
         end
     end
diff --git a/test/spatial/ABC.jl b/test/spatial/ABC.jl
index 558a701a..7acda4f9 100644
--- a/test/spatial/ABC.jl
+++ b/test/spatial/ABC.jl
@@ -24,6 +24,21 @@ netstoch = [[1 => -1, 2 => -1, 3 => 1], [1 => 1, 2 => 1, 3 => -1]]
 rates = [0.1 / mesh_size, 1.0]
 majumps = MassActionJump(rates, reactstoch, netstoch)
 
+# equivalent constant rate jumps
+rate1(u,p,t,site) = u[1,site]*u[2,site] / 2
+rate2(u,p,t,site) = u[3,site]
+affect1!(integrator) = begin
+    integrator.u[1] -= 1
+    integrator.u[2] -= 1
+    integrator.u[3] += 1
+end
+affect2!(integrator) = begin
+    integrator.u[1] += 1
+    integrator.u[2] += 1
+    integrator.u[3] -= 1
+end
+crjumps = JumpSet([ConstantRateJump(rate1, affect1!), ConstantRateJump(rate2, affect2!)])
+
 # spatial system setup
 hopping_rate = diffusivity * (linear_size / domain_size)^2
 
@@ -54,6 +69,9 @@ jump_problems = JumpProblem[JumpProblem(prob, NSM(), majumps,
 push!(jump_problems,
       JumpProblem(prob, DirectCRDirect(), majumps, hopping_constants = hopping_constants,
                   spatial_system = grids[1], save_positions = (false, false)))
+push!(jump_problems,
+JumpProblem(prob, DirectCRDirect(), crjumps, hopping_constants = hopping_constants,
+            spatial_system = grids[1], save_positions = (false, false)))
 # setup flattenned jump prob
 push!(jump_problems,
       JumpProblem(prob, NRM(), majumps, hopping_constants = hopping_constants,

From f0997cd2ecf7142b8555efc16479ef30f62e6df4 Mon Sep 17 00:00:00 2001
From: Vasily Ilin <vasilin97@gmail.com>
Date: Sun, 3 Sep 2023 01:56:11 -0400
Subject: [PATCH 09/12] Add constant rate jump problems to the `ABC.jl` test.

---
 src/spatial/directcrdirect.jl |  8 +++++++-
 src/spatial/nsm.jl            |  8 +++++++-
 src/spatial/reaction_rates.jl |  2 +-
 test/spatial/ABC.jl           | 26 +++++++++++++++-----------
 4 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/src/spatial/directcrdirect.jl b/src/spatial/directcrdirect.jl
index 52f76640..baf081f2 100644
--- a/src/spatial/directcrdirect.jl
+++ b/src/spatial/directcrdirect.jl
@@ -39,6 +39,9 @@ function DirectCRDirectJumpAggregation(nj::SpatialJump{J}, njt::T, et::T, rx_rat
 
     # a dependency graph is needed
     if dep_graph === nothing
+        if length(rx_rates.cr_jumps) != 0
+            error("Provide a dependency graph to use DirectCRDirect with constant rate jumps.")
+        end
         dg = make_dependency_graph(num_specs, rx_rates.ma_jumps)
     else
         dg = dep_graph
@@ -54,6 +57,9 @@ function DirectCRDirectJumpAggregation(nj::SpatialJump{J}, njt::T, et::T, rx_rat
     end
 
     if jumptovars_map === nothing
+        if length(rx_rates.cr_jumps) != 0
+            error("Provide a jump-to-species dependency graph to use DirectCRDirect with constant rate jumps.")
+        end
         jtov_map = jump_to_vars_map(rx_rates.ma_jumps)
     else
         jtov_map = jumptovars_map
@@ -94,7 +100,7 @@ function aggregate(aggregator::DirectCRDirect, starting_state, p, t, end_time,
 
     next_jump = SpatialJump{Int}(typemax(Int), typemax(Int), typemax(Int)) #a placeholder
     next_jump_time = typemax(typeof(end_time))
-    rx_rates = RxRates(num_sites(spatial_system), majumps)
+    rx_rates = RxRates(num_sites(spatial_system), majumps, constant_jumps)
     hop_rates = HopRates(hopping_constants, spatial_system)
     site_rates = zeros(typeof(end_time), num_sites(spatial_system))
 
diff --git a/src/spatial/nsm.jl b/src/spatial/nsm.jl
index a79c46fa..2ae29805 100644
--- a/src/spatial/nsm.jl
+++ b/src/spatial/nsm.jl
@@ -32,6 +32,9 @@ function NSMJumpAggregation(nj::SpatialJump{J}, njt::T, et::T, rx_rates::RX, hop
 
     # a dependency graph is needed
     if dep_graph === nothing
+        if length(rx_rates.cr_jumps) != 0
+            error("Provide a dependency graph to use NSM with constant rate jumps.")
+        end
         dg = make_dependency_graph(num_specs, rx_rates.ma_jumps)
     else
         dg = dep_graph
@@ -47,6 +50,9 @@ function NSMJumpAggregation(nj::SpatialJump{J}, njt::T, et::T, rx_rates::RX, hop
     end
 
     if jumptovars_map === nothing
+        if length(rx_rates.cr_jumps) != 0
+            error("Provide a jump-to-species dependency graph to use NSM with constant rate jumps.")
+        end
         jtov_map = jump_to_vars_map(rx_rates.ma_jumps)
     else
         jtov_map = jumptovars_map
@@ -83,7 +89,7 @@ function aggregate(aggregator::NSM, starting_state, p, t, end_time, constant_jum
 
     next_jump = SpatialJump{Int}(typemax(Int), typemax(Int), typemax(Int)) #a placeholder
     next_jump_time = typemax(typeof(end_time))
-    rx_rates = RxRates(num_sites(spatial_system), majumps)
+    rx_rates = RxRates(num_sites(spatial_system), majumps, constant_jumps)
     hop_rates = HopRates(hopping_constants, spatial_system)
 
     NSMJumpAggregation(next_jump, next_jump_time, end_time, rx_rates, hop_rates,
diff --git a/src/spatial/reaction_rates.jl b/src/spatial/reaction_rates.jl
index 070e73c4..7912a80d 100644
--- a/src/spatial/reaction_rates.jl
+++ b/src/spatial/reaction_rates.jl
@@ -89,7 +89,7 @@ function execute_rx_at_site!(integrator, rx_rates::RxRates, rx, site)
                              rx_rates.ma_jumps)
     else
         cr_jump = rx_rates.cr_jumps[rx - get_num_majumps(rx_rates.ma_jumps)]
-        cr_jump.affect!(integrator)
+        cr_jump.affect!(integrator, site)
     end
 end
 
diff --git a/test/spatial/ABC.jl b/test/spatial/ABC.jl
index 7acda4f9..05c88272 100644
--- a/test/spatial/ABC.jl
+++ b/test/spatial/ABC.jl
@@ -27,17 +27,19 @@ majumps = MassActionJump(rates, reactstoch, netstoch)
 # equivalent constant rate jumps
 rate1(u,p,t,site) = u[1,site]*u[2,site] / 2
 rate2(u,p,t,site) = u[3,site]
-affect1!(integrator) = begin
-    integrator.u[1] -= 1
-    integrator.u[2] -= 1
-    integrator.u[3] += 1
+affect1!(integrator,site) = begin
+    integrator.u[1, site] -= 1
+    integrator.u[2, site] -= 1
+    integrator.u[3, site] += 1
 end
-affect2!(integrator) = begin
-    integrator.u[1] += 1
-    integrator.u[2] += 1
-    integrator.u[3] -= 1
+affect2!(integrator,site) = begin
+    integrator.u[1, site] += 1
+    integrator.u[2, site] += 1
+    integrator.u[3, site] -= 1
 end
 crjumps = JumpSet([ConstantRateJump(rate1, affect1!), ConstantRateJump(rate2, affect2!)])
+dep_graph = [[1,2],[1,2]]
+jumptovars_map = [[1,2,3],[1,2,3]]
 
 # spatial system setup
 hopping_rate = diffusivity * (linear_size / domain_size)^2
@@ -69,9 +71,11 @@ jump_problems = JumpProblem[JumpProblem(prob, NSM(), majumps,
 push!(jump_problems,
       JumpProblem(prob, DirectCRDirect(), majumps, hopping_constants = hopping_constants,
                   spatial_system = grids[1], save_positions = (false, false)))
-push!(jump_problems,
-JumpProblem(prob, DirectCRDirect(), crjumps, hopping_constants = hopping_constants,
-            spatial_system = grids[1], save_positions = (false, false)))
+# setup constant rate jump problems
+push!(jump_problems, JumpProblem(prob, NSM(), crjumps, hopping_constants = hopping_constants,
+            spatial_system = CartesianGrid(dims), save_positions = (false, false), dep_graph = dep_graph, jumptovars_map = jumptovars_map))
+push!(jump_problems, JumpProblem(prob, DirectCRDirect(), crjumps, hopping_constants = hopping_constants,
+            spatial_system = CartesianGrid(dims), save_positions = (false, false), dep_graph = dep_graph, jumptovars_map = jumptovars_map))
 # setup flattenned jump prob
 push!(jump_problems,
       JumpProblem(prob, NRM(), majumps, hopping_constants = hopping_constants,

From 0614fd95d5652da3805c2c8a140023cf49d97b85 Mon Sep 17 00:00:00 2001
From: Vasily Ilin <vasilin97@gmail.com>
Date: Sun, 3 Sep 2023 02:01:56 -0400
Subject: [PATCH 10/12] Add an error message that flattening is not supported
 for costant rate jumps.

---
 src/problem.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/problem.jl b/src/problem.jl
index 109e2459..3ad21375 100644
--- a/src/problem.jl
+++ b/src/problem.jl
@@ -204,7 +204,9 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS
         if is_spatial(aggregator)
             kwargs = merge((; hopping_constants, spatial_system), kwargs)
         else
-            # TODO(vilin97): Handle flattening of constant rate jumps.
+            if num_crjs(jumps) != 0
+                error("Use a spatial SSA, e.g. DirectCRDirect in order to use ConstantRateJumps.")
+            end 
             prob, maj = flatten(maj, prob, spatial_system, hopping_constants; kwargs...)
         end
     end

From ea26ce298ba30cb4bc64d50d18ef421d65664a79 Mon Sep 17 00:00:00 2001
From: Vasily Ilin <vasilin97@gmail.com>
Date: Sat, 9 Sep 2023 00:16:39 -0400
Subject: [PATCH 11/12] Change empty spatial massaction jump to use concrete
 types.

---
 src/spatial/reaction_rates.jl          | 2 +-
 src/spatial/spatial_massaction_jump.jl | 7 +++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/spatial/reaction_rates.jl b/src/spatial/reaction_rates.jl
index 7912a80d..d2c26c4d 100644
--- a/src/spatial/reaction_rates.jl
+++ b/src/spatial/reaction_rates.jl
@@ -29,7 +29,7 @@ function RxRates(num_sites::Int, ma_jumps::M, cr_jumps::C) where {M, C}
     RxRates{Float64, M, C}(rates, vec(sum(rates, dims = 1)), ma_jumps, cr_jumps)
 end
 RxRates(num_sites::Int, ma_jumps::M) where {M<:AbstractMassActionJump} = RxRates(num_sites, ma_jumps, ConstantRateJump[])
-RxRates(num_sites::Int, cr_jumps::C) where {C} = RxRates(num_sites, SpatialMassActionJump(nothing), cr_jumps)
+RxRates(num_sites::Int, cr_jumps::C) where {C} = RxRates(num_sites, SpatialMassActionJump(), cr_jumps)
 
 num_rxs(rx_rates::RxRates) = get_num_majumps(rx_rates.ma_jumps) + length(rx_rates.cr_jumps)
 
diff --git a/src/spatial/spatial_massaction_jump.jl b/src/spatial/spatial_massaction_jump.jl
index 559595a3..2a5dd437 100644
--- a/src/spatial/spatial_massaction_jump.jl
+++ b/src/spatial/spatial_massaction_jump.jl
@@ -89,8 +89,11 @@ function SpatialMassActionJump(ma_jumps::MassActionJump{T, S, U, V}; scale_rates
                           scale_rates = scale_rates, useiszero = useiszero, nocopy = nocopy)
 end
 
-function SpatialMassActionJump(nothing)
-    SpatialMassActionJump([], [], [])
+function SpatialMassActionJump()
+    empty_majump = MassActionJump(Vector{Float64}(),
+            Vector{Vector{Pair{Int, Int}}}(),
+            Vector{Vector{Pair{Int, Int}}}())
+    SpatialMassActionJump(empty_majump)
 end
 ##############################################
 

From 6f1ca89580e0e235ece750c7a9a871c397a50ca9 Mon Sep 17 00:00:00 2001
From: Vasily Ilin <vasilin97@gmail.com>
Date: Mon, 11 Sep 2023 21:39:56 -0400
Subject: [PATCH 12/12] Fix two typos.

---
 src/spatial/reaction_rates.jl  | 2 +-
 test/spatial/reaction_rates.jl | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/spatial/reaction_rates.jl b/src/spatial/reaction_rates.jl
index d2c26c4d..5653fe72 100644
--- a/src/spatial/reaction_rates.jl
+++ b/src/spatial/reaction_rates.jl
@@ -112,4 +112,4 @@ function is_massaction(rx_rates::RxRates, rx)
 end
 
 eval_massaction_rate(u, rx, ma_jumps::M, site) where {M <: SpatialMassActionJump} = evalrxrate(u, rx, ma_jumps, site)
-eval_massaction_rate(u, rx, ma_jumps::M, site) where {M <: MassActionJump} = evalrxrate((@view u[:, site]), rx, ma_jumps)
\ No newline at end of file
+eval_massaction_rate(u, rx, ma_jumps::M, site) where {M <: MassActionJump} = evalrxrate((@view u[:, site]), rx, ma_jumps)
diff --git a/test/spatial/reaction_rates.jl b/test/spatial/reaction_rates.jl
index a985f604..a95d04e0 100644
--- a/test/spatial/reaction_rates.jl
+++ b/test/spatial/reaction_rates.jl
@@ -63,4 +63,4 @@ for rx_rates in rx_rates_list
     for site in 1:num_nodes
         @test JP.total_site_rx_rate(rx_rates, site) == 0.0
     end
-end
\ No newline at end of file
+end