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
  
      infection_rate = 2.3
      recovery_rate = 0.5
      population_size = 1000
      detection_rate = 0.5
  
      observed_infections = np.array(
          [
              15.0,
              41.0,
              104.0,
              209.0,
              269.0,
              226.0,
              160.0,
              97.0,
              69.0,
              34.0,
              19.0,
              20.0,
              13.0,
              5.0,
              6.0,
              3.0,
              2.0,
              1.0,
              1.0,
              1.0,
          ]
      )
  
      # Add batch dimensions and an event dimension.
      observed_infections = np.tile(
          observed_infections, (2, 3, 1)
      ).astype(np.float32)
      observed_infections = np.swapaxes(observed_infections, -1, 0)[
          ..., np.newaxis
      ]
  
      # The state `x` is defined by `[num_susceptible, num_infected]`.
      def transition_fn(x):
          p_s = 1.0 - tf.math.exp(
              -infection_rate * x[..., 1] / population_size
          )
          p_r = 1.0 - tf.math.exp(-recovery_rate)
          d_s = -x[..., 0] * p_s
          return mvn_diag.MultivariateNormalDiag(
              x + tf.stack([d_s, -d_s - x[..., 1] * p_r], axis=-1),
              scale_diag=tf.stack(
                  [
                      x[..., 0] * p_s * (1 - p_s),
                      x[..., 1] * p_r * (1 - p_r),
                  ],
                  axis=-1,
              ),
              validate_args=True,
          )
  
      def transition_jacobian_fn(x):
          shape_out = prefer_static.concat(
              [prefer_static.shape(x)[:-1], (2, 2)], axis=0
          )
          return tf.reshape(
              tf.stack(
                  [
                      1.0
                      - infection_rate * x[..., 1] / population_size,
                      -infection_rate * x[..., 1] / population_size,
                      infection_rate * x[..., 0] / population_size,
                      1
                      + infection_rate * x[..., 0] / population_size
                      - recovery_rate,
                  ],
                  axis=-1,
              ),
              shape=shape_out,
          )
  
      # We use the Binomial mean and covariance, assuming observed infections are
      # detected with `prob=detection_rate`.
      def observation_fn(x):
          return mvn_diag.MultivariateNormalDiag(
              detection_rate * x[..., 1:],
              scale_diag=x[..., 1:]
              * detection_rate
              * (1 - detection_rate),
              validate_args=True,
          )
  
      observation_jacobian_fn = lambda _: [[0.0, detection_rate]]
  
      # Start with an incorrect guess for susceptible and infected.
      x = np.array([population_size - 25, 25]).astype(np.float32)
      initial_state_prior = mvn_diag.MultivariateNormalDiag(
          x, scale_diag=0.1 * x, validate_args=True
      )
  
      results = self.evaluate(
          extended_kalman_filter.extended_kalman_filter(
              observations=observed_infections,
              initial_state_prior=initial_state_prior,
              transition_fn=transition_fn,
              observation_fn=observation_fn,
              transition_jacobian_fn=transition_jacobian_fn,
              observation_jacobian_fn=observation_jacobian_fn,
          )
      )
      self.assertAllEqual(results[0].shape, [20, 3, 2, 2])
      self.assertAllEqual(results[1].shape, [20, 3, 2, 2, 2])
  
file_path: probability/tensorflow_probability/python/experimental/sequential/extended_kalman_filter_test.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.experimental.sequential import extended_kalman_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 ExtendedKalmanFilterTest obj's underlying member method code with example usage: obj.test_epidemiological_model() and returns: mvn_diag, tf, mvn_diag