diff --git a/docs/src/torch/reference/index.rst b/docs/src/torch/reference/index.rst index 7cb577e4..f3d4cb61 100644 --- a/docs/src/torch/reference/index.rst +++ b/docs/src/torch/reference/index.rst @@ -11,6 +11,7 @@ API reference units wrappers o3 + symmetrized-model ase misc diff --git a/docs/src/torch/reference/symmetrized-model.rst b/docs/src/torch/reference/symmetrized-model.rst new file mode 100644 index 00000000..3555d1b8 --- /dev/null +++ b/docs/src/torch/reference/symmetrized-model.rst @@ -0,0 +1,147 @@ +.. _symmetrized-model: + +O(3)-symmetrized models +======================= + +The :py:class:`metatomic.torch.SymmetrizedModel` class wraps an exported +:py:class:`~metatomic.torch.AtomisticModel` with finite-quadrature O(3) +averaging and equivariance diagnostics. Pre-existing outputs of the model are +averaged over rotated and inverted copies of each input. +:py:class:`~metatomic.torch.SymmetrizedModel` also adds extra outputs to +compute the equivariance variance or squared character-projection +contributions of the model response. + +Output requests +--------------- + +The requested output name selects both the source output and the calculation: + +.. list-table:: + :header-rows: 1 + + * - Requested and returned name + - Result + * - ```` + - O(3) average of the underlying ```` output + * - ``o3::variance::`` + - component-averaged equivariance variance of ```` + * - ``o3::character_projection::`` + - unnormalized squared character-projection contributions of ```` + +```` is preserved verbatim. It can therefore be a standard quantity, a +variant such as ``energy/pbe``, or a custom name such as +``mtt::feature::node``. For example, +``o3::variance::energy/pbe`` evaluates the underlying ``energy/pbe`` output. + +Average and variance +-------------------- + +For an input :math:`x`, an O(3) operation :math:`g`, and the target +representation :math:`\rho_\alpha`, define the response transformed back to the +input frame as + +.. math:: + + z_\alpha(g;x) = \rho_\alpha(g^{-1}) f(gx). + +The ordinary result is the normalized Haar average + +.. math:: + + \Pi_\alpha(f,x) + = \int_{\mathrm{O}(3)} z_\alpha(g;x)\,\mathrm{d}\mu(g). + +For a TensorMap block with component multiplicity :math:`d`, the corresponding +variance output contains + +.. math:: + + v_\alpha(f,x) + = \frac{1}{d}\left[ + \int_{\mathrm{O}(3)} \lVert z_\alpha(g;x) \rVert_2^2\, + \mathrm{d}\mu(g) + - \lVert \Pi_\alpha(f,x) \rVert_2^2 + \right]. + +This value is returned separately for every sample and property. It has no +component axes, and it is not reduced across samples or square-rooted. A +weighted mean of these values over a group of samples, followed by a square +root, gives a block-wise equivariance RMSE. + +TensorMap representation +------------------------ + +An averaged output retains the physical schema declared by the source model. +For diagnostics, the standard quantities are represented as follows: + +.. list-table:: + :header-rows: 1 + + * - Source quantity + - Diagnostic target keys + * - ``energy``, ``energy_ensemble``, ``energy_uncertainty`` + - ``o3_lambda=0``, ``o3_sigma=1`` + * - ``non_conservative_force`` + - ``o3_lambda=1``, ``o3_sigma=1`` + * - ``non_conservative_stress`` + - ``(o3_lambda, o3_sigma)=(0,1)``, ``(1,-1)``, and ``(2,1)`` + +Variants after ``/`` use the same representation as their base quantity. + +Energy-like scalars acquire an ``o3_mu`` component of size one for diagnostics. +Cartesian force components are reordered into the real spherical +:math:`\ell=1` basis described in :ref:`o3-conventions`. Stress diagnostics +cover the full matrix: the scalar trace, the antisymmetric (axial pseudovector, +:math:`\ell=1` with ``o3_sigma=-1``) part, and the symmetric-traceless sector. +For a symmetric stress the pseudovector sector is exactly zero; a model +producing a non-symmetric stress (before any downstream symmetrization) sees +its antisymmetric response in this sector. + +Already-spherical outputs retain their ``o3_lambda`` and ``o3_sigma`` keys and +``o3_mu`` components, and other semantic source keys are preserved. The wrapper +does not infer the physical meaning of a custom output from its shape; in +particular, a custom Cartesian :math:`3\times3` output is not treated as a +symmetric stress. + +Character projections +--------------------- + +Character projections analyze the direct response :math:`u(g;x)=f(gx)`, rather +than the back-transformed response used for averaging. For the character sector +:math:`\beta=(\lambda,\sigma)` with :math:`d_\beta=2\lambda+1`, the squared +projection norm is + +.. math:: + + B_\beta(u,x) + = d_\beta \iint_{\mathrm{O}(3)} + u(g_1;x)^\dagger + \chi_\beta(g_1g_2^{-1})u(g_2;x)\, + \mathrm{d}\mu(g_1)\,\mathrm{d}\mu(g_2). + +Character results append ``chi_lambda`` and ``chi_sigma`` to the TensorMap +keys. These labels describe the O(3) dependence of the response over the +rotation orbit. They are distinct from ``o3_lambda`` and ``o3_sigma``, which +describe the target representation of the output itself. Target component axes +are retained; summing over them gives the complete component norm in the +equation above. + +Quadrature +---------- + +The deterministic grid combines a Lebedev rule on the sphere, uniformly spaced +in-plane rotations, and both O(3) cosets: O(3) splits into two cosets of SO(3), +the proper rotations, and the improper ones (a rotation composed with +inversion). Its weights are normalized to sum to one. A general +machine-learning model need not be band-limited, so a finite grid is not +automatically exact. ``max_angular_momentum_grid`` controls the quadrature +resolution, not the representation: increase it until the averages, variances, +and character projections of interest converge. + +Reference +--------- + +.. py:currentmodule:: metatomic.torch + +.. autoclass:: SymmetrizedModel + :members: diff --git a/metatomic-torch/CHANGELOG.md b/metatomic-torch/CHANGELOG.md index b066d1ad..3163eaac 100644 --- a/metatomic-torch/CHANGELOG.md +++ b/metatomic-torch/CHANGELOG.md @@ -16,8 +16,18 @@ a changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows ### Removed --> +### Added + +- Added `metatomic.torch.SymmetrizedModel` for finite-quadrature O(3) + averaging, equivariance variances, and character projections of existing + atomistic models. + ### Changed +- `O3Transformation` now holds a batch of one or more operations, can be + constructed from precomputed tensors inside scripted models, and gained + `inverse`, `with_inversion`, `transform_systems`, and `transform_tensormap`; + `SymmetrizedModel` shares this single implementation. - Renamed `O3Transformation.is_inverted` to `is_improper`. - `wigners >= 0.4.0` is now required. diff --git a/python/metatomic_torch/metatomic/torch/__init__.py b/python/metatomic_torch/metatomic/torch/__init__.py index 06a9ae9c..c76f2764 100644 --- a/python/metatomic_torch/metatomic/torch/__init__.py +++ b/python/metatomic_torch/metatomic/torch/__init__.py @@ -61,6 +61,7 @@ is_atomistic_model, load_atomistic_model, ) +from .o3._symmetrized import SymmetrizedModel # noqa: F401 from .serialization import ( # noqa: F401 load_system, load_system_buffer, diff --git a/python/metatomic_torch/metatomic/torch/_quantities.py b/python/metatomic_torch/metatomic/torch/_quantities.py new file mode 100644 index 00000000..d0800858 --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/_quantities.py @@ -0,0 +1,79 @@ +""" +Python-side mirror of ``metatomic-torch/src/quantities.cpp`` (``KNOWN_QUANTITIES`` +and the per-quantity checks), holding the metadata for standard quantities: their +category (Cartesian layout and spherical character) and their deprecated-name +aliases. + +This module must not import anything from ``metatomic``, so that any other module +can import it without creating an import cycle. +""" + +from typing import Dict + + +def standard_quantity_categories() -> Dict[str, str]: + """Return the Cartesian layout and spherical character for standard quantities. + + This is the single source of truth for which outputs and inputs are + decomposed; it mirrors ``KNOWN_QUANTITIES`` in + ``metatomic-torch/src/quantities.cpp``, minus ``feature``. Only the current + (singular) spellings appear here: deprecated names are normalized before + they reach the code using this table. + + TorchScript cannot read a module-level dictionary from a compiled function, + so the table is built by this function and bound to + :py:data:`STANDARD_QUANTITY_CATEGORIES` for Python callers. + """ + return { + # scalars: l = 0 + "charge": "scalar", + "energy": "scalar", + "energy_ensemble": "scalar", + "energy_uncertainty": "scalar", + "mass": "scalar", + "spin_multiplicity": "scalar", + # Cartesian vectors: l = 1 + "heat_flux": "cartesian_vector", + "momentum": "cartesian_vector", + "non_conservative_force": "cartesian_vector", + "position": "cartesian_vector", + "velocity": "cartesian_vector", + # symmetric 3x3 matrices: l = 0 and l = 2 + "non_conservative_stress": "symmetric_matrix", + } + + +STANDARD_QUANTITY_CATEGORIES: Dict[str, str] = standard_quantity_categories() + +#: maximum angular momentum carried by each category above +MAX_ANGULAR_MOMENTUM_PER_CATEGORY: Dict[str, int] = { + "scalar": 0, + "cartesian_vector": 1, + "symmetric_matrix": 2, +} + + +def _new_quantity_names() -> Dict[str, str]: + """Return the map from deprecated quantity names to their current name. + + TorchScript cannot read a module-level dictionary from a compiled function, + so the table is built by this function and bound to + :py:data:`NEW_QUANTITY_NAMES` for Python callers. + """ + return { + "features": "feature", + "non_conservative_forces": "non_conservative_force", + "positions": "position", + "momenta": "momentum", + "masses": "mass", + "velocities": "velocity", + "charges": "charge", + } + + +NEW_QUANTITY_NAMES: Dict[str, str] = _new_quantity_names() + +#: mapping from current quantity names to the corresponding deprecated name +DEPRECATED_QUANTITY_NAMES: Dict[str, str] = { + new: deprecated for deprecated, new in NEW_QUANTITY_NAMES.items() +} diff --git a/python/metatomic_torch/metatomic/torch/model.py b/python/metatomic_torch/metatomic/torch/model.py index 561c1004..15fc498e 100644 --- a/python/metatomic_torch/metatomic/torch/model.py +++ b/python/metatomic_torch/metatomic/torch/model.py @@ -1,3 +1,4 @@ +import copy import datetime import json import math @@ -25,6 +26,7 @@ ) from . import __version__ as metatomic_version from ._extensions import _collect_extensions +from ._quantities import DEPRECATED_QUANTITY_NAMES, NEW_QUANTITY_NAMES def load_atomistic_model(path, extensions_directory=None) -> "AtomisticModel": @@ -394,27 +396,13 @@ def __init__( else: raise ValueError(f"unknown dtype in capabilities: {capabilities.dtype}") - # mapping from deprecated output/input names to their new name - self._new_names = { - "features": "feature", - "non_conservative_forces": "non_conservative_force", - "positions": "position", - "momenta": "momentum", - "masses": "mass", - "velocities": "velocity", - "charges": "charge", - } + # mapping from deprecated output/input names to their new name, copied + # onto the instance because TorchScript methods cannot read module-level + # dictionaries + self._new_names = copy.deepcopy(NEW_QUANTITY_NAMES) # mapping from new names to the corresponding deprecated name - self._deprecated_names = { - "feature": "features", - "non_conservative_force": "non_conservative_forces", - "position": "positions", - "momentum": "momenta", - "mass": "masses", - "velocity": "velocities", - "charge": "charges", - } + self._deprecated_names = copy.deepcopy(DEPRECATED_QUANTITY_NAMES) # Pretend that the model can output either the new or deprecated names new_outputs = {} diff --git a/python/metatomic_torch/metatomic/torch/o3/__init__.py b/python/metatomic_torch/metatomic/torch/o3/__init__.py index 40cb50ba..e27c7fbc 100644 --- a/python/metatomic_torch/metatomic/torch/o3/__init__.py +++ b/python/metatomic_torch/metatomic/torch/o3/__init__.py @@ -7,7 +7,7 @@ spherical components in a :py:class:`~metatensor.torch.TensorBlock`. """ -from ._tranformations import ( +from ._transformations import ( O3Transformation, random_transformations, transform_block, diff --git a/python/metatomic_torch/metatomic/torch/o3/_decompose.py b/python/metatomic_torch/metatomic/torch/o3/_decompose.py new file mode 100644 index 00000000..a3e3827d --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/o3/_decompose.py @@ -0,0 +1,249 @@ +""" +Decomposition of standard Cartesian outputs into O(3) irreducible components. + +Standard scalars, vectors and matrices are re-expressed with explicit +``o3_lambda`` and ``o3_sigma`` keys, so that the variance and +character-projection machinery of the O(3)-symmetrized model treats them exactly +like natively spherical outputs. +""" + +import math +from typing import List, Tuple + +import torch +from metatensor.torch import Labels, TensorBlock, TensorMap + +from .._quantities import standard_quantity_categories + + +def _o3_mu_labels(o3_lambda: int, device: torch.device) -> Labels: + """Return ``o3_mu`` labels from ``-o3_lambda`` through ``o3_lambda``.""" + return Labels( + "o3_mu", + torch.arange( + -o3_lambda, + o3_lambda + 1, + dtype=torch.int32, + device=device, + ).reshape(-1, 1), + ) + + +def _cartesian_vectors_to_spherical( + values: torch.Tensor, + component_axis: int, +) -> torch.Tensor: + """Reorder ``(x, y, z)`` as ``(mu=-1, 0, 1) = (y, z, x)``.""" + return values.roll(-1, dims=component_axis) + + +def _matrices_to_spherical( + values: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return orthonormal l=0, l=1, and l=2 components of a Cartesian matrix. + + Standard matrix quantities are symmetric, so their antisymmetric (l=1, + pseudovector) part is zero by construction; it is still returned so that + the diagnostics of a model producing a non-symmetric output (before any + downstream symmetrization) capture the antisymmetric response as well. + + The three stacks together preserve the Frobenius norm of the matrix. + """ + assert values.dim() == 4 and values.size(1) == 3 and values.size(2) == 3 + + l0 = (values[:, 0, 0, :] + values[:, 1, 1, :] + values[:, 2, 2, :]).unsqueeze( + 1 + ) / math.sqrt(3.0) + + sqrt_two = math.sqrt(2.0) + # the antisymmetric part packs into the axial pseudovector + # v = ((M_21 - M_12) / 2, (M_02 - M_20) / 2, (M_10 - M_01) / 2), listed here + # as sqrt(2) * v in the real spherical (mu=-1, 0, 1) = (y, z, x) order + l1 = torch.stack( + [ + (values[:, 0, 2, :] - values[:, 2, 0, :]) / sqrt_two, + (values[:, 1, 0, :] - values[:, 0, 1, :]) / sqrt_two, + (values[:, 2, 1, :] - values[:, 1, 2, :]) / sqrt_two, + ], + dim=1, + ) + + l2 = torch.stack( + [ + (values[:, 0, 1, :] + values[:, 1, 0, :]) / sqrt_two, + (values[:, 1, 2, :] + values[:, 2, 1, :]) / sqrt_two, + (2.0 * values[:, 2, 2, :] - values[:, 0, 0, :] - values[:, 1, 1, :]) + / math.sqrt(6.0), + (values[:, 0, 2, :] + values[:, 2, 0, :]) / sqrt_two, + (values[:, 0, 0, :] - values[:, 1, 1, :]) / sqrt_two, + ], + dim=1, + ) + + return l0, l1, l2 + + +def decompose_quantity( + name: str, + tensor: TensorMap, +) -> TensorMap: + """Decompose standard quantities for variance and character projection. + + This takes the standard Cartesian or scalar quantities (inputs or outputs + of a model) and re-expresses them in the usual O(3) spherical convention, + i.e. as blocks labelled by ``o3_lambda``/``o3_sigma`` with ``o3_mu`` + components. + + ``feature`` is excluded from the decomposition table: features are not an + irreducible representation of O(3), so they are passed through unchanged and + their variance measures the deviation from invariance. + """ + quantity = name.split("/", 1)[0] + categories = standard_quantity_categories() + if quantity not in categories: + return tensor + category = categories[quantity] + + if category == "scalar": + scalar_blocks: List[TensorBlock] = [] + for block in tensor.blocks(): + assert len(block.components) == 0, ( + f"'{quantity}' outputs must not have components" + ) + scalar_blocks.append( + TensorBlock( + values=block.values.unsqueeze(1), + samples=block.samples, + components=[_o3_mu_labels(0, block.values.device)], + properties=block.properties, + ) + ) + result = TensorMap( + _add_o3_irrep_to_keys(tensor.keys, 0, 1), + scalar_blocks, + ) + + elif category == "cartesian_vector": + vector_blocks: List[TensorBlock] = [] + for block in tensor.blocks(): + assert ( + len(block.components) == 1 + and block.components[0].names == ["xyz"] + and len(block.components[0]) == 3 + ), f"'{quantity}' must have one 'xyz' component axis of size 3" + vector_blocks.append( + TensorBlock( + values=_cartesian_vectors_to_spherical(block.values, 1), + samples=block.samples, + components=[_o3_mu_labels(1, block.values.device)], + properties=block.properties, + ) + ) + result = TensorMap( + _add_o3_irrep_to_keys(tensor.keys, 1, 1), + vector_blocks, + ) + + else: + assert category == "symmetric_matrix" + blocks_l0: List[TensorBlock] = [] + blocks_l1: List[TensorBlock] = [] + blocks_l2: List[TensorBlock] = [] + for block in tensor.blocks(): + assert ( + len(block.components) == 2 + and block.components[0].names == ["xyz_1"] + and block.components[1].names == ["xyz_2"] + and len(block.components[0]) == 3 + and len(block.components[1]) == 3 + ), f"'{quantity}' must have 'xyz_1' and 'xyz_2' component axes of size 3" + + values_l0, values_l1, values_l2 = _matrices_to_spherical(block.values) + blocks_l0.append( + TensorBlock( + values=values_l0, + samples=block.samples, + components=[_o3_mu_labels(0, block.values.device)], + properties=block.properties, + ) + ) + blocks_l1.append( + TensorBlock( + values=values_l1, + samples=block.samples, + components=[_o3_mu_labels(1, block.values.device)], + properties=block.properties, + ) + ) + blocks_l2.append( + TensorBlock( + values=values_l2, + samples=block.samples, + components=[_o3_mu_labels(2, block.values.device)], + properties=block.properties, + ) + ) + + keys_l0 = _add_o3_irrep_to_keys(tensor.keys, 0, 1) + # a Cartesian rank-2 tensor is inversion-even, so its antisymmetric + # (axial-vector) part is an l=1 pseudovector: o3_sigma = -1 + keys_l1 = _add_o3_irrep_to_keys(tensor.keys, 1, -1) + keys_l2 = _add_o3_irrep_to_keys(tensor.keys, 2, 1) + result = TensorMap( + Labels( + list(keys_l0.names), + torch.cat([keys_l0.values, keys_l1.values, keys_l2.values], dim=0), + ), + blocks_l0 + blocks_l1 + blocks_l2, + ) + + for info_name, info_value in tensor.info().items(): + result.set_info(info_name, info_value) + return result + + +def _add_o3_irrep_to_keys( + keys: Labels, + o3_lambda: int, + o3_sigma: int, +) -> Labels: + """Add or validate the ``o3_lambda`` and ``o3_sigma`` key columns.""" + names = list(keys.names) + values = keys.values + + if names == ["_"]: + if len(keys) != 1 or int(values[0, 0]) != 0: + raise ValueError( + "the '_' placeholder must contain exactly one key with value 0" + ) + names = [] + values = values[:, :0] + + for name, expected in ( + ("o3_lambda", o3_lambda), + ("o3_sigma", o3_sigma), + ): + if name in names: + column = values[:, names.index(name)] + if not bool(torch.all(column == expected).item()): + raise ValueError( + f"the existing '{name}' key column must contain only " + f"{expected} to assign O(3) irrep " + f"({o3_lambda}, {o3_sigma})" + ) + else: + names.append(name) + values = torch.cat( + [ + values, + torch.full( + (len(keys), 1), + expected, + dtype=values.dtype, + device=values.device, + ), + ], + dim=1, + ) + + return Labels(names, values) diff --git a/python/metatomic_torch/metatomic/torch/o3/_projections.py b/python/metatomic_torch/metatomic/torch/o3/_projections.py new file mode 100644 index 00000000..e4cd8e51 --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/o3/_projections.py @@ -0,0 +1,194 @@ +"""Character-projection helpers. + +The projected quantity is defined in the :py:class:`SymmetrizedModel` class +docstring. The projections are accumulated separately over the two cosets of +SO(3) in O(3): the proper rotations, and the improper ones (a rotation composed +with inversion). +""" + +from typing import List, Tuple + +import torch +from metatensor.torch import Labels, TensorBlock, TensorMap + +from ._utils import ( + group_samples_by_rotated_copy, + restore_input_system_to_samples, +) + + +def _character_projection_coefficients_from_rotation_batch( + values: torch.Tensor, + weights: torch.Tensor, + inverse_wigner_matrices: torch.Tensor, +) -> torch.Tensor: + """Compute one rotation batch's character-projection coefficients.""" + weighted_wigner_matrices = weights.to( + dtype=values.dtype, + device=values.device, + ).view(-1, 1, 1) * inverse_wigner_matrices.to( + dtype=values.dtype, + device=values.device, + ) + return torch.einsum( + "gmn,gs...->smn...", + weighted_wigner_matrices, + values, + ) + + +def _character_projections_from_proper_and_improper_coefficients( + proper_coefficients: torch.Tensor, + improper_coefficients: torch.Tensor, + chi_lambda: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Return squared character projections for ``chi_sigma=+1`` and ``-1``.""" + dimension = 2 * chi_lambda + 1 + parity = (-1) ** chi_lambda + sigma_plus = proper_coefficients + parity * improper_coefficients + sigma_minus = proper_coefficients - parity * improper_coefficients + factor = float(dimension) / 4.0 + return ( + factor * sigma_plus.square().sum(dim=(1, 2)), + factor * sigma_minus.square().sum(dim=(1, 2)), + ) + + +def character_projection_coefficients_from_batch( + tensor: TensorMap, + weights: torch.Tensor, + inverse_wigner_matrices: List[torch.Tensor], + input_system_index: int, +) -> TensorMap: + """Accumulate all angular momenta of the character projection for one + rotation batch.""" + key_names = list(tensor.keys.names) + key_values = tensor.keys.values + if key_names == ["_"]: + if len(tensor.keys) != 1 or int(key_values[0, 0]) != 0: + raise ValueError( + "the '_' placeholder must contain exactly one key with value 0" + ) + key_names = [] + key_values = key_values[:, :0] + + if "chi_lambda" in key_names or "chi_sigma" in key_names: + raise ValueError( + "source output keys must not contain 'chi_lambda' or 'chi_sigma'" + ) + + blocks: List[TensorBlock] = [] + output_key_values: List[torch.Tensor] = [] + n_rotated_copies = weights.numel() + for key_index in range(len(tensor.keys)): + block = tensor.block(key_index) + values, sample_names, sample_values = group_samples_by_rotated_copy( + block, + n_rotated_copies, + ) + samples = restore_input_system_to_samples( + sample_names, + sample_values, + input_system_index, + device=block.samples.device, + ) + + for chi_lambda in range(len(inverse_wigner_matrices)): + coefficients = _character_projection_coefficients_from_rotation_batch( + values, + weights, + inverse_wigner_matrices[chi_lambda], + ) + dimension = 2 * chi_lambda + 1 + character_indices = torch.arange( + dimension, + dtype=torch.int32, + device=coefficients.device, + ).reshape(-1, 1) + components = [ + Labels("chi_m", character_indices), + Labels("chi_n", character_indices), + ] + for component in block.components: + components.append(component) + blocks.append( + TensorBlock( + values=coefficients, + samples=samples, + components=components, + properties=block.properties, + ) + ) + output_key_values.append( + torch.cat( + [ + key_values[key_index], + torch.tensor( + [chi_lambda], + dtype=key_values.dtype, + device=key_values.device, + ), + ] + ) + ) + + if len(output_key_values) == 0: + values = key_values.new_empty((0, len(key_names) + 1)) + else: + values = torch.stack(output_key_values) + return TensorMap(Labels(key_names + ["chi_lambda"], values), blocks) + + +def character_projection_tensormap_from_cosets( + proper_coefficients: TensorMap, + improper_coefficients: TensorMap, +) -> TensorMap: + """Combine proper and improper coefficient TensorMaps into O(3) irreps.""" + key_names = list(proper_coefficients.keys.names) + if "chi_lambda" not in key_names: + raise ValueError("character coefficients must contain a 'chi_lambda' key") + if "chi_sigma" in key_names: + raise ValueError("source output keys must not contain 'chi_sigma'") + chi_lambda_column = key_names.index("chi_lambda") + + blocks: List[TensorBlock] = [] + output_key_values: List[torch.Tensor] = [] + for key_index in range(len(proper_coefficients.keys)): + proper_block = proper_coefficients.block(key_index) + improper_block = improper_coefficients.block(key_index) + chi_lambda = int(proper_coefficients.keys.values[key_index, chi_lambda_column]) + sigma_plus, sigma_minus = ( + _character_projections_from_proper_and_improper_coefficients( + proper_block.values, + improper_block.values, + chi_lambda, + ) + ) + target_components = proper_block.components[2:] + for chi_sigma, values in ((1, sigma_plus), (-1, sigma_minus)): + blocks.append( + TensorBlock( + values=values, + samples=proper_block.samples, + components=target_components, + properties=proper_block.properties, + ) + ) + output_key_values.append( + torch.cat( + [ + proper_coefficients.keys.values[key_index], + torch.tensor( + [chi_sigma], + dtype=proper_coefficients.keys.values.dtype, + device=proper_coefficients.keys.values.device, + ), + ] + ) + ) + + if len(output_key_values) == 0: + values = proper_coefficients.keys.values.new_empty((0, len(key_names) + 1)) + else: + values = torch.stack(output_key_values) + return TensorMap(Labels(key_names + ["chi_sigma"], values), blocks) diff --git a/python/metatomic_torch/metatomic/torch/o3/_quadrature.py b/python/metatomic_torch/metatomic/torch/o3/_quadrature.py new file mode 100644 index 00000000..3daeacdf --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/o3/_quadrature.py @@ -0,0 +1,175 @@ +""" +Quadrature rules used to integrate over the rotation group SO(3), and over O(3) when +inversion is included. Lebedev rules on the unit sphere are combined with uniformly +spaced in-plane rotations, giving rotations and weights that integrate spherical +harmonics exactly up to a requested maximum angular momentum. +""" + +import numpy as np + +from ._utils import validate_integer + + +_LEBEDEV_ORDERS = ( + 3, + 5, + 7, + 9, + 11, + 13, + 15, + 17, + 19, + 21, + 23, + 25, + 27, + 29, + 31, + 35, + 41, + 47, + 53, + 59, + 65, + 71, + 77, + 83, + 89, + 95, + 101, + 107, + 113, + 119, + 125, + 131, +) + + +def _import_scipy(): + """Import the SciPy functions required to construct a rotation quadrature.""" + try: + from scipy.integrate import lebedev_rule + from scipy.spatial.transform import Rotation + except ImportError as e: + raise ImportError( + "scipy >= 1.15 is required for SymmetrizedModel quadrature construction " + "(scipy.integrate.lebedev_rule); install it with `pip install scipy`." + ) from e + return lebedev_rule, Rotation + + +def choose_quadrature(max_angular_momentum: int) -> tuple[int, int]: + """ + Choose a Lebedev quadrature order and number of in-plane rotations to integrate + spherical harmonics up to ``max_angular_momentum``. + + :param max_angular_momentum: maximum spherical harmonic degree + :return: (lebedev_order, n_inplane_rotations) + """ + max_angular_momentum = validate_integer( + "max_angular_momentum", max_angular_momentum, 0 + ) + if max_angular_momentum > _LEBEDEV_ORDERS[-1]: + raise ValueError( + "the requested quadrature degree " + f"max_angular_momentum={max_angular_momentum} exceeds the largest " + f"available Lebedev order ({_LEBEDEV_ORDERS[-1]})" + ) + # pick smallest order >= max_angular_momentum + n = min(o for o in _LEBEDEV_ORDERS if o >= max_angular_momentum) + # minimal gamma count + K = max_angular_momentum + 1 + return n, K + + +def get_euler_angles_quadrature( + lebedev_order: int, n_rotations: int +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """ + Get the Euler angles and weights for a Lebedev quadrature combined with in-plane + rotations for SO(3) integration. + + :param lebedev_order: order of the Lebedev quadrature on the unit sphere + :param n_rotations: positive integer number of in-plane rotations per Lebedev node + :return: alpha, beta, gamma, w arrays, each of shape (M*K,), where M is the + number of Lebedev nodes and K the number of in-plane rotations; entries + are paired elementwise, one per rotation of the grid. + """ + + lebedev_order = validate_integer("lebedev_order", lebedev_order, 1) + n_rotations = validate_integer("n_rotations", n_rotations, 1) + if lebedev_order not in _LEBEDEV_ORDERS: + raise ValueError( + f"unsupported Lebedev order {lebedev_order}; supported orders are " + f"{list(_LEBEDEV_ORDERS)}" + ) + + lebedev_rule, _ = _import_scipy() + # Lebedev nodes (X: (3, M)) + X, w = lebedev_rule(lebedev_order) # w sums to 4*pi + x, y, z = X + alpha = np.arctan2(y, x) # (M,) + beta = np.arccos(np.clip(z, -1.0, 1.0)) # (M,) + gamma = np.linspace(0.0, 2 * np.pi, n_rotations, endpoint=False) # (K,) + + w_so3 = np.repeat(w / (4 * np.pi * n_rotations), repeats=gamma.size) # (M*K,) + + A = np.repeat(alpha, gamma.size) # (N,) + B = np.repeat(beta, gamma.size) # (N,) + G = np.tile(gamma, alpha.size) # (N,) + + return A, B, G, w_so3 + + +def _rotations_from_euler_angles( + alpha: np.ndarray, beta: np.ndarray, gamma: np.ndarray +): + """ + Construct one active ZYZ rotation from each Euler-angle triple. + + The rotation at index ``i`` is + ``Rz(alpha[i]) @ Ry(beta[i]) @ Rz(gamma[i])``. + + :param alpha: array of alpha angles (N,) + :param beta: array of beta angles (N,) + :param gamma: array of gamma angles (N,) + :return: Rotation object containing the N rotations + """ + + _, Rotation = _import_scipy() + rotations = ( + Rotation.from_euler("z", alpha.reshape(-1, 1)) + * Rotation.from_euler("y", beta.reshape(-1, 1)) + * Rotation.from_euler("z", gamma.reshape(-1, 1)) + ) + + return rotations + + +def get_rotation_quadrature( + lebedev_order: int, n_rotations: int, include_inversion: bool = False +) -> tuple[np.ndarray, np.ndarray]: + """ + Construct rotation matrices and weights for normalized group integration. + + The SO(3) grid combines a Lebedev rule on the sphere with uniformly spaced + in-plane rotations, with weights normalized to sum to one. If + ``include_inversion`` is ``True``, each proper rotation is paired with an + improper one and the original weight is divided equally between the pair. + + :param lebedev_order: order of the Lebedev quadrature on the unit sphere; + must be one of the orders supported by ``scipy.integrate.lebedev_rule`` + :param n_rotations: positive integer number of in-plane rotations per Lebedev node + :param include_inversion: whether to extend the quadrature from SO(3) to O(3) + :return: float64 rotations of shape ``(N, 3, 3)`` and weights of shape + ``(N,)``, summing to 1 + """ + alpha, beta, gamma, weights = get_euler_angles_quadrature( + lebedev_order, n_rotations + ) + rotations = _rotations_from_euler_angles(alpha, beta, gamma).as_matrix() + if include_inversion: + rotations = np.concatenate([rotations, -rotations], axis=0) + weights = np.concatenate([0.5 * weights, 0.5 * weights], axis=0) + return rotations, weights diff --git a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py new file mode 100644 index 00000000..a4816267 --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py @@ -0,0 +1,1204 @@ +""" +:py:class:`SymmetrizedModel`, which averages another model's outputs over a finite +O(3) quadrature. The wrapped model is evaluated on rotated and inverted copies of +each system, and the results are transformed back to the input frame to build the +O(3) average together with the equivariance diagnostics of the requested outputs. +""" + +import warnings +from typing import Dict, List, Optional, Tuple + +import metatensor.torch as mts +import torch +from metatensor.torch import Labels, TensorBlock, TensorMap, dtype_name + +from metatomic.torch import ( + AtomisticModel, + ModelCapabilities, + ModelInterface, + ModelOutput, + NeighborListOptions, + System, +) + +from .._quantities import ( + MAX_ANGULAR_MOMENTUM_PER_CATEGORY, + NEW_QUANTITY_NAMES, + STANDARD_QUANTITY_CATEGORIES, +) +from ._decompose import decompose_quantity +from ._projections import ( + character_projection_coefficients_from_batch, + character_projection_tensormap_from_cosets, +) +from ._quadrature import choose_quadrature, get_rotation_quadrature +from ._transformations import O3Transformation, _max_o3_lambda_in_tensor +from ._utils import ( + group_samples_by_rotated_copy, + map_selected_atoms_to_rotated_copies, + restore_input_system_to_samples, + validate_integer, +) +from ._wigner import ( + build_packed_wigner_matrices, + wigner_matrices_for_lambda, +) + + +def _use_new_quantity_name(name: str) -> str: + """Replace a deprecated base quantity in ``name`` with its current name. + + Public ``AtomisticModel`` capabilities advertise the deprecated alias of + each standard output next to its current name; normalizing here lets the + angular-momentum inference in :py:meth:`SymmetrizedModel.wrap` categorize + both spellings. Requests to the wrapper itself must use current names. + """ + parts = name.split("/") + if parts[0] in NEW_QUANTITY_NAMES: + parts[0] = NEW_QUANTITY_NAMES[parts[0]] + return "/".join(parts) + return name + + +def _check_o3_lambda_limit( + tensor: TensorMap, + tensor_description: str, + max_angular_momentum: int, + limit_name: str, +) -> None: + """Check a TensorMap's spherical component ranks against one limit.""" + tensor_max_o3_lambda = _max_o3_lambda_in_tensor(tensor) + if tensor_max_o3_lambda > max_angular_momentum: + raise ValueError( + f"{tensor_description} contains o3_lambda={tensor_max_o3_lambda}, " + f"exceeding {limit_name}={max_angular_momentum}" + ) + + +def _parse_output_request(requested_name: str) -> Tuple[str, str]: + """Return the underlying output name and requested calculation.""" + variance_prefix = "o3::variance::" + character_projection_prefix = "o3::character_projection::" + + if requested_name.startswith(variance_prefix): + source_name = requested_name[len(variance_prefix) :] + calculation = "variance" + elif requested_name.startswith(character_projection_prefix): + source_name = requested_name[len(character_projection_prefix) :] + calculation = "character_projection" + else: + if requested_name.startswith("o3::"): + raise ValueError( + f"requested output '{requested_name}' uses the 'o3::' prefix " + "reserved by SymmetrizedModel, but is neither a variance nor a " + "character-projection request" + ) + source_name = requested_name + calculation = "average" + + if len(source_name) == 0: + raise ValueError( + f"requested output '{requested_name}' does not identify an " + "underlying model output" + ) + + return source_name, calculation + + +def _group_output_requests( + outputs: Dict[str, ModelOutput], +) -> Tuple[ + Dict[str, str], + Dict[str, str], + Dict[str, str], + Dict[str, str], +]: + """Group public requests by underlying output and calculation. + + The returned dictionaries map each source name to the exact spelling the + caller requested it under. + """ + source_sample_kinds: Dict[str, str] = {} + average_names: Dict[str, str] = {} + variance_names: Dict[str, str] = {} + character_projection_names: Dict[str, str] = {} + + for requested_name, output in outputs.items(): + source_name, calculation = _parse_output_request(requested_name) + sample_kind = output.sample_kind + if source_name in source_sample_kinds: + previous_sample_kind = source_sample_kinds[source_name] + if sample_kind != previous_sample_kind: + raise ValueError( + f"all requests derived from '{source_name}' must use the same " + f"sample_kind; got '{previous_sample_kind}' and '{sample_kind}'" + ) + else: + source_sample_kinds[source_name] = sample_kind + + if calculation == "average": + average_names[source_name] = requested_name + elif calculation == "variance": + variance_names[source_name] = requested_name + else: + character_projection_names[source_name] = requested_name + + return ( + source_sample_kinds, + average_names, + variance_names, + character_projection_names, + ) + + +def _infer_max_angular_momentum( + names: Dict[str, ModelOutput], + kind: str, + argument: str, +) -> int: + """Guess an angular-momentum limit from standard quantity names.""" + max_angular_momentum = 0 + found_standard = False + custom_names: List[str] = [] + for name in names.keys(): + quantity = _use_new_quantity_name(name).split("/")[0] + if quantity == "feature": + # features are not an irreducible representation of O(3): they are + # passed through unchanged and never rotated back + found_standard = True + continue + if quantity not in STANDARD_QUANTITY_CATEGORIES: + # a custom name says nothing about its angular momenta, so it is + # skipped: if it turns out to carry a larger one and is requested, + # _check_o3_lambda_limit rejects it at forward time, naming the limit + custom_names.append(name) + continue + found_standard = True + category = STANDARD_QUANTITY_CATEGORIES[quantity] + max_angular_momentum = max( + max_angular_momentum, MAX_ANGULAR_MOMENTUM_PER_CATEGORY[category] + ) + + if not found_standard and len(custom_names) != 0: + raise ValueError( + f"no standard quantities were found among the {kind}s " + f"{custom_names}, please set {argument} explicitly" + ) + + return max_angular_momentum + + +def _reduce_weighted_centered_batch( + tensor: TensorMap, + weights: torch.Tensor, + input_system_index: int, + reference: Optional[TensorMap], + compute_second_moments: bool, +) -> Tuple[ + TensorMap, + Optional[TensorMap], + Optional[TensorMap], + TensorMap, +]: + """Accumulate one rotation batch's weighted moments, centered on a reference. + + The variance is later formed as ``E[X^2] - E[X]^2``; when the mean response + is much larger than its variation, both terms are huge and nearly equal, and + their difference loses most significant digits. Subtracting a fixed + per-output reference (the first rotated copy) from every response first + leaves the variance unchanged but keeps both moments of the order of the + variation itself, so the subtraction is numerically safe. + """ + n_rotated_copies = weights.numel() + centered_first_moment_blocks: List[TensorBlock] = [] + second_moment_blocks: List[TensorBlock] = [] + absolute_second_moment_blocks: List[TensorBlock] = [] + reference_blocks: List[TensorBlock] = [] + + for key, block in tensor.items(): + values, sample_names, sample_values = group_samples_by_rotated_copy( + block, n_rotated_copies + ) + if reference is None: + # clone so the reference does not keep the full batch tensor alive + reference_values = values[0].clone() + else: + reference_values = reference.block(key).values + matching_shape = reference_values.dim() + 1 == values.dim() + if matching_shape: + for axis in range(reference_values.dim()): + if reference_values.size(axis) != values.size(axis + 1): + matching_shape = False + if not matching_shape: + raise ValueError( + "reference and batch block shapes do not match: reference is " + f"{list(reference_values.shape)}, batch is {list(values.shape)}" + ) + centered_values = values - reference_values.unsqueeze(0) + + batch_weights = weights.to( + dtype=centered_values.dtype, + device=centered_values.device, + ) + weight_shape = [centered_values.shape[0]] + [1] * (centered_values.ndim - 1) + centered_first_moment_values = torch.sum( + batch_weights.view(weight_shape) * centered_values, + dim=0, + ) + + samples = restore_input_system_to_samples( + sample_names, + sample_values, + input_system_index, + device=block.samples.values.device, + ) + centered_first_moment_blocks.append( + TensorBlock( + values=centered_first_moment_values, + samples=samples, + components=block.components, + properties=block.properties, + ) + ) + + if compute_second_moments: + squared_norms = centered_values**2 + if len(block.components) != 0: + n_components = 1 + for component in block.components: + n_components *= len(component) + squared_norms = squared_norms.reshape( + centered_values.shape[0], + centered_values.shape[1], + n_components, + centered_values.shape[-1], + ).sum(dim=2) + moment_weight_shape = [squared_norms.shape[0]] + [1] * ( + squared_norms.ndim - 1 + ) + second_moment_values = torch.sum( + batch_weights.view(moment_weight_shape) * squared_norms, + dim=0, + ) + absolute_second_moment_values = torch.sum( + torch.abs(batch_weights).view(moment_weight_shape) * squared_norms, + dim=0, + ) + second_moment_blocks.append( + TensorBlock( + values=second_moment_values, + samples=samples, + components=[], + properties=block.properties, + ) + ) + absolute_second_moment_blocks.append( + TensorBlock( + values=absolute_second_moment_values, + samples=samples, + components=[], + properties=block.properties, + ) + ) + if reference is None: + reference_blocks.append( + TensorBlock( + values=reference_values, + samples=samples, + components=block.components, + properties=block.properties, + ) + ) + + if reference is None: + reference = TensorMap(tensor.keys, reference_blocks) + + second_moment: Optional[TensorMap] = None + absolute_second_moment: Optional[TensorMap] = None + if compute_second_moments: + second_moment = TensorMap(tensor.keys, second_moment_blocks) + absolute_second_moment = TensorMap( + tensor.keys, + absolute_second_moment_blocks, + ) + + return ( + TensorMap(tensor.keys, centered_first_moment_blocks), + second_moment, + absolute_second_moment, + reference, + ) + + +def _add_tensormap_contribution( + accumulator: Dict[str, TensorMap], + output_name: str, + contribution: TensorMap, +) -> None: + """Add a TensorMap contribution to the running sum for one output.""" + if output_name in accumulator: + accumulator[output_name] = mts.add(accumulator[output_name], contribution) + else: + accumulator[output_name] = contribution + + +def _copy_tensormap_info(source: TensorMap, result: TensorMap) -> TensorMap: + """Copy global information from ``source`` to ``result``.""" + for info_name, info_value in source.info().items(): + result.set_info(info_name, info_value) + return result + + +def _component_norm_squared(tensor: TensorMap) -> TensorMap: + """Return squared values summed over all component axes.""" + blocks: List[TensorBlock] = [] + for block in tensor.blocks(): + values = block.values.square() + if len(block.components) != 0: + values = values.flatten(start_dim=1, end_dim=-2).sum(dim=1) + blocks.append( + TensorBlock( + values=values, + samples=block.samples, + components=[], + properties=block.properties, + ) + ) + return TensorMap(tensor.keys, blocks) + + +def _clamp_roundoff_negative_diagnostic( + tensor: TensorMap, + scale: TensorMap, + *, + n_grid_points: int, + quantity: str, + max_angular_momentum_grid: int, +) -> TensorMap: + """Clamp round-off negatives and reject invalid or materially negative values. + + The variance and character projections are non-negative by construction, but + the finite quadrature evaluates them as differences of large accumulated + sums, so exact zeros come out as tiny values of either sign. Values within + the accumulated round-off bound (estimated from ``scale``) are clamped to + zero; more negative values mean the quadrature did not resolve the response, + which is reported instead of silently returned. + """ + blocks: List[TensorBlock] = [] + for key, block in tensor.items(): + scale_values = scale.block(key).values + if bool(torch.any(~torch.isfinite(block.values)).item()): + raise ValueError(f"O(3) {quantity} is not finite for block ({key.print()})") + if bool(torch.any(~torch.isfinite(scale_values) | (scale_values < 0)).item()): + raise ValueError( + f"round-off scale of the O(3) {quantity} is negative or not " + f"finite for block ({key.print()})" + ) + + # TorchScript does not support torch.finfo; use the IEEE-754 values for + # the floating-point dtypes supported by metatomic models. + if block.values.dtype == torch.float64: + epsilon = 2.220446049250313e-16 + tiny = 2.2250738585072014e-308 + elif block.values.dtype == torch.float32: + epsilon = 1.1920928955078125e-07 + tiny = 1.1754943508222875e-38 + else: + raise TypeError( + "O(3) diagnostics require float32 or float64 values, got " + f"{dtype_name(block.values.dtype)}" + ) + + n_epsilon = n_grid_points * epsilon + gamma = n_epsilon / (1.0 - n_epsilon) + tolerance = ( + 64.0 + * gamma + * torch.clamp( + scale_values, + min=tiny, + ) + ) + if bool(torch.any(block.values < -tolerance).item()): + raise ValueError( + f"finite O(3) {quantity} is materially negative; the quadrature " + "does not resolve this response. Increase max_angular_momentum_grid " + f"above {max_angular_momentum_grid} and check convergence" + ) + + blocks.append( + TensorBlock( + values=torch.clamp(block.values, min=0.0), + samples=block.samples, + components=block.components, + properties=block.properties, + ) + ) + return TensorMap(tensor.keys, blocks) + + +def _variance_from_centered_moments( + centered_first_moment: TensorMap, + centered_second_moment: TensorMap, + absolute_centered_second_moment: TensorMap, + *, + n_grid_points: int, + max_angular_momentum_grid: int, +) -> TensorMap: + """Compute a validated component-summed variance from centered moments.""" + centered_first_moment_norm_squared = _component_norm_squared(centered_first_moment) + variance = mts.subtract( + centered_second_moment, + centered_first_moment_norm_squared, + ) + roundoff_scale = mts.add( + absolute_centered_second_moment, + centered_first_moment_norm_squared, + ) + return _clamp_roundoff_negative_diagnostic( + variance, + roundoff_scale, + n_grid_points=n_grid_points, + quantity="variance", + max_angular_momentum_grid=max_angular_momentum_grid, + ) + + +def _mean_variance_over_components( + variance: TensorMap, + component_layout: TensorMap, +) -> TensorMap: + """Average component-summed variance over each block's components.""" + if variance.keys != component_layout.keys: + raise ValueError("variance and component-layout keys do not match") + + blocks: List[TensorBlock] = [] + for key, block in variance.items(): + if len(block.components) != 0: + raise ValueError("component-summed variance must not have components") + + layout_block = component_layout.block(key) + if ( + layout_block.samples != block.samples + or layout_block.properties != block.properties + ): + raise ValueError("variance and component-layout metadata do not match") + + n_components = 1 + for component in layout_block.components: + n_components *= len(component) + + blocks.append( + TensorBlock( + values=block.values / n_components, + samples=block.samples, + components=[], + properties=block.properties, + ) + ) + + return TensorMap(variance.keys, blocks) + + +class SymmetrizedModel(torch.nn.Module): + """ + Wrap a model with finite-quadrature O(3) averaging and equivariance + diagnostics. + + Requesting an output declared by the wrapped model returns its O(3) + average, evaluated over rotated and inverted copies of the input and + transformed back to the input frame. Requests named + ``o3::variance::`` return the component-averaged equivariance + variance of the ```` output and, when ``max_angular_momentum_character`` is + set, ``o3::character_projection::`` requests return its unnormalized + squared character-projection contributions. The definition of these + quantities, their TensorMap representation, and convergence guidance for + the quadrature are documented in :ref:`symmetrized-model`. + + Only CPU and CUDA execution is supported, and requests for explicit + TensorBlock gradients are rejected. When an input requires gradients, + differentiating an averaged result through PyTorch autograd retains the + source-model activations from all quadrature batches; ``batch_size`` does + not bound their total size. Use :py:func:`torch.inference_mode` or + :py:func:`torch.no_grad` when derivatives are not required. + + :param model: underlying :py:class:`ModelInterface`. The :py:meth:`wrap` method + obtains this module from :py:attr:`AtomisticModel.module`. + :param max_angular_momentum_target: maximum angular momentum that can be transformed + back to the input frame when an average or variance of an + already-spherical output is requested. Cartesian outputs and + character-only requests are not limited by this value. + :param max_angular_momentum_input: maximum angular momentum that can be rotated in + already-spherical custom System data. The default of zero still allows + Cartesian custom inputs. The ``ModelOutput`` declarations returned by a + model's ``requested_inputs()`` do not specify which angular momenta + may occur in the corresponding TensorMaps, so this limit must be + supplied before export for all required Wigner-D matrices to be + serialized. + :param max_angular_momentum_character: maximum angular momentum included in + character projections. ``None`` disables character-projection outputs; zero + enables the scalar (``o3_lambda = 0``) contribution only. + :param max_angular_momentum_grid: quadrature integration degree. If ``None``, use + the larger of ``2 * max_angular_momentum_target + 1`` and + ``2 * max_angular_momentum_character`` when character projections are enabled. + An explicit value must be non-negative and no larger than the highest available + Lebedev order, 131; a value below ``2 * max_angular_momentum_character`` is + rejected. + :param batch_size: positive number of transformed systems evaluated in one call to + ``model``. The default is 32. + """ + + max_angular_momentum_character: Optional[int] + _requested_inputs: Dict[str, ModelOutput] + _requested_neighbor_lists: List[NeighborListOptions] + + def __init__( + self, + model: ModelInterface, + *, + max_angular_momentum_target: int, + max_angular_momentum_input: int = 0, + max_angular_momentum_character: Optional[int] = None, + max_angular_momentum_grid: Optional[int] = None, + batch_size: int = 32, + ): + super().__init__() + + self._model = model + self._requested_inputs = {} + self._requested_neighbor_lists = [] + self.max_angular_momentum_target = validate_integer( + "max_angular_momentum_target", max_angular_momentum_target, 0 + ) + self.max_angular_momentum_input = validate_integer( + "max_angular_momentum_input", max_angular_momentum_input, 0 + ) + if max_angular_momentum_character is not None: + max_angular_momentum_character = validate_integer( + "max_angular_momentum_character", max_angular_momentum_character, 0 + ) + self.max_angular_momentum_character = max_angular_momentum_character + self.batch_size = validate_integer("batch_size", batch_size, 1) + + if max_angular_momentum_grid is None: + max_angular_momentum_grid = 2 * self.max_angular_momentum_target + 1 + if self.max_angular_momentum_character is not None: + max_angular_momentum_grid = max( + max_angular_momentum_grid, + 2 * self.max_angular_momentum_character, + ) + else: + max_angular_momentum_grid = validate_integer( + "max_angular_momentum_grid", max_angular_momentum_grid, 0 + ) + if ( + self.max_angular_momentum_character is not None + and max_angular_momentum_grid < 2 * self.max_angular_momentum_character + ): + raise ValueError( + "max_angular_momentum_grid must be at least twice " + "max_angular_momentum_character" + ) + self.max_angular_momentum_grid = max_angular_momentum_grid + + device = torch.device("cpu") + for parameter in model.parameters(): + device = parameter.device + break + else: + for buffer in model.buffers(): + device = buffer.device + break + if device.type != "cpu" and device.type != "cuda": + # the quadrature buffers are stored in float64, which other + # accelerators (e.g. MPS) do not support + raise ValueError("SymmetrizedModel supports CPU and CUDA execution") + + lebedev_order, n_rotations = choose_quadrature(self.max_angular_momentum_grid) + rotations, weights = get_rotation_quadrature( + lebedev_order, + n_rotations, + ) + rotation_matrices = torch.from_numpy(rotations).to( + dtype=torch.float64, + device=device, + ) + rotation_weights = torch.from_numpy(weights).to( + dtype=torch.float64, + device=device, + ) + + max_angular_momentum_wigner = max( + self.max_angular_momentum_input, + self.max_angular_momentum_target, + 0 + if self.max_angular_momentum_character is None + else self.max_angular_momentum_character, + ) + packed_wigner_matrices = build_packed_wigner_matrices( + rotation_matrices, + max_angular_momentum_wigner, + ) + self._max_angular_momentum_wigner = max_angular_momentum_wigner + + self.register_buffer("_rotation_matrices", rotation_matrices) + self.register_buffer("_rotation_weights", rotation_weights) + self.register_buffer("_packed_wigner_matrices", packed_wigner_matrices) + + @staticmethod + def wrap( + model: AtomisticModel, + *, + max_angular_momentum_target: Optional[int] = None, + max_angular_momentum_input: Optional[int] = None, + max_angular_momentum_character: Optional[int] = None, + max_angular_momentum_grid: Optional[int] = None, + batch_size: int = 32, + ) -> AtomisticModel: + """ + Wrap an exported model with O(3) averaging and diagnostics. + + The returned model retains every output declared by ``model`` under its + original name. Requesting such an output evaluates its O(3) average. + Additional outputs named ``o3::variance::`` provide the + component-averaged equivariance variance. If ``max_angular_momentum_character`` + is set, ``o3::character_projection::`` outputs provide squared + character projections through that angular momentum. + + The original metadata, requested inputs, neighbor lists, and compatible + capabilities are preserved. + + Constructing a wrapper requires SciPy 1.15 or newer for its Lebedev + quadrature. SciPy is not required to evaluate a wrapper that has + already been saved. + + :param model: the :py:class:`AtomisticModel` to wrap + :param max_angular_momentum_target: maximum angular momentum accepted in + already-spherical model outputs requested for averaging or variance. + When ``None``, it is guessed as the largest angular momentum of the + standard quantities declared by ``model``; non-standard outputs are + skipped, and an explicit value is required if ``model`` declares + outputs but none of them is a standard quantity. + :param max_angular_momentum_input: maximum angular momentum accepted in custom + System data. When ``None``, it is guessed the same way from the + quantities in ``model.requested_inputs()``. + :param max_angular_momentum_character: maximum angular momentum in character + projections, or ``None`` to disable them + :param max_angular_momentum_grid: quadrature integration degree, selected + automatically when ``None`` + :param batch_size: number of transformed Systems evaluated in one model call + """ + if not isinstance(model, AtomisticModel): + raise TypeError("model must be an AtomisticModel") + + capabilities = model.capabilities() + supported_devices = [ + device + for device in capabilities.supported_devices + if device == "cpu" or device == "cuda" + ] + if len(supported_devices) == 0: + raise ValueError( + "SymmetrizedModel supports CPU and CUDA execution, but the " + "wrapped model declares " + str(capabilities.supported_devices) + ) + + if max_angular_momentum_target is None: + max_angular_momentum_target = _infer_max_angular_momentum( + capabilities.outputs, + "output", + "max_angular_momentum_target", + ) + if max_angular_momentum_input is None: + max_angular_momentum_input = _infer_max_angular_momentum( + model.requested_inputs(use_new_names=True), + "input", + "max_angular_momentum_input", + ) + + outputs: Dict[str, ModelOutput] = {} + # private field: the as-declared output names, deliberately without the + # deprecation aliases added by the public accessors + for name in model._model_capabilities_outputs_names: + if name.startswith("o3::"): + raise ValueError( + "the wrapped model output '" + + name + + "' uses a prefix reserved by SymmetrizedModel" + ) + + source_output = capabilities.outputs[name] + average_description = "O(3) average of the '" + name + "' output." + if source_output.description != "": + average_description += " " + source_output.description + outputs[name] = ModelOutput( + unit=source_output.unit, + sample_kind=source_output.sample_kind, + explicit_gradients=[], + description=average_description, + ) + + squared_unit = "" + if source_output.unit != "": + squared_unit = "(" + source_output.unit + ")^2" + outputs["o3::variance::" + name] = ModelOutput( + unit=squared_unit, + sample_kind=source_output.sample_kind, + explicit_gradients=[], + description=( + "O(3) equivariance variance of the '" + + name + + "' output for each sample, averaged over components." + ), + ) + if max_angular_momentum_character is not None: + outputs["o3::character_projection::" + name] = ModelOutput( + unit=squared_unit, + sample_kind=source_output.sample_kind, + explicit_gradients=[], + description=( + "Unnormalized squared O(3) character-projection " + "contributions of the '" + + name + + "' output, resolved by chi_lambda and chi_sigma." + ), + ) + + wrapper = SymmetrizedModel( + model.module, + max_angular_momentum_target=max_angular_momentum_target, + max_angular_momentum_input=max_angular_momentum_input, + max_angular_momentum_character=max_angular_momentum_character, + batch_size=batch_size, + max_angular_momentum_grid=max_angular_momentum_grid, + ) + # private field: the as-declared inputs, deliberately without deprecation + # aliases + wrapper._requested_inputs = { + name: requested_input + for name, requested_input in model._requested_inputs.items() + } + # copy the options: constructing the AtomisticModel below mutates them by + # adding requestors and setting the length unit + for options in model.requested_neighbor_lists(): + copied_options = NeighborListOptions( + options.cutoff, + options.full_list, + options.strict, + ) + for requestor in options.requestors(): + copied_options.add_requestor(requestor) + wrapper._requested_neighbor_lists.append(copied_options) + new_capabilities = ModelCapabilities( + outputs=outputs, + atomic_types=capabilities.atomic_types, + interaction_range=capabilities.interaction_range, + length_unit=capabilities.length_unit, + supported_devices=supported_devices, + dtype=capabilities.dtype, + ) + return AtomisticModel( + wrapper.eval(), + model.metadata(), + capabilities=new_capabilities, + ) + + def requested_neighbor_lists(self) -> List[NeighborListOptions]: + """Return the neighbor lists requested by the wrapped model.""" + return self._requested_neighbor_lists + + def requested_inputs(self) -> Dict[str, ModelOutput]: + """Return the custom System data requested by the wrapped model.""" + return self._requested_inputs + + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels], + ) -> Dict[str, TensorMap]: + """Evaluate the requested O(3) averages and diagnostics.""" + if len(outputs) == 0: + empty: Dict[str, TensorMap] = {} + return empty + if len(systems) == 0: + # the metadata of the outputs (keys, sample and property labels) only + # becomes known by evaluating the wrapped model on at least one + # system, so there is no way to build correctly-labelled empty results + raise ValueError("SymmetrizedModel requires at least one System") + + for requested_name, output in outputs.items(): + if len(output.explicit_gradients) != 0: + raise ValueError( + "SymmetrizedModel does not support explicit gradients for " + f"output '{requested_name}'" + ) + + ( + source_sample_kinds, + average_names, + variance_names, + character_projection_names, + ) = _group_output_requests(outputs) + if ( + len(character_projection_names) != 0 + and self.max_angular_momentum_character is None + ): + raise ValueError( + "max_angular_momentum_character must be set to request " + "character projections" + ) + + source_outputs: Dict[str, ModelOutput] = {} + for source_name in source_sample_kinds: + source_outputs[source_name] = ModelOutput( + sample_kind=source_sample_kinds[source_name], + ) + + per_output_results: Dict[str, List[TensorMap]] = {} + for requested_name in outputs: + empty_results: List[TensorMap] = [] + per_output_results[requested_name] = empty_results + + for input_system_index, system in enumerate(systems): + system_results = self._evaluate_system( + system, + input_system_index, + source_outputs, + average_names, + variance_names, + character_projection_names, + selected_atoms, + ) + for requested_name in outputs: + per_output_results[requested_name].append( + system_results[requested_name] + ) + + results: Dict[str, TensorMap] = {} + for requested_name in outputs: + results[requested_name] = mts.join( + per_output_results[requested_name], + "samples", + different_keys="union", + ) + return results + + def _evaluate_system( + self, + system: System, + input_system_index: int, + source_outputs: Dict[str, ModelOutput], + average_names: Dict[str, str], + variance_names: Dict[str, str], + character_projection_names: Dict[str, str], + selected_atoms: Optional[Labels], + ) -> Dict[str, TensorMap]: + """Stream all quadrature batches for one input System.""" + work_dtype = system.positions.dtype + work_device = system.positions.device + if work_dtype != torch.float32 and work_dtype != torch.float64: + raise TypeError( + "SymmetrizedModel requires float32 or float64 Systems, got " + f"{dtype_name(work_dtype)}" + ) + if work_device.type != "cpu" and work_device.type != "cuda": + raise ValueError("SymmetrizedModel supports CPU and CUDA execution") + integration_dtype = self._rotation_matrices.dtype + if integration_dtype != torch.float32 and integration_dtype != torch.float64: + raise TypeError( + "SymmetrizedModel integration buffers must use float32 or " + f"float64, got {dtype_name(integration_dtype)}" + ) + if integration_dtype != torch.float64: + warnings.warn( + "SymmetrizedModel integration buffers were downcast from " + "float64; averages and diagnostics will be less accurate", + stacklevel=2, + ) + if ( + self._rotation_matrices.device != work_device + or self._rotation_weights.device != work_device + or self._packed_wigner_matrices.device != work_device + ): + raise ValueError( + "SymmetrizedModel and input Systems must use the same device" + ) + + for data_name in system.known_data(): + _check_o3_lambda_limit( + system.get_data(data_name), + f"custom input '{data_name}'", + self.max_angular_momentum_input, + "max_angular_momentum_input", + ) + + character_max = 0 + configured_character_max = self.max_angular_momentum_character + if configured_character_max is not None: + character_max = configured_character_max + + average_references: Dict[str, TensorMap] = {} + average_first_moments: Dict[str, TensorMap] = {} + variance_references: Dict[str, TensorMap] = {} + variance_first_moments: Dict[str, TensorMap] = {} + variance_second_moments: Dict[str, TensorMap] = {} + variance_absolute_second_moments: Dict[str, TensorMap] = {} + proper_character_coefficients: Dict[str, TensorMap] = {} + improper_character_coefficients: Dict[str, TensorMap] = {} + + n_rotations = self._rotation_matrices.size(0) + needs_backrotation = len(average_names) != 0 or len(variance_names) != 0 + + # per-``ell`` views into the packed Wigner-D buffer, shared by every batch + wigner_views: List[torch.Tensor] = [] + for o3_lambda in range(self._max_angular_momentum_wigner + 1): + wigner_views.append( + wigner_matrices_for_lambda( + self._packed_wigner_matrices, + n_rotations, + o3_lambda, + ) + ) + + for batch_start in range(0, n_rotations, self.batch_size): + batch_stop = min(batch_start + self.batch_size, n_rotations) + n_rotated_copies = batch_stop - batch_start + proper_matrices = self._rotation_matrices[batch_start:batch_stop] + so3_weights = self._rotation_weights[batch_start:batch_stop] + o3_weights = 0.5 * so3_weights + local_selected_atoms = map_selected_atoms_to_rotated_copies( + selected_atoms, + input_system_index, + n_rotated_copies, + ) + + batch_wigner: List[torch.Tensor] = [] + for wigner_view in wigner_views: + batch_wigner.append(wigner_view[batch_start:batch_stop]) + all_proper = torch.zeros( + n_rotated_copies, + dtype=torch.bool, + device=work_device, + ) + + # the batch of proper rotations applied to the input System, in the + # working dtype of the wrapped model + input_wigner: List[torch.Tensor] = [] + for o3_lambda in range(self.max_angular_momentum_input + 1): + input_wigner.append( + batch_wigner[o3_lambda].to(dtype=work_dtype, device=work_device) + ) + input_rotations = O3Transformation( + proper_matrices.to(dtype=work_dtype, device=work_device), + self.max_angular_momentum_input, + _improper=all_proper, + _wigner_D=input_wigner, + ) + + # the same rotations in the float64 integration dtype, used to + # transform outputs back to the input frame + backrotation_rotations: Optional[O3Transformation] = None + if needs_backrotation: + target_wigner: List[torch.Tensor] = [] + for o3_lambda in range(self.max_angular_momentum_target + 1): + target_wigner.append(batch_wigner[o3_lambda]) + backrotation_rotations = O3Transformation( + proper_matrices, + self.max_angular_momentum_target, + _improper=all_proper, + _wigner_D=target_wigner, + ) + + inverse_character_wigner_matrices: List[torch.Tensor] = [] + if len(character_projection_names) != 0: + for chi_lambda in range(character_max + 1): + inverse_character_wigner_matrices.append( + batch_wigner[chi_lambda].transpose(1, 2) + ) + + for coset_index in range(2): + is_improper = coset_index == 1 + if is_improper: + input_transformations = input_rotations.with_inversion() + else: + input_transformations = input_rotations + transformed_systems = input_transformations.transform_systems(system) + raw_outputs = self._model( + transformed_systems, + source_outputs, + local_selected_atoms, + ) + + for source_name in source_outputs: + if source_name not in raw_outputs: + raise ValueError( + "underlying model did not return requested output " + f"'{source_name}'" + ) + + for source_name in source_outputs: + raw_tensor = raw_outputs[source_name] + for block in raw_tensor.blocks(): + gradient_names = block.gradients_list() + if len(gradient_names) != 0: + raise ValueError( + f"underlying output '{source_name}' contains " + f"unsupported explicit gradient '{gradient_names[0]}'" + ) + + tensor = raw_tensor.to( + dtype=integration_dtype, + device=work_device, + ) + if source_name in average_names or source_name in variance_names: + # the component metadata does not change across batches: + # check it once per output + if batch_start == 0 and not is_improper: + _check_o3_lambda_limit( + tensor, + f"output '{source_name}'", + self.max_angular_momentum_target, + "max_angular_momentum_target", + ) + if backrotation_rotations is None: + raise RuntimeError( + "backrotation transformations were not prepared" + ) + if is_improper: + backrotation = ( + backrotation_rotations.with_inversion().inverse() + ) + else: + backrotation = backrotation_rotations.inverse() + backrotated = backrotation.transform_tensormap(tensor) + + if source_name in average_names: + has_average_reference = source_name in average_references + average_reference: Optional[TensorMap] = None + if has_average_reference: + average_reference = average_references[source_name] + ( + first_moment, + _, + _, + updated_average_reference, + ) = _reduce_weighted_centered_batch( + backrotated, + o3_weights, + input_system_index, + average_reference, + compute_second_moments=False, + ) + if not has_average_reference: + updated_average_reference = _copy_tensormap_info( + backrotated, + updated_average_reference, + ) + average_references[source_name] = updated_average_reference + _add_tensormap_contribution( + average_first_moments, + source_name, + first_moment, + ) + + if source_name in variance_names: + diagnostic_tensor = decompose_quantity( + source_name, + backrotated, + ) + variance_reference: Optional[TensorMap] = None + if source_name in variance_references: + variance_reference = variance_references[source_name] + ( + first_moment, + second_moment, + absolute_second_moment, + variance_reference, + ) = _reduce_weighted_centered_batch( + diagnostic_tensor, + o3_weights, + input_system_index, + variance_reference, + compute_second_moments=True, + ) + if second_moment is None or absolute_second_moment is None: + raise RuntimeError("variance moments were not computed") + variance_references[source_name] = variance_reference + _add_tensormap_contribution( + variance_first_moments, + source_name, + first_moment, + ) + _add_tensormap_contribution( + variance_second_moments, + source_name, + second_moment, + ) + _add_tensormap_contribution( + variance_absolute_second_moments, + source_name, + absolute_second_moment, + ) + + if source_name in character_projection_names: + direct_tensor = decompose_quantity(source_name, tensor) + contribution = character_projection_coefficients_from_batch( + direct_tensor, + so3_weights, + inverse_character_wigner_matrices, + input_system_index, + ) + if is_improper: + _add_tensormap_contribution( + improper_character_coefficients, + source_name, + contribution, + ) + else: + _add_tensormap_contribution( + proper_character_coefficients, + source_name, + contribution, + ) + + results: Dict[str, TensorMap] = {} + for source_name, requested_name in average_names.items(): + mean = mts.add( + average_references[source_name], + average_first_moments[source_name], + ) + mean = _copy_tensormap_info(average_references[source_name], mean) + results[requested_name] = mean.to( + dtype=work_dtype, + device=work_device, + ) + + for source_name, requested_name in variance_names.items(): + variance = _variance_from_centered_moments( + variance_first_moments[source_name], + variance_second_moments[source_name], + variance_absolute_second_moments[source_name], + n_grid_points=2 * n_rotations, + max_angular_momentum_grid=self.max_angular_momentum_grid, + ) + variance = _mean_variance_over_components( + variance, + variance_references[source_name], + ) + results[requested_name] = variance.to( + dtype=work_dtype, + device=work_device, + ) + + for source_name, requested_name in character_projection_names.items(): + projection = character_projection_tensormap_from_cosets( + proper_character_coefficients[source_name], + improper_character_coefficients[source_name], + ) + results[requested_name] = projection.to( + dtype=work_dtype, + device=work_device, + ) + + return results diff --git a/python/metatomic_torch/metatomic/torch/o3/_tranformations.py b/python/metatomic_torch/metatomic/torch/o3/_tranformations.py deleted file mode 100644 index c0a4736b..00000000 --- a/python/metatomic_torch/metatomic/torch/o3/_tranformations.py +++ /dev/null @@ -1,979 +0,0 @@ -""" -Rotate systems and tensor maps under O(3) transformations, routing rows of -multi-system tensors by their ``"system"`` sample label. -""" - -from numbers import Integral - -import torch -from metatensor.torch import Labels, LabelsEntry, TensorBlock, TensorMap - -from .. import System, register_autograd_neighbors -from ._wigner import build_wigner_D_cache - - -_INTEGER_DTYPES = ( - torch.uint8, - torch.uint16, - torch.uint32, - torch.uint64, - torch.int8, - torch.int16, - torch.int32, - torch.int64, -) - - -def _validate_nonnegative_integer(name: str, value: int) -> int: - """Validate a non-negative integer and return it as a Python int.""" - if torch.jit.is_scripting(): - integer_value = value - else: - if isinstance(value, bool) or not isinstance(value, Integral): - raise TypeError( - f"{name} must be a non-negative integer, got {type(value).__name__}." - ) - integer_value = int(value) - if integer_value < 0: - raise ValueError(f"{name} must be a non-negative integer, got {integer_value}.") - - return integer_value - - -def _spherical_parity_factor( - ell: int, - sigma: int, - is_improper: bool, -) -> int: - """Return ``sigma * (-1) ** ell`` for an improper transformation, else ``1``.""" - if torch.jit.is_scripting(): - integer_sigma = sigma - else: - if isinstance(sigma, bool) or not isinstance(sigma, Integral): - raise TypeError(f"sigma must be an integer, got {type(sigma).__name__}.") - integer_sigma = int(sigma) - if integer_sigma not in (-1, 1): - raise ValueError(f"sigma must be either -1 or +1, got {integer_sigma}.") - - if is_improper: - return integer_sigma * int((-1) ** ell) - - return 1 - - -def _validate_system_ids( - systems: list[System], - transformations: list["O3Transformation"], - system_ids: list[int] | torch.Tensor | None, - *, - expected_device: torch.device | None, -) -> torch.Tensor: - """Check and normalize the ``system_ids`` argument of ``transform_tensor``. - - ``system_ids[i]`` is the value in a block's ``"system"`` sample column that - selects ``transformations[i]``. This checks that systems and transformations - pair up one-to-one and that there is one distinct integer id per system, - returning the ids as a ``torch.long`` tensor (``0..n_systems - 1`` when - ``system_ids`` is ``None``). - """ - n_systems = len(systems) - n_transformations = len(transformations) - if n_systems != n_transformations: - raise ValueError( - "Expected one transformation per system, but got " - f"len(systems)={n_systems} and " - f"len(transformations)={n_transformations}." - ) - - if system_ids is None: - return torch.arange(n_systems, dtype=torch.long, device=expected_device) - - if isinstance(system_ids, torch.Tensor): - if system_ids.ndim != 1: - raise ValueError( - "system_ids must be one-dimensional, but got a tensor with shape " - f"{tuple(system_ids.shape)}." - ) - if system_ids.dtype not in _INTEGER_DTYPES: - raise ValueError( - "system_ids must contain integers, but got a tensor with dtype " - f"{system_ids.dtype}." - ) - if expected_device is not None and system_ids.device != expected_device: - raise ValueError( - f"system_ids are on device {system_ids.device}, but the values to " - f"transform are on device {expected_device}." - ) - validated_ids = system_ids.to(dtype=torch.long) - else: - python_ids: list[int] = [] - for system_id in system_ids: - if isinstance(system_id, bool) or not isinstance(system_id, Integral): - raise ValueError("system_ids must contain integers.") - python_ids.append(int(system_id)) - validated_ids = torch.tensor( - python_ids, - dtype=torch.long, - device=expected_device, - ) - - if len(validated_ids) != n_systems: - raise ValueError( - "system_ids must contain exactly one entry per system, but got " - f"len(system_ids)={len(validated_ids)} and len(systems)={n_systems}." - ) - if torch.unique(validated_ids).numel() != n_systems: - raise ValueError( - "system_ids must contain one distinct entry per system, but got " - f"{validated_ids.tolist()}." - ) - - return validated_ids - - -def _validate_transformations_dtype_device( - transformations: list["O3Transformation"], - *, - expected_dtype: torch.dtype, - expected_device: torch.device, -) -> None: - """Check that every transformation has the expected dtype and device.""" - for index, transformation in enumerate(transformations): - if ( - transformation.dtype != expected_dtype - or transformation.device != expected_device - ): - raise ValueError( - f"Transformation at index {index} has dtype/device " - f"({transformation.dtype}, {transformation.device}), differing from " - f"the values to transform ({expected_dtype}, {expected_device})." - ) - - -class O3Transformation: - """ - A single O(3) transformation, represented by a (3, 3) rotation or improper-rotation - matrix. - - The constructor stores a copy of ``matrix``. - """ - - def __init__(self, matrix: torch.Tensor, max_angular_momentum: int): - """ - :param matrix: (3, 3) rotation or improper-rotation matrix - :param max_angular_momentum: non-negative maximum angular momentum for - which Wigner-D matrices are available - """ - max_angular_momentum = _validate_nonnegative_integer( - "max_angular_momentum", 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)." - ) - - # Keep an independent copy so modifying the input tensor later cannot make - # the matrix disagree with the cached parity and Wigner-D matrices. - self._matrix = matrix.clone() - self._max_angular_momentum = max_angular_momentum - self._is_improper = bool(torch.det(self._matrix) < 0) - - self._wigner_D_cache: dict[int, torch.Tensor] | None = None - - @classmethod - def _create_no_checks( - cls, - matrix: torch.Tensor, - max_angular_momentum: int, - *, - is_improper: bool, - ) -> "O3Transformation": - """Create a transformation after validation in ``random_transformations``. - - The random factory validates its arguments and matrices before calling this - method. This avoids repeating the public constructor's checks, matrix copy, - and determinant calculation for every matrix. ``is_improper`` must match - ``matrix``. - """ - transformation = cls.__new__(cls) - transformation._matrix = matrix - transformation._max_angular_momentum = max_angular_momentum - transformation._is_improper = is_improper - transformation._wigner_D_cache = None - return transformation - - def _ensure_wigner_D_cache(self) -> dict[int, torch.Tensor]: - """Ensure that the Wigner-D cache has been built and return it.""" - if self._wigner_D_cache is None: - self._wigner_D_cache = build_wigner_D_cache( - self._max_angular_momentum, - self._matrix, - device=self._matrix.device, - dtype=self._matrix.dtype, - ) - - return self._wigner_D_cache - - def _wigner_D_cache_entry(self, ell: int) -> torch.Tensor: - """Return the internal cache entry for ``ell`` without copying it.""" - ell = self._validate_ell_range(ell) - - D = self._ensure_wigner_D_cache().get(ell) - if D is None: - raise ValueError(f"Wigner-D matrix for ell={ell} not found in cache.") - - return D - - @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_improper(self) -> bool: - """Whether this transformation is improper, with negative determinant.""" - return self._is_improper - - 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 _validate_ell_range(self, ell: int) -> int: - """Check that ``ell`` is an integer in ``[0, max_angular_momentum]``.""" - ell = _validate_nonnegative_integer("ell", ell) - - if ell > self._max_angular_momentum: - raise ValueError( - f"ell={ell} exceeds max_angular_momentum={self._max_angular_momentum}." - ) - - return ell - - 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 in ``[0, max_angular_momentum]`` - :param sigma: ``+1`` for a proper spherical representation or ``-1`` for - a pseudo one. Under an improper transformation, the representation - acquires the factor ``sigma * (-1) ** ell``. - :return: (..., 2*ell+1) tensor of transformed spherical values - """ - ell = self._validate_ell_range(ell) - parity_factor = _spherical_parity_factor( - ell, - sigma, - is_improper=self.is_improper, - ) - - D = self._wigner_D_cache_entry(ell) - transformed = values @ D.T - if parity_factor != 1: - transformed = transformed * parity_factor - - return transformed - - def wigner_D_matrix(self, ell: int) -> torch.Tensor: - """Return the proper-part Wigner-D matrix for ``ell``. - - For an improper transformation, :meth:`transform_spherical` applies the - inversion-parity factor separately. - - :param ell: angular momentum in ``[0, max_angular_momentum]`` - :return: (2*ell+1, 2*ell+1) Wigner-D matrix - """ - return self._wigner_D_cache_entry(ell) - - -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`` transformations uniformly from SO(3), or from O(3) when - inversions are included. - - 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: non-negative number of transformations to generate - :param max_angular_momentum: non-negative maximum angular momentum for - Wigner-D matrices - :param device: target device for the output tensors - :param dtype: target dtype for the output tensors; must be - :attr:`torch.float32` or :attr:`torch.float64` - :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`` :class:`O3Transformation` objects - """ - n = _validate_nonnegative_integer("n", n) - max_angular_momentum = _validate_nonnegative_integer( - "max_angular_momentum", max_angular_momentum - ) - - if dtype not in (torch.float32, torch.float64): - raise ValueError(f"dtype must be torch.float32 or torch.float64, got {dtype}.") - - 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) - - matrices_are_improper = [False] * n - - 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) - matrices_are_improper = (signs < 0).tolist() - - identity = torch.eye( - 3, - device=R.device, - dtype=R.dtype, - ).expand(n, 3, 3) - if not torch.allclose( - R @ R.transpose(-1, -2), - identity, - atol=1e-5, - ): - raise ValueError("Generated transformations are not orthogonal.") - - return [ - O3Transformation._create_no_checks( - matrix, - max_angular_momentum, - is_improper=is_improper, - ) - for matrix, is_improper in zip( - R.unbind(0), - matrices_are_improper, - strict=True, - ) - ] - - -def _value_row_indices_by_system( - block: TensorBlock, - system_ids: torch.Tensor, -) -> list[torch.Tensor]: - """Return value-row indices in ``system_ids`` order, or all rows for one system.""" - if len(system_ids) == 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_labels = block.samples.column("system").to(dtype=torch.long) - unique_labels = torch.unique(system_labels) - labels_are_known = torch.isin(unique_labels, system_ids) - if not labels_are_known.all(): - unknown_labels = unique_labels[~labels_are_known] - raise ValueError( - f"Block samples contain system labels {unknown_labels.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_labels == system_id, as_tuple=False).reshape(-1) - for system_id in system_ids - ] - - -def _gradient_row_indices_by_system( - grad_block: TensorBlock, - parent_block: TensorBlock, - system_ids: torch.Tensor, -) -> list[torch.Tensor]: - """Group gradient rows by the system of their referenced value row.""" - if len(system_ids) == 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_labels = parent_block.samples.column("system").to(dtype=torch.long) - parent_value_rows = grad_block.samples.column("sample").to(dtype=torch.long) - gradient_system_labels = parent_system_labels[parent_value_rows] - - return [ - torch.nonzero( - gradient_system_labels == system_id, - as_tuple=False, - ).reshape(-1) - for system_id in system_ids - ] - - -def transform_system(system: System, transformation: O3Transformation) -> System: - """Apply an O(3) transformation to a single System. - - Positions, cell vectors, neighbor-list displacements, and custom data following - :ref:`o3-conventions` are transformed. Atomic types and periodic-boundary flags - are preserved. - - :param system: input system - :param transformation: O(3) transformation to apply, matching - ``system.positions`` in dtype and device - :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 - # Detach the input graph before registering the rotated values below. - neighbors_values = neighbors.values.detach().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 - """ - # Reserve einsum indices for all ten component axes supported by Metatomic. - _EINSUM_IN = "abcdefghjk" - _EINSUM_OUT = "ABCDEFGHIJ" - - if len(matrices) == 0: - return values - n_axes = len(matrices) - if n_axes > len(_EINSUM_IN): - raise ValueError(f"can not transform a tensor with {n_axes} component axes") - 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 _component_axis_suffix(axis_name: str, prefix: str) -> tuple[bool, str]: - """Match a component-axis name and return its supported suffix.""" - suffixes = ["", "_1", "_2", "_3", "_4", "_5", "_6", "_7", "_8", "_9"] - for suffix in suffixes: - if axis_name == prefix + suffix: - return True, suffix - return False, "" - - -def _validate_component_axis_metadata( - components: list[Labels], - key: LabelsEntry, -) -> list[tuple[bool, int, int]]: - """Validate component axes and return ``(is_spherical, ell, sigma)`` metadata.""" - if len(components) > 10: - raise ValueError( - f"can not transform a tensor with {len(components)} component axes; " - "at most 10 are supported" - ) - - metadata: list[tuple[bool, int, int]] = [] - for component in components: - axis_name = component.names[0] - is_cartesian, _ = _component_axis_suffix(axis_name, "xyz") - is_spherical, suffix = _component_axis_suffix(axis_name, "o3_mu") - if is_cartesian: - expected_labels = torch.arange( - 3, - device=component.values.device, - dtype=component.values.dtype, - ) - if not torch.equal(component.values[:, 0], expected_labels): - raise ValueError( - f"Cartesian component axis '{axis_name}' must use labels " - "[0, 1, 2] in x, y, z order." - ) - metadata.append((False, 0, 1)) - elif is_spherical: - ell = _validate_nonnegative_integer( - "ell", - int(key["o3_lambda" + suffix]), - ) - sigma = int(key["o3_sigma" + suffix]) - _spherical_parity_factor(ell, sigma, is_improper=False) - - expected_labels = torch.arange( - -ell, - ell + 1, - device=component.values.device, - dtype=component.values.dtype, - ) - if not torch.equal(component.values[:, 0], expected_labels): - raise ValueError( - f"Spherical component axis '{axis_name}' for ell={ell} must use " - f"labels from {-ell} through {ell} in ascending order." - ) - metadata.append((True, 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 metadata - - -def _max_o3_lambda_in_tensor(tensor: TensorMap) -> int: - """Return the largest spherical rank in block values or attached gradients. - - A TensorMap containing only scalar or Cartesian component axes returns ``-1``. - """ - max_o3_lambda = -1 - for key, block in tensor.items(): - metadata = _validate_component_axis_metadata(block.components, key) - for is_spherical, ell, _sigma in metadata: - if is_spherical and ell > max_o3_lambda: - max_o3_lambda = ell - - for _gradient_name, gradient in block.gradients(): - gradient_metadata = _validate_component_axis_metadata( - gradient.components, - key, - ) - for is_spherical, ell, _sigma in gradient_metadata: - if is_spherical and ell > max_o3_lambda: - max_o3_lambda = ell - - return max_o3_lambda - - -def _axis_matrices_and_parity( - metadata: list[tuple[bool, int, int]], - transformation: O3Transformation, -) -> tuple[list[torch.Tensor], int]: - """Return the axis matrices and their combined spherical parity factor.""" - matrices: list[torch.Tensor] = [] - parity = 1 - for is_spherical, ell, sigma in metadata: - if is_spherical: - matrices.append(transformation._wigner_D_cache_entry(ell)) - parity *= _spherical_parity_factor( - ell, - sigma, - transformation.is_improper, - ) - else: - matrices.append(transformation._matrix) - - 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 value or gradient rows with their assigned transformation.""" - metadata = _validate_component_axis_metadata(components, key) - new_values = values.clone() - for system_index, rows in enumerate(row_indices): - if len(rows) == 0: - continue - matrices, parity = _axis_matrices_and_parity( - metadata, - 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: - """Apply per-system O(3) transformations to a block and its gradients. - - With one system, the ``"system"`` sample label is optional and ignored, as in - :py:func:`transform_tensor`. - - :param key: parent block key, supplying the O(3) labels required by spherical - component axes - :param block: block to transform - :param systems: systems corresponding positionally to ``transformations`` - :param transformations: one O(3) transformation per system, matching - ``block.values`` in dtype and device - :param system_ids: one distinct integer ``"system"`` sample label per system; - entry ``i`` is paired with ``transformations[i]``. A tensor argument must - be one-dimensional and use the same device as ``block.values``. Defaults - to ``range(len(systems))`` - :return: block with transformed values and gradients and unchanged labels; when - ``systems`` is empty, the block is unchanged - """ - system_ids = _validate_system_ids( - systems, - transformations, - system_ids, - expected_device=block.values.device, - ) - if len(systems) == 0: - return block - - _validate_transformations_dtype_device( - transformations, - expected_dtype=block.values.dtype, - expected_device=block.values.device, - ) - - return _transform_block_impl(key, block, transformations, system_ids) - - -def _transform_block_impl( - key: LabelsEntry, - block: TensorBlock, - transformations: list[O3Transformation], - system_ids: torch.Tensor, -) -> TensorBlock: - """Transform block values and gradients using validated system assignments.""" - value_sample_indices = _value_row_indices_by_system(block, system_ids) - new_block = TensorBlock( - values=_transform_component_values( - block.values, - block.components, - key, - value_sample_indices, - transformations, - ), - samples=block.samples, - components=block.components, - properties=block.properties, - ) - for gradient_name, gradient in block.gradients(): - gradient_sample_indices = _gradient_row_indices_by_system( - gradient, - block, - system_ids, - ) - new_block.add_gradient( - gradient_name, - TensorBlock( - values=_transform_component_values( - gradient.values, - gradient.components, - key, - gradient_sample_indices, - transformations, - ), - samples=gradient.samples, - components=gradient.components, - properties=gradient.properties, - ), - ) - return new_block - - -def transform_tensor( - tensor: TensorMap, - systems: list[System], - transformations: list[O3Transformation], - system_ids: list[int] | torch.Tensor | None = None, -) -> TensorMap: - """Apply per-system O(3) transformations to a TensorMap and its gradients. - - Scalar, Cartesian, and spherical data are identified by their component-axis - names, following :ref:`o3-conventions`; one :py:class:`TensorMap` may contain - all three kinds of data. At most ten component axes are supported in one - value or gradient block. - - With multiple systems, the ``"system"`` sample label assigns each value sample - to a transformation: samples labelled ``system_ids[i]`` use - ``transformations[i]``. A block may contain samples for only some of the - systems, but every ``"system"`` label present in the block must appear in - ``system_ids``. A gradient sample uses the same transformation as the parent - value sample referenced by its ``"sample"`` label. With one system, the - ``"system"`` label is optional and ignored. - - :param tensor: TensorMap to transform - :param systems: systems corresponding positionally to ``transformations`` - :param transformations: one O(3) transformation per system, matching the tensor - values in dtype and device when present - :param system_ids: one distinct integer ``"system"`` sample label per system; - entry ``i`` is paired with ``transformations[i]``. A tensor argument must - be one-dimensional and use the same device as the tensor values. Defaults - to ``range(len(systems))`` - :return: transformed TensorMap with the same keys and global information; when - ``systems`` is empty, the tensor is unchanged - """ - if len(tensor) != 0: - system_ids_device = tensor.block(0).values.device - elif len(transformations) != 0: - system_ids_device = transformations[0].device - else: - system_ids_device = None - - system_ids = _validate_system_ids( - systems, - transformations, - system_ids, - expected_device=system_ids_device, - ) - if len(systems) == 0: - return tensor - - if len(tensor) != 0: - values = tensor.block(0).values - _validate_transformations_dtype_device( - transformations, - expected_dtype=values.dtype, - expected_device=values.device, - ) - - new_blocks = [ - _transform_block_impl(key, block, transformations, system_ids) - for key, block in tensor.items() - ] - transformed = TensorMap(keys=tensor.keys, blocks=new_blocks) - for info_key, info_value in tensor.info().items(): - transformed.set_info(info_key, info_value) - - return transformed - - -def _transformation_indices( - samples: Labels, - n_transformations: int, -) -> torch.Tensor: - """Map sample rows to local transformation indices.""" - if n_transformations <= 0: - raise ValueError("n_transformations must be positive") - if n_transformations == 1: - return torch.zeros( - len(samples), - dtype=torch.long, - device=samples.device, - ) - if "system" not in samples.names: - raise ValueError("multiple transformations require a 'system' sample dimension") - - indices = samples.column("system").to(dtype=torch.long) - if bool(torch.any((indices < 0) | (indices >= n_transformations)).item()): - raise ValueError("sample system indices exceed the transformation batch") - return indices - - -def _transform_component_values_with_precomputed_matrices( - values: torch.Tensor, - components: list[Labels], - key: LabelsEntry, - transformation_indices: torch.Tensor, - matrices: torch.Tensor, - wigner_matrices: list[torch.Tensor], - is_improper: bool, -) -> torch.Tensor: - """Transform component axes with precomputed O(3) matrices.""" - metadata = _validate_component_axis_metadata(components, key) - if len(metadata) == 0: - return values.clone() - - transformed = values - parity = 1 - for component_index, (is_spherical, ell, sigma) in enumerate(metadata): - if is_spherical: - if ell >= len(wigner_matrices): - raise ValueError("spherical rank exceeds the Wigner-D storage") - axis_matrices = wigner_matrices[ell] - parity *= _spherical_parity_factor(ell, sigma, is_improper) - else: - axis_matrices = matrices - - component_axis = component_index + 1 - moved = torch.movedim(transformed, component_axis, -1) - moved_shape = moved.shape - flattened = moved.flatten(start_dim=1, end_dim=-2) - matrices_for_rows = axis_matrices.index_select( - 0, - transformation_indices, - ) - transformed = torch.bmm( - flattened, - matrices_for_rows.transpose(1, 2), - ) - transformed = transformed.reshape(moved_shape) - transformed = torch.movedim(transformed, -1, component_axis) - - if parity != 1: - transformed = transformed * parity - return transformed - - -def _transform_tensor_with_precomputed_matrices( - tensor: TensorMap, - matrices: torch.Tensor, - wigner_matrices: list[torch.Tensor], - is_improper: bool, -) -> TensorMap: - """Transform a TensorMap using precomputed matrices from one O(3) coset. - - ``matrices[i]`` is the actual Cartesian operation for local system ``i``, - while ``wigner_matrices[ell][i]`` is the Wigner-D matrix for its proper - rotational part. Every operation in the batch must be either proper or - improper, as selected by ``is_improper``. - - With multiple operations, ``"system"`` sample labels are local indices into - the matrix batch. A singleton batch does not require this sample dimension. - The caller chooses the transformation direction by supplying either the - forward matrices or their inverses. - """ - if ( - matrices.dim() != 3 - or matrices.size(0) == 0 - or matrices.size(1) != 3 - or matrices.size(2) != 3 - ): - raise ValueError("matrices must have shape (N, 3, 3) with N > 0") - if matrices.dtype != torch.float32 and matrices.dtype != torch.float64: - raise TypeError("matrices must use float32 or float64") - if len(tensor) != 0: - reference_values = tensor.block(0).values - if ( - matrices.dtype != reference_values.dtype - or matrices.device != reference_values.device - ): - raise ValueError("tensor and matrices must have the same dtype and device") - - blocks: list[TensorBlock] = [] - for key, block in tensor.items(): - value_indices = _transformation_indices( - block.samples, - matrices.size(0), - ) - new_block = TensorBlock( - values=_transform_component_values_with_precomputed_matrices( - block.values, - block.components, - key, - value_indices, - matrices, - wigner_matrices, - is_improper, - ), - samples=block.samples, - components=block.components, - properties=block.properties, - ) - - for gradient_name, gradient in block.gradients(): - parent_rows = gradient.samples.column("sample").to(dtype=torch.long) - gradient_indices = value_indices.index_select(0, parent_rows) - new_block.add_gradient( - gradient_name, - TensorBlock( - values=_transform_component_values_with_precomputed_matrices( - gradient.values, - gradient.components, - key, - gradient_indices, - matrices, - wigner_matrices, - is_improper, - ), - samples=gradient.samples, - components=gradient.components, - properties=gradient.properties, - ), - ) - blocks.append(new_block) - - transformed = TensorMap(tensor.keys, blocks) - for info_name, info_value in tensor.info().items(): - transformed.set_info(info_name, info_value) - return transformed diff --git a/python/metatomic_torch/metatomic/torch/o3/_transformations.py b/python/metatomic_torch/metatomic/torch/o3/_transformations.py new file mode 100644 index 00000000..ee7955cf --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/o3/_transformations.py @@ -0,0 +1,1181 @@ +""" +Rotate systems and tensor maps under O(3) transformations, routing rows of +multi-system tensors by their ``"system"`` sample label. + +:py:class:`O3Transformation` holds a batch of one or more operations. The +tensor-transformation kernel in this module is TorchScript compatible, so a +scripted model can construct transformations from precomputed tensors inside +``forward`` and share one implementation with the eager public functions. +""" + +from numbers import Integral +from typing import Optional + +import torch +from metatensor.torch import Labels, LabelsEntry, TensorBlock, TensorMap + +from .. import System, register_autograd_neighbors +from ._wigner import build_packed_wigner_matrices, wigner_matrices_for_lambda + + +_INTEGER_DTYPES = ( + torch.uint8, + torch.uint16, + torch.uint32, + torch.uint64, + torch.int8, + torch.int16, + torch.int32, + torch.int64, +) + + +def _validate_nonnegative_integer(name: str, value: int) -> int: + """Validate a non-negative integer and return it as a Python int.""" + if torch.jit.is_scripting(): + integer_value = value + else: + if isinstance(value, bool) or not isinstance(value, Integral): + raise TypeError( + f"{name} must be a non-negative integer, got {type(value).__name__}." + ) + integer_value = int(value) + if integer_value < 0: + raise ValueError(f"{name} must be a non-negative integer, got {integer_value}.") + + return integer_value + + +def _spherical_parity_factor( + ell: int, + sigma: int, + is_improper: bool, +) -> int: + """Return ``sigma * (-1) ** ell`` for an improper transformation, else ``1``.""" + if torch.jit.is_scripting(): + integer_sigma = sigma + else: + if isinstance(sigma, bool) or not isinstance(sigma, Integral): + raise TypeError(f"sigma must be an integer, got {type(sigma).__name__}.") + integer_sigma = int(sigma) + if integer_sigma not in (-1, 1): + raise ValueError(f"sigma must be either -1 or +1, got {integer_sigma}.") + + if is_improper: + return integer_sigma * int((-1) ** ell) + + return 1 + + +def _determinants_3x3(matrices: torch.Tensor) -> torch.Tensor: + """Return the determinants of a ``(N, 3, 3)`` batch of matrices. + + Written out explicitly because ``torch.linalg`` is not available in + TorchScript. + """ + a = matrices[:, 0, 0] + b = matrices[:, 0, 1] + c = matrices[:, 0, 2] + d = matrices[:, 1, 0] + e = matrices[:, 1, 1] + f = matrices[:, 1, 2] + g = matrices[:, 2, 0] + h = matrices[:, 2, 1] + i = matrices[:, 2, 2] + return a * (e * i - f * h) - b * (d * i - f * g) + c * (d * h - e * g) + + +class O3Transformation: + """ + A batch of one or more O(3) transformations, represented by ``(N, 3, 3)`` + rotation or improper-rotation matrices. A single ``(3, 3)`` matrix is + stored as a batch of one. + + The constructor stores a copy of ``matrix`` and builds the Wigner-D + matrices lazily on first use, which requires the eager ``wigners`` package. + Scripted models construct transformations from precomputed tensors instead, + through the private constructor arguments. + """ + + def __init__( + self, + matrix: torch.Tensor, + max_angular_momentum: int, + _improper: Optional[torch.Tensor] = None, + _wigner_D: Optional[list[torch.Tensor]] = None, + ): + """ + :param matrix: ``(3, 3)`` or ``(N, 3, 3)`` rotation or + improper-rotation matrices + :param max_angular_momentum: non-negative maximum angular momentum for + which Wigner-D matrices are available + :param _improper: private trusted path used for internal batching: + an ``(N,)`` boolean mask of the negative-determinant operations, + paired with already-validated ``(N, 3, 3)`` matrices which are + stored without copying or checks + :param _wigner_D: private, only together with ``_improper``: one + ``(N, 2*ell+1, 2*ell+1)`` stack of proper-part Wigner-D matrices + per ``ell`` through ``max_angular_momentum``; ``None`` defers to + the lazy eager build + """ + if _improper is not None: + matrices = matrix + improper = _improper + else: + max_angular_momentum = _validate_nonnegative_integer( + "max_angular_momentum", max_angular_momentum + ) + + if matrix.dim() == 2: + matrices = matrix.unsqueeze(0) + else: + matrices = matrix + if ( + matrices.dim() != 3 + or matrices.size(0) == 0 + or matrices.size(1) != 3 + or matrices.size(2) != 3 + ): + if torch.jit.is_scripting(): + raise ValueError( + "transformation matrices must have shape (3, 3) or " + "(N, 3, 3) with N > 0" + ) + else: + raise ValueError( + f"Transformation has shape {tuple(matrix.shape)}; " + "expected (3, 3) or (N, 3, 3) with N > 0." + ) + + identity = torch.eye(3, device=matrices.device, dtype=matrices.dtype) + if not torch.allclose( + matrices @ matrices.transpose(1, 2), + identity, + atol=1e-5, + ): + raise ValueError( + "Transformation is not orthogonal (R @ R.T deviates from I)." + ) + + # Keep an independent copy so modifying the input tensor later cannot + # make the matrices disagree with the parity and Wigner-D matrices. + matrices = matrices.clone() + improper = _determinants_3x3(matrices) < 0.0 + + self._matrices = matrices + self._max_angular_momentum = max_angular_momentum + self._improper = improper + self._wigner_D = _wigner_D + + @property + def matrices(self) -> torch.Tensor: + """The ``(N, 3, 3)`` batch of rotation or improper-rotation matrices.""" + return self._matrices + + @property + def matrix(self) -> torch.Tensor: + """The ``(3, 3)`` matrix of a single transformation. + + Raises for a batch of more than one operation; use :py:attr:`matrices` + there. + """ + if self._matrices.size(0) != 1: + raise ValueError( + f"this O3Transformation holds {self._matrices.size(0)} " + "operations; use .matrices" + ) + return self._matrices[0] + + @property + def max_angular_momentum(self) -> int: + """The maximum angular momentum with available Wigner-D matrices.""" + return self._max_angular_momentum + + @property + def improper(self) -> torch.Tensor: + """Boolean mask marking the improper operations in the batch.""" + return self._improper + + @property + def is_improper(self) -> bool: + """Whether the transformations are improper, with negative determinant. + + Raises for a batch mixing proper and improper operations; use + :py:attr:`improper` there. + """ + n_improper = int(self._improper.to(dtype=torch.long).sum().item()) + if n_improper == 0: + return False + if n_improper == self._improper.numel(): + return True + raise ValueError( + "this O3Transformation mixes proper and improper operations; use .improper" + ) + + @property + @torch.jit.unused + def dtype(self) -> torch.dtype: + """The dtype of the transformation matrices.""" + return self._matrices.dtype + + @property + @torch.jit.unused + def device(self) -> torch.device: + """The device of the transformation matrices.""" + return self._matrices.device + + def _validate_ell_range(self, ell: int) -> int: + """Check that ``ell`` is an integer in ``[0, max_angular_momentum]``.""" + ell = _validate_nonnegative_integer("ell", ell) + + if ell > self._max_angular_momentum: + raise ValueError( + f"ell={ell} exceeds max_angular_momentum={self._max_angular_momentum}." + ) + + return ell + + def _wigner_D_matrices(self) -> list[torch.Tensor]: + """Return the per-``ell`` Wigner-D stacks, building them on first use.""" + wigner_D = self._wigner_D + if wigner_D is None: + wigner_D = self._build_wigner_D() + self._wigner_D = wigner_D + return wigner_D + + @torch.jit.unused + def _build_wigner_D(self) -> list[torch.Tensor]: + """Build the Wigner-D stacks with the eager numpy-based path.""" + packed = build_packed_wigner_matrices( + self._matrices, + self._max_angular_momentum, + ) + n_matrices = self._matrices.size(0) + return [ + wigner_matrices_for_lambda(packed, n_matrices, ell) + for ell in range(self._max_angular_momentum + 1) + ] + + def wigner_D_matrices(self, ell: int) -> torch.Tensor: + """Return the proper-part Wigner-D matrices for ``ell``. + + For improper operations, the inversion-parity factor + ``sigma * (-1) ** ell`` is applied separately when transforming + spherical values. + + :param ell: angular momentum in ``[0, max_angular_momentum]`` + :return: ``(N, 2*ell+1, 2*ell+1)`` stack of Wigner-D matrices + """ + ell = self._validate_ell_range(ell) + return self._wigner_D_matrices()[ell] + + def wigner_D_matrix(self, ell: int) -> torch.Tensor: + """Return the proper-part Wigner-D matrix of a single transformation. + + Raises for a batch of more than one operation; use + :py:meth:`wigner_D_matrices` there. + + :param ell: angular momentum in ``[0, max_angular_momentum]`` + :return: (2*ell+1, 2*ell+1) Wigner-D matrix + """ + if self._matrices.size(0) != 1: + raise ValueError( + f"this O3Transformation holds {self._matrices.size(0)} " + "operations; use .wigner_D_matrices" + ) + return self.wigner_D_matrices(ell)[0] + + def inverse(self) -> "O3Transformation": + """Return the batch of inverse transformations. + + The inverse of an orthogonal matrix is its transpose, and the Wigner-D + matrices of the inverse are the transposed Wigner-D matrices, so this + returns transposed views of the existing storage without copying. + """ + wigner_D = self._wigner_D + inverse_wigner: Optional[list[torch.Tensor]] = None + if wigner_D is not None: + inverse_wigner = [D.transpose(1, 2) for D in wigner_D] + return O3Transformation( + self._matrices.transpose(1, 2), + self._max_angular_momentum, + _improper=self._improper, + _wigner_D=inverse_wigner, + ) + + def with_inversion(self) -> "O3Transformation": + """Return the batch composed with the inversion. + + Composing with the inversion negates the matrices and flips their + parity, while the proper rotational part -- and with it the Wigner-D + matrices -- is unchanged and shared with this batch. + """ + return O3Transformation( + -self._matrices, + self._max_angular_momentum, + _improper=torch.logical_not(self._improper), + _wigner_D=self._wigner_D, + ) + + def transform_cartesian(self, vectors: torch.Tensor) -> torch.Tensor: + """Apply the transformations to Cartesian vectors. + + :param vectors: ``(..., 3)`` tensor of Cartesian vectors + :return: transformed vectors, with the input shape for a single + transformation or a leading batch axis (``(N, ..., 3)``) for a + batch of more than one + """ + if self._matrices.size(0) == 1: + return vectors @ self._matrices[0].transpose(0, 1) + + flattened = vectors.reshape(1, -1, 3) + transformed = flattened @ self._matrices.transpose(1, 2) + output_shape: list[int] = [self._matrices.size(0)] + for size in vectors.shape: + output_shape.append(size) + return transformed.reshape(output_shape) + + def transform_spherical( + self, values: torch.Tensor, ell: int, sigma: int + ) -> torch.Tensor: + """Apply the transformations to spherical values. + + :param values: (..., 2*ell+1) tensor of spherical values + :param ell: angular momentum in ``[0, max_angular_momentum]`` + :param sigma: ``+1`` for a proper spherical representation or ``-1`` for + a pseudo one. Under an improper transformation, the representation + acquires the factor ``sigma * (-1) ** ell``. + :return: transformed values, with the input shape for a single + transformation or a leading batch axis for a batch of more than one + """ + ell = self._validate_ell_range(ell) + # the parity factor acquired by the improper operations in the batch; + # this also validates sigma + parity = _spherical_parity_factor(ell, sigma, True) + D = self.wigner_D_matrices(ell) + + if self._matrices.size(0) == 1: + transformed = values @ D[0].transpose(0, 1) + if parity != 1 and bool(torch.any(self._improper).item()): + transformed = transformed * float(parity) + return transformed + + dimension = 2 * ell + 1 + flattened = values.reshape(1, -1, dimension) + transformed = flattened @ D.transpose(1, 2) + if parity != 1 and bool(torch.any(self._improper).item()): + factors = torch.where( + self._improper, + torch.tensor(float(parity), dtype=values.dtype, device=values.device), + torch.tensor(1.0, dtype=values.dtype, device=values.device), + ) + transformed = transformed * factors.view(-1, 1, 1) + output_shape: list[int] = [self._matrices.size(0)] + for size in values.shape: + output_shape.append(size) + return transformed.reshape(output_shape) + + def transform_systems(self, system: System) -> list[System]: + """Apply every transformation in the batch to one System. + + Positions, cell vectors, neighbor-list displacements, and custom data + following :ref:`o3-conventions` are transformed. Atomic types and + periodic-boundary flags are preserved. + + :param system: input system, matching the transformation matrices in + dtype and device + :return: one transformed System per operation in the batch + """ + matrices = self._matrices + if ( + matrices.dtype != system.positions.dtype + or matrices.device != system.positions.device + ): + raise ValueError( + "system and transformation matrices must have the same dtype and device" + ) + + positions = system.positions.unsqueeze(0) @ matrices.transpose(1, 2) + cells = system.cell.unsqueeze(0) @ matrices.transpose(1, 2) + + transformed_systems: list[System] = [] + for index in range(matrices.size(0)): + transformed_systems.append( + System( + positions=positions[index], + types=system.types, + cell=cells[index], + pbc=system.pbc, + ) + ) + + for options in system.known_neighbor_lists(): + neighbors = system.get_neighbor_list(options) + # neighbor vectors are stored as (n_pairs, 3, 1); squeeze/unsqueeze + # around the matmul. Detach the input graph before registering the + # rotated values below. + source_values = neighbors.values.detach().squeeze(-1) + neighbor_values = source_values.unsqueeze(0) @ matrices.transpose(1, 2) + for index in range(matrices.size(0)): + rotated_neighbors = TensorBlock( + values=neighbor_values[index].unsqueeze(-1), + samples=neighbors.samples, + components=neighbors.components, + properties=neighbors.properties, + ) + register_autograd_neighbors( + transformed_systems[index], + rotated_neighbors, + ) + transformed_systems[index].add_neighbor_list( + options, + rotated_neighbors, + ) + + for data_name in system.known_data(): + data = system.get_data(data_name) + wigner_matrices: list[torch.Tensor] = [] + if _max_o3_lambda_in_tensor(data) >= 0: + wigner_matrices = self._wigner_D_matrices() + for index in range(matrices.size(0)): + index_wigner: list[torch.Tensor] = [] + for D in wigner_matrices: + index_wigner.append(D[index : index + 1]) + transformed_systems[index].add_data( + data_name, + _transform_tensormap_batched( + data, + matrices[index : index + 1], + index_wigner, + self._improper[index : index + 1], + None, + ), + ) + + return transformed_systems + + def transform_tensormap( + self, + tensor: TensorMap, + system_ids: Optional[torch.Tensor] = None, + ) -> TensorMap: + """Apply the transformations to a TensorMap and its gradients. + + Scalar, Cartesian, and spherical data are identified by their + component-axis names, following :ref:`o3-conventions`. With a batch of + more than one operation, the ``"system"`` sample label assigns each + value row to an operation: when ``system_ids`` is ``None``, the labels + index the batch directly, and otherwise rows labelled ``system_ids[i]`` + use operation ``i``. Gradient rows use the operation of the value row + referenced by their ``"sample"`` label. With a single operation, the + ``"system"`` label is optional and ignored. + + :param tensor: TensorMap to transform, matching the transformation + matrices in dtype and device + :param system_ids: optional one-dimensional tensor with one distinct + ``"system"`` sample label per operation in the batch + :return: transformed TensorMap with the same metadata and global + information + """ + wigner_matrices: list[torch.Tensor] = [] + if _max_o3_lambda_in_tensor(tensor) >= 0: + wigner_matrices = self._wigner_D_matrices() + return _transform_tensormap_batched( + tensor, + self._matrices, + wigner_matrices, + self._improper, + system_ids, + ) + + +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`` transformations uniformly from SO(3), or from O(3) when + inversions are included. + + 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: non-negative number of transformations to generate + :param max_angular_momentum: non-negative maximum angular momentum for + Wigner-D matrices + :param device: target device for the output tensors + :param dtype: target dtype for the output tensors; must be + :attr:`torch.float32` or :attr:`torch.float64` + :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`` single-operation :class:`O3Transformation` objects + """ + n = _validate_nonnegative_integer("n", n) + max_angular_momentum = _validate_nonnegative_integer( + "max_angular_momentum", max_angular_momentum + ) + + if dtype not in (torch.float32, torch.float64): + raise ValueError(f"dtype must be torch.float32 or torch.float64, got {dtype}.") + + 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) + + improper = torch.zeros(n, dtype=torch.bool, device=device) + + 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) + improper = signs < 0 + + identity = torch.eye( + 3, + device=R.device, + dtype=R.dtype, + ).expand(n, 3, 3) + if not torch.allclose( + R @ R.transpose(-1, -2), + identity, + atol=1e-5, + ): + raise ValueError("Generated transformations are not orthogonal.") + + return [ + O3Transformation( + matrix.unsqueeze(0), + max_angular_momentum, + _improper=improper[index : index + 1], + ) + for index, matrix in enumerate(R.unbind(0)) + ] + + +def _validate_system_ids( + systems: list[System], + transformations: list[O3Transformation], + system_ids: list[int] | torch.Tensor | None, + *, + expected_device: torch.device | None, +) -> torch.Tensor | None: + """Check and normalize the ``system_ids`` argument of ``transform_tensor``. + + ``system_ids[i]`` is the value in a block's ``"system"`` sample column that + selects ``transformations[i]``. This checks that systems and transformations + pair up one-to-one and that there is one distinct integer id per system, + returning the ids as a ``torch.long`` tensor, or ``None`` when + ``system_ids`` is ``None`` and the labels index the transformations + directly. + """ + n_systems = len(systems) + n_transformations = len(transformations) + if n_systems != n_transformations: + raise ValueError( + "Expected one transformation per system, but got " + f"len(systems)={n_systems} and " + f"len(transformations)={n_transformations}." + ) + + if system_ids is None: + return None + + if isinstance(system_ids, torch.Tensor): + if system_ids.ndim != 1: + raise ValueError( + "system_ids must be one-dimensional, but got a tensor with shape " + f"{tuple(system_ids.shape)}." + ) + if system_ids.dtype not in _INTEGER_DTYPES: + raise ValueError( + "system_ids must contain integers, but got a tensor with dtype " + f"{system_ids.dtype}." + ) + if expected_device is not None and system_ids.device != expected_device: + raise ValueError( + f"system_ids are on device {system_ids.device}, but the values to " + f"transform are on device {expected_device}." + ) + validated_ids = system_ids.to(dtype=torch.long) + else: + python_ids: list[int] = [] + for system_id in system_ids: + if isinstance(system_id, bool) or not isinstance(system_id, Integral): + raise ValueError("system_ids must contain integers.") + python_ids.append(int(system_id)) + validated_ids = torch.tensor( + python_ids, + dtype=torch.long, + device=expected_device, + ) + + if len(validated_ids) != n_systems: + raise ValueError( + "system_ids must contain exactly one entry per system, but got " + f"len(system_ids)={len(validated_ids)} and len(systems)={n_systems}." + ) + if torch.unique(validated_ids).numel() != n_systems: + raise ValueError( + "system_ids must contain one distinct entry per system, but got " + f"{validated_ids.tolist()}." + ) + + return validated_ids + + +def _validate_transformations_dtype_device( + transformations: list[O3Transformation], + *, + expected_dtype: torch.dtype, + expected_device: torch.device, +) -> None: + """Check that every transformation has the expected dtype and device.""" + for index, transformation in enumerate(transformations): + if ( + transformation.dtype != expected_dtype + or transformation.device != expected_device + ): + raise ValueError( + f"Transformation at index {index} has dtype/device " + f"({transformation.dtype}, {transformation.device}), differing from " + f"the values to transform ({expected_dtype}, {expected_device})." + ) + + +def _combine_transformations( + transformations: list[O3Transformation], + max_o3_lambda: int, +) -> O3Transformation: + """Concatenate per-system transformations into one batch. + + Each entry must hold a single operation. The combined batch carries + Wigner-D stacks through ``max_o3_lambda``; entries whose + ``max_angular_momentum`` cannot cover it raise the usual range error. + """ + for index, transformation in enumerate(transformations): + if transformation.matrices.size(0) != 1: + raise ValueError( + f"transformations[{index}] holds " + f"{transformation.matrices.size(0)} operations; pass one " + "single-operation O3Transformation per system" + ) + + if len(transformations) == 1: + return transformations[0] + + matrices = torch.cat( + [transformation.matrices for transformation in transformations], + dim=0, + ) + improper = torch.cat( + [transformation.improper for transformation in transformations], + dim=0, + ) + wigner_D: Optional[list[torch.Tensor]] = None + if max_o3_lambda >= 0: + wigner_D = [ + torch.cat( + [ + transformation.wigner_D_matrices(ell) + for transformation in transformations + ], + dim=0, + ) + for ell in range(max_o3_lambda + 1) + ] + return O3Transformation( + matrices, + max(max_o3_lambda, 0), + _improper=improper, + _wigner_D=wigner_D, + ) + + +def transform_system(system: System, transformation: O3Transformation) -> System: + """Apply an O(3) transformation to a single System. + + Positions, cell vectors, neighbor-list displacements, and custom data following + :ref:`o3-conventions` are transformed. Atomic types and periodic-boundary flags + are preserved. + + :param system: input system + :param transformation: single-operation O(3) transformation to apply, matching + ``system.positions`` in dtype and device + :return: new System with transformed geometry + """ + if transformation.matrices.size(0) != 1: + raise ValueError( + "transform_system expects a single operation; use " + "O3Transformation.transform_systems for batches" + ) + 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})." + ) + + return transformation.transform_systems(system)[0] + + +def _component_axis_suffix(axis_name: str, prefix: str) -> tuple[bool, str]: + """Match a component-axis name and return its supported suffix.""" + suffixes = ["", "_1", "_2", "_3", "_4", "_5", "_6", "_7", "_8", "_9"] + for suffix in suffixes: + if axis_name == prefix + suffix: + return True, suffix + return False, "" + + +def _validate_component_axis_metadata( + components: list[Labels], + key: LabelsEntry, +) -> list[tuple[bool, int, int]]: + """Validate component axes and return ``(is_spherical, ell, sigma)`` metadata.""" + if len(components) > 10: + raise ValueError( + f"can not transform a tensor with {len(components)} component axes; " + "at most 10 are supported" + ) + + metadata: list[tuple[bool, int, int]] = [] + for component in components: + axis_name = component.names[0] + is_cartesian, _ = _component_axis_suffix(axis_name, "xyz") + is_spherical, suffix = _component_axis_suffix(axis_name, "o3_mu") + if is_cartesian: + expected_labels = torch.arange( + 3, + device=component.values.device, + dtype=component.values.dtype, + ) + if not torch.equal(component.values[:, 0], expected_labels): + raise ValueError( + f"Cartesian component axis '{axis_name}' must use labels " + "[0, 1, 2] in x, y, z order." + ) + metadata.append((False, 0, 1)) + elif is_spherical: + ell = _validate_nonnegative_integer( + "ell", + int(key["o3_lambda" + suffix]), + ) + sigma = int(key["o3_sigma" + suffix]) + _spherical_parity_factor(ell, sigma, is_improper=False) + + expected_labels = torch.arange( + -ell, + ell + 1, + device=component.values.device, + dtype=component.values.dtype, + ) + if not torch.equal(component.values[:, 0], expected_labels): + raise ValueError( + f"Spherical component axis '{axis_name}' for ell={ell} must use " + f"labels from {-ell} through {ell} in ascending order." + ) + metadata.append((True, 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 metadata + + +def _max_o3_lambda_in_block(key: LabelsEntry, block: TensorBlock) -> int: + """Return the largest angular momentum in one block's values or gradients. + + A block containing only scalar or Cartesian component axes returns ``-1``. + """ + max_o3_lambda = -1 + metadata = _validate_component_axis_metadata(block.components, key) + for is_spherical, ell, _sigma in metadata: + if is_spherical and ell > max_o3_lambda: + max_o3_lambda = ell + + for _gradient_name, gradient in block.gradients(): + gradient_metadata = _validate_component_axis_metadata( + gradient.components, + key, + ) + for is_spherical, ell, _sigma in gradient_metadata: + if is_spherical and ell > max_o3_lambda: + max_o3_lambda = ell + + return max_o3_lambda + + +def _max_o3_lambda_in_tensor(tensor: TensorMap) -> int: + """Return the largest angular momentum in block values or attached gradients. + + A TensorMap containing only scalar or Cartesian component axes returns ``-1``. + """ + max_o3_lambda = -1 + for key, block in tensor.items(): + block_max = _max_o3_lambda_in_block(key, block) + if block_max > max_o3_lambda: + max_o3_lambda = block_max + + return max_o3_lambda + + +def transform_block( + key: LabelsEntry, + block: TensorBlock, + systems: list[System], + transformations: list[O3Transformation], + system_ids: list[int] | torch.Tensor | None = None, +) -> TensorBlock: + """Apply per-system O(3) transformations to a block and its gradients. + + With one system, the ``"system"`` sample label is optional and ignored, as in + :py:func:`transform_tensor`. + + :param key: parent block key, supplying the O(3) labels required by spherical + component axes + :param block: block to transform + :param systems: systems corresponding positionally to ``transformations`` + :param transformations: one single-operation O(3) transformation per system, + matching ``block.values`` in dtype and device + :param system_ids: one distinct integer ``"system"`` sample label per system; + entry ``i`` is paired with ``transformations[i]``. A tensor argument must + be one-dimensional and use the same device as ``block.values``. Defaults + to ``range(len(systems))`` + :return: block with transformed values and gradients and unchanged labels; when + ``systems`` is empty, the block is unchanged + """ + validated_ids = _validate_system_ids( + systems, + transformations, + system_ids, + expected_device=block.values.device, + ) + if len(systems) == 0: + return block + + _validate_transformations_dtype_device( + transformations, + expected_dtype=block.values.dtype, + expected_device=block.values.device, + ) + + block_max_o3_lambda = _max_o3_lambda_in_block(key, block) + combined = _combine_transformations(transformations, block_max_o3_lambda) + wigner_matrices: list[torch.Tensor] = [] + if block_max_o3_lambda >= 0: + wigner_matrices = combined._wigner_D_matrices() + return _transform_block_batched( + key, + block, + combined.matrices, + wigner_matrices, + combined.improper, + validated_ids, + ) + + +def transform_tensor( + tensor: TensorMap, + systems: list[System], + transformations: list[O3Transformation], + system_ids: list[int] | torch.Tensor | None = None, +) -> TensorMap: + """Apply per-system O(3) transformations to a TensorMap and its gradients. + + Scalar, Cartesian, and spherical data are identified by their component-axis + names, following :ref:`o3-conventions`; one :py:class:`TensorMap` may contain + all three kinds of data. At most ten component axes are supported in one + value or gradient block. + + With multiple systems, the ``"system"`` sample label assigns each value sample + to a transformation: samples labelled ``system_ids[i]`` use + ``transformations[i]``. A block may contain samples for only some of the + systems, but every ``"system"`` label present in the block must appear in + ``system_ids``. A gradient sample uses the same transformation as the parent + value sample referenced by its ``"sample"`` label. With one system, the + ``"system"`` label is optional and ignored. + + :param tensor: TensorMap to transform + :param systems: systems corresponding positionally to ``transformations`` + :param transformations: one single-operation O(3) transformation per system, + matching the tensor values in dtype and device when present + :param system_ids: one distinct integer ``"system"`` sample label per system; + entry ``i`` is paired with ``transformations[i]``. A tensor argument must + be one-dimensional and use the same device as the tensor values. Defaults + to ``range(len(systems))`` + :return: transformed TensorMap with the same keys and global information; when + ``systems`` is empty, the tensor is unchanged + """ + if len(tensor) != 0: + system_ids_device = tensor.block(0).values.device + elif len(transformations) != 0: + system_ids_device = transformations[0].device + else: + system_ids_device = None + + validated_ids = _validate_system_ids( + systems, + transformations, + system_ids, + expected_device=system_ids_device, + ) + if len(systems) == 0: + return tensor + + if len(tensor) != 0: + values = tensor.block(0).values + _validate_transformations_dtype_device( + transformations, + expected_dtype=values.dtype, + expected_device=values.device, + ) + + combined = _combine_transformations( + transformations, + _max_o3_lambda_in_tensor(tensor), + ) + return combined.transform_tensormap(tensor, validated_ids) + + +def _transformation_local_indices( + samples: Labels, + n_transformations: int, + system_ids: Optional[torch.Tensor], +) -> torch.Tensor: + """Map sample rows to local indices into the transformation batch. + + With ``system_ids``, rows are matched to operations by their ``"system"`` + label; without, the labels are used as batch indices directly. + """ + if n_transformations <= 0: + raise ValueError("n_transformations must be positive") + if n_transformations == 1: + return torch.zeros( + len(samples), + dtype=torch.long, + device=samples.device, + ) + if "system" not in samples.names: + raise ValueError("multiple transformations require a 'system' sample dimension") + + labels = samples.column("system").to(dtype=torch.long) + if system_ids is None: + if bool(torch.any((labels < 0) | (labels >= n_transformations)).item()): + raise ValueError("sample system indices exceed the transformation batch") + return labels + + sorted_ids, sort_order = torch.sort(system_ids) + positions = torch.searchsorted(sorted_ids, labels) + positions = torch.clamp(positions, min=0, max=int(sorted_ids.numel()) - 1) + matched = sorted_ids.index_select(0, positions) == labels + if not bool(torch.all(matched).item()): + if torch.jit.is_scripting(): + raise ValueError( + "block samples contain system labels that are not in system_ids" + ) + else: + unknown_labels = torch.unique(labels[~matched]) + raise ValueError( + f"Block samples contain system labels {unknown_labels.tolist()} " + f"that are not in system_ids={system_ids.tolist()}. Every sample " + "must be assigned to a system in the transformation." + ) + return sort_order.index_select(0, positions) + + +def _transform_component_values_batched( + values: torch.Tensor, + components: list[Labels], + key: LabelsEntry, + local_indices: torch.Tensor, + matrices: torch.Tensor, + wigner_matrices: list[torch.Tensor], + improper: torch.Tensor, +) -> torch.Tensor: + """Transform the component axes of one values tensor. + + ``local_indices[i]`` selects the operation applied to row ``i``. A batch + of one operation skips the per-row matrix gather entirely. + """ + metadata = _validate_component_axis_metadata(components, key) + if len(metadata) == 0: + return values.clone() + + n_transformations = matrices.size(0) + transformed = values + spherical_parity = 1 + for component_index, (is_spherical, ell, sigma) in enumerate(metadata): + if is_spherical: + if ell >= len(wigner_matrices): + raise ValueError( + f"ell={ell} exceeds " + f"max_angular_momentum={len(wigner_matrices) - 1}." + ) + axis_matrices = wigner_matrices[ell] + # the factor acquired by improper operations; applied per row below + spherical_parity *= _spherical_parity_factor(ell, sigma, True) + else: + axis_matrices = matrices + + component_axis = component_index + 1 + moved = torch.movedim(transformed, component_axis, -1) + moved_shape = moved.shape + flattened = moved.flatten(start_dim=1, end_dim=-2) + if n_transformations == 1: + transformed = flattened @ axis_matrices[0].transpose(0, 1) + else: + matrices_for_rows = axis_matrices.index_select(0, local_indices) + transformed = torch.bmm( + flattened, + matrices_for_rows.transpose(1, 2), + ) + transformed = transformed.reshape(moved_shape) + transformed = torch.movedim(transformed, -1, component_axis) + + if spherical_parity != 1 and bool(torch.any(improper).item()): + if n_transformations == 1: + transformed = transformed * float(spherical_parity) + else: + factors = torch.where( + improper.index_select(0, local_indices), + torch.tensor( + float(spherical_parity), + dtype=values.dtype, + device=values.device, + ), + torch.tensor(1.0, dtype=values.dtype, device=values.device), + ) + factors_shape: list[int] = [-1] + for _axis in range(values.dim() - 1): + factors_shape.append(1) + transformed = transformed * factors.view(factors_shape) + return transformed + + +def _transform_block_batched( + key: LabelsEntry, + block: TensorBlock, + matrices: torch.Tensor, + wigner_matrices: list[torch.Tensor], + improper: torch.Tensor, + system_ids: Optional[torch.Tensor], +) -> TensorBlock: + """Transform one block and its gradients with a batch of operations.""" + value_indices = _transformation_local_indices( + block.samples, + matrices.size(0), + system_ids, + ) + new_block = TensorBlock( + values=_transform_component_values_batched( + block.values, + block.components, + key, + value_indices, + matrices, + wigner_matrices, + improper, + ), + samples=block.samples, + components=block.components, + properties=block.properties, + ) + for gradient_name, gradient in block.gradients(): + parent_rows = gradient.samples.column("sample").to(dtype=torch.long) + gradient_indices = value_indices.index_select(0, parent_rows) + new_block.add_gradient( + gradient_name, + TensorBlock( + values=_transform_component_values_batched( + gradient.values, + gradient.components, + key, + gradient_indices, + matrices, + wigner_matrices, + improper, + ), + samples=gradient.samples, + components=gradient.components, + properties=gradient.properties, + ), + ) + return new_block + + +def _transform_tensormap_batched( + tensor: TensorMap, + matrices: torch.Tensor, + wigner_matrices: list[torch.Tensor], + improper: torch.Tensor, + system_ids: Optional[torch.Tensor], +) -> TensorMap: + """Transform a TensorMap with a batch of O(3) operations. + + This is the single implementation behind every tensor-transformation entry + point: ``matrices[i]`` is the Cartesian operation of local index ``i``, + ``wigner_matrices[ell][i]`` the Wigner-D matrix of its proper rotational + part, and ``improper[i]`` whether it includes the inversion. Row routing + follows :py:meth:`O3Transformation.transform_tensormap`. + """ + if ( + matrices.dim() != 3 + or matrices.size(0) == 0 + or matrices.size(1) != 3 + or matrices.size(2) != 3 + ): + raise ValueError("matrices must have shape (N, 3, 3) with N > 0") + if matrices.dtype != torch.float32 and matrices.dtype != torch.float64: + raise TypeError("matrices must use float32 or float64") + if len(tensor) != 0: + reference_values = tensor.block(0).values + if ( + matrices.dtype != reference_values.dtype + or matrices.device != reference_values.device + ): + raise ValueError("tensor and matrices must have the same dtype and device") + + blocks: list[TensorBlock] = [] + for key, block in tensor.items(): + blocks.append( + _transform_block_batched( + key, + block, + matrices, + wigner_matrices, + improper, + system_ids, + ) + ) + + transformed = TensorMap(tensor.keys, blocks) + for info_name, info_value in tensor.info().items(): + transformed.set_info(info_name, info_value) + return transformed diff --git a/python/metatomic_torch/metatomic/torch/o3/_utils.py b/python/metatomic_torch/metatomic/torch/o3/_utils.py new file mode 100644 index 00000000..646ba840 --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/o3/_utils.py @@ -0,0 +1,146 @@ +""" +Shared helpers for the O(3)-symmetrized model machinery: argument validation, and the +sample-label bookkeeping that maps one input system onto its rotated copies and back +again once the outputs have been reduced over those copies. +""" + +from numbers import Integral +from typing import List, Optional, Tuple + +import torch +from metatensor.torch import Labels, TensorBlock + + +def validate_integer(name: str, value, minimum: int) -> int: + """Check that ``value`` is an integer at least ``minimum``. + + Return it as a Python ``int``. + """ + if isinstance(value, bool) or not isinstance(value, Integral): + raise TypeError(f"{name} must be an integer, got {type(value).__name__}") + integer_value = int(value) + if integer_value < minimum: + if minimum == 0: + qualifier = "non-negative" + elif minimum == 1: + qualifier = "positive" + else: + qualifier = f"larger or equal to {minimum}" + raise ValueError(f"{name} must be {qualifier}, got {integer_value}") + return integer_value + + +def map_selected_atoms_to_rotated_copies( + selected_atoms: Optional[Labels], + input_system_index: int, + n_rotated_copies: int, +) -> Optional[Labels]: + """Map one input system's selected atoms to each rotated copy.""" + if selected_atoms is None: + return None + + input_system_mask = ( + selected_atoms.column("system").to(dtype=torch.long) == input_system_index + ) + selected_atoms_for_input_system = selected_atoms.values[input_system_mask] + if selected_atoms_for_input_system.shape[0] == 0: + return Labels( + list(selected_atoms.names), + selected_atoms.values.new_empty((0, len(selected_atoms.names))), + ) + + rotated_values = selected_atoms_for_input_system.repeat((n_rotated_copies, 1)) + rotated_values[:, list(selected_atoms.names).index("system")] = torch.arange( + n_rotated_copies, + dtype=rotated_values.dtype, + device=rotated_values.device, + ).repeat_interleave(len(selected_atoms_for_input_system)) + return Labels(list(selected_atoms.names), rotated_values) + + +def group_samples_by_rotated_copy( + block: TensorBlock, n_rotated_copies: int +) -> Tuple[torch.Tensor, List[str], torch.Tensor]: + """Group samples from rotated copies along a leading copy axis.""" + sample_names = list(block.samples.names) + system_column = sample_names.index("system") + copy_indices = block.samples.column("system").to(dtype=torch.long) + sample_values_without_system = torch.cat( + [ + block.samples.values[:, :system_column], + block.samples.values[:, system_column + 1 :], + ], + dim=1, + ) + if len(copy_indices) != 0 and bool( + torch.any((copy_indices < 0) | (copy_indices >= n_rotated_copies)).item() + ): + raise ValueError( + "encountered output samples with out-of-range rotated-copy indices: " + f"the system column spans [{int(copy_indices.min())}, " + f"{int(copy_indices.max())}], expected [0, {n_rotated_copies - 1}]" + ) + + if len(copy_indices) % n_rotated_copies != 0: + raise ValueError( + "SymmetrizedModel expects every rotated copy to produce the same " + "sample labels in the same order." + ) + n_samples_per_copy = len(copy_indices) // n_rotated_copies + order = torch.argsort(copy_indices, stable=True) + expected_copy_indices = torch.arange( + n_rotated_copies, + dtype=copy_indices.dtype, + device=copy_indices.device, + ).repeat_interleave(n_samples_per_copy) + if not torch.equal(copy_indices[order], expected_copy_indices): + raise ValueError( + "SymmetrizedModel expects every rotated copy to produce the same " + "sample labels in the same order." + ) + + values_shape = [n_rotated_copies, n_samples_per_copy] + for axis in range(1, block.values.dim()): + values_shape.append(block.values.shape[axis]) + values_by_copy = block.values[order].reshape(values_shape) + sample_values_by_copy = sample_values_without_system[order].reshape( + n_rotated_copies, + n_samples_per_copy, + sample_values_without_system.shape[1], + ) + shared_sample_values = sample_values_by_copy[0] + if not torch.equal( + sample_values_by_copy, + shared_sample_values.unsqueeze(0).expand_as(sample_values_by_copy), + ): + raise ValueError( + "SymmetrizedModel expects every rotated copy to produce the same " + "sample labels in the same order." + ) + + return ( + values_by_copy, + sample_names[:system_column] + sample_names[system_column + 1 :], + shared_sample_values, + ) + + +def restore_input_system_to_samples( + sample_names: List[str], + sample_values: torch.Tensor, + input_system_index: int, + *, + device: torch.device, +) -> Labels: + """Restore the input-system label after reducing over rotated copies.""" + sample_values = sample_values.to(device=device) + system_values = torch.full( + (sample_values.shape[0], 1), + input_system_index, + dtype=sample_values.dtype, + device=device, + ) + return Labels( + ["system"] + sample_names, + torch.cat([system_values, sample_values], dim=1), + ) diff --git a/python/metatomic_torch/metatomic/torch/o3/_wigner.py b/python/metatomic_torch/metatomic/torch/o3/_wigner.py index 8631a71d..7052adcd 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_wigner.py +++ b/python/metatomic_torch/metatomic/torch/o3/_wigner.py @@ -99,12 +99,12 @@ def _rotation_to_angles( def build_wigner_D_cache( - o3_lambda_max: int, + max_angular_momentum: int, matrix: torch.Tensor, device: torch.device, dtype: torch.dtype, ) -> dict[int, torch.Tensor]: - """Return real Wigner-D matrices for ``ell = 0, ..., o3_lambda_max``. + """Return real Wigner-D matrices for ``ell = 0, ..., max_angular_momentum``. If ``matrix`` has negative determinant, ``-matrix`` is a proper rotation. Build the Wigner-D matrices for this proper rotation; the caller restores @@ -114,8 +114,82 @@ def build_wigner_D_cache( angles = _rotation_to_angles(matrix) complex_to_real = { ell: _complex_to_real_spherical_harmonics_transform(ell) - for ell in range(o3_lambda_max + 1) + for ell in range(max_angular_momentum + 1) } - cache = _compute_real_wigner_d_matrices(o3_lambda_max, angles, complex_to_real) + cache = _compute_real_wigner_d_matrices( + max_angular_momentum, angles, complex_to_real + ) return {ell: tensor.to(device=device, dtype=dtype) for ell, tensor in cache.items()} + + +def build_packed_wigner_matrices( + matrices: torch.Tensor, + max_angular_momentum: int, +) -> torch.Tensor: + """Build and pack proper Wigner-D matrices through ``max_angular_momentum``. + + :param matrices: ``(n_matrices, 3, 3)`` stack of O(3) matrices + :param max_angular_momentum: maximum angular momentum to include + :return: flat tensor holding every Wigner-D matrix, laid out for + :py:func:`wigner_matrices_for_lambda`, with the dtype and device of + ``matrices`` + """ + output_device = matrices.device + output_dtype = matrices.dtype + calculation_matrices = matrices.detach().to(device="cpu") + cpu = torch.device("cpu") + n_matrices = matrices.size(0) + n_elements_per_matrix = ( + (max_angular_momentum + 1) + * (2 * max_angular_momentum + 1) + * (2 * max_angular_momentum + 3) + // 3 + ) + packed = torch.empty( + n_matrices * n_elements_per_matrix, + dtype=output_dtype, + device="cpu", + ) + + for matrix_index, matrix in enumerate(calculation_matrices.unbind(0)): + cache = build_wigner_D_cache( + max_angular_momentum, + matrix, + device=cpu, + dtype=output_dtype, + ) + for o3_lambda in range(max_angular_momentum + 1): + dimension = 2 * o3_lambda + 1 + elements_before = o3_lambda * (4 * o3_lambda * o3_lambda - 1) // 3 + offset = n_matrices * elements_before + matrix_index * dimension * dimension + packed[offset : offset + dimension * dimension].copy_( + cache[o3_lambda].reshape(-1) + ) + + return packed.to( + device=output_device, + dtype=output_dtype, + ) + + +def wigner_matrices_for_lambda( + packed: torch.Tensor, + n_matrices: int, + o3_lambda: int, +) -> torch.Tensor: + """Return the packed Wigner-D stack for one ``o3_lambda`` as a view.""" + # the packed layout is rank-major then matrix-major: all matrices for + # o3_lambda=0 come first, then all matrices for o3_lambda=1, and so on + dimension = 2 * o3_lambda + 1 + elements_before = o3_lambda * (4 * o3_lambda * o3_lambda - 1) // 3 + offset = n_matrices * elements_before + length = n_matrices * dimension * dimension + if offset + length > packed.numel(): + raise ValueError("o3_lambda exceeds the packed Wigner-D storage") + + return packed[offset : offset + length].view( + n_matrices, + dimension, + dimension, + ) diff --git a/python/metatomic_torch/tests/o3.py b/python/metatomic_torch/tests/o3.py index 2009e946..e66a1bd5 100644 --- a/python/metatomic_torch/tests/o3.py +++ b/python/metatomic_torch/tests/o3.py @@ -21,9 +21,9 @@ # These private helpers back exported symmetrized-model wrappers and have no public # entry point yet; their tests below compare them against the public transform_tensor. -from metatomic.torch.o3._tranformations import ( +from metatomic.torch.o3._transformations import ( _max_o3_lambda_in_tensor, - _transform_tensor_with_precomputed_matrices, + _transform_tensormap_batched, ) # The complex-to-real spherical harmonics conversion is defined only here for now. @@ -110,17 +110,23 @@ def _single_block_tensor_map( def _stack_o3_matrices(transformations, max_angular_momentum): - """Stack Cartesian and Wigner matrices from O3 transformations.""" - matrices = torch.stack( - [transformation.matrix for transformation in transformations] + """Stack Cartesian, Wigner, and parity tensors from O3 transformations.""" + matrices = torch.cat( + [transformation.matrices for transformation in transformations] ) wigner_matrices = [ - torch.stack( - [transformation.wigner_D_matrix(ell) for transformation in transformations] + torch.cat( + [ + transformation.wigner_D_matrices(ell) + for transformation in transformations + ] ) for ell in range(max_angular_momentum + 1) ] - return matrices, wigner_matrices + improper = torch.cat( + [transformation.improper for transformation in transformations] + ) + return matrices, wigner_matrices, improper def test_max_o3_lambda_in_tensor(): @@ -344,8 +350,10 @@ def test_transformation_validation(): with pytest.raises(ValueError, match=f"^{message}$"): random_transformations(0, device=torch.device("cpu"), dtype=torch.float16) - # matrices must be (3, 3) and orthogonal - message = re.escape("Transformation has shape (2, 2); expected (3, 3).") + # matrices must be (3, 3) or (N, 3, 3) and orthogonal + message = re.escape( + "Transformation has shape (2, 2); expected (3, 3) or (N, 3, 3) with N > 0." + ) with pytest.raises(ValueError, match=f"^{message}$"): O3Transformation(torch.eye(2, dtype=torch.float64), max_angular_momentum=0) matrix = torch.eye(3, dtype=torch.float64) @@ -999,10 +1007,7 @@ def test_system_ids_validation(): samples=Labels(["atom"], torch.tensor([[0]])), components=[], ) - message = re.escape( - "Rotational augmentation expects output samples to include a 'system' " - "dimension when transforming multiple systems." - ) + message = re.escape("multiple transformations require a 'system' sample dimension") with pytest.raises(ValueError, match=f"^{message}$"): transform_tensor(no_system_column, systems, transformations) @@ -1209,16 +1214,16 @@ def test_pair_samples_routing(): @pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE) -@pytest.mark.parametrize("is_improper", [False, True]) -def test_precomputed_tensor_transform_matches_transform_tensor( +@pytest.mark.parametrize("parities", [(1.0, 1.0), (-1.0, -1.0), (1.0, -1.0)]) +def test_batched_tensor_transform_matches_transform_tensor( device, dtype, - is_improper, + parities, ): - """The scripted precomputed-matrices path matches ``transform_tensor``.""" + """The scripted batched kernel matches ``transform_tensor``, including for + batches mixing proper and improper operations.""" dtype = getattr(torch, dtype) atol = 1.0e-5 if dtype == torch.float32 else 1.0e-12 - sign = -1.0 if is_improper else 1.0 proper_matrices = [ _rotation_90_degrees_around_z().to(device=device, dtype=dtype), torch.tensor( @@ -1229,9 +1234,9 @@ def test_precomputed_tensor_transform_matches_transform_tensor( ] transformations = [ O3Transformation(sign * matrix, max_angular_momentum=2) - for matrix in proper_matrices + for sign, matrix in zip(parities, proper_matrices, strict=True) ] - matrices, wigner_matrices = _stack_o3_matrices( + matrices, wigner_matrices, improper = _stack_o3_matrices( transformations, max_angular_momentum=2, ) @@ -1300,12 +1305,13 @@ def test_precomputed_tensor_transform_matches_transform_tensor( ], transformations, ) - scripted_transform = torch.jit.script(_transform_tensor_with_precomputed_matrices) + scripted_transform = torch.jit.script(_transform_tensormap_batched) result = scripted_transform( tensor, matrices, wigner_matrices, - is_improper, + improper, + None, ) mts.allclose_raise(result, expected, rtol=0.0, atol=atol) @@ -1330,13 +1336,13 @@ def test_precomputed_tensor_transform_matches_transform_tensor( ) -def test_precomputed_tensor_transform_single_transformation_is_scriptable(): +def test_batched_tensor_transform_single_transformation_is_scriptable(): """The scripted singleton path should not require a ``system`` sample label.""" transformation = O3Transformation( _rotation_90_degrees_around_z(), max_angular_momentum=1, ) - matrices, wigner_matrices = _stack_o3_matrices( + matrices, wigner_matrices, improper = _stack_o3_matrices( [transformation], max_angular_momentum=1, ) @@ -1355,12 +1361,13 @@ def test_precomputed_tensor_transform_single_transformation_is_scriptable(): ], ) - scripted_transform = torch.jit.script(_transform_tensor_with_precomputed_matrices) + scripted_transform = torch.jit.script(_transform_tensormap_batched) result = scripted_transform( tensor, matrices, wigner_matrices, - False, + improper, + None, ) expected = transform_tensor( tensor, @@ -1376,13 +1383,13 @@ def test_precomputed_tensor_transform_single_transformation_is_scriptable(): ) -def test_precomputed_tensor_transform_rejects_invalid_routing_and_wigner_rank(): +def test_batched_tensor_transform_rejects_invalid_routing_and_wigner_rank(): """Ambiguous routing or missing Wigner-D ranks fail instead of misrotating.""" transformations = [ O3Transformation(torch.eye(3, dtype=torch.float64), 1), O3Transformation(_rotation_90_degrees_around_z(), 1), ] - matrices, wigner_matrices = _stack_o3_matrices( + matrices, wigner_matrices, improper = _stack_o3_matrices( transformations, max_angular_momentum=1, ) @@ -1394,11 +1401,12 @@ def test_precomputed_tensor_transform_rejects_invalid_routing_and_wigner_rank(): ) message = re.escape("multiple transformations require a 'system' sample dimension") with pytest.raises(ValueError, match=f"^{message}$"): - _transform_tensor_with_precomputed_matrices( + _transform_tensormap_batched( missing_system, matrices, wigner_matrices, - False, + improper, + None, ) for system_index in (-1, 2): @@ -1409,11 +1417,12 @@ def test_precomputed_tensor_transform_rejects_invalid_routing_and_wigner_rank(): ) message = re.escape("sample system indices exceed the transformation batch") with pytest.raises(ValueError, match=f"^{message}$"): - _transform_tensor_with_precomputed_matrices( + _transform_tensormap_batched( out_of_range, matrices, wigner_matrices, - False, + improper, + None, ) unavailable_rank = _single_block_tensor_map( @@ -1427,11 +1436,194 @@ def test_precomputed_tensor_transform_rejects_invalid_routing_and_wigner_rank(): Labels("o3_mu", torch.arange(-1, 2).reshape(-1, 1)), ], ) - message = re.escape("spherical rank exceeds the Wigner-D storage") + message = re.escape("ell=1 exceeds max_angular_momentum=0.") with pytest.raises(ValueError, match=f"^{message}$"): - _transform_tensor_with_precomputed_matrices( + _transform_tensormap_batched( unavailable_rank, matrices, wigner_matrices[:1], - False, + improper, + None, + ) + + +def test_batched_transformation_matches_singles(): + """One batched O3Transformation behaves like its per-operation singles.""" + singles = random_transformations( + 4, + max_angular_momentum=2, + device=torch.device("cpu"), + dtype=torch.float64, + include_inversions=True, + generator=torch.Generator().manual_seed(20260805), + ) + batch = O3Transformation( + torch.cat([single.matrices for single in singles]), + max_angular_momentum=2, + ) + + assert batch.matrices.shape == (4, 3, 3) + determinant_signs = torch.stack( + [torch.det(single.matrix) < 0 for single in singles] + ) + assert torch.equal(batch.improper, determinant_signs) + + # Cartesian and spherical actions gain a leading batch axis + vectors = torch.randn(5, 3, dtype=torch.float64) + cartesian = batch.transform_cartesian(vectors) + spherical = batch.transform_spherical(vectors, ell=1, sigma=-1) + for index, single in enumerate(singles): + assert torch.allclose( + cartesian[index], + single.transform_cartesian(vectors), + atol=1e-12, + ) + assert torch.allclose( + spherical[index], + single.transform_spherical(vectors, ell=1, sigma=-1), + atol=1e-12, + ) + + # a batched System transformation matches transform_system per operation + system = _make_system( + [1, 8], + positions=torch.randn(2, 3, dtype=torch.float64), + cell=torch.eye(3, dtype=torch.float64), + pbc=torch.tensor([True, True, True]), + ) + batch_systems = batch.transform_systems(system) + assert len(batch_systems) == 4 + for index, single in enumerate(singles): + expected_system = transform_system(system, single) + assert torch.allclose( + batch_systems[index].positions, + expected_system.positions, + atol=1e-12, + ) + assert torch.allclose( + batch_systems[index].cell, + expected_system.cell, + atol=1e-12, + ) + + +def test_inverse_and_with_inversion_views(): + """``inverse`` and ``with_inversion`` return views composing correctly.""" + matrix = torch.tensor(_axis_angle([1.0, 2.0, 3.0], 0.7), dtype=torch.float64) + transformation = O3Transformation(matrix, max_angular_momentum=2) + + inverse = transformation.inverse() + assert torch.allclose( + inverse.matrix @ transformation.matrix, + torch.eye(3, dtype=torch.float64), + atol=1e-12, + ) + for ell in range(3): + assert torch.allclose( + inverse.wigner_D_matrix(ell), + transformation.wigner_D_matrix(ell).T, + atol=1e-12, ) + + flipped = transformation.with_inversion() + assert flipped.is_improper + assert torch.equal(flipped.matrices, -transformation.matrices) + # the proper part -- and with it the Wigner-D matrices -- is unchanged + for ell in range(3): + assert torch.equal( + flipped.wigner_D_matrix(ell), + transformation.wigner_D_matrix(ell), + ) + + # (-R)^-1 = -R^T, still improper + inverse_flipped = flipped.inverse() + assert inverse_flipped.is_improper + assert torch.allclose( + inverse_flipped.matrix, + -matrix.T, + atol=1e-12, + ) + + # values round-trip through a transformation and its inverse + values = torch.randn(4, 5, dtype=torch.float64) + roundtrip = inverse.transform_spherical( + transformation.transform_spherical(values, ell=2, sigma=-1), + ell=2, + sigma=-1, + ) + assert torch.allclose(roundtrip, values, atol=1e-12) + + +def test_o3_transformation_scriptable_in_forward(): + """A scripted model can build transformations from buffers inside forward + and still be saved and loaded.""" + + class BackRotate(torch.nn.Module): + wigner_D: list[torch.Tensor] + + def __init__(self, matrices, wigner_D): + super().__init__() + self.register_buffer("matrices", matrices) + self.wigner_D = list(wigner_D) + + def forward(self, tensor: TensorMap) -> TensorMap: + proper = torch.zeros( + self.matrices.size(0), + dtype=torch.bool, + device=self.matrices.device, + ) + transformation = O3Transformation( + self.matrices, + len(self.wigner_D) - 1, + _improper=proper, + _wigner_D=self.wigner_D, + ) + inverse = transformation.with_inversion().inverse() + return inverse.transform_tensormap(tensor) + + singles = random_transformations( + 3, + max_angular_momentum=2, + device=torch.device("cpu"), + dtype=torch.float64, + generator=torch.Generator().manual_seed(3), + ) + matrices, wigner_matrices, _improper = _stack_o3_matrices( + singles, + max_angular_momentum=2, + ) + module = BackRotate(matrices, wigner_matrices) + scripted = torch.jit.script(module) + + values = torch.randn(6, 5, 2, dtype=torch.float64) + tensor = _single_block_tensor_map( + keys=Labels(["o3_lambda", "o3_sigma"], torch.tensor([[2, 1]])), + values=values, + samples=Labels( + ["system", "sample"], + torch.stack([torch.arange(6) % 3, torch.arange(6)], dim=1), + ), + components=[Labels(["o3_mu"], torch.arange(-2, 3).reshape(-1, 1))], + properties=Labels(["p"], torch.tensor([[0], [1]])), + ) + + eager_result = module(tensor) + scripted_result = scripted(tensor) + assert torch.allclose( + eager_result.block().values, + scripted_result.block().values, + atol=1e-12, + ) + + import io + + buffer = io.BytesIO() + torch.jit.save(scripted, buffer) + buffer.seek(0) + loaded = torch.jit.load(buffer) + loaded_result = loaded(tensor) + assert torch.allclose( + eager_result.block().values, + loaded_result.block().values, + atol=1e-12, + ) diff --git a/python/metatomic_torch/tests/symmetrized_model.py b/python/metatomic_torch/tests/symmetrized_model.py new file mode 100644 index 00000000..98605114 --- /dev/null +++ b/python/metatomic_torch/tests/symmetrized_model.py @@ -0,0 +1,2325 @@ +import re +from typing import Dict, List, Optional + +import metatensor.torch as mts +import numpy as np +import pytest +import torch +from metatensor.torch import Labels, TensorBlock, TensorMap + +from metatomic.torch import ( + AtomisticModel, + ModelCapabilities, + ModelEvaluationOptions, + ModelMetadata, + ModelOutput, + NeighborListOptions, + SymmetrizedModel, + System, + load_atomistic_model, +) +from metatomic.torch.o3 import O3Transformation +from metatomic.torch.o3._decompose import ( + _cartesian_vectors_to_spherical, + _matrices_to_spherical, + _o3_mu_labels, + decompose_quantity, +) +from metatomic.torch.o3._quadrature import ( + _rotations_from_euler_angles, + choose_quadrature, + get_euler_angles_quadrature, + get_rotation_quadrature, +) +from metatomic.torch.o3._utils import map_selected_atoms_to_rotated_copies + + +def _make_single_block_tensor_map( + values: torch.Tensor, sample_name: str = "sample" +) -> TensorMap: + """Create a one-block TensorMap test input from ``values``.""" + device = values.device + components = [ + Labels.range(f"component_{axis}", size).to(device=device) + for axis, size in enumerate(values.shape[1:-1]) + ] + return TensorMap( + Labels("_", torch.tensor([[0]], dtype=torch.int64, device=device)), + [ + TensorBlock( + values=values, + samples=Labels.range(sample_name, values.shape[0]).to(device=device), + components=components, + properties=Labels.range("property", values.shape[-1]).to(device=device), + ) + ], + ) + + +def _tensor_map_with_components( + values: torch.Tensor, + component_names, +) -> TensorMap: + """Create a one-block TensorMap with the requested component-axis names.""" + components = [ + Labels.range(name, values.shape[axis + 1]) + for axis, name in enumerate(component_names) + ] + return TensorMap( + Labels("_", torch.tensor([[0]], dtype=torch.int64)), + [ + TensorBlock( + values=values, + samples=Labels.range("system", values.shape[0]), + components=components, + properties=Labels.range("property", values.shape[-1]), + ) + ], + ) + + +class _EmptyModel(torch.nn.Module): + """Provide the model interface without producing any outputs.""" + + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels], + ) -> Dict[str, TensorMap]: + return {} + + +def _forward_test_system( + positions: List[List[float]], + dtype: torch.dtype = torch.float64, + requires_grad: bool = False, +) -> System: + """Create a non-periodic input System with configurable dtype and autograd.""" + position_values = torch.tensor( + positions, + dtype=dtype, + requires_grad=requires_grad, + ) + return System( + types=torch.ones(len(position_values), dtype=torch.int64), + positions=position_values, + cell=torch.zeros((3, 3), dtype=dtype), + pbc=torch.tensor([False, False, False]), + ) + + +def _system_scalar_tensor_map( + values: torch.Tensor, + property_name: str = "property", +) -> TensorMap: + """Package one scalar response for each System in a model call.""" + device = values.device + return TensorMap( + Labels("_", torch.tensor([[0]], dtype=torch.int64, device=device)), + [ + TensorBlock( + values=values, + samples=Labels( + "system", + torch.arange( + len(values), + dtype=torch.int64, + device=device, + ).reshape(-1, 1), + ), + components=[], + properties=Labels( + property_name, + torch.arange( + values.shape[-1], + dtype=torch.int64, + device=device, + ).reshape(-1, 1), + ), + ) + ], + ) + + +class _LinearEnergyModel(torch.nn.Module): + """Return the first atom's x coordinate as every requested scalar output.""" + + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels], + ) -> Dict[str, TensorMap]: + values = torch.stack([system.positions[0, 0] for system in systems]).reshape( + -1, 1 + ) + result = torch.jit.annotate(Dict[str, TensorMap], {}) + for output_name in outputs: + result[output_name] = _system_scalar_tensor_map(values) + return result + + +class _CountingLinearEnergyModel(_LinearEnergyModel): + """Record how often ``forward`` is called and the requests it receives.""" + + def __init__(self): + super().__init__() + self.call_count = 0 + self.requested_names: List[List[str]] = [] + self.requested_units: List[str] = [] + self.requested_sample_kinds: List[str] = [] + self.requested_explicit_gradients: List[List[str]] = [] + self.requested_descriptions: List[str] = [] + + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels], + ) -> Dict[str, TensorMap]: + self.call_count += 1 + self.requested_names.append(list(outputs.keys())) + for output in outputs.values(): + self.requested_units.append(output.unit) + self.requested_sample_kinds.append(output.sample_kind) + self.requested_explicit_gradients.append(list(output.explicit_gradients)) + self.requested_descriptions.append(output.description) + return super().forward(systems, outputs, selected_atoms) + + +class _OffsetLinearEnergyModel(_LinearEnergyModel): + """Add a large invariant offset to the linear response.""" + + def __init__(self, offset: float): + super().__init__() + self._offset = offset + + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels], + ) -> Dict[str, TensorMap]: + values = self._offset + torch.stack( + [system.positions[0, 0] for system in systems] + ).reshape(-1, 1) + result = torch.jit.annotate(Dict[str, TensorMap], {}) + for output_name in outputs: + result[output_name] = _system_scalar_tensor_map(values) + return result + + +class _InconsistentSampleModel(torch.nn.Module): + """Label every returned sample as system 0, whatever the input batch.""" + + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels], + ) -> Dict[str, TensorMap]: + values = torch.stack([system.positions[0, 0] for system in systems]).reshape( + -1, 1 + ) + sample_values = torch.stack( + [ + torch.zeros(len(systems), dtype=torch.int64), + torch.arange(len(systems), dtype=torch.int64), + ], + dim=1, + ) + result = torch.jit.annotate(Dict[str, TensorMap], {}) + for output_name in outputs: + result[output_name] = TensorMap( + Labels("_", torch.tensor([[0]])), + [ + TensorBlock( + values=values, + samples=Labels(["system", "atom"], sample_values), + components=[], + properties=Labels.range("property", 1), + ) + ], + ) + return result + + +class _LinearModelWithRequirements(torch.nn.Module): + """Provide a scalar output while requesting custom data and a neighbor list.""" + + def __init__(self): + super().__init__() + self._neighbor_list = NeighborListOptions( + 2.5, + False, + True, + "linear model", + ) + + def requested_neighbor_lists(self) -> List[NeighborListOptions]: + return [self._neighbor_list] + + def requested_inputs(self) -> Dict[str, ModelOutput]: + return { + "mtt::field": ModelOutput( + unit="eV", + sample_kind="atom", + description="Cartesian field used by the model.", + ) + } + + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels], + ) -> Dict[str, TensorMap]: + values: List[torch.Tensor] = [] + for system in systems: + neighbors = system.get_neighbor_list(self._neighbor_list) + field = system.get_data("mtt::field") + invariant_input = ( + neighbors.values.square().sum() + field.block().values.square().sum() + ) + values.append(system.positions[0, 0] + 0.01 * invariant_input) + scalar_values = torch.stack(values).reshape(-1, 1) + + result = torch.jit.annotate(Dict[str, TensorMap], {}) + for output_name in outputs: + property_name = "energy" if output_name == "energy" else "property" + result[output_name] = _system_scalar_tensor_map( + scalar_values, + property_name, + ) + return result + + +def _system_with_linear_model_requirements( + neighbor_options: NeighborListOptions, + device: torch.device, +) -> System: + """Create the float32 System required by ``_LinearModelWithRequirements``.""" + system = _forward_test_system( + [[1.0, 2.0, 3.0], [1.2, 2.1, 3.1]], + dtype=torch.float32, + ) + system.add_neighbor_list( + neighbor_options, + TensorBlock( + values=(system.positions[1] - system.positions[0]).reshape(1, 3, 1), + samples=Labels( + [ + "first_atom", + "second_atom", + "cell_shift_a", + "cell_shift_b", + "cell_shift_c", + ], + torch.tensor([[0, 1, 0, 0, 0]], dtype=torch.int64), + ), + components=[Labels.range("xyz", 3)], + properties=Labels.range("distance", 1), + ), + ) + field = TensorMap( + Labels( + "_", + torch.tensor([[0]], dtype=torch.int64), + ), + [ + TensorBlock( + values=system.positions.unsqueeze(-1), + samples=Labels.range("atom", len(system)), + components=[Labels.range("xyz", 3)], + properties=Labels.range("field", 1), + ) + ], + ) + field.set_info("unit", "eV") + system.add_data("mtt::field", field) + return system.to(device=device) + + +class _O3PolynomialSectorModel(torch.nn.Module): + """Return one analytic polynomial response in every O(3) sector to lambda=3.""" + + # the polynomials 1, x, x*y, and x*y*z transform purely in lambda=0..3 with + # sigma=+1; multiplying by det(positions) flips the parity to sigma=-1 + + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels], + ) -> Dict[str, TensorMap]: + device = systems[0].positions.device + sectors = [ + (o3_lambda, o3_sigma) for o3_lambda in range(4) for o3_sigma in (1, -1) + ] + values: List[torch.Tensor] = [] + for system in systems: + x, y, z = system.positions[0] + sigma_plus_values = [ + x.new_ones(()), + x, + x * y, + x * y * z, + ] + pseudoscalar = torch.det(system.positions) + system_values: List[torch.Tensor] = [] + for o3_lambda, o3_sigma in sectors: + value = sigma_plus_values[o3_lambda] + if o3_sigma == -1: + value = pseudoscalar * value + system_values.append(value) + values.append(torch.stack(system_values)) + + tensor = TensorMap( + Labels("_", torch.tensor([[0]], dtype=torch.int64, device=device)), + [ + TensorBlock( + values=torch.stack(values), + samples=Labels( + "system", + torch.arange( + len(systems), + dtype=torch.int64, + device=device, + ).reshape(-1, 1), + ), + components=[], + properties=Labels( + ["source_lambda", "source_sigma"], + torch.tensor(sectors, dtype=torch.int64, device=device), + ), + ) + ], + ) + result = torch.jit.annotate(Dict[str, TensorMap], {}) + for output_name in outputs: + result[output_name] = tensor + return result + + +class _AtomFeatureModel(torch.nn.Module): + """Return one component-less per-atom feature that is not O(3) invariant.""" + + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels], + ) -> Dict[str, TensorMap]: + device = systems[0].positions.device + values: List[torch.Tensor] = [] + samples: List[torch.Tensor] = [] + for system_index, system in enumerate(systems): + for atom_index in range(len(system)): + values.append(system.positions[atom_index, 0].reshape(1)) + samples.append( + torch.tensor( + [system_index, atom_index], + dtype=torch.int64, + device=device, + ) + ) + + tensor = TensorMap( + Labels("_", torch.tensor([[0]], dtype=torch.int64, device=device)), + [ + TensorBlock( + torch.stack(values), + Labels(["system", "atom"], torch.stack(samples)), + [], + Labels.range("feature", 1).to(device=device), + ) + ], + ) + result = torch.jit.annotate(Dict[str, TensorMap], {}) + for output_name in outputs: + result[output_name] = tensor + return result + + +class _DegreeSevenEnergyModel(torch.nn.Module): + """Return an odd degree-seven response with a degree-fourteen square.""" + + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels], + ) -> Dict[str, TensorMap]: + energies: List[torch.Tensor] = [] + for system in systems: + x, y, z = system.positions[0] + fourth_order = 0.625 * (x**4 + y**4 + z**4) + mixed = x**2 * y**2 + x**2 * z**2 + y**2 * z**2 + energies.append(1000.0 * x * y * z * (fourth_order - mixed)) + return { + "energy": _system_scalar_tensor_map(torch.stack(energies).reshape(-1, 1)) + } + + +class _EquivariantOutputModel(torch.nn.Module): + """Provide exactly equivariant scalar, Cartesian, and spherical test outputs.""" + + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels], + ) -> Dict[str, TensorMap]: + device = systems[0].positions.device + result: Dict[str, TensorMap] = {} + system_samples = Labels( + "system", + torch.arange(len(systems), dtype=torch.int64, device=device).reshape(-1, 1), + ) + placeholder = Labels( + "_", + torch.tensor([[0]], dtype=torch.int64, device=device), + ) + properties = Labels.range("property", 1).to(device=device) + + if "energy" in outputs: + energy = torch.stack( + [system.positions.square().sum() for system in systems] + ).reshape(-1, 1) + result["energy"] = TensorMap( + placeholder, + [TensorBlock(energy, system_samples, [], properties)], + ) + + if "non_conservative_force" in outputs: + force_values: List[torch.Tensor] = [] + force_samples: List[torch.Tensor] = [] + if selected_atoms is None: + for system_index, system in enumerate(systems): + for atom_index in range(len(system)): + force_values.append(system.positions[atom_index]) + force_samples.append( + torch.tensor( + [system_index, atom_index], + dtype=torch.int64, + device=device, + ) + ) + else: + system_indices = selected_atoms.column("system").to(dtype=torch.long) + atom_indices = selected_atoms.column("atom").to(dtype=torch.long) + for row in range(len(selected_atoms)): + system_index = int(system_indices[row]) + atom_index = int(atom_indices[row]) + force_values.append(systems[system_index].positions[atom_index]) + force_samples.append(selected_atoms.values[row]) + + if len(force_values) == 0: + force = torch.empty( + (0, 3, 1), + dtype=systems[0].positions.dtype, + device=device, + ) + samples = torch.empty((0, 2), dtype=torch.int64, device=device) + else: + force = torch.stack(force_values).unsqueeze(-1) + samples = torch.stack(force_samples) + result["non_conservative_force"] = TensorMap( + placeholder, + [ + TensorBlock( + force, + Labels(["system", "atom"], samples), + [Labels.range("xyz", 3).to(device=device)], + properties, + ) + ], + ) + + if "non_conservative_stress" in outputs: + stress = torch.stack( + [system.positions.T @ system.positions for system in systems] + ).unsqueeze(-1) + result["non_conservative_stress"] = TensorMap( + placeholder, + [ + TensorBlock( + stress, + system_samples, + [ + Labels.range("xyz_1", 3).to(device=device), + Labels.range("xyz_2", 3).to(device=device), + ], + properties, + ) + ], + ) + + if "mtt::spherical_vector" in outputs: + spherical = torch.stack( + [system.positions[0].roll(-1) for system in systems] + ).unsqueeze(-1) + result["mtt::spherical_vector"] = TensorMap( + Labels( + ["o3_lambda", "o3_sigma"], + torch.tensor([[1, 1]], dtype=torch.int64, device=device), + ), + [ + TensorBlock( + spherical, + system_samples, + [_o3_mu_labels(1, device)], + properties, + ) + ], + ) + + if "mtt::spherical_quadrupole" in outputs: + matrices = torch.stack( + [ + torch.outer(system.positions[0], system.positions[0]) + for system in systems + ] + ).unsqueeze(-1) + _, _, spherical = _matrices_to_spherical(matrices) + result["mtt::spherical_quadrupole"] = TensorMap( + Labels( + ["o3_lambda", "o3_sigma"], + torch.tensor([[2, 1]], dtype=torch.int64, device=device), + ), + [ + TensorBlock( + spherical, + system_samples, + [_o3_mu_labels(2, device)], + properties, + ) + ], + ) + + return result + + +class TestQuadrature: + """Test quadrature weights and grid properties.""" + + def test_weights_sum(self): + """Quadrature weights should sum to 1 (normalized Haar measure on SO(3)).""" + for L_max in [3, 5, 7]: + lebedev_order, n_inplane = choose_quadrature(L_max) + _, _, _, w = get_euler_angles_quadrature(lebedev_order, n_inplane) + # The weights are w_i / (4*pi*K) repeated K times, where w_i sum to 4*pi + # So total sum = sum(w_i)/(4*pi*K) * K = sum(w_i)/(4*pi) = 1 + assert np.allclose(w.sum(), 1.0, atol=1e-12), ( + f"Weights don't sum to 1 for L_max={L_max}: sum={w.sum()}" + ) + + def test_euler_angle_rotations_are_in_so3(self): + """Euler-angle matrices should be orthogonal with determinant +1.""" + lebedev_order, n_inplane = choose_quadrature(5) + alpha, beta, gamma, _ = get_euler_angles_quadrature(lebedev_order, n_inplane) + rotations = _rotations_from_euler_angles(alpha, beta, gamma) + matrices = rotations.as_matrix() + + identity = np.broadcast_to(np.eye(3), matrices.shape) + assert np.allclose( + matrices @ matrices.transpose(0, 2, 1), + identity, + rtol=0.0, + atol=1e-12, + ) + assert np.allclose( + np.linalg.det(matrices), + 1.0, + rtol=0.0, + atol=1e-12, + ) + + def test_quadrature_validation(self): + """Quadrature construction rejects invalid degrees, counts, and orders.""" + message = ( + "the requested quadrature degree max_angular_momentum=132 exceeds the " + "largest available Lebedev order (131)" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + choose_quadrature(132) + + message = "max_angular_momentum must be non-negative, got -1" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + choose_quadrature(-1) + + message = "max_angular_momentum must be an integer, got float" + with pytest.raises(TypeError, match=f"^{re.escape(message)}$"): + choose_quadrature(1.5) + + message = "n_rotations must be positive, got 0" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + get_rotation_quadrature(3, 0) + + message = "n_rotations must be an integer, got float" + with pytest.raises(TypeError, match=f"^{re.escape(message)}$"): + get_rotation_quadrature(3, 1.5) + + supported_orders = [ + *range(3, 32, 2), + *range(35, 132, 6), + ] + message = ( + f"unsupported Lebedev order 4; supported orders are {supported_orders}" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + get_rotation_quadrature(4, 3) + + def test_degree_two_grid_resolves_l1_products(self): + order, n_rotations = choose_quadrature(2) + rotations, weights = get_rotation_quadrature(order, n_rotations) + function = rotations[:, 2, 0] + + norm = np.sum(weights * function**2) + projection_matrix = np.einsum("g,gij,g->ij", weights, rotations, function) + projected_norm = 3.0 * np.sum(projection_matrix**2) + assert np.isclose(norm, 1.0 / 3.0, atol=1e-12) + assert np.isclose(projected_norm, 1.0 / 3.0, atol=1e-12) + + def test_rotation_quadrature_matrices(self): + """Inversion should pair every proper rotation with an improper partner.""" + rotations, _ = get_rotation_quadrature(11, 5) + o3_rotations, _ = get_rotation_quadrature(11, 5, include_inversion=True) + + assert len(o3_rotations) == 2 * len(rotations) + dets = np.linalg.det(o3_rotations) + assert np.allclose(np.sort(dets), np.repeat([-1.0, 1.0], len(rotations))) + + +class TestSymmetrizedModelConstruction: + """Test construction of the quadrature and persistent Wigner-D storage.""" + + def test_character_limit_controls_default_grid(self): + """Character sectors should raise the default grid degree when necessary.""" + model = SymmetrizedModel( + _EmptyModel(), + max_angular_momentum_target=0, + max_angular_momentum_character=2, + ) + + assert model.max_angular_momentum_grid == 4 + + def test_rejects_grid_too_small_for_character_sectors(self): + """An explicit grid must resolve products for every requested sector.""" + message = ( + "max_angular_momentum_grid must be at least twice " + "max_angular_momentum_character" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + SymmetrizedModel( + _EmptyModel(), + max_angular_momentum_target=0, + max_angular_momentum_character=2, + max_angular_momentum_grid=3, + ) + + def test_rejects_invalid_constructor_arguments(self): + """Every integer constructor argument should enforce its documented range.""" + message = "max_angular_momentum_target must be non-negative, got -1" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + SymmetrizedModel(_EmptyModel(), max_angular_momentum_target=-1) + + message = "max_angular_momentum_target must be an integer, got bool" + with pytest.raises(TypeError, match=f"^{re.escape(message)}$"): + SymmetrizedModel(_EmptyModel(), max_angular_momentum_target=True) + + message = "max_angular_momentum_input must be an integer, got float" + with pytest.raises(TypeError, match=f"^{re.escape(message)}$"): + SymmetrizedModel( + _EmptyModel(), + max_angular_momentum_target=0, + max_angular_momentum_input=1.5, + ) + + message = "max_angular_momentum_character must be non-negative, got -1" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + SymmetrizedModel( + _EmptyModel(), + max_angular_momentum_target=0, + max_angular_momentum_character=-1, + ) + + message = "batch_size must be positive, got 0" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + SymmetrizedModel( + _EmptyModel(), + max_angular_momentum_target=0, + batch_size=0, + ) + + message = "max_angular_momentum_grid must be non-negative, got -1" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + SymmetrizedModel( + _EmptyModel(), + max_angular_momentum_target=0, + max_angular_momentum_grid=-1, + ) + + def test_rejects_a_model_stored_on_an_unsupported_device(self): + """Reject direct construction from a model outside CPU or CUDA.""" + base_model = _EmptyModel() + base_model.register_buffer("_device_marker", torch.empty(0, device="meta")) + + message = "SymmetrizedModel supports CPU and CUDA execution" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + SymmetrizedModel(base_model, max_angular_momentum_target=0) + + +class TestSymmetrizedModelForward: + """Test how requested averages and diagnostics are computed and returned.""" + + def test_character_projection_separates_sectors_through_lambda_three(self): + """Character projection separates the eight analytic sectors to lambda=3.""" + source_name = "mtt::o3_polynomial_sectors" + requested_name = "o3::character_projection::" + source_name + sectors = [ + (chi_lambda, chi_sigma) for chi_lambda in range(4) for chi_sigma in (1, -1) + ] + model = SymmetrizedModel( + _O3PolynomialSectorModel(), + max_angular_momentum_target=0, + max_angular_momentum_character=3, + max_angular_momentum_grid=6, + batch_size=17, + ) + system = _forward_test_system(torch.eye(3, dtype=torch.float64).tolist()) + + result = model( + [system], + {requested_name: ModelOutput(sample_kind="system")}, + None, + )[requested_name] + + assert result.keys.names == ["chi_lambda", "chi_sigma"] + assert result.keys.values.tolist() == [list(sector) for sector in sectors] + expected_properties = Labels( + ["source_lambda", "source_sigma"], + torch.tensor(sectors, dtype=torch.int64), + ) + # O(3) averages on the unit sphere: =1/3, <(xy)^2>=1/15, + # <(xyz)^2>=1/105 + expected_norms = [1.0, 1.0 / 3.0, 1.0 / 15.0, 1.0 / 105.0] + for key, block in result.items(): + assert block.samples == Labels("system", torch.tensor([[0]])) + assert block.components == [] + assert block.properties == expected_properties + + expected = torch.zeros((1, len(sectors)), dtype=torch.float64) + source_index = sectors.index( + (int(key["chi_lambda"]), int(key["chi_sigma"])) + ) + expected[0, source_index] = expected_norms[int(key["chi_lambda"])] + assert torch.allclose( + block.values, + expected, + rtol=0.0, + atol=1.0e-11, + ) + + def test_energy_results_match_analytic_values_and_reuse_predictions(self): + """Reuse each energy prediction for its average and both diagnostics.""" + base_model = _CountingLinearEnergyModel() + batch_size = 5 + model = SymmetrizedModel( + base_model, + max_angular_momentum_target=0, + max_angular_momentum_character=1, + max_angular_momentum_grid=2, + batch_size=batch_size, + ) + system = _forward_test_system([[1.0, 2.0, 3.0]]) + outputs = { + "energy": ModelOutput(sample_kind="system"), + "o3::variance::energy": ModelOutput(sample_kind="system"), + "o3::character_projection::energy": ModelOutput(sample_kind="system"), + } + + with torch.inference_mode(): + result = model([system], outputs, None) + + assert set(result) == set(outputs) + n_rotations = len(model._rotation_matrices) + assert base_model.call_count == 2 * ( + (n_rotations + batch_size - 1) // batch_size + ) + assert all(names == ["energy"] for names in base_model.requested_names) + + assert torch.allclose( + result["energy"].block().values, + torch.zeros((1, 1), dtype=torch.float64), + atol=1.0e-12, + ) + # under O(3) rotations of r=(1, 2, 3): |r|^2 / 3 = 14/3 + expected_variance = torch.tensor([[14.0 / 3.0]], dtype=torch.float64) + assert torch.allclose( + result["o3::variance::energy"].block().values, + expected_variance, + atol=1.0e-12, + ) + + projection = result["o3::character_projection::energy"] + assert projection.keys.names == [ + "o3_lambda", + "o3_sigma", + "chi_lambda", + "chi_sigma", + ] + vector_projection = projection.block( + { + "o3_lambda": 0, + "o3_sigma": 1, + "chi_lambda": 1, + "chi_sigma": 1, + } + ) + assert torch.allclose( + vector_projection.values.squeeze(1), + expected_variance, + atol=1.0e-12, + ) + for key, block in projection.items(): + if int(key["chi_lambda"]) == 1 and int(key["chi_sigma"]) == 1: + continue + assert torch.allclose( + block.values, + torch.zeros_like(block.values), + atol=1.0e-12, + ) + + def test_stress_character_projection_combines_target_and_character_sectors(self): + """Keep the stress irreps separate from its O(3) character sectors.""" + requested_name = "o3::character_projection::non_conservative_stress" + model = SymmetrizedModel( + _EquivariantOutputModel(), + max_angular_momentum_target=2, + max_angular_momentum_character=2, + max_angular_momentum_grid=4, + batch_size=17, + ) + system = _forward_test_system([[1.0, 2.0, 3.0], [-0.5, 0.25, 1.0]]) + + result = model( + [system], + {requested_name: ModelOutput(sample_kind="system")}, + None, + ) + + assert set(result) == {requested_name} + projection = result[requested_name] + assert projection.keys.names == [ + "o3_lambda", + "o3_sigma", + "chi_lambda", + "chi_sigma", + ] + assert { + tuple(int(value) for value in key.values) for key in projection.keys + } == { + (o3_lambda, o3_sigma, chi_lambda, chi_sigma) + for o3_lambda, o3_sigma in ((0, 1), (1, -1), (2, 1)) + for chi_lambda in range(3) + for chi_sigma in (1, -1) + } + + for key, block in projection.items(): + o3_lambda = int(key["o3_lambda"]) + o3_sigma = int(key["o3_sigma"]) + chi_lambda = int(key["chi_lambda"]) + chi_sigma = int(key["chi_sigma"]) + assert block.components == [_o3_mu_labels(o3_lambda, block.values.device)] + + # the stress of this model is exactly symmetric, so its l=1 + # pseudovector sector is zero + if o3_sigma == 1 and chi_lambda == o3_lambda and chi_sigma == 1: + assert bool(torch.any(block.values > 1.0e-12)) + else: + assert torch.allclose( + block.values, + torch.zeros_like(block.values), + rtol=0.0, + atol=1.0e-11, + ) + + @pytest.mark.parametrize( + ("requested_name", "unit"), + [ + ("energy", "eV"), + ("o3::variance::energy", "(eV)^2"), + ("o3::character_projection::energy", "(eV)^2"), + ], + ) + def test_source_request_contains_only_the_shared_sample_kind( + self, + requested_name, + unit, + ): + """Do not pass diagnostic metadata to the underlying source output.""" + base_model = _CountingLinearEnergyModel() + model = SymmetrizedModel( + base_model, + max_angular_momentum_target=0, + max_angular_momentum_character=1, + max_angular_momentum_grid=2, + ) + + model( + [_forward_test_system([[1.0, 2.0, 3.0]])], + { + requested_name: ModelOutput( + unit=unit, + sample_kind="system", + description="Metadata for the public result.", + ) + }, + None, + ) + + assert all(names == ["energy"] for names in base_model.requested_names) + assert set(base_model.requested_sample_kinds) == {"system"} + assert set(base_model.requested_units) == {""} + assert base_model.requested_explicit_gradients == [ + [] for _ in base_model.requested_explicit_gradients + ] + assert set(base_model.requested_descriptions) == {""} + + def test_rejects_an_output_above_the_declared_target_rank(self): + """Reject a rank-two spherical output when the declared limit is one.""" + model = SymmetrizedModel( + _EquivariantOutputModel(), + max_angular_momentum_target=1, + ) + + message = ( + "output 'mtt::spherical_quadrupole' contains o3_lambda=2, " + "exceeding max_angular_momentum_target=1" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + model( + [_forward_test_system([[1.0, 2.0, 3.0]])], + { + "mtt::spherical_quadrupole": ModelOutput( + sample_kind="system", + ) + }, + None, + ) + + def test_rejects_a_negative_quadrature_error_and_converges(self): + """A degree-12 grid rejects the degree-14 response; degree 14 is exact.""" + position = torch.tensor( + [[-1.12984253e-2, 3.64940445e-4, -9.99936104e-1]], + dtype=torch.float64, + ) + position = position / torch.linalg.norm(position) + system = _forward_test_system(position.tolist()) + variance_name = "o3::variance::energy" + variance_request = { + variance_name: ModelOutput(sample_kind="system"), + } + + underresolved = SymmetrizedModel( + _DegreeSevenEnergyModel(), + max_angular_momentum_target=0, + max_angular_momentum_grid=12, + batch_size=64, + ) + message = ( + "finite O(3) variance is materially negative; the quadrature does " + "not resolve this response. Increase max_angular_momentum_grid above 12 " + "and check convergence" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + underresolved([system], variance_request, None) + + resolved = SymmetrizedModel( + _DegreeSevenEnergyModel(), + max_angular_momentum_target=0, + max_angular_momentum_grid=14, + batch_size=64, + ) + outputs = { + "energy": ModelOutput(sample_kind="system"), + variance_name: ModelOutput(sample_kind="system"), + } + result = resolved([system], outputs, None) + # closed-form Haar variance of the degree-seven response at |r| = 1 + expected_variance = 1.0e6 * 17.0 / 137280.0 + assert result["energy"].block().values.item() == pytest.approx( + 0.0, + abs=1.0e-12, + ) + assert result[variance_name].block().values.item() == pytest.approx( + expected_variance, + rel=1.0e-12, + ) + + @pytest.mark.parametrize( + "source_name", + [ + "energy/pbe", + "mtt::feature::node", + # the reserved prefix is stripped exactly once, keeping "mtt::aux::" + "mtt::aux::features", + ], + ) + def test_preserves_variant_and_custom_output_names(self, source_name): + """Return variants and custom outputs under their exact requested names.""" + base_model = _CountingLinearEnergyModel() + model = SymmetrizedModel( + base_model, + max_angular_momentum_target=0, + max_angular_momentum_grid=2, + ) + variance_name = "o3::variance::" + source_name + outputs = { + source_name: ModelOutput(sample_kind="system"), + variance_name: ModelOutput(sample_kind="system"), + } + + result = model( + [_forward_test_system([[1.0, 2.0, 3.0]])], + outputs, + None, + ) + + assert set(result) == set(outputs) + assert all(names == [source_name] for names in base_model.requested_names) + assert torch.allclose( + result[variance_name].block().values, + torch.tensor([[14.0 / 3.0]], dtype=torch.float64), + atol=1.0e-12, + ) + + def test_component_less_output_averages_and_measures_invariance(self): + """Features have no spherical character: plain mean, invariance variance.""" + system = _forward_test_system([[1.0, 2.0, 3.0], [0.0, 1.0, 0.0]]) + model = SymmetrizedModel( + _AtomFeatureModel(), + max_angular_momentum_target=0, + max_angular_momentum_grid=2, + batch_size=5, + ) + outputs = { + "feature": ModelOutput(sample_kind="atom"), + "o3::variance::feature": ModelOutput(sample_kind="atom"), + } + + result = model([system], outputs, None) + + mean = result["feature"].block() + assert mean.samples.values.tolist() == [[0, 0], [0, 1]] + # the mean of x over O(3) is zero, without any back-rotation + assert torch.allclose(mean.values, torch.zeros_like(mean.values), atol=1.0e-12) + + variance = result["o3::variance::feature"] + # the tensor is passed through undecomposed, keeping its original keys + assert variance.keys.names == ["_"] + assert torch.allclose( + variance.block().values, + # - ^2 = |r|^2 / 3 for each atom + (system.positions.square().sum(dim=1) / 3.0).reshape(-1, 1), + atol=1.0e-12, + ) + + def test_selected_atoms_excludes_unselected_input_systems(self): + """Selecting only from System 1 must not create samples for System 0.""" + systems = [ + _forward_test_system([[1.0, 0.0, 0.0], [0.0, 2.0, 0.0]]), + _forward_test_system([[0.0, 0.0, 3.0], [4.0, 5.0, 6.0]]), + ] + model = SymmetrizedModel( + _EquivariantOutputModel(), + max_angular_momentum_target=1, + max_angular_momentum_grid=2, + batch_size=5, + ) + outputs = { + "non_conservative_force": ModelOutput(sample_kind="atom"), + "o3::variance::non_conservative_force": ModelOutput(sample_kind="atom"), + } + selected_atoms = Labels( + ["system", "atom"], + torch.tensor([[1, 1]], dtype=torch.int64), + ) + + result = model(systems, outputs, selected_atoms) + + mean = result["non_conservative_force"].block() + assert mean.samples.values.tolist() == [[1, 1]] + assert torch.allclose( + mean.values.squeeze(-1), + systems[1].positions[1].reshape(1, 3), + atol=1.0e-12, + ) + variance = result["o3::variance::non_conservative_force"] + assert variance.keys.values.tolist() == [[1, 1]] + assert variance.block().samples.values.tolist() == [[1, 1]] + assert torch.allclose( + variance.block().values, + torch.zeros_like(variance.block().values), + atol=1.0e-12, + ) + + def test_empty_selected_atoms_returns_empty_outputs(self): + """A fully empty atom selection must not create artificial samples.""" + systems = [ + _forward_test_system([[1.0, 0.0, 0.0]]), + _forward_test_system([[0.0, 2.0, 0.0]]), + ] + model = SymmetrizedModel( + _EquivariantOutputModel(), + max_angular_momentum_target=1, + max_angular_momentum_grid=2, + batch_size=5, + ) + outputs = { + "non_conservative_force": ModelOutput(sample_kind="atom"), + "o3::variance::non_conservative_force": ModelOutput(sample_kind="atom"), + } + selected_atoms = Labels( + ["system", "atom"], + torch.empty((0, 2), dtype=torch.int64), + ) + + result = model(systems, outputs, selected_atoms) + + assert set(result) == set(outputs) + mean = result["non_conservative_force"].block() + assert mean.samples.names == ["system", "atom"] + assert len(mean.samples) == 0 + assert mean.values.shape == (0, 3, 1) + + variance = result["o3::variance::non_conservative_force"] + assert variance.keys.names == ["o3_lambda", "o3_sigma"] + assert variance.keys.values.tolist() == [[1, 1]] + variance_block = variance.block() + assert variance_block.samples.names == ["system", "atom"] + assert len(variance_block.samples) == 0 + assert variance_block.components == [] + assert variance_block.values.shape == (0, 1) + + def test_multiple_systems_keep_per_system_rows_in_order(self): + """Joined outputs should keep one correct row per input System, in order.""" + systems = [ + _forward_test_system([[1.0, 2.0, 3.0]]), + _forward_test_system([[-0.5, 0.25, 1.0], [0.5, -1.0, 2.0]]), + ] + model = SymmetrizedModel( + _EquivariantOutputModel(), + max_angular_momentum_target=0, + max_angular_momentum_grid=2, + batch_size=5, + ) + outputs = { + "energy": ModelOutput(sample_kind="system"), + "o3::variance::energy": ModelOutput(sample_kind="system"), + } + + result = model(systems, outputs, None) + + energy = result["energy"].block() + assert energy.samples.values.tolist() == [[0], [1]] + expected = torch.stack( + [system.positions.square().sum() for system in systems] + ).reshape(-1, 1) + assert torch.allclose(energy.values, expected, atol=1.0e-12) + variance = result["o3::variance::energy"].block() + assert variance.samples.values.tolist() == [[0], [1]] + assert torch.allclose( + variance.values, + torch.zeros_like(variance.values), + atol=1.0e-12, + ) + + def test_equivariant_outputs_preserve_values_metadata_and_zero_variance(self): + """Return exact equivariant outputs unchanged and report zero variance.""" + system = _forward_test_system([[1.0, 2.0, 3.0], [-0.5, 0.25, 1.0]]) + sources = [ + "energy", + "non_conservative_force", + "non_conservative_stress", + "mtt::spherical_vector", + "mtt::spherical_quadrupole", + ] + outputs = { + name: ModelOutput( + sample_kind="atom" if name == "non_conservative_force" else "system" + ) + for name in sources + } + for name in sources: + outputs["o3::variance::" + name] = outputs[name] + model = SymmetrizedModel( + _EquivariantOutputModel(), + max_angular_momentum_target=2, + max_angular_momentum_grid=2, + batch_size=7, + ) + + result = model([system], outputs, None) + + assert set(result) == set(outputs) + assert torch.allclose( + result["energy"].block().values, + system.positions.square().sum().reshape(1, 1), + atol=1.0e-12, + ) + assert torch.allclose( + result["non_conservative_force"].block().values.squeeze(-1), + system.positions, + atol=1.0e-12, + ) + assert torch.allclose( + result["non_conservative_stress"].block().values.squeeze(-1), + (system.positions.T @ system.positions).unsqueeze(0), + atol=1.0e-12, + ) + assert torch.allclose( + result["mtt::spherical_vector"].block().values.squeeze(-1), + system.positions[0].roll(-1).reshape(1, 3), + atol=1.0e-12, + ) + quadrupole = result["mtt::spherical_quadrupole"] + assert quadrupole.keys.values.tolist() == [[2, 1]] + _, _, expected_quadrupole = _matrices_to_spherical( + torch.outer(system.positions[0], system.positions[0]).reshape(1, 3, 3, 1) + ) + assert torch.allclose( + quadrupole.block().values, + expected_quadrupole, + atol=1.0e-12, + ) + + expected_target_keys = { + "o3::variance::energy": [[0, 1]], + "o3::variance::non_conservative_force": [[1, 1]], + "o3::variance::non_conservative_stress": [[0, 1], [1, -1], [2, 1]], + "o3::variance::mtt::spherical_vector": [[1, 1]], + "o3::variance::mtt::spherical_quadrupole": [[2, 1]], + } + for name, expected_keys in expected_target_keys.items(): + variance = result[name] + assert variance.keys.names == ["o3_lambda", "o3_sigma"] + assert variance.keys.values.tolist() == expected_keys + for block in variance.blocks(): + assert block.components == [] + assert torch.allclose( + block.values, + torch.zeros_like(block.values), + atol=1.0e-12, + ) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) + def test_dtype_and_implicit_autograd(self, dtype): + """Variance should preserve the model dtype and its implicit backward path.""" + system = _forward_test_system( + [[1.0, 2.0, 3.0]], + dtype=dtype, + requires_grad=True, + ) + model = SymmetrizedModel( + _LinearEnergyModel(), + max_angular_momentum_target=0, + max_angular_momentum_grid=2, + ) + outputs = { + "o3::variance::energy": ModelOutput(sample_kind="system"), + } + + result = model([system], outputs, None) + variance = result["o3::variance::energy"].block().values + + assert variance.dtype == dtype + gradient = torch.autograd.grad(variance.sum(), system.positions)[0] + tolerance = 2.0e-5 if dtype == torch.float32 else 1.0e-12 + assert torch.allclose( + gradient, + 2.0 * system.positions / 3.0, + rtol=0.0, + atol=tolerance, + ) + + def test_average_output_preserves_implicit_autograd(self): + """The averaged output should keep the implicit backward path to positions.""" + system = _forward_test_system( + [[1.0, 2.0, 3.0], [-0.5, 0.25, 1.0]], + requires_grad=True, + ) + model = SymmetrizedModel( + _EquivariantOutputModel(), + max_angular_momentum_target=0, + max_angular_momentum_grid=2, + ) + + result = model([system], {"energy": ModelOutput(sample_kind="system")}, None) + + gradient = torch.autograd.grad( + result["energy"].block().values.sum(), + system.positions, + )[0] + assert torch.allclose( + gradient, + 2.0 * system.positions, + rtol=0.0, + atol=1.0e-12, + ) + + def test_rejects_unknown_o3_requests(self): + """Every unrecognized 'o3::' request is reserved, not a source output.""" + model = SymmetrizedModel( + _LinearEnergyModel(), + max_angular_momentum_target=0, + max_angular_momentum_grid=2, + ) + message = ( + "requested output 'o3::variance_extra::energy' uses the 'o3::' " + "prefix reserved by SymmetrizedModel, but is neither a variance nor " + "a character-projection request" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + model( + [_forward_test_system([[1.0, 2.0, 3.0]])], + {"o3::variance_extra::energy": ModelOutput(sample_kind="system")}, + None, + ) + + def test_input_limit_distinguishes_spherical_from_cartesian(self): + """A zero angular-momentum input limit still allows Cartesian custom data.""" + model = SymmetrizedModel( + _LinearEnergyModel(), + max_angular_momentum_target=0, + max_angular_momentum_grid=2, + ) + outputs = {"energy": ModelOutput(sample_kind="system")} + + cartesian_system = _forward_test_system([[1.0, 2.0, 3.0]]) + cartesian_system.add_data( + "mtt::field", + TensorMap( + Labels("_", torch.tensor([[0]])), + [ + TensorBlock( + values=torch.ones((1, 3, 1), dtype=torch.float64), + samples=Labels.range("atom", 1), + components=[Labels.range("xyz", 3)], + properties=Labels.range("property", 1), + ) + ], + ), + ) + model([cartesian_system], outputs, None) + + spherical_system = _forward_test_system([[1.0, 2.0, 3.0]]) + spherical_system.add_data( + "mtt::field", + TensorMap( + Labels(["o3_lambda", "o3_sigma"], torch.tensor([[1, 1]])), + [ + TensorBlock( + values=torch.ones((1, 3, 1), dtype=torch.float64), + samples=Labels.range("atom", 1), + components=[_o3_mu_labels(1, torch.device("cpu"))], + properties=Labels.range("property", 1), + ) + ], + ), + ) + message = ( + "custom input 'mtt::field' contains o3_lambda=1, exceeding " + "max_angular_momentum_input=0" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + model([spherical_system], outputs, None) + + def test_variance_is_stable_with_large_mean_offset(self): + """A huge invariant offset must not destroy the variance numerically.""" + model = SymmetrizedModel( + _OffsetLinearEnergyModel(1.0e8), + max_angular_momentum_target=0, + max_angular_momentum_grid=2, + ) + system = _forward_test_system([[1.0, 2.0, 3.0]]) + + result = model( + [system], + {"o3::variance::energy": ModelOutput(sample_kind="system")}, + None, + ) + + # - ^2 = |r|^2 / 3, unchanged by the constant offset + assert result["o3::variance::energy"].block().values.item() == pytest.approx( + 14.0 / 3.0, + rel=1.0e-5, + ) + + def test_inconsistent_sample_labels_from_wrapped_model_are_rejected(self): + """A wrapped model mislabelling its per-copy samples fails loudly.""" + model = SymmetrizedModel( + _InconsistentSampleModel(), + max_angular_momentum_target=0, + max_angular_momentum_grid=2, + ) + message = ( + "SymmetrizedModel expects every rotated copy to produce the same " + "sample labels in the same order." + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + model( + [_forward_test_system([[1.0, 2.0, 3.0]])], + {"energy": ModelOutput(sample_kind="atom")}, + None, + ) + + def test_rejects_invalid_requests_before_model_evaluation(self): + """Invalid public requests should fail without running the source model.""" + base_model = _CountingLinearEnergyModel() + model = SymmetrizedModel(base_model, max_angular_momentum_target=0) + system = _forward_test_system([[1.0, 2.0, 3.0]]) + + assert model([], {}, None) == {} + message = "SymmetrizedModel requires at least one System" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + model([], {"energy": ModelOutput(sample_kind="system")}, None) + message = ( + "max_angular_momentum_character must be set to request " + "character projections" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + model( + [system], + {"o3::character_projection::energy": ModelOutput(sample_kind="system")}, + None, + ) + message = ( + "SymmetrizedModel does not support explicit gradients for output 'energy'" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + model( + [system], + { + "energy": ModelOutput( + sample_kind="system", + explicit_gradients=["positions"], + ) + }, + None, + ) + message = ( + "all requests derived from 'energy' must use the same sample_kind; " + "got 'system' and 'atom'" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + model( + [system], + { + "energy": ModelOutput(sample_kind="system"), + "o3::variance::energy": ModelOutput(sample_kind="atom"), + }, + None, + ) + assert base_model.call_count == 0 + + def test_downcast_integration_buffers_warn_and_run(self): + """Calling .float() on the module warns and loses accuracy, not correctness.""" + model = SymmetrizedModel( + _LinearEnergyModel(), + max_angular_momentum_target=0, + max_angular_momentum_grid=2, + ).float() + system = _forward_test_system([[1.0, 2.0, 3.0]], dtype=torch.float32) + + with pytest.warns(UserWarning, match="downcast from float64"): + result = model( + [system], + {"o3::variance::energy": ModelOutput(sample_kind="system")}, + None, + ) + + variance = result["o3::variance::energy"].block().values + assert variance.dtype == torch.float32 + # under O(3) rotations of r=(1, 2, 3): |r|^2 / 3 = 14/3 + assert variance.item() == pytest.approx(14.0 / 3.0, rel=1.0e-3) + + def test_rejects_a_model_that_omits_the_requested_output(self): + """Fail loudly when the underlying model does not return a source.""" + model = SymmetrizedModel( + _EmptyModel(), + max_angular_momentum_target=0, + max_angular_momentum_grid=2, + ) + + message = "underlying model did not return requested output 'energy'" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + model( + [_forward_test_system([[1.0, 2.0, 3.0]])], + {"energy": ModelOutput(sample_kind="system")}, + None, + ) + + def test_rejects_a_non_finite_variance(self): + """A NaN model response should fail the variance finiteness check.""" + model = SymmetrizedModel( + _LinearEnergyModel(), + max_angular_momentum_target=0, + max_angular_momentum_grid=2, + ) + + message = "O(3) variance is not finite for block ((o3_lambda=0, o3_sigma=1))" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + model( + [_forward_test_system([[float("nan"), 2.0, 3.0]])], + {"o3::variance::energy": ModelOutput(sample_kind="system")}, + None, + ) + + def test_is_scriptable_and_serializable(self, tmp_path): + """The complete forward path should execute after scripting and reloading.""" + constructor_arguments = { + "max_angular_momentum_target": 0, + "max_angular_momentum_character": 1, + "max_angular_momentum_grid": 2, + "batch_size": 5, + } + eager = SymmetrizedModel(_LinearEnergyModel(), **constructor_arguments) + scripted = torch.jit.script( + SymmetrizedModel(_LinearEnergyModel(), **constructor_arguments) + ) + path = tmp_path / "symmetrized-model.pt" + torch.jit.save(scripted, str(path)) + loaded = torch.jit.load(str(path)) + system = _forward_test_system([[1.0, 2.0, 3.0]]) + outputs = { + "energy": ModelOutput(sample_kind="system"), + "o3::variance::energy": ModelOutput(sample_kind="system"), + "o3::character_projection::energy": ModelOutput(sample_kind="system"), + } + + expected = eager([system], outputs, None) + actual = loaded([system], outputs, None) + + assert set(actual) == set(expected) + for name in expected: + mts.allclose_raise(actual[name], expected[name], rtol=0.0, atol=1.0e-12) + + # a SymmetrizedModel can also wrap a scripted + reloaded inner model + inner_path = tmp_path / "inner-model.pt" + torch.jit.save(torch.jit.script(_LinearEnergyModel()), str(inner_path)) + rewrapped = SymmetrizedModel( + torch.jit.load(str(inner_path)), + **constructor_arguments, + ) + rewrapped_result = rewrapped([system], outputs, None) + for name in expected: + mts.allclose_raise( + rewrapped_result[name], + expected[name], + rtol=0.0, + atol=1.0e-12, + ) + + +class TestSymmetrizedModelWrap: + """Test exported-model capabilities, dependencies, and execution.""" + + @pytest.mark.parametrize("max_angular_momentum_character", [None, 1]) + def test_wrap_declares_capabilities(self, max_angular_momentum_character): + """Wrapping declares averages and diagnostics with squared units.""" + source_outputs = { + "energy": ModelOutput( + unit="eV", + sample_kind="system", + explicit_gradients=["positions"], + ), + "mass": ModelOutput(unit="u", sample_kind="atom"), + "mtt::pair": ModelOutput(sample_kind="atom_pair"), + } + base = AtomisticModel( + _EmptyModel().eval(), + ModelMetadata(name="wrapped source model"), + ModelCapabilities( + outputs=source_outputs, + atomic_types=[1, 6, 8], + interaction_range=4.5, + length_unit="A", + supported_devices=["cuda", "mps", "cpu"], + dtype="float32", + ), + ) + + wrapped = SymmetrizedModel.wrap( + base, + max_angular_momentum_target=0, + max_angular_momentum_character=max_angular_momentum_character, + max_angular_momentum_grid=2, + ) + + capabilities = wrapped.capabilities() + assert wrapped.metadata().name == "wrapped source model" + assert capabilities.atomic_types == [1, 6, 8] + assert capabilities.interaction_range == 4.5 + assert capabilities.length_unit == "A" + assert capabilities.supported_devices == ["cuda", "cpu"] + + expected_names = set(source_outputs) + expected_names.update("o3::variance::" + name for name in source_outputs) + if max_angular_momentum_character is not None: + expected_names.update( + "o3::character_projection::" + name for name in source_outputs + ) + # "masses" is a compatibility alias added by AtomisticModel; it must not + # become another declared source with its own diagnostics + assert set(capabilities.outputs) == expected_names | {"masses"} + assert "o3::variance::masses" not in capabilities.outputs + assert "o3::character_projection::masses" not in capabilities.outputs + + for name, source_output in source_outputs.items(): + squared_unit = ( + "" if source_output.unit == "" else f"({source_output.unit})^2" + ) + assert capabilities.outputs["o3::variance::" + name].unit == squared_unit + character_name = "o3::character_projection::" + name + if max_angular_momentum_character is None: + assert character_name not in capabilities.outputs + else: + assert capabilities.outputs[character_name].unit == squared_unit + + def test_guesses_limits_from_standard_quantities(self): + """Both limits default to what the standard quantities require.""" + + class _VelocityInputModel(_EmptyModel): + def requested_inputs(self) -> Dict[str, ModelOutput]: + # the custom input is skipped by the guess as well + return { + "velocity": ModelOutput(sample_kind="atom"), + "mtt::field": ModelOutput(sample_kind="atom"), + } + + def guessed_target(outputs): + base = AtomisticModel( + _VelocityInputModel().eval(), + ModelMetadata(), + ModelCapabilities( + outputs=outputs, + atomic_types=[1], + interaction_range=0.0, + length_unit="A", + supported_devices=["cpu"], + dtype="float64", + ), + ) + wrapped = SymmetrizedModel.wrap(base, max_angular_momentum_grid=2) + # velocity is a Cartesian vector + assert wrapped.module.max_angular_momentum_input == 1 + return wrapped.module.max_angular_momentum_target + + energy = ModelOutput(unit="eV", sample_kind="system") + force = ModelOutput(unit="eV/A", sample_kind="atom") + stress = ModelOutput(unit="eV/A^3", sample_kind="system") + + assert guessed_target({"energy": energy}) == 0 + assert guessed_target({"feature": ModelOutput(sample_kind="atom")}) == 0 + assert guessed_target({"energy": energy, "non_conservative_force": force}) == 1 + assert guessed_target({"non_conservative_stress": stress}) == 2 + # a custom output is skipped, the standard ones still set the limit + custom = ModelOutput(sample_kind="system") + assert guessed_target({"energy": energy, "mtt::custom": custom}) == 0 + + def test_rejects_guessing_a_limit_without_standard_outputs(self): + """Only non-standard outputs leave nothing to guess the limit from.""" + base = AtomisticModel( + _EmptyModel().eval(), + ModelMetadata(), + ModelCapabilities( + outputs={ + "mtt::custom": ModelOutput(sample_kind="system"), + "mtt::other": ModelOutput(sample_kind="system"), + }, + atomic_types=[1], + interaction_range=0.0, + length_unit="A", + supported_devices=["cpu"], + dtype="float64", + ), + ) + + message = ( + "no standard quantities were found among the outputs " + "['mtt::custom', 'mtt::other'], please set max_angular_momentum_target " + "explicitly" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + SymmetrizedModel.wrap(base) + + @pytest.mark.parametrize( + "source_name", + [ + "o3::variance::mtt::source", + "o3::character_projection::mtt::source", + ], + ) + def test_rejects_reserved_source_names(self, source_name): + """Reject source names that look like wrapper-generated diagnostics.""" + base = AtomisticModel( + _EmptyModel().eval(), + ModelMetadata(), + ModelCapabilities( + outputs={source_name: ModelOutput(sample_kind="system")}, + atomic_types=[1], + interaction_range=0.0, + length_unit="A", + supported_devices=["cpu"], + dtype="float64", + ), + ) + + message = ( + f"the wrapped model output '{source_name}' uses a prefix reserved " + "by SymmetrizedModel" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + SymmetrizedModel.wrap(base, max_angular_momentum_target=0) + + def test_rejects_models_without_a_supported_device(self): + """Reject models whose declared devices contain neither CPU nor CUDA.""" + base = AtomisticModel( + _EmptyModel().eval(), + ModelMetadata(), + ModelCapabilities( + outputs={"mtt::value": ModelOutput(sample_kind="system")}, + atomic_types=[1], + interaction_range=0.0, + length_unit="A", + supported_devices=["mps"], + dtype="float64", + ), + ) + + message = ( + "SymmetrizedModel supports CPU and CUDA execution, but the " + "wrapped model declares ['mps']" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + SymmetrizedModel.wrap(base, max_angular_momentum_target=0) + + def test_preserves_requirements_and_runs_after_save_load(self, tmp_path): + """Preserve model requirements through wrapping, saving, and reloading.""" + metadata = ModelMetadata(name="model with requirements") + base = AtomisticModel( + _LinearModelWithRequirements().eval(), + metadata, + ModelCapabilities( + outputs={ + "mtt::linear": ModelOutput( + unit="eV", + sample_kind="system", + description="First Cartesian coordinate.", + ) + }, + atomic_types=[1], + interaction_range=2.5, + length_unit="A", + supported_devices=["cpu"], + dtype="float32", + ), + ) + base_path = tmp_path / "base-model.pt" + base.save(base_path) + loaded_base = load_atomistic_model(base_path) + base_requestors = set(loaded_base.requested_neighbor_lists()[0].requestors()) + + wrapped = SymmetrizedModel.wrap( + loaded_base, + max_angular_momentum_target=0, + # 'mtt::linear' and 'mtt::field' are not standard quantities, so + # both limits have to be given explicitly + max_angular_momentum_input=0, + max_angular_momentum_character=1, + max_angular_momentum_grid=2, + batch_size=5, + ) + assert ( + set(loaded_base.requested_neighbor_lists()[0].requestors()) + == base_requestors + ) + + wrapped_path = tmp_path / "symmetrized-model.pt" + wrapped.save(wrapped_path) + loaded = load_atomistic_model(wrapped_path) + + requested_inputs = loaded.requested_inputs(use_new_names=True) + assert set(requested_inputs) == {"mtt::field"} + assert requested_inputs["mtt::field"].unit == "eV" + assert requested_inputs["mtt::field"].sample_kind == "atom" + assert ( + requested_inputs["mtt::field"].description + == "Cartesian field used by the model." + ) + + requested_neighbor_lists = loaded.requested_neighbor_lists() + assert len(requested_neighbor_lists) == 1 + neighbor_options = requested_neighbor_lists[0] + assert neighbor_options.cutoff == 2.5 + assert neighbor_options.full_list is False + assert neighbor_options.strict is True + assert base_requestors.issubset(set(neighbor_options.requestors())) + + system = _system_with_linear_model_requirements( + neighbor_options, + torch.device("cpu"), + ) + + requested_outputs = { + "mtt::linear": ModelOutput( + unit="meV", + sample_kind="system", + ), + "o3::variance::mtt::linear": ModelOutput( + unit="(meV)^2", + sample_kind="system", + ), + "o3::character_projection::mtt::linear": ModelOutput( + unit="(meV)^2", + sample_kind="system", + ), + } + evaluation_options = ModelEvaluationOptions( + length_unit="A", + outputs=requested_outputs, + ) + eager = wrapped([system], evaluation_options, check_consistency=True) + reloaded = loaded([system], evaluation_options, check_consistency=True) + + assert set(reloaded) == set(requested_outputs) + for name in eager: + mts.allclose_raise( + reloaded[name], + eager[name], + rtol=0.0, + atol=0.0, + ) + for block in reloaded[name].blocks(): + assert block.values.dtype == torch.float32 + + expected_variance = torch.tensor( + [[14.0 / 3.0 * 1.0e6]], + dtype=torch.float32, + ) + assert torch.allclose( + reloaded["o3::variance::mtt::linear"].block().values, + expected_variance, + rtol=2.0e-5, + atol=1.0, + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") + def test_saved_wrapper_runs_on_cuda(self, tmp_path): + """Match CPU results after moving a saved float32 wrapper to CUDA. + + The wrapper's model dtype is float32, while its integration buffers stay + float64 throughout; nothing here downcasts the module itself. + """ + base = AtomisticModel( + _LinearModelWithRequirements().eval(), + ModelMetadata(name="CUDA source model"), + ModelCapabilities( + outputs={ + "energy": ModelOutput( + unit="eV", + sample_kind="system", + ) + }, + atomic_types=[1], + interaction_range=2.5, + length_unit="A", + supported_devices=["cpu", "cuda"], + dtype="float32", + ), + ) + wrapped = SymmetrizedModel.wrap( + base, + max_angular_momentum_target=0, + max_angular_momentum_input=0, + max_angular_momentum_character=1, + max_angular_momentum_grid=2, + batch_size=5, + ) + path = tmp_path / "cuda-symmetrized-model.pt" + wrapped.save(path) + + cpu_model = load_atomistic_model(path) + cuda_device = torch.device("cuda", torch.cuda.current_device()) + cuda_model = load_atomistic_model(path).to(device=cuda_device) + neighbor_options = cpu_model.requested_neighbor_lists()[0] + cpu_system = _system_with_linear_model_requirements( + neighbor_options, + torch.device("cpu"), + ) + cuda_system = cpu_system.to(device=cuda_device) + + cpu_module = SymmetrizedModel( + _LinearEnergyModel(), max_angular_momentum_target=0 + ) + message = "SymmetrizedModel and input Systems must use the same device" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + cpu_module( + [cuda_system], + {"energy": ModelOutput(sample_kind="system")}, + None, + ) + + requested_outputs = { + "energy": ModelOutput( + unit="meV", + sample_kind="system", + ), + "o3::variance::energy": ModelOutput( + unit="(meV)^2", + sample_kind="system", + ), + "o3::character_projection::energy": ModelOutput( + unit="(meV)^2", + sample_kind="system", + ), + } + evaluation_options = ModelEvaluationOptions( + length_unit="A", + outputs=requested_outputs, + ) + with torch.inference_mode(): + expected = cpu_model( + [cpu_system], + evaluation_options, + check_consistency=True, + ) + actual = cuda_model( + [cuda_system], + evaluation_options, + check_consistency=True, + ) + + assert set(actual) == set(requested_outputs) + assert actual["energy"].block().values.device.type == "cuda" + for name, tensor in actual.items(): + mts.allclose_raise( + tensor.to(device="cpu"), + expected[name], + rtol=2.0e-5, + atol=2.0e-5, + ) + + +def test_selected_atoms_system_column_found_by_name(): + """The rotated-copy index goes into the "system" column wherever it sits.""" + selection = Labels(["atom", "system"], torch.tensor([[3, 0], [5, 0]])) + rotated = map_selected_atoms_to_rotated_copies(selection, 0, 2) + assert rotated.names == ["atom", "system"] + assert rotated.values[:, 0].tolist() == [3, 5, 3, 5] + assert rotated.values[:, 1].tolist() == [0, 0, 1, 1] + + +def test_cartesian_vectors_to_spherical(): + """Map Cartesian components to the real spherical l=1 ordering.""" + values = torch.tensor( + [[[1.0, 10.0], [2.0, 20.0], [3.0, 30.0]]], + dtype=torch.float64, + ) + + result = _cartesian_vectors_to_spherical(values, component_axis=1) + + assert torch.equal( + result, + torch.tensor( + [[[2.0, 20.0], [3.0, 30.0], [1.0, 10.0]]], + dtype=torch.float64, + ), + ) + + +@pytest.mark.parametrize("inversion", [1.0, -1.0]) +def test_cartesian_vectors_to_spherical_commutes_with_o3(inversion): + """Converting before or after an O(3) transformation should give the same result.""" + proper_rotation = torch.tensor( + [ + [-2.0 / 3.0, 2.0 / 15.0, 11.0 / 15.0], + [2.0 / 3.0, -1.0 / 3.0, 2.0 / 3.0], + [1.0 / 3.0, 14.0 / 15.0, 2.0 / 15.0], + ], + dtype=torch.float64, + ) + transformation = O3Transformation( + inversion * proper_rotation, + max_angular_momentum=1, + ) + cartesian = torch.tensor( + [[1.2, -0.7, 2.3], [-0.4, 1.1, 0.8]], + dtype=torch.float64, + ) + + transformed_cartesian = _cartesian_vectors_to_spherical( + transformation.transform_cartesian(cartesian), + component_axis=1, + ) + transformed_spherical = transformation.transform_spherical( + _cartesian_vectors_to_spherical(cartesian, component_axis=1), + ell=1, + sigma=1, + ) + + assert torch.allclose( + transformed_cartesian, + transformed_spherical, + rtol=0.0, + atol=1.0e-12, + ) + + +def test_matrices_to_spherical_known_components(): + """Known matrices map as expected and the Frobenius norm is preserved.""" + matrices = torch.zeros((3, 3, 3, 1), dtype=torch.float64) + matrices[0, :, :, 0] = torch.eye(3, dtype=torch.float64) + matrices[1, 0, 0, 0] = 1.0 + matrices[1, 1, 1, 0] = -1.0 + matrices[2, 0, 1, 0] = 2.0 + matrices[2, 1, 0, 0] = -2.0 + + l0, l1, l2 = _matrices_to_spherical(matrices) + + expected_l0 = torch.zeros((3, 1, 1), dtype=torch.float64) + expected_l0[0, 0, 0] = 3.0**0.5 + expected_l1 = torch.zeros((3, 3, 1), dtype=torch.float64) + # antisymmetric part of matrices[2]: axial vector (0, 0, -2), times sqrt(2) + expected_l1[2, 1, 0] = -2.0 * 2.0**0.5 + expected_l2 = torch.zeros((3, 5, 1), dtype=torch.float64) + expected_l2[1, 4, 0] = 2.0**0.5 + assert torch.allclose(l0, expected_l0, rtol=0.0, atol=1.0e-12) + assert torch.allclose(l1, expected_l1, rtol=0.0, atol=1.0e-12) + assert torch.allclose(l2, expected_l2, rtol=0.0, atol=1.0e-12) + + generator = torch.Generator().manual_seed(1234) + random_matrices = torch.randn( + (4, 3, 3, 2), + dtype=torch.float64, + generator=generator, + ) + + l0, l1, l2 = _matrices_to_spherical(random_matrices) + + spherical_norm_squared = ( + l0.square().sum(dim=1) + l1.square().sum(dim=1) + l2.square().sum(dim=1) + ) + cartesian_norm_squared = random_matrices.square().sum(dim=(1, 2)) + assert torch.allclose( + spherical_norm_squared, + cartesian_norm_squared, + rtol=0.0, + atol=1.0e-12, + ) + + +@pytest.mark.parametrize("inversion", [1.0, -1.0]) +def test_matrices_to_spherical_commutes_with_o3(inversion): + """Cartesian and spherical transformations should give the same components.""" + proper_rotation = torch.tensor( + [ + [-2.0 / 3.0, 2.0 / 15.0, 11.0 / 15.0], + [2.0 / 3.0, -1.0 / 3.0, 2.0 / 3.0], + [1.0 / 3.0, 14.0 / 15.0, 2.0 / 15.0], + ], + dtype=torch.float64, + ) + transformation = O3Transformation( + inversion * proper_rotation, + max_angular_momentum=2, + ) + # deliberately non-symmetric, so the l=1 pseudovector part is non-zero + matrices = torch.tensor( + [ + [[1.2, -0.7, 2.3], [0.9, 1.1, 0.8], [-1.6, 0.3, -0.4]], + [[-0.2, 1.4, 0.5], [0.6, 0.9, -1.1], [1.7, -0.3, 2.0]], + ], + dtype=torch.float64, + ).unsqueeze(-1) + + matrix = transformation.matrix + transformed_matrices = torch.einsum( + "ia,sabp,jb->sijp", + matrix, + matrices, + matrix, + ) + transformed_l0, transformed_l1, transformed_l2 = _matrices_to_spherical( + transformed_matrices + ) + l0, l1, l2 = _matrices_to_spherical(matrices) + + expected_l0 = transformation.transform_spherical( + l0[..., 0], ell=0, sigma=1 + ).unsqueeze(-1) + expected_l1 = transformation.transform_spherical( + l1[..., 0], ell=1, sigma=-1 + ).unsqueeze(-1) + expected_l2 = transformation.transform_spherical( + l2[..., 0], ell=2, sigma=1 + ).unsqueeze(-1) + assert torch.allclose(transformed_l0, expected_l0, rtol=0.0, atol=1.0e-12) + assert torch.allclose(transformed_l1, expected_l1, rtol=0.0, atol=1.0e-12) + assert torch.allclose(transformed_l2, expected_l2, rtol=0.0, atol=1.0e-12) + + +@pytest.mark.parametrize( + "source_name", + [ + "energy", + "energy/pbe", + "energy_ensemble/member", + "energy_uncertainty/direct", + "charge", + ], +) +def test_decompose_quantity_scalar_quantities(source_name): + """Scalar quantities and their variants become one l=0 spherical block.""" + values = torch.tensor([[1.0, 2.0]], dtype=torch.float64) + tensor = _tensor_map_with_components(values, []) + tensor.set_info("unit", "eV") + + result = decompose_quantity(source_name, tensor) + + assert result.keys.names == ["o3_lambda", "o3_sigma"] + assert result.keys.values.tolist() == [[0, 1]] + assert torch.equal(result.block().values, values.unsqueeze(1)) + assert result.block().samples == tensor.block().samples + assert result.block().components == [_o3_mu_labels(0, values.device)] + assert result.block().properties == tensor.block().properties + assert result.info() == tensor.info() + + +@pytest.mark.parametrize( + "source_name", + [ + "non_conservative_force/direct", + "velocity", + ], +) +def test_decompose_quantity_cartesian_vectors_preserve_autograd(source_name): + """Cartesian vectors should become l=1 and preserve implicit autograd.""" + values = torch.tensor( + [[[1.0], [2.0], [3.0]]], + dtype=torch.float64, + requires_grad=True, + ) + tensor = _tensor_map_with_components(values, ["xyz"]) + + result = decompose_quantity(source_name, tensor) + + assert result.keys.names == ["o3_lambda", "o3_sigma"] + assert result.keys.values.tolist() == [[1, 1]] + assert result.block().components == [_o3_mu_labels(1, values.device)] + assert torch.equal( + result.block().values, + torch.tensor([[[2.0], [3.0], [1.0]]], dtype=torch.float64), + ) + + result.block().values.sum().backward() + assert torch.equal(values.grad, torch.ones_like(values)) + + +def test_decompose_quantity_non_conservative_stress_combines_irreps(): + """Stress returns l=0, l=1 (pseudovector), and l=2 blocks.""" + values = torch.zeros((2, 3, 3, 1), dtype=torch.float64) + values[0, :, :, 0] = torch.eye(3, dtype=torch.float64) + values[1, 0, 1, 0] = 2.0 + values[1, 1, 0, 0] = -2.0 + tensor = _tensor_map_with_components(values, ["xyz_1", "xyz_2"]) + + result = decompose_quantity("non_conservative_stress/direct", tensor) + + assert result.keys.names == ["o3_lambda", "o3_sigma"] + assert result.keys.values.tolist() == [[0, 1], [1, -1], [2, 1]] + block_l0 = result.block({"o3_lambda": 0, "o3_sigma": 1}) + block_l1 = result.block({"o3_lambda": 1, "o3_sigma": -1}) + block_l2 = result.block({"o3_lambda": 2, "o3_sigma": 1}) + assert block_l0.components == [_o3_mu_labels(0, values.device)] + assert block_l1.components == [_o3_mu_labels(1, values.device)] + assert block_l2.components == [_o3_mu_labels(2, values.device)] + assert torch.allclose( + block_l0.values, + torch.tensor([[[3.0**0.5]], [[0.0]]], dtype=torch.float64), + rtol=0.0, + atol=1.0e-12, + ) + # antisymmetric part of the second matrix: axial vector (0, 0, -2) + expected_l1 = torch.zeros((2, 3, 1), dtype=torch.float64) + expected_l1[1, 1, 0] = -2.0 * 2.0**0.5 + assert torch.allclose(block_l1.values, expected_l1, rtol=0.0, atol=1.0e-12) + assert torch.equal(block_l2.values, torch.zeros((2, 5, 1), dtype=torch.float64)) + for block in (block_l0, block_l1, block_l2): + assert block.samples == tensor.block().samples + assert block.properties == tensor.block().properties + + +def test_decompose_quantity_does_not_infer_custom_cartesian_semantics(): + """A generic 3x3 output should pass through unchanged.""" + tensor = _tensor_map_with_components( + torch.rand((1, 3, 3, 1), dtype=torch.float64), + ["xyz_1", "xyz_2"], + ) + + result = decompose_quantity("mtt::custom", tensor) + + mts.equal_raise(result, tensor) + + +def test_forward_rejects_outputs_with_attached_gradients(): + """The wrapper should not silently discard explicit TensorBlock gradients.""" + + class _AttachedGradientModel(torch.nn.Module): + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels], + ) -> Dict[str, TensorMap]: + properties = Labels.range("property", 1) + block = TensorBlock( + values=torch.ones((len(systems), 1), dtype=torch.float64), + samples=Labels.range("system", len(systems)), + components=[], + properties=properties, + ) + block.add_gradient( + "positions", + TensorBlock( + values=torch.ones((1, 3, 1), dtype=torch.float64), + samples=Labels("sample", torch.tensor([[0]], dtype=torch.int64)), + components=[Labels.range("xyz", 3)], + properties=properties, + ), + ) + return { + "energy": TensorMap( + Labels("_", torch.tensor([[0]], dtype=torch.int64)), + [block], + ) + } + + model = SymmetrizedModel( + _AttachedGradientModel(), + max_angular_momentum_target=0, + max_angular_momentum_grid=2, + ) + + message = ( + "underlying output 'energy' contains unsupported explicit gradient 'positions'" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + model( + [_forward_test_system([[1.0, 2.0, 3.0]])], + {"energy": ModelOutput(sample_kind="system")}, + None, + )