Time-warped autoregressive HMM (TWARHMM)

TWARHMM Models

Model classes for time warped ARHMMs.

class ssm.twarhmm.models.GaussianTWARHMM(num_discrete_states, time_constants, num_emission_dims=None, initial_state_probs=None, discrete_state_transition_matrix=None, time_constant_stay_probability=0.98, emission_weights=None, emission_biases=None, emission_covariances=None, emission_prior=None, seed=None)

Gaussian time-warped autoregressive HMM.

This model has the following likelihood,

x_t - x_{t-1} \mid x_{t-1}, z_t \sim \mathcal{N}\left(\tau_t^{-1} (A_{z_t} x_{t-1} + b_{z_t}), \tau_t^{-1} Q_{z_t} \right).

When the time constant is small, dynamics are faster. When the time constant is large, dynamics are slower.

Technically, the time-warped ARHMM is a special case of a factorial HMM, since we model the time-constants as taking values in a discrete set and following a Markov transition model,

p(z_t, \tau_t \mid z_{t-1}, \tau_{t-1}) = p(z_t \mid z_{t-1}) \times p(\tau_t \mid \tau_{t-1})

Parameters

  • num_discrete_states (int) – number of discrete latent states

  • time_constants (np.ndarray) – discrete non-negative values for the possible time constants.

  • num_emission_dims (int, optional) – emissions dimensionality. Defaults to None.

  • initial_state_probs (np.ndarray, optional) – initial discrete state probabilities with shape (\text{num\_discrete\_states},). Defaults to None.

  • discrete_state_transition_matrix (np.ndarray, optional) – transition matrix for discrete states with shape (\text{num\_discrete\_states}, \text{num\_discrete\_states}). Defaults to None.

  • time_constant_stay_probability (float, optional) – defines transition matrix for time constants. Probability of using the same time constant as in the previous time step. Other transition probabilities are calculated uniformly. Defaults to 0.98.

  • emission_weights (np.ndarray, optional) – emission weights A_{z_t} with shape (\text{num\_discrete\_states}, \text{emissions\_dim}, \text{emissions\_dim} * \text{num\_lags}). Defaults to None.

  • emission_biases (np.ndarray, optional) – emission biases b_{z_t} with shape (\text{num\_discrete\_states}, \text{emissions\_dim}). Defaults to None.

  • emission_covariances (np.ndarray, optional) – emission covariance Q_{z_t} with shape (\text{num\_discrete\_states}, \text{emissions\_dim}, \text{emissions\_dim}). Defaults to None.

  • emission_prior (GaussianLinearRegressionPrior, optional) – prior on emissions. Defaults to None.

  • seed (jr.PRNGKey, optional) – random seed. Defaults to None.

initialize(key, data, covariates=None, metadata=None, method='kmeans')

Initialize the model parameters by performing an M-step with state assignments determined by the specified method (random or kmeans).

Parameters

  • dataset (np.ndarray) – array of observed data of shape ( ext{batch} , ext{num\_timesteps} , ext{emissions\_dim})

  • key (jr.PRNGKey) – random seed

  • method (str, optional) – state assignment method. One of “random” or “kmeans”. Defaults to “kmeans”.

  • data (Array) –

Raises

ValueError – when initialize method is not recognized

Return type

None

TWARHMM Components

TWARHMM Emissions

class ssm.twarhmm.emissions.TimeWarpedAutoregressiveEmissions(num_discrete_states, time_constants, weights=None, biases=None, scale_trils=None, emissions_distribution_prior=None)

Emissions class for a Time-Warped Autoregressive HMM (TWARHMM).

Note: We (currently) only support lag-1 ARHMMs.

Optionally takes in a prior distribution.

x_t - x_{t-1} \mid x_{t-1}, z_t \sim \mathcal{N}\left(\tau_t^{-1} (A_{z_t} x_{t-1} + b_{z_t}), \tau_t^{-1} Q_{z_t} \right).

Parameters

  • num_discrete_states (int) – number of discrete latent states

  • time_constants (np.ndarray) – discrete non-negative values for the possible time constants

  • weights (np.ndarray, optional) – state-based weight matrix for Gaussian linear regression of shape (\text{num\_discrete\_states}, \text{emissions\_dim}, \text{emissions\_dim}). Note: dynamics are I + \frac{1}{\tau} A where A are the weights. Defaults to None.

  • biases (np.ndarray, optional) – state-based bias vector for Gaussian linear regression of shape (\text{num\_discrete\_states}, \text{emissions\_dim}). Note: dynamics are \frac{1}{\tau} b where b are the biases. Defaults to None.

  • scale_trils (np.ndarray, optional) – state-based scale_trils for Gaussian linear regression of shape (\text{num\_discrete\_states}, \text{emissions\_dim}, \text{emissions\_dim}). Note: scale_trils are \frac{1}{\tau} \texttt{chol}(Q) where Q are the covariance matrices. Defaults to None.

  • emissions_distribution_prior (ssmd.MatrixNormalInverseWishart, optional) – emissions prior distribution. Defaults to None.

distribution(state, covariates=None, metadata=None, history=None)

Returns the emissions distribution conditioned on a given state.

Parameters

  • state (int) – latent state

  • covariates (np.ndarray, optional) – optional covariates. Not yet supported. Defaults to None.

  • history (Optional[Array]) –

Returns

emissions_distribution (GaussianLinearRegression) – the emissions distribution

Return type

GaussianLinearRegression

log_likelihoods(data, covariates=None, metadata=None)

Compute log p(x_t | z_t=k) for all t and k.

m_step(dataset, posteriors, covariates=None, metadata=None)

Update the distribution with an M step.

The parameters are (A_k, b_k, Q_k) for each of the discrete states. The key idea is that once we fix the time-warping constant tau, the likelihood is just a Gaussian linear regression where

dx_t \sim \mathcal{N}(\frac{1}{\tau_t}(A_{z_t} x_{t-1} + b_{z_t}), \frac{1}{\tau_t} Q_{z_t})

and dx_t = x_{t+1} - x_t.

To update the parameters of the linear regression, we need the expected sufficient statistics, which we obtain by expanding the log likelihood above and writing it as a sum of inner products between parameters (A, b, Q) and sufficient statistics, which are functions of tau, x, and dx.

(1, E[x_{t-1} x_{t-1}^\top / \tau_t], E[x_{t-1} / \tau_t], E[1 / \tau_t], E[dx_t x_{t-1}^\top], E[dx_t], E[\tau_t dx_t dx_t^\top])

These are easy to compute because \tau_t is the only random variable and it is constrained to a finite set of values.

Parameters

  • dataset (np.ndarray) – observed data of shape (\text{batch\_dim}, \text{num\_timesteps}, \text{emissions\_dim}).

  • posteriors (StationaryHMMPosterior) – HMM posterior object with batch_dim to match dataset.

Returns

emissions (TimeWarpedAutoregressiveEmissions) – updated emissions object

Return type

TimeWarpedAutoregressiveEmissions