Make O3 transformations batched and torchscript compatible

docs/src/torch/reference/o3.rst

@@ -11,14 +11,15 @@ augmentation.
 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:
+To transform a :py:class:`~metatensor.torch.TensorMap`,
+:py:meth:`O3Transformations.transform_tensormap` needs 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
+  :py:class:`O3Transformations`, 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
@@ -78,13 +79,7 @@ ignored: every row is rotated with the (single) given transformation.
 Reference
 ---------
 
-.. autoclass:: metatomic.torch.o3.O3Transformation
+.. autoclass:: metatomic.torch.o3.O3Transformations
     :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

metatomic-torch/CHANGELOG.md

@@ -18,12 +18,20 @@ a changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows
 
 ### Changed
 
-- Renamed `O3Transformation.is_inverted` to `is_improper`.
-- `wigners >= 0.4.0` is now required.
+- `O3Transformation` has been renamed to `O3Transformations` and is now a
+  `torch.nn.Module` holding a batch of one or more O(3) operations. It is
+  compatible with TorchScript and has member functions to transform systems, TensorMap, Cartesian tensors, and spherical tensors; as well as functions to appply the reverse of the transformations.
+- Renamed `O3Transformations.is_inverted` to `improper_mask`.
+
+### Removed
+
+- Removed the `transform_system`, `transform_block`, and `transform_tensor` free
+  functions from `metatomic.torch.o3`; use the `O3Transformations` methods
+  instead.
 
 ### Fixed
 
-- `O3Transformation.transform_spherical` no longer applies the `(-1)^ell`
+- `O3Transformations.transform_spherical` no longer applies the `(-1)^ell`
   parity factor for proper transformations with `sigma = -1`.
 - Wigner-D evaluation is now stable near the ZYZ Euler-angle poles.
 - `transform_system` now preserves autograd for registered neighbor lists.

python/metatomic_torch/metatomic/torch/o3/__init__.py

@@ -7,19 +7,10 @@
 spherical components in a :py:class:`~metatensor.torch.TensorBlock`.
 """
 
-from ._tranformations import (
-    O3Transformation,
-    random_transformations,
-    transform_block,
-    transform_system,
-    transform_tensor,
-)
+from ._transformations import O3Transformations, random_transformations
 
 
 __all__ = [
-    "O3Transformation",
+    "O3Transformations",
     "random_transformations",
-    "transform_system",
-    "transform_tensor",
-    "transform_block",
 ]