Source code for matmelt.core.stages.structure_runner

#!/usr/bin/env python3
"""MatMelt structure runner: PyXtal-generate candidates for ONE composition.

Self-contained by design -- imports **stdlib + pyxtal 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 structure_runner.py --composition composition.json
    python structure_runner.py --composition composition.json --dims 3D 2D
    python structure_runner.py --help

Given one composition (``formula`` + ``elements`` + ``stoichiometry``), it
asks PyXtal for one random symmetry-constrained crystal in every group of each
requested **dimensionality**. Each dimensionality is swept over its full
symmetry family -- 3D = 230 space groups, 2D = 80 layer groups, 1D = 75 rod
groups (``--dims``, default ``3D``). Groups whose constraints cannot be met are
skipped -- but **counted and reported**, never swallowed silently. The fan-out
result lands in the current directory (or ``--out-dir``)::

    children.json                                # JSON list of child ids (may be [])
    children/<Formula>_<DIM>-<SG>/structure.cif  # the candidate crystal
    children/<Formula>_<DIM>-<SG>/structure.json # provenance: dim, SG, symbol, lattice

A child id is ``<formula>_<DIM>-<space-group-number>`` (``Cl1Na1_3D-225``):
the dimensionality tag and group *number* are joined with ``-`` so the last
underscore stays the formula boundary, and space-group *symbols* (``P2_1/c``,
which carry path separators) never appear. Deterministic and filesystem-safe.

Exit codes: ``0`` = completed, **including zero valid structures** (an empty
fan-out is a valid scientific result, not an error); ``2`` = bad arguments or
an unreadable/malformed composition file; ``3`` = pyxtal is not importable in
this python environment (stderr says exactly what to install); ``1`` =
unexpected error (traceback on stderr).
"""

from __future__ import annotations

import argparse
import json
import sys
from collections.abc import Callable, Sequence
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

MANIFEST_NAME = "children.json"
CHILDREN_DIR_NAME = "children"

#: Dimensionality tag -> (PyXtal ``dim`` argument, highest group number in that
#: family). 3D = 230 space groups, 2D = 80 layer groups, 1D = 75 rod groups.
#: Each dimensionality is swept over its own full family (a group number above a
#: family's count does not exist for it).
DIM_GROUPS: dict[str, tuple[int, int]] = {"3D": (3, 230), "2D": (2, 80), "1D": (1, 75)}

#: Keys one composition file must carry (the compose stage's child seed shape).
REQUIRED_KEYS = ("formula", "elements", "stoichiometry")

#: At most this many per-group error messages are echoed in the summary.
_MAX_ERROR_EXAMPLES = 3


def _parse_args(argv: list[str] | None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        prog="structure_runner",
        description="PyXtal-generate symmetry-constrained candidate crystals for one composition.",
    )
    parser.add_argument(
        "--composition",
        default="composition.json",
        metavar="PATH",
        help="JSON file holding one composition object with keys "
        "formula/elements/stoichiometry (default: composition.json).",
    )
    parser.add_argument(
        "--out-dir",
        default=".",
        metavar="PATH",
        help="Directory to write children.json and children/ into (default: cwd).",
    )
    parser.add_argument(
        "--dims",
        nargs="+",
        default=["3D"],
        choices=["3D", "2D", "1D"],
        metavar="DIM",
        help="Dimensionalities to generate (any of 3D/2D/1D; default 3D). Each is "
        "swept over its full symmetry family: 3D=230 space, 2D=80 layer, 1D=75 rod groups.",
    )
    parser.add_argument(
        "--space-groups",
        nargs=2,
        type=int,
        default=None,
        metavar=("FIRST", "LAST"),
        help="Optional debug/smoke limiter: restrict each dimensionality's sweep to the "
        "inclusive group range [FIRST, LAST] (capped to that family's size). Omit for the "
        "full family. Handy for debugging one group by hand or a quick real-PyXtal run.",
    )
    args = parser.parse_args(argv)
    if args.space_groups is not None:
        first, last = args.space_groups
        if not (1 <= first <= last):
            parser.error(f"--space-groups must satisfy 1 <= FIRST <= LAST, got {first} {last}")
    return args


[docs] def load_composition(path: Path) -> dict[str, Any]: """Read and validate one composition file. The same rules ``StructureStage.composition_of`` enforces, repeated here because the runner must stay debuggable standalone -- it cannot import the stage that normally checks its input. Parameters ---------- path The ``composition.json`` to read. Returns ------- dict The decoded composition. Raises ------ ValueError Naming the defect: unreadable, not valid JSON, not an object, a missing required key, a formula that is not alphanumeric (it names child directories), or elements and stoichiometry that do not line up. """ try: data = json.loads(path.read_text(encoding="utf-8")) except OSError as exc: raise ValueError(f"cannot read composition file {path}: {exc}") from exc except json.JSONDecodeError as exc: raise ValueError(f"composition file {path} is not valid JSON: {exc}") from exc if not isinstance(data, dict): raise ValueError( f"composition file {path} must hold one composition object, got {type(data).__name__}" ) missing = [key for key in REQUIRED_KEYS if key not in data] if missing: raise ValueError(f"composition is missing required key(s): {', '.join(missing)}") formula = data["formula"] if not isinstance(formula, str) or not formula or not formula.isalnum(): raise ValueError( 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 ValueError(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 ValueError( f"stoichiometry must be a list of integers >= 1, one per element, " f"got {stoichiometry!r} for elements {elements!r}" ) return data
[docs] def child_id(formula: str, dim_tag: str, space_group: int) -> str: """Build the child id for one generated structure. Parameters ---------- formula The parent composition's formula. dim_tag ``"3D"``, ``"2D"`` or ``"1D"``. space_group The group **number**. Group *symbols* never appear in an id: ``P2_1/c`` carries a path separator, and these ids become directory names. Returns ------- str ``<formula>_<dim>-<group>``. The dim tag and group number join with ``-`` so the last underscore stays the formula boundary. Examples -------- >>> child_id("Cl1Na1", "3D", 225) 'Cl1Na1_3D-225' """ return f"{formula}_{dim_tag}-{space_group}"
def _group_range(dim_tag: str, limit: tuple[int, int] | None) -> range: """Give the group numbers to sweep for one dimensionality. Parameters ---------- dim_tag Which symmetry family: ``"3D"`` (230 space groups), ``"2D"`` (80 layer groups) or ``"1D"`` (75 rod groups). limit An inclusive ``(first, last)`` cap, clamped to the family's real range. ``None`` sweeps the whole family. Returns ------- range The group numbers, ascending. """ _, max_group = DIM_GROUPS[dim_tag] if limit is None: return range(1, max_group + 1) first, last = limit return range(max(1, first), min(max_group, last) + 1)
[docs] @dataclass class GenerationReport: """What happened across the swept groups -- nothing is silent. Attributes ---------- children : list of str The ids of the structures actually written, in generation order. n_rejected : int Groups where a crystal was generated but its symmetry constraints were not met (``crystal.valid`` false). A normal outcome, not an error. n_errors : int Groups where PyXtal raised -- an incompatible stoichiometry, for instance. error_examples : list of str Up to a few example messages, so a systematic failure is visible in the job's stdout instead of being inferred from a zero count. """ children: list[str] = field(default_factory=list) n_rejected: int = 0 # crystal generated but symmetry constraints not met n_errors: int = 0 # PyXtal raised (incompatible stoichiometry, ...) error_examples: list[str] = field(default_factory=list)
[docs] def generate( composition: dict[str, Any], out_dir: Path, crystal_factory: Callable[[], Any], dims: Sequence[str] = ("3D",), group_limit: tuple[int, int] | None = None, ) -> GenerationReport: """Try each group of each requested dimensionality once, writing valid crystals. Parameters ---------- composition The validated composition to build structures for. out_dir The job's working directory; children land under ``children/<child_id>/``. crystal_factory Returns a fresh pyxtal-like object per attempt. This is the seam that keeps the loop testable without pyxtal installed. dims Which dimensionalities to sweep. group_limit An optional inclusive group range, applied per family. Useful for hand-debugging, where sweeping 230 groups is not the point. Returns ------- GenerationReport The children written plus the counted failures. Each generated structure gets a ``structure.cif`` and a ``structure.json`` recording its group, lattice type, and volume. Notes ----- A failing group never stops the sweep: it is counted, separately for symmetry rejections and raised errors, and the loop continues. The report exists precisely so that no failure is silent. """ report = GenerationReport() for dim_tag in dims: pdim, _ = DIM_GROUPS[dim_tag] for sg in _group_range(dim_tag, group_limit): crystal = crystal_factory() try: crystal.from_random( dim=pdim, group=sg, species=composition["elements"], numIons=composition["stoichiometry"], ) if not crystal.valid: report.n_rejected += 1 continue cid = child_id(composition["formula"], dim_tag, sg) child_dir = out_dir / CHILDREN_DIR_NAME / cid child_dir.mkdir(parents=True, exist_ok=True) crystal.to_file(str(child_dir / "structure.cif")) metadata = { "formula": composition["formula"], "elements": composition["elements"], "stoichiometry": composition["stoichiometry"], "dimension": dim_tag, "space_group": sg, "space_group_symbol": str(crystal.group.symbol), "lattice_type": str(crystal.group.lattice_type), "volume_A3": float(crystal.lattice.volume), } (child_dir / "structure.json").write_text( json.dumps(metadata, indent=2), encoding="utf-8" ) report.children.append(cid) except Exception as exc: # counted and reported, never a silent skip report.n_errors += 1 if len(report.error_examples) < _MAX_ERROR_EXAMPLES: report.error_examples.append(f"{dim_tag} SG {sg}: {exc}") return report
[docs] def write_manifest(children: list[str], out_dir: Path) -> None: """Write the fan-out manifest last, atomically. Parameters ---------- children The child ids to name, in generation order -- ascending dimensionality then group number, which is deterministic by construction. out_dir The job's working directory. Notes ----- The manifest is the driver's signal that the fan-out result is complete, so it lands only after every child directory exists, via tmp-then-rename. A crash mid-run leaves child directories but no manifest, never a manifest naming children that do not exist. """ manifest = out_dir / MANIFEST_NAME tmp = manifest.with_name(MANIFEST_NAME + ".tmp") tmp.write_text(json.dumps(children, indent=2), encoding="utf-8") tmp.replace(manifest)
def _load_pyxtal_factory() -> Callable[[], Any]: """Import PyXtal -- the single touchpoint, deliberately lazy. Returns ------- callable The ``pyxtal`` class itself, used as a factory of fresh crystal objects. Raises ------ ImportError When pyxtal is not installed in this interpreter. ``main`` catches it and exits 3 with the exact install command. """ from pyxtal import pyxtal factory: Callable[[], Any] = pyxtal return factory
[docs] def main(argv: list[str] | None = None) -> int: """Run the structure stage for one composition. Parameters ---------- argv Command-line arguments. ``None`` reads ``sys.argv``. Returns ------- int The process exit code: ``0`` when the sweep completed, including a fan-out into zero children, with the per-group tally printed; ``2`` for a usage error or an unusable composition file; ``3`` when pyxtal is missing, with the exact ``pip install`` command and this interpreter's path printed to stderr. """ args = _parse_args(argv) try: composition = load_composition(Path(args.composition)) except ValueError as exc: print(str(exc), file=sys.stderr) return 2 try: crystal_factory = _load_pyxtal_factory() except ImportError as exc: print( "structure_runner needs the 'pyxtal' package, which is not importable in " f"this python environment ({sys.executable}).\n" f"Install it there with: pip install pyxtal\n(import error: {exc})", file=sys.stderr, ) return 3 out_dir = Path(args.out_dir) dims = tuple(args.dims) group_limit = tuple(args.space_groups) if args.space_groups is not None else None report = generate(composition, out_dir, crystal_factory, dims, group_limit) write_manifest(report.children, out_dir) n_groups = sum(len(_group_range(d, group_limit)) for d in dims) print( f"{len(report.children)} candidate structure(s) for {composition['formula']} " f"across {len(dims)} dimensionalit(y/ies) [{', '.join(dims)}], " f"{n_groups} group(s) tried " f"({report.n_rejected} rejected by symmetry, {report.n_errors} errored)" ) for example in report.error_examples: print(f" e.g. {example}") return 0
if __name__ == "__main__": # pragma: no cover — exercised via subprocess sys.exit(main())