code: def ensemble_kalman_filter_log_marginal_likelihood(
      state,
      observation,
      observation_fn,
      perturbed_observations=True,
      low_rank_ensemble=False,
      seed=None,
      name=None,
  ):
      """Ensemble Kalman filter log marginal likelihood.
  
      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])
      ```
  
      This method estimates (logarithm of) the marginal likelihood of the
      observation at step `t`, `Y[t]`, given `state`. Typically, `state` is the
      predictive ensemble at time `t`. In that case, this function approximates
       `Log[p(Y[t] | Y[t-1], Y[t-2],...)]`
      The approximation is correct under a Linear Gaussian state space model
      assumption, as ensemble size --> infinity.
  
      Args:
        state: Instance of `EnsembleKalmanFilterState` at step `k`,
          conditioned on previous observations `Y_{1:k}`. Typically this is the
          output of `ensemble_kalman_filter_predict`.
        observation: `Tensor` representing the observation at step `k`.
        observation_fn: callable returning an instance of
          `tfd.MultivariateNormalLinearOperator` along with an extra information
          to be returned in the `EnsembleKalmanFilterState`.
        perturbed_observations: Whether the marginal distribution `p(Y[t] | ...)`
          is estimated using samples from the `observation_fn`'s distribution. If
          `False`, the distribution's covariance matrix is used directly. This
          latter choice is less common in the literature, but works even if the
          ensemble size is smaller than the number of observations.
        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_log_marginal_likelihood'`).
  
      Returns:
        log_marginal_likelihood: `Tensor` with same dtype as `state`.
  
      Raises:
        InsufficientEnsembleSizeError: If `perturbed_observations=True` and the
          ensemble size is not at least one greater than the number of observations.
  
      #### 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_log_marginal_likelihood"
      ):
          observation_particles_dist, unused_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
          )
  
          if low_rank_ensemble and perturbed_observations:
              raise ValueError(
                  "A low rank update cannot be used with `perturbed_observations=True`"
              )
  
          if perturbed_observations:
              # With G the observation operator and B the batch shape,
              # observation_particles = G(X) + η, where η ~ Normal(0, Γ).
              # Both are shape [n_ensemble] + B + [n_observations]
              observation_particles = observation_particles_dist.sample(
                  seed=seed
              )
              n_observations = observation_particles_dist.event_shape[0]
              n_ensemble = observation_particles_dist.batch_shape[0]
              if (
                  n_ensemble is not None
                  and n_observations is not None
                  and n_ensemble < n_observations + 1
              ):
                  raise InsufficientEnsembleSizeError(
                      f"When `perturbed_observations=True`, ensemble size ({n_ensemble}) "
                      "must be at least one greater than the number of observations "
                      f"({n_observations}), but it was not."
                  )
              observation_dist = mvn_tril.MultivariateNormalTriL(
                  loc=tf.reduce_mean(observation_particles, axis=0),
                  # Cholesky(Cov(G(X) + η)), where Cov(..) is the ensemble covariance.
                  scale_tril=tf.linalg.cholesky(
                      _covariance(observation_particles)
                  ),
              )
          else:
              if (
                  low_rank_ensemble
              ):  # low_rank_ensemble and not perturbed_observations.
                  predicted_observation_particles = (
                      observation_particles_dist.mean()
                  )
                  ensemble_mean_observations = tf.reduce_mean(
                      predicted_observation_particles, axis=0
                  )
                  observation_dist = MVNLowRankCov(
                      loc=ensemble_mean_observations,
                      cov_operator=_observation_particles_cov_linop(
                          predicted_observation_particles=predicted_observation_particles,
                          ensemble_mean_observations=ensemble_mean_observations,
                          observation_cov=_linop_covariance(
                              observation_particles_dist
                          ),
                      ),
                  )
              else:  # not low_rank_ensemble and not perturbed_observations.
                  # predicted_observation = G(X),
                  # and is shape [n_ensemble] + B.
                  predicted_observations = (
                      observation_particles_dist.mean()
                  )
                  observation_dist = mvn_tril.MultivariateNormalTriL(
                      loc=tf.reduce_mean(
                          predicted_observations, axis=0
                      ),  # ensemble mean
                      # Cholesky(Cov(G(X)) + Γ), where Cov(..) is the ensemble covariance.
                      scale_tril=tf.linalg.cholesky(
                          _covariance(predicted_observations)
                          + _linop_covariance(
                              observation_particles_dist
                          ).to_dense()
                      ),
                  )
  
          # Above we computed observation_dist, the distribution of observations given
          # the predictive distribution of states (e.g. states from previous time).
          # Here we evaluate the log_prob on the actual observations.
          return observation_dist.log_prob(observation)
  
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 5 for calling a global function with example usage: ensemble_kalman_filter_log_marginal_likelihood(state, observation, observation_fn, perturbed_observations, low_rank_ensemble, seed, name) and returns: observation_dist