Hidden Markov Model (HMM)

Contents

Hidden Markov Model (HMM)
- HMM Models
- HMM Components

HMM Models 

Module defining base model behavior for Hidden Markov Models (HMMs).

class ssm.hmm.base.HMM(num_states, initial_condition, transitions, emissions)

The Hidden Markov Model base class.

The HMM base class.

Parameters

num_states (int) – number of discrete states
initial_condition (initial.InitialCondition) – initial condition object defining $p(z_1)$
transitions (transitions.Transitions) – transitions object defining $p(z_t|z_{t-1})$
emissions (emissions.Emissions) – emissions ojbect defining $p(x_t|z_t)$

property emissions_shape

Returns the shape of a single emission, $y_t$ .

Returns: A tuple or tree of tuples giving the emission shape (s)

initial_distribution(covariates=None, metadata=None)

The distribution over the initial state of the SSM.

$p(x_1)$

Parameters

covariates (PyTree, optional) – optional covariates with leaf shape [B, T, …]. Defaults to None.
metadata (PyTree, optional) – optional metadata with leaf shape [B, …]. Defaults to None.

Returns

initial_distribution (tfp.distributions.Distribution) – A distribution over initial states in the SSM.

dynamics_distribution(state, covariates=None, metadata=None)

The dynamics (or state-transition) distribution conditioned on the current state.

$p(x_{t+1} | x_t, u_{t+1})$

Parameters: state – The current state on which to condition the dynamics.
Returns: dynamics_distribution (tfp.distributions.Distribution) – The distribution over states conditioned on the current state.

emissions_distribution(state, covariates=None, metadata=None)

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

$p(y_t | x_t, u_t)$

Parameters: state – The current state on which to condition the emissions.
Returns: emissions_distribution (tfp.distributions.Distribution) – The emissions distribution conditioned on the provided state.

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

key (jr.PRNGKey) – random seed
data (np.ndarray) – array of observed data of shape $(\text{batch} , \text{num\_timesteps} , \text{emissions\_dim})$
covariates (PyTree, optional) – optional covariates with leaf shape (B, T, …). Defaults to None.
metadata (PyTree, optional) – optional metadata with leaf shape (B, …). Defaults to None.
method (str, optional) – state assignment method. One of “random” or “kmeans”. Defaults to “kmeans”.

Raises

ValueError – when initialize method is not recognized

Return type

None

fit(data, covariates=None, metadata=None, method='em', num_iters=100, tol=0.0001, initialization_method='kmeans', key=None, verbosity=Verbosity.DEBUG)

Fit the HMM to a dataset using the specified method and initialization.

Parameters

data (np.ndarray) – observed data of shape $(\text{[batch]} , \text{num\_timesteps} , \text{emissions\_dim})$
covariates (PyTree, optional) – optional covariates with leaf shape (B, T, …). Defaults to None.
metadata (PyTree, optional) – optional metadata with leaf shape (B, …). Defaults to None.
method (str, optional) – model fit method. Must be one of [“em”]. Defaults to “em”.
num_iters (int, optional) – number of fit iterations. Defaults to 100.
tol (float, optional) – tolerance in log probability to determine convergence. Defaults to 1e-4.
initialization_method (str, optional) – method to initialize latent states. Defaults to “kmeans”.
key (jr.PRNGKey, optional) – Random seed. Defaults to None.
verbosity (Verbosity, optional) – print verbosity. Defaults to Verbosity.DEBUG.

Raises

ValueError – if fit method is not reocgnized

Returns

log_probs (np.ndarray) – log probabilities at each fit iteration
model (HMM) – the fitted model
posteriors (StationaryHMMPosterior) – the fitted posteriors

class ssm.hmm.models.BernoulliHMM(num_states, num_emission_dims=None, initial_state_probs=None, transition_matrix=None, emission_probs=None, seed=None)

HMM with conditionally independent Bernoulli emissions.

$p(x_t | z_t = k) \sim \prod_{d=1}^D \mathrm{Bern}(x_{td} \mid p_{kd})$

where $p_{kd}$ is the probability of seeing a one in dimension $d$ given discrete latent state $k$ .

The BernoulliHMM can be initialized by specifying each parameter explicitly, or you can simply specify the num_states, num_emission_dims, and seed to create a BernoulliHMM 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.
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_probs] (np.ndarray, optional) – specifies emission probabilities with shape $(\text{num\_states}, \text{emission\_dims})$ . Defaults to None.
seed (jr.PRNGKey, optional) – random seed. Defaults to None.
emission_probs (Array) –

class ssm.hmm.models.GaussianHMM(num_states, num_emission_dims=None, initial_state_probs=None, transitions=None, transition_matrix=None, emission_means=None, emission_covariances=None, seed=None)

HMM with Gaussian emissions.

$p(x_t | z_t = k) \sim \mathcal{N}(\mu_k, \Sigma_k)$

The GaussianHMM can be initialized by specifying each parameter explicitly, or you can simply specify the num_states, num_emission_dims, and seed to create a GaussianHMM 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.
initial_state_probs (np.ndarray, optional) – initial state probabilities with shape $(\text{num\_states},)$ . Defaults to None.
transitions (Transitions, optional) – object specifying transitions Defaults to None. If specified, then transition_matrix is ignored.
transition_matrix (np.ndarray, optional) – transition matrix with shape $(\text{num\_states}, \text{num\_states})$ . Defaults to None. Only used if transitions is None.
emission_means (np.ndarray, optional) – specifies emission means with shape $(\text{num\_states}, \text{emission\_dims})$ . Defaults to None.
emission_covariances (np.ndarray, optional) – specifies emissions covariances with shape $(\text{num\_states}, \text{emission\_dims}, \text{emission\_dims})$ . Defaults to None.
seed (jr.PRNGKey, optional) – random seed. Defaults to None.

class ssm.hmm.models.PoissonHMM(num_states, num_emission_dims=None, initial_state_probs=None, transition_matrix=None, emission_rates=None, seed=None)

HMM with Poisson emissions.

$p(x_t | z_t = k) \sim \text{Po}(\lambda=\lambda_k)$

The PoissonHMM can be initialized by specifying each parameter explicitly, or you can simply specify the num_states, num_emission_dims, and seed to create a GaussianHMM 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.
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_rates (np.ndarray, optional) – specifies Poisson emission rates with shape $(\text{num\_states}, \text{emission\_dims})$ . Defaults to None.
seed (jr.PRNGKey, optional) – random seed. Defaults to None.

HMM Components 

HMM Initials 

class ssm.hmm.initial.InitialCondition(num_states)

Base class for initial state distributions of an HMM.

$p(z_1 \mid u_t)$

where u_t are optional covariates at time t.

Parameters: num_states (int) –

distribution(covariates=None, metadata=None)

Return the distribution of z_1.

Parameters

covariates (PyTree, optional) – optional covariates with leaf shape (B, T, …). Defaults to None.
metadata (PyTree, optional) – optional metadata with leaf shape (B, …). Defaults to None.

Returns

distribution (tfd.Distribution) – distribution of z_1

log_initial_probs(data, covariates=None, metadata=None): Return [log Pr(z_1 = k) for k in range(num_states)]

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

Update the initial distribution in an M step given posteriors over the latent states.

Parameters

dataset (np.ndarray) – the observed dataset with shape (B, T, D)
posteriors (HMMPosterior) – posteriors over the latent states with leaf shape (B, …)
covariates (PyTree, optional) – optional covariates with leaf shape (B, T, …). Defaults to None.
metadata (PyTree, optional) – optional metadata with leaf shape (B, …). Defaults to None.

Returns

initial_condition (InitialCondition) – updated initial condition object

Return type

InitialCondition

class ssm.hmm.initial.StandardInitialCondition(num_states, initial_probs=None, initial_distribution=None, initial_distribution_prior=None)

The standard model is a categorical distribution.

Parameters

num_states (int) –
initial_distribution (ssmd.Categorical) –
initial_distribution_prior (ssmd.Dirichlet) –

distribution(covariates=None, metadata=None)

Return the distribution of z_1.

Parameters

covariates (PyTree, optional) – optional covariates with leaf shape (B, T, …). Defaults to None.
metadata (PyTree, optional) – optional metadata with leaf shape (B, …). Defaults to None.

Returns

distribution (tfd.Distribution) – distribution of z_1

log_initial_probs(data, covariates=None, metadata=None): Return [log Pr(z_1 = k) for k in range(num_states)]

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

Update the initial distribution in an M step given posteriors over the latent states.

Here, an exact M-step is performed.

Parameters

dataset (np.ndarray) – the observed dataset with shape (B, T, D)
posteriors (HMMPosterior) – posteriors over the latent states with leaf shape (B, …)
covariates (PyTree, optional) – optional covariates with leaf shape (B, T, …). Defaults to None.
metadata (PyTree, optional) – optional metadata with leaf shape (B, …). Defaults to None.

Returns

initial_condition (StandardInitialCondition) – updated initial condition object

Return type

StandardInitialCondition

HMM Transitions 

class ssm.hmm.transitions.Transitions(num_states)

Base class for HMM transitions models,

$p_t(z_t \mid z_{t-1}, u_t)$

where u_t are optional covariates at time t.

Parameters: num_states (int) –

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

Return the conditional distribution of z_t given state z_{t-1}

Parameters

state (int) – state z_{t-1}
covariates (PyTree, optional) – optional covariates with leaf shape (B, T, …). Defaults to None.
metadata (PyTree, optional) – optional metadata with leaf shape (B, …). Defaults to None.

Returns

distribution (tfd.Distribution) – conditional distribution of z_t given state z_{t-1}.

log_transition_matrices(data, covariates=None, metadata=None)

Returns the log probability of data where

$\texttt{log\_P}[i, j] = \log \Pr(z_{t+1} = j | z_t = i)$

if the transition probabilities are stationary or

$\texttt{log\_P}[t, i, j] = \log \Pr(z_{t+1} = j | z_t = i)$

if they are nonstationary.

Parameters: data (np.ndarray) – observed data
Returns: log probs (np.ndarray) – log probability as defined above

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

Update the transition parameters in an M step given posteriors over the latent states.

Parameters

dataset (np.ndarray) – the observed dataset with shape (B, T, D)
posteriors (HMMPosterior) – posteriors over the latent states with leaf shape (B, …)
covariates (PyTree, optional) – optional covariates with leaf shape (B, T, …). Defaults to None.
metadata (PyTree, optional) – optional metadata with leaf shape (B, …). Defaults to None.

Returns

transitions (Transitions) – updated transitions object

Return type

Transitions

class ssm.hmm.transitions.StationaryTransitions(num_states, transition_matrix=None, transition_distribution=None, transition_distribution_prior=None)

Basic transition model with a fixed transition matrix.

Parameters

num_states (int) –
transition_distribution (ssmd.Categorical) –
transition_distribution_prior (ssmd.Dirichlet) –

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

Return the conditional distribution of z_t given state z_{t-1}

Parameters

state (int) – state z_{t-1}
covariates (PyTree, optional) – optional covariates with leaf shape (B, T, …). Defaults to None.
metadata (PyTree, optional) – optional metadata with leaf shape (B, …). Defaults to None.

Returns

distribution (tfd.Distribution) – conditional distribution of z_t given state z_{t-1}.

log_transition_matrices(data, covariates=None, metadata=None)

Returns the log probability of data where

$\texttt{log\_P}[i, j] = \log \Pr(z_{t+1} = j | z_t = i)$

if the transition probabilities are stationary or

$\texttt{log\_P}[t, i, j] = \log \Pr(z_{t+1} = j | z_t = i)$

if they are nonstationary.

Parameters: data (np.ndarray) – observed data
Returns: log probs (np.ndarray) – log probability as defined above

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

Update the transition parameters in an M step given posteriors over the latent states.

Parameters

dataset (np.ndarray) – the observed dataset with shape (B, T, D)
posteriors (HMMPosterior) – posteriors over the latent states with leaf shape (B, …)
covariates (PyTree, optional) – optional covariates with leaf shape (B, T, …). Defaults to None.
metadata (PyTree, optional) – optional metadata with leaf shape (B, …). Defaults to None.

Returns

transitions (Transitions) – updated transitions object

Return type

StationaryTransitions

class ssm.hmm.transitions.SimpleStickyTransitions(num_states, stay_probability=None, transition_distribution=None, transition_distribution_prior=None)

HMM transition model where diagonal elements are a learned parameter (stay_probability) and off-diagonal elements are uniform. That is, for a model with n hidden states, the diagonal entries of the transition matrix are stay_probability, and the off-diagonal entries are (1 - stay_probability) / (n - 1).

Parameters

num_states (int) –
stay_probability (float) –
transition_distribution (ssmd.Categorical) –
transition_distribution_prior (ssmd.Beta) –

log_transition_matrices(data, covariates=None, metadata=None)

Returns the log probability of data where

$\texttt{log\_P}[i, j] = \log \Pr(z_{t+1} = j | z_t = i)$

if the transition probabilities are stationary or

$\texttt{log\_P}[t, i, j] = \log \Pr(z_{t+1} = j | z_t = i)$

if they are nonstationary.

Parameters: data (np.ndarray) – observed data
Returns: log probs (np.ndarray) – log probability as defined above

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

Return the conditional distribution of z_t given state z_{t-1}

Parameters

state (int) – state z_{t-1}
covariates (PyTree, optional) – optional covariates with leaf shape (B, T, …). Defaults to None.
metadata (PyTree, optional) – optional metadata with leaf shape (B, …). Defaults to None.

Returns

distribution (tfd.Distribution) – conditional distribution of z_t given state z_{t-1}.

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

Update the transition parameters in an M step given posteriors over the latent states.

Parameters

dataset (np.ndarray) – the observed dataset with shape (B, T, D)
posteriors (HMMPosterior) – posteriors over the latent states with leaf shape (B, …)
covariates (PyTree, optional) – optional covariates with leaf shape (B, T, …). Defaults to None.
metadata (PyTree, optional) – optional metadata with leaf shape (B, …). Defaults to None.

Returns

transitions (Transitions) – updated transitions object

Return type

SimpleStickyTransitions

HMM Emissions 

class ssm.hmm.emissions.Emissions(num_states)

Base class of emission distribution of an HMM

$p_t(x_t \mid z_t, u_t)$

where u_t are optional covariates.

Parameters: num_states (int) –

distribution(state, covariates=None, metadata=None): Return the conditional distribution of emission x_t given state z_t and (optionally) covariates u_t.

log_likelihoods(data, covariates=None, metadata=None): Compute log p(x_t | z_t=k) for all t and k.

m_step(data, posterior, covariates=None, metadata=None)

By default, try to optimize the emission distribution via generic gradient-based optimization of the expected log likelihood.

This function assumes that the Emissions subclass is a PyTree and that all of its leaf nodes are unconstrained parameters.

Parameters

data (np.ndarray) – the observed data
posterior (HMMPosterior) – the HMM posterior
covariates (PyTree, optional) – optional covariates with leaf shape (B, T, …). Defaults to None.
metadata (PyTree, optional) – optional metadata with leaf shape (B, …). Defaults to None.

Returns

emissions (ExponentialFamilyEmissions) – updated emissions object

Return type

Emissions

class ssm.hmm.emissions.ExponentialFamilyEmissions(num_states, emissions_distribution=None, emissions_distribution_prior=None)

Exponential Family Emissions for HMM.

Can be initialized by specifying parameters or by passing in a pre-initialized emissions_distribution object.

Parameters

num_states (int) – number of discrete states
means (np.ndarray, optional) – state-dependent emission means. Defaults to None.
covariances (np.ndarray, optional) – state-dependent emission covariances. Defaults to None.
emissions_distribution (ssmd.MultivariateNormalTriL, optional) – initialized emissions distribution. Defaults to None.
emissions_distribution_prior (ssmd.NormalInverseWishart, optional) – initialized emissions distribution prior. Defaults to None.

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

Get the emissions distribution at the provided state.

Parameters

state (int) – discrete state
covariates (PyTree, optional) – optional covariates with leaf shape (B, T, …). Defaults to None.
metadata (PyTree, optional) – optional metadata with leaf shape (B, …). Defaults to None.

Returns

emissions distribution (tfd.MultivariateNormalLinearOperator) – emissions distribution at given state

Return type

MultivariateNormalTriL

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

Update the emissions distribution using an M-step.

Operates over a batch of data (posterior must have the same batch dim).

Parameters

dataset (np.ndarray) – the observed dataset
posteriors (HMMPosterior) – the HMM posteriors
covariates (PyTree, optional) – optional covariates with leaf shape (B, T, …). Defaults to None.
metadata (PyTree, optional) – optional metadata with leaf shape (B, …). Defaults to None.

Returns

emissions (ExponentialFamilyEmissions) – updated emissions object

Return type

ExponentialFamilyEmissions

class ssm.hmm.emissions.BernoulliEmissions(num_states, probs=None, emissions_distribution=None, emissions_distribution_prior=None)

Gaussian Emissions for HMM.

Can be initialized by specifying parameters or by passing in a pre-initialized emissions_distribution object.

Parameters

num_states (int) – number of discrete states
probs (np.ndarray, optional) – state-dependent emission probabilities. Defaults to None.
covariances (np.ndarray, optional) – state-dependent emission covariances. Defaults to None.
emissions_distribution (ssmd.MultivariateNormalTriL, optional) – initialized emissions distribution. Defaults to None.
emissions_distribution_prior (ssmd.NormalInverseWishart, optional) – initialized emissions distribution prior. Defaults to None.

class ssm.hmm.emissions.GaussianEmissions(num_states, means=None, covariances=None, emissions_distribution=None, emissions_distribution_prior=None)

Gaussian Emissions for HMM.

Can be initialized by specifying parameters or by passing in a pre-initialized emissions_distribution object.

Parameters

num_states (int) – number of discrete states
means (np.ndarray, optional) – state-dependent emission means. Defaults to None.
covariances (np.ndarray, optional) – state-dependent emission covariances. Defaults to None.
emissions_distribution (ssmd.MultivariateNormalTriL, optional) – initialized emissions distribution. Defaults to None.
emissions_distribution_prior (ssmd.NormalInverseWishart, optional) – initialized emissions distribution prior. Defaults to None.

class ssm.hmm.emissions.PoissonEmissions(num_states, rates=None, emissions_distribution=None, emissions_distribution_prior=None)

Poisson Emissions for HMM.

Can be initialized by specifying parameters or by passing in a pre-initialized emissions_distribution object.

Parameters

num_states (int) – number of discrete states
rates (np.ndarray, optional) – state-dependent Poisson rates. Defaults to None.
emissions_distribution (tfd.Distribution, optional) – pre-initialized emissions distribution. Defaults to None.
emissions_distribution_prior (tfd.Gamma, optional) – pre-initialized emissions distribution prior. Defaults to None.

Hidden Markov Model (HMM)

HMM Models

HMM Components

HMM Initials

HMM Transitions

HMM Emissions

HMM Models 

HMM Components 

HMM Initials 

HMM Transitions 

HMM Emissions 