Sticky Poisson HMM Python

Python utilities for discovering latent states in neural population activity with Hidden Markov Models.

The package includes Python counterparts of the original MATLAB demo families plus graph-initialized extensions:

• standard Poisson HMM
• sticky Poisson HMM
• Poisson HMM with Dirichlet prior over transition rows
• graph-initialized Poisson and Gaussian HMMs
• graph-smoothed Multinoulli HMM
• Multinoulli HMM
• sticky Gaussian HMM for continuous signals

The Poisson and Multinoulli models are designed for count-like observations:

• binned spike-count arrays
• arrays of spike trains
• a single spike-time array
• fiber photometry signals after converting transients to event counts

Brutally honest modeling rule: sticky-Poisson HMM is a Poisson count model. It is principled for spike counts and detected event counts. Do not feed raw continuous fiber photometry fluorescence directly into the Poisson model unless you first convert it to event counts. If you want to model raw continuous amplitudes, use fit_sticky_gaussian_hmm instead.

Install

From this repository:

python -m pip install -e .

Then any script in the same Python environment can import:

from hmm_spikes import fit_sticky_poisson_hmm, fit_sticky_gaussian_hmm

HMM Studio Desktop GUI

This repository also includes hmm_studio, a PySide6 desktop interface for loading datasets, fitting Poisson, Gaussian, Multinoulli, and graph-initialized HMMs, comparing runs, and saving complete projects.

Install the GUI extras:

python -m pip install -e ".[studio]"

Launch from the repository root:

.\Launch_HMM_Studio.bat

HMM Studio project files use the .hmmstudio extension and store the dataset, fitted runs, BIC scan result, active run, and UI state so fitted models can be reopened without recomputing EM fits. Treat project files as trusted local files.

Minimal Use With A Count Array

from hmm_spikes import fit_sticky_poisson_hmm, state_probabilities, viterbi_decode

dt = 0.05
counts = your_counts_array  # shape: n_neurons x n_time_bins

result = fit_sticky_poisson_hmm(
    [counts],
    n_states=3,
    bin_size=dt,
    threshold=0.8,
    max_iter=1000,
    random_state=3456,
)

posterior = state_probabilities(counts, result.means, result.gamma)
posterior_states = posterior.argmax(axis=0)
viterbi_states = viterbi_decode(counts, result.means, result.gamma)

Algorithm Map

from hmm_spikes import (
    fit_poisson_hmm,
    fit_sticky_poisson_hmm,
    fit_dirichlet_poisson_hmm,
    fit_sticky_graph_poisson_hmm,
    fit_sticky_graph_gaussian_hmm,
    fit_graph_multinoulli_hmm,
    fit_multinoulli_hmm,
    fit_sticky_gaussian_hmm,
)

Use fit_poisson_hmm for the standard PHMM, fit_sticky_poisson_hmm for the recommended sticky count model, fit_dirichlet_poisson_hmm when you want a soft transition prior rather than a hard sticky reset, fit_multinoulli_hmm for categorical symbols, and fit_sticky_gaussian_hmm for continuous traces such as raw photometry. Use the graph-initialized Poisson and Gaussian variants when graph clustering over time bins should provide the starting state labels.

Graph-Initialized HMM Family

The graph-initialized Poisson and Gaussian variants build a weighted k-nearest-neighbor graph over time bins, spectral-cluster that graph into rough state labels, then use those labels to initialize the HMM emissions, transitions, and trial-start probabilities. After initialization, EM is the ordinary sticky Poisson or sticky Gaussian HMM update. There is no Poisson or Gaussian emission smoothing step.

The graph is built from smoothed, standardized observation geometry. For Poisson counts the graph uses log1p(counts) for geometry, while the HMM itself is still fit on raw nonnegative counts. For Gaussian observations the graph uses the continuous observation geometry directly.

from hmm_spikes import (
    fit_sticky_graph_poisson_hmm,
    fit_sticky_graph_gaussian_hmm,
    fit_graph_multinoulli_hmm,
    infer_symbol_transition_graph,
    infer_time_bin_graph,
)

graph = infer_time_bin_graph(
    trial_counts,
    top_k=15,
    temporal_weight=0.2,
)

result = fit_sticky_graph_poisson_hmm(
    trial_counts,
    n_states=3,
    bin_size=0.05,
    adjacency=graph,
    graph_strength=0.2,
    threshold=0.8,
    max_iter=1000,
    random_state=3456,
)

gaussian_graph = infer_time_bin_graph(continuous_trials, top_k=15)
gaussian_result = fit_sticky_graph_gaussian_hmm(
    continuous_trials,
    n_states=3,
    adjacency=gaussian_graph,
    graph_strength=0.2,
    threshold=0.8,
    max_iter=1000,
)

symbol_graph = infer_symbol_transition_graph(symbol_trials, n_symbols=n_symbols)
symbol_result = fit_graph_multinoulli_hmm(
    symbol_trials,
    n_states=3,
    n_symbols=n_symbols,
    adjacency=symbol_graph,
    graph_strength=0.1,
    max_iter=1000,
)

Brutal honesty: graph initialization is a starting point, not proof that graph communities are true neural states. If the graph geometry is garbage, the HMM starts from garbage. Use multiple seeds, held-out likelihood, BIC, posterior diagnostics, and biological interpretability before trusting it. The Multinoulli graph variant remains a graph-smoothed emission model because symbols have a natural similarity graph.

Reliability Note

The implementation does not depend on an external HMM backend. The forward-backward and EM updates are implemented in this package. Transition expectations are accumulated without constructing a large n_states x n_states x n_time_bins tensor, which is important for long recordings.

For unstable model-selection sweeps, use the crash-isolated scanner:

python scripts\run_bic_scan.py `
  --dataset data\python\exampledata_dt50ms.npz `
  --output-dir figures\bic_scan `
  --states 2..6 `
  --restarts 20 `
  --max-iter 1000

Each restart runs in a fresh Python subprocess. A native crash is recorded as crashed and the scan continues. Fits that satisfy the sticky threshold but do not hit strict EM convergence are labeled diagnostic.

Citation

This package is based on the sticky Poisson HMM methodology from:

Li T, La Camera G (2025) A sticky Poisson Hidden Markov Model for solving the problem of over-segmentation and rapid state switching in cortical datasets. PLOS One 20(7): e0325979. https://doi.org/10.1371/journal.pone.0325979

Tutorial

See tutorial.md for spike trains, single spike-time arrays, binned count arrays, fiber photometry event counts, and state-number selection.