Refine a trajectory's barriers after its MD when the campaign asks for it

README.md

@@ -59,6 +59,10 @@ it asks about GPUs to run on a CPU. A campaign with several trajectories runs on
 and one GPU per trajectory. The [campaign guide](docs/campaigns.md) describes the file, and
 [pathway refinement](docs/pathway-refinement.md) lists the refinement settings.
 
+By default a trajectory pauses for each reaction and refines it before MD goes on. With
+`"refine_after_md": true` in the campaign, MD runs to its last step without pausing, and the events
+wait in a queue until then ([when barriers are refined](docs/campaigns.md#when-barriers-are-refined)).
+
 Two more commands work on finished output. `reactionflow scan` replays the saved frames through
 bond detection, with the run's settings or new ones, and tabulates every bond change and species
 ([offline scan](docs/offline-scan.md)). `reactionflow check campaign.json band.traj` tests whether

docs/campaigns.md

@@ -218,6 +218,28 @@ remain independent of Torch. The complete four-GPU
 ANI-1xnr example is in
 [`examples/perlmutter/acn_20gpa_ani1xnr`](../examples/perlmutter/acn_20gpa_ani1xnr/README.md).
 
+## When barriers are refined
+
+By default a trajectory pauses at each confirmed reaction, refines it, and continues MD from the
+same checkpoint, so a slow refinement holds up that trajectory's MD. Set `"refine_after_md": true`
+at the top of the campaign to run MD straight through to `total_steps`. The events wait in a queue
+and are refined when MD ends, one at a time in the order MD found them, by the same process on the
+same GPU. Resubmit a job that ends partway, as usual: finished results are kept, and the next
+unfinished event is refined from the start.
+
+If a refinement gives no barrier, the model is evaluated once on the last checkpoint. If that fails
+or returns NaN, the job stops with the rest of the queue intact, so a lost GPU costs one event, not
+the queue. The check reloads the model at that checkpoint, so a calculator that needs its previous
+call to converge there, as an SCF can, may stop the job and hold the queue. If MD itself cannot go
+on, set the option to `false` and resubmit: the queue is refined before MD is touched. A trajectory
+marked failed refines its queue when resubmitted, then reports the failure again.
+
+MD restarts exactly after a pause, and refinement reads only saved frames, so with a model that
+keeps no state between evaluations both orders give the same frames, events, and barriers. A model
+with internal state, such as ASE's EMT and its neighbor list, can differ at rounding level, as it
+does after any restart. The setting is not part of a trajectory's contract, so you can change it
+when you resubmit. A trajectory already paused for refinement finishes that refinement first.
+
 ## Campaign status
 
 `reactionflow status campaign.json` summarizes a campaign without changing it, so it is safe to

docs/how-it-works.md

@@ -17,7 +17,9 @@ after it persists for three observations, and it becomes a reaction candidate on
 within two bonds of it has been stable for three observations. The trajectory then writes a
 checkpoint, relaxes the reactant and product, runs NEB and CI-NEB, checks the saddle, and records
 the result. MD resumes from the checkpoint, not from a relaxed structure, so it continues the
-trajectory that would have run without the pause.
+trajectory that would have run without the pause. With `"refine_after_md": true` in the campaign,
+the trajectory does not pause, and its events are refined when MD ends
+([when barriers are refined](campaigns.md#when-barriers-are-refined)).
 
 ## Bond detection
 
@@ -67,9 +69,11 @@ could not be mapped onto the relaxed endpoints, because a bond stayed between th
 the change disappeared, or relaxation changed bonds elsewhere in the cell. `stepwise` means the
 event was refined in steps, each with its own outcome. `failed` means an unexpected error. The
 saddle check's four verdicts are described in [pathway
-refinement](pathway-refinement.md#saddle-check). In every case the result is saved, the
-checkpoint is restored, and MD continues.
-A trajectory itself stops only if it cannot save its own state.
+refinement](pathway-refinement.md#saddle-check). In every case the result is saved. If MD has not
+reached `total_steps`, the checkpoint is restored and MD continues.
+A trajectory is marked failed only if its own bookkeeping fails: the bond monitor, the reaction
+store, or a save. If its model cannot load or raises an error, the job stops with a non-zero exit,
+and the next job continues from the last checkpoint.
 
 Results are never overwritten. Every resolved occurrence is refined, including one that repeats an
 earlier reaction, because the reactions that have happened around it can change its barrier. The
@@ -231,7 +235,8 @@ directory.
 
 To continue after a job ends, submit the same command again. Unfinished trajectories resume from
 their last checkpoint and finished ones are left as they are. Each output directory is bound to
-its structure, settings, and model, and a changed campaign is rejected instead of resumed.
+its structure, trajectory conditions, reaction settings, and model, and a change to those is
+rejected instead of resumed.
 
 [`examples/perlmutter/run-campaign.sbatch`](../examples/perlmutter/run-campaign.sbatch) is a job
 script for Perlmutter's four-GPU nodes. Thirty-two trajectories on eight nodes are submitted like

src/reactionflow/campaign.py

@@ -168,6 +168,7 @@ class CampaignConfig:
     trajectory_profiles: tuple[str, ...]
     trajectory_structures: tuple[Path, ...]
     require_gpu: bool = True
+    refine_after_md: bool = False
 
     @classmethod
     def load(cls, path: str | Path) -> CampaignConfig:
@@ -185,6 +186,7 @@ def load(cls, path: str | Path) -> CampaignConfig:
             "reaction_run",
             "trajectories",
             "require_gpu",
+            "refine_after_md",
             "adapter_profiles",
         }
         if unknown:
@@ -253,6 +255,9 @@ def structure_path(raw: object, name: str) -> Path:
         require_gpu = value.get("require_gpu", True)
         if type(require_gpu) is not bool:
             raise TypeError("campaign.require_gpu must be a boolean")
+        refine_after_md = value.get("refine_after_md", False)
+        if type(refine_after_md) is not bool:
+            raise TypeError("campaign.refine_after_md must be a boolean")
         return cls(
             source=source,
             output_root=(source.parent / output_value).resolve(),
@@ -262,6 +267,7 @@ def structure_path(raw: object, name: str) -> Path:
             trajectory_profiles=tuple(trajectory_profiles),
             trajectory_structures=tuple(trajectory_structures),
             require_gpu=require_gpu,
+            refine_after_md=refine_after_md,
         )
 
     def trajectory(self, index: int) -> TrajectorySpec:

src/reactionflow/cli.py

@@ -179,6 +179,7 @@ def run_selected_trajectory(
         pathway_calculator_provider=adapter.calculator,
         total_steps=trajectory.total_steps,
         pressure_GPa=trajectory.pressure_GPa,
+        refine_after_md=campaign.refine_after_md,
     )
 
 

src/reactionflow/run.py

@@ -124,7 +124,8 @@ class ReactionRun:
     Every observation publishes one exact checkpoint: the dynamics and calculator snapshot plus
     the reaction tracker, named by ``state.json`` together with the detector state. A confirmed
     reaction releases the MD runtime, refines the queued pathways, and restores that same
-    checkpoint into the next trajectory generation, exactly as a resubmitted job does.
+    checkpoint into the next trajectory generation, exactly as a resubmitted job does. With
+    ``refine_after_md``, confirmed reactions instead wait in the queue until MD reaches its target.
     """
 
     def __init__(self, root: Path, config: ReactionRunConfig) -> None:
@@ -148,6 +149,7 @@ def __init__(self, root: Path, config: ReactionRunConfig) -> None:
         self._active_checkpoint: Path | None = None
         self._snapshot: ExactRestartSnapshot | None = None
         self._created_with: str | None = __version__
+        self._refine_after_md = False
 
     @classmethod
     def create(
@@ -296,7 +298,8 @@ def _fatal(self, stage: str) -> Iterator[None]:
             yield
         except Exception as error:
             self._phase = "failed"
-            self._failure = {"stage": stage, "type": type(error).__name__, "message": str(error)}
+            failure = {"stage": stage, "type": type(error).__name__, "message": str(error)}
+            self._failure = self._failure or failure  # the first is why MD stopped
             self._write_state()
             raise
 
@@ -342,7 +345,7 @@ def _register(self, candidates: tuple[ReactionCandidate, ...], *, label: str) ->
             )
             self._queue_pathway(record, candidate)
 
-    def _observe(self, atoms: Atoms, global_step: int) -> None:
+    def _observe(self, atoms: Atoms, global_step: int, total_steps: int) -> None:
         """Advance the monitor one frame and register every completed occurrence."""
 
         assert self._detector is not None and self._tracker is not None
@@ -358,7 +361,7 @@ def _observe(self, atoms: Atoms, global_step: int) -> None:
             self._register(candidates, label=f"{frame:08d}")
             self._global_step = global_step
             self._global_frame = frame
-            if self._pending:
+            if self._pending and (not self._refine_after_md or global_step >= total_steps):
                 self._phase = "refining"
 
     def _write_frame(self, trajectory: Trajectory, atoms: Atoms, written: int) -> None:
@@ -372,7 +375,7 @@ def _write_frame(self, trajectory: Trajectory, atoms: Atoms, written: int) -> No
         trajectory.write(frame)
 
     def _run_segment(self, total_steps: int, runtime_provider: ExactRuntimeProvider) -> None:
-        """Run MD from the current checkpoint until a reaction is confirmed or the run ends."""
+        """Run MD from the current checkpoint until refinement is due or the run ends."""
 
         manager = (
             runtime_provider.start(self._initial)
@@ -408,7 +411,7 @@ def _run_segment(self, total_steps: int, runtime_provider: ExactRuntimeProvider)
                         raise RuntimeError(
                             f"exact runtime advanced {advanced} steps; expected {steps}"
                         )
-                    self._observe(runtime.atoms, self._global_step + steps)
+                    self._observe(runtime.atoms, self._global_step + steps, total_steps)
                     self._write_frame(trajectory, runtime.atoms, written)
                     # A checkpoint survives a crash, so the frames written before it must too.
                     flush_to_disk(path)
@@ -490,11 +493,15 @@ def _refine_pending(
         calculator_provider: CalculatorProvider,
         pressure_GPa: float | None,
     ) -> None:
-        """Refine each queued occurrence serially and publish its immutable result."""
+        """Refine each queued occurrence serially and publish its immutable result.
+
+        Returns early after a refinement with no barrier, so that the caller can check the model.
+        """
 
         with self._fatal("refine"):
             while self._pending:
                 occurrence_id = self._pending[0]
+                no_barrier = False
                 if (self.pathways / occurrence_id).is_dir():
                     # Published before an interruption; make its name durable and move on.
                     sync_directory(self.pathways)
@@ -512,8 +519,28 @@ def _refine_pending(
                         ),
                     )
                     self._publish_outcome(occurrence_id, outcome)
+                    no_barrier = outcome.barrier is None
                 self._pending.pop(0)
                 self._write_state()
+                if no_barrier:
+                    return
+
+    def _check_model(self, runtime_provider: ExactRuntimeProvider) -> None:
+        """Restore the last checkpoint, which loads and evaluates the model, or raise.
+
+        Reads a separate snapshot so that validation cannot mutate the run's checkpoint.
+        """
+
+        assert self._active_checkpoint is not None
+        snapshot = ExactRestartSnapshot.read(self._active_checkpoint / "exact-restart")
+        with runtime_provider.restore(snapshot) as runtime:
+            if runtime.atoms.calc is not None and not np.isfinite(runtime.atoms.get_forces()).all():
+                raise ValueError("the model returned non-finite forces at the last checkpoint")
+
+    def _refine_due(self) -> bool:
+        """Whether to refine now: the queue waits only for running MD under ``refine_after_md``."""
+
+        return bool(self._pending) and (self._phase != "running" or not self._refine_after_md)
 
     def _complete(self) -> None:
         """Drain unresolved terminal candidates and mark the run complete."""
@@ -546,6 +573,7 @@ def run(
         pathway_calculator_provider: CalculatorProvider,
         total_steps: int,
         pressure_GPa: float | None = None,
+        refine_after_md: bool = False,
     ) -> RunSummary:
         """Run or resume live detection, serial NEB/CI-NEB, and exact MD to a step target."""
 
@@ -559,25 +587,25 @@ def run(
             self._start(atoms)
         elif atoms is not None:
             raise ValueError("initial atoms may only be supplied to a new run")
+        self._refine_after_md = refine_after_md
 
-        if self._phase == "refining":
+        if self._phase == "refining" or self._refine_due():
             # A reopened adapter has not checked its model against the saved MD state yet.
             # Verify before any pathway result can be published, including terminal work
-            # that will never resume MD. Read a separate snapshot so validation cannot
-            # mutate the run's checkpoint.
-            assert self._active_checkpoint is not None
-            snapshot = ExactRestartSnapshot.read(self._active_checkpoint / "exact-restart")
-            with runtime_provider.restore(snapshot):
-                pass
+            # that will never resume MD.
+            self._check_model(runtime_provider)
 
         # Only failures to maintain the durable record mark the run failed. Any other error, such
         # as a runtime that this environment refuses to restore, leaves the last exact checkpoint
         # in place, so the run resumes once the cause is fixed.
         while self._phase != "completed":
-            if self._phase == "failed":
-                raise RuntimeError(f"ReactionRun failed: {self._failure}")
-            if self._pending:
+            if self._refine_due():
                 self._refine_pending(pathway_calculator_provider, pressure_GPa)
+                if self._pending:
+                    # No barrier. If the model is gone, stop before it fails the rest of the queue.
+                    self._check_model(runtime_provider)
+            elif self._phase == "failed":
+                raise RuntimeError(f"ReactionRun failed: {self._failure}")
             elif self._global_step >= total_steps:
                 self._complete()
             else:

tests/test_run.py

@@ -24,7 +24,7 @@
     assign_atom_ids,
 )
 from reactionflow.campaign import CampaignConfig
-from reactionflow.cli import main
+from reactionflow.cli import main, run_selected_trajectory
 from reactionflow.run import refine_pathway
 from reactionflow.status import _trajectory, campaign_status, format_status
 
@@ -735,3 +735,69 @@ def test_missing_tracker_checkpoint_does_not_reset_pending_history(tmp_path, mon
     shutil.rmtree(tmp_path / state["active_checkpoint"] / "tracker")
     with pytest.raises(FileNotFoundError):
         ReactionRun.open(tmp_path)
+
+
+def test_refine_after_md_runs_md_through_and_then_refines(tmp_path, monkeypatch) -> None:
+    path = pair_campaign(tmp_path, "pair")
+    value = {
+        **json.loads(path.read_text()),
+        "refine_after_md": True,
+        "reaction_run": run_config().to_dict(),
+    }
+    path.write_text(json.dumps(value))
+    leases = LeaseCounter()
+    adapter = ScriptedRuntimeProvider(leases)
+    adapter.calculator = leases
+    monkeypatch.setattr("reactionflow.cli.load_mlip_adapter", lambda *_: adapter)
+
+    summary = run_selected_trajectory(CampaignConfig.load(path), index=0, environment={})
+
+    # One unbroken segment, and the MD model is gone before refinement loads its own.
+    assert (summary.phase, summary.generation, summary.pathways) == ("completed", 0, 1)
+    assert leases.stages == ["md", "relax_reactant", "relax_product", "neb"]
+    assert leases.live == 0 and leases.max_live == 1
+    result = json.loads(next((tmp_path / "runs/pair/pathways").glob("*/result.json")).read_text())
+    assert result["barrier_eV"] == pytest.approx(0.0625, abs=2e-4)
+
+
+def test_refine_after_md_keeps_its_queue_across_interruptions(tmp_path, monkeypatch) -> None:
+    def advance(leases, **provider):
+        state = tmp_path / "state.json"
+        config = run_config(observation_interval=1)
+        run = (
+            ReactionRun.open(tmp_path)
+            if state.exists()
+            else ReactionRun.create(tmp_path, config=config)
+        )
+        return run.run(
+            pair(0.65) if run.phase == "new" else None,
+            # Break the bond, re-form it, and break it again.
+            runtime_provider=ScriptedRuntimeProvider(
+                leases, script={1: 1.55, 2: 0.65, 3: 1.55}, **provider
+            ),
+            pathway_calculator_provider=leases,
+            total_steps=3,
+            refine_after_md=True,
+        )
+
+    def saved_state():
+        state = json.loads((tmp_path / "state.json").read_text())
+        return state["phase"], state["global_step"], len(state["pending_pathway_ids"])
+
+    with pytest.raises(KeyboardInterrupt, match="scheduler interruption"):
+        advance(LeaseCounter(), interrupt_after_calls=2)
+    assert saved_state() == ("running", 2, 2)
+
+    with monkeypatch.context() as patch:
+        patch.setattr("reactionflow.run.refine_pathway", interrupt_refinement)
+        with pytest.raises(KeyboardInterrupt, match="before refinement"):
+            advance(LeaseCounter())
+    assert saved_state() == ("refining", 3, 3)
+
+    leases = LeaseCounter()
+    summary = advance(leases)
+
+    # The model is checked against the checkpoint once, MD does not run again, and each event is
+    # refined once.
+    assert leases.stages == ["md", *["relax_reactant", "relax_product", "neb"] * 3]
+    assert (summary.phase, summary.generation, summary.pathways) == ("completed", 0, 3)