tensorflow_probability/python/distributions/linear_gaussian_ssm.py
def build_kalman_filter_step(get_transition_matrix_for_timestep,
                             get_transition_noise_for_timestep,
                             get_observation_matrix_for_timestep,
                             get_observation_noise_for_timestep):
  """Build a callable that performs one step of Kalman filtering.

  Args:
    get_transition_matrix_for_timestep: callable taking a timestep
      as an integer `Tensor` argument, and returning a `LinearOperator`
      of shape `[latent_size, latent_size]`.
    get_transition_noise_for_timestep: callable taking a timestep as
      an integer `Tensor` argument, and returning a
      `MultivariateNormalLinearOperator` of event shape
      `[latent_size]`.
    get_observation_matrix_for_timestep: callable taking a timestep
      as an integer `Tensor` argument, and returning a `LinearOperator`
      of shape `[observation_size, observation_size]`.
    get_observation_noise_for_timestep: callable taking a timestep as
      an integer `Tensor` argument, and returning a
      `MultivariateNormalLinearOperator` of event shape
      `[observation_size]`.

  Returns:
    kalman_filter_step: a callable that updates a KalmanFilterState
      from timestep `t-1` to `t`.
  """

  def kalman_filter_step(state, elems_t):
    """Run a single step of Kalman filtering.

    Args:
      state: A `KalmanFilterState` object representing the previous
        filter state at time `t-1`.
      elems_t: A tuple of Tensors `(x_t, mask_t)`, or a `Tensor` `x_t`.
        `x_t` is a `Tensor` with rightmost shape dimensions
        `[observation_size, 1]` representing the vector observed at time `t`,
        and `mask_t` is a `Tensor` with rightmost dimensions`[1, 1]`
        representing the observation mask at time `t`. Both `x_t` and `mask_t`
        may have batch dimensions, which must be compatible with the batch
        dimensions of `state.predicted_mean` and `state.predictived_cov`
        respectively. If `mask_t` is not provided, it is assumed to be `None`.

    Returns:
      new_state: A `KalmanFilterState` object representing the new
        filter state at time `t`.
    """

    if isinstance(elems_t, tuple):
      x_t, mask_t = elems_t
    else:
      x_t = elems_t
      mask_t = None

    observation_matrix = get_observation_matrix_for_timestep(state.timestep)
    observation_noise = get_observation_noise_for_timestep(state.timestep)
    if mask_t is not None:
      # Before running the update, fill in masked observations using the prior
      # expectation. The precise filled value shouldn't matter since updates
      # from masked elements will not be selected below, but we need to ensure
      # that any results we incidently compute on masked values are at least
      # finite (not inf or NaN) so that they don't screw up gradient propagation
      # through `tf.where`, as described in
      #  https://github.com/tensorflow/tensorflow/issues/2540.
      # We fill with the prior expectation because any fixed value such as zero
      # might be arbitrarily unlikely under the prior, leading to overflow in
      # the updates, but the prior expectation should always be a
      # 'reasonable' observation.
      x_expected = _propagate_mean(state.predicted_mean,
                                   observation_matrix,
                                   observation_noise) * tf.ones_like(x_t)
      x_t = tf.where(
          tf.broadcast_to(mask_t, tf.shape(input=x_expected)), x_expected,
          tf.broadcast_to(x_t, tf.shape(input=x_expected)))

    # Given predicted mean u_{t|t-1} and covariance P_{t|t-1} from the
    # previous step, incorporate the observation x_t, producing the
    # filtered mean u_t and covariance P_t.
    (filtered_mean,
     filtered_cov,
     observation_dist) = linear_gaussian_update(
         state.predicted_mean, state.predicted_cov,
         observation_matrix, observation_noise,
         x_t)

    # Compute the marginal likelihood p(x_{t} | x_{:t-1}) for this
    # observation.
    log_marginal_likelihood = observation_dist.log_prob(x_t[..., 0])

    if mask_t is not None:
      filtered_mean = tf.where(
          tf.broadcast_to(mask_t, tf.shape(input=filtered_mean)),
          state.predicted_mean, filtered_mean)
      filtered_cov = tf.where(
          tf.broadcast_to(mask_t, tf.shape(input=filtered_cov)),
          state.predicted_cov, filtered_cov)
      log_marginal_likelihood = tf.where(
          tf.broadcast_to(mask_t[..., 0, 0],
                          tf.shape(input=log_marginal_likelihood)),
          tf.zeros_like(log_marginal_likelihood),
          log_marginal_likelihood)

    # Run the filtered posterior through the transition
    # model to predict the next time step:
    #  u_{t|t-1} = F_t u_{t-1} + b_t
    #  P_{t|t-1} = F_t P_{t-1} F_t' + Q_t
    predicted_mean, predicted_cov = kalman_transition(
        filtered_mean,
        filtered_cov,
        get_transition_matrix_for_timestep(state.timestep),
        get_transition_noise_for_timestep(state.timestep))

    return KalmanFilterState(
        filtered_mean, filtered_cov,
        predicted_mean, predicted_cov,
        observation_dist.mean()[..., tf.newaxis],
        observation_dist.covariance(),
        log_marginal_likelihood,
        state.timestep+1)

  return kalman_filter_step
def build_kalman_filter_step(get_transition_matrix_for_timestep,
                             get_transition_noise_for_timestep,
                             get_observation_matrix_for_timestep,
                             get_observation_noise_for_timestep):
  def kalman_filter_step(state, elems_t):
    """Run a single step of Kalman filtering.

    Args:
      state: A `KalmanFilterState` object representing the previous
        filter state at time `t-1`.
      elems_t: A tuple of Tensors `(x_t, mask_t)`, or a `Tensor` `x_t`.
        `x_t` is a `Tensor` with rightmost shape dimensions
        `[observation_size, 1]` representing the vector observed at time `t`,
        and `mask_t` is a `Tensor` with rightmost dimensions`[1, 1]`
        representing the observation mask at time `t`. Both `x_t` and `mask_t`
        may have batch dimensions, which must be compatible with the batch
        dimensions of `state.predicted_mean` and `state.predictived_cov`
        respectively. If `mask_t` is not provided, it is assumed to be `None`.

    Returns:
      new_state: A `KalmanFilterState` object representing the new
        filter state at time `t`.
    """

    if isinstance(elems_t, tuple):
      x_t, mask_t = elems_t
    else:
      x_t = elems_t
      mask_t = None

    observation_matrix = get_observation_matrix_for_timestep(state.timestep)
    observation_noise = get_observation_noise_for_timestep(state.timestep)
    if mask_t is not None:
      # Before running the update, fill in masked observations using the prior
      # expectation. The precise filled value shouldn't matter since updates
      # from masked elements will not be selected below, but we need to ensure
      # that any results we incidently compute on masked values are at least
      # finite (not inf or NaN) so that they don't screw up gradient propagation
      # through `tf.where`, as described in
      #  https://github.com/tensorflow/tensorflow/issues/2540.
      # We fill with the prior expectation because any fixed value such as zero
      # might be arbitrarily unlikely under the prior, leading to overflow in
      # the updates, but the prior expectation should always be a
      # 'reasonable' observation.
      x_expected = _propagate_mean(state.predicted_mean,
                                   observation_matrix,
                                   observation_noise) * tf.ones_like(x_t)
      x_t = tf.where(
          tf.broadcast_to(mask_t, tf.shape(input=x_expected)), x_expected,
          tf.broadcast_to(x_t, tf.shape(input=x_expected)))

    # Given predicted mean u_{t|t-1} and covariance P_{t|t-1} from the
    # previous step, incorporate the observation x_t, producing the
    # filtered mean u_t and covariance P_t.
    (filtered_mean,
     filtered_cov,
     observation_dist) = linear_gaussian_update(
         state.predicted_mean, state.predicted_cov,
         observation_matrix, observation_noise,
         x_t)

    # Compute the marginal likelihood p(x_{t} | x_{:t-1}) for this
    # observation.
    log_marginal_likelihood = observation_dist.log_prob(x_t[..., 0])

    if mask_t is not None:
      filtered_mean = tf.where(
          tf.broadcast_to(mask_t, tf.shape(input=filtered_mean)),
          state.predicted_mean, filtered_mean)
      filtered_cov = tf.where(
          tf.broadcast_to(mask_t, tf.shape(input=filtered_cov)),
          state.predicted_cov, filtered_cov)
      log_marginal_likelihood = tf.where(
          tf.broadcast_to(mask_t[..., 0, 0],
                          tf.shape(input=log_marginal_likelihood)),
          tf.zeros_like(log_marginal_likelihood),
          log_marginal_likelihood)

    # Run the filtered posterior through the transition
    # model to predict the next time step:
    #  u_{t|t-1} = F_t u_{t-1} + b_t
    #  P_{t|t-1} = F_t P_{t-1} F_t' + Q_t
    predicted_mean, predicted_cov = kalman_transition(
        filtered_mean,
        filtered_cov,
        get_transition_matrix_for_timestep(state.timestep),
        get_transition_noise_for_timestep(state.timestep))

    return KalmanFilterState(
        filtered_mean, filtered_cov,
        predicted_mean, predicted_cov,
        observation_dist.mean()[..., tf.newaxis],
        observation_dist.covariance(),
        log_marginal_likelihood,
        state.timestep+1)

  return kalman_filter_step
def build_kalman_filter_step ( get_transition_matrix_for_timestep , get_transition_noise_for_timestep , get_observation_matrix_for_timestep , get_observation_noise_for_timestep ) : def kalman_filter_step ( state , elems_t ) : """Run a single step of Kalman filtering.

    Args:
      state: A `KalmanFilterState` object representing the previous
        filter state at time `t-1`.
      elems_t: A tuple of Tensors `(x_t, mask_t)`, or a `Tensor` `x_t`.
        `x_t` is a `Tensor` with rightmost shape dimensions
        `[observation_size, 1]` representing the vector observed at time `t`,
        and `mask_t` is a `Tensor` with rightmost dimensions`[1, 1]`
        representing the observation mask at time `t`. Both `x_t` and `mask_t`
        may have batch dimensions, which must be compatible with the batch
        dimensions of `state.predicted_mean` and `state.predictived_cov`
        respectively. If `mask_t` is not provided, it is assumed to be `None`.

    Returns:
      new_state: A `KalmanFilterState` object representing the new
        filter state at time `t`.
    """ if isinstance ( elems_t , tuple ) : x_t , mask_t = elems_t else : x_t = elems_t mask_t = None observation_matrix = get_observation_matrix_for_timestep ( state . timestep ) observation_noise = get_observation_noise_for_timestep ( state . timestep ) if mask_t is not None : # Before running the update, fill in masked observations using the prior # expectation. The precise filled value shouldn't matter since updates # from masked elements will not be selected below, but we need to ensure # that any results we incidently compute on masked values are at least # finite (not inf or NaN) so that they don't screw up gradient propagation # through `tf.where`, as described in #  https://github.com/tensorflow/tensorflow/issues/2540. # We fill with the prior expectation because any fixed value such as zero # might be arbitrarily unlikely under the prior, leading to overflow in # the updates, but the prior expectation should always be a # 'reasonable' observation. x_expected = _propagate_mean ( state . predicted_mean , observation_matrix , observation_noise ) * tf . ones_like ( x_t ) x_t = tf . where ( tf . broadcast_to ( mask_t , tf . shape ( input = x_expected ) ) , x_expected , tf . broadcast_to ( x_t , tf . shape ( input = x_expected ) ) ) # Given predicted mean u_{t|t-1} and covariance P_{t|t-1} from the # previous step, incorporate the observation x_t, producing the # filtered mean u_t and covariance P_t. ( filtered_mean , filtered_cov , observation_dist ) = linear_gaussian_update ( state . predicted_mean , state . predicted_cov , observation_matrix , observation_noise , x_t ) # Compute the marginal likelihood p(x_{t} | x_{:t-1}) for this # observation. log_marginal_likelihood = observation_dist . log_prob ( x_t [ ... , 0 ] ) if mask_t is not None : filtered_mean = tf . where ( tf . broadcast_to ( mask_t , tf . shape ( input = filtered_mean ) ) , state . predicted_mean , filtered_mean ) filtered_cov = tf . where ( tf . broadcast_to ( mask_t , tf . shape ( input = filtered_cov ) ) , state . predicted_cov , filtered_cov ) log_marginal_likelihood = tf . where ( tf . broadcast_to ( mask_t [ ... , 0 , 0 ] , tf . shape ( input = log_marginal_likelihood ) ) , tf . zeros_like ( log_marginal_likelihood ) , log_marginal_likelihood ) # Run the filtered posterior through the transition # model to predict the next time step: #  u_{t|t-1} = F_t u_{t-1} + b_t #  P_{t|t-1} = F_t P_{t-1} F_t' + Q_t predicted_mean , predicted_cov = kalman_transition ( filtered_mean , filtered_cov , get_transition_matrix_for_timestep ( state . timestep ) , get_transition_noise_for_timestep ( state . timestep ) ) return KalmanFilterState ( filtered_mean , filtered_cov , predicted_mean , predicted_cov , observation_dist . mean ( ) [ ... , tf . newaxis ] , observation_dist . covariance ( ) , log_marginal_likelihood , state . timestep + 1 ) return kalman_filter_step
Build a callable that performs one step of Kalman filtering.

  Args:
    get_transition_matrix_for_timestep: callable taking a timestep
      as an integer `Tensor` argument, and returning a `LinearOperator`
      of shape `[latent_size, latent_size]`.
    get_transition_noise_for_timestep: callable taking a timestep as
      an integer `Tensor` argument, and returning a
      `MultivariateNormalLinearOperator` of event shape
      `[latent_size]`.
    get_observation_matrix_for_timestep: callable taking a timestep
      as an integer `Tensor` argument, and returning a `LinearOperator`
      of shape `[observation_size, observation_size]`.
    get_observation_noise_for_timestep: callable taking a timestep as
      an integer `Tensor` argument, and returning a
      `MultivariateNormalLinearOperator` of event shape
      `[observation_size]`.

  Returns:
    kalman_filter_step: a callable that updates a KalmanFilterState
      from timestep `t-1` to `t`.
Build a callable that performs one step of Kalman filtering .