tfp.distributions.LinearGaussianStateSpaceModel

Class LinearGaussianStateSpaceModel

Inherits From: Distribution

Observation distribution from a linear Gaussian state space model.

The state space model, sometimes called a Kalman filter, posits a latent state vector z_t of dimension latent_size that evolves over time following linear Gaussian transitions,

z_{t+1} = F * z_t + N(b; Q)

for transition matrix F, bias b and covariance matrix Q. At each timestep, we observe a noisy projection of the latent state x_t = H * z_t + N(c; R). The transition and observation models may be fixed or may vary between timesteps.

This Distribution represents the marginal distribution on observations, p(x). The marginal log_prob is computed by Kalman filtering [1], and sample by an efficient forward recursion. Both operations require time linear in T, the total number of timesteps.

Shapes

The event shape is [num_timesteps, observation_size], where observation_size is the dimension of each observation x_t. The observation and transition models must return consistent shapes.

This implementation supports vectorized computation over a batch of models. All of the parameters (prior distribution, transition and observation operators and noise models) must have a consistent batch shape.

Time-varying processes

Any of the model-defining parameters (prior distribution, transition and observation operators and noise models) may be specified as a callable taking an integer timestep t and returning a time-dependent value. The dimensionality (latent_size and observation_size) must be the same at all timesteps.

Importantly, the timestep is passed as a Tensor, not a Python integer, so any conditional behavior must occur inside the TensorFlow graph. For example, suppose we want to use a different transition model on even days than odd days. It does not work to write

def transition_matrix(t):
  if t % 2 == 0:
    return even_day_matrix
  else:
    return odd_day_matrix

since the value of t is not fixed at graph-construction time. Instead we need to write

def transition_matrix(t):
  return tf.cond(tf.equal(tf.mod(t, 2), 0),
                 lambda : even_day_matrix,
                 lambda : odd_day_matrix)

so that TensorFlow can switch between operators appropriately at runtime.

Examples

Consider a simple tracking model. The two-dimensional latent state represents the true position of a vehicle, and at each timestep we see a noisy observation of this position (e.g., a GPS reading). The vehicle is assumed to move by a random walk with standard deviation step_std at each step, and observation noise level std. We build the distribution over noisy observations by

ndims = 2
step_std = 1.0
noise_std = 5.0
model = LinearGaussianStateSpaceModel(
  num_timesteps=100,
  transition_matrix=tfl.LinearOperatorIdentity(ndims),
  transition_noise=tfd.MultivariateNormalDiag(
   scale_diag=step_std**2 * tf.ones([ndims])),
  observation_matrix=tfl.LinearOperatorIdentity(ndims),
  observation_noise=tfd.MultivariateNormalDiag(
   scale_diag=noise_std**2 * tf.ones([ndims])),
  initial_state_prior=tfd.MultivariateNormalDiag(
   scale_diag=tf.ones([ndims])))
)

using the identity matrix for the transition and observation operators. We can then use this model to generate samples, compute marginal likelihood of observed sequences, and perform posterior inference.

x = model.sample(5) # Sample from the prior on sequences of observations.
lp = model.log_prob(x) # Marginal likelihood of a (batch of) observations.

# Compute the filtered posterior on latent states given observations,
# and extract the mean and covariance for the current (final) timestep.
_, filtered_means, filtered_covs, _, _ = model.forward_filter(x)
final_step = tfd.MultivariateNormalFullCovariance(
              loc=filtered_means[..., -1, :],
              scale=filtered_covs[..., -1, :])
  • TODO(davmre): implement and describe full posterior inference / smoothing.

  • TODO(davmre): show example of fitting parameters.

__init__

__init__(
    num_timesteps,
    transition_matrix,
    transition_noise,
    observation_matrix,
    observation_noise,
    initial_state_prior,
    initial_step=0,
    validate_args=False,
    allow_nan_stats=True,
    name='LinearGaussianStateSpaceModel'
)

Initialize a `LinearGaussianStateSpaceModel.

Args:

  • num_timesteps: Integer Tensor total number of timesteps.
  • transition_matrix: A transition operator, represented by a Tensor or LinearOperator of shape [latent_size, latent_size], or by a callable taking as argument a scalar integer Tensor t and returning a Tensor or LinearOperator representing the transition operator from latent state at time t to time t + 1.
  • transition_noise: An instance of tfd.MultivariateNormalLinearOperator with event shape [latent_size], representing the mean and covariance of the transition noise model, or a callable taking as argument a scalar integer Tensor t and returning such a distribution representing the noise in the transition from time t to time t + 1.
  • observation_matrix: An observation operator, represented by a Tensor or LinearOperator of shape [observation_size, latent_size], or by a callable taking as argument a scalar integer Tensor t and returning a timestep-specific Tensor or LinearOperator.
  • observation_noise: An instance of tfd.MultivariateNormalLinearOperator with event shape [observation_size], representing the mean and covariance of the observation noise model, or a callable taking as argument a scalar integer Tensor t and returning a timestep-specific noise model.
  • initial_state_prior: An instance of MultivariateNormalLinearOperator representing the prior distribution on latent states; must have event shape [latent_size].
  • initial_step: optional int specifying the time of the first modeled timestep. This is added as an offset when passing timesteps t to (optional) callables specifying timestep-specific transition and observation models.
  • validate_args: Python bool, default False. Whether to validate input with asserts. If validate_args is False, and the inputs are invalid, correct behavior is not guaranteed.
  • allow_nan_stats: Python bool, default True. If False, raise an exception if a statistic (e.g. mean/mode/etc...) is undefined for any batch member If True, batch members with valid parameters leading to undefined statistics will return NaN for this statistic.
  • name: The name to give Ops created by the initializer.

Properties

allow_nan_stats

Python bool describing behavior when a stat is undefined.

Stats return +/- infinity when it makes sense. E.g., the variance of a Cauchy distribution is infinity. However, sometimes the statistic is undefined, e.g., if a distribution's pdf does not achieve a maximum within the support of the distribution, the mode is undefined. If the mean is undefined, then by definition the variance is undefined. E.g. the mean for Student's T for df = 1 is undefined (no clear way to say it is either + or - infinity), so the variance = E[(X - mean)**2] is also undefined.

Returns:

  • allow_nan_stats: Python bool.

batch_shape

Shape of a single sample from a single event index as a TensorShape.

May be partially defined or unknown.

The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.

Returns:

  • batch_shape: TensorShape, possibly unknown.

dtype

The DType of Tensors handled by this Distribution.

event_shape

Shape of a single sample from a single batch as a TensorShape.

May be partially defined or unknown.

Returns:

  • event_shape: TensorShape, possibly unknown.

name

Name prepended to all ops created by this Distribution.

parameters

Dictionary of parameters used to instantiate this Distribution.

reparameterization_type

Describes how samples from the distribution are reparameterized.

Currently this is one of the static instances tfd.FULLY_REPARAMETERIZED or tfd.NOT_REPARAMETERIZED.

Returns:

An instance of ReparameterizationType.

validate_args

Python bool indicating possibly expensive checks are enabled.

Methods

backward_smoothing_pass

backward_smoothing_pass(
    filtered_means,
    filtered_covs,
    predicted_means,
    predicted_covs
)

Run the backward pass in Kalman smoother.

The backward smoothing is using Rauch, Tung and Striebel smoother as as discussed in section 18.3.2 of Kevin P. Murphy, 2012, Machine Learning: A Probabilistic Perspective, The MIT Press. The inputs are returned by forward_filter function.

Args:

  • filtered_means: Means of the per-timestep filtered marginal distributions p(zt | x{:t}), as a Tensor of shape sample_shape(x) + batch_shape + [num_timesteps, latent_size].
  • filtered_covs: Covariances of the per-timestep filtered marginal distributions p(zt | x{:t}), as a Tensor of shape batch_shape + [num_timesteps, latent_size, latent_size].
  • predicted_means: Means of the per-timestep predictive distributions over latent states, p(z{t+1} | x{:t}), as a Tensor of shape sample_shape(x) + batch_shape + [num_timesteps, latent_size].
  • predicted_covs: Covariances of the per-timestep predictive distributions over latent states, p(z{t+1} | x{:t}), as a Tensor of shape batch_shape + [num_timesteps, latent_size, latent_size].

Returns:

  • posterior_means: Means of the smoothed marginal distributions p(zt | x{1:T}), as a Tensor of shape sample_shape(x) + batch_shape + [num_timesteps, latent_size], which is of the same shape as filtered_means.
  • posterior_covs: Covariances of the smoothed marginal distributions p(zt | x{1:T}), as a Tensor of shape batch_shape + [num_timesteps, latent_size, latent_size]. which is of the same shape as filtered_covs.

batch_shape_tensor

batch_shape_tensor(name='batch_shape_tensor')

Shape of a single sample from a single event index as a 1-D Tensor.

The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.

Args:

  • name: name to give to the op

Returns:

  • batch_shape: Tensor.

cdf

cdf(
    value,
    name='cdf'
)

Cumulative distribution function.

Given random variable X, the cumulative distribution function cdf is:

cdf(x) := P[X <= x]

Args:

  • value: float or double Tensor.
  • name: Python str prepended to names of ops created by this function.

Returns:

  • cdf: a Tensor of shape sample_shape(x) + self.batch_shape with values of type self.dtype.

copy

copy(**override_parameters_kwargs)

Creates a deep copy of the distribution.

Args:

  • **override_parameters_kwargs: String/value dictionary of initialization arguments to override with new values.

Returns:

  • distribution: A new instance of type(self) initialized from the union of self.parameters and override_parameters_kwargs, i.e., dict(self.parameters, **override_parameters_kwargs).

covariance

covariance(name='covariance')

Covariance.

Covariance is (possibly) defined only for non-scalar-event distributions.

For example, for a length-k, vector-valued distribution, it is calculated as,

Cov[i, j] = Covariance(X_i, X_j) = E[(X_i - E[X_i]) (X_j - E[X_j])]

where Cov is a (batch of) k x k matrix, 0 <= (i, j) < k, and E denotes expectation.

Alternatively, for non-vector, multivariate distributions (e.g., matrix-valued, Wishart), Covariance shall return a (batch of) matrices under some vectorization of the events, i.e.,

Cov[i, j] = Covariance(Vec(X)_i, Vec(X)_j) = [as above]

where Cov is a (batch of) k' x k' matrices, 0 <= (i, j) < k' = reduce_prod(event_shape), and Vec is some function mapping indices of this distribution's event dimensions to indices of a length-k' vector.

Args:

  • name: Python str prepended to names of ops created by this function.

Returns:

  • covariance: Floating-point Tensor with shape [B1, ..., Bn, k', k'] where the first n dimensions are batch coordinates and k' = reduce_prod(self.event_shape).

cross_entropy

cross_entropy(
    other,
    name='cross_entropy'
)

Computes the (Shannon) cross entropy.

Denote this distribution (self) by P and the other distribution by Q. Assuming P, Q are absolutely continuous with respect to one another and permit densities p(x) dr(x) and q(x) dr(x), (Shannon) cross entropy is defined as:

H[P, Q] = E_p[-log q(X)] = -int_F p(x) log q(x) dr(x)

where F denotes the support of the random variable X ~ P.

Args:

Returns:

  • cross_entropy: self.dtype Tensor with shape [B1, ..., Bn] representing n different calculations of (Shannon) cross entropy.

entropy

entropy(name='entropy')

Shannon entropy in nats.

event_shape_tensor

event_shape_tensor(name='event_shape_tensor')

Shape of a single sample from a single batch as a 1-D int32 Tensor.

Args:

  • name: name to give to the op

Returns:

  • event_shape: Tensor.

forward_filter

forward_filter(x)

Run a Kalman filter over a provided sequence of outputs.

Note that the returned values filtered_means, predicted_means, and observation_means depend on the observed time series x, while the corresponding covariances are independent of the observed series; i.e., they depend only on the model itself. This means that the mean values have shape concat([sample_shape(x), batch_shape, [num_timesteps, {latent/observation}_size]]), while the covariances have shape concat[(batch_shape, [num_timesteps, {latent/observation}_size, {latent/observation}_size]]), which does not depend on the sample shape.

Args:

  • x: a float-type Tensor with rightmost dimensions [num_timesteps, observation_size] matching self.event_shape. Additional dimensions must match or be broadcastable to self.batch_shape; any further dimensions are interpreted as a sample shape.

Returns:

  • log_likelihoods: Per-timestep log marginal likelihoods log p(x_t | x_{:t-1}) evaluated at the input x, as a Tensor of shape sample_shape(x) + batch_shape + [num_timesteps].
  • filtered_means: Means of the per-timestep filtered marginal distributions p(zt | x{:t}), as a Tensor of shape sample_shape(x) + batch_shape + [num_timesteps, latent_size].
  • filtered_covs: Covariances of the per-timestep filtered marginal distributions p(zt | x{:t}), as a Tensor of shape batch_shape + [num_timesteps, latent_size, latent_size].
  • predicted_means: Means of the per-timestep predictive distributions over latent states, p(z{t+1} | x{:t}), as a Tensor of shape sample_shape(x) + batch_shape + [num_timesteps, latent_size].
  • predicted_covs: Covariances of the per-timestep predictive distributions over latent states, p(z{t+1} | x{:t}), as a Tensor of shape batch_shape + [num_timesteps, latent_size, latent_size].
  • observation_means: Means of the per-timestep predictive distributions over observations, p(x{t} | x{:t-1}), as a Tensor of shape sample_shape(x) + batch_shape + [num_timesteps, observation_size].
  • observation_covs: Covariances of the per-timestep predictive distributions over observations, p(x{t} | x{:t-1}), as a Tensor of shape batch_shape + [num_timesteps, observation_size, observation_size].

is_scalar_batch

is_scalar_batch(name='is_scalar_batch')

Indicates that batch_shape == [].

Args:

  • name: Python str prepended to names of ops created by this function.

Returns:

  • is_scalar_batch: bool scalar Tensor.

is_scalar_event

is_scalar_event(name='is_scalar_event')

Indicates that event_shape == [].

Args:

  • name: Python str prepended to names of ops created by this function.

Returns:

  • is_scalar_event: bool scalar Tensor.

kl_divergence

kl_divergence(
    other,
    name='kl_divergence'
)

Computes the Kullback--Leibler divergence.

Denote this distribution (self) by p and the other distribution by q. Assuming p, q are absolutely continuous with respect to reference measure r, the KL divergence is defined as:

KL[p, q] = E_p[log(p(X)/q(X))]
         = -int_F p(x) log q(x) dr(x) + int_F p(x) log p(x) dr(x)
         = H[p, q] - H[p]

where F denotes the support of the random variable X ~ p, H[., .] denotes (Shannon) cross entropy, and H[.] denotes (Shannon) entropy.

Args:

Returns:

  • kl_divergence: self.dtype Tensor with shape [B1, ..., Bn] representing n different calculations of the Kullback-Leibler divergence.

log_cdf

log_cdf(
    value,
    name='log_cdf'
)

Log cumulative distribution function.

Given random variable X, the cumulative distribution function cdf is:

log_cdf(x) := Log[ P[X <= x] ]

Often, a numerical approximation can be used for log_cdf(x) that yields a more accurate answer than simply taking the logarithm of the cdf when x << -1.

Args:

  • value: float or double Tensor.
  • name: Python str prepended to names of ops created by this function.

Returns:

  • logcdf: a Tensor of shape sample_shape(x) + self.batch_shape with values of type self.dtype.

log_prob

log_prob(
    value,
    name='log_prob'
)

Log probability density/mass function.

Args:

  • value: float or double Tensor.
  • name: Python str prepended to names of ops created by this function.

Returns:

  • log_prob: a Tensor of shape sample_shape(x) + self.batch_shape with values of type self.dtype.

log_survival_function

log_survival_function(
    value,
    name='log_survival_function'
)

Log survival function.

Given random variable X, the survival function is defined:

log_survival_function(x) = Log[ P[X > x] ]
                         = Log[ 1 - P[X <= x] ]
                         = Log[ 1 - cdf(x) ]

Typically, different numerical approximations can be used for the log survival function, which are more accurate than 1 - cdf(x) when x >> 1.

Args:

  • value: float or double Tensor.
  • name: Python str prepended to names of ops created by this function.

Returns:

Tensor of shape sample_shape(x) + self.batch_shape with values of type self.dtype.

mean

mean(name='mean')

Mean.

mode

mode(name='mode')

Mode.

param_shapes

param_shapes(
    cls,
    sample_shape,
    name='DistributionParamShapes'
)

Shapes of parameters given the desired shape of a call to sample().

This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample().

Subclasses should override class method _param_shapes.

Args:

  • sample_shape: Tensor or python list/tuple. Desired shape of a call to sample().
  • name: name to prepend ops with.

Returns:

dict of parameter name to Tensor shapes.

param_static_shapes

param_static_shapes(
    cls,
    sample_shape
)

param_shapes with static (i.e. TensorShape) shapes.

This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample(). Assumes that the sample's shape is known statically.

Subclasses should override class method _param_shapes to return constant-valued tensors when constant values are fed.

Args:

  • sample_shape: TensorShape or python list/tuple. Desired shape of a call to sample().

Returns:

dict of parameter name to TensorShape.

Raises:

  • ValueError: if sample_shape is a TensorShape and is not fully defined.

posterior_marginals

posterior_marginals(x)

Run a Kalman smoother to return posterior mean and cov.

Note that the returned values smoothed_means depend on the observed time series x, while the smoothed_covs are independent of the observed series; i.e., they depend only on the model itself. This means that the mean values have shape concat([sample_shape(x), batch_shape, [num_timesteps, {latent/observation}_size]]), while the covariances have shape concat[(batch_shape, [num_timesteps, {latent/observation}_size, {latent/observation}_size]]), which does not depend on the sample shape.

This function only performs smoothing. If the user wants the intermediate values, which are returned by filtering pass forward_filter, one could get it by:

(log_likelihoods,
 filtered_means, filtered_covs,
 predicted_means, predicted_covs,
 observation_means, observation_covs) = model.forward_filter(x)
smoothed_means, smoothed_covs = model.backward_smoothing_pass(x)

where x is an observation sequence.

Args:

  • x: a float-type Tensor with rightmost dimensions [num_timesteps, observation_size] matching self.event_shape. Additional dimensions must match or be broadcastable to self.batch_shape; any further dimensions are interpreted as a sample shape.

Returns:

  • smoothed_means: Means of the per-timestep smoothed distributions over latent states, p(x{t} | x{:T}), as a Tensor of shape sample_shape(x) + batch_shape + [num_timesteps, observation_size].
  • smoothed_covs: Covariances of the per-timestep smoothed distributions over latent states, p(x{t} | x{:T}), as a Tensor of shape batch_shape + [num_timesteps, observation_size, observation_size].

prob

prob(
    value,
    name='prob'
)

Probability density/mass function.

Args:

  • value: float or double Tensor.
  • name: Python str prepended to names of ops created by this function.

Returns:

  • prob: a Tensor of shape sample_shape(x) + self.batch_shape with values of type self.dtype.

quantile

quantile(
    value,
    name='quantile'
)

Quantile function. Aka "inverse cdf" or "percent point function".

Given random variable X and p in [0, 1], the quantile is:

quantile(p) := x such that P[X <= x] == p

Args:

  • value: float or double Tensor.
  • name: Python str prepended to names of ops created by this function.

Returns:

  • quantile: a Tensor of shape sample_shape(x) + self.batch_shape with values of type self.dtype.

sample

sample(
    sample_shape=(),
    seed=None,
    name='sample'
)

Generate samples of the specified shape.

Note that a call to sample() without arguments will generate a single sample.

Args:

  • sample_shape: 0D or 1D int32 Tensor. Shape of the generated samples.
  • seed: Python integer seed for RNG
  • name: name to give to the op.

Returns:

  • samples: a Tensor with prepended dimensions sample_shape.

stddev

stddev(name='stddev')

Standard deviation.

Standard deviation is defined as,

stddev = E[(X - E[X])**2]**0.5

where X is the random variable associated with this distribution, E denotes expectation, and stddev.shape = batch_shape + event_shape.

Args:

  • name: Python str prepended to names of ops created by this function.

Returns:

  • stddev: Floating-point Tensor with shape identical to batch_shape + event_shape, i.e., the same shape as self.mean().

survival_function

survival_function(
    value,
    name='survival_function'
)

Survival function.

Given random variable X, the survival function is defined:

survival_function(x) = P[X > x]
                     = 1 - P[X <= x]
                     = 1 - cdf(x).

Args:

  • value: float or double Tensor.
  • name: Python str prepended to names of ops created by this function.

Returns:

Tensor of shape sample_shape(x) + self.batch_shape with values of type self.dtype.

variance

variance(name='variance')

Variance.

Variance is defined as,

Var = E[(X - E[X])**2]

where X is the random variable associated with this distribution, E denotes expectation, and Var.shape = batch_shape + event_shape.

Args:

  • name: Python str prepended to names of ops created by this function.

Returns:

  • variance: Floating-point Tensor with shape identical to batch_shape + event_shape, i.e., the same shape as self.mean().