Skip to content

Commit

Permalink
Merge pull request #756 from SciML/spatial_internal_update
Browse files Browse the repository at this point in the history 
[WIP] LatticeReactionSystem internal update
  • Loading branch information
@TorkelE
TorkelE authored Jul 9, 2024
2 parents a4c75fc + e19d67d commit 2228734
Show file tree
Hide file tree
Showing 16 changed files with 2,827 additions and 1,371 deletions.
17 changes: 12 additions & 5 deletions src/Catalyst.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ using LaTeXStrings, Latexify, Requires
using LinearAlgebra, Combinatorics
using JumpProcesses: JumpProcesses, JumpProblem,
MassActionJump, ConstantRateJump, VariableRateJump,
SpatialMassActionJump
SpatialMassActionJump, CartesianGrid, CartesianGridRej

# ModelingToolkit imports and convenience functions we use
using ModelingToolkit
Expand Down Expand Up @@ -171,18 +171,25 @@ 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

# Various utility functions
include("spatial_reaction_systems/utility.jl")
export CartesianGrid, CartesianGridReJ # (Implemented in JumpProcesses)
export has_cartesian_lattice, has_masked_lattice, has_grid_lattice, has_graph_lattice,
grid_dims, grid_size
export make_edge_p_values, make_directed_edge_values
include("spatial_reaction_systems/lattice_solution_interfacing.jl")
export get_lrs_vals

# Specific spatial problem types.
include("spatial_reaction_systems/spatial_ODE_systems.jl")
export rebuild_lat_internals!
include("spatial_reaction_systems/lattice_jump_systems.jl")

# General spatial modelling utility functions.
include("spatial_reaction_systems/utility.jl")

### ReactionSystem Serialisation ###
# Has to be at the end (because it uses records of all metadata declared by Catalyst).
include("reactionsystem_serialisation/serialisation_support.jl")
Expand Down
124 changes: 67 additions & 57 deletions src/spatial_reaction_systems/lattice_jump_systems.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,123 +7,123 @@ function DiffEqBase.DiscreteProblem(lrs::LatticeReactionSystem, u0_in, tspan,
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.
# Converts potential symmaps to varmaps.
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)
p_in = symmap_to_varmap(lrs, p_in)

# Converts u0 and p to their internal forms.
# u0 is simply a vector with all the species' initial condition values across all vertices.
# 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).
u0 = lattice_process_u0(u0_in, species(lrs), lrs)
# vert_ps and `edge_ps` are vector maps, taking each parameter's Symbolics representation to its value(s).
# vert_ps values are vectors. Here, index (i) is a parameter's value in vertex i.
# edge_ps values are sparse matrices. Here, index (i,j) is a parameter's value in the edge from vertex i to vertex j.
# Uniform vertex/edge parameters store only a single value (a length 1 vector, or size 1x1 sparse matrix).
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(u0, tspan, [vert_ps, edge_ps], args...; kwargs...)
# Returns a DiscreteProblem (which basically just stores the processed input).
return DiscreteProblem(u0, tspan, [vert_ps; edge_ps], args...; kwargs...)
end

# Builds a spatial JumpProblem from a DiscreteProblem containing a Lattice Reaction System.
function JumpProcesses.JumpProblem(lrs::LatticeReactionSystem, dprob, aggregator,
args...; name = nameof(lrs.rs),
combinatoric_ratelaws = get_combinatoric_ratelaws(lrs.rs), kwargs...)
# Builds a spatial JumpProblem from a DiscreteProblem containing a `LatticeReactionSystem`.
function JumpProcesses.JumpProblem(lrs::LatticeReactionSystem, dprob, aggregator, args...;
combinatoric_ratelaws = get_combinatoric_ratelaws(reactionsystem(lrs)),
name = nameof(reactionsystem(lrs)), kwargs...)
# Error checks.
if !isnothing(dprob.f.sys)
error("Unexpected `DiscreteProblem` passed into `JumpProblem`. Was a `LatticeReactionSystem` used as input to the initial `DiscreteProblem`?")
throw(ArgumentError("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).
# Hence the parameters of this one does nto actually matter. If at some point JumpProcess can
# Hence the parameters of this one do not 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]))
non_spat_dprob = DiscreteProblem(reshape(dprob.u0, num_species(lrs), num_verts(lrs)),
dprob.tspan, first.(dprob.p[1]))

# Creates and returns a spatial JumpProblem (masked lattices are not supported by these).
spatial_system = has_masked_lattice(lrs) ? get_lattice_graph(lrs) : lattice(lrs)
return JumpProblem(non_spat_dprob, aggregator, sma_jumps;
hopping_constants, spatial_system = lrs.lattice, name, kwargs...)
hopping_constants, spatial_system, 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))
# If the transport rate is uniform for one species, the vector has a single element, else one for each edge.
spatial_rates_dict = Dict(compute_all_transport_rates(Dict(dprob.p), 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
# Creates an array (of the same size as the hopping constant array) containing all edges.
# First the array is a NxM matrix (number of species x number of vertices). Each element is a
# vector containing all edges leading out from that vertex (sorted by destination index).
edge_array = [Pair{Int64, Int64}[] for _1 in 1:num_species(lrs), _2 in 1:num_verts(lrs)]
for e in edge_iterator(lrs), s_idx in 1:num_species(lrs)
push!(edge_array[s_idx, e[1]], e)
end
foreach(e_vec -> sort!(e_vec; by = e -> e[2]), edge_array)

# Creates the hopping constants array. It has the same shape as the edge array, but each
# element is that species transportation rate along that edge
hopping_constants = [[Catalyst.get_edge_value(all_diff_rates[s_idx], e)
for e in edge_array[s_idx, src_idx]]
for s_idx in 1:num_species(lrs), src_idx in 1:num_verts(lrs)]
return hopping_constants
end

# 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.
# Could implement a version which, if all reactions' rates are uniform, returns a MassActionJump.
# Not sure if there is any form of performance improvement from that though. Likely 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)]
is_spatials = [has_spatial_vertex_component(rx.rate, dprob.p)
for rx in reactions(reactionsystem(lrs))]

# 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)))
num_rxs = length(reactions(reactionsystem(lrs)))
u_rates = Vector{Float64}(undef, num_rxs - count(is_spatials))
s_rates = Matrix{Float64}(undef, count(is_spatials), num_verts(lrs))
reactant_stoich = Vector{Vector{Pair{Int64, Int64}}}(undef, num_rxs)
net_stoich = Vector{Vector{Pair{Int64, Int64}}}(undef, num_rxs)

# 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))
for (is_spat, rx) in zip(is_spatials, reactions(reactionsystem(lrs)))
is_spat && continue
u_rates[cur_rx] = compute_vertex_value(rx.rate, lrs; vert_ps = dprob.p[1])[1]
u_rates[cur_rx] = compute_vertex_value(rx.rate, lrs; 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)
reactant_stoich[cur_rx] = int_map(substoich_map, reactionsystem(lrs))
net_stoich[cur_rx] = int_map(rx.netstoich, reactionsystem(lrs))
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))
for (is_spat, rx) in zip(is_spatials, reactions(reactionsystem(lrs)))
is_spat || continue
s_rates[cur_rx - length(u_rates), :] = compute_vertex_value(rx.rate, lrs;
vert_ps = dprob.p[1])
s_rates[cur_rx - length(u_rates), :] .= compute_vertex_value(rx.rate, lrs;
ps = dprob.p)
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)
reactant_stoich[cur_rx] = int_map(substoich_map, reactionsystem(lrs))
net_stoich[cur_rx] = int_map(rx.netstoich, reactionsystem(lrs))
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)
return SpatialMassActionJump(u_rates, s_rates, reactant_stoich, net_stoich, nothing)
end

### Extra ###

# Temporary. Awaiting implementation in SII, or proper implementation withinCatalyst (with more general functionality).
# Temporary. Awaiting implementation in SII, or proper implementation within Catalyst (with
# more general functionality).
function int_map(map_in, sys)
return [ModelingToolkit.variable_index(sys, pair[1]) => pair[2] for pair in map_in]
end
Expand All @@ -133,7 +133,7 @@ end
# 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)))
# statetoid = Dict(ModelingToolkit.value(state) => i for (i, state) in enumerate(unknowns(rs)))
# eqs = equations(js)
# invttype = non_spat_dprob.tspan[1] === nothing ? Float64 : typeof(1 / non_spat_dprob.tspan[2])
#
Expand All @@ -142,3 +142,13 @@ end
# majpmapper = ModelingToolkit.JumpSysMajParamMapper(js, p; jseqs = eqs, rateconsttype = invttype)
# return ModelingToolkit.assemble_maj(eqs.x[1], statetoid, majpmapper)
# end

### Problem & Integrator Rebuilding ###

# Currently not implemented.
function rebuild_lat_internals!(dprob::DiscreteProblem)
error("Modification and/or rebuilding of `DiscreteProblem`s is currently not supported. Please create a new problem instead.")
end
function rebuild_lat_internals!(jprob::JumpProblem)
error("Modification and/or rebuilding of `JumpProblem`s is currently not supported. Please create a new problem instead.")
end
Loading

0 comments on commit 2228734

Please sign in to comment.