Initial commit

Signed-off-by: 周唤海 <albus.zhouhh@gmail.com>

Author: fliingelephant

.github/workflows/ci.yml

@@ -0,0 +1,25 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+
+      - name: Set up Python
+        run: uv python install 3.12
+
+      - name: Install dependencies
+        run: uv sync
+
+      - name: Run tests
+        run: uv run pytest -v

.gitignore

@@ -0,0 +1,6 @@
+.DS_Store
+
+data/*.npz
+
+*egg-info/
+*.pyc

AGENTS.md

@@ -0,0 +1,26 @@
+# Agent Guidelines
+
+Follow [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html), with JAX patterns in this document taking precedence.
+
+**Virtual environment**: Use `uv` for environment management. Run commands with `uv run`.
+
+## Design Decisions
+
+- **Match the paper.** Implementation should follow equations in `notes/Whitelam - 2026 - Generative Thermodynamic Computing.pdf`.
+- **DRY (Don't Repeat Yourself).** Consolidate duplicated implementations into a single source of truth.
+- **Let it crash.** Avoid defensive parameter checks; assume correct wiring and let errors surface.
+- **Uncertain correctness.** For uncertain behavior, refer to `notes/` or ask the user.
+- **Julia-style defaults.** Put defaults in function signature: `def foo(x, y=10):` not `def foo(x, y=None): y = y or 10`.
+- **No unnecessary intermediate variables.** Return directly: `return expr` not `result = expr; return result`.
+- **No unused imports/variables.** Remove any defined but unreferenced code.
+
+## JAX Patterns
+
+- **Use `jax.lax.scan`** over Python loops.
+- **Use `jax.vmap`** for batch operations (e.g., multiple denoising trajectories).
+- **No `jax.block_until_ready` in hot paths.** It breaks XLA fusion.
+- **Configurable dtype** via `config.dtype_name` (float32 or float64).
+
+## Logging
+
+- Use Python `logging` module, not `print()`.

README.md

@@ -0,0 +1,129 @@
+A JAX replication of generative thermodynamic computing[^1][^2], which uses Langevin dynamics for visual generation. This repository provides MNIST digit synthesis as a minimal working example.
+
+## How it works
+
+Standard diffusion models use neural networks for denoising. Here, denoising is done by Langevin dynamics of a physical system with trained couplings — no neural network at inference time.
+
+Training maximizes the probability that the system generates the reverse of noising trajectories, which is equivalent to minimizing heat dissipation. In hardware, this could be $\gt 10^{10}\times$ more efficient than digital computation.
+
+**Noising** (image → noise):
+
+![Noising trajectory](outputs/demo/fig1a_noising.png)
+
+**Denoising & generation** (noise → image):
+
+![Denoising and samples](outputs/demo/fig2.png)
+
+## Installation
+
+This project uses [uv](https://github.com/astral-sh/uv) for environment management.
+
+```bash
+uv sync
+```
+
+Verify the JAX backend:
+
+```bash
+uv run python -c "import jax; print(jax.__version__); print(jax.default_backend())"
+```
+
+## Data
+
+Download MNIST:
+
+```bash
+uv run download-mnist            # saves to data/mnist.npz
+uv run download-mnist --out path/to/mnist.npz  # custom path
+```
+
+## Quick Start
+
+Train a model and generate figures:
+
+```bash
+uv run whitelam-2026 --mnist data/mnist.npz --out outputs/demo
+```
+
+By default, the model trains on digits **(0, 1, 2)** with **512 hidden units**, matching the paper's setup.
+
+## Configuration
+
+Override any config field with `--set key=value` (repeatable). Values are parsed as JSON.
+
+```bash
+# More training trajectories for better samples
+uv run whitelam-2026 --mnist data/mnist.npz --out outputs/demo \
+  --set n_training_trajectories=1000
+
+# Set random seed for reproducibility (default is 42)
+uv run whitelam-2026 --mnist data/mnist.npz --out outputs/demo --seed 123
+
+# Higher DPI for publication-quality figures
+uv run whitelam-2026 --mnist data/mnist.npz --out outputs/demo --set fig_dpi=600
+```
+
+### Scaling to All 10 Digits
+
+The default model (512 hidden units) works well for 3 digit classes. To train on all digits (0-9), increase model capacity proportionally:
+
+```bash
+uv run whitelam-2026 --mnist data/mnist.npz --out outputs/full_mnist \
+  --set n_h=2048 \
+  --set 'train_digits=[0,1,2,3,4,5,6,7,8,9]' \
+  --set n_training_trajectories=1000
+```
+
+**Scaling considerations:**
+- Training time scales with `n_h²` (hidden-hidden couplings)
+- Memory scales with `n_v × n_h` (visible-hidden couplings)
+
+### Re-rendering Figures
+
+To regenerate figures from saved parameters without retraining:
+
+```bash
+uv run whitelam-2026 \
+  --mnist data/mnist.npz \
+  --params outputs/demo/params_learned.npz \
+  --out outputs/demo_rerender
+```
+
+## Outputs
+
+Each run produces:
+
+| File | Description |
+|------|-------------|
+| `fig1.png`, `fig2.png` | Composite figures |
+| `fig1a_noising.png` | Noising trajectory |
+| `fig1b_training_set.png` | Training digits |
+| `fig2a_denoising_trajectories.png` | Denoising trajectories |
+| `fig2b_samples.png` | Generated samples |
+| `fig2c_receptive_fields.png` | Learned hidden unit couplings |
+| `params_learned.npz` | Trained model parameters |
+| `config.json` | Configuration used |
+| `metrics.json` | Heat dissipation metrics |
+
+## Development
+
+Run tests:
+
+```bash
+uv run pytest -v
+```
+
+Use as a library:
+
+```python
+from generative_langevin.whitelam_2026.config import Whitelam2026Config
+from generative_langevin.whitelam_2026.model import init_params
+from generative_langevin.whitelam_2026.train import train_many_noising_trajectories
+from generative_langevin.whitelam_2026.sample import run_denoising_trajectory
+```
+
+## Reference
+
+[^1]: S. Whitelam. **Generative Thermodynamic Computing**. *Physical Review Letters* 136(3):037101, 2026. https://doi.org/10.1103/kwyy-1xln
+
+[^2]: https://github.com/swhitelam/generative_thermodynamic_computing