#!/usr/bin/env python3
"""MatMelt screen runner: MACE vc-relax + classify ONE candidate structure.
Self-contained by design -- imports **stdlib + mace/ase only, never
matmelt** -- because it is staged to the compute backend per job and
executed there by the user's own python; remotes never receive our package
code. It is equally runnable by hand for isolated debugging::
python screen_runner.py --cif structure.cif
python screen_runner.py --help
It reads one CIF, runs a variable-cell relaxation (LBFGS over a cell filter
-- forces and stress minimized together, the vc-relax equivalent), then
classifies the result on physical-sanity grounds only::
screen_result.json # category RELAXED|FAILED + failure_reasons + metrics
relaxed.cif # written on RELAXED only (a stale one is removed)
Classification is convergence-based (energy window, finite energy, bond
overlap, volume change, residual stress, final fmax, energy rise) -- NOT a
stability threshold: a high-energy metastable structure is information, not
a rejection. ``float64`` is the default dtype -- required for a reliable
stress tensor in vc-relax.
A ``FAILED`` category is a normal completed screening (exit 0) -- the stage
turns it into a terminal science verdict. Exit codes: ``0`` = screening
completed (either category); ``2`` = bad arguments or unreadable ``--cif``;
``3`` = mace/ase not importable in this python environment (stderr says
exactly what to install); ``1`` = unexpected error, e.g. a crashed
relaxation (traceback on stderr -- the stage treats it as a retryable
fault, never a science rejection).
"""
from __future__ import annotations
import argparse
import json
import math
import sys
from pathlib import Path
from typing import Any
VERDICT_NAME = "screen_result.json"
RELAXED_NAME = "relaxed.cif"
#: eV/A^3 -> GPa.
EV_PER_A3_TO_GPA = 160.21766208
#: Physical-sanity bounds. Deliberately constants, not CLI knobs (only the
#: model and the main relaxation settings are configurable): they flag MACE
#: extrapolation / unphysical geometries, and loosening them per run would
#: make verdicts incomparable across a database.
SANITY = {
"min_bond_A": 1.2,
"max_stress_GPa": 2.0,
"max_vol_change_pct": 60.0,
"energy_floor_eV_atom": -35.0,
"energy_ceil_eV_atom": 10.0,
"fmax_acceptance_eV_A": 0.5,
"max_energy_rise_eV_atom": 0.01,
}
def _parse_args(argv: list[str] | None) -> argparse.Namespace:
parser = argparse.ArgumentParser(
prog="screen_runner",
description="MACE vc-relax one candidate structure and classify it "
"RELAXED or FAILED on physical-sanity grounds.",
)
parser.add_argument(
"--cif",
default="structure.cif",
metavar="PATH",
help="Candidate structure to screen (default: structure.cif).",
)
parser.add_argument(
"--out-dir",
default=".",
metavar="PATH",
help="Directory to write screen_result.json / relaxed.cif into (default: cwd).",
)
parser.add_argument(
"--model-path",
default="",
metavar="PATH",
help="Local MACE model file; empty (default) uses the MACE-MP foundation "
"model from the mace cache (downloaded on first use, disclosed below).",
)
parser.add_argument(
"--model-size",
default="medium",
choices=["small", "medium", "large"],
help="MACE-MP foundation model size when --model-path is empty.",
)
parser.add_argument(
"--device",
default="cpu",
help='Compute device: "cpu" | "cuda" | "mps" | "auto".',
)
parser.add_argument(
"--dtype",
default="float64",
choices=["float64", "float32"],
help="Torch dtype; float64 (default) is required for a reliable stress tensor.",
)
parser.add_argument(
"--fmax",
type=float,
default=0.05,
metavar="EV_PER_A",
help="Force convergence target for the relaxation (eV/A).",
)
parser.add_argument(
"--max-steps",
type=int,
default=200,
metavar="N",
help="Maximum optimizer steps for the relaxation.",
)
args = parser.parse_args(argv)
if args.fmax <= 0:
parser.error(f"--fmax must be > 0, got {args.fmax}")
if args.max_steps < 1:
parser.error(f"--max-steps must be >= 1, got {args.max_steps}")
return args
# ---------------------------------------------------------------------------
# Pure classification (no science imports -- fully unit-testable)
# ---------------------------------------------------------------------------
[docs]
def classify(initial: dict[str, Any], final: dict[str, Any]) -> tuple[str, list[str]]:
"""Classify a finished relaxation as ``RELAXED`` or ``FAILED``.
A pure function over plain metric dicts, so every bound is unit-testable
with no science stack installed.
Parameters
----------
initial
Observables before relaxation. Needs ``e_eV_atom`` and ``volume_A3``.
final
Observables after relaxation. Needs ``e_eV_atom``, ``fmax_eV_A`` and
``volume_A3``; ``stress_max_GPa`` and ``min_bond_A`` may be ``None``,
which **skips** that check rather than failing the structure -- an
observable the calculator could not supply is not evidence against a
crystal.
Returns
-------
tuple
``(category, reasons)`` where category is ``"RELAXED"`` when no check
fired and ``"FAILED"`` otherwise, and ``reasons`` holds one
human-readable sentence per failed check.
Notes
-----
The checks are convergence and physical sanity: the energy window, NaN
energies, overlapping atoms, runaway volume change, residual stress, the
fmax acceptance gate, and an optimizer that walked uphill. Thresholds live
in the module's ``SANITY`` constants rather than in CLI options, so
verdicts stay comparable across a whole database.
Deliberately absent: any rejection on *absolute* energy ranking.
Metastable is information, not failure.
"""
reasons: list[str] = []
e_final = float(final["e_eV_atom"])
fmax_final = float(final["fmax_eV_A"])
floor = SANITY["energy_floor_eV_atom"]
ceil = SANITY["energy_ceil_eV_atom"]
if not (floor < e_final < ceil): # NaN lands here too: comparisons are False
reasons.append(
f"energy {e_final:.4f} eV/atom outside the physical window "
f"[{floor}, {ceil}] eV/atom -- the model is almost certainly extrapolating"
)
if not math.isfinite(e_final):
reasons.append("energy is NaN or Inf -- numerical failure in the model")
min_bond = final.get("min_bond_A")
if min_bond is not None and float(min_bond) < SANITY["min_bond_A"]:
reasons.append(
f"minimum interatomic distance {float(min_bond):.4f} A < "
f"{SANITY['min_bond_A']} A -- atoms are overlapping (unphysical geometry)"
)
vol_initial = float(initial["volume_A3"])
vol_final = float(final["volume_A3"])
vol_pct = (vol_final - vol_initial) / vol_initial * 100.0
if abs(vol_pct) > SANITY["max_vol_change_pct"]:
reasons.append(
f"volume changed by {vol_pct:+.1f}% during relaxation "
f"(threshold +/-{SANITY['max_vol_change_pct']}%) -- "
"structure collapsed or exploded"
)
stress_max = final.get("stress_max_GPa")
if stress_max is not None and float(stress_max) > SANITY["max_stress_GPa"]:
reasons.append(
f"residual stress {float(stress_max):.3f} GPa > "
f"{SANITY['max_stress_GPa']} GPa -- unit cell not mechanically equilibrated"
)
if fmax_final > SANITY["fmax_acceptance_eV_A"]:
reasons.append(
f"final fmax {fmax_final:.5f} eV/A > acceptance gate "
f"{SANITY['fmax_acceptance_eV_A']} eV/A -- no local minimum reached"
)
e_drop = float(initial["e_eV_atom"]) - e_final
if e_drop < -SANITY["max_energy_rise_eV_atom"]:
reasons.append(
f"energy rose by {abs(e_drop):.5f} eV/atom during relaxation "
"-- the optimizer walked uphill, suggesting gradient confusion"
)
return ("FAILED" if reasons else "RELAXED"), reasons
[docs]
def write_verdict(verdict: dict[str, Any], out_dir: Path) -> None:
"""Write the verdict atomically, scrubbing any stale relaxed structure.
Parameters
----------
verdict
The document to write.
out_dir
The job's working directory.
Notes
-----
A non-``RELAXED`` verdict must never sit next to a leftover
``relaxed.cif`` from an earlier attempt, because everything downstream
trusts that pairing -- so the stale file is removed first.
"""
if verdict.get("category") != "RELAXED":
(out_dir / RELAXED_NAME).unlink(missing_ok=True)
target = out_dir / VERDICT_NAME
tmp = target.with_name(VERDICT_NAME + ".tmp")
tmp.write_text(json.dumps(verdict, indent=2), encoding="utf-8")
tmp.replace(target)
# ---------------------------------------------------------------------------
# The science (lazy imports -- everything below needs mace/ase installed)
# ---------------------------------------------------------------------------
def _load_science() -> dict[str, Any]:
"""Import mace and ase -- the single touchpoint, deliberately lazy.
Returns
-------
dict
The symbols the rest of the module uses, passed around explicitly so
nothing else in this file references a science package by name.
Raises
------
ImportError
When mace-torch is not installed in this interpreter. ``main`` catches
it and exits 3 with the exact install command.
Notes
-----
The cell filter is looked up with a fallback: ``FrechetCellFilter`` on ASE
3.23 and newer, ``ExpCellFilter`` on older releases.
"""
from ase.io import read as ase_read
from ase.io import write as ase_write
from ase.neighborlist import neighbor_list
from ase.optimize import LBFGS
from mace.calculators import MACECalculator, mace_mp
try:
from ase.filters import FrechetCellFilter as CellFilter # ASE >= 3.23
except ImportError:
from ase.constraints import ExpCellFilter as CellFilter # older ASE
return {
"ase_read": ase_read,
"ase_write": ase_write,
"neighbor_list": neighbor_list,
"LBFGS": LBFGS,
"MACECalculator": MACECalculator,
"mace_mp": mace_mp,
"CellFilter": CellFilter,
}
def _load_calculator(sci: dict[str, Any], args: argparse.Namespace) -> Any:
"""Build the MACE calculator, from a local model or the foundation model.
Parameters
----------
sci
The loaded science symbols.
args
Parsed arguments; ``model_path``, ``model_size``, ``device`` and
``dtype`` are read.
Returns
-------
Any
A MACE calculator ready to attach to an ``Atoms`` object.
Notes
-----
With no ``model_path``, the MACE-MP foundation model is used and its first
use downloads it into the mace cache. That is **disclosed on stdout, never
silent**; the preflight offers the up-front consent.
"""
if args.model_path:
print(f"loading local MACE model: {args.model_path}")
return sci["MACECalculator"](
model_paths=args.model_path, device=args.device, default_dtype=args.dtype
)
# Disclosed, never silent: first use downloads the foundation model into
# the mace cache (the 3.6 preflight adds the up-front consent offer).
print(
f"no model_path configured: using the MACE-MP foundation model "
f"(size={args.model_size}); first use downloads it to the mace cache"
)
return sci["mace_mp"](
model=args.model_size,
device=args.device,
default_dtype=args.dtype,
dispersion=False,
)
def _evaluate(sci: dict[str, Any], atoms: Any) -> dict[str, Any]:
"""Measure the current energy, forces, stress and geometry.
Parameters
----------
sci
The loaded science symbols.
atoms
The structure to measure, with a calculator already attached.
Returns
-------
dict
Plain floats -- ``e_eV``, ``e_eV_atom``, ``fmax_eV_A``,
``stress_max_GPa``, ``volume_A3``, ``min_bond_A`` -- ready for
``classify``. ``stress_max_GPa`` is ``None`` when the calculator cannot
supply a stress tensor or returns a non-finite one, and ``min_bond_A``
is ``None`` for a structure with no neighbour pairs within the cutoff.
Neither is fatal: an unavailable observable simply skips its check.
"""
e_total = float(atoms.get_potential_energy())
n_atoms = len(atoms)
forces = atoms.get_forces()
fmax = float(max((sum(f * f for f in row)) ** 0.5 for row in forces.tolist()))
stress_max: float | None
try:
stress = atoms.get_stress(voigt=True)
stress_max = float(max(abs(float(c)) for c in stress)) * EV_PER_A3_TO_GPA
if not math.isfinite(stress_max):
stress_max = None
except Exception: # stress unavailable -> observable skipped, never fatal
stress_max = None
distances = sci["neighbor_list"]("d", atoms, cutoff=8.0)
min_bond = float(min(distances)) if len(distances) else None
return {
"e_eV": e_total,
"e_eV_atom": e_total / n_atoms,
"fmax_eV_A": fmax,
"stress_max_GPa": stress_max,
"volume_A3": float(atoms.get_volume()),
"min_bond_A": min_bond,
}
[docs]
def screen(sci: dict[str, Any], args: argparse.Namespace, out_dir: Path) -> dict[str, Any]:
"""Read, variable-cell relax, and classify one structure.
Parameters
----------
sci
The loaded science symbols.
args
Parsed arguments: the input CIF, the model settings, and the
convergence targets.
out_dir
The job's working directory, where ``relaxed.cif`` is written on a
``RELAXED`` verdict.
Returns
-------
dict
The verdict document: category, failure reasons, whether the optimizer
converged, the step count, the initial and final observables, and the
settings used -- so a stored verdict is reproducible from its own
record.
Notes
-----
The relaxation is LBFGS over a cell filter, i.e. atoms *and* cell relax
together. Convergence is not itself the verdict: a run that stops at the
step limit can still pass ``classify``'s acceptance gate, and one that
converges can still be rejected as unphysical.
"""
atoms = sci["ase_read"](args.cif)
atoms.calc = _load_calculator(sci, args)
initial = _evaluate(sci, atoms)
print(
f"initial: E = {initial['e_eV_atom']:.6f} eV/atom, "
f"fmax = {initial['fmax_eV_A']:.5f} eV/A, V = {initial['volume_A3']:.2f} A^3"
)
cell_filter = sci["CellFilter"](atoms)
optimizer = sci["LBFGS"](cell_filter, logfile=None, memory=10)
converged = bool(optimizer.run(fmax=args.fmax, steps=args.max_steps))
n_steps = int(optimizer.nsteps)
final = _evaluate(sci, atoms)
print(
f"relaxed ({n_steps} steps, converged={converged}): "
f"E = {final['e_eV_atom']:.6f} eV/atom, fmax = {final['fmax_eV_A']:.5f} eV/A, "
f"V = {final['volume_A3']:.2f} A^3"
)
category, reasons = classify(initial, final)
if category == "RELAXED":
sci["ase_write"](str(out_dir / RELAXED_NAME), atoms)
return {
"category": category,
"failure_reasons": reasons,
"converged": converged,
"n_steps": n_steps,
"initial": initial,
"final": final,
"settings": {
"model": args.model_path or f"mace-mp:{args.model_size}",
"device": args.device,
"dtype": args.dtype,
"fmax": args.fmax,
"max_steps": args.max_steps,
},
}
[docs]
def main(argv: list[str] | None = None) -> int:
"""Screen one candidate structure.
Parameters
----------
argv
Command-line arguments. ``None`` reads ``sys.argv``.
Returns
-------
int
The process exit code: ``0`` when screening *completed*, whichever way
the verdict went -- a ``FAILED`` verdict is a successful screening, and
the science lives in the written document; ``2`` for a usage error or
an unreadable structure file; ``3`` when mace-torch is missing, with
the exact ``pip install`` command and this interpreter's path printed
to stderr.
Notes
-----
A *crashed* relaxation deliberately propagates as an unhandled exception
(exit 1 with a traceback). The stage reads that as a retryable fault, which
is what keeps an infrastructure failure from ever being recorded as a
science rejection.
"""
args = _parse_args(argv)
cif_path = Path(args.cif)
if not cif_path.is_file():
print(f"cannot read structure file {cif_path}: not a file", file=sys.stderr)
return 2
try:
sci = _load_science()
except ImportError as exc:
print(
"screen_runner needs the 'mace-torch' package (which brings ase and "
f"torch), not importable in this python environment ({sys.executable}).\n"
f"Install it there with: pip install mace-torch\n(import error: {exc})",
file=sys.stderr,
)
return 3
out_dir = Path(args.out_dir)
# A crashed relaxation propagates (exit 1, traceback): the stage treats it
# as a retryable fault -- it must never be recorded as a science FAILED.
verdict = screen(sci, args, out_dir)
write_verdict(verdict, out_dir)
print(f"{cif_path} -> {verdict['category']}")
return 0
if __name__ == "__main__": # pragma: no cover — exercised via subprocess
sys.exit(main())