code: def generate_data(args):
      extended_data = [None] * (args.duration + args.forecast)
      model = Model(args, extended_data)
      logging.info("Simulating from a {}".format(type(model).__name__))
      for attempt in range(100):
          samples = model.generate(
              {
                  "R0": args.basic_reproduction_number,
                  "rho": args.response_rate,
                  "k": args.concentration,
                  "od": args.overdispersion,
              }
          )
          obs = samples["obs"][: args.duration]
          new_I = samples.get("S2I", samples.get("E2I"))
  
          obs_sum = int(obs.sum())
          new_I_sum = int(new_I[: args.duration].sum())
          assert 0 <= args.min_obs_portion < args.max_obs_portion <= 1
          min_obs = int(
              math.ceil(args.min_obs_portion * args.population)
          )
          max_obs = int(
              math.floor(args.max_obs_portion * args.population)
          )
          if min_obs <= obs_sum <= max_obs:
              logging.info(
                  "Observed {:d}/{:d} infections:\n{}".format(
                      obs_sum,
                      new_I_sum,
                      " ".join(str(int(x)) for x in obs),
                  )
              )
              return {"new_I": new_I, "obs": obs}
  
      if obs_sum < min_obs:
          raise ValueError(
              "Failed to generate >={} observations. "
              "Try decreasing --min-obs-portion (currently {}).".format(
                  min_obs, args.min_obs_portion
              )
          )
      else:
          raise ValueError(
              "Failed to generate <={} observations. "
              "Try increasing --max-obs-portion (currently {}).".format(
                  max_obs, args.max_obs_portion
              )
          )
  
file_path: pyro/examples/contrib/epidemiology/sir.py
imports: import argparse
  import logging
  import math
  import torch
  from torch.distributions import biject_to, constraints
  import pyro
  from pyro.contrib.epidemiology.models import HeterogeneousSIRModel, OverdispersedSEIRModel, OverdispersedSIRModel, SimpleSEIRModel, SimpleSIRModel, SuperspreadingSEIRModel, SuperspreadingSIRModel
label_desc: Use image node_id 2 for calling a global function with example usage: generate_data(args) and returns: dict