From 3f26335d578144d030a612b424f26008ce2352a0 Mon Sep 17 00:00:00 2001 From: Jonathan Schultz Date: Sun, 16 Aug 2026 17:10:40 +0000 Subject: [PATCH] M-CHECKPOINT CHK.7: checkpoint Reorganization Energy's geometry optimizations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reorganization Energy runs 2-3 independent geometry optimizations under one Calculate-tab run (the neutral reference plus one per ion channel) and was dispatched with no checkpoint at all — the run most likely to be interrupted (two full geometry opts plus four SCF energies) had nothing to resume into. - checkpoint.py: new Checkpoint.sub() derives a nested child checkpoint (own resume_key, dir nested under /legs/) for one leg of a multi-stage run — a single shared checkpoint would let one leg's trajectory/Hessian overwrite another's. has_progress() now also checks nested legs, since the parent itself never writes a trajectory directly; without this the run would never be offered as resumable no matter how much of a leg had completed. - reorganization_energy.py: run_reorganization_energy() takes checkpoint/resume params, builds+begins a leg checkpoint for the neutral optimization and each ion-channel optimization, and marks the parent complete once every required single point has converged (mirrors pes_scan.py's "every point attempted" convention) — without this a successful run would linger forever in the "unfinished calculations" list. - app.py: the Reorganization Energy branch of _do_run now passes checkpoint=_ckpt, resume=_resume like the PES Scan and Single Point branches already did. - app_runflow.py: fixed a real bug found while wiring this up — _CALC_TYPE_KEYS (the dropdown-label -> checkpoint/estimator key map) never listed "Reorganization Energy", so checkpoint_identity() silently built every reorg checkpoint under the "single_point" key, colliding with actual single-point checkpoints on the same molecule/method/basis, and the runtime estimator was scoring reorg runs against single-point history. Verified against real PySCF (RHF/STO-3G, H2 hole channel): two leg checkpoints written with distinct resume keys, correct steps_done, parent correctly marked complete and no longer offered as resumable. CHK.4 (frequency displacement restart) stays deferred, per the roadmap's own scope note — logged as a grounded design proposal in the planning repo rather than implemented this pass. Full suite: 2557 passed, 14 failed (pre-existing NMR failures, unrelated), 23 skipped. ruff + black + mypy (pre-push hook-stage) clean. Contributions: - Claude (Opus 4.8): implementation, tests, verification - Jonathan Schultz: direction and review Co-authored-by: Jonathan Schultz Co-authored-by: Claude --- quantui/app.py | 2 + quantui/app_runflow.py | 1 + quantui/checkpoint.py | 59 ++++- quantui/reorganization_energy.py | 70 +++++- tests/test_checkpoint_reorg_legs.py | 356 ++++++++++++++++++++++++++++ 5 files changed, 486 insertions(+), 2 deletions(-) create mode 100644 tests/test_checkpoint_reorg_legs.py diff --git a/quantui/app.py b/quantui/app.py index edd5c0b..fa12d9e 100644 --- a/quantui/app.py +++ b/quantui/app.py @@ -5186,6 +5186,8 @@ def _run_required_final_single_point(target_mol, reason: str): steps=self.max_steps_si.value, progress_stream=log, # type: ignore[arg-type] solvent=_solvent, + checkpoint=_ckpt, + resume=_resume, ) result_html = self._format_reorg_result(result) save_spectra = result.to_spectra() diff --git a/quantui/app_runflow.py b/quantui/app_runflow.py index 440e5b7..490819d 100644 --- a/quantui/app_runflow.py +++ b/quantui/app_runflow.py @@ -1570,6 +1570,7 @@ def _update_open_shell_hint(app: Any) -> None: "UV-Vis (TD-DFT)": "tddft", "NMR Shielding": "nmr", "PES Scan": "pes_scan", + "Reorganization Energy": "reorganization_energy", } diff --git a/quantui/checkpoint.py b/quantui/checkpoint.py index d581931..52f6b50 100644 --- a/quantui/checkpoint.py +++ b/quantui/checkpoint.py @@ -429,9 +429,66 @@ def has_progress(self) -> bool: return True traj = self.trajectory_path try: - return traj.is_file() and traj.stat().st_size > 0 + if traj.is_file() and traj.stat().st_size > 0: + return True + except OSError: + pass + return self._has_leg_progress() + + def _has_leg_progress(self) -> bool: + """True when any nested leg checkpoint (see :meth:`sub`) has progress. + + A multi-stage run like Reorganization Energy writes nothing directly + into its own checkpoint — all the real optimizer state lives in its + legs. Without this, the top-level checkpoint would always look empty + and the run would never be offered as resumable, no matter how much + of a leg had actually completed. + """ + try: + entries = list((self.dir / "legs").iterdir()) except OSError: return False + return any(entry.is_dir() and _has_progress_in(entry) for entry in entries) + + # ── Multi-stage runs (CHK.7) ──────────────────────────────────────────── + + def sub( + self, + tag: str, + *, + charge: int, + multiplicity: int, + coords: Sequence[Sequence[float]], + ) -> Checkpoint: + """A nested checkpoint for one leg of a multi-stage run. + + Reorganization Energy runs several independent geometry optimizations + under a single top-level checkpoint — the neutral reference plus one + per ion channel. Sharing one ``resume_key`` across all of them would + let one leg's trajectory/Hessian overwrite another's, so each leg + needs its own identity. This derives one from *tag* (a short label + like ``"neutral_opt"`` or ``"hole_opt"``), *charge*, *multiplicity* + and *coords* — the actual starting point of that leg's optimization — + while keeping method/basis/atoms from the parent. + + The child's directory is nested under this checkpoint's own + (``/legs/``) rather than living as a sibling in + the top-level checkpoint root. That keeps per-leg internals out of + the general "unfinished calculations" listing, which has no useful + way to present or restore half of a reorg run on its own. + """ + leg_identity = CalcIdentity( + calc_type=f"{self.identity.calc_type}:{tag}", + method=self.identity.method, + basis=self.identity.basis, + charge=int(charge), + multiplicity=int(multiplicity), + atom_symbols=self.identity.atom_symbols, + coords=tuple(tuple(float(v) for v in row) for row in coords), + ) + return Checkpoint( + leg_identity, root=self.dir / "legs", log_stream=self._log_stream + ) # ── Scan points (CHK.3) ───────────────────────────────────────────────── diff --git a/quantui/reorganization_energy.py b/quantui/reorganization_energy.py index fd3f5f0..9654818 100644 --- a/quantui/reorganization_energy.py +++ b/quantui/reorganization_energy.py @@ -37,7 +37,7 @@ import sys from dataclasses import dataclass, field -from typing import IO, List, Optional +from typing import IO, Any, List, Optional from .molecule import Molecule from .optimizer import DEFAULT_FMAX, DEFAULT_OPT_STEPS, optimize_geometry @@ -351,6 +351,8 @@ def run_reorganization_energy( steps: int = DEFAULT_OPT_STEPS, progress_stream: Optional[IO[str]] = None, solvent: Optional[str] = None, + checkpoint: Optional[Any] = None, + resume: bool = False, ) -> ReorganizationEnergyResult: """Compute the 4-point Marcus reorganization energy for a molecule. @@ -366,6 +368,27 @@ def run_reorganization_energy( progress_stream: Writable stream for live log output (Jupyter widget stream in the app, ``sys.stdout`` otherwise). solvent: Optional PCM solvent name for the single-point evaluations. + checkpoint: Optional :class:`~quantui.checkpoint.Checkpoint` for the + whole run (M-CHECKPOINT CHK.7). This run is really 2-3 independent + geometry optimizations — the neutral reference plus one per ion + channel — and each needs its own resume identity, or one leg's + trajectory/Hessian would overwrite another's. So *checkpoint* + itself is never handed to ``optimize_geometry`` directly; instead + each leg gets its own nested checkpoint via + :meth:`~quantui.checkpoint.Checkpoint.sub`. A leg that fails to + open its own checkpoint (disk full, permissions) simply runs + uncheckpointed — never the reason a leg doesn't run. + resume: Continue every leg from its own checkpoint where one exists. + A leg with nothing usable to resume just starts fresh — this is + not an error, it is the ordinary case for legs added after the + interrupted run's progress. Note: a leg that had already + *completed* before the interruption has no resumable state + either (by design — see ``Checkpoint.resumable_state``), so it is + re-optimized rather than instantly reused. In practice this + re-optimization converges in a step or two, since it starts + already at the minimum, and — because BFGS is deterministic — + reproduces the same geometry, so a later leg seeded from it still + finds its own checkpoint. Returns: :class:`ReorganizationEnergyResult`. @@ -399,6 +422,24 @@ def _single_point(mol: Molecule, mth: str, tag: str) -> float: raise RuntimeError(f"Single point did not converge: {tag}") return float(res.energy_hartree) + def _leg_checkpoint( + tag: str, *, charge: int, multiplicity: int, coords: Any + ) -> Optional[Any]: + """Nested checkpoint for one geometry-opt leg (CHK.7), or ``None``. + + ``checkpoint`` (the whole-run checkpoint) is never passed straight to + ``optimize_geometry`` — see the docstring above for why each leg needs + its own. A leg whose checkpoint fails to open behaves exactly like no + checkpoint at all: the optimization still runs, it just isn't + resumable. + """ + if checkpoint is None: + return None + leg = checkpoint.sub( + tag, charge=charge, multiplicity=multiplicity, coords=coords + ) + return leg if leg.begin() else None + _emit( stream, "\n" @@ -411,6 +452,12 @@ def _single_point(mol: Molecule, mth: str, tag: str) -> float: # ── Step 1: optimize the neutral reference geometry ────────────────────── _emit(stream, "\n── Optimizing neutral geometry (R_neutral) ──────────\n") + neutral_leg = _leg_checkpoint( + "neutral_opt", + charge=base_charge, + multiplicity=base_mult, + coords=molecule.coordinates, + ) neutral_opt = optimize_geometry( molecule=molecule, method=neutral_method, @@ -420,6 +467,8 @@ def _single_point(mol: Molecule, mth: str, tag: str) -> float: progress_stream=stream, # type: ignore[arg-type] status_label="Reorg: optimizing neutral geometry", report_fraction=False, # Don't let sub-opt 0→1 resets oscillate ETA + checkpoint=neutral_leg, + resume=resume, ) neutral_mol = neutral_opt.molecule n_total_steps = neutral_opt.n_steps @@ -455,6 +504,12 @@ def _single_point(mol: Molecule, mth: str, tag: str) -> float: charge=ion_charge, multiplicity=ion_mult, ) + ion_leg = _leg_checkpoint( + f"{kind}_opt", + charge=ion_charge, + multiplicity=ion_mult, + coords=ion_seed.coordinates, + ) ion_opt = optimize_geometry( molecule=ion_seed, method=ion_method, @@ -464,6 +519,8 @@ def _single_point(mol: Molecule, mth: str, tag: str) -> float: progress_stream=stream, # type: ignore[arg-type] status_label=f"Reorg: optimizing {kind} ion geometry", report_fraction=False, # See neutral-opt note above + checkpoint=ion_leg, + resume=resume, ) ion_mol = ion_opt.molecule n_total_steps += ion_opt.n_steps @@ -518,6 +575,17 @@ def _single_point(mol: Molecule, mth: str, tag: str) -> float: f"({lambda_total * HARTREE_TO_KCAL:.2f} kcal/mol)\n", ) + if checkpoint is not None: + # Reaching here means every required single point converged (a + # non-convergent one raises and never gets this far) — the run is + # done, whether or not an individual leg itself fully converged. + # Mirrors pes_scan.py: "every point was attempted, so there is + # nothing left to resume." A leg's own checkpoint (see _leg_checkpoint) + # only marks itself complete on convergence, independently of this — + # a leg that hit max steps without converging stays resumable at its + # own level even after the overall run completes. + checkpoint.mark_complete() + result = ReorganizationEnergyResult( formula=molecule.get_formula(), method=method, diff --git a/tests/test_checkpoint_reorg_legs.py b/tests/test_checkpoint_reorg_legs.py new file mode 100644 index 0000000..9d9c45b --- /dev/null +++ b/tests/test_checkpoint_reorg_legs.py @@ -0,0 +1,356 @@ +"""M-CHECKPOINT CHK.7 — Reorganization Energy checkpointing. + +Reorganization Energy is really 2-3 independent geometry optimizations (the +neutral reference plus one per ion channel) run under a single Calculate-tab +"calculation". A single shared checkpoint would let one leg's trajectory and +BFGS Hessian overwrite another's, so each leg needs its own resume identity — +that's what ``Checkpoint.sub()`` exists for. + +No PySCF/ASE here: ``optimize_geometry``/``run_in_session`` are monkeypatched +with fakes so the wiring is tested directly, mirroring +``tests/test_checkpoint_wiring.py``'s no-SCF style. +""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from quantui import checkpoint as C +from quantui import reorganization_energy as R +from quantui.molecule import Molecule + +# ══ Fixtures ═════════════════════════════════════════════════════════════════ + + +@pytest.fixture +def root(tmp_path, monkeypatch) -> Path: + monkeypatch.setenv("QUANTUI_CHECKPOINT_DIR", str(tmp_path / "ckpt")) + return tmp_path / "ckpt" + + +def _identity(**overrides) -> C.CalcIdentity: + base = dict( + calc_type="reorganization_energy", + method="B3LYP", + basis="6-31G*", + charge=0, + multiplicity=1, + atom_symbols=("O", "H", "H"), + coords=((0.0, 0.0, 0.0), (0.76, 0.59, 0.0), (-0.76, 0.59, 0.0)), + ) + base.update(overrides) + return C.CalcIdentity(**base) + + +def _neutral_molecule() -> Molecule: + return Molecule( + atoms=["O", "H", "H"], + coordinates=[[0.0, 0.0, 0.0], [0.76, 0.59, 0.0], [-0.76, 0.59, 0.0]], + charge=0, + multiplicity=1, + ) + + +# ══ Checkpoint.sub() — nested leg identity ══════════════════════════════════ + + +class TestCheckpointSub: + def test_a_leg_has_its_own_resume_key(self, root): + parent = C.Checkpoint(_identity()) + leg = parent.sub( + "neutral_opt", charge=0, multiplicity=1, coords=[[0.0, 0.0, 0.0]] + ) + assert leg.identity.resume_key != parent.identity.resume_key + + def test_two_different_tags_never_collide(self, root): + parent = C.Checkpoint(_identity()) + hole = parent.sub("hole_opt", charge=1, multiplicity=2, coords=[[0, 0, 0]]) + electron = parent.sub( + "electron_opt", charge=-1, multiplicity=2, coords=[[0, 0, 0]] + ) + assert hole.identity.resume_key != electron.identity.resume_key + + def test_the_same_tag_and_inputs_give_the_same_key(self, root): + parent = C.Checkpoint(_identity()) + a = parent.sub("neutral_opt", charge=0, multiplicity=1, coords=[[0, 0, 0]]) + b = parent.sub("neutral_opt", charge=0, multiplicity=1, coords=[[0, 0, 0]]) + assert a.identity.resume_key == b.identity.resume_key + + def test_a_leg_nests_under_the_parent_directory(self, root): + parent = C.Checkpoint(_identity()) + leg = parent.sub("neutral_opt", charge=0, multiplicity=1, coords=[[0, 0, 0]]) + assert leg.dir.parent.parent == parent.dir + assert leg.dir.parent.name == "legs" + + def test_method_basis_and_atoms_are_inherited_from_the_parent(self, root): + parent = C.Checkpoint(_identity(method="B3LYP", basis="6-31G*")) + leg = parent.sub("hole_opt", charge=1, multiplicity=2, coords=[[0, 0, 0]]) + assert leg.identity.method == "B3LYP" + assert leg.identity.basis == "6-31G*" + assert leg.identity.atom_symbols == ("O", "H", "H") + + def test_a_leg_writes_a_real_usable_checkpoint(self, root): + parent = C.Checkpoint(_identity()) + leg = parent.sub("neutral_opt", charge=0, multiplicity=1, coords=[[0, 0, 0]]) + assert leg.begin() is True + assert leg.exists() + + +class TestParentProgressReflectsLegs: + """The parent checkpoint never writes a trajectory of its own — all the + real state lives in its legs. Without _has_leg_progress, the parent would + always look empty and the run would never be offered as resumable.""" + + def test_no_progress_when_no_leg_has_run(self, root): + parent = C.Checkpoint(_identity()) + parent.begin() + assert parent.has_progress() is False + + def test_progress_appears_once_a_leg_has_a_trajectory(self, root): + parent = C.Checkpoint(_identity()) + parent.begin() + leg = parent.sub("neutral_opt", charge=0, multiplicity=1, coords=[[0, 0, 0]]) + leg.begin() + leg.trajectory_path.write_bytes(b"not empty") + assert parent.has_progress() is True + + def test_an_empty_leg_trajectory_is_not_progress(self, root): + parent = C.Checkpoint(_identity()) + parent.begin() + leg = parent.sub("neutral_opt", charge=0, multiplicity=1, coords=[[0, 0, 0]]) + leg.begin() + leg.trajectory_path.write_bytes(b"") + assert parent.has_progress() is False + + def test_a_leg_with_no_begin_never_created_a_directory(self, root): + # begin() never called -> no legs/ dir at all -> must not raise. + parent = C.Checkpoint(_identity()) + parent.begin() + assert parent.has_progress() is False + + +# ══ run_reorganization_energy wiring ════════════════════════════════════════ + + +class _FakeOptResult: + def __init__(self, molecule, n_steps=3, converged=True): + self.molecule = molecule + self.n_steps = n_steps + self.converged = converged + + +def _make_optimize_geometry(calls: list): + """A fake optimize_geometry that records every call's checkpoint/resume + and returns the input geometry unchanged (converged).""" + + def _fake(*, molecule, checkpoint=None, resume=False, **kw): + calls.append( + { + "molecule": molecule, + "checkpoint": checkpoint, + "resume": resume, + "status_label": kw.get("status_label"), + } + ) + return _FakeOptResult(molecule) + + return _fake + + +def _fake_run_in_session(*, molecule, method, basis, **kw): + return SimpleNamespace(converged=True, energy_hartree=-1.0) + + +class TestReorgCheckpointWiring: + def test_without_a_checkpoint_every_leg_runs_uncheckpointed( + self, root, monkeypatch + ): + calls: list = [] + monkeypatch.setattr(R, "optimize_geometry", _make_optimize_geometry(calls)) + monkeypatch.setattr(R, "run_in_session", _fake_run_in_session) + + R.run_reorganization_energy( + _neutral_molecule(), mode="both", method="B3LYP", basis="6-31G*" + ) + + assert len(calls) == 3 # neutral + hole + electron + assert all(c["checkpoint"] is None for c in calls) + assert all(c["resume"] is False for c in calls) + + def test_each_leg_gets_its_own_begun_checkpoint(self, root, monkeypatch): + calls: list = [] + monkeypatch.setattr(R, "optimize_geometry", _make_optimize_geometry(calls)) + monkeypatch.setattr(R, "run_in_session", _fake_run_in_session) + + parent = C.Checkpoint(_identity()) + R.run_reorganization_energy( + _neutral_molecule(), + mode="both", + method="B3LYP", + basis="6-31G*", + checkpoint=parent, + ) + + assert len(calls) == 3 + legs = [c["checkpoint"] for c in calls] + assert all(leg is not None for leg in legs) + assert all(leg.exists() for leg in legs) + # No two legs share a resume key. + keys = {leg.identity.resume_key for leg in legs} + assert len(keys) == 3 + + def test_leg_calc_type_tags_are_distinct_and_prefixed_by_the_parent( + self, root, monkeypatch + ): + calls: list = [] + monkeypatch.setattr(R, "optimize_geometry", _make_optimize_geometry(calls)) + monkeypatch.setattr(R, "run_in_session", _fake_run_in_session) + + parent = C.Checkpoint(_identity(calc_type="reorganization_energy")) + R.run_reorganization_energy( + _neutral_molecule(), + mode="both", + method="B3LYP", + basis="6-31G*", + checkpoint=parent, + ) + + tags = {c["checkpoint"].identity.calc_type for c in calls} + assert tags == { + "reorganization_energy:neutral_opt", + "reorganization_energy:hole_opt", + "reorganization_energy:electron_opt", + } + + def test_the_parent_checkpoint_is_marked_complete_on_success( + self, root, monkeypatch + ): + # Without this, a fully successful run would linger forever in the + # "unfinished calculations" listing — resumable_state() only excludes + # STATUS_COMPLETE, and nothing else ever marks the parent done. + calls: list = [] + monkeypatch.setattr(R, "optimize_geometry", _make_optimize_geometry(calls)) + monkeypatch.setattr(R, "run_in_session", _fake_run_in_session) + + parent = C.Checkpoint(_identity()) + parent.begin() + R.run_reorganization_energy( + _neutral_molecule(), + mode="hole", + method="B3LYP", + basis="6-31G*", + checkpoint=parent, + ) + + assert parent.load_state()["status"] == C.STATUS_COMPLETE + assert parent.resumable_state() is None + + def test_the_parent_is_not_marked_complete_when_a_single_point_fails( + self, root, monkeypatch + ): + calls: list = [] + monkeypatch.setattr(R, "optimize_geometry", _make_optimize_geometry(calls)) + + def _failing_run_in_session(*, molecule, method, basis, **kw): + return SimpleNamespace(converged=False, energy_hartree=0.0) + + monkeypatch.setattr(R, "run_in_session", _failing_run_in_session) + + parent = C.Checkpoint(_identity()) + parent.begin() + with pytest.raises(RuntimeError): + R.run_reorganization_energy( + _neutral_molecule(), + mode="hole", + method="B3LYP", + basis="6-31G*", + checkpoint=parent, + ) + + assert parent.load_state()["status"] != C.STATUS_COMPLETE + + def test_hole_only_mode_only_checkpoints_two_legs(self, root, monkeypatch): + calls: list = [] + monkeypatch.setattr(R, "optimize_geometry", _make_optimize_geometry(calls)) + monkeypatch.setattr(R, "run_in_session", _fake_run_in_session) + + parent = C.Checkpoint(_identity()) + R.run_reorganization_energy( + _neutral_molecule(), + mode="hole", + method="B3LYP", + basis="6-31G*", + checkpoint=parent, + ) + + tags = {c["checkpoint"].identity.calc_type for c in calls} + assert tags == { + "reorganization_energy:neutral_opt", + "reorganization_energy:hole_opt", + } + + def test_resume_is_threaded_to_every_leg(self, root, monkeypatch): + calls: list = [] + monkeypatch.setattr(R, "optimize_geometry", _make_optimize_geometry(calls)) + monkeypatch.setattr(R, "run_in_session", _fake_run_in_session) + + parent = C.Checkpoint(_identity()) + R.run_reorganization_energy( + _neutral_molecule(), + mode="both", + method="B3LYP", + basis="6-31G*", + checkpoint=parent, + resume=True, + ) + + assert all(c["resume"] is True for c in calls) + + def test_a_leg_that_fails_to_open_runs_uncheckpointed_not_broken( + self, root, monkeypatch + ): + # A checkpoint is an optimisation for the failure case, never the + # reason a leg doesn't run — mirrors Checkpoint.begin()'s own contract. + calls: list = [] + monkeypatch.setattr(R, "optimize_geometry", _make_optimize_geometry(calls)) + monkeypatch.setattr(R, "run_in_session", _fake_run_in_session) + monkeypatch.setattr(C.Checkpoint, "begin", lambda self, **kw: False) + + parent = C.Checkpoint(_identity()) + result = R.run_reorganization_energy( + _neutral_molecule(), + mode="hole", + method="B3LYP", + basis="6-31G*", + checkpoint=parent, + ) + + assert result.converged + assert all(c["checkpoint"] is None for c in calls) + + +# ══ calc_type_key regression (found while implementing CHK.7) ══════════════ + + +class TestReorgCalcTypeKey: + """``_CALC_TYPE_KEYS`` (app_runflow.py) is the label->key map that feeds + both the checkpoint identity and the runtime estimator. It listed every + calc type except Reorganization Energy, which silently fell back to + "single_point" — colliding a reorg run's checkpoint identity with an + actual single-point run on the same molecule/method/basis, and feeding + the estimator single-point history for a run that is 2-3 full geometry + optimizations. Fixed alongside CHK.7 since the leg tags + (``f"{parent.identity.calc_type}:{tag}"``) are meaningless if the parent + key itself is wrong. + """ + + def test_the_dropdown_label_maps_to_the_canonical_key(self): + from quantui import app_runflow + + app = SimpleNamespace( + calc_type_dd=SimpleNamespace(value="Reorganization Energy") + ) + assert app_runflow.calc_type_key(app) == "reorganization_energy"