Exact EM Documentation

General EM routines

ssm.inference.em.em(model, data, covariates=None, metadata=None, num_iters=100, tol=0.0001, verbosity=Verbosity.DEBUG)

Fit a model using EM.

Assumes the model has the following methods for EM:

  • model.e_step(data) (i.e. E-step)

  • model.m_step(dataset, posteriors)

  • model.marginal_likelihood(data, posterior)

Parameters

  • model (ssm.base.SSM) – the model to be fit

  • data (PyTree) – the observed data with leaf shape (B, T, D).

  • 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.

  • num_iters (int, optional) – number of iterations of EM fit. Defaults to 100.

  • tol (float, optional) – tolerance in marginal lp to declare convergence. Defaults to 1e-4.

  • verbosity (ssm.utils.Verbosity, optional) – verbosity of fit. Defaults to Verbosity.DEBUG.

Returns

  • log_probs – log probabilities across EM iterations

  • model – the fitted model

  • posterior – the posterior over the inferred latent states