saddleflow 0.1.0__py3-none-any.whl

saddleflow/__init__.py

@@ -0,0 +1,14 @@
+__version__ = "0.1.0"
+
+__all__ = ["SaddleFlow", "SaddleFlowSearch", "find_saddles"]
+
+
+def __getattr__(name):
+    # Lazy, so `import saddleflow` stays light for the training code paths.
+    if name == "SaddleFlow":
+        from .pretrained import SaddleFlow
+        return SaddleFlow
+    if name in ("SaddleFlowSearch", "find_saddles"):
+        from . import search
+        return getattr(search, name)
+    raise AttributeError(f"module 'saddleflow' has no attribute {name!r}")

saddleflow/cli.py

@@ -0,0 +1,59 @@
+"""Command line: `saddleflow propose` and `saddleflow search`.
+
+    saddleflow propose reactant.traj -o proposals.traj --n 8 --seed 0
+    saddleflow search  reactant.traj -o saddles.traj   --n 8 --seed 0 --task omat
+
+`propose` writes the flow endpoints; `search` also converges them with Sella on UMA-S-1.2 (fairchem's
+calculator, gated access to facebook/UMA) and keeps the distinct first-order saddles.
+"""
+
+from __future__ import annotations
+
+import argparse
+
+from ase.io import read, write
+
+
+def main(argv=None):
+    p = argparse.ArgumentParser(prog="saddleflow", description=__doc__.split("\n")[0])
+    sub = p.add_subparsers(dest="cmd", required=True)
+    for name in ("propose", "search"):
+        s = sub.add_parser(name)
+        s.add_argument("input", help="ASE-readable structure file (every frame is used)")
+        s.add_argument("-o", "--output", required=True)
+        s.add_argument("--n", type=int, default=8, help="noise draws per structure")
+        s.add_argument("--seed", type=int, default=None)
+        s.add_argument("--sigma", type=float, default=None, help="start noise (A), at most the training 0.3")
+        s.add_argument("--task", default=None, help="omat (bulk), oc20 or oc22 (surfaces); default omat")
+        s.add_argument("--model", default="SciLM/SaddleFlow", help="Hub repository or local directory")
+        s.add_argument("--device", default=None)
+        if name == "search":
+            s.add_argument("--uma", default="uma-s-1p2", help="UMA model for the Sella refinement")
+            s.add_argument("--fmax", type=float, default=0.01)
+            s.add_argument("--steps", type=int, default=300)
+    a = p.parse_args(argv)
+
+    from .pretrained import SaddleFlow
+    model = SaddleFlow.from_pretrained(a.model, device=a.device)
+    frames = read(a.input, index=":")
+    out = []
+    if a.cmd == "propose":
+        for props in model.propose(frames, n=a.n, seed=a.seed, sigma=a.sigma, task_name=a.task):
+            out.extend(props)
+    else:
+        from fairchem.core import FAIRChemCalculator
+        from fairchem.core.calculate.pretrained_mlip import get_predict_unit
+
+        from .search import find_saddles
+        unit = get_predict_unit(a.uma, device=model.device.type)          # fairchem takes "cuda" or "cpu"
+        for frame in frames:
+            task = a.task or frame.info.get("task_name") or "omat"
+            calc = FAIRChemCalculator(unit, task_name=task)
+            out.extend(find_saddles(frame, calc, n=a.n, model=model, seed=a.seed, fmax=a.fmax, steps=a.steps,
+                                    sigma=a.sigma, task_name=task))
+    write(a.output, out)
+    print(f"wrote {len(out)} structures to {a.output}")
+
+
+if __name__ == "__main__":
+    main()

saddleflow/data/__init__.py

@@ -0,0 +1,31 @@
+from .core import (
+    atoms_to_sample_dict,
+    extract_fixed_mask,
+    iter_triplets_from_traj_paths,
+    load_validated_triplets,
+    mic_unwrap,
+    resolve_paths,
+    triplet_to_pair_records,
+    validate_triplet,
+)
+from .db_dataset import AseDbSaddleDataset
+from .materials_saddles_dataset import MaterialsSaddlesDataset
+from .traj_dataset import TrajTripletDataset
+from .transforms import gaussian_perturbation, mic_displacement, wrap_positions
+
+__all__ = [
+    "AseDbSaddleDataset",
+    "MaterialsSaddlesDataset",
+    "TrajTripletDataset",
+    "atoms_to_sample_dict",
+    "extract_fixed_mask",
+    "gaussian_perturbation",
+    "iter_triplets_from_traj_paths",
+    "load_validated_triplets",
+    "mic_displacement",
+    "mic_unwrap",
+    "resolve_paths",
+    "triplet_to_pair_records",
+    "validate_triplet",
+    "wrap_positions",
+]

saddleflow/data/convert_to_db.py

@@ -0,0 +1,124 @@
+"""
+(Backend B) Convert .traj triplets [R, S, P, R, S, P, ...] into a single ASE-DB.
+
+Output backend is chosen by file extension:
+    *.db / *.sqlite   → ASE's default SQLite backend (good for debugging, moderate scale)
+    *.aselmdb         → LMDB-backed ASE-DB (fairchem's production format; fast at 30M+ scale)
+
+Each triplet writes two rows, one per (start, saddle) pair, keyed by role ∈ {"R2S", "P2S"}.
+Saddle positions are stored minimum-image-unwrapped relative to their start.
+
+Usage:
+    python -m saddleflow.data.convert_to_db \\
+        --traj 'data/raw/*.traj' \\
+        --out  data/processed/li_c.aselmdb \\
+        --default-task-name omat
+"""
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+from ase import Atoms
+from ase.constraints import FixAtoms
+from ase.db import connect
+from tqdm import tqdm
+
+from .core import iter_triplets_from_traj_paths, resolve_paths, triplet_to_pair_records
+
+
+def _start_atoms(record: dict) -> Atoms:
+    atoms = Atoms(
+        positions=record["start_pos"],
+        numbers=record["Z"],
+        cell=record["cell"],
+        pbc=True,
+    )
+    fixed_idx = np.where(record["fixed"])[0]
+    if len(fixed_idx):
+        atoms.set_constraint(FixAtoms(indices=fixed_idx.tolist()))
+    return atoms
+
+
+def convert(
+    traj_paths,
+    out_path: str,
+    default_task_name: str = "omat",
+    default_charge: int = 0,
+    default_spin: int = 0,
+    progress: bool = True,
+) -> dict:
+    paths = resolve_paths(traj_paths)
+    out = Path(out_path)
+    out.parent.mkdir(parents=True, exist_ok=True)
+
+    db = connect(str(out))
+    deltas: list[float] = []
+    triplet_id = 0
+    it = iter_triplets_from_traj_paths(paths)
+    if progress:
+        it = tqdm(it, desc=f"converting → {out.name}")
+
+    for R, S, P in it:
+        for r in triplet_to_pair_records(
+            R, S, P, triplet_id,
+            default_task_name, default_charge, default_spin,
+        ):
+            db.write(
+                _start_atoms(r),
+                key_value_pairs={
+                    "role": r["role"],
+                    "task_name": r["task_name"],
+                    "triplet_id": int(r["triplet_id"]),
+                    "charge": int(r["charge"]),
+                    "spin": int(r["spin"]),
+                    "delta_norm": float(r["delta_norm"]),
+                },
+                data={
+                    "saddle_un_pos": r["saddle_un_pos"],
+                    "partner_un_pos": r["partner_un_pos"],
+                    "metadata": r["metadata"],
+                },
+            )
+            deltas.append(float(r["delta_norm"]))
+        triplet_id += 1
+
+    stats = {
+        "num_triplets": triplet_id,
+        "num_records": 2 * triplet_id,
+        "delta_norm_mean": float(np.mean(deltas)) if deltas else 0.0,
+        "delta_norm_std": float(np.std(deltas)) if deltas else 0.0,
+        "source_paths": paths,
+        "out_path": str(out),
+    }
+    Path(str(out) + ".stats.json").write_text(json.dumps(stats, indent=2))
+    return stats
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--traj", required=True,
+                   help="Glob, directory, single path, or comma-separated list of .traj paths")
+    p.add_argument("--out", required=True,
+                   help="Output ASE-DB path (*.db, *.sqlite, or *.aselmdb)")
+    p.add_argument("--default-task-name", default="omat")
+    p.add_argument("--default-charge", type=int, default=0)
+    p.add_argument("--default-spin", type=int, default=0)
+    p.add_argument("--no-progress", action="store_true")
+    args = p.parse_args()
+
+    traj_arg = args.traj.split(",") if "," in args.traj else args.traj
+    stats = convert(
+        traj_arg, args.out,
+        default_task_name=args.default_task_name,
+        default_charge=args.default_charge,
+        default_spin=args.default_spin,
+        progress=not args.no_progress,
+    )
+    print(f"Wrote {stats['num_records']} records from {stats['num_triplets']} triplets → {args.out}")
+    print(f"<||Δ||> = {stats['delta_norm_mean']:.4f} Å   (std {stats['delta_norm_std']:.4f} Å)")
+
+
+if __name__ == "__main__":
+    main()

saddleflow/data/core.py

@@ -0,0 +1,200 @@
+"""
+Shared helpers for both data backends.
+
+A *triplet* is a contiguous (reactant, saddle, product) block of ASE Atoms, as
+found in the source .traj files (with `atoms.info['side']` ∈ {-1, 0, 1}).
+
+A *pair record* is one training example: a start structure (either R or P) and
+its MIC-unwrapped saddle. Every triplet produces two pair records (R→S and
+P→S) by microscopic reversibility — the saddle is the same saddle whether
+approached from R or P; only the MIC image choice differs.
+"""
+
+import os
+from glob import glob as _glob
+from typing import Iterator
+
+import numpy as np
+import torch
+from ase import Atoms
+from ase.constraints import FixAtoms
+from ase.io import Trajectory
+
+SIDE_REACTANT = -1
+SIDE_SADDLE = 0
+SIDE_PRODUCT = 1
+
+
+def mic_unwrap(start_pos: np.ndarray, target_pos: np.ndarray, cell: np.ndarray) -> np.ndarray:
+    """Return `target_pos` minimum-image-unwrapped relative to `start_pos`.
+
+    Result may lie outside the unit cell — that is intentional; the flow
+    interpolation is done in unwrapped space and wrapped only for UMA forward.
+    """
+    cell = np.asarray(cell, dtype=np.float64)
+    delta = np.asarray(target_pos, dtype=np.float64) - np.asarray(start_pos, dtype=np.float64)
+    frac = delta @ np.linalg.inv(cell)
+    frac -= np.round(frac)
+    return np.asarray(start_pos, dtype=np.float64) + frac @ cell
+
+
+def extract_fixed_mask(atoms: Atoms) -> np.ndarray:
+    mask = np.zeros(len(atoms), dtype=bool)
+    for c in getattr(atoms, "constraints", []) or []:
+        if isinstance(c, FixAtoms):
+            mask[c.index] = True
+    return mask
+
+
+def validate_triplet(R: Atoms, S: Atoms, P: Atoms) -> None:
+    assert np.array_equal(R.numbers, S.numbers) and np.array_equal(S.numbers, P.numbers), \
+        "atom ordering differs across triplet frames"
+    assert np.allclose(np.asarray(R.cell), np.asarray(S.cell)) \
+        and np.allclose(np.asarray(S.cell), np.asarray(P.cell)), \
+        "cell differs across triplet frames"
+    assert all(R.pbc) and all(S.pbc) and all(P.pbc), \
+        "triplet is not 3D-periodic"
+    fR, fS, fP = extract_fixed_mask(R), extract_fixed_mask(S), extract_fixed_mask(P)
+    assert np.array_equal(fR, fS) and np.array_equal(fS, fP), \
+        "FixAtoms mask differs across triplet frames"
+    sides = [R.info.get("side"), S.info.get("side"), P.info.get("side")]
+    if all(s is not None for s in sides):
+        assert sides == [SIDE_REACTANT, SIDE_SADDLE, SIDE_PRODUCT], \
+            f"unexpected info['side'] ordering {sides}, want [-1, 0, 1]"
+
+
+def _sanitize_info(info: dict) -> dict:
+    """Drop non-serializable entries from an ASE info dict so ASE-DB can store it."""
+    out = {}
+    for k, v in info.items():
+        if isinstance(v, (str, int, float, bool, type(None))):
+            out[k] = v
+        elif isinstance(v, np.ndarray):
+            out[k] = v
+        elif isinstance(v, (list, tuple)):
+            try:
+                out[k] = np.asarray(v)
+            except Exception:
+                continue
+        elif isinstance(v, dict):
+            out[k] = _sanitize_info(v)
+    return out
+
+
+def triplet_to_pair_records(
+    R: Atoms,
+    S: Atoms,
+    P: Atoms,
+    triplet_id: int,
+    default_task_name: str = "omat",
+    default_charge: int = 0,
+    default_spin: int = 0,
+) -> list[dict]:
+    validate_triplet(R, S, P)
+    Z = R.numbers.astype(np.int32)
+    cell = np.asarray(R.cell[:], dtype=np.float32)
+    fixed = extract_fixed_mask(R)
+    task_name = R.info.get("task_name", default_task_name)
+    charge = int(R.info.get("charge", default_charge))
+    spin = int(R.info.get("spin", default_spin))
+    metadata = {
+        "reactant_info": _sanitize_info(R.info),
+        "saddle_info": _sanitize_info(S.info),
+        "product_info": _sanitize_info(P.info),
+    }
+
+    S_un_from_R = mic_unwrap(R.positions, S.positions, cell).astype(np.float32)
+    S_un_from_P = mic_unwrap(P.positions, S.positions, cell).astype(np.float32)
+
+    base = dict(Z=Z, cell=cell, fixed=fixed, task_name=task_name,
+                charge=charge, spin=spin, metadata=metadata,
+                triplet_id=int(triplet_id))
+
+    # Mode-1 partner: the OTHER local minimum (the one we are flowing AWAY from).
+    # Stored MIC-unwrapped relative to start so it lives in the same frame as
+    # `saddle_un_pos`. Frozen atoms have identical positions in R and P, so on
+    # those atoms `partner_un_pos == start_pos`, i.e. delta_partner = 0 — the
+    # head sees no spurious frozen-atom signal.
+    P_un_from_R = mic_unwrap(R.positions, P.positions, cell).astype(np.float32)
+    R_un_from_P = mic_unwrap(P.positions, R.positions, cell).astype(np.float32)
+
+    return [
+        dict(base,
+             start_pos=R.positions.astype(np.float32),
+             saddle_un_pos=S_un_from_R,
+             partner_un_pos=P_un_from_R,
+             delta_norm=np.float32(np.linalg.norm(S_un_from_R - R.positions)),
+             role="R2S"),
+        dict(base,
+             start_pos=P.positions.astype(np.float32),
+             saddle_un_pos=S_un_from_P,
+             partner_un_pos=R_un_from_P,
+             delta_norm=np.float32(np.linalg.norm(S_un_from_P - P.positions)),
+             role="P2S"),
+    ]
+
+
+def iter_triplets_from_traj_paths(paths: list[str]) -> Iterator[tuple[Atoms, Atoms, Atoms]]:
+    for path in paths:
+        t = Trajectory(path, "r")
+        n = len(t)
+        assert n % 3 == 0, f"{path}: {n} frames, not a multiple of 3"
+        for i in range(0, n, 3):
+            yield t[i], t[i + 1], t[i + 2]
+        t.close()
+
+
+def resolve_paths(pattern_or_list) -> list[str]:
+    """Accept a single path, a glob pattern, a directory, or a list of any of those."""
+    if isinstance(pattern_or_list, (list, tuple)):
+        out = []
+        for p in pattern_or_list:
+            out.extend(resolve_paths(p))
+        return sorted(set(out))
+    p = str(pattern_or_list)
+    if any(c in p for c in "*?["):
+        return sorted(_glob(p))
+    if os.path.isdir(p):
+        return sorted(_glob(os.path.join(p, "*.traj")))
+    if os.path.isfile(p):
+        return [p]
+    raise FileNotFoundError(f"no .traj files matched {pattern_or_list!r}")
+
+
+def load_validated_triplets(paths) -> list[tuple[Atoms, Atoms, Atoms]]:
+    """Read one or more .traj files and return the validated list of (R, S, P) tuples.
+
+    Convenience wrapper over `iter_triplets_from_traj_paths` + `validate_triplet`
+    for callers that need all triplets up-front (e.g. evaluation scripts that
+    group by reactant).
+    """
+    triplets = list(iter_triplets_from_traj_paths(resolve_paths(paths)))
+    for R, S, P in triplets:
+        validate_triplet(R, S, P)
+    return triplets
+
+
+def atoms_to_sample_dict(
+    atoms: Atoms,
+    default_task_name: str = "omat",
+    default_charge: int = 0,
+    default_spin: int = 0,
+) -> dict:
+    """Build an inference-ready sample dict from a single ASE `Atoms`.
+
+    Matches the subset of fields that `sample_saddles` and
+    `FlowMatchingLoss.build_atomic_data` consume (no `saddle_un_pos` /
+    `delta_norm` / `metadata`, which are training-only). `task_name`,
+    `charge`, and `spin` are pulled from `atoms.info` if present, else from
+    the `default_*` arguments.
+    """
+    fixed = extract_fixed_mask(atoms)
+    return {
+        "start_pos": torch.tensor(atoms.get_positions(), dtype=torch.float32),
+        "Z": torch.tensor(atoms.numbers, dtype=torch.long),
+        "cell": torch.tensor(np.asarray(atoms.cell[:]), dtype=torch.float32),
+        "fixed": torch.tensor(fixed, dtype=torch.bool),
+        "task_name": atoms.info.get("task_name", default_task_name),
+        "charge": int(atoms.info.get("charge", default_charge)),
+        "spin": int(atoms.info.get("spin", default_spin)),
+    }

saddleflow/data/db_dataset.py

@@ -0,0 +1,80 @@
+"""
+(Backend B) Read (start, saddle) pair records from an ASE-DB written by
+`convert_to_db`. Yields the same dict schema as `TrajTripletDataset` so
+downstream training/inference code is backend-agnostic.
+"""
+
+import json
+import os
+from pathlib import Path
+
+import numpy as np
+import torch
+from ase.db import connect
+from torch.utils.data import Dataset
+
+
+class AseDbSaddleDataset(Dataset):
+    def __init__(self, path: str, select_args: dict | None = None):
+        self._path = str(path)
+        self._select_args = dict(select_args or {})
+
+        db = connect(self._path)
+        rows = list(db.select(**self._select_args))
+        self._ids = [row.id for row in rows]
+        self.delta_norm_mean = float(np.mean([row.delta_norm for row in rows])) if rows else 0.0
+        del rows
+        db.close()
+
+        stats_file = Path(self._path + ".stats.json")
+        if stats_file.is_file():
+            stats = json.loads(stats_file.read_text())
+            self.delta_norm_mean = float(stats.get("delta_norm_mean", self.delta_norm_mean))
+
+        self._db_cache: dict[int, "object"] = {}
+
+    def _get_db(self):
+        key = os.getpid()
+        d = self._db_cache.get(key)
+        if d is None:
+            d = connect(self._path)
+            self._db_cache[key] = d
+        return d
+
+    def __len__(self) -> int:
+        return len(self._ids)
+
+    def __getitem__(self, idx: int) -> dict:
+        row = self._get_db().get(id=self._ids[idx])
+        atoms = row.toatoms()
+        kvp = row.key_value_pairs
+        data = row.data
+
+        fixed = np.zeros(len(atoms), dtype=bool)
+        for c in atoms.constraints or []:
+            if type(c).__name__ == "FixAtoms":
+                fixed[c.index] = True
+
+        # `partner_un_pos` is required by the released training scheme.
+        # Older DBs without it fall back to a copy of `start_pos`
+        # (delta_partner = 0), but those DBs cannot be used for training.
+        if "partner_un_pos" in data:
+            partner_un = torch.tensor(np.asarray(data["partner_un_pos"]), dtype=torch.float32)
+        else:
+            partner_un = torch.tensor(atoms.positions, dtype=torch.float32)
+
+        return {
+            "start_pos": torch.tensor(atoms.positions, dtype=torch.float32),
+            "saddle_un_pos": torch.tensor(np.asarray(data["saddle_un_pos"]), dtype=torch.float32),
+            "partner_un_pos": partner_un,
+            "Z": torch.tensor(atoms.numbers, dtype=torch.long),
+            "cell": torch.tensor(np.asarray(atoms.cell), dtype=torch.float32),
+            "fixed": torch.tensor(fixed, dtype=torch.bool),
+            "task_name": str(kvp["task_name"]),
+            "charge": int(kvp["charge"]),
+            "spin": int(kvp["spin"]),
+            "delta_norm": torch.tensor(float(kvp["delta_norm"])),
+            "role": str(kvp["role"]),
+            "triplet_id": int(kvp["triplet_id"]),
+            "metadata": data.get("metadata", {}),
+        }