Source code for matmelt.core.stages.extract_runner

#!/usr/bin/env python3
"""MatMelt extract runner: assemble ONE material's OPTIMADE-shaped record.

Self-contained AND **stdlib only** -- extraction is bookkeeping, not science,
so unlike the other runners there is nothing to install and no exit-3 path:
it runs under any python, local or remote. It is equally runnable by hand::

    python extract_runner.py --id Cl1Na1_225
    python extract_runner.py --help

Inputs (all in the cwd, or under ``--in-dir``):

* ``relaxed.cif``        -- the screened structure (required); parsed by a
  minimal built-in CIF reader targeting ASE-written files. Geometry is
  best-effort: a parse failure degrades to a geometry-less record with a
  stderr warning -- never a crash (composition then comes from the metadata).
* ``screen_result.json`` -- must say ``category: RELAXED``; extract serves
  only screened-in materials (exit 2 otherwise, stated clearly).
* ``structure.json``     -- optional provenance from the structure stage
  (space group, lattice type); folded in when present.
* ``calc_result.json``   -- optional DFTB+ outcome from the calculate stage
  (WorkChain status + the full healing-attempt history, SKF paths already
  stripped by its runner); folded under ``_mcli.calculate`` when present.

Output: ``record.json`` (written atomically) -- one OPTIMADE ``structures``
-shaped document: standard attributes (elements, formulas, sites, lattice)
plus the non-standard screening/provenance extras namespaced under
``_mcli``, per the OPTIMADE custom-field convention.

Exit codes: ``0`` = record written; ``2`` = bad arguments, a missing or
malformed input, a non-RELAXED verdict, or no composition source at all;
``1`` = unexpected error (traceback on stderr).
"""

from __future__ import annotations

import argparse
import json
import math
import shlex
import string
import sys
from datetime import datetime, timezone
from functools import reduce
from pathlib import Path
from typing import Any

RECORD_NAME = "record.json"
RELAXED_NAME = "relaxed.cif"
SCREEN_NAME = "screen_result.json"
META_NAME = "structure.json"
CALC_NAME = "calc_result.json"


def _parse_args(argv: list[str] | None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        prog="extract_runner",
        description="Assemble one screened-in material's OPTIMADE-shaped record.json.",
    )
    parser.add_argument(
        "--id",
        required=True,
        metavar="ID",
        help="The material id to stamp on the record (e.g. Cl1Na1_225).",
    )
    parser.add_argument(
        "--in-dir",
        default=".",
        metavar="PATH",
        help="Directory holding relaxed.cif / screen_result.json and the optional "
        "structure.json / calc_result.json (default: cwd).",
    )
    parser.add_argument(
        "--out-dir",
        default=".",
        metavar="PATH",
        help="Directory to write record.json into (default: cwd).",
    )
    return parser.parse_args(argv)


# ---------------------------------------------------------------------------
# Minimal CIF reading (pure stdlib; targets ASE-written CIFs, degrades safely)
# ---------------------------------------------------------------------------

_CELL_TAGS = (
    "_cell_length_a",
    "_cell_length_b",
    "_cell_length_c",
    "_cell_angle_alpha",
    "_cell_angle_beta",
    "_cell_angle_gamma",
)


def _cif_number(token: str) -> float:
    """Read one CIF numeric value.

    Parameters
    ----------
    token
        The raw token, possibly carrying a standard-uncertainty suffix.

    Returns
    -------
    float
        The value, with any ``(5)``-style uncertainty stripped.

    Examples
    --------
    >>> _cif_number("1.234(5)")
    1.234
    """
    return float(token.split("(")[0])


[docs] def parse_cif(text: str) -> dict[str, Any]: """Parse cell parameters and atom sites out of one CIF document. A minimal reader targeting the shape ASE writes: scalar ``_cell_*`` tags plus one ``loop_`` whose ``_atom_site_*`` columns may come in any order. Extra loops are skipped, and the site symbol falls back to ``_atom_site_label`` when no type symbol column exists. Parameters ---------- text The CIF document. Returns ------- dict ``{"cell": (a, b, c, alpha, beta, gamma), "sites": [(symbol, fx, fy, fz), ...]}``. Raises ------ ValueError Naming what is missing -- a cell tag, the coordinate columns, a short row, or any atom sites at all. The caller decides whether that is fatal; ``main`` degrades loudly rather than failing the record. """ cell: dict[str, float] = {} sites: list[tuple[str, float, float, float]] = [] lines = [line.strip() for line in text.splitlines()] i = 0 while i < len(lines): line = lines[i] if not line or line.startswith("#"): i += 1 continue if line.split()[0] in _CELL_TAGS: tag, *rest = line.split() if rest: cell[tag] = _cif_number(rest[0]) i += 1 continue if line == "loop_": headers: list[str] = [] i += 1 while i < len(lines) and lines[i].startswith("_"): headers.append(lines[i].split()[0]) i += 1 if not any(h.startswith("_atom_site") for h in headers): continue # some other loop; its data rows fail the header test below columns = {name: idx for idx, name in enumerate(headers)} sym_col = columns.get("_atom_site_type_symbol", columns.get("_atom_site_label")) fx_col = columns.get("_atom_site_fract_x") fy_col = columns.get("_atom_site_fract_y") fz_col = columns.get("_atom_site_fract_z") if sym_col is None or fx_col is None or fy_col is None or fz_col is None: raise ValueError( "CIF atom-site loop lacks a symbol/label or fractional-coordinate column" ) while i < len(lines): row = lines[i] if not row or row.startswith(("_", "loop_", "data_", "#")): break tokens = shlex.split(row) if len(tokens) < len(headers): raise ValueError(f"CIF atom-site row has too few columns: {row!r}") symbol = tokens[sym_col].rstrip("0123456789+-") sites.append( ( symbol, _cif_number(tokens[fx_col]), _cif_number(tokens[fy_col]), _cif_number(tokens[fz_col]), ) ) i += 1 continue i += 1 missing = [tag for tag in _CELL_TAGS if tag not in cell] if missing: raise ValueError(f"CIF is missing cell tag(s): {', '.join(missing)}") if not sites: raise ValueError("CIF holds no atom sites") return {"cell": tuple(cell[tag] for tag in _CELL_TAGS), "sites": sites}
[docs] def cell_to_lattice( a: float, b: float, c: float, alpha: float, beta: float, gamma: float ) -> list[list[float]]: """Convert the six cell parameters into lattice vectors. Parameters ---------- a, b, c Cell edge lengths, in angstrom. alpha, beta, gamma Cell angles, in degrees. Returns ------- list Three lattice vectors (rows) in angstrom, in the standard crystallographic convention: ``a`` along x, ``b`` in the xy-plane. Right-angle float noise is snapped to honest zeros, since ``cos(90 deg)`` is ~6e-17 rather than 0. Examples -------- >>> cell_to_lattice(4.0, 4.0, 4.0, 90.0, 90.0, 90.0) [[4.0, 0.0, 0.0], [0.0, 4.0, 0.0], [0.0, 0.0, 4.0]] """ al, be, ga = (math.radians(angle) for angle in (alpha, beta, gamma)) cx = c * math.cos(be) cy = c * (math.cos(al) - math.cos(be) * math.cos(ga)) / math.sin(ga) cz = math.sqrt(max(c * c - cx * cx - cy * cy, 0.0)) def snap(value: float) -> float: # cos(90 deg) is ~6e-17, not 0 -- zero-snap the float noise so right # angles serialize as honest zeros in the record. return 0.0 if abs(value) < 1e-10 else value return [ [snap(v) for v in row] for row in ( [a, 0.0, 0.0], [b * math.cos(ga), b * math.sin(ga), 0.0], [cx, cy, cz], ) ]
[docs] def frac_to_cart(frac: tuple[float, float, float], lattice: list[list[float]]) -> list[float]: """Convert fractional coordinates to cartesian ones. Parameters ---------- frac An ``(fx, fy, fz)`` triple. lattice The three lattice vectors to express it against. Returns ------- list of float The cartesian position, in angstrom. """ return [sum(frac[row] * lattice[row][axis] for row in range(3)) for axis in range(3)]
# --------------------------------------------------------------------------- # OPTIMADE assembly (pure -- fully unit-testable) # ---------------------------------------------------------------------------
[docs] def chemical_formulas(counts_by_element: dict[str, int]) -> dict[str, Any]: """Build the OPTIMADE chemistry fields for one composition. Parameters ---------- counts_by_element ``{symbol: count}`` for the whole structure. Returns ------- dict The OPTIMADE fields: ``elements`` (sorted) and ``nelements``; ``elements_ratios`` as fractions of the total, per sorted element; ``chemical_formula_reduced``, gcd-reduced with a count of 1 omitted; ``chemical_formula_descriptive``, the same but keeping every count; and ``chemical_formula_anonymous``, reduced counts descending and lettered A, B, ... Raises ------ ValueError When there are no elements -- a record with no chemistry serves nobody. Examples -------- >>> fields = chemical_formulas({"Na": 4, "Cl": 4}) >>> fields["chemical_formula_reduced"], fields["chemical_formula_anonymous"] ('ClNa', 'AB') """ if not counts_by_element: raise ValueError("no elements to build formulas from") symbols = sorted(counts_by_element) counts = [counts_by_element[sym] for sym in symbols] divisor = reduce(math.gcd, counts) reduced = [count // divisor for count in counts] total = sum(counts) def fmt(symbol: str, count: int) -> str: return symbol if count == 1 else f"{symbol}{count}" # strict=False: the alphabet is deliberately longer than the element list letter_counts = zip(string.ascii_uppercase, sorted(reduced, reverse=True), strict=False) anonymous = "".join(fmt(letter, count) for letter, count in letter_counts) return { "elements": symbols, "nelements": len(symbols), "elements_ratios": [count / total for count in counts], "chemical_formula_reduced": "".join( fmt(s, n) for s, n in zip(symbols, reduced, strict=True) ), "chemical_formula_descriptive": "".join( f"{s}{n}" for s, n in zip(symbols, reduced, strict=True) ), "chemical_formula_anonymous": anonymous, }
[docs] def structure_block(parsed: dict[str, Any]) -> dict[str, Any]: """Build the OPTIMADE geometry attributes from a parsed CIF. Parameters ---------- parsed What ``parse_cif`` returned. Returns ------- dict ``lattice_vectors``, ``cartesian_site_positions``, ``nsites``, ``species_at_sites``, ``species``, and the periodicity fields. Every site is treated as fully occupied by a single element, which is what a generated-and-relaxed structure always is. """ lattice = cell_to_lattice(*parsed["cell"]) symbols = [site[0] for site in parsed["sites"]] positions = [frac_to_cart((fx, fy, fz), lattice) for _, fx, fy, fz in parsed["sites"]] return { "lattice_vectors": lattice, "cartesian_site_positions": positions, "nsites": len(symbols), "species_at_sites": symbols, "species": [ {"name": sym, "chemical_symbols": [sym], "concentration": [1.0]} for sym in sorted(set(symbols)) ], "dimension_types": [1, 1, 1], "nperiodic_dimensions": 3, "structure_features": [], }
[docs] def build_record( material_id: str, screen: dict[str, Any], parsed_cif: dict[str, Any] | None, meta: dict[str, Any], calc: dict[str, Any] | None = None, ) -> dict[str, Any]: """Assemble the OPTIMADE ``structures``-shaped document. Parameters ---------- material_id The record's id. screen The screening verdict, whose category, convergence, step count, final observables and settings are carried into the record. parsed_cif The parsed relaxed structure, or ``None`` when the CIF could not be read. Geometry is best-effort: without it the record is still written, with ``_mcli.geometry_source`` set to ``null``. meta The structure stage's provenance -- space group, dimensionality, lattice type -- and the fallback composition source. calc The calculate stage's result document, folded **verbatim** under ``_mcli.calculate`` when given, healing history included. Left ``None``, the record is byte-for-byte what a calculate-less funnel always produced. Returns ------- dict The complete record, with everything MatMelt-specific namespaced under ``_mcli`` per the OPTIMADE custom-field convention. Raises ------ ValueError When neither the CIF nor the metadata can supply a composition. """ attributes: dict[str, Any] = {} if parsed_cif is not None: attributes.update(structure_block(parsed_cif)) counts: dict[str, int] = {} for symbol, *_ in parsed_cif["sites"]: counts[symbol] = counts.get(symbol, 0) + 1 elif meta.get("elements") and meta.get("stoichiometry"): # strict: a mismatched metadata pair is a defect, not a composition counts = dict(zip(meta["elements"], meta["stoichiometry"], strict=True)) else: raise ValueError( "no composition source: the CIF is unusable and structure.json " "carries no elements/stoichiometry" ) attributes.update(chemical_formulas(counts)) attributes["last_modified"] = datetime.now(timezone.utc).isoformat(timespec="seconds") mcli: dict[str, Any] = { "screen": { key: screen.get(key) for key in ("category", "converged", "n_steps", "final", "settings") if key in screen }, "geometry_source": RELAXED_NAME if parsed_cif is not None else None, } provenance = { key: meta[key] for key in ("dimension", "space_group", "space_group_symbol", "lattice_type", "formula") if key in meta } if provenance: mcli["provenance"] = provenance if calc is not None: mcli["calculate"] = calc attributes["_mcli"] = mcli return {"id": material_id, "type": "structures", "attributes": attributes}
[docs] def write_record(record: dict[str, Any], out_dir: Path) -> None: """Write ``record.json`` atomically. Parameters ---------- record The assembled document. out_dir The job's working directory. """ target = out_dir / RECORD_NAME tmp = target.with_name(RECORD_NAME + ".tmp") tmp.write_text(json.dumps(record, indent=2), encoding="utf-8") tmp.replace(target)
# --------------------------------------------------------------------------- # Entry point # --------------------------------------------------------------------------- def _load_json(path: Path, what: str) -> dict[str, Any]: """Read a required JSON object. Parameters ---------- path The file to read. what A human name for the file, used in error messages. Returns ------- dict The decoded object. Raises ------ ValueError Naming the defect: unreadable, not valid JSON, or not an object. """ try: data = json.loads(path.read_text(encoding="utf-8")) except OSError as exc: raise ValueError(f"cannot read {what} {path}: {exc}") from exc except json.JSONDecodeError as exc: raise ValueError(f"{what} {path} is not valid JSON: {exc}") from exc if not isinstance(data, dict): raise ValueError(f"{what} {path} must hold a JSON object, got {type(data).__name__}") return data
[docs] def main(argv: list[str] | None = None) -> int: """Assemble one screened-in material's record. Parameters ---------- argv Command-line arguments. ``None`` reads ``sys.argv``. Returns ------- int The process exit code: ``0`` when the record was written; ``2`` for a usage error, an unusable screening verdict, a non-``RELAXED`` verdict, or a structure that yields no composition from any source. There is no exit 3 here -- the runner is stdlib-only, so it has no dependency to be missing. Notes ----- Geometry is best-effort and degrades **loudly**: an unparseable CIF prints a warning to stderr and the record is still written from the structure stage's metadata, with ``_mcli.geometry_source`` set to ``null``. The optional ``structure.json`` and ``calc_result.json`` inputs degrade the same way. Only a total absence of composition is fatal. """ args = _parse_args(argv) in_dir = Path(args.in_dir) try: screen = _load_json(in_dir / SCREEN_NAME, "screening verdict") except ValueError as exc: print(str(exc), file=sys.stderr) return 2 if screen.get("category") != "RELAXED": print( f"refusing to extract: {SCREEN_NAME} says category " f"{screen.get('category')!r}, but extract serves only RELAXED " f"(screened-in) materials", file=sys.stderr, ) return 2 cif_path = in_dir / RELAXED_NAME if not cif_path.is_file(): print(f"cannot read structure file {cif_path}: not a file", file=sys.stderr) return 2 parsed_cif: dict[str, Any] | None try: parsed_cif = parse_cif(cif_path.read_text(encoding="utf-8")) except (OSError, ValueError) as exc: # Geometry is best-effort: degrade loudly, never crash the record. print( f"WARNING: could not parse {cif_path} ({exc}); record will carry no geometry", file=sys.stderr, ) parsed_cif = None meta: dict[str, Any] = {} if (in_dir / META_NAME).is_file(): try: meta = _load_json(in_dir / META_NAME, "structure metadata") except ValueError as exc: print(f"WARNING: ignoring {META_NAME}: {exc}", file=sys.stderr) calc: dict[str, Any] | None = None if (in_dir / CALC_NAME).is_file(): try: calc = _load_json(in_dir / CALC_NAME, "calculate result") except ValueError as exc: print(f"WARNING: ignoring {CALC_NAME}: {exc}", file=sys.stderr) try: record = build_record(args.id, screen, parsed_cif, meta, calc) except ValueError as exc: print(str(exc), file=sys.stderr) return 2 write_record(record, Path(args.out_dir)) nsites = record["attributes"].get("nsites", "no") print( f"{args.id} -> {RECORD_NAME} " f"({record['attributes']['chemical_formula_reduced']}, {nsites} site(s))" ) return 0
if __name__ == "__main__": # pragma: no cover — exercised via subprocess sys.exit(main())