update

Author: TorkelE

src/spatial_reaction_systems/spatial_ODE_systems.jl

@@ -13,15 +13,17 @@ struct LatticeTransportODEf{R,S,T}
     """Temporary vector. For parameters which values are identical across the lattice, at some point these have to be converted of a length num_verts vector. To avoid re-allocation they are written to this vector."""
     work_vert_ps::Vector{S}    
     """For each parameter in vert_ps, its value is a vector with length either num_verts or 1. To know whenever a parameter's value need expanding to the work_vert_ps array, its length needs checking. This check is done once, and the value stored to this array. This field (specifically) is an enumerate over that array."""
-    enum_v_ps_idx_types::Base.Iterators.Enumerate{Vector{Bool}}
+    v_ps_idx_types::Vector{Bool}
     """A vector of pairs, with a value for each species with transportation. The first value is the species index (in the species(::ReactionSystem) vector), and the second is a vector with its transport rate values. If the transport rate is uniform (across all edges), that value is the only value in the vector. Else, there is one value for each edge in the lattice."""
     transport_rates::Vector{Pair{Int64, Vector{S}}}
     """A matrix, NxM, where N is the number of species with transportation and M the number of vertexes. Each value is the total rate at which that species leaves that vertex (e.g. for a species with constant diffusion rate D, in a vertex with n neighbours, this value is n*D)."""
     leaving_rates::Matrix{S}
     """An (enumerate'ed) iterator over all the edges of the lattice."""
-    enum_edges::T
+    edges::Graphs.SimpleGraphs.SimpleEdgeIter{SimpleDiGraph{Int64}}
+    """The edge parameters used to create the spatial ODEProblem. Currently unused, but will be needed to support changing these (e.g. due to events). Contain one vector for each edge parameter (length one if uniform, else one value for each edge)."""
+    edge_ps::Vector{Vector{T}}
 
-    function LatticeTransportODEf(ofunc::R, vert_ps::Vector{Vector{S}}, transport_rates::Vector{Pair{Int64, Vector{S}}}, lrs::LatticeReactionSystem) where {R,S}
+    function LatticeTransportODEf(ofunc::R, vert_ps::Vector{Vector{S}}, transport_rates::Vector{Pair{Int64, Vector{S}}}, edge_ps::Vector{Vector{T}}, lrs::LatticeReactionSystem) where {R,S,T}
         leaving_rates = zeros(length(transport_rates), lrs.num_verts)
         for (s_idx, trpair) in enumerate(transport_rates)
             rates = last(trpair)
@@ -31,37 +33,38 @@ struct LatticeTransportODEf{R,S,T}
             end
         end
         work_vert_ps = zeros(lrs.num_verts)
-        # 1 if ps are constant across the graph, 0 else
-        enum_v_ps_idx_types = enumerate(map(vp -> length(vp) == 1, vert_ps))
-        enum_edges = deepcopy(enumerate(edges(lrs.lattice)))                    # Creates an iterator over all the edges. Again, this is always used in the enumerated form. 
-        new{R,S,typeof(enum_edges)}(ofunc, lrs.num_verts, lrs.num_species, vert_ps, work_vert_ps, enum_v_ps_idx_types, transport_rates, leaving_rates, enum_edges)
+        # 1 if ps are constant across the graph, 0 else.
+        v_ps_idx_types = map(vp -> length(vp) == 1, vert_ps)
+        eds = edges(lrs.lattice)                  
+        new{R,S,T}(ofunc, lrs.num_verts, lrs.num_species, vert_ps, work_vert_ps, v_ps_idx_types, transport_rates, leaving_rates, eds, edge_ps)
     end
 end
 
 # Functor structure containing the information for the forcing function of a spatial ODE with spatial movement on a lattice.
-struct LatticeTransportODEjac{R,S,T}
+struct LatticeTransportODEjac{Q,R,S,T}
     """The ODEFunction of the (non-spatial) reaction system which generated this function."""
-    ofunc::R
+    ofunc::Q
     """The number of vertices."""
     num_verts::Int64
     """The number of species."""
     num_species::Int64
     """The values of the parameters which values are tied to vertexes."""
-    vert_ps::Vector{Vector{S}}
+    vert_ps::Vector{Vector{R}}
     """Temporary vector. For parameters which values are identical across the lattice, at some point these have to be converted of a length(num_verts) vector. To avoid re-allocation they are written to this vector."""
-    work_vert_ps::Vector{S} 
+    work_vert_ps::Vector{R} 
     """For each parameter in vert_ps, it either have length num_verts or 1. To know whenever a parameter's value need expanding to the work_vert_ps array, its length needs checking. This check is done once, and the value stored to this array. This field (specifically) is an enumerate over that array."""
-    enum_v_ps_idx_types::Base.Iterators.Enumerate{Vector{Bool}}
+    v_ps_idx_types::Vector{Bool}
     """Whether the Jacobian is sparse or not."""
     sparse::Bool
     """The transport rates. Can be a dense matrix (for non-sparse) or as the "nzval" field if sparse."""
-    jac_values::T
+    jac_transport::S
+    """The edge parameters used to create the spatial ODEProblem. Currently unused, but will be needed to support changing these (e.g. due to events). Contain one vector for each edge parameter (length one if uniform, else one value for each edge)."""
+    edge_ps::Vector{Vector{T}}
 
-    function LatticeTransportODEjac(ofunc::R, vert_ps::Vector{Vector{S}}, lrs::LatticeReactionSystem, jac_prototype::Union{Nothing, SparseMatrixCSC{Float64, Int64}}, sparse::Bool) where {R,S}
+    function LatticeTransportODEjac(ofunc::R, vert_ps::Vector{Vector{S}}, lrs::LatticeReactionSystem, jac_transport::Union{Nothing, SparseMatrixCSC{Float64, Int64}}, edge_ps::Vector{Vector{T}}, sparse::Bool) where {R,S,T}
         work_vert_ps = zeros(lrs.num_verts)
-        enum_v_ps_idx_types = enumerate(map(vp -> length(vp) == 1, vert_ps))
-        jac_values = sparse ? jac_prototype.nzval : Matrix(jac_prototype)   # Retrieves the diffusion values (form depending on Jacobian sparsity).
-        new{R,S,typeof(jac_values)}(ofunc, lrs.num_verts, lrs.num_species, vert_ps, work_vert_ps, enum_v_ps_idx_types, sparse, jac_values)
+        v_ps_idx_types = map(vp -> length(vp) == 1, vert_ps)
+        new{R,S,typeof(jac_transport),T}(ofunc, lrs.num_verts, lrs.num_species, vert_ps, work_vert_ps, v_ps_idx_types, sparse, jac_transport, edge_ps)
     end
 end
 
@@ -70,7 +73,10 @@ end
 # Creates an ODEProblem from a LatticeReactionSystem.
 function DiffEqBase.ODEProblem(lrs::LatticeReactionSystem, u0_in, tspan,
                                p_in = DiffEqBase.NullParameters(), args...;
-                               jac = true, sparse = jac, kwargs...)
+                               jac = true, sparse = jac, 
+                               name = nameof(lrs), include_zero_odes = true,
+                               combinatoric_ratelaws = get_combinatoric_ratelaws(lrs.rs),
+                               remove_conserved = false, checks = false, kwargs...)
     is_transport_system(lrs) || error("Currently lattice ODE simulations are only supported when all spatial reactions are TransportReactions.")
 
     # Converts potential symmaps to varmaps (parameter conversion is more involved since the vertex and edge parameters may be given in a tuple, or in a common vector).
@@ -85,26 +91,49 @@ function DiffEqBase.ODEProblem(lrs::LatticeReactionSystem, u0_in, tspan,
     vert_ps, edge_ps = lattice_process_p(p_in, vertex_parameters(lrs), edge_parameters(lrs), lrs)   
 
     # Creates ODEProblem.
-    ofun = build_odefunction(lrs, vert_ps, edge_ps, jac, sparse)        # Builds the ODEFunction.
-    return ODEProblem(ofun, u0, tspan, vert_ps, args...; kwargs...)     # Creates a normal ODEProblem.
+    ofun = build_odefunction(lrs, vert_ps, edge_ps, jac, sparse, name, include_zero_odes, combinatoric_ratelaws, remove_conserved, checks)
+    return ODEProblem(ofun, u0, tspan, vert_ps, args...; kwargs...) 
 end
 
 # Builds an ODEFunction for a spatial ODEProblem.
 function build_odefunction(lrs::LatticeReactionSystem, vert_ps::Vector{Vector{T}},
-                           edge_ps::Vector{Vector{T}}, use_jac::Bool, sparse::Bool) where {T}
-    # Prepares (non-spatial) ODE functions and list of spatially moving species and their rates.
-    ofunc = ODEFunction(convert(ODESystem, lrs.rs); jac = use_jac, sparse = false)                  # Creates the (non-spatial) ODEFunction corresponding to the (non-spatial) reaction network.
-    ofunc_sparse = ODEFunction(convert(ODESystem, lrs.rs); jac = use_jac, sparse = true)            # Creates the same function, but sparse. Could insert so this is only computed for sparse cases.
-    transport_rates_speciesmap = compute_all_transport_rates(vert_ps, edge_ps, lrs)                 # Creates a map (Vector{Pair}), mapping each species that is transported to a vector with its transportation rate. If the rate is uniform across all edges, the vector will be length 1 (with this value), else there will be a separate value for each edge.
-    transport_rates = Pair{Int64, Vector{T}}[findfirst(isequal(spat_rates[1]), species(lrs)) => spat_rates[2]
-                       for spat_rates in transport_rates_speciesmap]                                # Remakes "transport_rates_speciesmap". Rates are identical, but the species are represented as their index (in the species(::ReactionSystem) vector). In "transport_rates_speciesmap" they instead were Symbolics. Pair{Int64, Vector{T}}[] is required in case vector is empty (otherwise it becomes Any[], causing type error later).
-
-    f = LatticeTransportODEf(ofunc, vert_ps, transport_rates, lrs)
-    jac_prototype = (use_jac || sparse) ?
-                    build_jac_prototype(ofunc_sparse.jac_prototype, transport_rates,
-                                        lrs; set_nonzero = use_jac) : nothing                       # Computes the Jacobian prototype (nothing if `jac=false`). 
-    jac = use_jac ? LatticeTransportODEjac(ofunc, vert_ps, lrs, jac_prototype, sparse) : nothing    # (Potentially) Creates a functor for the ODE Jacobian function (incorporating spatial and non-spatial reactions).
-    return ODEFunction(f; jac = jac, jac_prototype = (sparse ? jac_prototype : nothing))            # Creates the ODEFunction used in the ODEProblem.
+                           edge_ps::Vector{Vector{T}}, jac::Bool, sparse::Bool,
+                           name, include_zero_odes, combinatoric_ratelaws, remove_conserved, checks) where {T}
+    println()
+    remove_conserved && error("Removal of conserved quantities is currently not supported for `LatticeReactionSystem`s")
+
+    # Creates a map, taking (the index in species(lrs) each species (with transportation) to its transportation rate (uniform or one value for each edge).
+    transport_rates = make_sidxs_to_transrate_map(vert_ps, edge_ps, lrs)    
+
+    # Prepares the Jacobian and forcing functions (depending on jacobian and sparsity selection).
+    if jac
+        ofunc_dense = ODEFunction(convert(ODESystem, lrs.rs; name, combinatoric_ratelaws, include_zero_odes, checks); jac = true, sparse = false) # Always used for build_jac_prototype.
+        ofunc_sparse = ODEFunction(convert(ODESystem, lrs.rs; name, combinatoric_ratelaws, include_zero_odes, checks); jac = true, sparse = true) # Always used for LatticeTransportODEjac.
+        jac_vals = build_jac_prototype(ofunc_sparse.jac_prototype, transport_rates, lrs; set_nonzero = true)
+        if sparse
+            f = LatticeTransportODEf(ofunc_sparse, vert_ps, transport_rates, edge_ps, lrs)
+            jac_vals = build_jac_prototype(ofunc_sparse.jac_prototype, transport_rates, lrs; set_nonzero = true)
+            J = LatticeTransportODEjac(ofunc_dense, vert_ps, lrs, jac_vals, edge_ps, true)
+            jac_prototype = jac_vals
+        else
+            f = LatticeTransportODEf(ofunc_dense, vert_ps, transport_rates, edge_ps, lrs)
+            J = LatticeTransportODEjac(ofunc_dense, vert_ps, lrs, jac_vals, edge_ps, false)
+            jac_prototype = nothing
+        end
+    else
+        if sparse
+            ofunc_sparse = ODEFunction(convert(ODESystem, lrs.rs; name, combinatoric_ratelaws, include_zero_odes, checks); jac = false, sparse = true)
+            f = LatticeTransportODEf(ofunc_sparse, vert_ps, transport_rates, edge_ps, lrs)
+            jac_prototype = build_jac_prototype(ofunc_sparse.jac_prototype, transport_rates, lrs; set_nonzero = false)
+        else
+            ofunc_dense = ODEFunction(convert(ODESystem, lrs.rs; name, combinatoric_ratelaws, include_zero_odes, checks); jac = false, sparse = false)
+            f = LatticeTransportODEf(ofunc_dense, vert_ps, transport_rates, edge_ps, lrs)
+            jac_prototype = nothing
+        end
+        J = nothing
+    end
+
+    return ODEFunction(f; jac = J, jac_prototype = jac_prototype)           
 end
 
 # Builds a jacobian prototype. If requested, populate it with the Jacobian's (constant) values as well.
@@ -120,7 +149,7 @@ function build_jac_prototype(ns_jac_prototype::SparseMatrixCSC{Float64, Int64},
     ns_j_idxs = ns_jac_prototype_idxs[2]
 
     # List the indexes of all non-zero Jacobian terms.
-    non_spat_terms = [[get_index(vert, s_i, lrs.num_species), get_index(vert, s_j, lrs.num_species)] for vert in 1:(lrs.num_verts) for (s_i, s_j) in zip(ns_i_idxs,ns_j_idxs)]               # Indexes of elements due to non-spatial dynamics.
+    non_spat_terms = [[get_index(vert, s_i, lrs.num_species), get_index(vert, s_j, lrs.num_species)] for vert in 1:(lrs.num_verts) for (s_i, s_j) in zip(ns_i_idxs,ns_j_idxs)]        # Indexes of elements due to non-spatial dynamics.
     trans_only_leaving_terms = [[get_index(e.src, s_idx, lrs.num_species), get_index(e.src, s_idx, lrs.num_species)] for e in edges(lrs.lattice) for s_idx in trans_only_species]     # Indexes due to terms for a species leaves its current vertex (but does not have non-spatial dynamics). If the non-spatial Jacobian is fully dense, these would already be accounted for. 
     trans_arriving_terms = [[get_index(e.src, s_idx, lrs.num_species), get_index(e.dst, s_idx, lrs.num_species)] for e in edges(lrs.lattice) for s_idx in trans_species]              # Indexes due to terms for species arriving into a new vertex.
     all_terms = [non_spat_terms; trans_only_leaving_terms; trans_arriving_terms]
@@ -130,7 +159,7 @@ function build_jac_prototype(ns_jac_prototype::SparseMatrixCSC{Float64, Int64},
 
     # Set element values.
     if set_nonzero
-        for (s, rates) in trans_rates, (e_idx, e) in enumerate(edges(lrs.lattice))    # Loops through all species with transportation and all edges along which the can be transported.
+        for (s, rates) in trans_rates, (e_idx, e) in enumerate(edges(lrs.lattice))
             jac_prototype[get_index(e.src, s, lrs.num_species), get_index(e.src, s, lrs.num_species)] -= get_component_value(rates, e_idx)    # Term due to species leaving source vertex.
             jac_prototype[get_index(e.src, s, lrs.num_species), get_index(e.dst, s, lrs.num_species)] += get_component_value(rates, e_idx)    # Term due to species arriving to destination vertex.
         end
@@ -147,7 +176,7 @@ function (f_func::LatticeTransportODEf)(du, u, p, t)
         idxs = get_indexes(vert_i, f_func.num_species)
 
         # vector of vertex ps at vert_i
-        vert_i_ps = view_vert_ps_vector!(f_func.work_vert_ps, p, vert_i, f_func.enum_v_ps_idx_types)
+        vert_i_ps = view_vert_ps_vector!(f_func.work_vert_ps, p, vert_i, enumerate(f_func.v_ps_idx_types))
 
         # evaluate reaction contributions to du at vert_i
         f_func.ofunc((@view du[idxs]),  (@view u[idxs]), vert_i_ps, t)
@@ -156,13 +185,13 @@ function (f_func::LatticeTransportODEf)(du, u, p, t)
     # s_idx is species index among transport species, s is index among all species
     # rates are the species' transport rates
     for (s_idx, (s, rates)) in enumerate(f_func.transport_rates)  
-        # rate for leaving vert_i
+        # Rate for leaving vert_i
         for vert_i in 1:(f_func.num_verts)                                 
             idx = get_index(vert_i, s, f_func.num_species)
             du[idx] -= f_func.leaving_rates[s_idx, vert_i] * u[idx]
         end
-        # add rates for entering a given vertex via an incoming edge
-        for (e_idx, e) in f_func.enum_edges 
+        # Add rates for entering a given vertex via an incoming edge
+        for (e_idx, e) in enumerate(f_func.edges)
             idx_dst = get_index(e.dst, s, f_func.num_species)
             idx_src = get_index(e.src, s, f_func.num_species)
             du[idx_dst] += get_component_value(rates, e_idx) * u[idx_src] 
@@ -177,14 +206,10 @@ function (jac_func::LatticeTransportODEjac)(J, u, p, t)
     # Update the Jacobian from reaction terms
     for vert_i in 1:(jac_func.num_verts)
         idxs = get_indexes(vert_i, jac_func.num_species)
-        vert_ps = view_vert_ps_vector!(jac_func.work_vert_ps, p, vert_i, jac_func.enum_v_ps_idx_types)
+        vert_ps = view_vert_ps_vector!(jac_func.work_vert_ps, p, vert_i, enumerate(jac_func.v_ps_idx_types))
         jac_func.ofunc.jac((@view J[idxs, idxs]), (@view u[idxs]), vert_ps, t)
     end
 
     # Updates for the spatial reactions (adds the Jacobian values from the diffusion reactions).
-    add_spat_J_vals!(J, jac_func)
-end
-
-# Updates the jacobian matrix with the diffusion values. Separate for spatial and non-spatial cases.
-add_spat_J_vals!(J::SparseMatrixCSC, jac_func::LatticeTransportODEjac) = (J.nzval .+= jac_func.jac_values)
-add_spat_J_vals!(J::Matrix, jac_func::LatticeTransportODEjac) = (J .+= jac_func.jac_values)
+    J .+= jac_func.jac_transport
+end

src/spatial_reaction_systems/spatial_reactions.jl

@@ -105,7 +105,7 @@ function find_parameters_in_rate!(parameters, rateex::ExprValues)
             push!(parameters, rateex)
         end
     elseif rateex isa Expr
-        # note, this (correctly) skips $(...) expressions
+        # Note, this (correctly) skips $(...) expressions
         for i in 2:length(rateex.args)
             find_parameters_in_rate!(parameters, rateex.args[i])
         end

src/spatial_reaction_systems/utility.jl

@@ -131,6 +131,15 @@ function compute_transport_rates(rate_law::Num,
                             for p in relevant_parameters)) for idxE in 1:num_edges]
 end
 
+# Creates a map, taking each species (with transportation) to its transportation rate.
+# The species is represented by its index (in species(lrs). 
+# If the rate is uniform across all edges, the vector will be length 1 (with this value), else there will be a separate value for each edge.
+# Pair{Int64, Vector{T}}[] is required in case vector is empty (otherwise it becomes Any[], causing type error later).
+function make_sidxs_to_transrate_map(vert_ps::Vector{Vector{Float64}}, edge_ps::Vector{Vector{T}}, lrs::LatticeReactionSystem) where T
+    transport_rates_speciesmap = compute_all_transport_rates(vert_ps, edge_ps, lrs)
+    return Pair{Int64, Vector{T}}[speciesmap(lrs.rs)[spat_rates[1]] => spat_rates[2] for spat_rates in transport_rates_speciesmap]
+end
+
 ### Accessing State & Parameter Array Values ###
 
 # Gets the index in the u array of species s in vertex vert (when their are num_species species).