"""The compose stage: one element-combination seed -> N candidate compositions (3.1).
The thin ``Stage`` half of the pair; the science lives entirely in
``compose_runner.py`` (self-contained: stdlib + smact, 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 seed needs no input file: its *id* is the data. Seed ids come from
``core/seeds.py`` (``"Cl-Na"``, ``"Cl-Na-O"`` -- sorted element symbols joined
by ``-``), and ``build_jobspec`` decodes the id back into ``--elements``
arguments. A malformed id is refused clearly by ``validate_inputs``.
Fan-out contract: the runner writes ``children.json`` (declared in
``produces``) plus one ``children/<formula>/composition.json`` per
charge-balanced composition, all in the job's cwd; the orchestrator adopts
the children and retires the seed as ``EXPANDED``. Zero children is a valid
result -- some combinations simply have no charge-balanced compound.
"""
from __future__ import annotations
import re
from pathlib import Path
from typing import TYPE_CHECKING
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 compose 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 = "compose_runner.py"
#: One element symbol: a capital letter plus up to two lowercase letters.
_SYMBOL = re.compile(r"[A-Z][a-z]{0,2}")
#: How many characters of a failed runner's stderr to surface in the verdict.
_REASON_TAIL = 300
[docs]
class ComposeStage(Stage):
"""Turn one element-combination seed into candidate-composition children.
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.
threshold
Maximum stoichiometric coefficient SMACT explores. Higher finds more
exotic stoichiometries and costs combinatorially more.
max_unary_atoms
Cap for single-element seeds: the runner emits the elemental
compositions ``X1..Xn`` (allotrope stoichiometries -- SMACT's charge
balance cannot apply to one element). ``None`` (the default) follows
``threshold``; set it separately to reach e.g. ``C60`` without
exploding SMACT's multi-element search.
Raises
------
StageError
For any invalid option. Every value is checked at construction, so a
config mistake surfaces when the run is wired rather than as N failed
jobs later.
"""
fan_out = True
runtime_requirements = (
Requirement(
module="smact",
pip_name="smact",
purpose="charge-balance filtering of element combinations (compose runner)",
),
)
def __init__(
self,
python: str = "python3",
threshold: int = 2,
max_unary_atoms: int | None = None,
) -> None:
if not isinstance(python, str) or not python.strip():
raise StageError(f"compose: python must be a non-empty string, got {python!r}")
if isinstance(threshold, bool) or not isinstance(threshold, int) or threshold < 1:
raise StageError(f"compose: threshold must be an integer >= 1, got {threshold!r}")
if max_unary_atoms is not None and (
isinstance(max_unary_atoms, bool)
or not isinstance(max_unary_atoms, int)
or max_unary_atoms < 1
):
raise StageError(
f"compose: max_unary_atoms must be an integer >= 1 (or omitted to "
f"follow threshold), got {max_unary_atoms!r}"
)
self.runner_python = python
self._threshold = threshold
self._max_unary_atoms = max_unary_atoms
@property
def name(self) -> str:
return "compose"
@property
def requires(self) -> StageInputSpec:
return StageInputSpec(
files=(),
description="no input files: the seed id itself encodes the element combination",
)
@property
def produces(self) -> StageOutputSpec:
return StageOutputSpec(
files=(CHILDREN_MANIFEST,),
description="fan-out manifest + children/<formula>/composition.json per child",
)
[docs]
@staticmethod
def elements_of(item: WorkItem) -> list[str]:
"""Decode a seed id into its element symbols.
Parameters
----------
item
The seed, whose ``id`` looks like ``"Cl-Na-O"``.
Returns
-------
list of str
The symbols, order preserved as-is -- seed ids from
``core/seeds.py`` are already sorted.
Raises
------
StageInputError
Naming exactly what is wrong: an empty or non-symbol part (symbols
are ``[A-Z][a-z]{0,2}``), or a duplicated element.
"""
parts = item.id.split("-")
bad = [part for part in parts if not _SYMBOL.fullmatch(part)]
if bad:
raise StageInputError(
f"stage 'compose' cannot run on item {item.id!r}: a seed id must be "
f"element symbols joined by '-' (e.g. 'Cl-Na'); invalid part(s) {bad}"
)
if len(set(parts)) != len(parts):
raise StageInputError(
f"stage 'compose' cannot run on item {item.id!r}: duplicate element symbols"
)
return parts
[docs]
def build_jobspec(self, item: WorkItem) -> JobSpec:
elements = self.elements_of(item)
runner = Path(__file__).with_name(RUNNER_NAME)
command = [
self.runner_python,
RUNNER_NAME,
"--elements",
*elements,
"--threshold",
str(self._threshold),
]
if self._max_unary_atoms is not None:
command += ["--max-unary-atoms", str(self._max_unary_atoms)]
return JobSpec(
command=command,
input_artifacts={RUNNER_NAME: str(runner)},
# 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 compose job.
Parameters
----------
item
The seed 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 compose failure is an environment problem --
missing smact, a broken interpreter -- that a retry can genuinely
resolve once the user fixes the venv. The runner reserves science
judgments for exit 0.
"""
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"compose runner exited {code}: {tail}")
register("compose", ComposeStage)