fix for LatticeReactionSystems

Author: TorkelE

src/spatial_reaction_systems/lattice_reaction_systems.jl

@@ -25,17 +25,17 @@ struct LatticeReactionSystem{S,T} # <: MT.AbstractTimeDependentSystem
     All parameters related to the lattice reaction system
     (both with spatial and non-spatial effects).
     """
-    parameters::Vector{BasicSymbolic{Real}}
+    parameters::Vector{Any}
     """
     Parameters which values are tied to vertexes (adjacencies), 
     e.g. (possibly) have an unique value at each vertex of the system.
     """
-    vertex_parameters::Vector{BasicSymbolic{Real}}
+    vertex_parameters::Vector{Any}
     """
     Parameters which values are tied to edges (adjacencies), 
     e.g. (possibly) have an unique value at each edge of the system.
     """
-    edge_parameters::Vector{BasicSymbolic{Real}}
+    edge_parameters::Vector{Any}
 
     function LatticeReactionSystem(rs::ReactionSystem{S}, spatial_reactions::Vector{T},
                                    lattice::DiGraph; init_digraph = true) where {S, T}
@@ -52,7 +52,7 @@ struct LatticeReactionSystem{S,T} # <: MT.AbstractTimeDependentSystem
         num_species = length(unique([species(rs); spat_species]))
         rs_edge_parameters = filter(isedgeparameter, parameters(rs))
         if isempty(spatial_reactions)
-            srs_edge_parameters = Vector{BasicSymbolic{Real}}[]
+            srs_edge_parameters = Vector{Any}[]
         else
             srs_edge_parameters = setdiff(reduce(vcat, [parameters(sr) for sr in spatial_reactions]), parameters(rs))
         end

src/spatial_reaction_systems/spatial_reactions.jl

@@ -95,7 +95,7 @@ function check_spatial_reaction_validity(rs::ReactionSystem, tr::TransportReacti
     if any([isequal(tr.species, s) && !isequivalent(tr.species, s) for s in species(rs)]) 
         error("A transport reaction used a species, $(tr.species), with metadata not matching its lattice reaction system. Please fetch this species from the reaction system and used in transport reaction creation.")
     end
-    if any([isequal(rs_p, tr_p) && !equivalent_metadata(rs_p, tr_p) 
+    if any([isequal(rs_p, tr_p) && !isequivalent(rs_p, tr_p) 
             for rs_p in parameters(rs), tr_p in Symbolics.get_variables(tr.rate)]) 
         error("A transport reaction used a parameter with metadata not matching its lattice reaction system. Please fetch this parameter from the reaction system and used in transport reaction creation.")
     end
@@ -105,15 +105,16 @@ function check_spatial_reaction_validity(rs::ReactionSystem, tr::TransportReacti
         error("Edge paramter(s) were found as a rate of a non-spatial reaction.")
     end
 end
-equivalent_metadata(p1, p2) = isempty(setdiff(p1.metadata, p2.metadata, [Catalyst.EdgeParameter => true]))
 
 # Since MTK's "isequal" ignores metadata, we have to use a special function that accounts for this.
 # This is important because whether something is an edge parameter is defined in metadata.
 function isequivalent(sym1, sym2)
-    !isequal(sym1, sym2) && (return false)
-    (sym1.metadata != sym2.metadata) && (return false)
+    isequal(sym1, sym2) || (return false)
+    (ignore_ep_metadata(sym1.metadata) != ignore_ep_metadata(sym2.metadata)) && (return false)
+    (typeof(sym1) != typeof(sym2)) && (return false)
     return true
 end
+ignore_ep_metadata(metadata) = setdiff(metadata, [Catalyst.EdgeParameter => true])
 
 # Implements custom `==`.
 """

test/dsl/dsl_options.jl

@@ -4,6 +4,7 @@
 
 # Fetch packages.
 using Catalyst, ModelingToolkit, OrdinaryDiffEq, Plots, Test
+using Symbolics: unwrap
 
 # Sets the default `t` to use.
 t = default_t()
@@ -344,6 +345,40 @@ let
     @test merge(ModelingToolkit.defaults(rn28), defs28) == ModelingToolkit.defaults(rn27)
 end
 
+# Tests that parameter type designation works.
+let
+    rn = @reaction_network begin
+        @parameters begin
+            k1
+            l1
+            k2::Float64 = 2.0
+            l2::Float64
+            k3::Int64 = 2, [description="A parameter"]
+            l3::Int64
+            k4::Float32, [description="Another parameter"]
+            l4::Float32
+            k5::Rational{Int64}
+            l5::Rational{Int64}
+        end
+        (k1,l1), X1 <--> Y1
+        (k2,l2), X2 <--> Y2
+        (k3,l3), X3 <--> Y3
+        (k4,l4), X4 <--> Y4
+        (k5,l5), X5 <--> Y5
+    end
+    
+    @test unwrap(rn.k1) isa SymbolicUtils.BasicSymbolic{Real}
+    @test unwrap(rn.l1) isa SymbolicUtils.BasicSymbolic{Real}
+    @test unwrap(rn.k2) isa SymbolicUtils.BasicSymbolic{Float64}
+    @test unwrap(rn.l2) isa SymbolicUtils.BasicSymbolic{Float64}
+    @test unwrap(rn.k3) isa SymbolicUtils.BasicSymbolic{Int64}
+    @test unwrap(rn.l3) isa SymbolicUtils.BasicSymbolic{Int64}
+    @test unwrap(rn.k4) isa SymbolicUtils.BasicSymbolic{Float32}
+    @test unwrap(rn.l4) isa SymbolicUtils.BasicSymbolic{Float32}
+    @test unwrap(rn.k5) isa SymbolicUtils.BasicSymbolic{Rational{Int64}}
+    @test unwrap(rn.l5) isa SymbolicUtils.BasicSymbolic{Rational{Int64}}
+end
+
 ### Observables ###
 
 # Test basic functionality.

test/reactionsystem_structure/designating_parameter_types.jl

@@ -0,0 +1,106 @@
+### Fetch Packages and Set Global Variables ###
+
+# Fetch packages.
+using Catalyst, JumpProcesses, NonlinearSolve, OrdinaryDiffEq, StochasticDiffEq, Test
+using Symbolics: unwrap
+
+# Sets stable rng number.
+using StableRNGs
+rng = StableRNG(12345)
+seed = rand(rng, 1:100)
+
+# Sets the default `t` to use.
+t = default_t()
+
+### Basic Tests ###
+
+# Declares a simple model to run tests on.
+begin
+    t = default_t()
+    @parameters p1 p2::Float64 p3::Int64 p4::Float32 p5::Rational{Int64}
+    @parameters d1 d2::Float64 = 1.2 d3::Int64 = 2 [description = "A parameter"] d4::Rational{Int64} d5::Float32
+    @species X1(t) X2(t) X3(t) X4(t) X5(t)
+
+    rxs = [
+        Reaction(p1, nothing, [X1]),
+        Reaction(p2, nothing, [X2]),
+        Reaction(p3, nothing, [X3]),
+        Reaction(p4, nothing, [X4]),
+        Reaction(p5, nothing, [X5]),
+        Reaction(d1, [X1], nothing),
+        Reaction(d2, [X2], nothing),
+        Reaction(d3, [X3], nothing),
+        Reaction(d4, [X4], nothing),
+        Reaction(d5, [X5], nothing)
+    ]
+    @named rs = ReactionSystem(rxs)
+    rs = complete(rs)
+
+    # Declares initial condition and potential parameter sets.
+    u0 = [X1 => 0.1, X2 => 0.2, X3 => 0.3, X4 => 0.4, X5 => 0.5]
+    p_alts = [
+        [p1 => 1.0, d1 => 1.0, p2 => 1.2, p3 => 2, p4 => 0.5, d4 => 1//2, p5 => 3//2, d5 => 1.5],
+        (p1 => 1.0, d1 => 1.0, p2 => 1.2, p3 => 2, p4 => 0.5, d4 => 1//2, p5 => 3//2, d5 => 1.5),
+        Dict([p1 => 1.0, d1 => 1.0, p2 => 1.2, p3 => 2, p4 => 0.5, d4 => 1//2, p5 => 3//2, d5 => 1.5])
+    ]
+end
+
+# Tests that parameters stored in the system have the correct type.
+let
+    @test Symbolics.unwrap(rs.p1) isa SymbolicUtils.BasicSymbolic{Real}
+    @test Symbolics.unwrap(rs.d1) isa SymbolicUtils.BasicSymbolic{Real}
+    @test Symbolics.unwrap(rs.p2) isa SymbolicUtils.BasicSymbolic{Float64}
+    @test Symbolics.unwrap(rs.d2) isa SymbolicUtils.BasicSymbolic{Float64}
+    @test Symbolics.unwrap(rs.p3) isa SymbolicUtils.BasicSymbolic{Int64}
+    @test Symbolics.unwrap(rs.d3) isa SymbolicUtils.BasicSymbolic{Int64}
+    @test Symbolics.unwrap(rs.p4) isa SymbolicUtils.BasicSymbolic{Float32}
+    @test Symbolics.unwrap(rs.d4) isa SymbolicUtils.BasicSymbolic{Rational{Int64}}
+    @test Symbolics.unwrap(rs.p5) isa SymbolicUtils.BasicSymbolic{Rational{Int64}}
+    @test Symbolics.unwrap(rs.d5) isa SymbolicUtils.BasicSymbolic{Float32}
+end
+
+# Tests that simulations with differentially typed variables yields correct results.
+let
+    for p in p_alts
+        oprob = ODEProblem(rs, u0, (0.0, 1000.0), p; abstol = 1e-10, reltol = 1e-10)
+        sol = solve(oprob, Tsit5())
+        @test all(sol[end] .≈ 1.0)
+    end
+end
+
+# Test that the various structures stores the parameters using the correct type.
+let
+    # Creates problems, integrators, and solutions.
+    oprob = ODEProblem(rs, u0, (0.0, 1.0), p_alts[1])
+    sprob = SDEProblem(rs, u0, (0.0, 1.0), p_alts[1])
+    dprob = DiscreteProblem(rs, u0, (0.0, 1.0), p_alts[1])
+    jprob = JumpProblem(rs, dprob, Direct(); rng)
+    nprob = NonlinearProblem(rs, u0, p_alts[1])
+
+    oinit = init(oprob, Tsit5()) 
+    sinit = init(sprob, ImplicitEM()) 
+    jinit = init(jprob, SSAStepper()) 
+    ninit = init(nprob, NewtonRaphson()) 
+
+    osol = solve(oprob, Tsit5())
+    ssol = solve(sprob, ImplicitEM(); seed)
+    jsol = solve(jprob, SSAStepper(); seed)
+    nsol = solve(nprob, NewtonRaphson())
+
+    # Checks the types of all stored parameter values.
+    for mtk_struct in [oprob, sprob, dprob, jprob, nprob, oinit, sinit, jinit, osol, ssol, jsol, nsol]
+        @test unwrap(mtk_struct.ps[p1]) isa Float64
+        @test unwrap(mtk_struct.ps[d1]) isa Float64
+        @test unwrap(mtk_struct.ps[p2]) isa Float64
+        @test unwrap(mtk_struct.ps[d2]) isa Float64
+        @test unwrap(mtk_struct.ps[p3]) isa Int64
+        @test unwrap(mtk_struct.ps[d3]) isa Int64
+        @test unwrap(mtk_struct.ps[p4]) isa Float32
+        @test unwrap(mtk_struct.ps[d4]) isa Rational{Int64}
+        @test unwrap(mtk_struct.ps[p5]) isa Rational{Int64}
+        @test unwrap(mtk_struct.ps[d5]) isa Float32
+    end
+
+    # Indexing currently broken for NonlinearSystem integrators (MTK intend to support this though).
+    @test_broken unwrap(ninit.ps[p1]) isa Float64
+end

test/runtests.jl

@@ -8,6 +8,7 @@ using SafeTestsets
     @time @safetestset "Reactions" begin include("reactionsystem_structure/reactions.jl") end
     @time @safetestset "ReactionSystem" begin include("reactionsystem_structure/reactionsystem.jl") end
     @time @safetestset "Higher Order Reactions" begin include("reactionsystem_structure/higher_order_reactions.jl") end
+    @time @safetestset "Designation of Parameter Types" begin include("reactionsystem_structure/designating_parameter_types.jl") end
 
     ### Tests model creation via the @reaction_network DSL. ###
     @time @safetestset "Basic DSL" begin include("dsl/dsl_basics.jl") end

test/spatial_reaction_systems/lattice_reaction_systems.jl

@@ -2,12 +2,15 @@
 
 # Fetch packages.
 using Catalyst, Graphs, Test
+using Symbolics: unwrap
 t = default_t()
 
 # Pre declares a grid.
 grid = Graphs.grid([2, 2])
 
+
 ### Tests LatticeReactionSystem Getters Correctness ###
+
 # Test case 1.
 let 
     rs = @reaction_network begin
@@ -272,3 +275,90 @@ let
     @test_throws ErrorException LatticeReactionSystem(rs, [tr], grid)
 end
 
+### Test Designation of Parameter Types ###
+
+# Checks that parameter types designated in the non-spatial `ReactionSystem` is handled correctly.
+let
+    # Declares LatticeReactionSystem with designated parameter types.
+    rs = @reaction_network begin
+        @parameters begin
+            k1
+            l1
+            k2::Float64 = 2.0
+            l2::Float64
+            k3::Int64 = 2, [description="A parameter"]
+            l3::Int64
+            k4::Float32, [description="Another parameter"]
+            l4::Float32
+            k5::Rational{Int64}
+            l5::Rational{Int64}
+            D1::Float32
+            D2, [edgeparameter=true]
+            D3::Rational{Int64}, [edgeparameter=true]
+        end
+        (k1,l1), X1 <--> Y1
+        (k2,l2), X2 <--> Y2
+        (k3,l3), X3 <--> Y3
+        (k4,l4), X4 <--> Y4
+        (k5,l5), X5 <--> Y5
+    end
+    tr1 = @transport_reaction $(rs.D1) X1 
+    tr2 = @transport_reaction $(rs.D2) X2 
+    tr3 = @transport_reaction $(rs.D3) X3    
+    lrs = LatticeReactionSystem(rs, [tr1, tr2, tr3], grid)
+    
+    # Loops through all parameters, ensuring that they have the correct type
+    p_types = Dict([ModelingToolkit.nameof(p) => typeof(unwrap(p)) for p in parameters(lrs)])
+    @test p_types[:k1] == SymbolicUtils.BasicSymbolic{Real}
+    @test p_types[:l1] == SymbolicUtils.BasicSymbolic{Real}
+    @test p_types[:k2] == SymbolicUtils.BasicSymbolic{Float64}
+    @test p_types[:l2] == SymbolicUtils.BasicSymbolic{Float64}
+    @test p_types[:k3] == SymbolicUtils.BasicSymbolic{Int64}
+    @test p_types[:l3] == SymbolicUtils.BasicSymbolic{Int64}
+    @test p_types[:k4] == SymbolicUtils.BasicSymbolic{Float32}
+    @test p_types[:l4] == SymbolicUtils.BasicSymbolic{Float32}
+    @test p_types[:k5] == SymbolicUtils.BasicSymbolic{Rational{Int64}}
+    @test p_types[:l5] == SymbolicUtils.BasicSymbolic{Rational{Int64}}
+    @test p_types[:D1] == SymbolicUtils.BasicSymbolic{Float32}
+    @test p_types[:D2] == SymbolicUtils.BasicSymbolic{Real}
+    @test p_types[:D3] == SymbolicUtils.BasicSymbolic{Rational{Int64}}
+end
+
+# Checks that programmatically declared parameters (with types) can be used in `TransportReaction`s.
+# Checks that LatticeReactionSystem with non-default parameter types can be simulated.
+let
+    rs = @reaction_network begin
+        @parameters p::Float32
+        (p,d), 0 <--> X
+    end
+    @parameters D::Rational{Int64}
+    tr = TransportReaction(D, rs.X)  
+    lrs = LatticeReactionSystem(rs, [tr], grid)
+    
+    p_types = Dict([ModelingToolkit.nameof(p) => typeof(unwrap(p)) for p in parameters(lrs)])
+    @test p_types[:p] == SymbolicUtils.BasicSymbolic{Float32}
+    @test p_types[:d] == SymbolicUtils.BasicSymbolic{Real}
+    @test p_types[:D] == SymbolicUtils.BasicSymbolic{Rational{Int64}}
+
+    u0 = [:X => [0.25, 0.5, 2.0, 4.0]]
+    ps = [rs.p => 2.0, rs.d => 1.0, D => 1//2]
+
+    # Currently broken. This requires some non-trivial reworking of internals.
+    # However, spatial internals have already been reworked (and greatly improved) in an unmerged PR.
+    # This will be sorted out once that has finished.
+    @test_broken false 
+    # oprob = ODEProblem(lrs, u0, (0.0, 10.0), ps)
+    # sol = solve(oprob, Tsit5())
+    # @test sol[end] == [1.0, 1.0, 1.0, 1.0]
+end
+
+# Tests that LatticeReactionSystem cannot be generated where transport reactions depend on parameters
+# that have a type designated in the non-spatial `ReactionSystem`.
+let
+    rs = @reaction_network begin
+        @parameters D::Int64
+        (p,d), 0 <--> X
+    end
+    tr = @transport_reaction D X 
+    @test_throws Exception LatticeReactionSystem(rs, tr, grid)
+end