Repository: jhuggins/viabel

Issue:
Model abstract class
Currently models are just log density functions that support autograd. Instead, they should inherit from the Model abstract class:

```
class Model:
    // compute log density; must be differentiable via autograd 
    def __call__(model_param)
    // get the constrained version of the model parameter; optional
    def constrain(model_param)
    def set_inverse_temp(temp)
    property supports_tempering``` 

Discussion:

Patch:
diff --git a/docs/source/api/api.rst b/docs/source/api/api.rst
index ccabe7f..707008c 100644
--- a/docs/source/api/api.rst
+++ b/docs/source/api/api.rst
@@ -6,4 +6,5 @@ API
 
    vi
    approximations
+   models
    bounds
diff --git a/docs/source/api/models.rst b/docs/source/api/models.rst
new file mode 100644
index 0000000..13e64dd
--- /dev/null
+++ b/docs/source/api/models.rst
@@ -0,0 +1,5 @@
+Models
+======
+
+.. automodule:: viabel.models
+  :members:
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 42ec620..16e0b45 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -14,7 +14,7 @@ development) that provides two types of functionality:
 #. Methods for computing bounds on the errors of the mean, standard deviation, and variance estimates produced by a continuous approximation to an (unnormalized) distribution. A canonical application is a variational approximation to a Bayesian posterior distribution.
 
 .. toctree::
-   :maxdepth: 1
+   :maxdepth: 2
 
    installation
    quickstart
diff --git a/viabel/approximations.py b/viabel/approximations.py
index 7ffd845..183fd19 100644
--- a/viabel/approximations.py
+++ b/viabel/approximations.py
@@ -56,7 +56,7 @@ class ApproximationFamily(ABC):
 
     @abstractmethod
     def sample(self, var_param, n_samples, seed=None):
-        """Return n_samples from variational distribution
+        """Generate samples from the variational distribution
 
         Parameters
         ----------
@@ -67,12 +67,22 @@ class ApproximationFamily(ABC):
 
         Returns
         -------
-        samples : `numpy.ndarray` of shape (n_samples, var_param_dim)
+        samples : `numpy.ndarray`, shape (n_samples, var_param_dim)
         """
         pass
 
     def entropy(self, var_param):
-        """Compute entropy of variational distribution."""
+        """Compute entropy of variational distribution.
+
+        Parameters
+        ----------
+        var_param : `numpy.ndarray`, shape (var_param_dim,)
+            The variational parameter.
+
+        Raises
+        ------
+        NotImplementedError
+            If entropy computation is not supported."""
         if self.supports_entropy:
             return self._entropy(var_param)
         raise NotImplementedError()
@@ -86,7 +96,18 @@ class ApproximationFamily(ABC):
         return self._supports_entropy
 
     def kl(self, var_param0, var_param1):
-        """Compute the Kullback-Leibler divergence KL(var_param0 | var_param1)"""
+        """Compute the Kullback-Leibler (KL) divergence.
+
+        Parameters
+        ----------
+        var_param0, var_param1 : `numpy.ndarray`, shape (var_param_dim,)
+            The variational parameters.
+
+        Raises
+        ------
+        NotImplementedError
+            If KL divergence computation is not supported.
+        """
         if self.supports_kl:
             return self._kl(var_param0, var_param1)
         raise NotImplementedError()
@@ -101,18 +122,42 @@ class ApproximationFamily(ABC):
 
     @abstractmethod
     def log_density(self, var_param, x):
-        """Compute the log density of the variational distribution at `x`"""
+        """The log density of the variational distribution.
+
+        Parameters
+        ----------
+        var_param : `numpy.ndarray`, shape (var_param_dim,)
+            The variational parameter.
+        x : `numpy.ndarray`, shape (dim,)
+            Value at which to evaluate the density."""
         pass
 
     @abstractmethod
     def mean_and_cov(self, var_param):
-        """Compute mean and covariance of the variational distribution"""
+        """The mean and covariance of the variational distribution.
+
+        Parameters
+        ----------
+        var_param : `numpy.ndarray`, shape (var_param_dim,)
+            The variational parameter.
+        """
         pass
 
     def pth_moment(self, var_param, p):
-        """Compute the absolute pth moment :math:`\\mathbb{E}[|X|^p]` of the variational distribution
+        """The absolute pth moment of the variational distribution.
+
+        The absolute pth moment is given by :math:`\\mathbb{E}[|X|^p]`.
+
+        Parameters
+        ----------
+        var_param : `numpy.ndarray`, shape (var_param_dim,)
+            The variational parameter.
+        p : `int`
 
-        Raises ``ValueError`` if `p` value not supported"""
+        Raises
+        ------
+        ValueError
+            If `p` value not supported"""
         if self.supports_pth_moment(p):
             return self._pth_moment(var_param, p)
         raise ValueError('p = {} is not a supported moment'.format(p))
diff --git a/viabel/convenience.py b/viabel/convenience.py
index 518a4b2..04ccf8f 100644
--- a/viabel/convenience.py
+++ b/viabel/convenience.py
@@ -1,7 +1,7 @@
 import numpy as np
 
 from viabel.approximations import MFGaussian
-from viabel.models import make_stan_log_density
+from viabel.models import Model, StanModel
 from viabel.objectives import black_box_klvi
 from viabel.optimization import adagrad_optimize
 
@@ -21,15 +21,17 @@ def bbvi(dimension, n_iters=10000, n_samples=10, log_density=None, approx=None,
         Number of Monte Carlo samples to use for estimating the gradient of
         the objective.
     log_density : `function`
-        Log density of the model. It must be provided unless ``fit`` is specified.
+        (Unnormalized) log density of the model. Must support automatic
+        differentiation with ``autograd``. Either ``log_density`` or ``fit``
+        must be provided.
     approx : `ApproximationFamily` object
         The approximation family. The default is to use ``viabel.approximations.MFGaussian``.
     objective : `function`
         Function for constructing the objective and gradient function. The default is
         to use ``viabel.objectives.black_box_klvi``.
     fit : `StanFit4model` object
-        If provided, ``log_density`` will be constructed using ``viabel.models.make_stan_log_density``.
-        Both ``fit`` and ``log_density`` cannot be given.
+        If provided, a ``StanModel`` will be used. Both ``fit`` and
+        ``log_density`` cannot be given.
     **kwargs
         Keyword arguments to pass to ``adagrad_optimize``.
 
@@ -43,9 +45,12 @@ def bbvi(dimension, n_iters=10000, n_samples=10, log_density=None, approx=None,
             raise ValueError('either log_density or fit must be specified')
         if objective_and_grad is not None:
             raise ValueError('objective_and_grad can only be specified if log_density is too')
-        log_density = make_stan_log_density(fit)
-    elif fit is not None:
+        model = StanModel(fit)
+    elif fit is None:
+        model = Model(log_density)
+    else:
         raise ValueError('log_density and fit cannot both be specified')
+
     if approx is None:
         if objective_and_grad is not None:
             raise ValueError('objective_and_grad can only be specified if approx is too')
@@ -56,7 +61,7 @@ def bbvi(dimension, n_iters=10000, n_samples=10, log_density=None, approx=None,
     var_param, var_param_history, _, _ = adagrad_optimize(n_iters, objective_and_grad, init_param, **kwargs)
     results = dict(var_param=var_param,
                    var_param_history=var_param_history,
-                   log_density=log_density,
+                   model=model,
                    approx=approx,
                    objective_and_grad=objective_and_grad)
     return results
diff --git a/viabel/models.py b/viabel/models.py
index e58954f..5b10485 100644
--- a/viabel/models.py
+++ b/viabel/models.py
@@ -2,10 +2,79 @@ from autograd.extend import primitive, defvjp
 import autograd.numpy as np
 
 __all__ = [
-    'make_stan_log_density'
+    'Model',
+    'StanModel'
 ]
 
 
+class Model(object):
+    """Base class for representing a model.
+
+    Does not support tempering. It can be overridden in part or in whole by
+    classes that inherit it. See ``StanModel`` for an example."""
+    def __init__(self, log_density):
+        """
+        Parameters
+        ----------
+        log_density : `function`
+            Function for computing the (unnormalized) log density of the model.
+            Must support automatic differentiation with ``autograd``.
+        """
+        self._log_density = log_density
+
+    def __call__(self, model_param):
+        """Compute (unnormalized) log density of the model.
+
+        Parameters
+        ----------
+        model_param : `numpy.ndarray`, shape (dim,)
+            Model parameter value
+
+        Returns
+        -------
+        log_density : `float`
+        """
+        return self._log_density(model_param)
+
+    def constrain(self, model_param):
+        """Construct dictionary of constrained parameters.
+
+        Parameters
+        ----------
+        model_param : `numpy.ndarray`, shape (dim,)
+            Model parameter value
+
+        Returns
+        -------
+        constrained_params : `dict`
+
+        Raises
+        ------
+        NotImplementedError
+            If constrained parameterization is not supported.
+        """
+        raise NotImplementedError()
+
+    @property
+    def supports_tempering(self):
+        """Whether the model supports tempering."""
+        return False
+
+    def set_inverse_temperature(self, inverse_temp):
+        """If tempering supported, set inverse temperature.
+
+        Parameters
+        ----------
+        inverse_temp : `float`
+
+        Raises
+        ------
+        NotImplementedError
+            If tempering is not supported.
+        """
+        raise NotImplementedError()
+
+
 def _vectorize_if_needed(f, a, axis=-1):
     if a.ndim > 1:
         return np.apply_along_axis(f, axis, a)
@@ -14,12 +83,14 @@ def _vectorize_if_needed(f, a, axis=-1):
 
 
 def _ensure_2d(a):
+    if a.ndim == 0:
+        return a
     while a.ndim < 2:
         a = a[:,np.newaxis]
     return a
 
 
-def make_stan_log_density(fitobj):
+def _make_stan_log_density(fitobj):
     @primitive
     def log_density(x):
         return _vectorize_if_needed(fitobj.log_prob, x)
@@ -27,3 +98,18 @@ def make_stan_log_density(fitobj):
         return lambda g: _ensure_2d(g) * _vectorize_if_needed(fitobj.grad_log_prob, x)
     defvjp(log_density, log_density_vjp)
     return log_density
+
+
+class StanModel(Model):
+    """Class that encapsulates a PyStan model."""
+    def __init__(self, fit):
+        """
+        Parameters
+        ----------
+        fit : `StanFit4model` object
+        """
+        self._fit = fit
+        super().__init__(_make_stan_log_density(fit))
+
+    def constrain(self, model_param):
+        return self._fit.constrain_pars(model_param)


Test patch:
diff --git a/viabel/tests/test_approximations.py b/viabel/tests/test_approximations.py
index a503279..a0da26e 100644
--- a/viabel/tests/test_approximations.py
+++ b/viabel/tests/test_approximations.py
@@ -56,16 +56,16 @@ def _test_family(vf, var_param0, var_param1, should_support=[], entropy_offset=0
     # These tests check that the variational family is defined self-consistently
     if vf.supports_entropy:
         _test_entropy(vf, var_param0, entropy_offset)
-    else:
+    else: # pragma: no cover
         with pytest.raises(NotImplementedError):
             vf.entropy(var_param0)
     if vf.supports_kl:
         _test_kl(vf, var_param0, var_param1)
-    else:
+    else: # pragma: no cover
         with pytest.raises(NotImplementedError):
             vf.kl(var_param0, var_param1)
     _test_mean_and_cov(vf, var_param0)
-    for p in set([2,4]) | set(should_support):
+    for p in set([1, 2,4]) | set(should_support):
         if p in should_support:
             assert vf.supports_pth_moment(p)
         if vf.supports_pth_moment(p):
diff --git a/viabel/tests/test_bounds.py b/viabel/tests/test_bounds.py
index 8f50772..5668e1c 100644
--- a/viabel/tests/test_bounds.py
+++ b/viabel/tests/test_bounds.py
@@ -11,7 +11,7 @@ def _gaussian_alpha_divergence(alpha, var1, var2):
     """Compute D_alpha(N(0, var1) | N(0, var2))"""
     tmp = alpha*var2 - (alpha - 1)*var1
     print('tmp =', tmp)
-    if tmp < 0:
+    if tmp < 0: # pragma: no cover
         return np.inf
     return -0.5 / (alpha - 1) * np.log(tmp) + .5*alpha/(alpha - 1)*np.log(var2)  - .5*np.log(var1)
 
diff --git a/viabel/tests/test_convenience.py b/viabel/tests/test_convenience.py
index bd4ff5f..b2f164b 100644
--- a/viabel/tests/test_convenience.py
+++ b/viabel/tests/test_convenience.py
@@ -4,6 +4,8 @@ import autograd.numpy as anp
 import numpy as np
 from autograd.scipy.stats import norm
 
+import pytest
+
 
 def test_bbvi():
     np.random.seed(851)
@@ -18,3 +20,12 @@ def test_bbvi():
     est_stdev = np.sqrt(np.diag(est_cov))
     np.testing.assert_almost_equal(mean.squeeze(), est_mean, decimal=2)
     np.testing.assert_almost_equal(stdev.squeeze(), est_stdev, decimal=2)
+
+    with pytest.raises(ValueError):
+        convenience.bbvi(2)
+    with pytest.raises(ValueError):
+        convenience.bbvi(2, objective_and_grad=True, fit=True)
+    with pytest.raises(ValueError):
+        convenience.bbvi(2, log_density=True, fit=True)
+    with pytest.raises(ValueError):
+        convenience.bbvi(2, log_density=True, objective_and_grad=True)
diff --git a/viabel/tests/test_models.py b/viabel/tests/test_models.py
index 899f732..7bac113 100644
--- a/viabel/tests/test_models.py
+++ b/viabel/tests/test_models.py
@@ -1,3 +1,5 @@
+import pytest
+
 from viabel import models
 
 import pickle
@@ -5,11 +7,26 @@ import pickle
 import numpy as np
 import pystan
 
-from autograd.test_util import check_grads
+import autograd.numpy as anp
+from autograd.scipy.stats import norm
+from autograd.test_util import check_vjp
+from autograd import elementwise_grad
 
 
-def _test_model(m, x):
-    check_grads(m, modes=['rev'], order=1)(x)
+def _test_model(m, x, supports_tempering, supports_constrain):
+    check_vjp(m, x)
+    check_vjp(m, x[0])
+    assert supports_tempering == m.supports_tempering
+    if supports_tempering: # pragma: no cover
+        m.set_inverse_temperature(.5)
+    else:
+        with pytest.raises(NotImplementedError):
+            m.set_inverse_temperature(.5)
+    if supports_constrain:
+        m.constrain(x[0]) == supports_constrain
+    else:
+        with pytest.raises(NotImplementedError):
+            m.constrain(x[0])
 
 
 test_model = """data {
@@ -29,12 +46,21 @@ model {
 }"""
 
 
-def test_stan_model():
+def test_Model():
+    mean = np.array([1.,-1.])[np.newaxis,:]
+    stdev = np.array([2.,5.])[np.newaxis,:]
+    log_p = lambda x: anp.sum(norm.logpdf(x, loc=mean, scale=stdev), axis=1)
+    model = models.Model(log_p)
+    x = 4*np.random.randn(10,2)
+    _test_model(model, x, False, False)
+
+
+def test_StanModel():
     compiled_model_file = 'robust_reg_model.pkl'
     try:
         with open(compiled_model_file, 'rb') as f:
             regression_model = pickle.load(f)
-    except:
+    except: # pragma: no cover
         regression_model = pystan.StanModel(model_code=test_model,
                                             model_name='regression_model')
         with open('robust_reg_model.pkl', 'wb') as f:
@@ -48,7 +74,7 @@ def test_stan_model():
 
     data = dict(N=N, x=x, y=y, df=40)
     fit = regression_model.sampling(data=data, iter=10, thin=1, chains=1)
-    stan_log_density = models.make_stan_log_density(fit)
+    model = models.StanModel(fit)
 
     x = 4*np.random.randn(10,2)
-    _test_model(stan_log_density, x)
+    _test_model(model, x, False, dict(beta=x[0]))
diff --git a/viabel/tests/test_objectives.py b/viabel/tests/test_objectives.py
index 9fd2039..a752c3c 100644
--- a/viabel/tests/test_objectives.py
+++ b/viabel/tests/test_objectives.py
@@ -31,7 +31,3 @@ def test_klvi():
 
 def test_chivi():
     _test_vi(black_box_chivi, 100, alpha=2)
-
-
-if __name__ == '__main__':
-    test_klvi()