Repository: pyro-ppl/pyro

Issue:
Resampling bug in SMCFilter?
@tuananhle7 [points out](https://github.com/pyro-ppl/pyro/pull/1947#issuecomment-551885868)

> I was looking at [this example](https://github.com/pyro-ppl/pyro/blob/dev/examples/smcfilter.py) and the [source code](https://github.com/pyro-ppl/pyro/blob/dev/pyro/infer/smcfilter.py) in order to understand the implementation.

> I don't understand how resampling works:

> Assume that at time `t - 1`, we have particles `{z_(t - 1)k}` and weights `{w_(t - 1)k}` (`k = 1, ..., K`) which I think (?) are stored in [`_values`](https://github.com/pyro-ppl/pyro/blob/dev/pyro/infer/smcfilter.py#L42) and [`_log_weights`](https://github.com/pyro-ppl/pyro/blob/dev/pyro/infer/smcfilter.py#L43).

> The [resampling step](https://github.com/pyro-ppl/pyro/blob/dev/pyro/infer/smcfilter.py#L73) produces ancestral indices `a_(t - 1)k` (`k = 1, ..., K`) which are marginally distributed according to a Categorical distribution parameterized by the normalized weights ([here](https://github.com/pyro-ppl/pyro/blob/dev/pyro/infer/smcfilter.py#L123) I think). Then `z_(t - 1)k` which is stored in `_values` is updated to be `z_(t - 1)(a_(t - 1)k)` [here](https://github.com/pyro-ppl/pyro/blob/dev/pyro/infer/smcfilter.py#L124).

> At time `t`, we should propose `z_tk` from `guide(z_t(k - 1))` where `z_t(k - 1)` is already resampled. However, [here](https://github.com/pyro-ppl/pyro/blob/dev/examples/smcfilter.py#L60-L61), `z_tk` is proposed from a guide which is conditioned on `self.z` (attribute of guide) which doesn't seem to interact with `_values` which stores the resampled values.

cc @chanjed @karalets 


Discussion:
Thanks for the careful reading, @tuananhle7. It's been a while since we wrote this, and yes I agree that it appears `._values` is not affecting `SimpleHarmonicGuide.z` when that value is later read. We'd happily accept a fix if you see one. In the meantime I'll think about possible fixes.

Resolving patch:
diff --git a/examples/smcfilter.py b/examples/smcfilter.py
--- a/examples/smcfilter.py
+++ b/examples/smcfilter.py
@@ -7,7 +7,7 @@
 import pyro.distributions as dist
 from pyro.infer import SMCFilter
 
-logging.basicConfig(format='%(relativeCreated) 9d %(message)s', level=logging.INFO)
+logging.basicConfig(format="%(relativeCreated) 9d %(message)s", level=logging.INFO)
 
 """
 This file demonstrates how to use the SMCFilter algorithm with
@@ -28,20 +28,19 @@ def __init__(self, process_noise, measurement_noise):
         self.sigma_z = torch.tensor(process_noise)
         self.sigma_y = torch.tensor(measurement_noise)
 
-    def init(self, initial):
+    def init(self, state, initial):
         self.t = 0
-        self.z = initial
-        self.y = None
+        state["z"] = pyro.sample("z_init", dist.Delta(initial, event_dim=1))
 
-    def step(self, y=None):
+    def step(self, state, y=None):
         self.t += 1
-        self.z = pyro.sample("z_{}".format(self.t),
-                             dist.Normal(self.z.matmul(self.A), self.B*self.sigma_z).to_event(1))
-        self.y = pyro.sample("y_{}".format(self.t),
-                             dist.Normal(self.z[..., 0], self.sigma_y),
-                             obs=y)
-
-        return self.z, self.y
+        state["z"] = pyro.sample(
+            "z_{}".format(self.t),
+            dist.Normal(state["z"].matmul(self.A), self.B*self.sigma_z).to_event(1))
+        y = pyro.sample("y_{}".format(self.t),
+                        dist.Normal(state["z"][..., 0], self.sigma_y),
+                        obs=y)
+        return state["z"], y
 
 
 class SimpleHarmonicModel_Guide:
@@ -49,27 +48,29 @@ class SimpleHarmonicModel_Guide:
     def __init__(self, model):
         self.model = model
 
-    def init(self, initial):
+    def init(self, state, initial):
         self.t = 0
-        self.z = initial
+        pyro.sample("z_init", dist.Delta(initial, event_dim=1))
 
-    def step(self, y=None):
+    def step(self, state, y=None):
         self.t += 1
 
         # Proposal distribution
-        self.z = pyro.sample("z_{}".format(self.t),
-                             dist.Normal(self.z.matmul(self.model.A), torch.tensor([1., 1.])).to_event(1))
+        pyro.sample(
+            "z_{}".format(self.t),
+            dist.Normal(state["z"].matmul(self.model.A), torch.tensor([1., 1.])).to_event(1))
 
 
 def generate_data(args):
     model = SimpleHarmonicModel(args.process_noise, args.measurement_noise)
 
+    state = {}
     initial = torch.tensor([1., 0.])
-    model.init(initial=initial)
+    model.init(state, initial=initial)
     zs = [initial]
     ys = [None]
     for t in range(args.num_timesteps):
-        z, y = model.step()
+        z, y = model.step(state)
         zs.append(z)
         ys.append(y)
 
@@ -85,24 +86,25 @@ def main(args):
 
     smc = SMCFilter(model, guide, num_particles=args.num_particles, max_plate_nesting=0)
 
-    logging.info('Generating data')
+    logging.info("Generating data")
     zs, ys = generate_data(args)
 
-    logging.info('Filtering')
+    logging.info("Filtering")
+
     smc.init(initial=torch.tensor([1., 0.]))
     for y in ys[1:]:
         smc.step(y)
 
-    logging.info('Marginals')
-    empirical = smc.get_empirical()
-    for t in range(1, 1+args.num_timesteps):
-        z = empirical["z_{}".format(t)]
-        logging.info("{}\t{}\t{}\t{}".format(t, zs[t], z.mean, z.variance))
+    logging.info("At final time step:")
+    z = smc.get_empirical()["z"]
+    logging.info("truth: {}".format(zs[-1]))
+    logging.info("mean: {}".format(z.mean))
+    logging.info("std: {}".format(z.variance ** 0.5))
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Simple Harmonic Oscillator w/ SMC Filtering Inference")
-    parser.add_argument("-n", "--num-timesteps", default=50, type=int)
+    parser.add_argument("-n", "--num-timesteps", default=500, type=int)
     parser.add_argument("-p", "--num-particles", default=100, type=int)
     parser.add_argument("--process-noise", default=1., type=float)
     parser.add_argument("--measurement-noise", default=1., type=float)
diff --git a/pyro/infer/smcfilter.py b/pyro/infer/smcfilter.py
--- a/pyro/infer/smcfilter.py
+++ b/pyro/infer/smcfilter.py
@@ -1,31 +1,35 @@
+import contextlib
+
 import torch
 
 import pyro
 import pyro.distributions as dist
 import pyro.poutine as poutine
+from pyro.infer.util import is_validation_enabled
 from pyro.poutine.util import prune_subsample_sites
 
 
-def _extract_samples(trace):
-    return {name: site["value"]
-            for name, site in trace.nodes.items()
-            if site["type"] == "sample"
-            if not site["is_observed"]
-            if type(site["fn"]).__name__ != "_Subsample"}
-
-
-class SMCFilter(object):
+class SMCFilter:
     """
     :class:`SMCFilter` is the top-level interface for filtering via sequential
     monte carlo.
 
-    The model and guide should be objects with two methods: ``.init()`` and
-    ``.step()``. These two methods should have the same signature as :meth:`init`
-    and :meth:`step` of this class. These methods are intended to be called first
-    with :meth:`init`, then with :meth:`step` repeatedly.
-
-    :param object model: probabilistic model defined as a function
-    :param object guide: guide used for sampling defined as a function
+    The model and guide should be objects with two methods:
+    ``.init(state, ...)`` and ``.step(state, ...)``, intended to be called
+    first with :meth:`init` , then with :meth:`step` repeatedly.  These two
+    methods should have the same signature as :class:`SMCFilter` 's
+    :meth:`init` and :meth:`step` of this class, but with an extra first
+    argument ``state`` that should be used to store all tensors that depend on
+    sampled variables. The ``state`` will be a dict-like object,
+    :class:`SMCState` , with arbitrary keys and :class:`torch.Tensor` values.
+    Models can read and write ``state`` but guides can only read from it.
+
+    Inference complexity is ``O(len(state) * num_time_steps)``, so to avoid
+    quadratic complexity in Markov models, ensure that ``state`` has fixed size.
+
+    :param object model: probabilistic model with ``init`` and ``step`` methods
+    :param object guide: guide used for sampling,  with ``init`` and ``step``
+        methods
     :param int num_particles: The number of particles used to form the
         distribution.
     :param int max_plate_nesting: Bound on max number of nested
@@ -38,9 +42,9 @@ def __init__(self, model, guide, num_particles, max_plate_nesting):
         self.num_particles = num_particles
         self.max_plate_nesting = max_plate_nesting
 
-        # Equivalent to an empirical distribution.
-        self._values = {}
-        self._log_weights = torch.zeros(self.num_particles)
+        # Equivalent to an empirical distribution, but allows a
+        # user-defined dynamic collection of tensors.
+        self.state = SMCState(self.num_particles)
 
     def init(self, *args, **kwargs):
         """
@@ -49,12 +53,12 @@ def init(self, *args, **kwargs):
         """
         self.particle_plate = pyro.plate("particles", self.num_particles, dim=-1-self.max_plate_nesting)
         with poutine.block(), self.particle_plate:
-            guide_trace = poutine.trace(self.guide.init).get_trace(*args, **kwargs)
+            with self.state._lock():
+                guide_trace = poutine.trace(self.guide.init).get_trace(self.state, *args, **kwargs)
             model = poutine.replay(self.model.init, guide_trace)
-            model_trace = poutine.trace(model).get_trace(*args, **kwargs)
+            model_trace = poutine.trace(model).get_trace(self.state, *args, **kwargs)
 
         self._update_weights(model_trace, guide_trace)
-        self._values.update(_extract_samples(model_trace))
         self._maybe_importance_resample()
 
     def step(self, *args, **kwargs):
@@ -64,31 +68,22 @@ def step(self, *args, **kwargs):
         Any args or kwargs are passed to the model and guide
         """
         with poutine.block(), self.particle_plate:
-            guide_trace = poutine.trace(self.guide.step).get_trace(*args, **kwargs)
+            with self.state._lock():
+                guide_trace = poutine.trace(self.guide.step).get_trace(self.state, *args, **kwargs)
             model = poutine.replay(self.model.step, guide_trace)
-            model_trace = poutine.trace(model).get_trace(*args, **kwargs)
+            model_trace = poutine.trace(model).get_trace(self.state, *args, **kwargs)
 
         self._update_weights(model_trace, guide_trace)
-        self._values.update(_extract_samples(model_trace))
         self._maybe_importance_resample()
 
-    def get_values_and_log_weights(self):
-        """
-        Returns the particles and its (unnormalized) log weights.
-        :returns: the values and unnormalized log weights.
-        :rtype: tuple of dict and floats where the dict is a key of name of latent to value of latent.
-        """
-        # TODO: Be clear that these are unnormalized weights. May want to normalize later.
-        return self._values, self._log_weights
-
     def get_empirical(self):
         """
-        :returns: a marginal distribution over every latent variable.
+        :returns: a marginal distribution over all state tensors.
         :rtype: a dictionary with keys which are latent variables and values
             which are :class:`~pyro.distributions.Empirical` objects.
         """
-        return {name: dist.Empirical(value, self._log_weights)
-                for name, value in self._values.items()}
+        return {key: dist.Empirical(value, self.state._log_weights)
+                for key, value in self.state.items()}
 
     @torch.no_grad()
     def _update_weights(self, model_trace, guide_trace):
@@ -105,21 +100,63 @@ def _update_weights(self, model_trace, guide_trace):
                 model_site = model_trace.nodes[name]
                 log_p = model_site["log_prob"].reshape(self.num_particles, -1).sum(-1)
                 log_q = guide_site["log_prob"].reshape(self.num_particles, -1).sum(-1)
-                self._log_weights += log_p - log_q
+                self.state._log_weights += log_p - log_q
 
         for site in model_trace.nodes.values():
             if site["type"] == "sample" and site["is_observed"]:
                 log_p = site["log_prob"].reshape(self.num_particles, -1).sum(-1)
-                self._log_weights += log_p
+                self.state._log_weights += log_p
 
-        self._log_weights -= self._log_weights.max()
+        self.state._log_weights -= self.state._log_weights.max()
 
     def _maybe_importance_resample(self):
-        if True:  # TODO check perplexity
+        if self.state:  # TODO check perplexity
             self._importance_resample()
 
     def _importance_resample(self):
-        # TODO: Turn quadratic algo -> linear algo by being lazier
-        index = dist.Categorical(logits=self._log_weights).sample(sample_shape=(self.num_particles,))
-        self._values = {name: value[index].contiguous() for name, value in self._values.items()}
+        index = dist.Categorical(logits=self.state._log_weights).sample(sample_shape=(self.num_particles,))
+        self.state._resample(index)
+
+
+class SMCState(dict):
+    """
+    Dictionary-like object to hold a vectorized collection of tensors to
+    represent all state during inference with :class:`SMCFilter`. During
+    inference, the :class:`SMCFilter` resample these tensors.
+
+    Keys may have arbitrary hashable type.
+    Values must be :class:`torch.Tensor` s.
+
+    :param int num_particles:
+    """
+    def __init__(self, num_particles):
+        assert isinstance(num_particles, int) and num_particles > 0
+        super().__init__()
+        self._num_particles = num_particles
+        self._log_weights = torch.zeros(num_particles)
+        self._locked = False
+
+    @contextlib.contextmanager
+    def _lock(self):
+        self._locked = True
+        try:
+            yield
+        finally:
+            self._locked = False
+
+    def __setitem__(self, key, value):
+        if self._locked:
+            raise RuntimeError("Guide cannot write to SMCState")
+        if is_validation_enabled():
+            if not isinstance(value, torch.Tensor):
+                raise TypeError("Only Tensors can be stored in an SMCState, but got {}"
+                                .format(type(value).__name__))
+            if value.dim() == 0 or value.size(0) != self._num_particles:
+                raise ValueError("Expected leading dim of size {} but got shape {}"
+                                 .format(self._num_particles, value.shape))
+        super().__setitem__(key, value)
+
+    def _resample(self, index):
+        for key, value in self.items():
+            self[key] = value[index].contiguous()
         self._log_weights.fill_(0.)


Test patch:
diff --git a/tests/infer/test_smcfilter.py b/tests/infer/test_smcfilter.py
--- a/tests/infer/test_smcfilter.py
+++ b/tests/infer/test_smcfilter.py
@@ -5,6 +5,7 @@
 import pyro.distributions as dist
 import pyro.poutine as poutine
 from pyro.infer import SMCFilter
+from tests.common import assert_close
 
 
 class SmokeModel:
@@ -13,23 +14,23 @@ def __init__(self, state_size, plate_size):
         self.state_size = state_size
         self.plate_size = plate_size
 
-    def init(self):
+    def init(self, state):
         self.t = 0
-        self.x_mean = pyro.sample("x_mean", dist.Normal(0., 1.))
-        self.y_mean = pyro.sample("y_mean",
-                                  dist.MultivariateNormal(torch.zeros(self.state_size), torch.eye(self.state_size)))
+        state["x_mean"] = pyro.sample("x_mean", dist.Normal(0., 1.))
+        state["y_mean"] = pyro.sample("y_mean",
+                                      dist.MultivariateNormal(torch.zeros(self.state_size),
+                                                              torch.eye(self.state_size)))
 
-    def step(self, x=None, y=None):
+    def step(self, state, x=None, y=None):
         v = pyro.sample("v_{}".format(self.t), dist.Normal(0., 1.))
         with pyro.plate("plate", self.plate_size):
             w = pyro.sample("w_{}".format(self.t), dist.Normal(v, 1.))
             x = pyro.sample("x_{}".format(self.t),
-                            dist.Normal(self.x_mean + w, 1), obs=x)
+                            dist.Normal(state["x_mean"] + w, 1), obs=x)
             y = pyro.sample("y_{}".format(self.t),
-                            dist.MultivariateNormal(self.y_mean + w.unsqueeze(-1), torch.eye(self.state_size)),
+                            dist.MultivariateNormal(state["y_mean"] + w.unsqueeze(-1), torch.eye(self.state_size)),
                             obs=y)
         self.t += 1
-
         return x, y
 
 
@@ -39,18 +40,17 @@ def __init__(self, state_size, plate_size):
         self.state_size = state_size
         self.plate_size = plate_size
 
-    def init(self):
+    def init(self, state):
         self.t = 0
-        self.x_mean = pyro.sample("x_mean",
-                                  dist.Normal(0., 2.))
-        self.y_mean = pyro.sample("y_mean",
-                                  dist.MultivariateNormal(torch.zeros(self.state_size), 2.*torch.eye(self.state_size)))
+        pyro.sample("x_mean", dist.Normal(0., 2.))
+        pyro.sample("y_mean",
+                    dist.MultivariateNormal(torch.zeros(self.state_size),
+                                            2.*torch.eye(self.state_size)))
 
-    def step(self, x=None, y=None):
+    def step(self, state, x=None, y=None):
         v = pyro.sample("v_{}".format(self.t), dist.Normal(0., 2.))
         with pyro.plate("plate", self.plate_size):
             pyro.sample("w_{}".format(self.t), dist.Normal(v, 2.))
-
         self.t += 1
 
 
@@ -66,13 +66,15 @@ def test_smoke(max_plate_nesting, state_size, plate_size, num_steps):
 
     true_model = SmokeModel(state_size, plate_size)
 
-    true_model.init()
-    truth = [true_model.step() for t in range(num_steps)]
+    state = {}
+    true_model.init(state)
+    truth = [true_model.step(state) for t in range(num_steps)]
 
     smc.init()
-    for xy in truth:
-        smc.step(*xy)
-    smc.get_values_and_log_weights()
+    assert set(smc.state) == {"x_mean", "y_mean"}
+    for x, y in truth:
+        smc.step(x, y)
+    assert set(smc.state) == {"x_mean", "y_mean"}
     smc.get_empirical()
 
 
@@ -85,20 +87,23 @@ def __init__(self):
         self.sigma_z = torch.tensor(1.)
         self.sigma_y = torch.tensor(1.)
 
-    def init(self):
+    def init(self, state):
         self.t = 0
-        self.z = torch.tensor([1., 0.])
-        self.y = None
+        state["z"] = pyro.sample("z_init",
+                                 dist.Delta(torch.tensor([1., 0.]), event_dim=1))
 
-    def step(self, y=None):
+    def step(self, state, y=None):
         self.t += 1
-        self.z = pyro.sample("z_{}".format(self.t),
-                             dist.Normal(self.z.matmul(self.A), self.B*self.sigma_z).to_event(1))
-        self.y = pyro.sample("y_{}".format(self.t),
-                             dist.Normal(self.z[..., 0], self.sigma_y),
-                             obs=y)
+        state["z"] = pyro.sample("z_{}".format(self.t),
+                                 dist.Normal(state["z"].matmul(self.A),
+                                             self.B*self.sigma_z).to_event(1))
+        y = pyro.sample("y_{}".format(self.t),
+                        dist.Normal(state["z"][..., 0], self.sigma_y),
+                        obs=y)
+
+        state["z_{}".format(self.t)] = state["z"]  # saved for testing
 
-        return self.z, self.y
+        return state["z"], y
 
 
 class HarmonicGuide:
@@ -106,26 +111,28 @@ class HarmonicGuide:
     def __init__(self):
         self.model = HarmonicModel()
 
-    def init(self):
+    def init(self, state):
         self.t = 0
-        self.z = torch.tensor([1., 0.])
+        pyro.sample("z_init", dist.Delta(torch.tensor([1., 0.]), event_dim=1))
 
-    def step(self, y=None):
+    def step(self, state, y=None):
         self.t += 1
 
         # Proposal distribution
-        self.z = pyro.sample("z_{}".format(self.t),
-                             dist.Normal(self.z.matmul(self.model.A), torch.tensor([2., 2.])).to_event(1))
+        pyro.sample("z_{}".format(self.t),
+                    dist.Normal(state["z"].matmul(self.model.A),
+                                torch.tensor([2., 2.])).to_event(1))
 
 
 def generate_data():
     model = HarmonicModel()
 
-    model.init()
+    state = {}
+    model.init(state)
     zs = [torch.tensor([1., 0.])]
     ys = [None]
     for t in range(50):
-        z, y = model.step()
+        z, y = model.step(state)
         zs.append(z)
         ys.append(y)
 
@@ -136,9 +143,10 @@ def score_latent(zs, ys):
     model = HarmonicModel()
     with poutine.trace() as trace:
         with poutine.condition(data={"z_{}".format(t): z for t, z in enumerate(zs)}):
-            model.init()
+            state = {}
+            model.init(state)
             for y in ys[1:]:
-                model.step(y)
+                model.step(state, y)
 
     return trace.trace.log_prob_sum()
 
@@ -155,13 +163,66 @@ def test_likelihood_ratio():
     smc.init()
     for y in ys_true[1:]:
         smc.step(y)
-    values, logweights = smc.get_values_and_log_weights()
-    i = logweights.max(0)[1]
-    values = {k: v[i] for k, v in values.items()}
+    i = smc.state._log_weights.max(0)[1]
+    values = {k: v[i] for k, v in smc.state.items()}
 
     zs_pred = [torch.tensor([1., 0.])]
-    zs_pred += [values.get("z_{}".format(t)) for t in range(1, 51)]
+    zs_pred += [values["z_{}".format(t)] for t in range(1, 51)]
 
     assert(score_latent(zs_true, ys_true) > score_latent(zs, ys_true))
     assert(score_latent(zs_pred, ys_true) > score_latent(zs_pred, ys))
     assert(score_latent(zs_pred, ys_true) > score_latent(zs, ys_true))
+
+
+def test_gaussian_filter():
+    dim = 4
+    init_dist = dist.MultivariateNormal(torch.zeros(dim), scale_tril=torch.eye(dim) * 10)
+    trans_mat = torch.eye(dim)
+    trans_dist = dist.MultivariateNormal(torch.zeros(dim), scale_tril=torch.eye(dim))
+    obs_mat = torch.eye(dim)
+    obs_dist = dist.MultivariateNormal(torch.zeros(dim), scale_tril=torch.eye(dim) * 2)
+    hmm = dist.GaussianHMM(init_dist, trans_mat, trans_dist, obs_mat, obs_dist)
+
+    class Model:
+        def init(self, state):
+            state["z"] = pyro.sample("z_init", init_dist)
+            self.t = 0
+
+        def step(self, state, datum=None):
+            state["z"] = pyro.sample("z_{}".format(self.t),
+                                     dist.MultivariateNormal(state["z"], scale_tril=trans_dist.scale_tril))
+            datum = pyro.sample("obs_{}".format(self.t),
+                                dist.MultivariateNormal(state["z"], scale_tril=obs_dist.scale_tril),
+                                obs=datum)
+            self.t += 1
+            return datum
+
+    class Guide:
+        def init(self, state):
+            pyro.sample("z_init", init_dist)
+            self.t = 0
+
+        def step(self, state, datum):
+            pyro.sample("z_{}".format(self.t),
+                        dist.MultivariateNormal(state["z"], scale_tril=trans_dist.scale_tril * 2))
+            self.t += 1
+
+    # Generate data.
+    num_steps = 20
+    model = Model()
+    state = {}
+    model.init(state)
+    data = torch.stack([model.step(state) for _ in range(num_steps)])
+
+    # Perform inference.
+    model = Model()
+    guide = Guide()
+    smc = SMCFilter(model, guide, num_particles=1000, max_plate_nesting=0)
+    smc.init()
+    for t, datum in enumerate(data):
+        smc.step(datum)
+        expected = hmm.filter(data[:1+t])
+        actual = smc.get_empirical()["z"]
+        assert_close(actual.variance ** 0.5, expected.variance ** 0.5, atol=0.1, rtol=0.5)
+        sigma = actual.variance.max().item() ** 0.5
+        assert_close(actual.mean, expected.mean, atol=3 * sigma)