Augmentation

class metatrain.utils.augmentation.O3Augmenter(target_info_dict: Dict[str, TargetInfo], extra_data_info_dict: Dict[str, TargetInfo] | None = None, group: str = 'O3')

Bases: object

Applies random O(3) transformations 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.

• extra_data_info_dict (Dict[str, TargetInfo] | None) – An optional dictionary mapping extra data names to their corresponding TargetInfo objects.

• group (str) – which transformations apply_random_augmentations() samples from: "O3" (the default) draws uniform rotations and improper rotations, while "inversions" draws only the identity and the inversion, for architectures that are already rotation-equivariant by construction.

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]]

Applies random O(3) augmentations to systems, targets, and optional extra data.

Parameters:

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

• targets (Dict[str, TensorMap]) – A dictionary mapping target names to TensorMap objects.

• extra_data (Dict[str, TensorMap] | None) – An optional dictionary of additional TensorMap objects to augment alongside targets.

Returns:

A tuple of augmented systems, targets, and extra data.

Return type:

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

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

Applies the given O(3) transformations to systems, targets, and optional extra data.

Parameters:

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

• targets (Dict[str, TensorMap]) – A dictionary mapping target names to TensorMap objects.

• transformations (List[Tensor]) – A list of 3x3 orthogonal torch.Tensor matrices, one per system. Matrices with determinant -1 are improper rotations.

• extra_data (Dict[str, TensorMap] | None) – An optional dictionary of additional TensorMap objects to augment alongside targets.

Returns:

A tuple of augmented systems, targets, and extra data.

Return type:

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