Skip to main content

SIMPL logo

Tests Colab demo Paper Docs PyPI Downloads

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, use the platform-specific instructions in GPU Acceleration.




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. Load your data
Y = ...     # spike counts (T, N_neurons)
Xb = ...    # behavioural initialisation positions (T, D)
time = ...  # timestamps (T,) or use time = np.linspace(0, T_end, len(Y)) if length is known but timestamps aren't

# 2. 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
)

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

# 4. 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.

# 5. 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.

Y_new = ...  # new spike counts (T_new, N_neurons)
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 — to create figures from scratch, 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=60)  # first 60 seconds
model.plot_latent_trajectory(time_range=(0, 60), iterations=(0, 3, 5))  # 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, # Xt_test is optional ground truth for the prediction data
  )


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, pass time=None when calling fit():

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

When time=None, 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 marking 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. Most typical neural recordings 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

Choose exactly one installation path:

# CPU (recommended; macOS, Linux, or Windows)
pip install simpl-neuro

# NVIDIA CUDA 12 (Linux/Windows only)
pip install -U "jax[cuda12]"
pip install simpl-neuro

# Apple Silicon Metal (experimental)
pip install "simpl-neuro[metal]"

Do not install the CUDA extra on macOS.

pip install --force-reinstall "simpl-neuro[metal]"
# Local editable checkout:
pip install --force-reinstall -e ".[metal]"
model = SIMPL()               # use a GPU when available (default)
model = SIMPL(use_gpu=False)  # force CPU
model = SIMPL(use_gpu=True)   # require a GPU

With use_gpu=False, SIMPL selects CPU directly without initialising installed GPU or Metal backends.

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

Open In Colab




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}
}

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

simpl_neuro-0.10.5.tar.gz (12.0 MB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

simpl_neuro-0.10.5-py3-none-any.whl (84.7 kB view details)

Uploaded Python 3

File details

Details for the file simpl_neuro-0.10.5.tar.gz.

File metadata

  • Download URL: simpl_neuro-0.10.5.tar.gz
  • Upload date:
  • Size: 12.0 MB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.5

File hashes

Hashes for simpl_neuro-0.10.5.tar.gz
Algorithm Hash digest
SHA256 74ca274969177b842532b0a36b42ba24541a50ab49d36a2742605eaeabae1e59
MD5 3a8ede54db248ade1c01792a37738243
BLAKE2b-256 b07b56c8ef1c0c5e34b6f18e4d73b48e87ef29f386bbb70dc6f5304613dadec6

See more details on using hashes here.

File details

Details for the file simpl_neuro-0.10.5-py3-none-any.whl.

File metadata

  • Download URL: simpl_neuro-0.10.5-py3-none-any.whl
  • Upload date:
  • Size: 84.7 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.5

File hashes

Hashes for simpl_neuro-0.10.5-py3-none-any.whl
Algorithm Hash digest
SHA256 80413af1deaa38b2542590610ac3325e40e0d4cb395fbf91eab8bbe7c6f10d69
MD5 7d497448b9df4886c36ea980d45c6aad
BLAKE2b-256 48f53eff8e093644b152dc47a47b6e816e4b086cfc852466768c46e79ef39417

See more details on using hashes here.

Release history Release notifications | RSS feed

0.11.1

2 files

0.11.0

2 files

This release

0.10.5 This release

2 files

0.10.4

2 files

0.10.2

2 files

0.10.1

2 files

0.10.0

2 files

0.9.3

2 files

0.9.2

2 files

0.9.1

2 files

0.9.0

2 files

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page