Custom omegaconf functions

Resolvers to handle special fields in our configs as well as the expansion/completion of the dataset section.

metatrain.utils.omegaconf.default_device(_root_: BaseContainer) → str

Custom OmegaConf resolver to find the default device of an architecture.

Device is found using the :py:func:metatrain.utils.devices.pick_devices` function.

Parameters:

_root_ (BaseContainer)

Return type:

str

metatrain.utils.omegaconf.default_precision(_root_: BaseContainer) → int

Custom OmegaConf resolver to find the default precision of an architecture.

File format is obtained based on the architecture name and its first entry in the supported_dtypes list.

Parameters:

_root_ (BaseContainer)

Return type:

int

metatrain.utils.omegaconf.default_random_seed() → int

Return session seed in the range [0, 2**32).

Return type:

int

metatrain.utils.omegaconf.check_dataset_options(dataset_config: ListConfig) → None

Perform consistency checks within one dataset config.

This is useful if the dataset config is made of several datasets.

• The function checks if length_units in each system section are known and the

  same.

• For unknown quantities a warning is given.

• If the names of the "targets" sections are the same between the elements of

  the list of datasets also the units must be the same.

• Two targets with the names {target} and mtt::{target} are not allowed.

Parameters:

dataset_config (ListConfig) – A List of configuration to be checked. In the list contains only one element no checks are performed.

Raises:

ValueError – If the units are not consistent between the dataset options or if two different targets have the {target} and mtt::{target} names.

Return type:

None

metatrain.utils.omegaconf.expand_dataset_config(conf: str | DictConfig | ListConfig) → ListConfig

Expands shorthand notations in a dataset configuration to its full format.

This function takes a dataset configuration, either as a :py:class:str, omegaconf.DictConfig or a omegaconf.ListConfig, and expands it into a detailed configuration format. It processes systems, targets, and gradient sections, setting default values and inferring missing information. Unknown keys are ignored, allowing for flexibility.

If the dataset configuration is either a str or a omegaconf.DictConfig

The function performs the following steps for each config

• Merges and interpolates the input configuration with the base configurations.

• Expands shorthand notations like file paths or simple true/false settings to full dictionary systems. This includes setting the units to the base units of "angstrom" and "eV".

• Handles special cases, such as the mandatory nature of the “energy” section for MD simulations and the mutual exclusivity of ‘stress’ and ‘virial’ sections. Additionally the gradient sections for “forces” are enables by default.

Parameters:

conf (str | DictConfig | ListConfig) – The dataset configuration, either as a file path string or a DictConfig object.

Raises:

ValueError – If both virial and stress sections are enabled in the “energy” target, as this is not permissible for training.

Returns:

List of datasets configurations. If conf was a str or a omegaconf.DictConfig the list contains only a single element.

Return type:

ListConfig

metatrain.utils.omegaconf.check_units(actual_options: DictConfig | ListConfig, desired_options: DictConfig | ListConfig) → None

Perform consistency checks between two dataset configs.

Parameters:

• actual_options (DictConfig | ListConfig) – The dataset options that you want to test.

• desired_options (DictConfig | ListConfig) – The dataset options actual_options is tested against.

Raises:

• ValueError – If the length units are not consistent between the system in the dataset options.

• ValueError – If a target is present only in desider_option and not in actual_option.

• ValueError – If the unit of a target quantity is not consistent between the dataset option.

Return type:

None