# DIMA — Diffusion–Isocoder Manifold-Autoencoder

**DIMA** combines three components into a practical, scalable manifold autoencoder:

- **DMAP**: Diffusion Maps encoder (builds a kNN graph in ambient space and embeds points into diffusion coordinates).
- **GPLM**: Nyström / inducing-point kernel ridge decoder (maps latent diffusion coordinates back to ambient space).
- **DDPM**: Latent diffusion model (optional) that learns a generative prior over *normalized* latents and can “refine” latents.

> Philosophy: keep **DMAP + GPLM** fast on CPU/RAM (NumPy/SciPy, sparse ops), and optionally run **DDPM** on JAX (CPU or GPU).

---

## Installation

### CPU-only (recommended to start)
```bash
pip install dima
```

### Optional extras

* Faster kNN search (FAISS):

```bash
pip install dima[faiss]
```

* Hugging Face upload/load:

```bash
pip install dima[hf]
```

> **GPU note (JAX):** `pip install dima` installs CPU `jaxlib` by default.
> For GPU, install the correct JAX wheel for your CUDA/ROCm setup first (per JAX docs), then install `dima`.

---

## Quickstart

```python
import numpy as np
import jax
from dima import DIMA

R_iX = np.random.randn(5000, 16).astype(np.float32)

dima = DIMA(R_iX)                 # trains DMAP, GPLM, DDPM with defaults

Z = dima.encode(R_iX[:10])        # ambient -> normalized latent (jnp)
X_hat = dima.decode(Z, refine=False)  # latent -> ambient (np), no DDPM
X_ref = dima.decode(Z, refine=True, t_start=10)  # DDPM refinement

X_gen = dima.sample(1000)         # unconditional samples -> ambient (np)
```

---

## What DIMA trains

### 1) DMAP encoder (Diffusion Maps)

DMAP builds a kNN graph over ambient data $R_{iX}\in\mathbb{R}^{N\times D}$ and returns diffusion coordinates
$R_{ix}\in\mathbb{R}^{N\times d}$.

Key idea (no theory): DMAP produces a **geometry-aware** latent space where nearby points on the manifold remain nearby
in diffusion distance.

**Main knobs (DMAP):**

* `d` *(int)*: latent dimension.
* `k` *(int)*: neighbors in the kNN graph. If `None`, a heuristic is used.
* `beta` / `β` *(float)*: kernel sharpness for the RBF affinity
  $K_{ij}=\exp{-\beta |x_i-x_j|^2/\varepsilon}$.
* `eps` / `ε` *(float or None)*: kernel bandwidth. If `None`, estimated from kNN distances (median heuristic).
* `alpha` / `α` *(float)*: density normalization exponent. Common values: `0.0` (none) or `1.0` (often robust).
* `t` *(float)*: diffusion time exponent (scales eigenvalues as $\lambda^t$). Often `0.5` or `1.0`.
* `drop_trivial` *(bool)*: drops the top eigenvector/eigenvalue (the constant mode).
* `sym` *(str)*: symmetrization mode for sparse kNN kernel graph: `"max"` or `"mean"`.
* `ann_backend` *(str)*: `"auto" | "faiss" | "pynndescent" | "sklearn" | "brute"`.

**Typical DMAP presets**

* Fast-ish and stable: `k=128..512`, `alpha=1.0`, `t=0.5`, `drop_trivial=True`.
* If your data is very noisy, try larger `k` and/or larger `eps_mul`.

---

### 2) GPLM decoder (Nyström kernel ridge / inducing GP)

GPLM learns a mapping from latents back to ambient:

* Inputs: latent training points $R_{ix}\in\mathbb{R}^{N\times d}$
* Targets: ambient training points $R_{iX}\in\mathbb{R}^{N\times D}$

Instead of a full $N\times N$ kernel solve, GPLM uses **inducing points** $Z_{mx}$ with $m\ll N$ and solves a reduced system:

* Build affinities $C_{im}=\exp{-\beta|R_{ix}-Z_{mx}|^2/\varepsilon}$
* Solve a stabilized kernel ridge / GP mean system to obtain weights $M_{mX}$
* Predict: $\hat R_{aX}=C_{am}M_{mX}+\mu_X$

**Main knobs (GPLM):**

* `m` *(int)*: number of inducing points. Bigger → better accuracy, more compute.
* `inducing` *(str)*: inducing strategy:

  * `"kmeans_medoids"` (default): kmeans centers snapped to nearest training latent.
  * `"fps"`: farthest-point sampling (space-filling).
  * `"random_subset"`: fastest.
  * `"given"`: use provided `Z_mx`.
* `sigma2` / `σ2` *(float)*: ridge regularization. Too small can overfit / cause instability; too large blurs reconstructions.
* `jitter` *(float)*: tiny diagonal stabilizer for Cholesky.
* `eps` / `ε` *(float or None)*: RBF bandwidth in latent space. If `None`, estimated from latent kNN distances.
* `k_eps` / `κ_eps` *(int)*: neighbors used for the $\varepsilon$ heuristic.
* `pred_k` / `pred_κ` *(int or None)*: prediction-time inducing neighbors:

  * `None` means use all inducing points (best accuracy).
  * a small number (e.g. `128` or `256`) speeds inference (slightly lower accuracy).
* `whiten_latent` *(bool)*: optionally standardize latent dimensions before kernel computation.
* `center_X` *(bool)*: subtract and re-add ambient mean (usually helpful).
* `fit_block` *(int)*: block size for streaming $C^T C$ accumulation (memory/perf knob).
* `ann_backend` *(str)*: same options as DMAP.

**Typical GPLM presets**

* Accurate: `m=1024..4096`, `pred_k=None`, `sigma2=1e-5` (tune).
* Faster inference: set `pred_k=128..512`.

---

### 3) DDPM latent diffusion (optional generative prior)

DDPM learns a distribution over **normalized** latents:
[
Z = \frac{R_x - \mu}{\sigma}
]
and can:

* **sample** new latents,
* **refine** a given latent by projecting it onto the learned latent manifold/prior.

In DIMA, DDPM operates purely in latent space (dimension `d`), so it’s lightweight compared to image DDPMs.

**Main knobs (DDPM):**

* `T` *(int)*: number of diffusion steps. Common: `100..1000`. DIMA default: `200`.
* `hidden_dim` *(int)*: MLP width.
* `t_embed_dim` *(int)*: time embedding dimension.
* `n_iter` *(int)*: training iterations (more is better for sampling quality).
* `batch_size` *(int)*: training batch size.
* `learning_rate` *(float)*: Adam learning rate.
* `ema_decay` *(float)*: EMA smoothing for stable sampling (typical: `0.999`).
* `beta_max` *(float)*: caps noise schedule (stability knob).
* `eps` *(float)*: numerical stabilizer.
* `verbose_every` *(int)*: progress prints.

**Refinement knobs (during decoding):**

* `refine` *(bool)*: enable/disable DDPM refinement.
* `t_start` *(int)*: how strongly to “project” using reverse diffusion.

  * small (`1..10`) = gentle projection
  * larger (`20..100`) = stronger projection (can oversmooth or drift if DDPM undertrained)
* `add_noise` *(bool)*: whether to forward-noise before reverse steps.

---

## API

### Core class

```python
from dima import DIMA
dima = DIMA(R_iX)
```

### Encode / decode

```python
Z = dima.encode(X)                      # (B,d) normalized latent (jnp)
X_hat = dima.decode(Z, refine=False)    # (B,D) reconstruction (np)
X_ref = dima.decode(Z, refine=True, t_start=10)  # refined decode
```

### Polymorphic call

```python
Z = dima(X)     # if X.shape[-1] == D
X = dima(Z)     # if Z.shape[-1] == d
```

### Sampling

```python
X_gen = dima.sample(1000)            # ambient samples (np)
Z_gen = dima.sample(1000, decode=False)  # latent samples (jnp)
```

---

## Configuration patterns

### 1) Use defaults (cleanest)

```python
dima = DIMA(R_iX)
```

### 2) Pass only a few knobs

```python
dima = DIMA(
    R_iX,
    d=64,
    dmap_kwargs=dict(k=256, alpha=1.0, t=0.5),
    gplm_kwargs=dict(m=2048, sigma2=1e-5, pred_k=256),
    ddpm_kwargs=dict(n_iter=50_000, hidden_dim=256),
)
```

---

## Saving / Loading

### Local

```python
dima.save_local("dima.msgpack", "config.json")
dima2 = DIMA.load_local("dima.msgpack", ddpm_device="auto")
```

### Hugging Face (optional)

```python
dima.upload_to_huggingface(repo_id="username/dima-model", hf_token="...")

dima3 = DIMA.load_from_huggingface("username/dima-model", ddpm_device="auto")
```

---

## Performance notes

* **DMAP**: main costs are kNN search + sparse eigensolve.

  * kNN is faster with FAISS.
  * Increasing `k` increases graph density and compute.
* **GPLM**: training cost roughly scales with $N\cdot m$ (streamed by `fit_block`).

  * Inference cost is $B\cdot m$ if `pred_k=None`, or $B\cdot \text{pred_k}$ if using inducing kNN.
* **DDPM**: scales with latent dimension `d`, steps `T`, and training iterations.

  * GPU helps but CPU works for smaller `d` and fewer iterations.

---

## Troubleshooting

### “Unknown backend: gpu”

Your JAX install only has CPU. Use:

```python
dima = DIMA(R_iX, ddpm_device="cpu")
```

or install a GPU-enabled JAX build.

### Reconstructions are blurry / low quality

* Increase `gplm_kwargs["m"]`
* Decrease `gplm_kwargs["sigma2"]` slightly (careful: too small can destabilize)
* Set `pred_k=None` (use all inducing points)

### DDPM refinement makes reconstructions worse

* Reduce `t_start` (try `3..10`)
* Increase DDPM training (`ddpm_kwargs["n_iter"]`)
* Disable `add_noise` for gentler behavior

---

## Citations

* Diffusion Maps / DMAE inspiration: *Diffusion Map AutoEncoder (DMAE)*.
* Latent diffusion prior: *Denoising Diffusion Probabilistic Models (DDPM)*.

---

## License

MIT