#!/usr/bin/env python3
"""MatMelt calculate runner: DFTB+ via the user's own AiiDA, ONE material.
Self-contained by design -- imports **stdlib + aiida/aiida_mydftbp/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 calculate_runner.py --id Cl1Na1_225 --code-label dftbplus@cluster \\
--skf-prefix /opt/skf/
python calculate_runner.py --help
The user's environment provides everything AiiDA: aiida-core, the
aiida-mydftbp plugin (whose ``mydftbp.single`` WorkChain does the
self-healing DFTB+ retries itself), a configured profile + code, and a
running daemon. This runner is the launcher/extractor around that, per
material:
1. prepare -- ``relaxed.cif`` -> ``.gen`` geometry + ``dftb_in.hsd`` from an
embedded template (SKF prefix, k-mesh from the options);
2. ensure the labelled Slater-Koster ``FolderData`` exists (find-or-create
from ``--skf-local-path``; missing and no local path = loud failure,
never a silent stall);
3. submit one ``mydftbp.single`` WorkChain and write
``calc_submission.json`` BEFORE polling, so a killed run re-attaches to
the same WorkChain instead of resubmitting;
4. poll to terminal, then extract every healing attempt (parameters with
SKF absolute paths stripped, per-attempt diffs) into an atomic
``calc_result.json``.
Exit codes: ``0`` = calculation completed and ``calc_result.json`` written
(converged or not -- science lives in the JSON, the stage judges); ``2`` =
bad arguments or unreadable ``--cif``; ``3`` = aiida/aiida_mydftbp/ase not
importable in this python environment (stderr says exactly what to
install); ``1`` = unexpected error (traceback on stderr -- the stage treats
it as a retryable fault).
Every function below takes plain arguments (the AiiDA symbols ride in an
``api`` dict built by ``_load_aiida``), so the whole orchestration is
testable with fakes and no aiida installed -- the dev-venv hard rule.
"""
from __future__ import annotations
import argparse
import copy
import json
import re
import sys
import time
from pathlib import Path
from typing import Any
RESULT_NAME = "calc_result.json"
SUBMISSION_NAME = "calc_submission.json"
GEN_NAME = "geometry.gen"
HSD_NAME = "dftb_in.hsd"
#: What DFTB+ writes the relaxed geometry to, given the template's
#: ``OutputPrefix = "geom.out"``; retrieved by the plugin, cached into
#: ``calc_result.json`` so exports need no AiiDA round-trip.
GEOMETRY_NAME = "geom.out.gen"
def _parse_args(argv: list[str] | None) -> argparse.Namespace:
parser = argparse.ArgumentParser(
prog="calculate_runner",
description="Run one material through DFTB+ via the user's own AiiDA "
"(self-healing mydftbp.single WorkChain) and extract the full "
"attempt history.",
)
parser.add_argument(
"--id",
required=True,
metavar="MATERIAL_ID",
help="Material identifier, recorded in calc_result.json.",
)
parser.add_argument(
"--cif",
default="relaxed.cif",
metavar="PATH",
help="Relaxed structure to calculate (default: relaxed.cif).",
)
parser.add_argument(
"--code-label",
required=True,
metavar="LABEL",
help="Label of the configured AiiDA code for DFTB+ (see 'verdi code list').",
)
parser.add_argument(
"--profile",
default=None,
metavar="NAME",
help="AiiDA profile to load (default: the user's default profile).",
)
parser.add_argument(
"--skf-label",
default="ptbp_complete_set",
metavar="LABEL",
help="Label of the stored Slater-Koster FolderData node (find-or-create).",
)
parser.add_argument(
"--skf-prefix",
required=True,
metavar="DIR",
help="Absolute .skf directory on the machine where DFTB+ executes "
"(baked into the HSD input).",
)
parser.add_argument(
"--skf-local-path",
default=None,
metavar="DIR",
help="Local .skf directory used to CREATE the labelled FolderData when "
"it does not exist yet; omitted = find-only, missing node is a loud "
"failure naming this option.",
)
parser.add_argument(
"--kpoints",
type=int,
nargs=3,
default=[4, 4, 4],
metavar=("K1", "K2", "K3"),
help="SupercellFolding k-mesh diagonal (default: 4 4 4).",
)
parser.add_argument(
"--max-wallclock-seconds",
type=int,
default=3600,
metavar="S",
help="Walltime request per WorkChain attempt, in seconds.",
)
parser.add_argument(
"--queue-name",
default=None,
metavar="NAME",
help="Scheduler queue/partition (default: the code's computer default).",
)
parser.add_argument(
"--poll-interval",
type=float,
default=30.0,
metavar="S",
help="Seconds between WorkChain status polls.",
)
parser.add_argument(
"--out-dir",
default=".",
metavar="PATH",
help="Directory to write calc_submission.json / calc_result.json into (default: cwd).",
)
args = parser.parse_args(argv)
if any(k < 1 for k in args.kpoints):
parser.error(f"--kpoints values must be >= 1, got {args.kpoints}")
if args.max_wallclock_seconds < 1:
parser.error(f"--max-wallclock-seconds must be >= 1, got {args.max_wallclock_seconds}")
if args.poll_interval <= 0:
parser.error(f"--poll-interval must be > 0, got {args.poll_interval}")
return args
# ---------------------------------------------------------------------------
# Prepare: .gen geometry + dftb_in.hsd from the embedded template (6.6)
# ---------------------------------------------------------------------------
#: Highest occupied angular momentum per element, by periodic-table block:
#: s-block "s", d-block "d", lanthanides/actinides "f"; everything else --
#: the whole p-block and any unknown symbol -- defaults to "p".
_S_BLOCK = "H He Li Be Na Mg K Ca Rb Sr Cs Ba Fr Ra"
_D_BLOCK = "Sc Ti V Cr Mn Fe Co Ni Cu Zn Y Zr Nb Mo Tc Ru Rh Pd Ag Cd Lu Hf Ta W Re Os Ir Pt Au Hg"
_F_BLOCK = "La Ce Pr Nd Pm Sm Eu Gd Tb Dy Ho Er Tm Yb Ac Th Pa U Np Pu Am Cm Bk Cf Es Fm Md No"
_ANGULAR_MOMENTUM: dict[str, str] = (
dict.fromkeys(_S_BLOCK.split(), "s")
| dict.fromkeys(_D_BLOCK.split(), "d")
| dict.fromkeys(_F_BLOCK.split(), "f")
)
[docs]
def max_angular_momentum(symbol: str) -> str:
"""Give the ``MaxAngularMomentum`` shell for one element.
Parameters
----------
symbol
The element symbol.
Returns
-------
str
``"s"`` for the s-block, ``"d"`` for the d-block, ``"f"`` for the
lanthanides and actinides, and ``"p"`` for the p-block -- which is also
the fallback for an unknown symbol, since it is the commonest case and
DFTB+ would otherwise refuse to start.
Examples
--------
>>> max_angular_momentum("Na"), max_angular_momentum("Fe")
('s', 'd')
"""
return _ANGULAR_MOMENTUM.get(symbol, "p")
[docs]
def species_from_gen(gen_text: str) -> list[str]:
"""Read the element symbols from DFTB+ ``.gen`` text.
Parameters
----------
gen_text
The file's contents. The species row is line 2 of the header.
Returns
-------
list of str
The symbols, in the order the ``.gen`` file indexes them.
Raises
------
ValueError
When the text is too short to hold a header, or the species row does
not look like element symbols.
"""
lines = [line for line in gen_text.splitlines() if line.strip()]
if len(lines) < 2:
raise ValueError("gen text has no species line (expected the 2-line .gen header)")
species = lines[1].split()
if not species or not all(symbol.isalpha() for symbol in species):
raise ValueError(f"gen species line does not look like element symbols: {lines[1]!r}")
return species
[docs]
def build_hsd(gen_text: str, skf_prefix: str, kpoints: tuple[int, int, int]) -> str:
"""Assemble ``dftb_in.hsd`` from the embedded template.
Parameters
----------
gen_text
The geometry, embedded inline as ``GenFormat``.
skf_prefix
Directory of the ``.skf`` files **on the machine where DFTB+ runs**.
Normalized to a trailing slash, because DFTB+ concatenates it with
``El1-El2.skf`` verbatim.
kpoints
The ``SupercellFolding`` k-mesh diagonal.
Returns
-------
str
The complete HSD input.
Notes
-----
``MaxAngularMomentum`` is derived per species. The document's shape is not
arbitrary: it must round-trip through the aiida-mydftbp plugin's HSD
parser, so the keys that plugin's self-healing ladders tune -- ``MaxSteps``,
``Optimizer``, ``Convergence``, ``SCC``, ``KPointsAndWeights`` -- are all
written in the forms it parses explicitly.
"""
species = species_from_gen(gen_text)
momenta = "\n".join(f' {symbol} = "{max_angular_momentum(symbol)}"' for symbol in species)
prefix = skf_prefix if skf_prefix.endswith("/") else skf_prefix + "/"
k1, k2, k3 = kpoints
return f"""Geometry = GenFormat {{
{gen_text.rstrip()}
}}
Driver = GeometryOptimization {{
Optimizer = Rational {{}}
MovedAtoms = 1:-1
MaxSteps = 100
OutputPrefix = "geom.out"
Convergence {{ GradElem = 1E-4 }}
}}
Hamiltonian = DFTB {{
SCC = Yes
SlaterKosterFiles = Type2FileNames {{
Prefix = "{prefix}"
Separator = "-"
Suffix = ".skf"
}}
MaxAngularMomentum = {{
{momenta}
}}
KPointsAndWeights = SupercellFolding {{
{k1} 0 0
0 {k2} 0
0 0 {k3}
0.5 0.5 0.5
}}
}}
Options {{}}
Analysis {{
PrintForces = Yes
}}
ParserOptions {{
ParserVersion = 12
}}
"""
[docs]
def write_gen(cif_path: str | Path, gen_path: str | Path) -> None:
"""Convert a CIF to a DFTB+ ``.gen`` file, in fractional coordinates.
Parameters
----------
cif_path
The structure to read.
gen_path
Where to write the ``.gen`` file.
Notes
-----
The runner's only ase touchpoint besides ``_load_aiida``, imported lazily
so everything above stays runnable anywhere. A missing ase has already
surfaced as exit 3 before this is ever called.
"""
from ase.io import read as ase_read
from ase.io import write as ase_write
atoms = ase_read(str(cif_path))
ase_write(str(gen_path), atoms, format="dftb", fractional=True)
# ---------------------------------------------------------------------------
# Extraction helpers: pure functions over plain dicts (no ORM types cross
# these seams -- everything is JSON-safe by construction)
# ---------------------------------------------------------------------------
#: What replaces the SKF directory in extracted parameters. The only place
#: an absolute machine path leaks into the plugin's parameter dicts is the
#: SlaterKosterFiles passthrough string under Hamiltonian.DFTB.
SKF_PATH_REDACTED = '"<stripped>"'
_SKF_PREFIX_RE = re.compile(r'(Prefix\s*=\s*)("[^"]*"|\S+)')
[docs]
def strip_skf_paths(params: dict[str, Any]) -> dict[str, Any]:
"""Copy ``params`` with the SKF directory path redacted.
Parameters
----------
params
A parameter dict as the plugin's HSD parser produced it.
Returns
-------
dict
A deep copy with the ``Prefix`` value replaced. The plugin keeps
``SlaterKosterFiles`` as a ``_raw_N`` passthrough string inside
``Hamiltonian.DFTB``, and that is the one place an absolute machine
path leaks into a record -- extracted documents must never carry it.
Anything unexpected is left untouched: redaction is best-effort over a
copy, never a failure.
"""
out = copy.deepcopy(params)
hamiltonian = out.get("Hamiltonian")
if not isinstance(hamiltonian, dict):
return out
dftb = hamiltonian.get("DFTB")
if not isinstance(dftb, dict):
return out
for key, value in dftb.items():
if key.startswith("_raw") and isinstance(value, str) and "SlaterKosterFiles" in value:
dftb[key] = _SKF_PREFIX_RE.sub(rf"\g<1>{SKF_PATH_REDACTED}", value)
return out
[docs]
def param_diff(old: dict[str, Any], new: dict[str, Any], _prefix: str = "") -> dict[str, Any]:
"""Diff two parameter dicts -- what the plugin's healing actually changed.
Parameters
----------
old, new
The two parameter dicts to compare.
_prefix
Internal, used to build dotted paths while recursing.
Returns
-------
dict
``{dotted.path: change}``, where change is ``{"from": ..., "to": ...}``
for a changed value and ``{"added": ...}`` or ``{"removed": ...}`` for
a key present on one side only. Two dict values recurse; a type change
compares whole.
Notes
-----
Keys starting with ``_`` are skipped at every level: ``_raw_*`` passthrough
strings and ``_healing`` metadata are not parameter changes.
Examples
--------
>>> param_diff({"Driver": {"MaxSteps": 100}}, {"Driver": {"MaxSteps": 500}})
{'Driver.MaxSteps': {'from': 100, 'to': 500}}
"""
changes: dict[str, Any] = {}
for key in sorted(set(old) | set(new), key=str):
if isinstance(key, str) and key.startswith("_"):
continue
path = f"{_prefix}{key}"
if key not in old:
changes[path] = {"added": new[key]}
elif key not in new:
changes[path] = {"removed": old[key]}
elif isinstance(old[key], dict) and isinstance(new[key], dict):
changes.update(param_diff(old[key], new[key], f"{path}."))
elif old[key] != new[key]:
changes[path] = {"from": old[key], "to": new[key]}
return changes
[docs]
def shape_attempt(
pk: int,
attempt: int,
exit_status: int | None,
params: dict[str, Any] | None,
previous_params: dict[str, Any] | None,
result: dict[str, Any] | None,
) -> dict[str, Any]:
"""Shape one WorkChain attempt as a plain JSON-safe dict.
Parameters
----------
pk
The CalcJob's primary key in the user's AiiDA database.
attempt
Its 1-based position in the healing chain.
exit_status
How that attempt ended, or ``None`` when unavailable.
params
The parameters it ran with, or ``None`` when unavailable.
previous_params
The previous attempt's parameters, for the diff.
result
Its parsed output, or ``None``.
Returns
-------
dict
The attempt record. ``healing`` carries the plugin's ``_healing``
metadata, ``None`` when the attempt predates it or parameters are
unavailable; ``params`` are stored SKF-stripped; and
``param_changes_from_previous`` is ``None`` for the first attempt and a
``param_diff`` otherwise -- which is what makes a healing chain
readable as "it raised MaxSteps, then loosened convergence".
"""
return {
"pk": pk,
"attempt": attempt,
"exit_status": exit_status,
"healing": params.get("_healing") if isinstance(params, dict) else None,
"params": strip_skf_paths(params) if isinstance(params, dict) else None,
"param_changes_from_previous": (
param_diff(previous_params, params)
if isinstance(previous_params, dict) and isinstance(params, dict)
else None
),
"result": result,
}
[docs]
def shape_result(
material_id: str,
single_wc_pk: int,
uuid: str,
converged: bool,
exit_status: int | None,
attempts: list[dict[str, Any]],
final_geometry: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Assemble the top-level ``calc_result.json`` document.
Parameters
----------
material_id
The material this calculation belongs to.
single_wc_pk
The ``mydftbp.single`` WorkChain's primary key.
uuid
Its uuid, which is what a re-attach verifies against.
converged
Whether the WorkChain reached a converged result.
exit_status
Its exit status, which the stage maps to a verdict.
attempts
The shaped attempt records, in chronological order.
final_geometry
The geometry DFTB+ actually relaxed **to**.
Returns
-------
dict
The result document.
Notes
-----
``final_result`` is the last attempt's parsed output only when the
WorkChain converged, and ``None`` otherwise: failed attempt history is data
too, but only a converged result may be consumed as one.
``final_geometry`` is cached here so consumers never need to reach into the
AiiDA repository -- the attempts' own ``params`` hold only what each
attempt *started* from. It is captured whether or not the WorkChain
converged (the top-level ``converged`` flag is what says it may be consumed
as a result), and is ``None`` when it could not be read.
"""
return {
"material_id": material_id,
"single_wc_pk": single_wc_pk,
"uuid": uuid,
"converged": bool(converged),
"exit_status": exit_status,
"n_attempts": len(attempts),
"attempts": list(attempts),
"final_result": attempts[-1]["result"] if converged and attempts else None,
"final_geometry": final_geometry,
}
# ---------------------------------------------------------------------------
# The AiiDA stack (lazy imports -- everything below needs the user's aiida env)
# ---------------------------------------------------------------------------
def _load_aiida() -> dict[str, Any]:
"""Import the user's AiiDA stack -- the single touchpoint, deliberately lazy.
Returns
-------
dict
Exactly the symbols the orchestration below uses, handed around as a
plain dict so every function is testable with a fake api and no aiida
installed anywhere.
Raises
------
ImportError
When aiida-core, aiida-mydftbp or ase is missing from this
interpreter. ``main`` catches it and exits 3 with the exact install
command.
Notes
-----
``ase`` is imported here purely so its absence is caught by the same exit-3
path; ``write_gen`` imports it again where it is used.
"""
import ase # noqa: F401 (write_gen imports it lazily; probed here so exit 3 covers it)
from aiida import load_profile, orm
from aiida.common import LinkType
from aiida.engine import submit
from aiida.plugins import WorkflowFactory
from aiida_mydftbp.hsd_parser import load_initial_parameters
return {
"load_profile": load_profile,
"orm": orm,
"LinkType": LinkType,
"submit": submit,
"WorkflowFactory": WorkflowFactory,
"load_initial_parameters": load_initial_parameters,
}
# ---------------------------------------------------------------------------
# Orchestration: ensure SKF -> submit (or re-attach) -> poll -> extract (6.7)
# ---------------------------------------------------------------------------
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 for a
later run to misread.
"""
tmp = path.with_name(path.name + ".tmp")
tmp.write_text(json.dumps(doc, indent=2), encoding="utf-8")
tmp.replace(path)
[docs]
def ensure_skf_folder(orm: Any, label: str, local_path: str | None) -> Any:
"""Find the labelled Slater-Koster ``FolderData``, or create it -- loudly.
Parameters
----------
orm
The AiiDA ``orm`` module.
label
The node label to find or claim.
local_path
A local directory of ``.skf`` files used to *create* the node when none
exists. ``None`` means find-only.
Returns
-------
Any
The stored ``FolderData`` node. An existing node is reused, so a second
run never duplicates it.
Raises
------
RuntimeError
When no node carries the label and either no ``local_path`` was given
or it holds no ``.skf`` files. The message names exactly what to
configure.
Notes
-----
A missing SKF node is **never a quiet skip** -- that silent stall is the
documented defect this function exists to fix. The failure surfaces as a
clean exit-1 message rather than a traceback or a run that appears to be
working.
"""
query = orm.QueryBuilder()
query.append(orm.FolderData, filters={"label": label})
existing = query.first(flat=True)
if existing is not None:
return existing
if local_path is None:
raise RuntimeError(
f"no FolderData labelled {label!r} exists in this AiiDA profile, and no "
"--skf-local-path was given to create it. Configure the calculate "
"stage's skf_local_path option (a local directory of .skf files) once, "
"or store a labelled FolderData yourself"
)
tree = Path(local_path)
if not tree.is_dir() or not any(tree.glob("*.skf")):
raise RuntimeError(
f"--skf-local-path {local_path!r} is not a directory containing .skf "
f"files; cannot create the {label!r} FolderData from it"
)
folder = orm.FolderData(tree=str(tree))
folder.label = label
folder.store()
print(f"created SKF FolderData {label!r} from {local_path}")
return folder
[docs]
def submit_workchain(api: dict[str, Any], args: argparse.Namespace, hsd_path: Path) -> Any:
"""Submit one self-healing ``mydftbp.single`` WorkChain for this material.
Parameters
----------
api
The loaded AiiDA symbols.
args
Parsed arguments: the code label, SKF settings, walltime and queue.
hsd_path
The generated ``dftb_in.hsd``, parsed by the plugin into the
WorkChain's initial parameters.
Returns
-------
Any
The submitted WorkChain node. Submission is non-blocking: the node is
handed to the daemon, which is what lets the runner be killed without
losing the calculation.
Notes
-----
``skf_files`` is passed even though the WorkChain's default
``use_remote_skf_path=True`` means the bytes are never uploaded -- it is a
required input, and it records provenance.
"""
orm = api["orm"]
code = orm.load_code(args.code_label)
single = api["WorkflowFactory"]("mydftbp.single")
skf = ensure_skf_folder(orm, args.skf_label, args.skf_local_path)
params = api["load_initial_parameters"](str(hsd_path))
options: dict[str, Any] = {
"resources": {"num_machines": 1, "num_mpiprocs_per_machine": 1},
"max_wallclock_seconds": args.max_wallclock_seconds,
"withmpi": False,
}
if args.queue_name:
options["queue_name"] = args.queue_name
return api["submit"](
single,
code=code,
formula=orm.Str(args.id),
material_dir=orm.Str(str(Path(args.out_dir).resolve())),
initial_parameters=orm.Dict(dict=params),
# Required input even though the WorkChain's use_remote_skf_path=True
# default means the bytes are never uploaded -- provenance only.
skf_files=skf,
options=orm.Dict(dict=options),
)
[docs]
def load_submission(out_dir: Path) -> dict[str, Any] | None:
"""Read the recorded submission, if this material has one.
Parameters
----------
out_dir
The job's working directory.
Returns
-------
dict or None
``{"pk": ..., "uuid": ...}``, or ``None`` when no record exists or it
is unusable -- which the caller treats as "submit a fresh one". A
record that exists but is corrupt warns on stderr first: silently
double-submitting a DFTB+ calculation is expensive.
"""
path = out_dir / SUBMISSION_NAME
try:
doc = json.loads(path.read_text(encoding="utf-8"))
except OSError:
return None
except ValueError:
print(
f"warning: {SUBMISSION_NAME} is not valid JSON; submitting a fresh WorkChain",
file=sys.stderr,
)
return None
if isinstance(doc, dict) and isinstance(doc.get("pk"), int):
return doc
print(
f"warning: {SUBMISSION_NAME} has no usable pk; submitting a fresh WorkChain",
file=sys.stderr,
)
return None
[docs]
def attach_or_submit(api: dict[str, Any], args: argparse.Namespace, hsd_path: Path) -> Any:
"""Re-attach to the recorded WorkChain, or submit and record a fresh one.
Parameters
----------
api
The loaded AiiDA symbols.
args
Parsed arguments.
hsd_path
The generated HSD input, used only on a fresh submission.
Returns
-------
Any
The WorkChain node this run will follow.
Notes
-----
``calc_submission.json`` is written **immediately after submission, before
any polling**. That ordering is the whole point: a killed run re-attaches
to the same WorkChain on rerun and never double-submits, which was proven
live by SIGKILLing the CLI mid-poll and finding exactly one WorkChain in
the database afterwards.
A record that no longer loads, or whose uuid disagrees with the loaded
node, is treated as stale -- warned about, then replaced. A quiet
double-submit is never acceptable here.
"""
out_dir = Path(args.out_dir)
recorded = load_submission(out_dir)
if recorded is not None:
try:
node = api["orm"].load_node(recorded["pk"])
# Broad by intent: aiida raises NotExistent and friends here, and any
# unloadable record means the same thing -- stale, so submit fresh.
except Exception as exc:
print(
f"warning: recorded WorkChain PK {recorded['pk']} is not loadable "
f"({exc}); submitting a fresh one",
file=sys.stderr,
)
else:
if recorded.get("uuid") in (None, str(node.uuid)):
print(f"re-attached to WorkChain PK {node.pk}")
return node
print(
f"warning: recorded uuid for PK {recorded['pk']} does not match the "
f"loaded node; submitting a fresh one",
file=sys.stderr,
)
node = submit_workchain(api, args, hsd_path)
_write_json_atomic(out_dir / SUBMISSION_NAME, {"pk": node.pk, "uuid": str(node.uuid)})
print(f"submitted WorkChain PK {node.pk}")
return node
[docs]
def poll_until_terminated(api: dict[str, Any], pk: int, interval_s: float) -> Any:
"""Reload the WorkChain node until it reports terminated.
Parameters
----------
api
The loaded AiiDA symbols.
pk
The WorkChain's primary key.
interval_s
Seconds to sleep between reloads.
Returns
-------
Any
The terminated node.
Notes
-----
The daemon does the work; this only watches. Killing the runner mid-poll
therefore costs nothing -- the WorkChain keeps running, and the next run
re-attaches to it.
"""
while True:
node = api["orm"].load_node(pk)
if node.is_terminated:
return node
time.sleep(interval_s)
def _linked_dict(namespace: Any, name: str) -> dict[str, Any] | None:
"""Read one linked ``Dict`` node's contents.
Parameters
----------
namespace
A node's ``inputs`` or ``outputs`` namespace.
name
The link label to read.
Returns
-------
dict or None
The contents, or ``None`` when the link is absent or does not hold a
dict. Absence is data here, not an error: an attempt that never
produced output parameters is a normal part of a healing chain.
"""
try:
node = getattr(namespace, name)
result = node.get_dict()
except AttributeError:
return None
return dict(result) if isinstance(result, dict) else None
def _retrieved_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, e.g. ``list_object_names``.
Returns
-------
Any
The bound method, taken from ``node.base.repository`` when present and
from the node itself otherwise.
Raises
------
RuntimeError
When neither exposes it -- an unexpected aiida-core API.
Notes
-----
The same three lines live in ``cif_runner.py`` on purpose: self-contained
runners ship alone and may never import each other or our package.
"""
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 run(api: dict[str, Any], args: argparse.Namespace) -> int:
"""Drive one material end to end, with the profile already loaded.
Prepare the geometry and HSD input, attach to or submit the WorkChain, poll
it to termination, extract every attempt, and write the result atomically.
Parameters
----------
api
The loaded AiiDA symbols.
args
Parsed arguments.
Returns
-------
int
``0``. The science lives in the written document, not in the exit code:
a WorkChain that failed still produced a *successful* extraction, and
the stage reads the verdict out of ``calc_result.json``.
"""
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
gen_path = out_dir / GEN_NAME
write_gen(args.cif, gen_path)
hsd_path = out_dir / HSD_NAME
hsd_path.write_text(
build_hsd(gen_path.read_text(encoding="utf-8"), args.skf_prefix, tuple(args.kpoints)),
encoding="utf-8",
)
node = attach_or_submit(api, args, hsd_path)
node = poll_until_terminated(api, node.pk, args.poll_interval)
attempts = extract_attempts(api, node)
converged = bool(node.is_finished_ok)
document = shape_result(
material_id=args.id,
single_wc_pk=node.pk,
uuid=str(node.uuid),
converged=converged,
exit_status=node.exit_status,
attempts=attempts,
final_geometry=extract_final_geometry(api, node),
)
_write_json_atomic(out_dir / RESULT_NAME, document)
outcome = "converged" if converged else f"exit_status {node.exit_status}"
print(f"{args.id} -> WorkChain PK {node.pk}: {outcome} ({len(attempts)} attempt(s))")
return 0
[docs]
def main(argv: list[str] | None = None) -> int:
"""Calculate one screened-in material with DFTB+ through the user's AiiDA.
Parameters
----------
argv
Command-line arguments. ``None`` reads ``sys.argv``.
Returns
-------
int
The process exit code: ``0`` when a result document was written,
whatever the WorkChain's own outcome; ``1`` for an expected, actionable
failure such as the SKF ensure refusing, printed as a message rather
than a traceback; ``2`` for a usage error or an unreadable structure
file; ``3`` when the AiiDA stack is missing, with the exact install
command and this interpreter's path printed to stderr. Anything
unexpected propagates with its traceback.
"""
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:
api = _load_aiida()
except ImportError as exc:
print(
"calculate_runner needs the user-owned AiiDA stack (aiida-core + the "
"aiida-mydftbp plugin) plus ase, not importable in this python "
f"environment ({sys.executable}).\n"
"Install them there with: pip install aiida-core aiida-mydftbp ase\n"
f"(import error: {exc})",
file=sys.stderr,
)
return 3
api["load_profile"](args.profile)
try:
return run(api, args)
except RuntimeError as exc:
# Expected, actionable failures (e.g. the SKF ensure refusing) --
# message without a traceback; anything else propagates as exit 1.
print(f"calculate_runner: {exc}", file=sys.stderr)
return 1
if __name__ == "__main__": # pragma: no cover — exercised via subprocess
sys.exit(main())