Clean up contrib.epidemiology docs

docs/source/contrib.epidemiology.rst

@@ -24,8 +24,6 @@ Base Compartmental Model
 Example Models
 --------------
 .. automodule:: pyro.contrib.epidemiology.models
-    :members:
-    :member-order: bysource
 
 Distributions
 -------------

pyro/contrib/epidemiology/compartmental.py

@@ -151,53 +151,6 @@ def _clear_plates(self):
     series = ()
     full_mass = False
 
-    @torch.no_grad()
-    @set_approx_log_prob_tol(0.1)
-    @set_approx_sample_thresh(100)  # This is robust to gross approximation.
-    def heuristic(self, num_particles=1024, ess_threshold=0.5, retries=10):
-        """
-        Finds an initial feasible guess of all latent variables, consistent
-        with observed data. This is needed because not all hypotheses are
-        feasible and HMC needs to start at a feasible solution to progress.
-
-        The default implementation attempts to find a feasible state using
-        :class:`~pyro.infer.smcfilter.SMCFilter` with proprosals from the
-        prior.  However this method may be overridden in cases where SMC
-        performs poorly e.g. in high-dimensional models.
-
-        :param int num_particles: Number of particles used for SMC.
-        :param float ess_threshold: Effective sample size threshold for SMC.
-        :returns: A dictionary mapping sample site name to tensor value.
-        :rtype: dict
-        """
-        # Run SMC.
-        model = _SMCModel(self)
-        guide = _SMCGuide(self)
-        for attempt in range(1, 1 + retries):
-            smc = SMCFilter(model, guide, num_particles=num_particles,
-                            ess_threshold=ess_threshold,
-                            max_plate_nesting=self.max_plate_nesting)
-            try:
-                smc.init()
-                for t in range(1, self.duration):
-                    smc.step()
-                break
-            except SMCFailed as e:
-                if attempt == retries:
-                    raise
-                logger.info("{}. Retrying...".format(e))
-                continue
-
-        # Select the most probable hypothesis.
-        i = int(smc.state._log_weights.max(0).indices)
-        init = {key: value[i] for key, value in smc.state.items()}
-
-        # Fill in sample site values.
-        init = self.generate(init)
-        aux = torch.stack([init[name] for name in self.compartments], dim=0)
-        init["auxiliary"] = clamp(aux, min=0.5, max=self.population - 0.5)
-        return init
-
     def global_model(self):
         """
         Samples and returns any global parameters.
@@ -462,6 +415,53 @@ def predict(self, forecast=0):
         self._concat_series(samples, forecast)
         return samples
 
+    @torch.no_grad()
+    @set_approx_log_prob_tol(0.1)
+    @set_approx_sample_thresh(100)  # This is robust to gross approximation.
+    def heuristic(self, num_particles=1024, ess_threshold=0.5, retries=10):
+        """
+        Finds an initial feasible guess of all latent variables, consistent
+        with observed data. This is needed because not all hypotheses are
+        feasible and HMC needs to start at a feasible solution to progress.
+
+        The default implementation attempts to find a feasible state using
+        :class:`~pyro.infer.smcfilter.SMCFilter` with proprosals from the
+        prior.  However this method may be overridden in cases where SMC
+        performs poorly e.g. in high-dimensional models.
+
+        :param int num_particles: Number of particles used for SMC.
+        :param float ess_threshold: Effective sample size threshold for SMC.
+        :returns: A dictionary mapping sample site name to tensor value.
+        :rtype: dict
+        """
+        # Run SMC.
+        model = _SMCModel(self)
+        guide = _SMCGuide(self)
+        for attempt in range(1, 1 + retries):
+            smc = SMCFilter(model, guide, num_particles=num_particles,
+                            ess_threshold=ess_threshold,
+                            max_plate_nesting=self.max_plate_nesting)
+            try:
+                smc.init()
+                for t in range(1, self.duration):
+                    smc.step()
+                break
+            except SMCFailed as e:
+                if attempt == retries:
+                    raise
+                logger.info("{}. Retrying...".format(e))
+                continue
+
+        # Select the most probable hypothesis.
+        i = int(smc.state._log_weights.max(0).indices)
+        init = {key: value[i] for key, value in smc.state.items()}
+
+        # Fill in sample site values.
+        init = self.generate(init)
+        aux = torch.stack([init[name] for name in self.compartments], dim=0)
+        init["auxiliary"] = clamp(aux, min=0.5, max=self.population - 0.5)
+        return init
+
     # Internal helpers ########################################
 
     def _concat_series(self, samples, forecast=0):

pyro/contrib/epidemiology/models.py

@@ -1,6 +1,8 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
+import re
+
 import torch
 from torch.nn.functional import pad
 
@@ -1101,3 +1103,20 @@ def transition_bwd(self, params, prev, curr, t):
             pyro.sample("obs_{}".format(t),
                         dist.ExtendedBinomial(S2I, rho),
                         obs=self.data[t])
+
+
+# Create sphinx documentation.
+__all__ = []
+for _name, _Model in list(locals().items()):
+    if isinstance(_Model, type) and issubclass(_Model, CompartmentalModel):
+        if _Model is not CompartmentalModel:
+            __all__.append(_name)
+__all__.sort(key=lambda name, vals=locals(): vals[name].__init__.__code__.co_firstlineno)
+__doc__ = "\n\n".join([
+    """
+    {}
+    ----------------------------------------------------------------
+    .. autoclass:: pyro.contrib.epidemiology.models.{}
+    """.format(re.sub("([A-Z][a-z]+)", r"\1 ", _name[:-5]), _name)
+    for _name in __all__
+])