Source code for matmelt.commands.cif

"""The ``cif`` command: export CIF files for calculated materials (M10).

A material that passed the ``calculate`` stage has its DFTB+ result in
``calc_result.json`` -- but not the geometry DFTB+ relaxed *to*. That
geometry (``geom.out.gen``) lives only in the user's AiiDA repository, so
this command:

1. picks the materials that passed calculate (state ``CALCULATED`` or
   beyond, with a ``converged`` ``calc_result.json``), optionally narrowed
   by ``--materials`` / ``--materials-file``;
2. fetches their relaxed geometries in **one** batched call to
   ``cif_runner.py``, run under the calculate stage's own ``python`` -- the
   interpreter the user's AiiDA lives in (the ``aiida_status`` precedent;
   this package never imports aiida);
3. writes ``<out-dir>/<id>.cif`` -- P1, with the WorkChain/CalcJob pks and
   the headline results as ``_mcli_*`` provenance items.

Default output is ``<work_dir>/cifs/``; item directories keep holding stage
artifacts only, and the exported set is one folder to open in VESTA or hand
to a database.

Materials that never converged (or never reached calculate) are **skipped
and counted**, not failed -- they simply have no calculated structure yet.
A material whose geometry cannot be fetched is a **failure**, reported with
its reason.

Exit codes: 0 = every eligible material exported (including "nothing
matched"), 1 = no run state / the fetch failed / some materials could not be
exported, 2 = config or usage error.
"""

from __future__ import annotations

import json
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any

from matmelt.commands import STATE_DB_FILENAME, resolve_config_path
from matmelt.config import ConfigError, load_config
from matmelt.core.cif import gen_to_cif
from matmelt.core.stages import cif_runner
from matmelt.core.stages.calculate import RESULT_NAME
from matmelt.core.state import MaterialState, StateError, StateStore
from matmelt.ui.console import (
    print_json_payload,
    render_cif_export,
    render_config_error,
    render_error,
    render_note,
    render_warning,
)

#: The self-contained fetch runner, run under the user's AiiDA interpreter.
RUNNER_PATH = Path(cif_runner.__file__)

#: Where a calculated structure can be found: the states at or past calculate.
_ELIGIBLE_STATES = (MaterialState.CALCULATED, MaterialState.EXTRACTED, MaterialState.DONE)

#: Headline numbers copied from ``final_result`` into the CIF header. Absent
#: keys become ``?`` -- an honest unknown beats a fabricated zero.
_RESULT_FIELDS = (
    "total_energy_eV",
    "band_gap_eV",
    "fermi_energy_eV",
    "max_force_eV_Ang",
)


[docs] def run_cif( config_path: str, *, materials: str | None = None, materials_file: str | None = None, out_dir: str | None = None, as_json: bool = False, ) -> int: """Export CIF files for the materials that passed the calculate stage. Parameters ---------- config_path The config file naming the workspace. It must configure a ``calculate`` stage: without one, no material has a DFTB+ relaxed geometry to export. materials Comma-separated ids to export. materials_file A file of ids, one per line. Combines with ``materials``. Selection is by id only -- an export is a flat set of files, not a lineage traversal. out_dir Where the CIFs go. Defaults to ``<work_dir>/cifs``. as_json Emit a JSON payload on stdout instead of Rich output. Returns ------- int ``0`` when the export completed, even with per-material skips or failures; ``1`` when nothing could be exported at all; ``2`` for a config error, a missing calculate stage, or an unusable selection. Notes ----- Eligibility is state ``CALCULATED`` or beyond **and** a converged ``calc_result.json``. Anything else is a counted skip with a stated reason, never a failure: a material that did not converge simply has no calculated structure to export. Geometries come from the result document's cache when a recent enough calculate run wrote one, and only the remainder is fetched through AiiDA -- so a fully-cached workspace spawns no AiiDA process at all. The store is opened read-only. """ config_path = resolve_config_path(config_path) try: config = load_config(config_path) except ConfigError as exc: render_config_error(str(exc)) return 2 stage_cfg = config.stages.get("calculate") if stage_cfg is None: render_error( "cif: this config has no 'calculate' stage, so no material has a " "DFTB+ relaxed geometry to export. Add a calculate stage (see " "'matmelt init' for a commented example) and run it first." ) return 2 options: dict[str, Any] = dict(stage_cfg.options) python = str(options.get("python") or sys.executable) profile = options.get("profile") selection = _parse_selection(materials, materials_file) if isinstance(selection, int): # error path returned an exit code return selection work_dir = Path(config.run.work_dir) db_path = work_dir / STATE_DB_FILENAME if not db_path.is_file(): render_error(f"cif: no run state at {db_path} -- start one with 'matmelt run'.") return 1 try: store = StateStore(db_path, read_only=True) except StateError as exc: render_error(f"cif: {exc}") return 1 try: candidates = sorted( ( record.material_id for state in _ELIGIBLE_STATES for record in store.query({"state": state}) ), ) finally: store.close() unmatched: list[str] = [] if selection is not None: unmatched = sorted(selection - set(candidates)) candidates = [material_id for material_id in candidates if material_id in selection] if unmatched and not as_json: render_warning( f"cif: {len(unmatched)} selected id(s) have not passed calculate (or are " f"unknown), skipped: {', '.join(unmatched)}" ) requests, skipped = _eligible_requests(work_dir, candidates) if not requests: if as_json: print_json_payload( _payload( config.run.name, _target_dir(work_dir, out_dir), [], skipped, [], [], unmatched ) ) else: render_note("cif: no calculated materials to export -- nothing written.") return 0 target = _target_dir(work_dir, out_dir) try: target.mkdir(parents=True, exist_ok=True) except OSError as exc: render_error(f"cif: cannot create output directory {target}: {exc}") return 1 # Geometries cached by a newer calculate run need no AiiDA at all. rows = [_cached_row(request) for request in requests if request["_gen"] is not None] to_fetch = [request for request in requests if request["_gen"] is None] if not as_json: render_note(_fetch_plan(len(rows), len(to_fetch))) if to_fetch: fetched = _fetch_geometries(python, profile, to_fetch) if fetched is None: return 1 rows += fetched order = {request["id"]: index for index, request in enumerate(requests)} rows.sort(key=lambda row: order.get(str(row.get("id", "")), len(order))) written, failed, notes = _write_cif_files(rows, requests, target) if as_json: print_json_payload( _payload(config.run.name, target, written, skipped, failed, notes, unmatched) ) else: for note in notes: render_warning(f"cif: {note['id']}: {note['note']}") render_cif_export(config.run.name, str(target), written, skipped, failed) return 1 if failed else 0
def _payload( run_name: str, target: Path, written: list[dict[str, Any]], skipped: list[dict[str, str]], failed: list[dict[str, str]], notes: list[dict[str, str]], unmatched: list[str], ) -> dict[str, Any]: """Build the ``--json`` document. Parameters ---------- run_name The configured run name. target Where the CIFs were written. written One entry per exported file. skipped Ineligible materials, each with its reason. failed Materials whose geometry could not be fetched or parsed. notes Per-material remarks, such as a stale recorded CalcJob pk. unmatched Selected ids the workspace does not know. Returns ------- dict Everything the human view prints has a key here, so a machine caller never has to scrape rendered text -- and nothing human-facing is printed to stdout alongside it. """ return { "run": run_name, "out_dir": str(target), "written": written, "skipped": skipped, "failed": failed, "notes": notes, "unmatched": unmatched, } # --------------------------------------------------------------------------- # Selection + eligibility # --------------------------------------------------------------------------- def _parse_selection( materials: str | None, materials_file: str | None ) -> frozenset[str] | int | None: """Combine the two selection options into one set of ids. Parameters ---------- materials Comma-separated ids, or ``None``. materials_file A file of ids, or ``None``. Same conventions as ``run`` and ``retry``: one per line, blanks and ``#`` comments skipped. Returns ------- frozenset or int or None The union of both sources; ``None`` when neither was given, meaning every eligible material; or the exit code ``2`` when the file cannot be read or an explicit selection came out empty. Notes ----- Unlike ``run``, this selects by id alone. An export is a flat set of files, not a lineage traversal, so a parent's id does not pull in its children. """ if materials is None and materials_file is None: return None picked = {part.strip() for part in (materials or "").split(",") if part.strip()} if materials_file is not None: try: lines = Path(materials_file).read_text(encoding="utf-8").splitlines() except OSError as exc: render_error(f"cif: cannot read --materials-file {materials_file}: {exc}") return 2 picked.update( line.strip() for line in lines if line.strip() and not line.lstrip().startswith("#") ) if not picked: render_error( "cif: --materials/--materials-file given but the selection is empty -- " "name at least one material id." ) return 2 return frozenset(picked) def _target_dir(work_dir: Path, out_dir: str | None) -> Path: """Decide where the exported CIFs go. Parameters ---------- work_dir The run's workspace root. out_dir An explicit destination, or ``None``. Returns ------- pathlib.Path ``out_dir`` when given, else ``<work_dir>/cifs``. """ return Path(out_dir) if out_dir is not None else work_dir / "cifs" def _eligible_requests( work_dir: Path, candidates: list[str] ) -> tuple[list[dict[str, Any]], list[dict[str, str]]]: """Split candidate materials into fetch requests and counted skips. Parameters ---------- work_dir The workspace root, under which each material's directory lives. candidates Material ids whose state makes them worth inspecting. Returns ------- tuple ``(requests, skipped)``. A material qualifies when its ``calc_result.json`` parses, says ``converged``, and names the WorkChain that produced it. Everything else is a skip with a stated reason -- never a failure, because a material that did not converge simply has no calculated structure to export. """ requests: list[dict[str, Any]] = [] skipped: list[dict[str, str]] = [] for material_id in candidates: path = work_dir / material_id / RESULT_NAME if not path.is_file(): skipped.append({"id": material_id, "reason": f"no {RESULT_NAME}"}) continue try: document = json.loads(path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError) as exc: skipped.append({"id": material_id, "reason": f"unreadable {RESULT_NAME}: {exc}"}) continue if not isinstance(document, dict): skipped.append({"id": material_id, "reason": f"{RESULT_NAME} is not a JSON object"}) continue if document.get("converged") is not True: skipped.append({"id": material_id, "reason": "did not converge"}) continue wc_pk = document.get("single_wc_pk") if not isinstance(wc_pk, int) or isinstance(wc_pk, bool): skipped.append({"id": material_id, "reason": f"{RESULT_NAME} names no WorkChain pk"}) continue attempts = document.get("attempts") calc_pk = None if isinstance(attempts, list) and attempts and isinstance(attempts[-1], dict): last = attempts[-1].get("pk") calc_pk = last if isinstance(last, int) and not isinstance(last, bool) else None final = document.get("final_result") cached_gen, cached_pk = _cached_geometry(document) requests.append( { "id": material_id, "wc_pk": wc_pk, "uuid": document.get("uuid"), "calc_pk": calc_pk, "_result": final if isinstance(final, dict) else {}, "_n_attempts": document.get("n_attempts"), "_gen": cached_gen, "_gen_calc_pk": cached_pk, } ) return requests, skipped def _cached_geometry(document: dict[str, Any]) -> tuple[str | None, int | None]: """Read the relaxed geometry a calculate run cached into its result. Parameters ---------- document A parsed ``calc_result.json``. Returns ------- tuple ``(gen_text, source_calc_pk)``, either of which is ``None`` when the document carries no usable cache. Results written before this cache existed simply have none, which is why the AiiDA fetch path stays. """ cached = document.get("final_geometry") if not isinstance(cached, dict): return None, None gen = cached.get("gen") source_pk = cached.get("source_calc_pk") return ( gen if isinstance(gen, str) and gen.strip() else None, source_pk if isinstance(source_pk, int) and not isinstance(source_pk, bool) else None, ) # --------------------------------------------------------------------------- # Fetch (one batched subprocess) + write # --------------------------------------------------------------------------- def _cached_row(request: dict[str, Any]) -> dict[str, Any]: """Synthesize a fetch result from a geometry already cached on disk. Parameters ---------- request The request whose cached geometry is used. Returns ------- dict Deliberately the same shape the runner returns, so everything downstream -- provenance, failure handling, the written CIF -- is identical whether the geometry came from disk or from AiiDA. It *is* the same file either way. """ return { "id": request["id"], "wc_pk": request["wc_pk"], "calc_pk": request["_gen_calc_pk"] if request["_gen_calc_pk"] is not None else request["calc_pk"], "gen": request["_gen"], "error": None, "note": None, } def _fetch_plan(cached: int, to_fetch: int) -> str: """Phrase one line saying where the geometries are coming from. Parameters ---------- cached How many will be served from disk. to_fetch How many need an AiiDA lookup. Returns ------- str The message, which says plainly when no AiiDA lookup is needed at all. """ if not to_fetch: plural = "y" if cached == 1 else "ies" return ( f"Using {cached} cached geometr{plural} from {RESULT_NAME} -- no AiiDA lookup needed." ) plural = "y" if to_fetch == 1 else "ies" cached_note = f"{cached} cached, " if cached else "" return f"{cached_note}fetching {to_fetch} relaxed geometr{plural} from AiiDA..." def _fetch_geometries( python: str, profile: Any, requests: list[dict[str, Any]] ) -> list[dict[str, Any]] | None: """Run the fetch runner once for the whole batch. Parameters ---------- python The calculate stage's interpreter -- the one whose AiiDA install holds the data. profile The AiiDA profile to load, or ``None`` for the default. requests The materials to fetch. Returns ------- list or None One result row per request, or ``None`` when the runner itself failed. A *per-material* failure is a row carrying an ``error``, not a ``None`` return: one unreadable geometry must not cost the whole export. Notes ----- The runner's stderr is deliberately **not** captured. Its progress lines and its own actionable messages -- a missing aiida names the pip command -- reach the user live, which matters when thousands of geometries are being read. """ entries = [ {key: request[key] for key in ("id", "wc_pk", "uuid", "calc_pk")} for request in requests ] with tempfile.TemporaryDirectory(prefix="matmelt-cif-") as scratch: batch_path = Path(scratch) / "batch.json" out_path = Path(scratch) / "geometries.json" batch_path.write_text(json.dumps(entries), encoding="utf-8") command = [python, str(RUNNER_PATH), "--batch", str(batch_path), "--out", str(out_path)] if profile: command += ["--profile", str(profile)] try: proc = subprocess.run(command, stdout=subprocess.PIPE, text=True, check=False) except FileNotFoundError: render_error( f"cif: cannot run the calculate stage's python {python!r} -- fix " f"stages.calculate.options.python in your config." ) return None except OSError as exc: render_error(f"cif: cannot run {python!r}: {exc}") return None if proc.returncode != 0: render_error( f"cif: fetching geometries failed (exit {proc.returncode}); see the " f"message above. The runner ran as: {' '.join(command)}" ) return None try: document = json.loads(out_path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError) as exc: render_error(f"cif: the geometry fetch produced no readable result: {exc}") return None results = document.get("results") if isinstance(document, dict) else None if not isinstance(results, list): render_error("cif: the geometry fetch produced no 'results' list.") return None return [row for row in results if isinstance(row, dict)] def _write_cif_files( fetched: list[dict[str, Any]], requests: list[dict[str, Any]], target: Path, ) -> tuple[list[dict[str, Any]], list[dict[str, str]], list[dict[str, str]]]: """Turn fetched geometries into CIF files, reporting per material. Parameters ---------- fetched The runner's result rows. requests The originating requests, for provenance. target The export directory. Returns ------- tuple ``(written, failed, notes)``. A geometry that will not parse, or a file that will not write, fails only its own material -- the rest of the export still lands. Notes ----- Notes, such as a recorded CalcJob pk that disagrees with the database, are **returned rather than printed**, so ``--json`` output stays pure JSON. That was a real defect: progress and warnings once polluted the payload. """ by_id = {request["id"]: request for request in requests} written: list[dict[str, Any]] = [] failed: list[dict[str, str]] = [] notes: list[dict[str, str]] = [] for row in fetched: material_id = str(row.get("id", "")) request = by_id.get(material_id, {}) if row.get("error"): failed.append({"id": material_id, "error": str(row["error"])}) continue gen = row.get("gen") if not isinstance(gen, str) or not gen.strip(): failed.append({"id": material_id, "error": "the fetch returned no geometry"}) continue if row.get("note"): notes.append({"id": material_id, "note": str(row["note"])}) try: text = gen_to_cif(gen, data_name=material_id, extra=_provenance(row, request)) except ValueError as exc: failed.append({"id": material_id, "error": f"cannot build a CIF: {exc}"}) continue path = target / f"{material_id}.cif" try: tmp = path.with_name(path.name + ".tmp") tmp.write_text(text, encoding="utf-8") tmp.replace(path) except OSError as exc: failed.append({"id": material_id, "error": f"cannot write {path}: {exc}"}) continue written.append( { "id": material_id, "path": str(path), "wc_pk": row.get("wc_pk"), "calc_pk": row.get("calc_pk"), } ) return written, failed, notes def _provenance(row: dict[str, Any], request: dict[str, Any]) -> dict[str, Any]: """Build the provenance stamped into one exported CIF. Parameters ---------- row The fetch result row. request Its originating request, carrying the parsed result document. Returns ------- dict The items written as ``_mcli_*`` CIF tags: the material id, where the geometry came from, the WorkChain and CalcJob pks, the attempt count, and the recorded energetics -- so an exported file can be traced back to the calculation that produced it without the workspace. """ result: dict[str, Any] = request.get("_result") or {} provenance: dict[str, Any] = { "material_id": row.get("id"), "geometry_source": "geom.out.gen (DFTB+ relaxed)", "single_wc_pk": row.get("wc_pk"), "calcjob_pk": row.get("calc_pk"), "n_attempts": request.get("_n_attempts"), } for field in _RESULT_FIELDS: provenance[field] = result.get(field) return provenance