#!/usr/bin/env python3
"""MatMelt cif runner: fetch DFTB+ relaxed geometries out of AiiDA (M10).
Self-contained by design -- imports **stdlib + aiida only, never
matmelt** -- because it is executed by the user's own python (the one
their AiiDA lives in), exactly like the stage runners beside it. It is the
one exception to the "one runner per funnel stage" rule: it belongs to the
``cif`` command, not to a stage, and is never staged to a compute backend.
Why it exists: ``calc_result.json`` records every healing attempt's *input*
parameters, so the geometry embedded there is what DFTB+ *started* from
(the MACE-relaxed cell from the screen stage). The geometry DFTB+ relaxed
*to* is ``geom.out.gen`` in the final CalcJob's retrieved folder, which
lives only in the AiiDA repository. This runner reads it back.
It works in batches -- one profile load serves thousands of materials::
python cif_runner.py --batch batch.json --out geometries.json
python cif_runner.py --id Cl1Na1_225 --wc-pk 13524 --uuid abc-123
python cif_runner.py --help
``--batch`` takes a JSON list of ``{"id", "wc_pk", "uuid"?, "calc_pk"?}``
entries. The output document is ``{"results": [{"id", "wc_pk", "calc_pk",
"gen", "error", "note"}, ...]}`` -- one row per requested material, in the
requested order. A material that cannot be resolved gets an ``error``
string and a ``null`` geometry: one bad material never costs the batch, and
the caller decides what a failure means.
Resolution mirrors the calculate runner's own extraction: load the
WorkChain, verify its uuid when one was recorded (pks are per-profile --
this is what stops a profile switch from exporting a stranger's structure),
walk its ``CALL_CALC`` CalcJob children, and take the most recent by
creation time -- the attempt that converged. A recorded ``calc_pk`` that
disagrees with it is reported in ``note``; the database, not the record,
decides.
Exit codes: ``0`` = the batch ran and the document was written (per-material
failures live inside it); ``2`` = bad arguments or an unreadable/malformed
batch file; ``3`` = aiida not importable in this python environment (stderr
says exactly what to install); ``1`` = the profile could not be loaded or
the output could not be written.
"""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
from typing import Any
#: What DFTB+ writes the relaxed geometry to, given our HSD template's
#: ``OutputPrefix = "geom.out"`` -- and what the plugin retrieves.
GEOMETRY_NAME = "geom.out.gen"
def _parse_args(argv: list[str] | None) -> argparse.Namespace:
parser = argparse.ArgumentParser(
prog="cif_runner",
description="Fetch the DFTB+ relaxed geometry (geom.out.gen) of one or many "
"finished mydftbp.single WorkChains from the user's AiiDA profile.",
)
parser.add_argument(
"--batch",
default=None,
metavar="PATH",
help='JSON list of {"id", "wc_pk", "uuid"?, "calc_pk"?} entries to fetch.',
)
parser.add_argument(
"--id",
default=None,
metavar="MATERIAL_ID",
help="Single-material mode: the material id (needs --wc-pk).",
)
parser.add_argument(
"--wc-pk",
type=int,
default=None,
metavar="PK",
help="Single-material mode: pk of the mydftbp.single WorkChain.",
)
parser.add_argument(
"--uuid",
default=None,
metavar="UUID",
help="Single-material mode: the WorkChain uuid to verify against the pk.",
)
parser.add_argument(
"--calc-pk",
type=int,
default=None,
metavar="PK",
help="Single-material mode: the expected final CalcJob pk (cross-checked).",
)
parser.add_argument(
"--geometry-name",
default=GEOMETRY_NAME,
metavar="FILE",
help=f"Retrieved file holding the relaxed geometry (default: {GEOMETRY_NAME}).",
)
parser.add_argument(
"--profile",
default=None,
metavar="NAME",
help="AiiDA profile to load (default: the user's default profile).",
)
parser.add_argument(
"--out",
default=None,
metavar="PATH",
help="Write the result document here (default: stdout).",
)
args = parser.parse_args(argv)
if (args.batch is None) == (args.id is None):
parser.error("give either --batch or --id/--wc-pk, not both and not neither")
if args.id is not None and args.wc_pk is None:
parser.error("--id needs --wc-pk (the WorkChain to read the geometry from)")
return args
# ---------------------------------------------------------------------------
# The batch file (pure -- no aiida needed to validate a request)
# ---------------------------------------------------------------------------
[docs]
def load_batch(path: str | Path) -> list[dict[str, Any]]:
"""Read and validate the batch request.
Parameters
----------
path
The JSON request file the ``cif`` command wrote.
Returns
-------
list of dict
One normalised entry per material, each with ``id``, ``wc_pk``,
``uuid`` and ``calc_pk`` keys.
Raises
------
ValueError
Naming the defect: unreadable, not valid JSON, not a list, or an entry
without a non-empty string ``id`` and an integer ``wc_pk``. ``uuid``
and ``calc_pk`` are optional.
Notes
-----
Validation happens before AiiDA is touched at all, so a typo in the request
costs a message rather than a profile load.
"""
try:
raw = json.loads(Path(path).read_text(encoding="utf-8"))
except OSError as exc:
raise ValueError(f"cannot read batch file {path}: {exc}") from exc
except json.JSONDecodeError as exc:
raise ValueError(f"batch file {path} is not valid JSON: {exc}") from exc
if not isinstance(raw, list):
raise ValueError(f"batch file {path} must hold a JSON list, got {type(raw).__name__}")
entries: list[dict[str, Any]] = []
for index, entry in enumerate(raw):
where = f"batch entry {index}"
if not isinstance(entry, dict):
raise ValueError(f"{where} must be an object, got {type(entry).__name__}")
material_id = entry.get("id")
if not isinstance(material_id, str) or not material_id:
raise ValueError(f"{where} needs a non-empty string 'id', got {material_id!r}")
wc_pk = entry.get("wc_pk")
if not isinstance(wc_pk, int) or isinstance(wc_pk, bool):
raise ValueError(f"{where} ({material_id}) needs an integer 'wc_pk', got {wc_pk!r}")
entries.append(
{
"id": material_id,
"wc_pk": wc_pk,
"uuid": entry.get("uuid"),
"calc_pk": entry.get("calc_pk"),
}
)
return entries
# ---------------------------------------------------------------------------
# The AiiDA stack (the single lazy import touchpoint)
# ---------------------------------------------------------------------------
def _load_aiida() -> dict[str, Any]:
"""Import the user's AiiDA stack -- the single touchpoint, deliberately lazy.
Returns
-------
dict
The symbols used below, handed around explicitly so every function is
testable with a fake api.
Raises
------
ImportError
When aiida-core is missing from this interpreter. ``main`` catches it
and exits 3 with the exact install command.
Notes
-----
No ase here, unlike the calculate runner: reading a stored text file needs
nothing but aiida-core.
"""
from aiida import load_profile, orm
from aiida.common import LinkType
return {"load_profile": load_profile, "orm": orm, "LinkType": LinkType}
# ---------------------------------------------------------------------------
# Repository access (small compatibility seam over aiida-core versions)
# ---------------------------------------------------------------------------
def _repository_member(node: Any, method: str) -> Any:
"""Resolve one repository method across aiida-core API generations.
Parameters
----------
node
The node whose repository is wanted.
method
The method name.
Returns
-------
Any
The bound method. aiida-core 2.x keeps the repository behind
``node.base.repository``; older nodes exposed the same calls directly,
and trying both keeps this runner working across the versions a user
may have installed.
Raises
------
RuntimeError
When neither exposes it -- said loudly rather than guessed at.
"""
repository = getattr(getattr(node, "base", None), "repository", None)
member = getattr(repository, method, None)
if not callable(member):
member = getattr(node, method, None)
if not callable(member):
raise RuntimeError(f"node exposes no {method} (unexpected aiida-core API)")
return member
[docs]
def list_object_names(node: Any) -> list[str]:
"""List the files stored on a node.
Parameters
----------
node
A node with a repository, typically a retrieved ``FolderData``.
Returns
-------
list of str
The stored file names.
"""
return [str(name) for name in _repository_member(node, "list_object_names")()]
[docs]
def read_object(node: Any, name: str) -> str:
"""Read one stored file as text.
Parameters
----------
node
The node holding it.
name
The file to read.
Returns
-------
str
Its contents, decoded when stored as bytes.
Raises
------
RuntimeError
When the file is absent. The message lists what *is* there, which is
what turns "no geometry" into a diagnosable answer.
"""
names = list_object_names(node)
if name not in names:
found = ", ".join(sorted(names)) or "nothing"
raise RuntimeError(f"{name!r} is not among the retrieved files (found: {found})")
content = _repository_member(node, "get_object_content")(name)
return content.decode("utf-8") if isinstance(content, bytes) else str(content)
# ---------------------------------------------------------------------------
# Resolution: WorkChain -> final CalcJob -> geom.out.gen
# ---------------------------------------------------------------------------
[docs]
def final_calcjob(api: dict[str, Any], workchain: Any) -> Any:
"""Find the WorkChain's most recent CalcJob -- its last healing attempt.
Parameters
----------
api
The loaded AiiDA symbols.
workchain
The WorkChain to traverse.
Returns
-------
Any
The newest ``CALL_CALC`` CalcJob by creation time.
Raises
------
RuntimeError
When the WorkChain has no attempts to read a geometry from.
Notes
-----
The same traversal the calculate runner extracts attempts with, so "the
final attempt" means the same thing in both documents.
"""
calcs = workchain.base.links.get_outgoing(
node_class=api["orm"].CalcJobNode, link_type=api["LinkType"].CALL_CALC
).all_nodes()
if not calcs:
raise RuntimeError("the WorkChain has no CalcJob attempts to read a geometry from")
return max(calcs, key=lambda node: node.ctime)
[docs]
def fetch_geometry(
api: dict[str, Any], entry: dict[str, Any], geometry_name: str
) -> dict[str, Any]:
"""Fetch one material's relaxed geometry as a result row. Never raises.
Parameters
----------
api
The loaded AiiDA symbols.
entry
One normalised batch entry.
geometry_name
The retrieved file to read.
Returns
-------
dict
A row with ``id``, ``wc_pk``, ``calc_pk``, ``gen``, ``error`` and
``note``. Any failure -- an unloadable pk, a uuid mismatch, no
attempts, no retrieved folder, an absent geometry file -- comes back as
``error`` text with a ``null`` geometry, so **the batch always
completes** and the caller reports per material.
Notes
-----
A recorded uuid that disagrees with the loaded node is a **refusal**, not a
warning: exporting a different node's geometry under this material's id
would silently corrupt the export. A recorded ``calc_pk`` that disagrees is
only a ``note``, and the database wins.
"""
row: dict[str, Any] = {
"id": entry["id"],
"wc_pk": entry["wc_pk"],
"calc_pk": None,
"gen": None,
"error": None,
"note": None,
}
try:
workchain = api["orm"].load_node(entry["wc_pk"])
recorded_uuid = entry.get("uuid")
if recorded_uuid is not None and str(workchain.uuid) != str(recorded_uuid):
raise RuntimeError(
f"pk {entry['wc_pk']} holds uuid {workchain.uuid} but the record says "
f"{recorded_uuid} -- refusing to export a different node's geometry"
)
calc = final_calcjob(api, workchain)
row["calc_pk"] = calc.pk
expected = entry.get("calc_pk")
if expected is not None and expected != calc.pk:
row["note"] = (
f"recorded final attempt pk {expected} but the WorkChain's last CalcJob "
f"is pk {calc.pk}; used the database"
)
retrieved = getattr(calc.outputs, "retrieved", None)
if retrieved is None:
raise RuntimeError(f"CalcJob pk {calc.pk} has no retrieved folder")
row["gen"] = read_object(retrieved, geometry_name)
# Broad by intent: aiida raises its own exception family here, and every
# failure means the same thing to the caller -- no geometry for this one.
except Exception as exc:
row["error"] = f"{type(exc).__name__}: {exc}"
return row
# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
def _write_json_atomic(path: Path, doc: dict[str, Any]) -> None:
"""Write a JSON document atomically.
Parameters
----------
path
Where to write.
doc
The document.
Notes
-----
tmp-then-rename, so a crash never leaves a half-written document.
"""
tmp = path.with_name(path.name + ".tmp")
tmp.write_text(json.dumps(doc, indent=2), encoding="utf-8")
tmp.replace(path)
[docs]
def run(api: dict[str, Any], args: argparse.Namespace) -> int:
"""Fetch every requested geometry and emit the result document.
Parameters
----------
api
The loaded AiiDA symbols.
args
Parsed arguments: a batch file or a single material, plus where to
write.
Returns
-------
int
``0`` when the document was emitted, even if some rows carry errors --
a per-material failure is data, not a process failure. ``1`` only when
the output file itself could not be written.
Notes
-----
Progress and the final tally go to **stderr**, leaving stdout clean for the
result document when no ``--out`` is given.
"""
if args.batch is not None:
entries = load_batch(args.batch)
else:
entries = [{"id": args.id, "wc_pk": args.wc_pk, "uuid": args.uuid, "calc_pk": args.calc_pk}]
results: list[dict[str, Any]] = []
for index, entry in enumerate(entries, start=1):
results.append(fetch_geometry(api, entry, args.geometry_name))
if args.batch is not None and index % 200 == 0:
print(f"fetched {index}/{len(entries)} geometries", file=sys.stderr)
document = {"results": results}
if args.out is not None:
try:
_write_json_atomic(Path(args.out), document)
except OSError as exc:
print(f"cif_runner: cannot write {args.out}: {exc}", file=sys.stderr)
return 1
else:
print(json.dumps(document, indent=2))
failed = sum(1 for row in results if row["error"] is not None)
print(
f"fetched {len(results) - failed}/{len(results)} geometr"
f"{'y' if len(results) == 1 else 'ies'}" + (f", {failed} failed" if failed else ""),
file=sys.stderr,
)
return 0
[docs]
def main(argv: list[str] | None = None) -> int:
"""Fetch relaxed DFTB+ geometries out of the user's AiiDA repository.
Parameters
----------
argv
Command-line arguments. ``None`` reads ``sys.argv``.
Returns
-------
int
The process exit code: ``0`` when the result document was emitted;
``1`` when the output could not be written; ``2`` for a usage error or
an unusable batch file; ``3`` when aiida-core is missing, with the
exact install command and this interpreter's path printed to stderr.
"""
args = _parse_args(argv)
try:
api = _load_aiida()
except ImportError as exc:
print(
"cif_runner needs the user-owned AiiDA stack (aiida-core), not importable "
f"in this python environment ({sys.executable}).\n"
"Install it there with: pip install aiida-core\n"
f"(import error: {exc})",
file=sys.stderr,
)
return 3
if args.batch is not None and not Path(args.batch).is_file():
print(f"cannot read batch file {args.batch}: not a file", file=sys.stderr)
return 2
try:
api["load_profile"](args.profile)
# Broad by intent: a missing/misconfigured profile raises aiida's own
# exceptions; the user needs the message, not a traceback.
except Exception as exc:
print(
f"cif_runner: cannot load AiiDA profile {args.profile or '(default)'}: {exc}",
file=sys.stderr,
)
return 1
try:
return run(api, args)
except ValueError as exc:
print(f"cif_runner: {exc}", file=sys.stderr)
return 2
if __name__ == "__main__": # pragma: no cover — exercised via subprocess
sys.exit(main())