Automatically round proposal values for discrete variables in SMC

Author: ricardoV94

pymc3/smc/smc.py

@@ -12,15 +12,13 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
-import warnings
-
 from collections import OrderedDict
 
 import aesara.tensor as at
 import numpy as np
 
-from aesara import config
 from aesara import function as aesara_function
+from aesara.graph.basic import clone_replace
 from scipy.special import logsumexp
 from scipy.stats import multivariate_normal
 
@@ -34,6 +32,7 @@
 from pymc3.blocking import DictToArrayBijection
 from pymc3.model import Point, modelcontext
 from pymc3.sampling import sample_prior_predictive
+from pymc3.vartypes import discrete_types
 
 
 class SMC:
@@ -273,9 +272,15 @@ def posterior_to_trace(self):
         for i in range(lenght_pos):
             value = []
             size = 0
-            for var in varnames:
-                shape, new_size = self.var_info[var]
-                value.append(self.posterior[i][size : size + new_size].reshape(shape))
+            for varname in varnames:
+                shape, new_size = self.var_info[varname]
+                var_samples = self.posterior[i][size : size + new_size]
+                # Round discrete variable samples. The rounded values were the ones
+                # actually used in the logp evaluations (see logp_forw)
+                var = self.model[varname]
+                if var.dtype in discrete_types:
+                    var_samples = np.round(var_samples).astype(var.dtype)
+                value.append(var_samples.reshape(shape))
                 size += new_size
             strace.record(point={k: v for k, v in zip(varnames, value)})
         return strace
@@ -294,20 +299,32 @@ def logp_forw(point, out_vars, vars, shared):
         containing :class:`aesara.tensor.Tensor` for depended shared data
     """
 
+    # Convert expected input of discrete variables to (rounded) floats
+    if any(var.dtype in discrete_types for var in vars):
+        replace_int_to_float = {}
+        replace_float_to_round = {}
+        new_vars = []
+        for var in vars:
+            if var.dtype in discrete_types:
+                float_var = at.TensorType("floatX", var.broadcastable)(var.name)
+                replace_int_to_float[var] = float_var
+                new_vars.append(float_var)
+
+                round_float_var = at.round(float_var)
+                round_float_var.name = var.name
+                replace_float_to_round[float_var] = round_float_var
+            else:
+                new_vars.append(var)
+
+        replace_int_to_float.update(shared)
+        replace_float_to_round.update(shared)
+        out_vars = clone_replace(out_vars, replace_int_to_float, strict=False)
+        out_vars = clone_replace(out_vars, replace_float_to_round)
+        vars = new_vars
+
     out_list, inarray0 = join_nonshared_inputs(point, out_vars, vars, shared)
-    # TODO: Figure out how to safely accept float32 (floatX) input when there are
-    # discrete variables of int64 dtype in `vars`.
-    # See https://github.com/pymc-devs/pymc3/pull/4769#issuecomment-861494080
-    if config.floatX == "float32" and any(var.dtype == "int64" for var in vars):
-        warnings.warn(
-            "SMC sampling may run slower due to the presence of discrete variables "
-            "together with aesara.config.floatX == `float32`",
-            UserWarning,
-        )
-        f = aesara_function([inarray0], out_list[0], allow_input_downcast=True)
-    else:
-        f = aesara_function([inarray0], out_list[0])
-        f.trust_input = True
+    f = aesara_function([inarray0], out_list[0])
+    f.trust_input = True
     return f
 
 

pymc3/tests/test_smc.py

@@ -16,12 +16,15 @@
 import aesara.tensor as at
 import numpy as np
 import pytest
+import scipy.stats as st
 
 from arviz.data.inference_data import InferenceData
 
 import pymc3 as pm
 
+from pymc3.aesaraf import floatX
 from pymc3.backends.base import MultiTrace
+from pymc3.smc.smc import SMC
 from pymc3.tests.helpers import SeededTest
 
 
@@ -64,10 +67,6 @@ def two_gaussians(x):
             x = pm.Normal("x", 0, 1)
             y = pm.Normal("y", x, 1, observed=0)
 
-        with pm.Model() as self.slow_model:
-            x = pm.Normal("x", 0, 1)
-            y = pm.Normal("y", x, 1, observed=100)
-
     def test_sample(self):
         with self.SMC_test:
             mtrace = pm.sample_smc(draws=self.samples, return_inferencedata=False)
@@ -76,12 +75,43 @@ def test_sample(self):
         mu1d = np.abs(x).mean(axis=0)
         np.testing.assert_allclose(self.muref, mu1d, rtol=0.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_smc(draws=10)
+    def test_discrete_rounding_proposal(self):
+        """
+        Test that discrete variable values are automatically rounded
+        in SMC logp functions
+        """
+
+        with pm.Model() as m:
+            z = pm.Bernoulli("z", p=0.7)
+            like = pm.Potential("like", z * 1.0)
+
+        smc = SMC(model=m)
+        smc.initialize_population()
+        smc.setup_kernel()
+        smc.initialize_logp()
+
+        assert smc.prior_logp_func(floatX(np.array([-0.51]))) == -np.inf
+        assert np.isclose(smc.prior_logp_func(floatX(np.array([-0.49]))), np.log(0.3))
+        assert np.isclose(smc.prior_logp_func(floatX(np.array([0.49]))), np.log(0.3))
+        assert np.isclose(smc.prior_logp_func(floatX(np.array([0.51]))), np.log(0.7))
+        assert smc.prior_logp_func(floatX(np.array([1.51]))) == -np.inf
+
+    def test_unobserved_discrete(self):
+        n = 10
+        rng = self.get_random_state()
+
+        z_true = np.zeros(n, dtype=int)
+        z_true[int(n / 2) :] = 1
+        y = st.norm(np.array([-1, 1])[z_true], 0.25).rvs(random_state=rng)
+
+        with pm.Model() as m:
+            z = pm.Bernoulli("z", p=0.5, size=n)
+            mu = pm.math.switch(z, 1.0, -1.0)
+            like = pm.Normal("like", mu=mu, sigma=0.25, observed=y)
+
+            trace = pm.sample_smc(chains=1, return_inferencedata=False)
+
+        assert np.all(np.median(trace["z"], axis=0) == z_true)
 
     def test_ml(self):
         data = np.repeat([1, 0], [50, 50])
@@ -109,14 +139,6 @@ def test_start(self):
             }
             trace = pm.sample_smc(500, chains=1, start=start)
 
-    def test_slowdown_warning(self):
-        with aesara.config.change_flags(floatX="float32"):
-            with pytest.warns(UserWarning, match="SMC sampling may run slower due to"):
-                with pm.Model() as model:
-                    a = pm.Poisson("a", 5)
-                    y = pm.Normal("y", a, 5, observed=[1, 2, 3, 4])
-                    trace = pm.sample_smc(draws=100, chains=2, cores=1)
-
     @pytest.mark.parametrize("chains", (1, 2))
     def test_return_datatype(self, chains):
         draws = 10