Skip to content

Mv cleanup #3008

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Jun 11, 2018
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 7 additions & 61 deletions pymc3/distributions/timeseries.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand All @@ -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):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This only works if x has rank 1, right? Maybe we should intercept the shape in __init__ (one of the kwargs) and throw an error if that doesn't have length 1, just to give people a better error message.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Er, wait, I would hope not, wouldn't rank 1 defeat the purpose of Mv?
I'll run the demo to get a better grasp. More tests for the timeseries would be nice wrt shapes #2859 (comment)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I checked to be sure: this only accepts x and cov/chol with ndim 2, and correctly controls that the shapes are compatible.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Related observations:

  • To be more precise, it controls/catches any glitch between shape and cov upon sampling, not at declaration time. That is also the behaviour of MvN/MvT. While that could be improved, I think it should then happen in those classes.
  • Another question entirely (for later, perhaps) is wether we want to intercept any such errors here, as to make them more contextual for the timeseries classes.

However this is all probably not what your concern about the rank was?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I should have said rank 2, not 1. I didn't actually think about the fact that mvnormal doesn't support higher rank input either, but I still think we should check this in this class. If at some point in the future we extend mvnormal, people wouldn't get an error then, but I'm not sure the results are what we want. (eg which dimension is the one with the correlations? In mvnormal it is the last one....)
Here is the code in mvnormal that I have in mind:
https://github.com/pymc-devs/pymc3/blob/master/pymc3/distributions/multivariate.py#L37

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right, gotcha. If my understanding is correct, in the case of a random walk, ndim=1 doesn't make sense either, right? So shape can only be two-dimensional for these distributions.

So, as it stands, I would simply add this control in __init__:

    if len(self.shape) != 2:
        raise ValueError("Only 2 dimensions are allowed.")

Perhaps for a separate PR — when given 3 dims, we could (I think?) diff the obs separately along the time axis, concatenate/stack the thus obtained innovations to two dims, and pass that on to Mv.

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:
Expand All @@ -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

Expand All @@ -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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice simplification!

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.innov = multivariate.MvStudentT.dist(self.nu, *self.innovArgs)

def _repr_latex_(self, name=None, dist=None):
if dist is None:
Expand Down
40 changes: 18 additions & 22 deletions pymc3/gp/gp.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,13 +8,14 @@
from pymc3.gp.cov import Covariance, Constant
from pymc3.gp.mean import Zero
from pymc3.gp.util import (conditioned_vars,
infer_shape, stabilize, cholesky, solve_lower, solve_upper)
infer_shape, stabilize, solve_lower, solve_upper)
from pymc3.distributions import draw_values
from pymc3.distributions.dist_math import eigh
from ..math import cartesian, kron_dot, kron_diag

__all__ = ['Latent', 'Marginal', 'TP', 'MarginalSparse', 'MarginalKron']

cholesky = tt.slinalg.cholesky

class Base(object):
R"""
Expand Down Expand Up @@ -107,13 +108,13 @@ def __init__(self, mean_func=Zero(), cov_func=Constant(0.0)):

def _build_prior(self, name, X, reparameterize=True, **kwargs):
mu = self.mean_func(X)
chol = cholesky(stabilize(self.cov_func(X)))
cov = stabilize(self.cov_func(X))
shape = infer_shape(X, kwargs.pop("shape", None))
if reparameterize:
v = pm.Normal(name + "_rotated_", mu=0.0, sd=1.0, shape=shape, **kwargs)
f = pm.Deterministic(name, mu + tt.dot(chol, v))
f = pm.Deterministic(name, mu + cholesky(cov).dot(v))
else:
f = pm.MvNormal(name, mu=mu, chol=chol, shape=shape, **kwargs)
f = pm.MvNormal(name, mu=mu, cov=cov, shape=shape, **kwargs)
return f

def prior(self, name, X, reparameterize=True, **kwargs):
Expand Down Expand Up @@ -203,9 +204,8 @@ def conditional(self, name, Xnew, given=None, **kwargs):
"""
givens = self._get_given_vals(given)
mu, cov = self._build_conditional(Xnew, *givens)
chol = cholesky(stabilize(cov))
shape = infer_shape(Xnew, kwargs.pop("shape", None))
return pm.MvNormal(name, mu=mu, chol=chol, shape=shape, **kwargs)
return pm.MvNormal(name, mu=mu, cov=cov, shape=shape, **kwargs)


@conditioned_vars(["X", "f", "nu"])
Expand Down Expand Up @@ -249,14 +249,14 @@ def __add__(self, other):

def _build_prior(self, name, X, reparameterize=True, **kwargs):
mu = self.mean_func(X)
chol = cholesky(stabilize(self.cov_func(X)))
cov = stabilize(self.cov_func(X))
shape = infer_shape(X, kwargs.pop("shape", None))
if reparameterize:
chi2 = pm.ChiSquared("chi2_", self.nu)
v = pm.Normal(name + "_rotated_", mu=0.0, sd=1.0, shape=shape, **kwargs)
f = pm.Deterministic(name, (tt.sqrt(self.nu) / chi2) * (mu + tt.dot(chol, v)))
f = pm.Deterministic(name, (tt.sqrt(self.nu) / chi2) * (mu + cholesky(cov).dot(v)))
else:
f = pm.MvStudentT(name, nu=self.nu, mu=mu, chol=chol, shape=shape, **kwargs)
f = pm.MvStudentT(name, nu=self.nu, mu=mu, cov=cov, shape=shape, **kwargs)
return f

def prior(self, name, X, reparameterize=True, **kwargs):
Expand Down Expand Up @@ -321,10 +321,9 @@ def conditional(self, name, Xnew, **kwargs):

X = self.X
f = self.f
nu2, mu, covT = self._build_conditional(Xnew, X, f)
chol = cholesky(stabilize(covT))
nu2, mu, cov = self._build_conditional(Xnew, X, f)
shape = infer_shape(Xnew, kwargs.pop("shape", None))
return pm.MvStudentT(name, nu=nu2, mu=mu, chol=chol, shape=shape, **kwargs)
return pm.MvStudentT(name, nu=nu2, mu=mu, cov=cov, shape=shape, **kwargs)


@conditioned_vars(["X", "y", "noise"])
Expand Down Expand Up @@ -418,15 +417,15 @@ def marginal_likelihood(self, name, X, y, noise, is_observed=True, **kwargs):
if not isinstance(noise, Covariance):
noise = pm.gp.cov.WhiteNoise(noise)
mu, cov = self._build_marginal_likelihood(X, noise)
chol = cholesky(stabilize(cov))
cov = stabilize(cov)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure the stabilisation is needed here since noise is already being added to the diagonal in _build_marginal_likelihood.

self.X = X
self.y = y
self.noise = noise
if is_observed:
return pm.MvNormal(name, mu=mu, chol=chol, observed=y, **kwargs)
return pm.MvNormal(name, mu=mu, cov=cov, observed=y, **kwargs)
else:
shape = infer_shape(X, kwargs.pop("shape", None))
return pm.MvNormal(name, mu=mu, chol=chol, shape=shape, **kwargs)
return pm.MvNormal(name, mu=mu, cov=cov, shape=shape, **kwargs)

def _get_given_vals(self, given):
if given is None:
Expand Down Expand Up @@ -504,9 +503,8 @@ def conditional(self, name, Xnew, pred_noise=False, given=None, **kwargs):

givens = self._get_given_vals(given)
mu, cov = self._build_conditional(Xnew, pred_noise, False, *givens)
chol = cholesky(cov)
shape = infer_shape(Xnew, kwargs.pop("shape", None))
return pm.MvNormal(name, mu=mu, chol=chol, shape=shape, **kwargs)
return pm.MvNormal(name, mu=mu, cov=cov, shape=shape, **kwargs)

def predict(self, Xnew, point=None, diag=False, pred_noise=False, given=None):
R"""
Expand Down Expand Up @@ -797,9 +795,8 @@ def conditional(self, name, Xnew, pred_noise=False, given=None, **kwargs):

givens = self._get_given_vals(given)
mu, cov = self._build_conditional(Xnew, pred_noise, False, *givens)
chol = cholesky(cov)
shape = infer_shape(Xnew, kwargs.pop("shape", None))
return pm.MvNormal(name, mu=mu, chol=chol, shape=shape, **kwargs)
return pm.MvNormal(name, mu=mu, cov=cov, shape=shape, **kwargs)


@conditioned_vars(["Xs", "y", "sigma"])
Expand Down Expand Up @@ -959,7 +956,7 @@ def _build_conditional(self, Xnew, pred_noise, diag):
cov = Km - Asq
if pred_noise:
cov += sigma * np.eye(cov.shape)
return mu, cov
return mu, stabilize(cov)
Copy link
Contributor Author

@gBokiau gBokiau Jun 5, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ditto, actually almost certain this is redundant. (Moved from line 999, BTW)
Maybe not redundant when pred_noise == False.


def conditional(self, name, Xnew, pred_noise=False, **kwargs):
"""
Expand Down Expand Up @@ -996,9 +993,8 @@ def conditional(self, name, Xnew, pred_noise=False, **kwargs):
constructor.
"""
mu, cov = self._build_conditional(Xnew, pred_noise, False)
chol = cholesky(stabilize(cov))
shape = infer_shape(Xnew, kwargs.pop("shape", None))
return pm.MvNormal(name, mu=mu, chol=chol, shape=shape, **kwargs)
return pm.MvNormal(name, mu=mu, cov=cov, shape=shape, **kwargs)

def predict(self, Xnew, point=None, diag=False, pred_noise=False):
R"""
Expand Down