Conversion from SunODE to pure PyMC 4

Questions v5
1

Hello all,

I am working on comparing SunODE to just raw PyMC 4, and am having trouble translating a SunODE Lotka-Volterra model (Hudson Bay), to PyMC 4. Mainly just want to compare the 2 and see which one we want to use for a bigger project. Below is the code from my Jupyter Notebook:

import numpy as np
import sunode
import sunode.wrappers.as_aesara
import pymc as pm
import arviz as az

times = np.arange(1900,1921,1)
lynx_data = np.array([
4.0, 6.1, 9.8, 35.2, 59.4, 41.7, 19.0, 13.0, 8.3, 9.1, 7.4,
8.0, 12.3, 19.5, 45.7, 51.1, 29.7, 15.8, 9.7, 10.1, 8.6
])
hare_data = np.array([
30.0, 47.2, 70.2, 77.4, 36.3, 20.6, 18.1, 21.4, 22.0, 25.4,
27.1, 40.3, 57.0, 76.6, 52.3, 19.5, 11.2, 7.6, 14.6, 16.2, 24.7
])

def lotka_volterra(t, y, p):
“”"Right hand side of Lotka-Volterra equation.

All inputs are dataclasses of sympy variables, or in the case
of non-scalar variables numpy arrays of sympy variables.
"""
return {
    'hares': p.alpha * y.hares - p.beta * y.lynx * y.hares,
    'lynx': p.delta * y.hares * y.lynx - p.gamma * y.lynx,
}

with pm.Model() as model:
hares_start = pm.HalfNormal(‘hares_start’, sigma=50)
lynx_start = pm.HalfNormal(‘lynx_start’, sigma=50)

ratio = pm.Beta('ratio', alpha=0.5, beta=0.5)
    
fixed_hares = pm.HalfNormal('fixed_hares', sigma=50)
fixed_lynx = pm.Deterministic('fixed_lynx', ratio * fixed_hares)

period = pm.Gamma('period', mu=10, sigma=1)
freq = pm.Deterministic('freq', 2 * np.pi / period)

log_speed_ratio = pm.Normal('log_speed_ratio', mu=0, sigma=0.1)
speed_ratio = np.exp(log_speed_ratio)

# Compute the parameters of the ode based on our prior parameters
alpha = pm.Deterministic('alpha', freq * speed_ratio * ratio)
beta = pm.Deterministic('beta', freq * speed_ratio / fixed_hares)
gamma = pm.Deterministic('gamma', freq / speed_ratio / ratio)
delta = pm.Deterministic('delta', freq / speed_ratio / fixed_hares / ratio)

y_hat, _, problem, solver, _, _ = sunode.wrappers.as_aesara.solve_ivp(
    y0={
    # The initial conditions of the ode. Each variable
    # needs to specify a theano or numpy variable and a shape.
    # This dict can be nested.
        'hares': (hares_start, ()),
        'lynx': (lynx_start, ()),
    },
    params={
    # Each parameter of the ode. sunode will only compute derivatives
    # with respect to theano variables. The shape needs to be specified
    # as well. It it infered automatically for numpy variables.
    # This dict can be nested.
        'alpha': (alpha, ()),
        'beta': (beta, ()),
        'gamma': (gamma, ()),
        'delta': (delta, ()),
        'extra': np.zeros(1),
    },
# A functions that computes the right-hand-side of the ode using
# sympy variables.
    rhs=lotka_volterra,
# The time points where we want to access the solution
    tvals=times,
    t0=times[0],
)

# We can access the individual variables of the solution using the
# variable names.
pm.Deterministic('hares_mu', y_hat['hares'])
pm.Deterministic('lynx_mu', y_hat['lynx'])

sd = pm.HalfNormal('sd')
pm.LogNormal('hares', mu=y_hat['hares'], sigma=sd, observed=hare_data)
pm.LogNormal('lynx', mu=y_hat['lynx'], sigma=sd, observed=lynx_data)

with model:
idata = pm.sample(tune=1000, draws=1000, chains=6, cores=6)

*This is the example from GitHub - aseyboldt/sunode: Solve ODEs fast, with support for PyMC3

2

Maybe see here for pm.ode.DifferentialEquation examples :point_right: GSoC 2019: Introduction of pymc3.ode API — PyMC example gallery

But I can already tell you that you’ll definitely want to use sunode for the bigger project. It has more capable solvers, adjoints gradients and is faster.

Depending on what kinds of models you are building, maybe also check our our murefi package which uses sunode internally (if installed). You can also use just parts of it, for example to predict a murefi.Dataset containing Replicates of Timeseries with TensorVariables that are already conveniently sliced to match to the shapes of observed data. In other words, it’s not technically necessary to combine it with our other package calibr8 for likelihoods.

1 Like