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.

_write_trajectory_to_db(traj, db, keys_to_write)

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 (DataSeries) – 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:

xr.Dataset

_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. We expect trees to have been converted to common units before this is invoked.

Parameters:

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

  • keys_to_write (Collection[str], optional) – The keys of variables to write to the db.

Returns:

The resulting metadata dictionary.

Return type:

dict[str, Any]

write_ase_db(traj, db_path, db_format=None, keys_to_write=None, preprocess=True, force=False)#

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

Parameters:

  • traj (xr.Dataset | DataSeries | TreeNode[Any, DataSeries | xr.Dataset]) – The Dataset to be written to an ASE db style database. Can also be in tree format. If provided as a tree, the data must be consistent with each other, i.e. all coordinates except for the leading dimension must match. Inconsistencies

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

  • db_format (Literal["schnet", "spainn";] | None, optional) – 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) – Whether to apply preprocessing of the data. Defaults to True.

  • force (bool, optional) – A flag to force overwriting of an existing database at the position denoted by db_path.

Raises:

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

  • ValueError – If an unsupported data type was provided as an input.

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

  • ValueError – If the data in a provided tree is inconsistent.

Notes

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

_write_trajectory_to_db(traj, db, keys_to_write)#
Parameters: