Document pinned inference and reproducible training recipes; add paper citation

CITATION.cff

@@ -0,0 +1,35 @@
+cff-version: 1.2.0
+message: "If you use Morph, please cite the accompanying paper."
+type: software
+title: Morph
+authors:
+  - family-names: Franke
+    given-names: Malte
+  - family-names: Schmid
+    given-names: Stefan P.
+  - family-names: Ivkovic
+    given-names: Zarko
+  - family-names: Jorner
+    given-names: Kjell
+  - family-names: Krause
+    given-names: Andreas
+repository-code: "https://github.com/maltefranke/morph-chem"
+preferred-citation:
+  type: generic
+  title: "Generative Molecular Morphing for Flexible-Size Design via Unbalanced Optimal Transport"
+  authors:
+    - family-names: Franke
+      given-names: Malte
+    - family-names: Schmid
+      given-names: Stefan P.
+    - family-names: Ivkovic
+      given-names: Zarko
+    - family-names: Jorner
+      given-names: Kjell
+    - family-names: Krause
+      given-names: Andreas
+  year: 2026
+  url: "https://arxiv.org/abs/2606.07239"
+  identifiers:
+    - type: other
+      value: "arXiv:2606.07239"

README.md

@@ -7,6 +7,7 @@ tags:
   - chembl3d
   - geom-drugs
   - 3d
+  - arxiv:2606.07239
 ---
 
 # Morph: 3D molecular generation
@@ -16,6 +17,8 @@ Unconditional molecular generation checkpoints trained on **ChEMBL3D** and
 needed for inference. No training dataset is required to load a model or generate
 molecules.
 
+See our [paper](https://arxiv.org/abs/2606.07239), *Generative Molecular Morphing for Flexible-Size Design via Unbalanced Optimal Transport*.
+
 ## Available checkpoints
 
 | Dataset | Model name | Hydrogen representation | Completed epochs | Files |
@@ -42,7 +45,8 @@ discovery for all models in the table. Install it in a Python environment:
 ```bash
 python -m pip install 'huggingface_hub>=1.0,<2'
 hf auth login
-hf download maltefranke/morph --local-dir morph-runtime \
+hf download maltefranke/morph --revision da1c241d83f2628b172fc510e1bc724ef4429db2 \
+  --local-dir morph-runtime \
   --include 'geom/explicit_h/runtime/*'
 python -m pip install ./morph-runtime/geom/explicit_h/runtime
 morph models list
@@ -132,6 +136,86 @@ The root `sample.py` remains available as the archived ChEMBL3D sampling entry
 point with its corresponding bundled runtimes. The shared interface above
 supports all three published model variants.
 
+## Reproducibility
+
+### Generate with a fixed release
+
+Pin the **runtime code and dependencies** as well as the **model bundle**.
+This installs the archived shared inference runtime, independently of changes
+to this repository or the model catalog. Install [uv](https://docs.astral.sh/uv/)
+and authenticate with an account that can access the model repository:
+
+```bash
+uv tool install 'huggingface_hub==1.33.0'
+hf auth login
+hf download maltefranke/morph \
+  --revision da1c241d83f2628b172fc510e1bc724ef4429db2 \
+  --include 'geom/explicit_h/runtime/*' --local-dir morph-runtime
+cd morph-runtime/geom/explicit_h/runtime
+uv sync --frozen --no-dev --python 3.13.7
+source .venv/bin/activate
+
+# Run on a compute worker. Save all attempts and their validity flags.
+MORPH_COMPILE=0 MORPH_ATTENTION_BACKEND=sdpa MORPH_FLASH_ATTN=0 \
+morph generate --model maltefranke/morph --subfolder geom/explicit_h \
+  --revision b839af108eaade757a5d2d668d05498e8c733c1e --epoch 50 \
+  --device cuda --ema --steps 100 --seed 42 --batch 64 --batch-size 16 \
+  --save-trajs --output geom-e50-seed42.sdf
+uv pip freeze > environment.txt
+nvidia-smi --query-gpu=name,driver_version --format=csv > gpu.csv
+```
+
+Choose another published run by changing these three options:
+
+| `--subfolder` | `--revision` | `--epoch` |
+| --- | --- | --- |
+| `chembl3d/explicit_h` | `ab9bea92f499f03450ebc99459fce04461504a41` | `200`, `300`, `400` |
+| `chembl3d/implicit_h` | `ab9bea92f499f03450ebc99459fce04461504a41` | `100`, `200`, `300` |
+| `geom/explicit_h` | `b839af108eaade757a5d2d668d05498e8c733c1e` | `50` |
+
+Keep the seed, batch size, integration steps, EMA selection, precision, attention
+backend and software/hardware environment fixed when comparing samples. A seed
+alone does not guarantee identical samples across GPU models or library versions.
+The checkpoint SHA-256 hashes are in each bundle's `metadata.json`.
+
+### Repeat a training recipe
+
+The [training archives](https://huggingface.co/maltefranke/morph/tree/main/reproduction)
+contain the original source (including `run.py`), saved training config, launcher,
+dependency versions, and a portable `train.sh`. They use **2 RTX 4090s, seed 42,
+and bf16 mixed precision**. Run installation and training in a compute allocation:
+
+```bash
+# Start from a directory outside the inference environment above.
+run=geom-explicit-h  # Or chembl3d-explicit-h, chembl3d-implicit-h, qm9-explicit-h
+hf download maltefranke/morph --revision d0332f1c022345956470bb948ada79e5702ed8dc \
+  "reproduction/$run.tar.gz" reproduction/SHA256SUMS --local-dir morph-training
+cd morph-training/reproduction
+sha256sum --ignore-missing -c SHA256SUMS
+tar -xzf "$run.tar.gz"
+cd "$run"
+uv venv --python 3.13.7 .venv
+uv pip install --python .venv/bin/python -r pyproject.toml \
+  -c reproduction/requirements-environment.txt
+source .venv/bin/activate
+export MORPH_DATASET_ROOT=/absolute/path/to/geom  # Contains processed/
+export MORPH_OUTPUT_DIR="$PWD/outputs/repeat-01"
+bash train.sh
+```
+
+All portable recipes evaluate every epoch and retain every fifth epoch; QM9
+retains every **25th** epoch. The archives start fresh and stop at 400 epochs for
+explicit-H ChEMBL3D, 300 for implicit-H ChEMBL3D, 50 for GEOM, or 3000 for QM9.
+Override the stopping point with `bash train.sh trainer.trainer.max_epochs=100`.
+QM9 currently has a training recipe but no published checkpoint.
+
+These commands reproduce the documented setup, not an exact replay of every
+historical run: explicit-H ChEMBL3D resumed at epoch 5 and GEOM at epoch 48;
+the older explicit-H ChEMBL3D run also evaluated every five epochs. Data caches
+and those intermediate resume checkpoints are not included in the archives.
+See [reproducibility details](reproduction/README.md) for the data requirements,
+original settings, environment provenance and resume instructions.
+
 ## Hydrogen handling and training data
 
 | Dataset | Training filter | Representation |
@@ -181,3 +265,21 @@ historical model card. GEOM has its own
 [bundle manifest](https://huggingface.co/maltefranke/morph/blob/b839af108eaade757a5d2d668d05498e8c733c1e/geom/explicit_h/manifest.json).
 Original training configs contain historical cluster paths; use `config.yaml`
 through the inference loader for portable generation.
+
+## Citation
+
+If you use Morph, please cite our [paper](https://arxiv.org/abs/2606.07239):
+
+```bibtex
+@misc{franke2026generativemolecularmorphingflexiblesize,
+      title={Generative Molecular Morphing for Flexible-Size Design via Unbalanced Optimal Transport},
+      author={Malte Franke and Stefan P. Schmid and Zarko Ivkovic and Kjell Jorner and Andreas Krause},
+      year={2026},
+      eprint={2606.07239},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2606.07239},
+}
+```
+
+Machine-readable citation metadata is available in [CITATION.cff](CITATION.cff).

reproduction/README.md

@@ -0,0 +1,178 @@
+# Reproducing models and training runs
+
+The [README](../README.md#reproducibility) contains complete commands for pinned
+inference and training. The model repository is
+[maltefranke/morph](https://huggingface.co/maltefranke/morph).
+
+## Published inference
+
+Pin both the installed source and the checkpoint/assets. The shared inference
+runtime at Hub revision `da1c241d83f2628b172fc510e1bc724ef4429db2` supports all
+published models; its `uv.lock` pins dependencies. The README installs that
+runtime with Python 3.13.7. Installing a different checkout of Morph changes the
+implementation even when `load_pretrained(..., revision=...)` selects the same
+weights: that argument pins the model bundle, not the installed Python code.
+
+With that environment active, this is the Python equivalent of the pinned
+GEOM example. Set the environment variables before importing Morph:
+
+```python
+import os
+
+os.environ["MORPH_COMPILE"] = "0"
+os.environ["MORPH_ATTENTION_BACKEND"] = "sdpa"
+os.environ["MORPH_FLASH_ATTN"] = "0"
+
+from morph.inference import load_pretrained
+
+model = load_pretrained(
+    "maltefranke/morph",
+    subfolder="geom/explicit_h",
+    revision="b839af108eaade757a5d2d668d05498e8c733c1e",
+    epoch=50,
+    device="cuda",
+    steps=100,
+    use_ema=True,
+)
+trajectories, validity, vocab, coordinate_std = model.generate(
+    n_mols=64, batch_size=16, seed=42, return_traj=True,
+)
+print("Model revision:", model.bundle.revision)
+print("Valid attempts:", sum(validity), "/", len(validity))
+```
+
+Direct repository IDs bypass catalog lookup. To use a short model name while
+freezing discovery, pass a pinned catalog:
+
+```python
+from morph.inference import load_catalog, load_pretrained
+
+catalog = load_catalog(
+    repo_id="maltefranke/morph",
+    revision="377e8b52505d564df8529fe6fef2da9c77f57661",
+)
+model = load_pretrained(
+    "chembl3d-implicit-h", catalog=catalog, epoch=300,
+    device="cuda", steps=100, use_ema=True,
+)
+```
+
+Changing batch size can change random-number consumption and generated samples.
+Record the generation count, batch size, seed, steps, EMA/current choice,
+precision, backend, source revision, package versions, GPU and driver. The CLI's
+SDF includes valid molecules only; `--save-trajs` preserves every attempt and
+its validity flag. Coordinates are already in Angstrom.
+
+Seeds and fixed releases provide controlled comparisons, but CUDA kernels and
+changes to hardware, dependencies or precision can prevent bitwise equality.
+The historical training configuration uses `deterministic: warn`, which is not
+a guarantee that every operation has a deterministic implementation.
+
+## Training archives
+
+Download from `reproduction/` at immutable Hub revision `d0332f1c022345956470bb948ada79e5702ed8dc`.
+Each archive contains:
+
+- The original `run.py`, `src/`, `configs/`, `pyproject.toml`, and `uv.lock`
+  where present in the original source snapshot. Original source files are
+  checked against their saved SHA-256 manifest.
+- `reproduction/original_training_config.yaml` and
+  `reproduction/original_launcher.sbatch` for historical settings and paths.
+  For QM9, the config is from the fresh training run after its successful GPU
+  preflight; the launcher records the preflight reuse and allocation time limit.
+- `configs/reproduce.yaml` and `train.sh` for portable fresh training.
+- `reproduction/provenance.json`, `dataset_reference.json`,
+  `original_source_manifest.json`, and `manifest.json` for auditability.
+- `reproduction/requirements-environment.txt`: version constraints captured
+  from the shared training environment on 2026-10-08. This is not a historical
+  per-job environment capture; Python 3.13.7 is the reference interpreter.
+
+The archives are roughly 0.2–0.5 MB each and contain no datasets or weights.
+`reproduction/SHA256SUMS` verifies the archives; each archive's manifest verifies
+its files. Preserve the archive and immutable revision with your results.
+
+The new launcher records package versions, GPU/driver information, requested
+overrides, recipe provenance and data references in `MORPH_OUTPUT_DIR`.
+Hydra records the effective configuration in `.hydra/`. Choose a new output
+directory for each run, and retain these records alongside checkpoints.
+
+| Archive | Batch per GPU | Accumulation | Effective batch | Original validation cadence | Portable validation / save cadence |
+| --- | ---: | ---: | ---: | --- | --- |
+| `chembl3d-explicit-h.tar.gz` | 320 | 1 | 640 | Every 5 epochs | Every 1 / 5 epochs |
+| `chembl3d-implicit-h.tar.gz` | 320 | 1 | 640 | Every epoch | Every 1 / 5 epochs |
+| `geom-explicit-h.tar.gz` | 320 | 1 | 640 | Every epoch | Every 1 / 5 epochs |
+| `qm9-explicit-h.tar.gz` | 1024 | 2 | 4096 | Every epoch | Every 1 / 25 epochs |
+
+All recipes use two RTX 4090 GPUs, seed 42, bf16 mixed precision, SDPA attention,
+compiled networks and 4096 generation attempts per validation epoch. GEOM uses
+Muon with CUDA graph execution; ChEMBL3D and QM9 use native Muon execution with
+DDP compiler graph splitting disabled and compilation fallback enabled.
+Environment flags are explicit in each archive's `reproduction/environment.sh`.
+
+The portable configs preserve the saved model, optimizer, schedule and batch
+settings. They change data/output paths, disable W&B, remove the historical
+resume path and allocation timer, and use a finite epoch limit. They also apply
+the current validation and checkpoint policy from the table. Those changes are
+listed in `provenance.json`; the historical config is preserved separately.
+
+## Data needed for training
+
+Set `MORPH_DATASET_ROOT` to a dataset directory containing `processed/`:
+
+| Archive | Dataset directory | Selection / representation |
+| --- | --- | --- |
+| ChEMBL3D explicit H | `chembl3d_max100` | At most 100 total atoms, explicit H |
+| ChEMBL3D implicit H | `chembl3d_max100_implicit_h` | Same source molecules, then remove H |
+| GEOM explicit H | `geom` | At most 72 total atoms, explicit H |
+| QM9 explicit H | `qm9` | Original QM9 split, explicit H |
+
+The ChEMBL3D cap counts hydrogens **before removal**. Both variants have
+1,474,008 / 184,295 / 184,242 train/validation/test molecules. The implicit-H cache
+also repairs stale cached charges and bond representations from the canonical
+source SMILES, then recomputes heavy-atom statistics. It is not interchangeable
+with the explicit-H cache. QM9 has 100,000 / 20,000 / 12,947 molecules.
+
+Keep the original processed split files, their molecule/conformer order,
+`*_smiles.txt` files, ordered vocabularies and `distributions.pt` together.
+ChEMBL3D and GEOM use `<split>_data.lmdb/data.mdb`; QM9 uses `<split>.pt`.
+Training uses the training SMILES for novelty metrics as well as the graph cache.
+
+The launcher checks vocabulary/statistics SHA-256 hashes and processed split
+file sizes before starting distributed training. **File sizes do not verify
+dataset contents.** The historical release did not record complete dataset
+content hashes, and these caches are not hosted in the model repository.
+Rebuilding from upstream data therefore requires checking split membership,
+conformer selection and ordering; matching an atom cap alone is insufficient.
+For new datasets, use the repository's preprocessing workflow before a
+distributed launch, and preserve the resulting cache. Run dataset hashing on
+a CPU worker if you need a full content manifest for future comparisons.
+
+## Fresh training versus resuming
+
+The explicit-H ChEMBL3D release resumed after 5 completed epochs. GEOM epoch 50
+resumed from the best checkpoint after 48 completed epochs. Starting these
+recipes from seed 42 reproduces their setup but does not replay that history.
+The intermediate epoch-5 and epoch-48 checkpoints are not in the Hub release.
+The implicit-H ChEMBL3D run and the QM9 recipe start fresh. QM9 has no published
+checkpoint in this release.
+
+To continue from an available full training checkpoint, pass its absolute path
+and the desired **total** completed epoch count to the matching archive:
+
+```bash
+bash train.sh \
+  "trainer.checkpoint_path='/absolute/path/to/epoch_050.ckpt'" \
+  trainer.trainer.max_epochs=100
+```
+
+Use `trainer.checkpoint_path` to restore optimizer, scheduler, EMA and loop state;
+loading inference weights alone does not resume training. Periodic checkpoint
+filenames use Lightning's zero-based epoch index: index 49 is completed epoch
+50. A `final.ckpt` written **after** `fit` may already have advanced that counter;
+inspect the completed loop count rather than adding one unconditionally.
+
+The portable recipes disable the wall-clock timer. If you add `max_time` when
+resuming, Lightning also restores elapsed timer state: account for that elapsed
+time when setting the new budget. Keep the original source, data order, rank
+count and microbatch settings for the closest continuation; state restoration
+does not guarantee a bitwise replay of all RNG and data-worker behavior.