class_docstr: Abstract base class for discrete-time discrete-value stochastic
  compartmental models.
  
  Derived classes must implement methods :meth:`initialize` and
  :meth:`transition`. Derived classes may optionally implement
  :meth:`global_model`, :meth:`compute_flows`, and :meth:`heuristic`.
  
  Example usage::
  
      # First implement a concrete derived class.
      class MyModel(CompartmentalModel):
          def __init__(self, ...): ...
          def global_model(self): ...
          def initialize(self, params): ...
          def transition(self, params, state, t): ...
  
      # Run inference to fit the model to data.
      model = MyModel(...)
      model.fit_svi(num_samples=100)  # or .fit_mcmc(...)
      R0 = model.samples["R0"]  # An example parameter.
      print("R0 = {:0.3g} ± {:0.3g}".format(R0.mean(), R0.std()))
  
      # Predict latent variables.
      samples = model.predict()
  
      # Forecast forward.
      samples = model.predict(forecast=30)
  
      # You can assess future interventions (applied after ``duration``) by
      # storing them as attributes that are read by your derived methods.
      model.my_intervention = False
      samples1 = model.predict(forecast=30)
      model.my_intervention = True
      samples2 = model.predict(forecast=30)
      effect = samples2["my_result"].mean() - samples1["my_result"].mean()
      print("average effect = {:0.3g}".format(effect))
  
  An example workflow is to use cheaper approximate inference while finding
  good model structure and priors, then move to more accurate but more
  expensive inference once the model is plausible.
  
  1.  Start with ``.fit_svi(guide_rank=0, num_steps=2000)`` for cheap
      inference while you search for a good model.
  2.  Additionally infer long-range correlations by moving to a low-rank
      multivariate normal guide via ``.fit_svi(guide_rank=None,
      num_steps=5000)``.
  3.  Optionally additionally infer non-Gaussian posterior by moving to the
      more expensive (but still approximate via moment matching)
      ``.fit_mcmc(num_quant_bins=1, num_samples=10000, num_chains=2)``.
  4.  Optionally improve fit around small counts by moving the the more
      expensive enumeration-based algorithm ``.fit_mcmc(num_quant_bins=4,
      num_samples=10000, num_chains=2)`` (GPU recommended).
  
  :ivar dict samples: Dictionary of posterior samples.
  :param list compartments: A list of strings of compartment names.
  :param int duration: The number of discrete time steps in this model.
  :param population: Either the total population of a single-region model or
      a tensor of each region's population in a regional model.
  :type population: int or torch.Tensor
  :param tuple approximate: Names of compartments for which pointwise
      approximations should be provided in :meth:`transition`, e.g. if you
      specify ``approximate=("I")`` then the ``state["I_approx"]`` will be a
      continuous-valued non-enumerated point estimate of ``state["I"]``.
      Approximations are useful to reduce computational cost. Approximations
      are continuous-valued with support ``(-0.5, population + 0.5)``.
class_docstr_tok: ["Abstract","base","class","for","discrete-time","discrete-value","stochastic","compartmental","models",".","Derived","classes","must","implement","methods",":","meth",":","`","initialize","`","and",":","meth",":","`","transition","`",".","Derived","classes","may","optionally","implement",":","meth",":","`","global_model","`",",",":","meth",":","`","compute_flows","`",",","and",":","meth",":","`","heuristic","`",".","Example","usage",":",":","#","First","implement","a","concrete","derived","class",".","class","MyModel","(","CompartmentalModel",")",":","def","__init__","(","self",",","...",")",":","...","def","global_model","(","self",")",":","...","def","initialize","(","self",",","params",")",":","...","def","transition","(","self",",","params",",","state",",","t",")",":","...","#","Run","inference","to","fit","the","model","to","data",".","model","=","MyModel","(","...",")","model.fit_svi","(","num_samples=100",")","#","or",".fit_mcmc","(","...",")","R0","=","model.samples","[","``","R0","''","]","#","An","example","parameter",".","print","(","``","R0","=","{",":0.3g","}","\u00b1","{",":0.3g","}","''",".format","(","R0.mean","(",")",",","R0.std","(",")",")",")","#","Predict","latent","variables",".","samples","=","model.predict","(",")","#","Forecast","forward",".","samples","=","model.predict","(","forecast=30",")","#","You","can","assess","future","interventions","(","applied","after","``","duration","``",")","by","#","storing","them","as","attributes","that","are","read","by","your","derived","methods",".","model.my_intervention","=","False","samples1","=","model.predict","(","forecast=30",")","model.my_intervention","=","True","samples2","=","model.predict","(","forecast=30",")","effect","=","samples2","[","``","my_result","''","]",".mean","(",")","-","samples1","[","``","my_result","''","]",".mean","(",")","print","(","``","average","effect","=","{",":0.3g","}","''",".format","(","effect",")",")","An","example","workflow","is","to","use","cheaper","approximate","inference","while","finding","good","model","structure","and","priors",",","then","move","to","more","accurate","but","more","expensive","inference","once","the","model","is","plausible",".","1",".","Start","with","``",".fit_svi","(","guide_rank=0",",","num_steps=2000",")","``","for","cheap","inference","while","you","search","for","a","good","model",".","2",".","Additionally","infer","long-range","correlations","by","moving","to","a","low-rank","multivariate","normal","guide","via","``",".fit_svi","(","guide_rank=None",",","num_steps=5000",")","``",".","3",".","Optionally","additionally","infer","non-Gaussian","posterior","by","moving","to","the","more","expensive","(","but","still","approximate","via","moment","matching",")","``",".fit_mcmc","(","num_quant_bins=1",",","num_samples=10000",",","num_chains=2",")","``",".","4",".","Optionally","improve","fit","around","small","counts","by","moving","the","the","more","expensive","enumeration-based","algorithm","``",".fit_mcmc","(","num_quant_bins=4",",","num_samples=10000",",","num_chains=2",")","``","(","GPU","recommended",")",".",":","ivar","dict","samples",":","Dictionary","of","posterior","samples",".",":","param","list","compartments",":","A","list","of","strings","of","compartment","names",".",":","param","int","duration",":","The","number","of","discrete","time","steps","in","this","model",".",":","param","population",":","Either","the","total","population","of","a","single-region","model","or","a","tensor","of","each","region","'s","population","in","a","regional","model",".",":","type","population",":","int","or","torch.Tensor",":","param","tuple","approximate",":","Names","of","compartments","for","which","pointwise","approximations","should","be","provided","in",":","meth",":","`","transition","`",",","e.g",".","if","you","specify","``","approximate=","(","``","I","''",")","``","then","the","``","state","[","``","I_approx","''","]","``","will","be","a","continuous-valued","non-enumerated","point","estimate","of","``","state","[","``","I","''","]","``",".","Approximations","are","useful","to","reduce","computational","cost",".","Approximations","are","continuous-valued","with","support","``","(","-0.5",",","population","+","0.5",")","``","."]
docstr: Forward generative model used for simulation and forecasting.
docstr_tok: ["Forward","generative","model","used","for","simulation","and","forecasting","."]
code: def _generative_model(self, forecast=0):
      """
      Forward generative model used for simulation and forecasting.
      """
      # Sample global parameters.
      params = self.global_model()
  
      # Sample initial values.
      state = self.initialize(params)
      state = {
          k: v
          if isinstance(v, torch.Tensor)
          else torch.tensor(float(v))
          for k, v in state.items()
      }
  
      # Sequentially transition.
      for t in range(self.duration + forecast):
          for name in self.approximate:
              state[name + "_approx"] = state[name]
          self.transition(params, state, t)
          with self.region_plate:
              for name in self.compartments:
                  pyro.deterministic(
                      "{}_{}".format(name, t), state[name], event_dim=0
                  )
  
      self._clear_plates()
  
imports: import functools
  import logging
  import operator
  import re
  import warnings
  from abc import ABC, abstractmethod
  from collections import OrderedDict
  from contextlib import ExitStack, contextmanager
  from functools import reduce
  from timeit import default_timer
  import torch
  from torch.distributions import biject_to, constraints
  from torch.distributions.utils import lazy_property
  import pyro.distributions
  import pyro.distributions.hmm
  import pyro.poutine
  from pyro.distributions.transforms import HaarTransform
  from pyro.infer import MCMC, NUTS, SVI, JitTrace_ELBO, SMCFilter, Trace_ELBO, infer_discrete
  from pyro.infer.autoguide import AutoLowRankMultivariateNormal, AutoMultivariateNormal, AutoNormal, init_to_generated, init_to_value
  from pyro.infer.mcmc import ArrowheadMassMatrix
  from pyro.infer.reparam import HaarReparam, SplitReparam
  from pyro.infer.smcfilter import SMCFailed
  from pyro.infer.util import is_validation_enabled
  from pyro.optim import ClippedAdam
  from pyro.poutine.util import site_is_factor, site_is_subsample
  from pyro.util import warn_if_nan
  from .distributions import set_approx_log_prob_tol, set_approx_sample_thresh, set_relaxed_distributions
  from .util import align_samples, cat2, clamp, quantize, quantize_enumerate
label_desc: Use image node_id 22 for calling the CompartmentalModel obj's underlying member method code with example usage: obj._generative_model(forecast) without return types