Merge pull request #3557 from fonnesbeck/alt_vi_backend

Added backend option for VI sample trace

Author: fonnesbeck

RELEASE-NOTES.md

@@ -8,6 +8,7 @@
 - Sequential Monte Carlo - Approximate Bayesian Computation step method is now available. The implementation is in an experimental stage and will be further improved.
 - Added `Matern12` covariance function for Gaussian processes. This is the Matern kernel with nu=1/2.
 - Progressbar reports number of divergences in real time, when available [#3547](https://github.com/pymc-devs/pymc3/pull/3547).
+- Sampling from variational approximation now allows for alternative trace backends [#3550].
 
 ### Maintenance
 - Moved math operations out of `Rice`, `TruncatedNormal`, `Triangular` and `ZeroInflatedNegativeBinomial` `random` methods. Math operations on values returned by `draw_values` might not broadcast well, and all the `size` aware broadcasting is left to `generate_samples`. Fixes [#3481](https://github.com/pymc-devs/pymc3/issues/3481) and [#3508](https://github.com/pymc-devs/pymc3/issues/3508)

pymc3/tests/test_variational_inference.py

@@ -151,14 +151,22 @@ def three_var_approx(three_var_model, three_var_groups):
 def three_var_approx_single_group_mf(three_var_model):
     return MeanField(model=three_var_model)
 
-
-def test_sample_simple(three_var_approx):
-    trace = three_var_approx.sample(500)
-    assert set(trace.varnames) == {'one', 'one_log__', 'three', 'two'}
-    assert len(trace) == 500
-    assert trace[0]['one'].shape == (10, 2)
-    assert trace[0]['two'].shape == (10, )
-    assert trace[0]['three'].shape == (10, 1, 2)
+@pytest.fixture(
+    params = [
+        ('ndarray', None),
+        ('text', 'test'),
+        ('sqlite', 'test.sqlite'),
+        ('hdf5', 'test.h5')
+    ]
+)
+def test_sample_simple(three_var_approx, request):
+        backend, name = request.param
+        trace = three_var_approx.sample(100, backend=backend, name=name)
+        assert set(trace.varnames) == {'one', 'one_log__', 'three', 'two'}
+        assert len(trace) == 100
+        assert trace[0]['one'].shape == (10, 2)
+        assert trace[0]['two'].shape == (10, )
+        assert trace[0]['three'].shape == (10, 1, 2)
 
 
 @pytest.fixture

pymc3/variational/opvi.py

@@ -45,6 +45,7 @@
 from ..blocking import (
     ArrayOrdering, DictToArrayBijection, VarMap
 )
+from ..backends import NDArray, Text, SQLite, HDF5
 from ..model import modelcontext
 from ..theanof import tt_rng, change_flags, identity
 from ..util import get_default_varnames
@@ -1569,7 +1570,8 @@ def inner(draws=100):
 
         return inner
 
-    def sample(self, draws=500, include_transformed=True):
+    def sample(self, draws=500, include_transformed=True, backend='ndarray', 
+               name=None):
         """Draw samples from variational posterior.
 
         Parameters
@@ -1578,6 +1580,11 @@ def sample(self, draws=500, include_transformed=True):
             Number of random samples.
         include_transformed : `bool`
             If True, transformed variables are also sampled. Default is False.
+        backend : `str`
+            Trace backend type to use. Valid entries include: 'ndarray' (default),
+            'text', 'sqlite', 'hdf5'.
+        name : `str`
+            Name for backend (required for non-NDArray backends). Default is None.
 
         Returns
         -------
@@ -1588,8 +1595,10 @@ def sample(self, draws=500, include_transformed=True):
                                             include_transformed=include_transformed)
         samples = self.sample_dict_fn(draws)  # type: dict
         points = ({name: records[i] for name, records in samples.items()} for i in range(draws))
-        trace = pm.sampling.NDArray(model=self.model, vars=vars_sampled, test_point={
-            name: records[0] for name, records in samples.items()
+        _backends = dict(ndarray=NDArray, text=Text, hdf5=HDF5, sqlite=SQLite)
+
+        trace = _backends[backend](name=name, model=self.model, vars=vars_sampled, test_point={
+                name: records[0] for name, records in samples.items()
         })
         try:
             trace.setup(draws=draws, chain=0)