Added base class for variational methods

pymc3/distributions/discrete.py

@@ -6,7 +6,7 @@
 from scipy import stats
 
 from .dist_math import bound, factln, binomln, betaln, logpow
-from .distribution import Discrete, draw_values, generate_samples
+from .distribution import Discrete, draw_values, generate_samples, reshape_sampled
 
 __all__ = ['Binomial',  'BetaBinomial',  'Bernoulli',  'Poisson',
            'NegativeBinomial', 'ConstantDist', 'Constant', 'ZeroInflatedPoisson',
@@ -250,7 +250,7 @@ def random(self, point=None, size=None, repeat=None):
                              dist_shape=self.shape,
                              size=size)
         g[g == 0] = np.finfo(float).eps  # Just in case
-        return stats.poisson.rvs(g)
+        return reshape_sampled(stats.poisson.rvs(g), size, self.shape)
 
     def logp(self, value):
         mu = self.mu
@@ -441,9 +441,11 @@ def logp(self, value):
         c = self.c
         return bound(0, tt.eq(value, c))
 
+
 def ConstantDist(*args, **kwargs):
+    import warnings
     warnings.warn("ConstantDist has been deprecated. In future, use Constant instead.",
-                DeprecationWarning)
+                  DeprecationWarning)
     return Constant(*args, **kwargs)
 
 
@@ -489,7 +491,8 @@ def random(self, point=None, size=None, repeat=None):
         g = generate_samples(stats.poisson.rvs, theta,
                              dist_shape=self.shape,
                              size=size)
-        return g * (np.random.random(np.squeeze(g.shape)) < psi)
+        sampled = g * (np.random.random(np.squeeze(g.shape)) < psi)
+        return reshape_sampled(sampled, size, self.shape)
 
     def logp(self, value):
         return tt.switch(value > 0,
@@ -543,7 +546,8 @@ def random(self, point=None, size=None, repeat=None):
                              dist_shape=self.shape,
                              size=size)
         g[g == 0] = np.finfo(float).eps  # Just in case
-        return stats.poisson.rvs(g) * (np.random.random(np.squeeze(g.shape)) < psi)
+        sampled = stats.poisson.rvs(g) * (np.random.random(np.squeeze(g.shape)) < psi)
+        return reshape_sampled(sampled, size, self.shape)
 
     def logp(self, value):
         return tt.switch(value > 0,

pymc3/distributions/dist_math.py

@@ -10,6 +10,8 @@
 
 from .special import gammaln, multigammaln
 
+c = - 0.5 * np.log(2 * np.pi)
+
 
 def bound(logp, *conditions):
     """
@@ -77,3 +79,39 @@ def i1(x):
                      x**9 / 1474560 + x**11 / 176947200 + x**13 / 29727129600,
                      np.e**x / (2 * np.pi * x)**0.5 * (1 - 3 / (8 * x) + 15 / (128 * x**2) + 315 / (3072 * x**3)
                                                        + 14175 / (98304 * x**4)))
+
+
+def sd2rho(sd):
+    """sd -> rho
+    theano converter
+    mu + sd*e = mu + log(1+exp(rho))*e"""
+    return tt.log(tt.exp(sd) - 1)
+
+
+def rho2sd(rho):
+    """rho -> sd
+    theano converter
+    mu + sd*e = mu + log(1+exp(rho))*e"""
+    return tt.log1p(tt.exp(rho))
+
+
+def kl_divergence_normal_pair(mu1, mu2, sd1, sd2):
+    elemwise_kl = (tt.log(sd2/sd1) +
+                   (sd2**2 + (mu1-mu2)**2)/(2.*sd2**2) -
+                   0.5)
+    return tt.sum(elemwise_kl)
+
+
+def kl_divergence_normal_pair3(mu1, mu2, rho1, rho2):
+    sd1, sd2 = rho2sd(rho1), rho2sd(rho2)
+    return kl_divergence_normal_pair(mu1, mu2, sd1, sd2)
+
+
+def log_normal(x, mean, std, eps=0.0):
+    std += eps
+    return c - tt.log(tt.abs_(std)) - (x - mean) ** 2 / (2 * std ** 2)
+
+
+def log_normal3(x, mean, rho, eps=0.0):
+    std = rho2sd(rho)
+    return log_normal(x, mean, std, eps)

pymc3/distributions/distribution.py

@@ -10,9 +10,11 @@
 __all__ = ['DensityDist', 'Distribution', 'Continuous',
            'Discrete', 'NoDistribution', 'TensorType', 'draw_values']
 
+
 class _Unpickling(object):
     pass
 
+
 class Distribution(object):
     """Statistical distribution"""
     def __new__(cls, name, *args, **kwargs):
@@ -129,12 +131,10 @@ def __init__(self, logp, shape=(), dtype='float64', testval=0, *args, **kwargs):
 
 
 class MultivariateContinuous(Continuous):
-
     pass
 
 
 class MultivariateDiscrete(Discrete):
-
     pass
 
 
@@ -265,6 +265,22 @@ def broadcast_shapes(*args):
     return tuple(x)
 
 
+def infer_shape(shape):
+    try:
+        shape = tuple(shape or ())
+    except TypeError:  # If size is an int
+        shape = tuple((shape,))
+    except ValueError:  # If size is np.array
+        shape = tuple(shape)
+    return shape
+
+
+def reshape_sampled(sampled, size, dist_shape):
+    dist_shape = infer_shape(dist_shape)
+    repeat_shape = infer_shape(size)
+    return np.reshape(sampled, repeat_shape + dist_shape)
+
+
 def replicate_samples(generator, size, repeats, *args, **kwargs):
     n = int(np.prod(repeats))
     if n == 1:
@@ -326,10 +342,7 @@ def generate_samples(generator, *args, **kwargs):
     else:
         prefix_shape = tuple(dist_shape)
 
-    try:
-        repeat_shape = tuple(size or ())
-    except TypeError:  # If size is an int
-        repeat_shape = tuple((size,))
+    repeat_shape = infer_shape(size)
 
     if broadcast_shape == (1,) and prefix_shape == ():
         if size is not None:
@@ -342,13 +355,9 @@ def generate_samples(generator, *args, **kwargs):
                                         broadcast_shape,
                                         repeat_shape + prefix_shape,
                                         *args, **kwargs)
-            if broadcast_shape == (1,) and not prefix_shape == ():
-                samples = np.reshape(samples, repeat_shape + prefix_shape)
         else:
             samples = replicate_samples(generator,
                                         broadcast_shape,
                                         prefix_shape,
                                         *args, **kwargs)
-            if broadcast_shape == (1,):
-                samples = np.reshape(samples, prefix_shape)
-    return samples
+    return reshape_sampled(samples, size, dist_shape)

pymc3/distributions/timeseries.py

@@ -1,10 +1,18 @@
 import theano.tensor as tt
 from theano import scan
 
+from .multivariate import get_tau_cov, MvNormal, MvStudentT
 from .continuous import Normal, Flat
 from .distribution import Continuous
 
-__all__ = ['AR1', 'GaussianRandomWalk', 'GARCH11', 'EulerMaruyama']
+__all__ = [
+    'AR1',
+    'GaussianRandomWalk',
+    'GARCH11',
+    'EulerMaruyama',
+    'MvGaussianRandomWalk',
+    'MvStudentTRandomWalk'
+]
 
 
 class AR1(Continuous):
@@ -108,7 +116,8 @@ def __init__(self, omega=None, alpha_1=None, beta_1=None,
         self.initial_vol = initial_vol
         self.mean = 0
 
-    def _get_volatility(self, x):
+    def get_volatility(self, x):
+        x = x[:-1]
 
         def volatility_update(x, vol, w, a, b):
             return tt.sqrt(w + a * tt.square(x) + b * tt.square(vol))
@@ -118,12 +127,11 @@ def volatility_update(x, vol, w, a, b):
                       outputs_info=[self.initial_vol],
                       non_sequences=[self.omega, self.alpha_1,
                                      self.beta_1])
-        return vol
+        return tt.concatenate(self.initial_vol, vol)
 
     def logp(self, x):
-        vol = self._get_volatility(x[:-1])
-        return (Normal.dist(0., sd=self.initial_vol).logp(x[0]) +
-                tt.sum(Normal.dist(0, sd=vol).logp(x[1:])))
+        vol = self.get_volatility(x)
+        return tt.sum(Normal.dist(0, sd=vol).logp(x))
 
 
 class EulerMaruyama(Continuous):
@@ -151,3 +159,79 @@ def logp(self, x):
         mu = xt + self.dt * f
         sd = tt.sqrt(self.dt) * g
         return tt.sum(Normal.dist(mu=mu, sd=sd).logp(x[1:]))
+
+
+class MvGaussianRandomWalk(Continuous):
+    """
+    Multivariate Random Walk with Normal innovations
+
+    Parameters
+    ----------
+    mu : tensor
+        innovation drift, defaults to 0.0
+    cov : tensor
+        pos def matrix, innovation covariance matrix
+    tau : tensor
+        pos def matrix, innovation precision (alternative to specifying cov)
+    init : distribution
+        distribution for initial value (Defaults to Flat())
+    """
+    def __init__(self, mu=0., cov=None, tau=None, init=Flat.dist(),
+                 *args, **kwargs):
+        super(MvGaussianRandomWalk, self).__init__(*args, **kwargs)
+        tau, cov = get_tau_cov(mu, tau=tau, cov=cov)
+        self.tau = tau
+        self.cov = cov
+        self.mu = mu
+        self.init = init
+        self.mean = 0.
+
+    def logp(self, x):
+        tau = self.tau
+        mu = self.mu
+        init = self.init
+
+        x_im1 = x[:-1]
+        x_i = x[1:]
+
+        innov_like = MvNormal.dist(mu=x_im1 + mu, tau=tau).logp(x_i)
+        return init.logp(x[0]) + tt.sum(innov_like)
+
+
+class MvStudentTRandomWalk(Continuous):
+    """
+    Multivariate Random Walk with StudentT innovations
+
+    Parameters
+    ----------
+    nu : degrees of freedom
+    mu : tensor
+        innovation drift, defaults to 0.0
+    cov : tensor
+        pos def matrix, innovation covariance matrix
+    tau : tensor
+        pos def matrix, innovation precision (alternative to specifying cov)
+    init : distribution
+        distribution for initial value (Defaults to Flat())
+    """
+    def __init__(self, nu, mu=0., cov=None, tau=None, init=Flat.dist(),
+                 *args, **kwargs):
+        super(MvStudentTRandomWalk, self).__init__(*args, **kwargs)
+        tau, cov = get_tau_cov(mu, tau=tau, cov=cov)
+        self.tau = tau
+        self.cov = cov
+        self.mu = mu
+        self.nu = nu
+        self.init = init
+        self.mean = 0.
+
+    def logp(self, x):
+        cov = self.cov
+        mu = self.mu
+        nu = self.nu
+        init = self.init
+
+        x_im1 = x[:-1]
+        x_i = x[1:]
+        innov_like = MvStudentT.dist(nu, cov, mu=x_im1 + mu).logp(x_i)
+        return init.logp(x[0]) + tt.sum(innov_like)

pymc3/math.py

@@ -22,3 +22,7 @@ def invlogit(x, eps=sys.float_info.epsilon):
 
 def logit(p):
     return tt.log(p / (1 - p))
+
+
+def flatten_list(tensors):
+    return tt.concatenate([var.ravel() for var in tensors])