shnitsel.io.ase.write
=====================

.. py:module:: shnitsel.io.ase.write


Functions
---------

.. autoapisummary::

   shnitsel.io.ase.write._prepare_for_write_schnetpack
   shnitsel.io.ase.write._ndarray_to_json_ser
   shnitsel.io.ase.write._collect_metadata
   shnitsel.io.ase.write.write_ase_db


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

.. py:function:: _prepare_for_write_schnetpack(traj, leading_dim_name)

   Helper function to perform some preprocessing on the dataset before writing to a SchnetPack compatible database.

   Combines the dipole variables into one entry.

   :param traj: The Dataset to transform into a SchnetPack conforming format.
   :type traj: Trajectory
   :param leading_dim_name: The name of the leading dimension identifying different frames within the dataset. Depending on the setup, this should be 'frame' or 'time'.
   :type leading_dim_name: Literal['frame', 'time']

   :returns: The transformed dataset
   :rtype: Trajectory


.. py:function:: _ndarray_to_json_ser(value)

.. py:function:: _collect_metadata(traj, keys_to_write)

   Helper function to generate the SPaiNN Metadata dict from a Trajectory struct.

   Extracts info from attributes and variables to set up the dict.

   :param traj: The Dataset to extract the metadata from.
   :type traj: Trajectory

   :returns: The resulting metadata dictionary.
   :rtype: dict[str, Any]


.. py:function:: write_ase_db(traj, db_path, db_format, keys_to_write = None, preprocess = True)

   Function to write a Dataset into a ASE db in either SchNet or SPaiNN format.

   :param traj: The Dataset to be written to an ASE db style database
   :type traj: Trajectory
   :param db_path: Path to write the database to
   :type db_path: str
   :param db_format: Format of the target database. Used to control order of dimensions in data arrays. Can be either "schnet" or "spainn".
   :type db_format: Literal["schnet", "spainn";] | None
   :param keys_to_write: Optional parameter to restrict which data variables to . Defaults to None.
   :type keys_to_write: Collection | None, optional
   :param preprocess: _description_. Defaults to True.
   :type preprocess: bool, optional

   :raises ValueError: If neither `frame` nor `time` dimension is present on the dataset.
   :raises ValueError: If the `db_format` is neither `schnet`, `spainn` nor None

   .. rubric:: Notes

   See `https://spainn-md.readthedocs.io/en/latest/userguide/data_pipeline.html#generate-a-spainn-database` for details on SPaiNN format.