Skip to content

Latest commit

 

History

History
575 lines (432 loc) · 15.6 KB

File metadata and controls

575 lines (432 loc) · 15.6 KB

Tutorial: Sticky Poisson HMM For Neural State Discovery

The sticky Poisson HMM is for segmenting neural activity into discrete latent states while discouraging unrealistically fast state switching. The key idea is simple: each hidden state has a vector of expected event counts across neurons, units, or channels, and the transition matrix is constrained to keep self-transition probabilities high.

Use the Poisson and Multinoulli models when your observations are counts or symbols in time bins.

Do not use raw continuous signals as Poisson observations. For fiber photometry, either convert transients to detected events and use sticky-Poisson HMM, or model the raw continuous trace with sticky Gaussian HMM.

Install

From the package folder:

python -m pip install -e .

Then your other scripts can call:

from hmm_spikes import fit_sticky_poisson_hmm

1. Binned Spike-Count Array

If you already have a count matrix:

import numpy as np
from hmm_spikes import fit_sticky_poisson_hmm, state_probabilities, viterbi_decode

dt = 0.05  # seconds
counts = np.asarray(your_counts)  # 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)

Outputs:

  • posterior: probability of each state in each time bin, shape n_states x n_time_bins
  • posterior_states: most probable state in each bin
  • viterbi_states: most likely complete state sequence
  • result.rates_hz: state firing rates or event rates, shape n_neurons x n_states
  • result.gamma: transition probability matrix

State labels are zero-based in Python. Add 1 if you want MATLAB-style labels.

1b. The MATLAB-Style HMM Families

The Python package exposes direct public entry points for the four MATLAB demo families.

Standard Poisson HMM:

from hmm_spikes import fit_poisson_hmm

standard = fit_poisson_hmm(
    trial_counts,
    n_states=3,
    bin_size=0.05,
    max_iter=1000,
)

Sticky Poisson HMM:

from hmm_spikes import fit_sticky_poisson_hmm

sticky = fit_sticky_poisson_hmm(
    trial_counts,
    n_states=3,
    bin_size=0.05,
    threshold=0.8,
    max_iter=1000,
)

Poisson HMM with Dirichlet transition prior:

from hmm_spikes import fit_dirichlet_poisson_hmm

dirichlet = fit_dirichlet_poisson_hmm(
    trial_counts,
    n_states=3,
    bin_size=0.05,
    mode_self_transition=0.9,
    offdiag_alpha=1.1,
    max_iter=1000,
)

Multinoulli HMM:

from hmm_spikes import counts_to_multinoulli_symbols, fit_multinoulli_hmm

symbols = [
    counts_to_multinoulli_symbols(counts, random_state=0)
    for counts in trial_counts
]

multi = fit_multinoulli_hmm(
    symbols,
    n_states=3,
    n_symbols=trial_counts[0].shape[0] + 1,
    max_iter=1000,
)

Multinoulli warning: use short bins. The model assumes one categorical symbol per time bin: one selected unit fired, or no unit fired. If many neurons fire in the same bin, the symbol representation is throwing information away.

1c. Graph-Initialized HMM Family

The graph-initialized Poisson and Gaussian HMM family is added on top of the MATLAB-style models. It builds a weighted graph over time bins, clusters that graph into rough state labels, and uses those labels to initialize emissions, transitions, and trial-start probabilities. After initialization, EM is the ordinary sticky Poisson or sticky Gaussian HMM. There is no Poisson or Gaussian graph-smoothing step.

The graph variants now mean:

  • Poisson: graph labels initialize state count-rate vectors
  • Gaussian: graph labels initialize state mean and variance vectors
  • Multinoulli: symbol graph still smooths categorical emission probabilities

Use Poisson or Gaussian graph initialization when time-bin geometry gives a better starting segmentation than random starts:

  • smoothed and normalized population activity
  • PCA or another low-dimensional embedding
  • k-nearest-neighbor similarity among time bins
  • optional temporal edges between neighboring bins
from hmm_spikes import (
    counts_to_multinoulli_symbols,
    fit_graph_multinoulli_hmm,
    fit_sticky_graph_gaussian_hmm,
    fit_sticky_graph_poisson_hmm,
    infer_symbol_transition_graph,
    infer_time_bin_graph,
)

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

graph_model = 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,
)

print(graph_model.rates_hz)
print(graph_model.graph_summary["density"])
print(graph_model.graph_smoothness)

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

symbol_trials = [
    counts_to_multinoulli_symbols(counts, random_state=0)
    for counts in trial_counts
]
symbol_graph = infer_symbol_transition_graph(
    symbol_trials,
    n_symbols=trial_counts[0].shape[0] + 1,
)
symbol_model = fit_graph_multinoulli_hmm(
    symbol_trials,
    n_states=3,
    n_symbols=trial_counts[0].shape[0] + 1,
    adjacency=symbol_graph,
    graph_strength=0.1,
    max_iter=1000,
)

If you do not provide adjacency, the Poisson and Gaussian graph fitters infer a time-bin graph from the training observations, and the Multinoulli graph fitter infers a short-lag symbol transition graph. Be careful. That is convenient, but it can leak task structure into the initialization if you estimate the graph on the same condition you are trying to interpret. For serious analysis, compare against random restarts, report sensitivity to graph parameters, and treat the graph-initialized fit as exploratory unless it survives those checks.

2. Multiple Spike Trains

Use this when each neuron has its own spike-time array.

import numpy as np
from hmm_spikes import spike_trains_to_counts, fit_sticky_poisson_hmm

dt = 0.05
time_edges = np.arange(0.0, 30.0 + dt, dt)

spike_trains = [
    neuron_1_spike_times,
    neuron_2_spike_times,
    neuron_3_spike_times,
]

counts = spike_trains_to_counts(spike_trains, time_edges)

result = fit_sticky_poisson_hmm(
    [counts],
    n_states=3,
    bin_size=dt,
    threshold=0.8,
)

This gives a count matrix with shape:

n_neurons x n_time_bins

3. One Single Spike-Time Array

You can train and decode one unit, although it is statistically weak compared with population data.

import numpy as np
from hmm_spikes import spike_times_to_counts, fit_sticky_poisson_hmm, viterbi_decode

dt = 0.05
time_edges = np.arange(0.0, 60.0 + dt, dt)
spike_times = single_unit_spike_times

counts = spike_times_to_counts(spike_times, time_edges)

result = fit_sticky_poisson_hmm(
    [counts],
    n_states=2,
    bin_size=dt,
    threshold=0.8,
)

states = viterbi_decode(counts, result.means, result.gamma)

Brutal honesty: a single spike train can support simple low-rate versus high-rate segmentation. It cannot reliably support a large number of hidden states.

4. Two-Column Firing Matrix

If your spikes are stored as [spike_time, neuron_id], with one-based neuron IDs:

import numpy as np
from hmm_spikes import firings_to_counts, fit_sticky_poisson_hmm

dt = 0.05
time_edges = np.arange(-2.0, 5.0 + dt, dt)

counts = firings_to_counts(
    firings,
    time_edges,
    n_neurons=12,
)

result = fit_sticky_poisson_hmm(
    [counts],
    n_states=3,
    bin_size=dt,
    threshold=0.8,
)

5. Multiple Trials

Training is much better when you provide multiple trials:

trial_counts = [
    counts_trial_1,
    counts_trial_2,
    counts_trial_3,
]

result = fit_sticky_poisson_hmm(
    trial_counts,
    n_states=3,
    bin_size=0.05,
    threshold=0.8,
    max_iter=1000,
)

Then decode each trial:

from hmm_spikes import state_probabilities, viterbi_decode

for counts in trial_counts:
    posterior = state_probabilities(counts, result.means, result.gamma)
    states = viterbi_decode(counts, result.means, result.gamma)

6. Fiber Photometry

Fiber photometry is continuous. A Poisson HMM does not model raw fluorescence values. To use sticky-Poisson HMM, detect transients and turn them into event counts.

Single photometry signal:

from hmm_spikes import signal_to_event_counts, fit_sticky_poisson_hmm, viterbi_decode

counts, time_edges, event_times = signal_to_event_counts(
    photometry_signal,
    sampling_rate=1000.0,
    bin_size=0.1,
    threshold_z=2.5,
    refractory=0.2,
)

result = fit_sticky_poisson_hmm(
    [counts],
    n_states=2,
    bin_size=0.1,
    threshold=0.8,
)

states = viterbi_decode(counts, result.means, result.gamma)

Array of photometry signals:

counts, time_edges, event_times = signal_to_event_counts(
    photometry_signals,  # shape: n_signals x n_samples
    sampling_rate=1000.0,
    bin_size=0.1,
    threshold_z=2.5,
    refractory=0.2,
)

result = fit_sticky_poisson_hmm(
    [counts],
    n_states=3,
    bin_size=0.1,
    threshold=0.8,
)

Interpretation:

  • each photometry channel becomes one row
  • each detected transient contributes event counts
  • result.rates_hz is an event rate, not a fluorescence amplitude

This is a modeling compromise. It is useful when transients are the event-like objects of interest. It is not a replacement for modeling raw fluorescence dynamics.

6b. Sticky Gaussian HMM For Raw Photometry Or Continuous Signals

If you want to model the raw continuous signal itself, use the sticky Gaussian HMM. This uses a diagonal Gaussian emission model, so each state has a mean and variance for each signal.

Single continuous signal:

from hmm_spikes import fit_sticky_gaussian_hmm, gaussian_viterbi_decode

result = fit_sticky_gaussian_hmm(
    [photometry_signal],
    n_states=3,
    threshold=0.8,
    max_iter=1000,
    random_state=3456,
)

states = gaussian_viterbi_decode(
    photometry_signal,
    result.means,
    result.variances,
    result.gamma,
)

Array of continuous signals:

from hmm_spikes import fit_sticky_gaussian_hmm, gaussian_state_probabilities

signals = photometry_signals  # shape: n_signals x n_time_bins

result = fit_sticky_gaussian_hmm(
    [signals],
    n_states=3,
    threshold=0.8,
    max_iter=1000,
)

posterior = gaussian_state_probabilities(
    signals,
    result.means,
    result.variances,
    result.gamma,
)

Interpretation:

  • result.means: state-dependent signal amplitude means
  • result.variances: state-dependent signal variances
  • result.gamma: transition matrix
  • result.threshold_satisfied: whether the final sticky constraint holds

Use sticky Gaussian HMM for raw continuous photometry. Use sticky Poisson HMM only after event detection and binning.

7. Choosing The Number Of States

Do not choose the state number by raw likelihood. Raw likelihood usually rewards too many states.

Use:

  1. Candidate state numbers, for example m = 2, 3, ..., 12
  2. Many random restarts per m, ideally 20 to 100
  3. Keep only converged models with threshold_satisfied == True
  4. Compute BIC
  5. Inspect rasters, posterior confidence, state rates, and dwell times
import numpy as np
from hmm_spikes import fit_sticky_poisson_hmm

def bic(log_likelihood, m, n_neurons, total_bins):
    n_params = m * (m - 1) + n_neurons * m
    return -2 * log_likelihood + n_params * np.log(total_bins)

results = []

for m in range(2, 13):
    best = None

    for seed in range(30):
        result = fit_sticky_poisson_hmm(
            trial_counts,
            n_states=m,
            bin_size=0.05,
            threshold=0.8,
            max_iter=1000,
            random_state=seed,
        )

        if not result.converged:
            continue
        if not result.threshold_satisfied:
            continue

        if best is None or result.log_likelihood > best.log_likelihood:
            best = result

    if best is None:
        print(f"m={m}: no valid model")
        continue

    n_neurons = trial_counts[0].shape[0]
    total_bins = sum(x.shape[1] for x in trial_counts)
    score = bic(best.log_likelihood, m, n_neurons, total_bins)
    results.append((m, score, best))
    print(f"m={m}, BIC={score:.2f}, LL={best.log_likelihood:.2f}")

best_m, best_bic, best_model = min(results, key=lambda x: x[1])
print("Best state number:", best_m)

For long or unstable recordings, prefer the crash-isolated scanner:

from hmm_spikes import run_sticky_poisson_bic_scan_isolated

scan = run_sticky_poisson_bic_scan_isolated(
    trial_counts,
    candidate_states=range(2, 7),
    bin_size=0.05,
    n_restarts=20,
    threshold=0.8,
    max_iter=1000,
    output_dir="bic_scan_outputs",
)

print(scan.best_strict)
print(scan.best_diagnostic)

The scanner classifies each restart:

  • strict: converged and sticky threshold satisfied
  • diagnostic: sticky threshold satisfied but strict EM convergence was not reached
  • invalid: returned a model that did not satisfy the sticky threshold
  • timeout: exceeded the per-restart timeout
  • crashed: subprocess exited abnormally

For a case like NAc where the sticky threshold is satisfied but the strict convergence flag is false at 100 iterations, call it a diagnostic fit and increase max_iter before interpreting it as final.

Command-line version:

python scripts\run_bic_scan.py `
  --dataset path\to\data.npz `
  --output-dir bic_scan_outputs `
  --states 2..6 `
  --restarts 20 `
  --max-iter 1000

If only K=4 is stable and biologically interpretable, report K=4 as the prespecified stable model and label the broader K scan as unstable. Do not pretend an unstable K sweep provides strong model-selection evidence.

7b. Condition Metrics And Paired Tests

The package includes helpers for the table-level quantities:

from hmm_spikes import (
    state_sequence_metrics,
    paired_condition_table,
)

metrics = state_sequence_metrics(viterbi_states, bin_size=0.05)
print(metrics["switches_per_min"])
print(metrics["mean_state_duration_s"])
print(metrics["state_entropy_bits"])

For paired condition comparisons:

condition_a = {
    "switches_per_min": switches_control,
    "mean_state_duration_s": duration_control,
    "state_entropy_bits": entropy_control,
}

condition_b = {
    "switches_per_min": switches_treatment,
    "mean_state_duration_s": duration_treatment,
    "state_entropy_bits": entropy_treatment,
}

rows = paired_condition_table(condition_a, condition_b)
for row in rows:
    print(row)

With n=3 mice, exact sign-flip p values are coarse. No FDR-surviving effect should be described as a reliable condition effect. Treat it as descriptive and hypothesis-generating.

8. Citation

This package implements the sticky Poisson HMM workflow introduced by:

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

PLOS article page:

https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0325979

Wrap-Up

Use sticky-Poisson HMM when your observations are event counts. Spike trains naturally become counts. Single spike-time arrays can be counted too, but support only simple models. Fiber photometry must be converted into event counts first. The sticky transition constraint protects you from one of the most common HMM failures in neural data: fake rapid switching caused by weak self-transition probabilities.