Spatial modelling integration planning

[TorkelE]

Basically, I think there are 3 things we need to coordinate so the same stuff is used across spatial ODE/Jump simulations:

• What the content of the LatticeReactionNetwork needs to be.
• What the content of the DiffusionReaction structure needs to be.
• What formats to allow for providing input u0 and p to the ODEProblems and DiscreteProblems.

Lattice reaction system

The LatticeReactionNetwork I think we were happy with. It currently contains:

• rs: the Catalyst ReactionSystem.
• spatial_reactions: A vector of DiffusionReactions.
• lattice: A directed graph (over which we wish to simulate the system). If an undirected graph is provided as the lattice, it is converted to a directed one (with edges in both directions).
  It also contains some auxiliary information (which can be derived from those above), such as the number of compartments or edges, or whenever the lattice was initially provided as a directed or undirected graph. Potentially we might want to give some additional information (such as if the lattice has a special structure, like being a grid, would this be useful for jump simulators?)

Diffusion reactions

Currently have the following fields:

• rate: a Num. Can potentially be a number (e.g. 1.0), a single parameter (e.g. dX), or an algebraic expression of parameters (e.g. dE*dX).
• species: The species transported in this diffusion reaction (a single Num).
  It also contains some auxiliary information (e.g. the Symbol representation of the transported species, since this is easier to compare to various u0 inputs). I think the main discussion here is how to do with the rate, when to store it in what form and when to convert it into what etc.

Initial condition inputs

I will start with this one since it is easier than the parameters. For each species, it can either have an initial condition that is the same across all compartments, or a vector (with one value for each compartment). Consider a system with species [X, Y], and 3 lattice points. We can have the following possible types inputs:

# Inputs where order matters:
u0 = [1.0, 0.2]    # Vector of scalars. X=1.0 on all lattice points, Y=2.0 on all lattice points.
u0 = [1.0, [0.1, 0.2, 0.3]]    # Vector of scalars and vectors.  X=1.0 on all lattice points, Y=0.1 on the first lattice point, 0.2 on the 2nd, and 0.3 on the third.
u0 = [[1.0, 2.0, 3.0], [0.1, 0.2, 0.3]]       # Vector of vectors.  X=1.0 and Y=0.1 on the first lattice point, X=2.0 and Y=0.2 on the 2nd, and X=3.0 and Y=0.3 on the third.
u0 = [1.0 2.0 3.0; 0.1 0.2 0.3]    # Matrix. X=1.0 and Y=0.1 on the first lattice point, X=2.0 and Y=0.2 on the 2nd, and X=3.0 and Y=0.3 on the third.

# Inputs where order does not matter:
u0 = [X = > 1.0, Y => 0.2]    # Vector of pairs to scalars. X=1.0 on all lattice points, Y=2.0 on all lattice points.
u0 = [X = > 1.0, Y => [0.1, 0.2, 0.3]]    # Vector of pairs to scalars or vectors.  X=1.0 on all lattice points, Y=0.1 on the first lattice point, 0.2 on the 2nd, and 0.3 on the third.
u0 = [X = > [1.0, 2.0, 3.0], Y => [0.1, 0.2, 0.3]]       # Vector of pairs to vectors.  X=1.0 and Y=0.1 on the first lattice point, X=2.0 and Y=0.2 on the 2nd, and X=3.0 and Y=0.3 on the third.
# Here the numeric symbolics for X and Y are used, but Symbols :X and :Y can also be used.

Parameter inputs

Parameters are messier for two reasons:

• After starting the simulation, we can forget about the initial conditions. I.e. for these, we just need to convert the input into the initial u vector and can then forget about it. For parameters we must figure out what is the best ways to store the values (and potentially update them due to callbacks). We might want to consider several storage alternatives as well (could have different for ODEs and Jumps). However, storage is internal, so less important now.
• Parameters exist in both a compartment and edge form. This can matter when we have parameters with values varying between different compartments/edges.

Currently, parameters can either be given as a single vector:

p= [...] # Contain all parameter values.

or as a as a tuple, containing the compartment and edge parameters separately:

p = (pC, pE)

let us consider our previous system (X and Y spread across 3 compartments). If the compartments are on a line (with connections between the first and second and second and third), we have two compartment parameters (A and B) and a single edge parameter (D). Using the first parameter input we have the following options (for simplicity I do not split the edges into edges into two directions here):

# Inputs where order matters:
p = [1.0, 10.0, 0.1]    # Vector of scalars. A=1.0 on all lattice points, B=10.0 on all lattice points, and D=0.1 on all edges.
p = [1.0, [10.0, 20.0, 30.0], [0.1, 0.2]]     # Vector of scalars and vectors.  A=1.0 on all lattice points, B=10.0, 20.0, and 30.0 on the three lattice points. D=0.1, and 0.2 along the two edges. 
p = [[1.0, 2.0, 3.0], [10.0, 20.0, 30.0], [0.1, 0.2]]       # Vector of vectors.  A=1.0, 2.0, and 3.0 and B=10.0, 20.0, and 30.0 on the three lattice points. D=0.1, and 0.2 along the two edges.

# Inputs where order does not matter:
p = [A=>1.0, B=>10.0, D=>0.1]    # Vector of scalars. A=1.0 on all lattice points, B=10.0 on all lattice points, and D=0.1 on all edges.
p = [A=>1.0, B=>[10.0, 20.0, 30.0], D=>[0.1, 0.2]]     # Vector of scalars and vectors.  A=1.0 on all lattice points, B=10.0, 20.0, and 30.0 on the three lattice points. D=0.1, and 0.2 along the two edges.
p = [A=>[1.0, 2.0, 3.0], B=>[10.0, 20.0, 30.0], D=>[0.1, 0.2]]       # Vector of vectors.  A=1.0, 2.0, and 3.0 and B=10.0, 20.0, and 30.0 on the three lattice points. D=0.1, and 0.2 along the two edges.
# Here the numeric symbolics for A, B, and D are used, but Symbols :A, :B, and :D can also be used.

Using the second version (p = (pC, pE)) the advantage is that we can input values in Matrix for as well (pC and pE being separate matrices). If it wasn't because of this, I'd contemplate not allowing this form (simplifies things).

Next, when we consider that each edge might have different values in both directions, that adds another layer of complication

Currently, internally, everything is converted to the p = [1.0, [10.0, 20.0, 30.0], [0.1, 0.2]] # Vector of scalars and vectors. form (because it is the most compact, but as discussed other storage forms might be desired).

On a general note, I am thinking of adding metadata to parameters whenever they a compartment or edge parameter. As default all parameters that appear in the reaction system will be compartment ones. Those only appearing in DiffusionReactionss, or specified as edge parameters in the reaction system, will be edge parameters.

[Vilin97]

@TorkelE, I will describe what kind of input is expected for spatial solvers in JumpProcesses.
First of all, this tutorial is a good example. I will summarize its contents here. To create a spatial JumpProblem we must pass in 5 variables:

JumpProblem(prob, alg, majumps, hopping_constants = hopping_constants, spatial_system = grid

where

• prob is a DiscreteProblem with u0 given by a matrix whose (s,i) entry is the number of species s at site i.
• alg is either NSM() or DirectCRDirect() or any ordinary SSA (in case an ordinary SSA is passed, we flatten).
• majumps is MassActionJump or SpatialMassActionJump
• hopping_constants can be one of many things: a vector, a matrix, a matrix of vectors, a (vector, vector-of-vectors) pair or a (matrix, vector-of-vectors) pair. Let's first get rates of form D_{s,i,j} working. To do this, hopping_constants must be of type Matrix{Vector{F}} where F <: Number (usually F is Float64) with hopping_constants[s,i][j] being the hopping rate of species s at site i to the neighbor at index j. Look at the first example here if the description above is unclear. Alternativaly, an easier form of hopping rates is D_{s,i}, where you just pass a matrix whose (s,i) entry is the hopping rate of species s from site i.
• spatial_system can be either a CartesianGrid of dimension 1, 2 or 3, or a Graph.

If a MassActionJump is passed as the third argument, the reaction rates are assumed to be uniform in space. This is the most memory-efficient way. If reaction rates vary between sites in the spatial_system, we must construct a SpatialMassActionJump.

I will now describe how to build a SpatialMassActionJump. Here is the constructor code and here is the unit test, which is a great example. The constructor is
SpatialMassActionJump(urates::A, srates::B, rs, ns), where

• urates is the vector of rates for reactions whose rates are uniform in space.
• srates is the matrix of rates for reactions whose rates are vary in space. Its (r,i) entry is the rate of reaction r at site i.
• rs and ns are reaction stoichiometry and net stoichiometry. These are the same as for a MassActionJump Here is the example from the same test.

Let's focus on the most general case when all reaction rates vary in space so that we use nothing for urates in SpatialMassActionJump, and hopping rates are of the form D_{s,i,j}. This is the most general case, and it will be a good fallback. All the other variations are just optimizations and can be done later.

Also, let's not focus too much on keeping allocations low. Let's just get something working first and then iterate to make it better. Feel free to continue this thread if anything is unclear!

Tagging @isaacsas for visibility.

[Vilin97]

Potential roadmap:

• Get the uniform massaction case with hopping_rates of form $D_{s,i,j}$ and reaction rates of form (r,i). If it makes life easier, we can assume that the SSA is either NSM() or DirectCRDirect(). Start with Graph.
• Expand to the case when some reaction rates are non-uniform in space.
• Add CartesianGrid support.
• Expand to the case of hopping_rates of form $D_s$.
• Expand to the case of hopping_rates of form $D_{s,i}$.
• Expand to the other two hopping_rates cases: $D_{s} * L_{i,j}$ and $D_{s,i} * L_{i,j}$.
• Add support of ConstantRateJump.
• Add support of VariableRateJump?
• Make a tutorial with the new functionality.
• Create a helper package for plotting/animations. Something like AnimateCatalyst.jl.

The constant and variable rate jumps will require changes in JumpProcesses.jl as well as Catalyst.jl. The other items should only require PRs to Catalyst.jl.

[TorkelE]

So, for a starter what I think we need to do is to create a dispatch for when you call

JumpProblem(lrs::LatticeReactionSystem, prob::DiscreteProblem, alg)

so that the

• majumpss are extracted from lrs.rs (the reaction system),
• hopping_constantss are extracted from prob.p (the parameters)
• spatial_system is extracted from lrs.lattice (the spatial graph).
  and a call to

JumpProblem(prob, alg, majumps, hopping_constants = hopping_constants, spatial_system = grid)

is made.

If possible, the p that the users pass into DiscrteProblem can remain unchanged, and then we can revamp them into a desired form once passed into JumpProblem.

As long as we decide how we want parameters to be given when creating ODEProblems or DiscreteProblems, we don't need to coordinate the internals of the two approaches (ODe and Jump).

For now, would it be ok to limit ourselves to Graphs, and not CartesianGrids. Internally, is the structure of a cartesian grid utilised, or it is treated as a graph?

[Vilin97]

If it helps, we can only work with graphs for now and extend to Cartesian grids later. Graphs are more general but use more memory.

Do you want any help with this? I'm not familiar with the internals of catalyst though.

[TorkelE]

Good point, yes, let's start with graphs and take it from there.

[TorkelE]

I was going to make a connection from Catalyst to LatticeReactionSystem's to JumpProblem, for spatial simulations, I think I should be able to do that and we can move from there. I think right now the most important point is opinions really, what kind of inputs should we allow and not, and then I know what to limit myself to. However, we probably need Sam's opinion here as well. @isaacsas you are busy right now though?

[isaacsas]

Hi both, feel free to go ahead without me. I'll take a look and give feedback next week or when you open a PR. I unfortunately don't have time at the moment.

[TorkelE]

I've done another revamp of the SPatial ODE code, started on building a link from Lattice Reaction NEtworks to the spatial Jump stuff.

Would it make sense to have the hopping_constants input into the DiscreteProblem rather than the JumpProblem? These are essential parameters to the simulation, so would it make sense to give these at the same time the parameter values are give, i.e. the DiscreteProblem.

Currently if you have a diffusion reaction of species X at rate dX, the value of dX is part of the parameter array that is passed into the DiscreteProblem. I will later use this to calculate the hopping constants. I could save this information in some intermediary structure within DiscreteProblem as well and produced the hopping_constants later, but to me calculating them in the DiscreteProblem seems more natural.
(but then I never really understood why we separate Discrete and JumpProblems in the first place anyway, so there might be some things I'm missing)

[TorkelE]

Relating this last issue, I don't think solving it is urgent, but something we could have a think about.

[TorkelE]

Also, am I correct in that for spatial JumpProcesses, rather than having the hopping constants/diffusion rates (whatever we call them) tied to a edge (or a combination of edges and compartment), it depends on the compartment only. That is the rate to make a jump [comp_1 -> comp_2] and [comp_1 -> comp_3] are the same?

This might be implicitly stated in your previous message @Vilin97 , but I was a bit unsure about the D and L notation on what index meant what)

[Vilin97]

There are different forms of hopping constants (btw, I do not call them diffusion rate because that is kind of an established term for the constant D in the heat equation, which depends on mesh size). Let's first get rates of form D_{s,i,j} working. To do this, hopping_constants must be of type Matrix{Vector{F}} where F <: Number (usually F is Float64) with hopping_constants[s,i][j] being the hopping rate of species s at site i to the neighbor at index j. Look at the first example here if the description above is unclear.

The notation D_{s,i,j} just means that the hopping constant D depends on three indices: s,i,j. Note that whilei is the index of the node/site i in the graph, j is the index of ith neighbor. For example, if j=1, it's the hopping constant to the first neighbor, which is likely not the first node/site in the graph. The ordering of the neighbors is given by the graph itself via its neighbors function.

[TorkelE]

Got it, will integrate this as well. Basically:

• If you give a Matrix{<:Number} as hopping_constants each value will be the rate from a compartment to its neighbours.
• If you give a Matrix{Vector{<:Number}} as hopping_constants each value will be the rate from a compartment to a specific neighbour.
  ?

You are probably right we use different terminology. In this case I image the spatial system to be e.g. a plant where we have a rigid grid of cells (compartments), each with a well-mixed system within it. Then, if you have undirected movement between neighbouring cells (compartment) I call that a diffusion reaction (and for something more fancy, a transportation reaction).

[TorkelE]

Ok, that clarification is good. I will implement this and push today.

[TorkelE]

I have added support for converting LatticeReactionSystems to JumpProblems with spatial structure. The examples in https://docs.sciml.ai/JumpProcesses/stable/tutorials/spatial/#Hopping-rates now become:

Uniform diffusion rates

# Uniform diffusion.
using Catalyst, JumpProcesses, Graphs

rs = @reaction_network begin
    (k1,k2), A + B <--> C
end
srs = diffusion_reactions([(:D, :A),(:D, :B)])
lattice = Graphs.grid([5,5])
lrs = LatticeReactionSystem(rs, srs, lattice)

u0 = [:A => [25; fill(0, 24)], :B => [fill(0, 24); 25], :C =>0]
tspan = (0.0, 3.0)
p = [:k1 => 6.0, :k2 => 0.05, :D => 1.0]

prob = DiscreteProblem(lrs, u0, tspan, p)
jprob = JumpProblem(lrs, prob, NSM())
sol = solve(jprob, SSAStepper())

Non-uniform diffusion rates (only diffusion up and right)

using Catalyst, JumpProcesses, Graphs

rs = @reaction_network begin
    (k1,k2), A + B <--> C
end
srs = diffusion_reactions([(:D, :A),(:D, :B)])
lattice = DiGraph(Graphs.grid([5,5]))
lrs = LatticeReactionSystem(rs, srs, lattice)

u0 = [:A => [25; fill(0, 24)], :B => [fill(0, 24); 25], :C =>0]
tspan = (0.0, 3.0)
D_values =  [(e.dst>e.src) ? 1.0 : 0.0 for e in edges(lattice)]
p = [:k1 => 6.0, :k2 => 0.05, :D => 1.0]

prob = DiscreteProblem(lrs, u0, tspan, p)
jprob = JumpProblem(lrs, prob, NSM())
sol = solve(jprob, SSAStepper())

Details

Currently, I am using the same interface for defining parameters and initial conditions as for the ODE system (with these being passed into the DiscreteProblem. Internally, in the DiscreteProblem, u0 is stored as a Vector{Vector{<:Number}} (where the first vector is the species, the second its value in all compartments, if identical this will be a vector of length 1 with this single value). p is similarly stored, but as a Tuple{Vector{Vector{<:Number}},Vector{Vector{<:Number}}} with parameters tied to compartments and edges stored separately. When JumpProblem is called these are then converted into the format you use.

Generally, you won't have to worry about this, I am happy to maintain the link for converting this internal format to whatever your spatial jump algorithms want in JumpProblem. There are some cases where the JumpSimulator takes parameter values as vectors and then uses products of these. In this cases, my preferred approach is a symbolic representation when the diffusion reaction is created:

dr = DiffusionReaction(D*L, :X)
p = [:D => [...], :L => [...] ]

I think this is intuitive, and also means that these cases will not appear as specific cases to the end-user. This currently works well, however, I will have to think of which cases your notation allows more compact storage. In the internal representation, things are compact, but it might be that I expand this unnecessarily when the JumpProblem is created. If this is the case, my proposal is that we have a special routine that checks when a rate is given in this form, and then uses your internal more compact representation. That way we won't have to worry about how the user inputs things.

[TorkelE]

@Vilin97 Do you want to meet sometime Monday-Wednesday and discuss through again? Except for the slowdown for large ODEs we are mostly done. If you are happy with the current internal formats for initial conditions and parameters in the JumpProblems, then I think I can handle all conversions to them. You shouldn't need do do anything more. On the Jump side I think the biggest issues is supporting more general forms of systems (as you mentioned, e.g. mass action jumps with rates depending on the compartment and constant rate jumps).

One small inconsistency we have (which need not matter, but we might wish to coordinate) is that the output solution, at each time point, is given as a nxm matrix for JumpProblems, and a vector of length n*m for ODEProblem. The first makes more sense (and was how I initially did it for ODEs), but it turned out weird when I needed to calculate the Jacobian. @isaacsas do you have any advice on how to handle this?

[Vilin97]

@TorkelE and I met offline today. Some notes:

• reshape is essentially free. Use it to make the output of ODEProblem consistent with the output of JumpProblem.
• The PR implementing the JumpProcesses spatial stuff will be separate from the ODE PR.
• We will start with one simple case (Graphs, uniform reaction rates, one form of hopping_constants) and add support for the rest in future PRs. See roadmap.
• Tests can be essentially copied from the spatial tests in JumpProcesses. The two relevant ones are ABC.jl and diffusion.jl.
• I will start working on ConstantRateJumps once the first spatial PR gets merged into Catalyst.

[TorkelE]

(Simple) spatial models are currently possible. I have also created an issue for gathering future steps for improving spatial implementations at: #726. Would it make sense to close this and gather future spatial discussions there?

[TorkelE]

Stuff from here is now intergrated to https://github.com/SciML/CatalyStuff from here is now intergrated to https://github.com/SciML/Catalyst.jl/issues/726st.jl/issues/726