pyRBM/CLAUDE.md

# CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

## What this is

pyRBM is a from-scratch, hand-derived implementation of Restricted Boltzmann Machines (RBMs) and
Recurrent Temporal RBMs (RTRBMs) — no autodiff, no PyTorch training loop. Every gradient (Contrastive
Divergence) is coded by hand in `src/rbm/train.py`. It's a research/educational codebase, not a
competitive ML framework — see `README.rbm.md` for a full Q&A on the design rationale, why a
Transformer can't be bolted on without adding autodiff, and how the `src/compat/torch.py` bridge
relates to that.

## Environment

- Python >=3.12, project uses a `.venv` at the repo root (`.venv/bin/python`, `.venv/bin/pytest`).
- Package is installed editable (`pip install -e .`), src-layout: `src/{rbm,stack,model,label,image,compat}`
  import as top-level packages (`from rbm.entity import Entity`, `from stack.rnn import StackRnn`, etc.)
  — there is no `pyrbm.` prefix.
- GPU (CUDA) is used by default — see "CPU vs GPU" below.
- No linter/formatter config exists (no ruff/black/flake8/mypy) and no pytest config file — match
  existing style (tabs for indentation in most `src/rbm` and `src/stack` files, spaces elsewhere;
  check the file you're editing).

## Running tests

```bash
.venv/bin/pytest tests/                     # whole suite
.venv/bin/pytest tests/test_rnn.py -q       # single file
.venv/bin/pytest tests/test_rnn.py::test_rnn_shared -q      # single test
```

`tests/README.md` documents each test's purpose and is worth reading before adding new tests — it
also has a fix-log at the bottom explaining several non-obvious bugs already found and fixed in
`train.py`/`entity.py` (e.g. why L1/L2 regularization must NOT be scaled by `1/mini_batch_size`, why
CD-k>1 Gibbs loops must sample rather than use means). Some `tests/*.py` files are pytest suites
(`def test_...`), others (`test_faces_sub_image.py`, `test_norbs.py`, `test_rbm.py`, etc.) are
standalone scripts invoked directly (`python tests/test_rbm.py <project_name>`) that need external
data/images not in the repo (see the table in `tests/README.md`).

## CPU vs GPU (important gotcha)

`src/rbm/matrix.py` is the single backend switch for the entire codebase:

```python
USE_CUDA = 1
if USE_CUDA:
    import cupy as np
else:
    import numpy as np
```

Every module imports the array module as `np` from here (`from .matrix import np`), so flipping
`USE_CUDA` to `0` switches the whole codebase to CPU-only numpy with no other code changes. This
machine has cupy-cuda12x + an RTX 4070 Ti, so GPU is the default/expected path. If you see
`cupy`-related errors on a machine without a GPU, this flag is the first place to check. Note that
`cupy-cuda12x` is a de-facto hard runtime dependency (nothing works with `USE_CUDA=1` without it) but
is *not* listed in `pyproject.toml`'s `dependencies` — it's only present because it was installed
manually into this `.venv`. `torch` is also installed but only backs `src/compat/torch.py` /
notebook use, not core training.

## Architecture

### Core primitives (`src/rbm/`)

- **`matrix.py`** — numpy/cupy backend switch (see above) + shared math helpers (`sample`, `prob`
  (sigmoid), `gaussian`, `rms_error`, `read_armadillo` for loading legacy Armadillo `.dat` files).
- **`state.py`** (`RbmState`) — holds `w_hv`, `b_v`, `b_h`; pure linear algebra (`v_to_h`, `h_to_v`)
  plus `.npz` save/load.
- **`entity.py`** (`Entity`) — one RBM unit. `Entity.Type` is one of `BB_RBM`/`BG_RBM`/`GB_RBM`/`GG_RBM`
  (binary/Gaussian visible × binary/Gaussian hidden), derived from `EntityParams`. Owns a `RbmState`
  (weights) and a `grad` (momentum accumulator). `forward`/`reconstruct` do (optionally multi-step)
  Gibbs sampling; `h_given_v`/`v_given_h` are the raw (unsampled) linear step. `grad_compute` /
  `state_adjust` apply momentum + L1/L2/weight-decay (regularization is applied at full strength,
  NOT scaled by mini-batch size — see the fix-log note above).
- **`train.py`** — the four CD variants (`cd_binary_binary`, `cd_gaussian_binary`,
  `cd_gaussian_gaussian`, `cd_binary_gaussian`), each a positive-phase/negative-phase Gibbs pass
  mirroring the classic Hinton Matlab RBM code (comments in the file reference the original Matlab
  pseudocode). `train()` is the standard epoch/mini-batch loop for a single `Entity`, dispatching to
  the right CD function via `entity.type` and reporting through a `Status`.
- **`status.py`** — pluggable progress reporting/checkpointing (`Status`, `CheckpointStatus`) driven
  by a callback (`on_report`) that can also halt training early.
- **`layer.py`** (`Layer`) — thin wrapper: one `Entity` + name + save/load, sized from
  `(numVisibleX, numVisibleY, numContext, numHidden)`.

### Composition patterns — two different ways to build multi-entity models

1. **`Stack` family (`src/stack/`)** — an ordered `list[Layer]` with save/load-by-index into
   `{name}-{i}-state.npz` files under `work_dir`.
   - `StackDeep` (`deep.py`) — classic layer-wise-pretrained deep belief net: `train()` trains layer 0,
     forwards its output as layer 1's input, etc. `pass_up`/`pass_down`/`pass_down_up` for
     encode/decode/round-trip through the stack.
   - `StackRnn` (`rnn.py`) — the RTRBM. Visible = `[context h_{t-1} | sensory x_t]`. Two modes based on
     how many layers you `append()`:
     - **Shared weights** (1 layer): the same `Entity` is reused every time step; sequences can be any
       length T.
     - **Unrolled** (N layers): `layers[t % N]` handles time step t, each with its own weights;
       training requires sequences of length exactly N (uses "temporal-shift padding" — see the
       docstring in `_train_unrolled`).
     Use `StackRnn.make_layer(...)` / `StackRnn.make_unrolled(...)` to construct layers, `.step(x)` /
     `.reconstruct(h)` for inference, `.train(sequences)` for training.
   - `factory.py` (`StackFactory`) — builds a `Stack` from a JSON project dict/file (legacy `.prj`
     format from a prior C++ implementation this was ported from — field names like `numVisibleX`,
     `doSampleBatch` reflect that origin).
   - `rnn_helper.py` — one-hot vocab encode/decode (`ch2idx`/`idx2ch`/`vec2idx`/`idx2vec`) and
     visible-vector `concat`/`split`/`shift_*` helpers used by the RNN scripts/notebooks. `VOCAB` here
     is a small ad-hoc set (currently space + digits) — the top-level scripts/notebooks each define
     their own, larger vocab inline rather than importing this one.

2. **`Model` (`src/model/model.py`)** — an alternate, reflection-based composition: subclass `Model`,
   assign `Entity` instances as plain attributes in `__init__`, implement `forward()`. Base class
   discovers entities via `self.objects(Entity)` (walks `__dict__`) for `train`/`save`/`load`/`init`,
   training each entity in sequence and feeding its output forward — used by `test_model.py`'s
   `DeepModel` and by `src/label/label.py`'s `Label.Fitter`.

### Other modules

- **`src/label/label.py`** (`Label`) — encodes discrete labels as binary or one-hot vectors, then fits
  a `Model` (an RBM) on the vectors so labels can be embedded/decoded through learned features.
- **`src/image/sub_image.py`** — patch-based image tiling: `SubImageExtract` cuts overlapping patches
  (with stride) for training on image patches instead of whole images; `SubImageCombine` blends
  overlapping reconstructed patches back into a full image (pixel-averaged); also a small `conv2d`.
- **`src/compat/torch.py`** — a thin PyTorch-`optim`-shaped wrapper (`LossFunction`, `Optimizer` with
  `zero_grad()`/`step()`) around the *same* hand-coded CD functions from `train.py` — it mimics the
  PyTorch API surface but is not autodiff; see `README.rbm.md` for why real backprop would require
  replacing the numpy/cupy backend with torch tensors throughout.

### Top-level scripts and notebooks

These are the char-level RTRBM demos/experiments, each self-contained and runnable directly
(`python Rnn.py`, no install needed — they `sys.path.insert(0, "src")` themselves):

- **`Rnn.py`** — minimal unrolled-mode demo matching `docs/Rnn.drawio.png` (sliding window + context,
  small hard-coded text).
- **`JayRnn.py`** / `README_JayRnn.md` — larger char-level language model, single shared-weight RBM,
  fixed vocabulary, checkpoint/resume via `results/JayRnn-*-state.npz`.
- **`moby_rnn.ipynb`** / `README_moby_rnn.md` — the most fully worked example: trains on
  `data/moby.txt`, unrolled `StackRnn` with `TEMPORAL_DEPTH` positions, includes evaluation
  (teacher-forced reconstruction vs. honest next-step prediction), clamped-Gibbs text generation at
  multiple temperatures, and hidden-state visualization.
- **`RNN-RBM.ipynb`** — the notebook the top-level `README.md` documents (context/sensory concat,
  `UNROLL_DEPTH` modes).
- **`mnist_test*.ipynb`, `cifar_test*.ipynb`, `layer_test.ipynb`** — image-domain (non-recurrent)
  RBM/DBN experiments using `data/MNIST` and `data/cifar-10-batches-py`.

Checkpoints from all of the above land in `results/{name}-{index}-state.npz`, loaded automatically on
the next run if present (`RbmState.load` silently no-ops on missing file/shape mismatch, so deleting a
checkpoint is the way to force retraining from scratch).