ODE with matrix calculation

Questions v3
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