code: def test_epidemiological_model_docstring_example(self):
      # Toy SIR model
      # (https://en.wikipedia.org/wiki/Compartmental_models_in_epidemiology)
      population_size = 1000
      parameter_prior = jdn.JointDistributionNamed(
          {
              "infection_rate": uniform.Uniform(low=0.0, high=3.0),
              "recovery_rate": uniform.Uniform(low=0.0, high=3.0),
          }
      )
  
      initial_state_prior_fn = lambda parameters: jdn.JointDistributionNamed(
          {  # pylint: disable=g-long-lambda
              "new_infections": poisson.Poisson(
                  parameters["infection_rate"]
              ),
              "new_recoveries": deterministic.Deterministic(
                  tf.broadcast_to(
                      0.0, ps.shape(parameters["recovery_rate"])
                  )
              ),
              "susceptible": (
                  lambda new_infections: deterministic.Deterministic(  # pylint: disable=g-long-lambda
                      population_size - new_infections
                  )
              ),
              "infected": (
                  lambda new_infections: deterministic.Deterministic(  # pylint: disable=unnecessary-lambda, g-long-lambda
                      new_infections
                  )
              ),
          }
      )
  
      # Dynamics model: new infections and recoveries are given by the SIR
      # model with Poisson noise.
      def parameterized_infection_dynamics(
          _, previous_state, parameters
      ):
          new_infections = poisson.Poisson(
              parameters["infection_rate"]
              * previous_state["infected"]
              * previous_state["susceptible"]
              / population_size
          )
          new_recoveries = poisson.Poisson(
              previous_state["infected"] * parameters["recovery_rate"]
          )
          return jdn.JointDistributionNamed(
              {
                  "new_infections": new_infections,
                  "new_recoveries": new_recoveries,
                  "susceptible": lambda new_infections: deterministic.Deterministic(  # pylint: disable=g-long-lambda
                      tf.maximum(
                          0.0,
                          previous_state["susceptible"]
                          - new_infections,
                      )
                  ),
                  "infected": lambda new_infections, new_recoveries: deterministic.Deterministic(  # pylint:disable=g-long-lambda
                      tf.maximum(
                          0.0,
                          (
                              previous_state["infected"]
                              + new_infections
                              - new_recoveries
                          ),
                      )
                  ),
              }
          )
  
      # Observation model: each day we detect new cases and recoveries, noisily.
      def parameterized_infection_observations(_, state, parameters):
          del parameters
          return jdn.JointDistributionNamed(
              {
                  "new_infections": poisson.Poisson(
                      state["new_infections"] + 0.1
                  ),
                  "new_recoveries": poisson.Poisson(
                      state["new_recoveries"] + 0.1
                  ),
              }
          )
  
      iterated_filter = iterated_filter_lib.IteratedFilter(
          parameter_prior=parameter_prior,
          parameterized_initial_state_prior_fn=initial_state_prior_fn,
          parameterized_transition_fn=parameterized_infection_dynamics,
          parameterized_observation_fn=parameterized_infection_observations,
      )
  
      observed_values = {
          "new_infections": tf.convert_to_tensor(
              [
                  2.0,
                  7.0,
                  14.0,
                  24.0,
                  45.0,
                  93.0,
                  160.0,
                  228.0,
                  252.0,
                  158.0,
                  17.0,
                  0.0,
                  0.0,
                  0.0,
                  0.0,
                  0.0,
                  0.0,
                  0.0,
                  0.0,
                  0.0,
                  0.0,
                  0.0,
                  0.0,
                  0.0,
                  0.0,
                  0.0,
              ]
          ),
          "new_recoveries": tf.convert_to_tensor(
              [
                  0.0,
                  0.0,
                  3.0,
                  4.0,
                  3.0,
                  8.0,
                  12.0,
                  31.0,
                  49.0,
                  73.0,
                  85.0,
                  65.0,
                  71.0,
                  58.0,
                  42.0,
                  65.0,
                  36.0,
                  31.0,
                  32.0,
                  27.0,
                  31.0,
                  20.0,
                  19.0,
                  19.0,
                  14.0,
                  27.0,
              ]
          ),
      }
      parameter_particles = self.evaluate(
          iterated_filter.estimate_parameters(
              observations=observed_values,
              num_iterations=20,
              num_particles=4096,
              initial_perturbation_scale=1.0,
              cooling_schedule=(
                  iterated_filter_lib.geometric_cooling_schedule(
                      0.001, k=20
                  )
              ),
              seed=test_util.test_seed(),
          )
      )
      parameter_estimates = tf.nest.map_structure(
          lambda x: np.mean(x[-1], axis=0), parameter_particles
      )
      # Check that we recovered the generating parameters.
      self.assertAllClose(
          parameter_estimates["infection_rate"], 1.2, atol=0.1
      )
      self.assertAllClose(
          parameter_estimates["recovery_rate"], 0.1, atol=0.05
      )
  
file_path: probability/tensorflow_probability/python/experimental/sequential/iterated_filter_test.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.experimental.sequential import iterated_filter
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 2 for calling the IteratedFilterTest obj's underlying member method code with example usage: obj.test_epidemiological_model_docstring_example() and returns: jdn, jdn