ODE with matrix calculation

Questions v3
1

Hi, I’m trying to infer two parameters in an ODE system called p0 and p1 in pymc3. Here A is a known 3 * 3 matrix. y is the observation, which is a vector containing 3 elements. The idea is to infer p0 and p1 given the observation y.

In my case the derivative dy/dt includes both :

  1. product between matrix and column vextor : A @ y, which returns a column vector
  2. elementwise product between two column vectors: y * (3 - 2y), which returns a column vector

In scipy, we can define such an ODE system by:

import numpy as np
from scipy.integrate import odeint

A = np.array([[-0.25,    0,    0],
               [ 0.25, -0.2,    0],
               [    0,  0.2, -0.1]])

def deriv(y, t, p0, p1):
    return p0 * A @ y + p1 * y * (3 - 2 *y)

p0 = 2
p1 = 4

time = np.arange(0,100,0.1)
y0 = np.array([10, 20, 30])

sol = odeint(deriv, y0, time, args=(p0,p1))

However, in pymc, when I tried to use DifferentialEquation:

def odefunc(y, t, p):
    return p[0] * pm.math.dot(A , y) + p[1] * y * (3 -  2 * y) 

times = np.arange(0,100,0.1)
ode_model = DifferentialEquation(func=odefunc, times=times, n_states=3, n_theta=2, t0=0)`

I got the following error:

AssertionError: theano.tensor.jacobian expects a 1 dimensional variable as `expression`. If not use flatten to make it a vector

I also found the sunode package difficult to implement such kind of idea. (Maybe it’s because I’m not so familiar with sunode).

Thus, I’m wondering if there’s a way of implementing such ODE system with both calculations: (product between matrix and column vextor) and (elementwise product between two vectors) mentioned above in pymc? I need these two calculations because in reality my matrix A and vector y are large.

Also, I’ve noticed that the odeint function in jax.experimental.ode is able to do that. Is it possible to use the jax odeint solver in pymc?

Thank you so much for your help!

2

I am not familiar with the DifferentialEquation module, but if you would like to try it with sunode, your example might look like this:

import numpy as np
import pymc as pm
from sunode.wrappers.as_theano import solve_ivp

def deriv(t, y, p):
    return {"y": p.p0 * p.A @ y.y + p.p1 * y.y * (3 - 2 * y.y)}


A = np.array([[-0.25, 0, 0], [0.25, -0.2, 0], [0, 0.2, -0.1]])
time = np.arange(0, 100, 0.1)
y0 = np.array([10.0, 20.0, 30.0])

coords = {"state": [0, 1, 2], "time": time}

with pm.Model(coords=coords) as model:
    # Create variables for the parameters you want to estimate
    p0 = pm.LogNormal("p0", mu=np.log(2), sigma=1)
    p1 = pm.LogNormal("p1", mu=np.log(4), sigma=1)

    y_hat, _, problem, solver, _, _ = solve_ivp(
        y0={
            "y": (y0, ("state",)),
        },
        params={
            "p0": (p0, ()),
            "p1": (p1, ()),
            "A": A,
        },
        rhs=deriv,
        tvals=time,
        t0=time[0],
        coords=coords,
        solver_kwargs=dict(abstol=1e-6, reltol=1e-8),
    )
    # Save the solution for y as Deterministic
    pm.Deterministic("y", y_hat["y"], dims=("time", "state"))

# Print the symbolic representation of the right-hand-side
problem._sym_dydt

# Sample the prior distribution
with model:
    prior = pm.sample_prior_predictive()
1 Like
3

Thank you so much for your detailed reply and patient help! I’m so happy to find out that sunode can give a solution!

One challenge I’m currently facing is that: instead of the 3*3 matrix A and the vector y with length 3, in reality, my matrix is actually 85 * 85 and len(y) = 85. Then I find that it takes quite a long time to run your code above (about 10hours on my laptop). I’m wondering what has caused this and if there’s any way to speed up? Thank you so much!

1 Like
4

I think the problem is that the automatic code generation from sympy expressions becomes very inefficient when the number of states becomes large. I have been facing this issue too when I tried to use sunode for a discretized version of a PDE, and there is also an issue (opened by someone else) on github.
I have been discussing this with @aseyboldt (the main developer of sunode). He thinks that it might be possible to solve this by using the new numba backend in aesara. I had started experimenting with this last November but then I did not have the time to get it to work. I would like to get back to this but currently I am still very busy with other stuff.

1 Like