Add AGENTS.md and trim the README

AGENTS.md

@@ -0,0 +1,224 @@
+# AGENTS.md
+
+Guidance for AI coding agents working in this repository.
+
+## What this is
+
+STAMP (Solid Tumor Associative Modeling in Pathology) is an end-to-end, weakly
+supervised deep-learning pipeline for computational pathology: gigapixel
+whole-slide images (WSIs) in, biomarker predictions out. It ships as a single
+CLI (`stamp`) driven by one YAML config file, plus an optional MCP server.
+
+The pipeline stages, in order, map 1:1 onto CLI subcommands and onto top-level
+keys in the config file:
+
+```
+WSIs → preprocess → [encode_slides | encode_patients] → train / crossval
+                                                       → deploy → statistics
+                                                       → heatmaps
+```
+
+Feature files are HDF5 (`.h5`) with a `feats` dataset plus metadata attributes
+(`extractor`, `encoder`, `stamp_version`, `code_hash`, `feat_type`). Downstream
+stages read those attributes to validate that features came from a compatible
+extractor.
+
+## Branches and PRs
+
+**Never commit directly to `main`.** STAMP is developed branch-first: `main` is
+the integration branch and only receives changes through pull requests.
+
+- Branch off an up-to-date `main` (`git fetch origin && git switch -c <branch>
+  origin/main`).
+- Name branches with a topic prefix, as in the existing history:
+  `fix/…`, `feature/…`, `dev/…` (e.g. `fix/build-6`, `dev/multitask`,
+  `feature/random-seed`).
+- Open a PR against `main`. CI is wired to `push` and `pull_request` on `main`
+  only, so a branch gets its first full check when the PR opens — expect the
+  dependency-config, ruff and Linux test jobs to gate the merge.
+- Recent history is squash-merged, giving one commit per PR with a `(#NNN)`
+  suffix. Write the commit message accordingly: imperative mood, one line
+  describing the change ("Pin every dependency to an exact version", "Run the
+  test suite on Linux only").
+
+If you are an agent asked to "commit" work, check what branch you are on first
+and create one off `main` if the answer is `main`.
+
+## Repository layout
+
+```
+src/stamp/
+  __main__.py            CLI entry point (`stamp`); argparse + dispatch on config
+  config.yaml            Factory-settings config, copied by `stamp init`
+  types.py               Shared type aliases (Microns, SlidePixels, Task, Bags, …)
+  utils/                 config.py (StampConfig root model), seed.py, cache.py, path.py
+  preprocessing/         Tiling and tile-level feature extraction
+    __init__.py            `extract_()` — dispatches ExtractorName → extractor
+    config.py              `ExtractorName` enum + `PreprocessingConfig`
+    tiling.py              WSI → tiles, background rejection
+    extractor/             One module per foundation model (uni2, virchow2, …)
+  encoding/              Tile features → slide-/patient-level embeddings
+    config.py              `EncoderName` enum + encoding configs
+    encoder/__init__.py    `Encoder` ABC; subclasses in encoder/{titan,cobra,…}.py
+  modeling/              Training, cross-validation, deployment
+    data.py                Dataloaders, clini/slide table parsing, feature typing
+    train.py, crossval.py, deploy.py
+    registry.py            `ModelName` enum + (feature_type, task) → Lightning class
+    models/                VisionTransformer, MLP, TransMIL, Barspoon, Cox heads
+  statistics/            AUROC/AUPRC/regression/survival metrics and plots
+  heatmaps/              Attention heatmaps and top/bottom-scoring tile exports
+tests/                   pytest suite; `random_data.py` builds synthetic inputs
+scripts/                 check_dependency_config.py, verify_installed_stack.py,
+                         remote_gpu_check.sh (manual GPU-host harness, not CI)
+mcp/                     FastMCP server exposing the CLI stages as MCP tools
+.github/workflows/build.yml   CI
+```
+
+`README.md` covers installation, `getting-started.md` is the user-facing tutorial.
+
+## Environment and commands
+
+The project uses **uv** (pinned to `>=0.12,<0.13`) and Python 3.14 (3.13 also
+supported). uv is not optional — the CUDA wheel indexes and build exclusions
+live in `[tool.uv.*]` and are invisible to pip. See "Dependency rules" below.
+
+```bash
+uv sync --extra cpu --dev          # what you normally want locally
+uv sync --locked --extra cpu --dev # what CI runs; never rewrites uv.lock
+uv sync --extra gpu                # Linux + CUDA 13.0
+uv sync --extra gpu_all            # + conchv1_5, gigapath, musk
+```
+
+`cpu`, `gpu` and `gpu_all` are declared as mutually conflicting extras; pick one.
+
+```bash
+# Tests. The full suite is slow (model downloads, real training runs).
+uv run pytest tests/ --ignore=tests/test_feature_extractors.py   # what CI runs
+uv run pytest -m "not slow" tests/                               # quick pass
+uv run pytest tests/test_data.py -k some_case --verbose          # one case
+uv run pytest tests/test_feature_extractors.py -k uni2           # one extractor
+
+# Lint, format, types
+uv run ruff check --target-version=py313
+uv run ruff format --target-version=py313      # add --diff to only check
+uv run pyright
+
+# Dependency-configuration checks (stdlib only, no install needed)
+python scripts/check_dependency_config.py
+uv run python scripts/verify_installed_stack.py --expect-no-cuda
+```
+
+Extractor tests download gated model weights from Hugging Face and need
+`HF_TOKEN` in the environment; without it they are expected to fail, not to be
+"fixed". `tests/conftest.py` adds a `--extractor` option to select which ones
+run.
+
+CI (`.github/workflows/build.yml`) runs: dependency-config check, ruff
+lint+format, the test suite on Linux (3.13 and 3.14), a macOS install/import
+smoke test, GPU-extra installs on a machine with no CUDA toolkit (x86_64 and
+aarch64), and one job per extractor. macOS runs installation and imports only —
+it is a supported development platform, not a tested one.
+
+## Dependency rules
+
+These are enforced by `scripts/check_dependency_config.py` (and mirrored in
+`tests/test_dependency_config.py`). Breaking one turns CI red:
+
+- **Every requirement is pinned with `==`.** No floors, no ranges. A loose
+  specifier lets an unreviewed release into a fresh resolution.
+- **Git dependencies must be HTTPS + a full 40-char commit SHA.** No tags, no
+  branches.
+- **`flash-attn`, `mamba-ssm` and `causal-conv1d` come from the Astral index**
+  (`https://wheels.astral.sh/simple/cu130/`) as pre-built wheels, never from a
+  pinned wheel URL and never from source. `[tool.uv.no-build-package]` forbids
+  building them; their local versions (`+cu.13.0.torch.2.11`) must agree with
+  the PyTorch minor and CUDA channel the project selects.
+- **CPU installs must contain zero CUDA packages.** That is what
+  `[tool.uv.exclude-dependencies]` is for — upstream forks (UNI, GigaPath,
+  COBRA) declare CUDA extras unconditionally, so STAMP drops them and
+  re-declares them with platform markers in its own extras.
+- Several pins carry a comment explaining *why* they are held back (numpy 2.4
+  vs beartype, transformers 4.57 vs `trust_remote_code` models). Read the
+  comment before bumping; if you bump anyway, update or remove the comment.
+- Changing a dependency means regenerating `uv.lock` (`uv lock`) in the same
+  commit.
+
+## Code conventions
+
+- **Runtime type checking is on.** `src/stamp/__init__.py` calls
+  `beartype_this_package()`, so annotations are checked at runtime (violations
+  surface as `UserWarning`). Annotate accurately; a sloppy annotation is a bug,
+  not cosmetic.
+- **Array shapes are annotated with jaxtyping**, e.g.
+  `Float[Tensor, "batch tile feature"]`. Ruff's `F722` is globally ignored
+  because of this — don't re-enable it.
+- **Domain units are `NewType`s** in `stamp/types.py`: `Microns`,
+  `SlidePixels`, `TilePixels`, `SlideMPP`. Use them rather than bare `float`/
+  `int` so slide-space and tile-space pixels can't be mixed up.
+- **A trailing underscore means the function has side effects** (writes files,
+  mutates its arguments): `extract_()`, `train_categorical_model_()`,
+  `encode_slides_()`, `filter_complete_patient_data_()`. Keep the convention
+  when adding functions.
+- **Configs are Pydantic models with `extra="forbid"`**, composed under
+  `StampConfig` in `stamp/utils/config.py`. A new CLI option means: add the
+  field to the right config model, document it in `src/stamp/config.yaml`, and
+  pass it through in `__main__.py`.
+- **Heavy imports are deferred.** `__main__.py` imports torch/pydantic-heavy
+  modules inside the `match` arms so `stamp init` and `stamp --help` stay fast;
+  `extract_()` imports each extractor inside its `case`. Keep new imports local
+  in the same way.
+- **Optional backends fail loudly at import.** Every `extractor/*.py` and
+  `encoder/*.py` wraps its third-party imports in `try: … except
+  ModuleNotFoundError` and re-raises with the extra to install
+  (`pip install 'stamp[uni2]'`). Follow that pattern; the macOS CI job relies on
+  it, treating only `cobra` and `gigapath` as legitimately absent.
+- **Logging, not printing.** Use `_logger = logging.getLogger("stamp")`. Each
+  command also attaches a file handler writing `logfile.log` into its
+  `output_dir`.
+- Ruff is the formatter and linter (line length and style are its defaults).
+  Import sorting is on in the VS Code config; run `ruff format` before
+  committing.
+
+## Adding things
+
+**A tile-level feature extractor**: add a module under
+`preprocessing/extractor/` returning an `Extractor(model=…, transform=…,
+identifier=…)`; add a variant to `ExtractorName`; add a `case` to `extract_()`;
+declare its dependencies as an optional extra in `pyproject.toml` and add it to
+the `cpu`/`gpu`/`gpu_all` blanket extras as appropriate; list it in
+`src/stamp/config.yaml`'s comment and in the `test_extractors` CI matrix.
+
+**A slide/patient encoder**: subclass `Encoder` in `encoding/encoder/`,
+implementing `_generate_slide_embedding` and `_generate_patient_embedding`
+(override `encode_slides_`/`encode_patients_` only if you need tile
+coordinates); add a variant to `EncoderName` and wire it into
+`init_slide_encoder_`/`init_patient_encoder_`. `required_extractors` gates which
+tile features the encoder accepts.
+
+**A model**: add it under `modeling/models/`, add a `ModelName` variant, and
+wire it into `load_model_class()` in `modeling/registry.py`. The Lightning
+wrapper is chosen by `(feature_type, task)` from `MODEL_REGISTRY`; feature type
+is detected from the `.h5` metadata, task is one of `classification`,
+`regression`, `survival`. Only `barspoon` supports multi-target classification.
+
+## Gotchas
+
+- `identifier` on an `Extractor` must uniquely identify model *and* weights —
+  it is written into every feature file and checked downstream.
+- `generate_hash: True` appends a hash of the preprocessing source (see
+  `get_processing_code_hash`) to output directory names. Editing any file in
+  `preprocessing/` therefore changes where features land; that is deliberate.
+- Reproducibility goes through `Seed.set(seed)`, called once from `__main__.py`.
+  `torch.use_deterministic_algorithms()` is intentionally *not* enabled (large
+  performance cost), so same-seed runs follow the same trajectory but are not
+  bit-identical. `Seed.get_loader_worker_init()` raises if the seed was never
+  set.
+- Tests force the `spawn` multiprocessing start method (`tests/conftest.py`).
+  Don't switch to `fork`; it warns and misbehaves with the threaded loaders.
+- `deploy` accepts multiple checkpoints and majority-votes across them — that's
+  how cross-validation folds are combined.
+- Slides that raise during feature extraction are skipped, not fatal;
+  `extract_()` is deliberately fail-safe. Check `logfile.log` when output looks
+  short.
+- Don't commit `config.yaml` at the repo root — it's gitignored on purpose
+  (it's a user's local, path-laden copy of `src/stamp/config.yaml`).

CLAUDE.md

@@ -0,0 +1,4 @@
+# CLAUDE.md
+
+See [AGENTS.md](AGENTS.md) for all repository guidance — project layout, build
+and test commands, dependency rules, code conventions, and gotchas.

README.md

@@ -48,25 +48,13 @@ To setup STAMP you need [uv](https://docs.astral.sh/uv/) 0.12 or newer.
 | CUDA (`--extra gpu`, `--extra gpu_all`) | 3.13, 3.14 | Linux x86_64, Linux aarch64 |
 
 Python **3.14 is recommended** and is what `.python-version` selects; 3.13 remains
-supported. The CUDA builds are pinned to one ABI stack: **CUDA 13.0 with PyTorch
-2.11.0 and TorchVision 0.26.0**.
-
-CI runs the test suite on Linux. macOS is checked for installation and imports
-only, so it stays usable for development but is not a tested target.
-
-There is no CUDA build for macOS or Windows. `flash-attn`, `mamba-ssm` and
-`causal-conv1d` are only published as pre-built wheels for Linux, and STAMP
-refuses to compile them (see [below](#why-uv-is-required)), so a GPU extra on
-those platforms fails with a clear resolution error rather than a compiler error.
+supported. The CUDA builds are pinned to **CUDA 13.0 with PyTorch 2.11.0 and
+TorchVision 0.26.0**. There is no CUDA build for macOS or Windows — use
+`--extra cpu` there.
 
 > [!IMPORTANT]
-> uv is required for the GPU workflow. The PyTorch and Astral wheel indexes are
-> configured through `[tool.uv.sources]` in `pyproject.toml`, which is
-> uv-specific and invisible to other installers. Installing with plain `pip`
-> would mean pointing it at
-> `https://download.pytorch.org/whl/cu130` and
-> `https://wheels.astral.sh/simple/cu130/` yourself, and re-deriving the exact
-> pins by hand. See [Why uv is required](#why-uv-is-required).
+> Install with uv, not pip. The wheel indexes STAMP needs are configured in
+> `[tool.uv.sources]` in `pyproject.toml`, which pip does not read.
 
 ### Install or Update uv:
 
@@ -107,17 +95,13 @@ uv sync --extra gpu_all
 source .venv/bin/activate
 ```
 
-Both GPU extras install **pre-built wheels by default**. `flash-attn`,
-`mamba-ssm` and `causal-conv1d` come from the
-[Astral wheel index](https://wheels.astral.sh/simple/cu130/) already compiled
-against CUDA 13.0 and PyTorch 2.11, so nothing is compiled locally and no CUDA
-toolkit needs to be installed to *install* STAMP. You still need an NVIDIA
-driver to *run* on a GPU.
+Both GPU extras install pre-built wheels, so nothing is compiled locally and no
+CUDA toolkit is needed to install STAMP. You still need an NVIDIA driver to run
+on a GPU.
 
 > [!NOTE]
-> `--extra gpu_prebuilt` is deprecated and now just an alias for
-> `--extra gpu_all`, which is pre-built anyway. It will be removed in the next
-> breaking release.
+> `--extra gpu_prebuilt` is a deprecated alias for `--extra gpu_all` and will be
+> removed in the next breaking release.
 
 If you encounter errors during installation please read Installation Troubleshooting [below](#installation-troubleshooting).
 
@@ -137,31 +121,6 @@ If you encounter errors during installation please read Installation Troubleshoo
 > apt update && apt install -y libgl1 libglx-mesa0 libglib2.0-0
 > ```
 
-### Why uv is required
-
-The GPU workflow depends on configuration that only uv reads:
-
-* `[tool.uv.sources]` routes `torch` and `torchvision` to
-  `https://download.pytorch.org/whl/cu130` (or `.../cpu` for the `cpu` extra)
-  and the three compiled extensions to `https://wheels.astral.sh/simple/cu130/`.
-* `[tool.uv.exclude-dependencies]` drops the unconditional CUDA requirements
-  that the UNI, GigaPath and COBRA forks declare, so a CPU install stays free of
-  CUDA-only packages.
-* `[tool.uv.no-build-package]` forbids source builds of `flash-attn`,
-  `mamba-ssm` and `causal-conv1d`.
-
-None of this is visible to `pip`, which reads only `[project]`. Installing with
-pip would resolve `flash-attn` from PyPI and try to compile it — which is
-exactly what this configuration exists to prevent. If you must use pip, you have
-to add both indexes yourself and pin the extensions to the same
-`+cu.13.0.torch.2.11` local versions listed in `pyproject.toml`.
-
-Source builds of the three extensions are refused on purpose: they are
-ABI-locked to one PyTorch build, take a long time, need a matching CUDA toolkit,
-and were the most common cause of broken installs. An unsupported platform now
-fails during resolution with a clear message instead of part-way through a
-compile.
-
 ## Basic Usage
 
 If the installation was successful, running `stamp` in your terminal should yield the following output:
@@ -228,8 +187,7 @@ please consider citing our [Nature Protocols publication](https://www.nature.com
 
 These only exist as pre-built wheels for Linux x86_64 and Linux aarch64 on
 Python 3.13/3.14. On any other platform uv reports that no compatible version
-was found. This is deliberate — STAMP does not fall back to compiling them.
-Use `--extra cpu` on macOS and Windows.
+was found. Use `--extra cpu` on macOS and Windows.
 
 #### Triton Errors
 
@@ -258,10 +216,8 @@ was compiled against:
 E       ImportError: [...]/.venv/lib/python3.14/site-packages/flash_attn_2_cuda.cpython-314-x86_64-linux-gnu.so: undefined symbol: _ZN3c105ErrorC2ENS_14SourceLocationENSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
 ```
 
-This should no longer happen, because the extension versions carry the PyTorch
-build in their version string (`+cu.13.0.torch.2.11`) and the lockfile pins them
-together with torch. If you do hit it, you are almost certainly in an
-environment left over from an older STAMP release. Check what is installed:
+This usually means the environment is left over from an older STAMP release.
+Check what is installed:
 
 ```bash
 uv run python scripts/verify_installed_stack.py --expect-cuda