"""The structure stage: one composition -> M candidate crystal structures (3.2).
The thin ``Stage`` half of the pair; the science lives entirely in
``structure_runner.py`` (self-contained: stdlib + pyxtal, zero
``matmelt`` imports), which this class ships to the backend as an input
artifact of every job. This module itself imports **no science** -- it builds
job descriptions and judges results, nothing more, so it loads on any machine.
The input is the ``composition.json`` a compose child carries (formula,
elements, stoichiometry). ``validate_inputs`` goes beyond file presence and
checks the file actually decodes into that shape -- someone starting the
pipeline mid-way with a hand-made composition is refused clearly, before a
job is built, never with silent garbage or a confusing runner crash.
Fan-out contract: the runner tries every group of each requested
dimensionality (``--dims``; 3D=230 space, 2D=80 layer, 1D=75 rod groups) and
writes ``children.json`` (declared in ``produces``) plus one
``children/<Formula>_<DIM>-<SG>/`` directory per valid candidate (holding
``structure.cif`` + ``structure.json``), all in the job's cwd; the
orchestrator adopts the children and retires the composition as ``EXPANDED``.
Zero children is a valid result -- some compositions satisfy no symmetry
constraints at all.
"""
from __future__ import annotations
import json
from pathlib import Path
from typing import TYPE_CHECKING, Any
from matmelt.core.jobspec import JobSpec
from matmelt.core.stages.base import (
CHILDREN_DIR,
CHILDREN_MANIFEST,
Requirement,
Stage,
StageError,
StageInputError,
StageInputSpec,
StageOutputSpec,
StageVerdict,
)
from matmelt.core.stages.registry import register
if TYPE_CHECKING:
from matmelt.core.jobspec import JobResult
from matmelt.core.stages.base import WorkItem
#: The self-contained script shipped with every structure job. Lives next to
#: this module inside the installed package; staged into the job's cwd by the
#: backend, so the command references it by bare name.
RUNNER_NAME = "structure_runner.py"
#: The input file a compose child carries -- declared as an input artifact so
#: the Phase-4 staging transport knows to push it (it already sits in the
#: item dir locally; backends stage in-place inputs as a no-op).
COMPOSITION_NAME = "composition.json"
#: Keys ``composition.json`` must carry for the runner to consume it.
_REQUIRED_KEYS = ("formula", "elements", "stoichiometry")
#: How many characters of a failed runner's stderr to surface in the verdict.
_REASON_TAIL = 300
#: The dimensionalities the runner can sweep (mirrors ``structure_runner.DIM_GROUPS``
#: keys; kept as a local constant so this stage stays science-free -- the runner's
#: argparse ``choices`` is the executing enforcement, this is an early, clear guard).
_VALID_DIMENSIONS = ("3D", "2D", "1D")
[docs]
class StructureStage(Stage):
"""PyXtal-generate symmetry-constrained candidates for one composition.
Parameters
----------
python
Interpreter to invoke on the backend this stage is wired to (the
user's own environment holds the science stack -- never bundled).
Plain ``python3`` by default; a config option, never a hardcoded
compute detail.
dimensions
Which dimensionalities the runner generates -- any of ``3D`` (230 space
groups), ``2D`` (80 layer groups), ``1D`` (75 rod groups). A single tag
string or a list; defaults to ``("3D",)`` (the pre-existing behaviour).
A config option, forwarded to the runner as ``--dims``.
Raises
------
StageError
For any invalid option, so a config mistake surfaces when the run is
wired rather than as N failed jobs later.
"""
fan_out = True
runtime_requirements = (
Requirement(
module="pyxtal",
pip_name="pyxtal",
purpose="symmetry-constrained random crystal generation (structure runner)",
),
)
def __init__(self, python: str = "python3", dimensions: Any = ("3D",)) -> None:
if not isinstance(python, str) or not python.strip():
raise StageError(f"structure: python must be a non-empty string, got {python!r}")
self.runner_python = python
self.dimensions = self._validate_dimensions(dimensions)
@staticmethod
def _validate_dimensions(dimensions: Any) -> tuple[str, ...]:
"""Normalise and validate the ``dimensions`` option.
Parameters
----------
dimensions
A single tag string such as ``"3D"``, or a list or tuple of tags.
Returns
-------
tuple of str
The tags, in the given order.
Raises
------
StageError
For an empty value, a wrong type, or an unknown tag, with the valid
set named -- refusing the misconfigured stage before any job.
"""
if isinstance(dimensions, str):
dimensions = (dimensions,)
if not isinstance(dimensions, (list, tuple)) or not dimensions:
raise StageError(
f"structure: dimensions must be a non-empty list of "
f"{list(_VALID_DIMENSIONS)}, got {dimensions!r}"
)
dims = tuple(dimensions)
bad = [d for d in dims if d not in _VALID_DIMENSIONS]
if bad:
raise StageError(
f"structure: unknown dimensionalit(y/ies) {bad}; valid: {list(_VALID_DIMENSIONS)}"
)
return dims
@property
def name(self) -> str:
return "structure"
@property
def requires(self) -> StageInputSpec:
return StageInputSpec(
files=(COMPOSITION_NAME,),
description="one charge-balanced composition from the compose stage",
)
@property
def produces(self) -> StageOutputSpec:
return StageOutputSpec(
files=(CHILDREN_MANIFEST,),
description=(
"fan-out manifest + children/<Formula>_<DIM>-<SG>/structure.cif per candidate"
),
)
[docs]
@staticmethod
def composition_of(item: WorkItem) -> dict[str, Any]:
"""Read and check the item's ``composition.json``.
The same rules are enforced by the runner, which must stay debuggable
standalone; checking here refuses a bad item before any job is built or
submitted.
Parameters
----------
item
The item whose directory holds the file.
Returns
-------
dict
The decoded composition.
Raises
------
StageInputError
Naming exactly what is wrong: unreadable or malformed JSON, a
missing required key, a bad formula (it names child directories, so
it must be alphanumeric), or elements and stoichiometry that do not
line up.
"""
def refuse(defect: str) -> StageInputError:
return StageInputError(
f"stage 'structure' cannot run on item {item.id!r}: {COMPOSITION_NAME} {defect}"
)
path = item.dir / COMPOSITION_NAME
try:
data = json.loads(path.read_text(encoding="utf-8"))
except OSError as exc:
raise refuse(f"is unreadable: {exc}") from exc
except ValueError as exc:
raise refuse(f"is not valid JSON: {exc}") from exc
if not isinstance(data, dict):
raise refuse(f"must hold one composition object, got {type(data).__name__}")
missing = [key for key in _REQUIRED_KEYS if key not in data]
if missing:
raise refuse(f"is missing required key(s): {', '.join(missing)}")
formula = data["formula"]
if not isinstance(formula, str) or not formula or not formula.isalnum():
raise refuse(
f"formula must be a non-empty alphanumeric string "
f"(it names child directories), got {formula!r}"
)
elements = data["elements"]
if (
not isinstance(elements, list)
or not elements
or not all(isinstance(el, str) and el for el in elements)
):
raise refuse(f"elements must be a non-empty list of element symbols, got {elements!r}")
stoichiometry = data["stoichiometry"]
if (
not isinstance(stoichiometry, list)
or len(stoichiometry) != len(elements)
or not all(
isinstance(n, int) and not isinstance(n, bool) and n >= 1 for n in stoichiometry
)
):
raise refuse(
f"stoichiometry must be a list of integers >= 1, one per element, "
f"got {stoichiometry!r} for elements {elements!r}"
)
return data
[docs]
def build_jobspec(self, item: WorkItem) -> JobSpec:
runner = Path(__file__).with_name(RUNNER_NAME)
return JobSpec(
command=[
self.runner_python,
RUNNER_NAME,
"--composition",
COMPOSITION_NAME,
"--dims",
*self.dimensions,
],
input_artifacts={
RUNNER_NAME: str(runner),
COMPOSITION_NAME: str(item.dir / COMPOSITION_NAME),
},
# The manifest AND the children/ subtree it names -- both must be
# staged back so a remote (ssh_exec) fan-out job's children survive.
# produces/validate_outputs stay on the manifest alone (see produces).
output_artifacts=[CHILDREN_MANIFEST, CHILDREN_DIR],
cwd=str(item.dir),
scratch_inputs=frozenset({RUNNER_NAME}),
)
[docs]
def interpret_result(self, item: WorkItem, result: JobResult) -> StageVerdict:
"""Judge a finished structure job.
Parameters
----------
item
The composition whose job finished.
result
What the backend reported.
Returns
-------
StageVerdict
``ok`` on exit 0, which includes a fan-out into zero children.
Anything else is a fault carrying the runner's stderr tail, and
**never terminal**: a structure failure is an environment problem
-- missing pyxtal, a broken interpreter -- that a retry can
genuinely resolve. Per-space-group science rejections are absorbed
inside the runner's exit 0 and counted in its report.
"""
if result.exit_code == 0:
return StageVerdict(ok=True)
tail = (result.stderr or result.stdout).strip()[-_REASON_TAIL:]
code = "unknown" if result.exit_code is None else str(result.exit_code)
return StageVerdict(ok=False, reason=f"structure runner exited {code}: {tail}")
register("structure", StructureStage)