simpl-neuro 0.10.0

SIMPL: Scalable and hassle-free optimisation of neural representations from behaviour

SIMPL is a JAX-python package for optimising latent representations and neural tuning curves from spike data. It does this by iteratively decoding the latent and fitting the tuning curves, starting from behavior or stimuli. It is lightweight, scalable, and very fast. Published at ICLR 2025.

Install | Quickstart | Examples | Cite

---

Key Features

• Fast — fits 200 neurons over 1 hour of data in under 10 seconds on CPU. GPU optional but rarely needed.
• Scalable - scales to state-of-the-art size neural datasets (1000s of neurons, millions of time points, billions of spikes) on CPU.
• Simple — scikit-learn API. Minimal hyperparameters. Get started in <10 lines of code.
• Flexible — 1D, 1D angular (e.g. head direction), 2D spatial (e.g. place/grid cells), and higher dimensional data all permitted. Trial-structure aware. Temporal or non-temporal data.
• Rich outputs — results stored as xarray.Dataset with per-iteration variables, metrics and units.
• Visual — built-in plotting for trajectories, receptive fields, spike rasters, and fitting summaries.

Neural data analysis in < 5 seconds

Installation

pip install simpl-neuro

To run the demo notebook locally (or else )

pip install "simpl-neuro[demos]"
simpl demo                # downloads demo notebook into the cwd

If you need GPU, see the Advanced Usage section.

---

Quickstart

SIMPL follows sklearn conventions: configure hyperparameters at init, pass data to fit(). For hyperparameter units, see the Model / Maths section.

from simpl import SIMPL

# 1. Configure the model (no data, no computation)
model = SIMPL(
    speed_prior=0.4,        # "0.4m/s" prior on latent speed
    behavior_prior=None,    # (optional) soft tether to the initial behaviour/stimulus
    kernel_bandwidth=0.02,  # "2 cm" kernel bandwidth for KDE spike smoothing
    bin_size=0.02,          # "2 cm" spatial bin size for environment discretisation
)

# 2. Fit
model.fit(
    Y,                      # spike counts (T, N_neurons)
    Xb,                     # behavioural initialisation positions (T, D)
    time,                   # timestamps (T,)
    n_iterations=5,
    )

# 3. Access results
model.X_           # final decoded latent positions, shape (T, D)
model.F_           # final receptive fields, shape (N_neurons, *env_dims)
model.results_     # full xarray.Dataset with metrics, likelihoods, and baselines, across iterations.

# 4. Plot results 
model.plot_fitting_summary()  # Shows bits-per-spike metric and spike-latent mutual information. 

# (optional) Resume training if not yet converged
model.fit(Y, Xb, time, n_iterations=5, resume=True)

Prediction

Decode new spikes using the fitted receptive fields (no behavioural input needed). The new data must be binned at the same dt as the training data.

X_decoded = model.predict(Y_new)
model.prediction_results_  # xr.Dataset with rich results (mu_s, sigma_s, log-likelihoods, etc.)

---

Model, Maths and Notation

Notation

SIMPL tries to stick to the following notations:

• Latent trajectory: $X \in \mathbb{R}^{T \times D}$ in maths, X / model.X_ in code. $X_t$ is the inferred latent at time bin $t$.
  • Latent-space coordinate: $x$ in maths, a grid point in model.xF_ in code. This is a possible position/location, not a whole trajectory.
  • Behavioral initialisation: $X_b$ in maths, Xb in code. This starts the fit and can optionally tether the latent through behavior_prior. $X_t$/Xt is ground truth (if known).
• Spike counts: $Y \in \mathbb{R}^{T \times N}$. $y_t$ is one time-bin vector, and $y_{t,n}$ is one neuron's count in one time bin.
• Receptive fields: $F \in \mathbb{R}^{N \times N_{\textrm{bins}}}$. F / model.F_ in code are reshaped to the environment size, e.g. F.shape = (N, N_x_bins, N_y_bins, ...). $F_n(x)$ is neuron $n$'s expected spike count at latent-space point $x$. Thus $F_n(X_t)$ is neuron $n$'s expected spike count at the decoded latent position and is the Poisson rate parameter.

Full model

This is only a summary, see ICLR paper for full details.

At its heart SIMPL approximately optimises a latent trajectory $X_{1:T}$ and receptive fields $F(x)$ under:

$$ p(X_{1:T}, Y \mid F) ;\propto; \prod_t \underbrace{{\color{A92E5E}p(y_t \mid X_t, F)}}{{\color{A92E5E}\mathrm{observation\ model}}}, \underbrace{{\color{1D5C84}p(X_t \mid X{t-\Delta t})}}_{{\color{1D5C84}\mathrm{dynamics\ model}}} $$

Dynamics model

The temporal prior is a Gaussian random-walk model controlled by $\sigma_v$ (speed_prior):

$$ {\color{1D5C84}p(X_t \mid X_{t-\Delta t})} \approx {\color{1D5C84}\mathcal{N}!\left(X_t; X_{t-\Delta t}, (\sigma_v,\Delta t)^2 I\right)} $$

Smaller speed_prior values enforce smoother decoded trajectories; larger values let each time bin follow the spike likelihood more freely. Set speed_prior=None disables Kalman smoothing. For data without meaningful temporal structure, see Temporal vs. non-temporal datasets.

Optional (behavior_prior)
SIMPL can also include a soft Gaussian tether to whatever the latent was initialised to (typically behavior), controlled by $\sigma_b$ (behavior_prior):

$$ {\color{1D5C84}p(X_t \mid X_{t-\Delta t})} \propto \underbrace{{\color{1D5C84}\mathcal{N}!\left(X_t; X_{t-\Delta t}, (\sigma_v,\Delta t)^2 I\right)}}{{\color{1D5C84}\mathrm{latent\ close\ to\ previous\ latent}}} , \cdot \underbrace{{\color{A3CC90}\mathcal{N}!\left(X_t; X_t^{(0)}, \sigma_b^2 I\right)}}{{\color{A3CC90}\mathrm{latent\ close\ to\ initialisation}}} $$

Smaller behavior_prior values enforce decoded trajectory to be closer to the behavior passed into Xb; larger values let each time bin follow the spike likelihood more freely (potentially moving far from behavior). Set behavior_prior=None to strictly disable behavior tether.

Once fields are estimated an approximation (see paper) converts Poisson-nonlinear observations to linear observations, allowing these dynamics to be inferred with fast Kalman smoothing.

Observation model

The spike likelihood comes from the fitted tuning curves:

$$ {\color{A92E5E}p(y_t \mid X_t, F)} = {\color{A92E5E}\prod_n \mathrm{Poisson}!\left(y_{t,n}; F_n(X_t)\right)} $$

where, for neuron $n$, $F_n(X_t)$ is the expected spike count in that time bin, i.e. its tuning curve evaluated at the decoded latent position. The tuning curve itself is estimated by the standard KDE equation from the current latent:

$$ {\color{A92E5E}F_n(x) = \frac{\sum_t y_{t,n},K(x, X_t)}{\sum_t K(x, X_t)}} $$

$K$ is a Gaussian kernel with bandwidth kernel_bandwidth. The denominator corrects for non-uniform occupancy. Receptive fields are evaluated on a spatial grid with bin size $\Delta x$, but decoded positions are not restricted to those grid points.

Units and Discretisation

All hyperparameters (e.g. speed_prior, kernel_bandwidth, bin_size etc.) are defined in data units (e.g. typically [m/s], [m], [m] but these depend on your data of course), not arbitrary time/spatial-bin units.

---

Plotting and Metrics

Metrics

The four headline fitting metrics are:

• Spike log-likelihood (logPYXF, logPYXF_val) — the mean Poisson log-likelihood of the observed spike counts under the fitted receptive fields evaluated along the decoded trajectory.

  We strongly suggest using bits_per_spike over logPYXF for comparisons. Its zero point is interpretable, and the normalisation makes it easier to compare datasets with different neuron counts or recording lengths.

$$ \mathcal{L}

\sum_t \sum_n \log \mathrm{Poisson}!\left(y_{t,n}; F_n(X_t)\right) $$

• Bits per spike (bits_per_spike, bits_per_spike_val) — how much better the fitted tuning curves explain spikes than a mean-rate baseline, in bits per observed spike. This is useful for comparing fits across datasets with different spike counts or bin sizes:

$$ \mathrm{BPS} = \frac{\mathcal{L}(\hat{F}) - \mathcal{L}(\bar{F})}{N_{\mathrm{spk}} \ln 2} $$

• Skaggs spatial information (spatial_information) — the standard spatial-information rate for each neuron, in bits/sec. It measures how informative the receptive field is about latent position in the small-time-bin limit:

$$ I_{\mathrm{Skaggs}} = \frac{1}{\Delta t}\sum_x p_X(x),F_n(x),\log_2 \frac{F_n(x)}{\bar{F}_n} $$

Other metrics available in model.results_ include:

• mutual_information — the exact finite-time-bin analog of spatial information, $I(X;Y)$, in bits/s. It is the same idea as Skaggs spatial information, but computed from the full spike-count distribution rather than the small-bin approximation.
• X_R2, X_err — latent-position agreement with ground truth, when Xt is registered with add_baselines.
• F_err — receptive-field error against ground-truth fields, when Ft is registered.
• stability — correlation between fields estimated from odd and even minutes.
• field_change, trajectory_change — per-iteration changes in tuning curves and decoded trajectory.
• negative_entropy, sparsity — compactness/sparsity summaries of the fitted tuning curves.

For 2D environments, model.analyse_place_fields() adds morphology metrics to model.results_, including place-field count, size, position, roundness, and peak firing rate. This uses connected-component analysis on receptive fields and is not run automatically during fit().

Plotting

Built-in plotting methods provide quick diagnostics. All methods return matplotlib Axes for further customisation — for publication-quality figures, use model.results_ (an xarray.Dataset) to access the data directly.

# Log-likelihood and spatial information across iterations
model.plot_fitting_summary()

# Decoded trajectory (all iterations by default)
model.plot_latent_trajectory()
model.plot_latent_trajectory(time_range=(0, 60))  # zoom in, specific iterations

# Receptive fields (iteration 0 + last by default)
model.plot_receptive_fields(neurons=[0, 5, 10])

# Spike raster heatmap (time × neurons)
model.plot_spikes()
model.plot_spikes(time_range=(0, 60))

# Auto-discover and plot all per-iteration metrics
model.plot_all_metrics(show_neurons=False)

# 2D place-field morphology metrics
model.analyse_place_fields()

# Prediction on held-out data
model.predict(Y_test)
model.plot_prediction(Xb=Xb_test, Xt=Xt_test)

Synthetic grid cell tuning curves optimised from a noisy behavioural initialisation

True latent trajectory recovered by SIMPL

Bits-per-spike and mutual-information metrics improve across epochs and exceed naive ML

---

Advanced Usage

Saving and loading

model.save_results("results.nc")

# Load results as an xr.Dataset for custom analysis
from simpl import load_results
results = load_results("results.nc")

# Or rehydrate a full model for plotting, prediction, or resumed training
# (constructor arguments must exactly match the original training run)
model = SIMPL(speed_prior=0.4, kernel_bandwidth=0.025, bin_size=0.02)
model.load("results.nc")
model.fit(Y, Xb, time, n_iterations=5, resume=True)  # pick up where you left off

Ground truth baselines

If you have ground truth positions (and optionally ground truth receptive fields), register them before fitting so that baseline metrics (latent R2, field error, etc.) are computed at each iteration:

model.add_baselines(Xt=Xt, Ft=Ft, Ft_coords_dict={"y": ybins, "x": xbins})
model.fit(Y, Xb, time, n_iterations=5)  # baselines computed automatically

Temporal vs. non-temporal datasets

Temporal structure enters SIMPL only through the Kalman smoothing prior. A temporal dataset might be hippocampal navigation, where neighboring samples are adjacent moments in an animal's trajectory and smoothness over time is meaningful. A non-temporal dataset might be neural responses to visual stimuli, where each point is a separate trial and neighboring rows do not imply temporal continuity.

Once Kalman smoothing is disabled, SIMPL is, in effect, iterative maximum likelihood: each observation is decoded independently from the current receptive fields, then the receptive fields are refit from those decoded positions.

If your samples do not have meaningful time stamps or temporal ordering, omit time when calling fit():

model = SIMPL(speed_prior=None)
model.fit(Y, Xb)

When time is omitted, SIMPL treats the data as non-temporal, replaces the missing time coordinate with np.arange(T), and forcefully disables Kalman smoothing. In this setting, the saved time coordinate is better interpreted as trial/sample index rather than physical time.

1D angular / circular data

SIMPL supports 1D circular latent variables (e.g. head direction) via the is_1D_angular flag. When enabled, the environment is fixed to [-π, π), angular KDE is used for receptive fields, and the Kalman filter wraps its state to [-π, π) after every predict, update, and smooth step.

model = SIMPL(
    is_1D_angular=True,
    bin_size=np.pi / 32,
    env_pad=0.0,
    speed_prior=0.1,
    kernel_bandwidth=0.3,
)
model.fit(Y, Xb, time, n_iterations=5)  # Xb should be in radians, [-pi, pi)

Note: The wrapped Kalman filter assumes a tight posterior (σ ≪ 2π). If posterior uncertainty is large relative to the circular domain, decoding accuracy may degrade.

Trial boundaries

When data comes from multiple recording sessions or trials, you don't want the Kalman smoother blending across discontinuities. Pass trial_boundaries — an array of time-bin indices where each new trial starts — and SIMPL will run the filter/smoother independently within each segment. The initial state for each trial is estimated from the likelihood modes within that trial.

# Three trials starting at time-bins 0, 5000, and 12000
model.fit(Y, Xb, time, n_iterations=5, trial_boundaries=[0, 5000, 12000])

If your timestamps have gaps (e.g. concatenated sessions), SIMPL will warn you and suggest using trial_boundaries to avoid smoothing across the jumps.

GPU Acceleration

SIMPL auto-detects and offloads compute-heavy steps to GPU when available. Typical neural recordings (< 2 hrs) fit in under 60 s on CPU alone, so a GPU is rarely needed.

200 neurons, dt=0.02s (50Hz), dx=2cm (2,500 bins), 5 iterations, includes JIT overheads

pip install -U "jax[cuda12]"   # NVIDIA GPU (CUDA)
pip install ".[metal]"           # Apple Silicon GPU (experimental and not recommended, pins JAX to 0.4.35)

model = SIMPL(use_gpu=False)   # force CPU

Data preprocessing utilities

from simpl import accumulate_spikes, coarsen_dt

# Roll up spikes into wider time bins (e.g. sum every 2 bins)
Y_coarse, Xb_coarse, time_coarse = coarsen_dt(Y, Xb, time, dt_multiplier=2)

# Accumulate spikes with a causal sliding window
Y_accum = accumulate_spikes(Y, window=3)

---

Examples/Demos

The examples/simpl_demo.ipynb notebook walks through the full SIMPL workflow across four datasets:

1. Synthetic grid cells — fits SIMPL on artificial grid cell data with known ground truth, demonstrating decoded trajectories, receptive field recovery, log-likelihood improvements, and prediction on held-out data.
2. Real place cells — fits SIMPL on real hippocampal place cell recordings from Tanni et al. (2022), where no ground truth is available.
3. Real head direction cells — fits SIMPL in 1D angular mode on head direction cell recordings from Vollan et al. (2025), demonstrating circular latent variable decoding and polar receptive field plots.
4. Motor cortex hand reaching — fits SIMPL on somatosensory cortex recordings from Chowdhury et al. (2020), demonstrating higher-dimensional latent variables (2D and 4D) and model comparison across different behavioural initialisations (position vs velocity vs combined).

---

Code and Development

Package Structure

src/simpl/
├── __init__.py        # Top-level exports: SIMPL, load_datafile, ...
├── simpl.py           # Core SIMPL class (EM algorithm, fit/predict)
├── kde.py             # KDE, Poisson log-likelihood, gaussian_kernel
├── kalman.py          # KalmanFilter class + Kalman functions
├── environment.py     # Environment class (spatial discretisation)
├── plotting.py        # Built-in diagnostic plots (trajectory, fields, metrics)
└── utils.py           # Gaussian helpers, CCA, data prep, I/O

Development

# Install for development
pip install -e ".[dev]"

# Lint
ruff check src/
ruff format --check src/

# Run tests
pytest

---

Cite

If you use SIMPL in your work, please cite it as:

Tom George, Pierre Glaser, Kim Stachenfeld, Caswell Barry, & Claudia Clopath (2025). SIMPL: Scalable and hassle-free optimisation of neural representations from behaviour. In The Thirteenth International Conference on Learning Representations.

@inproceedings{
    george2025simpl,
    title={{SIMPL}: Scalable and hassle-free optimisation of neural representations from behaviour},
    author={Tom George and Pierre Glaser and Kim Stachenfeld and Caswell Barry and Claudia Clopath},
    booktitle={The Thirteenth International Conference on Learning Representations},
    year={2025},
    url={https://openreview.net/forum?id=9kFaNwX6rv}
}