Implement spatial ODE simulations #571

TorkelE · 2022-12-02T20:35:15Z

New implementation of spatial ODE simulations. Instead of being based on MTK we now just write a new f function for our given spatial system, that loops through all compartments and connections. The advantage of this is that we no longer need to store a big equation, which helps with simulations over a large number of cells.

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 A B
srs = [SpatialReaction(:D, ([:X], []), ([], [:X]))]
lattice = grid([20, 20])
lrs = LatticeReactionSystem(rs, srs, 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, Tsit5());

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)

which show spatial effects in the Brusselator:
https://user-images.githubusercontent.com/18099310/205381134-b72f02bd-c24a-4264-a238-41d0d8fa2420.mp4
(I've never used Makie before, so still some things to iron out there to make the animation prettier)

isaacsas · 2022-12-02T21:36:20Z

Cool! Will give a review tomorrow!

isaacsas

I'm ok with merging something like this, though I'd prefer to drop the Symbol usage and make SpatialReactions look more like normal Reactions in their interface.

But I think we will need to iterate on this interface over time to get it both performant and convenient to use.

@ChrisRackauckas any thoughts?

src/spatial_reaction_network.jl

isaacsas · 2022-12-05T16:56:52Z

src/spatial_reaction_network.jl

+# A set of spatial reactions (denoting interaction between comaprtments).
+# A network of compartments (a meta graph that can contain some additional infro for each compartment).
+# The lattice is a DiGraph, normals graphs are converted to DiGraphs (with one edge in each direction).
+struct LatticeReactionSystem # <: MT.AbstractTimeDependentSystem # Adding this part messes up show, disabling me from creating LRSs


Define a custom show then?

Ahh, yeah, that makes sense.

isaacsas · 2022-12-05T16:57:25Z

src/spatial_reaction_network.jl

+                               kwargs...)
+    @unpack rs, spatial_reactions, lattice = lrs       
+
+    spatial_params = unique(getfield.(spatial_reactions, :rate))


Again, why not use symbolic variables here?

I'm happy to use symbolic variables, but was a bit uncertain of the proper way to use them here. Also, given that we don't actually use MTK I presume there was no advantage?

The advantage is consistency with how non-spatial Catalyst works and with MTK (particularly since we will hopefully be able to switch to MTK/Symbolics for the code generation at some point). Also, wouldn't we potentially want to allow more complicated rate expressions at some point that are symbolic expressions?

Basically we should use symbolic variables/parameters as the standard objects -- Symbols in parameter mappings and initial conditions are just provided as a convenience (and should be the only place we currently use Symbols in the non-spatial code).

isaacsas · 2022-12-05T16:59:30Z

src/spatial_reaction_network.jl

+    @unpack rs, spatial_reactions, lattice = lrs       
+
+    spatial_params = unique(getfield.(spatial_reactions, :rate))
+    pV_in, pE_in = split_parameters(p, spatial_params)


Longer-term this interface will probably need changes. It is pretty common that transport rates for diffusion all have the form D_S k_{i,j} for D_S the diffusivity of a species, and the k_{i,j} the same across all species. As such, we'd want to avoid users having to essentially instantiate these matrices for every species to save memory (i.e. we'd want to support different forms for these rate structures as JumpProcesses attempts to do). But that is an optmization we can work towards later on as we refine the interace.

I'm not sure if I fully follow.

Right now the parameter interface requires a user to give the full set of weights for all spatial transitions for each species across the graph edges right? In many cases though, a transition from node i to node j has the form I wrote above, where there is a species-dependent (spatially-constant) diffusion rate that multiplies a species-independent transition rate k_{i,j}. We can discuss more on Monday if that is still confusing!

Ok, I think I get it. Basically, typically there are N+M diffusion parameters (N is the number of species and M is the number of edges on the lattice). Then the diffusion rate along an edge is the product of a parameter for the species and a parameter for the edge.

The current interface doesn't directly permit for this. You could declare NxM different parameters, and set their values to follow this rule. However, a direct way to do this would be convenient?

I think this should be doable. I already implemented (for the old spatial reaction, but will re-implement) a DiffusionReaction function, that creates a diffusion reaction without lots of empty [] having to be passed to the SpatialReaction. I could expand on this.

I think we can follow up on this after thinking a bit more about how to specify the full spatial parameters anyways, so what you have for now is probably fine for this PR. (In an effort to avoid it growing really big and/or sitting forever.)

src/spatial_reaction_network.jl

isaacsas · 2022-12-05T17:06:02Z

src/spatial_reaction_network.jl

+matrix_form(input::Vector, n, index_dict) = matrix_form(map(i -> (i[2] isa Vector) ? i[1] => i[2] : i[1] => fill(i[2], n), input), n, index_dict);
+
+# Creates a function for simulating the spatial ODE with spatial reactions.
+function build_f(lrs::LatticeReactionSystem, u_idxs::Dict{Symbol,Int64}, pE_idxes::Dict{Symbol,Int64})


Since your generated function is capturing variables, why not use a functor (i.e. make a structure, SpatialReactionODEFun that stores the variables and then create a general functor SpatialReactionODEFun(du, u, p, t))? This should avoid any need to dynamically generate a function right? Then you just return an instance of the function.

That's probably the way to go, I'm just unfamiliar with actually using the approach and its advantages,

https://docs.julialang.org/en/v1/manual/methods/#Function-like-objects

mutable struct GraphODEFunWrapper p1 p2 end # in-place ODE derivative function function (g::GraphODEFunWrapper)(du, u, p, t) # can now use g.p1 and g.p2 end

So you create a new GraphODEFunWrapper to store the needed parameters that aren't in p and return that to the user. It can still be called like a function in the ODE solvers.

isaacsas · 2022-12-05T17:07:32Z

src/spatial_reaction_network.jl

+end
+
+# Get the rate of a specific reaction.
+function get_rate(sr, pE, u_src, u_dst, u_idxs, pE_idxes)


Why not reuse the current ODE rate generation code for non-spatial ReactionSystems?

Will check if that works, good point.

I had a look at this, but didn't really see how I could reuse it given that here we have two sets of substrates (one for each of the two neighbouring compartments.

Also, for normal ODEs we just generate a symbolic expression with oderatelaw, while here we actually have to calculate a number to plug into our ODE function. Similarily I was a bit uncertain how to utilise symbolic rates in the SpatrialReaction, given that we at some point have to go back to just dealing with numbers, and have to get a value from the symbolic expression and the p and u matrices.

TorkelE · 2022-12-14T20:43:25Z

I had to revert to specifying things as Symbols and Int64s in the SpatialReaction. Turns out, due to how the f function is created, not specifying these leads to type instability.
(why I am not entirely certain of, have been trying to figure out for a while, but not getting anywhere)

isaacsas · 2022-12-14T20:55:56Z

That sounds like an issue with the f implementation. It really shouldn't be using either symbols or symbolics internally, just the index of a species.

TorkelE · 2022-12-14T21:04:31Z

Yeah, there's definitely something with the f function. I guess it is something with a function (ODEProblem) creating a function (the f function) that means the second function must have all these type specifications without crashing. What it is, I am uncertain of. I will probably need to find an expert to help have a look at it for me.

Right now I don't really use symbols on the inside either. However, given that parameter inputs are given on forms like:

params = [k1 => 1.0, p1 =>3.0]

each parameter has to have some kind of name/tag so that we can figure out which parameter should have which value from the input. Internally, we use Int64s to index the array though. Currently avoiding Symbolics entierly.

TorkelE · 2023-01-17T15:14:17Z

Gone through the code, and now the profiler looks good. Still a problem that:

using Catalyst, OrdinaryDiffEq, Graphs

rs = @reaction_network begin
    A, ∅ → X
    1, 2X + Y → 3X
    B, X → Y
    1, X → ∅
end A B
srs = [SpatialReaction(:D, ([:X], []), ([], [:X]), ([1], Vector{Int64}()), (Vector{Int64}(), [1]))]
lattice = grid([25, 10])
lrs = LatticeReactionSystem(rs, srs, 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]

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

e.g. the FBDF solver here doesn't work well (although I can do autodiff now). Tried a lot of implicit solvers, but they all scale poorly with lattice size. For non-stiff problems using the explicit ones should be fine, but unsure what to do for stiff problems.

ChrisRackauckas · 2023-01-17T15:15:48Z

Preconditioned GMRES?

TorkelE · 2023-01-20T20:19:02Z

I can run GMRES (small speedup), but haven't managed to get preconditioners working yet.

However, GMRES is Jacobian free-right? And currently the problem is that getting a Jacobian scales poorly with the problem size? I don't think I had this problem when I created the problem through MTK, but MTK automatically produces an jacobian for you. However, it should be possible to create the Jacobian at problem creation from the input (just as the f function is created).

Does this reasoning make sense? In which case trying to generate a Jacobian when the ODEProblem is created might solve thing, and allow the implicit methods to scale with the explicit ones?

ChrisRackauckas · 2023-01-20T20:20:20Z

GMRES is Jacobain-free. Preconditioners need Jacobians but if you just do sparse=true it can just use coloring to build the Jacobian for the preconditioner.

TorkelE · 2023-01-20T20:24:17Z

Figured so. But since we are not using MTK then we would have to create the Jacobian from scratch anyway? Which would mean that there is no way around writing a function that can generate the jacobian.

ChrisRackauckas · 2023-01-20T20:25:12Z

No, you'd just need to determine the sparsity pattern.

TorkelE · 2023-01-20T20:25:42Z

For the preconditioner? Got it.

TorkelE · 2023-05-03T23:10:00Z

Ok, we are back in action. In addition to, for a

Reaction network
Set of spatial reaction events
Network of adjacent compartments
computing the f function, we now also compute the jacobian and the jac_prototype with the sparsity. Using this, implicit solvers scale well with the number of compartments, hence we can simulate large stiff systems.

There are still some minor tuning to do, as well as some deciding how to set up the inputs. However, the big hurdle should be done now.

TorkelE · 2023-05-04T19:19:07Z

Since it is a big thing, I simply reuploaded it as #644, trying to ensure that it is properly built on the latest Catalyst version.

isaacsas reviewed Dec 5, 2022

View reviewed changes

TorkelE mentioned this pull request Dec 6, 2022

Lattice reaction network #558

Closed

TorkelE added 11 commits May 4, 2023 14:28

Implement spatial ODE simulations

c1b1f44

Use Symbolics.getname

c97ae85

Change SpatialReaction structure to similar of normal reaction.

bb7de6d

Add simple test

8b4186c

Basic formatting

d04abf3

Make function for simple spatial reactions

e663cd2

Revert Symbol specification

c58d27e

Fix type instability and garbage collection error

c42ebb6

Add automatic, sparse, jacobian

759c2de

minor updates

c60b98b

update sparsity usage

1b2e152

TorkelE force-pushed the ODE_n_LatticeReactionSystem_dec2022 branch from 17f64f4 to 1b2e152 Compare May 4, 2023 18:30

TorkelE mentioned this pull request May 4, 2023

Spatial Reaction Network Implementation #644

Merged

TorkelE closed this May 4, 2023

TorkelE deleted the ODE_n_LatticeReactionSystem_dec2022 branch June 8, 2024 18:39

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement spatial ODE simulations #571

Implement spatial ODE simulations #571

TorkelE commented Dec 2, 2022

isaacsas commented Dec 2, 2022

isaacsas left a comment

isaacsas Dec 5, 2022

TorkelE Dec 5, 2022

isaacsas Dec 5, 2022

TorkelE Dec 5, 2022

isaacsas Dec 7, 2022

isaacsas Dec 5, 2022

TorkelE Dec 5, 2022

isaacsas Dec 7, 2022 •

edited

Loading

TorkelE Dec 7, 2022

isaacsas Dec 7, 2022

isaacsas Dec 5, 2022

TorkelE Dec 5, 2022

isaacsas Dec 7, 2022

isaacsas Dec 5, 2022

TorkelE Dec 5, 2022

TorkelE Dec 9, 2022

TorkelE commented Dec 14, 2022

isaacsas commented Dec 14, 2022

TorkelE commented Dec 14, 2022

TorkelE commented Jan 17, 2023

ChrisRackauckas commented Jan 17, 2023

TorkelE commented Jan 20, 2023

ChrisRackauckas commented Jan 20, 2023

TorkelE commented Jan 20, 2023

ChrisRackauckas commented Jan 20, 2023

TorkelE commented Jan 20, 2023

TorkelE commented May 3, 2023

TorkelE commented May 4, 2023

Implement spatial ODE simulations #571

Implement spatial ODE simulations #571

Conversation

TorkelE commented Dec 2, 2022

isaacsas commented Dec 2, 2022

isaacsas left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

isaacsas Dec 7, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

TorkelE commented Dec 14, 2022

isaacsas commented Dec 14, 2022

TorkelE commented Dec 14, 2022

TorkelE commented Jan 17, 2023

ChrisRackauckas commented Jan 17, 2023

TorkelE commented Jan 20, 2023

ChrisRackauckas commented Jan 20, 2023

TorkelE commented Jan 20, 2023

ChrisRackauckas commented Jan 20, 2023

TorkelE commented Jan 20, 2023

TorkelE commented May 3, 2023

TorkelE commented May 4, 2023

isaacsas Dec 7, 2022 •

edited

Loading