diff --git a/CHANGELOG.md b/CHANGELOG.md index a4213e9..24898d2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- Runtime-selectable, agent-free MLIP APIs and CLI commands for direct ASE + MACE, Rootstock, and NVIDIA ALCHEMI MACE calculations +- Ordered MLIP batch execution with per-item JSON results, calculator/model + reuse, partial-failure handling, and a persistent batch manifest +- Polaris installation and live smoke-test recipes for the three MLIP paths - Zeo++ module (`matkit.zeopp`) for pore geometry analysis: pore diameters (Di/Df/Dif), accessible surface area, accessible volume, pore size distribution, and channel identification - CLI `matkit zeopp run` and `matkit zeopp analyze` subcommands with support for high accuracy mode (`-ha`), custom radii files (`-r UFF.rad`), and configurable probe parameters - CLI interface (`matkit` command) with subcommands for graspa, graspa_sycl, raspa2, and tobacco diff --git a/README.md b/README.md index 4c00825..d7855aa 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ - **RASPA2** -- Classical GCMC simulations - **RASPA3** -- Force field format conversion from RASPA2 - **Zeo++** -- Pore geometry analysis (pore diameters, surface area, volume, channels) -- **MACE-MP** -- ML interatomic potential geometry/cell optimization +- **MLIPs** -- direct MACE, Rootstock, and NVIDIA ALCHEMI execution - **ORCA** -- Quantum chemistry (planned) ## Features @@ -44,6 +44,12 @@ pip install -e ".[rdkit]" # For ML interatomic potentials (MACE) pip install -e ".[mlip]" +# Lightweight access to cluster-managed Rootstock models +pip install -e ".[rootstock]" + +# NVIDIA ALCHEMI MACE support (install a matching CUDA extra too) +pip install -e ".[nvalchemi_mace]" + # All optional dependencies pip install -e ".[all]" @@ -75,8 +81,42 @@ matkit zeopp run --cif structure.cif --analysis res --analysis sa --radii UFF.ra # Parse existing Zeo++ output files matkit zeopp analyze --path output_dir/ + +# Run MACE directly through ASE +matkit mlip run --input structure.cif --backend ase-mace \ + --checkpoint medium --device cuda --dtype float32 + +# Run a Rootstock checkpoint already deployed on Polaris +matkit mlip run --input structure.cif --backend rootstock \ + --checkpoint mace-mp-0-medium --cluster polaris --device cuda + +# Run a native NVIDIA ALCHEMI batch +matkit mlip run-batch --input-dir cifs --backend nvalchemi-mace \ + --checkpoint medium --device cuda --batch-size 16 +``` + +### GPU examples + +[`examples/mlip_gpu.py`](examples/mlip_gpu.py) runs one backend per Python +process so GPU runtime state is isolated. It accepts one or more ASE-readable +structure files and writes a manifest plus one JSON result per input. + +```bash +# Direct MACE calculator through ASE +python examples/mlip_gpu.py --backend ase-mace structure.cif + +# Rootstock-managed MACE checkpoint on Polaris +python examples/mlip_gpu.py --backend rootstock \ + --cluster polaris --checkpoint mace-mp-0-medium structure.cif + +# NVIDIA ALCHEMI MACE with native GPU batching +python examples/mlip_gpu.py --backend nvalchemi-mace \ + --checkpoint medium --batch-size 16 structures/*.cif ``` +All three commands force `device="cuda"` and use `float32` where the backend +exposes a dtype. Pass `--driver opt` for a fixed-cell geometry optimization. + ## Python API ```python @@ -111,6 +151,24 @@ print(result["results"]["sa"]) # {'ASA': 4004.7, 'ASA_m2_g': 3918.3, ...} # Parse existing Zeo++ output files result = get_output_data("output_dir/") + +# Agent-free, runtime-selectable MLIP execution +from matkit.mlip import ( + ASEMACEConfig, + MLIPCalculationConfig, + run_mlip, +) + +result = run_mlip( + "structure.cif", + ASEMACEConfig( + checkpoint="medium", + device="cuda", + dtype="float32", + ), + MLIPCalculationConfig(driver="energy"), + output_file="mace_result.json", +) ``` ## License diff --git a/alcf/polaris/mlip/README.md b/alcf/polaris/mlip/README.md new file mode 100644 index 0000000..36bf946 --- /dev/null +++ b/alcf/polaris/mlip/README.md @@ -0,0 +1,59 @@ +# MLIP playground on Polaris + +This directory installs and smoke-tests the three agent-free MLIP paths in +MatKit: + +- direct MACE through ASE; +- a cluster-managed MACE checkpoint through Rootstock; +- native batched MACE through NVIDIA ALCHEMI Toolkit. + +These scripts validate that the installations work. They are not performance +benchmarks and do not collect repeated timing or parity statistics. + +## Prerequisites + +Rootstock is deployed on Polaris, but ALCF users need access to its shared +installation. Follow the current Polaris instructions in the +[Matter Model Almanac](https://garden-ai.github.io/almanac/clusters/) before +running the Rootstock smoke test. + +Run the installer from the MatKit checkout on a Polaris login node: + +```bash +export MATKIT_MLIP_ENV=/lus/eagle/projects/PROJECT/USER/envs/matkit-mlip +bash alcf/polaris/mlip/install.sh +``` + +The installer creates an isolated Python 3.12 environment, installs the CUDA +12 and MACE extras for ALCHEMI, installs Rootstock, and installs this checkout +in editable mode. Override `MATKIT_MLIP_ENV`; the default is `.venv` in the +repository. + +## Smoke test + +Edit the `#PBS -A` project in `smoke.pbs`, then submit it while the checkout is +your working directory: + +```bash +qsub -v MATKIT_MLIP_ENV="$MATKIT_MLIP_ENV" \ + alcf/polaris/mlip/smoke.pbs +``` + +The defaults use the small periodic structure in `tests/data`, direct and +ALCHEMI checkpoint alias `medium`, and Rootstock checkpoint +`mace-mp-0-medium`. Override paths or checkpoint names when submitting: + +```bash +qsub -v MATKIT_MLIP_ENV="$MATKIT_MLIP_ENV",\ +MATKIT_SMOKE_INPUT=/path/to/input.cif,\ +MACE_CHECKPOINT=/path/to/model.pt,\ +ROOTSTOCK_CHECKPOINT=mace-mp-0-medium \ + alcf/polaris/mlip/smoke.pbs +``` + +Results are written under `projects/mlip_smoke_$PBS_JOBID` by default. Set +`MATKIT_SMOKE_OUTPUT` to choose another persistent directory. + +For compute-node downloads, the PBS script exports the ALCF HTTP proxy. Model +weights should be allowed to finish downloading before treating later timings +as performance measurements. diff --git a/alcf/polaris/mlip/install.sh b/alcf/polaris/mlip/install.sh new file mode 100644 index 0000000..1ceebcf --- /dev/null +++ b/alcf/polaris/mlip/install.sh @@ -0,0 +1,44 @@ +#!/bin/bash -l + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +MLIP_ENV="${MATKIT_MLIP_ENV:-${REPO_ROOT}/.venv}" + +module use /soft/modulefiles +module load conda/2025-09-25 + +if [[ ! -x "${MLIP_ENV}/bin/python" ]]; then + python -m venv "${MLIP_ENV}" +fi + +source "${MLIP_ENV}/bin/activate" +python -m pip install --upgrade pip setuptools wheel + +python -m pip install \ + --extra-index-url https://download.pytorch.org/whl/cu126 \ + --extra-index-url https://pypi.nvidia.com \ + 'nvalchemi-toolkit[cu12,mace]>=0.2,<0.3' + +python -m pip install -e \ + "${REPO_ROOT}[mlip,rootstock,nvalchemi_mace]" + +python - <<'PY' +from importlib.metadata import version + +for package in ( + "matkit", + "ase", + "mace-torch", + "rootstock", + "nvalchemi-toolkit", + "torch", +): + print(f"{package}=={version(package)}") +PY + +echo +echo "Environment installed at ${MLIP_ENV}" +echo "Check Rootstock access with: rootstock resolve --cluster polaris --json" +echo "Submit alcf/polaris/mlip/smoke.pbs from the MatKit checkout next." diff --git a/alcf/polaris/mlip/smoke.pbs b/alcf/polaris/mlip/smoke.pbs new file mode 100644 index 0000000..0b2c08c --- /dev/null +++ b/alcf/polaris/mlip/smoke.pbs @@ -0,0 +1,85 @@ +#!/bin/bash -l +#PBS -N matkit-mlip-smoke +#PBS -l select=1:system=polaris +#PBS -l place=scatter +#PBS -l walltime=00:30:00 +#PBS -l filesystems=home:eagle +#PBS -q debug +#PBS -A PROJECT + +set -euo pipefail + +: "${MATKIT_MLIP_ENV:?Submit with -v MATKIT_MLIP_ENV=/path/to/env}" + +MATKIT_REPO="${MATKIT_REPO:-${PBS_O_WORKDIR}}" +INPUT_FILE="${MATKIT_SMOKE_INPUT:-${MATKIT_REPO}/tests/data/test_structure.cif}" +OUTPUT_DIR="${MATKIT_SMOKE_OUTPUT:-${MATKIT_REPO}/projects/mlip_smoke_${PBS_JOBID}}" +MACE_CHECKPOINT="${MACE_CHECKPOINT:-medium}" +ROOTSTOCK_CHECKPOINT="${ROOTSTOCK_CHECKPOINT:-mace-mp-0-medium}" + +module use /soft/modulefiles +module load conda/2025-09-25 +source "${MATKIT_MLIP_ENV}/bin/activate" + +export HTTP_PROXY="http://proxy.alcf.anl.gov:3128" +export HTTPS_PROXY="http://proxy.alcf.anl.gov:3128" +export http_proxy="${HTTP_PROXY}" +export https_proxy="${HTTPS_PROXY}" +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}" + +mkdir -p "${OUTPUT_DIR}" +cd "${MATKIT_REPO}" +nvidia-smi +rootstock resolve --cluster polaris --json + +matkit mlip run \ + --input "${INPUT_FILE}" \ + --output "${OUTPUT_DIR}/ase_mace.json" \ + --backend ase-mace \ + --checkpoint "${MACE_CHECKPOINT}" \ + --device cuda \ + --dtype float32 + +matkit mlip run \ + --input "${INPUT_FILE}" \ + --output "${OUTPUT_DIR}/rootstock_mace.json" \ + --backend rootstock \ + --checkpoint "${ROOTSTOCK_CHECKPOINT}" \ + --cluster polaris \ + --device cuda \ + --timeout 1200 + +matkit mlip run-batch \ + --input "${INPUT_FILE}" \ + --outdir "${OUTPUT_DIR}/nvalchemi_mace" \ + --backend nvalchemi-mace \ + --checkpoint "${MACE_CHECKPOINT}" \ + --device cuda \ + --dtype float32 \ + --batch-size 1 + +python - \ + "${OUTPUT_DIR}/ase_mace.json" \ + "${OUTPUT_DIR}/rootstock_mace.json" \ + "${OUTPUT_DIR}/nvalchemi_mace" <<'PY' +import json +import math +import sys +from pathlib import Path + +names = sys.argv[1:3] +batch_results = sorted(Path(sys.argv[3]).glob("[0-9][0-9][0-9][0-9][0-9]_*.json")) +assert len(batch_results) == 1, batch_results +names.append(str(batch_results[0])) + +for name in names: + data = json.loads(Path(name).read_text()) + assert data["success"], data["error"] + assert math.isfinite(data["energy"]) + assert data["forces"] + assert all(math.isfinite(value) for row in data["forces"] for value in row) + print(f"validated {name}: energy={data['energy']} eV") +PY + +echo "MLIP smoke test complete: ${OUTPUT_DIR}" diff --git a/examples/mlip_gpu.py b/examples/mlip_gpu.py new file mode 100644 index 0000000..556c998 --- /dev/null +++ b/examples/mlip_gpu.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python3 +"""Run one MatKit MLIP backend on one or more structures using a GPU.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from matkit.mlip import ( + ASEMACEConfig, + MLIPCalculationConfig, + NVAlchemiMACEConfig, + RootstockConfig, + run_mlip_batch, +) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Run direct ASE MACE, Rootstock, or NVIDIA ALCHEMI MACE on a GPU. " + "Use one backend per process to keep GPU runtime state isolated." + ) + ) + parser.add_argument( + "inputs", + nargs="+", + type=Path, + help="Structure files readable by ASE.", + ) + parser.add_argument( + "--backend", + choices=("ase-mace", "rootstock", "nvalchemi-mace"), + required=True, + ) + parser.add_argument( + "--checkpoint", + help=( + "Model alias or checkpoint path. Defaults to 'medium' for MACE and " + "ALCHEMI, or 'mace-mp-0-medium' for Rootstock." + ), + ) + parser.add_argument( + "--output-dir", + type=Path, + help="Result directory (default: mlip_gpu_results/).", + ) + parser.add_argument("--driver", choices=("energy", "opt"), default="energy") + parser.add_argument("--fmax", type=float, default=0.01) + parser.add_argument("--steps", type=int, default=1000) + parser.add_argument( + "--batch-size", + type=int, + default=16, + help="Native batch size for ALCHEMI; accepted by all backends.", + ) + parser.add_argument( + "--cluster", + help="Rootstock cluster name (defaults to 'polaris').", + ) + parser.add_argument( + "--root", + type=Path, + help="Rootstock deployment root instead of a named cluster.", + ) + parser.add_argument( + "--compile-model", + action="store_true", + help="Enable model compilation for NVIDIA ALCHEMI MACE.", + ) + parser.add_argument( + "--enable-cueq", + action="store_true", + help="Enable CuEquivariance for NVIDIA ALCHEMI MACE.", + ) + return parser + + +def main() -> int: + parser = build_parser() + args = parser.parse_args() + + if args.cluster and args.root: + parser.error("--cluster and --root are mutually exclusive") + + if args.backend == "ase-mace": + backend = ASEMACEConfig( + checkpoint=args.checkpoint or "medium", + device="cuda", + dtype="float32", + ) + elif args.backend == "rootstock": + backend = RootstockConfig( + checkpoint=args.checkpoint or "mace-mp-0-medium", + cluster=None if args.root else (args.cluster or "polaris"), + root=str(args.root) if args.root else None, + device="cuda", + ) + else: + backend = NVAlchemiMACEConfig( + checkpoint=args.checkpoint or "medium", + device="cuda", + dtype="float32", + compile_model=args.compile_model, + enable_cueq=args.enable_cueq, + ) + + calculation = MLIPCalculationConfig( + driver=args.driver, + fmax=args.fmax, + steps=args.steps, + ) + output_dir = args.output_dir or Path("mlip_gpu_results") / args.backend + summary = run_mlip_batch( + args.inputs, + backend, + calculation, + output_dir=output_dir, + batch_size=args.batch_size, + ) + + print( + json.dumps( + { + "status": summary["status"], + "backend": args.backend, + "total": summary["total"], + "succeeded": summary["succeeded"], + "failed": summary["failed"], + "manifest_file": summary["manifest_file"], + }, + indent=2, + ) + ) + return 0 if summary["failed"] == 0 else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pyproject.toml b/pyproject.toml index 711c5bd..bce10c1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,11 +18,13 @@ dependencies = [ [project.optional-dependencies] rdkit = ["rdkit"] mlip = ["mace-torch"] +rootstock = ["rootstock>=1.6,<2"] +nvalchemi_mace = ["nvalchemi-toolkit[mace]>=0.2,<0.3"] plot = ["matplotlib>=3.5"] -pacmof2 = ["pacmof2"] +pacmof2 = ["pacmof2 @ git+https://github.com/snurr-group/pacmof2.git"] graspa = ["pyyaml>=6.0"] pygraspa = ["pyyaml>=6.0"] -all = ["rdkit", "mace-torch", "matplotlib>=3.5", "pacmof2", "pyyaml>=6.0"] +all = ["rdkit", "mace-torch", "matplotlib>=3.5", "pacmof2 @ git+https://github.com/snurr-group/pacmof2.git", "pyyaml>=6.0"] dev = ["pytest>=7.0", "ruff>=0.4"] [project.scripts] diff --git a/skills.md b/skills.md index fff7ab2..66eaf17 100644 --- a/skills.md +++ b/skills.md @@ -19,7 +19,7 @@ src/matkit/ raspa3/ # RASPA2 -> RASPA3 format conversion zeopp/ # Zeo++ pore geometry analysis (wraps network binary) tobacco/ # SMILES -> CIF linker generation for ToBaCCo - mlip/ # MACE-MP ML interatomic potential optimization + mlip/ # Direct, Rootstock, and ALCHEMI MLIP execution orca/ # ORCA quantum chemistry (stub) io/ # File format converters (SMILES, PubChem JSON) utils/ # Shared utilities (unit cell calc, solvent removal, CIF sampling) @@ -50,7 +50,8 @@ matkit - **Language**: Python >= 3.10 - **Core deps**: ase (atomic simulation), click (CLI), networkx (graph analysis), numpy -- **Optional deps**: rdkit (SMILES), mace-torch (MLIP), openbabel CLI (obabel) +- **Optional deps**: rdkit, mace-torch, rootstock, nvalchemi-toolkit, + openbabel CLI (obabel) - **Build**: setuptools via pyproject.toml (PEP 621) - **Linting**: ruff (E, F rules, 80 char line length) - **Testing**: pytest (tests/ directory) diff --git a/src/matkit/cli.py b/src/matkit/cli.py index 6449d4b..2d86da7 100644 --- a/src/matkit/cli.py +++ b/src/matkit/cli.py @@ -1012,6 +1012,311 @@ def mlip_cli(): pass +def _mlip_options(function): + """Add runtime-neutral MLIP options to a Click command.""" + decorators = [ + click.option( + "--backend", + required=True, + type=click.Choice(["ase-mace", "rootstock", "nvalchemi-mace"]), + help="MLIP execution backend.", + ), + click.option( + "--checkpoint", + required=True, + help="Model alias, canonical Rootstock ID, or checkpoint path.", + ), + click.option( + "--device", default=None, help="Device such as cpu or cuda." + ), + click.option( + "--dtype", + default=None, + type=click.Choice(["float32", "float64"]), + help="Floating-point precision.", + ), + click.option( + "--driver", + default="energy", + show_default=True, + type=click.Choice(["energy", "opt"]), + ), + click.option( + "--optimizer", + default="fire", + show_default=True, + type=click.Choice(["bfgs", "lbfgs", "gpmin", "fire", "mdmin"]), + ), + click.option("--fmax", default=0.01, show_default=True, type=float), + click.option("--steps", default=1000, show_default=True, type=int), + click.option( + "--calculator-type", + default="mace_mp", + show_default=True, + type=click.Choice(["mace_mp", "mace_off", "mace_anicc"]), + help="Direct ASE MACE calculator factory.", + ), + click.option( + "--dispersion/--no-dispersion", + default=False, + help="Enable MACE-MP D3 dispersion.", + ), + click.option("--cluster", default=None, help="Rootstock cluster ID."), + click.option( + "--root", + "root_path", + default=None, + type=click.Path(), + help="Custom Rootstock installation root.", + ), + click.option( + "--cache-root", + default=None, + type=click.Path(), + help="Custom Rootstock cache root.", + ), + click.option( + "--setup-kwarg", + multiple=True, + metavar="KEY=JSON", + help="Rootstock setup keyword; repeat as needed.", + ), + click.option( + "--timeout", + default=600.0, + show_default=True, + type=float, + help="Rootstock worker startup timeout.", + ), + click.option( + "--weights", + default=None, + type=click.Path(), + help="Rootstock custom checkpoint weights.", + ), + click.option( + "--dt", + default=0.1, + show_default=True, + type=float, + help="ALCHEMI FIRE timestep.", + ), + click.option( + "--compile-model", + is_flag=True, + help="Compile the ALCHEMI MACE model.", + ), + click.option( + "--enable-cueq", + is_flag=True, + help="Enable cuEquivariance in ALCHEMI MACE.", + ), + ] + for decorator in reversed(decorators): + function = decorator(function) + return function + + +def _parse_setup_kwargs(values): + parsed = {} + for value in values: + if "=" not in value: + raise click.BadParameter( + "must use KEY=JSON syntax", param_hint="--setup-kwarg" + ) + key, raw = value.split("=", 1) + if not key: + raise click.BadParameter( + "key must not be empty", param_hint="--setup-kwarg" + ) + try: + parsed[key] = json.loads(raw) + except json.JSONDecodeError: + parsed[key] = raw + return parsed + + +def _build_mlip_configs(options): + from matkit.mlip import ( + ASEMACEConfig, + MLIPCalculationConfig, + NVAlchemiMACEConfig, + RootstockConfig, + ) + + backend_name = options["backend"] + checkpoint = options["checkpoint"] + device = options["device"] + dtype = options["dtype"] + if backend_name == "ase-mace": + backend = ASEMACEConfig( + checkpoint=checkpoint, + device=device or "cpu", + dtype=dtype or "float64", + calculator_type=options["calculator_type"], + dispersion=options["dispersion"], + ) + elif backend_name == "rootstock": + backend = RootstockConfig( + checkpoint=checkpoint, + cluster=options["cluster"], + root=options["root_path"], + cache_root=options["cache_root"], + setup_kwargs=_parse_setup_kwargs(options["setup_kwarg"]), + timeout=options["timeout"], + weights=options["weights"], + device=device or "cpu", + ) + else: + backend = NVAlchemiMACEConfig( + checkpoint=checkpoint, + device=device or "cuda", + dtype=dtype or "float32", + dt=options["dt"], + compile_model=options["compile_model"], + enable_cueq=options["enable_cueq"], + ) + calculation = MLIPCalculationConfig( + driver=options["driver"], + optimizer=options["optimizer"], + fmax=options["fmax"], + steps=options["steps"], + ) + return backend, calculation + + +@mlip_cli.command("run") +@click.option( + "--input", + "input_file", + required=True, + type=click.Path(exists=True, dir_okay=False), + help="Input structure readable by ASE.", +) +@click.option( + "--output", + default="output.json", + show_default=True, + type=click.Path(dir_okay=False), + help="JSON result file.", +) +@_mlip_options +def mlip_run_cmd(input_file, output, **options): + """Run one MLIP energy calculation or fixed-cell optimization.""" + from matkit.mlip import run_mlip + + try: + backend, calculation = _build_mlip_configs(options) + result = run_mlip( + input_file, + backend, + calculation=calculation, + output_file=output, + ) + except (ImportError, OSError, ValueError) as exc: + raise click.ClickException(str(exc)) from exc + if not result["success"]: + raise click.ClickException(result["error"]) + click.echo( + json.dumps( + { + "status": "success", + "energy": result["energy"], + "unit": result["energy_unit"], + "converged": result["converged"], + "output_results_file": str(output), + }, + indent=2, + ) + ) + + +@mlip_cli.command("run-batch") +@click.option( + "--input", + "input_files", + multiple=True, + type=click.Path(exists=True, dir_okay=False), + help="Input structure; repeat for an ordered list.", +) +@click.option( + "--input-dir", + default=None, + type=click.Path(exists=True, file_okay=False), + help="Directory containing input structures.", +) +@click.option( + "--pattern", + default="*.cif", + show_default=True, + help="Glob used with --input-dir.", +) +@click.option( + "--outdir", + default="mlip_results", + show_default=True, + type=click.Path(file_okay=False), +) +@click.option("--batch-size", default=16, show_default=True, type=int) +@click.option("--max-atoms", default=None, type=int) +@_mlip_options +def mlip_run_batch_cmd( + input_files, + input_dir, + pattern, + outdir, + batch_size, + max_atoms, + **options, +): + """Run an ordered MLIP batch and write a JSON manifest.""" + from pathlib import Path + + from matkit.mlip import run_mlip_batch + + if bool(input_files) == bool(input_dir): + raise click.UsageError("Specify exactly one of --input or --input-dir.") + if input_dir: + files = [ + str(path) + for path in sorted(Path(input_dir).glob(pattern)) + if path.is_file() + ] + if not files: + raise click.UsageError( + f"No files matching {pattern!r} in {input_dir}." + ) + else: + files = list(input_files) + + try: + backend, calculation = _build_mlip_configs(options) + summary = run_mlip_batch( + files, + backend, + calculation=calculation, + output_dir=outdir, + batch_size=batch_size, + max_atoms=max_atoms, + ) + except (ImportError, OSError, ValueError) as exc: + raise click.ClickException(str(exc)) from exc + click.echo( + json.dumps( + { + "status": summary["status"], + "manifest_file": summary["manifest_file"], + "total": summary["total"], + "succeeded": summary["succeeded"], + "failed": summary["failed"], + }, + indent=2, + ) + ) + if summary["status"] == "failure": + raise click.ClickException("All MLIP batch items failed.") + + @mlip_cli.command("mace-opt") @click.option( "--fname", diff --git a/src/matkit/mlip/__init__.py b/src/matkit/mlip/__init__.py index 83231a6..897e32e 100644 --- a/src/matkit/mlip/__init__.py +++ b/src/matkit/mlip/__init__.py @@ -1,4 +1,24 @@ -__all__ = [] +from matkit.mlip.config import ( + ASEMACEConfig, + MLIPBackendConfig, + MLIPCalculationConfig, + NVAlchemiMACEConfig, + RootstockConfig, +) +from matkit.mlip.runner import run_mlip, run_mlip_batch +from matkit.types import MLIPBatchSummary, MLIPResult + +__all__ = [ + "ASEMACEConfig", + "MLIPBackendConfig", + "MLIPBatchSummary", + "MLIPCalculationConfig", + "MLIPResult", + "NVAlchemiMACEConfig", + "RootstockConfig", + "run_mlip", + "run_mlip_batch", +] try: from matkit.mlip.mace_opt import run_opt_mace diff --git a/src/matkit/mlip/config.py b/src/matkit/mlip/config.py new file mode 100644 index 0000000..25f8cad --- /dev/null +++ b/src/matkit/mlip/config.py @@ -0,0 +1,124 @@ +"""Configuration objects for runtime-selectable MLIP calculations.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from typing import Any, Literal, TypeAlias + + +_DTYPES = {"float32", "float64"} +_OPTIMIZERS = {"bfgs", "lbfgs", "gpmin", "fire", "mdmin"} + + +@dataclass(frozen=True) +class ASEMACEConfig: + """Run a MACE calculator directly through ASE.""" + + checkpoint: str = "medium" + device: str = "cpu" + dtype: Literal["float32", "float64"] = "float64" + calculator_type: Literal["mace_mp", "mace_off", "mace_anicc"] = "mace_mp" + dispersion: bool = False + damping: str = "bj" + dispersion_xc: str = "pbe" + dispersion_cutoff: float = 21.167088422553647 + type: Literal["ase-mace"] = field(default="ase-mace", init=False) + + def __post_init__(self) -> None: + if not self.checkpoint: + raise ValueError("checkpoint must not be empty") + if self.dtype not in _DTYPES: + raise ValueError(f"Unsupported dtype: {self.dtype}") + if self.calculator_type not in { + "mace_mp", + "mace_off", + "mace_anicc", + }: + raise ValueError( + f"Unsupported MACE calculator type: {self.calculator_type}" + ) + if self.dispersion_cutoff <= 0: + raise ValueError("dispersion_cutoff must be positive") + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class RootstockConfig: + """Run a Rootstock-managed checkpoint through its ASE calculator.""" + + checkpoint: str + cluster: str | None = None + root: str | None = None + cache_root: str | None = None + setup_kwargs: dict[str, Any] = field(default_factory=dict) + timeout: float = 600.0 + weights: str | None = None + device: str = "cpu" + type: Literal["rootstock"] = field(default="rootstock", init=False) + + def __post_init__(self) -> None: + if not self.checkpoint: + raise ValueError("checkpoint must not be empty") + if self.cluster is not None and self.root is not None: + raise ValueError("Rootstock cannot specify both cluster and root") + if self.timeout <= 0: + raise ValueError("timeout must be positive") + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class NVAlchemiMACEConfig: + """Run MACE using NVIDIA ALCHEMI Toolkit.""" + + checkpoint: str + device: str = "cuda" + dtype: Literal["float32", "float64"] = "float32" + dt: float = 0.1 + compile_model: bool = False + enable_cueq: bool = False + type: Literal["nvalchemi-mace"] = field( + default="nvalchemi-mace", init=False + ) + + def __post_init__(self) -> None: + if not self.checkpoint: + raise ValueError("checkpoint must not be empty") + if self.dtype not in _DTYPES: + raise ValueError(f"Unsupported dtype: {self.dtype}") + if self.dt <= 0: + raise ValueError("dt must be positive") + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +MLIPBackendConfig: TypeAlias = ( + ASEMACEConfig | RootstockConfig | NVAlchemiMACEConfig +) + + +@dataclass(frozen=True) +class MLIPCalculationConfig: + """Calculation settings shared by all MLIP backends.""" + + driver: Literal["energy", "opt"] = "energy" + optimizer: Literal["bfgs", "lbfgs", "gpmin", "fire", "mdmin"] = "fire" + fmax: float = 0.01 + steps: int = 1000 + + def __post_init__(self) -> None: + if self.driver not in {"energy", "opt"}: + raise ValueError(f"Unsupported MLIP driver: {self.driver}") + if self.optimizer not in _OPTIMIZERS: + raise ValueError(f"Unsupported ASE optimizer: {self.optimizer}") + if self.fmax <= 0: + raise ValueError("fmax must be positive") + if self.steps < 1: + raise ValueError("steps must be at least 1") + + def to_dict(self) -> dict[str, Any]: + return asdict(self) diff --git a/src/matkit/mlip/runner.py b/src/matkit/mlip/runner.py new file mode 100644 index 0000000..f0647e8 --- /dev/null +++ b/src/matkit/mlip/runner.py @@ -0,0 +1,608 @@ +"""Plain-Python execution core for runtime-selectable MLIPs.""" + +from __future__ import annotations + +import json +import time +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Iterator, Sequence + +from ase.io import read as ase_read + +from matkit.mlip.config import ( + ASEMACEConfig, + MLIPBackendConfig, + MLIPCalculationConfig, + NVAlchemiMACEConfig, + RootstockConfig, +) +from matkit.types import MLIPBatchSummary, MLIPResult + + +def _atoms_payload(atoms) -> dict[str, Any]: + """Convert the portable portion of an ASE Atoms object to JSON data.""" + return { + "atomic_numbers": atoms.get_atomic_numbers().tolist(), + "positions": atoms.get_positions().tolist(), + "cell": atoms.cell.array.tolist(), + "pbc": atoms.pbc.tolist(), + } + + +def _optional_stress(atoms): + if not atoms.pbc.any() or atoms.cell.rank != 3: + return None + try: + return atoms.get_stress(voigt=False) + except Exception: + return None + + +def _synchronize_device(device: str) -> None: + if not device.startswith("cuda"): + return + try: + import torch + + torch.cuda.synchronize(torch.device(device)) + except (ImportError, RuntimeError, ValueError): + return + + +def _create_mace_calculator(config: ASEMACEConfig): + try: + import mace.calculators as mace_calculators + except ImportError as exc: + raise ImportError( + "Direct MACE requires the 'mlip' extra: pip install matkit[mlip]" + ) from exc + + try: + factory = getattr(mace_calculators, config.calculator_type) + except AttributeError as exc: + raise ImportError( + f"Installed mace-torch does not provide {config.calculator_type}." + ) from exc + + if config.calculator_type == "mace_anicc": + if config.dispersion: + raise ValueError( + "Dispersion options are supported only by " + "calculator_type='mace_mp'." + ) + return factory(device=config.device, model_path=config.checkpoint) + + kwargs: dict[str, Any] = { + "model": config.checkpoint, + "device": config.device, + "default_dtype": config.dtype, + } + if config.calculator_type == "mace_mp": + kwargs["dispersion"] = config.dispersion + if config.dispersion: + kwargs.update( + { + "damping": config.damping, + "dispersion_xc": config.dispersion_xc, + "dispersion_cutoff": config.dispersion_cutoff, + } + ) + elif config.dispersion: + raise ValueError( + "Dispersion options are supported only by calculator_type=" + "'mace_mp'." + ) + return factory(**kwargs) + + +@contextmanager +def _ase_backend_context( + config: ASEMACEConfig | RootstockConfig, +) -> Iterator[Any]: + """Create one ASE calculator and retain it for the whole request.""" + if isinstance(config, ASEMACEConfig): + yield _create_mace_calculator(config) + return + + try: + from rootstock import RootstockCalculator + except ImportError as exc: + raise ImportError( + "Rootstock requires its optional extra: " + "pip install matkit[rootstock]" + ) from exc + + kwargs = { + "checkpoint": config.checkpoint, + "cluster": config.cluster, + "root": config.root, + "cache_root": config.cache_root, + "device": config.device, + "setup_kwargs": config.setup_kwargs, + "timeout": config.timeout, + "weights": config.weights, + } + kwargs = {key: value for key, value in kwargs.items() if value is not None} + with RootstockCalculator(**kwargs) as calculator: + yield calculator + + +def _optimizer_class(name: str): + from ase.optimize import BFGS, FIRE, GPMin, LBFGS, MDMin + + return { + "bfgs": BFGS, + "lbfgs": LBFGS, + "gpmin": GPMin, + "fire": FIRE, + "mdmin": MDMin, + }[name] + + +def _success_result( + input_file: str, + backend: MLIPBackendConfig, + calculation: MLIPCalculationConfig, + atoms, + energy: float, + forces, + stress, + converged: bool, + n_steps: int | None, + calculation_time: float, +) -> dict[str, Any]: + return { + "schema_version": 1, + "success": True, + "error": "", + "input_structure_file": input_file, + "backend_info": backend.to_dict(), + "calculation_input": calculation.to_dict(), + "energy": float(energy), + "energy_unit": "eV", + "forces": None if forces is None else forces.tolist(), + "force_unit": "eV/angstrom", + "stress": None if stress is None else stress.tolist(), + "stress_unit": "eV/angstrom^3", + "converged": bool(converged), + "n_steps": n_steps, + "final_structure": _atoms_payload(atoms), + "calculation_time_s": calculation_time, + } + + +def _failure_result( + input_file: str, + backend: MLIPBackendConfig, + calculation: MLIPCalculationConfig, + error: Exception | str, + calculation_time: float = 0.0, +) -> dict[str, Any]: + return { + "schema_version": 1, + "success": False, + "error": str(error), + "input_structure_file": input_file, + "backend_info": backend.to_dict(), + "calculation_input": calculation.to_dict(), + "energy": None, + "energy_unit": "eV", + "forces": None, + "force_unit": "eV/angstrom", + "stress": None, + "stress_unit": "eV/angstrom^3", + "converged": False, + "n_steps": None, + "final_structure": None, + "calculation_time_s": calculation_time, + } + + +def _run_ase_item( + input_file: str, + atoms, + calculator, + backend: ASEMACEConfig | RootstockConfig, + calculation: MLIPCalculationConfig, +) -> dict[str, Any]: + atoms.calc = calculator + _synchronize_device(backend.device) + started = time.perf_counter() + converged = True + n_steps = 0 + if calculation.driver == "opt" and len(atoms) > 1: + optimizer = _optimizer_class(calculation.optimizer)(atoms, logfile=None) + converged = bool( + optimizer.run(fmax=calculation.fmax, steps=calculation.steps) + ) + n_steps = optimizer.nsteps + + energy = atoms.get_potential_energy() + forces = atoms.get_forces() + stress = _optional_stress(atoms) + _synchronize_device(backend.device) + elapsed = time.perf_counter() - started + return _success_result( + input_file, + backend, + calculation, + atoms, + energy, + forces, + stress, + converged, + n_steps, + elapsed, + ) + + +def _load_nvalchemi_model(config: NVAlchemiMACEConfig): + try: + import torch + from nvalchemi.models.mace import MACEWrapper + except ImportError as exc: + raise ImportError( + "NVIDIA ALCHEMI MACE requires the 'nvalchemi_mace' " + "extra and a matching CUDA extra." + ) from exc + + model = MACEWrapper.from_checkpoint( + config.checkpoint, + device=torch.device(config.device), + dtype=getattr(torch, config.dtype), + enable_cueq=config.enable_cueq, + compile_model=config.compile_model, + ) + model.eval() + return model + + +def _atoms_to_nvalchemi_data(atoms, config: NVAlchemiMACEConfig): + try: + import torch + from nvalchemi.data import AtomicData + except ImportError as exc: + raise ImportError( + "NVIDIA ALCHEMI MACE requires the 'nvalchemi_mace' " + "extra and a matching CUDA extra." + ) from exc + + dtype = getattr(torch, config.dtype) + data = AtomicData.from_atoms( + atoms, + device=torch.device(config.device), + dtype=dtype, + ) + data.forces = torch.zeros( + data.num_nodes, 3, device=data.device, dtype=dtype + ) + data.energy = torch.zeros(1, 1, device=data.device, dtype=dtype) + data.velocities = torch.zeros( + data.num_nodes, 3, device=data.device, dtype=dtype + ) + return data + + +def _nvalchemi_result( + input_file: str, + original_atoms, + data, + backend: NVAlchemiMACEConfig, + calculation: MLIPCalculationConfig, + converged: bool, + started: float, +) -> dict[str, Any]: + final_atoms = original_atoms.copy() + final_atoms.positions = data.positions.detach().cpu().numpy() + if data.cell is not None: + final_atoms.cell = data.cell.squeeze(0).detach().cpu().numpy() + if data.pbc is not None: + final_atoms.pbc = data.pbc.squeeze(0).detach().cpu().numpy() + + energy = data.energy.detach().cpu().reshape(-1)[0].item() + forces = None + if data.forces is not None: + forces = data.forces.detach().cpu().numpy() + stress = None + if data.stress is not None: + stress = data.stress.detach().cpu().reshape(-1, 3, 3)[0].numpy() + return _success_result( + input_file, + backend, + calculation, + final_atoms, + energy, + forces, + stress, + converged, + None if calculation.driver == "opt" else 0, + time.perf_counter() - started, + ) + + +def _run_nvalchemi_chunk( + model, + entries: Sequence[tuple[int, str, Any]], + backend: NVAlchemiMACEConfig, + calculation: MLIPCalculationConfig, +) -> list[tuple[int, dict[str, Any]]]: + try: + from nvalchemi.data import Batch + from nvalchemi.dynamics import BaseDynamics, ConvergenceHook, FIRE + except ImportError as exc: + raise ImportError( + "NVIDIA ALCHEMI MACE requires the 'nvalchemi_mace' " + "extra and a matching CUDA extra." + ) from exc + + started = time.perf_counter() + data_list = [ + _atoms_to_nvalchemi_data(atoms, backend) for _, _, atoms in entries + ] + batch = Batch.from_data_list(data_list) + hooks = model.make_neighbor_hooks() + convergence = None + if calculation.driver == "energy": + dynamics = BaseDynamics(model=model, hooks=hooks, n_steps=1) + else: + convergence = ConvergenceHook.from_fmax(calculation.fmax) + dynamics = FIRE( + model=model, + dt=backend.dt, + hooks=hooks, + convergence_hook=convergence, + n_steps=calculation.steps, + ) + + _synchronize_device(backend.device) + with dynamics: + batch = dynamics.run(batch) + _synchronize_device(backend.device) + + converged_indices = set(range(len(entries))) + if convergence is not None: + indices = convergence.evaluate(batch) + converged_indices = ( + set() if indices is None else set(indices.detach().cpu().tolist()) + ) + + return [ + ( + original_index, + _nvalchemi_result( + input_file, + atoms, + data, + backend, + calculation, + chunk_index in converged_indices, + started, + ), + ) + for chunk_index, ((original_index, input_file, atoms), data) in ( + enumerate(zip(entries, batch.to_data_list())) + ) + ] + + +def _chunks_by_capacity( + entries: Sequence[tuple[int, str, Any]], + batch_size: int, + max_atoms: int | None, +) -> Iterator[list[tuple[int, str, Any]]]: + chunk: list[tuple[int, str, Any]] = [] + atom_count = 0 + for entry in entries: + n_atoms = len(entry[2]) + exceeds_atoms = ( + max_atoms is not None + and bool(chunk) + and atom_count + n_atoms > max_atoms + ) + if len(chunk) >= batch_size or exceeds_atoms: + yield chunk + chunk = [] + atom_count = 0 + chunk.append(entry) + atom_count += n_atoms + if chunk: + yield chunk + + +def _read_inputs( + input_files: Sequence[str | Path], + backend: MLIPBackendConfig, + calculation: MLIPCalculationConfig, +) -> tuple[list[tuple[int, str, Any]], list[dict[str, Any] | None]]: + prepared = [] + results: list[dict[str, Any] | None] = [None] * len(input_files) + for index, value in enumerate(input_files): + input_file = str(Path(value).expanduser().resolve()) + try: + if not Path(input_file).is_file(): + raise FileNotFoundError( + f"Input structure file does not exist: {value}" + ) + prepared.append((index, input_file, ase_read(input_file))) + except Exception as exc: + results[index] = _failure_result( + input_file, backend, calculation, exc + ) + return prepared, results + + +def _execute_inputs( + input_files: Sequence[str | Path], + backend: MLIPBackendConfig, + calculation: MLIPCalculationConfig, + batch_size: int, + max_atoms: int | None, +) -> tuple[list[dict[str, Any]], float]: + if not input_files: + raise ValueError("At least one input structure file is required") + if batch_size < 1: + raise ValueError("batch_size must be at least 1") + if max_atoms is not None and max_atoms < 1: + raise ValueError("max_atoms must be at least 1") + if ( + isinstance(backend, NVAlchemiMACEConfig) + and calculation.optimizer != "fire" + ): + raise ValueError("NVIDIA ALCHEMI supports only the FIRE optimizer") + + prepared, results = _read_inputs(input_files, backend, calculation) + if not prepared: + return [result for result in results if result is not None], 0.0 + + setup_started = time.perf_counter() + if isinstance(backend, (ASEMACEConfig, RootstockConfig)): + try: + with _ase_backend_context(backend) as calculator: + _synchronize_device(backend.device) + setup_time = time.perf_counter() - setup_started + for index, input_file, atoms in prepared: + try: + results[index] = _run_ase_item( + input_file, + atoms, + calculator, + backend, + calculation, + ) + except Exception as exc: + results[index] = _failure_result( + input_file, backend, calculation, exc + ) + except Exception as exc: + setup_time = time.perf_counter() - setup_started + for index, input_file, _ in prepared: + results[index] = _failure_result( + input_file, backend, calculation, exc + ) + else: + try: + model = _load_nvalchemi_model(backend) + _synchronize_device(backend.device) + setup_time = time.perf_counter() - setup_started + for chunk in _chunks_by_capacity(prepared, batch_size, max_atoms): + try: + chunk_results = _run_nvalchemi_chunk( + model, chunk, backend, calculation + ) + if len(chunk_results) != len(chunk): + raise RuntimeError( + "NVIDIA ALCHEMI returned a different number of " + "results than input structures." + ) + for index, result in chunk_results: + results[index] = result + except Exception as exc: + for index, input_file, _ in chunk: + results[index] = _failure_result( + input_file, backend, calculation, exc + ) + except Exception as exc: + setup_time = time.perf_counter() - setup_started + for index, input_file, _ in prepared: + results[index] = _failure_result( + input_file, backend, calculation, exc + ) + + final_results = [] + for result in results: + if result is None: + raise RuntimeError("Internal MLIP result ordering error") + final_results.append(result) + return final_results, setup_time + + +def _write_json(path: Path, data: dict[str, Any]) -> str: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(data, indent=2), encoding="utf-8") + return str(path.resolve()) + + +def run_mlip( + input_file: str | Path, + backend: MLIPBackendConfig, + calculation: MLIPCalculationConfig | None = None, + output_file: str | Path | None = None, +) -> MLIPResult: + """Run one energy calculation or fixed-cell optimization.""" + calculation = calculation or MLIPCalculationConfig() + results, setup_time = _execute_inputs( + [input_file], backend, calculation, batch_size=1, max_atoms=None + ) + result = results[0] + result["setup_time_s"] = setup_time + if output_file is not None: + result["output_results_file"] = _write_json( + Path(output_file).expanduser(), result + ) + return result + + +def run_mlip_batch( + input_files: Sequence[str | Path], + backend: MLIPBackendConfig, + calculation: MLIPCalculationConfig | None = None, + output_dir: str | Path = "mlip_results", + batch_size: int = 16, + max_atoms: int | None = None, +) -> MLIPBatchSummary: + """Run an ordered MLIP batch and persist results plus a manifest.""" + calculation = calculation or MLIPCalculationConfig() + started = time.perf_counter() + output_path = Path(output_dir).expanduser() + output_path.mkdir(parents=True, exist_ok=True) + results, setup_time = _execute_inputs( + input_files, + backend, + calculation, + batch_size=batch_size, + max_atoms=max_atoms, + ) + + items = [] + for index, result in enumerate(results): + stem = Path(result["input_structure_file"]).stem + result_path = output_path / f"{index:05d}_{stem}.json" + result["setup_time_s"] = setup_time + result_file = _write_json(result_path, result) + items.append( + { + "index": index, + "input_structure_file": result["input_structure_file"], + "status": "success" if result["success"] else "failure", + "result_file": result_file, + "error": result["error"], + } + ) + + succeeded = sum(item["status"] == "success" for item in items) + failed = len(items) - succeeded + if items and failed == 0: + status = "completed" + elif succeeded: + status = "partial" + else: + status = "failure" + manifest = { + "schema_version": 1, + "status": status, + "backend_info": backend.to_dict(), + "calculation_input": calculation.to_dict(), + "setup_time_s": setup_time, + "wall_time_s": time.perf_counter() - started, + "total": len(items), + "succeeded": succeeded, + "failed": failed, + "items": items, + } + manifest_file = _write_json(output_path / "batch_manifest.json", manifest) + return { + **manifest, + "manifest_file": manifest_file, + "results": results, + } diff --git a/src/matkit/types.py b/src/matkit/types.py index b463428..19dc838 100644 --- a/src/matkit/types.py +++ b/src/matkit/types.py @@ -73,3 +73,42 @@ class UMABatchResult(TypedDict): final_energy: Optional[float] n_steps: Optional[int] error_message: Optional[str] + + +class MLIPResult(TypedDict): + """Runtime-neutral result from ``matkit.mlip.run_mlip``.""" + + schema_version: int + success: bool + error: str + input_structure_file: str + backend_info: dict + calculation_input: dict + energy: Optional[float] + energy_unit: str + forces: Optional[list[list[float]]] + force_unit: str + stress: Optional[list[list[float]]] + stress_unit: str + converged: bool + n_steps: Optional[int] + final_structure: Optional[dict] + calculation_time_s: float + setup_time_s: float + + +class MLIPBatchSummary(TypedDict): + """Persistent summary from ``matkit.mlip.run_mlip_batch``.""" + + schema_version: int + status: str + backend_info: dict + calculation_input: dict + setup_time_s: float + wall_time_s: float + total: int + succeeded: int + failed: int + items: list[dict] + manifest_file: str + results: list[MLIPResult] diff --git a/tests/test_cli.py b/tests/test_cli.py index 342a4d2..6d30282 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,5 +1,7 @@ """Tests for matkit CLI.""" +import json + from click.testing import CliRunner from matkit.cli import main @@ -59,3 +61,80 @@ def test_unknown_command(self): runner = CliRunner() result = runner.invoke(main, ["nonexistent"]) assert result.exit_code != 0 + + def test_mlip_group_exposes_runtime_neutral_commands(self): + runner = CliRunner() + result = runner.invoke(main, ["mlip", "--help"]) + assert result.exit_code == 0 + assert "run" in result.output + assert "run-batch" in result.output + + def test_mlip_run_delegates_to_python_api( + self, sample_cif, tmp_path, monkeypatch + ): + import matkit.mlip + + received = {} + + def fake_run(input_file, backend, calculation, output_file): + received.update( + { + "input": input_file, + "backend": backend, + "calculation": calculation, + "output": output_file, + } + ) + return { + "success": True, + "energy": -1.25, + "energy_unit": "eV", + "converged": True, + } + + monkeypatch.setattr(matkit.mlip, "run_mlip", fake_run) + output = tmp_path / "result.json" + runner = CliRunner() + result = runner.invoke( + main, + [ + "mlip", + "run", + "--input", + sample_cif, + "--output", + str(output), + "--backend", + "rootstock", + "--checkpoint", + "mace-mp-0-medium", + "--cluster", + "polaris", + "--device", + "cuda", + "--setup-kwarg", + 'default_dtype="float32"', + ], + ) + + assert result.exit_code == 0 + assert json.loads(result.output)["energy"] == -1.25 + assert received["backend"].cluster == "polaris" + assert received["backend"].setup_kwargs == {"default_dtype": "float32"} + assert received["output"] == str(output) + + def test_mlip_batch_requires_one_input_source(self): + runner = CliRunner() + result = runner.invoke( + main, + [ + "mlip", + "run-batch", + "--backend", + "ase-mace", + "--checkpoint", + "medium", + ], + ) + assert result.exit_code != 0 + assert "exactly one" in result.output diff --git a/tests/test_mlip_runner.py b/tests/test_mlip_runner.py new file mode 100644 index 0000000..495b645 --- /dev/null +++ b/tests/test_mlip_runner.py @@ -0,0 +1,312 @@ +"""Tests for runtime-selectable, agent-free MLIP execution.""" + +import json +import sys +from contextlib import contextmanager +from types import SimpleNamespace + +import numpy as np +import pytest +from ase import Atoms +from ase.calculators.emt import EMT +from ase.io import write + +from matkit.mlip import ( + ASEMACEConfig, + MLIPCalculationConfig, + NVAlchemiMACEConfig, + RootstockConfig, + run_mlip, + run_mlip_batch, +) +from matkit.mlip import runner as mlip_runner + + +def _write_copper(path, distance=2.5): + atoms = Atoms( + "Cu2", + positions=[[0.0, 0.0, 0.0], [distance, 0.0, 0.0]], + cell=[8.0, 8.0, 8.0], + pbc=False, + ) + write(path, atoms) + return atoms + + +def _emt_context(counter=None): + @contextmanager + def calculator_context(_config): + if counter is not None: + counter["entered"] += 1 + yield EMT() + + return calculator_context + + +def test_backend_config_validation(): + with pytest.raises(ValueError, match="both cluster and root"): + RootstockConfig( + checkpoint="mace-mp-0-medium", + cluster="polaris", + root="/shared/rootstock", + ) + with pytest.raises(ValueError, match="dt must be positive"): + NVAlchemiMACEConfig(checkpoint="medium", dt=0) + with pytest.raises(ValueError, match="steps must be at least 1"): + MLIPCalculationConfig(steps=0) + + +def test_mace_anicc_uses_model_path_signature(monkeypatch): + received = {} + + def mace_anicc(**kwargs): + received.update(kwargs) + return "calculator" + + monkeypatch.setitem( + sys.modules, + "mace.calculators", + SimpleNamespace(mace_anicc=mace_anicc), + ) + monkeypatch.setitem( + sys.modules, + "mace", + SimpleNamespace(calculators=sys.modules["mace.calculators"]), + ) + + calculator = mlip_runner._create_mace_calculator( + ASEMACEConfig( + checkpoint="ani.model", + calculator_type="mace_anicc", + device="cuda", + ) + ) + + assert calculator == "calculator" + assert received == {"device": "cuda", "model_path": "ani.model"} + + +def test_run_mlip_energy_writes_runtime_neutral_result(tmp_path, monkeypatch): + input_file = tmp_path / "copper.xyz" + output_file = tmp_path / "result.json" + _write_copper(input_file) + monkeypatch.setattr(mlip_runner, "_ase_backend_context", _emt_context()) + + result = run_mlip( + input_file, + ASEMACEConfig(checkpoint="unused"), + output_file=output_file, + ) + + assert result["success"] is True + assert result["energy"] is not None + assert len(result["forces"]) == 2 + assert result["stress"] is None + assert result["n_steps"] == 0 + stored = json.loads(output_file.read_text()) + assert stored["backend_info"]["type"] == "ase-mace" + assert stored["final_structure"]["atomic_numbers"] == [29, 29] + + +def test_periodic_energy_includes_full_stress(tmp_path, monkeypatch): + input_file = tmp_path / "periodic.xyz" + atoms = Atoms( + "Cu", + positions=[[0.0, 0.0, 0.0]], + cell=[3.6, 3.6, 3.6], + pbc=True, + ) + write(input_file, atoms) + monkeypatch.setattr(mlip_runner, "_ase_backend_context", _emt_context()) + + result = run_mlip( + input_file, + ASEMACEConfig(checkpoint="unused"), + ) + + assert result["success"] is True + assert np.asarray(result["stress"]).shape == (3, 3) + + +def test_run_mlip_fixed_cell_optimization(tmp_path, monkeypatch): + input_file = tmp_path / "copper.xyz" + initial = _write_copper(input_file, distance=3.0) + monkeypatch.setattr(mlip_runner, "_ase_backend_context", _emt_context()) + + result = run_mlip( + input_file, + ASEMACEConfig(checkpoint="unused"), + MLIPCalculationConfig(driver="opt", steps=2), + ) + + assert result["success"] is True + assert result["n_steps"] <= 2 + assert result["final_structure"]["cell"] == initial.cell.tolist() + + +def test_batch_reuses_calculator_and_preserves_failures(tmp_path, monkeypatch): + first = tmp_path / "first.xyz" + second = tmp_path / "second.xyz" + missing = tmp_path / "missing.xyz" + _write_copper(first) + _write_copper(second, distance=2.7) + counter = {"entered": 0} + monkeypatch.setattr( + mlip_runner, + "_ase_backend_context", + _emt_context(counter), + ) + + summary = run_mlip_batch( + [first, missing, second], + ASEMACEConfig(checkpoint="unused"), + output_dir=tmp_path / "results", + ) + + assert summary["status"] == "partial" + assert summary["succeeded"] == 2 + assert summary["failed"] == 1 + assert counter["entered"] == 1 + assert [item["index"] for item in summary["items"]] == [0, 1, 2] + assert [item["status"] for item in summary["items"]] == [ + "success", + "failure", + "success", + ] + manifest = json.loads( + (tmp_path / "results" / "batch_manifest.json").read_text() + ) + assert manifest["status"] == "partial" + assert all( + (tmp_path / "results" / name).exists() + for name in ( + "00000_first.json", + "00001_missing.json", + "00002_second.json", + ) + ) + + +def test_rootstock_context_forwards_options_and_closes(monkeypatch): + events = [] + + class FakeRootstockCalculator: + def __init__(self, **kwargs): + events.append(("init", kwargs)) + + def __enter__(self): + events.append(("enter", None)) + return "calculator" + + def __exit__(self, exc_type, exc, traceback): + events.append(("exit", None)) + + monkeypatch.setitem( + sys.modules, + "rootstock", + SimpleNamespace(RootstockCalculator=FakeRootstockCalculator), + ) + config = RootstockConfig( + checkpoint="mace-mp-0-medium", + cluster="polaris", + setup_kwargs={"default_dtype": "float32"}, + timeout=1200, + device="cuda", + ) + + with mlip_runner._ase_backend_context(config) as calculator: + assert calculator == "calculator" + + assert [event[0] for event in events] == ["init", "enter", "exit"] + kwargs = events[0][1] + assert kwargs["checkpoint"] == "mace-mp-0-medium" + assert kwargs["cluster"] == "polaris" + assert kwargs["setup_kwargs"] == {"default_dtype": "float32"} + assert kwargs["timeout"] == 1200 + + +def test_nvalchemi_loads_once_and_chunks_in_order(tmp_path, monkeypatch): + input_files = [] + for index in range(3): + input_file = tmp_path / f"input_{index}.xyz" + _write_copper(input_file, distance=2.5 + index * 0.1) + input_files.append(input_file) + + loaded = [] + chunks = [] + sentinel_model = object() + + def load_model(config): + loaded.append(config.checkpoint) + return sentinel_model + + def run_chunk(model, entries, backend, calculation): + assert model is sentinel_model + chunks.append([entry[1] for entry in entries]) + return [ + ( + index, + mlip_runner._success_result( + input_file, + backend, + calculation, + atoms, + 1.25, + np.zeros((len(atoms), 3)), + None, + True, + 0, + 0.01, + ), + ) + for index, input_file, atoms in entries + ] + + monkeypatch.setattr(mlip_runner, "_load_nvalchemi_model", load_model) + monkeypatch.setattr(mlip_runner, "_run_nvalchemi_chunk", run_chunk) + summary = run_mlip_batch( + input_files, + NVAlchemiMACEConfig(checkpoint="medium"), + output_dir=tmp_path / "results", + batch_size=3, + max_atoms=4, + ) + + assert summary["status"] == "completed" + assert loaded == ["medium"] + assert [len(chunk) for chunk in chunks] == [2, 1] + assert chunks[0] + chunks[1] == [ + str(path.resolve()) for path in input_files + ] + + +def test_nvalchemi_missing_backend_is_persisted(tmp_path, monkeypatch): + input_file = tmp_path / "input.xyz" + output_file = tmp_path / "failure.json" + _write_copper(input_file) + + def missing_backend(_config): + raise ImportError("Install the nvalchemi_mace extra") + + monkeypatch.setattr(mlip_runner, "_load_nvalchemi_model", missing_backend) + result = run_mlip( + input_file, + NVAlchemiMACEConfig(checkpoint="medium"), + output_file=output_file, + ) + + assert result["success"] is False + assert "nvalchemi_mace" in result["error"] + assert json.loads(output_file.read_text())["success"] is False + + +def test_nvalchemi_rejects_non_fire_optimizer(tmp_path): + input_file = tmp_path / "input.xyz" + _write_copper(input_file) + + with pytest.raises(ValueError, match="only the FIRE optimizer"): + run_mlip( + input_file, + NVAlchemiMACEConfig(checkpoint="medium"), + MLIPCalculationConfig(driver="opt", optimizer="bfgs"), + )