diff --git a/ThermoScreening/cli/thermo.py b/ThermoScreening/cli/thermo.py index c8cb066..59ddd8e 100644 --- a/ThermoScreening/cli/thermo.py +++ b/ThermoScreening/cli/thermo.py @@ -122,6 +122,11 @@ def _command_parser(): "--method", default="GFN2-xTB", choices=["GFN2-xTB", "GFN1-xTB"], help="GFN-xTB parametrisation for the xtb engines (default 'GFN2-xTB').", ) + screen_parser.add_argument( + "--resume", action="store_true", + help="Reuse successful results from a prior .json and only (re)run " + "the missing/failed molecules.", + ) return parser @@ -207,6 +212,7 @@ def run_screen(parser_args): quasi_rrho=parser_args.quasi_rrho, engine=parser_args.engine, method=parser_args.method, + resume=parser_args.resume, ) failed = sum(1 for record in results if record["status"] != "ok") diff --git a/ThermoScreening/thermo/screening.py b/ThermoScreening/thermo/screening.py index 9c95fab..f949180 100644 --- a/ThermoScreening/thermo/screening.py +++ b/ThermoScreening/thermo/screening.py @@ -10,6 +10,7 @@ import csv import json import logging +import os from dataclasses import dataclass from pathlib import Path @@ -114,20 +115,56 @@ def _thermo_summary(thermo): } +def _load_completed(out): + """ + Load the successfully-completed records from a prior ``.json``. + + Returns a ``{name: record}`` mapping of the records whose ``status`` is + ``"ok"``, so a resumed screen can skip them. A missing or unreadable file + yields an empty mapping (nothing to resume). + """ + json_path = Path(str(out)).with_suffix(".json") + if not json_path.exists(): + return {} + try: + prior = json.loads(json_path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError): + 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 + } + + +def _atomic_write(path, write_fn): + """ + Write via a temp file + ``os.replace`` so an interrupted write never leaves a + truncated (unresumable) file behind. + """ + tmp = path.with_name(path.name + ".tmp") + with open(tmp, "w", newline="", encoding="utf-8") as handle: + write_fn(handle) + os.replace(tmp, path) + + def _write_results(results, out): stem = Path(str(out)).with_suffix("") if stem.parent != Path(""): stem.parent.mkdir(parents=True, exist_ok=True) - csv_path = stem.with_suffix(".csv") - with open(csv_path, "w", newline="", encoding="utf-8") as handle: + def write_csv(handle): writer = csv.DictWriter(handle, fieldnames=_RESULT_FIELDS) writer.writeheader() for record in results: writer.writerow({field: record.get(field, "") for field in _RESULT_FIELDS}) + csv_path = stem.with_suffix(".csv") json_path = stem.with_suffix(".json") - json_path.write_text(json.dumps(results, indent=2), encoding="utf-8") + _atomic_write(csv_path, write_csv) + _atomic_write(json_path, lambda handle: handle.write(json.dumps(results, indent=2))) return csv_path, json_path @@ -146,6 +183,7 @@ def screen( quasi_rrho=False, engine="dftb+", method="GFN2-xTB", + resume=False, ): """ Run a thermochemistry screen over a set of molecules. @@ -192,6 +230,11 @@ def screen( applies only to DFTB+; ``solvent`` applies to DFTB+ and xtb-cli. method : str GFN-xTB parametrisation for the xtb engines (``"GFN2-xTB"`` default). + resume : bool + If True, reuse the successful (``status="ok"``) records from a prior + ``.json`` and skip those molecules; molecules that previously failed + (or were never run) are (re-)run. Results are written incrementally after + every molecule regardless, so an interrupted screen can be resumed. Returns ------- @@ -210,8 +253,16 @@ def screen( jobs = _load_jobs(source, charge, spin) root = Path(directory) + completed = _load_completed(out) if resume else {} + results = [] + csv_path = json_path = None for job in jobs: + if job.name in completed: + logger.info(f"Skipping {job.name} (already completed)") + results.append(completed[job.name]) + csv_path, json_path = _write_results(results, out) + continue record = { "name": job.name, "path": str(job.path), @@ -267,7 +318,8 @@ def screen( # we must keep screening the remaining molecules logger.warning(f"Screening failed for {job.name}: {exc}") results.append(record) + # write after every molecule so an interrupted run stays resumable + csv_path, json_path = _write_results(results, out) - csv_path, json_path = _write_results(results, out) logger.info(f"Wrote {csv_path} and {json_path}") return results diff --git a/tests/thermo/test_screening.py b/tests/thermo/test_screening.py index 5d62281..c51ec99 100644 --- a/tests/thermo/test_screening.py +++ b/tests/thermo/test_screening.py @@ -174,6 +174,70 @@ def fake_dftb(*args, **kwargs): assert captured["method"] == "GFN1-xTB" +def test_screen_resume_skips_completed_and_reruns_failed(monkeypatch, tmp_path): + from pathlib import Path + + _write_xyz(tmp_path / "mol_a.xyz") + _write_xyz(tmp_path / "mol_b.xyz") + out = tmp_path / "out" + + calls = [] + + def thermo_fail_b(atoms, directory=None, **kwargs): + name = Path(directory).name + calls.append(name) + if name == "mol_b": + raise RuntimeError("boom") + return _FakeThermo() + + # first run: mol_a ok, mol_b fails + monkeypatch.setattr(screening, "dftbplus_thermo", thermo_fail_b) + first = screening.screen(str(tmp_path), out=str(out), directory=str(tmp_path / "r1")) + assert {r["name"]: r["status"] for r in first} == {"mol_a": "ok", "mol_b": "error"} + + # resume: mol_a (ok) skipped, only mol_b (was error) re-run + calls.clear() + monkeypatch.setattr( + screening, "dftbplus_thermo", + lambda atoms, directory=None, **kw: (calls.append(Path(directory).name), _FakeThermo())[1], + ) + second = screening.screen( + str(tmp_path), out=str(out), directory=str(tmp_path / "r2"), resume=True + ) + assert calls == ["mol_b"] # only the previously-failed molecule re-run + assert {r["name"]: r["status"] for r in second} == {"mol_a": "ok", "mol_b": "ok"} + + +def test_load_completed_handles_missing_and_corrupt(tmp_path): + # no prior file -> nothing to resume + assert screening._load_completed(str(tmp_path / "nope")) == {} + # unreadable/corrupt json -> empty, does not crash + (tmp_path / "bad.json").write_text("{ not valid json", encoding="utf-8") + assert screening._load_completed(str(tmp_path / "bad")) == {} + # valid JSON but not a list of records -> empty, does not crash + (tmp_path / "scalar.json").write_text("42", encoding="utf-8") + assert screening._load_completed(str(tmp_path / "scalar")) == {} + + +def test_screen_writes_results_incrementally(monkeypatch, tmp_path): + for name in ("mol_a", "mol_b", "mol_c"): + _write_xyz(tmp_path / f"{name}.xyz") + + writes = [] + real_write = screening._write_results + + def spy_write(results, out): + writes.append(len(results)) + return real_write(results, out) + + monkeypatch.setattr(screening, "_write_results", spy_write) + monkeypatch.setattr(screening, "dftbplus_thermo", lambda atoms, **kw: _FakeThermo()) + + screening.screen(str(tmp_path), out=str(tmp_path / "out"), directory=str(tmp_path / "r")) + # written after each molecule with a growing result set (durable/resumable) + assert writes == [1, 2, 3] + + def test_screen_dispatches_to_xtb_cli_engine(monkeypatch, tmp_path): _write_xyz(tmp_path / "mol.xyz") @@ -345,6 +409,9 @@ def test_cli_parse_args_routes_screen(): cli_args = cli.parse_args(["screen", "molecules.csv", "--engine", "xtb-cli"]) assert cli_args.engine == "xtb-cli" + assert args.resume is False # default + assert cli.parse_args(["screen", "molecules.csv", "--resume"]).resume is True + def test_cli_run_screen_returns_failure_count(monkeypatch): import ThermoScreening.cli.thermo as cli @@ -357,6 +424,7 @@ def test_cli_run_screen_returns_failure_count(monkeypatch): source="x", out="res", charge=0.0, temperature=298.15, pressure=101325.0, directory="screening", parameter_set="3ob", solvent=None, quasi_rrho=False, engine="dftb+", method="GFN2-xTB", + resume=False, ) assert cli.run_screen(args) == 1 # one molecule failed