Add change_dist_size

Add tests.

Author: junpenglao

pymc/distributions/timeseries.py

@@ -27,7 +27,6 @@
 from aesara.tensor.random.op import RandomVariable
 
 from pymc.aesaraf import constant_fold, floatX, intX
-from pymc.distributions import distribution
 from pymc.distributions.continuous import Normal, get_tau_sigma
 from pymc.distributions.distribution import (
     Distribution,
@@ -461,7 +460,7 @@ class AR(Distribution):
         process.
     init_dist : unnamed distribution, optional
         Scalar or vector distribution for initial values. Unnamed refers to distributions
-         created with the ``.dist()`` API. Distributions should have shape (*shape[:-1], ar_order).
+        created with the ``.dist()`` API. Distributions should have shape (*shape[:-1], ar_order).
         If not, it will be automatically resized. Defaults to pm.Normal.dist(0, 100, shape=...).
 
         .. warning:: init_dist will be cloned, rendering it independent of the one passed as input.
@@ -914,7 +913,7 @@ class EulerMaruyama(Distribution):
         parameters of the SDE, passed as ``*args`` to ``sde_fn``
     init_dist : unnamed distribution, optional
         Scalar or vector distribution for initial values. Unnamed refers to distributions
-         created with the ``.dist()`` API. Distributions should have shape (*shape[:-1], ar_order).
+        created with the ``.dist()`` API. Distributions should have shape (*shape[:-1]).
         If not, it will be automatically resized. Defaults to pm.Normal.dist(0, 100, shape=...).
 
         .. warning:: init_dist will be cloned, rendering it independent of the one passed as input.
@@ -953,7 +952,7 @@ def dist(cls, dt, sde_fn, sde_pars, *, init_dist=None, steps=None, **kwargs):
                     f"Init dist must be a distribution created via the `.dist()` API, "
                     f"got {type(init_dist)}"
                 )
-                check_dist_not_registered(init_dist)
+            check_dist_not_registered(init_dist)
             if init_dist.owner.op.ndim_supp > 1:
                 raise ValueError(
                     "Init distribution must have a scalar or vector support dimension, ",
@@ -970,17 +969,15 @@ def dist(cls, dt, sde_fn, sde_pars, *, init_dist=None, steps=None, **kwargs):
         # Tell Aeppl to ignore init_dist, as it will be accounted for in the logp term
         init_dist = ignore_logprob(init_dist)
 
-        return super().dist([dt, sde_fn, sde_pars, init_dist, steps], **kwargs)
+        return super().dist([init_dist, steps, sde_pars, dt, sde_fn], **kwargs)
 
     @classmethod
-    def rv_op(cls, dt, sde_fn, sde_pars, init_dist, steps, size=None):
-        # Init dist should have shape (*size, ar_order)
+    def rv_op(cls, init_dist, steps, sde_pars, dt, sde_fn, size=None):
+        # Init dist should have shape (*size,)
         if size is not None:
             batch_size = size
         else:
-            # In this case the size of the init_dist depends on the parameters shape
-            # The last dimension of rho and init_dist does not matter
-            batch_size = at.broadcast_shape(*sde_pars, at.atleast_1d(init_dist)[..., 0])
+            batch_size = at.broadcast_shape(*sde_pars, init_dist)
         init_dist = change_dist_size(init_dist, batch_size)
 
         # Create OpFromGraph representing random draws form AR process
@@ -1024,6 +1021,24 @@ def step(*prev_args):
         return eulermaruyama
 
 
+@_change_dist_size.register(EulerMaruyamaRV)
+def change_eulermaruyama_size(op, dist, new_size, expand=False):
+
+    if expand:
+        old_size = dist.shape[:-1]
+        new_size = tuple(new_size) + tuple(old_size)
+
+    init_dist, steps, *sde_pars, _ = dist.owner.inputs
+    return EulerMaruyama.rv_op(
+        init_dist,
+        steps,
+        sde_pars,
+        dt=op.dt,
+        sde_fn=op.sde_fn,
+        size=new_size,
+    )
+
+
 @_logprob.register(EulerMaruyamaRV)
 def eulermaruyama_logp(op, values, init_dist, steps, *sde_pars_noise_arg, **kwargs):
     (x,) = values

pymc/tests/distributions/test_timeseries.py

@@ -831,22 +831,82 @@ def test_change_dist_size(self):
 
 
 class TestEulerMaruyama:
+    @pytest.mark.parametrize("batched_param", [1, 2])
+    @pytest.mark.parametrize("explicit_shape", (True, False))
+    def test_batched_size(self, explicit_shape, batched_param):
+        steps, batch_size = 100, 5
+        param_val = np.square(np.random.randn(batch_size))
+        if explicit_shape:
+            kwargs = {"shape": (batch_size, steps)}
+        else:
+            kwargs = {"steps": steps - 1}
+
+        def sde_fn(x, k, d, s):
+            return (k - d * x, s)
+
+        sde_pars = [1.0, 2.0, 0.1]
+        sde_pars[batched_param] = sde_pars[batched_param] * param_val
+        with Model() as t0:
+            y = EulerMaruyama("y", dt=0.02, sde_fn=sde_fn, sde_pars=sde_pars, **kwargs)
+
+        y_eval = draw(y, draws=2)
+        assert y_eval[0].shape == (batch_size, steps)
+        assert not np.any(np.isclose(y_eval[0], y_eval[1]))
+
+        if explicit_shape:
+            kwargs["shape"] = steps
+        with Model() as t1:
+            for i in range(batch_size):
+                sde_pars_slice = sde_pars.copy()
+                sde_pars_slice[batched_param] = sde_pars[batched_param][i]
+                EulerMaruyama(f"y_{i}", dt=0.02, sde_fn=sde_fn, sde_pars=sde_pars, **kwargs)
+
+        np.testing.assert_allclose(
+            t0.compile_logp()(t0.initial_point()),
+            t1.compile_logp()(t1.initial_point()),
+        )
+
+    def test_change_dist_size1(self):
+        def sde1(x, k, d, s):
+            return (k - d * x, s)
+
+        base_dist = EulerMaruyama.dist(dt=0.01, sde_fn=sde1, sde_pars=(1, 2, 0.1), shape=(5, 10))
+
+        new_dist = change_dist_size(base_dist, (4,))
+        assert new_dist.eval().shape == (4, 10)
+
+        new_dist = change_dist_size(base_dist, (4,), expand=True)
+        assert new_dist.eval().shape == (4, 5, 10)
 
-    def _gen_sde_path(self, sde, pars, dt, n, x0):
-        xs = [x0]
-        wt = np.random.normal(size=(n,) if isinstance(x0, float) else (n, x0.size))
-        for i in range(n):
-            f, g = sde(xs[-1], *pars)
-            xs.append(xs[-1] + f * dt + np.sqrt(dt) * g * wt[i])
-        return np.array(xs)
+    def test_change_dist_size2(self):
+        def sde2(p, s):
+            N = 500.0
+            return s * p * (1 - p) / (1 + s * p), pm.math.sqrt(p * (1 - p) / N)
 
-    def test_linear(self):
+        base_dist = EulerMaruyama.dist(dt=0.01, sde_fn=sde2, sde_pars=(0.1,), shape=(3, 10))
+
+        new_dist = change_dist_size(base_dist, (4,))
+        assert new_dist.eval().shape == (4, 10)
+
+        new_dist = change_dist_size(base_dist, (4,), expand=True)
+        assert new_dist.eval().shape == (4, 3, 10)
+
+    def test_linear_model(self):
         lam = -0.78
         sig2 = 5e-3
         N = 300
         dt = 1e-1
+
+        def _gen_sde_path(sde, pars, dt, n, x0):
+            xs = [x0]
+            wt = np.random.normal(size=(n,) if isinstance(x0, float) else (n, x0.size))
+            for i in range(n):
+                f, g = sde(xs[-1], *pars)
+                xs.append(xs[-1] + f * dt + np.sqrt(dt) * g * wt[i])
+            return np.array(xs)
+
         sde = lambda x, lam: (lam * x, sig2)
-        x = floatX(self._gen_sde_path(sde, (lam,), dt, N, 5.0))
+        x = floatX(_gen_sde_path(sde, (lam,), dt, N, 5.0))
         z = x + np.random.randn(x.size) * sig2
         # build model
         with Model() as model: