TrajectoryReporter

class torch_sim.trajectory.TrajectoryReporter(filenames, state_frequency=100, *, prop_calculators=None, state_kwargs=None, metadata=None, trajectory_kwargs=None)

Bases: object

Trajectory reporter for saving simulation data at specified intervals.

This class manages writing multiple trajectory files simultaneously. It handles periodic saving of full system states and custom property calculations.

Variables:

• state_frequency (int) – How often to save full states (in simulation steps)

• prop_calculators (dict) – Map of frequencies to property calculators

• state_kwargs (dict) – Additional arguments for state writing

• metadata (dict) – Metadata to save in trajectory files

• trajectories (list) – TorchSimTrajectory instances

• filenames (list) – Trajectory file paths

• array_registry (dict) – Map of array names to (shape, dtype) tuples

Parameters:

• filenames (str | Path | Sequence[str | Path] | None)

• state_frequency (int)

• prop_calculators (dict[int, dict[str, Callable]])

• state_kwargs (dict[str, Any])

• metadata (dict[str, str] | None)

• trajectory_kwargs (dict[str, Any])

Examples

>>> reporter = TrajectoryReporter(
...     ["system1.h5", "system2.h5"],
...     state_frequency=100,
...     prop_calculators={10: {"energy": calculate_energy}},
... )
>>> for step in range(1000):
...     # Run simulation step
...     state = step_fn(state)
...     reporter.report(state, step, model)
>>> reporter.close()

property filenames: list[str] | None

Get the list of trajectory filenames.

Returns:

List of trajectory file paths,

or None if no trajectories are loaded.

Return type:

list[str] | None

reopen_trajectories(filenames)

Closes any existing trajectory files and reopens new ones given by filenames.

Parameters:

filenames (str | pathlib.Path | list[str | pathlib.Path]) – Path(s) to save trajectory file(s)

Raises:

ValueError – If filenames are not unique

Return type:

None

property array_registry: dict[str, tuple[tuple[int, ...], dtype]]

Registry of array shapes and dtypes.

truncate_to_step(step)

Truncate all trajectory files to the specified step. WARNING: This operation is irreversible and will remove data from the trajectory files.

Parameters:

step (int) – The step to truncate to.

Return type:

None

report(state, step, model=None, *, force=False)

Report a state and step to the trajectory files.

Writes states and calculated properties to all trajectory files at the specified frequencies. Splits multi-system states across separate trajectory files. The number of systems must match the number of trajectory files.

Parameters:

• state (SimState) – Current system state with n_systems equal to len(filenames)

• step (int | list[int]) – Current simulation step per system, setting step to 0 will write the state and all properties. If a list is provided, it must have length equal to n_systems. Otherwise, a single integer step is broadcast to all systems.

• 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:

Map of property names to tensors for each

system.

Return type:

list[dict[str, Tensor]]

Raises:

ValueError – If number of systems doesn’t match number of trajectory files

report_final_frame(index, state, step, model=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.

Parameters:

• 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

Return type:

None

finish()

Finish writing the trajectory files.

Closes all open trajectory files.

Return type:

None

close()

Close all trajectory files.

Ensures all data is written to disk and releases the file handles.

Return type:

None

property mode: Literal['r', 'w', 'a']

Get the mode of the first trajectory file.

Returns:

Mode from the trajectory_kwargs used during initialization.

Return type:

”r” | ”w” | ”a”

property last_steps: list[int | None]

Get the last logged step across all trajectory files.

This is useful for resuming optimizations from where they left off.

Returns:

The last step number for each trajectory, or None if

the trajectory is empty. Returns empty list if no trajectories exist.

Return type:

list[int | None]

property last_written_steps: list[int | None]

Get the largest step recorded across all arrays of each trajectory file.

Unlike 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:

The largest recorded step number for each trajectory,

or None if the trajectory is empty. Returns an empty list if no trajectories exist.

Return type:

list[int | None]