# /// script # requires-python = ">=3.11,<3.14" # dependencies = [ # "mepd[gxtb,aam] @ git+https://github.com/atomsforhumanity/mepd.git@da31dcb295a067b433a890a08542cd7bc2cabc03", # "h5py>=3.10", # "matplotlib>=3.8", # ] # [tool.uv] # # mepd's own locked versions at the pinned commit (its uv.lock, via `uv export`): the # # environment mepd is tested in. Regenerate with update_pin.sh when the pin moves. # constraint-dependencies = [ # "absl-py==2.5.0 ; python_full_version >= '3.11'", # "aimnet==0.2.0 ; python_full_version >= '3.11'", # "annotated-doc==0.0.5", # "annotated-types==0.8.0", # "antlr4-python3-runtime==4.9.3 ; python_full_version >= '3.11'", # "anyio==4.15.1", # "asciitree==0.3.3", # "ase==3.29.0", # "ase-db-backends==0.11.0 ; python_full_version >= '3.11'", # "asttokens==3.0.2", # "backoff==2.2.1 ; python_full_version >= '3.11'", # "certifi==2026.7.22", # "cffi==2.1.1", # "charset-normalizer==3.5.2 ; python_full_version >= '3.11'", # "chemcloud==0.17.0", # "click==8.5.0", # "cloudpickle==3.1.2 ; python_full_version >= '3.11'", # "clusterscope==0.0.18 ; python_full_version >= '3.11'", # "colorama==0.4.6 ; sys_platform == 'win32'", # "contourpy==1.3.2 ; python_full_version < '3.11'", # "contourpy==1.3.3 ; python_full_version == '3.11.*'", # "contourpy==1.4.0 ; python_full_version >= '3.12'", # "cryptography==50.0.2 ; python_full_version >= '3.11'", # "cuda-bindings==13.4.3 ; sys_platform == 'linux'", # "cuda-pathfinder==1.8.3 ; sys_platform == 'linux'", # "cuda-toolkit==13.0.3.0 ; sys_platform == 'linux'", # "cycler==0.12.1", # "decorator==5.3.1 ; python_full_version < '3.11'", # "deprecated==1.3.1 ; python_full_version == '3.11.*'", # "deprecated==3.0.0 ; python_full_version >= '3.12'", # "e3nn==0.6.0 ; python_full_version >= '3.11'", # "eval-type-backport==0.4.0", # "exceptiongroup==1.3.1 ; python_full_version < '3.11'", # "executing==2.2.1", # "fairchem-core==2.23.0 ; python_full_version >= '3.11'", # "fastapi==0.143.0", # "fasteners==0.20 ; sys_platform != 'emscripten'", # "filelock==4.0.12", # "flatbuffers==25.12.19", # "flexcache==0.3", # "flexparser==0.4", # "fonttools==4.65.0 ; python_full_version < '3.11'", # "fonttools==4.66.1 ; python_full_version >= '3.11'", # "fsspec==2026.9.0", # "geometric==1.1.1", # "googleapis-common-protos==1.75.5 ; python_full_version >= '3.11'", # "grpcio==1.84.0 ; python_full_version >= '3.11'", # "h11==0.16.0", # "h5py==3.16.0", # "hf-xet==1.7.0 ; platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'", # "httpcore==1.0.9", # "httpx==0.28.1", # "huggingface-hub==1.33.0", # "hydra-core==1.3.7 ; python_full_version >= '3.11'", # "idna==3.20", # "ipython==8.39.0 ; python_full_version < '3.11'", # "ipython==9.17.1 ; python_full_version >= '3.11'", # "ipython-pygments-lexers==1.1.1 ; python_full_version >= '3.11'", # "jax==0.6.2 ; python_full_version < '3.11'", # "jax==0.10.2 ; python_full_version == '3.11.*'", # "jax==0.11.2 ; python_full_version >= '3.12'", # "jaxlib==0.6.2 ; python_full_version < '3.11'", # "jaxlib==0.10.2 ; python_full_version == '3.11.*'", # "jaxlib==0.11.2 ; python_full_version >= '3.12'", # "jedi==0.20.0", # "jinja2==3.1.6", # "kiwisolver==1.5.1", # "legacy-cgi==2.6.4 ; python_full_version >= '3.13'", # "llvmlite==0.50.0", # "lmdb==1.7.3 ; python_full_version >= '3.11'", # "markdown==3.10.3 ; python_full_version < '3.11'", # "markdown==3.11 ; python_full_version >= '3.11'", # "markdown-it-py==4.2.0", # "markupsafe==3.0.4", # "matplotlib==3.10.9 ; python_full_version < '3.11'", # "matplotlib==3.11.2 ; python_full_version >= '3.11'", # "matplotlib-inline==0.2.2", # "mdurl==0.1.2", # "ml-dtypes==0.6.0", # "monty==2026.7.16 ; python_full_version >= '3.11'", # "mpmath==1.3.0", # "mypy-extensions==1.1.0 ; python_full_version >= '3.11'", # "networkx==3.4.2 ; python_full_version < '3.11'", # "networkx==3.6.1 ; python_full_version == '3.11.*'", # "networkx==3.7 ; python_full_version >= '3.12'", # "ninja==1.13.2", # "numba==0.68.0", # "numcodecs==0.13.1 ; python_full_version < '3.11'", # "numcodecs==0.15.1 ; python_full_version >= '3.11'", # "numpy==2.2.6 ; python_full_version < '3.11'", # "numpy==2.4.6 ; python_full_version == '3.11.*'", # "numpy==2.5.3 ; python_full_version >= '3.12'", # "nvalchemi-toolkit-ops==0.4.1 ; python_full_version >= '3.11'", # "nvidia-cublas==13.1.1.3 ; sys_platform == 'linux'", # "nvidia-cuda-cupti==13.0.85 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')", # "nvidia-cuda-nvrtc==13.0.88 ; sys_platform == 'linux'", # "nvidia-cuda-runtime==13.0.96 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')", # "nvidia-cudnn-cu13==9.20.0.48 ; sys_platform == 'linux'", # "nvidia-cufft==12.0.0.61 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')", # "nvidia-cufile==1.15.1.6 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')", # "nvidia-curand==10.4.0.35 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')", # "nvidia-cusolver==12.0.4.66 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')", # "nvidia-cusparse==12.6.3.3 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')", # "nvidia-cusparselt-cu13==0.8.1 ; sys_platform == 'linux'", # "nvidia-nccl-cu13==2.29.7 ; sys_platform == 'linux'", # "nvidia-nvjitlink==13.4.92 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')", # "nvidia-nvshmem-cu13==3.4.5 ; sys_platform == 'linux'", # "nvidia-nvtx==13.0.85 ; (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')", # "omegaconf==2.3.1 ; python_full_version >= '3.11'", # "onnxruntime==1.24.3 ; python_full_version < '3.11'", # "onnxruntime==1.30.0 ; python_full_version >= '3.11'", # "openbabel-wheel==3.1.1.23", # "openmm==8.6.1", # "opentelemetry-api==1.45.1", # "opentelemetry-exporter-http-transport==0.66b1 ; python_full_version >= '3.11'", # "opentelemetry-exporter-otlp-common==0.66b1 ; python_full_version >= '3.11'", # "opentelemetry-exporter-otlp-proto-common==1.45.1 ; python_full_version >= '3.11'", # "opentelemetry-exporter-otlp-proto-http==1.45.1 ; python_full_version >= '3.11'", # "opentelemetry-proto==1.45.1 ; python_full_version >= '3.11'", # "opentelemetry-sdk==1.45.1 ; python_full_version >= '3.11'", # "opentelemetry-semantic-conventions==0.66b1 ; python_full_version >= '3.11'", # "opt-einsum==3.4.0", # "opt-einsum-fx==0.1.4 ; python_full_version >= '3.11'", # "orjson==3.13.0 ; python_full_version >= '3.11'", # "packaging==26.3", # "pandas==2.3.3 ; python_full_version < '3.11'", # "pandas==3.0.6 ; python_full_version >= '3.11'", # "parso==0.8.7", # "pdbfixer==1.12.0", # "pexpect==4.9.0 ; sys_platform != 'emscripten' and sys_platform != 'win32'", # "pillow==12.3.0", # "pint==0.24.4 ; python_full_version < '3.11'", # "pint==0.25.3 ; python_full_version == '3.11.*'", # "pint==0.26.1 ; python_full_version >= '3.12'", # "platformdirs==4.12.4", # "prompt-toolkit==3.0.53", # "protobuf==7.36.2", # "psutil==7.2.2", # "psycopg2-binary==2.9.13 ; python_full_version >= '3.11'", # "ptyprocess==0.7.0 ; sys_platform != 'emscripten' and sys_platform != 'win32'", # "pure-eval==0.2.4", # "py-cpuinfo==9.0.0", # "py3dmol==2.5.5", # "pyarrow==25.0.1", # "pycparser==3.1 ; implementation_name != 'PyPy'", # "pydantic==2.14.0", # "pydantic-core==2.50.0", # "pydantic-settings==2.15.0", # "pygments==2.21.0", # "pymysql==1.2.3 ; python_full_version >= '3.11'", # "pynauty==2.8.8.1", # "pyparsing==3.3.3", # "pyre-extensions==0.0.32 ; python_full_version >= '3.11'", # "pysmiles==2.1.1", # "python-dateutil==2.9.0.post0", # "python-dotenv==1.2.4", # "python-multipart==0.0.32", # "pytz==2026.5 ; python_full_version < '3.11'", # "pyyaml==6.0.3", # "qccodec==0.12.0", # "qccompute==0.14.0", # "qcconst==0.2.2", # "qcdata==0.19.0", # "qcelemental==0.51.2", # "qcengine==0.51.0", # "qcinf==0.4.1", # "qcref==0.3.0", # "rdchiral==1.1.0", # "rdkit==2026.3.6", # "regex==2026.9.29", # "requests==2.34.2 ; python_full_version >= '3.11'", # "rich==15.0.0", # "ruamel-yaml==0.19.1 ; python_full_version >= '3.11'", # "rustworkx==0.18.1", # "safetensors==0.8.0", # "scipy==1.15.3 ; python_full_version < '3.11'", # "scipy==1.17.1 ; python_full_version == '3.11.*'", # "scipy==1.18.1 ; python_full_version >= '3.12'", # "sella==2.6.0", # "setuptools==80.10.2", # "shellingham==1.5.4", # "six==1.17.0", # "stack-data==0.6.3", # "starlette==1.7.0", # "submitit==1.5.4 ; python_full_version >= '3.11'", # "sympy==1.14.0", # "tabulate==0.10.0 ; python_full_version >= '3.11'", # "tblite==0.7.0", # "tensorboard==2.21.0 ; python_full_version >= '3.11'", # "tensorboard-data-server==0.7.2 ; python_full_version >= '3.11'", # "timeout-timer==0.2.0", # "tokenizers==0.23.2", # "toml==0.10.2", # "tomli==2.5.0", # "tomli-w==1.2.0", # "torch==2.13.0", # "torchani==2.9.0", # "torchtnt==0.2.4 ; python_full_version >= '3.11'", # "tqdm==4.70.1", # "traitlets==5.16.1", # "transformers==5.19.0", # "triton==3.7.1 ; sys_platform == 'linux'", # "typer==0.27.3", # "typing-extensions==4.16.0", # "typing-inspect==0.9.0 ; python_full_version >= '3.11'", # "typing-inspection==0.4.4", # "tzdata==2026.5 ; python_full_version < '3.11' or sys_platform == 'emscripten' or sys_platform == 'win32'", # "urllib3==2.8.0 ; python_full_version >= '3.11'", # "uvicorn==0.54.0", # "vina==1.2.7", # "wandb==0.30.0 ; python_full_version >= '3.11'", # "warp-lang==1.18.0 ; python_full_version >= '3.11'", # "wcwidth==0.9.2", # "websockets==17.2 ; python_full_version >= '3.11'", # "werkzeug==3.1.9 ; python_full_version >= '3.11'", # "wheel==0.48.0", # "wrapt==2.5.0 ; python_full_version >= '3.11'", # "xxhash==4.0.1 ; python_full_version >= '3.11'", # "zarr==2.18.3 ; python_full_version < '3.11'", # "zarr==2.18.7 ; python_full_version >= '3.11'", # ] # /// """Benchmark a reaction dataset with mepd: re-converge its transition states at a cheap level, keep the consistent set, search for lower channels, and rescore what is lower with DFT. uv run benchmark.py import-manifest WS reactions.csv # or: import-reaction-qm WS B3LYPD3_TZVP.h5 uv run benchmark.py screen WS # TS opt + IRC of every reference TS uv run benchmark.py search WS --profile gsm.toml # mepd.channels on the kept set uv run benchmark.py verify WS --profile dft.toml # DFT on the reactions with a lower route uv run benchmark.py report WS # table.csv, summary.json, figures Every stage works in the workspace folder WS, records each reaction's result as it finishes (WS//results.jsonl) and skips finished ones when run again, so a stage can be stopped and resumed. See README.md. """ from __future__ import annotations import csv import json import os import shutil import subprocess import sys import tarfile import time from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path from typing import Optional import typer app = typer.Typer(add_completion=False, help=__doc__.split("\n\n")[0], no_args_is_help=True) SCRIPT = Path(__file__).resolve() STAGES = ("screen", "search", "verify") # --- workspace ------------------------------------------------------------------------------------- FIELDS = ["id", "source", "charge", "multiplicity", "n_atoms", "n_heavy", "start", "end", "ts", "reference_barrier_kcal", "smiles"] def _reactions(ws: Path) -> list[dict]: fp = ws / "reactions.csv" if not fp.exists(): raise typer.BadParameter(f"{ws} has no reactions.csv: import a dataset first (import-manifest / " "import-reaction-qm / import-transition1x / import-transition1x-tmc).") return list(csv.DictReader(open(fp, newline=""))) def _write_reactions(ws: Path, rows: list[dict]) -> None: ws.mkdir(parents=True, exist_ok=True) fp = ws / "reactions.csv" old = {r["id"]: r for r in csv.DictReader(open(fp, newline=""))} if fp.exists() else {} old.update({r["id"]: r for r in rows}) with open(fp, "w", newline="") as fh: w = csv.DictWriter(fh, fieldnames=FIELDS) w.writeheader() w.writerows({k: r.get(k, "") for k in FIELDS} for r in old.values()) def _results(ws: Path, stage: str) -> dict[str, dict]: fp = ws / stage / "results.jsonl" out = {} if fp.exists(): for line in fp.read_text().splitlines(): if line.strip(): rec = json.loads(line) out[rec["id"]] = rec return out def _settings(ws: Path, stage: str, given: Optional[dict] = None) -> dict: """A stage's settings, saved on its first run; a rerun must not change them (results would mix settings) unless the stage folder is cleared first.""" fp = ws / stage / "settings.json" if given is None: return json.loads(fp.read_text()) if fp.exists(): saved = json.loads(fp.read_text()) if saved != given: changed = sorted(k for k in set(saved) | set(given) if saved.get(k) != given.get(k)) raise typer.BadParameter( f"{stage} already ran in {ws} with other settings ({', '.join(changed)}). Use the same ones to " f"resume, or move {ws / stage} aside to start the stage over.") return saved fp.parent.mkdir(parents=True, exist_ok=True) fp.write_text(json.dumps(given, indent=1) + "\n") try: # provenance of the stage: code, versions, binaries (mepd.provenance) from mepd.provenance import run_manifest, write_manifest profile = given.get("profile") write_manifest(fp.parent, run_manifest(sys.argv, Path(profile) if profile else None, extra={"stage": stage, "settings": given})) except Exception as exc: # provenance is a record, never a reason to stop typer.echo(f"(provenance not recorded: {type(exc).__name__}: {exc})") return given def _engine(settings: dict): import mepd if settings.get("profile"): return mepd.engine_from_profile(settings["profile"]) return mepd.engine(settings.get("engine") or "gxtb") def _run_all(ws: Path, stage: str, ids: list[str], jobs: int, timeout: float) -> None: """Run `stage` on each id in its own process (a timeout kills it; one crash stops nothing else), `jobs` at a time, appending each result as it lands.""" done = _results(ws, stage) todo = [i for i in ids if i not in done] typer.echo(f"{stage}: {len(ids)} reaction(s), {len(ids) - len(todo)} done before, {len(todo)} to run " f"({jobs} at a time).") out = ws / stage / "results.jsonl" out.parent.mkdir(parents=True, exist_ok=True) logs = ws / stage / "logs" logs.mkdir(exist_ok=True) def one(rid: str) -> dict: t0 = time.time() env = dict(os.environ, OMP_NUM_THREADS=os.environ.get("OMP_NUM_THREADS", "1")) with open(logs / f"{rid}.log", "w") as log: try: proc = subprocess.run([sys.executable, str(SCRIPT), "_one", stage, str(ws), rid], stdout=log, stderr=subprocess.STDOUT, timeout=timeout or None, env=env) except subprocess.TimeoutExpired: return {"id": rid, "status": "timeout", "seconds": round(time.time() - t0, 1)} result_fp = ws / stage / rid / "result.json" if proc.returncode == 0 and result_fp.exists(): return json.loads(result_fp.read_text()) tail = (logs / f"{rid}.log").read_text().strip().splitlines()[-1:] or [""] return {"id": rid, "status": "error", "message": tail[0][:300], "seconds": round(time.time() - t0, 1)} with ThreadPoolExecutor(max_workers=max(1, jobs)) as pool, open(out, "a") as fh: futures = {pool.submit(one, rid): rid for rid in todo} for n, fut in enumerate(as_completed(futures), 1): rec = fut.result() fh.write(json.dumps(rec) + "\n") fh.flush() typer.echo(f" [{n}/{len(todo)}] {rec['id']}: {rec['status']}") # --- cluster backend ------------------------------------------------------------------------------ # `--backend afh` runs a stage's reactions on the AFH k3s cluster through the `afh` client # (github.com/atomsforhumanity/job-runner, client/install.sh) instead of local processes: one # cluster task per reaction, each running `_one` in node-local scratch. The run folder # (WS//afh/) holds a copy of what the stage reads; results come back into WS// # exactly as a local run leaves them, so resuming, mixing backends and `report` work the same. CLUSTER_DIR = SCRIPT.parent / "cluster" # What each stage reads from the workspace, per reaction (beyond reactions.csv). _READS = {"screen": ("inputs",), "search": ("screen",), "verify": ("screen", "search"), "reirc": ("screen",), "rescore": ()} _TASK_SH = r"""set -eo pipefail stage=$(cat "$RUN_DIR/STAGE") mkdir -p "ws/$stage" for x in "$RUN_DIR"/ws/*; do [ "$(basename "$x")" = "$stage" ] || ln -s "$x" "ws/$(basename "$x")"; done sed "s|__RUN_DIR__|$RUN_DIR|g" "$RUN_DIR/ws/$stage/settings.json" > "ws/$stage/settings.json" # What the reaction wrote is kept, in output/ or (on failure) errors/, as one tar: a search # folder has ~10k small files, and copying those to NFS one by one loads the shared server. # The path optimizers' per-iteration histories (node_*_history/, ~98% of a search folder; # training data for ML potentials) go into their own history.tar.zst (~20x smaller), unless # KEEP_HISTORY=0 (--afh keep_history=false) drops them. pack() { d="ws/$stage/$TASK_ID"; [ -d "$d" ] || return 0 if [ "${KEEP_HISTORY:-1}" != 0 ] && [ -n "$(cd "$d" && find . -type d -name '*_history' -print -quit)" ]; then (cd "$d" && find . -type d -name '*_history' -print0 | tar --null -cf - -T - | zstd -q -10 -T"${NTHREADS:-1}" -o "$OUT_DIR/history.tar.zst") fi tar -cf "$OUT_DIR/reaction.tar" --exclude='*_history' -C "$d" . cp -f "$d/result.json" "$OUT_DIR/" 2>/dev/null || true } trap pack EXIT python3 "$RUN_DIR/benchmark.py" _one "$stage" ws "$TASK_ID" """ def _afh(*args: str, capture: bool = False) -> str: if shutil.which("afh") is None: raise typer.BadParameter("--backend afh needs the `afh` client: clone " "github.com/atomsforhumanity/job-runner and run client/install.sh.") proc = subprocess.run(["afh", *args], text=True, capture_output=capture) if proc.returncode == 24 and args[0] == "pull": # rsync: files vanished mid-copy (tasks still writing) return "" if proc.returncode: raise RuntimeError(f"afh {' '.join(args)} failed ({proc.returncode}): {(proc.stderr or '').strip()[-500:]}") return proc.stdout if capture else "" def _afh_spec(stage: str, cluster: dict, timeout: float, n_tasks: int) -> dict: image = cluster.get("image") or ((CLUSTER_DIR / "IMAGE").read_text().strip() if (CLUSTER_DIR / "IMAGE").exists() else "") if not image: raise typer.BadParameter("No cluster image: build one with cluster/build_image.sh, or give --afh image=...") # Small pods: a pod holds all its CPUs until its slowest reaction ends, so with 16 slots one # long search kept 15 idle CPUs reserved on a shared cluster (kaist500-endmin: 304 reserved, 14 used). spec = {"image": image, "command": "bash $RUN_DIR/task.sh", "threads_per_task": 1, "cpus_per_pod": 4, "memory_per_pod": "8Gi", "task_timeout": int(timeout) if timeout else 0, "max_attempts": 2} spec.update({k: v for k, v in cluster.items() if k not in ("image", "run", "poll", "retry_failed", "keep_history")}) if stage in ("screen", "search", "reirc", "rescore"): # g-xTB, xtb and GSM must run single-threaded (they slow down badly otherwise); afh only # forces that when the image or command names xtb, and ours don't. (verify, e.g. DFT, # keeps afh's default: as many threads as the task reserves.) spec["env"] = {"OMP_NUM_THREADS": "1", "MKL_NUM_THREADS": "1", "OPENBLAS_NUM_THREADS": "1", **spec.get("env", {})} if "keep_history" in cluster and not cluster["keep_history"]: spec["env"] = {**spec.get("env", {}), "KEEP_HISTORY": "0"} spec["image"] = image if "shards" not in cluster and stage not in ("screen", "reirc", "rescore"): # searches and DFT take half an hour or more: one round per pod, so a pod frees its CPUs when # its reactions end. The screen (about a minute a reaction) keeps afh's `auto`, which sizes # pods by measured task time to a few rounds each; a fixed cap made each pod run ~500 # screen tasks for hours and kept every other run, `high` ones too, off the cluster. per_pod = max(1, int(spec["cpus_per_pod"]) // int(spec["threads_per_task"])) spec["shards"] = max(1, -(-n_tasks // per_pod)) return spec def _afh_pull_new(ws: Path, stage: str, run_name: str, local: Path, status: dict) -> None: """Pull only the finished tasks not yet in the workspace, plus logs. (`afh pull` re-syncs the whole run, which walks every file of every finished task.)""" remote = f"/srv/jobs/runs/{status['run']}" listing = _afh("ssh", "sh", "-c", f"cd {remote} && for d in output errors; do [ -d $d ] && ls $d | " f"sed \"s|^|$d/|\"; done; true", capture=True) done = _results(ws, stage) want = [ln for ln in listing.split() if not ln.split("/")[-1].startswith(".tmp-") and ln.split("/")[-1] not in done] if not want: return lst = local / ".pull-list" lst.write_text("\n".join(want) + "\n") proc = subprocess.run(["rsync", "-rlt", "--files-from", str(lst), f"{os.environ['AFH_HOST']}:{remote}/", f"{local}/"]) if proc.returncode not in (0, 23, 24): # 23/24: a task folder appeared or vanished mid-copy raise RuntimeError(f"rsync of {len(want)} task folder(s) from {remote} failed ({proc.returncode})") def _afh_collect(ws: Path, stage: str, local: Path) -> int: """Move finished cluster tasks from the pulled run folder into the workspace; returns how many were new. Tasks lost to a dead pod stay pending.""" done = _results(ws, stage) new = 0 with open(ws / stage / "results.jsonl", "a") as fh: for kind in ("output", "errors"): d = local / kind if not d.is_dir(): continue for task in sorted(d.iterdir()): rid = task.name if rid.startswith(".") or rid in done or not (task / "_meta.json").exists(): continue # .tmp--: a task still being published on the cluster meta = json.loads((task / "_meta.json").read_text()) reason = str(meta.get("reason", "")) if kind == "errors" and reason.startswith("abandoned"): continue dest = ws / stage / rid if dest.exists(): shutil.rmtree(dest) if (task / "reaction.tar").exists(): dest.mkdir(parents=True) with tarfile.open(task / "reaction.tar") as tf: tf.extractall(dest, filter="data") if (task / "history.tar.zst").exists(): # kept packed: read with `tar --zstd -xf` shutil.copyfile(task / "history.tar.zst", dest / "history.tar.zst") else: # runs from before the tar (one file per file) shutil.copytree(task, dest, ignore=shutil.ignore_patterns("_meta.json", "task.log")) if (task / "task.log").exists(): shutil.copyfile(task / "task.log", ws / stage / "logs" / f"{rid}.log") rec = None if kind == "output" and (dest / "result.json").exists(): rec = json.loads((dest / "result.json").read_text()) elif reason.startswith("timeout"): rec = {"id": rid, "status": "timeout", "seconds": meta.get("seconds")} if rec is None: log = (task / "task.log").read_text(errors="replace").strip().splitlines() \ if (task / "task.log").exists() else [] rec = {"id": rid, "status": "error", "message": (log[-1:] or [reason])[0][:300], "seconds": meta.get("seconds")} rec["host"] = meta.get("host") for f in task.iterdir(): # the workspace has it now; keep only the timing record if f.name != "_meta.json": shutil.rmtree(f) if f.is_dir() and not f.is_symlink() else f.unlink() fh.write(json.dumps(rec) + "\n") fh.flush() done[rid] = rec new += 1 typer.echo(f" {rid}: {rec['status']}") return new def _afh_run_all(ws: Path, stage: str, ids: list[str], timeout: float, cluster: dict) -> None: """`_run_all` on the cluster: stage the run folder, push, submit, then collect results as they land until the cluster run finishes.""" done = _results(ws, stage) todo = [i for i in ids if i not in done] run_name = cluster.get("run") or f"{ws.resolve().name}-{stage}".lower() run_name = "".join(c if c.isalnum() or c == "-" else "-" for c in run_name.lower())[:32].strip("-") typer.echo(f"{stage}: {len(ids)} reaction(s), {len(ids) - len(todo)} done before, {len(todo)} to run " f"on the cluster (run {run_name}).") (ws / stage / "logs").mkdir(parents=True, exist_ok=True) local = ws / stage / "afh" try: status = json.loads(_afh("status", run_name, "--json", capture=True) or "{}") except RuntimeError: # not pushed yet status = {} if status.get("run"): # a previous submit's finished tasks first _afh_pull_new(ws, stage, run_name, local, status) if (local / "output").exists() or (local / "errors").exists(): _afh_collect(ws, stage, local) done = _results(ws, stage) todo = [i for i in ids if i not in done] if not todo: return attach = status.get("state") == "running" if attach: # an earlier call submitted it and stopped waiting: pick it up, change nothing typer.echo(f" cluster run {run_name} is still running: collecting its results (`afh cancel " f"{run_name}` first to change or resubmit it).") if not attach: # stage the inputs: only what the remaining reactions read wsd = local / "ws" if wsd.exists(): shutil.rmtree(wsd) (wsd / stage).mkdir(parents=True) shutil.copyfile(ws / "reactions.csv", wsd / "reactions.csv") for sub in _READS[stage]: (wsd / sub).mkdir() if (ws / sub / "settings.json").exists(): shutil.copyfile(ws / sub / "settings.json", wsd / sub / "settings.json") for rid in todo: if (ws / sub / rid).is_dir(): shutil.copytree(ws / sub / rid, wsd / sub / rid, symlinks=True) if stage == "rescore": _rescore_stage_routes(ws, wsd, todo) settings = _settings(ws, stage) if settings.get("profile"): shutil.copyfile(settings["profile"], local / "profile.toml") settings = {**settings, "profile": "__RUN_DIR__/profile.toml"} (wsd / stage / "settings.json").write_text(json.dumps(settings, indent=1) + "\n") shutil.copyfile(SCRIPT, local / "benchmark.py") (local / "task.sh").write_text(_TASK_SH) (local / "STAGE").write_text(stage + "\n") known = set() if (local / "tasks.jsonl").exists(): # append only: never reorder a run's tasks known = {json.loads(ln)["id"] for ln in (local / "tasks.jsonl").read_text().splitlines() if ln.strip()} with open(local / "tasks.jsonl", "a") as fh: for rid in todo: if rid not in known: fh.write(json.dumps({"id": rid}) + "\n") cluster = {"description": f"mepd dataset-benchmark {stage} of {ws.resolve().name} " f"({len(_reactions(ws))} reactions, one task each)", **cluster} spec = _afh_spec(stage, cluster, timeout, len(todo)) (local / "run.yaml").write_text("".join(f"{k}: {json.dumps(v)}\n" for k, v in spec.items())) _afh("push", str(local), run_name) _afh("submit", run_name, *(["--retry-failed"] if cluster.get("retry_failed") else [])) poll = float(cluster.get("poll", 60)) while True: time.sleep(poll) status = json.loads(_afh("status", run_name, "--json", capture=True) or "{}") _afh_pull_new(ws, stage, run_name, local, status) _afh_collect(ws, stage, local) typer.echo(f" [{time.strftime('%H:%M')}] {status.get('state')}: {status.get('done', 0)} done, " f"{status.get('failed', 0)} failed, {status.get('pending', '?')} pending" + (f", ETA {status['eta_hours']} h" if status.get("eta_hours") else "")) if status.get("state") != "running": break finished = _results(ws, stage) # once: it parses the whole results file left = [i for i in todo if i not in finished] if left: typer.echo(f"{stage}: {len(left)} reaction(s) not finished on the cluster (lost pods); rerun the " "stage to submit them again.") def _cluster_opts(backend: str, afh: list[str]) -> Optional[dict]: if backend == "local": return None if backend != "afh": raise typer.BadParameter("--backend must be 'local' or 'afh'.") opts = {} for kv in afh or []: k, sep, v = kv.partition("=") if not sep: raise typer.BadParameter(f"--afh takes key=value (run.yaml fields, or run/image/poll), not {kv!r}.") try: opts[k.strip()] = json.loads(v) except json.JSONDecodeError: opts[k.strip()] = v return opts def _execute(ws: Path, stage: str, ids: list[str], jobs: int, timeout: float, cluster: Optional[dict]) -> None: if cluster is None: _run_all(ws, stage, ids, jobs, timeout) else: _afh_run_all(ws, stage, ids, timeout, cluster) BACKEND_HELP = "'local' (processes here) or 'afh' (the AFH cluster through `afh`; see cluster/README.md)." AFH_HELP = ("With --backend afh, key=value for the cluster run: run.yaml fields (cpus_per_pod, memory_per_pod, " "parallelism, priority, ...), run=, image=, poll=, retry_failed=true, " "keep_history=false (drop the path optimizers' per-iteration files; kept, compressed, by default). " "Repeatable.") # --- import ---------------------------------------------------------------------------------------- def _store_inputs(ws: Path, rid: str, start, end, ts) -> dict: """Copy one reaction's structures into WS/inputs// (xyz, Angstrom).""" import mepd d = ws / "inputs" / rid d.mkdir(parents=True, exist_ok=True) paths = {} for name, s in (("start", start), ("end", end), ("ts", ts)): mepd.write_xyz(d / f"{name}.xyz", s) paths[name] = str((d / f"{name}.xyz").relative_to(ws)) return paths def _row(rid: str, source: str, start, ts, paths: dict, barrier="", smiles="") -> dict: return {"id": rid, "source": source, "charge": int(ts.structure.charge), "multiplicity": int(ts.structure.multiplicity), "n_atoms": len(ts.symbols), "n_heavy": sum(s != "H" for s in ts.symbols), **paths, "reference_barrier_kcal": barrier, "smiles": smiles} @app.command("import-manifest") def import_manifest( ws: Path = typer.Argument(..., help="Workspace folder (created if needed)."), manifest: Path = typer.Argument(..., exists=True, help="CSV: id, start, end, ts (xyz paths), charge, " "multiplicity; optional reference_barrier_kcal, smiles. Column names " "start_xyz/end_xyz/ts_reference_xyz and reaction_id are accepted too."), root: Optional[Path] = typer.Option(None, help="Folder the xyz paths are relative to (default: the CSV's)."), source: str = typer.Option("", help="Dataset name to record (default: the CSV's file name)."), ): """Import reactions listed in a CSV with an xyz file per structure.""" import mepd root = root or manifest.parent rows = [] for r in csv.DictReader(open(manifest, newline="")): rid = r.get("id") or r.get("reaction_id") or r.get("idx") get = lambda *keys: next((r[k] for k in keys if r.get(k)), None) # noqa: E731 q, m = int(get("charge") or 0), int(get("multiplicity") or 1) start, end, ts = (mepd.load(root / get(*keys), charge=q, multiplicity=m) for keys in (("start", "start_xyz"), ("end", "end_xyz"), ("ts", "ts_reference_xyz", "ts_xyz"))) rows.append(_row(rid, source or manifest.stem, start, ts, _store_inputs(ws, rid, start, end, ts), get("reference_barrier_kcal", "barrier_kcal") or "", get("smiles") or "")) _write_reactions(ws, rows) typer.echo(f"Imported {len(rows)} reaction(s) into {ws}.") @app.command("import-reaction-qm") def import_reaction_qm( ws: Path = typer.Argument(..., help="Workspace folder (created if needed)."), h5: Path = typer.Argument(..., exists=True, help="Reaction-QM B3LYPD3_TZVP.h5 (doi:10.5281/zenodo.18551029)."), ids: Optional[Path] = typer.Option(None, help="File with one reaction id per line (RXN_... or a number)."), sample: int = typer.Option(0, help="Instead of --ids: this many reactions drawn at random."), seed: int = typer.Option(0, help="Random seed for --sample."), limit: int = typer.Option(0, help="Instead of --ids/--sample: the first N in file order."), ): """Import reactions from the Reaction-QM B3LYP-RXN subset (reactant fragments, product fragments and the TS, B3LYP-D3/TZVP).""" import random import h5py import mepd import numpy as np from qcdata import Structure from mepd.units import ANGSTROM_TO_BOHR, HARTREE_TO_KCAL symbols = {1: "H", 6: "C", 7: "N", 8: "O", 9: "F", 14: "Si", 15: "P", 16: "S", 17: "Cl", 35: "Br", 53: "I"} with h5py.File(h5, "r") as fh: index = {rid: chunk for chunk in fh for rid in fh[chunk]} if ids: wanted = [ln.strip() for ln in ids.read_text().splitlines() if ln.strip()] wanted = [f"RXN_{w.rjust(10, '0')}" if w.isdigit() else w for w in wanted] elif sample: wanted = random.Random(seed).sample(sorted(index), sample) else: wanted = list(index)[:limit] if limit else list(index) rows = [] for rid in wanted: g = fh[index[rid]][rid] def species(tag): s = g[tag] return Structure(symbols=[symbols[int(z)] for z in s["atomic_numbers"][()]], geometry=np.asarray(s["coordinates"][()]) * ANGSTROM_TO_BOHR, charge=int(s["charge"][()]), multiplicity=int(s["multiplicity"][()])) def side(prefix): parts = [species(t) for t in sorted(g) if t.startswith(prefix)] if len(parts) == 1: return parts[0] from mepd.complexes import packed # separate frames: pack the molecules together s = packed(parts, seed=0) return s.model_copy(update={"charge": sum(p.charge for p in parts)}) from mepd.structures import as_node ts = as_node(species("TS")) spin = {"multiplicity": int(ts.structure.multiplicity)} # a side's spin is the reaction's start = as_node(side("R").model_copy(update=spin)) end = as_node(side("P").model_copy(update=spin)) e_r = sum(float(g[t]["EHG"][()][0]) for t in g if t.startswith("R")) barrier = (float(g["TS"]["EHG"][()][0]) - e_r) * HARTREE_TO_KCAL smiles = g["TS"]["smiles"][()] rows.append(_row(rid, "reaction-qm", start, ts, _store_inputs(ws, rid, start, end, ts), f"{barrier:.4f}", smiles.decode() if isinstance(smiles, bytes) else str(smiles))) _write_reactions(ws, rows) typer.echo(f"Imported {len(rows)} Reaction-QM reaction(s) into {ws}.") @app.command("import-transition1x") def import_transition1x( ws: Path = typer.Argument(..., help="Workspace folder (created if needed)."), h5: Path = typer.Argument(..., exists=True, help="Transition1x.h5 (doi:10.6084/m9.figshare.19614657)."), split: str = typer.Option("data", help="Group of the file: data (all reactions), train, val or test."), ids: Optional[Path] = typer.Option(None, help="File with one reaction name per line (rxn0001 or T1X_rxn0001)."), sample: int = typer.Option(0, help="Instead of --ids: this many reactions drawn at random."), seed: int = typer.Option(0, help="Random seed for --sample."), limit: int = typer.Option(0, help="Instead of --ids/--sample: the first N in file order."), ): """Import reactions from Transition1x (M. Schreiner, A. Bhowmik, T. Vegge, J. Busk, O. Winther, "Transition1x - a dataset for building generalizable reactive machine learning potentials", Sci. Data 9, 779 (2022), doi:10.1038/s41597-022-01870-w; data doi:10.6084/m9.figshare.19614657, MIT): each reaction's reactant, product and transition state from its NEB at wB97x/6-31G(d), in one atom order. The reference barrier is the transition state's energy above the reactant's. All are neutral singlets.""" import random import h5py import numpy as np from qcdata import Structure from mepd.structures import as_node from mepd.units import ANGSTROM_TO_BOHR, HARTREE_TO_EV, HARTREE_TO_KCAL symbols = {1: "H", 6: "C", 7: "N", 8: "O", 9: "F"} key = "wB97x_6-31G(d).energy" with h5py.File(h5, "r") as fh: grp = fh[split] index = {rxn: formula for formula in grp for rxn in grp[formula]} if ids: wanted = [ln.strip().removeprefix("T1X_") for ln in ids.read_text().splitlines() if ln.strip()] elif sample: wanted = random.Random(seed).sample(sorted(index), sample) else: wanted = list(index)[:limit] if limit else list(index) rows = [] for rxn in wanted: g = grp[index[rxn]][rxn] def species(tag): s = g[tag] z = [int(v) for v in s["atomic_numbers"][()]] x = np.asarray(s["positions"][()], dtype=float).reshape(-1, len(z), 3)[0] st = Structure(symbols=[symbols[v] for v in z], geometry=x * ANGSTROM_TO_BOHR, charge=0, multiplicity=1) return as_node(st), float(np.asarray(s[key][()]).reshape(-1)[0]) (ts, e_ts), (start, e_r), (end, _) = species("transition_state"), species("reactant"), species("product") barrier = (e_ts - e_r) / HARTREE_TO_EV * HARTREE_TO_KCAL rid = f"T1X_{rxn}" rows.append(_row(rid, "transition1x", start, ts, _store_inputs(ws, rid, start, end, ts), f"{barrier:.4f}")) _write_reactions(ws, rows) typer.echo(f"Imported {len(rows)} Transition1x reaction(s) into {ws}.") def _extxyz_frames(path: Path): """(header dict, symbols, Angstrom coordinates) of each frame of an extended-xyz file.""" import re import numpy as np with open(path) as fh: while True: first = fh.readline() if not first.strip(): return n = int(first) head = dict(re.findall(r'(\w+)=("[^"]*"|\S+)', fh.readline())) rows = [fh.readline().split() for _ in range(n)] yield head, [r[0] for r in rows], np.array([[float(v) for v in r[1:4]] for r in rows]) @app.command("import-transition1x-tmc") def import_transition1x_tmc( ws: Path = typer.Argument(..., help="Workspace folder (created if needed)."), xyz: Path = typer.Argument(..., exists=True, help="t1x_with_tmc.xyz (AEFM.zip of doi:10.5281/zenodo.18338077)."), ids: Optional[Path] = typer.Option(None, help="File with one reaction id per line (TMC_00012_Pd)."), sample: int = typer.Option(0, help="Instead of --ids: this many reactions drawn at random."), seed: int = typer.Option(0, help="Random seed for --sample."), limit: int = typer.Option(0, help="Instead of --ids/--sample: the first N in file order."), ): """Import reactions from Transition1x-TMC (S. Darouich, J. W. Toney, W. Luo, J. Kastner, M. Niepert, H. J. Kulik, "Beyond the Training Domain: Robust Generative Transition State Models for Unseen Chemistry", arXiv:2601.16469 (2026); data in AEFM.zip of doi:10.5281/zenodo.18338077, CC BY 4.0): Transition1x reactions bound to a transition-metal centre (Ti, Zn, Ru, Rh, Pd, Ag, Os, Ir, Pt, Au), each with a reactant, product and transition state at GFN2-xTB in one atom order. The file lists them as consecutive reactant, transition_state, product frames; reaction k (0-based, in file order) becomes TMC__. Charge and spin come from the file (charge 0; uhf unpaired electrons -> multiplicity uhf + 1). The reference barrier is the transition state's energy above the reactant's (GFN2-xTB, eV in the file).""" import random from qcdata import Structure from mepd.structures import as_node from mepd.units import ANGSTROM_TO_BOHR, HARTREE_TO_EV, HARTREE_TO_KCAL frames = list(_extxyz_frames(xyz)) if len(frames) % 3: raise typer.BadParameter(f"{xyz}: {len(frames)} frames, not reactant/TS/product triples") reactions = {} for k in range(len(frames) // 3): trio = frames[3 * k:3 * k + 3] heads = [h for h, _, _ in trio] if [h["type"] for h in heads] != ["reactant", "transition_state", "product"] or \ len({(h["rxn"], h["tm"], h["charge"], h["uhf"]) for h in heads}) != 1 or \ len({tuple(s) for _, s, _ in trio}) != 1: raise typer.BadParameter(f"{xyz}: frames {3 * k}-{3 * k + 2} are not one reaction's reactant, TS and product") reactions[f"TMC_{k:05d}_{heads[0]['tm']}"] = trio if ids: wanted = [ln.strip() for ln in ids.read_text().splitlines() if ln.strip()] elif sample: wanted = random.Random(seed).sample(sorted(reactions), sample) else: wanted = list(reactions)[:limit] if limit else list(reactions) rows = [] for rid in wanted: (h_r, syms, x_r), (h_ts, _, x_ts), (_, _, x_p) = reactions[rid] q, m = int(h_ts["charge"]), int(h_ts["uhf"]) + 1 node = lambda x: as_node(Structure(symbols=syms, geometry=x * ANGSTROM_TO_BOHR, charge=q, multiplicity=m)) # noqa: E731 start, ts, end = node(x_r), node(x_ts), node(x_p) barrier = (float(h_ts["energy"]) - float(h_r["energy"])) / HARTREE_TO_EV * HARTREE_TO_KCAL rows.append(_row(rid, "transition1x-tmc", start, ts, _store_inputs(ws, rid, start, end, ts), f"{barrier:.4f}")) _write_reactions(ws, rows) typer.echo(f"Imported {len(rows)} Transition1x-TMC reaction(s) into {ws}.") # --- screen ---------------------------------------------------------------------------------------- def _side(c) -> dict: """What a mepd.Comparison found, for the record.""" out = {"graph": bool(c.graph), "stereo": c.stereo, "method": c.method, "reason": c.reason, "smiles": list(c.smiles)} if getattr(c, "metal_ligand_changed", None) is not None: out["metal_ligand_changed"] = bool(c.metal_ligand_changed) return out def _mlc(settings: dict) -> dict: """mepd comparison keywords from the screen settings: metal-ligand contacts are compared unless the settings say otherwise (older settings files have no such key). Leaving them out follows the dative zero-order bonds of Z. Li, H.-H. Hsu, B. M. Savoie, ChemRxiv (2026), doi:10.26434/chemrxiv.15006374/v1.""" return {} if settings.get("metal_ligand_contacts", True) else {"metal_ligand_contacts": False} def _decide(rec: dict, out: Path, e_ts: float, a, b, raw_a, raw_b, start, end, energy, settings: dict) -> dict: """From the two minimized IRC ends (a, b) and the raw IRC ends they came from: which dataset endpoints they reproduce, the kept ends (written to out/start.xyz, end.xyz) and the barriers. `energy(node)` gives a structure's energy in hartree. Shared by `screen` and `recompare`.""" import mepd from mepd.units import HARTREE_TO_KCAL kw = _mlc(settings) for e, raw, mn in zip(rec["irc_ends"], (raw_a, raw_b), (a, b)): e["same_species_as_irc_end"] = bool(mepd.same_molecule(raw, mn, **kw)) if mepd.same_molecule(a, b, **kw): return {**rec, "status": "degenerate_irc"} hit = mepd.connects([a, b], start, end, **kw) rec["sides"] = {"start": _side(hit.comparisons[0]), "end": _side(hit.comparisons[1])} rec["irc_reversed"] = hit.reversed start_side, end_side = (b, a) if hit.reversed else (a, b) kept_start = start_side if hit.start else start kept_end = end_side if hit.end else end mepd.write_xyz(out / "start.xyz", kept_start, comments=["irc minimum" if hit.start else "dataset"]) mepd.write_xyz(out / "end.xyz", kept_end, comments=["irc minimum" if hit.end else "dataset"]) undetermined = [k for k, u in zip(("start", "end"), hit.undetermined) if u] if hit.both: status = "match" elif hit.start or hit.end: status = "start_only" if hit.start else "end_only" else: status = "stereo_undetermined" if undetermined else "mismatch" return {**rec, "status": status, "undetermined_sides": undetermined, "start_from": "irc" if hit.start else "dataset", "end_from": "irc" if hit.end else "dataset", "barrier_kcal": float((e_ts - energy(kept_start)) * HARTREE_TO_KCAL) if hit.start else None, "reverse_barrier_kcal": float((e_ts - energy(kept_end)) * HARTREE_TO_KCAL) if hit.end else None} def _screen_one(ws: Path, row: dict, settings: dict) -> dict: """Re-converge one reference TS (converged, exactly one imaginary mode), follow its IRC to two minimized, frequency-checked minima, and decide which dataset endpoints they reproduce; a reproduced endpoint is replaced by that minimum. Every outcome is a status with the diagnostics behind it.""" import mepd from mepd.units import HARTREE_TO_KCAL t0 = time.time() q, m = int(row["charge"]), int(row["multiplicity"]) start, end, ts_ref = (mepd.load(ws / row[k], charge=q, multiplicity=m) for k in ("start", "end", "ts")) eng = _engine(settings) out = ws / "screen" / row["id"] rec = {"id": row["id"]} done = lambda **kw: {**rec, **kw, "seconds": round(time.time() - t0, 1)} # noqa: E731 try: ts = mepd.find_ts(ts_ref, eng, fmax=settings["ts_fmax"], maxiter=settings["ts_maxiter"], continue_unconverged=settings["ts_continue"]) except mepd.WrongStationaryPointError as exc: if exc.structure is not None: mepd.write_xyz(out / "ts_wrong_order.xyz", exc.structure) return done(status="ts_wrong_order", message=str(exc)[:300], ts_lowest_freqs_cm=exc.frequencies.lowest(4) if exc.frequencies is not None else None) except mepd.ConvergenceError as exc: last = getattr(exc, "structure", None) if last is not None: mepd.write_xyz(out / "ts_last.xyz", last) return done(status="ts_not_converged", message=str(exc)[:300]) mepd.write_xyz(out / "ts.xyz", ts) rec["ts_lowest_freqs_cm"] = ts.frequencies.lowest(4) try: path = mepd.irc(ts, eng) except mepd.ConvergenceError as exc: msg = str(exc) return done(status="irc_end_not_minimized" if "could not be minimized" in msg else "irc_failed", message=msg[:300]) mepd.write_xyz(out / "irc.xyz", path.irc_path.nodes) # the IRC itself mepd.write_xyz(out / "path.xyz", path.nodes) # minimum -> IRC -> minimum a, b = path.ends raw_a, raw_b = path.irc_path.nodes[0], path.irc_path.nodes[-1] e_ts, e_a, e_b, e_ra, e_rb = mepd.energies([ts, a, b, raw_a, raw_b], eng) rec["irc_ends"] = [{"drop_on_minimization_kcal": float((er - e) * HARTREE_TO_KCAL), "saddle_escapes": mn.saddle_escapes, "lowest_freqs_cm": mn.frequencies.lowest(3)} for mn, er, e in ((a, e_ra, e_a), (b, e_rb, e_b))] rec["ts_energy_hartree"] = float(e_ts) known = {id(a): e_a, id(b): e_b} energy = lambda n: known[id(n)] if id(n) in known else mepd.energies([n], eng)[0] # noqa: E731 return done(**_decide(rec, out, e_ts, a, b, raw_a, raw_b, start, end, energy, settings)) @app.command() def screen( ws: Path = typer.Argument(..., help="Workspace folder."), engine: str = typer.Option("gxtb", help="Engine by name (mepd.engine), unless --profile is given."), profile: Optional[Path] = typer.Option(None, exists=True, help="mepd profile .toml for the engine."), keep: str = typer.Option("both", help="Which reactions go on to the search: 'both' (the minimized IRC ends " "reproduce both dataset endpoints) or 'any' (at least one). Either way, an IRC end " "that changes species when minimized drops the reaction."), jobs: int = typer.Option(4, help="Reactions at a time (one process each, one thread each)."), timeout: float = typer.Option(600, help="Seconds per reaction; 0 = none."), ts_fmax: float = typer.Option(0.001, help="TS optimization converged below this largest force (eV/Angstrom)."), ts_maxiter: int = typer.Option(1000, help="TS optimizer steps per round."), ts_continue: int = typer.Option(2, help="Rounds that continue an unconverged TS optimization from where it " "stopped (0 = give up there)."), metal_ligand_contacts: bool = typer.Option( True, "--metal-ligand-contacts/--no-metal-ligand-contacts", help="Compare bonds to metal atoms too (default). Without, two structures are the same species when " "their covalent bonds and stereochemistry are, whatever the coordination at the metal (dative " "contacts as zero-order bonds: Li, Hsu, Savoie, ChemRxiv 2026, doi:10.26434/chemrxiv.15006374/v1); " "each side then records metal_ligand_changed."), backend: str = typer.Option("local", help=BACKEND_HELP), afh: Optional[list[str]] = typer.Option(None, help=AFH_HELP), ): """Re-converge every reference TS at the screen's level (TS optimization + IRC) and keep the consistent set: each dataset endpoint the IRC reproduces (same bond graph and stereochemistry) is replaced by that IRC end, so kept endpoints share the TS's atom order and frame.""" if keep not in ("any", "both"): raise typer.BadParameter("--keep must be 'any' or 'both'.") cluster = _cluster_opts(backend, afh) given = {"engine": engine, "profile": str(profile.resolve()) if profile else None, "keep": keep, "ts_fmax": ts_fmax, "ts_maxiter": ts_maxiter, "ts_continue": ts_continue} if not metal_ligand_contacts: # only when off, so earlier settings files still resume given["metal_ligand_contacts"] = False s = _settings(ws, "screen", given) _execute(ws, "screen", [r["id"] for r in _reactions(ws)], jobs, timeout, cluster) res = _results(ws, "screen") kept = [i for i, r in res.items() if _kept(r, s["keep"])] moved = sum(1 for r in res.values() if _kept({**r, "irc_ends": None}, s["keep"]) and not _irc_connects(r)) typer.echo(f"screen: {len(kept)} of {len(res)} kept ({s['keep']} endpoint(s) reproduced by the IRC; " f"{moved} more left out because minimizing an IRC end changed its species).") def _recompare_one(ws: Path, row: dict, old: dict, settings: dict) -> dict: """Re-decide one screened reaction from its stored geometries (`recompare`).""" import mepd from mepd.units import HARTREE_TO_KCAL out = ws / "screen" / row["id"] q, m = int(row["charge"]), int(row["multiplicity"]) from mepd.structures import read_frames path, irc = (read_frames(out / f, charge=q, multiplicity=m) for f in ("path.xyz", "irc.xyz")) a, b, raw_a, raw_b = path[0], path[-1], irc[0], irc[-1] start, end = (mepd.load(ws / row[k], charge=q, multiplicity=m) for k in ("start", "end")) e_ts = old["ts_energy_hartree"] known = {} # IRC-end energies the old record already implies rev = old.get("irc_reversed", False) if old.get("start_from") == "irc" and old.get("barrier_kcal") is not None: known[1 if rev else 0] = e_ts - old["barrier_kcal"] / HARTREE_TO_KCAL if old.get("end_from") == "irc" and old.get("reverse_barrier_kcal") is not None: known[0 if rev else 1] = e_ts - old["reverse_barrier_kcal"] / HARTREE_TO_KCAL eng = _engine(settings) def energy(n): k = 0 if n is a else 1 if k not in known: known[k] = float(mepd.energies([n], eng)[0]) return known[k] keep_keys = ("id", "ts_lowest_freqs_cm", "irc_ends", "ts_energy_hartree", "seconds", "host") rec = {k: old[k] for k in keep_keys if k in old} rec["irc_ends"] = [dict(e) for e in old["irc_ends"]] return _decide(rec, out, e_ts, a, b, raw_a, raw_b, start, end, energy, settings) @app.command() def recompare( ws: Path = typer.Argument(..., help="Workspace folder whose screen has run."), metal_ligand_contacts: bool = typer.Option( True, "--metal-ligand-contacts/--no-metal-ligand-contacts", help="As for `screen`: compare bonds to metal atoms too, or leave them out (Li, Hsu, Savoie, ChemRxiv " "2026, doi:10.26434/chemrxiv.15006374/v1)."), jobs: int = typer.Option(4, help="Reactions at a time."), force: bool = typer.Option(False, help="Re-decide even when the settings are unchanged (e.g. after a " "change in mepd's comparison)."), ): """Re-decide every finished screen result from its stored geometries (the TS, the IRC in irc.xyz and its minimized ends in path.xyz) under other comparison settings, without re-running a TS optimization or an IRC: statuses, sides, the IRC-connects check, kept ends and barriers (a newly kept IRC end gets one single point). The old results are kept as screen/results..jsonl, and screen/settings.json records the change ("recompared").""" import datetime from concurrent.futures import ProcessPoolExecutor import mepd fp = ws / "screen" / "settings.json" old_s = json.loads(fp.read_text()) new_s = {k: v for k, v in old_s.items() if k not in ("metal_ligand_contacts", "recompared")} if not metal_ligand_contacts: new_s["metal_ligand_contacts"] = False if {k: v for k, v in old_s.items() if k != "recompared"} == new_s and not force: typer.echo("screen already uses these comparison settings; nothing to do.") return rows = {r["id"]: r for r in _reactions(ws)} res = _results(ws, "screen") todo = [i for i, r in res.items() if r.get("irc_ends") and (ws / "screen" / i / "path.xyz").exists() and (ws / "screen" / i / "irc.xyz").exists() and r.get("ts_energy_hartree") is not None] stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") backup = ws / "screen" / f"results.{stamp}.jsonl" shutil.copyfile(ws / "screen" / "results.jsonl", backup) new = dict(res) with ProcessPoolExecutor(max(1, jobs)) as pool: futs = {i: pool.submit(_recompare_one, ws, rows[i], res[i], new_s) for i in todo} for k, (i, f) in enumerate(futs.items(), 1): new[i] = f.result() if k % 500 == 0: typer.echo(f" {k}/{len(todo)}") mepd_dir = Path(mepd.__file__).resolve().parent.parent commit = subprocess.run(["git", "-C", str(mepd_dir), "rev-parse", "HEAD"], capture_output=True, text=True).stdout.strip() or getattr(mepd, "__version__", "?") new_s["recompared"] = [*old_s.get("recompared", []), {"date": stamp, "from": {k: v for k, v in old_s.items() if k != "recompared"}, "backup": backup.name, "n": len(todo), "mepd": commit}] (ws / "screen" / "results.jsonl").write_text("".join(json.dumps(r) + "\n" for r in new.values())) fp.write_text(json.dumps(new_s, indent=1) + "\n") changed = sum(new[i].get("status") != res[i].get("status") for i in todo) kept = [i for i, r in new.items() if _kept(r, new_s["keep"])] typer.echo(f"recompare: {len(todo)} re-decided, {changed} changed status; {len(kept)} of {len(new)} kept " f"(was {sum(_kept(r, old_s['keep']) for r in res.values())}). Old results: {backup.name}.") def _irc_connects(rec: dict) -> bool: """The IRC itself reached each kept end: minimizing an IRC end kept it the same species. Otherwise the minimization did the chemistry (an IRC that barely moves off a soft saddle, then drops tens of kcal/mol into another molecule), and the "reference TS" is not this reaction's TS.""" return all(e.get("same_species_as_irc_end", True) for e in rec.get("irc_ends") or []) def _kept(rec: dict, keep: str) -> bool: ends_ok = rec.get("status") == "match" or (keep == "any" and rec.get("status") in ("start_only", "end_only")) return ends_ok and _irc_connects(rec) def _kept_ids(ws: Path) -> list[str]: """The reactions the screen keeps under the workspace's keep rule. The rule only selects (it changes no screen result), so it may be tightened after a search: the later stages and the report then leave out what it drops.""" keep = _settings(ws, "screen")["keep"] return [i for i, r in _results(ws, "screen").items() if _kept(r, keep)] # --- IRC retry --------------------------------------------------------------------------------------- # An IRC stops where the largest force falls below its fmax (0.1 eV/Angstrom by default): on a flat # shoulder that can be tens of kcal/mol above the minimum, and minimizing from there may cross into # another molecule. `reirc` follows the IRC again, tighter, from the screen's converged TS for each # reaction that the IRC-connects rule left out, and the tighter IRC decides. def _reirc_one(ws: Path, row: dict, settings: dict) -> dict: """The IRC again, at the retry's fmax, from the screen's TS: geometries and energies only (the comparison with the dataset endpoints is made by `reirc` afterwards, as in `recompare`).""" import mepd from mepd.units import HARTREE_TO_KCAL t0 = time.time() q, m = int(row["charge"]), int(row["multiplicity"]) ts = mepd.load(ws / "screen" / row["id"] / "ts.xyz", charge=q, multiplicity=m) eng = _engine(settings) out = ws / "reirc" / row["id"] out.mkdir(parents=True, exist_ok=True) rec = {"id": row["id"], "irc_fmax": settings["irc_fmax"], "irc_steps": settings["irc_steps"]} done = lambda **kw: {**rec, **kw, "seconds": round(time.time() - t0, 1)} # noqa: E731 try: path = mepd.irc(ts, eng, fmax=settings["irc_fmax"], steps=settings["irc_steps"]) except mepd.ConvergenceError as exc: msg = str(exc) return done(status="irc_end_not_minimized" if "could not be minimized" in msg else "irc_failed", message=msg[:300]) mepd.write_xyz(out / "irc.xyz", path.irc_path.nodes) mepd.write_xyz(out / "path.xyz", path.nodes) a, b = path.ends raw_a, raw_b = path.irc_path.nodes[0], path.irc_path.nodes[-1] e_ts, e_a, e_b, e_ra, e_rb = (float(e) for e in mepd.energies([ts, a, b, raw_a, raw_b], eng)) rec["irc_ends"] = [{"drop_on_minimization_kcal": float((er - e) * HARTREE_TO_KCAL), "saddle_escapes": mn.saddle_escapes, "lowest_freqs_cm": mn.frequencies.lowest(3)} for mn, er, e in ((a, e_ra, e_a), (b, e_rb, e_b))] rec["ts_energy_hartree"] = e_ts rec["end_energies_hartree"] = [e_a, e_b] rec["irc_frames"] = len(path.irc_path.nodes) return done(status="irc_done") def _reirc_decide(ws: Path, row: dict, old: dict, new: dict, settings: dict) -> dict: """The screen record from the retried IRC: `_decide` on its ends, the loose IRC's record kept.""" import mepd from mepd.structures import read_frames q, m = int(row["charge"]), int(row["multiplicity"]) d = ws / "reirc" / row["id"] path, irc = (read_frames(d / f, charge=q, multiplicity=m) for f in ("path.xyz", "irc.xyz")) a, b, raw_a, raw_b = path[0], path[-1], irc[0], irc[-1] start, end = (mepd.load(ws / row[k], charge=q, multiplicity=m) for k in ("start", "end")) known = {id(a): new["end_energies_hartree"][0], id(b): new["end_energies_hartree"][1]} rec = {"id": row["id"], "ts_lowest_freqs_cm": old.get("ts_lowest_freqs_cm"), "irc_ends": [dict(e) for e in new["irc_ends"]], "ts_energy_hartree": old["ts_energy_hartree"], "seconds": old.get("seconds"), "host": old.get("host")} rec = _decide(rec, d, old["ts_energy_hartree"], a, b, raw_a, raw_b, start, end, lambda n: known[id(n)], settings) loose = {k: v for k, v in old.items() if k not in ("id", "irc_retry", "irc_loose")} return {**rec, "irc_retry": {"fmax": new["irc_fmax"], "steps": new["irc_steps"], "seconds": new.get("seconds"), "host": new.get("host"), "ts_energy_hartree": new["ts_energy_hartree"]}, "irc_loose": loose} @app.command() def reirc( ws: Path = typer.Argument(..., help="Workspace folder whose screen has run."), fmax: float = typer.Option(0.01, help="The retried IRC stops below this largest force (eV/Angstrom)."), steps: int = typer.Option(3000, help="IRC steps per direction."), jobs: int = typer.Option(4, help="Reactions at a time."), timeout: float = typer.Option(1800, help="Seconds per reaction; 0 = none."), ids: Optional[Path] = typer.Option(None, exists=True, help="Retry these reactions instead (a file of ids): " "e.g. those an IRC end leaves under a stricter comparison than the screen's."), apply: bool = typer.Option(True, "--apply/--no-apply", help="Write the retried decisions into the screen " "results (default), or only run the IRCs."), backend: str = typer.Option("local", help=BACKEND_HELP), afh: Optional[list[str]] = typer.Option(None, help=AFH_HELP), ): """Retry the IRC, tighter, for every screened reaction with an IRC end that changed species when minimized (the IRC-connects rule): from the screen's converged TS, same engine. The retried IRC then decides the reaction (status, kept ends, barriers and the IRC-connects check) in place of the loose one, which stays in the record as `irc_loose`, with its files as screen//*.loose.xyz. The old results are kept as screen/results..jsonl. A reaction newly kept whose earlier search used other ends loses that search (moved to search/superseded-/), so the next `search` runs it again from the new ends.""" import datetime cluster = _cluster_opts(backend, afh) s_screen = _settings(ws, "screen") res = _results(ws, "screen") want = set(ids.read_text().split()) if ids else None todo = [i for i, r in res.items() if r.get("irc_ends") and r.get("ts_energy_hartree") is not None and (i in want if want is not None else not _irc_connects(r)) and "irc_retry" not in r and (ws / "screen" / i / "ts.xyz").exists()] s = _settings(ws, "reirc", {"engine": s_screen.get("engine"), "profile": s_screen.get("profile"), "irc_fmax": fmax, "irc_steps": steps}) typer.echo(f"reirc: {len(todo)} reaction(s) with an IRC end that changes species on minimization.") _execute(ws, "reirc", todo, jobs, timeout, cluster) if not apply: return rows = {r["id"]: r for r in _reactions(ws)} new = _results(ws, "reirc") todo = [i for i in todo if new.get(i, {}).get("status") == "irc_done"] failed = {i: new[i]["status"] for i in new if new[i].get("status") != "irc_done" and "irc_retry" not in res.get(i, {})} stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") backup = ws / "screen" / f"results.{stamp}.jsonl" shutil.copyfile(ws / "screen" / "results.jsonl", backup) out = dict(res) for k, i in enumerate(todo, 1): out[i] = _reirc_decide(ws, rows[i], res[i], new[i], s_screen) if k % 1000 == 0: typer.echo(f" {k}/{len(todo)} decided") for i, st in failed.items(): # the retry failed: the loose IRC's record stands (still left out) out[i] = {**res[i], "irc_retry": {"fmax": s["irc_fmax"], "steps": s["irc_steps"], "status": st, "message": new[i].get("message")}} for i in todo: # the retried IRC's files replace the loose ones sd, rd = ws / "screen" / i, ws / "reirc" / i for f in ("irc.xyz", "path.xyz", "start.xyz", "end.xyz"): if (sd / f).exists() and not (sd / f.replace(".xyz", ".loose.xyz")).exists(): shutil.move(sd / f, sd / f.replace(".xyz", ".loose.xyz")) if (rd / f).exists(): # a degenerate IRC keeps no start/end shutil.copyfile(rd / f, sd / f) (ws / "screen" / "results.jsonl").write_text("".join(json.dumps(r) + "\n" for r in out.values())) s_screen["irc_retry"] = [*s_screen.get("irc_retry", []), {"date": stamp, "fmax": s["irc_fmax"], "steps": s["irc_steps"], "n": len(todo), "failed": len(failed), "backup": backup.name}] (ws / "screen" / "settings.json").write_text(json.dumps(s_screen, indent=1) + "\n") rejoined = [i for i in todo if _kept(out[i], s_screen["keep"])] searched = _results(ws, "search") stale = [i for i in rejoined if i in searched] if stale: dest = ws / "search" / f"superseded-{stamp}" dest.mkdir(parents=True) shutil.copyfile(ws / "search" / "results.jsonl", dest / "results.jsonl") for i in stale: if (ws / "search" / i).exists(): shutil.move(ws / "search" / i, dest / i) gone = set(stale) (ws / "search" / "results.jsonl").write_text("".join(json.dumps(r) + "\n" for i, r in searched.items() if i not in gone)) changed = sum(out[i].get("status") != res[i].get("status") for i in todo) typer.echo(f"reirc: {len(todo)} re-decided ({changed} changed status), {len(failed)} retries failed; " f"{len(rejoined)} now kept ({len(stale)} earlier searches moved aside); " f"{sum(_kept(r, s_screen['keep']) for r in out.values())} of {len(out)} kept in all. " f"Old results: {backup.name}.") # --- search ---------------------------------------------------------------------------------------- def _search_one(ws: Path, row: dict, settings: dict) -> dict: """mepd.channels between the kept endpoints; compare the best route found with the re-converged reference TS against one barrier floor.""" import mepd t0 = time.time() q, m = int(row["charge"]), int(row["multiplicity"]) sd = ws / "screen" / row["id"] start, end, ref_ts = (mepd.load(sd / f"{n}.xyz", charge=q, multiplicity=m) for n in ("start", "end", "ts")) kw = {"profile": settings["profile"]} if settings.get("profile") else {} eng = _engine(settings) result = mepd.channels(start, end, output=ws / "search" / row["id"] / "run", engine=eng, **kw, **settings.get("options", {})) best = result.lowest() rec = {"id": row["id"], "n_direct": len(result.direct), "n_multistep": len(result.multistep), "n_offtarget": len(result.offtarget), "seconds": round(time.time() - t0, 1)} # one floor for both barriers: the lowest of every reactant structure there is floor = [start, *result.start_conformers] ref_barrier = mepd.barrier_kcal(ref_ts, floor, eng) rec["ref_barrier_kcal"] = ref_barrier if best is None: return {**rec, "status": "no_route"} best_barrier = mepd.barrier_kcal(best.ts, floor, eng) lower_by = ref_barrier - best_barrier tol = settings.get("tolerance_kcal", 1.0) outcome = "lower" if lower_by > tol else "higher" if lower_by < -tol else "matched" return {**rec, "status": "done", "outcome": outcome, "best_kind": best.kind, "best_folder": str(best.folder.relative_to(ws)), "best_barrier_kcal": best_barrier, "lower_by_kcal": lower_by} @app.command() def search( ws: Path = typer.Argument(..., help="Workspace folder."), profile: Optional[Path] = typer.Option(None, exists=True, help="mepd profile .toml: path-search method and " "settings, and the engine (default: mepd's defaults, g-xTB)."), engine: str = typer.Option("gxtb", help="Engine by name, when the profile does not set one."), options: Optional[Path] = typer.Option(None, exists=True, help="JSON or TOML with mepd.ChannelsOptions " "fields (the `mepd channels` options, underscores for dashes)."), tolerance: float = typer.Option(1.0, help="kcal/mol: a route counts as lower than the reference TS when " "its barrier is lower by more than this."), ids: Optional[Path] = typer.Option(None, exists=True, help="Search only these kept reactions (a file of ids, " "one per line)."), jobs: int = typer.Option(2, help="Reactions at a time. Each search also runs --workers path searches."), timeout: float = typer.Option(3600, help="Seconds per reaction (a hard stop; search_budget in --options " "ends a search gracefully); 0 = none."), backend: str = typer.Option("local", help=BACKEND_HELP), afh: Optional[list[str]] = typer.Option(None, help=AFH_HELP), ): """Search the kept reactions with mepd.channels and compare the lowest route found with the re-converged reference TS.""" cluster = _cluster_opts(backend, afh) opts = {"workers": 1, "parallel": False} # one core per search unless --options says otherwise if options: text = options.read_text() if options.suffix == ".json": opts.update(json.loads(text)) else: import tomllib opts.update(tomllib.loads(text)) import mepd mepd.ChannelsOptions(**opts).validate() # fail now on a bad option, not in every reaction s = _settings(ws, "search", {"profile": str(profile.resolve()) if profile else None, "engine": engine, "options": opts, "tolerance_kcal": tolerance}) kept = _kept_ids(ws) if ids: only = set(ids.read_text().split()) kept = [i for i in kept if i in only] _execute(ws, "search", kept, jobs, timeout, cluster) keep = set(kept) res = {i: r for i, r in _results(ws, "search").items() if i in keep} # searched under an earlier keep rule: out lower = [i for i, r in res.items() if r.get("outcome") == "lower"] typer.echo(f"search: {len(res)} searched, {len(lower)} with a route lower than the reference " f"(by > {s['tolerance_kcal']} kcal/mol).") # --- rescore --------------------------------------------------------------------------------------- # The screen's IRC-connects rule, applied to every step of the routes a search found: a step's IRC must # reach the minima it is filed under by itself (minimizing each raw IRC end keeps it the same species). # A step that fails gets its IRC again, tighter, from its TS; it stands when the tighter IRC passes and # reaches the same two species. The lowest route whose steps all stand is the reaction's best route. def _route_folders(run: Path) -> list[tuple[str, Path, list[Path]]]: """(kind, folder, step folders) of every route in a channels output folder, as mepd.read_channels orders them (direct, then multistep).""" def numbered(d: Path, prefix: str) -> list[Path]: return sorted((x for x in d.glob(f"{prefix}_*") if x.is_dir() and x.name.rsplit("_", 1)[1].isdigit()), key=lambda x: int(x.name.rsplit("_", 1)[1])) out = [("direct", d, [d]) for d in numbered(run / "channels", "channel")] out += [("multistep", d, numbered(d, "step")) for d in numbered(run / "alternate-channels", "alternate_channel")] return out def _route_energy(steps: list[Path]) -> float: """The highest TS energy along a route (hartree), as mepd's Channel.ts; inf when unreadable.""" try: return max(float((s / "ts.energies").read_text().split()[0]) for s in steps) except (OSError, ValueError, IndexError): return float("inf") def _routes_by_energy(ws: Path, rid: str) -> list[tuple[float, str, Path, list[Path]]]: routes = [(_route_energy(st), kind, d, st) for kind, d, st in _route_folders(ws / "search" / rid / "run")] return sorted(routes, key=lambda r: r[0]) # stable: ties keep read_channels' order, as lowest() does def _step_ok(step: Path, q: int, m: int) -> bool: """Does the stored IRC reach its minimized ends by itself? (channels' irc.xyz: minimized end, raw IRC, minimized end.)""" import mepd from mepd.structures import read_frames try: fr = read_frames(step / "irc.xyz", charge=q, multiplicity=m) except Exception: return False if len(fr) < 4: return False return bool(mepd.same_molecule(fr[1], fr[0]) and mepd.same_molecule(fr[-2], fr[-1])) def _step_retry(ws: Path, rid: str, step: Path, q: int, m: int, eng, settings: dict) -> dict: """The step's IRC again at the retry's fmax: does it pass, and reach the same two species?""" import mepd from mepd.structures import read_frames from mepd.units import HARTREE_TO_KCAL t0 = time.time() rel = step.relative_to(ws / "search" / rid / "run") out = {"step": str(rel)} try: stored = read_frames(step / "irc.xyz", charge=q, multiplicity=m) ts = mepd.load(step / "ts.xyz", charge=q, multiplicity=m) r = mepd.irc(ts, eng, fmax=settings["irc_fmax"], steps=settings["irc_steps"]) except Exception as exc: return {**out, "ok": False, "error": f"{type(exc).__name__}: {exc}"[:300], "seconds": round(time.time() - t0, 1)} raw, mn = r.irc_path.nodes, r.path.nodes keeps = [bool(mepd.same_molecule(raw[0], mn[0])), bool(mepd.same_molecule(raw[-1], mn[-1]))] a, b = stored[0], stored[-1] same = bool((mepd.same_molecule(mn[0], a) and mepd.same_molecule(mn[-1], b)) or (mepd.same_molecule(mn[0], b) and mepd.same_molecule(mn[-1], a))) e = [float(x) for x in mepd.energies([raw[0], mn[0], raw[-1], mn[-1]], eng)] d = ws / "rescore" / rid / rel d.mkdir(parents=True, exist_ok=True) mepd.write_xyz(d / "irc.xyz", list(mn)) # as channels writes it: minimized end, IRC, minimized end return {**out, "ok": all(keeps) and same, "ends_keep_species": keeps, "same_ends_as_stored": same, "drops_kcal": [round((e[0] - e[1]) * HARTREE_TO_KCAL, 3), round((e[2] - e[3]) * HARTREE_TO_KCAL, 3)], "irc_frames": len(raw), "seconds": round(time.time() - t0, 1)} def _rescore_one(ws: Path, row: dict, settings: dict, retry: bool = True) -> Optional[dict]: """A searched reaction's best route under the IRC-connects rule for route steps. Without `retry`, None when a step would need its IRC again (so the caller can send it to the cluster).""" from mepd.units import HARTREE_TO_KCAL t0 = time.time() rid = row["id"] q, m = int(row["charge"]), int(row["multiplicity"]) old = json.loads((ws / "search" / rid / "result.json").read_text()) base = {"id": rid, "search_status": old.get("status"), "search_outcome": old.get("outcome"), "search_best_folder": old.get("best_folder"), "search_lower_by_kcal": old.get("lower_by_kcal"), "ref_barrier_kcal": old.get("ref_barrier_kcal")} if old.get("status") != "done" or not old.get("best_folder"): return {**base, "status": old.get("status"), "seconds": 0.0} routes = _routes_by_energy(ws, rid) by_folder = {str(d.relative_to(ws)): e for e, _, d, _ in routes} e_best = by_folder.get(old["best_folder"]) if e_best is None or e_best == float("inf"): return {**base, "status": "error", "message": f"best route {old['best_folder']} missing or unreadable"} floor = e_best - old["best_barrier_kcal"] / HARTREE_TO_KCAL # the search's barrier floor (hartree) eng = _engine(settings) if retry else None checked, n_retry = [], 0 for e, kind, d, steps in routes: if e == float("inf"): checked.append({"route": str(d.relative_to(ws)), "ok": False, "unreadable": True}) continue entry = {"route": str(d.relative_to(ws)), "kind": kind, "barrier_kcal": float((e - floor) * HARTREE_TO_KCAL), "steps": []} ok = True for st in steps: if _step_ok(st, q, m): continue if not retry: return None n_retry += 1 r = _step_retry(ws, rid, st, q, m, eng, settings) entry["steps"].append(r) if not r["ok"]: ok = False break # this route is out; its other steps need no retry entry["ok"] = ok checked.append(entry) if ok: lower_by = old["ref_barrier_kcal"] - entry["barrier_kcal"] tol = settings["tolerance_kcal"] outcome = "lower" if lower_by > tol else "higher" if lower_by < -tol else "matched" return {**base, "status": "done", "outcome": outcome, "best_kind": kind, "best_folder": entry["route"], "best_barrier_kcal": entry["barrier_kcal"], "lower_by_kcal": lower_by, "routes_dropped": sum(not c["ok"] for c in checked), "irc_retries": n_retry, "checked": checked, "seconds": round(time.time() - t0, 1)} return {**base, "status": "no_route", "routes_dropped": len(checked), "irc_retries": n_retry, "checked": checked, "seconds": round(time.time() - t0, 1)} def _rescore_stage_routes(ws: Path, wsd: Path, todo: list[str]) -> None: """Cluster inputs for `rescore`: each reaction's search record and the files of its routes up to the lowest one whose stored steps all pass (higher routes are never needed).""" plan = json.loads((ws / "rescore" / "plan.json").read_text()) (wsd / "search").mkdir(exist_ok=True) for rid in todo: src, dst = ws / "search" / rid, wsd / "search" / rid dst.mkdir(parents=True) shutil.copyfile(src / "result.json", dst / "result.json") for folder in plan[rid]: d = ws / folder rel = d.relative_to(src) (dst / rel).mkdir(parents=True, exist_ok=True) if (d / "path.txt").exists(): shutil.copyfile(d / "path.txt", dst / rel / "path.txt") for st in [d] if not list(d.glob("step_*")) else d.glob("step_*"): srel = st.relative_to(src) (dst / srel).mkdir(parents=True, exist_ok=True) for f in ("ts.xyz", "ts.energies", "irc.xyz", "irc.energies"): if (st / f).exists(): shutil.copyfile(st / f, dst / srel / f) def _rescore_plan(ws: Path, row: dict, settings: dict) -> tuple[str, Optional[dict], list[str]]: """Local first pass: the rescored record when no step needs its IRC again, else the routes to stage.""" from mepd.units import HARTREE_TO_KCAL # noqa: F401 (imports mepd once per worker) q, m = int(row["charge"]), int(row["multiplicity"]) try: rec = _rescore_one(ws, row, settings, retry=False) except Exception as exc: return row["id"], {"id": row["id"], "status": "error", "message": f"{type(exc).__name__}: {exc}"[:300]}, [] if rec is not None: return row["id"], rec, [] need = [] for e, _, d, steps in _routes_by_energy(ws, row["id"]): need.append(str(d.relative_to(ws))) if e != float("inf") and all(_step_ok(st, q, m) for st in steps): break return row["id"], None, need @app.command() def rescore( ws: Path = typer.Argument(..., help="Workspace folder whose search has run."), fmax: float = typer.Option(0.01, help="A failing step's IRC is retried down to this largest force (eV/Angstrom)."), steps: int = typer.Option(3000, help="IRC steps per direction."), jobs: int = typer.Option(4, help="Reactions at a time (the local first pass, and local retries)."), timeout: float = typer.Option(3600, help="Seconds per reaction; 0 = none."), backend: str = typer.Option("local", help=BACKEND_HELP), afh: Optional[list[str]] = typer.Option(None, help=AFH_HELP), ): """Re-score every finished search under the IRC-connects rule for route steps (see the section comment): rescore/results.jsonl holds, per reaction, the best route that stands, its barrier against the search's own floor, the outcome, and every route checked on the way. A first pass reads the stored IRCs here; only reactions with a failing step below their lowest passing route are run (on --backend), with just those routes' files staged.""" from collections import Counter from concurrent.futures import ProcessPoolExecutor cluster = _cluster_opts(backend, afh) s_search = _settings(ws, "search") s = _settings(ws, "rescore", {"engine": s_search.get("engine"), "profile": s_search.get("profile"), "tolerance_kcal": s_search.get("tolerance_kcal", 1.0), "irc_fmax": fmax, "irc_steps": steps}) rows = {r["id"]: r for r in _reactions(ws)} kept = set(_kept_ids(ws)) searched = {i: r for i, r in _results(ws, "search").items() if i in kept and (ws / "search" / i / "result.json").exists()} done = _results(ws, "rescore") plan_fp = ws / "rescore" / "plan.json" plan = json.loads(plan_fp.read_text()) if plan_fp.exists() else {} first = [i for i in searched if i not in done and i not in plan] typer.echo(f"rescore: {len(searched)} searched, {len(done)} rescored before; first pass on {len(first)}.") (ws / "rescore").mkdir(exist_ok=True) with ProcessPoolExecutor(max(1, jobs)) as pool, open(ws / "rescore" / "results.jsonl", "a") as fh: futs = [pool.submit(_rescore_plan, ws, rows[i], s) for i in first] for k, f in enumerate(futs, 1): rid, rec, need = f.result() if rec is not None: fh.write(json.dumps(rec) + "\n") else: plan[rid] = need if k % 1000 == 0: fh.flush() plan_fp.write_text(json.dumps(plan) + "\n") typer.echo(f" {k}/{len(first)}") plan_fp.write_text(json.dumps(plan) + "\n") done = _results(ws, "rescore") todo = [i for i in plan if i not in done] typer.echo(f"rescore: {len(todo)} reaction(s) need a step's IRC again.") if todo: _execute(ws, "rescore", todo, jobs, timeout, cluster) res = {i: r for i, r in _results(ws, "rescore").items() if i in searched} count = lambda key: dict(sorted(Counter(str(r.get(key)) # noqa: E731 for r in res.values()).items())) changed = sum(r.get("best_folder") != r.get("search_best_folder") for r in res.values() if r.get("status") == "done") typer.echo(f"rescore: {len(res)} of {len(searched)} rescored; status {count('status')}; outcome " f"{count('outcome')}; best route changed for {changed}, " f"{sum(r.get('status') == 'no_route' and r.get('search_status') == 'done' for r in res.values())} " "now have no route.") # --- verify ---------------------------------------------------------------------------------------- def _dft_ts(guess, expected_ends, eng) -> dict: """TS opt + IRC at the verify level, and whether the IRC still joins the structures the cheap level's IRC joined.""" import mepd try: ts = mepd.find_ts(guess, eng) except mepd.ConvergenceError as exc: return {"status": "ts_failed", "message": str(exc)[:300]} rec = {"energy_hartree": float(mepd.energies(ts, eng)[0])} try: path = mepd.irc(ts, eng) except mepd.ConvergenceError as exc: return {**rec, "status": "irc_failed", "message": str(exc)[:300], "ts": ts} hit = mepd.connects(path, *expected_ends) return {**rec, "status": "connects" if hit.both else "other", "ts": ts, "irc": path} def _verify_one(ws: Path, row: dict, settings: dict) -> dict: """Re-optimize the reference TS and every TS of the lowest route at the verify level; the gap is their energy difference (same atoms, no floor).""" import mepd t0 = time.time() q, m = int(row["charge"]), int(row["multiplicity"]) eng = _engine(settings) out = ws / "verify" / row["id"] sd = ws / "screen" / row["id"] start, end, ref_guess = (mepd.load(sd / f"{n}.xyz", charge=q, multiplicity=m) for n in ("start", "end", "ts")) ref = _dft_ts(ref_guess, (start, end), eng) result = mepd.read_channels(ws / "search" / row["id"] / "run", charge=q, multiplicity=m) best = result.lowest() steps = [] for k, step in enumerate(best.steps): if step.irc is None: # the search kept no IRC for it: nothing to compare against steps.append({"status": "no_reference_irc"}) continue r = _dft_ts(step.ts, step.irc.ends, eng) if "ts" in r: mepd.write_xyz(out / f"route_step{k}_ts.xyz", r.pop("ts")) r.pop("irc", None) steps.append(r) if "ts" in ref: mepd.write_xyz(out / "reference_ts.xyz", ref.pop("ts")) ref.pop("irc", None) rec = {"id": row["id"], "reference": ref, "route": steps, "route_kind": best.kind, "seconds": round(time.time() - t0, 1)} ok = ref.get("status") == "connects" and steps and all(s.get("status") == "connects" for s in steps) if ok: from mepd.units import HARTREE_TO_KCAL gap = (ref["energy_hartree"] - max(s["energy_hartree"] for s in steps)) * HARTREE_TO_KCAL tol = _settings(ws, "search").get("tolerance_kcal", 1.0) return {**rec, "status": "verified", "lower_by_kcal": gap, "outcome": "lower" if gap > tol else "higher" if gap < -tol else "matched"} return {**rec, "status": "not_verified"} @app.command() def verify( ws: Path = typer.Argument(..., help="Workspace folder."), profile: Path = typer.Option(..., exists=True, help="mepd profile .toml for the verify level (e.g. DFT; " "see profiles/)."), which: str = typer.Option("lower", help="'lower': reactions whose search found a lower route; 'routes': " "every reaction with a route."), jobs: int = typer.Option(1, help="Reactions at a time."), timeout: float = typer.Option(0, help="Seconds per reaction; 0 = none."), backend: str = typer.Option("local", help=BACKEND_HELP), afh: Optional[list[str]] = typer.Option(None, help=AFH_HELP), ): """Rescore at a higher level: re-optimize the reference TS and every TS of the lowest route found, follow each IRC, and compare their energies.""" if which not in ("lower", "routes"): raise typer.BadParameter("--which must be 'lower' or 'routes'.") cluster = _cluster_opts(backend, afh) _settings(ws, "verify", {"profile": str(profile.resolve()), "which": which}) kept = set(_kept_ids(ws)) res = {i: r for i, r in _results(ws, "search").items() if i in kept} ids = [i for i, r in res.items() if (r.get("outcome") == "lower" if which == "lower" else r.get("status") == "done")] _execute(ws, "verify", ids, jobs, timeout, cluster) out = _results(ws, "verify") typer.echo(f"verify: {sum(r.get('status') == 'verified' for r in out.values())} of {len(out)} verified; " f"{sum(r.get('outcome') == 'lower' for r in out.values())} still lower at the verify level.") # --- report ---------------------------------------------------------------------------------------- @app.command() def report(ws: Path = typer.Argument(..., help="Workspace folder.")): """One table (report/table.csv), a summary (report/summary.json) and a figure of how much lower the routes found are.""" rows = _reactions(ws) keep = _settings(ws, "screen")["keep"] if (ws / "screen" / "settings.json").exists() else "both" kept = set(_kept_ids(ws)) if (ws / "screen" / "settings.json").exists() else set() st = {stage: _results(ws, stage) for stage in STAGES} for stage in STAGES[1:]: # searched or verified under an earlier, looser keep rule: not reported st[stage] = {i: r for i, r in st[stage].items() if i in kept} table = [] for r in rows: sc, se, ve = (st[s].get(r["id"], {}) for s in STAGES) table.append({ "id": r["id"], "source": r["source"], "n_atoms": r["n_atoms"], "charge": r["charge"], "multiplicity": r["multiplicity"], "reference_barrier_kcal": r["reference_barrier_kcal"], "screen": sc.get("status", ""), "start_from": sc.get("start_from", ""), "end_from": sc.get("end_from", ""), "screen_barrier_kcal": sc.get("barrier_kcal", ""), "kept": _kept(sc, keep) if sc else "", "search": se.get("status", ""), "routes": (se.get("n_direct", 0) + se.get("n_multistep", 0)) if se else "", "search_outcome": se.get("outcome", ""), "best_kind": se.get("best_kind", ""), "lower_by_kcal": se.get("lower_by_kcal", ""), "verify": ve.get("status", ""), "verify_outcome": ve.get("outcome", ""), "verify_lower_by_kcal": ve.get("lower_by_kcal", ""), }) out = ws / "report" out.mkdir(exist_ok=True) with open(out / "table.csv", "w", newline="") as fh: w = csv.DictWriter(fh, fieldnames=list(table[0]) if table else ["id"]) w.writeheader() w.writerows(table) def count(stage, key="status"): c: dict = {} for rec in st[stage].values(): c[rec.get(key) or "-"] = c.get(rec.get(key) or "-", 0) + 1 return c summary = {"reactions": len(rows), "screen": count("screen"), "kept": sum(t["kept"] is True for t in table), "keep_rule": keep, "search": count("search"), "search_outcome": count("search", "outcome"), "verify": count("verify"), "verify_outcome": count("verify", "outcome")} (out / "summary.json").write_text(json.dumps(summary, indent=1) + "\n") typer.echo(json.dumps(summary, indent=1)) lower = [float(t["lower_by_kcal"]) for t in table if t["search_outcome"] == "lower"] if lower: import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt fig, ax = plt.subplots(figsize=(4.5, 3), dpi=150) ax.hist(lower, bins=20, color="#2a5db0") ax.set_xlabel("route found lower than the reference TS by (kcal/mol)") ax.set_ylabel("reactions") ax.spines[["top", "right"]].set_visible(False) fig.tight_layout() fig.savefig(out / "lower_by.png", metadata={"Software": None}) typer.echo(f"Wrote {out / 'table.csv'} and {out / 'summary.json'}.") @app.command() def run( ws: Path = typer.Argument(..., help="Workspace folder (already imported)."), search_profile: Optional[Path] = typer.Option(None, exists=True, help="Profile for the search stage."), search_options: Optional[Path] = typer.Option(None, exists=True, help="ChannelsOptions for the search."), verify_profile: Optional[Path] = typer.Option(None, exists=True, help="Profile for verify; skipped if not given."), jobs: int = typer.Option(2, help="Reactions at a time in the screen and search stages."), ): """screen, search, verify (with --verify-profile) and report, with each stage's defaults otherwise. Resumes like the stages do.""" screen(ws, engine="gxtb", profile=None, keep="both", jobs=jobs, timeout=600, ts_fmax=0.001, ts_maxiter=1000, ts_continue=2, backend="local", afh=None) search(ws, profile=search_profile, engine="gxtb", options=search_options, tolerance=1.0, ids=None, jobs=jobs, timeout=3600, backend="local", afh=None) if verify_profile: verify(ws, profile=verify_profile, which="lower", jobs=1, timeout=0, backend="local", afh=None) report(ws) @app.command("_one", hidden=True) def _one(stage: str, ws: Path, rid: str): """One reaction of one stage, in its own process (used by the stages).""" import mepd mepd.set_threads(int(os.environ.get("OMP_NUM_THREADS", "1"))) row = next(r for r in _reactions(ws) if r["id"] == rid) fn = {"screen": _screen_one, "search": _search_one, "verify": _verify_one, "reirc": _reirc_one, "rescore": _rescore_one}[stage] rec = fn(ws, row, _settings(ws, stage)) rec.setdefault("status", "done") d = ws / stage / rid d.mkdir(parents=True, exist_ok=True) (d / "result.json").write_text(json.dumps(rec, indent=1, default=float) + "\n") @app.command() def clean(ws: Path = typer.Argument(...), stage: str = typer.Argument(..., help="screen, search, verify or report")): """Move a stage's results aside (to .old-