saddleflow 0.1.0__tar.gz

saddleflow-0.1.0/.claude/scheduled_tasks.lock

@@ -0,0 +1 @@
+{"sessionId":"94b4102a-2392-429c-bad0-9e0b075ba740","pid":2648357,"procStart":"49722193","acquiredAt":1778081663965}

saddleflow-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,36 @@
+# Build the package and upload it to PyPI when a GitHub release is published. Uses PyPI trusted publishing
+# (the `pypi` environment of this repository is registered as a publisher on PyPI), so no token is stored.
+name: Publish to PyPI
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - run: python -m pip install build twine
+      - run: python -m build
+      - run: python -m twine check dist/*
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish:
+    needs: build
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1

saddleflow-0.1.0/.gitignore

@@ -0,0 +1,75 @@
+# Python
+__pycache__/
+*.pyc
+*.egg-info/
+.pytest_cache/
+.venv/
+.eggs/
+build/
+dist/
+
+# PyTorch / checkpoints
+*.pt
+*.pth
+*.ckpt
+*.safetensors
+
+# Training outputs
+checkpoints/
+results/
+logs/
+wandb/
+lightning_logs/
+
+# Example-run output directories. Nothing under runs/ is tracked, so a user's
+# first training leaves a clean `git status`. To commit a reference figure
+# anyway, force it: `git add -f examples/<case>/runs/<run>/viz/foo.png`.
+**/runs/
+
+# Belt-and-braces for checkpoints written outside a runs/ directory: keep light
+# files (*.json, *.png, *.npz) but never per-checkpoint heavy tensors (model
+# weights, optimizer state, EMA shadow, RNG state, LR-scheduler state) or
+# intermediate epoch checkpoints.
+**/checkpoint_epoch_*/
+**/checkpoint_final/model.safetensors
+**/checkpoint_final/optimizer.bin
+**/checkpoint_final/scheduler.bin
+**/checkpoint_final/ema.pt
+**/checkpoint_final/random_states_*.pkl
+
+# Data directories (datasets not tracked in git) anywhere in the tree,
+# with an explicit exception for saddleflow/data/ (the source package).
+data/
+!saddleflow/data/
+!saddleflow/data/**
+# ...and for the recorded SaddleMill NEB input of the LiC_simpler example (10 KB).
+!examples/LiC_simpler/neb/data/
+!examples/LiC_simpler/neb/data/**
+# Re-exclude cache/build artifacts that the negation above would otherwise un-ignore.
+saddleflow/data/**/__pycache__/
+saddleflow/data/**/*.pyc
+*.lmdb
+*.h5
+*.hdf5
+
+# SLURM
+slurm*.log
+slurm-*.out
+*.err
+
+# OS / editors
+.DS_Store
+.idea/
+.vscode/
+
+# Reference codes — kept locally for design reference; will be deleted before public release
+other_codes/
+
+# Private experiments tree (its own git repo; not part of the public package)
+private/
+
+# NeurIPS rebuttal working directory (uploaded to OpenReview, not part of the package)
+rebuttal/
+
+# Personal / machine-specific Claude notes (research log, cluster paths); never commit
+CLAUDE.local.md

saddleflow-0.1.0/CLAUDE.md

@@ -0,0 +1,206 @@
+# CLAUDE.md
+
+Guide for anyone (person or coding agent) working on SaddleFlow. Machine- or person-specific notes (cluster paths,
+job logs, unpublished experiments) belong in an untracked `CLAUDE.local.md` next to this file; it is git-ignored and
+must never be committed.
+
+## What SaddleFlow is
+
+SaddleFlow generates transition-state (first-order saddle) structures for periodic materials by **flow matching**.
+The velocity field is Meta FAIR's pretrained **UMA** backbone (`uma-s-1p2` or `uma-m-1p1`) with an equivariant
+**time-FiLM** at selected backbone blocks (all of them in the paper model), an optional global-attention layer
+(`GlobalAttn`, off in every released recipe), and a small SO(3)-equivariant `VelocityHead` that outputs a per-atom
+velocity. Package import name `saddleflow`; written
+**SaddleFlow** in prose.
+
+Two ways to use it:
+
+- **Unconditional, single-ended (Dimer-like).** Only the reactant is known. Start from the reactant plus Gaussian noise
+  and flow to a nearby first-order saddle; the head never sees the product (`--delta-endpoint-channels 0`). This is
+  the model in the MaterialsSaddles paper (NeurIPS 2026, Evaluations & Datasets), trained in `examples/MaterialsSaddles`.
+- **Conditional, double-ended (NEB-like).** Both endpoints (R, P) are known. Start at their minimum-image midpoint and
+  give the head the per-atom displacement to each endpoint at every step (`--delta-endpoint-channels 32`). Recipe in
+  `examples/MP20Bat`.
+
+## Repository layout
+
+- `saddleflow/data/` — datasets (`TrajTripletDataset` for `.traj`, `AseDbSaddleDataset`, `MaterialsSaddlesDataset`
+  for the released ASE-LMDB shards), coordinate transforms (`wrap_positions`, `mic_displacement`).
+- `saddleflow/flow/` — `FlowMatchingConfig`, `FlowMatchingLoss`, start-distribution sampling, sampler.
+- `saddleflow/models/` — `TimeFiLMBackbone`, `VelocityHead`, `GlobalAttn`.
+- `saddleflow/utils/` — backbone loading, the `accelerate` training loop (`training.py`), EMA, evaluation helpers.
+- `examples/` — `MaterialsSaddles` (paper model, data download, evaluation), `MP20Bat` (conditional recipe),
+  `LiC` and `LiC_simpler` (small unconditioned examples with their own READMEs), `LematBulk`, `scaling_bench`.
+- `tests/` — `pytest tests`.
+
+## The paper model (unconditional, all four MaterialsSaddles subsets)
+
+**Data.** `python examples/MaterialsSaddles/data_prep.py` downloads MaterialsSaddles from Hugging Face
+(`SciLM/MaterialsSaddles`) into `$MATERIALSSADDLES_ROOT` (default `$SCRATCH/MaterialsSaddles_v2`). The release is
+deduplicated and split 90/5/5 by element set into `<subset>/{train,val,test}/` directories; no chemical system
+appears in two splits.
+
+**Training command** (one epoch over all subsets with oc20 ×2 and oc22, mp20bat ×8 = 67.7 M records; 64 GH200 GPUs,
+8 records per GPU, global batch 512, ~132 k steps). Launch one process per GPU with the SPMD pattern of
+`examples/MaterialsSaddles/run.sh`:
+
+```
+python examples/MaterialsSaddles/train.py --all-subsets --output-dir <out> \
+  --subset-repeat oc22=8,mp20bat=8,oc20=2 --task-name-map oc22=oc20 \
+  --backbone uma-m-1p1 --mixed-precision no --grad-clip-norm 1.0 --batch-size 8 --num-epochs 1 \
+  --learning-rate 1e-3 --uma-lr 3e-3 --warmup-steps $((250 * WORLD_SIZE)) --ema-decay 0.9996 \
+  --loss-type huber --huber-delta 0.05 \
+  --delta-endpoint-channels 0 --path-start-prob 1.0 --path-noise-sigma 0.3 --path-u-power 2.0 \
+  --save-every-steps 2000 --val-every-steps 2000 --max-val-records 4000 \
+  --unfreeze-uma-all --early-time-film-blocks 0,1,2,3,4,5,6,7,8,9 --head-depth 3 --attn-layers 0 \
+  --com-symmetric-loss --xt-target-correction --xt-perturb-sigma 0.05 \
+  --no-inject-force --no-dimer-residual --no-frozen-force-backbone --no-endpoint-features --eigenmode-aux-weight 0
+```
+
+**Why each choice.**
+- **Path starts** (`--path-start-prob 1`). Each sample starts on the line from an endpoint E (reactant or product) to its
+  saddle S plus noise: `x0 = E + u (S − E) + ε`, `ε ~ N(0, σ²)` on mobile atoms. Training only from `S + ε`
+  (transition-state denoising) fails for this task: in 3N dimensions an isotropic draw almost never points along the
+  reaction, so the model learns to undo noise and stays at the minimum.
+- **`u ~ U(0,1)²`** (`--path-u-power 2`) biases starts toward the endpoint, where inference begins.
+- **σ = 0.3 Å matches the inference noise.** Never deploy wider than the training σ (the model degrades sharply);
+  narrower is safe. Ship the training σ with the weights.
+- **x_t perturbation (0.05 Å)** is not redundant with the start noise: ε enters `x_t` scaled by `(1 − t)` and fades near
+  the saddle, where the model's own endpoints land off the line.
+- **Convergent target** (`--xt-target-correction`, `t_floor = 0.1`): `v = (S − x_t)/(1 − t)` for `t ≤ 0.9`, with S and
+  x_t unwrapped to the start (not a minimum-image difference); `S − x0` above (perturbation off there).
+- **Huber δ = 0.05** approximates the median of a multi-modal target where MSE averages the modes.
+- **Backbone LR 3e-3 on UMA-M** (1e-2, the UMA-S optimum, diverges on UMA-M); head/FiLM LR 1e-3.
+- **`--task-name-map oc22=oc20`**: `uma-m-1p1` has no `oc22` mixture-of-experts entry and raises `KeyError('oc22')`.
+
+**Inference: ×5 K4.** Spend 20 network evaluations as five passes of four Euler steps, each pass restarting at `t = 0`
+from the previous endpoint, rather than one pass of 20 (×1 K20). A single pass stops short of the saddle; because the
+training starts span the whole endpoint–saddle line, the same model also knows how to refine its own endpoint.
+
+```
+SF_NAPPLY=5 python examples/MaterialsSaddles/dump_throws_batched.py \
+  --data-glob "$MATERIALSSADDLES_ROOT/<subset>/*/<subset>_*.aselmdb" --subset <subset> --split test \
+  --ckpt <ckpt> --use-ema --task-name-map oc22=oc20 --sides 0 --sigma 0.3 --K 4 --seed 0 --tag ev --outdir <dir>
+```
+
+Then converge every endpoint with Sella (see Evaluation) and collect with `sm_collect.py`.
+
+## The conditional recipe (`examples/MP20Bat`)
+
+UMA-S-1.2 with all four blocks unfrozen (LR 1e-4) and time-FiLM at every block; `VelocityHead(depth=3,
+delta_endpoint_channels=32)` fed `(Δ_R, Δ_P) = MIC(R − x_t), MIC(P − x_t)` stacked as `(N, 2, 3)`; `x0 = (R + P)/2`
+(minimum-image midpoint); x_t perturbation 0.05 Å with the convergent target; CoM-symmetric loss. `run.sh` holds the
+exact hyperparameters. Later work found Huber loss and a much higher backbone LR (3e-3 to 1e-2) clearly better for
+this recipe. Score conditional models against the dataset label (they are asked for that specific saddle).
+
+## Architecture
+
+**Backbone (UMA).** Loaded by `saddleflow.utils.load_uma_backbone`; only `.backbone` is used, the energy/force heads are
+discarded. UMA-S-1.2: 4 blocks, lmax 2, 128 channels, embedding `(N, 9, 128)`. UMA-M-1.1: 10 blocks, lmax 4,
+128 channels, embedding `(N, 25, 128)`, 1.4 B parameters. With a 6 Å cutoff, even the 4 UMA-S blocks reach ~24 Å,
+beyond most released cells; this is why `GlobalAttn` stays off. Mixture-of-experts routing comes from each record's
+`task_name` (`omat`, `oc20`, `oc22`, ...) via `csd_embedding(charge, spin, task)`; `charge`/`spin` default to 0.
+
+**Time-FiLM.** A zero-initialised equivariant FiLM is hooked before each selected block (`TimeFiLMBackbone`) and in the
+head: an additive bias on `l = 0` channels and a `(1 + γ)` gate on `l ≥ 1` channels, both functions of `t` only, so
+equivariance is exact. The bias alone would never reach `l ≥ 1` outputs (`SO3_Linear` keeps the `l` paths separate);
+the gate is what lets time act on the vector channels. At init the wrapped model equals UMA bit for bit.
+
+**VelocityHead.** Built from UMA's `SO3_Linear` (not e3nn's `o3.Linear`; the two tensor layouts are not
+interchangeable) with `depth − 1` blocks of `SO3_Linear → UMAGate` before the `l = 1` readout. Optional zero-initialised
+inputs (endpoint displacements, forces) leave it identical to a plain force head when off.
+
+**Output projections, in order.** (1) `v[fixed] = 0`. (2) Subtract the mobile-atom mean velocity **only in systems with
+no frozen atoms** (`_com_projection_batched`). Do not make this unconditional: in a system with one mobile atom it
+zeroes that atom's velocity and silently kills training.
+
+## Flow formulation
+
+- **State:** Cartesian positions in Å; UMA always sees physical positions, wrapped into the cell just before each
+  forward pass. Interpolation happens in unwrapped space.
+- **Minimum-image unwrap** of every target relative to the start (once, when a record is built):
+  `Δs = inv(cell) @ (r_target − r_start); Δs −= round(Δs); r_target_un = r_start + cell @ Δs`.
+- **R↔P doubling:** every triplet gives two records, one per endpoint as start; the unwrap is recomputed per anchor.
+- **Start distributions** (`FlowMatchingConfig`): `(R+P)/2` midpoint (conditional), transition-state denoising
+  (`ts_denoise_sigma`), path starts (`path_start_prob`, `path_noise_sigma`, `path_u_power`, optional σ schedule
+  `path_noise_sigma_endpoint`/`_saddle`). All new options default off.
+- **Path:** `x_t = (1 − t) x0 + t S`, `t ~ U(0, 1)`; target as above.
+- **Inference:** forward Euler, `K` steps per pass, optional restarts (`SF_NAPPLY` in the dumpers). Higher-order
+  integrators and larger `K` do not help; the error is in the field, not the integration.
+
+## Data formats
+
+All sources reduce to the same per-record dict: `start_pos`, `partner_un_pos`, `saddle_un_pos` (both unwrapped to the
+start), `Z`, `cell`, `fixed`, `task_name`, `charge`, `spin`, `delta_norm`, `role` (`R2S`/`P2S`), `triplet_id`,
+`metadata`. Triplets are stored as consecutive `[R, S, P]` frames. Small problems read `.traj` directly
+(`TrajTripletDataset`); MaterialsSaddles reads ASE-LMDB shards (`MaterialsSaddlesDataset`, which also applies
+`task_name_map`). Neighbour lists are rebuilt every forward pass (`AtomicData.from_ase`), since bonds change along the
+flow.
+
+## Training infrastructure
+
+Plain PyTorch + Hugging Face `accelerate` (`saddleflow/utils/training.py`); no Hydra or Lightning.
+
+- **Multi-node: SPMD, one process per GPU** (`RANK=$SLURM_PROCID`, `WORLD_SIZE=$SLURM_NTASKS`, static
+  `init_process_group`), as in `examples/MaterialsSaddles/run.sh`. Do not use `accelerate launch --num_machines`
+  across many nodes: its elastic rendezvous strands stragglers. `examples/MP20Bat/run.sh` still uses the old pattern.
+- **Warmup and `--max-steps` count scheduler ticks, and accelerate ticks once per rank per optimizer step.** For a real
+  warmup of W optimizer steps pass `--warmup-steps W × world_size`. `--max-steps` is converted internally.
+- **Fixed global batch, free node count.** Changing nodes at fixed global batch changes nothing but speed; never
+  resume a run on a different node count (the schedule is per-rank ticks).
+- **Resume** (`--resume-from <checkpoint>`) continues mid-epoch with the same shuffle and skips trained batches.
+  `--init-weights` loads weights only (fresh optimizer and schedule) and re-snapshots the EMA.
+- **EMA:** choose `d = 1 − 1/window` with `window ≈ total_steps/50`. Too high a decay on a short run leaves the shadow
+  at the initial weights (evaluation then looks broken; compare with live weights).
+- **Large validation sets:** cap them with `--max-val-records`.
+- **Precision:** fp32 is the safe default (the paper model is fp32). bf16 autocast is acceptable for UMA-S training
+  because the coordinate helpers are pinned to fp32 (see bug log).
+
+## Evaluation
+
+- **Score with Sella, never with ASE Dimer.** Dimer wanders to a different saddle in a sizeable fraction of cases and
+  fabricates a tail. Run Sella (via SaddleMill, `method = Sella`; example `examples/Sella_OC22` in
+  `SciLM-ai/SaddleMill`) from each prediction: UMA-S-1.2 with the subset's own task (`omat` for lemat/mp20bat,
+  `oc20`, `oc22`), first order, `fmax = 0.01` eV/Å, ≤ 300 steps, `delta0 = 0.1`, index check (exactly one eigenvalue
+  below −0.01 eV/Ų). For SaddleFlow outputs use `dataset_type = bulk`: they carry their own `FixAtoms` and all-zero
+  tags, and `oc` would re-freeze every tag-0 atom.
+- **Check the task name against the subset before scoring.** Scoring through the wrong expert inflates errors.
+- **Metric:** maximum per-atom displacement (minimum image) from the prediction to the saddle Sella reaches, over cases
+  that converge and verify as first order; report mean and median and the Sella force calls. Unconditional models are
+  scored against the Sella saddle, conditional models against the dataset label.
+- **Collect strictly** (`sm_collect.py`): count only real Sella convergence and verified index-1; a desorption check
+  can mark an unoptimised input as converged at step 0; Sella crashes are failures, not dropped cases; a truncated
+  pass keeps the easy cases first and is biased.
+- **Compare arms on identical starts** with paired tests. Different noise draws move medians by several hundredths
+  of an Å, more than most training changes.
+
+## Latent-bug log — re-check these if behaviour looks wrong
+
+- **Autocast silently quantises coordinates.** `wrap_positions` and `mic_displacement` do a Cartesian→fractional→
+  Cartesian round trip; under `torch.autocast` those matmuls run in bf16 (spacing 0.06 Å at 12 Å), so the structure
+  fed to UMA is corrupted while the target is exact. Training still reaches a *lower* loss. Both helpers now force
+  autocast off (`tests/test_autocast_precision.py`). Rule: never do arithmetic on absolute coordinates under autocast.
+  fp16 hides the bug; do not test for it with fp16.
+- **UMA dropouts leak into training.** `.train()` enables UMA's composition and mixture dropouts; `FlowMatchingLoss`
+  re-calls `backbone.eval()`.
+- **Time embedding base.** The transformer base 10000 is ~constant over `t ∈ [0, 1]`; the head uses geometric
+  frequencies from 1 to `half` cycles per unit time.
+- **RMSD over all atoms dilutes error by √N.** All evaluation helpers take `mobile_mask`.
+- **CoM projection** must stay conditional on having no frozen atoms (see Architecture).
+- **AdaLN-zero:** the first time-MLP layer has zero gradient at step 0 only; if it persists, something is broken.
+- **Force injection uses the first sample's task name for the whole batch.** Fine for homogeneous batches or with
+  `--no-inject-force` (all released recipes); wrong for mixed-subset batches with force injection.
+- **Multi-node NCCL hang** on some clusters near an epoch boundary (a race between gradient all-reduce buckets):
+  `TORCH_DISTRIBUTED_DEBUG=DETAIL` serialises the collectives (~40 % slower) and fixes it; a per-step barrier does not.
+- **LiC:** the Li is atom 126 in `LiC` and 112 in `LiC_simpler`; the negative cell `z` is intentional vacuum.
+
+## Contributing
+
+- Work on a branch, open a pull request, merge it; do not push to `main` directly.
+- Commit messages: one sentence; add a short body only when the reason is not visible in the diff.
+- Figures as PDF.
+- Never commit data, checkpoints, or `CLAUDE.local.md`.
+
+## License
+
+MIT — see [`LICENSE`](LICENSE).

saddleflow-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Ilgar Baghishov, Sung Hoon Jung and Graeme Henkelman
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

saddleflow-0.1.0/PKG-INFO

@@ -0,0 +1,233 @@
+Metadata-Version: 2.5
+Name: saddleflow
+Version: 0.1.0
+Summary: Flow-matching generation of transition-state structures for periodic materials.
+Project-URL: Homepage, https://github.com/SciLM-ai/SaddleFlow
+Project-URL: Repository, https://github.com/SciLM-ai/SaddleFlow
+Project-URL: Model, https://huggingface.co/SciLM/SaddleFlow
+Project-URL: Dataset, https://huggingface.co/datasets/SciLM/MaterialsSaddles
+Author: Ilgar Baghishov, Sung Hoon Jung, Graeme Henkelman
+License: MIT
+License-File: LICENSE
+Requires-Python: <3.14,>=3.11
+Requires-Dist: ase>=3.26.0
+Requires-Dist: fairchem-core<2.24,>=2.19.0
+Requires-Dist: huggingface-hub>=0.27.1
+Requires-Dist: numpy<2.5,>=2.0
+Requires-Dist: safetensors>=0.4
+Requires-Dist: scipy>=1.15.0
+Provides-Extra: dev
+Requires-Dist: pre-commit; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: pytest-cov; extra == 'dev'
+Requires-Dist: ruff==0.5.1; extra == 'dev'
+Provides-Extra: eval
+Requires-Dist: sella>=2.5; extra == 'eval'
+Provides-Extra: extras
+Requires-Dist: pymatgen>=2025.1.9; extra == 'extras'
+Provides-Extra: search
+Requires-Dist: sella>=2.5; extra == 'search'
+Provides-Extra: train
+Requires-Dist: accelerate>=1.0.0; extra == 'train'
+Requires-Dist: lmdb<=1.7.3,>=1.6.2; extra == 'train'
+Requires-Dist: matplotlib>=3.7; extra == 'train'
+Requires-Dist: pyarrow; extra == 'train'
+Requires-Dist: pyyaml; extra == 'train'
+Requires-Dist: tqdm; extra == 'train'
+Requires-Dist: wandb; extra == 'train'
+Description-Content-Type: text/markdown
+
+# SaddleFlow
+
+Generative AI for **transition-state structures** in periodic materials.
+
+SaddleFlow uses flow matching on a fine-tuned [UMA](https://huggingface.co/facebook/UMA) backbone to propose
+saddle points. The released model is unconditional, like a Dimer search: from a single minimum plus noise it
+flows to a nearby first-order saddle, which Sella then converges in 25–61 force calls instead of 196–352
+(MaterialsSaddles paper). A conditional recipe that also takes the product, like NEB, is in
+[`examples/MP20Bat/`](https://github.com/SciLM-ai/SaddleFlow/tree/main/examples/MP20Bat/).
+
+## Use the pretrained model
+
+### Install
+
+```bash
+pip install "saddleflow[search]"
+python -c "import torch; print(torch.cuda.is_available())"   # should print True
+```
+
+**Linux on ARM (aarch64), e.g. NVIDIA GH200:** PyPI has only CPU builds of torch for aarch64, so the line above
+silently installs one. Install the CUDA build first, then SaddleFlow:
+
+```bash
+pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cu129
+pip install "saddleflow[search]"
+```
+
+`torch==2.13.0` matches the newest supported fairchem-core (2.23); with fairchem-core 2.19–2.21 use `torch==2.8.0`
+from the same index.
+
+**Access.** The weights are gated: accept the license at [`SciLM/SaddleFlow`](https://huggingface.co/SciLM/SaddleFlow)
+and, to converge proposals with Sella, at [`facebook/UMA`](https://huggingface.co/facebook/UMA); then run
+`hf auth login`. The first `from_pretrained()` downloads 5.6 GB.
+
+**Hardware.** One GPU. On one GH200 a proposal (20 network evaluations) takes about 1.8 s on its own, and 16–64
+proposals batched together run at 2–3 per second for 40–80-atom cells. GPU memory is about 6 GB for one proposal,
+12–15 GB for 16, and 26–34 GB for 64 (`batch_size` sets how many run together). CPU works but is slow. fairchem
+compiles kernels on first use and caches them under `~/.triton` and `~/.cache`; on clusters with a small home
+quota, set `TRITON_CACHE_DIR`, `TORCHINDUCTOR_CACHE_DIR`, and `XDG_CACHE_HOME` to a scratch directory.
+
+### Run
+
+[`examples/quickstart/`](https://github.com/SciLM-ai/SaddleFlow/tree/main/examples/quickstart/) runs all of this on a real 61-atom battery cell.
+
+```python
+from ase.io import read
+from saddleflow import SaddleFlow, SaddleFlowSearch, find_saddles
+
+sf = SaddleFlow.from_pretrained()                              # SciLM/SaddleFlow, GPU if available
+atoms = read("reactant.traj")
+guesses = sf.propose(atoms, n=8, seed=0, task_name="omat")     # 8 ase.Atoms, one per noise draw
+
+# Converge them with Sella on UMA-S-1.2 and keep the distinct first-order saddles.
+from fairchem.core import FAIRChemCalculator
+from fairchem.core.calculate.pretrained_mlip import get_predict_unit
+calc = FAIRChemCalculator(get_predict_unit("uma-s-1p2", device="cuda"), task_name="omat")
+saddles = find_saddles(atoms, calc, n=8, model=sf, seed=0)
+
+# Or use it like an ASE optimizer: atoms ends on a verified first-order saddle.
+atoms.calc = calc
+SaddleFlowSearch(atoms, sf, seed=0).run(fmax=0.01)
+```
+
+From the command line: `saddleflow propose reactant.traj -o proposals.traj --n 8` and
+`saddleflow search reactant.traj -o saddles.traj --n 8 --task omat`.
+
+Use `task_name="omat"` for bulk materials and `"oc20"` or `"oc22"` for adsorbates on catalyst surfaces. The start
+noise defaults to the training value (0.3 Å) and cannot be raised; each proposal is five flow passes of four
+Euler steps. `examples/MaterialsSaddles/export_pretrained.py` turns a training checkpoint into this format, and
+`check_pretrained.py` checks that it reproduces recorded evaluation endpoints.
+
+## Repository
+
+- `examples/MaterialsSaddles/`: training, evaluation and export of the released model on all four
+  MaterialsSaddles subsets.
+- `examples/MP20Bat/`: the conditional recipe on the `mp20bat` subset.
+- `examples/LiC/`, `examples/LiC_simpler/`: small Li-on-graphene examples for quick experiments.
+
+See [`CLAUDE.md`](https://github.com/SciLM-ai/SaddleFlow/blob/main/CLAUDE.md) for the method and the training recipe.
+
+## Install for training and development
+
+```bash
+conda create -n saddleflow python=3.12 -y
+conda activate saddleflow
+pip install -e '.[train,search]'
+```
+
+`train` adds what the training and evaluation scripts in `examples/` need (accelerate, wandb, matplotlib, pyarrow,
+lmdb); `search` adds Sella for the saddle-search evaluation below, which is not needed to train or to sample. The
+paper model was trained with fairchem-core 2.19 (`pip install 'fairchem-core==2.19.*'` to match it). On aarch64,
+install a CUDA build of torch first, as above.
+
+Use **Python 3.12** — fairchem-core builds cleanly on 3.12 but has had wheel issues on 3.11/3.13 on some platforms.
+
+First UMA load downloads the `uma-s-1p2` checkpoint from HuggingFace and requires a valid `HF_TOKEN` (or a
+`hf auth login`). If pip builds Sella from source with NVIDIA's `nvc` compiler, set `CC=gcc CXX=g++`.
+
+## Sample from your own conditional checkpoint
+
+```python
+import torch
+from ase.io import read
+
+from saddleflow.flow.sampler import sample_saddles
+# ... load_model is provided per-example; see examples/MP20Bat/sample_and_distance_eval.py
+
+# Load a trained checkpoint (architecture is rebuilt from the run's config.json).
+ckpt_dir = "$SCRATCH/SaddleFlow_mp20bat/runs/<RUN>/checkpoint_final"
+loss_module, _config = load_model(ckpt_dir, device="cuda", use_ema=True)
+
+# Read reactant and product geometries (ASE-readable formats).
+R = read("my_reactant.traj")
+P = read("my_product.traj")
+
+# Run deterministic Mode-1 sampling (sigma_inf=0, n_perturbations=1, K Euler steps).
+candidates = sample_saddles(
+    sample_dict_from_atoms(R, P),    # builds {start_pos, partner_un_pos, Z, cell, fixed, ...}
+    loss_module.backbone,
+    loss_module.global_attn,
+    loss_module.velocity_head,
+    sigma_inf=0.0,
+    n_perturbations=1,
+    K=50,
+    device="cuda",
+    partner_pos=...,                  # MIC-unwrapped partner positions
+    # Plus any optional heads / force backbones the run uses; see eval scripts.
+)
+```
+
+For end-to-end working examples, see `examples/MP20Bat/sample_and_distance_eval.py` (a 50-case eval that loads a checkpoint, samples saddles for held-out triplets, and writes parity plots and 4-frame `.traj` files for visualization).
+
+## Training your own model
+
+The production example is `examples/MP20Bat/`:
+
+```bash
+cd $SCRATCH/SaddleFlow_mp20bat            # so SLURM logs land on scratch
+sbatch $WORK/codes/SaddleFlow/examples/MP20Bat/run.sh
+```
+
+This trains on the full `mp20bat` subset (~34,742 triplets) for 60 epochs with 4 nodes × 3 A100s, then auto-runs a 50-case post-training eval. Heavy outputs (checkpoints, eval `.npz`, `.traj` files) all land under `$SCRATCH/SaddleFlow_mp20bat/runs/<TIMESTAMP>/`. See [`examples/MP20Bat/README.md`](https://github.com/SciLM-ai/SaddleFlow/blob/main/examples/MP20Bat/README.md) for the full pipeline (training → full-test-set eval at K=10 → K=10 vs K=50 stability → paper-style parity figures).
+
+**For full-dataset / large-scale training (all four MaterialsSaddles subsets — lemat + oc20 + oc22 + mp20bat = 34 M triplets — at 128+ nodes), use [`examples/MaterialsSaddles/`](https://github.com/SciLM-ai/SaddleFlow/tree/main/examples/MaterialsSaddles/) instead.** Its `run.sh` uses srun-native SPMD (each SLURM task = one rank, env-var DDP), which scales reliably past 128 nodes. The MP20Bat launcher above is the proven 4-node paper config; its `accelerate launch --num_machines=N` pattern is fragile at large scale — see [`CLAUDE.md`](https://github.com/SciLM-ai/SaddleFlow/blob/main/CLAUDE.md) "Training infrastructure" for why.
+
+For a quick sanity-check on small data:
+
+- [`examples/LiC_simpler/`](https://github.com/SciLM-ai/SaddleFlow/tree/main/examples/LiC_simpler/) — **start here.** One Li hop on pristine graphene (1 triplet), ~20 min on one GPU. The site is 6-fold symmetric and the data contains only *one* of the six equivalent saddles, so whether the model recovers all six is a visual pass/fail on the physics. **It passes:** with the backbone unfrozen and per-block time-FiLM, all 48 sampled trajectories end in one of the six saddle petals — a six-fold orbit from one training saddle, with hexatic order **0.98–1.00** over two independent runs. That recipe is the script's default; a frozen backbone reaches only 0.91 (see its [README](https://github.com/SciLM-ai/SaddleFlow/blob/main/examples/LiC_simpler/README.md)).
+- [`examples/LiC/`](https://github.com/SciLM-ai/SaddleFlow/tree/main/examples/LiC/) — defective graphene: 12 training saddles, 167 held-out ones (each Sella-reconverged at fmax 0.005 and verified index-1), and 10 bond midpoints that look like saddles but are not. An unconditioned model carries a Li thrown anywhere on the sheet to the nearest saddle; ~4 h on one GPU. Its [README](https://github.com/SciLM-ai/SaddleFlow/blob/main/examples/LiC/README.md) reproduces the trajectory figure.
+
+Both use an *unconditioned* model rather than the `(R, P)`-conditioned scheme below: on graphene the midpoint of R and P already sits on the saddle, so conditioning would leave the flow nothing to learn.
+
+## Method (in brief)
+
+- **Flow matching** over atomic Cartesian coordinates with straight-line Optimal Transport from `x_0` to `x_1 = saddle`.
+- **Product-conditional anchoring.** `x_0 = (R + P)/2` is the PBC-correct geodesic midpoint of the reactant–product pair, and the velocity head is conditioned on a per-atom Δ_partner = MIC(partner − x_t) at every flow step. This re-parameterises the L2-Bayes-optimal predictor from `E[saddle] − R ≈ midpoint − R` (large) to `E[saddle] − midpoint ≈ 0` (small) — the head only has to predict the residual deviation of the saddle from the midpoint, not rediscover the midpoint itself.
+- **Velocity field** = pretrained **UMA-S-1.2** (fairchem, 6.6M active params) as the equivariant backbone, equivariant time-FiLM injected at each backbone block, and a small `VelocityHead` on top. Selectively-unfrozen UMA blocks at low LR (1e-5) let the backbone adapt its layer-N l ≥ 1 outputs to the velocity-prediction task.
+- **R↔P doubling.** Each triplet contributes both `(start=R, partner=P)` and `(start=P, partner=R)` training pairs, so the head sees the saddle from both directions and the parity isn't a learned bias.
+- **Inference perturbation.** A small Gaussian `ε_inf` on the start position at `t=0` lets multiple integrations from the same `(R, P)` pair land in different angular wedges of the local environment, producing a small ensemble of saddle candidates that can be clustered downstream.
+
+## Evaluation — score against the saddle the prediction actually reaches
+
+RMSD to a stored label answers "is this near *that* saddle", which is not the question. Run a saddle optimiser from the prediction and measure how far it had to move:
+
+```bash
+# 1. Integrate the flow and dump candidates (--ckpt2 runs the two-stage cascade).
+python examples/MaterialsSaddles/dump_predictions.py \
+    --data-glob '<MaterialsSaddles>/mp20bat/*.aselmdb' \
+    --ckpt <stage1>/checkpoint_final --ckpt2 <stage2>/checkpoint_final \
+    --tag cascade --outdir preds --shard $i --nshards $N
+
+# 2. Optimise each candidate to a saddle and verify it is genuinely first-order.
+python examples/MaterialsSaddles/sella_eval.py \
+    --src preds/cascade_00.traj --out sella_00.json --check-index
+```
+
+Then keep records with `nneg == 1`. Reported per model: distance moved (rmsd/maxd), convergence rate, and force calls — a prediction that is close to a saddle *and* cheap to converge is the useful one.
+
+**Use Sella, not the ASE Dimer.** On identical structures the Dimer walked to a different saddle in ~8% of cases, which silently inflates the error tail (in the worst decile, predictions sat 0.17 Å from the label but 2.21 Å from where the Dimer stopped). Sella converged 100% vs 77%, in 44 vs 384 force calls, with 93% of results verified index-1.
+
+## Built on
+
+- [fairchem](https://github.com/facebookresearch/fairchem) (Meta FAIR) — UMA backbone, `AtomicData`, PBC-aware graph building, ASE Calculator integration.
+- [ASE](https://wiki.fysik.dtu.dk/ase/) — atomic simulation environment.
+- [PyTorch](https://pytorch.org) + [HuggingFace `accelerate`](https://huggingface.co/docs/accelerate) — training loop, multi-GPU / multi-node scaling.
+
+## Authors
+
+Ilgar Baghishov, Sung Hoon Jung and Graeme Henkelman (The University of Texas at Austin).
+
+## License
+
+The code is MIT — see [`LICENSE`](https://github.com/SciLM-ai/SaddleFlow/blob/main/LICENSE). The released weights are a fine-tuned version of Meta's UMA-M-1.1
+and are distributed under the FAIR Chemistry License v1 on the Hugging Face model page.