markovflow.gauss_markov

Module representing a Gauss-Markov chain.

Module Contents

class GaussMarkovDistribution(name=None)

Bases: tf.Module, abc.ABC

Abstract class for representing a Gauss-Markov chain. Classes that extend this one (such as StateSpaceModel) represent a different parameterisation of the joint Gaussian distribution.

property event_shape → tf.Tensor

Return the shape of the event in the Gauss-Markov chain that is [num_transitions + 1, state_dim].

property batch_shape → tf.TensorShape

Return the shape of any leading dimensions in the Gauss-Markov chain that come before event_shape.

property state_dim → int

Return the state dimension of the Gauss-Markov chain.

property num_transitions → tf.Tensor

Return the number of transitions in the Gauss-Markov chain.

abstract _build_precision() → markovflow.block_tri_diag.SymmetricBlockTriDiagonal

Compute the compact banded representation of the precision matrix.

property precision → markovflow.block_tri_diag.SymmetricBlockTriDiagonal

Return the precision matrix of the joint Gaussian.

property marginal_means → tf.Tensor

Return the marginal means of the joint Gaussian.

Returns

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

property marginal_covariances → tf.Tensor

Return the marginal covariances of the joint Gaussian.

Returns

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

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

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 marginals → Tuple[tf.Tensor, tf.Tensor]

Return the means \(μₖ\) and the covariances \(Σₖₖ\) of the marginal distributions over consecutive states \(xₖ\).

Returns

The means and covariances, with respective shapes batch_shape + [num_transitions + 1, state_dim], batch_shape + [num_transitions + 1, state_dim, state_dim].

abstract sample(sample_shape: markovflow.base.SampleShape) → tf.Tensor

Sample trajectories from the distribution.

Parameters

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

Returns

The state samples, with shape sample_shape + self.batch_shape + self.event_shape.

abstract log_det_precision() → tf.Tensor

Calculate the log determinant of the precision matrix.

Returns

A tensor with shape batch_shape.

abstract log_pdf(states) → tf.Tensor

Return the value of the log of the PDF evaluated at states.

That is:

\[log p(x) = log p(x₀) + Σₖ log p(xₖ₊₁|xₖ) \verb|(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.

abstract create_trainable_copy() → GaussMarkovDistribution

Create a trainable version.

This is primarily for use with variational approaches where we want to optimise the parameters of the Gauss-Markov distribution.

Returns

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

abstract create_non_trainable_copy() → GaussMarkovDistribution

Create a non-trainable version.

Convert a trainable version of this class back to being non-trainable.

Returns

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

abstract kl_divergence(dist: GaussMarkovDistribution) → tf.Tensor

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

\[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 and \(Pᵢ\) are the banded precisions.

The KL divergence is then 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

The KL divergences, with shape self.batch_shape.

check_compatible(dist_1: GaussMarkovDistribution, dist_2: GaussMarkovDistribution) → None

Check that two GaussMarkovDistribution objects are compatible.

If not, raise an exception.