code: def predict(args, model, truth):
      samples = model.predict(forecast=args.forecast)
      S2I = samples["S2I"]
      median = S2I.median(dim=0).values
      lines = [
          "Median prediction of new infections (starting on day 0):"
      ]
      for r in range(args.num_regions):
          lines.append(
              "Region {}: {}".format(
                  r, " ".join(map(str, map(int, median[:, r])))
              )
          )
      logging.info("\n".join(lines))
  
      # Optionally plot the latent and forecasted series of new infections.
      if args.plot:
          import matplotlib.pyplot as plt
  
          fig, axes = plt.subplots(
              args.num_regions,
              sharex=True,
              figsize=(6, 1 + args.num_regions),
          )
          time = torch.arange(args.duration + args.forecast)
          p05 = S2I.kthvalue(
              int(round(0.5 + 0.05 * args.num_samples)), dim=0
          ).values
          p95 = S2I.kthvalue(
              int(round(0.5 + 0.95 * args.num_samples)), dim=0
          ).values
          for r, ax in enumerate(axes):
              ax.fill_between(
                  time,
                  p05[:, r],
                  p95[:, r],
                  color="red",
                  alpha=0.3,
                  label="90% CI",
              )
              ax.plot(time, median[:, r], "r-", label="median")
              ax.plot(
                  time[: args.duration],
                  model.data[:, r],
                  "k.",
                  label="observed",
              )
              ax.plot(time, truth[:, r], "k--", label="truth")
              ax.axvline(args.duration - 0.5, color="gray", lw=1)
              ax.set_xlim(0, len(time) - 1)
              ax.set_ylim(0, None)
          axes[0].set_title(
              "New infections among {} regions each of size {}".format(
                  args.num_regions, args.population
              )
          )
          axes[args.num_regions // 2].set_ylabel("inf./day")
          axes[-1].set_xlabel("day after first infection")
          axes[-1].legend(loc="upper left")
          plt.tight_layout()
          plt.subplots_adjust(hspace=0)
  
file_path: pyro/examples/contrib/epidemiology/regional.py
imports: import argparse
  import logging
  import torch
  import pyro
  from pyro.contrib.epidemiology.models import RegionalSIRModel
label_desc: Use image node_id 5 for calling a global function with example usage: predict(args, model, truth) without return types