tensorflow_probability/python/distributions/linear_gaussian_ssm.py
def linear_gaussian_update(
    prior_mean, prior_cov, observation_matrix, observation_noise, x_observed):
  """Conjugate update for a linear Gaussian model.

  Given a normal prior on a latent variable `z`,
    `p(z) = N(prior_mean, prior_cov) = N(u, P)`,
  for which we observe a linear Gaussian transformation `x`,
    `p(x|z) = N(H * z + c, R)`,
  the posterior is also normal:
    `p(z|x) = N(u*, P*)`.

  We can write this update as
     x_expected = H * u + c # pushforward prior mean
     S = R + H * P * H'  # pushforward prior cov
     K = P * H' * S^{-1} # optimal Kalman gain
     u* = u + K * (x_observed - x_expected) # posterior mean
     P* = (I - K * H) * P (I - K * H)' + K * R * K' # posterior cov
  (see, e.g., https://en.wikipedia.org/wiki/Kalman_filter#Update)

  Args:
    prior_mean: `Tensor` with event shape `[latent_size, 1]` and
      potential batch shape `B = [b1, ..., b_n]`.
    prior_cov: `Tensor` with event shape `[latent_size, latent_size]`
      and batch shape `B` (matching `prior_mean`).
    observation_matrix: `LinearOperator` with shape
      `[observation_size, latent_size]` and batch shape broadcastable
      to `B`.
    observation_noise: potentially-batched
      `MultivariateNormalLinearOperator` instance with event shape
      `[observation_size]` and batch shape broadcastable to `B`.
    x_observed: potentially batched `Tensor` with event shape
      `[observation_size, 1]` and batch shape `B`.

  Returns:
    posterior_mean: `Tensor` with event shape `[latent_size, 1]` and
      batch shape `B`.
    posterior_cov: `Tensor` with event shape `[latent_size,
      latent_size]` and batch shape `B`.
    predictive_dist: the prior predictive distribution `p(x|z)`,
      as a `Distribution` instance with event
      shape `[observation_size]` and batch shape `B`. This will
      typically be `tfd.MultivariateNormalTriL`, but when
      `observation_size=1` we return a `tfd.Independent(tfd.Normal)`
      instance as an optimization.
  """

  # If observations are scalar, we can avoid some matrix ops.
  observation_size_is_static_and_scalar = (
      tf.compat.dimension_value(observation_matrix.shape[-2]) == 1)

  # Push the predicted mean for the latent state through the
  # observation model
  x_expected = _propagate_mean(prior_mean,
                               observation_matrix,
                               observation_noise)

  # Push the predictive covariance of the latent state through the
  # observation model:
  #  S = R + H * P * H'.
  # We use a temporary variable for H * P,
  # reused below to compute Kalman gain.
  tmp_obs_cov = observation_matrix.matmul(prior_cov)
  predicted_obs_cov = (
      observation_matrix.matmul(tmp_obs_cov, adjoint_arg=True)
      + observation_noise.covariance())

  # Compute optimal Kalman gain:
  #  K = P * H' * S^{-1}
  # Since both S and P are cov matrices, thus symmetric,
  # we can take the transpose and reuse our previous
  # computation:
  #      = (S^{-1} * H * P)'
  #      = (S^{-1} * tmp_obs_cov) '
  #      = (S \ tmp_obs_cov)'
  if observation_size_is_static_and_scalar:
    gain_transpose = tmp_obs_cov/predicted_obs_cov
  else:
    predicted_obs_cov_chol = tf.linalg.cholesky(predicted_obs_cov)
    gain_transpose = tf.linalg.cholesky_solve(predicted_obs_cov_chol,
                                              tmp_obs_cov)

  # Compute the posterior mean, incorporating the observation.
  #  u* = u + K (x_observed - x_expected)
  posterior_mean = (prior_mean +
                    tf.linalg.matmul(gain_transpose, x_observed - x_expected,
                                     adjoint_a=True))

  # For the posterior covariance, we could use the simple update
  #  P* = P - K * H * P
  # but this is prone to numerical issues because it subtracts a
  # value from a PSD matrix.  We choose instead to use the more
  # expensive Jordan form update
  #  P* = (I - K H) * P * (I - K H)' + K R K'
  # which always produces a PSD result. This uses
  #  tmp_term = (I - K * H)'
  # as an intermediate quantity.
  tmp_term = -observation_matrix.matmul(gain_transpose, adjoint=True)  # -K * H
  tmp_term = tf.linalg.set_diag(tmp_term, tf.linalg.diag_part(tmp_term) + 1)
  posterior_cov = (
      tf.linalg.matmul(
          tmp_term, tf.linalg.matmul(prior_cov, tmp_term), adjoint_a=True)
      + tf.linalg.matmul(gain_transpose,
                         tf.linalg.matmul(
                             observation_noise.covariance(), gain_transpose),
                         adjoint_a=True))

  if observation_size_is_static_and_scalar:
    # A plain Normal would have event shape `[]`; wrapping with Independent
    # ensures `event_shape=[1]` as required.
    predictive_dist = independent.Independent(
        normal.Normal(loc=x_expected[..., 0],
                      scale=tf.sqrt(predicted_obs_cov[..., 0])),
        reinterpreted_batch_ndims=1)

    # Minor hack to define the covariance, so that `predictive_dist` can pass as
    # an MVNTriL-like object.
    predictive_dist.covariance = lambda: predicted_obs_cov
  else:
    predictive_dist = mvn_tril.MultivariateNormalTriL(
        loc=x_expected[..., 0],
        scale_tril=predicted_obs_cov_chol)

  return posterior_mean, posterior_cov, predictive_dist
def linear_gaussian_update(
    prior_mean, prior_cov, observation_matrix, observation_noise, x_observed):
  # If observations are scalar, we can avoid some matrix ops.
  observation_size_is_static_and_scalar = (
      tf.compat.dimension_value(observation_matrix.shape[-2]) == 1)

  # Push the predicted mean for the latent state through the
  # observation model
  x_expected = _propagate_mean(prior_mean,
                               observation_matrix,
                               observation_noise)

  # Push the predictive covariance of the latent state through the
  # observation model:
  #  S = R + H * P * H'.
  # We use a temporary variable for H * P,
  # reused below to compute Kalman gain.
  tmp_obs_cov = observation_matrix.matmul(prior_cov)
  predicted_obs_cov = (
      observation_matrix.matmul(tmp_obs_cov, adjoint_arg=True)
      + observation_noise.covariance())

  # Compute optimal Kalman gain:
  #  K = P * H' * S^{-1}
  # Since both S and P are cov matrices, thus symmetric,
  # we can take the transpose and reuse our previous
  # computation:
  #      = (S^{-1} * H * P)'
  #      = (S^{-1} * tmp_obs_cov) '
  #      = (S \ tmp_obs_cov)'
  if observation_size_is_static_and_scalar:
    gain_transpose = tmp_obs_cov/predicted_obs_cov
  else:
    predicted_obs_cov_chol = tf.linalg.cholesky(predicted_obs_cov)
    gain_transpose = tf.linalg.cholesky_solve(predicted_obs_cov_chol,
                                              tmp_obs_cov)

  # Compute the posterior mean, incorporating the observation.
  #  u* = u + K (x_observed - x_expected)
  posterior_mean = (prior_mean +
                    tf.linalg.matmul(gain_transpose, x_observed - x_expected,
                                     adjoint_a=True))

  # For the posterior covariance, we could use the simple update
  #  P* = P - K * H * P
  # but this is prone to numerical issues because it subtracts a
  # value from a PSD matrix.  We choose instead to use the more
  # expensive Jordan form update
  #  P* = (I - K H) * P * (I - K H)' + K R K'
  # which always produces a PSD result. This uses
  #  tmp_term = (I - K * H)'
  # as an intermediate quantity.
  tmp_term = -observation_matrix.matmul(gain_transpose, adjoint=True)  # -K * H
  tmp_term = tf.linalg.set_diag(tmp_term, tf.linalg.diag_part(tmp_term) + 1)
  posterior_cov = (
      tf.linalg.matmul(
          tmp_term, tf.linalg.matmul(prior_cov, tmp_term), adjoint_a=True)
      + tf.linalg.matmul(gain_transpose,
                         tf.linalg.matmul(
                             observation_noise.covariance(), gain_transpose),
                         adjoint_a=True))

  if observation_size_is_static_and_scalar:
    # A plain Normal would have event shape `[]`; wrapping with Independent
    # ensures `event_shape=[1]` as required.
    predictive_dist = independent.Independent(
        normal.Normal(loc=x_expected[..., 0],
                      scale=tf.sqrt(predicted_obs_cov[..., 0])),
        reinterpreted_batch_ndims=1)

    # Minor hack to define the covariance, so that `predictive_dist` can pass as
    # an MVNTriL-like object.
    predictive_dist.covariance = lambda: predicted_obs_cov
  else:
    predictive_dist = mvn_tril.MultivariateNormalTriL(
        loc=x_expected[..., 0],
        scale_tril=predicted_obs_cov_chol)

  return posterior_mean, posterior_cov, predictive_dist
def linear_gaussian_update ( prior_mean , prior_cov , observation_matrix , observation_noise , x_observed ) : # If observations are scalar, we can avoid some matrix ops. observation_size_is_static_and_scalar = ( tf . compat . dimension_value ( observation_matrix . shape [ - 2 ] ) == 1 ) # Push the predicted mean for the latent state through the # observation model x_expected = _propagate_mean ( prior_mean , observation_matrix , observation_noise ) # Push the predictive covariance of the latent state through the # observation model: #  S = R + H * P * H'. # We use a temporary variable for H * P, # reused below to compute Kalman gain. tmp_obs_cov = observation_matrix . matmul ( prior_cov ) predicted_obs_cov = ( observation_matrix . matmul ( tmp_obs_cov , adjoint_arg = True ) + observation_noise . covariance ( ) ) # Compute optimal Kalman gain: #  K = P * H' * S^{-1} # Since both S and P are cov matrices, thus symmetric, # we can take the transpose and reuse our previous # computation: #      = (S^{-1} * H * P)' #      = (S^{-1} * tmp_obs_cov) ' #      = (S \ tmp_obs_cov)' if observation_size_is_static_and_scalar : gain_transpose = tmp_obs_cov / predicted_obs_cov else : predicted_obs_cov_chol = tf . linalg . cholesky ( predicted_obs_cov ) gain_transpose = tf . linalg . cholesky_solve ( predicted_obs_cov_chol , tmp_obs_cov ) # Compute the posterior mean, incorporating the observation. #  u* = u + K (x_observed - x_expected) posterior_mean = ( prior_mean + tf . linalg . matmul ( gain_transpose , x_observed - x_expected , adjoint_a = True ) ) # For the posterior covariance, we could use the simple update #  P* = P - K * H * P # but this is prone to numerical issues because it subtracts a # value from a PSD matrix.  We choose instead to use the more # expensive Jordan form update #  P* = (I - K H) * P * (I - K H)' + K R K' # which always produces a PSD result. This uses #  tmp_term = (I - K * H)' # as an intermediate quantity. tmp_term = - observation_matrix . matmul ( gain_transpose , adjoint = True ) # -K * H tmp_term = tf . linalg . set_diag ( tmp_term , tf . linalg . diag_part ( tmp_term ) + 1 ) posterior_cov = ( tf . linalg . matmul ( tmp_term , tf . linalg . matmul ( prior_cov , tmp_term ) , adjoint_a = True ) + tf . linalg . matmul ( gain_transpose , tf . linalg . matmul ( observation_noise . covariance ( ) , gain_transpose ) , adjoint_a = True ) ) if observation_size_is_static_and_scalar : # A plain Normal would have event shape `[]`; wrapping with Independent # ensures `event_shape=[1]` as required. predictive_dist = independent . Independent ( normal . Normal ( loc = x_expected [ ... , 0 ] , scale = tf . sqrt ( predicted_obs_cov [ ... , 0 ] ) ) , reinterpreted_batch_ndims = 1 ) # Minor hack to define the covariance, so that `predictive_dist` can pass as # an MVNTriL-like object. predictive_dist . covariance = lambda : predicted_obs_cov else : predictive_dist = mvn_tril . MultivariateNormalTriL ( loc = x_expected [ ... , 0 ] , scale_tril = predicted_obs_cov_chol ) return posterior_mean , posterior_cov , predictive_dist
Conjugate update for a linear Gaussian model.

  Given a normal prior on a latent variable `z`,
    `p(z) = N(prior_mean, prior_cov) = N(u, P)`,
  for which we observe a linear Gaussian transformation `x`,
    `p(x|z) = N(H * z + c, R)`,
  the posterior is also normal:
    `p(z|x) = N(u*, P*)`.

  We can write this update as
     x_expected = H * u + c # pushforward prior mean
     S = R + H * P * H'  # pushforward prior cov
     K = P * H' * S^{-1} # optimal Kalman gain
     u* = u + K * (x_observed - x_expected) # posterior mean
     P* = (I - K * H) * P (I - K * H)' + K * R * K' # posterior cov
  (see, e.g., https://en.wikipedia.org/wiki/Kalman_filter#Update)

  Args:
    prior_mean: `Tensor` with event shape `[latent_size, 1]` and
      potential batch shape `B = [b1, ..., b_n]`.
    prior_cov: `Tensor` with event shape `[latent_size, latent_size]`
      and batch shape `B` (matching `prior_mean`).
    observation_matrix: `LinearOperator` with shape
      `[observation_size, latent_size]` and batch shape broadcastable
      to `B`.
    observation_noise: potentially-batched
      `MultivariateNormalLinearOperator` instance with event shape
      `[observation_size]` and batch shape broadcastable to `B`.
    x_observed: potentially batched `Tensor` with event shape
      `[observation_size, 1]` and batch shape `B`.

  Returns:
    posterior_mean: `Tensor` with event shape `[latent_size, 1]` and
      batch shape `B`.
    posterior_cov: `Tensor` with event shape `[latent_size,
      latent_size]` and batch shape `B`.
    predictive_dist: the prior predictive distribution `p(x|z)`,
      as a `Distribution` instance with event
      shape `[observation_size]` and batch shape `B`. This will
      typically be `tfd.MultivariateNormalTriL`, but when
      `observation_size=1` we return a `tfd.Independent(tfd.Normal)`
      instance as an optimization.
Conjugate update for a linear Gaussian model .