Support HaarReparam for non-compartmental variables (#2523)

* Attempt to apply HaarReparam to non-compartmental time series * Factor _sample_auxiliary() out of three models * Fix tests * Remove debug statement * Fix docstring * Bump pytest version * Update pytest marks * Fix typo
pyro-ppl · Jun 13, 2020 · f7a5677 · f7a5677
1 parent d18fec8
commit f7a5677
Show file tree

Hide file tree

Showing 9 changed files with 322 additions and 123 deletions.
diff --git a/docs/source/contrib.epidemiology.rst b/docs/source/contrib.epidemiology.rst
@@ -1,5 +1,5 @@
-Epidemiological Models
-======================
+Epidemiology
+============
 .. automodule:: pyro.contrib.epidemiology
 
 .. warning:: Code in ``pyro.contrib.epidemiology`` is under development.
@@ -19,6 +19,7 @@ Base Compartmental Model
 ------------------------
 .. automodule:: pyro.contrib.epidemiology.compartmental
     :members:
+    :show-inheritance:
     :member-order: bysource
 
 Example Models
@@ -29,8 +30,11 @@ Distributions
 -------------
 .. automodule:: pyro.contrib.epidemiology.distributions
     :members:
+    :show-inheritance:
     :member-order: bysource
 
 .. autoclass:: pyro.distributions.CoalescentRateLikelihood
     :members:
+    :show-inheritance:
+    :member-order: bysource
     :special-members: __call__
diff --git a/examples/contrib/epidemiology/regional.py b/examples/contrib/epidemiology/regional.py
@@ -7,7 +7,7 @@
 import torch
 
 import pyro
-from pyro.contrib.epidemiology import RegionalSIRModel
+from pyro.contrib.epidemiology.models import RegionalSIRModel
 
 logging.basicConfig(format='%(message)s', level=logging.INFO)
 

diff --git a/examples/contrib/epidemiology/sir.py b/examples/contrib/epidemiology/sir.py
@@ -13,8 +13,9 @@
 from torch.distributions import biject_to, constraints
 
 import pyro
-from pyro.contrib.epidemiology import (HeterogeneousSIRModel, OverdispersedSEIRModel, OverdispersedSIRModel,
-                                       SimpleSEIRModel, SimpleSIRModel, SuperspreadingSEIRModel, SuperspreadingSIRModel)
+from pyro.contrib.epidemiology.models import (HeterogeneousSIRModel, OverdispersedSEIRModel, OverdispersedSIRModel,
+                                              SimpleSEIRModel, SimpleSIRModel, SuperspreadingSEIRModel,
+                                              SuperspreadingSIRModel)
 
 logging.basicConfig(format='%(message)s', level=logging.INFO)
 

diff --git a/pyro/contrib/epidemiology/__init__.py b/pyro/contrib/epidemiology/__init__.py
@@ -3,22 +3,9 @@
 
 from .compartmental import CompartmentalModel
 from .distributions import beta_binomial_dist, binomial_dist, infection_dist
-from .models import (HeterogeneousSIRModel, OverdispersedSEIRModel, OverdispersedSIRModel, RegionalSIRModel,
-                     SimpleSEIRModel, SimpleSIRModel, SparseSIRModel, SuperspreadingSEIRModel, SuperspreadingSIRModel,
-                     UnknownStartSIRModel)
 
 __all__ = [
     "CompartmentalModel",
-    "HeterogeneousSIRModel",
-    "OverdispersedSEIRModel",
-    "OverdispersedSIRModel",
-    "RegionalSIRModel",
-    "SimpleSEIRModel",
-    "SimpleSIRModel",
-    "SparseSIRModel",
-    "SuperspreadingSEIRModel",
-    "SuperspreadingSIRModel",
-    "UnknownStartSIRModel",
     "beta_binomial_dist",
     "binomial_dist",
     "infection_dist",