# /// 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).") 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")} _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"): # 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 != "screen": # 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) 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 left = [i for i in todo if i not in _results(ws, stage)] 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}.") # --- screen ---------------------------------------------------------------------------------------- def _side(c) -> dict: """What a mepd.Comparison found, for the record.""" return {"graph": bool(c.graph), "stereo": c.stereo, "method": c.method, "reason": c.reason, "smiles": list(c.smiles)} 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), "same_species_as_irc_end": bool(mepd.same_molecule(raw, mn)), "saddle_escapes": mn.saddle_escapes, "lowest_freqs_cm": mn.frequencies.lowest(3)} for raw, mn, er, e in ((raw_a, a, e_ra, e_a), (raw_b, b, e_rb, e_b))] rec["ts_energy_hartree"] = float(e_ts) if mepd.same_molecule(a, b): return done(status="degenerate_irc") hit = mepd.connects([a, b], start, end) 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"]) e_start, e_end = mepd.energies([kept_start, kept_end], eng) 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 done(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 - e_start) * HARTREE_TO_KCAL) if hit.start else None, reverse_barrier_kcal=float((e_ts - e_end) * HARTREE_TO_KCAL) if hit.end else None) @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)."), 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) s = _settings(ws, "screen", {"engine": engine, "profile": str(profile.resolve()) if profile else None, "keep": keep, "ts_fmax": ts_fmax, "ts_maxiter": ts_maxiter, "ts_continue": ts_continue}) _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 _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)] # --- 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."), 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) _execute(ws, "search", kept, jobs, timeout, cluster) res = {i: r for i, r in _results(ws, "search").items() if i in kept} # 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).") # --- 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, 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}[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-