#!/usr/bin/env python3
"""MatMelt compose runner: SMACT charge-balance ONE element combination.
Self-contained by design -- imports **stdlib + smact 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 compose_runner.py --elements Cl Na
python compose_runner.py --help
Given one element combination (any arity: binary, ternary, ...), it asks
SMACT for the charge-balanced compositions and writes the fan-out result into
the current directory (or ``--out-dir``)::
children.json # JSON list of child ids (may be [])
children/<formula>/composition.json # one seed file per child
A child id is the reduced formula with explicit counts (``Cl1Na1``,
``Cl2Na1O3``): deterministic, filesystem-safe, unambiguous. Compositions that
reduce to the same formula under different oxidation-state assignments merge
into one child carrying every assignment.
**Unary (single-element) seeds are special-cased.** SMACT's charge-neutrality
filter yields nothing for one element (no charge-balanced unary compound), so
a single ``--elements X`` emits the elemental compositions ``X1..Xn`` directly
(``n`` = ``--max-unary-atoms``, defaulting to ``--threshold``) -- allotrope
stoichiometries like ``C2`` or ``C60`` for the structure stage to find cells
for. Elemental compositions are deliberately **never gcd-reduced** (reduction
would collapse every ``Xn`` into ``X1``) and never touch smact at all.
Exit codes: ``0`` = completed, **including zero viable compositions** (an
empty fan-out is a valid scientific result, not an error); ``2`` = bad
arguments (argparse); ``3`` = smact 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 functools import reduce
from math import gcd
from pathlib import Path
from typing import Any
MANIFEST_NAME = "children.json"
CHILDREN_DIR_NAME = "children"
def _parse_args(argv: list[str] | None) -> argparse.Namespace:
parser = argparse.ArgumentParser(
prog="compose_runner",
description="SMACT charge-balance one element combination into candidate compositions.",
)
parser.add_argument(
"--elements",
nargs="+",
required=True,
metavar="SYMBOL",
help="Element symbols of one combination, e.g. --elements Cl Na O",
)
parser.add_argument(
"--threshold",
type=int,
default=2,
help="Maximum stoichiometric coefficient SMACT explores (default: 2).",
)
parser.add_argument(
"--max-unary-atoms",
type=int,
default=None,
metavar="N",
help=(
"Single-element seeds emit the elemental compositions X1..XN "
"(default: --threshold). Ignored for multi-element combinations."
),
)
parser.add_argument(
"--out-dir",
default=".",
metavar="PATH",
help="Directory to write children.json and children/ into (default: cwd).",
)
args = parser.parse_args(argv)
if args.threshold < 1:
parser.error(f"--threshold must be >= 1, got {args.threshold}")
if args.max_unary_atoms is not None and args.max_unary_atoms < 1:
parser.error(f"--max-unary-atoms must be >= 1, got {args.max_unary_atoms}")
return args
[docs]
def reduced_counts(counts: list[int]) -> list[int]:
"""Reduce a stoichiometry to its smallest whole-number form.
Parameters
----------
counts
Atom counts per element, all at least 1.
Returns
-------
list of int
The counts divided by their greatest common divisor.
Examples
--------
>>> reduced_counts([2, 2])
[1, 1]
>>> reduced_counts([4, 6, 2])
[2, 3, 1]
"""
divisor = reduce(gcd, counts)
return [count // divisor for count in counts]
[docs]
def merge_compositions(
raw: list[tuple[Any, Any, Any]],
) -> dict[str, dict[str, Any]]:
"""Fold SMACT's charge-balanced tuples into unique children.
Parameters
----------
raw
``(symbols, ox_states, counts)`` tuples as ``smact_filter`` yields
them.
Returns
-------
dict
``{formula_id: composition}``, where each composition carries
``formula``, ``elements``, ``stoichiometry`` and ``oxidation_states``.
Stoichiometries are reduced first, so entries landing on the same
formula -- the same compound reached through different oxidation-state
assignments, or an unreduced multiple -- merge into one composition
carrying every distinct assignment.
"""
compositions: dict[str, dict[str, Any]] = {}
for symbols_raw, ox_raw, counts_raw in raw:
symbols = [str(symbol) for symbol in symbols_raw]
counts = reduced_counts([int(count) for count in counts_raw])
fid = formula_id(symbols, counts)
entry = compositions.setdefault(
fid,
{
"formula": fid,
"elements": symbols,
"stoichiometry": counts,
"oxidation_states": [],
},
)
ox_states = [int(ox) for ox in ox_raw]
if ox_states not in entry["oxidation_states"]:
entry["oxidation_states"].append(ox_states)
return compositions
[docs]
def unary_compositions(element: str, max_atoms: int) -> dict[str, dict[str, Any]]:
"""Enumerate elemental compositions for a single-element seed.
Charge balance cannot apply to one element, so SMACT is never consulted --
and never imported -- on this path.
Parameters
----------
element
The element symbol.
max_atoms
The largest formula unit to emit.
Returns
-------
dict
``{formula_id: composition}`` for ``X1`` through ``X<max_atoms>``,
mirroring ``merge_compositions``'s shape so ``write_children`` is
reused unchanged. The stoichiometry is the atom count per formula unit
and is deliberately **not** gcd-reduced: ``C2`` and ``C60`` are
distinct allotrope candidates, not multiples of ``C1``. The oxidation
state of a neutral elemental solid is 0.
Examples
--------
>>> sorted(unary_compositions("C", 3))
['C1', 'C2', 'C3']
"""
return {
fid: {
"formula": fid,
"elements": [element],
"stoichiometry": [n],
"oxidation_states": [[0]],
}
for n in range(1, max_atoms + 1)
for fid in [formula_id([element], [n])]
}
[docs]
def discover(elements: list[str], threshold: int) -> list[tuple[Any, Any, Any]]:
"""Ask SMACT for the charge-balanced compositions of one element combination.
The single SMACT touchpoint in this file. The import is inside the function
body, which is what lets the module load -- and ``--help`` run -- in an
environment with no science stack at all.
Parameters
----------
elements
Two or more element symbols.
threshold
The largest stoichiometric coefficient SMACT explores.
Returns
-------
list of tuple
``(symbols, ox_states, counts)`` per charge-balanced candidate.
Raises
------
ImportError
When smact is not installed in this interpreter. ``main`` catches it
and exits 3 with the exact install command.
"""
import smact
from smact.screening import smact_filter
element_objects = [smact.Element(symbol) for symbol in elements]
return list(smact_filter(element_objects, threshold=threshold))
[docs]
def write_children(compositions: dict[str, dict[str, Any]], out_dir: Path) -> None:
"""Write one directory per composition, then the manifest.
Parameters
----------
compositions
``{formula_id: composition}`` to write.
out_dir
The job's working directory, which is the parent item's directory.
Notes
-----
The manifest is the driver's signal that the fan-out result is complete, so
it is written **after** every child directory and via tmp-then-rename. A
crash mid-run therefore leaves child directories but no manifest -- never a
manifest naming children that do not exist.
"""
for fid, composition in sorted(compositions.items()):
child_dir = out_dir / CHILDREN_DIR_NAME / fid
child_dir.mkdir(parents=True, exist_ok=True)
(child_dir / "composition.json").write_text(
json.dumps(composition, indent=2), encoding="utf-8"
)
manifest = out_dir / MANIFEST_NAME
tmp = manifest.with_name(MANIFEST_NAME + ".tmp")
tmp.write_text(json.dumps(sorted(compositions), indent=2), encoding="utf-8")
tmp.replace(manifest)
[docs]
def main(argv: list[str] | None = None) -> int:
"""Run the compose stage for one seed.
Parameters
----------
argv
Command-line arguments. ``None`` reads ``sys.argv``.
Returns
-------
int
The process exit code: ``0`` when the fan-out completed, including a
fan-out into zero children; ``2`` for a usage error (argparse's own);
``3`` when smact is missing, with the exact ``pip install`` command and
this interpreter's path printed to stderr.
"""
args = _parse_args(argv)
if len(args.elements) == 1:
# Unary seed: charge balance cannot apply, and smact is never imported
# -- elemental compositions are enumerated directly.
max_atoms = args.max_unary_atoms if args.max_unary_atoms is not None else args.threshold
compositions = unary_compositions(args.elements[0], max_atoms)
write_children(compositions, Path(args.out_dir))
print(
f"{len(compositions)} elemental composition(s) for "
f"{args.elements[0]} (max_unary_atoms={max_atoms})"
)
return 0
try:
raw = discover(args.elements, args.threshold)
except ImportError as exc:
print(
"compose_runner needs the 'smact' package, which is not importable in "
f"this python environment ({sys.executable}).\n"
f"Install it there with: pip install smact\n(import error: {exc})",
file=sys.stderr,
)
return 3
compositions = merge_compositions(raw)
write_children(compositions, Path(args.out_dir))
print(
f"{len(compositions)} charge-balanced composition(s) for "
f"{'-'.join(args.elements)} (threshold={args.threshold})"
)
return 0
if __name__ == "__main__": # pragma: no cover — exercised via subprocess
sys.exit(main())