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)

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.

• 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

system.

Return type:

list[dict[str, Tensor]]

Raises:

ValueError – If number of systems 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

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]