Augmentation

metatrain.utils.augmentation.get_random_rotation()

metatrain.utils.augmentation.get_random_inversion()

class metatrain.utils.augmentation.RotationalAugmenter(target_info_dict: Dict[str, TargetInfo], extra_data_info_dict: Dict[str, TargetInfo] | None = None)

Bases: object

A class to apply random rotations and inversions to a set of systems and their targets.

Parameters:

• target_info_dict (Dict[str, TargetInfo]) – A dictionary mapping target names to their corresponding TargetInfo objects. This is used to determine the type of targets and how to apply the augmentations.

• extra_data_info_dict (Dict[str, TargetInfo] | None)

apply_random_augmentations(systems: List[System], targets: Dict[str, TensorMap], extra_data: Dict[str, TensorMap] | None = None) → Tuple[List[System], Dict[str, TensorMap], Dict[str, TensorMap]]

Apply a random augmentation to a number of System objects and its targets.

Parameters:

• systems (List[System]) – A list of System objects to be augmented.

• targets (Dict[str, TensorMap]) – A dictionary mapping target names to their corresponding TensorMap objects. These are the targets to be augmented.

• extra_data (Dict[str, TensorMap] | None)

Returns:

A tuple containing the augmented systems and targets.

Return type:

Tuple[List[System], Dict[str, TensorMap], Dict[str, TensorMap]]