Add pairwise correlation plot to sir.py

examples/contrib/epidemiology/sir.py

@@ -99,10 +99,11 @@ def evaluate(args, samples):
         logging.info("{}: truth = {:0.3g}, estimate = {:0.3g} \u00B1 {:0.3g}"
                      .format(key, getattr(args, name), mean, std))
 
-    # Optionally plot histograms.
+    # Optionally plot histograms and pairwise correlations.
     if args.plot:
         import matplotlib.pyplot as plt
         import seaborn as sns
+
         fig, axes = plt.subplots(len(names), 1, figsize=(5, 2.5 * len(names)))
         axes[0].set_title("Posterior parameter estimates")
         for ax, (name, key) in zip(axes, names.items()):
@@ -114,6 +115,25 @@ def evaluate(args, samples):
             ax.legend(loc="best")
         plt.tight_layout()
 
+        covariates = [(name, samples[name]) for name in names.values()]
+        for i, aux in enumerate(samples["auxiliary"].unbind(-2)):
+            covariates.append(("aux[{},0]".format(i), aux[:, 0]))
+            covariates.append(("aux[{},-1]".format(i), aux[:, -1]))
+        N = len(covariates)
+        fig, axes = plt.subplots(N, N, figsize=(8, 8), sharex="col", sharey="row")
+        for i in range(N):
+            axes[i][0].set_ylabel(covariates[i][0])
+            axes[0][i].set_xlabel(covariates[i][0])
+            axes[0][i].xaxis.set_label_position("top")
+            for j in range(N):
+                ax = axes[i][j]
+                ax.set_xticks(())
+                ax.set_yticks(())
+                ax.scatter(covariates[j][1], -covariates[i][1],
+                           lw=0, color="darkblue", alpha=0.3)
+        plt.tight_layout()
+        plt.subplots_adjust(wspace=0, hspace=0)
+
 
 def predict(args, model, truth):
     samples = model.predict(forecast=args.forecast)

pyro/contrib/epidemiology/seir.py

@@ -52,7 +52,7 @@ def global_model(self):
         tau_e = self.incubation_time
         tau_i = self.recovery_time
         R0 = pyro.sample("R0", dist.LogNormal(0., 1.))
-        rho = pyro.sample("rho", dist.Uniform(0, 1))
+        rho = pyro.sample("rho", dist.Beta(2, 2))
         return R0, tau_e, tau_i, rho
 
     def initialize(self, params):
@@ -181,7 +181,7 @@ def global_model(self):
         tau_i = self.recovery_time
         R0 = pyro.sample("R0", dist.LogNormal(0., 1.))
         k = pyro.sample("k", dist.Exponential(1.))
-        rho = pyro.sample("rho", dist.Uniform(0, 1))
+        rho = pyro.sample("rho", dist.Beta(2, 2))
         return R0, k, tau_e, tau_i, rho
 
     def initialize(self, params):

pyro/contrib/epidemiology/sir.py

@@ -48,7 +48,7 @@ def __init__(self, population, recovery_time, data):
     def global_model(self):
         tau = self.recovery_time
         R0 = pyro.sample("R0", dist.LogNormal(0., 1.))
-        rho = pyro.sample("rho", dist.Uniform(0, 1))
+        rho = pyro.sample("rho", dist.Beta(2, 2))
         return R0, tau, rho
 
     def initialize(self, params):
@@ -162,7 +162,7 @@ def global_model(self):
         tau = self.recovery_time
         R0 = pyro.sample("R0", dist.LogNormal(0., 1.))
         k = pyro.sample("k", dist.Exponential(1.))
-        rho = pyro.sample("rho", dist.Uniform(0, 1))
+        rho = pyro.sample("rho", dist.Beta(2, 2))
         return R0, k, tau, rho
 
     def initialize(self, params):
@@ -263,7 +263,7 @@ def __init__(self, population, recovery_time, data, mask):
     def global_model(self):
         tau = self.recovery_time
         R0 = pyro.sample("R0", dist.LogNormal(0., 1.))
-        rho = pyro.sample("rho", dist.Uniform(0, 1))
+        rho = pyro.sample("rho", dist.Beta(2, 2))
         return R0, tau, rho
 
     def initialize(self, params):
@@ -390,8 +390,8 @@ def global_model(self):
         # Assume two different response rates: rho0 before any observations
         # were made (in pre_obs_window), followed by a higher response rate rho1
         # after observations were made (in post_obs_window).
-        rho0 = pyro.sample("rho0", dist.Uniform(0, 1))
-        rho1 = pyro.sample("rho1", dist.Uniform(0, 1))
+        rho0 = pyro.sample("rho0", dist.Beta(2, 2))
+        rho1 = pyro.sample("rho1", dist.Beta(2, 2))
         # Whereas each of rho0,rho1 are scalars (possibly batched over samples),
         # we construct a time series rho with an extra time dim on the right.
         rho = torch.cat([