Metadata-Version: 2.5
Name: nldisco
Version: 1.0.0
Summary: Neural latent discovery with sparse encoder-decoder models
Project-URL: Homepage, https://github.com/jkbhagatio/nldisco
Project-URL: Documentation, https://github.com/jkbhagatio/nldisco/blob/main/docs/README.md
Project-URL: Repository, https://github.com/jkbhagatio/nldisco
Project-URL: Issues, https://github.com/jkbhagatio/nldisco/issues
Author-email: Jai Bhagat <jkbhagatio@gmail.com>
License-Expression: PolyForm-Noncommercial-1.0.0
License-File: LICENSE.md
Requires-Python: <3.14,>=3.9
Requires-Dist: beartype<0.23,>=0.18
Requires-Dist: datasets>=2.14.6
Requires-Dist: einops
Requires-Dist: email-validator
Requires-Dist: huggingface-hub
Requires-Dist: hydra-core<1.4,>=1.3
Requires-Dist: ipython
Requires-Dist: ipywidgets
Requires-Dist: jax>=0.4.21
Requires-Dist: jaxtyping
Requires-Dist: matplotlib
Requires-Dist: numpy
Requires-Dist: openai
Requires-Dist: pandas
Requires-Dist: plotly<6,>=5.23
Requires-Dist: python-dotenv<2,>=1
Requires-Dist: rich
Requires-Dist: scikit-learn
Requires-Dist: scipy
Requires-Dist: seaborn
Requires-Dist: sortedcontainers
Requires-Dist: submitit<2,>=1.5
Requires-Dist: torch
Requires-Dist: torcheval
Requires-Dist: torchinfo
Requires-Dist: torchvision
Requires-Dist: tqdm
Requires-Dist: typeguard
Requires-Dist: wandb
Provides-Extra: analysis
Requires-Dist: cebra; (python_version < '3.10') and extra == 'analysis'
Requires-Dist: cebra>=0.6.1; (python_version >= '3.10') and extra == 'analysis'
Requires-Dist: setuptools<81; (python_version < '3.10') and extra == 'analysis'
Provides-Extra: churchland
Requires-Dist: brainsets; extra == 'churchland'
Requires-Dist: dandi; extra == 'churchland'
Requires-Dist: h5py; extra == 'churchland'
Requires-Dist: requests; extra == 'churchland'
Requires-Dist: temporaldata; extra == 'churchland'
Description-Content-Type: text/markdown

<div align="center">

# NLDisco

**Neural Latent Discovery**  
Find interpretable features in neural population activity.

[Quick start](#quick-start) · [User guide](#user-guide) · [Documentation](https://github.com/jkbhagatio/nldisco/blob/main/docs/README.md) · [Paper](https://github.com/jkbhagatio/nldisco/blob/paper/paper/iclr_paper/full_paper.pdf) · [Experiments](https://github.com/jkbhagatio/nldisco/blob/main/experiments/README.md)

</div>

![NLDisco pipeline: preprocess neural data, train sparse models, evaluate reconstruction, and interpret latents.](https://raw.githubusercontent.com/jkbhagatio/nldisco/main/docs/assets/figure-1.png)

NLDisco learns sparse, overcomplete representations of neural activity. Each latent
is a candidate feature to inspect alongside behavior, stimuli, or other metadata.
Training uses neural activity alone; behavioral labels enter during interpretation.

- **Discover at multiple scales.** Matryoshka dictionaries learn nested sets of latents.
- **Capture temporal context.** Encode single bins or complete windows with linear or
  transformer encoders, including an optional shift-equivariant mode.
- **Connect populations.** Reconstruct the input population or predict an aligned target population.
- **Inspect what you learn.** Evaluate reconstruction, latent ablations, and spectral fidelity.

![Example datasets: simulated rat navigation, macaque reaching, and mouse foraging.](https://raw.githubusercontent.com/jkbhagatio/nldisco/main/docs/assets/task_overview.png)

**NLDisco across datasets:** simulated navigation, macaque reaching, and mouse foraging.

## From activity to interpretable latents

![Sparse encoder-decoder architectures: single-bin, Matryoshka, and window models.](https://raw.githubusercontent.com/jkbhagatio/nldisco/main/docs/assets/figure-2.png)

**Figure 2 · Model architecture.** **(a)** Natural features are jointly encoded in neural
activity. **(b)** A sparse encoder-decoder reconstructs neural activity through a small
set of active latents; targets can be the input population or a separate, aligned
population. **(c)** Matryoshka levels each reconstruct the target using nested latent
sets. **(d)** Window models reconstruct every timebin using either a flattened linear
encoder (FW-SED) or a transformer encoder (TW-SED), with reconstruction scored per bin.

The Python API uses **[batch, timebin, unit]** tensors. Build a window model with:

```python
from nldisco import EncoderConfig, SedConfig, build_sed

model = build_sed(SedConfig(
    n_neurons=100,
    seq_len=8,
    dsed_topk_map={128: 4, 512: 8},
    encoder=EncoderConfig(type="TransformerWindow"),
))
```

Use `FlatWindow` for a linear encoder, or set `seq_len=1` for a single-bin SED.
For paired populations, set `n_output_neurons` and supply `targets=` to
`SpikeWindowDataset`; the CLI equivalents are `data.target_path` and
`data.target_normalization`. Inputs and targets share their time grid and splits,
while unit counts can differ. Inference needs only inputs.

See [models and inference](https://github.com/jkbhagatio/nldisco/blob/main/docs/model_names.md), the
[paired-population example](https://github.com/jkbhagatio/nldisco/blob/main/examples/paired_transcoder.py), and
[evaluation](https://github.com/jkbhagatio/nldisco/blob/main/docs/evaluation.md) for details.

## Why individual latents?

![Toy example comparing neural activity, a tangled two-dimensional latent space, and individual NLDisco latents.](https://raw.githubusercontent.com/jkbhagatio/nldisco/main/docs/assets/figure-s1.png)

**Figure S1 · Interpretable latents in a complex latent space.** A toy example with a
discrete variable (top) and a continuous variable (bottom). **(a)** Both are encoded
in three neurons' firing rates. **(b)** A two-dimensional projection creates tangled
trajectories; white and black dots mark their start and end. **(c)** Individual
NLDisco latents track the variables, illustrating the aim of interpretable feature
discovery. This is a conceptual illustration, not a benchmark result.

## Method comparison

[![Table S1: comparison of NLDisco with other neural latent-variable methods, including qualifications and footnotes.](https://raw.githubusercontent.com/jkbhagatio/nldisco/main/docs/assets/table-s1.png)](https://raw.githubusercontent.com/jkbhagatio/nldisco/main/docs/assets/table-s1.png)

The paper's qualitative comparison of 15 methodological features. Click the table
for full resolution; see the [paper](https://github.com/jkbhagatio/nldisco/blob/paper/paper/iclr_paper/full_paper.pdf) for context and references.

## Explore and develop

| Resource | Contents |
| --- | --- |
| [Documentation](https://github.com/jkbhagatio/nldisco/blob/main/docs/README.md) | Preprocessing, models, training, checkpoint reloads, and evaluation |
| [Examples](https://github.com/jkbhagatio/nldisco/blob/main/docs/tutorials/README.md) | Runnable walkthroughs using generated data |
| [Experiments](https://github.com/jkbhagatio/nldisco/blob/main/experiments/README.md) | Paper datasets, configurations, notebooks, and result artifacts |
| [Library](https://github.com/jkbhagatio/nldisco/blob/main/src/nldisco/) | Model, data, training, evaluation, plotting, and sweep code |

```console
uv run pytest                             # Core library tests
uv run pytest experiments/tests           # Paper-analysis tests
uv run jupyter lab
```

CEBRA comparisons require `uv sync --locked --extra analysis`; Churchland data
helpers additionally require `--extra churchland`. See [installation](https://github.com/jkbhagatio/nldisco/blob/main/docs/installation.md).
To register a Jupyter kernel, run `uv run python -m ipykernel install --user --name=nldisco`.

## Quick start

NLDisco requires **Python 3.9–3.13**.

To install the library into an existing, activated Python environment, use either:

```console
uv pip install nldisco
# Or:
python -m pip install nldisco
```

To work with the source, examples, or notebooks, install
[uv](https://docs.astral.sh/uv/getting-started/installation/) and clone the repository:

```console
git clone https://github.com/jkbhagatio/nldisco.git nldisco
cd nldisco
uv sync --locked
uv run python examples/paired_transcoder.py --layout flat
```

See [installation](https://github.com/jkbhagatio/nldisco/blob/main/docs/installation.md) for editable installs into existing
environments, optional dependencies, and development setup. Commands prefixed
with `uv run` below assume the source checkout; in an existing environment,
activate it and omit `uv run`.

The example generates its own data, trains a model, evaluates it, and demonstrates
inference. It also supports `single_bin`, `transformer`, and `shift_equivariant` layouts.

For your own data, supply a numeric `.npy` array shaped **[timebin, unit]**:

```console
# Train one model on nonnegative spike counts.
uv run python -m nldisco.sweep --config-name train data.path=/path/to/counts.npy training.epochs=10

# Sweep learning rates and seeds with two local workers.
uv run python -m nldisco.sweep data.path=/path/to/counts.npy execution.max_parallel=2
```

The defaults use MSLE loss and a ReLU decoder. For signed targets, including z-scores,
set `loss.type=mse model.decoder.output_activation=none`.
See [preprocessing](https://github.com/jkbhagatio/nldisco/blob/main/docs/preprocessing.md) for Kilosort/Phy loading and spike binning,
and [training and sweeps](https://github.com/jkbhagatio/nldisco/blob/main/docs/training_and_sweeps.md) for GPU, W&B, and Slurm execution.

## User guide

A Kilosort/Phy recording can move through the full pipeline in five steps.

**1. Load and bin spikes.** `load_kilosort` reads sorter exports and retains the
unit-ID mapping. Set the actual sampling rate and recording duration:

```python
import numpy as np
from nldisco.data import load_kilosort

binned = load_kilosort(
    "/path/to/kilosort", sampling_rate=30_000, bin_size=0.02,
    start_time=0, stop_time=600,
)
np.save("counts.npy", binned.counts)
np.save("timestamps.npy", binned.timestamps)
np.save("unit_ids.npy", binned.unit_ids)
```

**2. Preprocess and split.** Keep the saved counts raw. The runner below fits z-score
normalization on training rows only, then applies it to validation rows. It builds
8-bin windows with an 80/20 chronological split. Supply trial/session IDs and validity
masks when needed to prevent windows crossing boundaries; see [preprocessing](https://github.com/jkbhagatio/nldisco/blob/main/docs/preprocessing.md).

**3. Train with a W&B sweep.** Authenticate with `uv run wandb login` (or set
`WANDB_API_KEY` in `.env`), then launch a Bayesian search over learning rates and seeds:

```console
uv run python -m nldisco.sweep --config-name wandb \
  data.path=counts.npy data.timestamps=timestamps.npy data.expected_bin_size=0.02 \
  data.normalization=zscore loss.type=mse model.decoder.output_activation=none \
  training.epochs=20 search.max_runs=10 wandb.project=nldisco-kilosort
```

W&B minimizes `validation/loss`. Each trial saves its configuration, calibrated
checkpoint, and reconstruction metrics under `outputs/nldisco/`.
[Training and sweeps](https://github.com/jkbhagatio/nldisco/blob/main/docs/training_and_sweeps.md) covers search spaces and GPU/Slurm settings.

**4. Evaluate the model.** Pick a trial using validation results and reload its checkpoint.
With the original recording files unchanged, rebuild the same split and preprocessing:

```python
from nldisco.sweep.training import load_checkpoint, model_and_loss, prepare_data
from nldisco.train import evaluate_model
from nldisco.plot import plot_reconstruction_by_lag

saved = load_checkpoint("outputs/nldisco/<invocation>/run-0000/model.pt")
model = saved.model
_, loss = model_and_loss(saved.config, model.cfg.n_neurons, model.cfg.n_output_neurons)
_, validation, _ = prepare_data(saved.config, model.cfg)
result = evaluate_model(model, validation, loss)
print(result.metrics_by_lag)
plot_reconstruction_by_lag(result.metrics_by_lag)
```

Optional [ablation and spectral diagnostics](https://github.com/jkbhagatio/nldisco/blob/main/docs/evaluation.md) test individual latent
contributions and frequency fidelity. Keep a separate test set for final reporting.

**5. Evaluate the latents.** Inspect decoder patterns and align latent activations with
behavior. For example, plot latent 0 over the eligible validation positions, retaining
zeros where it was inactive:

```python
from nldisco.plot import plot_decoder_feature

plot_decoder_feature(model.decoder, latent_idx=0)
latent = result.evaluation_index[["source_time_idx"]].merge(
    result.activation_table.query("latent_idx == 0")[["source_time_idx", "activation_value"]],
    on="source_time_idx", how="left",
).fillna({"activation_value": 0})
latent["time_s"] = np.load("timestamps.npy")[latent.source_time_idx]
latent.plot(x="time_s", y="activation_value")
```

Synchronize behavioral measurements to the recording clock, then join them through
`source_time_idx`. Compare tuning and activating examples; quantify candidate features
with coverage, specificity, and AUROC over both active and inactive positions.
Decoder weights describe reconstructed activity; they are not causal attributions.
See the [paper experiments](https://github.com/jkbhagatio/nldisco/blob/main/experiments/README.md) for dataset-specific latent analyses.

## License

NLDisco's original software is source-available under the
[PolyForm Noncommercial License 1.0.0](https://github.com/jkbhagatio/nldisco/blob/main/LICENSE.md). It permits use, modification,
and redistribution for the purposes described in the license, including
noncommercial use and use by the educational and research organizations it covers.
Commercial uses outside those permissions require a separate license from the
copyright holders.

For commercial licensing enquiries, contact [Jai Bhagat](mailto:jkbhagatio@gmail.com).

Third-party software and datasets retain their respective licenses and terms.
In particular, the bundled LangevinFlow code retains its
[own license](https://github.com/jkbhagatio/nldisco/blob/main/experiments/churchland/environments/LangevinFlow_CCN/LICENSE).
