Add test

Implement momentum and fix some shape bug.

Author: junpenglao

pymc/distributions/timeseries.py

@@ -21,7 +21,6 @@
 
 from aeppl.abstract import _get_measurable_outputs
 from aeppl.logprob import _logprob
-from aesara import scan
 from aesara.graph import FunctionGraph, rewrite_graph
 from aesara.graph.basic import Node, clone_replace
 from aesara.raise_op import Assert
@@ -230,7 +229,7 @@ def random_walk_moment(op, rv, init_dist, innovation_dist, steps):
 
 @_logprob.register(RandomWalkRV)
 def random_walk_logp(op, values, *inputs, **kwargs):
-    # ALthough Aeppl can derive the logprob of random walks, it does not collapse
+    # Although Aeppl can derive the logprob of random walks, it does not collapse
     # what PyMC considers the core dimension of steps. We do it manually here.
     (value,) = values
     # Recreate RV and obtain inner graph
@@ -681,15 +680,15 @@ def dist(cls, omega, alpha_1, beta_1, initial_vol, *, steps=None, **kwargs):
 
         return super().dist([omega, alpha_1, beta_1, initial_vol, init_dist, steps], **kwargs)
 
-
     @classmethod
     def rv_op(cls, omega, alpha_1, beta_1, initial_vol, init_dist, steps, size=None):
         if size is not None:
             batch_size = size
         else:
             # In this case the size of the init_dist depends on the parameters shape
-            batch_size = at.broadcast_shape(omega, alpha_1, beta_1, init_dist)
+            batch_size = at.broadcast_shape(omega, alpha_1, beta_1, initial_vol)
         init_dist = change_dist_size(init_dist, batch_size)
+        # initial_vol = initial_vol * at.ones(batch_size)
 
         # Create OpFromGraph representing random draws form AR process
         # Variables with underscore suffix are dummy inputs into the OpFromGraph
@@ -712,13 +711,16 @@ def step(*args):
 
         (y_t, _), innov_updates_ = aesara.scan(
             fn=step,
-            outputs_info=[init_, initial_vol_],
+            outputs_info=[init_, initial_vol_ * at.ones(batch_size)],
             non_sequences=[omega_, alpha_1_, beta_1_, noise_rng],
             n_steps=steps_,
             strict=True,
         )
         (noise_next_rng,) = tuple(innov_updates_.values())
-        garch11_ = at.concatenate([at.atleast_1d(init_), y_t.T], axis=-1)
+
+        garch11_ = at.concatenate([init_[None, ...], y_t], axis=0).dimshuffle(
+            tuple(range(1, y_t.ndim)) + (0,)
+        )
 
         garch11_op = GARCH11RV(
             inputs=[omega_, alpha_1_, beta_1_, initial_vol_, init_, steps_],
@@ -748,22 +750,30 @@ def garch11_logp(
     op, values, omega, alpha_1, beta_1, initial_vol, init_dist, steps, noise_rng, **kwargs
 ):
     (value,) = values
+    value_dimswapped = value.dimshuffle((value.ndim - 1,) + tuple(range(0, value.ndim - 1)))
+    initial_vol = initial_vol * at.ones_like(value_dimswapped[0])
 
     def volatility_update(x, vol, w, a, b):
         return at.sqrt(w + a * at.square(x) + b * at.square(vol))
 
-    vol, _ = scan(
+    vol, _ = aesara.scan(
         fn=volatility_update,
-        sequences=[value[:-1]],
+        sequences=[value_dimswapped[:-1]],
         outputs_info=[initial_vol],
         non_sequences=[omega, alpha_1, beta_1],
     )
     sigma_t = at.concatenate([[initial_vol], vol])
     # Compute and collapse logp across time dimension
-    innov_logp = at.sum(logp(Normal.dist(0, sigma_t), value), axis=-1)
+    innov_logp = at.sum(logp(Normal.dist(0, sigma_t), value_dimswapped), axis=-1)
     return innov_logp
 
 
+@_moment.register(GARCH11RV)
+def garch11_moment(op, rv, omega, alpha_1, beta_1, initial_vol, init_dist, steps, noise_rng):
+    # GARCH(1,1) mean is zero
+    return at.zeros_like(rv)
+
+
 class EulerMaruyama(distribution.Continuous):
     r"""
     Stochastic differential equation discretized with the Euler-Maruyama method.

pymc/tests/distributions/test_timeseries.py

@@ -481,55 +481,120 @@ def test_change_dist_size(self):
         assert new_dist.eval().shape == (4, 3, 10)
 
 
-def test_GARCH11():
-    # test data ~ N(0, 1)
-    data = np.array(
+class TestGARCH11:
+    def test_logp(self):
+        # test data ~ N(0, 1)
+        data = np.array(
+            [
+                -1.35078362,
+                -0.81254164,
+                0.28918551,
+                -2.87043544,
+                -0.94353337,
+                0.83660719,
+                -0.23336562,
+                -0.58586298,
+                -1.36856736,
+                -1.60832975,
+                -1.31403141,
+                0.05446936,
+                -0.97213128,
+                -0.18928725,
+                1.62011258,
+                -0.95978616,
+                -2.06536047,
+                0.6556103,
+                -0.27816645,
+                -1.26413397,
+            ]
+        )
+        omega = 0.6
+        alpha_1 = 0.4
+        beta_1 = 0.5
+        initial_vol = np.float64(0.9)
+        vol = np.empty_like(data)
+        vol[0] = initial_vol
+        for i in range(len(data) - 1):
+            vol[i + 1] = np.sqrt(omega + beta_1 * vol[i] ** 2 + alpha_1 * data[i] ** 2)
+
+        with Model() as t:
+            y = GARCH11(
+                "y",
+                omega=omega,
+                alpha_1=alpha_1,
+                beta_1=beta_1,
+                initial_vol=initial_vol,
+                shape=data.shape,
+            )
+            z = Normal("z", mu=0, sigma=vol, shape=data.shape)
+        garch_like = t.compile_logp(y)({"y": data})
+        reg_like = t.compile_logp(z)({"z": data})
+        decimal = select_by_precision(float64=7, float32=4)
+        np.testing.assert_allclose(garch_like, reg_like, 10 ** (-decimal))
+
+    @pytest.mark.parametrize(
+        "arg_name",
+        ["omega", "alpha_1", "beta_1", "initial_vol"],
+    )
+    def test_batched_size(self, arg_name):
+        steps, batch_size = 100, 5
+        param_val = np.square(np.random.randn(batch_size))
+        init_kwargs = dict(
+            omega=1.25,
+            alpha_1=0.5,
+            beta_1=0.45,
+            initial_vol=2.5,
+        )
+        kwargs0 = init_kwargs.copy()
+        kwargs0[arg_name] = init_kwargs[arg_name] * param_val
+        with Model() as t0:
+            y = GARCH11("y", shape=(batch_size, steps), **kwargs0)
+
+        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]))
+
+        kwargs1 = init_kwargs.copy()
+        with Model() as t1:
+            for i in range(batch_size):
+                kwargs1[arg_name] = init_kwargs[arg_name] * param_val[i]
+                GARCH11(f"y_{i}", shape=steps, **kwargs1)
+
+        np.testing.assert_allclose(
+            t0.compile_logp()(t0.initial_point()),
+            t1.compile_logp()(t1.initial_point()),
+        )
+
+    @pytest.mark.parametrize(
+        "size, expected",
         [
-            -1.35078362,
-            -0.81254164,
-            0.28918551,
-            -2.87043544,
-            -0.94353337,
-            0.83660719,
-            -0.23336562,
-            -0.58586298,
-            -1.36856736,
-            -1.60832975,
-            -1.31403141,
-            0.05446936,
-            -0.97213128,
-            -0.18928725,
-            1.62011258,
-            -0.95978616,
-            -2.06536047,
-            0.6556103,
-            -0.27816645,
-            -1.26413397,
-        ]
+            (None, np.zeros((2, 8))),
+            ((5, 2), np.zeros((5, 2, 8))),
+        ],
     )
-    omega = 0.6
-    alpha_1 = 0.4
-    beta_1 = 0.5
-    initial_vol = np.float64(0.9)
-    vol = np.empty_like(data)
-    vol[0] = initial_vol
-    for i in range(len(data) - 1):
-        vol[i + 1] = np.sqrt(omega + beta_1 * vol[i] ** 2 + alpha_1 * data[i] ** 2)
-
-    with Model() as t:
-        y = GARCH11(
-            "y",
-            omega=omega,
-            alpha_1=alpha_1,
-            beta_1=beta_1,
-            initial_vol=initial_vol,
-            shape=data.shape,
+    def test_moment(self, size, expected):
+        with Model() as model:
+            GARCH11(
+                "x",
+                omega=1.25,
+                alpha_1=0.5,
+                beta_1=0.45,
+                initial_vol=np.ones(2),
+                steps=7,
+                size=size,
+            )
+        assert_moment_is_expected(model, expected, check_finite_logp=False)
+
+    def test_change_dist_size(self):
+        base_dist = pm.GARCH11.dist(
+            omega=1.25, alpha_1=0.5, beta_1=0.45, initial_vol=1.0, shape=(3, 10)
         )
-        z = Normal("z", mu=0, sigma=vol, shape=data.shape)
-    garch_like = t.compile_logp(y)({"y": data})
-    reg_like = t.compile_logp(z)({"z": data})
-    decimal = select_by_precision(float64=7, float32=4)
-    np.testing.assert_allclose(garch_like, reg_like, 10 ** (-decimal))
+
+        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 _gen_sde_path(sde, pars, dt, n, x0):