Add pearson correlation plot to examples/.../sir.py
@@ -10,6 +10,7 @@ import logging
import math
import torch
+from torch.distributions import biject_to, constraints
import pyro
from pyro.contrib.epidemiology import OverdispersedSEIRModel, OverdispersedSIRModel, SimpleSEIRModel, SimpleSIRModel
@@ -90,7 +91,7 @@ def infer(args, model):
return model.samples
-def evaluate(args, samples):
+def evaluate(args, model, samples):
# Print estimated values.
names = {"basic_reproduction_number": "R0",
"response_rate": "rho"}
@@ -107,6 +108,7 @@ def evaluate(args, samples):
import matplotlib.pyplot as plt
import seaborn as sns
+ # Plot individual histograms.
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()):
@@ -118,6 +120,7 @@ def evaluate(args, samples):
ax.legend(loc="best")
plt.tight_layout()
+ # Plot pairwise joint distributions for selected variables.
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]))
@@ -137,6 +140,36 @@ def evaluate(args, samples):
plt.tight_layout()
plt.subplots_adjust(wspace=0, hspace=0)
+ # Plot Pearson correlation for every pair of unconstrained variables.
+ def unconstrain(constraint, value):
+ value = biject_to(constraint).inv(value)
+ return value.reshape(args.num_samples, -1)
+
+ covariates = [
+ ("R1", unconstrain(constraints.positive, samples["R0"])),
+ ("rho", unconstrain(constraints.unit_interval, samples["rho"]))]
+ if "k" in samples:
+ covariates.append(
+ ("k", unconstrain(constraints.positive, samples["k"])))
+ constraint = constraints.interval(-0.5, model.population + 0.5)
+ for name, aux in zip(model.compartments, samples["auxiliary"].unbind(-2)):
+ covariates.append((name, unconstrain(constraint, aux)))
+ x = torch.cat([v for _, v in covariates], dim=-1)
+ x -= x.mean(0)
+ x /= x.std(0)
+ x = x.t().matmul(x)
+ x /= args.num_samples
+ x.clamp_(min=-1, max=1)
+ plt.figure(figsize=(8, 8))
+ plt.imshow(x, cmap="bwr")
+ ticks = torch.tensor([0] + [v.size(-1) for _, v in covariates]).cumsum(0)
+ ticks = (ticks[1:] + ticks[:-1]) / 2
+ plt.yticks(ticks, [name for name, _ in covariates])
+ plt.xticks(())
+ plt.tick_params(length=0)
+ plt.title("Pearson correlation (unconstrained coordinates)")
+ plt.tight_layout()
+
def predict(args, model, truth):
samples = model.predict(forecast=args.forecast)
@@ -183,7 +216,7 @@ def main(args):
samples = infer(args, model)
# Evaluate fit.
- evaluate(args, samples)
+ evaluate(args, model, samples)
# Predict latent time series.
if args.forecast:
