diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index fe1bae4..702d630 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -28,6 +28,10 @@ jobs: run: | ruff check tests/test_api*.py tests/test_foundations.py tests/test_operation_cli.py tests/test_mcp_api.py tests/test_packaged_resources.py tests/test_unified_smoke.py tests/fixtures/*.py examples/unified_smoke.py ruff format --check tests/test_api*.py tests/test_foundations.py tests/test_operation_cli.py tests/test_mcp_api.py tests/test_packaged_resources.py tests/test_unified_smoke.py tests/fixtures/*.py examples/unified_smoke.py + - name: Check calculation examples + run: | + ruff check tests/test_calculation_examples.py examples/zeopp.py examples/graspa.py examples/graspa_mixture_setup.py examples/mlip.py + ruff format --check tests/test_calculation_examples.py examples/zeopp.py examples/graspa.py examples/graspa_mixture_setup.py examples/mlip.py test: runs-on: ubuntu-latest diff --git a/README.md b/README.md index 697d758..109bc76 100644 --- a/README.md +++ b/README.md @@ -95,6 +95,14 @@ matkit mlip run-batch --input-dir cifs --backend nvalchemi-mace \ --checkpoint medium --device cuda --batch-size 16 ``` +### Calculation examples + +The [calculation examples guide](examples/README.md) covers Zeo++ pore analysis +and screening, pure-component gRASPA execution and pressure sweeps, mixture +preparation, and MLIP evaluation/relaxation/batches across all three current +backends. It includes runnable Python programs, equivalent CLI specifications, +and result inspection using public MatKit interfaces. + ### GPU examples [`examples/mlip_gpu.py`](examples/mlip_gpu.py) runs one backend per Python diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 0000000..44d08f4 --- /dev/null +++ b/examples/README.md @@ -0,0 +1,431 @@ +# Running calculations with MatKit + +These examples use public Python APIs and the existing MatKit CLI. MatKit +stages inputs, launches engines, parses outputs, and saves calculation bundles. +The Python programs contain no custom engine launchers, output parsers, +checkpoint writers, or charge-generation helpers. + +| Program | Scenarios | Main API | +| --- | --- | --- | +| [zeopp.py](zeopp.py) | Pore analyses; multiple structures | `analyze_pores`, `run_batch` | +| [graspa.py](graspa.py) | Pure-component preparation, execution, pressure sweep | `prepare_adsorption`, `run_adsorption` | +| [graspa_mixture_setup.py](graspa_mixture_setup.py) | CO2/N2 mixture preparation | `matkit.graspa.setup_simulation` | +| [mlip.py](mlip.py) | Energy/forces, fixed-cell relaxation, batches; three backends | `evaluate`, `relax`, `run_batch` | + +## Install and supply inputs + +Run commands from the repository root after installing MatKit in your selected +environment: + +```bash +python -m pip install -e . # Core, Zeo++, gRASPA preparation/CLI +python -m pip install -e '.[mlip]' # Direct MACE, when needed +python -m pip install -e '.[rootstock]' # Rootstock client, when needed +matkit capabilities --json +``` + +Zeo++'s `network` and CUDA gRASPA's `simulate` executables are separate +installations. For ALCHEMI, use the repository's +[compatible CUDA environment recipe](../alcf/polaris/mlip/README.md). +Capability discovery reports caller-side installation, not GPU or model +validation. Each worker interpreter must have MatKit and its selected backend +installed; Rootstock also needs an accessible deployment. + +Replace the example filenames with your own structures. Each file must contain +one structure. Unified CIF loading requires an unambiguous atom mapping; use +explicit P1 structures when symmetry expansion would change the site count. +Zeo++ requires a fully periodic cell. gRASPA additionally requires finite +`_atom_site_charge` values whose sum matches `--net-charge` (default zero). +Use charges and force-field definitions appropriate for your framework. +The repository's small uncharged test CIF is not a scientific reference or a +gRASPA input. + +For unified ALCHEMI, supply geometry-only extended XYZ with the correct cell +and PBC. The current implementation rejects atom arrays, constraints, and +bonds; even an ordinary CIF can carry ASE's `spacegroup_kinds` array and be +rejected. +The examples do not silently remove scientific metadata to bypass validation. + +Use a fresh output directory for every invocation. Subprocess execution keeps +engine output in bundle logs. The programs exit 0 for accepted calculations +(or successful preparation), 1 for execution/input-file failures or +unconverged relaxation, and 2 for argument parsing or request/profile +validation errors. +Use `--help` on any program to see its options. + +## Zeo++: pore geometry and screening + +With `network` on PATH, run all five analyses: + +```bash +python examples/zeopp.py framework.cif --outdir runs/pores +``` + +Select analyses and settings, or process several structures with the same +settings: + +```bash +python examples/zeopp.py framework.cif --outdir runs/area-volume \ + --analysis sa --analysis vol --probe-radius 1.2 --channel-radius 1.8 \ + --num-samples 100000 --radii /path/to/custom.rad + +python examples/zeopp.py framework_a.cif framework_b.cif \ + --analysis res --analysis sa --outdir runs/pore-screen +``` + +The default radii are MatKit's bundled UFF radii; high accuracy is enabled. +`res` reports Di/Df/Dif in angstrom, `sa` surface area, `vol` accessible volume, +`psd` a pore-size histogram, and `chan` channel dimensionalities. For sampled +analyses, `probe_radius` must not exceed `channel_radius`. Sample count and +probe choice affect the calculation and should be chosen for your study. +Unified Zeo++ batches run sequentially and retain one bundle per structure. + +For an executable outside PATH, save `zeopp-execution.json`: + +```json +{ + "executables": {"zeopp": ["/absolute/path/to/network"]}, + "timeout_s": 3600 +} +``` + +Pass `--execution zeopp-execution.json` to the program. A direct API call is: + +```python +from matkit.api import ExecutionConfig, PoreRequest, StructureRef, analyze_pores + +result = analyze_pores( + PoreRequest( + structure=StructureRef(path="framework.cif"), + analyses=["res", "sa"], + num_samples=100000, + ), + output_dir="runs/pores-api", + execution=ExecutionConfig(mode="subprocess"), +) +if not result.accepted: + raise RuntimeError(result.failure) +print(result.payload.results["res"]["Di"]) +print(result.payload.results["sa"]["ASA_m2_g"]) +``` + +For the CLI, save `pores.json` beside `framework.cif`: + +```json +{ + "operation": "pores", + "structure": {"path": "framework.cif"}, + "analyses": ["res", "sa", "vol", "psd", "chan"], + "probe_radius": 1.86, + "channel_radius": 1.86, + "num_samples": 100000 +} +``` + +```bash +matkit pores --spec pores.json --outdir runs/pores-cli \ + --execution zeopp-execution.json +matkit inspect runs/pores-cli + +# Alternatively, stage without an engine, then execute the same bundle. +matkit prepare --spec pores.json --outdir runs/pores-prepared +matkit execute runs/pores-prepared --execution zeopp-execution.json +``` + +CLI request paths are relative to their specification file. Python program +input paths are relative to the caller's working directory. Executable paths +in execution profiles should be absolute; commands are argument lists. + +## gRASPA: pure-component adsorption and pressure sweeps + +Preparation needs no CUDA or gRASPA executable: + +```bash +python examples/graspa.py charged_framework.cif \ + --adsorbate CO2 --pressure-pa 100000 --prepare-only \ + --outdir runs/co2-prepared +``` + +This creates `runs/co2-prepared/00000`. Every pressure point gets an indexed +directory, including a single point. The generated input uses 298 K, a 12.8 +angstrom cutoff, 1000 initialization cycles, 1000 equilibration cycles, +10000 production cycles, 5 blocks, and PR-EOS. These demonstration settings +do not establish adequate sampling. The CLI options expose temperature, +cycles, blocks, cutoff, net charge, uptake unit, and a custom template directory. + +Save `graspa-execution.json` with your executable: + +```json +{ + "executables": {"graspa": ["/absolute/path/to/gRASPA/bin/simulate"]}, + "timeout_s": 3600 +} +``` + +Execute the prepared point, or run a fresh pressure sweep: + +```bash +matkit execute runs/co2-prepared/00000 --execution graspa-execution.json +matkit adsorption analyze runs/co2-prepared/00000 + +python examples/graspa.py charged_framework.cif --adsorbate CO2 \ + --temperature-k 298 --pressure-pa 1000 --pressure-pa 10000 \ + --pressure-pa 100000 --outdir runs/co2-isotherm \ + --execution graspa-execution.json +``` + +Pressures above correspond to 0.01, 0.1, and 1 bar; **1 bar = 100000 Pa**. +Directories `00000`, `00001`, and `00002` preserve that input order. The +program prints one JSON run record per line, continues after a failed point, +and exits nonzero if any point fails. Each point uses `run_adsorption` because +`run_batch` requires identical pressures and other scientific settings. +For multiple frameworks at one pressure, a list of otherwise identical +`AdsorptionRequest` objects can use `run_batch`. + +The same single-point operation through the CLI uses `adsorption.json` beside +your charged CIF: + +```json +{ + "operation": "adsorption", + "structure": {"path": "charged_framework.cif"}, + "adsorbate": "CO2", + "temperature_K": 298, + "pressure_Pa": 100000, + "cutoff_angstrom": 12.8, + "initialization_cycles": 1000, + "equilibration_cycles": 1000, + "production_cycles": 10000, + "number_of_blocks": 5, + "fugacity_coefficient": "PR-EOS", + "net_charge": 0, + "unit": "mol/kg" +} +``` + +```bash +matkit adsorption run --spec adsorption.json --outdir runs/co2-cli \ + --execution graspa-execution.json + +# Or prepare now and execute later in an allocation with CUDA. +matkit adsorption prepare --spec adsorption.json --outdir runs/co2-staged +matkit execute runs/co2-staged --execution graspa-execution.json +``` + +Keep prepared bundles intact when copying them to another environment. +MatKit verifies staged inputs before execution; change the request or template +and prepare a new bundle when settings need changing. + +Inspect uptake, engine-reported uncertainty, and heat of adsorption: + +```python +from matkit.api import inspect_run + +result = inspect_run("runs/co2-isotherm/00000") +if not result.accepted: + raise RuntimeError(result.failure) +print(result.payload.uptake, result.payload.uncertainty, result.payload.unit) +print(result.payload.heat_of_adsorption, result.payload.heat_unit) +``` + +Uptake is absolute, per framework mass/volume; the default is mol/kg. +Heat is reported in kJ/mol with the engine's sign convention. Accepted output +does not establish equilibration or independent samples. + +### CO2/N2 mixture preparation + +```bash +python examples/graspa_mixture_setup.py charged_framework.cif \ + --outdir runs/co2-n2-inputs +``` + +The program directly calls `setup_simulation` with CO2/N2 mole fractions +0.15/0.85, 298 K, total pressure 100000 Pa, and the default template's +component placeholder. MatKit calculates unit-cell replication and copies +the input definitions. This legacy template uses 10000 initialization and +production cycles, zero equilibration cycles, and one block in this example. +Edit the explicit API settings in the program for another mixture. + +This directory contains engine inputs, not a unified run bundle. Unified +mixture execution and result parsing are not implemented. The legacy +setup API also does not perform the unified charged-CIF validation. Check +charge/force-field suitability before executing through your engine workflow; +do not pass this directory to `matkit execute` or the single-component parser. + +## MLIP: evaluation, relaxation, and batches + +The program uses `MLIPMethod` for the checkpoint and one of `MACEAdapter`, +`RootstockAdapter`, or `AlchemiAdapter` for the implementation. Energy mode +requests energy and forces by default; `--property potential_energy` requests +only energy, and repeated `--property` options can include stress when the +model supports it. Energy is in eV, forces in eV/angstrom, and stress in +eV/angstrom^3 using the ASE convention. + +Direct MACE defaults to CPU/float64: + +```bash +python examples/mlip.py framework.cif --backend ase-mace --checkpoint medium \ + --outdir runs/mace-energy + +python examples/mlip.py framework.cif --backend ase-mace --checkpoint medium \ + --driver relax --fmax 0.02 --steps 500 --outdir runs/mace-relax + +python examples/mlip.py framework_a.cif framework_b.cif \ + --backend ase-mace --checkpoint medium --outdir runs/mace-batch +``` + +Relaxation uses fixed-cell FIRE. The default tolerance is 0.01 eV/angstrom +and maximum steps is 1000. An unconverged calculation retains its numerical +results but exits 1. + +For CUDA, save `cuda-execution.json`: + +```json +{"device": "cuda"} +``` + +Run one backend per process in its installed environment: + +```bash +python examples/mlip.py framework.cif --backend ase-mace --checkpoint medium \ + --dtype float32 --execution cuda-execution.json --outdir runs/mace-cuda + +python examples/mlip.py framework.cif --backend rootstock \ + --cluster polaris --checkpoint mace-mp-0-medium \ + --execution cuda-execution.json --outdir runs/rootstock-energy + +python examples/mlip.py geometry_a.extxyz geometry_b.extxyz \ + --backend nvalchemi-mace --checkpoint medium --batch-size 16 \ + --execution cuda-execution.json --outdir runs/alchemi-energy + +python examples/mlip.py geometry_a.extxyz geometry_b.extxyz \ + --backend nvalchemi-mace --checkpoint medium --batch-size 16 \ + --driver relax --fmax 0.02 --steps 500 \ + --execution cuda-execution.json --outdir runs/alchemi-relax +``` + +Replace `polaris` with your deployment, or use `--root /path/to/deployment`. +Add `--driver relax` for Rootstock relaxation; multiple inputs form a batch +on any backend. Rootstock precision belongs to its deployment settings, so +the program rejects `--dtype` for Rootstock. Its worker owns CUDA; the caller +does not need a CUDA-enabled PyTorch. ALCHEMI defaults to CUDA/float32 and +supports native chunking; direct MACE and Rootstock batches reuse one +calculator while executing items sequentially. Different model aliases do +not imply identical weights or comparable scientific results. + +A direct Python relaxation has no example-specific helpers: + +```python +from matkit.api import ( + ExecutionConfig, MLIPMethod, RelaxRequest, StructureRef, relax, +) + +result = relax( + RelaxRequest( + structure=StructureRef(path="framework.cif"), + method=MLIPMethod(checkpoint="medium"), + fmax=0.02, + steps=500, + ), + output_dir="runs/relax-api", + execution=ExecutionConfig(mode="subprocess", device="cpu"), +) +print(result.accepted) +if result.payload is not None: + print(result.payload.potential_energy, result.payload.converged) + print(result.payload.final_structure) # Relative to runs/relax-api. +``` + +For the unified CLI, save `energy.json` beside the input: + +```json +{ + "operation": "evaluate", + "structure": {"path": "framework.cif"}, + "method": {"checkpoint": "medium"}, + "adapter": {"type": "ase-mace", "dtype": "float32"}, + "properties": ["potential_energy", "forces"] +} +``` + +Save a relaxation specification as `relaxation.json`: + +```json +{ + "operation": "relax", + "structure": {"path": "framework.cif"}, + "method": {"checkpoint": "medium"}, + "adapter": {"type": "ase-mace", "dtype": "float32"}, + "fmax": 0.02, + "steps": 500 +} +``` + +```bash +matkit evaluate --spec energy.json --outdir runs/energy-cli \ + --execution cuda-execution.json +matkit relax --spec relaxation.json --outdir runs/relax-cli \ + --execution cuda-execution.json +``` + +For Rootstock, change `adapter` to +`{"type": "rootstock", "cluster": "polaris"}` and use its checkpoint ID. +For ALCHEMI, use +`{"type": "nvalchemi-mace", "dtype": "float32", "batch_size": 16}` +and a geometry-only extended XYZ input. + +For CLI batches, save a JSON list of complete requests as `batch.json`. +Only the structure fields may differ within a batch. For example: + +```json +[ + { + "operation": "evaluate", + "structure": {"path": "framework_a.cif"}, + "method": {"checkpoint": "medium"}, + "properties": ["potential_energy", "forces"] + }, + { + "operation": "evaluate", + "structure": {"path": "framework_b.cif"}, + "method": {"checkpoint": "medium"}, + "properties": ["potential_energy", "forces"] + } +] +``` + +```bash +matkit batch --spec batch.json --outdir runs/batch-cli +matkit inspect runs/batch-cli +matkit inspect runs/batch-cli/00000 +``` + +## Read results and understand validation + +Each unified run retains `request.json`, original/supporting inputs under +`inputs/`, engine files under `work/`, and a `run.json` record. Completed or +failed executions normally also have `result.json`. Worker stdout/stderr logs +are at the bundle root. Zeo++ logs are `work/engine.stdout.log` and +`work/engine.stderr.log`; gRASPA stdout is `work/raspa.log`. + +MLIP final geometry is `work/final_structure.extxyz`; keep its adjacent +`.metadata.json` sidecar for atom correspondence and supported metadata. +Charges derived from the original geometry are invalidated after relaxation; +the relaxed output is not automatically a charged gRASPA input. + +Batch roots contain `batch_manifest.json` and numbered item bundles. A batch +uses one supervised worker, so its worker logs live at the batch root; +external-engine logs remain in each item's `work/` directory. A batch +may retain successful items alongside failures or unconverged optimizations. +Use `inspect_run` on an item's directory to read its scientific payload; +`matkit inspect` also reads batch manifests. Inspection exit 0 means the +record was readable, even if the calculation failed. `RunResult.accepted` +and `BatchResult.accepted` are computed Python properties; use the run +command's exit code and stored state/checks when consuming CLI JSON. + +The examples are checked with synthetic engines and calculator doubles. +No real Zeo++, gRASPA, MACE, Rootstock, or ALCHEMI execution was available +during local validation. For opt-in execution evidence use +[unified_smoke.py](unified_smoke.py) or the +[Polaris MLIP smoke recipe](../alcf/polaris/mlip/README.md). +See the [capability inventory](../docs/capabilities.md) for validation limits. diff --git a/examples/graspa.py b/examples/graspa.py new file mode 100644 index 0000000..9edb394 --- /dev/null +++ b/examples/graspa.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 +"""Prepare or run pure-component gRASPA at one or more pressures.""" + +import argparse +from pathlib import Path +import sys + +from matkit.api import ( + AdsorptionRequest, + ExecutionConfig, + StructureRef, + prepare_adsorption, + run_adsorption, +) + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("cif", help="Atom-mapped, charged, periodic CIF.") + parser.add_argument("--outdir", required=True, type=Path) + parser.add_argument( + "--execution", type=Path, help="Execution profile JSON." + ) + parser.add_argument("--prepare-only", action="store_true") + parser.add_argument("--adsorbate", default="CO2") + parser.add_argument("--temperature-k", type=float, default=298.0) + parser.add_argument( + "--pressure-pa", action="append", type=float, help="Repeat for a sweep." + ) + parser.add_argument("--cutoff", type=float, default=12.8) + parser.add_argument("--initialization-cycles", type=int, default=1000) + parser.add_argument("--equilibration-cycles", type=int, default=1000) + parser.add_argument("--production-cycles", type=int, default=10000) + parser.add_argument("--blocks", type=int, default=5) + parser.add_argument("--net-charge", type=float, default=0.0) + parser.add_argument( + "--unit", choices=("mol/kg", "mg/g", "g/L"), default="mol/kg" + ) + parser.add_argument( + "--template-dir", help="Complete custom template directory." + ) + args = parser.parse_args(argv) + if args.prepare_only and args.execution: + parser.error("Preparation does not use an execution profile") + + try: + execution = ( + ExecutionConfig.model_validate_json(args.execution.read_text()) + if args.execution + else ExecutionConfig() + ).model_copy(update={"mode": "subprocess"}) + requests = [ + AdsorptionRequest( + structure=StructureRef(path=args.cif), + adsorbate=args.adsorbate, + temperature_K=args.temperature_k, + pressure_Pa=pressure, + cutoff_angstrom=args.cutoff, + initialization_cycles=args.initialization_cycles, + equilibration_cycles=args.equilibration_cycles, + production_cycles=args.production_cycles, + number_of_blocks=args.blocks, + fugacity_coefficient="PR-EOS", + net_charge=args.net_charge, + unit=args.unit, + template_dir=args.template_dir, + ) + for pressure in args.pressure_pa or [100000.0] + ] + except (OSError, ValueError) as exc: + parser.error(str(exc)) + + try: + args.outdir.mkdir(parents=True, exist_ok=False) + except OSError as exc: + print(f"Use a fresh output directory: {exc}", file=sys.stderr) + return 1 + + succeeded = True + # Different pressures cannot share a homogeneous MatKit batch. + # Each public API call owns preparation, execution, logs, and results. + for index, request in enumerate(requests): + bundle = args.outdir / f"{index:05d}" + try: + if args.prepare_only: + result = prepare_adsorption(request, output_dir=bundle) + else: + result = run_adsorption( + request, output_dir=bundle, execution=execution + ) + print(result.model_dump_json()) # One run record per stdout line. + succeeded = succeeded and ( + result.state == "prepared" + if args.prepare_only + else result.accepted + ) + except Exception as exc: + print(f"{bundle}: {exc}", file=sys.stderr) + succeeded = False + return 0 if succeeded else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/graspa_mixture_setup.py b/examples/graspa_mixture_setup.py new file mode 100644 index 0000000..5481e4c --- /dev/null +++ b/examples/graspa_mixture_setup.py @@ -0,0 +1,40 @@ +#!/usr/bin/env python3 +"""Prepare a CO2/N2 mixture with explicit fractions; no engine execution.""" + +import argparse +from pathlib import Path +import sys + +from matkit.graspa import setup_simulation + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("cif", help="Charged periodic framework CIF.") + parser.add_argument("--outdir", required=True, type=Path) + args = parser.parse_args(argv) + + try: + args.outdir.mkdir(parents=True, exist_ok=False) + setup_simulation( + cif=args.cif, + outpath=str(args.outdir), + adsorbates=[ + {"MoleculeName": "CO2", "MolFraction": 0.15}, + {"MoleculeName": "N2", "MolFraction": 0.85}, + ], + temperature=298.0, + pressure=100000.0, + cutoff=12.8, + n_cycle=10000, + template_dir="template", + ) + except Exception as exc: + print(str(exc), file=sys.stderr) + return 1 + print(f"Prepared mixture input: {args.outdir / 'simulation.input'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/mlip.py b/examples/mlip.py new file mode 100644 index 0000000..524ab1e --- /dev/null +++ b/examples/mlip.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +"""Evaluate or relax structures with MACE, Rootstock, or ALCHEMI.""" + +import argparse +from pathlib import Path +import sys + +from matkit.api import ( + AlchemiAdapter, + EvaluateRequest, + ExecutionConfig, + MACEAdapter, + MLIPMethod, + RelaxRequest, + RootstockAdapter, + StructureRef, + evaluate, + relax, + run_batch, +) + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "inputs", nargs="+", help="One structure per input file." + ) + parser.add_argument("--outdir", required=True, type=Path) + parser.add_argument( + "--execution", type=Path, help="Execution profile JSON." + ) + parser.add_argument( + "--backend", + required=True, + choices=("ase-mace", "rootstock", "nvalchemi-mace"), + ) + parser.add_argument("--checkpoint", required=True) + parser.add_argument( + "--driver", choices=("energy", "relax"), default="energy" + ) + parser.add_argument("--dtype", choices=("float32", "float64")) + location = parser.add_mutually_exclusive_group() + location.add_argument("--cluster", help="Rootstock deployment cluster.") + location.add_argument("--root", help="Rootstock deployment directory.") + parser.add_argument( + "--batch-size", type=int, help="ALCHEMI native chunk size." + ) + parser.add_argument( + "--fmax", type=float, help="Relaxation force tolerance (eV/A)." + ) + parser.add_argument("--steps", type=int, help="Maximum relaxation steps.") + parser.add_argument( + "--property", + action="append", + choices=("potential_energy", "forces", "stress"), + ) + args = parser.parse_args(argv) + if args.backend != "rootstock" and (args.cluster or args.root): + parser.error("--cluster and --root require --backend rootstock") + if args.backend == "rootstock" and args.dtype: + parser.error("Rootstock precision is configured in its deployment") + if args.backend != "nvalchemi-mace" and args.batch_size is not None: + parser.error("--batch-size requires --backend nvalchemi-mace") + if args.driver == "energy" and ( + args.fmax is not None or args.steps is not None + ): + parser.error("--fmax and --steps require --driver relax") + if args.driver == "relax" and args.property: + parser.error("Relaxation requests energy and forces automatically") + + try: + execution = ( + ExecutionConfig.model_validate_json(args.execution.read_text()) + if args.execution + else ExecutionConfig() + ).model_copy(update={"mode": "subprocess"}) + if args.backend == "ase-mace": + adapter = MACEAdapter(dtype=args.dtype) + elif args.backend == "rootstock": + adapter = RootstockAdapter(cluster=args.cluster, root=args.root) + else: + adapter = AlchemiAdapter( + dtype=args.dtype or "float32", + batch_size=args.batch_size + if args.batch_size is not None + else 16, + ) + method = MLIPMethod(checkpoint=args.checkpoint) + if args.driver == "relax": + request_type, operation = RelaxRequest, relax + settings = { + "fmax": args.fmax if args.fmax is not None else 0.01, + "steps": args.steps if args.steps is not None else 1000, + } + else: + request_type, operation = EvaluateRequest, evaluate + settings = { + "properties": args.property or ["potential_energy", "forces"] + } + requests = [ + request_type( + structure=StructureRef(path=path), + method=method, + adapter=adapter, + **settings, + ) + for path in args.inputs + ] + except (OSError, ValueError) as exc: + parser.error(str(exc)) + + try: + if len(requests) == 1: + result = operation( + requests[0], output_dir=args.outdir, execution=execution + ) + else: + result = run_batch( + requests, output_dir=args.outdir, execution=execution + ) + except Exception as exc: + print(str(exc), file=sys.stderr) + return 1 + print(result.model_dump_json(indent=2)) + return 0 if result.accepted else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/zeopp.py b/examples/zeopp.py new file mode 100644 index 0000000..e3354fb --- /dev/null +++ b/examples/zeopp.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python3 +"""Run Zeo++ pore analyses for one structure or a homogeneous batch.""" + +import argparse +from pathlib import Path +import sys + +from matkit.api import ( + ExecutionConfig, + PoreRequest, + StructureRef, + analyze_pores, + run_batch, +) + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("inputs", nargs="+", help="Periodic structure files.") + parser.add_argument("--outdir", required=True, type=Path) + parser.add_argument( + "--execution", type=Path, help="Execution profile JSON." + ) + parser.add_argument( + "--analysis", + action="append", + choices=("res", "sa", "vol", "psd", "chan"), + ) + parser.add_argument("--probe-radius", type=float, default=1.86) + parser.add_argument("--channel-radius", type=float, default=1.86) + parser.add_argument("--num-samples", type=int, default=100000) + parser.add_argument("--radii", help="Defaults to MatKit's bundled UFF.rad.") + args = parser.parse_args(argv) + + try: + execution = ( + ExecutionConfig.model_validate_json(args.execution.read_text()) + if args.execution + else ExecutionConfig() + ).model_copy(update={"mode": "subprocess"}) + requests = [ + PoreRequest( + structure=StructureRef(path=path), + analyses=args.analysis or ["res", "sa", "vol", "psd", "chan"], + probe_radius=args.probe_radius, + channel_radius=args.channel_radius, + num_samples=args.num_samples, + radii_file=args.radii, + ) + for path in args.inputs + ] + except (OSError, ValueError) as exc: + parser.error(str(exc)) + + try: + if len(requests) == 1: + result = analyze_pores( + requests[0], output_dir=args.outdir, execution=execution + ) + else: + result = run_batch( + requests, output_dir=args.outdir, execution=execution + ) + except Exception as exc: + print(str(exc), file=sys.stderr) + return 1 + print(result.model_dump_json(indent=2)) + return 0 if result.accepted else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_calculation_examples.py b/tests/test_calculation_examples.py new file mode 100644 index 0000000..6c909dd --- /dev/null +++ b/tests/test_calculation_examples.py @@ -0,0 +1,329 @@ +"""Run the usage examples through public APIs and supervised test engines.""" + +import importlib.util +import json +import os +from pathlib import Path +import sys + +from ase import Atoms +from ase.io import write +import pytest + +from matkit.api import inspect_run + + +@pytest.fixture +def example(): + def load(name): + path = Path(__file__).parents[1] / "examples" / f"{name}.py" + spec = importlib.util.spec_from_file_location(f"example_{name}", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + return load + + +@pytest.fixture +def profile(tmp_path): + def create(engine, *arguments): + path = tmp_path / f"{engine}-execution.json" + fixture = Path(__file__).parent / "fixtures" / "fake_engine.py" + path.write_text( + json.dumps( + { + "executables": { + engine: [ + sys.executable, + str(fixture), + engine, + *arguments, + ] + } + } + ) + ) + return str(path) + + return create + + +@pytest.fixture +def charged_cif(sample_cif, tmp_path): + # Synthetic charges belong only to test fixtures, not the usage programs. + path = tmp_path / "charged.cif" + lines = [] + for line in Path(sample_cif).read_text().splitlines(): + if line.strip() == "_atom_site_occupancy": + lines.append(" _atom_site_charge") + fields = line.split() + if fields and fields[0] in {"Si", "O"}: + fields.insert(-1, "0.0") + line = " ".join(fields) + lines.append(line) + path.write_text("\n".join(lines) + "\n") + return str(path) + + +@pytest.mark.parametrize( + "name", ["zeopp", "graspa", "graspa_mixture_setup", "mlip"] +) +def test_help_requires_no_optional_engines(example, name, capsys): + with pytest.raises(SystemExit) as exc: + example(name).main(["--help"]) + assert exc.value.code == 0 + assert "--outdir" in capsys.readouterr().out + + +@pytest.mark.parametrize("count", [1, 2]) +def test_zeopp_all_analyses_and_batches( + example, profile, sample_cif, tmp_path, capsys, count +): + root = tmp_path / "pores" + code = example("zeopp").main( + [sample_cif] * count + + ["--outdir", str(root), "--execution", profile("zeopp")] + ) + assert code == 0 + stdout = json.loads(capsys.readouterr().out) + if count == 2: + assert [item["index"] for item in stdout["items"]] == [0, 1] + assert all(item["accepted"] for item in stdout["items"]) + assert (root / "worker.stdout.log").is_file() + for index in range(count): + bundle = root if count == 1 else root / f"{index:05d}" + result = inspect_run(bundle) + assert result.accepted + assert set(result.payload.results) == { + "res", + "sa", + "vol", + "psd", + "chan", + } + assert (bundle / "work" / "structure.psd_histo").is_file() + + +@pytest.mark.parametrize("failure", ["partial", "missing_executable"]) +def test_zeopp_failures_are_persisted_and_nonzero( + example, profile, sample_cif, tmp_path, failure +): + root = tmp_path / "pores" + path = Path(profile("zeopp", "--partial")) + if failure == "missing_executable": + path.write_text( + json.dumps({"executables": {"zeopp": [str(tmp_path / "missing")]}}) + ) + assert ( + example("zeopp").main( + [sample_cif, "--outdir", str(root), "--execution", str(path)] + ) + == 1 + ) + result = inspect_run(root) + assert result.state == "failed" + assert not result.accepted + assert result.failure is not None + + +def test_graspa_prepares_distinct_pressure_bundles( + example, charged_cif, tmp_path, capsys +): + root = tmp_path / "prepared" + arguments = [ + charged_cif, + "--outdir", + str(root), + "--prepare-only", + "--pressure-pa", + "100000.1", + "--pressure-pa", + "100000.2", + ] + assert example("graspa").main(arguments) == 0 + records = [ + json.loads(line) for line in capsys.readouterr().out.splitlines() + ] + assert [record["state"] for record in records] == ["prepared", "prepared"] + for index, pressure in enumerate([100000.1, 100000.2]): + bundle = root / f"{index:05d}" + record = inspect_run(bundle) + assert record.requested["pressure_Pa"] == pressure + assert ( + f"Pressure {pressure}" + in (bundle / "work/simulation.input").read_text() + ) + assert (bundle / "inputs/structure.cif").read_bytes() == Path( + charged_cif + ).read_bytes() + assert not (bundle / "worker.stdout.log").exists() + # A rerun must not silently replace the prepared cases. + assert example("graspa").main(arguments) == 1 + assert inspect_run(root / "00000").run_id == records[0]["run_id"] + + +@pytest.mark.parametrize("fails", [False, True]) +def test_graspa_executes_every_pressure_and_reports_outcomes( + example, profile, charged_cif, tmp_path, capsys, fails +): + root = tmp_path / "sweep" + execution = profile("graspa", *(["--fail"] if fails else [])) + code = example("graspa").main( + [ + charged_cif, + "--outdir", + str(root), + "--pressure-pa", + "1000", + "--pressure-pa", + "100000", + "--execution", + execution, + ] + ) + assert code == (1 if fails else 0) + records = [ + json.loads(line) for line in capsys.readouterr().out.splitlines() + ] + assert len(records) == 2 + for index in range(2): + result = inspect_run(root / f"{index:05d}") + assert result.accepted is not fails + if fails: + assert "code 7" in result.failure.message + else: + assert result.payload.uptake == 12 + assert result.payload.unit == "mol/kg" + + +def test_graspa_does_not_invent_charges(example, sample_cif, tmp_path): + root = tmp_path / "uncharged" + assert ( + example("graspa").main( + [sample_cif, "--outdir", str(root), "--prepare-only"] + ) + == 1 + ) + assert "_atom_site_charge" in inspect_run(root / "00000").failure.message + + +def test_mixture_preserves_input_and_explicit_composition( + example, charged_cif, tmp_path +): + root = tmp_path / "mixture" + arguments = [charged_cif, "--outdir", str(root)] + module = example("graspa_mixture_setup") + assert module.main(arguments) == 0 + content = (root / "simulation.input").read_text() + assert content.count("MoleculeName") == 2 + assert "CO2" in content and "N2" in content + assert [ + line.split()[-1] + for line in content.splitlines() + if "MolFraction" in line + ] == ["0.15", "0.85"] + assert (root / "charged.cif").read_bytes() == Path(charged_cif).read_bytes() + assert not (root / "run.json").exists() + assert module.main(arguments) == 1 + assert (root / "simulation.input").read_text() == content + + +@pytest.fixture +def mlip_inputs(tmp_path): + # A worker-local MACE factory double exercises the actual subprocess API. + # No example code or source calculator implementation is patched. + packages = tmp_path / "engine_packages" + package = packages / "mace" + package.mkdir(parents=True) + (package / "__init__.py").write_text("") + (package / "calculators.py").write_text( + "from ase.calculators.emt import EMT\n" + "def mace_mp(**kwargs):\n" + " return EMT()\n" + ) + profile = tmp_path / "mlip-execution.json" + profile.write_text( + json.dumps( + { + "device": "cpu", + "environment": { + "PYTHONPATH": os.pathsep.join([str(packages), *sys.path]) + }, + } + ) + ) + source = tmp_path / "copper.extxyz" + write( + source, + Atoms("Cu2", positions=[[0, 0, 0], [3, 0, 0]], cell=[10] * 3, pbc=True), + ) + return str(source), str(profile) + + +@pytest.mark.parametrize( + "driver,count", [("energy", 1), ("energy", 2), ("relax", 1)] +) +def test_mlip_worker_results_and_unconverged_relaxation( + example, mlip_inputs, tmp_path, capsys, driver, count +): + source, profile = mlip_inputs + root = tmp_path / "mlip" + arguments = [source] * count + [ + "--outdir", + str(root), + "--backend", + "ase-mace", + "--checkpoint", + "fixture", + "--execution", + profile, + "--driver", + driver, + ] + if driver == "relax": + arguments += ["--steps", "1", "--fmax", "1e-12"] + code = example("mlip").main(arguments) + assert code == (1 if driver == "relax" else 0) + stdout = json.loads(capsys.readouterr().out) + if count == 2: + assert [item["index"] for item in stdout["items"]] == [0, 1] + assert all(item["accepted"] for item in stdout["items"]) + for index in range(count): + bundle = root if count == 1 else root / f"{index:05d}" + result = inspect_run(bundle) + assert result.state == "completed" + assert result.payload.potential_energy is not None + assert len(result.payload.forces) == 2 + assert (bundle / result.payload.final_structure).exists() + assert result.accepted is (driver == "energy") + if driver == "relax": + assert result.payload.converged is False + + +@pytest.mark.parametrize( + "options", + [ + ["--backend", "rootstock", "--dtype", "float32"], + ["--backend", "ase-mace", "--batch-size", "16"], + ["--backend", "nvalchemi-mace", "--batch-size", "0"], + ["--backend", "ase-mace", "--steps", "100"], + ], +) +def test_invalid_mlip_options_fail_before_preparation( + example, tmp_path, options +): + root = tmp_path / "invalid" + with pytest.raises(SystemExit) as exc: + example("mlip").main( + [ + "input.extxyz", + "--outdir", + str(root), + "--checkpoint", + "medium", + *options, + ] + ) + assert exc.value.code == 2 + assert not root.exists()