Somewhat fixed MvNormal implementation

mvnormal-extension.py

from scipy import stats
from theano.tensor.nlinalg import det, matrix_inverse, trace, eigh

from pymc3 import transforms
from pymc3.distributions.distribution import Continuous, Discrete, draw_values, generate_samples
from pymc3.distributions.special import gammaln, multigammaln
from pymc3.distributions.dist_math import bound, logpow, factln

class MvNormal(Continuous):
    r"""
    This MvNormal can handle tensor arguments (to some degree).  Also, the sampling routine uses the correct covariance,
    but it's costly.
    """
    def __init__(self, mu, tau, *args, **kwargs):
        super(MvNormal, self).__init__(*args, **kwargs)
        self.mean = self.median = self.mode = self.mu = mu
        self.tau = tau

    def random(self, point=None, size=None):
        mu, tau = draw_values([self.mu, self.tau], point=point)

        def _random(mean, tau, size=None):
            supp_dim = mean.shape[-1]
            mus_collapsed = mean.reshape((-1, supp_dim))
            taus_collapsed = tau.reshape((-1, supp_dim, supp_dim))
            # FIXME: do something smarter about tau/cov
            covs_collapsed = np.apply_over_axes(lambda x,y: np.linalg.inv(x), taus_collapsed, 0)

            from functools import partial
            mvrvs = partial(stats.multivariate_normal.rvs, size=1)
            res = map(mvrvs, mus_collapsed, covs_collapsed)

            # FIXME: this is a hack; the PyMC sampling framework
            # will incorrectly set `size == Distribution.shape` when a single
            # sample is requested, implying that we want
            # `Distribution.shape`-many samples of a
            # `Distribution.shape` sized object: too many!  That's why
            # we're ignoring `size` right now and only ever asking
            # for a single sample.
            return np.asarray(res).reshape(mean.shape)

        samples = generate_samples(_random,
                                   mean=mu, tau=tau,
                                   dist_shape=self.shape,
                                   broadcast_shape=mu.shape,
                                   size=size)
        return samples

    def logp(self, value):
        mu = T.as_tensor_variable(self.mu)
        tau = T.as_tensor_variable(self.tau)

        reps_shape_T = tau.shape[:-2]
        reps_shape_prod = T.prod(reps_shape_T, keepdims=True)
        dist_shape_T = mu.shape[-1:]

        # collapse reps dimensions
        flat_supp_shape = T.concatenate((reps_shape_prod, dist_shape_T))
        mus_collapsed = mu.reshape(flat_supp_shape, ndim=2)
        taus_collapsed = tau.reshape(T.concatenate((reps_shape_prod,
            dist_shape_T, dist_shape_T)), ndim=3)

        # force value to conform to reps_shape
        value_reshape = T.ones_like(mu) * value
        values_collapsed = value_reshape.reshape(flat_supp_shape, ndim=2)

        def single_logl(_mu, _tau, _value, k):
            delta = _value - _mu
            result = k * T.log(2 * np.pi) + T.log(det(_tau))
            result += T.square(delta.dot(_tau)).sum(axis=-1)
            return -result/2

        from theano import scan
        res, _ = scan(fn=single_logl
                , sequences=[mus_collapsed, taus_collapsed, values_collapsed]
                , non_sequences=[dist_shape_T]
                , strict=True
                )
        return res.sum()