diff --git a/docs/src/torch/reference/index.rst b/docs/src/torch/reference/index.rst index a689e439b..fd5e15c21 100644 --- a/docs/src/torch/reference/index.rst +++ b/docs/src/torch/reference/index.rst @@ -9,6 +9,7 @@ API reference systems models/index misc + o3 ase serialization diff --git a/docs/src/torch/reference/o3.rst b/docs/src/torch/reference/o3.rst new file mode 100644 index 000000000..25e2920c3 --- /dev/null +++ b/docs/src/torch/reference/o3.rst @@ -0,0 +1,87 @@ +O(3) transformations +===================== + +The :py:mod:`metatomic.torch.o3` module rotates and inverts +:py:class:`~metatomic.torch.System` and :py:class:`~metatensor.torch.TensorMap` +objects, for example to generate randomly rotated copies of a structure for data +augmentation. + +.. _o3-conventions: + +Conventions +----------- + +To transform a :py:class:`~metatensor.torch.TensorBlock`, :py:func:`transform_block` +and :py:func:`transform_tensor` need to know, for each component axis, whether it +carries a Cartesian or a spherical tensor. This is inferred from the axis name: + +- Cartesian axes are named ``xyz``, or ``xyz_1``, ``xyz_2``, ... for blocks with + several Cartesian axes (e.g. rank-2 Cartesian tensors). These are rotated directly + with the (3, 3) transformation matrix :math:`R` passed to + :py:class:`O3Transformation`, following the usual column-vector convention + :math:`v' = R v` (equivalently, since values are stored as rows, + ``values_transformed = values @ R.T``). +- Spherical axes are named ``o3_mu``, or ``o3_mu_1``, ``o3_mu_2``, ... They are rotated + with the real Wigner-D matrix for the angular momentum ``o3_lambda`` (respectively + ``o3_lambda_1``, ``o3_lambda_2``, ...) found in the block's key. Real, rather than + complex, spherical harmonics are used throughout, matching the convention used + elsewhere in metatomic and metatensor for spherical targets. + +Wigner-D matrices +~~~~~~~~~~~~~~~~~ + +Complex Wigner-D matrices are computed by the `wigners `_ +package, using the convention + +.. math:: + + D^{\ell}_{m',m}(\alpha, \beta, \gamma) = \langle \ell, m' | + e^{-i J_z \alpha} e^{-i J_y \beta} e^{-i J_z \gamma} | \ell, m \rangle, + +with ZYZ Euler angles :math:`(\alpha, \beta, \gamma)` extracted from the proper part of +:math:`R`, i.e. from :math:`R` itself when :math:`\det R = 1`, or from :math:`-R` when +:math:`\det R = -1`, such that this proper part equals :math:`R_z(\alpha) R_y(\beta) +R_z(\gamma)`. + +These are then converted to real Wigner-D matrices through the unitary +change of basis :math:`T` mapping complex spherical harmonics +:math:`Y_{\ell}^{m}` to real ones: + +.. math:: + + Y_{\ell, m}^{\text{real}} = \begin{cases} + \dfrac{1}{\sqrt{2}} \left( Y_{\ell}^{-m} + (-1)^m Y_{\ell}^{m} \right) + & m > 0 \\[6pt] + Y_{\ell}^{0} & m = 0 \\[6pt] + \dfrac{i}{\sqrt{2}} \left( Y_{\ell}^{m} - (-1)^m Y_{\ell}^{-m} \right) + & m < 0 + \end{cases} + +so that the real Wigner-D matrix is :math:`D^{\ell}_{\text{real}} = T^{*} D^{\ell} T^{T}`. + +For an improper rotation (a rotation composed with an inversion), Cartesian axes are +flipped as part of the transformation matrix itself, while spherical axes pick up an +extra parity factor :math:`(-1)^{\ell} \sigma`, where :math:`\ell` is ``o3_lambda`` +and :math:`\sigma` is the block's ``o3_sigma`` (its behavior under inversion, +1 or +-1), also read from the key. + +A :py:class:`~metatensor.torch.TensorBlock` in general contains values associated with +several systems. This information is contained in the ``"system"`` samples dimension; +each row is rotated with the transformation of the system it belongs to. Gradient blocks +are routed the same way, via their parent block's ``"system"`` column. When only one +system is being transformed, the ``"system"`` column is optional and, if present, is +ignored: every row is rotated with the (single) given transformation. + +Reference +--------- + +.. autoclass:: metatomic.torch.o3.O3Transformation + :members: + +.. autofunction:: metatomic.torch.o3.random_transformations + +.. autofunction:: metatomic.torch.o3.transform_system + +.. autofunction:: metatomic.torch.o3.transform_tensor + +.. autofunction:: metatomic.torch.o3.transform_block diff --git a/python/metatomic_torch/metatomic/torch/__init__.py b/python/metatomic_torch/metatomic/torch/__init__.py index a8bf363aa..06a9ae9c5 100644 --- a/python/metatomic_torch/metatomic/torch/__init__.py +++ b/python/metatomic_torch/metatomic/torch/__init__.py @@ -51,7 +51,10 @@ pick_device = torch.ops.metatomic.pick_device pick_output = torch.ops.metatomic.pick_output -from . import ase_calculator # noqa: F401 +from . import ( # noqa: F401 + ase_calculator, + o3, +) from .model import ( # noqa: F401 AtomisticModel, ModelInterface, diff --git a/python/metatomic_torch/metatomic/torch/o3/__init__.py b/python/metatomic_torch/metatomic/torch/o3/__init__.py new file mode 100644 index 000000000..40cb50baa --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/o3/__init__.py @@ -0,0 +1,25 @@ +""" +Apply O(3) transformations (rotations and improper rotations) to +:py:class:`metatomic.torch.System` and :py:class:`metatensor.torch.TensorMap`, for +example to augment training data with randomly rotated copies of a structure. + +See :ref:`o3-conventions` for the naming conventions used to identify Cartesian and +spherical components in a :py:class:`~metatensor.torch.TensorBlock`. +""" + +from ._tranformations import ( + O3Transformation, + random_transformations, + transform_block, + transform_system, + transform_tensor, +) + + +__all__ = [ + "O3Transformation", + "random_transformations", + "transform_system", + "transform_tensor", + "transform_block", +] diff --git a/python/metatomic_torch/metatomic/torch/o3/_tranformations.py b/python/metatomic_torch/metatomic/torch/o3/_tranformations.py new file mode 100644 index 000000000..239d0df68 --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/o3/_tranformations.py @@ -0,0 +1,533 @@ +""" +Rotate systems and tensor maps under O(3) transformations, routing rows of +multi-system tensors by their ``"system"`` sample label. +""" + +import torch +from metatensor.torch import Labels, LabelsEntry, TensorBlock, TensorMap + +from .. import System, register_autograd_neighbors +from ._wigner import build_wigner_D_cache + + +# Component-axis names recognised by the augmentation machinery. +_SUFFIXES = [""] + [f"_{i}" for i in range(1, 10)] +_CARTESIAN_AXES = frozenset({f"xyz{s}" for s in _SUFFIXES}) +_SPHERICAL_AXIS_TO_LAMBDA = {f"o3_mu{s}": f"o3_lambda{s}" for s in _SUFFIXES} +_SPHERICAL_AXIS_TO_SIGMA = {f"o3_mu{s}": f"o3_sigma{s}" for s in _SUFFIXES} + + +class O3Transformation: + """ + A single O(3) transformation, represented by a (3, 3) rotation or improper-rotation + matrix. + """ + + def __init__(self, matrix: torch.Tensor, max_angular_momentum: int): + """ + :param matrix: (3, 3) rotation or improper-rotation matrix + :param max_angular_momentum: maximum angular momentum of any spherical + representation to be transformed by this transformation; Wigner-D matrices + will be precomputed for all ``ell <= max_angular_momentum`` + """ + if matrix.shape != (3, 3): + raise ValueError( + f"Transformation has shape {tuple(matrix.shape)}; expected (3, 3)." + ) + + identity = torch.eye(3, device=matrix.device, dtype=matrix.dtype) + if not torch.allclose(matrix @ matrix.T, identity, atol=1e-5): + raise ValueError( + "Transformation is not orthogonal (R @ R.T deviates from I)." + ) + + self._matrix = matrix + self._max_angular_momentum = max_angular_momentum + self._is_inverted = bool(torch.det(matrix) < 0) + + self._wigner_D_cache = build_wigner_D_cache( + max_angular_momentum, + matrix, + device=matrix.device, + dtype=matrix.dtype, + ) + + @property + def matrix(self) -> torch.Tensor: + """The (3, 3) rotation or improper-rotation matrix.""" + return self._matrix + + @property + def dtype(self) -> torch.dtype: + """The dtype of the transformation matrix.""" + return self._matrix.dtype + + @property + def device(self) -> torch.device: + """The device of the transformation matrix.""" + return self._matrix.device + + @property + def is_inverted(self) -> bool: + """Whether this transformation is an improper rotation (det < 0).""" + return self._is_inverted + + def transform_cartesian(self, vectors: torch.Tensor) -> torch.Tensor: + """Apply the transformation to Cartesian vectors. + + :param vectors: (..., 3) tensor of Cartesian vectors + :return: (..., 3) tensor of transformed vectors + """ + return vectors @ self._matrix.T + + def transform_spherical( + self, values: torch.Tensor, ell: int, sigma: int + ) -> torch.Tensor: + """Apply the transformation to spherical values. + + :param values: (..., 2*ell+1) tensor of spherical values + :param ell: angular momentum of the spherical representation + :param sigma: inversion parity of the spherical representation + :return: (..., 2*ell+1) tensor of transformed spherical values + """ + D = self._wigner_D_cache.get(ell) + if D is None: + raise ValueError(f"Wigner-D matrix for ell={ell} not found in cache.") + transformed = values @ D.T + if sigma == -1: + transformed = transformed * ((-1) ** ell) + return transformed + + def wigner_D_matrix(self, ell: int): + """Return the Wigner-D matrix for this transformation and angular momentum ell. + + :param ell: angular momentum of the spherical representation + :return: (2*ell+1, 2*ell+1) Wigner-D matrix + """ + D = self._wigner_D_cache.get(ell) + if D is None: + raise ValueError(f"Wigner-D matrix for ell={ell} not found in cache.") + return D + + +def random_transformations( + n: int, + max_angular_momentum: int = 0, + *, + device: torch.device, + dtype: torch.dtype, + include_inversions: bool = False, + generator: torch.Generator | None = None, +) -> list[O3Transformation]: + """Sample ``n`` uniformly distributed O(3) transformations. + + Rotations are sampled from the Haar measure on SO(3) via random unit quaternions. + When ``include_inversions`` is ``True``, each matrix is independently negated with + probability 0.5, giving a uniform distribution over the full O(3) group. + + :param n: number of transformations to generate + :param max_angular_momentum: maximum angular momentum for Wigner-D matrices + :param device: target device for the output tensors + :param dtype: target dtype for the output tensors + :param include_inversions: if ``True``, sample from O(3) instead of SO(3) + :param generator: optional :class:`torch.Generator` for reproducible sampling; when + ``None`` the global RNG is used + :return: list of ``n`` orthogonal (3, 3) tensors + """ + q = torch.randn(n, 4, device=device, dtype=dtype, generator=generator) + q = q / q.norm(dim=1, keepdim=True) + w, x, y, z = q.unbind(1) + # Quaternion to rotation matrix (standard formula) + R = torch.stack( + [ + 1 - 2 * (y * y + z * z), + 2 * (x * y - w * z), + 2 * (x * z + w * y), + 2 * (x * y + w * z), + 1 - 2 * (x * x + z * z), + 2 * (y * z - w * x), + 2 * (x * z - w * y), + 2 * (y * z + w * x), + 1 - 2 * (x * x + y * y), + ], + dim=1, + ).reshape(n, 3, 3) + + if include_inversions: + signs = torch.randint(0, 2, (n,), device=device, generator=generator) * 2 - 1 + R = R * signs.to(dtype=dtype).reshape(n, 1, 1) + + return [O3Transformation(r, max_angular_momentum) for r in R.unbind(0)] + + +def _block_row_indices_by_system( + block: TensorBlock, + system_ids: torch.Tensor, +) -> list[torch.Tensor]: + """Return row-index tensors into ``block.values``, one per system. + + With a single system every row belongs to it (any ``"system"`` label is ignored). + + With several systems the ``"system"`` column is required and each row is routed by + matching its ``"system"`` label against ``system_ids``. + + :param block: block whose ``samples`` may contain a ``"system"`` column + :param system_ids: one value per system; ``system_ids[i]`` is the value of the + ``"system"`` column identifying the rows belonging to the i-th entry in the + ``systems`` list passed to ``transform_tensor`` or ``transform_block`` + :return: list of length ``len(system_ids)``; entry ``i`` selects the rows of system + ``i`` + """ + n_systems = len(system_ids) + if n_systems == 1: + return [torch.arange(block.values.shape[0], device=block.values.device)] + + if "system" not in block.samples.names: + raise ValueError( + "Rotational augmentation expects output samples to include a 'system' " + "dimension when transforming multiple systems." + ) + system_column = block.samples.column("system").to(dtype=torch.long) + distinct = torch.unique(system_column) + covered = torch.isin(distinct, system_ids) + if not covered.all(): + uncovered = distinct[~covered] + raise ValueError( + f"Block samples contain system labels {uncovered.tolist()} that are " + f"not in system_ids={system_ids.tolist()}. Every sample must be " + f"assigned to a system in the transformation." + ) + return [ + torch.nonzero(system_column == label, as_tuple=False).reshape(-1) + for label in system_ids + ] + + +def _gradient_row_indices_by_system( + grad_block: TensorBlock, + parent_block: TensorBlock, + system_ids: torch.Tensor, +) -> list[torch.Tensor]: + """Return row-index tensors into a gradient block, one per system. + + Gradient samples carry a ``"sample"`` column indexing into the parent block rather + than their own ``"system"`` column, so the system of each gradient row is read from + the parent block's ``"system"`` column. + + :param grad_block: gradient block to route + :param parent_block: the value block this gradient is attached to + :param system_ids: one value per system; ``system_ids[i]`` is the value of the + ``"system"`` column identifying the rows belonging to the i-th entry in the + ``systems`` list passed to ``transform_tensor`` or ``transform_block`` + :return: list of length ``len(system_ids)``; entry ``i`` selects the gradient + rows of system ``i`` + """ + n_systems = len(system_ids) + if n_systems == 1: + return [ + torch.arange(grad_block.values.shape[0], device=grad_block.values.device) + ] + if "system" not in parent_block.samples.names: + raise ValueError( + "Rotational augmentation expects the values samples to include a 'system' " + "dimension when transforming gradients of multiple systems." + ) + parent_system = parent_block.samples.column("system").to(dtype=torch.long) + sample_index = grad_block.samples.column("sample").to(dtype=torch.long) + grad_system = parent_system[sample_index] + return [ + torch.nonzero(grad_system == label, as_tuple=False).reshape(-1) + for label in system_ids + ] + + +def transform_system(system: System, transformation: O3Transformation) -> System: + """Apply an O(3) transformation to a single System. + + This function will transform positions, cell vectors, neighbor-list displacement + vectors, and any custom data. + + :param system: input system + :param transformation: O(3) transformation to apply + :return: new System with transformed geometry + """ + if ( + system.positions.dtype != transformation.dtype + or system.positions.device != transformation.device + ): + raise ValueError( + f"System has positions with dtype/device " + f"({system.positions.dtype}, {system.positions.device}) differing " + f"from the transformations ({transformation.dtype}, " + f"{transformation.device})." + ) + + new_system = System( + positions=transformation.transform_cartesian(system.positions), + types=system.types, + cell=transformation.transform_cartesian(system.cell), + pbc=system.pbc, + ) + + for data_name in system.known_data(): + data = system.get_data(data_name) + new_system.add_data( + data_name, transform_tensor(data, [system], [transformation]) + ) + + for options in system.known_neighbor_lists(): + neighbors = system.get_neighbor_list(options) + # neighbor vectors are stored as (N, 3, 1); squeeze/unsqueeze around the matmul + neighbors_values = neighbors.values.squeeze(-1) + new_values = transformation.transform_cartesian(neighbors_values) + rotated_neighbors = TensorBlock( + values=new_values.unsqueeze(-1), + samples=neighbors.samples, + components=neighbors.components, + properties=neighbors.properties, + ) + register_autograd_neighbors(new_system, rotated_neighbors) + new_system.add_neighbor_list(options, rotated_neighbors) + + return new_system + + +def _contract_component_axes( + values: torch.Tensor, + matrices: list[torch.Tensor], +) -> torch.Tensor: + """Rotate each component axis of ``values`` by its matrix. + + ``values`` has shape ``(n_rows, d_1, ..., d_k, n_properties)`` and ``matrices[j]`` + (shape ``(d_j, d_j)``) is contracted with component axis ``j`` as + ``out[..., A, ...] = sum_a matrices[j][A, a] * values[..., a, ...]``. + + :param values: values tensor of a value or gradient block + :param matrices: one rotation matrix per component axis (empty for scalars) + :return: rotated values, same shape as the input + """ + # einsum index letters for the per-axis component contraction (input lower, output + # upper). Six axes is far more than any realistic block (value + gradient) needs. + _EINSUM_IN = "abcdef" + _EINSUM_OUT = "ABCDEF" + + if len(matrices) == 0: + return values + n_axes = len(matrices) + in_subscript = "i" + _EINSUM_IN[:n_axes] + "p" + out_subscript = "i" + _EINSUM_OUT[:n_axes] + "p" + matrix_subscripts = [_EINSUM_OUT[j] + _EINSUM_IN[j] for j in range(n_axes)] + equation = ",".join(matrix_subscripts + [in_subscript]) + "->" + out_subscript + return torch.einsum(equation, *matrices, values) + + +def _axis_matrices_and_parity( + components: list[Labels], + key: LabelsEntry, + transformation: O3Transformation, +) -> tuple[list[torch.Tensor], int]: + """Pick the rotation matrix for each component axis and the O(3) inversion parity. + + Cartesian axes (``xyz``/``xyz_1``/``xyz_2``/...) use the rotation directly, so + improper rotations flip vectors automatically. Spherical axes + (``o3_mu``/``o3_mu_1``/ ``o3_mu_2``) use the proper-rotation Wigner-D matrix of the + matching ``o3_lambda`` plus a ``(-1)^ell * sigma`` parity factor accumulated + whenever ``transformation`` is improper. + + :param name: TensorMap name, used only in error messages + :param components: component :class:`Labels` of the block (or gradient block) + :param key: the parent block's key, supplying ``o3_lambda``/``o3_sigma`` values + :param transformation: this system's O3 transformation + :return: ``(matrices, parity)`` with one matrix per component axis + """ + matrices: list[torch.Tensor] = [] + parity = 1 + for component in components: + axis_name = component.names[0] + if axis_name in _CARTESIAN_AXES: + matrices.append(transformation.matrix) + elif axis_name in _SPHERICAL_AXIS_TO_LAMBDA: + ell = int(key[_SPHERICAL_AXIS_TO_LAMBDA[axis_name]]) + matrices.append(transformation.wigner_D_matrix(ell)) + if transformation.is_inverted: + sigma = int(key[_SPHERICAL_AXIS_TO_SIGMA[axis_name]]) + parity *= ((-1) ** ell) * sigma + else: + raise ValueError( + f"Found a component axis '{axis_name}', which is neither a Cartesian " + "('xyz'/'xyz_1'/'xyz_2'/...) nor spherical ('o3_mu'/'o3_mu_1'/...) " + "axis; it can not be transformed." + ) + return matrices, parity + + +def _transform_component_values( + values: torch.Tensor, + components: list[Labels], + key: LabelsEntry, + row_indices: list[torch.Tensor], + transformations: list[O3Transformation], +) -> torch.Tensor: + """Rotate the values of a single value or gradient block, per system. + + :param name: TensorMap name, used only in error messages + :param values: the block's values tensor + :param components: the block's component :class:`Labels` + :param key: the parent block's key (for spherical ``o3_lambda``/``o3_sigma``) + :param row_indices: per-system row indices into ``values`` + :param transformations: per-system (3, 3) transformation matrices + :param wigner_D_matrices: ``{ell: [D_0, ..., D_{N-1}]}`` real Wigner-D matrices + :return: new values tensor with each system's rows rotated + """ + new_values = values.clone() + for system_index, rows in enumerate(row_indices): + if len(rows) == 0: + continue + matrices, parity = _axis_matrices_and_parity( + components, key, transformations[system_index] + ) + rotated = _contract_component_axes(values[rows], matrices) + if parity != 1: + rotated = rotated * parity + new_values[rows] = rotated + return new_values + + +def transform_block( + key: LabelsEntry, + block: TensorBlock, + systems: list[System], + transformations: list[O3Transformation], + system_ids: list[int] | torch.Tensor | None = None, +) -> TensorBlock: + """Rotate one block and all of its gradients. + + :param key: the block's key, containing ``o3_lambda``/``o3_sigma`` values for + spherical blocks + :param block: the block to rotate + :param systems: list of systems, one per transformation + :param transformations: per-system O(3) transformation matrices + :param system_ids: index of the ``systems`` used in the samples of ``block``; + defaults to ``list(range(len(systems)))`` + :return: new block with rotated values and gradients + """ + assert len(systems) == len(transformations) + if len(systems) == 0: + return block + + if system_ids is None: + system_ids = list(range(len(systems))) + + if not isinstance(system_ids, torch.Tensor): + system_ids = torch.tensor( + system_ids, dtype=torch.int32, device=block.values.device + ) + + if ( + block.values.dtype != transformations[0].dtype + or block.values.device != transformations[0].device + ): + raise ValueError( + f"TensorMap has values with dtype/device " + f"({block.values.dtype}, {block.values.device}) differing " + f"from the transformations ({transformations[0].dtype}, " + f"{transformations[0].device})." + ) + + return _transform_block_impl(key, block, systems, transformations, system_ids) + + +def _transform_block_impl( + key: LabelsEntry, + block: TensorBlock, + systems: list[System], + transformations: list[O3Transformation], + system_ids: torch.Tensor, +) -> TensorBlock: + """Implementation of ``transform_block`` without the initial checks""" + row_indices = _block_row_indices_by_system(block, system_ids) + new_block = TensorBlock( + values=_transform_component_values( + block.values, + block.components, + key, + row_indices, + transformations, + ), + samples=block.samples, + components=block.components, + properties=block.properties, + ) + for gradient_name in block.gradients_list(): + grad_block = block.gradient(gradient_name) + grad_rows = _gradient_row_indices_by_system(grad_block, block, system_ids) + new_block.add_gradient( + gradient_name, + TensorBlock( + values=_transform_component_values( + grad_block.values, + grad_block.components, + key, + grad_rows, + transformations, + ), + samples=grad_block.samples, + components=grad_block.components, + properties=grad_block.properties, + ), + ) + return new_block + + +def transform_tensor( + tensor: TensorMap, + systems: list[System], + transformations: list[O3Transformation], + system_ids: list[int] | torch.Tensor | None = None, +) -> TensorMap: + """Rotate every block (and its gradients) of a TensorMap. + + The tensor character of each component axis is inferred from its name, so a single + :py:class:`TensorMap` may freely mix scalar, Cartesian and spherical blocks, and + blocks may carry gradients. + + One of the samples dimensions must be ``"system"``, which is used to route each row + to the correct system and transformation. + + :param tensor: TensorMap to rotate + :param systems: input systems + :param transformations: per-system O(3) transformation matrices + :param system_ids: index of the ``systems`` used in the samples of ``tensor``; + defaults to ``list(range(len(systems)))`` + :return: new TensorMap with rotated values and gradients + """ + assert len(systems) == len(transformations) + if len(systems) == 0: + return tensor + + if system_ids is None: + system_ids = list(range(len(systems))) + + if not isinstance(system_ids, torch.Tensor): + system_ids = torch.tensor( + system_ids, dtype=torch.int32, device=transformations[0].device + ) + + if len(tensor) != 0: + block = tensor.block(0) + if ( + block.values.dtype != transformations[0].dtype + or block.values.device != transformations[0].device + ): + raise ValueError( + f"TensorMap has values with dtype/device " + f"({block.values.dtype}, {block.values.device}) differing " + f"from the transformations ({transformations[0].dtype}, " + f"{transformations[0].device})." + ) + + new_blocks = [ + _transform_block_impl(key, block, systems, transformations, system_ids) + for key, block in tensor.items() + ] + return TensorMap(keys=tensor.keys, blocks=new_blocks) diff --git a/python/metatomic_torch/metatomic/torch/o3/_wigner.py b/python/metatomic_torch/metatomic/torch/o3/_wigner.py new file mode 100644 index 000000000..6d6594b48 --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/o3/_wigner.py @@ -0,0 +1,213 @@ +"""Private helpers to build real Wigner-D matrices for augmentation. + +Complex Wigner-D matrices are delegated to :func:`wigners.wigner_D_array` using +ZYZ Euler angles. This module then reshapes/indexes these matrices into the flat +layout expected by the augmentation code and applies the complex-to-real +change-of-basis. + +The indexing convention used here matches ``wigners``: +``D[mp + ell, m + ell] == D^ell_{mp,m}``. +""" + +import functools + +import numpy as np +import torch +from wigners import wigner_D_array + + +def _wigner_d_size(ell_min: int, mp_max: int, ell_max: int) -> int: + if mp_max >= ell_max: + return ( + ell_max * (ell_max * (4 * ell_max + 12) + 11) + + ell_min * (1 - 4 * ell_min**2) + + 3 + ) // 3 + if mp_max > ell_min: + return ( + 3 * ell_max * (ell_max + 2) + + ell_min * (1 - 4 * ell_min**2) + + mp_max + * (3 * ell_max * (2 * ell_max + 4) + mp_max * (-2 * mp_max - 3) + 5) + + 3 + ) // 3 + + return (ell_max * (ell_max + 2) - ell_min**2) * (1 + 2 * mp_max) + 2 * mp_max + 1 + + +def _wigner_d_index(ell: int, mp: int, m: int, ell_min: int, mp_max: int) -> int: + idx = 0 + for ell_prev in range(ell_min, ell): + local_mp_max = mp_max if mp_max < ell_prev else ell_prev + idx += (2 * local_mp_max + 1) * (2 * ell_prev + 1) + + local_mp_max = mp_max if mp_max < ell else ell + idx += (mp + local_mp_max) * (2 * ell + 1) + idx += m + ell + return idx + + +def _compute_wigner_d_complex( + ell_max: int, alpha: float, beta: float, gamma: float +) -> np.ndarray: + """Compute complex Wigner-D matrix elements for all ell in ``[0, ell_max]``. + + Calls :func:`wigners.wigner_D_array` for each input angle triplet, then packs + values into a flat output array indexed by :func:`_wigner_d_index`. + + :param ell_max: maximum angular-momentum order + :param alpha: ZYZ first rotation angle, arbitrary shape + :param beta: ZYZ second rotation angle, same shape as ``alpha`` + :param gamma: ZYZ third rotation angle, same shape as ``alpha`` + :return: complex array of shape ``(*alpha.shape, dsize)`` + """ + + mp_max = ell_max + dsize = _wigner_d_size(0, mp_max, ell_max) + result = np.zeros(dsize, dtype=np.complex128) + + _wigner_D_array = wigner_D_array(ell_max, alpha, beta, gamma) + + for ell in range(0, ell_max + 1): + for mp in range(-ell, ell + 1): + i_d = _wigner_d_index(ell, mp, -ell, 0, mp_max) + for m in range(-ell, ell + 1): + result[i_d] = _wigner_D_array[ell][mp + ell, m + ell] + i_d += 1 + + return result + + +def _compute_complex_wigner_d_matrices( + ell_max: int, + angles: tuple[float, float, float], +) -> dict[int, np.ndarray]: + """Return complex Wigner-D matrices for ell in ``[0, ell_max]`` at the given ZYZ + angles. + + The returned matrices follow the standard ``wigners`` indexing convention: + ``matrix[..., mp + ell, m + ell] == D^ell_{mp,m}``. + + :param ell_max: maximum angular-momentum order + :param angles: ``(alpha, beta, gamma)`` ZYZ Euler-angles + :return: ``{ell: array of shape (2*ell+1, 2*ell+1)}`` + """ + alpha, beta, gamma = angles + raw = _compute_wigner_d_complex(ell_max, alpha, beta, gamma) + matrices: dict[int, np.ndarray] = {} + for ell in range(ell_max + 1): + block = np.zeros((2 * ell + 1, 2 * ell + 1), dtype=np.complex128) + for mp in range(-ell, ell + 1): + for m in range(-ell, ell + 1): + block[mp + ell, m + ell] = raw[_wigner_d_index(ell, mp, m, 0, ell_max)] + matrices[ell] = block + return matrices + + +def _compute_real_wigner_d_matrices( + ell_max: int, + angles: tuple[float, float, float], + complex_to_real: dict[int, np.ndarray], +) -> dict[int, torch.Tensor]: + """Convert complex Wigner-D matrices to real ones using the provided + change-of-basis. + + :param ell_max: maximum angular-momentum order + :param angles: ``(alpha, beta, gamma)`` ZYZ Euler-angles + :param complex_to_real: ``{ell: (2*ell+1, 2*ell+1)}`` unitary transform from complex + to real spherical harmonics + :return: ``{ell: real tensor of shape (*angles[0].shape, 2*ell+1, 2*ell+1)}`` + """ + complex_matrices = _compute_complex_wigner_d_matrices(ell_max, angles) + real_matrices: dict[int, torch.Tensor] = {} + for ell, matrix in complex_matrices.items(): + transform = complex_to_real[ell] + matrix = np.einsum("ij,...jk,kl->...il", transform.conj(), matrix, transform.T) + # The complex-to-real basis change can leave tiny imaginary residuals; + # scale the tolerance by matrix magnitude instead of using a fixed atol. + scale = float(np.max(np.abs(matrix.real))) if matrix.size else 1.0 + atol = max(1e-9, scale * 1e-10) + if not np.allclose(matrix.imag, 0.0, atol=atol): + raise ValueError("real Wigner matrix conversion produced complex values") + real_matrices[ell] = torch.from_numpy(matrix.real) + return real_matrices + + +@functools.lru_cache(maxsize=None) +def _complex_to_real_spherical_harmonics_transform(ell: int) -> np.ndarray: + """Return the complex-to-real spherical-harmonics transform for one ``ell``. + + The returned matrix has shape ``(2*ell+1, 2*ell+1)``. + """ + if ell < 0 or not isinstance(ell, int): + raise ValueError("ell must be a non-negative integer.") + + size = 2 * ell + 1 + T = np.zeros((size, size), dtype=complex) + + for m in range(-ell, ell + 1): + m_index = m + ell + if m > 0: + T[m_index, ell + m] = 1 / np.sqrt(2) * (-1) ** m + T[m_index, ell - m] = 1 / np.sqrt(2) + elif m < 0: + T[m_index, ell + abs(m)] = -1j / np.sqrt(2) * (-1) ** m + T[m_index, ell - abs(m)] = 1j / np.sqrt(2) + else: + T[m_index, ell] = 1 + + return T + + +def _rotation_to_angles( + rotation: torch.Tensor, +) -> tuple[float, float, float]: + """ + Decompose an O(3) rotation matrix into ZYZ Euler angles :math:`(\\alpha, \\beta, + \\gamma)`. + + For improper rotations (det < 0) the proper part ``-R`` is decomposed; the inversion + parity factor is handled separately when applying Wigner-D matrices. + """ + + rotation = rotation if torch.det(rotation) > 0 else -rotation + # R = Rz(alpha) Ry(beta) Rz(gamma): element [2,2] = cos(beta) + cos_beta = rotation[2, 2].clamp(-1.0, 1.0) + beta = torch.arccos(cos_beta) + sin_beta = torch.sin(beta) + beta = float(beta) + if abs(sin_beta) < 1e-10: + # Gimbal lock: only alpha +/- gamma is defined; fix gamma=0 + if cos_beta > 0: + alpha = float(torch.atan2(rotation[1, 0], rotation[0, 0])) + else: + alpha = float(torch.atan2(-rotation[1, 0], -rotation[0, 0])) + gamma = 0.0 + else: + # R[0,2]=cos(alpha)*sin(beta), R[1,2]=sin(alpha)*sin(beta): alpha via atan2 + # R[2,1]=sin(beta)*sin(gamma), R[2,0]=-sin(beta)*cos(gamma): gamma via atan2 + alpha = float(torch.atan2(rotation[1, 2], rotation[0, 2])) + gamma = float(torch.atan2(rotation[2, 1], -rotation[2, 0])) + + return alpha, beta, gamma + + +def build_wigner_D_cache( + o3_lambda_max: int, + matrix: torch.Tensor, + device: torch.device, + dtype: torch.dtype, +) -> dict[int, torch.Tensor]: + """ + Build the cache of real Wigner-D matrices for ``ell = 0..o3_lambda_max`` at the ZYZ + Euler angles corresponding to the given rotation matrix, using cached + complex-to-real transform per ell. + """ + angles = _rotation_to_angles(matrix) + complex_to_real = { + ell: _complex_to_real_spherical_harmonics_transform(ell) + for ell in range(o3_lambda_max + 1) + } + cache = _compute_real_wigner_d_matrices(o3_lambda_max, angles, complex_to_real) + + return {ell: tensor.to(device=device, dtype=dtype) for ell, tensor in cache.items()} diff --git a/python/metatomic_torch/setup.py b/python/metatomic_torch/setup.py index bc92b5b6e..01de7f38c 100644 --- a/python/metatomic_torch/setup.py +++ b/python/metatomic_torch/setup.py @@ -320,6 +320,7 @@ def create_version_number(version): f"torch {torch_version}", "metatensor-torch >=0.10.0,<0.11", "metatensor-operations >=0.5.0,<0.6", + "wigners", ] # when packaging a sdist for release, we should never use local dependencies diff --git a/python/metatomic_torch/tests/o3.py b/python/metatomic_torch/tests/o3.py new file mode 100644 index 000000000..6f0fea937 --- /dev/null +++ b/python/metatomic_torch/tests/o3.py @@ -0,0 +1,542 @@ +import re + +import numpy as np +import pytest +import torch +from metatensor.torch import Labels, TensorBlock, TensorMap + +from metatomic.torch import ( + NeighborListOptions, + System, +) +from metatomic.torch.o3 import ( + O3Transformation, + random_transformations, + transform_system, + transform_tensor, +) + +from ._tests_utils import can_use_mps_backend + + +ALL_DEVICE_DTYPE = [("cpu", "float64"), ("cpu", "float32")] + +if torch.cuda.is_available(): + ALL_DEVICE_DTYPE.append(("cuda", "float64")) + ALL_DEVICE_DTYPE.append(("cuda", "float32")) + +if can_use_mps_backend(): + ALL_DEVICE_DTYPE.append(("mps", "float32")) + + +def _make_system( + types, + positions=None, + cell=None, + pbc=None, + *, + device="cpu", + dtype=torch.float64, +): + n_atoms = len(types) + if positions is None: + positions = torch.zeros((n_atoms, 3), dtype=dtype, device=device) + if cell is None: + cell = torch.zeros((3, 3), dtype=dtype, device=device) + if pbc is None: + pbc = torch.tensor([False, False, False], device=device) + elif torch.is_tensor(pbc): + pbc = pbc.to(device=device) + else: + pbc = torch.tensor(pbc, device=device) + return System( + types=torch.tensor(types, device=device), + positions=positions, + cell=cell, + pbc=pbc, + ) + + +@pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE) +def test_transform_system(device, dtype): + dtype = getattr(torch, dtype) + atol = 1e-6 if dtype == torch.float32 else 1e-12 + + # create a system with neighbor list and custom data + system = _make_system( + [1, 1, 1], + positions=torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], + dtype=dtype, + device=device, + ), + cell=torch.eye(3, dtype=dtype, device=device) * 3.0, + pbc=torch.tensor([True, True, True]), + device=device, + dtype=dtype, + ) + + options = NeighborListOptions(cutoff=2.0, full_list=True, strict=False) + neighbors = TensorBlock( + values=torch.tensor([[[1.0], [0.0], [0.0]]], dtype=dtype, device=device), + samples=Labels( + [ + "first_atom", + "second_atom", + "cell_shift_a", + "cell_shift_b", + "cell_shift_c", + ], + torch.tensor([[0, 1, 0, 0, 0]], device=device), + ), + components=[Labels(["xyz"], torch.arange(3, device=device).reshape(-1, 1))], + properties=Labels(["distance"], torch.tensor([[0]], device=device)), + ) + system.add_neighbor_list(options, neighbors) + + # scalar per-atom data spread over two blocks: both must pass through untouched + samples = Labels( + ["system", "atom"], + torch.tensor([[0, 0], [0, 2], [0, 1]], device=device), + ) + scalar = TensorMap( + keys=Labels(["k"], torch.tensor([[0], [1]], device=device)), + blocks=[ + TensorBlock( + values=torch.tensor([[1.0], [2.0], [3.0]], dtype=dtype, device=device), + samples=samples, + components=[], + properties=Labels(["p"], torch.tensor([[0]], device=device)), + ), + TensorBlock( + values=torch.tensor([[3.0], [4.0], [5.0]], dtype=dtype, device=device), + samples=samples, + components=[], + properties=Labels(["p"], torch.tensor([[0]], device=device)), + ), + ], + ) + + # xyz vector per-atom data: must rotate by R on the component axis + vector_values = torch.tensor( + [[[1.0], [0.0], [0.0]], [[0.0], [2.0], [0.0]], [[0.0], [0.0], [3.0]]], + dtype=dtype, + device=device, + ) + vector = TensorMap( + keys=Labels(["_"], torch.tensor([[0]], device=device)), + blocks=[ + TensorBlock( + values=vector_values.clone(), + samples=samples, + components=[ + Labels(["xyz"], torch.arange(3, device=device).reshape(-1, 1)) + ], + properties=Labels(["p"], torch.tensor([[0]], device=device)), + ) + ], + ) + system.add_data("custom::scalar", scalar) + system.add_data("custom::vector", vector) + + # Apply a transformation to it + matrix = torch.tensor( + [ + [np.cos(np.pi / 3), -np.sin(np.pi / 3), 0.0], + [np.sin(np.pi / 3), np.cos(np.pi / 3), 0.0], + [0.0, 0.0, 1.0], + ], + dtype=dtype, + device=device, + ) + transformation = O3Transformation(matrix, max_angular_momentum=0) + + rotated = transform_system(system, transformation) + + assert torch.allclose(rotated.positions, system.positions @ matrix.T, atol=atol) + assert torch.allclose(rotated.cell, system.cell @ matrix.T, atol=atol) + assert torch.equal(rotated.types, system.types) + assert torch.equal(rotated.pbc, system.pbc) + + new_neighbors = rotated.get_neighbor_list(options).values + expected = ( + system.get_neighbor_list(options).values.squeeze(-1) @ matrix.T + ).unsqueeze(-1) + assert torch.allclose(new_neighbors, expected, atol=atol) + + new_scalar = rotated.get_data("custom::scalar") + for block_id in range(len(scalar.keys)): + assert torch.allclose( + new_scalar.block_by_id(block_id).values, + scalar.block_by_id(block_id).values, + ) + + new_vector = rotated.get_data("custom::vector").block().values + expected_vector = (vector_values.squeeze(-1) @ matrix.T).unsqueeze(-1) + assert torch.allclose(new_vector, expected_vector, atol=atol) + + +def test_random_rotations_are_orthogonal(): + transformations = random_transformations( + 20, device=torch.device("cpu"), dtype=torch.float64 + ) + assert len(transformations) == 20 + identity = torch.eye(3, dtype=torch.float64) + for transformation in transformations: + assert transformation.matrix.shape == (3, 3) + assert torch.allclose( + transformation.matrix @ transformation.matrix.T, identity, atol=1e-10 + ) + assert abs(float(torch.det(transformation.matrix)) - 1.0) < 1e-10 + + +def test_random_rotations_include_inversions(): + # With n=100 the probability that all determinants have the same sign is 2^{-99}. + transformations = random_transformations( + 100, + device="cpu", + dtype=torch.float64, + include_inversions=True, + ) + dets = torch.stack([torch.det(T.matrix) for T in transformations]) + assert torch.allclose(dets.abs(), torch.ones(100, dtype=torch.float64), atol=1e-10) + assert (dets > 0).any() and (dets < 0).any() + + +def _axis_angle(axis, theta): + """A general (non-degenerate, beta != 0) rotation matrix from axis and angle.""" + axis = np.asarray(axis, dtype=float) + axis = axis / np.linalg.norm(axis) + x, y, z = axis + c, s = np.cos(theta), np.sin(theta) + one_minus_c = 1.0 - c + return np.array( + [ + [ + c + x * x * one_minus_c, + x * y * one_minus_c - z * s, + x * z * one_minus_c + y * s, + ], + [ + y * x * one_minus_c + z * s, + c + y * y * one_minus_c, + y * z * one_minus_c - x * s, + ], + [ + z * x * one_minus_c - y * s, + z * y * one_minus_c + x * s, + c + z * z * one_minus_c, + ], + ] + ) + + +# Change of basis from Cartesian (x, y, z) to real ell=1 spherical harmonics, ordered +# (m=-1, 0, +1) = (y, z, x). The real ell=1 Wigner-D matrix satisfies D1 = C @ R @ C.T, +# which lets us cross-check the spherical path (Wigner-D) against the trivially-correct +# Cartesian path under arbitrary rotations. +CARTESIAN_TO_SPHERICAL_L1 = torch.tensor( + [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]], dtype=torch.float64 +) + +# Rotations exercising the full ZYZ decomposition path: generic, beta=0, beta=pi, +# and improper cases. +TRANSFORMATIONS = [ + pytest.param( + torch.tensor(_axis_angle([1.0, 2.0, 3.0], 0.7), dtype=torch.float64), + id="rotation-1", + ), + pytest.param( + torch.tensor(_axis_angle([-2.0, 1.0, 0.5], 2.4), dtype=torch.float64), + id="rotation-2", + ), + pytest.param( + torch.tensor(_axis_angle([0.0, 0.0, 1.0], 0.9), dtype=torch.float64), + id="gimbal-lock-beta-0", + ), + pytest.param( + torch.tensor(_axis_angle([1.0, 0.0, 0.0], np.pi), dtype=torch.float64), + id="gimbal-lock-beta-pi", + ), + pytest.param( + -torch.tensor(_axis_angle([0.3, -1.0, 2.0], 1.1), dtype=torch.float64), + id="improper-rotation", + ), +] + + +@pytest.mark.parametrize("matrix", TRANSFORMATIONS) +def test_general_rotation_L1_wigner_matches_cartesian(matrix): + """For ell=1, Wigner-D matches the Cartesian basis-change formula.""" + proper = matrix if torch.det(matrix) > 0 else -matrix + transformation = O3Transformation(matrix, max_angular_momentum=1) + + expected = CARTESIAN_TO_SPHERICAL_L1 @ proper @ CARTESIAN_TO_SPHERICAL_L1.T + assert torch.allclose(transformation._wigner_D_cache[1], expected, atol=1e-12) + + +@pytest.mark.parametrize("matrix", TRANSFORMATIONS) +@pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE) +def test_spherical_rotation_matches_cartesian(matrix, device, dtype): + dtype = getattr(torch, dtype) + atol = 1e-6 if dtype == torch.float32 else 1e-12 + + matrix = matrix.to(device=device, dtype=dtype) + + system = _make_system([1, 1], device=device, dtype=dtype) + + cartesian_vectors = torch.randn(2, 3, 1, dtype=dtype, device=device) + + cartesian = TensorMap( + Labels(["_"], torch.tensor([[0]], device=device)), + [ + TensorBlock( + values=cartesian_vectors, + samples=Labels( + ["system", "atom"], + torch.tensor([[0, 0], [0, 1]], device=device), + ), + components=[ + Labels(["xyz"], torch.arange(3, device=device).reshape(-1, 1)) + ], + properties=Labels(["p"], torch.tensor([[0]], device=device)), + ) + ], + ) + # spherical encoding: w = C @ v along the component axis + cart_to_sph = CARTESIAN_TO_SPHERICAL_L1.to(device=device, dtype=dtype) + spherical_values = torch.einsum("Aa,iap->iAp", cart_to_sph, cartesian_vectors) + spherical = TensorMap( + Labels(["o3_lambda", "o3_sigma"], torch.tensor([[1, 1]], device=device)), + [ + TensorBlock( + values=spherical_values, + samples=Labels( + ["system", "atom"], + torch.tensor([[0, 0], [0, 1]], device=device), + ), + components=[ + Labels(["o3_mu"], torch.arange(-1, 2, device=device).reshape(-1, 1)) + ], + properties=Labels(["p"], torch.tensor([[0]], device=device)), + ) + ], + ) + + transformation = O3Transformation(matrix, max_angular_momentum=1) + + cartesian_transformed = transform_tensor(cartesian, [system], [transformation]) + spherical_transformed = transform_tensor(spherical, [system], [transformation]) + + expected_spherical = torch.einsum( + "Aa,iap->iAp", cart_to_sph, cartesian_transformed.block().values + ) + + assert torch.allclose( + spherical_transformed.block().values, expected_spherical, atol=atol + ) + + +def test_gradients_are_rotated(): + systems = [ + _make_system([1, 1]), + _make_system([8, 8, 8]), + ] + R0 = torch.tensor(_axis_angle([1.0, 2.0, 3.0], 0.7), dtype=torch.float64) + R1 = torch.tensor(_axis_angle([0.0, 1.0, 1.0], 1.9), dtype=torch.float64) + + values = torch.tensor([[1.0], [2.0]], dtype=torch.float64) + pos_grad = torch.randn(5, 3, 1, dtype=torch.float64) # 2 + 3 atoms + strain_grad = torch.randn(2, 3, 3, 1, dtype=torch.float64) + + block = TensorBlock( + values=values, + samples=Labels(["system"], torch.tensor([[0], [1]])), + components=[], + properties=Labels(["energy"], torch.tensor([[0]])), + ) + block.add_gradient( + "positions", + TensorBlock( + values=pos_grad, + samples=Labels( + ["sample", "atom"], + # non-sorted samples values + torch.tensor([[0, 0], [1, 0], [0, 1], [1, 1], [1, 2]]), + ), + components=[Labels(["xyz"], torch.arange(3).reshape(-1, 1))], + properties=Labels(["energy"], torch.tensor([[0]])), + ), + ) + block.add_gradient( + "strain", + TensorBlock( + values=strain_grad, + samples=Labels(["sample"], torch.tensor([[0], [1]])), + components=[ + Labels(["xyz_1"], torch.arange(3).reshape(-1, 1)), + Labels(["xyz_2"], torch.arange(3).reshape(-1, 1)), + ], + properties=Labels(["energy"], torch.tensor([[0]])), + ), + ) + tensor = TensorMap(Labels(["_"], torch.tensor([[0]])), [block]) + + transformed = transform_tensor( + tensor, + systems, + [ + O3Transformation(R0, max_angular_momentum=1), + O3Transformation(R1, max_angular_momentum=1), + ], + ) + + # values should not change + assert torch.allclose(transformed.block().values, values) + + # positions gradients: rows 0,2 -> R0, rows 1,3,4 -> R1 + expected_pos = pos_grad.clone() + expected_pos[0] = R0 @ pos_grad[0] + expected_pos[1] = R1 @ pos_grad[1] + expected_pos[2] = R0 @ pos_grad[2] + expected_pos[3] = R1 @ pos_grad[3] + expected_pos[4] = R1 @ pos_grad[4] + + assert torch.allclose( + transformed.block().gradient("positions").values, expected_pos + ) + + # strain gradients: row 0 -> R0 S R0.T, row 1 -> R1 S R1.T + expected_strain = strain_grad.clone() + expected_strain[0] = torch.einsum("Aa,abp,Bb->ABp", R0, strain_grad[0], R0) + expected_strain[1] = torch.einsum("Aa,abp,Bb->ABp", R1, strain_grad[1], R1) + assert torch.allclose( + transformed.block().gradient("strain").values, expected_strain + ) + + +def test_unsupported_component_axis_raises(): + """Unknown component-axis names fail loudly.""" + systems = [_make_system([1])] + tensor = TensorMap( + Labels(["_"], torch.tensor([[0]])), + [ + TensorBlock( + values=torch.zeros(1, 3, 1, dtype=torch.float64), + samples=Labels(["system"], torch.tensor([[0]])), + components=[Labels(["direction"], torch.arange(3).reshape(-1, 1))], + properties=Labels(["p"], torch.tensor([[0]])), + ) + ], + ) + + message = ( + "Found a component axis 'direction', which is neither a Cartesian " + "('xyz'/'xyz_1'/'xyz_2'/...) nor spherical ('o3_mu'/'o3_mu_1'/...) axis; " + "it can not be transformed." + ) + with pytest.raises(ValueError, match=re.escape(message)): + transform_tensor( + tensor, systems, [O3Transformation(torch.eye(3, dtype=torch.float64), 1)] + ) + + +def test_system_ids_explicit_match(): + """Explicit ``system_ids`` correctly routes rows with non-contiguous labels.""" + systems = [_make_system([1, 8]), _make_system([1, 8])] + + R0 = O3Transformation(torch.eye(3, dtype=torch.float64), 1) + # 90-degree rotation around z: (x,y) -> (-y, x) + c, s = np.cos(np.pi / 2), np.sin(np.pi / 2) + R1 = O3Transformation( + torch.tensor([[c, -s, 0], [s, c, 0], [0, 0, 1]], dtype=torch.float64), 1 + ) + + # row 0 (label 92) -> R0 (identity): unchanged + # row 1 (label 38) -> R1 (z-90): (0,2,0) -> (-2,0,0) + vector_values = torch.tensor( + [[[1.0], [0.0], [0.0]], [[0.0], [2.0], [0.0]]], + dtype=torch.float64, + ) + tensor = TensorMap( + Labels(["_"], torch.tensor([[0]])), + [ + TensorBlock( + values=vector_values.clone(), + samples=Labels( + ["system", "atom"], + torch.tensor([[92, 0], [38, 0]]), + ), + components=[Labels(["xyz"], torch.arange(3).reshape(-1, 1))], + properties=Labels(["p"], torch.tensor([[0]])), + ) + ], + ) + + transformed = transform_tensor(tensor, systems, [R0, R1], system_ids=[92, 38]) + result = transformed.block().values + # row 0 (system 92) rotated by R0 (identity) + assert torch.allclose(result[0], vector_values[0]) + # row 1 (system 38) rotated by R1 (z-90) + expected = torch.tensor([[[-2.0], [0.0], [0.0]]], dtype=torch.float64) + assert torch.allclose(result[1], expected) + + +def test_system_ids_uncovered_samples_raises(): + """All distinct system indices in samples must appear in ``system_ids``.""" + systems = [_make_system([1, 8]), _make_system([1, 8])] + R = O3Transformation(torch.eye(3, dtype=torch.float64), 1) + + tensor = TensorMap( + Labels(["_"], torch.tensor([[0]])), + [ + TensorBlock( + values=torch.zeros(2, 3, 1, dtype=torch.float64), + samples=Labels( + ["system", "atom"], + torch.tensor([[92, 0], [38, 0]]), + ), + components=[Labels(["xyz"], torch.arange(3).reshape(-1, 1))], + properties=Labels(["p"], torch.tensor([[0]])), + ) + ], + ) + + message = ( + "Block samples contain system labels [38, 92] that are not in " + "system_ids=[99, 100]. Every sample must be assigned to a system " + "in the transformation." + ) + with pytest.raises(ValueError, match=re.escape(message)): + transform_tensor(tensor, systems, [R, R], system_ids=[99, 100]) + + message = ( + "Block samples contain system labels [38, 92] that are not in " + "system_ids=[0, 1]. Every sample must be assigned to a system " + "in the transformation." + ) + with pytest.raises(ValueError, match=re.escape(message)): + # default system_ids=[0, 1] also misses labels 92 and 38 + transform_tensor(tensor, systems, [R, R]) + + +def test_system_ids_single_system_ignores_label(): + """Single-system blocks return all rows regardless of the ``"system"`` label.""" + systems = [_make_system([1])] + values = torch.tensor([[[1.0], [2.0], [3.0]]], dtype=torch.float64) + + tensor = TensorMap( + Labels(["_"], torch.tensor([[0]])), + [ + TensorBlock( + values=values.clone(), + samples=Labels(["system", "atom"], torch.tensor([[4, 0]])), + components=[Labels(["xyz"], torch.arange(3).reshape(-1, 1))], + properties=Labels(["p"], torch.tensor([[0]])), + ) + ], + ) + + transformation = O3Transformation(torch.eye(3, dtype=torch.float64), 1) + transformed = transform_tensor(tensor, systems, [transformation]) + assert torch.allclose(transformed.block().values, values) diff --git a/tox.ini b/tox.ini index e623219c0..e64e8bd6a 100644 --- a/tox.ini +++ b/tox.ini @@ -39,9 +39,10 @@ testing_deps = pytest pytest-cov -metatensor_deps = +metatomic_deps = metatensor-torch >=0.10.0,<0.11 metatensor-operations >=0.5.0,<0.6 + wigners ################################################################################ @@ -52,7 +53,7 @@ metatensor_deps = description = Run the C++ tests for metatomic-torch deps = cmake - {[testenv]metatensor_deps} + {[testenv]metatomic_deps} torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.12}.* commands = @@ -75,7 +76,7 @@ commands = description = Run the C++ tests for metatomic-torch deps = cmake - {[testenv]metatensor_deps} + {[testenv]metatomic_deps} torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.12}.* commands = @@ -122,7 +123,7 @@ description = Run the tests of the metatomic-torch Python package deps = {[testenv]testing_deps} {[testenv]packaging_deps} - {[testenv]metatensor_deps} + {[testenv]metatomic_deps} torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.12}.* numpy @@ -150,7 +151,7 @@ description = Run the doctests defined in any metatomic package deps = {[testenv]testing_deps} {[testenv]packaging_deps} - {[testenv]metatensor_deps} + {[testenv]metatomic_deps} torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.12}.* numpy @@ -174,7 +175,7 @@ description = Run the tests of the metatomic-ase Python package deps = {[testenv]testing_deps} {[testenv]packaging_deps} - {[testenv]metatensor_deps} + {[testenv]metatomic_deps} torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.12}.* numpy @@ -210,7 +211,7 @@ description = Run the tests of the metatomic-torchsim Python package deps = {[testenv]testing_deps} {[testenv]packaging_deps} - {[testenv]metatensor_deps} + {[testenv]metatomic_deps} torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.12}.* numpy @@ -220,6 +221,7 @@ deps = torch-sim-atomistic nvalchemi-toolkit-ops >=0.3.0,<0.4 + changedir = python/metatomic_torchsim commands = pip install {[testenv]build_single_wheel} . @@ -264,7 +266,7 @@ setenv = # build the docs against the CPU only version of torch PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cpu {env:PIP_EXTRA_INDEX_URL:} deps = - {[testenv]metatensor_deps} + {[testenv]metatomic_deps} {[testenv]packaging_deps} {[testenv]testing_deps}