Simplify timeseries code, delegate to Mv

Author: gBokiau

pymc3/distributions/timeseries.py

@@ -3,7 +3,6 @@
 
 from pymc3.util import get_variable_name
 from .continuous import get_tau_sd, Normal, Flat
-from .dist_math import Cholesky
 from . import multivariate
 from . import distribution
 
@@ -280,48 +279,8 @@ def _repr_latex_(self, name=None, dist=None):
                                                 get_variable_name(dt))
 
 
-class _CovSet():
-    R"""
-    Convenience class to set Covariance, Inverse Covariance and Cholesky
-    descomposition of Covariance marrices.
-    """
-    def __initCov__(self, cov=None, tau=None, chol=None, lower=True):
-        if all([val is None for val in [cov, tau, chol]]):
-            raise ValueError('One of cov, tau or chol arguments must be provided.')
-
-        self.cov = self.tau = self.chol_cov = None
-
-        cholesky = Cholesky(nofail=True, lower=True)
-        if cov is not None:
-            self.k = cov.shape[0]
-            self._cov_type = 'cov'
-            cov = tt.as_tensor_variable(cov)
-            if cov.ndim != 2:
-                raise ValueError('cov must be two dimensional.')
-            self.chol_cov = cholesky(cov)
-            self.cov = cov
-            self._n = self.cov.shape[-1]
-        elif tau is not None:
-            self.k = tau.shape[0]
-            self._cov_type = 'tau'
-            tau = tt.as_tensor_variable(tau)
-            if tau.ndim != 2:
-                raise ValueError('tau must be two dimensional.')
-            self.chol_tau = cholesky(tau)
-            self.tau = tau
-            self._n = self.tau.shape[-1]
-        else:
-            if chol is not None and not lower:
-                chol = chol.T
-            self.k = chol.shape[0]
-            self._cov_type = 'chol'
-            if chol.ndim != 2:
-                raise ValueError('chol must be two dimensional.')
-            self.chol_cov = tt.as_tensor_variable(chol)
-            self._n = self.chol_cov.shape[-1]
 
-
-class MvGaussianRandomWalk(distribution.Continuous, _CovSet):
+class MvGaussianRandomWalk(distribution.Continuous):
     R"""
     Multivariate Random Walk with Normal innovations
 
@@ -346,19 +305,18 @@ class MvGaussianRandomWalk(distribution.Continuous, _CovSet):
     def __init__(self, mu=0., cov=None, tau=None, chol=None, lower=True, init=Flat.dist(),
                  *args, **kwargs):
         super(MvGaussianRandomWalk, self).__init__(*args, **kwargs)
-        super(MvGaussianRandomWalk, self).__initCov__(cov, tau, chol, lower)
 
         self.mu = mu = tt.as_tensor_variable(mu)
         self.init = init
         self.mean = tt.as_tensor_variable(0.)
+        self.innovArgs = (self.mu, cov, tau, chol, lower)
+        self.innov = multivariate.MvNormal.dist(*self.innovArgs)
 
     def logp(self, x):
         x_im1 = x[:-1]
         x_i = x[1:]
 
-        innov_like = multivariate.MvNormal.dist(mu=x_im1 + self.mu, cov=self.cov,
-                                                tau=self.tau, chol=self.chol_cov).logp(x_i)
-        return self.init.logp(x[0]) + tt.sum(innov_like)
+        return self.init.logp(x[0]) + self.innov.logp_sum(x_i - x_im1)
 
     def _repr_latex_(self, name=None, dist=None):
         if dist is None:
@@ -371,7 +329,7 @@ def _repr_latex_(self, name=None, dist=None):
                                                 get_variable_name(cov))
 
 
-class MvStudentTRandomWalk(distribution.Continuous, _CovSet):
+class MvStudentTRandomWalk(MvGaussianRandomWalk):
     R"""
     Multivariate Random Walk with StudentT innovations
 
@@ -389,22 +347,10 @@ class MvStudentTRandomWalk(distribution.Continuous, _CovSet):
     init : distribution
         distribution for initial value (Defaults to Flat())
     """
-    def __init__(self, nu, mu=0., cov=None, tau=None, chol=None, lower=True, init=Flat.dist(),
-                 *args, **kwargs):
+    def __init__(self, nu, *args, **kwargs):
         super(MvStudentTRandomWalk, self).__init__(*args, **kwargs)
-        super(MvStudentTRandomWalk, self).__initCov__(cov, tau, chol, lower)
-        self.mu = mu = tt.as_tensor_variable(mu)
         self.nu = nu = tt.as_tensor_variable(nu)
-        self.init = init
-        self.mean = tt.as_tensor_variable(0.)
-
-    def logp(self, x):
-        x_im1 = x[:-1]
-        x_i = x[1:]
-        innov_like = multivariate.MvStudentT.dist(self.nu, mu=x_im1 + self.mu,
-                                                  cov=self.cov, tau=self.tau,
-                                                  chol=self.chol_cov).logp(x_i)
-        return self.init.logp(x[0]) + tt.sum(innov_like)
+        self.inov = multivariate.MvStudentT.dist(self.nu, *self.innovArgs)
 
     def _repr_latex_(self, name=None, dist=None):
         if dist is None: