Add MLIP backends and GPU examples

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

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

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.

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."

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}"