Metadata-Version: 2.4
Name: egnlib
Version: 0.2.1
Summary: Equivariant Geodesic Networks: a universal SPD-manifold classifier for PyTorch
Author-email: Md Raihan Khan <kraihan918@gmail.com>
Maintainer-email: Md Raihan Khan <kraihan918@gmail.com>
License-Expression: MIT
Project-URL: Homepage, https://github.com/kraihan/egn
Project-URL: Repository, https://github.com/kraihan/egn
Project-URL: Issues, https://github.com/kraihan/egn/issues
Project-URL: Changelog, https://github.com/kraihan/egn/blob/main/CHANGELOG.md
Keywords: riemannian-geometry,spd-manifold,covariance,eeg,biosignal,deep-learning,pytorch,classification
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Mathematics
Classifier: Typing :: Typed
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0
Requires-Dist: numpy>=1.21
Provides-Extra: dev
Requires-Dist: pytest>=7; extra == "dev"
Requires-Dist: pytest-cov>=4; extra == "dev"
Requires-Dist: ruff>=0.4; extra == "dev"
Requires-Dist: build>=1.0; extra == "dev"
Requires-Dist: twine>=5.0; extra == "dev"
Provides-Extra: sklearn
Requires-Dist: scikit-learn>=1.0; extra == "sklearn"
Provides-Extra: docs
Requires-Dist: mkdocs>=1.5; extra == "docs"
Requires-Dist: mkdocs-material>=9.0; extra == "docs"
Requires-Dist: mkdocstrings[python]>=0.24; extra == "docs"
Dynamic: license-file

# EGN — Equivariant Geodesic Networks

A universal classifier on the SPD manifold, packaged the way a convolutional network is.

```python
import numpy as np
from egn import EGNClassifier

X = np.random.randn(512, 22, 256)      # 512 trials, 22 channels, 256 samples
y = np.random.randint(0, 4, 512)

clf = EGNClassifier(epochs=30).fit(X, y)
print(clf.score(X, y))
```

That is the whole API for the common case. The matrix size, the number of manifold
channels, the class count and the label vocabulary are all inferred from the data, so the
same object trains on EEG epochs, skeleton covariances, radar returns and image
descriptors without a per-dataset subclass.

```bash
pip install egn        # torch and numpy are the only dependencies
```

---

## What it does

Every intermediate representation is an exact symmetric positive definite matrix under the
affine-invariant Riemannian metric. There is no projection step in the forward pass, and
constrained parameters are updated by retraction, so a Stiefel frame stays orthonormal and
an SPD prototype stays positive definite regardless of the step size.

```
x ──ToSPD──► (B, K, d, d) ──EGNBlock ×L──► (B, K′, m, m) ──pool──► (B, m, m) ──head──► logits
```

The correspondence with a CNN is deliberate, and the intuitions transfer:

| CNN | EGN |
|---|---|
| `Conv2d(c_in, c_out, k)` | `BiMap(d_in, d_out, c_in, c_out)` |
| channels | manifold channels (branches) |
| stride / downsampling | matrix size reduction `d_in → d_out` |
| `BatchNorm2d` | `SPDBatchNorm` (whitens by a Fréchet mean) |
| `ReLU` | `SpectralActivation` (acts on eigenvalues) |
| `Dropout` | `GeodesicDropout` (interpolates towards the identity along a geodesic) |
| bias | `GeometricBias` (a metric isometry, not an addition) |
| global average pool | `RiemannianPool` (Fréchet or log-Euclidean barycentre) |
| linear classifier | `TangentHead` or `GeodesicPrototypeHead` |

---

## Input conventions

`ToSPD` accepts, and infers when `input_kind="auto"`:

| input | meaning | output |
|---|---|---|
| `(B, n, n)` | covariance / descriptor already on the manifold | `(B, 1, n, n)` |
| `(B, K, n, n)` | multi-branch SPD | `(B, K, n, n)` |
| `(B, C, T)` | multichannel signal | `(B, K, C, C)` |
| `(B, T, D)` | sequence of features (`input_kind="sequence"`) | `(B, K, D, D)` |
| `(B, C, H, W)` | feature map / image | `(B, K, C, C)` |

`branches=K` splits the sample axis into `K` windows and forms one covariance per window —
the manifold equivalent of a multi-channel stem.

Rank-2 input is rejected with an explanatory error rather than being turned into a singular
rank-one outer product, which would only fail later inside a logarithm.

---

## Why the GPU is fast now

The earlier research code ran the whole network in `float64` and was slower on GPU than on
CPU. Six things caused that, and all six are addressed:

1. **`float64` everywhere.** Consumer and inference-class GPUs execute double precision at
   1/32 of their `float32` rate; a T4 in `fp64` is genuinely slower than a decent CPU. The
   dtype policy is now `config.spectral_dtype = "auto"` — `float32` on CUDA, `float64` on
   CPU. Set it to `"float64"` only to reproduce a theory check.
2. **Host synchronisation inside the forward pass.** The old Fréchet mean called `.item()`
   on a residual every iteration and the spectral activation called `.item()` on its
   threshold, draining the CUDA queue several times per layer. Nothing in the forward pass
   calls `.item()` any more; the mean runs a fixed iteration budget.
3. **Repeated eigendecompositions.** `sqrtm_pair` returns `S^{1/2}` and `S^{-1/2}` from one
   decomposition; the prototype head whitens by each prototype once per forward instead of
   once per (sample, prototype) pair; distances use `eigvalsh`, which never forms
   eigenvectors.
4. **Batching.** Every operator takes an arbitrary leading shape and issues exactly one
   `eigh` per call. `config.eig_chunk` caps the batch when memory, not throughput, is the
   constraint.
5. **Per-step metric readback.** Training statistics accumulate on the device and are read
   once per epoch.
6. **Pooling cost.** The default pool is the closed-form log-Euclidean barycentre; the
   iterative Fréchet mean is one flag away (`pool="frechet"`) when the channel spread makes
   it worth the iterations.

Measure your own machine rather than trusting any of this:

```bash
python -m egn.benchmark --sizes 16 32 64 --batch 256 1024 4096
```

The output separates `eigh` throughput per dtype from end-to-end model throughput, which is
what tells you whether you are precision bound or launch-latency bound. If throughput is
flat in the batch size, the kernels are too small to saturate the device — raise the batch
before anything else.

**The honest caveat.** These are small-matrix, decomposition-heavy workloads. A GPU wins
decisively at large batch sizes and `float32`; at batch 32 with 8×8 matrices it may still
lose to a CPU, because the kernels never fill the device. That is a property of the
operation, not of this implementation.

---

## Scaling out

Distributed data parallel is a launch flag, not a rewrite:

```bash
torchrun --nproc_per_node=4 examples/train_ddp.py
```

`EGNClassifier` detects the process group, wraps the model in `DistributedDataParallel`,
installs a `DistributedSampler` and reduces metrics across ranks. The model is materialised
before wrapping, so the replicas have parameters to broadcast.

DDP works with the Riemannian optimiser without special handling: DDP all-reduces the
Euclidean gradients in its backward hook, and every rank then applies the same deterministic
retraction to the same parameters, so replicas stay identical.

`data_parallel=True` uses `nn.DataParallel` for a quick single-process multi-GPU run. It is
not recommended — it re-scatters the model every step, which on a network of short kernel
launches costs more than it saves.

---

## Building your own architecture

```python
import torch.nn as nn
from egn.nn import EGNBlock, GeodesicPrototypeHead, RiemannianPool, ToSPD

model = nn.Sequential(
    ToSPD(kind="signal", branches=4),
    EGNBlock(22, 16, in_channels=4, out_channels=8, mix=True),
    EGNBlock(16, 8,  in_channels=8, out_channels=8),
    RiemannianPool("frechet", iters=5),
    GeodesicPrototypeHead(8, num_classes=4),
)
```

Or use the factories, which follow the `torchvision` convention:

```python
from egn import egn_tiny, egn_small, egn_base
model = egn_base(num_classes=4)      # deeper trunk, mixed channels, geodesic head
```

### Choosing a head

`TangentHead` (default) takes one logarithm per sample and applies a linear classifier in
the tangent space at a learnable reference point. Its cost is independent of the class
count.

`GeodesicPrototypeHead` scores by squared geodesic distance to trainable SPD prototypes,
`p(c | Σ) = softmax(−d²(Σ, P_c)/τ)`. Fully geometric, and its cost grows with the number of
prototypes. Use it when you want prototypes you can inspect, or when classes are naturally
described as regions of the manifold.

Both are invariant in the *joint* sense: congruencing the input **and** the reference by the
same matrix leaves the logits unchanged. Invariance while holding the prototypes fixed is
false — congruencing only the input changes every distance.

---

## Geometry as a public API

`egn.geometry` and `egn.functional` are usable on their own, with no model involved:

```python
from egn.geometry import distance, frechet_mean, geodesic, riemannian_log

d = distance(A, B)                      # affine-invariant geodesic distance
M = frechet_mean(batch, dim=1)          # Riemannian barycentre
mid = geodesic(A, B, 0.5)               # midpoint on the manifold
```

Everything is differentiable, including through repeated eigenvalues: the backward pass uses
the Loewner divided-difference matrix, and clamped eigenvalues receive a zero subgradient
instead of an unbounded one.

---

## Tests

```bash
pytest -q
```

The suite asserts the geometry, not just the shapes: exp/log inverse to machine precision,
affine invariance of the distance, geodesic constant speed, the Fréchet mean as a fixed
point and its equivariance, `gradcheck` on every spectral operator including an exact
eigenvalue tie, and the closed-form prototype gradient against autograd. It also asserts the
counter-examples — that a convex combination is *not* a geodesic, and that a rectangular
`BiMap` is not output-congruent — because those are the claims that are easy to overstate.

---

## Citation

```bibtex
@unpublished{khan2026egn,
  author = {Md Raihan Khan and Airin Akter Tania},
  title  = {Equivariant Geodesic Networks: End-to-End Classification on the SPD Manifold},
  note   = {Under review at AAAI},
  year   = {2026},
  url    = {https://github.com/kraihan/egn}
}
```

## License

MIT. This package contains no third-party research code; see `MIGRATION.md` if you are
coming from the original research repository.
