Fix discrete sampling (#3240)

bijection only mapped continuous variables (not sure why) but it caused problems for discrete parameters in SMC.

Author: aloctavodia

pymc3/model.py

@@ -672,7 +672,7 @@ def isroot(self):
     @property
     @memoize(bound=True)
     def bijection(self):
-        vars = inputvars(self.cont_vars)
+        vars = inputvars(self.vars)
 
         bij = DictToArrayBijection(ArrayOrdering(vars),
                                    self.test_point)

pymc3/step_methods/smc.py

@@ -9,18 +9,18 @@
 
 from .arraystep import metrop_select
 from .metropolis import MultivariateNormalProposal
-from ..theanof import floatX, make_shared_replacements, join_nonshared_inputs
+from ..theanof import floatX, make_shared_replacements, join_nonshared_inputs, inputvars
 from ..model import modelcontext
 from ..backends.ndarray import NDArray
 from ..backends.base import MultiTrace
 
 
-__all__ = ['SMC', 'sample_smc']
+__all__ = ["SMC", "sample_smc"]
 
-proposal_dists = {'MultivariateNormal': MultivariateNormalProposal}
+proposal_dists = {"MultivariateNormal": MultivariateNormalProposal}
 
 
-class SMC():
+class SMC:
     """
     Sequential Monte Carlo step
 
@@ -59,8 +59,15 @@ class SMC():
         %282007%29133:7%28816%29>`__
     """
 
-    def __init__(self, n_steps=5, scaling=1., p_acc_rate=0.01, tune=True,
-                 proposal_name='MultivariateNormal', threshold=0.5):
+    def __init__(
+        self,
+        n_steps=25,
+        scaling=1.0,
+        p_acc_rate=0.01,
+        tune=True,
+        proposal_name="MultivariateNormal",
+        threshold=0.5,
+    ):
 
         self.n_steps = n_steps
         self.scaling = scaling
@@ -98,15 +105,15 @@ def sample_smc(draws=5000, step=None, progressbar=False, model=None, random_seed
     stage = 0
     acc_rate = 1
     model.marginal_likelihood = 1
-    variables = model.vars
+    variables = inputvars(model.vars)
     discrete = np.concatenate([[v.dtype in pm.discrete_types] * (v.dsize or 1) for v in variables])
     any_discrete = discrete.any()
     all_discrete = discrete.all()
     shared = make_shared_replacements(variables, model)
     prior_logp = logp_forw([model.varlogpt], variables, shared)
     likelihood_logp = logp_forw([model.datalogpt], variables, shared)
 
-    pm._log.info('Sample initial stage: ...')
+    pm._log.info("Sample initial stage: ...")
     posterior, var_info = _initial_population(draws, model, variables)
 
     while beta < 1:
@@ -127,15 +134,18 @@ def sample_smc(draws=5000, step=None, progressbar=False, model=None, random_seed
         # acceptance rate
         if step.tune and stage > 0:
             if acc_rate == 0:
-                acc_rate = 1. / step.n_steps
+                acc_rate = 1.0 / step.n_steps
             step.scaling = _tune(acc_rate)
             step.n_steps = 1 + int(np.log(step.p_acc_rate) / np.log(1 - acc_rate))
 
-        pm._log.info('Stage: {:d} Beta: {:f} Steps: {:d} Acc: {:f}'.format(stage, beta,
-                                                                           step.n_steps, acc_rate))
+        pm._log.info(
+            "Stage: {:d} Beta: {:f} Steps: {:d} Acc: {:f}".format(
+                stage, beta, step.n_steps, acc_rate
+            )
+        )
         # Apply Metropolis kernel (mutation)
-        proposed = 0.
-        accepted = 0.
+        proposed = 0.0
+        accepted = 0.0
         priors = np.array([prior_logp(sample) for sample in posterior]).squeeze()
         tempered_post = priors + likelihoods * beta
         for draw in tqdm(range(draws), disable=not progressbar):
@@ -147,12 +157,12 @@ def sample_smc(draws=5000, step=None, progressbar=False, model=None, random_seed
 
                 if any_discrete:
                     if all_discrete:
-                        delta = np.round(delta, 0).astype('int64')
-                        q_old = q_old.astype('int64')
-                        q_new = (q_old + delta).astype('int64')
+                        delta = np.round(delta, 0).astype("int64")
+                        q_old = q_old.astype("int64")
+                        q_new = (q_old + delta).astype("int64")
                     else:
                         delta[discrete] = np.round(delta[discrete], 0)
-                        q_new = (q_old + delta)
+                        q_new = floatX(q_old + delta)
                 else:
                     q_new = floatX(q_old + delta)
 
@@ -163,12 +173,12 @@ def sample_smc(draws=5000, step=None, progressbar=False, model=None, random_seed
                     accepted += accept
                     posterior[draw] = q_old
                     old_tempered_post = new_tempered_post
-                proposed += 1.
+                proposed += 1.0
 
         acc_rate = accepted / proposed
         stage += 1
 
-    trace = _posterior_to_trace(posterior, model, var_info)
+    trace = _posterior_to_trace(posterior, variables, model, var_info)
 
     return trace
 
@@ -177,6 +187,7 @@ def _initial_population(draws, model, variables):
     """
     Create an initial population from the prior
     """
+
     population = []
     var_info = {}
     start = model.test_point
@@ -188,7 +199,7 @@ def _initial_population(draws, model, variables):
         point = pm.Point({v.name: init_rnd[v.name][i] for v in variables}, model=model)
         population.append(model.dict_to_array(point))
 
-    return np.array(population), var_info
+    return np.array(floatX(population)), var_info
 
 
 def _calc_beta(beta, likelihoods, threshold=0.5):
@@ -219,11 +230,11 @@ def _calc_beta(beta, likelihoods, threshold=0.5):
         Partial marginal likelihood
     """
     low_beta = old_beta = beta
-    up_beta = 2.
+    up_beta = 2.0
     rN = int(len(likelihoods) * threshold)
 
     while up_beta - low_beta > 1e-6:
-        new_beta = (low_beta + up_beta) / 2.
+        new_beta = (low_beta + up_beta) / 2.0
         weights_un = np.exp((new_beta - old_beta) * (likelihoods - likelihoods.max()))
         weights = weights_un / np.sum(weights_un)
         ESS = int(1 / np.sum(weights ** 2))
@@ -241,11 +252,11 @@ def _calc_beta(beta, likelihoods, threshold=0.5):
     return new_beta, old_beta, weights, np.mean(sj)
 
 
-def _calc_covariance(posterior_array, weights):
+def _calc_covariance(posterior, weights):
     """
     Calculate trace covariance matrix based on importance weights.
     """
-    cov = np.cov(np.squeeze(posterior_array), aweights=weights.ravel(), bias=False, rowvar=0)
+    cov = np.cov(posterior, aweights=weights.ravel(), bias=False, rowvar=0)
     if np.isnan(cov).any() or np.isinf(cov).any():
         raise ValueError('Sample covariances not valid! Likely "chains" is too small!')
     return np.atleast_2d(cov)
@@ -264,18 +275,18 @@ def _tune(acc_rate):
     -------
     scaling: float
     """
-    # a and b after Muto & Beck 2008 .
-    a = 1. / 9
-    b = 8. / 9
+    # a and b after Muto & Beck 2008.
+    a = 1.0 / 9
+    b = 8.0 / 9
     return (a + b * acc_rate) ** 2
 
 
-def _posterior_to_trace(posterior, model, var_info):
+def _posterior_to_trace(posterior, variables, model, var_info):
     """
     Save results into a PyMC3 trace
     """
     lenght_pos = len(posterior)
-    varnames = [v.name for v in model.vars]
+    varnames = [v.name for v in variables]
 
     with model:
         strace = NDArray(model)
@@ -285,7 +296,7 @@ def _posterior_to_trace(posterior, model, var_info):
         size = 0
         for var in varnames:
             shape, new_size = var_info[var]
-            value.append(posterior[i][size:size+new_size].reshape(shape))
+            value.append(posterior[i][size : size + new_size].reshape(shape))
             size += new_size
         strace.record({k: v for k, v in zip(varnames, value)})
     return MultiTrace([strace])

pymc3/tests/test_smc.py

@@ -47,6 +47,15 @@ def test_sample(self):
         mu1d = np.abs(x).mean(axis=0)
         np.testing.assert_allclose(self.muref, mu1d, rtol=0., atol=0.03)
 
+
+    def test_discrete_continuous(self):
+        with pm.Model() as model:
+            a = pm.Poisson('a', 5)
+            b = pm.HalfNormal('b', 10)
+            y = pm.Normal('y', a, b, observed=[1, 2, 3, 4])
+            trace = pm.sample(step=pm.SMC())
+
+
     def test_ml(self):
         data = np.repeat([1, 0], [50, 50])
         marginals = []