diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py index 144b2c63..0feb1edc 100644 --- a/tests/test_trajectory.py +++ b/tests/test_trajectory.py @@ -9,6 +9,7 @@ import torch import torch_sim as ts +from torch_sim.autobatching import InFlightAutoBatcher from torch_sim.integrators import MDState from torch_sim.models.interface import ModelInterface from torch_sim.models.lennard_jones import LennardJonesModel @@ -1006,7 +1007,7 @@ def test_truncate_trajectory( ValueError, match=( r"Cannot truncate to a step greater than the last step\. " - r"self\.last_step=3 < step=10" + r"last_written_step=3 < step=10" ), ): traj.truncate_to_step(10) @@ -1047,8 +1048,10 @@ def test_truncate_trajectory_reporter( trajectory_reporter=trajectory_reporter, ) - trajectory_reporter.truncate_to_step(step=min(trajectory_reporter.last_steps)) - assert trajectory_reporter.last_steps == [5, 5] + trajectory_reporter.truncate_to_step( + step=min(trajectory_reporter.last_written_steps) + ) + assert trajectory_reporter.last_written_steps == [5, 5] with pytest.raises( ValueError, match=( @@ -1064,16 +1067,17 @@ def test_truncate_trajectory_reporter( trajectory_reporter.truncate_to_step(-2) # truncate to step 3 trajectory_reporter.truncate_to_step(3) - assert trajectory_reporter.last_steps == [3, 3] + assert trajectory_reporter.last_written_steps == [3, 3] def test_integrate_uneven_trajectory_append( si_double_sim_state: SimState, lj_model: LennardJonesModel ) -> None: """ - Test appending to an existing trajectory with uneven frames running ts.integrate. - Expected behavior: ts.integrate should first truncate all trajectories to the shortest - length, and then append new frames to all trajectories. + Test appending to trajectory files that disagree on their last step. + ts.integrate advances a single step counter for the whole batch, so its files + must agree before it can resume. Disagreement is an error the caller resolves + by truncating to a common step explicitly. """ # Create a temporary trajectory file @@ -1167,3 +1171,453 @@ def test_optimize_save_initial_state( steps = traj.get_steps("positions") # Should start at step 0 np.testing.assert_allclose(steps, [0, 1, 2, 3]) + + +# ---------------------------------------------------------------------------------- +# Resuming and truncating with mismatched array cadences +# ---------------------------------------------------------------------------------- + + +def _per_array_steps(traj_file: str | Path) -> dict[str, np.ndarray]: + """Map array name to its recorded steps, skipping global (step-0-only) arrays.""" + with TorchSimTrajectory(traj_file, mode="r") as traj: + return { + name: traj.get_steps(name) + for name in traj.array_registry + if set(traj.get_steps(name).tolist()) != {0} + } + + +def _run_integrate( + system: SimState, + model: LennardJonesModel, + traj_files: list[str], + *, + n_steps: int, + state_freq: int, + prop_freq: int, + mode: str = "w", +) -> SimState: + """Run ts.integrate with independent state and property cadences.""" + reporter = ts.TrajectoryReporter( + traj_files, + state_frequency=state_freq, + prop_calculators={prop_freq: {"potential_energy": lambda state: state.energy}}, + trajectory_kwargs=dict(mode=mode), + ) + return ts.integrate( + system=system, + model=model, + timestep=0.001, + n_steps=n_steps, + temperature=300.0, + integrator=ts.Integrator.nvt_langevin, + trajectory_reporter=reporter, + ) + + +def test_last_written_step_vs_last_step(test_file: Path, random_state: MDState) -> None: + """last_step tracks positions; last_written_step tracks all arrays.""" + with TorchSimTrajectory(test_file, mode="w") as traj: + # both are None for an empty trajectory + assert traj.last_step is None + assert traj.last_written_step is None + + traj.write_state(random_state, 0) + traj.write_state(random_state, 10) + for step in range(16): + traj.write_arrays({"potential_energy": torch.tensor([float(step)])}, step) + + assert traj.last_step == 10 + assert traj.last_written_step == 15 + assert traj.last_step_of("positions") == 10 + assert traj.last_step_of("potential_energy") == 15 + assert traj.last_step_of("not_an_array") is None + + +@pytest.mark.parametrize( + ("state_freq", "prop_freq"), + [(1, 1), (10, 1), (100, 10), (100, 30), (10, 100)], + ids=["matched", "sparse_state", "multiple", "non_multiple", "sparse_prop"], +) +def test_resume_preserves_monotonicity( + state_freq: int, + prop_freq: int, + si_double_sim_state: SimState, + lj_model: LennardJonesModel, + tmp_path: Path, +) -> None: + """Resumption preserves strict per-array step monotonicity for any combination + of cadences, including non-multiples.""" + traj_files = [str(tmp_path / f"monotonic_{idx}.h5") for idx in range(2)] + + state = _run_integrate( + si_double_sim_state, + lj_model, + traj_files, + n_steps=15, + state_freq=state_freq, + prop_freq=prop_freq, + ) + _run_integrate( + state, + lj_model, + traj_files, + n_steps=12, + state_freq=state_freq, + prop_freq=prop_freq, + mode="a", + ) + + for traj_file in traj_files: + for name, steps in _per_array_steps(traj_file).items(): + assert np.all(np.diff(steps) > 0), f"{name} steps not monotonic: {steps}" + + +def test_resume_does_not_discard_data( + si_double_sim_state: SimState, lj_model: LennardJonesModel, tmp_path: Path +) -> None: + """No array may lose rows across a close/reopen cycle.""" + traj_files = [str(tmp_path / f"no_discard_{idx}.h5") for idx in range(2)] + + state = _run_integrate( + si_double_sim_state, + lj_model, + traj_files, + n_steps=15, + state_freq=100, + prop_freq=10, + ) + before = {traj_file: _per_array_steps(traj_file) for traj_file in traj_files} + + _run_integrate( + state, + lj_model, + traj_files, + n_steps=12, + state_freq=100, + prop_freq=10, + mode="a", + ) + + for traj_file in traj_files: + after = _per_array_steps(traj_file) + for name, steps in before[traj_file].items(): + assert len(after[name]) >= len(steps) + assert after[name][-1] >= steps[-1] + # the pre-existing rows must be preserved verbatim + np.testing.assert_array_equal(after[name][: len(steps)], steps) + + +def test_resume_continues_on_absolute_grid( + si_double_sim_state: SimState, lj_model: LennardJonesModel, tmp_path: Path +) -> None: + """Each array resumes on its own absolute lattice.""" + traj_files = [str(tmp_path / f"abs_grid_{idx}.h5") for idx in range(2)] + state_freq, prop_freq = 10, 3 + + # both runs end on a step divisible by 10 and by 3, so no final frame is + # written off-grid and every step must lie on its array's own lattice + state = _run_integrate( + si_double_sim_state, + lj_model, + traj_files, + n_steps=30, + state_freq=state_freq, + prop_freq=prop_freq, + ) + _run_integrate( + state, + lj_model, + traj_files, + n_steps=30, + state_freq=state_freq, + prop_freq=prop_freq, + mode="a", + ) + + for traj_file in traj_files: + with TorchSimTrajectory(traj_file, mode="r") as traj: + pos_steps = traj.get_steps("positions") + energy_steps = traj.get_steps("potential_energy") + assert np.all(pos_steps % state_freq == 0) + assert np.all(energy_steps % prop_freq == 0) + # both grids span the interruption at step 30 + assert pos_steps[-1] == 60 + assert energy_steps[-1] == 60 + assert 30 in pos_steps.tolist() + assert 33 in energy_steps.tolist() + + +def test_truncate_aligns_all_arrays(test_file: Path, random_state: MDState) -> None: + """Explicit truncation trims every array, not just those on the positions grid.""" + with TorchSimTrajectory(test_file, mode="w") as traj: + for step in (0, 100): + traj.write_state(random_state, step) + for step in range(0, 121, 30): + traj.write_arrays({"potential_energy": torch.tensor([float(step)])}, step) + + assert traj.last_step == 100 + assert traj.last_written_step == 120 + + # align every array to the last positions frame + traj.truncate_to_step(100) + assert traj.last_written_step == 100 + for name in traj.array_registry: + assert traj.get_steps(name)[-1] <= 100 + + # a subsequent write at step 101 must now succeed + traj.write_state(random_state, 101) + traj.write_arrays({"potential_energy": torch.tensor([101.0])}, 101) + assert traj.last_written_step == 101 + + +def test_truncate_to_last_written_step_is_noop( + test_file: Path, random_state: MDState +) -> None: + """Truncating to a step no array exceeds leaves every array untouched.""" + with TorchSimTrajectory(test_file, mode="w") as traj: + for step in (0, 100): + traj.write_state(random_state, step) + for step in range(0, 121, 30): + traj.write_arrays({"potential_energy": torch.tensor([float(step)])}, step) + + before = {name: traj.get_steps(name).copy() for name in traj.array_registry} + traj.truncate_to_step(120) + for name, steps in before.items(): + np.testing.assert_array_equal(traj.get_steps(name), steps) + + +def test_mismatched_cadence_emits_no_warning( + si_double_sim_state: SimState, + lj_model: LennardJonesModel, + tmp_path: Path, + recwarn: pytest.WarningsRecorder, + caplog: pytest.LogCaptureFixture, +) -> None: + """Reopening a multi-cadence trajectory is not an error condition.""" + traj_files = [str(tmp_path / f"no_warn_{idx}.h5") for idx in range(2)] + _run_integrate( + si_double_sim_state, + lj_model, + traj_files, + n_steps=15, + state_freq=10, + prop_freq=1, + ) + + recwarn.clear() + with caplog.at_level("WARNING", logger="torch_sim.trajectory"): + for traj_file in traj_files: + TorchSimTrajectory(traj_file, mode="a").close() + + assert not [w for w in recwarn if "Inconsistent last steps" in str(w.message)] + assert not [rec for rec in caplog.records if "Inconsistent last steps" in rec.message] + + +def test_autobatcher_swap_does_not_warn( + lj_model: LennardJonesModel, + ar_supercell_sim_state: SimState, + tmp_path: Path, + recwarn: pytest.WarningsRecorder, + caplog: pytest.LogCaptureFixture, +) -> None: + """The InFlightAutoBatcher reopen path stays quiet under mismatched cadences.""" + torch.manual_seed(0) + # staggered perturbations make systems converge in different swaps, so that a + # surviving system's file is reopened mid-run with its property array ahead of + # its positions array - the only configuration that trips the removed check + substates = [] + for scale in (0.02, 0.05, 0.08, 0.12, 0.16, 0.2): + substate = ar_supercell_sim_state.clone() + substate.positions = ( + substate.positions + torch.randn_like(substate.positions) * scale + ) + substates.append(substate) + multi_state = ts.initialize_state(substates, lj_model.device, lj_model.dtype) + + traj_files = [ + str(tmp_path / f"swap_{idx}.h5") for idx in range(multi_state.n_systems) + ] + reporter = ts.TrajectoryReporter( + traj_files, + state_frequency=10, + prop_calculators={1: {"potential_energy": lambda state: state.energy}}, + ) + autobatcher = InFlightAutoBatcher( + model=lj_model, memory_scales_with="n_atoms", max_memory_scaler=70 + ) + + opened: set[str] = set() + reopened_mid_run: list[str] = [] + original_reopen = reporter.reopen_trajectories + + def spy(filenames: list[str]) -> None: + for name in map(str, filenames): + if name in opened: + reopened_mid_run.append(name) + opened.add(name) + original_reopen(filenames) + + reporter.reopen_trajectories = spy + + recwarn.clear() + with caplog.at_level("WARNING", logger="torch_sim.trajectory"): + ts.optimize( + system=multi_state, + model=lj_model, + optimizer=ts.Optimizer.fire, + convergence_fn=ts.generate_force_convergence_fn(force_tol=1e-1), + trajectory_reporter=reporter, + autobatcher=autobatcher, + max_steps=200, + steps_between_swaps=5, + ) + + assert reopened_mid_run, "no system survived a swap, scenario not exercised" + assert not [w for w in recwarn if "Inconsistent last steps" in str(w.message)] + assert not [rec for rec in caplog.records if "Inconsistent last steps" in rec.message] + + +# ---------------------------------------------------------------------------------- +# final frame recording +# ---------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("state_freq", "prop_freq", "n_steps"), + [(10, 10, 10), (10, 10, 13), (100, 10, 150), (10, 1, 15), (100, 30, 155)], + ids=["on_grid", "off_grid", "on_prop_grid", "dense_prop", "non_multiple"], +) +def test_final_frame_is_recorded( + state_freq: int, + prop_freq: int, + n_steps: int, + si_double_sim_state: SimState, + lj_model: LennardJonesModel, + tmp_path: Path, +) -> None: + """The last recorded frame is the final state, on-grid or not. + + The ``on_prop_grid`` and ``dense_prop`` cases end on the property cadence but + off the state cadence, so a guard keyed off the largest step across all arrays + would skip the final state write. + """ + traj_files = [str(tmp_path / f"final_frame_{idx}.h5") for idx in range(2)] + final_state = _run_integrate( + si_double_sim_state, + lj_model, + traj_files, + n_steps=n_steps, + state_freq=state_freq, + prop_freq=prop_freq, + ) + + for idx, traj_file in enumerate(traj_files): + with TorchSimTrajectory(traj_file, mode="r") as traj: + assert traj.get_steps("positions")[-1] == n_steps + assert traj.get_steps("potential_energy")[-1] == n_steps + last_positions = torch.tensor(traj.get_array("positions")[-1]) + torch.testing.assert_close( + last_positions.to(dtype=final_state.dtype), + final_state.split()[idx].positions, + ) + + +def test_no_duplicate_final_frame( + si_double_sim_state: SimState, lj_model: LennardJonesModel, tmp_path: Path +) -> None: + """A run ending exactly on-grid does not double-write its last frame.""" + traj_files = [str(tmp_path / f"no_dupe_{idx}.h5") for idx in range(2)] + _run_integrate( + si_double_sim_state, + lj_model, + traj_files, + n_steps=10, + state_freq=5, + prop_freq=5, + ) + + for traj_file in traj_files: + with TorchSimTrajectory(traj_file, mode="r") as traj: + np.testing.assert_array_equal(traj.get_steps("positions"), [0, 5, 10]) + np.testing.assert_array_equal(traj.get_steps("potential_energy"), [0, 5, 10]) + + +def test_optimize_records_converged_state( + si_double_sim_state: SimState, lj_model: LennardJonesModel, tmp_path: Path +) -> None: + """The converged geometry appears in its own trajectory.""" + traj_files = [str(tmp_path / f"opt_final_{idx}.h5") for idx in range(2)] + state = si_double_sim_state + state.positions += torch.randn_like(state.positions) * 0.05 + + reporter = ts.TrajectoryReporter(traj_files, state_frequency=10) + final_states = ts.optimize( + system=state, + model=lj_model, + optimizer=ts.Optimizer.fire, + convergence_fn=ts.generate_force_convergence_fn(force_tol=1e-1), + trajectory_reporter=reporter, + max_steps=100, + steps_between_swaps=5, + ) + + for idx, traj_file in enumerate(traj_files): + with TorchSimTrajectory(traj_file, mode="r") as traj: + last_positions = torch.tensor(traj.get_array("positions")[-1]) + torch.testing.assert_close( + last_positions.to(dtype=final_states.dtype), + final_states.split()[idx].positions, + ) + + +def test_optimize_final_frame_for_early_converging_systems( + lj_model: LennardJonesModel, + ar_supercell_sim_state: SimState, + fe_supercell_sim_state: SimState, + tmp_path: Path, +) -> None: + """Systems popped out mid-run still get their final frame written.""" + torch.manual_seed(0) + states = [ + ar_supercell_sim_state, + fe_supercell_sim_state, + ar_supercell_sim_state, + fe_supercell_sim_state, + ] + multi_state = ts.initialize_state(states, lj_model.device, lj_model.dtype) + for sub, scale in enumerate((0.01, 0.05, 0.1, 0.2)): + mask = multi_state.system_idx == sub + multi_state.positions[mask] += ( + torch.randn_like(multi_state.positions[mask]) * scale + ) + + traj_files = [ + str(tmp_path / f"early_conv_{idx}.h5") for idx in range(multi_state.n_systems) + ] + reporter = ts.TrajectoryReporter(traj_files, state_frequency=10) + autobatcher = InFlightAutoBatcher( + model=lj_model, memory_scales_with="n_atoms", max_memory_scaler=260 + ) + + final_states = ts.optimize( + system=multi_state, + model=lj_model, + optimizer=ts.Optimizer.fire, + convergence_fn=ts.generate_force_convergence_fn(force_tol=1e-1), + trajectory_reporter=reporter, + autobatcher=autobatcher, + max_steps=50, + steps_between_swaps=5, + ) + + # every file - not just those in the final batch - must end at its own + # system's final geometry + for idx, traj_file in enumerate(traj_files): + with TorchSimTrajectory(traj_file, mode="r") as traj: + last_positions = torch.tensor(traj.get_array("positions")[-1]) + torch.testing.assert_close( + last_positions.to(dtype=final_states.dtype), + final_states.split()[idx].positions, + ) diff --git a/torch_sim/runners.py b/torch_sim/runners.py index 7329db38..21eac389 100644 --- a/torch_sim/runners.py +++ b/torch_sim/runners.py @@ -115,12 +115,14 @@ def _determine_initial_step_for_integrate( check for resume information Returns: - int: The initial step to start from (1 if not resuming, otherwise last_step + 1) + int: The initial step to start from (1 if not resuming, otherwise the largest + step recorded in any array of any trajectory + 1) """ initial_step: int = 1 if trajectory_reporter is not None and trajectory_reporter.mode == "a": last_logged_steps = [ - step if step is not None else 0 for step in trajectory_reporter.last_steps + step if step is not None else 0 + for step in trajectory_reporter.last_written_steps ] last_logged_step = min(last_logged_steps) initial_step = initial_step + last_logged_step @@ -129,7 +131,7 @@ def _determine_initial_step_for_integrate( f"Trajectory files have different last steps: {set(last_logged_steps)} " "Cannot resume integration from inconsistent states." "You can truncate the trajectories to the same step using:\n\n" - " reporter.truncate_to_step(min(reporter.last_step))\n\n" + " reporter.truncate_to_step(min(reporter.last_written_steps))\n\n" "before calling integrate again." ) if last_logged_step > 0: @@ -161,7 +163,7 @@ def _determine_initial_step_for_optimize( size=(state.n_systems,), fill_value=1, dtype=torch.long, device=state.device ) if trajectory_reporter is not None and trajectory_reporter.mode == "a": - last_steps = trajectory_reporter.last_steps + last_steps = trajectory_reporter.last_written_steps last_steps = [step if step is not None else 0 for step in last_steps] last_logged_steps = torch.tensor( last_steps, dtype=torch.long, device=state.device @@ -243,10 +245,73 @@ def _write_initial_state( """ if trajectory_reporter: trajectories_empty = all( - traj.last_step is None for traj in trajectory_reporter.trajectories + traj.last_written_step is None for traj in trajectory_reporter.trajectories ) if trajectories_empty: - trajectory_reporter.report(state, 0, model=model) + trajectory_reporter.report(state, 0, model=model, force=True) + + +def _write_final_state( + trajectory_reporter: TrajectoryReporter | None, + state: SimState, + model: ModelInterface, + step: int, +) -> None: + """Write the final state if the run ended off the ``state_frequency`` grid. + + Args: + trajectory_reporter (TrajectoryReporter | None): Optional reporter + state (SimState): Final simulation state + model (ModelInterface): Model used for simulation + step (int): Final step of the run + """ + if not trajectory_reporter: + return + if all( + traj.last_step is not None and traj.last_step >= step + for traj in trajectory_reporter.trajectories + ): + return + trajectory_reporter.report(state, step, model=model, force=True) + + +def _write_final_states_for_converged( + trajectory_reporter: TrajectoryReporter | None, + converged_states: list[SimState], + og_indices: list[int], + og_filenames: list[str] | None, + model: ModelInterface, + step: torch.Tensor, +) -> None: + """Write the final frame of each newly converged system. + + ``optimize`` pops converged systems out of the batch mid-run and immediately + repoints the reporter at the remaining systems, so a single final write before + ``finish()`` would miss every system that converged earlier. This must therefore + be called while the reporter still holds open handles to the batch the systems + converged out of. + + Args: + trajectory_reporter (TrajectoryReporter | None): Optional reporter + converged_states (list[SimState]): Newly converged single-system states + og_indices (list[int]): Original index of each converged state + og_filenames (list[str] | None): Full list of trajectory filenames, indexed + by original system index + model (ModelInterface): Model used for optimization + step (torch.Tensor): Per-system step counter, indexed by original index. + Holds the *next* step, so the final written step is ``step - 1``. + """ + if not trajectory_reporter or not converged_states or og_filenames is None: + return + open_filenames = trajectory_reporter.filenames or [] + position = {str(name): idx for idx, name in enumerate(open_filenames)} + for og_idx, converged_state in zip(og_indices, converged_states, strict=True): + idx = position.get(str(og_filenames[og_idx])) + if idx is None: # file not currently open, cannot write without reopening + continue + trajectory_reporter.report_final_frame( + idx, converged_state, int(step[og_idx]) - 1, model=model + ) def integrate[T: SimState]( # noqa: C901, PLR0915 @@ -395,6 +460,12 @@ def integrate[T: SimState]( # noqa: C901, PLR0915 if trajectory_reporter: trajectory_reporter.report(state, step, model=model) + # ensure the final state is recorded even if it is off the cadence grid + if n_steps > 0: + _write_final_state( + trajectory_reporter, state, model, initial_step + n_steps - 1 + ) + # finish the trajectory reporter final_states.append(state) if tqdm_pbar: @@ -691,6 +762,23 @@ def optimize[T: OptimState]( # noqa: C901, PLR0915 while True: result = autobatcher.next_batch(state, convergence_tensor) + newly_converged = result[1] + # og indices of the states that just converged, in the same order + newly_converged_og_idx = ( + autobatcher.completed_idx_og_order[-len(newly_converged) :] + if newly_converged + else [] + ) + # must happen before reopen_trajectories, while the reporter still holds + # open handles to the batch these systems converged out of + _write_final_states_for_converged( + trajectory_reporter, + newly_converged, + newly_converged_og_idx, + og_filenames, + model, + step, + ) if result[0] is None: # All states have converged, collect the final converged states all_converged_states.extend(result[1]) diff --git a/torch_sim/trajectory.py b/torch_sim/trajectory.py index 338a827e..ef833a49 100644 --- a/torch_sim/trajectory.py +++ b/torch_sim/trajectory.py @@ -31,7 +31,6 @@ import inspect import logging import pathlib -import warnings from collections.abc import Callable, Mapping, Sequence from functools import partial from typing import TYPE_CHECKING, Any, Literal, Self @@ -227,7 +226,7 @@ def truncate_to_step(self, step: int) -> None: """ if step <= 0: raise ValueError(f"Step must be greater than 0. Got step={step}.") - last_steps = self.last_steps + last_steps = self.last_written_steps if any(s is None for s in last_steps): raise ValueError("Cannot truncate: one or more trajectories are empty.") if step > min(last_steps): @@ -265,7 +264,12 @@ def _add_model_arg_to_prop_calculators(self) -> None: self.prop_calculators[frequency][name] = new_fn def report( - self, state: SimState, step: int | list[int], model: ModelInterface | None = None + self, + state: SimState, + step: int | list[int], + model: ModelInterface | None = None, + *, + force: bool = False, ) -> list[dict[str, torch.Tensor]]: """Report a state and step to the trajectory files. @@ -283,6 +287,10 @@ def report( model (ModelInterface, optional): Model used for simulation. Defaults to None. Must be provided if any prop_calculators are provided. + force (bool): If True, bypass the frequency gates and write every array + whose last recorded step is below ``step``. Used to capture the + initial and final frames of a run, which need not lie on the + cadence grid. Defaults to False. Returns: list[dict[str, torch.Tensor]]: Map of property names to tensors for each @@ -311,33 +319,98 @@ def report( # Process each system separately for idx, substate in enumerate(split_states): sys_step = step[idx] if isinstance(step, list) else step - # Write state to trajectory if it's time - if self.state_frequency and sys_step % self.state_frequency == 0: - self.trajectories[idx].write_state( - substate, sys_step, **self.state_kwargs - ) + all_props.append( + self._report_system(idx, substate, sys_step, model, force=force) + ) + + return all_props + + def _report_system( # noqa: C901 + self, + idx: int, + substate: SimState, + sys_step: int, + model: ModelInterface | None = None, + *, + force: bool = False, + ) -> dict[str, torch.Tensor]: + """Write a single system's state and properties to its trajectory file. + + Args: + idx (int): Index of the trajectory file to write to + substate (SimState): Single-system state to write + sys_step (int): Step to label the written frames with + model (ModelInterface, optional): Model used for simulation + force (bool): Bypass the frequency gates, writing any array whose last + recorded step is below ``sys_step``. + + Returns: + dict[str, torch.Tensor]: Map of property names to tensors. + """ + trajectory = self.trajectories[idx] - all_state_props = {} - # Process property calculators for this system - for report_frequency, calculators in self.prop_calculators.items(): - if sys_step % report_frequency != 0 or report_frequency == 0: - continue - - # Calculate properties for this substate - props = {} - for prop_name, prop_fn in calculators.items(): - prop = prop_fn(substate, model) - if len(prop.shape) == 0: - prop = prop.unsqueeze(0) - props[prop_name] = prop - - # Write properties to this trajectory + # Write state to trajectory if it's time + if self.state_frequency and (force or sys_step % self.state_frequency == 0): + last_state_step = trajectory.last_step + already_written = ( + force and last_state_step is not None and last_state_step >= sys_step + ) + if not already_written: + trajectory.write_state(substate, sys_step, **self.state_kwargs) + + all_state_props: dict[str, torch.Tensor] = {} + # Process property calculators for this system + for report_frequency, calculators in self.prop_calculators.items(): + if report_frequency == 0: + continue + if not force and sys_step % report_frequency != 0: + continue + + # Calculate properties for this substate + props = {} + for prop_name, prop_fn in calculators.items(): + prop = prop_fn(substate, model) + if len(prop.shape) == 0: + prop = prop.unsqueeze(0) + props[prop_name] = prop + + # Write properties to this trajectory + if props: + all_state_props.update(props) + if force: + props = { + name: value + for name, value in props.items() + if (last := trajectory.last_step_of(name)) is None + or last < sys_step + } if props: - all_state_props.update(props) - self.trajectories[idx].write_arrays(props, sys_step) - all_props.append(all_state_props) + trajectory.write_arrays(props, sys_step) - return all_props + return all_state_props + + def report_final_frame( + self, + index: int, + state: SimState, + step: int, + model: ModelInterface | None = None, + ) -> None: + """Write the final frame of a single system, bypassing the cadence grid. + + Runs generally terminate at a step that is not a multiple of + ``state_frequency``, which would leave the final state absent from its own + trajectory. This writes it unconditionally, skipping any array that already + holds a step at or beyond ``step`` so a run that ends on-grid is not + double-written. + + Args: + index (int): Index of the trajectory file to write to + state (SimState): Single-system final state + step (int): Final step of that system + model (ModelInterface, optional): Model used for simulation + """ + self._report_system(index, state, step, model, force=True) def _extract_props_batched( self, @@ -442,6 +515,30 @@ def last_steps(self) -> list[int | None]: last_steps.append(traj.last_step) return last_steps + @property + def last_written_steps(self) -> list[int | None]: + """Get the largest step recorded across all arrays of each trajectory file. + + Unlike :attr:`last_steps`, which only tracks the ``positions`` cadence, this + is the correct basis for resuming a run when different arrays are written at + different frequencies. + + Returns: + list[int | None]: The largest recorded step number for each trajectory, + or None if the trajectory is empty. Returns an empty list if no + trajectories exist. + """ + if not self.trajectories: + return [] + last_steps = [] + for trajectory in self.trajectories: + if trajectory._file.isopen: + last_steps.append(trajectory.last_written_step) + else: + with TorchSimTrajectory(trajectory._file.filename, mode="r") as traj: + last_steps.append(traj.last_written_step) + return last_steps + def __enter__(self) -> Self: """Support the context manager protocol. @@ -537,18 +634,6 @@ def __init__( self.type_map = self._initialize_type_map( coerce_to_float32=coerce_to_float32, coerce_to_int32=coerce_to_int32 ) - if mode == "a" and self.last_step is not None: - inconsistent_step = any( - self.get_steps(name)[-1] > self.last_step for name in self.array_registry - ) - if inconsistent_step: - msg = ( - "Inconsistent last steps detected in trajectory arrays. " - "Truncating all arrays to the `positions` array's last step." - ) - warnings.warn(msg, UserWarning, stacklevel=2) - logger.warning(msg) - self.truncate_to_step(self.last_step) def _initialize_header(self, metadata: dict[str, str] | None = None) -> None: """Initialize the HDF5 file header with metadata. @@ -853,7 +938,48 @@ def last_step(self) -> int | None: """ if not self.array_registry or "positions" not in self.array_registry: return None - return self.get_steps("positions")[-1].item() + return self.last_step_of("positions") + + def last_step_of(self, name: str) -> int | None: + """Get the last recorded step of a single array. + + Uses a direct slice of the steps node rather than reading the whole + array, which is O(1) rather than O(rows). + + Args: + name (str): Name of the array + + Returns: + int | None: The last recorded step, or None if the array is absent + """ + if name not in self.array_registry: + return None + steps_node = self._file.get_node("/steps/", name=name) + if not isinstance(steps_node, tables.Array) or len(steps_node) == 0: + return None + return int(steps_node[-1]) + + @property + def last_written_step(self) -> int | None: + """Largest step recorded across *all* arrays. + + Unlike :attr:`last_step`, which only tracks the ``positions`` cadence, + this is the correct basis for resuming a run: writing strictly beyond it + guarantees the per-array step monotonicity that ``_validate_array`` + enforces, whatever cadence each individual array uses. + + Returns: + int | None: The largest recorded step number, or None if the + trajectory holds no data. + """ + if not self.array_registry: + return None + last_steps = [ + step + for step in (self.last_step_of(name) for name in self.array_registry) + if step is not None + ] + return max(last_steps) if last_steps else None def __str__(self) -> str: """Get a string representation of the trajectory. @@ -1200,17 +1326,18 @@ def truncate_to_step(self, step: int) -> None: Args: step (int): Desired last step of the trajectory after truncation """ - if self.last_step is None: + last_written_step = self.last_written_step + if last_written_step is None: raise ValueError( "Cannot truncate an empty trajectory (no data has been written)." ) - if self.last_step < step: + if last_written_step < step: raise ValueError( f"Cannot truncate to a step greater than the last step." - f" {self.last_step=} < {step=}" + f" {last_written_step=} < {step=}" ) - if self.last_step == step: - return # No truncation needed + if last_written_step == step: + return # no array holds steps beyond `step`, nothing to truncate if step <= 0: raise ValueError(f"Step must be larger than 0. Got {step=}") for name in self.array_registry: @@ -1221,9 +1348,10 @@ def truncate_to_step(self, step: int) -> None: if set(steps_data) == {0}: continue # skip global arrays # Find the index where the step is less than or equal to the desired step - # We know that it must be at least one index because of the earlier check. indices = np.where(steps_data <= step)[0] - length = indices[-1] + 1 # +1 because we want to include this index + # +1 because we want to include this index; an array written entirely + # after `step` is emptied. + length = int(indices[-1]) + 1 if len(indices) else 0 data_node = self._file.get_node("/data/", name=name) if isinstance(data_node, tables.EArray):