Metadata-Version: 2.4
Name: pyjae
Version: 0.2.2
Summary: Joint Autoencoders (JAE1 channel-split, JAE2 JEPA) for neural signal denoising and manifold learning
Author-email: Ege Altan <37878962+egealtan@users.noreply.github.com>
License: MIT
Project-URL: Homepage, https://github.com/egealtan/pyjae
Project-URL: Repository, https://github.com/egealtan/pyjae
Project-URL: Bug Tracker, https://github.com/egealtan/pyjae/issues
Keywords: neuroscience,neural signals,denoising,autoencoder,jepa,self-supervised-learning,manifold,pytorch
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0.0
Requires-Dist: numpy>=1.21.0
Requires-Dist: scikit-learn>=1.0.0
Requires-Dist: scipy>=1.7.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
Requires-Dist: jupyter>=1.0.0; extra == "dev"
Requires-Dist: matplotlib>=3.5.0; extra == "dev"
Requires-Dist: ruff>=0.1.0; extra == "dev"
Requires-Dist: build>=1.0.0; extra == "dev"
Requires-Dist: twine>=5.0.0; extra == "dev"
Dynamic: license-file

# `pyjae`: Joint Autoencoders (JAE) for Neural Signal Denoising

[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
[![PyTorch](https://img.shields.io/badge/PyTorch-2.0+-ee4c2c.svg)](https://pytorch.org/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)

`pyjae` (pronounced "pie-jay") is the official implementation of joint autoencoder
models for denoising high-dimensional neural population recordings and recovering
the low-dimensional manifold underneath, based on
[Altan et al. (2021)](https://journals.plos.org/ploscompbiol/article?id=10.1371/journal.pcbi.1008591).

## Installation

```bash
pip install pyjae
```

Or with [uv](https://docs.astral.sh/uv/):

```bash
uv add pyjae
```

For development from source:

```bash
git clone https://github.com/egealtan/pyjae.git
cd pyjae
uv sync            # installs the package with the dev extras
```

## Quick start

```python
from pyjae import JAE, simulate_neural_data

# Simulate a 6D nonlinear manifold observed on 64 noisy channels
clean, noisy, info = simulate_neural_data(
    n_samples=400, n_channels=64, n_timepoints=96,
    latent_dim=6, snr_db=10.0, nonlinear=True, alpha=3.0, seed=0,
)

# Train the channel-split model and denoise
model = JAE(latent_dim=6)          # backend="jae1" by default
model.fit(noisy, epochs=200)
denoised = model.denoise(noisy)
print("VAF:", model.score(clean, denoised))
```

Use the JEPA backend for representation learning plus a denoising readout:

```python
model = JAE(latent_dim=32, backend="jepa", patch_len=8, d_model=64)
model.fit(noisy, epochs=150)
denoised = model.denoise(noisy)
```

## Package layout

| Module | Purpose |
| --- | --- |
| `pyjae.api` | `JAE` facade (fit / denoise / score / save / load) over both backends |
| `pyjae.models` | `JAE1`, `JAE2`, and shared encoder building blocks |
| `pyjae.views` | Modular channel-split and JEPA-mask strategies |
| `pyjae.data` | Simulator (Altan et al. generative model) and evaluation controls |
| `pyjae.metrics` | Per-channel VAF plus a collapse-resistant latent-quality panel |
| `pyjae.baselines` | PCA, Factor Analysis, denoising autoencoder, Wiener oracle |
| `pyjae.eval` | Benchmark and evaluation harness |

## Requirements

- Python >= 3.10
- PyTorch >= 2.0, NumPy, scikit-learn, SciPy

## Citation

```bibtex
@article{altan2021jae,
  title={Estimating the dimensionality of the manifold underlying multi-electrode neural recordings},
  author={Altan, Ege and Solla, Sara A. and Miller, Lee E. and Perreault, Eric J.},
  journal={PLOS Computational Biology},
  year={2021},
  volume={17},
  number={11},
  pages={e1008591},
  doi={10.1371/journal.pcbi.1008591}
}
```

## License

MIT License. See [LICENSE](LICENSE).
