Spatial Reaction Network Implementation

[TorkelE]

A revamped version of #571. While we might be able to tune some stuff, this should work well. Jacobian can be built, so stiff systems can be implemented properly.

Syntax is similar to previously, see this example:

When making a spatial simulation:
u0 is a NxM matrix, where N is the number of species in the (non-spatial) reaction network, and M is the number of compartments.
p is a tupple (pV,pE),
pV is a LxM matrix, where L is the number of parameters in the (non-spatial) reaction network, and M is the number of compartments.
pE is a RxS matrix, where R is the number of parameters tied to spatial reactions, and S is the number of connections.

Sample code:

using Catalyst, OrdinaryDiffEq, Graphs

rs = @reaction_network begin
    A, ∅ → X
    1, 2X + Y → 3X
    B, X → Y
    1, X → ∅
end
tr = @transport_reaction D X
lattice = Graphs.grid([20, 20])
lrs = LatticeReactionSystem(rs, [tr], lattice);

u0_in = [:X => 10 * rand(nv(lattice)), :Y => 10 * rand(nv(lattice))]
tspan = (0.0, 100.0)
p_in = [:A => 1.0, :B => 4.0, :D => 0.2]

@time oprob = ODEProblem(lrs, u0_in, tspan, p_in)
@time sol = solve(oprob, QNDF());

Plotting is not added to Catalyst right now, but running this code:

using GLMakie
using Makie.Colors
using GraphMakie
function animate_spatial_sol(sol, x_pos, y_pos; animation_name="spatial_animation.mp4", framerate=100, min_val=nothing, max_val=nothing, var=1, node_size=2.0, node_marker=:circle, edge_width=0.0, timetitle=true)
    trajectories = map(vals -> vals[var, :], sol.u)
    mini, maxi = extrema(vcat(trajectories...))
    !isnothing(min_val) && (mini = min_val)
    !isnothing(max_val) && (maxi = min_val)
    fixed_layout(_) = map(i -> (x_pos[i], y_pos[i]), 1:length(trajectories[1]))

    idx = Observable(1)
    colors = @lift map(val -> RGB{Float64}(val, val, 0.0), (trajectories[$idx] .- mini) ./ maxi)
    title = timetitle ? @lift("t = $(sol.t[($idx)])") : ""
    fig = graphplot(lattice, layout=fixed_layout, node_size=node_size, node_marker=node_marker, edge_width=edge_width, node_color=colors,
        axis=(title=title,))

    record(fig, animation_name, 1:length(sol.t);
        framerate=framerate) do i
        idx[] = i
    end
end

enable us to create a plot through:

x_pos = vcat(fill(1:20, 20)...)
y_pos = vcat(map(i -> fill(i, 20), 1:20)...)
@time animate_spatial_sol(sol, x_pos, y_pos; node_size=40.0, node_marker=:rect)

[TorkelE]

Updated using improved jacobian, which should make scaling decent (although something still seems to happen around 10,000 compartments that I do not understand).

Syntax is basically the same, with the exception that the only type of spatial reaction is the normal diffusion reaction (where a species moves from one compartment to an adjacent one): DiffusionReaction(:dX, :X) (the first argument is the parameter symbol, the second the diffusing species).

[TorkelE]

This PR is probably as good as it can get now, the only possible addition is documentation (which I will add once we are happy with it. Around n=10,000 nodes, performance for implicit ODE solvers starts to drop (for the brusselator it can take about a minute, but it gets distinctly non-linearly worse as n increases). I have tried to track this down, but very little of the run time seems to be in the f function and Jacobian evaluation, and I think that the likely culprit is probably the linear solver.

I have modified the example in the initial post to work with the (slightly) modified new syntax. Where to incorporate the plotting functionality is still uncertain.

Currently, a LatticeReactionSystem is created by bundling

1. A normal ReactionSystem
2. A vector of SpatialReactions.
3. A graph.

Currently, the only type of SpatialReaction supported are diffusion reactions (a single component that diffuses linearly to its concentration to neighbouring compartments. These are created using:

DiffusionReaction(:dX, :X)

where dX is the parameter for the diffusion rate and :X is the diffusing species. One can also use

DiffusionReactions([(:dX, :X), (:dY, :Y), (:dZ, :Z)])

to create a vector of them directly. If only a single DiffusionReaction is given to LatticeReactionSystem, it is automatically converted to a vector.

The graph can be a directed on an undirected graph. If it is directed, species only diffuse in the direction of the connection. Internally, if you give an unidrected graph, it is connected to a directed graph with connections in both directions.

The initial condition u0 can either be:

1. Vector of numbers (ordered according to the species). Here, the nth number is the initial condition of the n'th species.
2. Vector of vectors (ordered according to the species). Here, each vector must have the same number of elements as the system's number of compartments. The m'th element of the n'th vector is the initial condition of the n'th species in the m'th compartment.
3. A vector of a combination of vectors of numbers (here, if the n'th vector is a number, the n'th species has identical initial conditions across all comaprtments, else the values in the vector is used).
4. A vector of pairs, where the first element in each pair is a species, and the second element is its initial condition (either a single number, of a vector with the same number of elements as the system have comaprtments.
5. An nxm matrix, where the n'th row is the values of the n'th species and the m'th column the values of the m'th compartment.

Parameters are similar, however, contain both spatial and normal parameters. Here, preferably the pair notation is used (listing all parameters in one vector. Alternatively, the parameter can be given as a tuple (pC, pE), where the first element is the non-spatial parameters, and the second the spatial ones. Note that the spatial parameters are spread out across edges (so if there are l edges, these can e.g. be given in a nxl matrix).

If an undirected graph with nE edges was given, internally, this is converted to a directed graph with 2*nE edges. Here, if a spatial parameter is given a vector of values of length nE, it is presumed that its n'th value is the values for the two (directed) edges corresponding to the n'th (undirected) edge.

[TorkelE]

The format checker fails, but it is the "docs/make.jl" file, but not sure what to do about that. Otherwise this should be all good.

[isaacsas]

Can you remind me why you need to use symbols instead of symbolics in specifying spatial reactions? This is a big break from the ModelingToolkit/Catalyst symbolic interface...

[TorkelE]

Since the modelingtoolkit/symbolic is never used for those it seemed weird to make them symbolics when I would just internally convert them to symbols/numbers internally. It should be possible to allow them to be both symbols/symbolics though, if that would be desirable?

[isaacsas]

I don't think the Catalyst tests should have timing tests. That is hardware dependent, will likely change when using Github CI. That kind of stuff should go in a SciMLBenchmark instead.

I also think we should switch to taking symbolics for spatial reactions to preserve the interface with how normal reactions are specified, and the more general ModelingToolkit approach of specifying things with symbolics instead of Symbols.

[TorkelE]

Sounds good,
The timers were mostly for me to keep track of whether my changes improved performance or not, I will remove them from tests, but save them somewhere (maybe just comment them out). I will make an update so that Symbolics can be used as well.

[TorkelE]

I have removed the timing tests, and also added the ability to use Nums. I also made another test checking that this works:

@parameters dS dI dR
    @variables t
    @species S(t) I(t) R(t)
    SIR_srs_numsym_1 = diffusion_reactions([(:dS, :S), (:dI, :I), (:dR, :R)])
    SIR_srs_numsym_2 = diffusion_reactions([(dS, :S), (dI, :I), (dR, :R)])
    SIR_srs_numsym_3 = diffusion_reactions([(:dS, S), (:dI, I), (:dR, R)])
    SIR_srs_numsym_4 = diffusion_reactions([(dS, S), (dI, I), (dR, R)])
    SIR_srs_numsym_5 = diffusion_reactions([(dS, :S), (:dI, I), (dR, :R)])
    SIR_srs_numsym_6 = diffusion_reactions([(:dS, :S), (:dI, I), (dR, R)])

    u0 = [:S => 990.0, :I => 20.0 * rand_v_vals(small_2d_grid), :R => 0.0]
    pV = SIR_p
    pE_1 = [:dS => 0.01, :dI => 0.01, :dR => 0.01]
    pE_2 = [dS => 0.01, dI => 0.01, dR => 0.01]
    pE_3 = [dS => 0.01, :dI => 0.01, :dR => 0.01]
    ss_explicit_base = solve(ODEProblem(LatticeReactionSystem(SIR_system, SIR_srs_numsym_1, small_2d_grid), u0, (0.0, 10.0), (pV, pE_1); jac = false), Tsit5()).u[end]
    ss_implicit_base = solve(ODEProblem(LatticeReactionSystem(SIR_system, SIR_srs_numsym_1, small_2d_grid), u0, (0.0, 10.0), (pV, pE_1); jac = true), Rosenbrock23()).u[end]

    for srs in [
            SIR_srs_numsym_1,
            SIR_srs_numsym_2,
            SIR_srs_numsym_3,
            SIR_srs_numsym_4,
            SIR_srs_numsym_5,
            SIR_srs_numsym_6,
        ], pE in [pE_1, pE_2, pE_3]
        lrs = LatticeReactionSystem(SIR_system, srs, small_2d_grid)
        ss_explicit = solve(ODEProblem(lrs, u0, (0.0, 10.0), (pV, pE); jac = false), Tsit5()).u[end]
        ss_implicit = solve(ODEProblem(lrs, u0, (0.0, 10.0), (pV, pE); jac = true), Rosenbrock23()).u[end]
        @test all(isapprox.(ss_explicit, ss_explicit_base))
        @test all(isapprox.(ss_implicit, ss_implicit_base))
    end

was it something like this you had in mind?

[isaacsas]

I'll try to get back for a review later in the week. Sorry for the delay, I'm dealing with a ton of deadlines right now...

[TorkelE]

Now worries, let's discuss on Wednesday, probably the quickest.

[TorkelE]

I tried to check allocation when using the Jacobian:

using Catalyst, OrdinaryDiffEq, BenchmarkTools

SIR_system = @reaction_network begin
    α, S + I --> 2I
    β, I --> R
end
SIR_p = [:α => 0.1 / 1000, :β => 0.01]
SIR_u0 = [:S => 999.0, :I => 1.0, :R => 0.0]

SIR_sr = diffusion_reactions([(:dS, :S), (:dI, :I), (:dR, :R)])
small_2d_grid = Graphs.grid([5, 5])
SIR_lrs = LatticeReactionSystem(SIR_system, SIR_sr, small_2d_grid)

oprob = ODEProblem(SIR_system, SIR_u0, (0.0, 500.0), SIR_p)
oprob_spat = ODEProblem(SIR_lrs, SIR_u0, (0.0, 500.0), (SIR_p, [:dS=>0.1, :dI=>0.2, :dR=>0.05]))



J_tmp = deepcopy(oprob_spat.f.jac_prototype)
@benchmark oprob_spat.f.jac(J_tmp, oprob_spat.u0, oprob_spat.p, 0.0)

It says allocs estimate: 5, would that be considered small?

[isaacsas]

If you merge master into your PR, I capped the formatter script to use the last version before the recent style change. Install that version (1.0.32) and you can format locally, and hopefully it will then only impact files you've modified.

[TorkelE]

Ok, this now based on symbolics. That mean it supports stuff like

@parameters d1 d2
@variable t 
@species X(t)
DiffusionReaction(d1*d2, X)

I also tried implementing a spatial f and jacobian for the Brusselator manually, to compare performance. For large system it is almost 2x faster. However, I still get the poor scaling where the simulation time approaches a minute when the number of comaprtments is 10000 (this is without rescaling the diffusion rate as number of compartment are increased, I should note).

[TorkelE]

I have recreated the example from https://docs.sciml.ai/JumpProcesses/stable/tutorials/spatial/, but going through a LatticeReactionSystem, will prettify the code and upload tomorrow.

[TorkelE]

Have made some improvements to the internals to make it easier to read. This should now also support spatial jump simulations, see #659.

I'm fairly happy with this now. What is currently lacking is:

• Some tweaking to better mesh with specialised forms of spatial Jump problem inputs.
• The performance of the ODE simulations still seems to suffer from similar slow downs at around 10000 compartments. I have gone through a lot, but really cannot find any issue with the code. To me it seems the runtime is all in the linear solve (but I am not sure how to improve this). There could definitely be something in the code here to fix it, but at this stage I do not know how to find it.
• The creation of the JumpProblem (by extracting mass action jumps) might be a bit unsmooth, if either of you (@isaacsas @Vilin97) could check if there's a more natural way to do it that could be useful.

Except for that, I think we just need to create some documentation and plotting interfaces.

[isaacsas]

@TorkelE can you put the jump updates in a separate PR to keep this more manageable? Then when this gets merged you can update it to master so it only shows the new commits.

[TorkelE]

Done. Tomorrow I will see if someone in the office can help me rebase this on the master branch as well

[TorkelE]

The latest master should now be part of this one. The format check still fails' @isaacsas , not sure what is going on though. Either way, when there's opportunity it would be good to deal with this one. If it is possible to make more optimization on the ODE side maybe we could do that later (so that we could move forward with the Jump stuff as well, which depend on this one).

[isaacsas]

@TorkelE here are some comments so far. I'm about halfway through the main lattice_reaction_system file, but not sure I'll be able to look at this more today.

[isaacsas]

Do we have examples of any spatial modeling packages and their notations? It might be helpful in further cleaning up / clarifying the interface here too.

[TorkelE]

Do you mean within Julia, or a spatial CRN modelling package?

[isaacsas]

I mean a spatial CRN modeling package. Like PyRDME, STEPS, the one that was mentioned in our session at JuliaCon Medyon I think, etc.

[isaacsas]

They've all already designed ways to specify lattice reactions, so it might be worth looking at their interface to see if they have nice ways for users to specify the transport model.

[TorkelE]

Just checking, if we update it like this, the intended workflow would be something like:

using Catalyst, Graphs
rn = @reaction_system begin
    (p,d), 0 <--> X
end
@parameters D
@variables t
@species X(t)
srs = [TransportReaction(D, X)]
lattice = Graphs.grid([5, 5])
lrs = LatticeReactionSystem(rn, srs, lattice)

?

I would have used unpack X = rn, however, if we use a macro (like @transport_reaction D X) or something else to create the transport reaction we have no idea of whether X comes from an already created ReactionSystem or not. In practise, for this case, one would use @unpack, but I think the code should work without.

[isaacsas]

That is what I added on to my earlier comment. I'd be ok with requiring all species to be declared within the attached ReactionSystem, so the only new information one can add is parameters representing transport-related information. Then one could use @unpack, or we make it work using namespacing, i.e. rn.X. For the former we just mimic the current macro behaviors

@unpack X = rn
@parameters D
@transport_rx $D $X

# or this creates the parameter if not defined
# then in LatticeReactionSystem we need to loop through and 
# pull out these parameters from the spatial reactions, as we do in ReactionSystem
@transport_rx D $X

We very much should not have users re-declare a species that is already in rn as that could lead to metadata being dropped and having two different versions of the same symbolic.

[isaacsas]

To be honest though, I find the current DiffusionReaction notation a bit confusing. A user is specifying the rates and the species, but not the actual transitions? I wonder if it might make more sense to just require attaching the transport rates as metadata to X, i.e.

@species X [transportrates = X]
rs = ...
lrs = LatticeReactionSystem(rs, grid)

[TorkelE]

You mean @species X [transportrates = D]? And then for more complicated expressions @species X [transportrates = Di*Dx/2]?

That could work. What if one would want to e.g. create a SIR model, and then from it create different spatial models (where different subsets of species move)? Also, in the future we might want to add spatial reactions that are no just one species moving compartment. It could be a species moving to another compartment while being converted to another one like
A_i --> B_j
Wouldn't defining the spatial dynamics separately help future-proof us for this case as well? The previous version of this PR did just that, before we simplified to only this case.
(Still happy to go with your approach if you think it is better)

[isaacsas]

I don't feel strongly about using metadata, so it is fine to keep the TransportReaction representation. I think that I just don't quite understand how I would encode a formula like the rate to hop from $i \to j$ for species $S$ is $D_s L_{i,j}$ where $D_s$ is a species-dependent diffusivity, and $L_{i,j}$ is a set of transition rates across the graph.

[TorkelE]

Bummer. Still some problem with SteadyStateDiffEq. Avik said that v2 was ready, but they forgot to release it. Should hopefully be out soon, might fix this.

[TorkelE]

OK, it have now switched to the latest SteadyStateDiffEq version (v2), but now NonlinearSolve and stuff seem to fail. Saw someone report the same issue on Slack, so might be a common thing?