statespacecheck

Local goodness-of-fit diagnostics for state space models: find the observations, down to individual spikes, where a decoder disagrees with the data.

A state space model decodes a latent state (for example an animal's position) by combining a prediction from the past with the evidence in each new observation. statespacecheck asks, for every observation, whether the two agree: does the observation fall where the model's one-step prediction expected it? Global scores such as the likelihood of a whole session cannot say when a model fails; these diagnostics can, so a misfit can be traced to a period, a behavior, or a part of the model.

The package implements the methods of the paper Local goodness-of-fit measures for neural decoding (Zeng, Comrie, Frank, Eden and Denovellis). Its analysis code and an interactive website are at statespacecheck-paper and https://edeno.github.io/statespacecheck-paper/.

Installation

pip install statespacecheck

Quick start

Given a decoder's one-step predictive distribution, each unit's firing rate at each position, and which unit fired in which time bin, compute three diagnostics for every spike and flag the poorly fit ones. Here a simulated decoder predicts the animal's position correctly in the first half of the recording and the mirror-image position in the second half:

import numpy as np
import statespacecheck as ssc

rng = np.random.default_rng(0)
position = np.linspace(0, 100, 51)  # position bins (cm)
n_time, n_units, dt = 2000, 20, 0.02  # time bins, units, bin width (s)

# Place fields: firing rate (Hz) of each unit at each position, (n_bins, n_units)
centers = np.linspace(0, 100, n_units)
place_fields = 0.1 + 20 * np.exp(-0.5 * ((position[:, None] - centers) / 8) ** 2)

# The animal runs back and forth; spikes follow the place fields
true_position = 50 + 45 * np.sin(np.arange(n_time) * dt * 0.6)
true_bin = np.abs(position[:, None] - true_position).argmin(axis=0)
spike_counts = rng.poisson(place_fields[true_bin] * dt)  # (n_time, n_units)
time_bin, unit = np.nonzero(spike_counts)
n_spikes = spike_counts[time_bin, unit]
time_ind, unit = np.repeat(time_bin, n_spikes), np.repeat(unit, n_spikes)  # one per spike

# A decoder's one-step predictive distribution, (n_time, n_bins): it tracks the
# animal in the first half and predicts the mirror-image position in the second
predicted = np.where(np.arange(n_time) < n_time // 2, true_position, 100 - true_position)
predictive = np.exp(-0.5 * ((position - predicted[:, None]) / 5) ** 2)
predictive /= predictive.sum(axis=1, keepdims=True)

# Diagnostics for every spike; thresholds from a baseline period; flags
diagnostics = ssc.event_diagnostics(predictive, place_fields, time_ind, unit)
baseline = time_ind < n_time // 4
hpd_threshold = ssc.baseline_threshold(diagnostics.hpd_overlap[baseline], 0.01)
flags = ssc.flag_events(diagnostics, hpd_overlap_threshold=hpd_threshold)
print(
    f"HPD overlap flagged at or below {hpd_threshold:.2f} (the baseline's 1st "
    f"percentile): {flags.hpd_overlap[baseline].mean():.1%} of baseline spikes"
)
misfit = time_ind >= n_time // 2
for name, flagged in [
    ("HPD overlap", flags.hpd_overlap),
    ("p-value", flags.predictive_pvalue),
]:
    print(
        f"{name:12s} flagged: {flagged[~misfit].mean():.0%} of spikes before, "
        f"{flagged[misfit].mean():.0%} during the misfit"
    )

HPD overlap flagged at or below 0.00 (the baseline's 1st percentile): 2.4% of baseline spikes
HPD overlap  flagged: 2% of spikes before, 81% during the misfit
p-value      flagged: 2% of spikes before, 85% during the misfit

More than 1% of baseline spikes have no overlap at all, so the 1st percentile is 0 and the inclusive rule flags all of them: 2.4% of the baseline, not 1%. That fraction describes the baseline; it is not a false-alarm rate to expect elsewhere.

The per-event tutorial walks through this workflow with a real filter, and the decoder guide shows how to get these inputs from your own decoder.

What your decoder provides

Every array is on the same grid of states, with bins of equal size: the diagnostics compare distributions over a discretized state space. A grid decoder already works this way. For a model whose distributions are continuous, such as a Kalman filter, evaluate the prediction and each observation's likelihood on a grid first; the Kalman filter tutorial shows how, and how to check that the grid is fine enough. The same diagnostics apply to any state space model, not only neural decoders.

Common pitfalls:

• Place fields stored one unit per row, (n_units, n_bins): pass place_fields.T.
• Positions outside the track marked NaN: the per-event functions need finite values, so keep only the valid bins in both arrays.
• Switching models (for example continuous and fragmented dynamics): sum the predictive distribution over the discrete states first, so every model is compared on the same position grid.
• Spike times: event_time_ind holds bin indices; convert times with np.digitize(spike_times, time_bin_edges) - 1.

Reading the results

The diagnostics measure consistency, not similarity: a spike is consistent with the prediction when it falls where the prediction put probability, even if the prediction is much broader. Passing a diagnostic (a high HPD overlap, a large p-value) means the spike is consistent with the prediction in that respect, not that the model is right: a prediction too broad to be useful passes both, a rare unit can get small p-values under a correct model, and counts, timing and the proportions of units across spikes are not checked. See Interpreting the diagnostics.

The paper's quantities in the package

The package also has tools the paper does not use: time-bin versions of the diagnostics for a whole-bin likelihood, run-based flagging of time series (statespacecheck.periods), a generic Monte Carlo predictive check (predictive_pvalue), and plot_diagnostics. The API reference marks which is which.

Continuous marks

With clusterless decoding, each spike's mark is a vector of waveform features rather than a unit, so the predictive p-value cannot be summed over marks. monte_carlo_mark_pvalue estimates it by simulation, as the paper describes: it draws a state from the event-weighted predictive distribution, draws a mark for a spike at that state, and ranks the observed mark's predictive density among the replicates'. It needs the observation model as a MarkModel: the log of its joint mark intensity (in log space, because densities of marks with many features underflow), a sampler of marks, and the total event rate at each state. The paper's Figure 2 example, with a 1-D mark:

import numpy as np
from scipy.stats import norm

import statespacecheck as ssc

position = np.linspace(0, 100, 200)  # state grid (cm)
predictive = norm.pdf(position, 35, 8)[np.newaxis]  # one spike's prediction, (1, n_bins)
sigma = 12.0  # spread of a spike's 1-D mark around the position


def log_mark_intensity(marks):  # log lambda(x, y) of marks (n, 1) at every position
    return norm.logpdf(marks, position, sigma)  # (n, n_bins)


def sample_marks(bins, rng):  # one mark for a spike at each position bin
    return rng.normal(position[bins], sigma)[:, np.newaxis]


model = ssc.MarkModel(
    log_intensity=log_mark_intensity,
    sample=sample_marks,
    ground_intensity=np.ones_like(position),  # the same total rate everywhere
)
check = ssc.monte_carlo_mark_pvalue(
    predictive,
    model,
    np.array([[60.0]]),  # the observed mark
    n_samples=5000,
    rng=0,
)
print(check.pvalue.round(2))

# All three diagnostics for every spike of a recording
spikes = ssc.clusterless_event_diagnostics(
    norm.pdf(position, [[35.0], [60.0]], 8),  # predictive of 2 time bins, (n_time, n_bins)
    model,
    event_time_ind=np.array([0, 0, 1]),  # each spike's time bin
    event_marks=np.array([[60.0], [36.0], [58.0]]),  # each spike's mark
    rng=0,
)
print(spikes.predictive_pvalue.round(2), spikes.hpd_overlap.round(2))

[0.08]
[0.08 0.95 0.87] [0.46 1.   1.  ]

The observed mark lies in the tail of the predictive mark distribution: the Monte Carlo p-value is about 0.08 (numerical integration gives 0.083), above the paper's 0.05 cutoff. clusterless_event_diagnostics computes HPD overlap, KL divergence and this p-value for every spike, as event_diagnostics does for sorted units. The first spike repeats the example above (mark 60 under a prediction centred at 35): its single-event likelihood has the lowest HPD overlap with the prediction (0.46), and its p-value is again about 0.08. Results are reproducible for a fixed seed and batch_size. The clusterless tutorial decodes a simulated clusterless recording and diagnoses a misspecified mark model.

Documentation

https://edeno.github.io/statespacecheck: tutorials, interpreting the diagnostics, using the package with your decoder, and the API reference.

Citation

If you use this package in your research, please cite it; CITATION.cff in the repository records the version and release date:

@software{statespacecheck,
  title={statespacecheck: Goodness-of-fit diagnostics for state space models},
  author={Denovellis, Eric and Zeng, Sirui and Eden, Uri T.},
  url={https://github.com/edeno/statespacecheck},
  doi={10.5281/zenodo.22999988}
}

Releases are archived on Zenodo. The DOI above, 10.5281/zenodo.22999988, covers all versions and resolves to the latest; to cite a specific version, use its DOI from the Zenodo record (version 0.3.0 is 10.5281/zenodo.22999989).

Please also cite the companion paper for the methods: Local goodness-of-fit measures for neural decoding (Zeng, Comrie, Frank, Eden and Denovellis); its analysis code is at https://github.com/edeno/statespacecheck-paper.

Contributing and license

Contributions are welcome; see CONTRIBUTING.md. MIT License.