# ODE with matrix calculation

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 * pm.math.dot(A , y) + p * 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!

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,
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

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

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