Implement augmentation routines in metatomic

docs/src/torch/reference/index.rst

@@ -9,6 +9,7 @@ API reference
     systems
     models/index
     misc
+    o3
     ase
     serialization
 

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 <https://github.com/luthaf/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

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,

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",
+]