refactor variational module, add histogram approximation (#1904)

* refactor module, add histogram

* add more tests

* refactor some code concerning AEVB histogram

* fix test for histogram

* use mean as deterministic point in Histogram

* remove unused import

* change names of shortcuts

* add names to shared params

* add new line at the end of `approximations.py`

Author: ferrine

pymc3/tests/test_variational_inference.py

@@ -1,13 +1,16 @@
-from six.moves import cPickle as pickle
+import pickle
 import unittest
 import numpy as np
 from theano import theano, tensor as tt
 import pymc3 as pm
 from pymc3 import Model, Normal
-from pymc3.variational.inference import (
-    KL, MeanField, ADVI, FullRankADVI,
+from pymc3.variational import (
+    ADVI, FullRankADVI,
+    Histogram,
     fit
 )
+from pymc3.variational.operators import KL
+from pymc3.variational.approximations import MeanField
 
 from pymc3.tests import models
 from pymc3.tests.helpers import SeededTest
@@ -186,7 +189,7 @@ def cb(*_):
                 mu_ = Normal('mu', mu=mu0, sd=sd0, testval=0)
                 Normal('x', mu=mu_, sd=sd, observed=data_t, total_size=n)
                 inf = self.inference()
-                approx = inf.fit(self.NITER, callbacks=[cb])
+                approx = inf.fit(self.NITER, callbacks=[cb], obj_n_mc=10)
                 trace = approx.sample_vp(10000)
             np.testing.assert_allclose(np.mean(trace['mu']), mu_post, rtol=0.4)
             np.testing.assert_allclose(np.std(trace['mu']), np.sqrt(1. / d), rtol=0.4)
@@ -199,22 +202,34 @@ def test_pickling(self):
             inference.fit(20)
 
         def test_aevb(self):
-            _, model, _ = models.exponential_beta()
+            _, model, _ = models.exponential_beta(n=2)
             x = model.x
             y = model.y
             mu = theano.shared(x.init_value) * 2
-            sd = theano.shared(x.init_value) * 3
+            rho = theano.shared(np.zeros_like(x.init_value))
             with model:
-                inference = self.inference(local_rv={y: (mu, sd)})
-                inference.fit(3)
+                inference = self.inference(local_rv={y: (mu, rho)})
+                approx = inference.fit(3, obj_n_mc=2)
+                approx.sample_vp(10)
+                approx.apply_replacements(
+                    y,
+                    more_replacements={x: np.asarray([1, 1], dtype=x.dtype)}
+                ).eval()
+
+        def test_profile(self):
+            with models.multidimensional_model()[1]:
+                self.inference().run_profiling(10)
 
 
 class TestMeanField(TestApproximates.Base):
     inference = ADVI
 
     def test_approximate(self):
         with models.multidimensional_model()[1]:
-            fit(10, method='advi')
+            meth = ADVI()
+            fit(10, method=meth)
+            self.assertRaises(KeyError, fit, 10, method='undefined')
+            self.assertRaises(TypeError, fit, 10, method=1)
 
 
 class TestFullRank(TestApproximates.Base):
@@ -234,11 +249,54 @@ def test_from_advi(self):
 
     def test_combined(self):
         with models.multidimensional_model()[1]:
+            self.assertRaises(ValueError, fit, 10, method='advi->fullrank_advi', frac=1)
             fit(10, method='advi->fullrank_advi', frac=.5)
 
     def test_approximate(self):
         with models.multidimensional_model()[1]:
             fit(10, method='fullrank_advi')
 
+
+class TestHistogram(SeededTest):
+    def test_sampling(self):
+        with models.multidimensional_model()[1]:
+            full_rank = FullRankADVI()
+            approx = full_rank.fit(20)
+            trace0 = approx.sample_vp(10000)
+            histogram = Histogram(trace0)
+        trace1 = histogram.sample_vp(100000)
+        np.testing.assert_allclose(trace0['x'].mean(0), trace1['x'].mean(0), atol=0.01)
+        np.testing.assert_allclose(trace0['x'].var(0), trace1['x'].var(0), atol=0.01)
+
+    def test_aevb_histogram(self):
+        _, model, _ = models.exponential_beta(n=2)
+        x = model.x
+        mu = theano.shared(x.init_value)
+        rho = theano.shared(np.zeros_like(x.init_value))
+        with model:
+            inference = ADVI(local_rv={x: (mu, rho)})
+            approx = inference.approx
+            trace0 = approx.sample_vp(1000)
+            histogram = Histogram(trace0, local_rv={x: (mu, rho)})
+            trace1 = histogram.sample_vp(10000)
+            histogram.random(no_rand=True)
+            histogram.random_fn(no_rand=True)
+        np.testing.assert_allclose(trace0['y'].mean(0), trace1['y'].mean(0), atol=0.02)
+        np.testing.assert_allclose(trace0['y'].var(0), trace1['y'].var(0), atol=0.02)
+        np.testing.assert_allclose(trace0['x'].mean(0), trace1['x'].mean(0), atol=0.02)
+        np.testing.assert_allclose(trace0['x'].var(0), trace1['x'].var(0), atol=0.02)
+
+    def test_random(self):
+        with models.multidimensional_model()[1]:
+            full_rank = FullRankADVI()
+            approx = full_rank.approx
+            trace0 = approx.sample_vp(10000)
+            histogram = Histogram(trace0)
+            histogram.randidx(None).eval()
+            histogram.randidx(1).eval()
+            histogram.random_fn(no_rand=True)
+            histogram.random_fn(no_rand=False)
+            histogram.histogram_logp.eval()
+
 if __name__ == '__main__':
     unittest.main()

pymc3/variational/__init__.py

@@ -23,3 +23,7 @@
     FullRankADVI,
     fit,
 )
+from .approximations import Histogram
+
+from . import approximations
+from . import operators