wishart-process-em

Finite-basis Wishart process models of neural noise covariance, fit by marginal-likelihood (QMC + Laplace) EM.

Estimate how the trial-to-trial covariance of a neural population changes smoothly across continuously-parameterised experimental conditions (grating orientation, reach angle, tone frequency, …), pooling statistical power from neighbouring conditions so that accurate covariance estimates are possible even with few trials per condition.

This package is a self-contained JAX implementation of the Wishart process model of Nejatbakhsh, Garon & Williams (2023; NeurIPS), with two deliberate differences from the reference implementation:

1. Finite (Fourier-feature) basis. Each latent Gaussian process is represented in weight space by a truncated Fourier expansion, turning the nonparametric GP into a finite, differentiable model.
2. Marginal-likelihood / EM-style inference instead of variational inference. Per-trial latents are integrated out with a randomised quasi-Monte-Carlo lattice and an optional Laplace-approximation importance-sampling correction; the marginal likelihood is then maximised with optax.

It supports Gaussian, Poisson, and negative-binomial observations and the full covariance family from the paper, Σ(x) = L (U(x)U(x)ᵀ + Λ(x)) Lᵀ.

---

Installation

pip install wishart-process-em            # from PyPI
# or, from source:
git clone https://github.com/neurostatslab/wishart-process-em
cd wishart-process-em
pip install -e ".[dev]"

Requires Python ≥ 3.12 (following nemos, on which the basis functions are built). For plotting (and the comparison baselines in examples/), install the extras: pip install "wishart-process-em[viz]".

Quickstart

import jax.numpy as jnp, jax.random as jxr
from wishart_process_em import (
    fourier_basis, WishartProcessModel, fit, squared_exponential,
)

# 1. A Fourier basis (a nemos FourierEval) over a 1-D periodic condition.
basis = fourier_basis(max_freq=8, num_dims=1)

# 2. A Gaussian Wishart process over N=25 neurons, rank P=3. The spectral
#    density sets the GP smoothness across conditions.
model = WishartProcessModel(basis, num_neurons=25, rank=3, likelihood="gaussian",
                            spectral_density=squared_exponential(lengthscale=0.2))

# 3. Simulate some data from the generative model.
true_params = model.init_params(jxr.PRNGKey(0))
X = jnp.linspace(0, 1, 400)                       # 400 trials, condition in [0,1]
Y, _ = model.sample(jxr.PRNGKey(1), true_params, X)

# 4. Fit and predict smooth condition-dependent covariance.
result = fit(model, Y, X, num_steps=500, learning_rate=1e-1)
Sigma = model.predict_cov(result.params, jnp.linspace(0, 1, 50))   # (50, N, N)

For spike counts, pass likelihood="poisson" (or "negative_binomial"): the latent integral is then estimated by QMC automatically.

model = WishartProcessModel(basis, num_neurons=25, rank=3, likelihood="poisson",
                            spectral_density=squared_exponential(lengthscale=0.2))
result = fit(model, counts, X, num_steps=1000)          # uses QMC each step

How it works

Component	Module	Notes
Weight-space GP basis	basis, kernels	nemos FourierEval features scaled by a kernel spectral density.
QMC integration	qmc	Randomised Korobov lattice → Gaussian latents.
Laplace approximation	optim	JAX-traceable damped Newton returning the Hessian Cholesky.
Observation models	likelihoods	Gaussian (conjugate), Poisson, negative-binomial.
Covariance Σ(x)=L(UUᵀ+Λ)Lᵀ	covariance, model	Rank-P factor, optional diagonal, optional scale matrix.
Marginal likelihood	inference	Conjugate / QMC / Laplace estimators.
Fitting	fit	optax-based, stochastic-EM for count models.
Metrics & diagnostics	diagnostics	Held-out log-likelihood, Fisher information, QDA decoding.

See docs/model.md for the full mathematical description.

Citation

If you use this software, please cite the paper:

@inproceedings{nejatbakhsh2023wishart,
  title     = {Estimating Noise Correlations Across Continuous Conditions With Wishart Processes},
  author    = {Nejatbakhsh, Amin and Garon, Isabel and Williams, Alex H.},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year      = {2023},
}

License

MIT — see LICENSE.