class_docstr: Compare various KF EnKF versions.
  
  If the model is linear and Gaussian the EnKF sample mean/cov and marginal
  likelihood converges to that of a KF in the large ensemble limit.
  This class tests that they are the same. It does that by implementing a
  one-step KF. It also does some simple checks on the KF, to make sure we didn't
  just replicate misunderstanding in the EnKF.
  
  This class also checks that various flavors of the EnKF are the same.
docstr: Solve one assimilation step using a KF.
docstr_tok: ["Solve","one","assimilation","step","using","a","KF","."]
code: def _kalman_filter_solve(self, observation, linear_model_params):
      """Solve one assimilation step using a KF."""
      # See http://screen/tnjSAEuo5nPKmYt for equations.
      # pylint: disable=unnecessary-lambda
      p = linear_model_params  # Simple & Sweet
  
      # With A, B matrices and x a vector, we define the operations...
      a_x = lambda a, x: tf.linalg.matvec(a, x)  # Ax
      a_b = lambda a, b: tf.linalg.matmul(a, b)  # AB
      a_bt = lambda a, b: tf.linalg.matmul(a, b, adjoint_b=True)  # ABᵀ
      a_b_at = lambda c, d: a_b(c, a_bt(d, c))  # ABAᵀ
  
      predictive_mean = a_x(p.transition_mat, p.prior_mean)
      predictive_cov = (
          a_b_at(p.transition_mat, p.prior_cov) + p.transition_cov
      )
  
      kalman_gain = a_b(
          a_bt(predictive_cov, p.observation_mat),
          tf.linalg.inv(
              a_b_at(p.observation_mat, predictive_cov)
              + p.observation_noise_cov
          ),
      )
      updated_mean = predictive_mean + a_x(
          kalman_gain,
          observation - a_x(p.observation_mat, predictive_mean),
      )
      updated_cov = a_b(
          tf.linalg.eye(p.n_states, dtype=p.dtype)
          - a_b(kalman_gain, p.observation_mat),
          predictive_cov,
      )
  
      # p(Y | X_{predictive})
      marginal_dist = mvn_tril.MultivariateNormalTriL(
          loc=a_x(p.observation_mat, predictive_mean),
          scale_tril=tf.linalg.cholesky(
              a_b_at(p.observation_mat, predictive_cov)
              + p.observation_noise_cov
          ),
      )
  
      return dict(
          predictive_mean=predictive_mean,
          predictive_cov=predictive_cov,
          predictive_stddev=tf.sqrt(
              tf.linalg.diag_part(predictive_cov)
          ),
          updated_mean=updated_mean,
          updated_cov=updated_cov,
          updated_stddev=tf.sqrt(tf.linalg.diag_part(updated_cov)),
          log_marginal_likelihood=marginal_dist.log_prob(observation),
      )
  
imports: import collections
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.experimental.sequential import ensemble_kalman_filter
  from tensorflow_probability.python.internal import test_combinations
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.stats import sample_stats
label_desc: Use image node_id 3 for calling the ComparingMethodsTest obj's underlying member method code with example usage: obj._kalman_filter_solve(observation, linear_model_params) and returns: dict