mythos.simulators.io

Common data structures for simulator I/O.

Classes

SimulatorTrajectory

A trajectory of a simulation run.

Functions

_merge_metadata(→ dict[str, jax.numpy.ndarray] | None)

Merge two metadata dictionaries for SimulatorTrajectory concatenation.

Module Contents

class mythos.simulators.io.SimulatorTrajectory

Bases: jax_md.rigid_body.RigidBody

A trajectory of a simulation run.

This class extends jax_md.rigid_body.RigidBody to include optional metadata associated with each state in the trajectory. This object can also store data associated with a single state, but in such a case certain methods do not make sense (e.g. filtering or slicing). Such single-state usage is primarily intended for use within mapping functions.

Parameters:

• center – The center of mass positions for each rigid body at each state in the trajectory.

• orientation – The orientations (as quaternions) for each rigid body at each state in the trajectory.

• box_size – Optional box size associated with each state in the trajectory.

• metadata – Optional metadata associated with each state in the trajectory. This must be a dictionary where each value is a numerical array whose first axis has length corresponding to number of states.

box_size: mythos.utils.types.Arr_Box | None = None

metadata: dict[str, jax.numpy.ndarray] | None = None

classmethod from_rigid_body(rigid_body: jax_md.rigid_body.RigidBody, **kwargs: Any) → SimulatorTrajectory

Create a SimulatorTrajectory from a RigidBody instance.

Parameters:

• rigid_body – The RigidBody instance to create the SimulatorTrajectory from.

• **kwargs – Additional keyword arguments to pass to the

• constructor. (SimulatorTrajectory)

Returns:

A SimulatorTrajectory instance.

with_state_metadata(**metadata: dict[str, mythos.utils.types.ARR_OR_SCALAR]) → SimulatorTrajectory

Set the same metadata for all states in the trajectory.

filter(filter_fn: collections.abc.Callable[[Any], bool]) → SimulatorTrajectory

Filter the trajectory based on metadata.

Parameters:

filter_fn – A function that takes in metadata tree and returns a boolean array of length equal to the number of states, indicating which states to keep.

Returns:

A new SimulatorTrajectory with only the states that pass the filter.

slice(key: int | slice | jax.numpy.ndarray | list) → SimulatorTrajectory

Slice the trajectory.

length() → int

Return the length of the trajectory.

Note, that this may have been more natural to implement as the built-in __len__ method. However, the chex.dataclass decorator overrides that method to be compatabile with the abc.Mapping interface

See here: https://github.com/google-deepmind/chex/blob/8af2c9e8a19f3a57d9bd283c2a34148aef952f60/chex/_src/dataclass.py#L50

__add__(other: SimulatorTrajectory) → SimulatorTrajectory

Concatenate two trajectories.

to_file(filepath: pathlib.Path, box_size: mythos.utils.types.Vector3D = (0, 0, 0)) → None

Write the trajectory to an oxDNA file.

Note that the SimulatorTrajectory does not store several of the fields necessary to fully reconstruct an oxDNA trajectory file (e.g. times, box size, velocities, angular momenta, and energies). Thus, times are filled with a monotonic sequence, while the rest of these fields are filled with 0’s. The resultant file can be used for inspection and visualization of non-time-dependent state-by-state spatial information only.

Parameters:

• filepath – The path to write the trajectory file to.

• box_size – The box size in 3 dimensions to write to the file. defaults to (0,0,0).

mythos.simulators.io._merge_metadata(left: dict[str, jax.numpy.ndarray] | None, len_left: int, right: dict[str, jax.numpy.ndarray] | None, len_right: int) → dict[str, jax.numpy.ndarray] | None

Merge two metadata dictionaries for SimulatorTrajectory concatenation.

If a key is missing in one of the dictionaries, it is filled with NaNs of the same shape (excluding leading axis which is num_states) as the corresponding array in the other dictionary. If a key is present in both dictionaries the shapes must be consistent beyond the leading axis.