"""Kalman Filter, Kalman Smoother, and EM Algorithm for Linear-Gaussian Systems.

=====================================

Inference for Linear-Gaussian Systems

=====================================

This module implements the Kalman Filter, Kalman Smoother, and

EM Algorithm for Linear-Gaussian state space models

"""

import warnings

import numpy as np

from scipy import linalg

from .utils import (

array1d,

array2d,

check_random_state,

get_params,

log_multivariate_normal_density,

preprocess_arguments,

)

from .utils_numpy import newbyteorder

# Dimensionality of each Kalman Filter parameter for a single time step

DIM = {

"transition_matrices": 2,

"transition_offsets": 1,

"observation_matrices": 2,

"observation_offsets": 1,

"transition_covariance": 2,

"observation_covariance": 2,

"initial_state_mean": 1,

"initial_state_covariance": 2,

}

def _arg_or_default(arg, default, dim, name):

if arg is None:

result = np.asarray(default)

else:

result = arg

if len(result.shape) > dim:

raise ValueError(

("%s is not constant for all time." + " You must specify it manually.")

% (name,)

)

return result

def _determine_dimensionality(variables, default):

"""Derive the dimensionality of the state space.

Parameters

----------

variables : list of ({None, array}, conversion function, index)

variables, functions to convert them to arrays, and indices in those

arrays to derive dimensionality from.

default : {None, int}

default dimensionality to return if variables is empty

Returns

-------

dim : int

dimensionality of state space as derived from variables or default.

"""

# gather possible values based on the variables

candidates = []

for v, converter, idx in variables:

if v is not None:

v = converter(v)

candidates.append(v.shape[idx])

# also use the manually specified default

if default is not None:

candidates.append(default)

# ensure consistency of all derived values

if len(candidates) == 0:

return 1

else:

if not np.all(np.array(candidates) == candidates[0]):

raise ValueError(

"The shape of all "

+ "parameters is not consistent. "

+ "Please re-check their values."

)

return candidates[0]

def _last_dims(X, t, ndims=2):

"""Extract the final dimensions of `X`.

Extract the final `ndim` dimensions at index `t` if `X` has >= `ndim` + 1

dimensions, otherwise return `X`.

Parameters

----------

X : array with at least dimension `ndims`

t : int

index to use for the `ndims` + 1th dimension

ndims : int, optional

number of dimensions in the array desired

Returns

-------

Y : array with dimension `ndims`

the final `ndims` dimensions indexed by `t`

"""

X = np.asarray(X)

if len(X.shape) == ndims + 1:

return X[t]

elif len(X.shape) == ndims:

return X

else:

raise ValueError(

("X only has %d dimensions when %d" + " or more are required")

% (len(X.shape), ndims)

)

def _loglikelihoods(

observation_matrices,

observation_offsets,

observation_covariance,

predicted_state_means,

predicted_state_covariances,

observations,

):

"""Calculate log-likelihood of all observations.

Parameters

----------

observation_matrices : [n_timesteps, n_dim_obs, n_dim_state] or [n_dim_obs,

n_dim_state] array

observation matrices for t in [0...n_timesteps-1]

observation_offsets : [n_timesteps, n_dim_obs] or [n_dim_obs] array

offsets for observations for t = [0...n_timesteps-1]

observation_covariance : [n_dim_obs, n_dim_obs] array

covariance matrix for all observations

predicted_state_means : [n_timesteps, n_dim_state] array

mean of state at time t given observations from times

[0...t-1] for t in [0...n_timesteps-1]

predicted_state_covariances : [n_timesteps, n_dim_state, n_dim_state] array

covariance of state at time t given observations from times

[0...t-1] for t in [0...n_timesteps-1]

observations : [n_dim_obs] array

All observations. If `observations[t]` is a masked array and any of

its values are masked, the observation will be ignored.

Returns

-------

loglikelihoods: [n_timesteps] array

`loglikelihoods[t]` is the probability density of the observation

generated at time step t

"""

n_timesteps = observations.shape[0]

loglikelihoods = np.zeros(n_timesteps)

for t in range(n_timesteps):

observation = observations[t]

if not np.any(np.ma.getmask(observation)):

observation_matrix = _last_dims(observation_matrices, t)

observation_offset = _last_dims(observation_offsets, t, ndims=1)

predicted_state_mean = _last_dims(predicted_state_means, t, ndims=1)

predicted_state_covariance = _last_dims(predicted_state_covariances, t)

predicted_observation_mean = (

np.dot(observation_matrix, predicted_state_mean) + observation_offset

)

predicted_observation_covariance = (

np.dot(

observation_matrix,

np.dot(predicted_state_covariance, observation_matrix.T),

)

+ observation_covariance

)

lmnd_arr = log_multivariate_normal_density(

observation[np.newaxis, :],

predicted_observation_mean[np.newaxis, :],

predicted_observation_covariance[np.newaxis, :, :],

)

loglikelihoods[t] = lmnd_arr[0, 0]

return loglikelihoods

def _filter_predict(

transition_matrix,

transition_covariance,

transition_offset,

current_state_mean,

current_state_covariance,

):

r"""Calculate the mean and covariance of :math:`P(x_{t+1} | z_{0:t})`.

Using the mean and covariance of :math:`P(x_t | z_{0:t})`, calculate the

mean and covariance of :math:`P(x_{t+1} | z_{0:t})`.

Parameters

----------

transition_matrix : [n_dim_state, n_dim_state} array

state transition matrix from time t to t+1

transition_covariance : [n_dim_state, n_dim_state] array

covariance matrix for state transition from time t to t+1

transition_offset : [n_dim_state] array

offset for state transition from time t to t+1

current_state_mean: [n_dim_state] array

mean of state at time t given observations from times

[0...t]

current_state_covariance: [n_dim_state, n_dim_state] array

covariance of state at time t given observations from times

[0...t]

Returns

-------

predicted_state_mean : [n_dim_state] array

mean of state at time t+1 given observations from times [0...t]

predicted_state_covariance : [n_dim_state, n_dim_state] array

covariance of state at time t+1 given observations from times

[0...t]

"""

predicted_state_mean = (

np.dot(transition_matrix, current_state_mean) + transition_offset

)

predicted_state_covariance = (

np.dot(transition_matrix, np.dot(current_state_covariance, transition_matrix.T))

+ transition_covariance

)

return (predicted_state_mean, predicted_state_covariance)

def _filter_correct(

observation_matrix,

observation_covariance,

observation_offset,

predicted_state_mean,

predicted_state_covariance,

observation,

):

r"""Correct a predicted state with a Kalman Filter update.

Incorporate observation `observation` from time `t` to turn

:math:`P(x_t | z_{0:t-1})` into :math:`P(x_t | z_{0:t})`

Parameters

----------

observation_matrix : [n_dim_obs, n_dim_state] array

observation matrix for time t

observation_covariance : [n_dim_obs, n_dim_obs] array

covariance matrix for observation at time t

observation_offset : [n_dim_obs] array

offset for observation at time t

predicted_state_mean : [n_dim_state] array

mean of state at time t given observations from times

[0...t-1]

predicted_state_covariance : [n_dim_state, n_dim_state] array

covariance of state at time t given observations from times

[0...t-1]

observation : [n_dim_obs] array

observation at time t. If `observation` is a masked array and any of

its values are masked, the observation will be ignored.

Returns

-------

kalman_gain : [n_dim_state, n_dim_obs] array

Kalman gain matrix for time t

corrected_state_mean : [n_dim_state] array

mean of state at time t given observations from times

[0...t]

corrected_state_covariance : [n_dim_state, n_dim_state] array

covariance of state at time t given observations from times

[0...t]

"""

if not np.any(np.ma.getmask(observation)):

predicted_observation_mean = (

np.dot(observation_matrix, predicted_state_mean) + observation_offset

)

predicted_observation_covariance = (

np.dot(

observation_matrix,

np.dot(predicted_state_covariance, observation_matrix.T),

)

+ observation_covariance

)

kalman_gain = np.dot(

predicted_state_covariance,

np.dot(observation_matrix.T, linalg.pinv(predicted_observation_covariance)),

)

corrected_state_mean = predicted_state_mean + np.dot(

kalman_gain, observation - predicted_observation_mean

)

corrected_state_covariance = predicted_state_covariance - np.dot(

kalman_gain, np.dot(observation_matrix, predicted_state_covariance)

)

else:

n_dim_state = predicted_state_covariance.shape[0]

n_dim_obs = observation_matrix.shape[0]

kalman_gain = np.zeros((n_dim_state, n_dim_obs))

corrected_state_mean = predicted_state_mean

corrected_state_covariance = predicted_state_covariance

return (kalman_gain, corrected_state_mean, corrected_state_covariance)

def _filter(

transition_matrices,

observation_matrices,

transition_covariance,

observation_covariance,

transition_offsets,

observation_offsets,

initial_state_mean,

initial_state_covariance,

observations,

):

"""Apply the Kalman Filter.

Calculate posterior distribution over hidden states given observations up

to and including the current time step.

Parameters

----------

transition_matrices : [n_timesteps-1,n_dim_state,n_dim_state] or

[n_dim_state,n_dim_state] array-like

state transition matrices

observation_matrices : [n_timesteps, n_dim_obs, n_dim_state] or [n_dim_obs, \

n_dim_state] array-like

observation matrix

transition_covariance : [n_timesteps-1,n_dim_state,n_dim_state] or

[n_dim_state,n_dim_state] array-like

state transition covariance matrix

observation_covariance : [n_timesteps, n_dim_obs, n_dim_obs] or [n_dim_obs,

n_dim_obs] array-like

observation covariance matrix

transition_offsets : [n_timesteps-1, n_dim_state] or [n_dim_state] \

array-like

state offset

observation_offsets : [n_timesteps, n_dim_obs] or [n_dim_obs] array-like

observations for times [0...n_timesteps-1]

initial_state_mean : [n_dim_state] array-like

mean of initial state distribution

initial_state_covariance : [n_dim_state, n_dim_state] array-like

covariance of initial state distribution

observations : [n_timesteps, n_dim_obs] array

observations from times [0...n_timesteps-1]. If `observations` is a

masked array and any of `observations[t]` is masked, then

`observations[t]` will be treated as a missing observation.

Returns

-------

predicted_state_means : [n_timesteps, n_dim_state] array

`predicted_state_means[t]` = mean of hidden state at time t given

observations from times [0...t-1]

predicted_state_covariances : [n_timesteps, n_dim_state, n_dim_state] array

`predicted_state_covariances[t]` = covariance of hidden state at time t

given observations from times [0...t-1]

kalman_gains : [n_timesteps, n_dim_state] array

`kalman_gains[t]` = Kalman gain matrix for time t

filtered_state_means : [n_timesteps, n_dim_state] array

`filtered_state_means[t]` = mean of hidden state at time t given

observations from times [0...t]

filtered_state_covariances : [n_timesteps, n_dim_state] array

`filtered_state_covariances[t]` = covariance of hidden state at time t

given observations from times [0...t]

"""

n_timesteps = observations.shape[0]

n_dim_state = len(initial_state_mean)

n_dim_obs = observations.shape[1]

predicted_state_means = np.zeros((n_timesteps, n_dim_state))

predicted_state_covariances = np.zeros((n_timesteps, n_dim_state, n_dim_state))

kalman_gains = np.zeros((n_timesteps, n_dim_state, n_dim_obs))

filtered_state_means = np.zeros((n_timesteps, n_dim_state))

filtered_state_covariances = np.zeros((n_timesteps, n_dim_state, n_dim_state))

for t in range(n_timesteps):

if t == 0:

predicted_state_means[t] = initial_state_mean

predicted_state_covariances[t] = initial_state_covariance

else:

transition_matrix = _last_dims(transition_matrices, t - 1)

transition_covariance_t = _last_dims(transition_covariance, t - 1)

transition_offset = _last_dims(transition_offsets, t - 1, ndims=1)

predicted_state_means[t], predicted_state_covariances[t] = _filter_predict(

transition_matrix,

transition_covariance_t,

transition_offset,

filtered_state_means[t - 1],

filtered_state_covariances[t - 1],

)

observation_matrix = _last_dims(observation_matrices, t)

observation_covariance_t = _last_dims(observation_covariance, t)

observation_offset = _last_dims(observation_offsets, t, ndims=1)

(

kalman_gains[t],

filtered_state_means[t],

filtered_state_covariances[t],

) = _filter_correct(

observation_matrix,

observation_covariance_t,

observation_offset,

predicted_state_means[t],

predicted_state_covariances[t],

observations[t],

)

return (

predicted_state_means,

predicted_state_covariances,

kalman_gains,

filtered_state_means,

filtered_state_covariances,

)

def _smooth_update(

transition_matrix,

filtered_state_mean,

filtered_state_covariance,

predicted_state_mean,

predicted_state_covariance,

next_smoothed_state_mean,

next_smoothed_state_covariance,

):

r"""Correct a predicted state with a Kalman Smoother update.

Calculates posterior distribution of the hidden state at time `t` given the

observations all observations via Kalman Smoothing.

Parameters

----------

transition_matrix : [n_dim_state, n_dim_state] array

state transition matrix from time t to t+1

filtered_state_mean : [n_dim_state] array

mean of filtered state at time t given observations from

times [0...t]

filtered_state_covariance : [n_dim_state, n_dim_state] array

covariance of filtered state at time t given observations from

times [0...t]

predicted_state_mean : [n_dim_state] array

mean of filtered state at time t+1 given observations from

times [0...t]

predicted_state_covariance : [n_dim_state, n_dim_state] array

covariance of filtered state at time t+1 given observations from

times [0...t]

next_smoothed_state_mean : [n_dim_state] array

mean of smoothed state at time t+1 given observations from

times [0...n_timesteps-1]

next_smoothed_state_covariance : [n_dim_state, n_dim_state] array

covariance of smoothed state at time t+1 given observations from

times [0...n_timesteps-1]

Returns

-------

smoothed_state_mean : [n_dim_state] array

mean of smoothed state at time t given observations from times

[0...n_timesteps-1]

smoothed_state_covariance : [n_dim_state, n_dim_state] array

covariance of smoothed state at time t given observations from

times [0...n_timesteps-1]

kalman_smoothing_gain : [n_dim_state, n_dim_state] array

correction matrix for Kalman Smoothing at time t

"""

kalman_smoothing_gain = np.dot(

filtered_state_covariance,

np.dot(transition_matrix.T, linalg.pinv(predicted_state_covariance)),

)

smoothed_state_mean = filtered_state_mean + np.dot(

kalman_smoothing_gain, next_smoothed_state_mean - predicted_state_mean

)

smoothed_state_covariance = filtered_state_covariance + np.dot(

kalman_smoothing_gain,

np.dot(

(next_smoothed_state_covariance - predicted_state_covariance),

kalman_smoothing_gain.T,

),

)

return (smoothed_state_mean, smoothed_state_covariance, kalman_smoothing_gain)

def _smooth(

transition_matrices,

filtered_state_means,

filtered_state_covariances,

predicted_state_means,

predicted_state_covariances,

):

"""Apply the Kalman Smoother.

Estimate the hidden state at time for each time step given all

observations.

Parameters

----------

transition_matrices : [n_timesteps-1, n_dim_state, n_dim_state] or \

[n_dim_state, n_dim_state] array

`transition_matrices[t]` = transition matrix from time t to t+1

filtered_state_means : [n_timesteps, n_dim_state] array

`filtered_state_means[t]` = mean state estimate for time t given

observations from times [0...t]

filtered_state_covariances : [n_timesteps, n_dim_state, n_dim_state] array

`filtered_state_covariances[t]` = covariance of state estimate for time

t given observations from times [0...t]

predicted_state_means : [n_timesteps, n_dim_state] array

`predicted_state_means[t]` = mean state estimate for time t given

observations from times [0...t-1]

predicted_state_covariances : [n_timesteps, n_dim_state, n_dim_state] array

`predicted_state_covariances[t]` = covariance of state estimate for

time t given observations from times [0...t-1]

Returns

-------

smoothed_state_means : [n_timesteps, n_dim_state]

mean of hidden state distributions for times [0...n_timesteps-1] given

all observations

smoothed_state_covariances : [n_timesteps, n_dim_state, n_dim_state] array

covariance matrix of hidden state distributions for times

[0...n_timesteps-1] given all observations

kalman_smoothing_gains : [n_timesteps-1, n_dim_state, n_dim_state] array

Kalman Smoothing correction matrices for times [0...n_timesteps-2]

"""

n_timesteps, n_dim_state = filtered_state_means.shape

smoothed_state_means = np.zeros((n_timesteps, n_dim_state))

smoothed_state_covariances = np.zeros((n_timesteps, n_dim_state, n_dim_state))

kalman_smoothing_gains = np.zeros((n_timesteps - 1, n_dim_state, n_dim_state))

smoothed_state_means[-1] = filtered_state_means[-1]

smoothed_state_covariances[-1] = filtered_state_covariances[-1]

for t in reversed(range(n_timesteps - 1)):

transition_matrix = _last_dims(transition_matrices, t)

(

smoothed_state_means[t],

smoothed_state_covariances[t],

kalman_smoothing_gains[t],

) = _smooth_update(

transition_matrix,

filtered_state_means[t],

filtered_state_covariances[t],

predicted_state_means[t + 1],

predicted_state_covariances[t + 1],

smoothed_state_means[t + 1],

smoothed_state_covariances[t + 1],

)

return (smoothed_state_means, smoothed_state_covariances, kalman_smoothing_gains)

def _smooth_pair(smoothed_state_covariances, kalman_smoothing_gain):

r"""Calculate pairwise covariance between hidden states.

Calculate covariance between hidden states at :math:`t` and :math:`t-1` for

all time step pairs

Parameters

----------

smoothed_state_covariances : [n_timesteps, n_dim_state, n_dim_state] array

covariance of hidden state given all observations

kalman_smoothing_gain : [n_timesteps-1, n_dim_state, n_dim_state]

Correction matrices from Kalman Smoothing

Returns

-------

pairwise_covariances : [n_timesteps, n_dim_state, n_dim_state] array

Covariance between hidden states at times t and t-1 for t =

[1...n_timesteps-1]. Time 0 is ignored.

"""

n_timesteps, n_dim_state, _ = smoothed_state_covariances.shape

pairwise_covariances = np.zeros((n_timesteps, n_dim_state, n_dim_state))

for t in range(1, n_timesteps):

pairwise_covariances[t] = np.dot(

smoothed_state_covariances[t], kalman_smoothing_gain[t - 1].T

)

return pairwise_covariances

def _em(

observations,

transition_offsets,

observation_offsets,

smoothed_state_means,

smoothed_state_covariances,

pairwise_covariances,

given=None,

):

"""Apply the EM Algorithm to the Linear-Gaussian model.

Estimate Linear-Gaussian model parameters by maximizing the expected log

likelihood of all observations.

Parameters

----------

observations : [n_timesteps, n_dim_obs] array

observations for times [0...n_timesteps-1]. If observations is a

masked array and any of observations[t] is masked, then it will be

treated as a missing observation.

transition_offsets : [n_dim_state] or [n_timesteps-1, n_dim_state] array

transition offset

observation_offsets : [n_dim_obs] or [n_timesteps, n_dim_obs] array

observation offsets

smoothed_state_means : [n_timesteps, n_dim_state] array

smoothed_state_means[t] = mean of state at time t given all

observations

smoothed_state_covariances : [n_timesteps, n_dim_state, n_dim_state] array

smoothed_state_covariances[t] = covariance of state at time t given all

observations

pairwise_covariances : [n_timesteps, n_dim_state, n_dim_state] array

pairwise_covariances[t] = covariance between states at times t and

t-1 given all observations. pairwise_covariances[0] is ignored.

given: dict

if one of the variables EM is capable of predicting is in given, then

that value will be used and EM will not attempt to estimate it. e.g.,

if 'observation_matrix' is in given, observation_matrix will not be

estimated and given['observation_matrix'] will be returned in its

place.

Returns

-------

transition_matrix : [n_dim_state, n_dim_state] array

estimated transition matrix

observation_matrix : [n_dim_obs, n_dim_state] array

estimated observation matrix

transition_offsets : [n_dim_state] array

estimated transition offset

observation_offsets : [n_dim_obs] array

estimated observation offset

transition_covariance : [n_dim_state, n_dim_state] array

estimated covariance matrix for state transitions

observation_covariance : [n_dim_obs, n_dim_obs] array

estimated covariance matrix for observations

initial_state_mean : [n_dim_state] array

estimated mean of initial state distribution

initial_state_covariance : [n_dim_state] array

estimated covariance of initial state distribution

"""

if given is None:

given = {}

if "observation_matrices" in given:

observation_matrix = given["observation_matrices"]

else:

observation_matrix = _em_observation_matrix(

observations,

observation_offsets,

smoothed_state_means,

smoothed_state_covariances,

)

if "observation_covariance" in given:

observation_covariance = given["observation_covariance"]

else:

observation_covariance = _em_observation_covariance(

observations,

observation_offsets,

observation_matrix,

smoothed_state_means,

smoothed_state_covariances,

)

if "transition_matrices" in given:

transition_matrix = given["transition_matrices"]

else:

transition_matrix = _em_transition_matrix(

transition_offsets,

smoothed_state_means,

smoothed_state_covariances,

pairwise_covariances,

)

if "transition_covariance" in given:

transition_covariance = given["transition_covariance"]

else:

transition_covariance = _em_transition_covariance(

transition_matrix,

transition_offsets,

smoothed_state_means,

smoothed_state_covariances,

pairwise_covariances,

)

if "initial_state_mean" in given:

initial_state_mean = given["initial_state_mean"]

else:

initial_state_mean = _em_initial_state_mean(smoothed_state_means)

if "initial_state_covariance" in given:

initial_state_covariance = given["initial_state_covariance"]

else:

initial_state_covariance = _em_initial_state_covariance(

initial_state_mean, smoothed_state_means, smoothed_state_covariances

)

if "transition_offsets" in given:

transition_offset = given["transition_offsets"]

else:

transition_offset = _em_transition_offset(

transition_matrix, smoothed_state_means

)

if "observation_offsets" in given:

observation_offset = given["observation_offsets"]

else:

observation_offset = _em_observation_offset(

observation_matrix, smoothed_state_means, observations

)

return (

transition_matrix,

observation_matrix,

transition_offset,

observation_offset,

transition_covariance,

observation_covariance,

initial_state_mean,

initial_state_covariance,

)

def _em_observation_matrix(

observations, observation_offsets, smoothed_state_means, smoothed_state_covariances

):

r"""Apply the EM algorithm to parameter `observation_matrix`.

Maximize expected log likelihood of observations with respect to the

observation matrix `observation_matrix`.

.. math::

C &= ( \sum_{t=0}^{T-1} (z_t - d_t) \mathbb{E}[x_t]^T )

( \sum_{t=0}^{T-1} \mathbb{E}[x_t x_t^T] )^-1

"""

_, n_dim_state = smoothed_state_means.shape

n_timesteps, n_dim_obs = observations.shape

res1 = np.zeros((n_dim_obs, n_dim_state))

res2 = np.zeros((n_dim_state, n_dim_state))

for t in range(n_timesteps):

if not np.any(np.ma.getmask(observations[t])):

observation_offset = _last_dims(observation_offsets, t, ndims=1)

res1 += np.outer(

observations[t] - observation_offset, smoothed_state_means[t]

)

res2 += smoothed_state_covariances[t] + np.outer(

smoothed_state_means[t], smoothed_state_means[t]

)

return np.dot(res1, linalg.pinv(res2))

def _em_observation_covariance(

observations,

observation_offsets,

observation_matrices,

smoothed_state_means,

smoothed_state_covariances,

):

r"""Apply the EM algorithm to parameter `observation_covariance`.

Maximize expected log likelihood of observations with respect to the

observation covariance matrix `observation_covariance`.

.. math::

R &= \frac{1}{T} \sum_{t=0}^{T-1}

[z_t - C_t \mathbb{E}[x_t] - d_t]

[z_t - C_t \mathbb{E}[x_t] - d_t]^T

+ C_t Var(x_t) C_t^T

"""

_, n_dim_state = smoothed_state_means.shape

n_timesteps, n_dim_obs = observations.shape

res = np.zeros((n_dim_obs, n_dim_obs))

n_obs = 0

for t in range(n_timesteps):

if not np.any(np.ma.getmask(observations[t])):

observation_matrix = _last_dims(observation_matrices, t)

observation_offset = _last_dims(observation_offsets, t, ndims=1)

err = (

observations[t]

- np.dot(observation_matrix, smoothed_state_means[t])

- observation_offset

)

res += np.outer(err, err) + np.dot(

observation_matrix,

np.dot(smoothed_state_covariances[t], observation_matrix.T),

)

n_obs += 1

if n_obs > 0:

return (1.0 / n_obs) * res

else:

return res

def _em_transition_matrix(

transition_offsets,

smoothed_state_means,

smoothed_state_covariances,

pairwise_covariances,

):

r"""Apply the EM algorithm to parameter `transition_matrix`.

Maximize expected log likelihood of observations with respect to the state

transition matrix `transition_matrix`.

.. math::

A &= ( \sum_{t=1}^{T-1} \mathbb{E}[x_t x_{t-1}^{T}]

- b_{t-1} \mathbb{E}[x_{t-1}]^T )

( \sum_{t=1}^{T-1} \mathbb{E}[x_{t-1} x_{t-1}^T] )^{-1}

"""

n_timesteps, n_dim_state, _ = smoothed_state_covariances.shape

res1 = np.zeros((n_dim_state, n_dim_state))

res2 = np.zeros((n_dim_state, n_dim_state))

for t in range(1, n_timesteps):

transition_offset = _last_dims(transition_offsets, t - 1, ndims=1)

res1 += (

pairwise_covariances[t]

+ np.outer(smoothed_state_means[t], smoothed_state_means[t - 1])

- np.outer(transition_offset, smoothed_state_means[t - 1])

)

res2 += smoothed_state_covariances[t - 1] + np.outer(

smoothed_state_means[t - 1], smoothed_state_means[t - 1]

)

return np.dot(res1, linalg.pinv(res2))

def _em_transition_covariance(

transition_matrices,

transition_offsets,

smoothed_state_means,

smoothed_state_covariances,

pairwise_covariances,

):

r"""Apply the EM algorithm to parameter `transition_covariance`.

Maximize expected log likelihood of observations with respect to the

transition covariance matrix `transition_covariance`.

.. math::

Q &= \frac{1}{T-1} \sum_{t=0}^{T-2}

(\mathbb{E}[x_{t+1}] - A_t \mathbb{E}[x_t] - b_t)

(\mathbb{E}[x_{t+1}] - A_t \mathbb{E}[x_t] - b_t)^T

+ A_t Var(x_t) A_t^T + Var(x_{t+1})

- Cov(x_{t+1}, x_t) A_t^T - A_t Cov(x_t, x_{t+1})

"""

n_timesteps, n_dim_state, _ = smoothed_state_covariances.shape

res = np.zeros((n_dim_state, n_dim_state))

for t in range(n_timesteps - 1):

transition_matrix = _last_dims(transition_matrices, t)

transition_offset = _last_dims(transition_offsets, t, ndims=1)

err = (

smoothed_state_means[t + 1]

- np.dot(transition_matrix, smoothed_state_means[t])

- transition_offset

)

Vt1t_A = np.dot(pairwise_covariances[t + 1], transition_matrix.T)

res += (

np.outer(err, err)

+ np.dot(

transition_matrix,

np.dot(smoothed_state_covariances[t], transition_matrix.T),

)

+ smoothed_state_covariances[t + 1]

- Vt1t_A

- Vt1t_A.T

)

if n_timesteps > 1:

return (1.0 / (n_timesteps - 1)) * res

else:

return res

def _em_initial_state_mean(smoothed_state_means):

r"""Apply the EM algorithm to parameter `initial_state_mean`.

Maximize expected log likelihood of observations with respect to the

initial state distribution mean `initial_state_mean`.

.. math::

\mu_0 = \mathbb{E}[x_0]

"""

return smoothed_state_means[0]

def _em_initial_state_covariance(

initial_state_mean, smoothed_state_means, smoothed_state_covariances

):

r"""Apply the EM algorithm to parameter `initial_state_covariance`.

Maximize expected log likelihood of observations with respect to the

covariance of the initial state distribution `initial_state_covariance`.

.. math::

\Sigma_0 = \mathbb{E}[x_0, x_0^T] - mu_0 \mathbb{E}[x_0]^T

- \mathbb{E}[x_0] mu_0^T + mu_0 mu_0^T

"""

x0 = smoothed_state_means[0]

x0_x0 = smoothed_state_covariances[0] + np.outer(x0, x0)

return (

x0_x0

- np.outer(initial_state_mean, x0)

- np.outer(x0, initial_state_mean)

+ np.outer(initial_state_mean, initial_state_mean)

)

def _em_transition_offset(transition_matrices, smoothed_state_means):

r"""Apply the EM algorithm to parameter `transition_offset`.

Maximize expected log likelihood of observations with respect to the

state transition offset `transition_offset`.

.. math::

b = \frac{1}{T-1} \sum_{t=1}^{T-1}

\mathbb{E}[x_t] - A_{t-1} \mathbb{E}[x_{t-1}]

"""

n_timesteps, n_dim_state = smoothed_state_means.shape

transition_offset = np.zeros(n_dim_state)

for t in range(1, n_timesteps):

transition_matrix = _last_dims(transition_matrices, t - 1)

transition_offset += smoothed_state_means[t] - np.dot(

transition_matrix, smoothed_state_means[t - 1]

)

if n_timesteps > 1:

return (1.0 / (n_timesteps - 1)) * transition_offset

else:

return np.zeros(n_dim_state)

def _em_observation_offset(observation_matrices, smoothed_state_means, observations):

r"""Apply the EM algorithm to parameter `observation_offset`.

Maximize expected log likelihood of observations with respect to the

observation offset `observation_offset`.

.. math::

d = \frac{1}{T} \sum_{t=0}^{T-1} z_t - C_{t} \mathbb{E}[x_{t}]

"""

n_timesteps, n_dim_obs = observations.shape

observation_offset = np.zeros(n_dim_obs)

n_obs = 0

for t in range(n_timesteps):

if not np.any(np.ma.getmask(observations[t])):

observation_matrix = _last_dims(observation_matrices, t)

observation_offset += observations[t] - np.dot(

observation_matrix, smoothed_state_means[t]

)

n_obs += 1

if n_obs > 0:

return (1.0 / n_obs) * observation_offset

else:

return observation_offset

class KalmanFilter:

r"""Implements the Kalman Filter, Kalman Smoother, and EM algorithm.

This class implements the Kalman Filter, Kalman Smoother, and EM Algorithm

for a Linear Gaussian model specified by,

.. math::

x_{t+1} &= A_{t} x_{t} + b_{t} + \\text{Normal}(0, Q_{t}) \\\\

z_{t} &= C_{t} x_{t} + d_{t} + \\text{Normal}(0, R_{t})

The Kalman Filter is an algorithm designed to estimate

:math:`P(x_t | z_{0:t})`. As all state transitions and observations are

linear with Gaussian distributed noise, these distributions can be

represented exactly as Gaussian distributions with mean

`filtered_state_means[t]` and covariances `filtered_state_covariances[t]`.

Similarly, the Kalman Smoother is an algorithm designed to estimate

:math:`P(x_t | z_{0:T-1})`.

The EM algorithm aims to find for

:math:`\\theta = (A, b, C, d, Q, R, \\mu_0, \\Sigma_0)`

.. math::

\\max_{\\theta} P(z_{0:T-1}; \\theta)

If we define :math:`L(x_{0:T-1},\\theta) = \\log P(z_{0:T-1}, x_{0:T-1};

\\theta)`, then the EM algorithm works by iteratively finding,

.. math::

P(x_{0:T-1} | z_{0:T-1}, \\theta_i)

then by maximizing,