code: def ensemble_kalman_filter_update(
      state,
      observation,
      observation_fn,
      damping=1.0,
      low_rank_ensemble=False,
      seed=None,
      name=None,
  ):
      """Ensemble Kalman filter update step.
  
      The [Ensemble Kalman Filter](
      https://en.wikipedia.org/wiki/Ensemble_Kalman_filter) is a Monte Carlo
      version of the traditional Kalman Filter. See also [2]. It assumes the model
  
      ```
      X[t] ~ transition_fn(X[t-1])
      Y[t] ~ observation_fn(X[t])
      ```
  
      Given the ensemble `state.particles` sampled from `P(X[t] | Y[t-1], ...)`,
      this function assimilates obervation `Y[t]` to produce the updated ensemble
      sampled from `P(X[t] | Y[t], ...)`.
  
      Typically, with `G` some deterministic observation mapping,
      `observation_fn(X)` returns a normal distribution centered at `G(X)`.
  
      Args:
        state: Instance of `EnsembleKalmanFilterState`.
        observation: `Tensor` representing the observation for this timestep.
        observation_fn: callable returning an instance of
          `tfd.MultivariateNormalLinearOperator` along with an extra information
          to be returned in the `EnsembleKalmanFilterState`.
        damping: Floating-point `Tensor` representing how much to damp the
          update by. Used to mitigate filter divergence. Default value: 1.
        low_rank_ensemble: Whether to use a LinearOperatorLowRankUpdate (rather than
          a dense Tensor) to represent the observation covariance. The "low rank" is
          the ensemble size. This is useful only if (i) the ensemble size is much
          less than the number of observations, and (ii) the LinearOperator
          associated with the observation_fn has an efficient inverse
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
        name: Python `str` name for ops created by this method.
          Default value: `None` (i.e., `'ensemble_kalman_filter_update'`).
  
      Returns:
        next_state: `EnsembleKalmanFilterState` representing particles at next
          timestep, after applying Kalman update equations.
  
      #### References
  
      [1] Geir Evensen. Sequential data assimilation with a nonlinear
          quasi-geostrophic model using Monte Carlo methods to forecast error
          statistics. Journal of Geophysical Research, 1994.
  
      [2] Matthias Katzfuss, Jonathan R. Stroud & Christopher K. Wikle
          Understanding the Ensemble Kalman Filter.
          The Americal Statistician, 2016.
      """
  
      with tf.name_scope(name or "ensemble_kalman_filter_update"):
          # In the example below, we let
          #  Y be the real observations, so Y.shape = [observation_size]
          #  X be an ensemble particles with X.shape = [ensemble_size, state_size]
          #  G be the observation function,
          #   so G(X).shape = [ensemble_size, state_size].
          # In practice, batch dims may appear between the ensemble and state dims.
  
          # In the traditional EnKF, observation_particles_dist ~ N(G(X), Γ).
          # However, our API would allow any Gaussian.
          observation_particles_dist, extra = observation_fn(
              state.step, state.particles, state.extra
          )
  
          common_dtype = dtype_util.common_dtype(
              [observation_particles_dist, observation],
              dtype_hint=tf.float32,
          )
  
          observation = tf.convert_to_tensor(
              observation, dtype=common_dtype
          )
          observation_size_is_static_and_scalar = (
              observation.shape[-1] == 1
          )
  
          if not isinstance(
              observation_particles_dist,
              mvn_linear_operator.MultivariateNormalLinearOperator,
          ):
              raise ValueError(
                  "Expected `observation_fn` to return an instance of "
                  "`MultivariateNormalLinearOperator`"
              )
  
          # predicted_observation_particles = G(X) + E[η],
          # and is shape [n_ensemble] + [observation_size]
          # Note that .mean() is the distribution mean, and the distribution is
          # centered at the predicted observations. This is *not* the ensemble mean.
          predicted_observation_particles = (
              observation_particles_dist.mean()
          )
  
          # covariance_between_state_and_predicted_observations
          # Cov(X, G(X))  = (X - μ(X))(G(X) - μ(G(X)))ᵀ
          covariance_between_state_and_predicted_observations = (
              tf.nest.map_structure(
                  lambda x: _covariance(
                      x, predicted_observation_particles
                  ),
                  state.particles,
              )
          )
  
          # observation_particles_diff = Y - G(X) - η
          observation_particles_diff = (
              observation - observation_particles_dist.sample(seed=seed)
          )
  
          # observation_particles_covariance ~ Cov(G(X)) + Γ
          if low_rank_ensemble:
              # observation_particles_covariance ~ LinearOperatorLowRankUpdate
              observation_particles_covariance = _observation_particles_cov_linop(
                  predicted_observation_particles=predicted_observation_particles,
                  ensemble_mean_observations=tf.reduce_mean(
                      predicted_observation_particles, axis=0
                  ),
                  observation_cov=_linop_covariance(
                      observation_particles_dist
                  ),
              )
          else:
              # observation_particles_covariance ~ LinearOperatorFullMatrix
              observation_particles_covariance_matrix = (
                  # With μ the ensemble average operator. For V a batch of column
                  # vectors, let Vᵀ be a batch of row vectors. Then
                  # _covariance(predicted_observation_particles)
                  #  = Cov(G(X)) = (G(X) - μ(G(X))) (G(X) - μ(G(X)))ᵀ
                  _covariance(predicted_observation_particles)
                  +
                  # Calling _linop_covariance(...).to_dense() rather than
                  # observation_particles_dist.covariance() means the shape is
                  # [observation_size, observation_size] rather than
                  # [ensemble_size] + [observation_size, observation_size].
                  # Both work, since this matrix is used to do mat-vecs with ensembles
                  # of vectors...however, doing things this way ensures we do an
                  # efficient batch-matmul and (more importantly) don't have to do a
                  # separate Cholesky for every ensemble member!
                  _linop_covariance(
                      observation_particles_dist
                  ).to_dense()
              )
              observation_particles_covariance = tf.linalg.LinearOperatorFullMatrix(
                  observation_particles_covariance_matrix,
                  # SPD because _linop_covariance(observation_particles_dist) is SPD
                  # and _covariance(predicted_observation_particles) is SSD
                  is_self_adjoint=True,
                  is_positive_definite=True,
              )
  
          # We specialize the univariate case.
          # TODO(srvasude): Refactor linear_gaussian_ssm, normal_conjugate_posteriors
          # and this code so we have a central place for normal conjugacy code.
          # Note that we do not use this code path if `low_rank_ensemble`, since our
          # API specifies we will use a LinearOperatorLowRankUpdate. This code path
          # would be a bit more efficient, but the user is warned in the docstring not
          # to set `low_rank_ensemble=True` if the observation dimension is low.
          if (
              observation_size_is_static_and_scalar
              and not low_rank_ensemble
          ):
              # In the univariate observation case, the Kalman gain is given by:
              # K = cov(X, Y) / (var(Y) + var_noise). That is we just divide
              # by the particle covariance plus the observation noise.
              kalman_gain = tf.nest.map_structure(
                  lambda x: x / observation_particles_covariance_matrix,
                  covariance_between_state_and_predicted_observations,
              )
              new_particles = tf.nest.map_structure(
                  lambda x, g: x
                  + damping
                  * tf.linalg.matvec(  # pylint:disable=g-long-lambda
                      g, observation_particles_diff
                  ),
                  state.particles,
                  kalman_gain,
              )
          else:
              # added_term = [Cov(G(X)) + Γ]⁻¹ [Y - G(X) - η]
              added_term = observation_particles_covariance.solvevec(
                  observation_particles_diff
              )
  
              # added_term
              #  = covariance_between_state_and_predicted_observations @ added_term
              #  = Cov(X, G(X)) [Cov(G(X)) + Γ]⁻¹ [Y - G(X) - η]
              #  = (X - μ(X))(G(X) - μ(G(X)))ᵀ [Cov(G(X)) + Γ]⁻¹ [Y - G(X) - η]
              added_term = tf.nest.map_structure(
                  lambda x: tf.linalg.matvec(x, added_term),
                  covariance_between_state_and_predicted_observations,
              )
  
              # new_particles = X + damping * added_term
              new_particles = tf.nest.map_structure(
                  lambda x, a: x + damping * a,
                  state.particles,
                  added_term,
              )
  
          return EnsembleKalmanFilterState(
              step=state.step + 1, particles=new_particles, extra=extra
          )
  
file_path: probability/tensorflow_probability/python/experimental/sequential/ensemble_kalman_filter.py
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import mvn_linear_operator
  from tensorflow_probability.python.distributions import mvn_low_rank_update_linear_operator_covariance
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
label_desc: Use image node_id 4 for calling a global function with example usage: ensemble_kalman_filter_update(state, observation, observation_fn, damping, low_rank_ensemble, seed, name) and returns: EnsembleKalmanFilterState