code: def test_epidemiological_model(self):
      # A toy, discrete version of an SIR (Susceptible, Infected, Recovered)
      # model (https://en.wikipedia.org/wiki/Compartmental_models_in_epidemiology)
  
      population_size = 1000
      infection_rate = tf.convert_to_tensor(1.1)
      infectious_period = tf.convert_to_tensor(8.0)
  
      initial_state_prior = jdn.JointDistributionNamed(
          {
              "susceptible": deterministic.Deterministic(999.0),
              "infected": deterministic.Deterministic(1.0),
              "new_infections": deterministic.Deterministic(1.0),
              "new_recoveries": deterministic.Deterministic(0.0),
          }
      )
  
      # Dynamics model: new infections and recoveries are given by the SIR
      # model with Poisson noise.
      def infection_dynamics(_, previous_state):
          new_infections = poisson.Poisson(
              infection_rate
              * previous_state["infected"]
              * previous_state["susceptible"]
              / population_size
          )
          new_recoveries = poisson.Poisson(
              previous_state["infected"] / infectious_period
          )
  
          def susceptible(new_infections):
              return deterministic.Deterministic(
                  ps.maximum(
                      0.0,
                      previous_state["susceptible"] - new_infections,
                  )
              )
  
          def infected(new_infections, new_recoveries):
              return deterministic.Deterministic(
                  ps.maximum(
                      0.0,
                      previous_state["infected"]
                      + new_infections
                      - new_recoveries,
                  )
              )
  
          return jdn.JointDistributionNamed(
              {
                  "new_infections": new_infections,
                  "new_recoveries": new_recoveries,
                  "susceptible": susceptible,
                  "infected": infected,
              }
          )
  
      # Observation model: each day we detect new cases, noisily.
      def infection_observations(_, state):
          return poisson.Poisson(state["infected"])
  
      # pylint: disable=bad-whitespace
      observations = tf.convert_to_tensor(
          [
              0.0,
              4.0,
              1.0,
              5.0,
              23.0,
              27.0,
              75.0,
              127.0,
              248.0,
              384.0,
              540.0,
              683.0,
              714.0,
              611.0,
              561.0,
              493.0,
              385.0,
              348.0,
              300.0,
              277.0,
              249.0,
              219.0,
              216.0,
              174.0,
              132.0,
              122.0,
              115.0,
              99.0,
              76.0,
              84.0,
              77.0,
              56.0,
              42.0,
              56.0,
              46.0,
              38.0,
              34.0,
              44.0,
              25.0,
              27.0,
          ]
      )
      # pylint: enable=bad-whitespace
  
      trajectories, _ = self.evaluate(
          particle_filter.infer_trajectories(
              observations=observations,
              initial_state_prior=initial_state_prior,
              transition_fn=infection_dynamics,
              observation_fn=infection_observations,
              num_particles=100,
              seed=test_util.test_seed(),
          )
      )
  
      # The susceptible population should decrease over time.
      self.assertAllLessEqual(
          trajectories["susceptible"][1:, ...]
          - trajectories["susceptible"][:-1, ...],
          0.0,
      )
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import linear_gaussian_ssm
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.experimental.mcmc import particle_filter
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 5 for calling the _ParticleFilterTest obj's underlying member method code with example usage: obj.test_epidemiological_model() and returns: jdn, poisson, deterministic, deterministic