Add check for discrete variables.

Author: taku-y

pymc3/tests/test_advi.py

@@ -1,10 +1,13 @@
 import numpy as np
-from pymc3 import Model, Normal
+from pymc3 import Model, Normal, HalfNormal, DiscreteUniform, Poisson, switch, Exponential
 from pymc3.theanof import inputvars
-from pymc3.variational.advi import variational_gradient_estimate
+from pymc3.variational.advi import variational_gradient_estimate, advi, advi_minibatch
 from theano import function
+import theano.tensor as tt
 
-def test_advi():
+from nose.tools import assert_raises
+
+def test_elbo():
     mu0 = 1.5
     sigma = 1.0
     y_obs = np.array([1.6, 1.4])
@@ -18,7 +21,7 @@ def test_advi():
 
     # Create variational gradient tensor
     grad, elbo, shared, uw = variational_gradient_estimate(
-        vars, model, n_mcsamples=1000000)
+        vars, model, n_mcsamples=10000)
 
     # Variational posterior parameters
     uw_ = np.array([1.88, np.log(1)])
@@ -34,3 +37,57 @@ def test_advi():
         0.5 * (np.log(2 * np.pi) + 1))
 
     np.testing.assert_allclose(elbo_mc, elbo_true, rtol=0, atol=1e-1)
+
+disaster_data = np.ma.masked_values([4, 5, 4, 0, 1, 4, 3, 4, 0, 6, 3, 3, 4, 0, 2, 6,
+                        3, 3, 5, 4, 5, 3, 1, 4, 4, 1, 5, 5, 3, 4, 2, 5,
+                        2, 2, 3, 4, 2, 1, 3, -999, 2, 1, 1, 1, 1, 3, 0, 0,
+                        1, 0, 1, 1, 0, 0, 3, 1, 0, 3, 2, 2, 0, 1, 1, 1,
+                        0, 1, 0, 1, 0, 0, 0, 2, 1, 0, 0, 0, 1, 1, 0, 2,
+                        3, 3, 1, -999, 2, 1, 1, 1, 1, 2, 4, 2, 0, 0, 1, 4,
+                        0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1], value=-999)
+year = np.arange(1851, 1962)
+
+def test_check_discrete():
+    with Model() as disaster_model:
+        switchpoint = DiscreteUniform('switchpoint', lower=year.min(), upper=year.max(), testval=1900)
+
+        # Priors for pre- and post-switch rates number of disasters
+        early_rate = Exponential('early_rate', 1)
+        late_rate = Exponential('late_rate', 1)
+
+        # Allocate appropriate Poisson rates to years before and after current
+        rate = switch(switchpoint >= year, early_rate, late_rate)
+
+        disasters = Poisson('disasters', rate, observed=disaster_data)
+
+    # This should raise ValueError
+    assert_raises(ValueError, advi, model=disaster_model, n=10)
+
+def test_check_discrete_minibatch():
+    disaster_data_t = tt.vector()
+    disaster_data_t.tag.test_value = np.zeros(len(disaster_data))
+
+    with Model() as disaster_model:
+
+        switchpoint = DiscreteUniform(
+            'switchpoint', lower=year.min(), upper=year.max(), testval=1900)
+
+        # Priors for pre- and post-switch rates number of disasters
+        early_rate = Exponential('early_rate', 1)
+        late_rate = Exponential('late_rate', 1)
+
+        # Allocate appropriate Poisson rates to years before and after current
+        rate = switch(switchpoint >= year, early_rate, late_rate)
+
+        disasters = Poisson('disasters', rate, observed=disaster_data_t)
+
+    def create_minibatch():
+        while True:
+            return disaster_data
+
+    # This should raise ValueError
+    assert_raises(
+        ValueError, advi_minibatch, model=disaster_model, n=10, 
+        minibatch_RVs=[disasters], minibatch_tensors=[disaster_data_t], 
+        minibatches=[create_minibatch()], verbose=False)
+    

pymc3/variational/advi.py

@@ -7,6 +7,9 @@
 from scipy import optimize
 import numpy as np
 from ..core import *
+from ..distributions import Discrete
+from ..distributions.transforms import TransformedDistribution
+from ..model import ObservedRV
 
 import theano
 from ..theanof import make_shared_replacements, join_nonshared_inputs, CallableTensor, gradient
@@ -19,6 +22,22 @@
 
 ADVIFit = namedtuple('ADVIFit', 'means, stds, elbo_vals')
 
+def is_discreteRV(var):
+    dist = var.distribution
+
+    # Transformed distributions are continuous. (is it true?)
+    if isinstance(dist, TransformedDistribution):
+        return False
+
+    return isinstance(dist, Discrete)
+
+def check_discrete_rvs(vars):
+    """Check that vars not include discrete variables, excepting ObservedRVs. 
+    """
+    vars_ = [var for var in vars if not isinstance(var, ObservedRV)]
+    if any([is_discreteRV(var) for var in vars_]):
+        raise ValueError('Model should not include discrete RVs for ADVI.')
+
 def advi(vars=None, start=None, model=None, n=5000, accurate_elbo=False, 
     learning_rate=.001, epsilon=.1, verbose=1):
     """Run ADVI. 
@@ -50,13 +69,13 @@ def advi(vars=None, start=None, model=None, n=5000, accurate_elbo=False,
     model = modelcontext(model)
     if start is None:
         start = model.test_point
-        import pdb
-        pdb.set_trace()
 
     if vars is None:
         vars = model.vars
     vars = inputvars(vars)
 
+    check_discrete_rvs(vars)
+
     n_mcsamples = 100 if accurate_elbo else 1
 
     # Create variational gradient tensor
@@ -130,6 +149,8 @@ def advi_minibatch(vars=None, start=None, model=None, n=5000, n_mcsamples=1,
 
     vars = set(inputvars(vars)) - set(minibatch_RVs)
 
+    check_discrete_rvs(vars)
+
     # Create variational gradient tensor
     grad, elbo, shared, uw = variational_gradient_estimate(
         vars, model, minibatch_RVs, minibatch_tensors, total_size,