Fix logp test and move it back to test_distributions.py

Author: ricardoV94

pymc/distributions/timeseries.py

@@ -23,6 +23,7 @@
 from pymc.aesaraf import change_rv_size, floatX, intX
 from pymc.distributions import distribution, logprob, multivariate
 from pymc.distributions.continuous import Flat, Normal, get_tau_sigma
+from pymc.distributions.dist_math import check_parameters
 from pymc.distributions.shape_utils import to_tuple
 
 __all__ = [
@@ -214,7 +215,11 @@ def logp(
         stationary_series = value[..., 1:] - value[..., :-1]
         series_logp = logprob.logp(Normal.dist(mu, sigma), stationary_series)
 
-        return init_logp + series_logp.sum(axis=-1)
+        return check_parameters(
+            init_logp + series_logp.sum(axis=-1),
+            steps > 0,
+            msg="steps > 0",
+        )
 
 
 class AR1(distribution.Continuous):

pymc/tests/test_distributions.py

@@ -2619,6 +2619,21 @@ def test_interpolated_transform(self, transform):
             assert not np.isfinite(m.compile_logp()({"x": -1.0}))
             assert not np.isfinite(m.compile_logp()({"x": 11.0}))
 
+    def test_grw(self):
+        def ref_logp(value, mu, sigma, steps):
+            # Relying on fact that init will be normal by default
+            return (
+                scipy.stats.norm.logpdf(value[0], mu, sigma)
+                + scipy.stats.norm.logpdf(np.diff(value), mu, sigma).sum()
+            )
+
+        self.check_logp(
+            pm.GaussianRandomWalk,
+            Vector(R, 4),
+            {"mu": R, "sigma": Rplus, "steps": Nat},
+            ref_logp,
+            decimal=select_by_precision(float64=6, float32=1),
+        )
 
 class TestBound:
     """Tests for pm.Bound distribution"""

pymc/tests/test_distributions_timeseries.py

@@ -89,31 +89,6 @@ def test_grw_inference(self):
         np.testing.assert_allclose([mu, sigma], [recovered_mu, recovered_sigma], atol=0.2)
 
 
-class TestGRWScipy(td.TestMatchesScipy):
-    # TODO: Test LogP for different inits in its own function
-
-    # TODO: Find issue that says GRW wont take vector 
-    def test_grw_logp(self):
-        def grw_logp(value, mu, sigma):
-            # Relying on fact that init will be normal
-            # Note: This means we're not testing 
-            stationary_series = np.diff(value)
-            logp = stats.norm.logpdf(value[0], mu, sigma) + \
-            stats.norm.logpdf(stationary_series, mu, sigma).sum(),
-            return logp
-             
-        # TODO: Make base class a static method
-        # TODO: Reuse this make this static so it doesnt run all other ones
-        self.check_logp(
-            pm.GaussianRandomWalk,
-            td.Vector(td.R, 10),
-            {"mu": td.R, "sigma": td.Rplus, "steps": td.Nat},
-            grw_logp,
-            decimal=select_by_precision(float64=6, float32=1),
-            n_samples=1,
-        )
-
-
 @pytest.mark.xfail(reason="Timeseries not refactored")
 def test_AR():
     # AR1