"""Pure-Python gRASPA simulation helpers (no LangChain / MCP decorators). Contains the core workflow functions for running gRASPA-SYCL simulations, parsing output, and mock simulations for testing. Used by the LangChain ``@tool`` wrapper in :mod:`graspa_tools` and the MCP/Parsl wrappers in :mod:`chemgraph.mcp.graspa_mcp_parsl`. """ from __future__ import annotations import glob import os import random import shlex import shutil import subprocess import time from pathlib import Path import ase import numpy as np from ase.io import read as ase_read from chemgraph.schemas.graspa_schema import graspa_input_schema # Template directory for gRASPA-SYCL input files _file_dir = Path(__file__).parent / "files" / "template_graspa_sycl" # Default gRASPA-SYCL command used on the original HPC environment. DEFAULT_GRASPA_CMD = ( "export OMP_NUM_THREADS=1; " "export ZE_FLAT_DEVICE_HIERARCHY=FLAT; " "/lus/flare/projects/IQC/thang/soft/gRASPA/graspa-sycl/bin/sycl.out" ) GRASPA_CMD_ENV = "CHEMGRAPH_GRASPA_CMD" def get_graspa_command() -> str: """Return the shell command used to launch gRASPA. The default command points at the original ALCF path and is not expected to exist on a laptop. Users can set ``CHEMGRAPH_GRASPA_CMD`` to a local command or executable path when running real gRASPA workflows. """ return os.environ.get(GRASPA_CMD_ENV, DEFAULT_GRASPA_CMD).strip() def _extract_graspa_executable(command: str) -> str | None: """Best-effort extraction of the executable from a shell command.""" segments = [segment.strip() for segment in command.split(";") if segment.strip()] if not segments: return None try: tokens = shlex.split(segments[-1]) except ValueError: return None return tokens[0] if tokens else None def get_graspa_runtime_status() -> dict: """Return a structured availability check for the configured gRASPA command.""" command = get_graspa_command() executable = _extract_graspa_executable(command) status = { "available": False, "env_var": GRASPA_CMD_ENV, "command": command, "executable": executable, "resolved_executable": None, "message": "", } if not executable: status["message"] = "No executable could be parsed from the gRASPA command." return status if os.path.isabs(executable): if os.path.isfile(executable) and os.access(executable, os.X_OK): status["available"] = True status["resolved_executable"] = executable status["message"] = "Configured gRASPA executable is available." else: status["message"] = f"Configured gRASPA executable is not available: {executable}" return status resolved = shutil.which(executable) if resolved: status["available"] = True status["resolved_executable"] = resolved status["message"] = "Configured gRASPA executable is available on PATH." else: status["message"] = f"Configured gRASPA executable is not on PATH: {executable}" return status # --------------------------------------------------------------------------- # Output parsing # --------------------------------------------------------------------------- def _read_graspa_sycl_output( output_path: str, adsorbate: str = "H2O", cifname: str = None, output_fname: str = "raspa.log", temperature: float = None, pressure: float = None, ) -> dict: """Parse gRASPA output and return uptake results. Parameters ---------- output_path : str Directory containing the gRASPA output files. adsorbate : str Name of the adsorbate molecule. cifname : str, optional Stem name of the CIF file (without extension). output_fname : str Name of the gRASPA log file. temperature : float, optional Simulation temperature in Kelvin. pressure : float, optional Simulation pressure in Pascal. Returns ------- dict Parsed results including uptake, status, and CIF path. """ result = { "status": "failure", "uptake_in_mol_kg": 0, "adsorbate": adsorbate, "temperature_in_K": None, "pressure_in_Pa": None, "cif_path": None, } target_file = Path(output_path) / Path(output_fname).name # --- Resolve CIF Path --- if cifname is None: cif_list = glob.glob(os.path.join(output_path, "*.cif")) if len(cif_list) != 1: cifpath = None else: cifpath = os.path.abspath(cif_list[0]) else: cifpath = os.path.abspath(os.path.join(output_path, f"{cifname}.cif")) result["cif_path"] = cifpath # --- Check Log Existence --- if not os.path.exists(target_file): return result # --- Parse Log --- unitcell_line = None uptake_line = None with open(target_file, "r") as rf: for line in rf: if "UnitCells" in line: unitcell_line = line.strip() elif "Overall: Average:" in line: uptake_line = line.strip() if unitcell_line is None or uptake_line is None: return result try: if cifpath is None: raise ValueError(f"Could not resolve CIF path in {output_path}") uptake_total_molecule = float(uptake_line.split()[2][:-1]) unitcell = unitcell_line.split()[4:] unitcell = [int(float(i)) for i in unitcell] atoms = ase_read(cifpath) framework_mass = ( sum(atoms.get_masses()) * unitcell[0] * unitcell[1] * unitcell[2] ) uptake_mol_kg = round((uptake_total_molecule / framework_mass) * 1000, 2) result["uptake_in_mol_kg"] = float(uptake_mol_kg) result["status"] = "success" result["temperature_in_K"] = temperature result["pressure_in_Pa"] = pressure except Exception as e: print(f"Error parsing results in {output_path}: {e}") result["status"] = "failure" return result # --------------------------------------------------------------------------- # Mock simulation (for testing) # --------------------------------------------------------------------------- def mock_graspa(params: graspa_input_schema) -> dict: """Return mock gRASPA results for testing without the SYCL runtime. Parameters ---------- params : graspa_input_schema Input parameters (only ``adsorbates`` is used to determine output shape). Returns ------- dict Simulated uptake results. """ def rand_uptake( low: float, high: float, ndigits: int = 3, min_positive: float | None = None ) -> float: """Generate a rounded mock uptake value. Parameters ---------- low : float Lower bound for the random value. high : float Upper bound for the random value. ndigits : int, optional Number of decimal places to round to. min_positive : float, optional Replacement value when rounding produces zero. Returns ------- float Mock uptake value. """ value = random.uniform(low, high) value = round(value, ndigits) if min_positive is not None and value == 0.0: value = min_positive return value time.sleep(random.uniform(20, 40)) n_ads = len(params.adsorbates) if n_ads == 1: uptake_co2 = rand_uptake(0, 2, ndigits=3) return {"co2_uptake_mol_per_kg": uptake_co2} elif n_ads == 2: uptake_co2 = rand_uptake(0, 2, ndigits=3) uptake_n2 = rand_uptake(0, 0.5, ndigits=3, min_positive=1e-3) try: selectivity = uptake_co2 / uptake_n2 except Exception: selectivity = 1e4 return { "co2_uptake_mol_per_kg": uptake_co2, "n2_uptake_mol_per_kg": uptake_n2, "co2_n2_selectivity": round(selectivity, 2), } elif n_ads == 3: uptake_co2 = rand_uptake(0, 2, ndigits=3) uptake_n2 = rand_uptake(0, 0.5, ndigits=3, min_positive=1e-3) uptake_h2o = rand_uptake(0, 5, ndigits=3) try: selectivity = uptake_co2 / uptake_n2 except Exception: selectivity = 1e4 return { "co2_uptake_mol_per_kg": uptake_co2, "n2_uptake_mol_per_kg": uptake_n2, "h2o_uptake_mol_per_kg": uptake_h2o, "co2_n2_selectivity": round(selectivity, 2), } else: raise ValueError("Only supports 1-3 adsorbates only.") # --------------------------------------------------------------------------- # Core simulation runner # --------------------------------------------------------------------------- def run_graspa_core(params: graspa_input_schema) -> dict: """Run a single gRASPA calculation using specified input parameters. Parameters ---------- params : graspa_input_schema Input parameters for the gRASPA calculation. Returns ------- dict Parsed simulation results including uptake and status. """ def _calculate_cell_size( atoms: ase.Atoms, cutoff: float = 12.8 ) -> list[int]: """Calculate unit-cell replication for GCMC with the given cutoff. Parameters ---------- atoms : ase.Atoms Unit-cell structure. cutoff : float, optional Minimum replicated cell length in angstrom. Returns ------- list[int] Replication factors along the three lattice vectors. """ unit_cell = atoms.cell[:] a = unit_cell[0] b = unit_cell[1] c = unit_cell[2] wa = np.divide( np.linalg.norm(np.dot(np.cross(b, c), a)), np.linalg.norm(np.cross(b, c)), ) wb = np.divide( np.linalg.norm(np.dot(np.cross(c, a), b)), np.linalg.norm(np.cross(c, a)), ) wc = np.divide( np.linalg.norm(np.dot(np.cross(a, b), c)), np.linalg.norm(np.cross(a, b)), ) uc_x = int(np.ceil(cutoff / (0.5 * wa))) uc_y = int(np.ceil(cutoff / (0.5 * wb))) uc_z = int(np.ceil(cutoff / (0.5 * wc))) return [uc_x, uc_y, uc_z] cif_path = Path(params.input_structure_file).resolve() if not cif_path.exists(): raise FileNotFoundError(f"CIF file does not exist: {cif_path}") runtime_status = get_graspa_runtime_status() if not runtime_status["available"]: return { "status": "failure", "error_type": "GraspaRuntimeUnavailable", "message": runtime_status["message"], "runtime": runtime_status, "uptake_in_mol_kg": None, "adsorbate": params.adsorbate, "temperature_in_K": params.temperature, "pressure_in_Pa": params.pressure, "cif_path": str(cif_path), } base_dir = cif_path.parent cifname = cif_path.stem temperature = params.temperature pressure = params.pressure adsorbate = params.adsorbate n_cycle = params.n_cycles folder_name = f"{cifname}--{adsorbate}-{temperature}-{pressure:g}" sim_dir = base_dir / folder_name sim_dir.mkdir(parents=True, exist_ok=True) for item in _file_dir.iterdir(): dest = sim_dir / item.name if item.is_dir(): if dest.exists(): shutil.rmtree(dest) shutil.copytree(item, dest) else: shutil.copy2(item, sim_dir) # Copy the specific CIF file shutil.copy2(cif_path, sim_dir / f"{cifname}.cif") atoms = ase_read(cif_path) [uc_x, uc_y, uc_z] = _calculate_cell_size(atoms) input_file = sim_dir / "simulation.input" temp_file = sim_dir / "simulation.input.tmp" with open(input_file, "r") as f_in, open(temp_file, "w") as f_out: for line in f_in: if "NCYCLE" in line: line = line.replace("NCYCLE", str(n_cycle)) if "ADSORBATE" in line: line = line.replace("ADSORBATE", adsorbate) if "TEMPERATURE" in line: line = line.replace("TEMPERATURE", str(temperature)) if "PRESSURE" in line: line = line.replace("PRESSURE", str(pressure)) if "UC_X UC_Y UC_Z" in line: line = line.replace("UC_X UC_Y UC_Z", f"{uc_x} {uc_y} {uc_z}") if "CUTOFF" in line: line = line.replace("CUTOFF", str(12.8)) if "CIFFILE" in line: line = line.replace("CIFFILE", cifname) f_out.write(line) shutil.move(temp_file, input_file) output_filename = Path(params.output_result_file).name with ( open(os.path.join(sim_dir, output_filename), "w") as fp, open(os.path.join(sim_dir, "raspa.err"), "w") as fe, ): subprocess.run(get_graspa_command(), cwd=sim_dir, stdout=fp, stderr=fe, shell=True) return _read_graspa_sycl_output( output_path=str(sim_dir), adsorbate=adsorbate, cifname=cifname, output_fname=params.output_result_file, temperature=temperature, pressure=pressure, )