markovflow.state_space_model

Module containing a state space model.

Module Contents

class StateSpaceModel(initial_mean: gpflow.base.TensorType, chol_initial_covariance: gpflow.base.TensorType, state_transitions: gpflow.base.TensorType, state_offsets: gpflow.base.TensorType, chol_process_covariances: gpflow.base.TensorType)[source]

Bases: markovflow.gauss_markov.GaussMarkovDistribution

Implements a state space model. This has the following form:

\[xₖ₊₁ = Aₖ xₖ + bₖ + qₖ\]

…where:

  • \(qₖ ~ 𝓝(0, Qₖ)\)

  • \(x₀ ~ 𝓝(μ₀, P₀)\)

  • \(xₖ ∈ ℝ^d\)

  • \(bₖ ∈ ℝ^d\)

  • \(Aₖ ∈ ℝ^{d × d}\)

  • \(Qₖ ∈ ℝ^{d × d}\)

  • \(μ₀ ∈ ℝ^{d × 1}\)

  • \(P₀ ∈ ℝ^{d × d}\)

The key reference is:

@inproceedings{grigorievskiy2017parallelizable,
    title={Parallelizable sparse inverse formulation Gaussian processes (SpInGP)},
    author={Grigorievskiy, Alexander and Lawrence, Neil and S{"a}rkk{"a}, Simo},
    booktitle={Int'l Workshop on Machine Learning for Signal Processing (MLSP)},
    pages={1--6},
    year={2017},
    organization={IEEE}
}

The model samples \(x₀\) with an initial Gaussian distribution in \(ℝ^d\) (in code \(d\) is state_dim).

The model then proceeds for \(n\) (num_transitions) to generate \([x₁, ... xₙ]\), according to the formula above. The marginal distribution of samples at a point \(k\) is a Gaussian with mean \(μₖ, Pₖ\).

This class allows the user to generate samples from this process as well as to calculate the marginal distributions for each transition.

Parameters

  • initial_mean – A TensorType containing the initial mean, with shape batch_shape + [state_dim].

  • chol_initial_covariance – A TensorType containing the Cholesky of the initial covariance, with shape batch_shape + [state_dim, state_dim]. That is, unless the initial covariance is zero, in which case it is zero.

  • state_transitions – A TensorType containing state transition matrices, with shape batch_shape + [num_transitions, state_dim, state_dim].

  • state_offsets – A TensorType containing the process means bₖ, with shape batch_shape + [num_transitions, state_dim].

  • chol_process_covariances – A TensorType containing the Cholesky of the noise covariance matrices, with shape batch_shape + [num_transitions, state_dim, state_dim]. That is, unless the noise covariance is zero, in which case it is zero.

property event_shape → tf.Tensor[source]

Return the shape of the event.

Returns

The shape is [num_transitions + 1, state_dim].

property batch_shape → tf.TensorShape[source]

Return the shape of any leading dimensions that come before event_shape.

property state_dim → int[source]

Return the state dimension.

property num_transitions → tf.Tensor[source]

Return the number of transitions.

property cholesky_process_covariances → gpflow.base.TensorType[source]

Return the Cholesky of \([Q₁, Q₂, ....]\).

Returns

A TensorType with shape [... num_transitions, state_dim, state_dim].

property cholesky_initial_covariance → gpflow.base.TensorType[source]

Return the Cholesky of \(P₀\).

Returns

A TensorType with shape [..., state_dim, state_dim].

property initial_covariance → tf.Tensor[source]

Return \(P₀\).

Returns

A TensorType with shape [..., state_dim, state_dim].

property concatenated_cholesky_process_covariance → tf.Tensor[source]

Return the Cholesky of \([P₀, Q₁, Q₂, ....]\).

Returns

A tensor with shape [... num_transitions + 1, state_dim, state_dim].

property state_offsets → gpflow.base.TensorType[source]

Return the state offsets \([b₁, b₂, ....]\).

Returns

A TensorType with shape [..., num_transitions, state_dim].

property initial_mean → gpflow.base.TensorType[source]

Return the initial mean \(μ₀\).

Returns

A TensorType with shape [..., state_dim].

property concatenated_state_offsets → tf.Tensor[source]

Return the concatenated state offsets \([μ₀, b₁, b₂, ....]\).

Returns

A tensor with shape [... num_transitions + 1, state_dim].

property state_transitions → gpflow.base.TensorType[source]

Return the concatenated state offsets \([A₀, A₁, A₂, ....]\).

Returns

A TensorType with shape [... num_transitions, state_dim, state_dim].

property marginal_means → tf.Tensor[source]

Return the mean of the marginal distributions at each time point. If:

\[xₖ ~ 𝓝(μₖ, Kₖₖ)\]

…then return \(μₖ\).

If we let the concatenated state offsets be \(m = [μ₀, b₁, b₂, ....]\) and \(A\) be defined as in equation (5) of the SpInGP paper (see class docstring), then:

\[μ = A m = (A⁻¹)⁻¹ m\]

…which we can do quickly using a_inv_block().

Returns

The marginal means of the joint Gaussian, with shape batch_shape + [num_transitions + 1, state_dim].

property marginal_covariances → tf.Tensor[source]

Return the ordered covariances \(Σₖₖ\) of the multivariate normal marginal distributions over consecutive states \(xₖ\).

Returns

The marginal covariances of the joint Gaussian, with shape batch_shape + [num_transitions + 1, state_dim, state_dim].

covariance_blocks() → Tuple[tf.Tensor, tf.Tensor][source]

Return the diagonal and lower off-diagonal blocks of the covariance.

Returns

A tuple of tensors with respective shapes batch_shape + [num_transitions + 1, state_dim], batch_shape + [num_transitions, state_dim, state_dim].

property a_inv_block → markovflow.block_tri_diag.LowerTriangularBlockTriDiagonal[source]

Return \(A⁻¹\).

This has the form:

A⁻¹ =  [ I             ]
       [-A₁, I         ]
       [    -A₂, I     ]
       [         ᨞  ᨞  ]
       [         -Aₙ, I]

…where \([A₁, ..., Aₙ]\) are the state transition matrices.

sample(sample_shape: markovflow.base.SampleShape) → tf.Tensor[source]

Return sample trajectories.

Parameters

sample_shape – The shape (and hence number of) trajectories to sample from the state space model.

Returns

A tensor containing state samples, with shape sample_shape + self.batch_shape + self.event_shape.

subsequent_covariances(marginal_covariances: tf.Tensor) → tf.Tensor[source]

For each pair of subsequent states \(xₖ, xₖ₊₁\), return the covariance of their joint distribution. That is:

\[Cov(xₖ₊₁, xₖ) = AₖPₖ\]
Parameters

marginal_covariances – The marginal covariances of each state in the model, with shape batch_shape + [num_transitions + 1, state_dim, state_dim].

Returns

The covariance between subsequent state, with shape batch_shape + [num_transitions, state_dim, state_dim].

log_det_precision() → tf.Tensor[source]

Calculate the log determinant of the precision matrix. This uses the precision as defined in the SpInGP paper (see class summary above).

Precision is defined as:

\[K⁻¹ = (AQAᵀ)⁻¹\]

so:

\[\begin{split}log |K⁻¹| &= log | Q⁻¹ | (since |A| = 1)\\ &= - log |P₀| - Σₜ log |Qₜ|\\ &= - 2 * (log |chol_P₀| + Σₜ log |chol_Qₜ|)\end{split}\]
Returns

A tensor with shape batch_shape.

create_non_trainable_copy() → StateSpaceModel[source]

Create a non-trainable version of GaussMarkovDistribution.

This is to convert a trainable version of this class back to being non-trainable.

Returns

A Gauss-Markov distribution that is a copy of this one.

create_trainable_copy() → StateSpaceModel[source]

Create a trainable version of this state space model.

This is primarily for use with variational approaches where we want to optimise the parameters of a state space model that is initialised from a prior state space model.

The initial mean and state transitions are the same.

The initial and process covariances are ‘flattened’. Since they are lower triangular, we only want to parametrise this part of the matrix. For this purpose we use the params.triangular constraint which is the tfp.bijectors.FillTriangular bijector that converts between a triangular matrix \([dim, dim]\) and a flattened vector of shape \([dim (dim + 1) / 2]\).

Returns

A state space model that is a copy of this one and a dataclass containing the variables that can be trained.

_build_precision() → markovflow.block_tri_diag.SymmetricBlockTriDiagonal[source]

Compute the compact banded representation of the Precision matrix using state space model parameters.

We construct matrix:

K⁻¹ = A⁻ᵀQ⁻¹A⁻¹

Using Q⁻¹ and A⁻¹ defined in equations (6) and (8) in the SpInGP paper (see class docstring)

It can be shown that

K⁻¹ = | P₀⁻¹ + A₁ᵀ Q₁⁻¹ A₁ | -A₁ᵀ Q₁⁻¹ | 0…
-Q₁⁻¹ A₁ | Q₁⁻¹ + A₂ᵀ Q₂⁻¹ A₂ | -A₂ᵀ Q₂⁻¹ | 0…
0 | -Q₂⁻¹ A₂ | Q₂⁻¹ + A₃ᵀ Q₃⁻¹ A₃ | -A₃ᵀ Q₃⁻¹| 0…
Returns

The precision as a SymmetricBlockTriDiagonal object

_log_pdf_factors(states: tf.Tensor) → tf.Tensor[source]

Return the value of the log of the factors of the probability density function (PDF) evaluated at a state trajectory:

[log p(x₀), log p(x₁|x₀), ..., log p(xₖ₊₁|xₖ)]

…with x₀ ~ 𝓝(μ₀, P₀) and xₖ₊₁|xₖ ~ 𝓝(Aₖ xₖ + bₖ, Qₖ)

States

The state trajectory has shape: sample_shape + self.batch_shape + self.event_shape

Returns

The log PDF of the factors with shape: sample_shape + self.batch_shape + [self.num_transitions + 1]

log_pdf(states) → tf.Tensor[source]

Return the value of the log of the probability density function (PDF) evaluated at states. That is:

\[log p(x) = log p(x₀) + Σₖ log p(xₖ₊₁|xₖ) (for 0 ⩽ k < n)\]
Parameters

states – The state trajectory, with shape sample_shape + self.batch_shape + self.event_shape.

Returns

The log PDF, with shape sample_shape + self.batch_shape.

kl_divergence(dist: markovflow.gauss_markov.GaussMarkovDistribution) → tf.Tensor[source]

Return the KL divergence of the current Gauss-Markov distribution from the specified input dist. That is:

\[KL(dist₁ ∥ dist₂)\]

To do so we first compute the marginal distributions from the Gauss-Markov form:

\[\begin{split}dist₁ = 𝓝(μ₁, P⁻¹₁)\\ dist₂ = 𝓝(μ₂, P⁻¹₂)\end{split}\]

…where:

  • \(μᵢ\) are the marginal means

  • \(Pᵢ\) are the banded precisions

The KL divergence is thus given by:

\[KL(dist₁ ∥ dist₂) = ½(tr(P₂P₁⁻¹) + (μ₂ - μ₁)ᵀP₂(μ₂ - μ₁) - N - log(|P₂|) + log(|P₁|))\]

…where \(N = (\verb |num_transitions| + 1) * \verb |state_dim|\) (that is, the dimensionality of the Gaussian).

Parameters

dist – Another similarly parameterised Gauss-Markov distribution.

Returns

A tensor of the KL divergences, with shape self.batch_shape.

normalizer()[source]

Conputes the normalizer Page 36 of Thang Bui :return:

state_space_model_from_covariances(initial_mean: tf.Tensor, initial_covariance: tf.Tensor, state_transitions: tf.Tensor, state_offsets: tf.Tensor, process_covariances: tf.Tensor) → StateSpaceModel[source]

Construct a state space model using the full covariance matrices for convenience.

Parameters

  • initial_mean – The initial mean, with shape batch_shape + [state_dim].

  • initial_covariance – Initial covariance, with shape batch_shape + [state_dim, state_dim].

  • state_transitions – State transition matrices, with shape batch_shape + [num_transitions, state_dim, state_dim].

  • state_offsets – The process means \(bₖ\), with shape batch_shape + [num_transitions, state_dim].

  • process_covariances – Noise covariance matrices, with shape batch_shape + [num_transitions, state_dim, state_dim].