diff --git a/.gitignore b/.gitignore index 2d459a7..622571f 100644 --- a/.gitignore +++ b/.gitignore @@ -168,3 +168,16 @@ __version__.py # Ignore temporary files tmp/ + +# Local calculation output +xtb_vib/ +/screening/ +/redox-screening/ +/results.csv +/results.json +/results-run.json +/redox-results*.csv +/redox-results*.json +/thermoscreening.slurm +/thermoscreening-*.out +/results-shards/ diff --git a/ThermoScreening/calculator/xtb.py b/ThermoScreening/calculator/xtb.py index b72a77f..0118eb9 100644 --- a/ThermoScreening/calculator/xtb.py +++ b/ThermoScreening/calculator/xtb.py @@ -56,8 +56,8 @@ def optimise_and_frequencies(atoms, calc, fmax=0.01): Parameters ---------- atoms : ase.Atoms - Initial geometry (with any ``info['charge']`` / ``info['spin']`` already - set for the calculator). + Initial geometry with total charge and unpaired electrons represented by + ASE initial charges and magnetic moments for the calculator. calc : ase.calculators.calculator.Calculator The ASE calculator to attach (e.g. a tblite ``TBLite`` instance). fmax : float @@ -79,17 +79,22 @@ def optimise_and_frequencies(atoms, calc, fmax=0.01): energy_hartree = _eV_to_hartree(atoms.get_potential_energy()) vibrations = Vibrations(atoms, name="xtb_vib") - vibrations.run() - frequencies = _real_frequencies_cm(vibrations.get_frequencies()) vibrations.clean() + try: + vibrations.run() + frequencies = _real_frequencies_cm(vibrations.get_frequencies()) + finally: + vibrations.clean() return atoms, energy_hartree, frequencies def xtb_calculator(method="GFN2-xTB"): """ - Build a tblite GFN-xTB ASE calculator (charge and spin are read from - ``atoms.info``). + Build a tblite GFN-xTB ASE calculator. + + tblite reads total charge and unpaired electrons from the attached atoms' + initial charges and magnetic moments, respectively. Parameters ---------- diff --git a/ThermoScreening/cli/slurm.py b/ThermoScreening/cli/slurm.py new file mode 100644 index 0000000..88c2d6b --- /dev/null +++ b/ThermoScreening/cli/slurm.py @@ -0,0 +1,182 @@ +"""Slurm job-array support for distributed thermochemistry screens.""" + +import os +import re +import shlex +import shutil +import subprocess +import sys +from pathlib import Path + +from ThermoScreening.exceptions import TSValueError + + +_SLURM_TOKEN = re.compile(r"^[A-Za-z0-9_.:,/-]+$") + + +def _positive_integer(value, name): + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise TSValueError(f"{name} must be an integer >= 1.") + return value + + +def _directive_value(value, name): + if value is None: + return None + if not _SLURM_TOKEN.fullmatch(str(value)): + raise TSValueError(f"{name} contains unsupported characters.") + return str(value) + + +def write_slurm_array_script( + command_args, + *, + tasks, + script, + local_jobs=1, + cpus_per_task=1, + job_name="thermoscreening", + walltime=None, + memory=None, + partition=None, + account=None, + preamble=None, + working_directory=None, + python_executable=None, +): + """Write a Slurm array that runs one deterministic screen shard per task.""" + tasks = _positive_integer(tasks, "tasks") + local_jobs = _positive_integer(local_jobs, "local_jobs") + cpus_per_task = _positive_integer(cpus_per_task, "cpus_per_task") + if local_jobs > cpus_per_task: + raise TSValueError("local --jobs cannot exceed cpus_per_task.") + if not command_args or command_args[0] != "screen": + raise TSValueError("Slurm arrays currently support the screen command.") + if "--shard-index" in command_args or "--shard-count" in command_args: + raise TSValueError("Do not pass shard options to the Slurm generator.") + + job_name = _directive_value(job_name, "job_name") + optional_directives = ( + ("time", _directive_value(walltime, "time")), + ("mem", _directive_value(memory, "memory")), + ("partition", _directive_value(partition, "partition")), + ("account", _directive_value(account, "account")), + ) + script = Path(script).resolve() + script.parent.mkdir(parents=True, exist_ok=True) + working_directory = Path(working_directory or os.getcwd()).resolve() + executable = str(python_executable or sys.executable) + threads_per_job = max(1, cpus_per_task // local_jobs) + + lines = [ + "#!/usr/bin/env bash", + f"#SBATCH --job-name={job_name}", + f"#SBATCH --array=0-{tasks - 1}", + f"#SBATCH --cpus-per-task={cpus_per_task}", + f"#SBATCH --output={job_name}-%A_%a.out", + ] + lines.extend( + f"#SBATCH --{option}={value}" + for option, value in optional_directives + if value is not None + ) + lines.extend( + [ + "", + "set -euo pipefail", + f"cd -- {shlex.quote(str(working_directory))}", + f"export OMP_NUM_THREADS={threads_per_job}", + f"export OPENBLAS_NUM_THREADS={threads_per_job}", + f"export MKL_NUM_THREADS={threads_per_job}", + f"export NUMEXPR_NUM_THREADS={threads_per_job}", + ] + ) + if preamble is not None: + preamble_text = Path(preamble).read_text(encoding="utf-8").rstrip() + if preamble_text: + lines.extend(["", preamble_text]) + + command = shlex.join([executable, "-m", "ThermoScreening", *command_args]) + lines.extend( + [ + "", + command + + ' --shard-index "$SLURM_ARRAY_TASK_ID"' + + ' --shard-count "$SLURM_ARRAY_TASK_COUNT"', + "", + ] + ) + script.write_text("\n".join(lines), encoding="utf-8") + script.chmod(script.stat().st_mode | 0o111) + return script + + +def submit_slurm_array( + script, + *, + shard_directory, + out, + job_name="thermoscreening", + partition=None, + account=None, + working_directory=None, + python_executable=None, +): + """Submit a Slurm array and a dependent result-collection job.""" + sbatch = shutil.which("sbatch") + if sbatch is None: + raise TSValueError("sbatch was not found on PATH.") + job_name = _directive_value(job_name, "job_name") + partition = _directive_value(partition, "partition") + account = _directive_value(account, "account") + working_directory = Path(working_directory or os.getcwd()).resolve() + executable = str(python_executable or sys.executable) + + try: + array = subprocess.run( + [sbatch, "--parsable", str(Path(script).resolve())], + check=True, + capture_output=True, + text=True, + ) + array_job_id = array.stdout.strip().split(";", maxsplit=1)[0] + if not array_job_id: + raise TSValueError("sbatch returned no array job ID.") + collect_command = shlex.join( + [ + executable, + "-m", + "ThermoScreening", + "collect", + str(shard_directory), + "-o", + str(out), + ] + ) + collector_args = [ + sbatch, + "--parsable", + f"--dependency=afterany:{array_job_id}", + f"--job-name={job_name}-collect", + f"--chdir={working_directory}", + f"--output={job_name}-collect-%j.out", + ] + if partition is not None: + collector_args.append(f"--partition={partition}") + if account is not None: + collector_args.append(f"--account={account}") + collector_args.append(f"--wrap={collect_command}") + collector = subprocess.run( + collector_args, + check=True, + capture_output=True, + text=True, + ) + except subprocess.CalledProcessError as exc: + detail = (exc.stderr or exc.stdout or str(exc)).strip() + raise TSValueError(f"Slurm submission failed: {detail}") from exc + + collector_job_id = collector.stdout.strip().split(";", maxsplit=1)[0] + if not collector_job_id: + raise TSValueError("sbatch returned no collector job ID.") + return array_job_id, collector_job_id diff --git a/ThermoScreening/cli/thermo.py b/ThermoScreening/cli/thermo.py index 55649b5..c1aa4f6 100644 --- a/ThermoScreening/cli/thermo.py +++ b/ThermoScreening/cli/thermo.py @@ -1,6 +1,6 @@ """The ``thermo`` command-line entry point.""" -from argparse import ArgumentParser +from argparse import ArgumentParser, REMAINDER import sys import time @@ -11,13 +11,29 @@ install_gbsa_param, install_slakos, ) +from ThermoScreening.exceptions import TSValueError +from ThermoScreening.cli.slurm import write_slurm_array_script, submit_slurm_array from ThermoScreening.thermo.api import execute -from ThermoScreening.thermo.screening import screen, rank_by_gibbs +from ThermoScreening.thermo.screening import ( + collect_screen_shards, + redox_screen, + screen, + screen_shard_directory, + rank_by_gibbs, +) from ThermoScreening.thermo.conformers import generate as generate_conformers, write_conformers from ThermoScreening.version import __version__ -SUBCOMMANDS = {"setup-dftb", "doctor", "screen", "conformers"} +SUBCOMMANDS = { + "setup-dftb", + "doctor", + "screen", + "collect", + "slurm", + "redox", + "conformers", +} def _run_parser(): @@ -138,6 +154,164 @@ def _command_parser(): help="Number of molecules to run concurrently (default 1). Each job runs " "in its own process and directory.", ) + screen_parser.add_argument( + "--shard-index", + type=int, + default=None, + help="Zero-based cluster shard to execute; requires --shard-count.", + ) + screen_parser.add_argument( + "--shard-count", + type=int, + default=None, + help="Total number of deterministic cluster shards.", + ) + + collect_parser = subparsers.add_parser( + "collect", + help="Validate and combine distributed screen shards.", + ) + collect_parser.add_argument( + "shard_directory", + help="Directory containing shard-*.json and shard-*-run.json files.", + ) + collect_parser.add_argument( + "-o", + "--out", + default="results", + help="Output stem for the combined CSV, JSON, and run metadata.", + ) + + slurm_parser = subparsers.add_parser( + "slurm", + help="Generate or submit a Slurm array for thermo screen.", + ) + slurm_parser.add_argument("--tasks", type=int, required=True) + slurm_parser.add_argument("--script", default="thermoscreening.slurm") + slurm_parser.add_argument("--job-name", default="thermoscreening") + slurm_parser.add_argument("--cpus-per-task", type=int, default=1) + slurm_parser.add_argument("--time", dest="walltime", default=None) + slurm_parser.add_argument("--mem", dest="memory", default=None) + slurm_parser.add_argument("--partition", default=None) + slurm_parser.add_argument("--account", default=None) + slurm_parser.add_argument( + "--preamble", + default=None, + help="Shell file inserted before the worker command for modules or variables.", + ) + slurm_parser.add_argument( + "--submit", + action="store_true", + help="Submit the array and a dependent collection job with sbatch.", + ) + slurm_parser.add_argument( + "command_args", + nargs=REMAINDER, + help="Screen command after '--', for example: -- screen molecules.csv -o out.", + ) + + redox_parser = subparsers.add_parser( + "redox", + help="Run a three-state redox screen from one starting geometry.", + ) + redox_parser.add_argument( + "source", + help="SMILES, one .xyz/.gen file, a directory, or a CSV containing " + "a path or SMILES per row.", + ) + redox_parser.add_argument( + "-o", + "--out", + default="redox-results", + help="Output stem for aggregate and per-state CSV/JSON results.", + ) + redox_parser.add_argument( + "--charge", type=int, default=None, + help="Oxidized-state charge. Ionic SMILES are inferred; structures default to 0.", + ) + redox_parser.add_argument( + "--reference", + default=None, + help="Reference molecule name, SMILES, or structure path.", + ) + redox_parser.add_argument( + "--reference-e1", + type=float, + default=None, + help="Measured first reduction potential of the reference in V.", + ) + redox_parser.add_argument( + "--reference-e2", + type=float, + default=None, + help="Measured second reduction potential of the reference in V.", + ) + redox_parser.add_argument( + "--reference-charge", + type=int, + default=None, + help="Oxidized-state charge of a separate reference, if not inferable.", + ) + redox_parser.add_argument( + "--potential-scale", + default=None, + help="Label for calibrated potentials, for example 'Fc/Fc+'.", + ) + redox_parser.add_argument( + "--max-conformers", + type=int, + default=20, + help="SMILES conformers to embed before selecting the lowest MMFF structure.", + ) + redox_parser.add_argument( + "--temperature", type=float, default=298.15, help="Temperature in K.", + ) + redox_parser.add_argument( + "--pressure", type=float, default=101325.0, help="Pressure in Pa.", + ) + redox_parser.add_argument( + "--directory", + default="redox-screening", + help="Working directory for generated inputs and state calculations.", + ) + redox_parser.add_argument( + "--parameter-set", default=None, choices=["3ob", "mio"], + help="DFTB+ Slater-Koster set; DFTB+ defaults to '3ob'.", + ) + redox_parser.add_argument( + "--solvent", + default=None, + help="Implicit-solvation solvent applied consistently to every state.", + ) + redox_parser.add_argument( + "--dispersion", default=None, choices=["d3-bj"], + help="DFTB+ dispersion correction. Default none.", + ) + redox_parser.add_argument( + "--quasi-rrho", + action="store_true", + help="Use quasi-RRHO vibrational entropy for every state.", + ) + redox_parser.add_argument( + "--engine", default="dftb+", choices=["dftb+", "xtb", "xtb-cli"], + help="Calculation engine (default 'dftb+').", + ) + redox_parser.add_argument( + "--method", default=None, choices=["GFN2-xTB", "GFN1-xTB"], + help="GFN parametrisation; xTB engines default to 'GFN2-xTB'.", + ) + redox_parser.add_argument( + "--resume", + action="store_true", + help="Reuse only state calculations with matching input fingerprints.", + ) + redox_parser.add_argument( + "-j", + "--jobs", + type=int, + default=1, + help="Number of charge-state calculations to run concurrently.", + ) conf_parser = subparsers.add_parser( "conformers", @@ -249,6 +423,8 @@ def run_screen(parser_args): method=parser_args.method, resume=parser_args.resume, jobs=parser_args.jobs, + shard_index=getattr(parser_args, "shard_index", None), + shard_count=getattr(parser_args, "shard_count", None), ) ranked = rank_by_gibbs(results) @@ -264,8 +440,124 @@ def run_screen(parser_args): print(f" {record['name']}: {record['error']}") print(f"Screened {len(results)} molecules ({len(failures)} failed).") + shard_index = getattr(parser_args, "shard_index", None) + if shard_index is None: + result_stem = parser_args.out + else: + result_stem = screen_shard_directory(parser_args.out) / f"shard-{shard_index:05d}" + print(f"Results: {result_stem}.csv, {result_stem}.json") + + return 1 if failures else 0 + + +def run_collect(parser_args): + """Combine and validate distributed screening results.""" + try: + results = collect_screen_shards(parser_args.shard_directory, out=parser_args.out) + except TSValueError as exc: + print(f"Collection failed: {exc}", file=sys.stderr) + return 1 + + failures = [record for record in results if record.get("status") != "ok"] + print(f"Collected {len(results)} molecules ({len(failures)} failed).") print(f"Results: {parser_args.out}.csv, {parser_args.out}.json") + return 1 if failures else 0 + + +def run_slurm(parser_args): + """Generate or submit a Slurm screening array.""" + command_args = list(parser_args.command_args) + if command_args and command_args[0] == "--": + command_args.pop(0) + try: + if not command_args or command_args[0] != "screen": + raise TSValueError("Slurm arrays currently support the screen command.") + nested = parse_args(command_args) + script = write_slurm_array_script( + command_args, + tasks=parser_args.tasks, + script=parser_args.script, + local_jobs=nested.jobs, + cpus_per_task=parser_args.cpus_per_task, + job_name=parser_args.job_name, + walltime=parser_args.walltime, + memory=parser_args.memory, + partition=parser_args.partition, + account=parser_args.account, + preamble=parser_args.preamble, + ) + print(f"Slurm script: {script}") + if parser_args.submit: + array_id, collector_id = submit_slurm_array( + script, + shard_directory=screen_shard_directory(nested.out), + out=nested.out, + job_name=parser_args.job_name, + partition=parser_args.partition, + account=parser_args.account, + ) + print(f"Submitted array job {array_id} and collector job {collector_id}.") + else: + print(f"Submit: sbatch {script}") + print( + "Collect after completion: thermo collect " + f"{screen_shard_directory(nested.out)} -o {nested.out}" + ) + except (OSError, TSValueError, ValueError) as exc: + print(f"Slurm setup failed: {exc}", file=sys.stderr) + return 1 + return 0 + + +def run_redox(parser_args): + """Run the three-state redox-screening workflow.""" + + try: + results = redox_screen( + parser_args.source, + out=parser_args.out, + charge=parser_args.charge, + temperature=parser_args.temperature, + pressure=parser_args.pressure, + directory=parser_args.directory, + parameter_set=parser_args.parameter_set, + solvent=parser_args.solvent, + dispersion=parser_args.dispersion, + quasi_rrho=parser_args.quasi_rrho, + engine=parser_args.engine, + method=parser_args.method, + resume=parser_args.resume, + jobs=parser_args.jobs, + reference=parser_args.reference, + reference_e1=parser_args.reference_e1, + reference_e2=parser_args.reference_e2, + reference_charge=parser_args.reference_charge, + potential_scale=parser_args.potential_scale, + max_conformers=parser_args.max_conformers, + ) + except (TSValueError, ValueError) as exc: + print(f"Redox screen failed: {exc}", file=sys.stderr) + return 1 + successful = [result for result in results if result["status"] == "ok"] + if successful: + print("Reduction potentials:") + for result in successful: + print( + f" {result['name']}: E1={result['E1_V']:.4f} V, " + f"E2={result['E2_V']:.4f} V, E2e={result['E2e_V']:.4f} V " + f"({result['potential_scale']})" + ) + + failures = [result for result in results if result["status"] != "ok"] + if failures: + print(f"Failed ({len(failures)}):") + for result in failures: + print(f" {result['name']}: {result['error']}") + print(f"Processed {len(results)} molecules ({len(failures)} failed).") + print(f"Results: {parser_args.out}.csv, {parser_args.out}.json") + print(f"State results: {parser_args.out}-states.csv, {parser_args.out}-states.json") + print(f"Run metadata: {parser_args.out}-run.json") return 1 if failures else 0 @@ -318,6 +610,15 @@ def main(): if command == "screen": return run_screen(parser_args) + if command == "collect": + return run_collect(parser_args) + + if command == "slurm": + return run_slurm(parser_args) + + if command == "redox": + return run_redox(parser_args) + if command == "conformers": return run_conformers(parser_args) diff --git a/ThermoScreening/thermo/__init__.py b/ThermoScreening/thermo/__init__.py index 1130f92..13aa2b9 100644 --- a/ThermoScreening/thermo/__init__.py +++ b/ThermoScreening/thermo/__init__.py @@ -6,7 +6,7 @@ from .inputFileReader import InputFileReader from .system import System from .thermo import Thermo -from .screening import screen, rank_by_gibbs +from .screening import collect_screen_shards, redox_screen, screen, rank_by_gibbs from .conformers import generate as generate_conformers, write_conformers, generate_thermo_ensemble from .reactions import ( calibrate_reduction_reference, diff --git a/ThermoScreening/thermo/api.py b/ThermoScreening/thermo/api.py index 3fe1cff..d86120e 100644 --- a/ThermoScreening/thermo/api.py +++ b/ThermoScreening/thermo/api.py @@ -977,8 +977,12 @@ def xtb_thermo( spin = 0.0 if electrons % 2 == 0 else 0.5 prepared = atoms.copy() - prepared.info["charge"] = int(round(charge)) - prepared.info["spin"] = int(round(2.0 * float(spin))) + atom_count = len(prepared) + unpaired = int(round(2.0 * float(spin))) + prepared.set_initial_charges(np.full(atom_count, charge / atom_count)) + prepared.set_initial_magnetic_moments( + np.full(atom_count, unpaired / atom_count) + ) with _run_in_directory(directory): optimized_atoms, potential_energy, frequencies = optimise_and_frequencies( diff --git a/ThermoScreening/thermo/conformers.py b/ThermoScreening/thermo/conformers.py index f32ecad..64cfabc 100644 --- a/ThermoScreening/thermo/conformers.py +++ b/ThermoScreening/thermo/conformers.py @@ -30,7 +30,9 @@ def _conformer_to_atoms(mol, conformer_id): conformer = mol.GetConformer(conformer_id) symbols = [atom.GetSymbol() for atom in mol.GetAtoms()] positions = np.asarray(conformer.GetPositions(), dtype=float) - return Atoms(symbols=symbols, positions=positions) + atoms = Atoms(symbols=symbols, positions=positions) + atoms.info["formal_charge"] = sum(atom.GetFormalCharge() for atom in mol.GetAtoms()) + return atoms def generate( diff --git a/ThermoScreening/thermo/screening.py b/ThermoScreening/thermo/screening.py index 80ca872..838fca9 100644 --- a/ThermoScreening/thermo/screening.py +++ b/ThermoScreening/thermo/screening.py @@ -9,8 +9,10 @@ import csv import functools +import hashlib import json import logging +import math import os from concurrent.futures import ProcessPoolExecutor, as_completed from dataclasses import dataclass @@ -22,8 +24,12 @@ from ThermoScreening.calculator.dftbplus import resolve_parameter_set from ThermoScreening.exceptions import TSValueError from ThermoScreening.utils.custom_logging import setup_logger +from ThermoScreening.version import __version__ from .api import dftbplus_thermo, xtb_thermo, xtb_cli_thermo +from ._units import HARTREE_TO_EV +from .conformers import generate as generate_conformers, write_conformers +from .reactions import SHE_ABSOLUTE_POTENTIAL logger = logging.getLogger(__package_name__).getChild("screening") logger = setup_logger(logger) @@ -41,8 +47,31 @@ "G_total_hartree", "S_cal_per_mol_K", "Cv_cal_per_mol_K", + "fingerprint", "error", ] +_REDOX_RESULT_FIELDS = [ + "name", + "path", + "charge", + "status", + "G_oxidized_hartree", + "G_reduced_once_hartree", + "G_reduced_twice_hartree", + "E1_V", + "E2_V", + "E2e_V", + "potential_gap_V", + "potential_inversion", + "potential_scale", + "run_fingerprint", + "error", +] +_REDOX_STATES = ( + ("oxidized", 0, 0), + ("reduced_once", -1, 1), + ("reduced_twice", -2, 2), +) @dataclass @@ -55,6 +84,16 @@ class ScreeningJob: spin: float | None = None +@dataclass +class _RedoxJob: + """One molecule and the three electronic states in a redox workflow.""" + + name: str + path: Path + charge: float + spins: tuple[float | None, float | None, float | None] + + def _jobs_from_directory(directory: Path, charge: float, spin=None): jobs = [ ScreeningJob(name=path.stem, path=path, charge=charge, spin=spin) @@ -155,7 +194,85 @@ def rank_by_gibbs(results): return sorted(ranked, key=lambda record: record["G_total_hartree"]) -def _load_completed(out): +def _stable_value(value): + """Return a deterministic JSON-compatible representation.""" + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, Path): + return str(value) + if isinstance(value, dict): + return {str(key): _stable_value(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_stable_value(item) for item in value] + if hasattr(value, "tolist"): + return _stable_value(value.tolist()) + return repr(value) + + +def _file_sha256(path): + """Hash a structure without making a missing input abort the whole screen.""" + digest = hashlib.sha256() + try: + with open(path, "rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + except OSError: + return None + return digest.hexdigest() + + +def _payload_fingerprint(payload): + encoded = json.dumps( + _stable_value(payload), sort_keys=True, separators=(",", ":"), allow_nan=False + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def _screen_settings( + *, + engine, + temperature, + pressure, + method, + solvent, + dispersion, + quasi_rrho, + parameter_set, + parameters, + spin_constants, +): + return { + "thermoscreening_version": __version__, + "engine": engine, + "temperature_K": temperature, + "pressure_Pa": pressure, + "method": method if engine in ("xtb", "xtb-cli") else None, + "solvent": solvent, + "dispersion": dispersion if engine == "dftb+" else None, + "quasi_rrho": quasi_rrho, + "parameter_set": parameter_set if engine == "dftb+" else None, + "parameters": parameters if engine == "dftb+" else None, + "spin_constants": spin_constants if engine == "dftb+" else None, + } + + +def _job_provenance(job, settings, input_index=None): + structure_hash = _file_sha256(job.path) + payload = { + "name": job.name, + "path": str(job.path), + "structure_sha256": structure_hash, + "charge": job.charge, + "spin": job.spin, + "settings": settings, + } + provenance = {**payload, "fingerprint": _payload_fingerprint(payload)} + if input_index is not None: + provenance["input_index"] = input_index + return provenance + + +def _load_completed(out, expected_fingerprints=None): """ Load the successfully-completed records from a prior ``.json``. @@ -172,11 +289,20 @@ def _load_completed(out): return {} if not isinstance(prior, list): return {} - return { - record["name"]: record - for record in prior - if isinstance(record, dict) and record.get("status") == "ok" and "name" in record - } + completed = {} + for record in prior: + if not ( + isinstance(record, dict) + and record.get("status") == "ok" + and "name" in record + ): + continue + if expected_fingerprints is not None: + expected = expected_fingerprints.get(record["name"]) + if expected is None or record.get("fingerprint") != expected: + continue + completed[record["name"]] = record + return completed def _atomic_write(path, write_fn): @@ -190,6 +316,65 @@ def _atomic_write(path, write_fn): os.replace(tmp, path) +def _sidecar_path(out, suffix): + stem = Path(str(out)).with_suffix("") + return stem.parent / f"{stem.name}{suffix}" + + +def _write_json(path, payload): + path.parent.mkdir(parents=True, exist_ok=True) + _atomic_write( + path, + lambda handle: handle.write( + json.dumps(_stable_value(payload), indent=2, sort_keys=True, allow_nan=False) + + "\n" + ), + ) + + +def _write_screen_run_metadata(out, source, settings, job_provenance, shard=None): + payload = { + "schema_version": 1, + "workflow": "thermochemistry_screen", + "source": str(source), + "settings": settings, + "jobs": job_provenance, + } + if shard is not None: + payload["shard"] = shard + path = _sidecar_path(out, "-run.json") + _write_json(path, payload) + return path + + +def _validate_shard(shard_index, shard_count): + if shard_index is None and shard_count is None: + return None + if shard_index is None or shard_count is None: + raise TSValueError("shard_index and shard_count must be supplied together.") + if isinstance(shard_index, bool) or not isinstance(shard_index, int): + raise TSValueError("shard_index must be an integer.") + if isinstance(shard_count, bool) or not isinstance(shard_count, int): + raise TSValueError("shard_count must be an integer.") + if shard_count < 1: + raise TSValueError("shard_count must be >= 1.") + if not 0 <= shard_index < shard_count: + raise TSValueError( + f"shard_index must satisfy 0 <= index < {shard_count}, got {shard_index}." + ) + return {"index": shard_index, "count": shard_count} + + +def screen_shard_directory(out): + """Return the directory containing distributed screen result shards.""" + stem = Path(str(out)).with_suffix("") + return stem.parent / f"{stem.name}-shards" + + +def _screen_shard_stem(out, shard_index): + return screen_shard_directory(out) / f"shard-{shard_index:05d}" + + def _write_results(results, out): stem = Path(str(out)).with_suffix("") if stem.parent != Path(""): @@ -212,6 +397,7 @@ def write_csv(handle): def _run_job( job, *, + fingerprint, engine, temperature, pressure, @@ -235,6 +421,7 @@ def _run_job( "path": str(job.path), "charge": job.charge, "status": "ok", + "fingerprint": fingerprint, "error": "", } try: @@ -300,6 +487,8 @@ def screen( method="GFN2-xTB", resume=False, jobs=1, + shard_index=None, + shard_count=None, ): """ Run a thermochemistry screen over a set of molecules. @@ -359,6 +548,11 @@ def screen( ``jobs > 1`` the molecules run in a process pool (process-based because each job changes the working directory). Results and output ordering are unchanged; only the wall-clock time differs. + shard_index, shard_count : int, optional + Select one deterministic zero-based shard of the input for execution on + a cluster. Both values are required together. Shards write to isolated + output and calculation directories and can be combined with + :func:`collect_screen_shards`. Returns ------- @@ -371,15 +565,52 @@ def screen( ) if jobs < 1: raise TSValueError(f"jobs must be >= 1, got {jobs}.") + if not math.isfinite(float(temperature)) or float(temperature) <= 0: + raise TSValueError("temperature must be finite and positive.") + if not math.isfinite(float(pressure)) or float(pressure) <= 0: + raise TSValueError("pressure must be finite and positive.") + shard = _validate_shard(shard_index, shard_count) spin_constants = None if engine == "dftb+": default_parameters, spin_constants = resolve_parameter_set(parameter_set) parameters = default_parameters if parameters is None else parameters - job_list = _load_jobs(source, charge, spin) + all_jobs = _load_jobs(source, charge, spin) root = Path(directory) - completed = _load_completed(out) if resume else {} + settings = _screen_settings( + engine=engine, + temperature=temperature, + pressure=pressure, + method=method, + solvent=solvent, + dispersion=dispersion, + quasi_rrho=quasi_rrho, + parameter_set=parameter_set, + parameters=parameters, + spin_constants=spin_constants, + ) + if shard is None: + selected = list(range(len(all_jobs))) + job_list = all_jobs + else: + selected = [ + index + for index in range(len(all_jobs)) + if index % shard["count"] == shard["index"] + ] + job_list = [all_jobs[index] for index in selected] + out = _screen_shard_stem(out, shard["index"]) + root = root / "shards" / f"shard-{shard['index']:05d}" + provenance = [ + _job_provenance(all_jobs[index], settings, input_index=index) + for index in selected + ] + fingerprints = {item["name"]: item["fingerprint"] for item in provenance} + _write_screen_run_metadata(out, source, settings, provenance, shard=shard) + completed = ( + _load_completed(out, expected_fingerprints=fingerprints) if resume else {} + ) run_one = functools.partial( _run_job, @@ -421,13 +652,19 @@ def _store(job, record): pending.append(job) if records: csv_path, json_path = _write_results(_ordered(), out) + elif not pending: + csv_path, json_path = _write_results([], out) if jobs > 1 and len(pending) > 1: with ProcessPoolExecutor(max_workers=min(jobs, len(pending))) as executor: future_to_job = {} for job in pending: logger.info(f"Screening {job.name} (charge {job.charge})") - future_to_job[executor.submit(run_one, job)] = job + future_to_job[ + executor.submit( + run_one, job, fingerprint=fingerprints[job.name] + ) + ] = job for future in as_completed(future_to_job): job = future_to_job[future] try: @@ -441,13 +678,785 @@ def _store(job, record): "path": str(job.path), "charge": job.charge, "status": "error", + "fingerprint": fingerprints[job.name], "error": str(exc), } _store(job, record) else: for job in pending: logger.info(f"Screening {job.name} (charge {job.charge})") - _store(job, run_one(job)) + _store(job, run_one(job, fingerprint=fingerprints[job.name])) logger.info(f"Wrote {csv_path} and {json_path}") return _ordered() + + +def _read_json(path, expected_type): + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError) as exc: + raise TSValueError(f"Could not read cluster result '{path}': {exc}") from exc + if not isinstance(payload, expected_type): + raise TSValueError(f"Cluster result '{path}' has an invalid JSON structure.") + return payload + + +def collect_screen_shards(shard_directory, out="results"): + """Validate and combine all shards from a distributed screen.""" + shard_directory = Path(shard_directory) + metadata_paths = sorted(shard_directory.glob("shard-*-run.json")) + if not metadata_paths: + raise TSValueError(f"No screen shards found in '{shard_directory}'.") + + expected_count = None + expected_settings = None + expected_source = None + seen_shards = set() + records_by_name = {} + provenance_by_name = {} + + for metadata_path in metadata_paths: + metadata = _read_json(metadata_path, dict) + if metadata.get("workflow") != "thermochemistry_screen": + raise TSValueError(f"'{metadata_path}' is not screen run metadata.") + shard = metadata.get("shard") + if not isinstance(shard, dict) or not {"index", "count"} <= set(shard): + raise TSValueError(f"'{metadata_path}' has no valid shard metadata.") + shard_index = shard["index"] + shard_count = shard["count"] + _validate_shard(shard_index, shard_count) + if shard_index in seen_shards: + raise TSValueError(f"Duplicate cluster shard index {shard_index}.") + seen_shards.add(shard_index) + + settings = metadata.get("settings") + source = metadata.get("source") + if expected_count is None: + expected_count = shard_count + expected_settings = settings + expected_source = source + elif shard_count != expected_count: + raise TSValueError("Cluster shards disagree on shard_count.") + elif settings != expected_settings or source != expected_source: + raise TSValueError("Cluster shards were produced from different inputs or settings.") + + result_name = metadata_path.name.removesuffix("-run.json") + ".json" + result_path = metadata_path.with_name(result_name) + if not result_path.is_file(): + raise TSValueError(f"Result file is missing for cluster shard {shard_index}.") + shard_records = _read_json(result_path, list) + shard_jobs = metadata.get("jobs") + if not isinstance(shard_jobs, list): + raise TSValueError(f"'{metadata_path}' has invalid job metadata.") + jobs_by_name = { + job.get("name"): job for job in shard_jobs if isinstance(job, dict) + } + if ( + len(jobs_by_name) != len(shard_jobs) + or any(not isinstance(name, str) or not name for name in jobs_by_name) + ): + raise TSValueError(f"'{metadata_path}' contains invalid or duplicate jobs.") + shard_records_by_name = { + record.get("name"): record + for record in shard_records + if isinstance(record, dict) + } + if ( + len(shard_records_by_name) != len(shard_records) + or set(shard_records_by_name) != set(jobs_by_name) + ): + raise TSValueError( + f"Cluster shard {shard_index} results do not match its job metadata." + ) + + for name, job in jobs_by_name.items(): + record = shard_records_by_name[name] + provenance_payload = { + key: job.get(key) + for key in ( + "name", + "path", + "structure_sha256", + "charge", + "spin", + "settings", + ) + } + if ( + job.get("settings") != settings + or _payload_fingerprint(provenance_payload) != job.get("fingerprint") + ): + raise TSValueError( + f"Invalid provenance fingerprint for {name!r} in " + f"cluster shard {shard_index}." + ) + if record.get("fingerprint") != job.get("fingerprint"): + raise TSValueError( + f"Fingerprint mismatch for {name!r} in cluster shard {shard_index}." + ) + if name in records_by_name: + raise TSValueError(f"Duplicate molecule {name!r} across cluster shards.") + if not isinstance(job.get("input_index"), int): + raise TSValueError(f"Missing input index for {name!r}.") + if job["input_index"] % shard_count != shard_index: + raise TSValueError( + f"Input position for {name!r} does not belong to " + f"cluster shard {shard_index}." + ) + records_by_name[name] = record + provenance_by_name[name] = job + + missing = sorted(set(range(expected_count)) - seen_shards) + if missing: + formatted = ", ".join(str(index) for index in missing) + raise TSValueError(f"Missing cluster shard(s): {formatted}.") + + provenance = sorted(provenance_by_name.values(), key=lambda job: job["input_index"]) + input_indices = [job["input_index"] for job in provenance] + if len(input_indices) != len(set(input_indices)): + raise TSValueError("Duplicate input positions across cluster shards.") + if input_indices != list(range(len(input_indices))): + raise TSValueError("Cluster shards do not cover every input position.") + results = [records_by_name[job["name"]] for job in provenance] + csv_path, json_path = _write_results(results, out) + _write_json( + _sidecar_path(out, "-run.json"), + { + "schema_version": 1, + "workflow": "thermochemistry_screen", + "source": expected_source, + "settings": expected_settings, + "jobs": provenance, + "collection": { + "shard_count": expected_count, + "shard_directory": str(shard_directory), + }, + }, + ) + logger.info(f"Collected {len(results)} molecules into {csv_path} and {json_path}") + return results + + +def _optional_float(value, default, field, row_number=None): + if value in (None, ""): + return default + try: + number = float(value) + except (TypeError, ValueError) as exc: + location = f" in row {row_number}" if row_number is not None else "" + raise TSValueError(f"{field}{location} must be numeric, got {value!r}.") from exc + if not math.isfinite(number): + location = f" in row {row_number}" if row_number is not None else "" + raise TSValueError(f"{field}{location} must be finite.") + return number + + +def _integer_charge(value, default, field, row_number=None): + number = _optional_float(value, default, field, row_number) + if number is None: + return None + if not float(number).is_integer(): + location = f" in row {row_number}" if row_number is not None else "" + raise TSValueError(f"{field}{location} must be an integer.") + return int(number) + + +def _valid_spin(value, default, field, row_number=None): + spin = _optional_float(value, default, field, row_number) + if spin is None: + return None + doubled = 2.0 * spin + if spin < 0 or not math.isclose(doubled, round(doubled), abs_tol=1e-9): + location = f" in row {row_number}" if row_number is not None else "" + raise TSValueError( + f"{field}{location} must be a non-negative integer or half-integer." + ) + return spin + + +def _validate_redox_name(name): + if not name or name in (".", "..") or Path(name).name != name: + raise TSValueError( + f"Invalid molecule name {name!r}; use a non-empty name without path separators." + ) + return name + + +def _redox_job_from_structure(path, name, charge, spins): + path = Path(path) + if not path.is_file(): + raise TSValueError(f"Structure file not found: '{path}'.") + if path.suffix.lower() not in _STRUCTURE_SUFFIXES: + raise TSValueError(f"Redox structures must be .xyz or .gen files, got '{path}'.") + charge = _integer_charge(charge, 0, "charge") + spins = tuple( + _valid_spin(spin, None, f"{state}_spin") + for (state, _offset, _index), spin in zip(_REDOX_STATES, spins) + ) + try: + atoms = ase.io.read(str(path)) + except Exception as exc: # pylint: disable=broad-except + raise TSValueError(f"Could not read redox structure '{path}': {exc}") from exc + nuclear_charge = int(sum(atoms.get_atomic_numbers())) + for (state, charge_offset, spin_index) in _REDOX_STATES: + spin = spins[spin_index] + if spin is None: + continue + electrons = nuclear_charge - (charge + charge_offset) + unpaired = int(round(2.0 * spin)) + if unpaired > electrons or (electrons - unpaired) % 2: + raise TSValueError( + f"{state}_spin={spin:g} is incompatible with {electrons} electrons " + f"for '{name}'." + ) + return _RedoxJob(_validate_redox_name(name), path.resolve(), charge, spins) + + +def _redox_job_from_smiles( + smiles, + name, + charge, + spins, + generated_directory, + max_conformers, +): + name = _validate_redox_name(name) + try: + conformers = generate_conformers(smiles, max_conformers=max_conformers) + except (ImportError, RuntimeError, ValueError) as exc: + raise TSValueError(f"Could not generate {name!r} from SMILES: {exc}") from exc + if not conformers: + raise TSValueError(f"Could not generate {name!r} from SMILES: no conformers returned.") + formal_charge = int(getattr(conformers[0], "info", {}).get("formal_charge", 0)) + charge = _integer_charge(charge, formal_charge, "charge") + if formal_charge and charge != formal_charge: + raise TSValueError( + f"charge={charge} conflicts with formal charge {formal_charge} in " + f"SMILES for {name!r}." + ) + path = write_conformers( + [conformers[0]], generated_directory, prefix=name + )[0] + return _redox_job_from_structure(path, name, charge, spins) + + +def _redox_jobs_from_manifest( + manifest, + charge, + spin, + generated_directory, + max_conformers, +): + jobs = [] + with open(manifest, newline="", encoding="utf-8") as handle: + reader = csv.DictReader(handle) + for row_number, row in enumerate(reader, start=2): + path_value = (row.get("path") or "").strip() + smiles = (row.get("smiles") or "").strip() + if bool(path_value) == bool(smiles): + raise TSValueError( + f"Row {row_number} of '{manifest}' needs exactly one of " + "'path' or 'smiles'." + ) + + fallback_name = Path(path_value).stem if path_value else f"molecule-{row_number - 1}" + name = _validate_redox_name((row.get("name") or fallback_name).strip()) + row_charge = _integer_charge( + row.get("charge"), charge, "charge", row_number + ) + oxidized_spin = _valid_spin( + row.get("oxidized_spin") or row.get("spin"), + spin, + "oxidized_spin", + row_number, + ) + spins = ( + oxidized_spin, + _valid_spin( + row.get("reduced_once_spin"), + None, + "reduced_once_spin", + row_number, + ), + _valid_spin( + row.get("reduced_twice_spin"), + None, + "reduced_twice_spin", + row_number, + ), + ) + + if path_value: + path = Path(path_value) + if not path.is_absolute(): + path = Path(manifest).parent / path + job = _redox_job_from_structure(path, name, row_charge, spins) + else: + job = _redox_job_from_smiles( + smiles, + name, + row_charge, + spins, + generated_directory, + max_conformers, + ) + jobs.append(job) + return jobs + + +def _load_redox_jobs( + source, + charge, + spin, + generated_directory, + max_conformers, + default_name="molecule", +): + source_path = Path(source) + try: + exists = source_path.exists() + except OSError: + exists = False + + if exists and source_path.is_dir(): + jobs = [ + _redox_job_from_structure(path, path.stem, charge, (spin, None, None)) + for path in sorted(source_path.iterdir()) + if path.suffix.lower() in _STRUCTURE_SUFFIXES + ] + elif exists and source_path.suffix.lower() == ".csv": + jobs = _redox_jobs_from_manifest( + source_path, + charge, + spin, + generated_directory, + max_conformers, + ) + elif exists: + jobs = [ + _redox_job_from_structure( + source_path, source_path.stem, charge, (spin, None, None) + ) + ] + else: + if isinstance(source, Path) or source_path.suffix.lower() in ( + ".csv", + *_STRUCTURE_SUFFIXES, + ): + raise TSValueError(f"Redox input not found: '{source}'.") + jobs = [ + _redox_job_from_smiles( + str(source), + default_name, + charge, + (spin, None, None), + generated_directory, + max_conformers, + ) + ] + + if not jobs: + raise TSValueError(f"No molecules found in redox input '{source}'.") + names = [job.name for job in jobs] + duplicates = sorted({name for name in names if names.count(name) > 1}) + if duplicates: + raise TSValueError( + f"Duplicate molecule name(s) in the redox screen: {', '.join(duplicates)}." + ) + return jobs + + +def _resolve_redox_reference( + reference, + candidates, + reference_charge, + generated_directory, + max_conformers, +): + for candidate in candidates: + if reference == candidate.name: + if ( + reference_charge is not None + and _integer_charge(reference_charge, None, "reference_charge") + != candidate.charge + ): + raise TSValueError( + f"reference_charge does not match candidate {candidate.name!r}." + ) + return candidate + + references = _load_redox_jobs( + reference, + reference_charge, + None, + generated_directory, + max_conformers, + default_name="reference", + ) + if len(references) != 1: + raise TSValueError("The redox reference must identify exactly one molecule.") + selected = references[0] + for candidate in candidates: + if ( + candidate.path == selected.path + and candidate.charge == selected.charge + and candidate.spins == selected.spins + ): + return candidate + return selected + + +def _write_redox_state_manifest(jobs, path): + mapping = {} + + def write(handle): + writer = csv.DictWriter(handle, fieldnames=("name", "path", "charge", "spin")) + writer.writeheader() + for job in jobs: + for state, charge_offset, spin_index in _REDOX_STATES: + state_name = f"{job.name}--{state}" + mapping[(job.name, state)] = state_name + writer.writerow( + { + "name": state_name, + "path": str(job.path), + "charge": job.charge + charge_offset, + "spin": "" if job.spins[spin_index] is None else job.spins[spin_index], + } + ) + + path.parent.mkdir(parents=True, exist_ok=True) + _atomic_write(path, write) + return mapping + + +def _redox_state_records(job, mapping, records): + state_records = {} + for state, _charge_offset, _spin_index in _REDOX_STATES: + name = mapping[(job.name, state)] + state_records[state] = records.get( + name, + {"status": "error", "error": "state calculation produced no result"}, + ) + return state_records + + +def _state_failure(states): + failures = [] + for state, record in states.items(): + if record.get("status") != "ok": + failures.append(f"{state}: {record.get('error') or 'calculation failed'}") + elif record.get("G_total_hartree") is None: + failures.append(f"{state}: total Gibbs free energy is missing") + else: + try: + finite = math.isfinite(float(record["G_total_hartree"])) + except (TypeError, ValueError): + finite = False + if not finite: + failures.append(f"{state}: total Gibbs free energy is not finite") + return "; ".join(failures) + + +def _absolute_stepwise_potentials(states): + oxidized = float(states["oxidized"]["G_total_hartree"]) + reduced_once = float(states["reduced_once"]["G_total_hartree"]) + reduced_twice = float(states["reduced_twice"]["G_total_hartree"]) + return ( + -(reduced_once - oxidized) * HARTREE_TO_EV, + -(reduced_twice - reduced_once) * HARTREE_TO_EV, + ) + + +def _write_redox_results(results, out): + stem = Path(str(out)).with_suffix("") + if stem.parent != Path(""): + stem.parent.mkdir(parents=True, exist_ok=True) + + def write_csv(handle): + writer = csv.DictWriter(handle, fieldnames=_REDOX_RESULT_FIELDS) + writer.writeheader() + for result in results: + writer.writerow( + {field: result.get(field, "") for field in _REDOX_RESULT_FIELDS} + ) + + csv_path = stem.with_suffix(".csv") + json_path = stem.with_suffix(".json") + _atomic_write(csv_path, write_csv) + _atomic_write(json_path, lambda handle: handle.write(json.dumps(results, indent=2))) + return csv_path, json_path + + +def _write_redox_run_metadata( + *, + out, + source, + candidates, + reference_job, + reference_e1, + reference_e2, + potential_scale, + settings, + max_conformers, +): + def describe(job): + if job is None: + return None + return { + "name": job.name, + "path": str(job.path), + "structure_sha256": _file_sha256(job.path), + "oxidized_charge": job.charge, + "spins": { + state: job.spins[index] + for state, _offset, index in _REDOX_STATES + }, + } + + payload = { + "schema_version": 1, + "workflow": "stepwise_reduction_screen", + "source": str(source), + "single_starting_geometry_approximation": True, + "smiles_embedding": { + "method": "ETKDGv3 followed by MMFF; lowest sampled conformer retained", + "max_conformers": max_conformers, + }, + "states": [ + {"name": state, "charge_offset": offset} + for state, offset, _index in _REDOX_STATES + ], + "settings": settings, + "candidates": [describe(job) for job in candidates], + "reference": { + "calculation": describe(reference_job), + "experimental_E1_V": reference_e1, + "experimental_E2_V": reference_e2, + "potential_scale": potential_scale, + }, + } + run_fingerprint = _payload_fingerprint(payload) + payload["run_fingerprint"] = run_fingerprint + _write_json(_sidecar_path(out, "-run.json"), payload) + return run_fingerprint + + +def redox_screen( + source, + out="redox-results", + charge=None, + temperature=298.15, + pressure=101325, + directory="redox-screening", + parameters=None, + spin=None, + parameter_set=None, + solvent=None, + dispersion=None, + quasi_rrho=False, + engine="dftb+", + method=None, + resume=False, + jobs=1, + reference=None, + reference_e1=None, + reference_e2=None, + reference_charge=None, + potential_scale=None, + max_conformers=20, +): + """ + Run a two-step molecular reduction screen from one starting geometry. + + The workflow prepares one lowest-MMFF starting conformer from SMILES when + needed, computes the oxidized, one-electron-reduced and two-electron-reduced + states, and reports the first, second and overall two-electron reduction + potentials. Each electronic state is optimized independently, but this is + not a charge-state conformer search or ensemble-free-energy calculation. A + reference molecule and measured E1/E2 may calibrate both steps. + """ + + calibration = (reference, reference_e1, reference_e2) + if any(value is not None for value in calibration) and not all( + value is not None for value in calibration + ): + raise TSValueError( + "reference, reference_e1 and reference_e2 must be supplied together." + ) + if potential_scale is not None and reference is None: + raise TSValueError("potential_scale requires reference calibration.") + if ( + isinstance(max_conformers, bool) + or not isinstance(max_conformers, int) + or max_conformers < 1 + ): + raise TSValueError("max_conformers must be an integer >= 1.") + if engine not in ("dftb+", "xtb", "xtb-cli"): + raise TSValueError( + f"Unknown engine {engine!r}; choose 'dftb+', 'xtb' or 'xtb-cli'." + ) + if engine == "xtb" and solvent is not None: + raise TSValueError("The tblite xTB engine does not support solvent.") + if engine != "dftb+" and dispersion is not None: + raise TSValueError("dispersion is supported only by the DFTB+ engine.") + if engine != "dftb+" and parameters is not None: + raise TSValueError("parameters are supported only by the DFTB+ engine.") + if engine == "dftb+" and method is not None: + raise TSValueError("method applies only to xTB engines.") + if engine != "dftb+" and parameter_set is not None: + raise TSValueError("parameter_set applies only to the DFTB+ engine.") + + resolved_parameter_set = parameter_set or "3ob" + resolved_method = method or "GFN2-xTB" + charge = _integer_charge(charge, None, "charge") + spin = _valid_spin(spin, None, "spin") + reference_charge = _integer_charge( + reference_charge, None, "reference_charge" + ) + reference_e1 = _optional_float(reference_e1, None, "reference_e1") + reference_e2 = _optional_float(reference_e2, None, "reference_e2") + root = Path(directory) + generated_directory = root / "inputs" + candidates = _load_redox_jobs( + source, + charge, + spin, + generated_directory, + max_conformers, + ) + + reference_job = None + calculation_jobs = list(candidates) + if reference is not None: + reference_job = _resolve_redox_reference( + reference, + candidates, + reference_charge, + generated_directory, + max_conformers, + ) + if reference_job not in calculation_jobs: + occupied_names = {job.name for job in calculation_jobs} + if reference_job.name in occupied_names: + index = 1 + reference_name = "reference" + while reference_name in occupied_names: + index += 1 + reference_name = f"reference-{index}" + reference_job = _RedoxJob( + reference_name, + reference_job.path, + reference_job.charge, + reference_job.spins, + ) + calculation_jobs.append(reference_job) + + scale = ( + "SHE" + if reference_job is None + else potential_scale or f"calibrated:{reference_job.name}" + ) + redox_settings = { + "thermoscreening_version": __version__, + "engine": engine, + "temperature_K": temperature, + "pressure_Pa": pressure, + "method": resolved_method if engine in ("xtb", "xtb-cli") else None, + "solvent": solvent, + "dispersion": dispersion if engine == "dftb+" else None, + "quasi_rrho": quasi_rrho, + "parameter_set": resolved_parameter_set if engine == "dftb+" else None, + "parameters": parameters if engine == "dftb+" else None, + } + run_fingerprint = _write_redox_run_metadata( + out=out, + source=source, + candidates=candidates, + reference_job=reference_job, + reference_e1=reference_e1, + reference_e2=reference_e2, + potential_scale=scale, + settings=redox_settings, + max_conformers=max_conformers, + ) + + state_manifest = root / "states.csv" + mapping = _write_redox_state_manifest(calculation_jobs, state_manifest) + state_out = f"{Path(str(out)).with_suffix('')}-states" + state_results = screen( + state_manifest, + out=state_out, + temperature=temperature, + pressure=pressure, + directory=root / "calculations", + parameters=parameters, + parameter_set=resolved_parameter_set, + solvent=solvent, + dispersion=dispersion, + quasi_rrho=quasi_rrho, + engine=engine, + method=resolved_method, + resume=resume, + jobs=jobs, + ) + records = {record["name"]: record for record in state_results} + + reference_error = "" + e1_reference = e2_reference = SHE_ABSOLUTE_POTENTIAL + if reference_job is not None: + reference_states = _redox_state_records(reference_job, mapping, records) + reference_error = _state_failure(reference_states) + if not reference_error: + reference_abs_e1, reference_abs_e2 = _absolute_stepwise_potentials( + reference_states + ) + e1_reference = reference_abs_e1 - reference_e1 + e2_reference = reference_abs_e2 - reference_e2 + + results = [] + for candidate in candidates: + states = _redox_state_records(candidate, mapping, records) + error = _state_failure(states) + if reference_error: + error = f"reference: {reference_error}" + (f"; {error}" if error else "") + result = { + "name": candidate.name, + "path": str(candidate.path), + "charge": candidate.charge, + "status": "error" if error else "ok", + "potential_scale": scale, + "run_fingerprint": run_fingerprint, + "error": error, + } + for state, field in ( + ("oxidized", "G_oxidized_hartree"), + ("reduced_once", "G_reduced_once_hartree"), + ("reduced_twice", "G_reduced_twice_hartree"), + ): + value = states[state].get("G_total_hartree") + try: + value = float(value) + except (TypeError, ValueError): + continue + if math.isfinite(value): + result[field] = value + + if not error: + absolute_e1, absolute_e2 = _absolute_stepwise_potentials(states) + e1 = absolute_e1 - e1_reference + e2 = absolute_e2 - e2_reference + result.update( + { + "E1_V": e1, + "E2_V": e2, + "E2e_V": (e1 + e2) / 2.0, + "potential_gap_V": e1 - e2, + "potential_inversion": e2 >= e1, + } + ) + results.append(result) + + csv_path, json_path = _write_redox_results(results, out) + logger.info(f"Wrote {csv_path}, {json_path} and the per-state results") + return results diff --git a/docs/api.rst b/docs/api.rst index 905226f..d82d80b 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -18,6 +18,8 @@ Screening --------- .. autofunction:: ThermoScreening.thermo.screening.screen +.. autofunction:: ThermoScreening.thermo.screening.collect_screen_shards +.. autofunction:: ThermoScreening.thermo.screening.redox_screen .. autofunction:: ThermoScreening.thermo.screening.rank_by_gibbs .. autoclass:: ThermoScreening.thermo.screening.ScreeningJob :members: diff --git a/docs/configuration.rst b/docs/configuration.rst index 3fe508a..c06ba03 100644 --- a/docs/configuration.rst +++ b/docs/configuration.rst @@ -64,4 +64,99 @@ Screening input Results are written to ``.csv`` and ``.json`` with the electronic energy (``Eelec_hartree``), the thermal corrections, the absolute Gibbs free -energy (``G_total_hartree``), entropy and heat capacity. +energy (``G_total_hartree``), entropy and heat capacity. ``-run.json`` +records the structure hashes and scientific settings used for safe resume. + +Cluster execution +----------------- + +Batch screening can be divided into deterministic, zero-based shards on any +scheduler. Every shard writes isolated result files under +``-shards/shard-NNNNN.*`` and calculations under +``/shards/shard-NNNNN``: + +.. code-block:: bash + + thermo screen molecules.csv -o results \ + --shard-index "$TASK_INDEX" --shard-count 32 + +After every shard has finished, validate and combine them in original input +order: + +.. code-block:: bash + + thermo collect results-shards -o results + +Collection fails on missing or duplicate shards, mixed settings, missing jobs, +or fingerprint mismatches. Scientific calculation failures remain in the +combined result and produce a non-zero command exit status. + +For Slurm, one command can generate an array script and optionally submit both +the array and a dependent collection job: + +.. code-block:: bash + + thermo slurm --tasks 32 --cpus-per-task 8 \ + --time 04:00:00 --mem 16G --partition compute --submit -- \ + screen molecules.csv -o results --engine dftb+ --jobs 2 + +``--cpus-per-task`` describes each allocation; ``--jobs`` controls independent +processes inside that allocation. The generated script divides available CPUs +between those processes through ``OMP_NUM_THREADS`` and common BLAS thread +variables. Use ``--preamble setup.sh`` for site-specific module loads or +environment variables. Omitting ``--submit`` only writes the script. + +All array tasks and the collector require the same shared working directory, +Python environment, executables, and parameter files. The shard interface also +works with PBS, LSF, and other schedulers by supplying their array index and the +chosen shard count directly. Cluster arrays currently distribute +``thermo screen``; ``thermo redox`` retains local ``--jobs`` parallelism. + +Redox workflow input +-------------------- + +``thermo redox`` accepts a SMILES string, one ``.xyz``/``.gen`` structure, a +directory of structures, or a CSV manifest. Each manifest row needs exactly one +of ``path`` or ``smiles``: + +.. code-block:: text + + name,path,smiles,charge,oxidized_spin,reduced_once_spin,reduced_twice_spin + AQ,aq.xyz,,0,0,0.5,0 + hydroxy-AQ,,OC1=CC2=C(C=C1)C(=O)C1=CC=CC(=O)C1=C2,0,0,0.5,0 + +``name`` and ``charge`` are optional. Spin columns are optional and otherwise +inferred from electron count for each state. Relative paths are resolved from +the manifest directory. Molecular charges must be integers and spins must be +non-negative integers or half-integers. Ionic SMILES supply their formal charge; +structure files without a charge default to zero. + +For SMILES input, ``--max-conformers`` controls how many ETKDG structures are +embedded. The lowest-MMFF-energy structure becomes the common starting geometry; +each charge state is then optimized independently by the selected backend. +This is a single-starting-geometry screening approximation, not a charge-state +conformer search or a conformational ensemble free energy. Flexible molecules +require separate conformer sampling for every charge state. + +The workflow writes: + +- ``.csv`` and ``.json`` with E1, E2, E2e, the potential gap, and the + three absolute Gibbs free energies; +- ``-states.csv`` and ``-states.json`` with the complete per-state + thermochemistry; and +- ``-run.json`` and ``-states-run.json`` with input hashes, + calculation settings, reference calibration, and reproducibility fingerprints; + and +- ``/states.csv`` as the generated, reproducible charge-state + manifest. + +``--jobs`` is the number of independent charge-state calculations executed in +parallel. A screen of *N* molecules contains ``3 * N`` jobs, plus three only when +a separately supplied reference is not already part of the candidate set. +``--resume`` reuses only successful states whose structure content, charge, +spin, engine, method, solvent, thermodynamic conditions, and other scientific +settings match the stored fingerprint. Missing, failed, legacy, or mismatched +states are rerun. On multi-core backends, set ``OMP_NUM_THREADS=1`` when using +several process jobs to avoid oversubscribing CPU cores. Launch process-based +parallel runs from the CLI or a script protected by ``if __name__ == "__main__"``; +use ``--jobs 1`` in notebooks and other interactive sessions. diff --git a/docs/index.rst b/docs/index.rst index 2b59452..2397d7d 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -73,8 +73,8 @@ Features :link: usage :link-type: doc - Reaction free energies and one-electron reduction potentials - (vs. SHE or a calibrated reference) from any two ``Thermo`` objects. + Reaction free energies, one-electron reduction potentials, and a parallel + three-state workflow with physical reference calibration and provenance. .. grid-item-card:: :octicon:`flame` Kinetics :link: usage @@ -87,9 +87,9 @@ Features :link: usage :link-type: doc - Screen a directory or CSV manifest in parallel (``--jobs``), with - per-molecule error isolation, ``--resume``, and ranking by Gibbs - free energy. + Screen a directory or CSV manifest with local processes or HPC job arrays, + per-molecule error isolation, validated collection, ``--resume``, and + ranking by Gibbs free energy. .. grid-item-card:: :octicon:`versions` Conformers & ensembles :link: usage diff --git a/docs/usage.rst b/docs/usage.rst index eb58215..744dc74 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -4,7 +4,8 @@ Usage Command line ------------ -The ``thermo`` command has four subcommands. +The ``thermo`` command provides calculation, setup, and batch-management +subcommands. ``thermo doctor`` Report which backends (dftb+, modes, DFTB_PREFIX + a Slater-Koster file, and @@ -41,6 +42,55 @@ The ``thermo`` command has four subcommands. # run 4 molecules at a time (each in its own process/directory) thermo screen molecules/ --jobs 4 + # submit 32 Slurm shards with 2 local workers per allocation + thermo slurm --tasks 32 --cpus-per-task 8 --submit -- \ + screen molecules.csv -o results --jobs 2 + + # scheduler-neutral collection after all shards finish + thermo collect results-shards -o results + +``thermo redox`` + Run input preparation, three successive charge states, thermochemistry, and + redox post-processing from one common starting geometry: + + .. code-block:: bash + + # A single molecule from SMILES; three states can run concurrently + thermo redox "O=C1C=CC(=O)C=C1" \ + --engine xtb-cli --solvent acetonitrile --jobs 3 + + # A directory or mixed path/SMILES manifest + thermo redox molecules.csv -o aq-screen \ + --parameter-set 3ob --solvent acetonitrile --quasi-rrho --jobs 12 + + Without an experimental reference, potentials are reported versus SHE using + the 4.44 V absolute convention. Semiempirical absolute redox potentials are + generally not quantitative. Calibrate both steps against one consistently + computed and measured reference molecule: + + .. code-block:: bash + + thermo redox molecules.csv -o aq-screen \ + --reference AQ --reference-e1 -0.75 --reference-e2 -1.40 \ + --potential-scale "chosen experimental scale" \ + --parameter-set 3ob --solvent acetonitrile --jobs 12 + + ``--reference`` may be a candidate name, a structure path, or a SMILES. When + it names a candidate, its existing three calculations are reused. E1 and E2 + receive separate physical reference calibrations; E2e is their arithmetic + mean. This is not a learned correction model. + + For SMILES, the lowest of the sampled MMFF conformers is retained as the + common starting geometry and each charge state is optimized independently. + This is deliberately a single-starting-geometry screen. It does not replace + charge-state-specific conformer searches or ensemble free energies for + flexible molecules. The ``*-run.json`` files record input hashes, settings, + calibration values, and the fingerprints used by safe ``--resume``. + + A result with ``potential_inversion=true`` has ``E2 >= E1``. Treat it as a + validation target: inspect conformers, minima, spin, and charge localization + before interpreting it as a merged two-electron wave. + Python API ---------- diff --git a/tests/calculator/test_xtb.py b/tests/calculator/test_xtb.py index 8a7907e..993da6f 100644 --- a/tests/calculator/test_xtb.py +++ b/tests/calculator/test_xtb.py @@ -44,6 +44,42 @@ def test_optimise_and_frequencies_with_emt(monkeypatch, tmp_path): assert frequencies[-1] > 0 +def test_optimise_and_frequencies_cleans_vibration_cache_on_failure( + monkeypatch, tmp_path +): + from ase import Atoms + from ase.calculators.emt import EMT + + clean_calls = [] + + class FakeOptimizer: + def __init__(self, atoms, logfile=None): + self.atoms = atoms + + def run(self, fmax): + return None + + class FakeVibrations: + def __init__(self, atoms, name): + self.name = name + + def clean(self): + clean_calls.append(self.name) + + def run(self): + raise RuntimeError("interrupted") + + monkeypatch.chdir(tmp_path) + monkeypatch.setattr("ase.optimize.BFGS", FakeOptimizer) + monkeypatch.setattr("ase.vibrations.Vibrations", FakeVibrations) + + atoms = Atoms("Cu", positions=[[0.0, 0.0, 0.0]]) + with pytest.raises(RuntimeError, match="interrupted"): + optimise_and_frequencies(atoms, EMT()) + + assert clean_calls == ["xtb_vib", "xtb_vib"] + + def test_xtb_thermo_has_no_solvation_parameter(): # regression guard: xTB-side implicit solvation is intentionally not wired up import inspect @@ -52,6 +88,36 @@ def test_xtb_thermo_has_no_solvation_parameter(): assert "solvent" not in inspect.signature(xtb_thermo).parameters +def test_xtb_thermo_passes_charge_and_unpaired_electrons_to_tblite( + monkeypatch, tmp_path +): + from ase import Atoms + import ThermoScreening.thermo.api as api + + captured = {} + sentinel = object() + + def fake_optimise(atoms, calculator, fmax): + captured["charge"] = atoms.get_initial_charges().sum() + captured["unpaired"] = atoms.get_initial_magnetic_moments().sum() + return atoms, -1.0, np.array([100.0, 200.0, 300.0]) + + monkeypatch.setattr(api, "optimise_and_frequencies", fake_optimise) + monkeypatch.setattr(api, "xtb_calculator", lambda method: object()) + monkeypatch.setattr(api, "run_thermo", lambda *args, **kwargs: sentinel) + + result = api.xtb_thermo( + Atoms("H2", positions=[[0, 0, 0], [0, 0, 0.75]]), + charge=-1, + spin=0.5, + directory=tmp_path, + ) + + assert result is sentinel + assert captured["charge"] == pytest.approx(-1.0) + assert captured["unpaired"] == pytest.approx(1.0) + + tblite_available = importlib.util.find_spec("tblite") is not None xtb_skip = pytest.mark.skipif(not tblite_available, reason="tblite (GFN-xTB) is not installed.") diff --git a/tests/cli/test_slurm.py b/tests/cli/test_slurm.py new file mode 100644 index 0000000..0dd571b --- /dev/null +++ b/tests/cli/test_slurm.py @@ -0,0 +1,172 @@ +from argparse import Namespace +from types import SimpleNamespace + +import pytest + +from ThermoScreening.cli import slurm +from ThermoScreening.exceptions import TSValueError + + +def test_write_slurm_array_script_sets_resources_and_shards(tmp_path): + preamble = tmp_path / "preamble.sh" + preamble.write_text("module load dftbplus\nsource env/bin/activate\n", encoding="utf-8") + + script = slurm.write_slurm_array_script( + ["screen", "molecules.csv", "-o", "results", "--jobs", "2"], + tasks=8, + script=tmp_path / "screen.slurm", + local_jobs=2, + cpus_per_task=8, + job_name="aq-screen", + walltime="04:00:00", + memory="16G", + partition="compute", + account="chemistry", + preamble=preamble, + working_directory=tmp_path, + python_executable="/shared/python", + ) + + text = script.read_text(encoding="utf-8") + assert "#SBATCH --array=0-7" in text + assert "#SBATCH --cpus-per-task=8" in text + assert "#SBATCH --time=04:00:00" in text + assert "#SBATCH --mem=16G" in text + assert "export OMP_NUM_THREADS=4" in text + assert "module load dftbplus" in text + assert "/shared/python -m ThermoScreening screen molecules.csv" in text + assert '--shard-index "$SLURM_ARRAY_TASK_ID"' in text + assert '--shard-count "$SLURM_ARRAY_TASK_COUNT"' in text + assert script.stat().st_mode & 0o111 + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"tasks": 0}, "tasks must be"), + ({"tasks": 2, "local_jobs": 3, "cpus_per_task": 2}, "cannot exceed"), + ({"tasks": 2, "job_name": "bad\nname"}, "unsupported characters"), + ], +) +def test_write_slurm_array_script_rejects_invalid_resources(tmp_path, kwargs, message): + options = {"tasks": 2, "script": tmp_path / "screen.slurm", **kwargs} + with pytest.raises(TSValueError, match=message): + slurm.write_slurm_array_script(["screen", "molecules.csv"], **options) + + +def test_write_slurm_array_script_rejects_manual_shard_options(tmp_path): + with pytest.raises(TSValueError, match="Do not pass shard options"): + slurm.write_slurm_array_script( + ["screen", "molecules.csv", "--shard-index", "0"], + tasks=2, + script=tmp_path / "screen.slurm", + ) + + +def test_submit_slurm_array_schedules_dependent_collector(monkeypatch, tmp_path): + calls = [] + responses = iter( + [SimpleNamespace(stdout="12345;cluster\n"), SimpleNamespace(stdout="12346\n")] + ) + + monkeypatch.setattr(slurm.shutil, "which", lambda command: "/usr/bin/sbatch") + + def fake_run(args, **kwargs): + calls.append((args, kwargs)) + return next(responses) + + monkeypatch.setattr(slurm.subprocess, "run", fake_run) + script = tmp_path / "screen.slurm" + script.write_text("#!/bin/bash\n", encoding="utf-8") + + result = slurm.submit_slurm_array( + script, + shard_directory=tmp_path / "results-shards", + out=tmp_path / "results", + job_name="screen", + partition="compute", + account="chemistry", + working_directory=tmp_path, + python_executable="/shared/python", + ) + + assert result == ("12345", "12346") + assert calls[0][0] == ["/usr/bin/sbatch", "--parsable", str(script.resolve())] + collector = calls[1][0] + assert "--dependency=afterany:12345" in collector + assert "--partition=compute" in collector + assert "--account=chemistry" in collector + assert any("ThermoScreening collect" in argument for argument in collector) + + +def test_submit_slurm_array_requires_sbatch(monkeypatch, tmp_path): + monkeypatch.setattr(slurm.shutil, "which", lambda command: None) + with pytest.raises(TSValueError, match="sbatch was not found"): + slurm.submit_slurm_array( + tmp_path / "screen.slurm", + shard_directory=tmp_path / "shards", + out=tmp_path / "results", + ) + + +def test_cli_generates_slurm_script(tmp_path, capsys): + import ThermoScreening.cli.thermo as cli + + args = Namespace( + tasks=4, + script=str(tmp_path / "screen.slurm"), + job_name="screen", + cpus_per_task=2, + walltime=None, + memory=None, + partition=None, + account=None, + preamble=None, + submit=False, + command_args=["--", "screen", "molecules.csv", "-o", "results"], + ) + + assert cli.run_slurm(args) == 0 + output = capsys.readouterr().out + assert "Slurm script:" in output + assert "thermo collect results-shards -o results" in output + + +def test_cli_collect_reports_failures(monkeypatch, capsys): + import ThermoScreening.cli.thermo as cli + + monkeypatch.setattr( + cli, + "collect_screen_shards", + lambda *args, **kwargs: [ + {"name": "ok", "status": "ok"}, + {"name": "bad", "status": "error", "error": "failed"}, + ], + ) + args = Namespace(shard_directory="results-shards", out="results") + + assert cli.run_collect(args) == 1 + assert "Collected 2 molecules (1 failed)" in capsys.readouterr().out + + +def test_slurm_parser_accepts_nested_screen_command(): + import ThermoScreening.cli.thermo as cli + + args = cli.parse_args( + [ + "slurm", + "--tasks", + "8", + "--cpus-per-task", + "4", + "--", + "screen", + "molecules.csv", + "--jobs", + "2", + ] + ) + + assert args.command == "slurm" + assert args.tasks == 8 + assert args.command_args[-3:] == ["molecules.csv", "--jobs", "2"] diff --git a/tests/thermo/test_api.py b/tests/thermo/test_api.py index 29fb58b..badba10 100644 --- a/tests/thermo/test_api.py +++ b/tests/thermo/test_api.py @@ -570,7 +570,8 @@ def fake_calculator(method): return "xtb-calc" def fake_optimise(atoms, calc, fmax=0.01): - seen["info"] = dict(atoms.info) + seen["charge"] = atoms.get_initial_charges().sum() + seen["unpaired"] = atoms.get_initial_magnetic_moments().sum() seen["calc"] = calc return "optimized-atoms", -5.0, np.array([1500.0, 3600.0, 3700.0]) @@ -597,8 +598,8 @@ def fake_run_thermo(frequencies, atoms=None, engine=None, spin=None, assert result == "thermo-result" assert seen["engine"] == "xtb" assert seen["spin"] == 0.5 # auto electron-count guess - assert seen["info"]["spin"] == 1 # unpaired electrons passed to xTB - assert seen["info"]["charge"] == 0 + assert seen["unpaired"] == pytest.approx(1.0) + assert seen["charge"] == pytest.approx(0.0) assert seen["method"] == "GFN1-xTB" assert seen["quasi_rrho"] is True assert seen["energy"] == -5.0 diff --git a/tests/thermo/test_conformers.py b/tests/thermo/test_conformers.py index a1d412b..64ac639 100644 --- a/tests/thermo/test_conformers.py +++ b/tests/thermo/test_conformers.py @@ -30,6 +30,12 @@ def test_generate_returns_ase_conformers(): assert all(atoms.get_chemical_formula() == "C4H10" for atoms in result) +def test_generate_preserves_smiles_formal_charge(): + result = conformers.generate("[NH4+]", max_conformers=1) + + assert result[0].info["formal_charge"] == 1 + + def test_generate_respects_max_conformers(): result = conformers.generate("CCCCCCCC", max_conformers=3) assert len(result) <= 3 diff --git a/tests/thermo/test_redox_workflow.py b/tests/thermo/test_redox_workflow.py new file mode 100644 index 0000000..0bcbefc --- /dev/null +++ b/tests/thermo/test_redox_workflow.py @@ -0,0 +1,737 @@ +import csv +import json +from argparse import Namespace +from concurrent.futures import Future + +import pytest + +from ThermoScreening.exceptions import TSValueError +from ThermoScreening.thermo import screening +from ThermoScreening.thermo._units import HARTREE_TO_EV +from ThermoScreening.thermo.reactions import SHE_ABSOLUTE_POTENTIAL + + +class _SyncExecutor: + def __init__(self, max_workers=None): + self.max_workers = max_workers + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def submit(self, function, *args, **kwargs): + future = Future() + future.set_result(function(*args, **kwargs)) + return future + + +class _ChargeThermo: + def __init__(self, charge): + self.charge = charge + + def electronic_energy(self): + return -100.0 + self.charge + + def total_energy(self, unit): + return -10.0 + + def total_enthalpy(self, unit): + return -9.0 + + def total_gibbs_free_energy(self, unit): + return -11.0 + + def total_EeGtot(self): + return -100.0 + 0.1 * self.charge + + def total_entropy(self, unit): + return 50.0 + + def total_heat_capacity(self, unit): + return 6.0 + + +def _write_xyz(path): + path.write_text("1\n\nH 0.0 0.0 0.0\n", encoding="utf-8") + + +def _fake_state_screen(energies, captured=None, failed=None): + failed = failed or {} + + def fake_screen(source, **kwargs): + if captured is not None: + captured.update(source=source, kwargs=kwargs) + with open(source, newline="", encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + results = [] + for row in rows: + molecule, state = row["name"].rsplit("--", 1) + error = failed.get((molecule, state)) + if error: + results.append({"name": row["name"], "status": "error", "error": error}) + else: + results.append( + { + "name": row["name"], + "status": "ok", + "G_total_hartree": energies[molecule][state], + } + ) + return results + + return fake_screen + + +def test_redox_screen_calculates_all_states_in_parallel(monkeypatch, tmp_path): + _write_xyz(tmp_path / "a.xyz") + _write_xyz(tmp_path / "b.xyz") + charges = [] + + def fake_thermo(atoms, charge=0.0, **kwargs): + charges.append(charge) + return _ChargeThermo(charge) + + monkeypatch.setattr(screening, "dftbplus_thermo", fake_thermo) + monkeypatch.setattr(screening, "ProcessPoolExecutor", _SyncExecutor) + + results = screening.redox_screen( + tmp_path, + out=tmp_path / "redox", + directory=tmp_path / "work", + jobs=4, + ) + + assert [result["name"] for result in results] == ["a", "b"] + assert sorted(charges) == [-2.0, -2.0, -1.0, -1.0, 0.0, 0.0] + assert all(result["status"] == "ok" for result in results) + assert all(result["potential_scale"] == "SHE" for result in results) + assert results[0]["E1_V"] == pytest.approx(0.1 * HARTREE_TO_EV - SHE_ABSOLUTE_POTENTIAL) + assert (tmp_path / "redox-states.csv").is_file() + assert (tmp_path / "redox.csv").is_file() + + +def test_redox_screen_resumes_all_completed_states(monkeypatch, tmp_path): + _write_xyz(tmp_path / "molecule.xyz") + calls = [] + + def fake_thermo(atoms, charge=0.0, **kwargs): + calls.append(charge) + return _ChargeThermo(charge) + + monkeypatch.setattr(screening, "dftbplus_thermo", fake_thermo) + arguments = { + "source": tmp_path / "molecule.xyz", + "out": tmp_path / "redox", + "directory": tmp_path / "work", + } + + first = screening.redox_screen(**arguments) + second = screening.redox_screen(**arguments, resume=True) + + assert calls == [0.0, -1.0, -2.0] + assert first == second + + +def test_redox_screen_calibrates_each_reduction_separately(monkeypatch, tmp_path): + _write_xyz(tmp_path / "target.xyz") + _write_xyz(tmp_path / "reference.xyz") + energies = { + "target": { + "oxidized": -200.0, + "reduced_once": -200.11, + "reduced_twice": -200.20, + }, + "reference": { + "oxidized": -100.0, + "reduced_once": -100.10, + "reduced_twice": -100.18, + }, + } + monkeypatch.setattr(screening, "screen", _fake_state_screen(energies)) + + result = screening.redox_screen( + tmp_path / "target.xyz", + out=tmp_path / "out", + directory=tmp_path / "work", + reference=tmp_path / "reference.xyz", + reference_e1=-0.75, + reference_e2=-1.40, + potential_scale="Fc/Fc+", + )[0] + + assert result["E1_V"] == pytest.approx(-0.75 + 0.01 * HARTREE_TO_EV) + assert result["E2_V"] == pytest.approx(-1.40 + 0.01 * HARTREE_TO_EV) + assert result["E2e_V"] == pytest.approx((result["E1_V"] + result["E2_V"]) / 2) + assert result["potential_scale"] == "Fc/Fc+" + + +def test_redox_screen_reuses_reference_from_candidate_set(monkeypatch, tmp_path): + _write_xyz(tmp_path / "target.xyz") + _write_xyz(tmp_path / "reference.xyz") + energies = { + "target": { + "oxidized": -200.0, + "reduced_once": -200.11, + "reduced_twice": -200.20, + }, + "reference": { + "oxidized": -100.0, + "reduced_once": -100.10, + "reduced_twice": -100.18, + }, + } + captured = {} + monkeypatch.setattr( + screening, "screen", _fake_state_screen(energies, captured=captured) + ) + + results = screening.redox_screen( + tmp_path, + out=tmp_path / "out", + directory=tmp_path / "work", + reference="reference", + reference_e1=-0.75, + reference_e2=-1.40, + ) + + with open(captured["source"], newline="", encoding="utf-8") as handle: + state_rows = list(csv.DictReader(handle)) + assert len(state_rows) == 6 + reference = next(result for result in results if result["name"] == "reference") + assert reference["E1_V"] == pytest.approx(-0.75) + assert reference["E2_V"] == pytest.approx(-1.40) + + +def test_redox_screen_rejects_candidate_reference_charge_mismatch(tmp_path): + _write_xyz(tmp_path / "reference.xyz") + + with pytest.raises(TSValueError, match="does not match candidate"): + screening.redox_screen( + tmp_path, + out=tmp_path / "out", + directory=tmp_path / "work", + reference="reference", + reference_e1=-0.75, + reference_e2=-1.40, + reference_charge=1, + ) + + +def test_redox_screen_reuses_reference_from_same_structure_path(monkeypatch, tmp_path): + _write_xyz(tmp_path / "reference.xyz") + energies = { + "reference": { + "oxidized": -100.0, + "reduced_once": -100.10, + "reduced_twice": -100.18, + } + } + captured = {} + monkeypatch.setattr( + screening, "screen", _fake_state_screen(energies, captured=captured) + ) + + result = screening.redox_screen( + tmp_path / "reference.xyz", + out=tmp_path / "out", + directory=tmp_path / "work", + reference=tmp_path / "reference.xyz", + reference_e1=-0.75, + reference_e2=-1.40, + )[0] + + with open(captured["source"], newline="", encoding="utf-8") as handle: + assert len(list(csv.DictReader(handle))) == 3 + assert result["E1_V"] == pytest.approx(-0.75) + + +def test_redox_screen_renames_a_distinct_reference_that_collides_with_candidate( + monkeypatch, tmp_path +): + candidate_directory = tmp_path / "candidates" + reference_directory = tmp_path / "references" + candidate_directory.mkdir() + reference_directory.mkdir() + _write_xyz(candidate_directory / "same.xyz") + _write_xyz(reference_directory / "same.xyz") + energies = { + "same": { + "oxidized": -200.0, + "reduced_once": -200.11, + "reduced_twice": -200.20, + }, + "reference": { + "oxidized": -100.0, + "reduced_once": -100.10, + "reduced_twice": -100.18, + }, + } + captured = {} + monkeypatch.setattr( + screening, "screen", _fake_state_screen(energies, captured=captured) + ) + + result = screening.redox_screen( + candidate_directory / "same.xyz", + out=tmp_path / "out", + directory=tmp_path / "work", + reference=reference_directory / "same.xyz", + reference_e1=-0.75, + reference_e2=-1.40, + )[0] + + with open(captured["source"], newline="", encoding="utf-8") as handle: + names = [row["name"] for row in csv.DictReader(handle)] + assert names == [ + "same--oxidized", + "same--reduced_once", + "same--reduced_twice", + "reference--oxidized", + "reference--reduced_once", + "reference--reduced_twice", + ] + assert result["status"] == "ok" + + +def test_redox_screen_propagates_state_and_reference_failures(monkeypatch, tmp_path): + _write_xyz(tmp_path / "target.xyz") + _write_xyz(tmp_path / "reference.xyz") + energies = { + "target": { + "oxidized": -200.0, + "reduced_once": -200.1, + "reduced_twice": -200.2, + }, + "reference": { + "oxidized": -100.0, + "reduced_once": -100.1, + "reduced_twice": -100.2, + }, + } + monkeypatch.setattr( + screening, + "screen", + _fake_state_screen( + energies, failed={("reference", "reduced_once"): "no minimum"} + ), + ) + + result = screening.redox_screen( + tmp_path / "target.xyz", + out=tmp_path / "out", + directory=tmp_path / "work", + reference=tmp_path / "reference.xyz", + reference_e1=-0.75, + reference_e2=-1.40, + )[0] + + assert result["status"] == "error" + assert "reference: reduced_once: no minimum" in result["error"] + assert "E1_V" not in result + + +def test_redox_screen_reports_missing_state_results_and_energies(monkeypatch, tmp_path): + _write_xyz(tmp_path / "molecule.xyz") + + def incomplete_screen(source, **kwargs): + return [ + { + "name": "molecule--oxidized", + "status": "ok", + "G_total_hartree": -100.0, + }, + {"name": "molecule--reduced_once", "status": "ok"}, + ] + + monkeypatch.setattr(screening, "screen", incomplete_screen) + + result = screening.redox_screen( + tmp_path / "molecule.xyz", + out=tmp_path / "out", + directory=tmp_path / "work", + )[0] + + assert result["status"] == "error" + assert "reduced_once: total Gibbs free energy is missing" in result["error"] + assert "reduced_twice: state calculation produced no result" in result["error"] + + +@pytest.mark.parametrize("invalid_energy", [float("nan"), float("inf"), "invalid"]) +def test_redox_screen_rejects_nonfinite_state_energies( + monkeypatch, tmp_path, invalid_energy +): + _write_xyz(tmp_path / "molecule.xyz") + energies = { + "molecule": { + "oxidized": -100.0, + "reduced_once": invalid_energy, + "reduced_twice": -100.2, + } + } + monkeypatch.setattr(screening, "screen", _fake_state_screen(energies)) + + result = screening.redox_screen( + tmp_path / "molecule.xyz", + out=tmp_path / "out", + directory=tmp_path / "work", + )[0] + + assert result["status"] == "error" + assert "reduced_once: total Gibbs free energy is not finite" in result["error"] + assert "G_reduced_once_hartree" not in result + + +def test_redox_screen_flags_potential_inversion(monkeypatch, tmp_path): + _write_xyz(tmp_path / "molecule.xyz") + energies = { + "molecule": { + "oxidized": -100.0, + "reduced_once": -100.05, + "reduced_twice": -100.12, + } + } + monkeypatch.setattr(screening, "screen", _fake_state_screen(energies)) + + result = screening.redox_screen( + tmp_path / "molecule.xyz", + out=tmp_path / "out", + directory=tmp_path / "work", + )[0] + + assert result["potential_inversion"] is True + assert result["potential_gap_V"] < 0 + + +def test_redox_screen_accepts_smiles_and_writes_generated_input(monkeypatch, tmp_path): + energies = { + "molecule": { + "oxidized": -100.0, + "reduced_once": -100.1, + "reduced_twice": -100.2, + } + } + monkeypatch.setattr(screening, "screen", _fake_state_screen(energies)) + + result = screening.redox_screen( + "CCO", + out=tmp_path / "out", + directory=tmp_path / "work", + max_conformers=2, + )[0] + + assert result["status"] == "ok" + assert (tmp_path / "work" / "inputs" / "molecule_0.xyz").is_file() + + +def test_redox_screen_infers_formal_charge_from_ionic_smiles(monkeypatch, tmp_path): + from ase import Atoms + + atoms = Atoms("H", positions=[[0.0, 0.0, 0.0]]) + atoms.info["formal_charge"] = -1 + energies = { + "molecule": { + "oxidized": -100.0, + "reduced_once": -100.1, + "reduced_twice": -100.2, + } + } + captured = {} + monkeypatch.setattr(screening, "generate_conformers", lambda *args, **kwargs: [atoms]) + monkeypatch.setattr( + screening, "screen", _fake_state_screen(energies, captured=captured) + ) + + result = screening.redox_screen( + "[H-]", + out=tmp_path / "out", + directory=tmp_path / "work", + )[0] + + with open(captured["source"], newline="", encoding="utf-8") as handle: + charges = [int(row["charge"]) for row in csv.DictReader(handle)] + assert result["charge"] == -1 + assert charges == [-1, -2, -3] + + +def test_redox_manifest_supports_paths_smiles_charges_and_spins(monkeypatch, tmp_path): + _write_xyz(tmp_path / "first.xyz") + manifest = tmp_path / "molecules.csv" + manifest.write_text( + "name,path,smiles,charge,spin,oxidized_spin,reduced_once_spin," + "reduced_twice_spin\n" + "first,first.xyz,,-1,0,,0.5,0\n" + "second,,CCO,0,,1,1.5,2\n", + encoding="utf-8", + ) + energies = { + "first": { + "oxidized": -100.0, + "reduced_once": -100.1, + "reduced_twice": -100.2, + }, + "second": { + "oxidized": -200.0, + "reduced_once": -200.1, + "reduced_twice": -200.2, + }, + } + captured = {} + monkeypatch.setattr( + screening, "screen", _fake_state_screen(energies, captured=captured) + ) + + results = screening.redox_screen( + manifest, + out=tmp_path / "out", + directory=tmp_path / "work", + max_conformers=2, + jobs=3, + ) + + with open(captured["source"], newline="", encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + assert [result["charge"] for result in results] == [-1, 0] + assert [row["charge"] for row in rows[:3]] == ["-1", "-2", "-3"] + assert [row["spin"] for row in rows[:3]] == ["0.0", "0.5", "0.0"] + assert captured["kwargs"]["jobs"] == 3 + + +@pytest.mark.parametrize( + ("manifest_text", "message"), + [ + ("name,path,smiles\na,,\n", "exactly one"), + ("name,path,smiles\na,a.xyz,CCO\n", "exactly one"), + ("name,path,smiles\na,,CCO\na,,CCC\n", "Duplicate molecule"), + ("name,path,charge\na,a.xyz,bad\n", "charge in row 2 must be numeric"), + ], +) +def test_redox_manifest_rejects_ambiguous_or_invalid_rows( + monkeypatch, tmp_path, manifest_text, message +): + _write_xyz(tmp_path / "a.xyz") + manifest = tmp_path / "molecules.csv" + manifest.write_text(manifest_text, encoding="utf-8") + monkeypatch.setattr(screening, "generate_conformers", lambda *args, **kwargs: [object()]) + monkeypatch.setattr( + screening, + "write_conformers", + lambda conformers, directory, prefix: [tmp_path / f"{prefix}.xyz"], + ) + + with pytest.raises(TSValueError, match=message): + screening.redox_screen( + manifest, + out=tmp_path / "out", + directory=tmp_path / "work", + ) + + +def test_redox_screen_rejects_missing_empty_and_unsupported_inputs(tmp_path): + empty = tmp_path / "empty" + empty.mkdir() + unsupported = tmp_path / "molecule.txt" + unsupported.write_text("not a structure", encoding="utf-8") + invalid_name = tmp_path / "invalid-name.csv" + invalid_name.write_text("name,path\n../escape,molecule.xyz\n", encoding="utf-8") + _write_xyz(tmp_path / "molecule.xyz") + + with pytest.raises(TSValueError, match="input not found"): + screening.redox_screen( + tmp_path / "missing.xyz", out=tmp_path / "out", directory=tmp_path / "work" + ) + with pytest.raises(TSValueError, match="No molecules found"): + screening.redox_screen(empty, out=tmp_path / "out", directory=tmp_path / "work") + with pytest.raises(TSValueError, match="must be .xyz or .gen"): + screening.redox_screen( + unsupported, out=tmp_path / "out", directory=tmp_path / "work" + ) + with pytest.raises(TSValueError, match="Invalid molecule name"): + screening.redox_screen( + invalid_name, out=tmp_path / "out", directory=tmp_path / "work" + ) + + +def test_redox_screen_requires_one_reference_molecule(tmp_path): + _write_xyz(tmp_path / "target.xyz") + references = tmp_path / "references" + references.mkdir() + _write_xyz(references / "one.xyz") + _write_xyz(references / "two.xyz") + + with pytest.raises(TSValueError, match="exactly one molecule"): + screening.redox_screen( + tmp_path / "target.xyz", + out=tmp_path / "out", + directory=tmp_path / "work", + reference=references, + reference_e1=-0.75, + reference_e2=-1.40, + ) + + +def test_redox_screen_reports_smiles_generation_failure(monkeypatch, tmp_path): + monkeypatch.setattr( + screening, + "generate_conformers", + lambda *args, **kwargs: (_ for _ in ()).throw(ValueError("embedding failed")), + ) + + with pytest.raises(TSValueError, match="Could not generate.*embedding failed"): + screening.redox_screen( + "invalid-smiles", + out=tmp_path / "out", + directory=tmp_path / "work", + ) + + +def test_redox_screen_reports_empty_smiles_generation(monkeypatch, tmp_path): + monkeypatch.setattr(screening, "generate_conformers", lambda *args, **kwargs: []) + + with pytest.raises(TSValueError, match="no conformers returned"): + screening.redox_screen( + "CCO", + out=tmp_path / "out", + directory=tmp_path / "work", + ) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"reference": "CCO"}, "supplied together"), + ({"reference_e1": -0.7}, "supplied together"), + ({"max_conformers": 0}, "integer >= 1"), + ({"max_conformers": 1.5}, "integer >= 1"), + ({"charge": float("inf")}, "charge must be finite"), + ({"charge": -0.5}, "charge must be an integer"), + ({"spin": 0.25}, "spin must be a non-negative integer or half-integer"), + ({"spin": 0.0}, "oxidized_spin=0 is incompatible"), + ({"potential_scale": "Fc/Fc+"}, "requires reference calibration"), + ({"engine": "xtb", "solvent": "water"}, "does not support solvent"), + ({"engine": "xtb", "dispersion": "d3-bj"}, "only by the DFTB"), + ({"engine": "xtb", "parameter_set": "3ob"}, "only to the DFTB"), + ({"engine": "dftb+", "method": "GFN2-xTB"}, "only to xTB"), + ], +) +def test_redox_screen_rejects_invalid_workflow_options(tmp_path, kwargs, message): + _write_xyz(tmp_path / "molecule.xyz") + with pytest.raises(TSValueError, match=message): + screening.redox_screen( + tmp_path / "molecule.xyz", + out=tmp_path / "out", + directory=tmp_path / "work", + **kwargs, + ) + + +def test_redox_screen_writes_machine_readable_aggregate(monkeypatch, tmp_path): + _write_xyz(tmp_path / "molecule.xyz") + energies = { + "molecule": { + "oxidized": -100.0, + "reduced_once": -100.1, + "reduced_twice": -100.18, + } + } + monkeypatch.setattr(screening, "screen", _fake_state_screen(energies)) + + results = screening.redox_screen( + tmp_path / "molecule.xyz", + out=tmp_path / "nested" / "redox", + directory=tmp_path / "work", + ) + + with open(tmp_path / "nested" / "redox.json", encoding="utf-8") as handle: + assert json.load(handle) == results + with open(tmp_path / "nested" / "redox.csv", newline="", encoding="utf-8") as handle: + row = next(csv.DictReader(handle)) + assert row["name"] == "molecule" + assert row["potential_inversion"] == "False" + metadata = json.loads( + (tmp_path / "nested" / "redox-run.json").read_text(encoding="utf-8") + ) + assert metadata["single_starting_geometry_approximation"] is True + assert metadata["candidates"][0]["structure_sha256"] + assert metadata["run_fingerprint"] == results[0]["run_fingerprint"] + + +def test_cli_parses_and_runs_redox(monkeypatch, capsys): + import ThermoScreening.cli.thermo as cli + + args = cli.parse_args( + [ + "redox", + "molecules.csv", + "--reference", + "AQ", + "--reference-e1", + "-0.75", + "--reference-e2", + "-1.40", + "--potential-scale", + "Fc/Fc+", + "--jobs", + "6", + "--resume", + ] + ) + captured = {} + + def fake_redox(source, **kwargs): + captured.update(source=source, kwargs=kwargs) + return [ + { + "name": "candidate", + "status": "ok", + "E1_V": -0.6, + "E2_V": -1.2, + "E2e_V": -0.9, + "potential_scale": "Fc/Fc+", + } + ] + + monkeypatch.setattr(cli, "redox_screen", fake_redox) + + assert args.command == "redox" + assert args.charge is None + assert args.parameter_set is None + assert args.method is None + assert cli.run_redox(args) == 0 + assert captured["kwargs"]["jobs"] == 6 + assert captured["kwargs"]["resume"] is True + assert captured["kwargs"]["reference_e1"] == -0.75 + assert "E1=-0.6000 V" in capsys.readouterr().out + + +def test_cli_redox_reports_workflow_errors(monkeypatch, capsys): + import ThermoScreening.cli.thermo as cli + + args = cli.parse_args(["redox", "missing.xyz"]) + monkeypatch.setattr( + cli, + "redox_screen", + lambda *args, **kwargs: (_ for _ in ()).throw(TSValueError("missing")), + ) + + assert cli.run_redox(args) == 1 + assert "Redox screen failed: missing" in capsys.readouterr().err + + +def test_cli_redox_reports_per_molecule_failures(monkeypatch, capsys): + import ThermoScreening.cli.thermo as cli + + args = cli.parse_args(["redox", "molecules.csv"]) + monkeypatch.setattr( + cli, + "redox_screen", + lambda *args, **kwargs: [ + {"name": "failed", "status": "error", "error": "anion failed"} + ], + ) + + assert cli.run_redox(args) == 1 + output = capsys.readouterr().out + assert "Failed (1)" in output + assert "failed: anion failed" in output diff --git a/tests/thermo/test_screening.py b/tests/thermo/test_screening.py index 31df9b2..ee12843 100644 --- a/tests/thermo/test_screening.py +++ b/tests/thermo/test_screening.py @@ -284,6 +284,179 @@ def fake_thermo(atoms, directory=None, **kwargs): assert by_name["good2"]["status"] == "ok" +def test_screen_shards_are_disjoint_and_collect_in_input_order(monkeypatch, tmp_path): + for name in ("mol_a", "mol_b", "mol_c", "mol_d", "mol_e"): + _write_xyz(tmp_path / f"{name}.xyz") + monkeypatch.setattr( + screening, "dftbplus_thermo", lambda atoms, **kwargs: _FakeThermo() + ) + out = tmp_path / "results" + work = tmp_path / "runs" + + first = screening.screen( + tmp_path, + out=out, + directory=work, + shard_index=0, + shard_count=2, + ) + second = screening.screen( + tmp_path, + out=out, + directory=work, + shard_index=1, + shard_count=2, + ) + combined = screening.collect_screen_shards( + screening.screen_shard_directory(out), out=out + ) + + assert [record["name"] for record in first] == ["mol_a", "mol_c", "mol_e"] + assert [record["name"] for record in second] == ["mol_b", "mol_d"] + assert [record["name"] for record in combined] == [ + "mol_a", + "mol_b", + "mol_c", + "mol_d", + "mol_e", + ] + assert (tmp_path / "results-shards" / "shard-00000.json").is_file() + assert (tmp_path / "results-shards" / "shard-00001-run.json").is_file() + metadata = json.loads((tmp_path / "results-run.json").read_text(encoding="utf-8")) + assert metadata["collection"]["shard_count"] == 2 + assert [job["input_index"] for job in metadata["jobs"]] == list(range(5)) + + +def test_screen_empty_shards_are_collectable(monkeypatch, tmp_path): + inputs = tmp_path / "inputs" + inputs.mkdir() + _write_xyz(inputs / "only.xyz") + monkeypatch.setattr( + screening, "dftbplus_thermo", lambda atoms, **kwargs: _FakeThermo() + ) + out = tmp_path / "results" + + shards = [ + screening.screen( + inputs, + out=out, + directory=tmp_path / "runs", + shard_index=index, + shard_count=3, + ) + for index in range(3) + ] + combined = screening.collect_screen_shards( + screening.screen_shard_directory(out), out=out + ) + + assert [len(shard) for shard in shards] == [1, 0, 0] + assert [record["name"] for record in combined] == ["only"] + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"shard_index": 0}, "supplied together"), + ({"shard_count": 2}, "supplied together"), + ({"shard_index": 0, "shard_count": 0}, "must be >= 1"), + ({"shard_index": 2, "shard_count": 2}, "0 <= index"), + ({"shard_index": 0.5, "shard_count": 2}, "must be an integer"), + ], +) +def test_screen_rejects_invalid_shards(tmp_path, kwargs, message): + _write_xyz(tmp_path / "mol.xyz") + with pytest.raises(TSValueError, match=message): + screening.screen(tmp_path, out=tmp_path / "out", **kwargs) + + +def test_collect_screen_shards_rejects_missing_and_tampered_shards( + monkeypatch, tmp_path +): + inputs = tmp_path / "inputs" + inputs.mkdir() + _write_xyz(inputs / "a.xyz") + _write_xyz(inputs / "b.xyz") + monkeypatch.setattr( + screening, "dftbplus_thermo", lambda atoms, **kwargs: _FakeThermo() + ) + out = tmp_path / "results" + screening.screen( + inputs, out=out, shard_index=0, shard_count=2, directory=tmp_path / "runs" + ) + + with pytest.raises(TSValueError, match="Missing cluster shard"): + screening.collect_screen_shards(screening.screen_shard_directory(out), out=out) + + screening.screen( + inputs, out=out, shard_index=1, shard_count=2, directory=tmp_path / "runs" + ) + shard_path = tmp_path / "results-shards" / "shard-00001.json" + records = json.loads(shard_path.read_text(encoding="utf-8")) + records[0]["fingerprint"] = "tampered" + shard_path.write_text(json.dumps(records), encoding="utf-8") + + with pytest.raises(TSValueError, match="Fingerprint mismatch"): + screening.collect_screen_shards(screening.screen_shard_directory(out), out=out) + + +def test_collect_screen_shards_rejects_wrong_assignment(monkeypatch, tmp_path): + inputs = tmp_path / "inputs" + inputs.mkdir() + _write_xyz(inputs / "a.xyz") + _write_xyz(inputs / "b.xyz") + monkeypatch.setattr( + screening, "dftbplus_thermo", lambda atoms, **kwargs: _FakeThermo() + ) + out = tmp_path / "results" + for index in range(2): + screening.screen( + inputs, + out=out, + shard_index=index, + shard_count=2, + directory=tmp_path / "runs", + ) + + metadata_path = tmp_path / "results-shards" / "shard-00001-run.json" + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + metadata["jobs"][0]["input_index"] = 2 + metadata_path.write_text(json.dumps(metadata), encoding="utf-8") + + with pytest.raises(TSValueError, match="does not belong"): + screening.collect_screen_shards(screening.screen_shard_directory(out), out=out) + + +def test_collect_screen_shards_rejects_mixed_settings(monkeypatch, tmp_path): + inputs = tmp_path / "inputs" + inputs.mkdir() + _write_xyz(inputs / "a.xyz") + _write_xyz(inputs / "b.xyz") + monkeypatch.setattr( + screening, "dftbplus_thermo", lambda atoms, **kwargs: _FakeThermo() + ) + out = tmp_path / "results" + screening.screen( + inputs, + out=out, + temperature=298.15, + shard_index=0, + shard_count=2, + directory=tmp_path / "runs", + ) + screening.screen( + inputs, + out=out, + temperature=310.0, + shard_index=1, + shard_count=2, + directory=tmp_path / "runs", + ) + + with pytest.raises(TSValueError, match="different inputs or settings"): + screening.collect_screen_shards(screening.screen_shard_directory(out), out=out) + + def test_screen_dispatches_to_xtb_engine(monkeypatch, tmp_path): _write_xyz(tmp_path / "mol.xyz") @@ -346,6 +519,37 @@ def thermo_fail_b(atoms, directory=None, **kwargs): assert {r["name"]: r["status"] for r in second} == {"mol_a": "ok", "mol_b": "ok"} +def test_screen_resume_invalidates_changed_scientific_inputs(monkeypatch, tmp_path): + _write_xyz(tmp_path / "mol.xyz") + out = tmp_path / "out" + calls = [] + + def fake_thermo(atoms, charge=0.0, **kwargs): + calls.append(charge) + return _FakeThermo() + + monkeypatch.setattr(screening, "dftbplus_thermo", fake_thermo) + first = screening.screen( + tmp_path, + out=out, + charge=0, + directory=tmp_path / "runs-a", + ) + second = screening.screen( + tmp_path, + out=out, + charge=1, + directory=tmp_path / "runs-b", + resume=True, + ) + + assert calls == [0, 1] + assert first[0]["fingerprint"] != second[0]["fingerprint"] + metadata = json.loads((tmp_path / "out-run.json").read_text(encoding="utf-8")) + assert metadata["jobs"][0]["charge"] == 1 + assert len(metadata["jobs"][0]["structure_sha256"]) == 64 + + def test_load_completed_handles_missing_and_corrupt(tmp_path): # no prior file -> nothing to resume assert screening._load_completed(str(tmp_path / "nope")) == {}