Metadata-Version: 2.4
Name: freegaussianizer
Version: 0.1.0
Summary: Matricial Free Energy loss for training Gaussianizing autoencoders.
Author-email: Thirulok Sundar <thiruloksundar278@gmail.com>
License: MIT
Project-URL: Homepage, https://github.com/Thiruloksundar/FreeGaussianizer
Project-URL: Repository, https://github.com/Thiruloksundar/FreeGaussianizer
Keywords: gaussianization,autoencoder,free-energy,random-matrix-theory,pytorch
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy>=1.21
Requires-Dist: scipy>=1.7
Requires-Dist: torch>=1.12
Provides-Extra: examples
Requires-Dist: matplotlib>=3.4; extra == "examples"
Provides-Extra: test
Requires-Dist: pytest>=7.0; extra == "test"
Dynamic: license-file

# FreeGaussianizer

A Python package implementing the **Matricial Free Energy** loss function for training Gaussianizing autoencoders, based on the paper:

> Sonthalia & Nadakuditi, *"Matricial Free Energy as a Gaussianizing Regularizer: Enhancing Autoencoders for Gaussian Code Generation"*, [arXiv:2510.17120](https://arxiv.org/abs/2510.17120)

The key idea: minimizing the Free Loss over mini-batches of encoder outputs drives the latent code distribution toward a Gaussian — without any explicit generative model.

This is the Python (PyTorch) port of [FreeGaussianizer.jl](https://github.com/Thiruloksundar/FreeGaussianizer.jl). The loss is implemented from the singular values of the code matrix (verified against the paper's reference value of −34.69; see [Testing](#testing)).

## Installation

```bash
pip install freegaussianizer
```

Or from source:

```bash
git clone https://github.com/Thiruloksundar/FreeGaussianizer
cd FreeGaussianizer
pip install -e .
```

## Quick start

```python
import torch
from freegaussianizer import (build_encoder, generate_data, train,
                              ks_statistic, delta_ot, rel_err_free_loss)

# 1. Generate synthetic data (Section 3 of the paper)
X_train, _ = generate_data(1280, p=2, mu=[5.0, 5.0], seed=42)
X_test,  _ = generate_data(1280, p=2, mu=[5.0, 5.0], seed=123)

# 2. Build the paper's MLP encoder (p=2 input -> d=32 embedding)
model = build_encoder(2, 32)

# 3. Train with mini-batched Adam, minimizing the Free Loss
result = train(model, X_train, epochs=2000, lr=1e-2, batch_size=256, X_test=X_test)

# 4. Evaluate Gaussianity of the learned embeddings
with torch.no_grad():
    Y = model(torch.as_tensor(X_train, dtype=torch.float64))  # (n x 32) code matrix
print("KS statistic :", ks_statistic(Y))
print("Delta_OT     :", delta_ot(Y))
print("Rel. err.    :", rel_err_free_loss(Y))
```

## Conventions

The code matrix `Y` is **batch-first**: shape `(b, d)` = (batch size × embedding dimension) with `d < b`. This matches PyTorch's standard `(batch, features)` layout, so the encoder is a plain `nn.Sequential` — no transposes needed. (The Julia package uses the transposed `d × b` convention; the mathematics is identical.)

## API

### Loss

| Function | Description |
|---|---|
| `free_loss(Y)` | Free Loss L_free(Ỹ) = −Φ̄_c(Ỹ) — minimize this to Gaussianize (Eq. 12) |
| `matricial_free_energy(Y)` | Matricial free energy Φ̄_c(Ỹ) (Eq. 9) |

`Y` is a `(b, d)` `torch.Tensor` with `d < b`. Both functions are autograd-differentiable for use inside a training loop.

### Data

| Function | Description |
|---|---|
| `generate_data(n_half, p, mu, seed)` | Generate balanced two-class χ²₁ data (Section 3) |

### Model

| Function | Description |
|---|---|
| `build_encoder(p, d, hidden=32)` | 5-layer MLP encoder matching the paper's architecture |

### Training

| Function | Description |
|---|---|
| `train(model, X_train, ...)` | Mini-batched Adam training loop; returns `TrainResult(train_losses, test_losses, snapshots)` |

### Evaluation metrics

| Function | Description |
|---|---|
| `ks_statistic(entries)` | Kolmogorov–Smirnov statistic vs N(0,1) |
| `delta_ot(Y, n_ref)` | Relative excess optimal transport cost Δ_OT (Eq. 13) |
| `rel_err_free_loss(Y, n_ref)` | Relative deviation of free loss from Gaussian reference |
| `ot_cost(A, B)` | Raw optimal transport cost between two matrices |

### Theoretical references

| Function | Description |
|---|---|
| `theoretical_min_free_loss(d, b, ...)` | Expected minimum free loss for a Gaussian `(b, d)` matrix |
| `marchenko_pastur_pdf(x, c)` | Marčenko–Pastur density at aspect ratio c = d/b |

## Examples

The [examples/](examples/) directory reproduces figures from the paper:

| Script | Reproduces |
|---|---|
| `free_loss_verification.py` | The three loss-verification checks (paper value, σ / mean minimizers) |
| `figure1.py` | Training data scatter + loss curves (Fig. 1) |
| `figure6_7.py` | Histogram / Q–Q / eigenvalue snapshots across training (Figs. 6–7) |
| `figure8.py` | Ridgeline of code-entry distributions over training (Fig. 8) |
| `figure9_10.py` | KS / Δ_OT / RelErr scaling experiments (Figs. 9–10) |

Run any example from the repo root (requires `pip install -e ".[examples]"`):

```bash
python examples/figure1.py
```

## Testing

Install the test extras and run the suite with `pytest`:

```bash
pip install -e ".[test]"
pytest -v
```

The test suite covers:

| Test file | What it checks |
|---|---|
| `test_exports.py` | All public functions are exported and accessible |
| `test_loss.py` | `free_loss` / `matricial_free_energy` values, autograd, and the paper's three verification checks (−34.69, σ = 1, mean = 0) |
| `test_data.py` | `generate_data` shape, label balance, reproducibility |
| `test_model.py` | `build_encoder` architecture, forward-pass shape, float64 weights |
| `test_metrics.py` | `ks_statistic`, `ot_cost`, `delta_ot` correctness |
| `test_theoretical.py` | `theoretical_min_free_loss` sign, `marchenko_pastur_pdf` integral |

## Citation

If you use this package, please cite the original paper:

```bibtex
@misc{sonthalia2025matricial,
  title   = {Matricial Free Energy as a Gaussianizing Regularizer:
             Enhancing Autoencoders for Gaussian Code Generation},
  author  = {Sonthalia, Rishi and Nadakuditi, Raj Rao},
  year    = {2025},
  eprint  = {2510.17120},
  archivePrefix = {arXiv},
}
```

## License

MIT — see [LICENSE](LICENSE).
