Multivariate EulerMaruyama

[fonnesbeck]

Describe the issue:

As described by #6953, this appears to be an actual bug that prevents vector-valued inputs for the sde_fn from working. As a result, the second example in the example notebook will not work under PyMC v5.

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

In the scan function, the value of xy will always be a scalar because of how the arguments are split.

        def step(*prev_args):
            prev_y, *prev_sde_pars, rng = prev_args
            f, g = sde_fn(prev_y, *prev_sde_pars)
            mu = prev_y + dt * f
            sigma = pt.sqrt(dt) * g
            next_rng, next_y = Normal.dist(mu=mu, sigma=sigma, rng=rng).owner.outputs
            return next_y, {rng: next_rng}

That is, prev_y will always be scalar.

Reproduceable code example:

import numpy as np
import pymc as pm

N, tau, a, m, s2 = 200, 3.0, 1.05, 0.2, 1e-1
xs, ys = [0.0], [1.0]
for i in range(N):
    x, y = xs[-1], ys[-1]
    dx = tau * (x - x**3.0 / 3.0 + y)
    dy = (1.0 / tau) * (a - x)
    xs.append(x + dt * dx + np.sqrt(dt) * s2 * np.random.randn())
    ys.append(y + dt * dy + np.sqrt(dt) * s2 * np.random.randn())
xs, ys = np.array(xs), np.array(ys)
zs = m * xs + (1 - m) * ys + np.random.randn(xs.size) * 0.1

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

xys = np.c_[xs, ys]

with pm.Model() as model:
    tau_h = pm.Uniform("tau_h", lower=0.1, upper=5.0)
    a_h = pm.Uniform("a_h", lower=0.5, upper=1.5)
    m_h = pm.Uniform("m_h", lower=0.0, upper=1.0)
    xy_h = pm.EulerMaruyama("xy_h", dt=dt, sde_fn=osc_sde, sde_pars=(tau_h, a_h), shape=xys.shape, initval=xys)
    zh = pm.Normal("zh", mu=m_h * xy_h[:, 0] + (1 - m_h) * xy_h[:, 1], sigma=0.1, observed=zs)

Error message:

Cell In[53], line 2
      1 def osc_sde(xy, tau, a):
----> 2     x, y = xy[:, 0], xy[:, 1]
      3     dx = tau * (x - x**3.0 / 3.0 + y)
      4     dy = (1.0 / tau) * (a - x)

File ~/repos/pymc-examples/.pixi/envs/default/lib/python3.12/site-packages/pytensor/tensor/variable.py:502, in _tensor_py_operators.__getitem__(self, args)
    500 # Check if the number of dimensions isn't too large.
    501 if self.ndim < index_dim_count:
--> 502     raise IndexError("too many indices for array")
    504 # Convert an Ellipsis if provided into an appropriate number of
    505 # slice(None).
    506 if len(ellipses) > 1:

IndexError: too many indices for array

PyMC version information:

PyMC 5.16.2

Context for the issue:

No response

[fonnesbeck]

is this a "wontfix" or is this still supposed to work for multivariate states?

[ricardoV94]

Yes the distribution is defined as a scalar timeseries in the core case. The error you hit immediately is because init_dist has shape=size, and not say shape=(*size, 2) if the core case was multivariate vector of length 2.

A multivariate version would be a feature request right now, even if it worked before in V3.

[tbosier]

Just got this to work:
Change
x, y = xy[:, 0], xy[:, 1]

to
x, y = xy[0], xy[1]

and use

init_dist = pm.Normal.dist([xs[0], ys[0]], 1e-3, shape=(2,))

xy_h = pm.EulerMaruyama("xy_h",dt=dt, sde_fn=osc_sde, sde_pars=(tau_h, a_h), steps=N, init_dist=init_dist)```

 # If N number of variables is desired, this should work:

N_vars = N 
initial_values = np.array([x0, y0, z0, ...])
initial_std = 1e-3  

init_dist = pm.Normal.dist(mu=initial_values, sigma=initial_std, shape=(N_vars,))```