test seems to pass

SciML · Feb 18, 2022 · 4ec4385 · 4ec4385
1 parent e35f8ef
commit 4ec4385
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 29 deletions.
diff --git a/src/spatial/flatten.jl b/src/spatial/flatten.jl
@@ -17,9 +17,11 @@ function flatten(ma_jump, prob::DiscreteProblem, spatial_system, hopping_constan
         rx_rates = ma_jump.scaled_rates
     elseif isa(ma_jump, SpatialMassActionJump)
         num_nodes = num_sites(spatial_system)
-        if isempty(ma_jump.uniform_rates)
+        if isnothing(ma_jump.uniform_rates) && isnothing(ma_jump.spatial_rates)
+            rx_rates = zeros(0,num_nodes)
+        elseif isnothing(ma_jump.uniform_rates)
             rx_rates = ma_jump.spatial_rates
-        elseif isempty(ma_jump.spatial_rates)
+        elseif isnothing(ma_jump.spatial_rates)
             rx_rates = reshape(repeat(ma_jump.uniform_rates, num_nodes), length(ma_jump.uniform_rates), num_nodes)
         else
             @assert size(ma_jump.spatial_rates, 2) == num_nodes

diff --git a/src/spatial/spatial_massaction_jump.jl b/src/spatial/spatial_massaction_jump.jl
@@ -1,51 +1,58 @@
-struct SpatialMassActionJump{F,S,U,V} <: AbstractMassActionJump
-    uniform_rates::Vector{F} # reactions that are uniform in space
-    spatial_rates::Matrix{F} # reactions whose rate depends on the site
+struct SpatialMassActionJump{A<:Union{AbstractVector,Nothing},B<:Union{AbstractMatrix,Nothing},S,U,V} <: AbstractMassActionJump
+    uniform_rates::A # reactions that are uniform in space
+    spatial_rates::B # reactions whose rate depends on the site
     reactant_stoch::S
     net_stoch::U
     param_mapper::V
 
-################ Constructors ##################
-
-"""
-uniform rates go first in ordering
-"""
-    function SpatialMassActionJump{F,S,U,V}(uniform_rates::Vector{F}, spatial_rates::Matrix{F}, reactant_stoch::S, net_stoch::U, param_mapper::V, scale_rates::Bool, useiszero::Bool, nocopy::Bool) where {F, S, U, V}
-        uniform_rates = nocopy ? uniform_rates : copy(uniform_rates)
-        spatial_rates = nocopy ? spatial_rates : copy(spatial_rates)
+    """
+    uniform rates go first in ordering
+    """
+    function SpatialMassActionJump{A,B,S,U,V}(uniform_rates::A, spatial_rates::B, reactant_stoch::S, net_stoch::U, param_mapper::V, scale_rates::Bool, useiszero::Bool, nocopy::Bool) where {A<:Union{AbstractVector,Nothing},B<:Union{AbstractMatrix,Nothing}, S, U, V}
+        uniform_rates = (nocopy || isnothing(uniform_rates)) ? uniform_rates : copy(uniform_rates)
+        spatial_rates = (nocopy || isnothing(spatial_rates)) ? spatial_rates : copy(spatial_rates)
         reactant_stoch = nocopy ? reactant_stoch : copy(reactant_stoch)
         for i in eachindex(reactant_stoch)
             if useiszero && (length(reactant_stoch[i]) == 1) && iszero(reactant_stoch[i][1][1])
-            reactant_stoch[i] = typeof(reactant_stoch[i])()
+                reactant_stoch[i] = typeof(reactant_stoch[i])()
             end
         end
-        if scale_rates && !isempty(uniform_rates)
+        num_unif_rates = isnothing(uniform_rates) ? 0 : length(uniform_rates)
+        if scale_rates && num_unif_rates > 0
             scalerates!(uniform_rates, reactant_stoch)
         end
-        if scale_rates && !isempty(spatial_rates)
-            scalerates!(spatial_rates, reactant_stoch[length(uniform_rates)+1:end])
+        if scale_rates && !isnothing(spatial_rates) && !isempty(spatial_rates)
+            scalerates!(spatial_rates, reactant_stoch[num_unif_rates+1:end])
         end
         new(uniform_rates, spatial_rates, reactant_stoch, net_stoch, param_mapper)
     end
 
 end
-SpatialMassActionJump(urates::Vector{F}, srates::Matrix{F}, rs::S, ns::U, pmapper::V; scale_rates = true, useiszero = true, nocopy=false) where {F,S,U,V} = SpatialMassActionJump{F,S,U,V}(urates, srates, rs, ns, pmapper, scale_rates, useiszero, nocopy)
-SpatialMassActionJump(urates::Vector{F}, srates::Matrix{F}, rs, ns; scale_rates = true, useiszero = true, nocopy=false) where {F} = SpatialMassActionJump(urates, srates, rs, ns, nothing; scale_rates=scale_rates, useiszero=useiszero, nocopy=nocopy)
 
-SpatialMassActionJump(srates::Matrix{F}, rs, ns, pmapper; scale_rates = true, useiszero = true, nocopy=false) where {F} = SpatialMassActionJump(zeros(F, 0), srates, rs, ns, pmapper; scale_rates = scale_rates, useiszero = useiszero, nocopy=nocopy)
-SpatialMassActionJump(srates::Matrix{F}, rs, ns; scale_rates = true, useiszero = true, nocopy=false) where {F} = SpatialMassActionJump(zeros(F, 0), srates, rs, ns, nothing; scale_rates = scale_rates, useiszero = useiszero, nocopy=nocopy)
+################ Constructors ##################
+
+SpatialMassActionJump(urates::A, srates::B, rs::S, ns::U, pmapper::V; scale_rates = true, useiszero = true, nocopy=false) where {A<:Union{AbstractVector,Nothing},B<:Union{AbstractMatrix,Nothing},S,U,V} = SpatialMassActionJump{A,B,S,U,V}(urates, srates, rs, ns, pmapper, scale_rates, useiszero, nocopy)
+SpatialMassActionJump(urates::A, srates::B, rs, ns; scale_rates = true, useiszero = true, nocopy=false) where {A<:Union{AbstractVector,Nothing},B<:Union{AbstractMatrix,Nothing}} = SpatialMassActionJump(urates, srates, rs, ns, nothing; scale_rates=scale_rates, useiszero=useiszero, nocopy=nocopy)
+
+SpatialMassActionJump(srates::B, rs, ns, pmapper; scale_rates = true, useiszero = true, nocopy=false) where {B<:Union{AbstractMatrix,Nothing}} = SpatialMassActionJump(nothing, srates, rs, ns, pmapper; scale_rates = scale_rates, useiszero = useiszero, nocopy=nocopy)
+SpatialMassActionJump(srates::B, rs, ns; scale_rates = true, useiszero = true, nocopy=false) where {B<:Union{AbstractMatrix,Nothing}} = SpatialMassActionJump(nothing, srates, rs, ns, nothing; scale_rates = scale_rates, useiszero = useiszero, nocopy=nocopy)
 
-SpatialMassActionJump(urates::Vector{F}, rs, ns, pmapper; scale_rates = true, useiszero = true, nocopy=false) where {F} = SpatialMassActionJump(urates, zeros(F, 0, 0), rs, ns, pmapper; scale_rates = scale_rates, useiszero = useiszero, nocopy=nocopy)
-SpatialMassActionJump(urates::Vector{F}, rs, ns; scale_rates = true, useiszero = true, nocopy=false) where {F} = SpatialMassActionJump(urates, zeros(F, 0, 0), rs, ns, nothing; scale_rates = scale_rates, useiszero = useiszero, nocopy=nocopy)
+SpatialMassActionJump(urates::A, rs, ns, pmapper; scale_rates = true, useiszero = true, nocopy=false) where {A<:Union{AbstractVector,Nothing}} = SpatialMassActionJump(urates, nothing, rs, ns, pmapper; scale_rates = scale_rates, useiszero = useiszero, nocopy=nocopy)
+SpatialMassActionJump(urates::A, rs, ns; scale_rates = true, useiszero = true, nocopy=false) where {A<:Union{AbstractVector,Nothing}} = SpatialMassActionJump(urates, nothing, rs, ns, nothing; scale_rates = scale_rates, useiszero = useiszero, nocopy=nocopy)
 
 SpatialMassActionJump(ma_jumps::MassActionJump{T,S,U,V}; scale_rates = true, useiszero = true, nocopy=false) where {T,S,U,V} = SpatialMassActionJump(ma_jumps.scaled_rates, ma_jumps.reactant_stoch, ma_jumps.net_stoch, ma_jumps.param_mapper; scale_rates = scale_rates, useiszero = useiszero, nocopy=nocopy)
 
 ##############################################
 
-get_num_majumps(spatial_majump::SpatialMassActionJump) = length(spatial_majump.uniform_rates) + size(spatial_majump.spatial_rates, 1)
+get_num_majumps(spatial_majump::SpatialMassActionJump{Nothing,Nothing,S,U,V}) where {S,U,V} = 0
+get_num_majumps(spatial_majump::SpatialMassActionJump{Nothing,B,S,U,V}) where {B,S,U,V} = size(spatial_majump.spatial_rates, 1)
+get_num_majumps(spatial_majump::SpatialMassActionJump{A,Nothing,S,U,V}) where {A,S,U,V} = length(spatial_majump.uniform_rates)
+get_num_majumps(spatial_majump::SpatialMassActionJump{A,B,S,U,V}) where {A<:AbstractVector,B<:AbstractMatrix,S,U,V} = length(spatial_majump.uniform_rates) + size(spatial_majump.spatial_rates, 1)
 using_params(spatial_majump::SpatialMassActionJump) = false
 
-function rate_at_site(rx, site, spatial_majump::SpatialMassActionJump)
+rate_at_site(rx, site, spatial_majump::SpatialMassActionJump{Nothing,B,S,U,V}) where {B,S,U,V} = spatial_majump.spatial_rates[rx, site]
+rate_at_site(rx, site, spatial_majump::SpatialMassActionJump{A,Nothing,S,U,V}) where {A,S,U,V} = spatial_majump.uniform_rates[rx]
+function rate_at_site(rx, site, spatial_majump::SpatialMassActionJump{A,B,S,U,V}) where {A<:AbstractVector,B<:AbstractMatrix,S,U,V}
     num_unif_rxs = length(spatial_majump.uniform_rates)
     rx <= num_unif_rxs ? spatial_majump.uniform_rates[rx] : spatial_majump.spatial_rates[rx-num_unif_rxs, site]
 end
@@ -60,6 +67,5 @@ function evalrxrate(speciesmat::AbstractMatrix{T}, rxidx::S, majump::SpatialMass
             val     *= specpop
         end
     end
-
     @inbounds return val * rate_at_site(rxidx, site, majump)
 end
diff --git a/test/spatial/spatial_majump.jl b/test/spatial/spatial_majump.jl
@@ -2,17 +2,17 @@ using DiffEqJump, DiffEqBase, OrdinaryDiffEq
 using Test, Graphs, LinearAlgebra
 
 reltol       = 0.05
-Nsims = 5*10^4
+Nsims = 10^4
 
 dim = 1
 linear_size = 5
 dims = Tuple(repeat([linear_size], dim))
 num_nodes = prod(dims)
 center_site = trunc(Int,(linear_size^dim + 1)/2)
 u0 = zeros(Int, 1, num_nodes)
-end_time = 1.0
+end_time = 100.0
 diffusivity = 1.0
-death_rate = 0.01
+death_rate = 0.1
 
 num_species = 1
 reactstoch = [Pair{Int64, Int64}[], [1 => 1]]