BUG: EulerMaruyama is not working

import arviz as az
import pymc as pm
import scipy
import pytensor.tensor as tt
from pymc import EulerMaruyama

%config InlineBackend.figure_format = 'retina'
az.style.use("arviz-darkgrid")



N, τ, a, m, σ2 = 200, 3.0, 1.05, 0.2, 1e-1

dt = 1e-1
xs, ys = [0.0], [1.0]
for i in range(N):
    x, y = xs[-1], ys[-1]
    dx = τ * (x - x ** 3.0 / 3.0 + y)
    dy = (1.0 / τ) * (a - x)
    xs.append(x + dt * dx + sqrt(dt) * σ2 * randn())
    ys.append(y + dt * dy + sqrt(dt) * σ2 * randn())
xs, ys = array(xs), array(ys)
zs = m * xs + (1 - m) * ys + randn(xs.size) * 0.1

figure(figsize=(10, 2))
plot(xs, label="$x(t)$")
plot(ys, label="$y(t)$")
plot(zs, label="$z(t)$")
legend()


def osc_sde(xy, τ, a):
    x, y = xy[:, 0], xy[:, 1]
    dx = τ * (x - x ** 3.0 / 3.0 + y)
    dy = (1.0 / τ) * (a - x)
    dxy = tt.stack([dx, dy], axis=0).T
    return dxy, σ2

xys = np.c_[xs, ys].T

with pm.Model() as model:
    τh = pm.Uniform("τh", lower=0.1, upper=5.0)
    ah = pm.Uniform("ah", lower=0.5, upper=1.5)
    mh = pm.Uniform("mh", lower=0.0, upper=1.0)
    xyh = EulerMaruyama("xyh", dt, osc_sde, (τh, ah), shape= N, testval=xys)
    zh = pm.Normal("zh", mu=mh * xyh[:, 0] + (1 - mh) * xyh[:, 1], sigma=0.1, observed=zs)

Error message:

/usr/local/lib/python3.10/dist-packages/pymc/distributions/timeseries.py:927: FutureWarning: The `testval` argument is deprecated; use `initval`.
  return super().__new__(cls, name, dt, sde_fn, *args, steps=steps, **kwargs)
/usr/local/lib/python3.10/dist-packages/pymc/distributions/timeseries.py:956: UserWarning: Initial distribution not specified, defaulting to `Normal.dist(0, 100, shape=...)`. You can specify an init_dist manually to suppress this warning.
  warnings.warn(
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-37-6153a081b667> in <cell line: 3>()
      5     ah = pm.Uniform("ah", lower=0.5, upper=1.5)
      6     mh = pm.Uniform("mh", lower=0.0, upper=1.0)
----> 7     xyh = EulerMaruyama("xyh", dt, osc_sde, (τh, ah), shape= N, testval=xys)
      8     zh = pm.Normal("zh", mu=mh * xyh[:, 0] + (1 - mh) * xyh[:, 1], sigma=0.1, observed=zs)

8 frames
/usr/local/lib/python3.10/dist-packages/pytensor/tensor/var.py in __getitem__(self, args)
    496         # Check if the number of dimensions isn't too large.
    497         if self.ndim < index_dim_count:
--> 498             raise IndexError("too many indices for array")
    499 
    500         # Convert an Ellipsis if provided into an appropriate number of

IndexError: too many indices for array

PyMC version information:

pymc version 5.7.2

I running my code on colab

Your shape doesn’t seem compatible with the shape of the initial value?

Hi Ricardo, Thank you for the quick answer!

using → shape = xys.shape

Leads to the same error

xyh = EulerMaruyama(“xyh”, dt, osc_sde, (τh, ah), shape= xys.shape, testval=xys)
52 zh = pm.Normal(“zh”, mu=mh * xyh[:, 0] + (1 - mh) * xyh[:, 1], sigma=0.1, observed=zs)
53

8 frames
/usr/local/lib/python3.10/dist-packages/pytensor/tensor/var.py in getitem(self, args)
496 # Check if the number of dimensions isn’t too large.
497 if self.ndim < index_dim_count:
→ 498 raise IndexError(“too many indices for array”)
499
500 # Convert an Ellipsis if provided into an appropriate number of

IndexError: too many indices for array

Am I missing something?