Getting solve_ivp to work with PyMC?

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.

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!