Merge e106eb5 into 7f8d7c3

pymc3/distributions/dist_math.py

@@ -10,6 +10,9 @@
 
 from .special import gammaln, multigammaln
 
+c = - 0.5 * np.log(2 * np.pi)
+
+
 def bound(logp, *conditions, **kwargs):
     """
     Bounds a log probability density with several conditions.
@@ -95,3 +98,105 @@ 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
+    :math:`mu + sd*e = mu + log(1+exp(rho))*e`"""
+    return tt.log(tt.exp(sd) - 1)
+
+
+def rho2sd(rho):
+    """
+    `rho -> sd` theano converter
+    :math:`mu + sd*e = mu + log(1+exp(rho))*e`"""
+    return tt.log1p(tt.exp(rho))
+
+
+def log_normal(x, mean, **kwargs):
+    """
+    Calculate logarithm of normal distribution at point `x`
+    with given `mean` and `std`
+    Parameters
+    ----------
+    x : Tensor
+        point of evaluation
+    mean : Tensor
+        mean of normal distribution
+    kwargs : one of parameters `{sd, tau, w, rho}`
+    Notes
+    -----
+    There are four variants for density parametrization.
+    They are:
+        1) standard deviation - `std`
+        2) `w`, logarithm of `std` :math:`w = log(std)`
+        3) `rho` that follows this equation :math:`rho = log(exp(std) - 1)`
+        4) `tau` that follows this equation :math:`tau = std^{-1}`
+    ----
+    """
+    sd = kwargs.get('sd')
+    w = kwargs.get('w')
+    rho = kwargs.get('rho')
+    tau = kwargs.get('tau')
+    eps = kwargs.get('eps', 0.0)
+    check = sum(map(lambda a: a is not None, [sd, w, rho, tau]))
+    if check > 1:
+        raise ValueError('more than one required kwarg is passed')
+    if check == 0:
+        raise ValueError('none of required kwarg is passed')
+    if sd is not None:
+        std = sd
+    elif w is not None:
+        std = tt.exp(w)
+    elif rho is not None:
+        std = rho2sd(rho)
+    else:
+        std = tau**(-1)
+    std += eps
+    return c - tt.log(tt.abs_(std)) - (x - mean) ** 2 / (2 * std ** 2)
+
+
+def log_normal_mv(x, mean, **kwargs):
+    """
+    Calculate logarithm of normal distribution at point `x`
+    with given `mean` and `sigma` matrix
+    Parameters
+    ----------
+    x : Tensor
+        point of evaluation
+    mean : Tensor
+        mean of normal distribution
+    kwargs : one of parameters `{cov, tau, chol}`
+    Notes
+    -----
+    There are three variants for density parametrization.
+    They are:
+        1) covariance matrix - `cov`
+        2) precision matrix - `tau`,
+        3) cholesky decomposition matrix  - `chol`
+    ----
+    """
+    T = kwargs.get('tau')
+    S = kwargs.get('cov')
+    L = kwargs.get('chol')
+    check = sum(map(lambda a: a is not None, [T, S, L]))
+    if check > 1:
+        raise ValueError('more than one required kwarg is passed')
+    if check == 0:
+        raise ValueError('none of required kwarg is passed')
+    # avoid unnecessary computations
+    if L is not None:
+        S = L.dot(L.T)
+        T = tt.nlinalg.matrix_inverse(S)
+        Det = tt.nlinalg.det(S)
+    elif T is not None:
+        Det = 1 / tt.nlinalg.det(T)
+    else:
+        T = tt.nlinalg.matrix_inverse(S)
+        Det = tt.nlinalg.det(S)
+    delta = x - mean
+    k = S.shape[0]
+    result = k * tt.log(2 * np.pi) + tt.log(Det)
+    result += delta.dot(T).dot(delta)
+    return -1 / 2. * result

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])

pymc3/tests/test_variational_inference.py

@@ -0,0 +1,123 @@
+import unittest
+import numpy as np
+import theano
+import theano.tensor as tt
+from theano.sandbox.rng_mrg import MRG_RandomStreams
+import pymc3 as pm
+from pymc3 import Model, Normal
+from pymc3.variational.opvi import KL, MeanField, FullRank, TestFunction
+from ..theanof import set_tt_rng
+from . import models
+set_tt_rng(MRG_RandomStreams(42))
+
+
+class TestApproximates:
+    class Base(unittest.TestCase):
+        op = KL
+        approx = MeanField
+        test_function = TestFunction
+        NITER = 30000
+
+        def test_elbo(self):
+            if self.approx is not MeanField:
+                raise unittest.SkipTest
+            mu0 = 1.5
+            sigma = 1.0
+            y_obs = np.array([1.6, 1.4])
+
+            post_mu = np.array([1.88])
+            post_sd = np.array([1])
+            # Create a model for test
+            with Model() as model:
+                mu = Normal('mu', mu=mu0, sd=sigma)
+                Normal('y', mu=mu, sd=1, observed=y_obs)
+
+            # Create variational gradient tensor
+            mean_field = self.approx(model=model)
+            elbo = -self.op(mean_field)(self.test_function())(mean_field.random())
+
+            mean_field.shared_params['mu'].set_value(post_mu)
+            mean_field.shared_params['rho'].set_value(np.log(np.exp(post_sd) - 1))
+
+            f = theano.function([], elbo)
+            elbo_mc = sum(f() for _ in range(10000))/10000
+
+            # Exact value
+            elbo_true = (-0.5 * (
+                3 + 3 * post_mu**2 - 2 * (y_obs[0] + y_obs[1] + mu0) * post_mu +
+                y_obs[0]**2 + y_obs[1]**2 + mu0**2 + 3 * np.log(2 * np.pi)) +
+                0.5 * (np.log(2 * np.pi) + 1))
+            np.testing.assert_allclose(elbo_mc, elbo_true, rtol=0, atol=1e-1)
+
+        @unittest.skip('Not ready')
+        def test_vary_samples(self):
+            _, model, _ = models.simple_model()
+            i = tt.iscalar('i')
+            i.tag.test_value = 1
+            with model:
+                mf = self.approx()
+                elbo = mf.elbo(i)
+            elbos = theano.function([i], elbo)
+            self.assertEqual(elbos(1).shape[0], 1)
+            self.assertEqual(elbos(10).shape[0], 10)
+
+        @unittest.skip('Not ready')
+        def test_vars_view(self):
+            _, model, _ = models.multidimensional_model()
+            with model:
+                mf = self.approx()
+                posterior = mf.posterior(10)
+                x_sampled = mf.view_from(posterior, 'x').eval()
+            self.assertEqual(x_sampled.shape, (10,) + model['x'].dshape)
+
+        @unittest.skip('Not ready')
+        def test_sample_vp(self):
+            n_samples = 100
+            xs = np.random.binomial(n=1, p=0.2, size=n_samples)
+            with pm.Model():
+                p = pm.Beta('p', alpha=1, beta=1)
+                pm.Binomial('xs', n=1, p=p, observed=xs)
+                mf = self.approx()
+                trace = mf.sample_vp(draws=1, hide_transformed=True)
+                self.assertListEqual(trace.varnames, ['p'])
+                trace = mf.sample_vp(draws=1, hide_transformed=False)
+                self.assertListEqual(sorted(trace.varnames), ['p', 'p_logodds_'])
+
+        def test_advi_optimizer(self):
+            n = 1000
+            sd0 = 2.
+            mu0 = 4.
+            sd = 3.
+            mu = -5.
+
+            data = sd * np.random.randn(n) + mu
+
+            d = n / sd ** 2 + 1 / sd0 ** 2
+            mu_post = (n * np.mean(data) / sd ** 2 + mu0 / sd0 ** 2) / d
+
+            with Model():
+                mu_ = Normal('mu', mu=mu0, sd=sd0, testval=0)
+                Normal('x', mu=mu_, sd=sd, observed=data)
+                pm.Deterministic('mu_sq', mu_**2)
+                approx = self.approx()
+                obj_f = self.op(approx)(self.test_function())
+                updates = obj_f.updates(obj_f.random())
+                step = theano.function([], [], updates=updates)
+                for _ in range(self.NITER):
+                    step()
+                trace = approx.sample_vp(10000)
+
+            np.testing.assert_allclose(np.mean(trace['mu']), mu_post, rtol=0.1)
+            np.testing.assert_allclose(np.std(trace['mu']), np.sqrt(1. / d), rtol=0.4)
+
+
+class TestMeanField(TestApproximates.Base):
+    pass
+
+
+class TestFullRank(TestApproximates.Base):
+    approx = FullRank
+
+
+if __name__ == '__main__':
+    unittest.main()

pymc3/theanof.py

@@ -4,16 +4,24 @@
 from theano import theano, scalar, tensor as tt
 from theano.gof.graph import inputs
 from theano.tensor import TensorVariable
+from theano.sandbox.rng_mrg import MRG_RandomStreams
 from theano.gof import Op
 from theano.configparser import change_flags
 from .memoize import memoize
 from .blocking import ArrayOrdering
 from .data import DataGenerator
 
-__all__ = ['gradient', 'hessian', 'hessian_diag', 'inputvars',
-           'cont_inputs', 'floatX', 'jacobian',
-           'CallableTensor', 'join_nonshared_inputs',
-           'make_shared_replacements', 'generator']
+__all__ = ['gradient',
+           'hessian',
+           'hessian_diag',
+           'inputvars',
+           'cont_inputs',
+           'floatX',
+           'jacobian',
+           'CallableTensor',
+           'join_nonshared_inputs',
+           'make_shared_replacements',
+           'generator']
 
 
 def inputvars(a):
@@ -305,5 +313,60 @@ def set_gen(self, gen):
 
 
 def generator(gen):
-    """shortcut for `GeneratorOp`"""
+    """
+    shortcut for `GeneratorOp`
+    """
     return GeneratorOp(gen)()
+
+
+@change_flags(compute_test_value='off')
+def launch_rng(rng):
+    """
+    Helper function for safe launch of theano random generator.
+    If not launched, there will be problems with test_value
+
+    Parameters
+    ----------
+    rng : `theano.sandbox.rng_mrg.MRG_RandomStreams` instance
+    """
+    state = rng.rstate
+    rng.inc_rstate()
+    rng.set_rstate(state)
+
+_tt_rng = MRG_RandomStreams()
+launch_rng(_tt_rng)
+
+
+def tt_rng():
+    """
+    Get the package-level random number generator.
+
+    Returns
+    -------
+    `theano.sandbox.rng_mrg.MRG_RandomStreams` instance
+        `theano.sandbox.rng_mrg.MRG_RandomStreams`
+        instance passed to the most recent call of `set_tt_rng`
+    """
+    return _tt_rng
+
+
+def set_tt_rng(new_rng):
+    """
+    Set the package-level random number generator.
+
+    Parameters
+    ----------
+    new_rng : `theano.sandbox.rng_mrg.MRG_RandomStreams` instance
+        The random number generator to use.
+    """
+    global _tt_rng
+    _tt_rng = new_rng
+    launch_rng(_tt_rng)
+
+
+class GradScale(theano.compile.ViewOp):
+    def __init__(self, multiplier):
+        self.multiplier = multiplier
+
+    def grad(self, args, g_outs):
+        return [self.multiplier * g_out for g_out in g_outs]