mythos.simulators.io
====================

.. py:module:: mythos.simulators.io

.. autoapi-nested-parse::

   Common data structures for simulator I/O.



Classes
-------

.. autoapisummary::

   mythos.simulators.io.SimulatorTrajectory


Functions
---------

.. autoapisummary::

   mythos.simulators.io._concat_optional_field
   mythos.simulators.io._merge_metadata


Module Contents
---------------

.. py:class:: SimulatorTrajectory

   Bases: :py:obj:`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.

   :param center: The center of mass positions for each rigid body at each
                  state in the trajectory.
   :param orientation: The orientations (as quaternions) for each rigid body at
                       each state in the trajectory.
   :param box_size: Optional box size associated with each state in the
                    trajectory.
   :param 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.
   :param temperature: Optional per-state temperature in kT (thermal energy in
                       simulation units). Shape ``(n_states,)``. When present,
                       ``beta = 1 / temperature`` can be used for reweighting. ``None``
                       indicates that the simulation temperature is unknown.


   .. py:attribute:: box_size
      :type:  mythos.utils.types.Arr_Box | None
      :value: None



   .. py:attribute:: temperature
      :type:  jax.numpy.ndarray | None
      :value: None



   .. py:attribute:: metadata
      :type:  dict[str, jax.numpy.ndarray] | None
      :value: None



   .. py:method:: from_rigid_body(rigid_body: jax_md.rigid_body.RigidBody, **kwargs: Any) -> SimulatorTrajectory
      :classmethod:


      Create a SimulatorTrajectory from a RigidBody instance.

      :param rigid_body: The RigidBody instance to create the SimulatorTrajectory from.
      :param \*\*kwargs: Additional keyword arguments to pass to the
      :param SimulatorTrajectory constructor.:

      :returns: A SimulatorTrajectory instance.



   .. py:method:: with_state_metadata(**metadata: dict[str, mythos.utils.types.ARR_OR_SCALAR]) -> SimulatorTrajectory

      Set the same metadata for all states in the trajectory.



   .. py:method:: filter(filter_fn: collections.abc.Callable[[Any], bool]) -> SimulatorTrajectory

      Filter the trajectory based on metadata.

      :param 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.



   .. py:method:: slice(key: int | slice | jax.numpy.ndarray | list) -> SimulatorTrajectory

      Slice the trajectory.



   .. py:method:: 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



   .. py:method:: concat(trajectories: list[SimulatorTrajectory]) -> SimulatorTrajectory
      :classmethod:


      Concatenate a list of SimulatorTrajectory instances.



   .. py:method:: __add__(other: SimulatorTrajectory) -> SimulatorTrajectory

      Concatenate two trajectories.



   .. py:method:: 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.

      :param filepath: The path to write the trajectory file to.
      :param box_size: The default box size in 3 dimensions to write to the file,
                       if trajectory box_size is not available. defaults to (0,0,0).



.. py:function:: _concat_optional_field(values: list[jax.numpy.ndarray | None], label: str) -> jax.numpy.ndarray | None

   Concatenate an optional per-trajectory field along axis 0.

   Returns ``None`` when all entries are ``None``, raises when the entries
   are a mix of ``None`` and non-``None``, and concatenates otherwise.


.. py:function:: _merge_metadata(metadata_list: list[dict[str, jax.numpy.ndarray] | None], lengths: list[int]) -> dict[str, jax.numpy.ndarray] | None

   Merge a list of metadata dictionaries for SimulatorTrajectory concatenation.

   If a key is missing in some dictionaries, it is filled with NaNs of
   the same shape (excluding leading axis which is num_states) as the
   corresponding array in the dictionaries that do have it. If a key is
   present in multiple dictionaries the shapes must be consistent beyond
   the leading axis.