Getting solve_ivp to work with PyMC?

,
1

Hi all,

I am trying to implement a simple exponential ODE model in the Bayesian framework using both odeint and solve_ivp methods from the scipy.integrate library.

import pymc as pm
import pytensor
import pytensor.tensor as pt
import arviz as az

import matplotlib.pyplot as plt
from pytensor.compile.ops import as_op
from scipy.integrate import odeint
from scipy.integrate import solve_ivp

The equation is very simple:

def exponential_decay(t, y): 
    
    dydt = -0.5 * y
    return dydt'

ODE-solution with solve_ivp without using pymc:

sol = solve_ivp(
    exponential_decay,
    t_span= [0, 10],
    y0 = [10],
    max_step=0.1)

fig, ax = plt.subplots()

ax.scatter(
    sol.t,
    sol.y[0],
    s=50,
    marker="^",
    label="ode-solution")

ax.grid()
ax.set_xlabel("Time [s]")
ax.set_ylabel("y [m]")
ax.legend()
plt.show()

Then, I proceed with generating synthetic data with Gaussian noise for the ODE model:

ode_array_size = len(sol.y[0])
mu, sigma = sol.y[0], 1

simulated_data_y = np.random.normal(mu, sigma, ode_array_size)

I learned from this example that we would need to convert the data-type to pytensor in order to use PyMC (Link: PyMC ODE Example):

@as_op(itypes=[pt.dvector], otypes=[pt.dvector])
def pytensor_ode_model(param):

    t_start, t_stop, y_init = param
    sol = solve_ivp(
        exponential_decay,
        t_span = [t_start, t_stop],
        y0 = [y_init],
        max_step=0.1)

    return sol.y[0]

Afterwards, one could get sample and get the distribution of the parameters:

t_start = 0.1
t_stop = 10
y_init = 10

with pm.Model() as model:
    
    y_init = pm.Normal("y_init", mu=y_init, sigma=1)
    t_start = pm.TruncatedNormal("t_start", mu=t_start, sigma=0.1, lower=0, initval=t_start)
    t_stop = pm.TruncatedNormal("t_stop", mu=t_stop, sigma=1, lower=0, initval=t_stop)
    
    sigma = pm.HalfNormal("sigma", 1)
    
    ode_solution = pytensor_ode_model(pm.math.stack([y_init, t_start, t_stop]))
    
    pm.Normal(
        "Y_obs",
        mu=ode_solution,
        sigma=sigma,
        observed=simulated_data_y)

    vars_list = list(model.values_to_rvs.keys())[:-1]
    
    trace = pm.sample(
            step=[pm.Slice(vars_list)],
            tune=2000,
            draws=2000,
            cores=1)

I get some errors about Input dimension mismatch:

ValueError: Input dimension mismatch: (input[0].shape[0] = 200, input[1].shape[0] = 101)
Apply node that caused the error: Composite{switch(i4, (((-0.5 * sqr(((i0 - i1) / i2))) - 0.9189385332046727) - i3), -inf)}(Y_obs{[ 1.022617 ... 63610e-02]}, FromFunctionOp{pytensor_ode_model}.0, ExpandDims{axis=0}.0, Log.0, Gt.0)

If I would solve the same equation using odeint, I would be able to proceed without errors.

debug_solve_ivp_ode
debug_solve_ivp_ode1374×490 77.2 KB

So, I am not sure what are the changes that I would need to make to resolve the input dimension problem using solve_ivp?

Thank you for any hints towards resolving this issue!