Autoregressive Hidden Markov Model (ARHMM)

ARHMM Models

Base class for Autoregressive HMM.

class ssm.arhmm.base.AutoregressiveHMM(num_states, initial_condition, transitions, emissions)

Base class for HMM with autoregressive dependencies.

Inherits from HMM base class.

The HMM base class.

Parameters
emissions_distribution(state, covariates=None, metadata=None, history=None)

The emissions (or observation) distribution conditioned on the current state.

p(y_t | x_t)

Parameters

state (float) – The current state on which to condition the emissions.

Returns

emissions_distribution (tfp.distributions.Distribution) – The emissions distribution conditioned on the provided state.

Return type

Distribution

log_probability(states, data, covariates=None, metadata=None, history=None)

Computes the log joint probability of a set of states and data (observations).

\log p(x, y) = \log p(x_1) + \sum_{t=1}^{T-1} \log p(x_{t+1} | x_t) + \sum_{t=1}^{T} \log p(y_t | x_t)

Parameters

  • states – latent states x_{1:T} of shape (\text{[batch]} , \text{num\_timesteps} , \text{latent\_dim})

  • data – observed data y_{1:T} of shape (\text{[batch]} , \text{num\_timesteps} , \text{emissions\_dim})

  • history – previous emissions to condition on of shape (\text{[batch]} , \text{num\_lags} , \text{emissions\_dim}). Default is None which will condition on zeros.

Returns

lp – log joint probability \log p(x, y) of shape (\text{batch]},)

sample(key, num_steps, initial_state=None, covariates=None, metadata=None, num_samples=1, history=None)

Sample from the joint distribution defined by the state space model.

x, y \sim p(x, y)

Parameters

  • key (jr.PRNGKey) – A JAX pseudorandom number generator key.

  • num_steps (int) – Number of steps for which to sample.

  • covariates – Optional covariates that may be needed for sampling. Default is None.

  • initial_state – Optional state on which to condition the sampled trajectory. Default is None which samples the intial state from the initial distribution.

  • num_samples (int) – Number of indepedent samples (defines a batch dimension).

  • prev_emissions – previous emissions to condition on of shape (\text{[batch]} , \text{num\_lags} , \text{emissions\_dim}). Default is None which will condition on zeros.

Returns

  • states – an array of latent states across time x_{1:T} of shape (\text{[batch]} , \text{num\_timesteps} , \text{latent\_dim})

  • emissions – an array of observations across time y_{1:T} of shape (\text{[batch]} , \text{num\_lags} , \text{emissions\_dim}).

Model classes for ARHMMs.

class ssm.arhmm.models.GaussianARHMM(num_states, num_emission_dims=None, num_lags=None, initial_state_probs=None, transition_matrix=None, emission_weights=None, emission_biases=None, emission_covariances=None, seed=None)

Gaussian autoregressive hidden markov Model (ARHMM).

Let $x_t$ denote the observation at time $t$. Let $z_t$ denote the corresponding discrete latent state. The autoregressive hidden Markov model (with ..math`text{lag}=1`) has the following likelihood,

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

The GaussianARHMM can be initialized by specifying each parameter explicitly, or you can simply specify the num_states, num_emission_dims, num_lags, and seed to create a GaussianARHMM with generic, randomly initialized parameters.

Parameters

  • num_states (int) – number of discrete latent states

  • num_emission_dims (int, optional) – number of emission dims. Defaults to None.

  • num_lags (int, optional) – number of previous timesteps on which to autoregress. Defaults to None.

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

  • transition_matrix (np.ndarray, optional) – transition matrix with shape (\text{num\_states}, \text{num\_states}). Defaults to None.

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

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

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

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

ARHMM Components

ARHMM Emissions

class ssm.arhmm.emissions.AutoregressiveEmissions(num_states, weights=None, biases=None, covariances=None, emissions_distribution=None, emissions_distribution_prior=None)

Gaussian linear regression emissions class for Autoregressive HMM.

Can be instantiated by specifying the parameters or you can pass in the initialized distribution object directly to emissions_distribution.

Optionally takes an emissions prior distribution.

Parameters

  • num_states (int) – number of discrete states

  • weights (np.ndarray, optional) – state-based weight matrix for Gaussian linear regression of shape (\text{num\_states}, \text{emissions\_dim}, \text{emissions\_dim} * \text{num\_lags}). Defaults to None.

  • biases (np.ndarray, optional) – state-based bias vector for Gaussian linear regression of shape (\text{num\_states}, \text{emissions\_dim}). Defaults to None.

  • covariances (np.ndarray, optional) – state-based covariances for Gaussian linear regression of shape (\text{num\_states}, \text{emissions\_dim}, \text{emissions\_dim}). Defaults to None.

  • emissions_distribution (ssmd.GaussianLinearRegression, optional) – initialized emissions distribution. 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 (ssmd.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.

Operates over a batch of data.

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 (AutoregressiveEmissions) – updated emissions object

Return type

AutoregressiveEmissions