shnitsel.io.ase.write

Functions

_prepare_for_write_schnetpack(traj, leading_dim_name)

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

_ndarray_to_json_ser(value)

_collect_metadata(traj, keys_to_write)

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

write_ase_db(traj, db_path, db_format[, ...])

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

Module Contents

_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.

Parameters:

  • traj (Trajectory) – The Dataset to transform into a SchnetPack conforming format.

  • leading_dim_name (Literal['frame', 'time']) – The name of the leading dimension identifying different frames within the dataset. Depending on the setup, this should be ‘frame’ or ‘time’.

Returns:

The transformed dataset

Return type:

Trajectory

_ndarray_to_json_ser(value)
_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.

Parameters:

  • traj (Trajectory) – The Dataset to extract the metadata from.

  • keys_to_write (Iterable[str])

Returns:

The resulting metadata dictionary.

Return type:

dict[str, Any]

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.

Parameters:

  • traj (Trajectory) – The Dataset to be written to an ASE db style database

  • db_path (str) – Path to write the database to

  • db_format (Literal["schnet", "spainn";] | None) – Format of the target database. Used to control order of dimensions in data arrays. Can be either “schnet” or “spainn”.

  • keys_to_write (Collection | None, optional) – Optional parameter to restrict which data variables to . Defaults to None.

  • preprocess (bool, optional) – _description_. Defaults to True.

Raises:

  • ValueError – If neither frame nor time dimension is present on the dataset.

  • ValueError – If the db_format is neither schnet, spainn nor None

Notes

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