Improve AR model (#3070)

As discussed in https://discourse.pymc.io/t/vectorized-autoregressive-model/1449/
improve AR implementation so that tensor parameter is allowed.

Author: Junpeng Lao

pymc3/distributions/timeseries.py

@@ -78,7 +78,7 @@ class AR(distribution.Continuous):
     Parameters
     ----------
     rho : tensor
-        Vector of autoregressive coefficients.
+        Tensor of autoregressive coefficients. The first dimension is the p lag.
     sd : float
         Standard deviation of innovation (sd > 0). (only required if tau is not specified)
     tau : float
@@ -100,29 +100,30 @@ def __init__(self, rho, sd=None, tau=None,
 
         self.mean = tt.as_tensor_variable(0.)
 
-        rho = tt.as_tensor_variable(rho, ndim=1)
+        if isinstance(rho, list):
+            p = len(rho)
+        else:
+            try:
+                p = rho.shape.tag.test_value[0]
+            except AttributeError:
+                p = rho.shape[0]
+
         if constant:
-            self.p = rho.shape[0] - 1
+            self.p = p - 1
         else:
-            self.p = rho.shape[0]
+            self.p = p
 
         self.constant = constant
-        self.rho = rho
+        self.rho = rho = tt.as_tensor_variable(rho)
         self.init = init
 
     def logp(self, value):
-
-        y = value[self.p:]
-        results, _ = scan(lambda l, obs, p: obs[p - l:-l],
-                          outputs_info=None, sequences=[tt.arange(1, self.p + 1)],
-                          non_sequences=[value, self.p])
-        x = tt.stack(results)
-
         if self.constant:
-            y = y - self.rho[0]
-            eps = y - self.rho[1:].dot(x)
+            x = tt.add(*[self.rho[i + 1] * value[self.p - (i + 1):-(i + 1)] for i in range(self.p)])
+            eps = value[self.p:] - self.rho[0] - x
         else:
-            eps = y - self.rho.dot(x)
+            x = tt.add(*[self.rho[i] * value[self.p - (i + 1):-(i + 1)] for i in range(self.p)])
+            eps = value[self.p:] - x
 
         innov_like = Normal.dist(mu=0.0, tau=self.tau).logp(eps)
         init_like = self.init.logp(value[:self.p])

pymc3/tests/test_distributions_timeseries.py

@@ -37,6 +37,29 @@ def test_AR():
     np.testing.assert_allclose(ar_like, reg_like)
 
 
+def test_AR_nd():
+    # AR2 multidimensional
+    p, T, n = 3, 100, 5
+    beta_tp = np.random.randn(p, n)
+    y_tp = np.random.randn(T, n)
+    with Model() as t0:
+        beta = Normal('beta', 0., 1.,
+                      shape=(p, n),
+                      testval=beta_tp)
+        AR('y', beta, sd=1.0,
+           shape=(T, n), testval=y_tp)
+
+    with Model() as t1:
+        beta = Normal('beta', 0., 1.,
+                      shape=(p, n),
+                      testval=beta_tp)
+        for i in range(n):
+            AR('y_%d' % i, beta[:, i], sd=1.0,
+               shape=T, testval=y_tp[:, i])
+
+    np.testing.assert_allclose(t0.logp(t0.test_point),
+                               t1.logp(t1.test_point))
+
 
 def test_GARCH11():
     # test data ~ N(0, 1)