trajectory_setup#

Attributes#

TrajType

Classes#

RequiredTrajectorySettings

Helper class to make sure all required settings are applied and to

OptionalTrajectorySettings

Helper class to make sure consistent keys are used in trajectory settings

Functions#

assign_required_settings(dataset, settings)

Function to assign all required settings to the dataset.

assign_optional_settings(dataset, settings)

Function to assign all assigned optional settings to a dataset.

get_statecomb_coordinate(states)

Helper function to create a statecombination coordinate if it is missing, based on the states registered.

create_initial_dataset(num_time_steps, num_states, ...)

Function to initialize an xr.Dataset with appropriate variables and coordinates to acommodate loaded data.

fill_missing_dataset_variables(ds)

Helper function to fill in missing variables in datasets that we need for initialization.

Module Contents#

TrajType#
class RequiredTrajectorySettings#

Helper class to make sure all required settings are applied and to avoid typos when assigning required settings to consistent keys in trajectories.

t_max: float#
delta_t: float#
max_ts: int#
completed: bool#
input_format: Literal['sharc', 'newtonx', 'ase', 'pyrai2md']#
input_type: Literal['static', 'dynamic', 'unknown']#
input_format_version: str#
num_singlets: int#
num_doublets: int#
num_triplets: int#
class OptionalTrajectorySettings#

Helper class to make sure consistent keys are used in trajectory settings for optional configuration information.

has_forces: bool | Literal['all', 'active_only'] | None = None#
trajectory_id: int | None = None#
is_multi_trajectory: bool | None = None#
trajectory_input_path: str | None = None#
theory_basis_set: str | None = None#
est_level: str | None = None#
misc_input_settings: Dict | None = None#
assign_required_settings(dataset, settings)#

Function to assign all required settings to the dataset. Some of the settings may be applied to a coordinate variable of zero-dimensional shape instead.

Just a handy tool so all values are assigned because all fields in settings should be assigned upon its creation

Parameters:

  • dataset (xr.Dataset | ShnitselDataset) – The dataset/trajectory to write the required settings into

  • settings (RequiredTrajectorySettings) – The fully assigned settings object containing all keys and values to be assigned.

Return type:

TrajType

assign_optional_settings(dataset, settings)#

Function to assign all assigned optional settings to a dataset.

Just a handy tool so we can be sure the settings are assigned with the correct keys.

Parameters:

  • dataset (xr.Dataset | Trajectory | Frames) – The dataset/trajectory to write the optional settings into

  • settings (OptionalTrajectorySettings) – The dataclass object that has all optional setting keys with optional values. Only assigned settings (not None) will be inserted.

Return type:

TrajType

get_statecomb_coordinate(states)#

Helper function to create a statecombination coordinate if it is missing, based on the states registered.

Parameters:

states (xr.DataArray) – The state coordinate

Returns:

The new coordinate having all non-ordered state combinations

Return type:

xr.Coordinates

create_initial_dataset(num_time_steps, num_states, num_atoms, format_name, loading_parameters, **kwargs)#

Function to initialize an xr.Dataset with appropriate variables and coordinates to acommodate loaded data.

All arguments are used to accurately size the dimensions of the dataset or assign. Also returns the default attributes associated with the variables in the dataset for later assignment of certain values like “time” which is not initialized with values.

Parameters:

  • num_time_steps (int) – The number of expected time steps in this trajectory. Set to 0 to not create a “time” dimension.

  • num_states (int) – The number of states within the datasets.

  • num_atoms (int) – _description_

  • format_name (Literal["sharc", "newtonx", "ase", "pyrai2md"]) – The number of atoms within the datasets. Set to 0 to remove all observables tied to an “atom” index.

  • loading_parameters (LoadingParameters | None) – Optional parameters to influence import and object creation. Used to set some input units and default handlers.

Returns:

First An xarray Dataset with appropriately sized DataArrays and coordinates also including default attributes for all variables. Second the key-value dict, where the key is the name of standard variables/coordinates in Shnitsel terminology and the value is the dict of default attributes associated with this variable in this format.

Return type:

Tuple[xr.Dataset, dict]

fill_missing_dataset_variables(ds)#

Helper function to fill in missing variables in datasets that we need for initialization.

Parameters:

ds (xr.Dataset) – The dataset to complete with state variables and atom variables

Returns:

The completed dataset with default variables applied.

Return type:

xr.Dataset