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

• shape_warned (bool) – Whether a shape warning has been issued

Parameters:

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

• state_frequency (int)

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

• state_kwargs (dict | None)

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

• trajectory_kwargs (dict | None)

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()

load_new_trajectories(filenames)

Load new trajectories into the reporter.

Closes any existing trajectory files and initializes new ones.

Parameters:

filenames (str | Path | list[str | 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.

report(state, step, model=None)

Report a state and step to the trajectory files.

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

Parameters:

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

• step (int) – Current simulation step, setting step to 0 will write the state and all properties.

• model (Module, optional) – Model used for simulation. Defaults to None. Must be provided if any prop_calculators are provided.

• write_to_file (bool, optional) – Whether to write the state to the trajectory files. Defaults to True. Should only be set to False if the props are being collected separately.

Returns:

Map of property names to tensors for each batch

Return type:

list[dict[str, Tensor]]

Raises:

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

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