From 9a4bcdb92f5a4f03fdbac8c93e39701daa0626d9 Mon Sep 17 00:00:00 2001 From: Michelangelo Domina Date: Thu, 23 Jul 2026 00:48:10 +0200 Subject: [PATCH 01/18] feat(torch): add reviewed SymmetrizedModel core --- .../torch/symmetrized_model/_decompose.py | 212 ++ .../torch/symmetrized_model/_model.py | 1137 ++++++++++ .../torch/symmetrized_model/_projections.py | 235 +++ .../torch/symmetrized_model/_quadrature.py | 169 ++ .../torch/symmetrized_model/_utils.py | 151 ++ .../symmetrized_model/_wigner_storage.py | 76 + .../tests/symmetrized_model.py | 1870 +++++++++++++++++ 7 files changed, 3850 insertions(+) create mode 100644 python/metatomic_torch/metatomic/torch/symmetrized_model/_decompose.py create mode 100644 python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py create mode 100644 python/metatomic_torch/metatomic/torch/symmetrized_model/_projections.py create mode 100644 python/metatomic_torch/metatomic/torch/symmetrized_model/_quadrature.py create mode 100644 python/metatomic_torch/metatomic/torch/symmetrized_model/_utils.py create mode 100644 python/metatomic_torch/metatomic/torch/symmetrized_model/_wigner_storage.py create mode 100644 python/metatomic_torch/tests/symmetrized_model.py diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_decompose.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_decompose.py new file mode 100644 index 00000000..50daca92 --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_decompose.py @@ -0,0 +1,212 @@ +import math +from typing import List + +import torch +from metatensor.torch import Labels, TensorBlock, TensorMap + + +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 _symmetric_matrices_to_spherical( + values: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return orthonormal l=0 and l=2 components of the symmetric matrix part.""" + l0 = (values[:, 0, 0, :] + values[:, 1, 1, :] + values[:, 2, 2, :]).unsqueeze( + 1 + ) / math.sqrt(3.0) + + sqrt_two = math.sqrt(2.0) + 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, l2 + + +def _decompose_output( + source_name: str, + tensor: TensorMap, +) -> TensorMap: + """Decompose standard outputs for variance and character projection.""" + quantity = source_name.split("/", 1)[0] + is_energy = quantity in ( + "energy", + "energy_ensemble", + "energy_uncertainty", + ) + is_force = quantity == "non_conservative_force" + is_stress = quantity == "non_conservative_stress" + if not (is_energy or is_force or is_stress): + return tensor + + for block in tensor.blocks(): + if len(block.gradients_list()) != 0: + raise ValueError( + "O(3) diagnostic decomposition does not support gradients " + "attached to '" + source_name + "'" + ) + + if is_energy: + energy_blocks: List[TensorBlock] = [] + for block in tensor.blocks(): + if len(block.components) != 0: + raise ValueError("energy-like outputs must not have components") + energy_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), + energy_blocks, + ) + + elif is_force: + force_blocks: List[TensorBlock] = [] + for block in tensor.blocks(): + if ( + len(block.components) != 1 + or block.components[0].names != ["xyz"] + or len(block.components[0]) != 3 + ): + raise ValueError( + "non_conservative_force must have one 'xyz' component axis " + "of size 3" + ) + force_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), + force_blocks, + ) + + else: + blocks_l0: List[TensorBlock] = [] + blocks_l2: List[TensorBlock] = [] + for block in tensor.blocks(): + if ( + len(block.components) != 2 + or block.components[0].names != ["xyz_1"] + or block.components[1].names != ["xyz_2"] + or len(block.components[0]) != 3 + or len(block.components[1]) != 3 + ): + raise ValueError( + "non_conservative_stress must have 'xyz_1' and 'xyz_2' " + "component axes of size 3" + ) + + values_l0, values_l2 = _symmetric_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_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) + keys_l2 = _add_o3_irrep_to_keys(tensor.keys, 2, 1) + result = TensorMap( + Labels( + list(keys_l0.names), + torch.cat([keys_l0.values, keys_l2.values], dim=0), + ), + blocks_l0 + 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/symmetrized_model/_model.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py new file mode 100644 index 00000000..6be842ca --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py @@ -0,0 +1,1137 @@ +from typing import Dict, List, Optional, Tuple + +import metatensor.torch as mts +import torch +from metatensor.torch import Labels, TensorBlock, TensorMap + +from metatomic.torch import ( + ModelInterface, + ModelOutput, + System, + register_autograd_neighbors, +) + +from ..o3._tranformations import ( + _max_o3_lambda_in_tensor, + _transform_tensor_with_precomputed_matrices, +) +from ._decompose import _decompose_output +from ._projections import ( + _character_projection_coefficients_from_batch, + _character_projection_tensormap_from_cosets, +) +from ._quadrature import _choose_quadrature, get_rotation_quadrature +from ._utils import ( + _group_samples_by_rotated_copy, + _map_selected_atoms_to_rotated_copies, + _restore_input_system_to_samples, + _validate_integer, +) +from ._wigner_storage import ( + _build_packed_wigner_matrices, + _wigner_matrices_for_lambda, +) + + +_DEFAULT_MAX_WIGNER_STORAGE_BYTES = 64 * 1024 * 1024 # 64 MiB + + +def _transform_system_geometry_batch( + system: System, + matrices: torch.Tensor, +) -> List[System]: + """Transform System geometry and neighbor lists with internal O(3) matrices.""" + 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 != system.positions.dtype + or matrices.device != system.positions.device + ): + raise ValueError("system and matrices must have the same dtype and device") + + if matrices.size(0) == 1: + positions = (system.positions @ matrices[0].transpose(0, 1)).unsqueeze(0) + cells = (system.cell @ matrices[0].transpose(0, 1)).unsqueeze(0) + else: + 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( + types=system.types, + positions=positions[index], + cell=cells[index], + pbc=system.pbc, + ) + ) + + for options in system.known_neighbor_lists(): + neighbors = system.get_neighbor_list(options) + source_values = neighbors.values.detach().squeeze(-1) + if matrices.size(0) == 1: + neighbor_values = (source_values @ matrices[0].transpose(0, 1)).unsqueeze(0) + else: + 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, + ) + + return transformed_systems + + +def _check_o3_lambda_limit( + tensor: TensorMap, + tensor_description: str, + max_o3_lambda: 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_o3_lambda: + raise ValueError( + tensor_description + + " contains o3_lambda=" + + str(tensor_max_o3_lambda) + + ", exceeding " + + limit_name + + "=" + + str(max_o3_lambda) + ) + + +def _transform_system_batch( + system: System, + matrices: torch.Tensor, + wigner_matrices: List[torch.Tensor], + max_o3_lambda_input: int, + is_improper: bool, +) -> List[System]: + """Transform a System batch, including its custom TensorMap data.""" + data_names = system.known_data() + for data_name in data_names: + _check_o3_lambda_limit( + system.get_data(data_name), + "custom input '" + data_name + "'", + max_o3_lambda_input, + "max_o3_lambda_input", + ) + + transformed_systems = _transform_system_geometry_batch(system, matrices) + if len(data_names) == 0: + return transformed_systems + + for index in range(len(transformed_systems)): + wigner_matrices_for_copy: List[torch.Tensor] = [] + for rank_matrices in wigner_matrices: + wigner_matrices_for_copy.append(rank_matrices[index : index + 1]) + + for data_name in data_names: + transformed_systems[index].add_data( + data_name, + _transform_tensor_with_precomputed_matrices( + system.get_data(data_name), + matrices[index : index + 1], + wigner_matrices_for_copy, + is_improper, + ), + ) + + return transformed_systems + + +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: + source_name = requested_name + calculation = "average" + + if len(source_name) == 0: + raise ValueError( + "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.""" + 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( + "all requests derived from '" + + source_name + + "' must use the same 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 _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 reference-centered weighted moments.""" + 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: + 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") + centered_values = values - reference_values.unsqueeze(0) + + # Any proper/improper weight split is applied by the caller. + 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 _join_per_system_tensormaps(tensors: List[TensorMap]) -> TensorMap: + """Join one TensorMap per input system along their sample axes.""" + if len(tensors) == 0: + raise ValueError("expected at least one per-system TensorMap") + + keys = tensors[0].keys + different_keys = "error" + for index in range(1, len(tensors)): + if tensors[index].keys != keys: + different_keys = "union" + break + + return mts.join(tensors, "samples", different_keys=different_keys) + + +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_o3_lambda_grid: int, +) -> TensorMap: + """Clamp round-off negatives and reject invalid or materially negative values.""" + blocks: List[TensorBlock] = [] + for key, block in tensor.items(): + scale_values = scale.block(key).values + invalid = ( + (~torch.isfinite(block.values)) + | (~torch.isfinite(scale_values)) + | (scale_values < 0) + ) + if bool(torch.any(invalid).item()): + raise ValueError(f"O(3) {quantity} or its round-off scale is invalid") + + # 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") + + 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_o3_lambda_grid " + f"above {max_o3_lambda_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_o3_lambda_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_o3_lambda_grid=max_o3_lambda_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): + r""" + Wrap a model with finite-quadrature O(3) averaging and equivariance + diagnostics. + + For a target representation :math:`\rho_\alpha`, define the model response + transformed back to the input frame as + + .. math:: + + z_\alpha(g;x) = \rho_\alpha(g^{-1}) f(gx). + + An ordinary requested output is the normalized Haar average + + .. math:: + + \Pi_\alpha(f,x) + = \int_{\mathrm{O}(3)} z_\alpha(g;x)\,\mathrm{d}\mu(g). + + The integrals are approximated by evaluating the underlying model on batches of + proper and improper transformations. For a TensorMap block with :math:`d` + component entries, ``o3::variance::`` returns + + .. 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] + = \frac{A_\alpha(f,x)^2}{d}. + + Here, :math:`A_\alpha` is the component-summed equivariance error defined in the + reference article. The returned value is instead a component-averaged variance for + every retained sample and property: this class neither takes its square root nor + aggregates it over samples. + + Character projections act on the direct response :math:`u(g;x) = f(gx)`. For a + character sector :math:`\beta=(\lambda,\sigma)` with + :math:`d_\beta=2\lambda+1`, the corresponding 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). + + Writing an O(3) operation as :math:`\Phi(R,s)`, with :math:`s=+1` for a proper + rotation and :math:`s=-1` for an improper operation, the character convention is + + .. math:: + + \chi_{\lambda,\sigma}(\Phi(R,s)) + = \left[\sigma(-1)^\lambda\right]^{(1-s)/2} + \operatorname{tr} D^\lambda(R). + + Requests named ``o3::character_projection::`` return the unnormalized + contributions to :math:`B_\beta`, labeled by ``chi_lambda`` and ``chi_sigma``. + Target component axes are retained; summing over them recovers the full + component norm in the equation above. + + The deterministic quadrature is exact only when it resolves the angular dependence + of the transformed model response. For unrestricted responses, convergence must be + checked by increasing ``max_o3_lambda_grid``. ``batch_size`` changes how many + transformed systems are evaluated in one model call, but does not change the grid + or the result. + + Rotation matrices, quadrature weights, and Wigner-D matrices are stored as float64 + buffers so they follow ordinary module device movement and serialization. The + packed Wigner-D allocation is checked against ``max_wigner_storage_bytes`` before + it is created. + + :param model: underlying :py:class:`ModelInterface`. The :py:meth:`wrap` method + obtains this module from :py:attr:`AtomisticModel.module`. + :param max_o3_lambda_target: largest ``o3_lambda`` accepted on an + already-spherical output component axis. Cartesian outputs are not limited by + this value. + :param max_o3_lambda_input: largest ``o3_lambda`` accepted on an + already-spherical component axis in custom System data. The default of zero + still allows Cartesian custom inputs. + :param max_o3_lambda_character: largest character sector included in character + projections. ``None`` disables character-projection outputs; zero enables the + scalar character sector only. + :param batch_size: positive number of transformed systems evaluated in one call to + ``model``. The default is 32. + :param max_o3_lambda_grid: quadrature integration degree. If ``None``, use the + larger of ``2 * max_o3_lambda_target + 1`` and + ``2 * max_o3_lambda_character`` when character projections are enabled. An + explicit value must be non-negative and no larger than the highest available + Lebedev order, 131. + :param max_wigner_storage_bytes: maximum number of bytes used by the serialized + packed Wigner-D matrices. Construction fails before allocation when this limit + would be exceeded. The default is 64 MiB. + """ + + max_o3_lambda_character: Optional[int] + + def __init__( + self, + model: ModelInterface, + max_o3_lambda_target: int, + max_o3_lambda_input: int = 0, + max_o3_lambda_character: Optional[int] = None, + batch_size: int = 32, + max_o3_lambda_grid: Optional[int] = None, + max_wigner_storage_bytes: int = _DEFAULT_MAX_WIGNER_STORAGE_BYTES, + ): + super().__init__() + + self._model = model + self.max_o3_lambda_target = _validate_integer( + "max_o3_lambda_target", max_o3_lambda_target, 0 + ) + self.max_o3_lambda_input = _validate_integer( + "max_o3_lambda_input", max_o3_lambda_input, 0 + ) + if max_o3_lambda_character is not None: + max_o3_lambda_character = _validate_integer( + "max_o3_lambda_character", max_o3_lambda_character, 0 + ) + self.max_o3_lambda_character = max_o3_lambda_character + self.batch_size = _validate_integer("batch_size", batch_size, 1) + self.max_wigner_storage_bytes = _validate_integer( + "max_wigner_storage_bytes", max_wigner_storage_bytes, 1 + ) + + if max_o3_lambda_grid is None: + max_o3_lambda_grid = 2 * self.max_o3_lambda_target + 1 + if self.max_o3_lambda_character is not None: + max_o3_lambda_grid = max( + max_o3_lambda_grid, + 2 * self.max_o3_lambda_character, + ) + else: + max_o3_lambda_grid = _validate_integer( + "max_o3_lambda_grid", max_o3_lambda_grid, 0 + ) + if ( + self.max_o3_lambda_character is not None + and max_o3_lambda_grid < 2 * self.max_o3_lambda_character + ): + raise ValueError( + "max_o3_lambda_grid must be at least twice max_o3_lambda_character" + ) + self.max_o3_lambda_grid = max_o3_lambda_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 == "mps": + raise ValueError("SymmetrizedModel supports CPU and CUDA execution") + + lebedev_order, n_rotations = _choose_quadrature(self.max_o3_lambda_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_o3_lambda_wigner = max( + self.max_o3_lambda_input, + self.max_o3_lambda_target, + 0 if self.max_o3_lambda_character is None else self.max_o3_lambda_character, + ) + n_wigner_elements_per_matrix = ( + (max_o3_lambda_wigner + 1) + * (2 * max_o3_lambda_wigner + 1) + * (2 * max_o3_lambda_wigner + 3) + // 3 + ) + required_wigner_storage_bytes = ( + len(rotation_matrices) + * n_wigner_elements_per_matrix + * rotation_matrices.element_size() + ) + if required_wigner_storage_bytes > self.max_wigner_storage_bytes: + raise ValueError( + "packed Wigner-D matrices require " + + str(required_wigner_storage_bytes) + + " bytes, exceeding max_wigner_storage_bytes=" + + str(self.max_wigner_storage_bytes) + ) + packed_wigner_matrices = _build_packed_wigner_matrices( + rotation_matrices, + max_o3_lambda_wigner, + ) + + self.register_buffer("_rotation_matrices", rotation_matrices) + self.register_buffer("_rotation_weights", rotation_weights) + self.register_buffer("_packed_wigner_matrices", packed_wigner_matrices) + + 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: + return torch.jit.annotate(Dict[str, TensorMap], {}) + if len(systems) == 0: + 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 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_o3_lambda_character is None + ): + raise ValueError( + "max_o3_lambda_character must be set to request character projections" + ) + + source_outputs = torch.jit.annotate(Dict[str, ModelOutput], {}) + for source_name in source_sample_kinds: + if source_name in average_names: + requested_name = average_names[source_name] + elif source_name in variance_names: + requested_name = variance_names[source_name] + else: + requested_name = character_projection_names[source_name] + source_outputs[source_name] = outputs[requested_name] + + per_output_results = torch.jit.annotate( + Dict[str, List[TensorMap]], + {}, + ) + for requested_name in outputs: + per_output_results[requested_name] = torch.jit.annotate(List[TensorMap], []) + + 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: + if requested_name not in system_results: + raise ValueError( + "SymmetrizedModel did not produce requested output '" + + requested_name + + "'" + ) + per_output_results[requested_name].append( + system_results[requested_name] + ) + + results = torch.jit.annotate(Dict[str, TensorMap], {}) + for requested_name in outputs: + results[requested_name] = _join_per_system_tensormaps( + per_output_results[requested_name] + ) + 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") + if ( + self._rotation_matrices.dtype != torch.float64 + or self._rotation_weights.dtype != torch.float64 + or self._packed_wigner_matrices.dtype != torch.float64 + ): + raise ValueError("SymmetrizedModel integration buffers must remain float64") + 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" + ) + + character_max = 0 + configured_character_max = self.max_o3_lambda_character + if configured_character_max is not None: + character_max = configured_character_max + + average_references = torch.jit.annotate(Dict[str, TensorMap], {}) + average_first_moments = torch.jit.annotate(Dict[str, TensorMap], {}) + variance_references = torch.jit.annotate(Dict[str, TensorMap], {}) + variance_first_moments = torch.jit.annotate(Dict[str, TensorMap], {}) + variance_second_moments = torch.jit.annotate(Dict[str, TensorMap], {}) + variance_absolute_second_moments = torch.jit.annotate( + Dict[str, TensorMap], + {}, + ) + proper_character_coefficients = torch.jit.annotate( + Dict[str, TensorMap], + {}, + ) + improper_character_coefficients = torch.jit.annotate( + Dict[str, TensorMap], + {}, + ) + + n_rotations = self._rotation_matrices.size(0) + needs_backrotation = len(average_names) != 0 or len(variance_names) != 0 + 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, + ) + + input_wigner_matrices: List[torch.Tensor] = [] + for o3_lambda in range(self.max_o3_lambda_input + 1): + input_wigner_matrices.append( + _wigner_matrices_for_lambda( + self._packed_wigner_matrices, + n_rotations, + o3_lambda, + )[batch_start:batch_stop].to( + dtype=work_dtype, + device=work_device, + ) + ) + + inverse_target_wigner_matrices: List[torch.Tensor] = [] + if needs_backrotation: + for o3_lambda in range(self.max_o3_lambda_target + 1): + inverse_target_wigner_matrices.append( + _wigner_matrices_for_lambda( + self._packed_wigner_matrices, + n_rotations, + o3_lambda, + )[batch_start:batch_stop].transpose(1, 2) + ) + + 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( + _wigner_matrices_for_lambda( + self._packed_wigner_matrices, + n_rotations, + chi_lambda, + )[batch_start:batch_stop].transpose(1, 2) + ) + + for coset_index in range(2): + is_improper = coset_index == 1 + sign = -1.0 if is_improper else 1.0 + matrices = (sign * proper_matrices).to( + dtype=work_dtype, + device=work_device, + ) + transformed_systems = _transform_system_batch( + system, + matrices, + input_wigner_matrices, + self.max_o3_lambda_input, + is_improper, + ) + 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 '" + + source_name + + "'" + ) + for returned_name in raw_outputs: + if returned_name not in source_outputs: + raise ValueError( + "underlying model returned unrequested output '" + + returned_name + + "'" + ) + + inverse_matrices = (sign * proper_matrices).transpose(1, 2) + 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( + "underlying output '" + + source_name + + "' contains unsupported explicit gradient '" + + gradient_names[0] + + "'" + ) + + tensor = raw_tensor.to( + dtype=torch.float64, + device=work_device, + ) + if source_name in average_names or source_name in variance_names: + _check_o3_lambda_limit( + tensor, + "output '" + source_name + "'", + self.max_o3_lambda_target, + "max_o3_lambda_target", + ) + backrotated = _transform_tensor_with_precomputed_matrices( + tensor, + inverse_matrices, + inverse_target_wigner_matrices, + is_improper, + ) + + 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_output( + 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_output(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 = torch.jit.annotate(Dict[str, TensorMap], {}) + for source_name, requested_name in average_names.items(): + if ( + source_name not in average_references + or source_name not in average_first_moments + ): + raise RuntimeError("average accumulation is incomplete") + 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(): + if ( + source_name not in variance_references + or source_name not in variance_first_moments + or source_name not in variance_second_moments + or source_name not in variance_absolute_second_moments + ): + raise RuntimeError("variance accumulation is incomplete") + 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_o3_lambda_grid=self.max_o3_lambda_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(): + if ( + source_name not in proper_character_coefficients + or source_name not in improper_character_coefficients + ): + raise RuntimeError("character-projection accumulation is incomplete") + 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/symmetrized_model/_projections.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_projections.py new file mode 100644 index 00000000..98bd147f --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_projections.py @@ -0,0 +1,235 @@ +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.""" + if ( + values.dim() < 3 + or weights.dim() != 1 + or inverse_wigner_matrices.dim() != 3 + or weights.size(0) == 0 + or values.size(0) != weights.size(0) + or inverse_wigner_matrices.size(0) != weights.size(0) + or inverse_wigner_matrices.size(1) != inverse_wigner_matrices.size(2) + ): + raise ValueError("incompatible values, weights, or Wigner-matrix shapes") + + 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 + if ( + chi_lambda < 0 + or proper_coefficients.dim() < 3 + or improper_coefficients.size() != proper_coefficients.size() + or proper_coefficients.size(1) != dimension + or proper_coefficients.size(2) != dimension + ): + raise ValueError("coefficient shapes do not match chi_lambda") + + 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 every character rank 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) sectors.""" + if proper_coefficients.keys != improper_coefficients.keys: + raise ValueError( + "proper and improper character coefficients must have same keys" + ) + + 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) + proper_components = proper_block.components + improper_components = improper_block.components + components_match = len(proper_components) == len(improper_components) + if components_match: + for component_index in range(len(proper_components)): + if ( + proper_components[component_index] + != improper_components[component_index] + ): + components_match = False + if ( + proper_block.samples != improper_block.samples + or not components_match + or proper_block.properties != improper_block.properties + ): + raise ValueError( + "proper and improper character coefficients must have same metadata" + ) + if ( + len(proper_block.components) < 2 + or proper_block.components[0].names != ["chi_m"] + or proper_block.components[1].names != ["chi_n"] + ): + raise ValueError("character coefficient component metadata is invalid") + + 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/symmetrized_model/_quadrature.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_quadrature.py new file mode 100644 index 00000000..1d7a9a5f --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_quadrature.py @@ -0,0 +1,169 @@ +from typing import Tuple + +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(L_max: int) -> Tuple[int, int]: + """ + Choose a Lebedev quadrature order and number of in-plane rotations to integrate + spherical harmonics up to degree ``L_max``. + + :param L_max: maximum spherical harmonic degree + :return: (lebedev_order, n_inplane_rotations) + """ + L_max = _validate_integer("L_max", L_max, 0) + if L_max > _LEBEDEV_ORDERS[-1]: + raise ValueError( + f"the requested quadrature degree L_max={L_max} exceeds the largest " + f"available Lebedev order ({_LEBEDEV_ORDERS[-1]})" + ) + # pick smallest order >= L_max + n = min(o for o in _LEBEDEV_ORDERS if o >= L_max) + # minimal gamma count + K = L_max + 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 +) -> "Rotation": # noqa: F821 (scipy is imported lazily) + """ + 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. SO(3) contains + proper rotations with determinant +1, while O(3) also contains improper + orthogonal transformations with determinant -1. 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 + :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. ``lebedev_order`` must be one of the orders + supported by ``scipy.integrate.lebedev_rule``. + """ + 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/symmetrized_model/_utils.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_utils.py new file mode 100644 index 00000000..4ef8558d --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_utils.py @@ -0,0 +1,151 @@ +import operator +from typing import List, Optional, Tuple + +import numpy as np +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, np.bool_)) or ( + isinstance(value, torch.Tensor) and value.dtype == torch.bool + ): + raise TypeError(f"{name} must be an integer, not a boolean") + try: + integer_value = int(operator.index(value)) + except TypeError as error: + raise TypeError( + f"{name} must be an integer, got {type(value).__name__}" + ) from error + if integer_value < minimum: + qualifier = "positive" if minimum == 1 else "non-negative" + 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." + ) + + # A single copy is already grouped; avoid sorting the common batch-size-one case. + if n_rotated_copies == 1: + return ( + block.values.unsqueeze(0), + sample_names[:system_column] + sample_names[system_column + 1 :], + sample_values_without_system, + ) + + 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/symmetrized_model/_wigner_storage.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_wigner_storage.py new file mode 100644 index 00000000..c841cb92 --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_wigner_storage.py @@ -0,0 +1,76 @@ +import torch + +from ..o3 import O3Transformation +from ._utils import _validate_integer + + +def _build_packed_wigner_matrices( + matrices: torch.Tensor, + max_o3_lambda: int, +) -> torch.Tensor: + """Build and pack proper Wigner-D matrices through ``max_o3_lambda``.""" + max_o3_lambda = _validate_integer("max_o3_lambda", max_o3_lambda, 0) + 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 not in (torch.float32, torch.float64): + raise TypeError("matrices must use float32 or float64") + + output_device = matrices.device + output_dtype = matrices.dtype + calculation_matrices = matrices.detach().to(device="cpu") + n_matrices = matrices.size(0) + n_elements_per_matrix = ( + (max_o3_lambda + 1) * (2 * max_o3_lambda + 1) * (2 * max_o3_lambda + 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)): + transformation = O3Transformation(matrix, max_o3_lambda) + for o3_lambda in range(max_o3_lambda + 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_( + transformation.wigner_D_matrix(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.""" + if packed.dim() != 1: + raise ValueError("packed Wigner-D storage must be one-dimensional") + if n_matrices <= 0: + raise ValueError("n_matrices must be positive") + if o3_lambda < 0: + raise ValueError("o3_lambda must be non-negative") + + 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/symmetrized_model.py b/python/metatomic_torch/tests/symmetrized_model.py new file mode 100644 index 00000000..e3f0d4af --- /dev/null +++ b/python/metatomic_torch/tests/symmetrized_model.py @@ -0,0 +1,1870 @@ +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 ModelOutput, NeighborListOptions, System +from metatomic.torch.o3 import O3Transformation, transform_system +from metatomic.torch.symmetrized_model._decompose import ( + _add_o3_irrep_to_keys, + _cartesian_vectors_to_spherical, + _decompose_output, + _o3_mu_labels, + _symmetric_matrices_to_spherical, +) +from metatomic.torch.symmetrized_model._model import ( + SymmetrizedModel, + _clamp_roundoff_negative_diagnostic, + _component_norm_squared, + _group_output_requests, + _join_per_system_tensormaps, + _mean_variance_over_components, + _parse_output_request, + _reduce_weighted_centered_batch, + _transform_system_batch, + _transform_system_geometry_batch, + _variance_from_centered_moments, +) +from metatomic.torch.symmetrized_model._projections import ( + _character_projection_coefficients_from_rotation_batch, + _character_projections_from_proper_and_improper_coefficients, +) +from metatomic.torch.symmetrized_model._quadrature import ( + _choose_quadrature, + _rotations_from_euler_angles, + get_euler_angles_quadrature, + get_rotation_quadrature, +) +from metatomic.torch.symmetrized_model._utils import ( + _group_samples_by_rotated_copy, + _map_selected_atoms_to_rotated_copies, + _restore_input_system_to_samples, +) +from metatomic.torch.symmetrized_model._wigner_storage import ( + _build_packed_wigner_matrices, + _wigner_matrices_for_lambda, +) + + +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 _system_with_neighbor_lists(dtype: torch.dtype) -> System: + """Create a test system with populated and empty neighbor lists.""" + positions = torch.tensor( + [[0.2, -0.1, 0.3], [1.1, 0.7, -0.4], [-0.3, 0.6, 1.2]], + dtype=dtype, + ) + cell = torch.tensor( + [[2.5, 0.1, 0.0], [0.0, 2.2, 0.2], [0.1, 0.0, 2.7]], + dtype=dtype, + ) + system = System( + types=torch.tensor([6, 1, 8]), + positions=positions, + cell=cell, + pbc=torch.tensor([True, True, True]), + ) + + samples = Labels( + [ + "first_atom", + "second_atom", + "cell_shift_a", + "cell_shift_b", + "cell_shift_c", + ], + torch.tensor([[0, 1, 0, 0, 0], [1, 2, 1, 0, 0]]), + ) + components = [Labels.range("xyz", 3)] + properties = Labels.range("distance", 1) + system.add_neighbor_list( + NeighborListOptions(3.0, False, True, "populated"), + TensorBlock( + values=torch.stack( + [ + positions[1] - positions[0], + positions[2] - positions[1] + cell[0], + ] + ).unsqueeze(-1), + samples=samples, + components=components, + properties=properties, + ), + ) + system.add_neighbor_list( + NeighborListOptions(1.0, True, False, "empty"), + TensorBlock( + values=torch.empty((0, 3, 1), dtype=dtype), + samples=Labels( + list(samples.names), + torch.empty((0, len(samples.names)), dtype=torch.int64), + ), + components=components, + properties=properties, + ), + ) + return system + + +class TestSystemGeometryBatch: + """Test batched O(3) transformation of System geometry.""" + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) + @pytest.mark.parametrize("n_matrices", [1, 3]) + def test_matches_individual_o3_transformations(self, dtype, n_matrices): + """Batched geometry should match one transformation at a time.""" + proper = torch.tensor( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], + dtype=dtype, + ) + if n_matrices == 1: + matrices = proper.unsqueeze(0) + else: + matrices = torch.stack([torch.eye(3, dtype=dtype), proper, -proper]) + system = _system_with_neighbor_lists(dtype) + + transformed = _transform_system_geometry_batch(system, matrices) + + assert len(transformed) == len(matrices) + for matrix, actual in zip(matrices, transformed, strict=True): + expected = transform_system( + system, + O3Transformation(matrix, max_angular_momentum=0), + ) + assert torch.equal(actual.positions, expected.positions) + assert torch.equal(actual.cell, expected.cell) + assert torch.equal(actual.types, expected.types) + assert torch.equal(actual.pbc, expected.pbc) + assert actual.known_neighbor_lists() == expected.known_neighbor_lists() + for options in expected.known_neighbor_lists(): + actual_neighbors = actual.get_neighbor_list(options) + expected_neighbors = expected.get_neighbor_list(options) + assert torch.equal(actual_neighbors.values, expected_neighbors.values) + assert actual_neighbors.samples == expected_neighbors.samples + assert actual_neighbors.components == expected_neighbors.components + assert actual_neighbors.properties == expected_neighbors.properties + + def test_preserves_neighbor_autograd(self): + """Rotated neighbor vectors should differentiate through positions and cell.""" + positions = torch.tensor( + [[0.2, -0.1, 0.3], [1.1, 0.7, -0.4]], + dtype=torch.float64, + requires_grad=True, + ) + cell = torch.tensor( + [[2.5, 0.1, 0.0], [0.0, 2.2, 0.2], [0.1, 0.0, 2.7]], + dtype=torch.float64, + requires_grad=True, + ) + system = System( + types=torch.tensor([6, 1]), + positions=positions, + cell=cell, + pbc=torch.tensor([True, True, True]), + ) + cell_shift = torch.tensor([1.0, -1.0, 0.0], dtype=torch.float64) + neighbor_vector = positions[1] - positions[0] + cell_shift @ cell + options = NeighborListOptions(4.0, False, True) + system.add_neighbor_list( + options, + TensorBlock( + values=neighbor_vector.reshape(1, 3, 1), + samples=Labels( + [ + "first_atom", + "second_atom", + "cell_shift_a", + "cell_shift_b", + "cell_shift_c", + ], + torch.tensor([[0, 1, 1, -1, 0]]), + ), + components=[Labels.range("xyz", 3)], + properties=Labels.range("distance", 1), + ), + ) + proper = torch.tensor( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], + dtype=torch.float64, + ) + matrices = torch.stack([proper, -proper]) + + transformed = _transform_system_geometry_batch(system, matrices) + loss = sum( + transformed_system.get_neighbor_list(options).values.square().sum() + for transformed_system in transformed + ) + position_gradient, cell_gradient = torch.autograd.grad( + loss, + (positions, cell), + ) + + vector_gradient = 2 * len(matrices) * neighbor_vector.detach() + assert torch.allclose( + position_gradient, + torch.stack([-vector_gradient, vector_gradient]), + ) + assert torch.allclose( + cell_gradient, + torch.outer(cell_shift, vector_gradient), + ) + + def test_rejects_invalid_matrix_batches(self): + """Matrix batches should have a non-empty shape and match the System.""" + system = _system_with_neighbor_lists(torch.float64) + invalid_shapes = [(3, 3), (0, 3, 3), (2, 2, 3), (2, 3, 2)] + for shape in invalid_shapes: + with pytest.raises(ValueError, match="shape \\(N, 3, 3\\)"): + _transform_system_geometry_batch( + system, + torch.empty(shape, dtype=torch.float64), + ) + + with pytest.raises(ValueError, match="same dtype and device"): + _transform_system_geometry_batch( + system, + torch.eye(3, dtype=torch.float32).unsqueeze(0), + ) + + def test_is_scriptable(self): + """The batched geometry transformation should compile and execute.""" + scripted = torch.jit.script(_transform_system_geometry_batch) + system = _system_with_neighbor_lists(torch.float64) + transformed = scripted( + system, + torch.eye(3, dtype=torch.float64).unsqueeze(0), + ) + + assert len(transformed) == 1 + assert torch.equal(transformed[0].positions, system.positions) + assert torch.equal(transformed[0].cell, system.cell) + + +class TestSystemBatch: + """Test batched O(3) transformation of complete Systems.""" + + @pytest.mark.parametrize("is_improper", [False, True]) + def test_transforms_spherical_custom_data(self, is_improper): + """Every transformed System should contain the corresponding custom data.""" + proper_matrices = torch.tensor( + [ + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], + [ + [-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, + ) + matrices = -proper_matrices if is_improper else proper_matrices + packed_wigner = _build_packed_wigner_matrices( + proper_matrices, + max_o3_lambda=1, + ) + wigner_matrices = [ + _wigner_matrices_for_lambda( + packed_wigner, + n_matrices=len(matrices), + o3_lambda=o3_lambda, + ) + for o3_lambda in range(2) + ] + + system = System( + types=torch.tensor([6, 8]), + positions=torch.tensor( + [[0.2, -0.1, 0.3], [1.1, 0.7, -0.4]], + dtype=torch.float64, + ), + cell=torch.eye(3, dtype=torch.float64) * 4.0, + pbc=torch.tensor([True, True, True]), + ) + values = torch.tensor( + [[[1.0], [2.0], [3.0]], [[-0.5], [1.5], [0.25]]], + dtype=torch.float64, + requires_grad=True, + ) + system.add_data( + "mtt::field", + TensorMap( + Labels( + ["o3_lambda", "o3_sigma"], + torch.tensor([[1, 1]]), + ), + [ + TensorBlock( + values=values, + samples=Labels.range("atom", 2), + components=[_o3_mu_labels(1, values.device)], + properties=Labels.range("property", 1), + ) + ], + ), + ) + + transformed = torch.jit.script(_transform_system_batch)( + system, + matrices, + wigner_matrices, + max_o3_lambda_input=1, + is_improper=is_improper, + ) + + assert len(transformed) == len(matrices) + for matrix, transformed_system in zip(matrices, transformed, strict=True): + expected_system = transform_system( + system, + O3Transformation(matrix, max_angular_momentum=1), + ) + assert "mtt::field" in transformed_system.known_data() + mts.allclose_raise( + transformed_system.get_data("mtt::field"), + expected_system.get_data("mtt::field"), + rtol=0.0, + atol=1.0e-12, + ) + + loss = sum( + transformed_system.get_data("mtt::field").block().values.square().sum() + for transformed_system in transformed + ) + gradient = torch.autograd.grad(loss, values)[0] + assert torch.allclose( + gradient, + 2 * len(matrices) * values, + rtol=0.0, + atol=1.0e-12, + ) + + def test_input_limit_distinguishes_spherical_from_cartesian(self): + """A zero spherical-rank limit should still allow Cartesian custom data.""" + matrix = torch.tensor( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], + dtype=torch.float64, + ).unsqueeze(0) + packed_wigner = _build_packed_wigner_matrices( + matrix, + max_o3_lambda=0, + ) + wigner_matrices = [ + _wigner_matrices_for_lambda( + packed_wigner, + n_matrices=1, + o3_lambda=0, + ) + ] + system = System( + types=torch.tensor([6]), + positions=torch.tensor([[0.2, -0.1, 0.3]], dtype=torch.float64), + cell=torch.eye(3, dtype=torch.float64) * 4.0, + pbc=torch.tensor([True, True, True]), + ) + cartesian = TensorMap( + Labels("_", torch.tensor([[0]])), + [ + TensorBlock( + values=torch.tensor( + [[[1.0], [2.0], [3.0]]], + dtype=torch.float64, + ), + samples=Labels.range("atom", 1), + components=[Labels.range("xyz", 3)], + properties=Labels.range("property", 1), + ) + ], + ) + system.add_data("mtt::field", cartesian) + scripted_transform = torch.jit.script(_transform_system_batch) + + transformed = scripted_transform( + system, + matrix, + wigner_matrices, + max_o3_lambda_input=0, + is_improper=False, + ) + expected = transform_system( + system, + O3Transformation(matrix[0], max_angular_momentum=0), + ) + mts.allclose_raise( + transformed[0].get_data("mtt::field"), + expected.get_data("mtt::field"), + rtol=0.0, + atol=1.0e-12, + ) + + spherical_system = System( + types=system.types, + positions=system.positions, + cell=system.cell, + pbc=system.pbc, + ) + 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), + ) + ], + ), + ) + with pytest.raises( + torch.jit.Error, + match=( + "custom input 'mtt::field' contains o3_lambda=1, exceeding " + "max_o3_lambda_input=0" + ), + ): + scripted_transform( + spherical_system, + matrix, + wigner_matrices, + max_o3_lambda_input=0, + is_improper=False, + ) + + +class TestCharacterProjections: + """Test construction of character projections from rotated model responses.""" + + @pytest.mark.parametrize("n_samples", [0, 2]) + def test_batch_coefficients_match_rotation_by_rotation_sum(self, n_samples): + """Batching should match summing the weighted rotations individually.""" + torch.manual_seed(7) + n_rotations = 4 + dimension = 3 + values = torch.randn( + (n_rotations, n_samples, 2, 3), + dtype=torch.float64, + ) + weights = torch.tensor( + [0.50, -0.25, 0.30, 0.45], + dtype=torch.float32, + ) + inverse_wigner_matrices = torch.randn( + (n_rotations, dimension, dimension), + dtype=torch.float32, + ) + + coefficients = _character_projection_coefficients_from_rotation_batch( + values, + weights, + inverse_wigner_matrices, + ) + + expected = torch.zeros( + (n_samples, dimension, dimension, 2, 3), + dtype=torch.float64, + ) + for rotation in range(n_rotations): + expected += ( + weights[rotation].to(torch.float64) + * inverse_wigner_matrices[rotation] + .to(torch.float64) + .reshape(1, dimension, dimension, 1, 1) + * values[rotation].reshape(n_samples, 1, 1, 2, 3) + ) + + assert torch.allclose(coefficients, expected, rtol=0.0, atol=1e-12) + + @pytest.mark.parametrize("chi_lambda", [0, 1, 2]) + def test_factorization_matches_all_rotation_pairs(self, chi_lambda): + """The factorization should match summing every pair of rotations.""" + torch.manual_seed(11 + chi_lambda) + n_rotations = 4 + dimension = 2 * chi_lambda + 1 + proper_values = torch.randn( + (n_rotations, 2, 2, 1), + dtype=torch.float64, + requires_grad=True, + ) + improper_values = torch.randn( + (n_rotations, 2, 2, 1), + dtype=torch.float64, + requires_grad=True, + ) + weights = torch.tensor( + [0.50, -0.25, 0.30, 0.45], + dtype=torch.float64, + ) + inverse_wigner_matrices = torch.randn( + (n_rotations, dimension, dimension), + dtype=torch.float64, + ) + proper_coefficients = _character_projection_coefficients_from_rotation_batch( + proper_values, + weights, + inverse_wigner_matrices, + ) + improper_coefficients = _character_projection_coefficients_from_rotation_batch( + improper_values, + weights, + inverse_wigner_matrices, + ) + + sigma_plus, sigma_minus = ( + _character_projections_from_proper_and_improper_coefficients( + proper_coefficients, + improper_coefficients, + chi_lambda, + ) + ) + + expected = [] + for chi_sigma in (1, -1): + combined_values = proper_values + ( + chi_sigma * (-1) ** chi_lambda * improper_values + ) + direct_sum = torch.zeros_like(combined_values[0]) + for first_rotation in range(n_rotations): + for second_rotation in range(n_rotations): + character = torch.sum( + inverse_wigner_matrices[first_rotation] + * inverse_wigner_matrices[second_rotation] + ) + direct_sum += ( + float(dimension) + / 4.0 + * weights[first_rotation] + * weights[second_rotation] + * character + * combined_values[first_rotation] + * combined_values[second_rotation] + ) + expected.append(direct_sum) + + assert torch.allclose(sigma_plus, expected[0], rtol=0.0, atol=1e-12) + assert torch.allclose(sigma_minus, expected[1], rtol=0.0, atol=1e-12) + assert sigma_plus.shape == proper_values.shape[1:] + assert sigma_minus.shape == improper_values.shape[1:] + assert torch.all(sigma_plus >= 0) + assert torch.all(sigma_minus >= 0) + + (sigma_plus.sum() + sigma_minus.sum()).backward() + assert torch.all(torch.isfinite(proper_values.grad)) + assert torch.all(torch.isfinite(improper_values.grad)) + + def test_rejects_mismatched_rotation_counts_and_coefficient_shapes(self): + """Reject unequal rotation counts or proper/improper coefficient shapes.""" + with pytest.raises(ValueError, match="incompatible values"): + _character_projection_coefficients_from_rotation_batch( + torch.zeros((3, 1, 1), dtype=torch.float64), + torch.ones(2, dtype=torch.float64), + torch.ones((3, 1, 1), dtype=torch.float64), + ) + + with pytest.raises(ValueError, match="chi_lambda"): + _character_projections_from_proper_and_improper_coefficients( + torch.zeros((1, 3, 3, 1), dtype=torch.float64), + torch.zeros((2, 3, 3, 1), dtype=torch.float64), + chi_lambda=1, + ) + + def test_is_scriptable(self): + """Both character-projection tensor operations should compile and run.""" + coefficient_function = torch.jit.script( + _character_projection_coefficients_from_rotation_batch + ) + projection_function = torch.jit.script( + _character_projections_from_proper_and_improper_coefficients + ) + values = torch.ones((1, 1, 1), dtype=torch.float64) + coefficients = coefficient_function( + values, + torch.ones(1, dtype=torch.float64), + torch.ones((1, 1, 1), dtype=torch.float64), + ) + sigma_plus, sigma_minus = projection_function( + coefficients, + coefficients, + 0, + ) + + assert sigma_plus.item() == 1.0 + assert sigma_minus.item() == 0.0 + + +class TestWignerStorage: + """Test persistent Wigner-D storage for the quadrature grid.""" + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) + def test_packed_matrices_match_o3(self, dtype): + """Packing and rank views should preserve the public O(3) matrices.""" + proper_rotation = torch.tensor( + [ + [0.0, -1.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 0.0, 1.0], + ], + dtype=dtype, + ) + matrices = torch.stack( + [ + torch.eye(3, dtype=dtype), + -proper_rotation, + ] + ) + max_o3_lambda = 2 + + packed = _build_packed_wigner_matrices(matrices, max_o3_lambda) + + assert packed.dim() == 1 + assert packed.numel() == len(matrices) * sum( + (2 * o3_lambda + 1) ** 2 for o3_lambda in range(max_o3_lambda + 1) + ) + assert packed.dtype == matrices.dtype + assert packed.device == matrices.device + + transformations = [ + O3Transformation(matrix, max_o3_lambda) for matrix in matrices.unbind(0) + ] + for o3_lambda in range(max_o3_lambda + 1): + actual = _wigner_matrices_for_lambda( + packed, + len(matrices), + o3_lambda, + ) + expected = torch.stack( + [ + transformation.wigner_D_matrix(o3_lambda) + for transformation in transformations + ] + ) + assert torch.equal(actual, expected) + + rank_one = _wigner_matrices_for_lambda(packed, len(matrices), 1) + previous = rank_one[0, 0, 0].clone() + rank_one[0, 0, 0] += 1 + assert packed[len(matrices)] == previous + 1 + + def test_builder_rejects_invalid_inputs(self): + """The builder should reject invalid ranks, shapes, and dtypes.""" + matrices = torch.eye(3, dtype=torch.float64).unsqueeze(0) + with pytest.raises(ValueError, match="non-negative"): + _build_packed_wigner_matrices(matrices, -1) + + for shape in ((0, 3, 3), (2, 3, 2)): + with pytest.raises(ValueError, match="shape \\(N, 3, 3\\)"): + _build_packed_wigner_matrices( + torch.empty(shape, dtype=torch.float64), + 1, + ) + + with pytest.raises(TypeError, match="float32 or float64"): + _build_packed_wigner_matrices(matrices.to(torch.float16), 1) + + def test_rank_view_rejects_invalid_inputs(self): + """Rank views should reject invalid storage, counts, and ranks.""" + with pytest.raises(ValueError, match="one-dimensional"): + _wigner_matrices_for_lambda(torch.empty((2, 2)), 1, 0) + with pytest.raises(ValueError, match="n_matrices must be positive"): + _wigner_matrices_for_lambda(torch.empty(1), 0, 0) + with pytest.raises(ValueError, match="o3_lambda must be non-negative"): + _wigner_matrices_for_lambda(torch.empty(1), 1, -1) + with pytest.raises(ValueError, match="exceeds the packed"): + _wigner_matrices_for_lambda(torch.empty(1), 1, 1) + + def test_rank_view_is_scriptable(self): + """The runtime rank accessor should compile and execute in TorchScript.""" + scripted = torch.jit.script(_wigner_matrices_for_lambda) + packed = torch.arange(70, dtype=torch.float64) + + assert torch.equal( + scripted(packed, 2, 2), + _wigner_matrices_for_lambda(packed, 2, 2), + ) + + +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_choose_quadrature_monotone(self): + """Higher L_max should give equal or larger quadrature grids.""" + prev_n = 0 + for L_max in [3, 5, 7, 11, 15]: + n, K = _choose_quadrature(L_max) + assert n >= prev_n + assert K == L_max + 1 + prev_n = n + + 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_choose_quadrature_too_large(self): + with pytest.raises(ValueError, match="exceeds the largest"): + _choose_quadrature(132) + + @pytest.mark.parametrize("value", [-1, -2]) + def test_choose_quadrature_rejects_negative_degree(self, value): + with pytest.raises(ValueError, match="non-negative"): + _choose_quadrature(value) + + @pytest.mark.parametrize("value", [1.5, True]) + def test_choose_quadrature_rejects_non_integer_degree(self, value): + with pytest.raises(TypeError, match="must be an integer"): + _choose_quadrature(value) + + @pytest.mark.parametrize("value", [0, -1]) + def test_rotation_quadrature_rejects_non_positive_rotation_count(self, value): + with pytest.raises(ValueError, match="positive"): + get_rotation_quadrature(3, value) + + @pytest.mark.parametrize("value", [1.5, True]) + def test_rotation_quadrature_rejects_non_integer_rotation_count(self, value): + with pytest.raises(TypeError, match="must be an integer"): + get_rotation_quadrature(3, value) + + def test_rotation_quadrature_rejects_unsupported_lebedev_order(self): + with pytest.raises(ValueError, match="unsupported Lebedev order"): + 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): + """Return normalized proper matrices and optional improper partners.""" + rotations, weights = get_rotation_quadrature(11, 5) + assert rotations.shape == (rotations.shape[0], 3, 3) + assert np.isclose(weights.sum(), 1.0) + assert np.allclose( + rotations @ rotations.transpose(0, 2, 1), + np.broadcast_to(np.eye(3), rotations.shape), + atol=1e-12, + ) + assert np.allclose(np.linalg.det(rotations), 1.0, atol=1e-12) + + o3_rotations, o3_weights = get_rotation_quadrature( + 11, 5, include_inversion=True + ) + assert len(o3_rotations) == 2 * len(rotations) + assert np.isclose(o3_weights.sum(), 1.0) + 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_constructs_registered_buffers(self): + """Constructor limits should determine the grid and Wigner-D storage.""" + model = SymmetrizedModel( + _EmptyModel(), + max_o3_lambda_target=1, + max_o3_lambda_input=2, + max_o3_lambda_character=1, + batch_size=7, + ) + + assert model.max_o3_lambda_target == 1 + assert model.max_o3_lambda_input == 2 + assert model.max_o3_lambda_character == 1 + assert model.max_o3_lambda_grid == 3 + assert model.batch_size == 7 + + buffers = dict(model.named_buffers()) + assert set(buffers) == { + "_rotation_matrices", + "_rotation_weights", + "_packed_wigner_matrices", + } + assert buffers["_rotation_matrices"].dtype == torch.float64 + assert buffers["_rotation_weights"].dtype == torch.float64 + assert buffers["_packed_wigner_matrices"].dtype == torch.float64 + assert torch.allclose( + buffers["_rotation_weights"].sum(), + torch.tensor(1.0, dtype=torch.float64), + ) + + n_rotations = len(buffers["_rotation_matrices"]) + expected_wigner_elements = n_rotations * sum( + (2 * o3_lambda + 1) ** 2 for o3_lambda in range(3) + ) + assert buffers["_packed_wigner_matrices"].numel() == expected_wigner_elements + + def test_character_limit_controls_default_grid(self): + """Character sectors should raise the default grid degree when necessary.""" + model = SymmetrizedModel( + _EmptyModel(), + max_o3_lambda_target=0, + max_o3_lambda_character=2, + ) + + assert model.max_o3_lambda_grid == 4 + + def test_rejects_grid_too_small_for_character_sectors(self): + """An explicit grid must resolve products for every requested sector.""" + with pytest.raises(ValueError, match="at least twice"): + SymmetrizedModel( + _EmptyModel(), + max_o3_lambda_target=0, + max_o3_lambda_character=2, + max_o3_lambda_grid=3, + ) + + @pytest.mark.parametrize( + ("argument", "value", "error", "message"), + [ + ("max_o3_lambda_target", -1, ValueError, "non-negative"), + ("max_o3_lambda_target", True, TypeError, "integer"), + ("max_o3_lambda_input", 1.5, TypeError, "integer"), + ("max_o3_lambda_character", -1, ValueError, "non-negative"), + ("batch_size", 0, ValueError, "positive"), + ("max_o3_lambda_grid", -1, ValueError, "non-negative"), + ("max_wigner_storage_bytes", 0, ValueError, "positive"), + ], + ) + def test_rejects_invalid_constructor_arguments( + self, + argument, + value, + error, + message, + ): + """Every integer constructor argument should enforce its documented range.""" + arguments = {"max_o3_lambda_target": 0, argument: value} + + with pytest.raises(error, match=message): + SymmetrizedModel(_EmptyModel(), **arguments) + + def test_checks_wigner_storage_limit_before_building(self, monkeypatch): + """An excessive Wigner-D allocation should be rejected before construction.""" + + def fail_if_called(*args, **kwargs): + raise AssertionError("Wigner-D construction should not have started") + + monkeypatch.setattr( + "metatomic.torch.symmetrized_model._model._build_packed_wigner_matrices", + fail_if_called, + ) + + with pytest.raises(ValueError, match="exceeding max_wigner_storage_bytes=1"): + SymmetrizedModel( + _EmptyModel(), + max_o3_lambda_target=0, + max_wigner_storage_bytes=1, + ) + + +class TestSelectedAtomsColumnOrder: + def test_system_column_found_by_name(self): + # the rotated-copy index must go into the "system" column wherever it + # is, not positionally into column 0 + 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] + + +@pytest.mark.parametrize( + ("sample_values", "message"), + [ + ([[0, 0], [2, 0]], "out-of-range rotated-copy indices"), + ([[0, 0], [0, 1], [1, 0]], "same sample labels"), + ([[0, 0], [0, 1], [0, 2], [1, 0]], "same sample labels"), + ([[0, 0], [0, 1], [1, 0], [1, 2]], "same sample labels"), + ], +) +def test_rotated_copy_layout_rejects_inconsistent_samples(sample_values, message): + """Samples from different rotated copies must never be mixed.""" + samples = Labels(["system", "atom"], torch.tensor(sample_values)) + block = TensorBlock( + values=torch.zeros((len(samples), 1), dtype=torch.float64), + samples=samples, + components=[], + properties=Labels.range("property", 1), + ) + + with pytest.raises(ValueError, match=message): + _group_samples_by_rotated_copy(block, n_rotated_copies=2) + + +@pytest.mark.parametrize( + ("sample_values", "values", "n_rotated_copies", "expected_values"), + [ + ( + [[3, 0], [5, 0]], + [3.0, 5.0], + 1, + [[[3.0], [5.0]]], + ), + ( + [[3, 1], [3, 0], [5, 1], [5, 0]], + [13.0, 3.0, 15.0, 5.0], + 2, + [[[3.0], [5.0]], [[13.0], [15.0]]], + ), + ], +) +def test_group_samples_by_rotated_copy( + sample_values, values, n_rotated_copies, expected_values +): + """Values and shared labels should remain aligned after grouping.""" + samples = Labels(["atom", "system"], torch.tensor(sample_values)) + block = TensorBlock( + values=torch.tensor(values, dtype=torch.float64).reshape(-1, 1), + samples=samples, + components=[], + properties=Labels.range("property", 1), + ) + + grouped_values, shared_names, shared_values = _group_samples_by_rotated_copy( + block, n_rotated_copies + ) + + assert torch.equal( + grouped_values, + torch.tensor(expected_values, dtype=torch.float64), + ) + assert shared_names == ["atom"] + assert shared_values.tolist() == [[3], [5]] + + +@pytest.mark.parametrize( + ("sample_names", "sample_values", "expected_names", "expected_values"), + [ + ([], [[]], ["system"], [[7]]), + (["atom"], [[3], [5]], ["system", "atom"], [[7, 3], [7, 5]]), + ], +) +def test_restore_input_system_to_samples( + sample_names, sample_values, expected_names, expected_values +): + """The original system index should be restored without changing samples.""" + samples = _restore_input_system_to_samples( + sample_names, + torch.tensor(sample_values, dtype=torch.int64), + input_system_index=7, + device=torch.device("cpu"), + ) + + assert samples.names == expected_names + assert samples.values.tolist() == expected_values + assert samples.device == torch.device("cpu") + + +@pytest.mark.parametrize("component_shape", [(), (2, 3)]) +def test_weighted_centered_batch_moments(component_shape): + """Compute weighted moments and reuse one fixed reference across batches.""" + n_rotated_copies = 3 + n_samples = 2 + n_properties = 2 + values = torch.arange( + n_rotated_copies * n_samples * int(np.prod(component_shape)) * n_properties, + dtype=torch.float64, + ).reshape(n_rotated_copies * n_samples, *component_shape, n_properties) + components = [ + Labels.range(name, size) + for name, size in zip(("a", "b"), component_shape, strict=False) + ] + tensor = TensorMap( + Labels("kind", torch.tensor([[0]])), + [ + TensorBlock( + values=values, + samples=Labels( + ["system", "item"], + torch.tensor( + [ + [copy, item] + for copy in range(n_rotated_copies) + for item in (5, 7) + ] + ), + ), + components=components, + properties=Labels.range("property", n_properties), + ) + ], + ) + weights = torch.tensor([0.2, -0.1, 0.4], dtype=torch.float64) + + moments = _reduce_weighted_centered_batch( + tensor, + weights, + input_system_index=4, + reference=None, + compute_second_moments=True, + ) + first_moment, second, absolute_second, reference = moments + + values_by_copy = values.reshape( + n_rotated_copies, n_samples, *component_shape, n_properties + ) + centered = values_by_copy - values_by_copy[0] + weight_shape = (n_rotated_copies,) + (1,) * (centered.ndim - 1) + assert torch.allclose( + first_moment.block().values, + torch.sum(weights.reshape(weight_shape) * centered, dim=0), + ) + squared_norms = centered**2 + if component_shape: + squared_norms = squared_norms.sum(dim=tuple(range(2, 2 + len(component_shape)))) + assert second is not None + assert absolute_second is not None + assert torch.allclose( + second.block().values, + torch.sum(weights.reshape(n_rotated_copies, 1, 1) * squared_norms, dim=0), + ) + assert torch.allclose( + absolute_second.block().values, + torch.sum( + torch.abs(weights).reshape(n_rotated_copies, 1, 1) * squared_norms, + dim=0, + ), + ) + expected_samples = Labels( + ["system", "item"], + torch.tensor([[4, 5], [4, 7]]), + ) + assert first_moment.keys == tensor.keys + assert first_moment.block().samples == expected_samples + assert first_moment.block().components == components + assert first_moment.block().properties == tensor.block().properties + assert second.block().samples == expected_samples + assert second.block().components == [] + assert second.block().properties == tensor.block().properties + assert absolute_second.block().samples == expected_samples + assert absolute_second.block().components == [] + assert absolute_second.block().properties == tensor.block().properties + + initial_reference_values = values_by_copy[0].clone() + assert torch.equal(reference.block().values, initial_reference_values) + + # Simulate a later batch with the same layout but different response values. + tensor.block().values.add_(10.0) + later_values_by_copy = tensor.block().values.reshape( + n_rotated_copies, n_samples, *component_shape, n_properties + ) + later_centered = later_values_by_copy - initial_reference_values.unsqueeze(0) + + later_moments = _reduce_weighted_centered_batch( + tensor, + weights, + input_system_index=4, + reference=reference, + compute_second_moments=False, + ) + first_moment, second, absolute_second, reused_reference = later_moments + assert torch.allclose( + first_moment.block().values, + torch.sum(weights.reshape(weight_shape) * later_centered, dim=0), + ) + assert second is None + assert absolute_second is None + assert reused_reference is reference + assert torch.equal(reference.block().values, initial_reference_values) + + +def test_join_per_system_tensormaps_with_matching_keys(monkeypatch): + """Systems with identical keys should be joined along samples.""" + tensors = [ + TensorMap( + Labels("kind", torch.tensor([[0]])), + [ + TensorBlock( + values=torch.tensor([[value]], dtype=torch.float64), + samples=Labels("system", torch.tensor([[system_index]])), + components=[], + properties=Labels.range("property", 1), + ) + ], + ) + for system_index, value in enumerate((1.0, 2.0)) + ] + + native_join = mts.join + different_keys_arguments = [] + + def record_join(tensors, axis, different_keys): + different_keys_arguments.append(different_keys) + return native_join(tensors, axis, different_keys=different_keys) + + monkeypatch.setattr(mts, "join", record_join) + joined = _join_per_system_tensormaps(tensors) + + assert different_keys_arguments == ["error"] + assert joined.keys == tensors[0].keys + assert joined.block().samples.values.tolist() == [[0], [1]] + assert joined.block().values.tolist() == [[1.0], [2.0]] + + +def test_join_per_system_tensormaps_with_different_keys(monkeypatch): + """System-dependent keys should be joined through their union.""" + tensors = [ + TensorMap( + Labels("kind", torch.tensor([[key]])), + [ + TensorBlock( + values=torch.tensor([[value]], dtype=torch.float64), + samples=Labels("system", torch.tensor([[system_index]])), + components=[], + properties=Labels.range("property", 1), + ) + ], + ) + for system_index, (key, value) in enumerate(((0, 1.0), (1, 2.0))) + ] + + native_join = mts.join + different_keys_arguments = [] + + def record_join(tensors, axis, different_keys): + different_keys_arguments.append(different_keys) + return native_join(tensors, axis, different_keys=different_keys) + + monkeypatch.setattr(mts, "join", record_join) + joined = _join_per_system_tensormaps(tensors) + + assert different_keys_arguments == ["union"] + assert joined.keys.values.tolist() == [[0], [1]] + assert joined.block(0).samples.values.tolist() == [[0]] + assert joined.block(0).values.tolist() == [[1.0]] + assert joined.block(1).samples.values.tolist() == [[1]] + assert joined.block(1).values.tolist() == [[2.0]] + + +@pytest.mark.parametrize( + ("component_shape", "n_samples"), + [((), 2), ((3,), 2), ((2, 3), 2), ((2, 3), 0)], +) +def test_component_norm_squared(component_shape, n_samples): + """All component axes should be contracted without changing metadata.""" + shape = (n_samples, *component_shape, 2) + values = torch.arange(int(np.prod(shape)), dtype=torch.float64).reshape(shape) + tensor = _make_single_block_tensor_map(values) + + result = _component_norm_squared(tensor) + + expected = values.square() + if component_shape: + expected = expected.sum(dim=tuple(range(1, 1 + len(component_shape)))) + assert torch.equal(result.block().values, expected) + assert result.keys == tensor.keys + assert result.block().samples == tensor.block().samples + assert result.block().components == [] + assert result.block().properties == tensor.block().properties + + +def test_variance_from_centered_moments(): + """Centered first and second moments should give component-summed variance.""" + component_shape = (2, 3) + shape = (2, *component_shape, 2) + centered_first_moment_values = ( + torch.arange(int(np.prod(shape)), dtype=torch.float64).reshape(shape) / 10 + ) + centered_first_moment = _make_single_block_tensor_map(centered_first_moment_values) + + norm_squared = centered_first_moment_values.square().sum(dim=(1, 2)) + expected_variance = torch.tensor([[0.25, 0.5], [0.75, 1.0]], dtype=torch.float64) + centered_second_moment = _make_single_block_tensor_map( + norm_squared + expected_variance + ) + absolute_centered_second_moment = _make_single_block_tensor_map( + norm_squared + expected_variance + 1.0 + ) + + variance = _variance_from_centered_moments( + centered_first_moment, + centered_second_moment, + absolute_centered_second_moment, + n_grid_points=12, + max_o3_lambda_grid=3, + ) + + assert torch.allclose(variance.block().values, expected_variance) + assert variance.keys == centered_first_moment.keys + assert variance.block().samples == centered_first_moment.block().samples + assert variance.block().components == [] + assert variance.block().properties == centered_first_moment.block().properties + + +def test_centered_variance_is_stable_with_large_offset(): + """A common offset should not cause cancellation in the variance.""" + values = torch.tensor( + [1.0e12, 1.0e12 + 1.0, 1.0e12 + 2.0, 1.0e12 + 3.0], + dtype=torch.float64, + ).reshape(-1, 1) + tensor = _make_single_block_tensor_map(values, sample_name="system") + weights = torch.tensor([0.125, 0.375, 0.375, 0.125], dtype=torch.float64) + + first, second, absolute_second, _ = _reduce_weighted_centered_batch( + tensor, + weights, + input_system_index=7, + reference=None, + compute_second_moments=True, + ) + assert second is not None + assert absolute_second is not None + variance = _variance_from_centered_moments( + first, + second, + absolute_second, + n_grid_points=4, + max_o3_lambda_grid=3, + ) + + assert torch.allclose( + variance.block().values, + torch.tensor([[0.75]], dtype=torch.float64), + rtol=0.0, + atol=1.0e-12, + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize("scale", [1.0e-12, 1.0e12]) +def test_roundoff_negative_diagnostic_uses_its_scale(dtype, scale): + """Only negative values within the summation tolerance should be clamped.""" + n_grid_points = 100 + n_epsilon = n_grid_points * torch.finfo(dtype).eps + gamma = n_epsilon / (1.0 - n_epsilon) + tolerance = 64.0 * gamma * scale + + cleaned = _clamp_roundoff_negative_diagnostic( + _make_single_block_tensor_map( + torch.tensor([[-0.5 * tolerance], [2.0]], dtype=dtype) + ), + _make_single_block_tensor_map(torch.tensor([[scale], [scale]], dtype=dtype)), + n_grid_points=n_grid_points, + quantity="variance", + max_o3_lambda_grid=3, + ) + assert cleaned.block().values[0, 0].item() == 0.0 + assert cleaned.block().values[1, 0].item() == 2.0 + + with pytest.raises(ValueError, match="materially negative"): + _clamp_roundoff_negative_diagnostic( + _make_single_block_tensor_map( + torch.tensor([[-2.0 * tolerance]], dtype=dtype) + ), + _make_single_block_tensor_map(torch.tensor([[scale]], dtype=dtype)), + n_grid_points=n_grid_points, + quantity="variance", + max_o3_lambda_grid=3, + ) + + +@pytest.mark.parametrize( + ("value", "scale"), + [ + (float("nan"), 1.0), + (0.0, float("inf")), + (0.0, -1.0), + ], +) +def test_roundoff_negative_diagnostic_rejects_invalid_input(value, scale): + """Values and their numerical scales should be finite and scales non-negative.""" + with pytest.raises(ValueError, match="round-off scale is invalid"): + _clamp_roundoff_negative_diagnostic( + _make_single_block_tensor_map(torch.tensor([[value]], dtype=torch.float64)), + _make_single_block_tensor_map(torch.tensor([[scale]], dtype=torch.float64)), + n_grid_points=100, + quantity="variance", + max_o3_lambda_grid=3, + ) + + +def test_roundoff_negative_diagnostic_rejects_unsupported_dtype(): + """Diagnostics should use one of the supported floating-point dtypes.""" + with pytest.raises(TypeError, match="float32 or float64"): + _clamp_roundoff_negative_diagnostic( + _make_single_block_tensor_map(torch.tensor([[0.0]], dtype=torch.float16)), + _make_single_block_tensor_map(torch.tensor([[1.0]], dtype=torch.float16)), + n_grid_points=100, + quantity="variance", + max_o3_lambda_grid=3, + ) + + +def test_variance_from_centered_moments_is_scriptable(): + """The complete centered-variance calculation should compile with TorchScript.""" + torch.jit.script(_variance_from_centered_moments) + + +@pytest.mark.parametrize( + ("component_shape", "n_samples"), + [((), 2), ((3,), 2), ((2, 3), 2), ((3,), 0)], +) +def test_mean_variance_over_components(component_shape, n_samples): + """Divide by component count without aggregating or creating samples.""" + variance_values = ( + torch.arange(n_samples * 2, dtype=torch.float64).reshape(n_samples, 2) + 1.0 + ) + variance = _make_single_block_tensor_map(variance_values, sample_name="atom") + component_layout = _make_single_block_tensor_map( + torch.zeros(n_samples, *component_shape, 2, dtype=torch.float64), + sample_name="atom", + ) + + result = _mean_variance_over_components(variance, component_layout) + + n_components = int(np.prod(component_shape)) if component_shape else 1 + assert torch.equal(result.block().values, variance_values / n_components) + assert result.keys == variance.keys + assert result.block().samples == variance.block().samples + assert result.block().components == [] + assert result.block().properties == variance.block().properties + + +@pytest.mark.parametrize( + ("requested_name", "source_name", "calculation"), + [ + ("energy", "energy", "average"), + ("energy/pbe", "energy/pbe", "average"), + ("mtt::aux::features", "mtt::aux::features", "average"), + ("o3::variance::energy/pbe", "energy/pbe", "variance"), + ( + "o3::variance::mtt::aux::features", + "mtt::aux::features", + "variance", + ), + ( + "o3::character_projection::mtt::feature::layer.0", + "mtt::feature::layer.0", + "character_projection", + ), + ( + "o3::variance_extra::energy", + "o3::variance_extra::energy", + "average", + ), + ], +) +def test_parse_output_request(requested_name, source_name, calculation): + """Recognize only complete prefixes and preserve the remaining name.""" + assert _parse_output_request(requested_name) == (source_name, calculation) + + +@pytest.mark.parametrize( + "requested_name", + ["", "o3::variance::", "o3::character_projection::"], +) +def test_parse_output_request_requires_source_name(requested_name): + """Every request should identify an underlying model output.""" + with pytest.raises(ValueError, match="does not identify"): + _parse_output_request(requested_name) + + +def test_group_output_requests_by_source_and_calculation(): + """Group requests while retaining each requested output name and sample kind.""" + outputs = { + "energy": ModelOutput(sample_kind="system"), + "o3::variance::energy": ModelOutput(sample_kind="system"), + "o3::character_projection::energy": ModelOutput(sample_kind="system"), + "o3::variance::mtt::aux::pairs": ModelOutput(sample_kind="atom_pair"), + } + + ( + source_sample_kinds, + average_names, + variance_names, + character_projection_names, + ) = _group_output_requests(outputs) + + assert source_sample_kinds == { + "energy": "system", + "mtt::aux::pairs": "atom_pair", + } + assert average_names == {"energy": "energy"} + assert variance_names == { + "energy": "o3::variance::energy", + "mtt::aux::pairs": "o3::variance::mtt::aux::pairs", + } + assert character_projection_names == {"energy": "o3::character_projection::energy"} + + +def test_group_output_requests_rejects_mixed_sample_kinds(): + """One source cannot share an evaluation at two sample resolutions.""" + outputs = { + "energy": ModelOutput(sample_kind="system"), + "o3::variance::energy": ModelOutput(sample_kind="atom"), + } + + with pytest.raises(ValueError, match="must use the same sample_kind"): + _group_output_requests(outputs) + + +@pytest.mark.parametrize( + ("o3_lambda", "expected"), + [ + (0, [0]), + (1, [-1, 0, 1]), + (2, [-2, -1, 0, 1, 2]), + ], +) +def test_o3_mu_labels(o3_lambda, expected): + """Spherical components should be ordered from -lambda to +lambda.""" + labels = _o3_mu_labels(o3_lambda, torch.device("cpu")) + + assert labels.names == ["o3_mu"] + assert labels.values[:, 0].tolist() == expected + assert labels.device == torch.device("cpu") + + +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_symmetric_matrices_to_spherical_known_components(): + """Identity, traceless diagonal, and skew matrices should map as expected.""" + 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, l2 = _symmetric_matrices_to_spherical(matrices) + + expected_l0 = torch.zeros((3, 1, 1), dtype=torch.float64) + expected_l0[0, 0, 0] = 3.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(l2, expected_l2, rtol=0.0, atol=1.0e-12) + + +def test_symmetric_matrices_to_spherical_preserves_norm(): + """The spherical norm should equal the symmetric-part Frobenius norm.""" + generator = torch.Generator().manual_seed(1234) + matrices = torch.randn( + (4, 3, 3, 2), + dtype=torch.float64, + generator=generator, + ) + symmetric = 0.5 * (matrices + matrices.transpose(1, 2)) + + l0, l2 = _symmetric_matrices_to_spherical(matrices) + + spherical_norm_squared = l0.square().sum(dim=1) + l2.square().sum(dim=1) + cartesian_norm_squared = symmetric.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_symmetric_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, + ) + matrices = torch.tensor( + [ + [[1.2, -0.7, 2.3], [-0.7, 1.1, 0.8], [2.3, 0.8, -0.4]], + [[-0.2, 1.4, 0.5], [1.4, 0.9, -1.1], [0.5, -1.1, 2.0]], + ], + dtype=torch.float64, + ).unsqueeze(-1) + + matrix = transformation.matrix + transformed_matrices = torch.einsum( + "ia,sabp,jb->sijp", + matrix, + matrices, + matrix, + ) + transformed_l0, transformed_l2 = _symmetric_matrices_to_spherical( + transformed_matrices + ) + l0, l2 = _symmetric_matrices_to_spherical(matrices) + + expected_l0 = transformation.transform_spherical( + l0[..., 0], ell=0, 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_l2, expected_l2, rtol=0.0, atol=1.0e-12) + + +@pytest.mark.parametrize( + ( + "names", + "values", + "o3_lambda", + "o3_sigma", + "expected_names", + "expected_values", + ), + [ + (["_"], [[0]], 2, 1, ["o3_lambda", "o3_sigma"], [[2, 1]]), + ( + ["channel"], + [[3], [7]], + 1, + -1, + ["channel", "o3_lambda", "o3_sigma"], + [[3, 1, -1], [7, 1, -1]], + ), + ( + ["channel", "o3_lambda"], + [[3, 1], [7, 1]], + 1, + -1, + ["channel", "o3_lambda", "o3_sigma"], + [[3, 1, -1], [7, 1, -1]], + ), + ], +) +def test_add_o3_irrep_to_keys( + names, + values, + o3_lambda, + o3_sigma, + expected_names, + expected_values, +): + """Preserve semantic keys while assigning one O(3) irrep.""" + result = _add_o3_irrep_to_keys( + Labels(names, torch.tensor(values)), + o3_lambda, + o3_sigma, + ) + + assert result.names == expected_names + assert result.values.tolist() == expected_values + + +@pytest.mark.parametrize( + ("names", "values", "message"), + [ + (["_"], [[1]], "placeholder"), + (["channel", "o3_lambda"], [[3, 1], [7, 2]], "o3_lambda"), + ], +) +def test_add_o3_irrep_to_keys_rejects_conflicting_metadata(names, values, message): + """Reject an invalid ``_`` placeholder or conflicting irrep key values.""" + with pytest.raises(ValueError, match=message): + _add_o3_irrep_to_keys( + Labels(names, torch.tensor(values)), + o3_lambda=1, + o3_sigma=1, + ) + + +@pytest.mark.parametrize( + "source_name", + [ + "energy", + "energy/pbe", + "energy_ensemble/member", + "energy_uncertainty/direct", + ], +) +def test_decompose_output_energy_like(source_name): + """Energy-like variants should become one scalar 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_output(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() + + +def test_decompose_output_non_conservative_force_preserves_autograd(): + """A force variant should become l=1 without breaking 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_output("non_conservative_force/direct", 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_output_non_conservative_stress_combines_irreps(): + """Stress should return l=0 and l=2 blocks and silently discard skew.""" + 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_output("non_conservative_stress/direct", tensor) + + assert result.keys.names == ["o3_lambda", "o3_sigma"] + assert result.keys.values.tolist() == [[0, 1], [2, 1]] + block_l0 = result.block({"o3_lambda": 0, "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_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, + ) + assert torch.equal(block_l2.values, torch.zeros((2, 5, 1), dtype=torch.float64)) + assert block_l0.samples == tensor.block().samples + assert block_l2.samples == tensor.block().samples + assert block_l0.properties == tensor.block().properties + assert block_l2.properties == tensor.block().properties + + +def test_decompose_output_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_output("mtt::custom", tensor) + + assert result is tensor + + +@pytest.mark.parametrize( + ("source_name", "shape", "component_names", "message"), + [ + ("energy", (1, 3, 1), ["xyz"], "must not have components"), + ( + "non_conservative_force", + (1, 3, 1), + ["component"], + "one 'xyz' component axis", + ), + ( + "non_conservative_stress", + (1, 3, 3, 1), + ["xyz_1", "component"], + "'xyz_1' and 'xyz_2' component axes", + ), + ], +) +def test_decompose_output_rejects_invalid_standard_components( + source_name, + shape, + component_names, + message, +): + """Standard quantities should use their required Cartesian component axes.""" + tensor = _tensor_map_with_components( + torch.zeros(shape, dtype=torch.float64), + component_names, + ) + + with pytest.raises(ValueError, match=message): + _decompose_output(source_name, tensor) + + +def test_decompose_output_rejects_attached_gradients(): + """Decomposition should not silently discard explicit TensorBlock gradients.""" + properties = Labels.range("property", 1) + block = TensorBlock( + values=torch.ones((1, 1), dtype=torch.float64), + samples=Labels.range("system", 1), + 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, + ), + ) + tensor = TensorMap( + Labels("_", torch.tensor([[0]], dtype=torch.int64)), + [block], + ) + + with pytest.raises(ValueError, match="gradients attached to 'energy'"): + _decompose_output("energy", tensor) + + +def test_decompose_output_is_scriptable(): + """The output decomposition should compile with TorchScript.""" + torch.jit.script(_decompose_output) From 5ebf4d5b1d028c64911bf881d4b4f30e02ac6e35 Mon Sep 17 00:00:00 2001 From: Michelangelo Domina Date: Thu, 23 Jul 2026 11:02:56 +0200 Subject: [PATCH 02/18] test(torch): cover SymmetrizedModel forward --- .../tests/symmetrized_model.py | 625 ++++++++++++++++++ 1 file changed, 625 insertions(+) diff --git a/python/metatomic_torch/tests/symmetrized_model.py b/python/metatomic_torch/tests/symmetrized_model.py index e3f0d4af..0143bc0a 100644 --- a/python/metatomic_torch/tests/symmetrized_model.py +++ b/python/metatomic_torch/tests/symmetrized_model.py @@ -105,6 +105,278 @@ def forward( 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) -> 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", + 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 which outputs it receives.""" + + def __init__(self): + super().__init__() + self.call_count = 0 + self.requested_names: List[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())) + return super().forward(systems, outputs, selected_atoms) + + +class _O3PolynomialSectorModel(torch.nn.Module): + """ + Return one analytic polynomial response in every O(3) sector through + ``lambda=3``. + + The homogeneous harmonic polynomials ``1``, ``x``, ``x*y``, and ``x*y*z`` + transform purely in the ``lambda=0``, ``1``, ``2``, and ``3`` sectors, + respectively. These responses have ``sigma=+1``. Multiplying each polynomial + by the determinant of the transformed Cartesian frame changes only its + inversion parity, producing the corresponding ``sigma=-1`` response. + + The eight responses are returned as properties of one scalar TensorMap block, + labeled by ``source_lambda`` and ``source_sigma``. + """ + + 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 _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, + ) + ], + ) + + return result + + def _system_with_neighbor_lists(dtype: torch.dtype) -> System: """Create a test system with populated and empty neighbor lists.""" positions = torch.tensor( @@ -949,6 +1221,359 @@ def fail_if_called(*args, **kwargs): ) +class TestSymmetrizedModelForward: + """Test how requested averages and diagnostics are computed and returned.""" + + def test_character_projection_separates_sectors_through_lambda_three(self): + """ + Separate every O(3) ``(lambda, sigma)`` sector through ``lambda=3``. + + Character projection of the eight analytic polynomial responses must + produce eight ``(chi_lambda, chi_sigma)`` blocks. Each block must contain + only the property belonging to the same sector, with squared norms + ``1``, ``1/3``, ``1/15``, and ``1/105`` for ``lambda=0``, ``1``, ``2``, + and ``3``. All projections onto the other seven sectors must vanish. + """ + 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_o3_lambda_target=0, + max_o3_lambda_character=3, + max_o3_lambda_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), + ) + 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_o3_lambda_target=0, + max_o3_lambda_character=1, + max_o3_lambda_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"), + } + + 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, + ) + 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, + ) + + @pytest.mark.parametrize("source_name", ["energy/pbe", "mtt::feature::node"]) + def test_preserves_variant_and_custom_output_names(self, source_name): + """Return variants and custom outputs under their exact requested names.""" + model = SymmetrizedModel( + _LinearEnergyModel(), + max_o3_lambda_target=0, + max_o3_lambda_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 torch.allclose( + result[variance_name].block().values, + torch.tensor([[14.0 / 3.0]], dtype=torch.float64), + 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_o3_lambda_target=1, + max_o3_lambda_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_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", + ] + 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_o3_lambda_target=2, + max_o3_lambda_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, + ) + + expected_target_keys = { + "o3::variance::energy": [[0, 1]], + "o3::variance::non_conservative_force": [[1, 1]], + "o3::variance::non_conservative_stress": [[0, 1], [2, 1]], + "o3::variance::mtt::spherical_vector": [[1, 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_o3_lambda_target=0, + max_o3_lambda_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, + ) + + with torch.no_grad(): + inference_result = model( + [system], + {"energy": ModelOutput(sample_kind="system")}, + None, + ) + assert not inference_result["energy"].block().values.requires_grad + + 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_o3_lambda_target=0) + system = _forward_test_system([[1.0, 2.0, 3.0]]) + + assert model([], {}, None) == {} + with pytest.raises(ValueError, match="at least one System"): + model([], {"energy": ModelOutput(sample_kind="system")}, None) + with pytest.raises(ValueError, match="max_o3_lambda_character must be set"): + model( + [system], + {"o3::character_projection::energy": ModelOutput(sample_kind="system")}, + None, + ) + with pytest.raises(ValueError, match="does not support explicit gradients"): + model( + [system], + { + "energy": ModelOutput( + sample_kind="system", + explicit_gradients=["positions"], + ) + }, + None, + ) + assert base_model.call_count == 0 + + def test_is_scriptable_and_serializable(self, tmp_path): + """The complete forward path should execute after scripting and reloading.""" + constructor_arguments = { + "max_o3_lambda_target": 0, + "max_o3_lambda_character": 1, + "max_o3_lambda_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) + + class TestSelectedAtomsColumnOrder: def test_system_column_found_by_name(self): # the rotated-copy index must go into the "system" column wherever it From 0056d07d0ffdb3f5513f66859495ecea01f3e60f Mon Sep 17 00:00:00 2001 From: Michelangelo Domina Date: Thu, 23 Jul 2026 12:05:48 +0200 Subject: [PATCH 03/18] feat(torch): wrap symmetrized models for export --- .../torch/symmetrized_model/__init__.py | 14 + .../torch/symmetrized_model/_model.py | 149 ++++++++ .../tests/symmetrized_model.py | 341 +++++++++++++++++- 3 files changed, 501 insertions(+), 3 deletions(-) create mode 100644 python/metatomic_torch/metatomic/torch/symmetrized_model/__init__.py diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/__init__.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/__init__.py new file mode 100644 index 00000000..0aa60093 --- /dev/null +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/__init__.py @@ -0,0 +1,14 @@ +""" +O(3) averaging and equivariance diagnostics for atomistic models. + +See :py:class:`SymmetrizedModel` for the method and public output conventions. +""" + +from ._model import SymmetrizedModel +from ._quadrature import get_rotation_quadrature + + +__all__ = [ + "SymmetrizedModel", + "get_rotation_quadrature", +] diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py index 6be842ca..4758bbb0 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py @@ -5,8 +5,11 @@ from metatensor.torch import Labels, TensorBlock, TensorMap from metatomic.torch import ( + AtomisticModel, + ModelCapabilities, ModelInterface, ModelOutput, + NeighborListOptions, System, register_autograd_neighbors, ) @@ -634,6 +637,8 @@ class SymmetrizedModel(torch.nn.Module): """ max_o3_lambda_character: Optional[int] + _requested_inputs: Dict[str, ModelOutput] + _requested_neighbor_lists: List[NeighborListOptions] def __init__( self, @@ -648,6 +653,8 @@ def __init__( super().__init__() self._model = model + self._requested_inputs = {} + self._requested_neighbor_lists = [] self.max_o3_lambda_target = _validate_integer( "max_o3_lambda_target", max_o3_lambda_target, 0 ) @@ -741,6 +748,148 @@ def __init__( self.register_buffer("_rotation_weights", rotation_weights) self.register_buffer("_packed_wigner_matrices", packed_wigner_matrices) + @staticmethod + def wrap( + model: AtomisticModel, + *, + max_o3_lambda_target: int, + max_o3_lambda_input: int = 0, + max_o3_lambda_character: Optional[int] = None, + batch_size: int = 32, + max_o3_lambda_grid: Optional[int] = None, + max_wigner_storage_bytes: int = _DEFAULT_MAX_WIGNER_STORAGE_BYTES, + ) -> 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_o3_lambda_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. + + :param model: the :py:class:`AtomisticModel` to wrap + :param max_o3_lambda_target: largest spherical rank accepted in model outputs + :param max_o3_lambda_input: largest spherical rank accepted in custom System + data + :param max_o3_lambda_character: largest character sector to report, or ``None`` + to disable character projections + :param batch_size: number of transformed Systems evaluated in one model call + :param max_o3_lambda_grid: quadrature integration degree, selected + automatically when ``None`` + :param max_wigner_storage_bytes: maximum size of the packed Wigner-D storage + """ + 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) + ) + + outputs: Dict[str, ModelOutput] = {} + for name in model._model_capabilities_outputs_names: + if name.startswith("o3::variance::") or name.startswith( + "o3::character_projection::" + ): + 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_o3_lambda_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_o3_lambda_target=max_o3_lambda_target, + max_o3_lambda_input=max_o3_lambda_input, + max_o3_lambda_character=max_o3_lambda_character, + batch_size=batch_size, + max_o3_lambda_grid=max_o3_lambda_grid, + max_wigner_storage_bytes=max_wigner_storage_bytes, + ) + wrapper._requested_inputs = { + name: requested_input + for name, requested_input in model._requested_inputs.items() + } + 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], diff --git a/python/metatomic_torch/tests/symmetrized_model.py b/python/metatomic_torch/tests/symmetrized_model.py index 0143bc0a..e21ebb55 100644 --- a/python/metatomic_torch/tests/symmetrized_model.py +++ b/python/metatomic_torch/tests/symmetrized_model.py @@ -6,8 +6,21 @@ import torch from metatensor.torch import Labels, TensorBlock, TensorMap -from metatomic.torch import ModelOutput, NeighborListOptions, System +from metatomic.torch import ( + AtomisticModel, + ModelCapabilities, + ModelEvaluationOptions, + ModelMetadata, + ModelOutput, + NeighborListOptions, + System, + load_atomistic_model, +) from metatomic.torch.o3 import O3Transformation, transform_system +from metatomic.torch.symmetrized_model import ( + SymmetrizedModel, + get_rotation_quadrature, +) from metatomic.torch.symmetrized_model._decompose import ( _add_o3_irrep_to_keys, _cartesian_vectors_to_spherical, @@ -16,7 +29,6 @@ _symmetric_matrices_to_spherical, ) from metatomic.torch.symmetrized_model._model import ( - SymmetrizedModel, _clamp_roundoff_negative_diagnostic, _component_norm_squared, _group_output_requests, @@ -36,7 +48,6 @@ _choose_quadrature, _rotations_from_euler_angles, get_euler_angles_quadrature, - get_rotation_quadrature, ) from metatomic.torch.symmetrized_model._utils import ( _group_samples_by_rotated_copy, @@ -191,6 +202,36 @@ def forward( return super().forward(systems, outputs, selected_atoms) +class _LinearModelWithRequirements(torch.nn.Module): + """Provide a scalar output while requesting custom data and a neighbor list.""" + + def requested_neighbor_lists(self) -> List[NeighborListOptions]: + return [NeighborListOptions(2.5, False, True, "linear model")] + + 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 = 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 _O3PolynomialSectorModel(torch.nn.Module): """ Return one analytic polynomial response in every O(3) sector through @@ -1574,6 +1615,300 @@ def test_is_scriptable_and_serializable(self, tmp_path): mts.allclose_raise(actual[name], expected[name], rtol=0.0, atol=1.0e-12) +class TestSymmetrizedModelWrap: + """Test exported-model capabilities, dependencies, and execution.""" + + @pytest.mark.parametrize("max_o3_lambda_character", [None, 1]) + def test_transfers_metadata_and_declared_capabilities( + self, + max_o3_lambda_character, + ): + """Publish truthful diagnostics without duplicating deprecated aliases.""" + metadata = ModelMetadata( + name="base model", + description="Metadata that should remain unchanged.", + authors=["A. Developer"], + references={"implementation": ["doi:10.0000/example"]}, + extra={"version": "test"}, + ) + source_outputs = { + "energy": ModelOutput( + unit="eV", + sample_kind="system", + explicit_gradients=["positions"], + description="Original energy description.", + ), + "mass": ModelOutput( + unit="u", + sample_kind="atom", + description="Original mass description.", + ), + "mtt::pair": ModelOutput( + sample_kind="atom_pair", + description="Original pair description.", + ), + } + base = AtomisticModel( + _EmptyModel().eval(), + metadata, + 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_o3_lambda_target=0, + max_o3_lambda_character=max_o3_lambda_character, + max_o3_lambda_grid=2, + ) + + actual_metadata = wrapped.metadata() + assert actual_metadata.name == metadata.name + assert actual_metadata.description == metadata.description + assert actual_metadata.authors == metadata.authors + assert actual_metadata.references == metadata.references + assert actual_metadata.extra == metadata.extra + + capabilities = wrapped.capabilities() + assert capabilities.atomic_types == [1, 6, 8] + assert capabilities.interaction_range == 4.5 + assert capabilities.length_unit == "A" + assert capabilities.supported_devices == ["cuda", "cpu"] + assert capabilities.dtype == "float32" + + declared_names = set(wrapped._model_capabilities_outputs_names) + expected_names = set(source_outputs) + expected_names.update("o3::variance::" + name for name in source_outputs) + if max_o3_lambda_character is not None: + expected_names.update( + "o3::character_projection::" + name for name in source_outputs + ) + assert declared_names == expected_names + + # ``AtomisticModel`` adds this compatibility alias, but it must not become + # another declared source with its own diagnostics. + assert "masses" in capabilities.outputs + assert "o3::variance::masses" not in capabilities.outputs + assert "o3::character_projection::masses" not in capabilities.outputs + + source_units = {"energy": "eV", "mass": "u", "mtt::pair": ""} + source_sample_kinds = { + "energy": "system", + "mass": "atom", + "mtt::pair": "atom_pair", + } + for name, source_output in source_outputs.items(): + average = capabilities.outputs[name] + assert average.unit == source_units[name] + assert average.sample_kind == source_sample_kinds[name] + assert average.explicit_gradients == [] + assert source_output.description in average.description + + squared_unit = ( + "" if source_output.unit == "" else f"({source_output.unit})^2" + ) + variance = capabilities.outputs["o3::variance::" + name] + assert variance.unit == squared_unit + assert variance.sample_kind == source_output.sample_kind + assert variance.explicit_gradients == [] + + character_name = "o3::character_projection::" + name + if max_o3_lambda_character is None: + assert character_name not in capabilities.outputs + else: + character = capabilities.outputs[character_name] + assert character.unit == squared_unit + assert character.sample_kind == source_output.sample_kind + assert character.explicit_gradients == [] + + @pytest.mark.parametrize( + "source_name", + [ + "o3::variance::mtt::source", + "o3::character_projection::mtt::source", + ], + ) + def test_rejects_reserved_source_names(self, source_name): + """A source name must not be ambiguous with a generated diagnostic.""" + 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", + ), + ) + + with pytest.raises(ValueError, match="prefix reserved"): + SymmetrizedModel.wrap(base, max_o3_lambda_target=0) + + def test_rejects_models_without_a_supported_device(self): + """The wrapper must not advertise a device on which it cannot run.""" + 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", + ), + ) + + with pytest.raises(ValueError, match="supports CPU and CUDA"): + SymmetrizedModel.wrap(base, max_o3_lambda_target=0) + + def test_preserves_requirements_and_runs_after_save_load(self, tmp_path): + """Wrap a loaded model, re-export it, and execute its declared contract.""" + 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_o3_lambda_target=0, + max_o3_lambda_character=1, + max_o3_lambda_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 = _forward_test_system( + [[1.0, 2.0, 3.0]], + dtype=torch.float32, + ) + system.add_neighbor_list( + neighbor_options, + TensorBlock( + values=torch.empty((0, 3, 1), dtype=torch.float32), + samples=Labels( + [ + "first_atom", + "second_atom", + "cell_shift_a", + "cell_shift_b", + "cell_shift_c", + ], + torch.empty((0, 5), 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", 1), + components=[Labels.range("xyz", 3)], + properties=Labels.range("field", 1), + ) + ], + ) + field.set_info("unit", "eV") + system.add_data("mtt::field", field) + + 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, + ) + + class TestSelectedAtomsColumnOrder: def test_system_column_found_by_name(self): # the rotated-copy index must go into the "system" column wherever it From 1daf33ab41c34663d57d96ee088c2f8fa3807c11 Mon Sep 17 00:00:00 2001 From: Michelangelo Domina Date: Thu, 23 Jul 2026 17:23:00 +0200 Subject: [PATCH 04/18] fix(torch): harden SymmetrizedModel contracts --- docs/src/torch/reference/index.rst | 1 + .../src/torch/reference/symmetrized-model.rst | 253 ++++++++ metatomic-torch/CHANGELOG.md | 6 + .../torch/symmetrized_model/_decompose.py | 5 +- .../torch/symmetrized_model/_model.py | 24 +- .../torch/symmetrized_model/_utils.py | 4 +- .../tests/symmetrized_model.py | 561 ++++++++++++++++-- 7 files changed, 791 insertions(+), 63 deletions(-) create mode 100644 docs/src/torch/reference/symmetrized-model.rst 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..6c19da11 --- /dev/null +++ b/docs/src/torch/reference/symmetrized-model.rst @@ -0,0 +1,253 @@ +.. _symmetrized-model: + +O(3)-symmetrized models +======================= + +The :py:mod:`metatomic.torch.symmetrized_model` module wraps an exported +:py:class:`~metatomic.torch.AtomisticModel` with finite-quadrature O(3) +averaging and equivariance diagnostics. Ordinary outputs are averaged over +rotated and inverted copies of each input. Additional output names request an +equivariance variance or squared character-projection contributions of the +model response. + +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. + +Wrapping and evaluating a model +------------------------------- + +Use +:py:meth:`~metatomic.torch.symmetrized_model.SymmetrizedModel.wrap` to retain +the model metadata, capabilities, requested neighbor lists, and requested +custom inputs: + +.. code-block:: python + + from metatomic.torch import ( + ModelEvaluationOptions, + ModelOutput, + load_atomistic_model, + ) + from metatomic.torch.symmetrized_model import SymmetrizedModel + + base_model = load_atomistic_model("model.pt") + model = SymmetrizedModel.wrap( + base_model, + max_o3_lambda_target=2, + max_o3_lambda_character=3, + ) + + options = ModelEvaluationOptions( + length_unit="angstrom", + outputs={ + "energy": ModelOutput(unit="eV", sample_kind="system"), + "o3::variance::energy": ModelOutput( + unit="(eV)^2", + sample_kind="system", + ), + "o3::character_projection::energy": ModelOutput( + unit="(eV)^2", + sample_kind="system", + ), + }, + ) + results = model(systems, options, check_consistency=True) + model.save("symmetrized-model.pt") + +The requested units must be compatible with the capabilities of the wrapped +model. The example assumes that its length and energy units are ``angstrom`` +and ``eV``. + +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. + +Several calculations for the same source output share the same model +predictions. Their :py:attr:`~metatomic.torch.ModelOutput.sample_kind` values +must agree. The returned dictionary contains exactly the requested names. +Character-projection requests are available only when +``max_o3_lambda_character`` was set during construction. + +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 later +evaluation operation can obtain a block-wise RMSE for a group of samples +:math:`G` as + +.. math:: + + \operatorname{RMSE}_{G} + = \sqrt{ + \frac{\sum_{s\in G} w_s v_\alpha(f,x_s)} + {\sum_{s\in G} w_s} + }. + +Different TensorMap blocks, including different irreducible sectors, remain +separate by default. Combining blocks with different component multiplicities +requires weighting each block by that multiplicity. + +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)`` 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`. Models should provide +symmetric ``non_conservative_stress`` tensors. Stress diagnostics retain only +the scalar trace and symmetric-traceless sectors, silently discarding any +antisymmetric part. A custom Cartesian :math:`3\times3` output is not assumed +to be a stress or to be symmetric. + +Already-spherical outputs retain their ``o3_lambda`` and ``o3_sigma`` keys and +``o3_mu`` components. Other semantic source keys are preserved. The wrapper +does not infer the physical meaning of a custom output from its shape. + +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 and angular-momentum limits +-------------------------------------- + +The deterministic grid combines a Lebedev rule on the sphere, uniformly spaced +in-plane rotations, and both O(3) cosets. 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. Increase ``max_o3_lambda_grid`` until the averages, +variances, and character projections of interest converge. A materially +negative or non-finite squared diagnostic is rejected instead of being reported +as a physical result. + +The constructor keeps four angular-momentum limits separate: + +- ``max_o3_lambda_input`` is the largest spherical rank accepted in custom + :py:class:`~metatomic.torch.System` data. Its default of zero still permits + Cartesian vectors and tensors; it restricts only already-spherical component + axes. +- ``max_o3_lambda_target`` is the largest spherical rank accepted in outputs + that must be transformed back to the input frame. +- ``max_o3_lambda_character`` is the largest character sector included in a + character-projection result. ``None`` disables these outputs. +- ``max_o3_lambda_grid`` controls the quadrature resolution, not an input or + output representation. + +The :py:class:`~metatomic.torch.ModelOutput` declarations returned by a model's +``requested_inputs()`` do not specify the spherical ranks that may occur in the +corresponding TensorMaps. The input limit must therefore be supplied before +export so that all required Wigner-D matrices can be serialized. At runtime, an +already-spherical custom input or an output requiring back-rotation is rejected +when its rank exceeds the corresponding declared limit; the error identifies +the offending name and rank. + +Execution, devices, and gradients +--------------------------------- + +The wrapper supports CPU and CUDA execution with float32 or float64 model +values. It stores quadrature and Wigner-D buffers in float64 and converts final +results back to the model dtype. Move a wrapped model between supported devices +without changing these buffer dtypes. MPS is not supported. + +``batch_size`` controls how many transformed copies are passed to the source +model in one call. It does not change the quadrature or its result. The +statistical accumulators are streamed, but the rotation grid and packed +Wigner-D matrices are persistent buffers. Construction rejects packed +Wigner-D storage larger than ``max_wigner_storage_bytes``. + +Explicit TensorBlock gradients are not supported in requests or source +outputs. Ordinary PyTorch autograd remains available through the returned +values. When an input requires gradients, differentiating an averaged result +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. + +Reference +--------- + +.. py:currentmodule:: metatomic.torch.symmetrized_model + +.. autoclass:: SymmetrizedModel + :members: + +.. autofunction:: get_rotation_quadrature diff --git a/metatomic-torch/CHANGELOG.md b/metatomic-torch/CHANGELOG.md index b066d1ad..2da208f3 100644 --- a/metatomic-torch/CHANGELOG.md +++ b/metatomic-torch/CHANGELOG.md @@ -16,6 +16,12 @@ a changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows ### Removed --> +### Added + +- Added `metatomic.torch.symmetrized_model` for finite-quadrature O(3) + averaging, equivariance variances, and character projections of exported + atomistic models. + ### Changed - Renamed `O3Transformation.is_inverted` to `is_improper`. diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_decompose.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_decompose.py index 50daca92..b28601d8 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_decompose.py +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_decompose.py @@ -61,7 +61,10 @@ def _decompose_output( "energy_ensemble", "energy_uncertainty", ) - is_force = quantity == "non_conservative_force" + is_force = quantity in ( + "non_conservative_force", + "non_conservative_forces", + ) is_stress = quantity == "non_conservative_stress" if not (is_energy or is_force or is_stress): return tensor diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py index 4758bbb0..3999956f 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py @@ -571,8 +571,8 @@ class SymmetrizedModel(torch.nn.Module): \right] = \frac{A_\alpha(f,x)^2}{d}. - Here, :math:`A_\alpha` is the component-summed equivariance error defined in the - reference article. The returned value is instead a component-averaged variance for + Thus, :math:`A_\alpha^2=d\,v_\alpha` is the squared component-summed + equivariance error. The returned value is the component-averaged variance for every retained sample and property: this class neither takes its square root nor aggregates it over samples. @@ -616,7 +616,8 @@ class SymmetrizedModel(torch.nn.Module): :param model: underlying :py:class:`ModelInterface`. The :py:meth:`wrap` method obtains this module from :py:attr:`AtomisticModel.module`. :param max_o3_lambda_target: largest ``o3_lambda`` accepted on an - already-spherical output component axis. Cartesian outputs are not limited by + already-spherical output component axis when an average or variance is + requested. Cartesian outputs and character-only requests are not limited by this value. :param max_o3_lambda_input: largest ``o3_lambda`` accepted on an already-spherical component axis in custom System data. The default of zero @@ -699,7 +700,7 @@ def __init__( for buffer in model.buffers(): device = buffer.device break - if device.type == "mps": + if device.type != "cpu" and device.type != "cuda": raise ValueError("SymmetrizedModel supports CPU and CUDA execution") lebedev_order, n_rotations = _choose_quadrature(self.max_o3_lambda_grid) @@ -773,7 +774,8 @@ def wrap( capabilities are preserved. :param model: the :py:class:`AtomisticModel` to wrap - :param max_o3_lambda_target: largest spherical rank accepted in model outputs + :param max_o3_lambda_target: largest spherical rank accepted in + already-spherical model outputs requested for averaging or variance :param max_o3_lambda_input: largest spherical rank accepted in custom System data :param max_o3_lambda_character: largest character sector to report, or ``None`` @@ -926,13 +928,9 @@ def forward( source_outputs = torch.jit.annotate(Dict[str, ModelOutput], {}) for source_name in source_sample_kinds: - if source_name in average_names: - requested_name = average_names[source_name] - elif source_name in variance_names: - requested_name = variance_names[source_name] - else: - requested_name = character_projection_names[source_name] - source_outputs[source_name] = outputs[requested_name] + source_outputs[source_name] = ModelOutput( + sample_kind=source_sample_kinds[source_name], + ) per_output_results = torch.jit.annotate( Dict[str, List[TensorMap]], @@ -984,6 +982,8 @@ def _evaluate_system( work_device = system.positions.device if work_dtype != torch.float32 and work_dtype != torch.float64: raise TypeError("SymmetrizedModel requires float32 or float64 Systems") + if work_device.type != "cpu" and work_device.type != "cuda": + raise ValueError("SymmetrizedModel supports CPU and CUDA execution") if ( self._rotation_matrices.dtype != torch.float64 or self._rotation_weights.dtype != torch.float64 diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_utils.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_utils.py index 4ef8558d..76fcf445 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_utils.py +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_utils.py @@ -70,9 +70,7 @@ def _group_samples_by_rotated_copy( dim=1, ) if len(copy_indices) != 0 and bool( - torch.any( - (copy_indices < 0) | (copy_indices >= n_rotated_copies) - ).item() + torch.any((copy_indices < 0) | (copy_indices >= n_rotated_copies)).item() ): raise ValueError( "Encountered output samples with out-of-range rotated-copy indices." diff --git a/python/metatomic_torch/tests/symmetrized_model.py b/python/metatomic_torch/tests/symmetrized_model.py index e21ebb55..953b7b0d 100644 --- a/python/metatomic_torch/tests/symmetrized_model.py +++ b/python/metatomic_torch/tests/symmetrized_model.py @@ -1,3 +1,4 @@ +import inspect from typing import Dict, List, Optional import metatensor.torch as mts @@ -135,7 +136,10 @@ def _forward_test_system( ) -def _system_scalar_tensor_map(values: torch.Tensor) -> TensorMap: +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( @@ -153,7 +157,7 @@ def _system_scalar_tensor_map(values: torch.Tensor) -> TensorMap: ), components=[], properties=Labels( - "property", + property_name, torch.arange( values.shape[-1], dtype=torch.int64, @@ -184,12 +188,16 @@ def forward( class _CountingLinearEnergyModel(_LinearEnergyModel): - """Record how often ``forward`` is called and which outputs it receives.""" + """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, @@ -199,14 +207,28 @@ def forward( ) -> 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 _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 [NeighborListOptions(2.5, False, True, "linear model")] + return [self._neighbor_list] def requested_inputs(self) -> Dict[str, ModelOutput]: return { @@ -223,15 +245,72 @@ def forward( 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 - ) + 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: - result[output_name] = _system_scalar_tensor_map(values) + 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 through @@ -302,6 +381,26 @@ def forward( 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.""" @@ -415,6 +514,29 @@ def forward( ], ) + if "mtt::spherical_quadrupole" in outputs: + matrices = torch.stack( + [ + torch.outer(system.positions[0], system.positions[0]) + for system in systems + ] + ).unsqueeze(-1) + _, spherical = _symmetric_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 @@ -1162,6 +1284,38 @@ def test_rotation_quadrature_matrices(self): class TestSymmetrizedModelConstruction: """Test construction of the quadrature and persistent Wigner-D storage.""" + def test_forward_has_the_exact_model_interface_signature(self): + """ + Require the canonical ``ModelInterface.forward`` signature. + + The wrapper must accept only ``systems``, ``outputs``, and + ``selected_atoms``, using the standard annotations and calling + convention without default values. + """ + signature = inspect.signature(SymmetrizedModel.forward) + parameters = list(signature.parameters.values()) + + assert [parameter.name for parameter in parameters] == [ + "self", + "systems", + "outputs", + "selected_atoms", + ] + assert all( + parameter.kind is inspect.Parameter.POSITIONAL_OR_KEYWORD + for parameter in parameters + ) + assert all( + parameter.default is inspect.Parameter.empty for parameter in parameters + ) + assert [parameter.annotation for parameter in parameters] == [ + inspect.Parameter.empty, + List[System], + Dict[str, ModelOutput], + Optional[Labels], + ] + assert signature.return_annotation == Dict[str, TensorMap] + def test_constructs_registered_buffers(self): """Constructor limits should determine the grid and Wigner-D storage.""" model = SymmetrizedModel( @@ -1261,6 +1415,14 @@ def fail_if_called(*args, **kwargs): max_wigner_storage_bytes=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")) + + with pytest.raises(ValueError, match="supports CPU and CUDA"): + SymmetrizedModel(base_model, max_o3_lambda_target=0) + class TestSymmetrizedModelForward: """Test how requested averages and diagnostics are computed and returned.""" @@ -1387,6 +1549,172 @@ def test_energy_results_match_analytic_values_and_reuse_predictions(self): 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_o3_lambda_target=2, + max_o3_lambda_character=2, + max_o3_lambda_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, 1, chi_lambda, chi_sigma) + for o3_lambda in (0, 2) + for chi_lambda in range(3) + for chi_sigma in (1, -1) + } + + for key, block in projection.items(): + o3_lambda = int(key["o3_lambda"]) + chi_lambda = int(key["chi_lambda"]) + chi_sigma = int(key["chi_sigma"]) + assert block.components == [_o3_mu_labels(o3_lambda, block.values.device)] + + if 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_o3_lambda_target=0, + max_o3_lambda_character=1, + max_o3_lambda_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_o3_lambda_target=1, + ) + + with pytest.raises( + ValueError, + match=( + "output 'mtt::spherical_quadrupole' contains o3_lambda=2, " + "exceeding max_o3_lambda_target=1" + ), + ): + 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): + """ + Reject a spurious negative variance caused by insufficient quadrature. + + The degree-12 grid can not integrate the degree-14 squared response and + yields a negative value. Raising the grid degree to 14 must recover the + exact variance. + """ + 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_o3_lambda_target=0, + max_o3_lambda_grid=12, + batch_size=64, + ) + with pytest.raises(ValueError, match="materially negative.*above 12"): + underresolved([system], variance_request, None) + + resolved = SymmetrizedModel( + _DegreeSevenEnergyModel(), + max_o3_lambda_target=0, + max_o3_lambda_grid=14, + batch_size=64, + ) + outputs = { + "energy": ModelOutput(sample_kind="system"), + variance_name: ModelOutput(sample_kind="system"), + } + result = resolved([system], outputs, None) + 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"]) def test_preserves_variant_and_custom_output_names(self, source_name): """Return variants and custom outputs under their exact requested names.""" @@ -1453,6 +1781,44 @@ def test_selected_atoms_excludes_unselected_input_systems(self): 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_o3_lambda_target=1, + max_o3_lambda_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_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]]) @@ -1735,7 +2101,7 @@ def test_transfers_metadata_and_declared_capabilities( ], ) def test_rejects_reserved_source_names(self, source_name): - """A source name must not be ambiguous with a generated diagnostic.""" + """Reject source names that look like wrapper-generated diagnostics.""" base = AtomisticModel( _EmptyModel().eval(), ModelMetadata(), @@ -1753,7 +2119,7 @@ def test_rejects_reserved_source_names(self, source_name): SymmetrizedModel.wrap(base, max_o3_lambda_target=0) def test_rejects_models_without_a_supported_device(self): - """The wrapper must not advertise a device on which it cannot run.""" + """Reject models whose declared devices contain neither CPU nor CUDA.""" base = AtomisticModel( _EmptyModel().eval(), ModelMetadata(), @@ -1771,7 +2137,7 @@ def test_rejects_models_without_a_supported_device(self): SymmetrizedModel.wrap(base, max_o3_lambda_target=0) def test_preserves_requirements_and_runs_after_save_load(self, tmp_path): - """Wrap a loaded model, re-export it, and execute its declared contract.""" + """Preserve model requirements through wrapping, saving, and reloading.""" metadata = ModelMetadata(name="model with requirements") base = AtomisticModel( _LinearModelWithRequirements().eval(), @@ -1829,41 +2195,10 @@ def test_preserves_requirements_and_runs_after_save_load(self, tmp_path): assert neighbor_options.strict is True assert base_requestors.issubset(set(neighbor_options.requestors())) - system = _forward_test_system( - [[1.0, 2.0, 3.0]], - dtype=torch.float32, - ) - system.add_neighbor_list( + system = _system_with_linear_model_requirements( neighbor_options, - TensorBlock( - values=torch.empty((0, 3, 1), dtype=torch.float32), - samples=Labels( - [ - "first_atom", - "second_atom", - "cell_shift_a", - "cell_shift_b", - "cell_shift_c", - ], - torch.empty((0, 5), 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", 1), - components=[Labels.range("xyz", 3)], - properties=Labels.range("field", 1), - ) - ], + torch.device("cpu"), ) - field.set_info("unit", "eV") - system.add_data("mtt::field", field) requested_outputs = { "mtt::linear": ModelOutput( @@ -1908,6 +2243,131 @@ def test_preserves_requirements_and_runs_after_save_load(self, tmp_path): 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.""" + 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_o3_lambda_target=0, + max_o3_lambda_character=1, + max_o3_lambda_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) + + assert cuda_system.positions.device.type == "cuda" + cuda_neighbors = cuda_system.get_neighbor_list(neighbor_options) + assert cuda_neighbors.values.device.type == "cuda" + assert cuda_neighbors.samples.device.type == "cuda" + cuda_field = cuda_system.get_data("mtt::field") + assert cuda_field.keys.device.type == "cuda" + assert cuda_field.block().values.device.type == "cuda" + assert cuda_field.block().samples.device.type == "cuda" + + 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 cuda_model.capabilities().dtype == "float32" + assert cuda_model.capabilities().supported_devices == ["cpu", "cuda"] + assert expected["energy"].block().values.item() == pytest.approx( + 295.2, + rel=2.0e-5, + ) + assert actual["energy"].keys.names == ["_"] + variance = actual["o3::variance::energy"] + assert variance.keys.names == ["o3_lambda", "o3_sigma"] + assert variance.keys.values.cpu().tolist() == [[0, 1]] + assert variance.block().components == [] + projection = actual["o3::character_projection::energy"] + assert projection.keys.names == [ + "o3_lambda", + "o3_sigma", + "chi_lambda", + "chi_sigma", + ] + assert projection.keys.values.cpu().tolist() == [ + [0, 1, 0, 1], + [0, 1, 0, -1], + [0, 1, 1, 1], + [0, 1, 1, -1], + ] + for block in projection.blocks(): + assert len(block.components) == 1 + assert block.components[0].names == ["o3_mu"] + assert len(block.components[0]) == 1 + for name, tensor in actual.items(): + assert tensor.keys.device.type == "cuda" + for block in tensor.blocks(): + assert block.values.device.type == "cuda" + assert block.values.dtype == torch.float32 + assert block.samples.device.type == "cuda" + assert block.properties.device.type == "cuda" + assert all( + component.device.type == "cuda" for component in block.components + ) + mts.allclose_raise( + tensor.to(device="cpu"), + expected[name], + rtol=2.0e-5, + atol=2.0e-5, + ) + class TestSelectedAtomsColumnOrder: def test_system_column_found_by_name(self): @@ -2700,8 +3160,15 @@ def test_decompose_output_energy_like(source_name): assert result.info() == tensor.info() -def test_decompose_output_non_conservative_force_preserves_autograd(): - """A force variant should become l=1 without breaking implicit autograd.""" +@pytest.mark.parametrize( + "source_name", + [ + "non_conservative_force/direct", + "non_conservative_forces/direct", + ], +) +def test_decompose_output_non_conservative_force_preserves_autograd(source_name): + """Both force spellings should become l=1 and preserve implicit autograd.""" values = torch.tensor( [[[1.0], [2.0], [3.0]]], dtype=torch.float64, @@ -2709,7 +3176,7 @@ def test_decompose_output_non_conservative_force_preserves_autograd(): ) tensor = _tensor_map_with_components(values, ["xyz"]) - result = _decompose_output("non_conservative_force/direct", tensor) + result = _decompose_output(source_name, tensor) assert result.keys.names == ["o3_lambda", "o3_sigma"] assert result.keys.values.tolist() == [[1, 1]] From 5058bd45d1f6d3e68d9c578da6193c0052505881 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Tue, 28 Jul 2026 12:05:58 +0200 Subject: [PATCH 05/18] Apply o3 review conventions to symmetrized model --- .../src/torch/reference/symmetrized-model.rst | 137 +-- .../torch/symmetrized_model/_decompose.py | 12 +- .../torch/symmetrized_model/_model.py | 334 ++---- .../torch/symmetrized_model/_projections.py | 56 +- .../torch/symmetrized_model/_quadrature.py | 24 +- .../torch/symmetrized_model/_utils.py | 28 +- .../symmetrized_model/_wigner_storage.py | 21 +- .../tests/symmetrized_model.py | 1044 +++++------------ 8 files changed, 419 insertions(+), 1237 deletions(-) diff --git a/docs/src/torch/reference/symmetrized-model.rst b/docs/src/torch/reference/symmetrized-model.rst index 6c19da11..897b9d02 100644 --- a/docs/src/torch/reference/symmetrized-model.rst +++ b/docs/src/torch/reference/symmetrized-model.rst @@ -10,54 +10,6 @@ rotated and inverted copies of each input. Additional output names request an equivariance variance or squared character-projection contributions of the model response. -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. - -Wrapping and evaluating a model -------------------------------- - -Use -:py:meth:`~metatomic.torch.symmetrized_model.SymmetrizedModel.wrap` to retain -the model metadata, capabilities, requested neighbor lists, and requested -custom inputs: - -.. code-block:: python - - from metatomic.torch import ( - ModelEvaluationOptions, - ModelOutput, - load_atomistic_model, - ) - from metatomic.torch.symmetrized_model import SymmetrizedModel - - base_model = load_atomistic_model("model.pt") - model = SymmetrizedModel.wrap( - base_model, - max_o3_lambda_target=2, - max_o3_lambda_character=3, - ) - - options = ModelEvaluationOptions( - length_unit="angstrom", - outputs={ - "energy": ModelOutput(unit="eV", sample_kind="system"), - "o3::variance::energy": ModelOutput( - unit="(eV)^2", - sample_kind="system", - ), - "o3::character_projection::energy": ModelOutput( - unit="(eV)^2", - sample_kind="system", - ), - }, - ) - results = model(systems, options, check_consistency=True) - model.save("symmetrized-model.pt") - -The requested units must be compatible with the capabilities of the wrapped -model. The example assumes that its length and energy units are ``angstrom`` -and ``eV``. - Output requests --------------- @@ -80,12 +32,6 @@ 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. -Several calculations for the same source output share the same model -predictions. Their :py:attr:`~metatomic.torch.ModelOutput.sample_kind` values -must agree. The returned dictionary contains exactly the requested names. -Character-projection requests are available only when -``max_o3_lambda_character`` was set during construction. - Average and variance -------------------- @@ -117,21 +63,9 @@ variance output contains \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 later -evaluation operation can obtain a block-wise RMSE for a group of samples -:math:`G` as - -.. math:: - - \operatorname{RMSE}_{G} - = \sqrt{ - \frac{\sum_{s\in G} w_s v_\alpha(f,x_s)} - {\sum_{s\in G} w_s} - }. - -Different TensorMap blocks, including different irreducible sectors, remain -separate by default. Combining blocks with different component multiplicities -requires weighting each block by that multiplicity. +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 ------------------------ @@ -158,12 +92,13 @@ Cartesian force components are reordered into the real spherical :math:`\ell=1` basis described in :ref:`o3-conventions`. Models should provide symmetric ``non_conservative_stress`` tensors. Stress diagnostics retain only the scalar trace and symmetric-traceless sectors, silently discarding any -antisymmetric part. A custom Cartesian :math:`3\times3` output is not assumed -to be a stress or to be symmetric. +antisymmetric part. Already-spherical outputs retain their ``o3_lambda`` and ``o3_sigma`` keys and -``o3_mu`` components. Other semantic source keys are preserved. The wrapper -does not infer the physical meaning of a custom output from its shape. +``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 --------------------- @@ -188,59 +123,15 @@ 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 and angular-momentum limits --------------------------------------- +Quadrature +---------- The deterministic grid combines a Lebedev rule on the sphere, uniformly spaced in-plane rotations, and both O(3) cosets. 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. Increase ``max_o3_lambda_grid`` until the averages, -variances, and character projections of interest converge. A materially -negative or non-finite squared diagnostic is rejected instead of being reported -as a physical result. - -The constructor keeps four angular-momentum limits separate: - -- ``max_o3_lambda_input`` is the largest spherical rank accepted in custom - :py:class:`~metatomic.torch.System` data. Its default of zero still permits - Cartesian vectors and tensors; it restricts only already-spherical component - axes. -- ``max_o3_lambda_target`` is the largest spherical rank accepted in outputs - that must be transformed back to the input frame. -- ``max_o3_lambda_character`` is the largest character sector included in a - character-projection result. ``None`` disables these outputs. -- ``max_o3_lambda_grid`` controls the quadrature resolution, not an input or - output representation. - -The :py:class:`~metatomic.torch.ModelOutput` declarations returned by a model's -``requested_inputs()`` do not specify the spherical ranks that may occur in the -corresponding TensorMaps. The input limit must therefore be supplied before -export so that all required Wigner-D matrices can be serialized. At runtime, an -already-spherical custom input or an output requiring back-rotation is rejected -when its rank exceeds the corresponding declared limit; the error identifies -the offending name and rank. - -Execution, devices, and gradients ---------------------------------- - -The wrapper supports CPU and CUDA execution with float32 or float64 model -values. It stores quadrature and Wigner-D buffers in float64 and converts final -results back to the model dtype. Move a wrapped model between supported devices -without changing these buffer dtypes. MPS is not supported. - -``batch_size`` controls how many transformed copies are passed to the source -model in one call. It does not change the quadrature or its result. The -statistical accumulators are streamed, but the rotation grid and packed -Wigner-D matrices are persistent buffers. Construction rejects packed -Wigner-D storage larger than ``max_wigner_storage_bytes``. - -Explicit TensorBlock gradients are not supported in requests or source -outputs. Ordinary PyTorch autograd remains available through the returned -values. When an input requires gradients, differentiating an averaged result -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. +one. A general machine-learning model need not be band-limited, so a finite +grid is not automatically exact. ``max_o3_lambda_grid`` controls the quadrature +resolution, not the representation: increase it until the averages, variances, +and character projections of interest converge. Reference --------- diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_decompose.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_decompose.py index b28601d8..a133b271 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_decompose.py +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_decompose.py @@ -29,7 +29,10 @@ def _cartesian_vectors_to_spherical( def _symmetric_matrices_to_spherical( values: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """Return orthonormal l=0 and l=2 components of the symmetric matrix part.""" + """Return orthonormal l=0 and l=2 components of the symmetric matrix part. + + The antisymmetric (l=1) part is silently discarded. + """ l0 = (values[:, 0, 0, :] + values[:, 1, 1, :] + values[:, 2, 2, :]).unsqueeze( 1 ) / math.sqrt(3.0) @@ -69,13 +72,6 @@ def _decompose_output( if not (is_energy or is_force or is_stress): return tensor - for block in tensor.blocks(): - if len(block.gradients_list()) != 0: - raise ValueError( - "O(3) diagnostic decomposition does not support gradients " - "attached to '" + source_name + "'" - ) - if is_energy: energy_blocks: List[TensorBlock] = [] for block in tensor.blocks(): diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py index 3999956f..da8e86f8 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py @@ -2,7 +2,7 @@ import metatensor.torch as mts import torch -from metatensor.torch import Labels, TensorBlock, TensorMap +from metatensor.torch import Labels, TensorBlock, TensorMap, dtype_name from metatomic.torch import ( AtomisticModel, @@ -36,9 +36,6 @@ ) -_DEFAULT_MAX_WIGNER_STORAGE_BYTES = 64 * 1024 * 1024 # 64 MiB - - def _transform_system_geometry_batch( system: System, matrices: torch.Tensor, @@ -57,12 +54,8 @@ def _transform_system_geometry_batch( ): raise ValueError("system and matrices must have the same dtype and device") - if matrices.size(0) == 1: - positions = (system.positions @ matrices[0].transpose(0, 1)).unsqueeze(0) - cells = (system.cell @ matrices[0].transpose(0, 1)).unsqueeze(0) - else: - positions = system.positions.unsqueeze(0) @ matrices.transpose(1, 2) - cells = system.cell.unsqueeze(0) @ matrices.transpose(1, 2) + 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)): @@ -78,10 +71,7 @@ def _transform_system_geometry_batch( for options in system.known_neighbor_lists(): neighbors = system.get_neighbor_list(options) source_values = neighbors.values.detach().squeeze(-1) - if matrices.size(0) == 1: - neighbor_values = (source_values @ matrices[0].transpose(0, 1)).unsqueeze(0) - else: - neighbor_values = source_values.unsqueeze(0) @ matrices.transpose(1, 2) + 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), @@ -111,13 +101,8 @@ def _check_o3_lambda_limit( tensor_max_o3_lambda = _max_o3_lambda_in_tensor(tensor) if tensor_max_o3_lambda > max_o3_lambda: raise ValueError( - tensor_description - + " contains o3_lambda=" - + str(tensor_max_o3_lambda) - + ", exceeding " - + limit_name - + "=" - + str(max_o3_lambda) + f"{tensor_description} contains o3_lambda={tensor_max_o3_lambda}, " + f"exceeding {limit_name}={max_o3_lambda}" ) @@ -125,19 +110,10 @@ def _transform_system_batch( system: System, matrices: torch.Tensor, wigner_matrices: List[torch.Tensor], - max_o3_lambda_input: int, is_improper: bool, ) -> List[System]: """Transform a System batch, including its custom TensorMap data.""" data_names = system.known_data() - for data_name in data_names: - _check_o3_lambda_limit( - system.get_data(data_name), - "custom input '" + data_name + "'", - max_o3_lambda_input, - "max_o3_lambda_input", - ) - transformed_systems = _transform_system_geometry_batch(system, matrices) if len(data_names) == 0: return transformed_systems @@ -178,9 +154,8 @@ def _parse_output_request(requested_name: str) -> Tuple[str, str]: if len(source_name) == 0: raise ValueError( - "requested output '" - + requested_name - + "' does not identify an underlying model output" + f"requested output '{requested_name}' does not identify an " + "underlying model output" ) return source_name, calculation @@ -207,13 +182,8 @@ def _group_output_requests( previous_sample_kind = source_sample_kinds[source_name] if sample_kind != previous_sample_kind: raise ValueError( - "all requests derived from '" - + source_name - + "' must use the same sample_kind; got '" - + previous_sample_kind - + "' and '" - + sample_kind - + "'" + 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 @@ -245,7 +215,8 @@ def _reduce_weighted_centered_batch( Optional[TensorMap], TensorMap, ]: - """Accumulate one rotation batch's reference-centered weighted moments.""" + """Accumulate one rotation batch's weighted moments, centered on a reference + value so the variance subtraction stays cancellation-safe.""" n_rotated_copies = weights.numel() centered_first_moment_blocks: List[TensorBlock] = [] second_moment_blocks: List[TensorBlock] = [] @@ -257,6 +228,7 @@ def _reduce_weighted_centered_batch( 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 @@ -266,10 +238,12 @@ def _reduce_weighted_centered_batch( 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") + 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) - # Any proper/improper weight split is applied by the caller. batch_weights = weights.to( dtype=centered_values.dtype, device=centered_values.device, @@ -383,21 +357,6 @@ def _copy_tensormap_info(source: TensorMap, result: TensorMap) -> TensorMap: return result -def _join_per_system_tensormaps(tensors: List[TensorMap]) -> TensorMap: - """Join one TensorMap per input system along their sample axes.""" - if len(tensors) == 0: - raise ValueError("expected at least one per-system TensorMap") - - keys = tensors[0].keys - different_keys = "error" - for index in range(1, len(tensors)): - if tensors[index].keys != keys: - different_keys = "union" - break - - return mts.join(tensors, "samples", different_keys=different_keys) - - def _component_norm_squared(tensor: TensorMap) -> TensorMap: """Return squared values summed over all component axes.""" blocks: List[TensorBlock] = [] @@ -428,13 +387,13 @@ def _clamp_roundoff_negative_diagnostic( blocks: List[TensorBlock] = [] for key, block in tensor.items(): scale_values = scale.block(key).values - invalid = ( - (~torch.isfinite(block.values)) - | (~torch.isfinite(scale_values)) - | (scale_values < 0) - ) - if bool(torch.any(invalid).item()): - raise ValueError(f"O(3) {quantity} or its round-off scale is invalid") + 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. @@ -445,7 +404,10 @@ def _clamp_roundoff_negative_diagnostic( epsilon = 1.1920928955078125e-07 tiny = 1.1754943508222875e-38 else: - raise TypeError("O(3) diagnostics require float32 or float64 values") + 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) @@ -539,89 +501,40 @@ def _mean_variance_over_components( class SymmetrizedModel(torch.nn.Module): - r""" + """ Wrap a model with finite-quadrature O(3) averaging and equivariance diagnostics. - For a target representation :math:`\rho_\alpha`, define the model response - transformed back to the input frame as - - .. math:: - - z_\alpha(g;x) = \rho_\alpha(g^{-1}) f(gx). - - An ordinary requested output is the normalized Haar average - - .. math:: - - \Pi_\alpha(f,x) - = \int_{\mathrm{O}(3)} z_\alpha(g;x)\,\mathrm{d}\mu(g). - - The integrals are approximated by evaluating the underlying model on batches of - proper and improper transformations. For a TensorMap block with :math:`d` - component entries, ``o3::variance::`` returns - - .. 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] - = \frac{A_\alpha(f,x)^2}{d}. - - Thus, :math:`A_\alpha^2=d\,v_\alpha` is the squared component-summed - equivariance error. The returned value is the component-averaged variance for - every retained sample and property: this class neither takes its square root nor - aggregates it over samples. - - Character projections act on the direct response :math:`u(g;x) = f(gx)`. For a - character sector :math:`\beta=(\lambda,\sigma)` with - :math:`d_\beta=2\lambda+1`, the corresponding 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). - - Writing an O(3) operation as :math:`\Phi(R,s)`, with :math:`s=+1` for a proper - rotation and :math:`s=-1` for an improper operation, the character convention is - - .. math:: - - \chi_{\lambda,\sigma}(\Phi(R,s)) - = \left[\sigma(-1)^\lambda\right]^{(1-s)/2} - \operatorname{tr} D^\lambda(R). - - Requests named ``o3::character_projection::`` return the unnormalized - contributions to :math:`B_\beta`, labeled by ``chi_lambda`` and ``chi_sigma``. - Target component axes are retained; summing over them recovers the full - component norm in the equation above. - - The deterministic quadrature is exact only when it resolves the angular dependence - of the transformed model response. For unrestricted responses, convergence must be - checked by increasing ``max_o3_lambda_grid``. ``batch_size`` changes how many - transformed systems are evaluated in one model call, but does not change the grid - or the result. - - Rotation matrices, quadrature weights, and Wigner-D matrices are stored as float64 - buffers so they follow ordinary module device movement and serialization. The - packed Wigner-D allocation is checked against ``max_wigner_storage_bytes`` before - it is created. + 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_o3_lambda_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_o3_lambda_target: largest ``o3_lambda`` accepted on an - already-spherical output component axis when an average or variance is - requested. Cartesian outputs and character-only requests are not limited by - this value. - :param max_o3_lambda_input: largest ``o3_lambda`` accepted on an - already-spherical component axis in custom System data. The default of zero - still allows Cartesian custom inputs. + :param max_o3_lambda_target: largest spherical rank 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_o3_lambda_input: largest spherical rank 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 the spherical ranks that + 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_o3_lambda_character: largest character sector included in character projections. ``None`` disables character-projection outputs; zero enables the scalar character sector only. @@ -631,10 +544,8 @@ class SymmetrizedModel(torch.nn.Module): larger of ``2 * max_o3_lambda_target + 1`` and ``2 * max_o3_lambda_character`` when character projections are enabled. An explicit value must be non-negative and no larger than the highest available - Lebedev order, 131. - :param max_wigner_storage_bytes: maximum number of bytes used by the serialized - packed Wigner-D matrices. Construction fails before allocation when this limit - would be exceeded. The default is 64 MiB. + Lebedev order, 131; a value below ``2 * max_o3_lambda_character`` is + rejected. """ max_o3_lambda_character: Optional[int] @@ -649,7 +560,6 @@ def __init__( max_o3_lambda_character: Optional[int] = None, batch_size: int = 32, max_o3_lambda_grid: Optional[int] = None, - max_wigner_storage_bytes: int = _DEFAULT_MAX_WIGNER_STORAGE_BYTES, ): super().__init__() @@ -668,9 +578,6 @@ def __init__( ) self.max_o3_lambda_character = max_o3_lambda_character self.batch_size = _validate_integer("batch_size", batch_size, 1) - self.max_wigner_storage_bytes = _validate_integer( - "max_wigner_storage_bytes", max_wigner_storage_bytes, 1 - ) if max_o3_lambda_grid is None: max_o3_lambda_grid = 2 * self.max_o3_lambda_target + 1 @@ -722,24 +629,6 @@ def __init__( self.max_o3_lambda_target, 0 if self.max_o3_lambda_character is None else self.max_o3_lambda_character, ) - n_wigner_elements_per_matrix = ( - (max_o3_lambda_wigner + 1) - * (2 * max_o3_lambda_wigner + 1) - * (2 * max_o3_lambda_wigner + 3) - // 3 - ) - required_wigner_storage_bytes = ( - len(rotation_matrices) - * n_wigner_elements_per_matrix - * rotation_matrices.element_size() - ) - if required_wigner_storage_bytes > self.max_wigner_storage_bytes: - raise ValueError( - "packed Wigner-D matrices require " - + str(required_wigner_storage_bytes) - + " bytes, exceeding max_wigner_storage_bytes=" - + str(self.max_wigner_storage_bytes) - ) packed_wigner_matrices = _build_packed_wigner_matrices( rotation_matrices, max_o3_lambda_wigner, @@ -758,7 +647,6 @@ def wrap( max_o3_lambda_character: Optional[int] = None, batch_size: int = 32, max_o3_lambda_grid: Optional[int] = None, - max_wigner_storage_bytes: int = _DEFAULT_MAX_WIGNER_STORAGE_BYTES, ) -> AtomisticModel: """ Wrap an exported model with O(3) averaging and diagnostics. @@ -773,6 +661,10 @@ def wrap( 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_o3_lambda_target: largest spherical rank accepted in already-spherical model outputs requested for averaging or variance @@ -783,7 +675,6 @@ def wrap( :param batch_size: number of transformed Systems evaluated in one model call :param max_o3_lambda_grid: quadrature integration degree, selected automatically when ``None`` - :param max_wigner_storage_bytes: maximum size of the packed Wigner-D storage """ if not isinstance(model, AtomisticModel): raise TypeError("model must be an AtomisticModel") @@ -801,6 +692,8 @@ def wrap( ) 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::variance::") or name.startswith( "o3::character_projection::" @@ -855,12 +748,15 @@ def wrap( max_o3_lambda_character=max_o3_lambda_character, batch_size=batch_size, max_o3_lambda_grid=max_o3_lambda_grid, - max_wigner_storage_bytes=max_wigner_storage_bytes, ) + # 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, @@ -907,9 +803,8 @@ def forward( for requested_name, output in outputs.items(): if len(output.explicit_gradients) != 0: raise ValueError( - "SymmetrizedModel does not support explicit gradients for output '" - + requested_name - + "'" + "SymmetrizedModel does not support explicit gradients for " + f"output '{requested_name}'" ) ( @@ -950,20 +845,16 @@ def forward( selected_atoms, ) for requested_name in outputs: - if requested_name not in system_results: - raise ValueError( - "SymmetrizedModel did not produce requested output '" - + requested_name - + "'" - ) per_output_results[requested_name].append( system_results[requested_name] ) results = torch.jit.annotate(Dict[str, TensorMap], {}) for requested_name in outputs: - results[requested_name] = _join_per_system_tensormaps( - per_output_results[requested_name] + results[requested_name] = mts.join( + per_output_results[requested_name], + "samples", + different_keys="union", ) return results @@ -981,15 +872,18 @@ def _evaluate_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") + 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") - if ( - self._rotation_matrices.dtype != torch.float64 - or self._rotation_weights.dtype != torch.float64 - or self._packed_wigner_matrices.dtype != torch.float64 - ): - raise ValueError("SymmetrizedModel integration buffers must remain float64") + if self._rotation_matrices.dtype != torch.float64: + raise ValueError( + "SymmetrizedModel integration buffers must remain float64, got " + f"{dtype_name(self._rotation_matrices.dtype)}; do not call " + ".float() or .half() on the module" + ) if ( self._rotation_matrices.device != work_device or self._rotation_weights.device != work_device @@ -999,6 +893,14 @@ def _evaluate_system( "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_o3_lambda_input, + "max_o3_lambda_input", + ) + character_max = 0 configured_character_max = self.max_o3_lambda_character if configured_character_max is not None: @@ -1082,7 +984,6 @@ def _evaluate_system( system, matrices, input_wigner_matrices, - self.max_o3_lambda_input, is_improper, ) raw_outputs = self._model( @@ -1094,16 +995,8 @@ def _evaluate_system( for source_name in source_outputs: if source_name not in raw_outputs: raise ValueError( - "underlying model did not return requested output '" - + source_name - + "'" - ) - for returned_name in raw_outputs: - if returned_name not in source_outputs: - raise ValueError( - "underlying model returned unrequested output '" - + returned_name - + "'" + "underlying model did not return requested output " + f"'{source_name}'" ) inverse_matrices = (sign * proper_matrices).transpose(1, 2) @@ -1113,11 +1006,8 @@ def _evaluate_system( gradient_names = block.gradients_list() if len(gradient_names) != 0: raise ValueError( - "underlying output '" - + source_name - + "' contains unsupported explicit gradient '" - + gradient_names[0] - + "'" + f"underlying output '{source_name}' contains " + f"unsupported explicit gradient '{gradient_names[0]}'" ) tensor = raw_tensor.to( @@ -1125,12 +1015,15 @@ def _evaluate_system( device=work_device, ) if source_name in average_names or source_name in variance_names: - _check_o3_lambda_limit( - tensor, - "output '" + source_name + "'", - self.max_o3_lambda_target, - "max_o3_lambda_target", - ) + # 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_o3_lambda_target, + "max_o3_lambda_target", + ) backrotated = _transform_tensor_with_precomputed_matrices( tensor, inverse_matrices, @@ -1229,11 +1122,6 @@ def _evaluate_system( results = torch.jit.annotate(Dict[str, TensorMap], {}) for source_name, requested_name in average_names.items(): - if ( - source_name not in average_references - or source_name not in average_first_moments - ): - raise RuntimeError("average accumulation is incomplete") mean = mts.add( average_references[source_name], average_first_moments[source_name], @@ -1245,13 +1133,6 @@ def _evaluate_system( ) for source_name, requested_name in variance_names.items(): - if ( - source_name not in variance_references - or source_name not in variance_first_moments - or source_name not in variance_second_moments - or source_name not in variance_absolute_second_moments - ): - raise RuntimeError("variance accumulation is incomplete") variance = _variance_from_centered_moments( variance_first_moments[source_name], variance_second_moments[source_name], @@ -1269,11 +1150,6 @@ def _evaluate_system( ) for source_name, requested_name in character_projection_names.items(): - if ( - source_name not in proper_character_coefficients - or source_name not in improper_character_coefficients - ): - raise RuntimeError("character-projection accumulation is incomplete") projection = _character_projection_tensormap_from_cosets( proper_character_coefficients[source_name], improper_character_coefficients[source_name], diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_projections.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_projections.py index 98bd147f..d635f1c7 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_projections.py +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_projections.py @@ -1,3 +1,9 @@ +"""Character-projection helpers. + +The projected quantity is defined in the :py:class:`SymmetrizedModel` class +docstring. +""" + from typing import List, Tuple import torch @@ -15,17 +21,6 @@ def _character_projection_coefficients_from_rotation_batch( inverse_wigner_matrices: torch.Tensor, ) -> torch.Tensor: """Compute one rotation batch's character-projection coefficients.""" - if ( - values.dim() < 3 - or weights.dim() != 1 - or inverse_wigner_matrices.dim() != 3 - or weights.size(0) == 0 - or values.size(0) != weights.size(0) - or inverse_wigner_matrices.size(0) != weights.size(0) - or inverse_wigner_matrices.size(1) != inverse_wigner_matrices.size(2) - ): - raise ValueError("incompatible values, weights, or Wigner-matrix shapes") - weighted_wigner_matrices = weights.to( dtype=values.dtype, device=values.device, @@ -47,15 +42,6 @@ def _character_projections_from_proper_and_improper_coefficients( ) -> Tuple[torch.Tensor, torch.Tensor]: """Return squared character projections for ``chi_sigma=+1`` and ``-1``.""" dimension = 2 * chi_lambda + 1 - if ( - chi_lambda < 0 - or proper_coefficients.dim() < 3 - or improper_coefficients.size() != proper_coefficients.size() - or proper_coefficients.size(1) != dimension - or proper_coefficients.size(2) != dimension - ): - raise ValueError("coefficient shapes do not match chi_lambda") - parity = (-1) ** chi_lambda sigma_plus = proper_coefficients + parity * improper_coefficients sigma_minus = proper_coefficients - parity * improper_coefficients @@ -155,11 +141,6 @@ def _character_projection_tensormap_from_cosets( improper_coefficients: TensorMap, ) -> TensorMap: """Combine proper and improper coefficient TensorMaps into O(3) sectors.""" - if proper_coefficients.keys != improper_coefficients.keys: - raise ValueError( - "proper and improper character coefficients must have same keys" - ) - key_names = list(proper_coefficients.keys.names) if "chi_lambda" not in key_names: raise ValueError("character coefficients must contain a 'chi_lambda' key") @@ -172,31 +153,6 @@ def _character_projection_tensormap_from_cosets( for key_index in range(len(proper_coefficients.keys)): proper_block = proper_coefficients.block(key_index) improper_block = improper_coefficients.block(key_index) - proper_components = proper_block.components - improper_components = improper_block.components - components_match = len(proper_components) == len(improper_components) - if components_match: - for component_index in range(len(proper_components)): - if ( - proper_components[component_index] - != improper_components[component_index] - ): - components_match = False - if ( - proper_block.samples != improper_block.samples - or not components_match - or proper_block.properties != improper_block.properties - ): - raise ValueError( - "proper and improper character coefficients must have same metadata" - ) - if ( - len(proper_block.components) < 2 - or proper_block.components[0].names != ["chi_m"] - or proper_block.components[1].names != ["chi_n"] - ): - raise ValueError("character coefficient component metadata is invalid") - chi_lambda = int(proper_coefficients.keys.values[key_index, chi_lambda_column]) sigma_plus, sigma_minus = ( _character_projections_from_proper_and_improper_coefficients( diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_quadrature.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_quadrature.py index 1d7a9a5f..ea5e4554 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_quadrature.py +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_quadrature.py @@ -1,5 +1,3 @@ -from typing import Tuple - import numpy as np from ._utils import _validate_integer @@ -54,7 +52,7 @@ def _import_scipy(): return lebedev_rule, Rotation -def _choose_quadrature(L_max: int) -> Tuple[int, int]: +def _choose_quadrature(L_max: int) -> tuple[int, int]: """ Choose a Lebedev quadrature order and number of in-plane rotations to integrate spherical harmonics up to degree ``L_max``. @@ -77,7 +75,7 @@ def _choose_quadrature(L_max: int) -> Tuple[int, int]: def get_euler_angles_quadrature( lebedev_order: int, n_rotations: int -) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: +) -> 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. @@ -116,7 +114,7 @@ def get_euler_angles_quadrature( def _rotations_from_euler_angles( alpha: np.ndarray, beta: np.ndarray, gamma: np.ndarray -) -> "Rotation": # noqa: F821 (scipy is imported lazily) +): """ Construct one active ZYZ rotation from each Euler-angle triple. @@ -141,23 +139,21 @@ def _rotations_from_euler_angles( def get_rotation_quadrature( lebedev_order: int, n_rotations: int, include_inversion: bool = False -) -> Tuple[np.ndarray, np.ndarray]: +) -> 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. SO(3) contains - proper rotations with determinant +1, while O(3) also contains improper - orthogonal transformations with determinant -1. If ``include_inversion`` - is ``True``, each proper rotation is paired with an improper one and the - original weight is divided equally between the pair. + 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 + :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. ``lebedev_order`` must be one of the orders - supported by ``scipy.integrate.lebedev_rule``. + ``(N,)``, summing to 1 """ alpha, beta, gamma, weights = get_euler_angles_quadrature( lebedev_order, n_rotations diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_utils.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_utils.py index 76fcf445..acae2e42 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_utils.py +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_utils.py @@ -1,7 +1,6 @@ -import operator +from numbers import Integral from typing import List, Optional, Tuple -import numpy as np import torch from metatensor.torch import Labels, TensorBlock @@ -11,16 +10,9 @@ def _validate_integer(name: str, value, minimum: int) -> int: Return it as a Python ``int``. """ - if isinstance(value, (bool, np.bool_)) or ( - isinstance(value, torch.Tensor) and value.dtype == torch.bool - ): - raise TypeError(f"{name} must be an integer, not a boolean") - try: - integer_value = int(operator.index(value)) - except TypeError as error: - raise TypeError( - f"{name} must be an integer, got {type(value).__name__}" - ) from error + 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: qualifier = "positive" if minimum == 1 else "non-negative" raise ValueError(f"{name} must be {qualifier}, got {integer_value}") @@ -73,15 +65,9 @@ def _group_samples_by_rotated_copy( torch.any((copy_indices < 0) | (copy_indices >= n_rotated_copies)).item() ): raise ValueError( - "Encountered output samples with out-of-range rotated-copy indices." - ) - - # A single copy is already grouped; avoid sorting the common batch-size-one case. - if n_rotated_copies == 1: - return ( - block.values.unsqueeze(0), - sample_names[:system_column] + sample_names[system_column + 1 :], - sample_values_without_system, + "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: diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_wigner_storage.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_wigner_storage.py index c841cb92..1980458a 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_wigner_storage.py +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_wigner_storage.py @@ -1,7 +1,6 @@ import torch from ..o3 import O3Transformation -from ._utils import _validate_integer def _build_packed_wigner_matrices( @@ -9,17 +8,6 @@ def _build_packed_wigner_matrices( max_o3_lambda: int, ) -> torch.Tensor: """Build and pack proper Wigner-D matrices through ``max_o3_lambda``.""" - max_o3_lambda = _validate_integer("max_o3_lambda", max_o3_lambda, 0) - 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 not in (torch.float32, torch.float64): - raise TypeError("matrices must use float32 or float64") - output_device = matrices.device output_dtype = matrices.dtype calculation_matrices = matrices.detach().to(device="cpu") @@ -55,13 +43,8 @@ def _wigner_matrices_for_lambda( o3_lambda: int, ) -> torch.Tensor: """Return the packed Wigner-D stack for one ``o3_lambda`` as a view.""" - if packed.dim() != 1: - raise ValueError("packed Wigner-D storage must be one-dimensional") - if n_matrices <= 0: - raise ValueError("n_matrices must be positive") - if o3_lambda < 0: - raise ValueError("o3_lambda must be non-negative") - + # 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 diff --git a/python/metatomic_torch/tests/symmetrized_model.py b/python/metatomic_torch/tests/symmetrized_model.py index 953b7b0d..8c1096d3 100644 --- a/python/metatomic_torch/tests/symmetrized_model.py +++ b/python/metatomic_torch/tests/symmetrized_model.py @@ -1,4 +1,4 @@ -import inspect +import re from typing import Dict, List, Optional import metatensor.torch as mts @@ -23,7 +23,6 @@ get_rotation_quadrature, ) from metatomic.torch.symmetrized_model._decompose import ( - _add_o3_irrep_to_keys, _cartesian_vectors_to_spherical, _decompose_output, _o3_mu_labels, @@ -32,10 +31,7 @@ from metatomic.torch.symmetrized_model._model import ( _clamp_roundoff_negative_diagnostic, _component_norm_squared, - _group_output_requests, - _join_per_system_tensormaps, _mean_variance_over_components, - _parse_output_request, _reduce_weighted_centered_batch, _transform_system_batch, _transform_system_geometry_batch, @@ -53,7 +49,6 @@ from metatomic.torch.symmetrized_model._utils import ( _group_samples_by_rotated_copy, _map_selected_atoms_to_rotated_copies, - _restore_input_system_to_samples, ) from metatomic.torch.symmetrized_model._wigner_storage import ( _build_packed_wigner_matrices, @@ -312,19 +307,10 @@ def _system_with_linear_model_requirements( class _O3PolynomialSectorModel(torch.nn.Module): - """ - Return one analytic polynomial response in every O(3) sector through - ``lambda=3``. + """Return one analytic polynomial response in every O(3) sector to lambda=3.""" - The homogeneous harmonic polynomials ``1``, ``x``, ``x*y``, and ``x*y*z`` - transform purely in the ``lambda=0``, ``1``, ``2``, and ``3`` sectors, - respectively. These responses have ``sigma=+1``. Multiplying each polynomial - by the determinant of the transformed Cartesian frame changes only its - inversion parity, producing the corresponding ``sigma=-1`` response. - - The eight responses are returned as properties of one scalar TensorMap block, - labeled by ``source_lambda`` and ``source_sigma``. - """ + # 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, @@ -602,17 +588,13 @@ class TestSystemGeometryBatch: """Test batched O(3) transformation of System geometry.""" @pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) - @pytest.mark.parametrize("n_matrices", [1, 3]) - def test_matches_individual_o3_transformations(self, dtype, n_matrices): + def test_matches_individual_o3_transformations(self, dtype): """Batched geometry should match one transformation at a time.""" proper = torch.tensor( [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=dtype, ) - if n_matrices == 1: - matrices = proper.unsqueeze(0) - else: - matrices = torch.stack([torch.eye(3, dtype=dtype), proper, -proper]) + matrices = torch.stack([torch.eye(3, dtype=dtype), proper, -proper]) system = _system_with_neighbor_lists(dtype) transformed = _transform_system_geometry_batch(system, matrices) @@ -629,12 +611,10 @@ def test_matches_individual_o3_transformations(self, dtype, n_matrices): assert torch.equal(actual.pbc, expected.pbc) assert actual.known_neighbor_lists() == expected.known_neighbor_lists() for options in expected.known_neighbor_lists(): - actual_neighbors = actual.get_neighbor_list(options) - expected_neighbors = expected.get_neighbor_list(options) - assert torch.equal(actual_neighbors.values, expected_neighbors.values) - assert actual_neighbors.samples == expected_neighbors.samples - assert actual_neighbors.components == expected_neighbors.components - assert actual_neighbors.properties == expected_neighbors.properties + assert torch.equal( + actual.get_neighbor_list(options).values, + expected.get_neighbor_list(options).values, + ) def test_preserves_neighbor_autograd(self): """Rotated neighbor vectors should differentiate through positions and cell.""" @@ -705,32 +685,21 @@ def test_rejects_invalid_matrix_batches(self): """Matrix batches should have a non-empty shape and match the System.""" system = _system_with_neighbor_lists(torch.float64) invalid_shapes = [(3, 3), (0, 3, 3), (2, 2, 3), (2, 3, 2)] + message = "matrices must have shape (N, 3, 3) with N > 0" for shape in invalid_shapes: - with pytest.raises(ValueError, match="shape \\(N, 3, 3\\)"): + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): _transform_system_geometry_batch( system, torch.empty(shape, dtype=torch.float64), ) - with pytest.raises(ValueError, match="same dtype and device"): + message = "system and matrices must have the same dtype and device" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): _transform_system_geometry_batch( system, torch.eye(3, dtype=torch.float32).unsqueeze(0), ) - def test_is_scriptable(self): - """The batched geometry transformation should compile and execute.""" - scripted = torch.jit.script(_transform_system_geometry_batch) - system = _system_with_neighbor_lists(torch.float64) - transformed = scripted( - system, - torch.eye(3, dtype=torch.float64).unsqueeze(0), - ) - - assert len(transformed) == 1 - assert torch.equal(transformed[0].positions, system.positions) - assert torch.equal(transformed[0].cell, system.cell) - class TestSystemBatch: """Test batched O(3) transformation of complete Systems.""" @@ -775,7 +744,6 @@ def test_transforms_spherical_custom_data(self, is_improper): values = torch.tensor( [[[1.0], [2.0], [3.0]], [[-0.5], [1.5], [0.25]]], dtype=torch.float64, - requires_grad=True, ) system.add_data( "mtt::field", @@ -799,7 +767,6 @@ def test_transforms_spherical_custom_data(self, is_improper): system, matrices, wigner_matrices, - max_o3_lambda_input=1, is_improper=is_improper, ) @@ -817,18 +784,6 @@ def test_transforms_spherical_custom_data(self, is_improper): atol=1.0e-12, ) - loss = sum( - transformed_system.get_data("mtt::field").block().values.square().sum() - for transformed_system in transformed - ) - gradient = torch.autograd.grad(loss, values)[0] - assert torch.allclose( - gradient, - 2 * len(matrices) * values, - rtol=0.0, - atol=1.0e-12, - ) - def test_input_limit_distinguishes_spherical_from_cartesian(self): """A zero spherical-rank limit should still allow Cartesian custom data.""" matrix = torch.tensor( @@ -867,13 +822,11 @@ def test_input_limit_distinguishes_spherical_from_cartesian(self): ], ) system.add_data("mtt::field", cartesian) - scripted_transform = torch.jit.script(_transform_system_batch) - transformed = scripted_transform( + transformed = _transform_system_batch( system, matrix, wigner_matrices, - max_o3_lambda_input=0, is_improper=False, ) expected = transform_system( @@ -910,19 +863,20 @@ def test_input_limit_distinguishes_spherical_from_cartesian(self): ], ), ) - with pytest.raises( - torch.jit.Error, - match=( - "custom input 'mtt::field' contains o3_lambda=1, exceeding " - "max_o3_lambda_input=0" - ), - ): - scripted_transform( - spherical_system, - matrix, - wigner_matrices, - max_o3_lambda_input=0, - is_improper=False, + model = SymmetrizedModel( + _LinearEnergyModel(), + max_o3_lambda_target=0, + max_o3_lambda_grid=2, + ) + message = ( + "custom input 'mtt::field' contains o3_lambda=1, exceeding " + "max_o3_lambda_input=0" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + model( + [spherical_system], + {"energy": ModelOutput(sample_kind="system")}, + None, ) @@ -932,12 +886,13 @@ class TestCharacterProjections: @pytest.mark.parametrize("n_samples", [0, 2]) def test_batch_coefficients_match_rotation_by_rotation_sum(self, n_samples): """Batching should match summing the weighted rotations individually.""" - torch.manual_seed(7) + generator = torch.Generator().manual_seed(7) n_rotations = 4 dimension = 3 values = torch.randn( (n_rotations, n_samples, 2, 3), dtype=torch.float64, + generator=generator, ) weights = torch.tensor( [0.50, -0.25, 0.30, 0.45], @@ -946,6 +901,7 @@ def test_batch_coefficients_match_rotation_by_rotation_sum(self, n_samples): inverse_wigner_matrices = torch.randn( (n_rotations, dimension, dimension), dtype=torch.float32, + generator=generator, ) coefficients = _character_projection_coefficients_from_rotation_batch( @@ -972,18 +928,18 @@ def test_batch_coefficients_match_rotation_by_rotation_sum(self, n_samples): @pytest.mark.parametrize("chi_lambda", [0, 1, 2]) def test_factorization_matches_all_rotation_pairs(self, chi_lambda): """The factorization should match summing every pair of rotations.""" - torch.manual_seed(11 + chi_lambda) + generator = torch.Generator().manual_seed(11 + chi_lambda) n_rotations = 4 dimension = 2 * chi_lambda + 1 proper_values = torch.randn( (n_rotations, 2, 2, 1), dtype=torch.float64, - requires_grad=True, + generator=generator, ) improper_values = torch.randn( (n_rotations, 2, 2, 1), dtype=torch.float64, - requires_grad=True, + generator=generator, ) weights = torch.tensor( [0.50, -0.25, 0.30, 0.45], @@ -992,6 +948,7 @@ def test_factorization_matches_all_rotation_pairs(self, chi_lambda): inverse_wigner_matrices = torch.randn( (n_rotations, dimension, dimension), dtype=torch.float64, + generator=generator, ) proper_coefficients = _character_projection_coefficients_from_rotation_batch( proper_values, @@ -1037,54 +994,9 @@ def test_factorization_matches_all_rotation_pairs(self, chi_lambda): assert torch.allclose(sigma_plus, expected[0], rtol=0.0, atol=1e-12) assert torch.allclose(sigma_minus, expected[1], rtol=0.0, atol=1e-12) - assert sigma_plus.shape == proper_values.shape[1:] - assert sigma_minus.shape == improper_values.shape[1:] assert torch.all(sigma_plus >= 0) assert torch.all(sigma_minus >= 0) - (sigma_plus.sum() + sigma_minus.sum()).backward() - assert torch.all(torch.isfinite(proper_values.grad)) - assert torch.all(torch.isfinite(improper_values.grad)) - - def test_rejects_mismatched_rotation_counts_and_coefficient_shapes(self): - """Reject unequal rotation counts or proper/improper coefficient shapes.""" - with pytest.raises(ValueError, match="incompatible values"): - _character_projection_coefficients_from_rotation_batch( - torch.zeros((3, 1, 1), dtype=torch.float64), - torch.ones(2, dtype=torch.float64), - torch.ones((3, 1, 1), dtype=torch.float64), - ) - - with pytest.raises(ValueError, match="chi_lambda"): - _character_projections_from_proper_and_improper_coefficients( - torch.zeros((1, 3, 3, 1), dtype=torch.float64), - torch.zeros((2, 3, 3, 1), dtype=torch.float64), - chi_lambda=1, - ) - - def test_is_scriptable(self): - """Both character-projection tensor operations should compile and run.""" - coefficient_function = torch.jit.script( - _character_projection_coefficients_from_rotation_batch - ) - projection_function = torch.jit.script( - _character_projections_from_proper_and_improper_coefficients - ) - values = torch.ones((1, 1, 1), dtype=torch.float64) - coefficients = coefficient_function( - values, - torch.ones(1, dtype=torch.float64), - torch.ones((1, 1, 1), dtype=torch.float64), - ) - sigma_plus, sigma_minus = projection_function( - coefficients, - coefficients, - 0, - ) - - assert sigma_plus.item() == 1.0 - assert sigma_minus.item() == 0.0 - class TestWignerStorage: """Test persistent Wigner-D storage for the quadrature grid.""" @@ -1134,48 +1046,12 @@ def test_packed_matrices_match_o3(self, dtype): ) assert torch.equal(actual, expected) - rank_one = _wigner_matrices_for_lambda(packed, len(matrices), 1) - previous = rank_one[0, 0, 0].clone() - rank_one[0, 0, 0] += 1 - assert packed[len(matrices)] == previous + 1 - - def test_builder_rejects_invalid_inputs(self): - """The builder should reject invalid ranks, shapes, and dtypes.""" - matrices = torch.eye(3, dtype=torch.float64).unsqueeze(0) - with pytest.raises(ValueError, match="non-negative"): - _build_packed_wigner_matrices(matrices, -1) - - for shape in ((0, 3, 3), (2, 3, 2)): - with pytest.raises(ValueError, match="shape \\(N, 3, 3\\)"): - _build_packed_wigner_matrices( - torch.empty(shape, dtype=torch.float64), - 1, - ) - - with pytest.raises(TypeError, match="float32 or float64"): - _build_packed_wigner_matrices(matrices.to(torch.float16), 1) - - def test_rank_view_rejects_invalid_inputs(self): - """Rank views should reject invalid storage, counts, and ranks.""" - with pytest.raises(ValueError, match="one-dimensional"): - _wigner_matrices_for_lambda(torch.empty((2, 2)), 1, 0) - with pytest.raises(ValueError, match="n_matrices must be positive"): - _wigner_matrices_for_lambda(torch.empty(1), 0, 0) - with pytest.raises(ValueError, match="o3_lambda must be non-negative"): - _wigner_matrices_for_lambda(torch.empty(1), 1, -1) - with pytest.raises(ValueError, match="exceeds the packed"): + def test_rank_view_rejects_out_of_range_lambda(self): + """Rank views should reject ranks beyond the packed storage.""" + message = "o3_lambda exceeds the packed Wigner-D storage" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): _wigner_matrices_for_lambda(torch.empty(1), 1, 1) - def test_rank_view_is_scriptable(self): - """The runtime rank accessor should compile and execute in TorchScript.""" - scripted = torch.jit.script(_wigner_matrices_for_lambda) - packed = torch.arange(70, dtype=torch.float64) - - assert torch.equal( - scripted(packed, 2, 2), - _wigner_matrices_for_lambda(packed, 2, 2), - ) - class TestQuadrature: """Test quadrature weights and grid properties.""" @@ -1191,15 +1067,6 @@ def test_weights_sum(self): f"Weights don't sum to 1 for L_max={L_max}: sum={w.sum()}" ) - def test_choose_quadrature_monotone(self): - """Higher L_max should give equal or larger quadrature grids.""" - prev_n = 0 - for L_max in [3, 5, 7, 11, 15]: - n, K = _choose_quadrature(L_max) - assert n >= prev_n - assert K == L_max + 1 - prev_n = n - 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) @@ -1221,32 +1088,39 @@ def test_euler_angle_rotations_are_in_so3(self): atol=1e-12, ) - def test_choose_quadrature_too_large(self): - with pytest.raises(ValueError, match="exceeds the largest"): + def test_quadrature_validation(self): + """Quadrature construction rejects invalid degrees, counts, and orders.""" + message = ( + "the requested quadrature degree L_max=132 exceeds the largest " + "available Lebedev order (131)" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): _choose_quadrature(132) - @pytest.mark.parametrize("value", [-1, -2]) - def test_choose_quadrature_rejects_negative_degree(self, value): - with pytest.raises(ValueError, match="non-negative"): - _choose_quadrature(value) + message = "L_max must be non-negative, got -1" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + _choose_quadrature(-1) - @pytest.mark.parametrize("value", [1.5, True]) - def test_choose_quadrature_rejects_non_integer_degree(self, value): - with pytest.raises(TypeError, match="must be an integer"): - _choose_quadrature(value) + message = "L_max must be an integer, got float" + with pytest.raises(TypeError, match=f"^{re.escape(message)}$"): + _choose_quadrature(1.5) - @pytest.mark.parametrize("value", [0, -1]) - def test_rotation_quadrature_rejects_non_positive_rotation_count(self, value): - with pytest.raises(ValueError, match="positive"): - get_rotation_quadrature(3, value) + message = "n_rotations must be positive, got 0" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + get_rotation_quadrature(3, 0) - @pytest.mark.parametrize("value", [1.5, True]) - def test_rotation_quadrature_rejects_non_integer_rotation_count(self, value): - with pytest.raises(TypeError, match="must be an integer"): - get_rotation_quadrature(3, value) + message = "n_rotations must be an integer, got float" + with pytest.raises(TypeError, match=f"^{re.escape(message)}$"): + get_rotation_quadrature(3, 1.5) - def test_rotation_quadrature_rejects_unsupported_lebedev_order(self): - with pytest.raises(ValueError, match="unsupported Lebedev order"): + 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): @@ -1261,22 +1135,11 @@ def test_degree_two_grid_resolves_l1_products(self): assert np.isclose(projected_norm, 1.0 / 3.0, atol=1e-12) def test_rotation_quadrature_matrices(self): - """Return normalized proper matrices and optional improper partners.""" - rotations, weights = get_rotation_quadrature(11, 5) - assert rotations.shape == (rotations.shape[0], 3, 3) - assert np.isclose(weights.sum(), 1.0) - assert np.allclose( - rotations @ rotations.transpose(0, 2, 1), - np.broadcast_to(np.eye(3), rotations.shape), - atol=1e-12, - ) - assert np.allclose(np.linalg.det(rotations), 1.0, atol=1e-12) + """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) - o3_rotations, o3_weights = get_rotation_quadrature( - 11, 5, include_inversion=True - ) assert len(o3_rotations) == 2 * len(rotations) - assert np.isclose(o3_weights.sum(), 1.0) dets = np.linalg.det(o3_rotations) assert np.allclose(np.sort(dets), np.repeat([-1.0, 1.0], len(rotations))) @@ -1284,38 +1147,6 @@ def test_rotation_quadrature_matrices(self): class TestSymmetrizedModelConstruction: """Test construction of the quadrature and persistent Wigner-D storage.""" - def test_forward_has_the_exact_model_interface_signature(self): - """ - Require the canonical ``ModelInterface.forward`` signature. - - The wrapper must accept only ``systems``, ``outputs``, and - ``selected_atoms``, using the standard annotations and calling - convention without default values. - """ - signature = inspect.signature(SymmetrizedModel.forward) - parameters = list(signature.parameters.values()) - - assert [parameter.name for parameter in parameters] == [ - "self", - "systems", - "outputs", - "selected_atoms", - ] - assert all( - parameter.kind is inspect.Parameter.POSITIONAL_OR_KEYWORD - for parameter in parameters - ) - assert all( - parameter.default is inspect.Parameter.empty for parameter in parameters - ) - assert [parameter.annotation for parameter in parameters] == [ - inspect.Parameter.empty, - List[System], - Dict[str, ModelOutput], - Optional[Labels], - ] - assert signature.return_annotation == Dict[str, TensorMap] - def test_constructs_registered_buffers(self): """Constructor limits should determine the grid and Wigner-D storage.""" model = SymmetrizedModel( @@ -1341,10 +1172,6 @@ def test_constructs_registered_buffers(self): assert buffers["_rotation_matrices"].dtype == torch.float64 assert buffers["_rotation_weights"].dtype == torch.float64 assert buffers["_packed_wigner_matrices"].dtype == torch.float64 - assert torch.allclose( - buffers["_rotation_weights"].sum(), - torch.tensor(1.0, dtype=torch.float64), - ) n_rotations = len(buffers["_rotation_matrices"]) expected_wigner_elements = n_rotations * sum( @@ -1364,7 +1191,8 @@ def test_character_limit_controls_default_grid(self): def test_rejects_grid_too_small_for_character_sectors(self): """An explicit grid must resolve products for every requested sector.""" - with pytest.raises(ValueError, match="at least twice"): + message = "max_o3_lambda_grid must be at least twice max_o3_lambda_character" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): SymmetrizedModel( _EmptyModel(), max_o3_lambda_target=0, @@ -1375,13 +1203,37 @@ def test_rejects_grid_too_small_for_character_sectors(self): @pytest.mark.parametrize( ("argument", "value", "error", "message"), [ - ("max_o3_lambda_target", -1, ValueError, "non-negative"), - ("max_o3_lambda_target", True, TypeError, "integer"), - ("max_o3_lambda_input", 1.5, TypeError, "integer"), - ("max_o3_lambda_character", -1, ValueError, "non-negative"), - ("batch_size", 0, ValueError, "positive"), - ("max_o3_lambda_grid", -1, ValueError, "non-negative"), - ("max_wigner_storage_bytes", 0, ValueError, "positive"), + ( + "max_o3_lambda_target", + -1, + ValueError, + "max_o3_lambda_target must be non-negative, got -1", + ), + ( + "max_o3_lambda_target", + True, + TypeError, + "max_o3_lambda_target must be an integer, got bool", + ), + ( + "max_o3_lambda_input", + 1.5, + TypeError, + "max_o3_lambda_input must be an integer, got float", + ), + ( + "max_o3_lambda_character", + -1, + ValueError, + "max_o3_lambda_character must be non-negative, got -1", + ), + ("batch_size", 0, ValueError, "batch_size must be positive, got 0"), + ( + "max_o3_lambda_grid", + -1, + ValueError, + "max_o3_lambda_grid must be non-negative, got -1", + ), ], ) def test_rejects_invalid_constructor_arguments( @@ -1394,33 +1246,16 @@ def test_rejects_invalid_constructor_arguments( """Every integer constructor argument should enforce its documented range.""" arguments = {"max_o3_lambda_target": 0, argument: value} - with pytest.raises(error, match=message): + with pytest.raises(error, match=f"^{re.escape(message)}$"): SymmetrizedModel(_EmptyModel(), **arguments) - def test_checks_wigner_storage_limit_before_building(self, monkeypatch): - """An excessive Wigner-D allocation should be rejected before construction.""" - - def fail_if_called(*args, **kwargs): - raise AssertionError("Wigner-D construction should not have started") - - monkeypatch.setattr( - "metatomic.torch.symmetrized_model._model._build_packed_wigner_matrices", - fail_if_called, - ) - - with pytest.raises(ValueError, match="exceeding max_wigner_storage_bytes=1"): - SymmetrizedModel( - _EmptyModel(), - max_o3_lambda_target=0, - max_wigner_storage_bytes=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")) - with pytest.raises(ValueError, match="supports CPU and CUDA"): + message = "SymmetrizedModel supports CPU and CUDA execution" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): SymmetrizedModel(base_model, max_o3_lambda_target=0) @@ -1428,15 +1263,7 @@ class TestSymmetrizedModelForward: """Test how requested averages and diagnostics are computed and returned.""" def test_character_projection_separates_sectors_through_lambda_three(self): - """ - Separate every O(3) ``(lambda, sigma)`` sector through ``lambda=3``. - - Character projection of the eight analytic polynomial responses must - produce eight ``(chi_lambda, chi_sigma)`` blocks. Each block must contain - only the property belonging to the same sector, with squared norms - ``1``, ``1/3``, ``1/15``, and ``1/105`` for ``lambda=0``, ``1``, ``2``, - and ``3``. All projections onto the other seven sectors must vanish. - """ + """Character projection separates the eight analytic sectors to lambda=3.""" source_name = "mtt::o3_polynomial_sectors" requested_name = "o3::character_projection::" + source_name sectors = [ @@ -1463,6 +1290,8 @@ def test_character_projection_separates_sectors_through_lambda_three(self): ["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]])) @@ -1513,6 +1342,7 @@ def test_energy_results_match_analytic_values_and_reuse_predictions(self): 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, @@ -1649,13 +1479,11 @@ def test_rejects_an_output_above_the_declared_target_rank(self): max_o3_lambda_target=1, ) - with pytest.raises( - ValueError, - match=( - "output 'mtt::spherical_quadrupole' contains o3_lambda=2, " - "exceeding max_o3_lambda_target=1" - ), - ): + message = ( + "output 'mtt::spherical_quadrupole' contains o3_lambda=2, " + "exceeding max_o3_lambda_target=1" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): model( [_forward_test_system([[1.0, 2.0, 3.0]])], { @@ -1667,13 +1495,7 @@ def test_rejects_an_output_above_the_declared_target_rank(self): ) def test_rejects_a_negative_quadrature_error_and_converges(self): - """ - Reject a spurious negative variance caused by insufficient quadrature. - - The degree-12 grid can not integrate the degree-14 squared response and - yields a negative value. Raising the grid degree to 14 must recover the - exact variance. - """ + """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, @@ -1691,7 +1513,12 @@ def test_rejects_a_negative_quadrature_error_and_converges(self): max_o3_lambda_grid=12, batch_size=64, ) - with pytest.raises(ValueError, match="materially negative.*above 12"): + message = ( + "finite O(3) variance is materially negative; the quadrature does " + "not resolve this response. Increase max_o3_lambda_grid above 12 " + "and check convergence" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): underresolved([system], variance_request, None) resolved = SymmetrizedModel( @@ -1705,6 +1532,7 @@ def test_rejects_a_negative_quadrature_error_and_converges(self): 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, @@ -1715,11 +1543,23 @@ def test_rejects_a_negative_quadrature_error_and_converges(self): rel=1.0e-12, ) - @pytest.mark.parametrize("source_name", ["energy/pbe", "mtt::feature::node"]) + @pytest.mark.parametrize( + "source_name", + [ + "energy/pbe", + "mtt::feature::node", + # "o3::variance_extra::" is not the reserved prefix: the full name + # is passed through as a source output + "o3::variance_extra::energy", + # 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( - _LinearEnergyModel(), + base_model, max_o3_lambda_target=0, max_o3_lambda_grid=2, ) @@ -1736,6 +1576,7 @@ def test_preserves_variant_and_custom_output_names(self, source_name): ) 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), @@ -1915,14 +1756,6 @@ def test_dtype_and_implicit_autograd(self, dtype): atol=tolerance, ) - with torch.no_grad(): - inference_result = model( - [system], - {"energy": ModelOutput(sample_kind="system")}, - None, - ) - assert not inference_result["energy"].block().values.requires_grad - def test_rejects_invalid_requests_before_model_evaluation(self): """Invalid public requests should fail without running the source model.""" base_model = _CountingLinearEnergyModel() @@ -1930,15 +1763,20 @@ def test_rejects_invalid_requests_before_model_evaluation(self): system = _forward_test_system([[1.0, 2.0, 3.0]]) assert model([], {}, None) == {} - with pytest.raises(ValueError, match="at least one System"): + message = "SymmetrizedModel requires at least one System" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): model([], {"energy": ModelOutput(sample_kind="system")}, None) - with pytest.raises(ValueError, match="max_o3_lambda_character must be set"): + message = "max_o3_lambda_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, ) - with pytest.raises(ValueError, match="does not support explicit gradients"): + message = ( + "SymmetrizedModel does not support explicit gradients for output 'energy'" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): model( [system], { @@ -1949,6 +1787,19 @@ def test_rejects_invalid_requests_before_model_evaluation(self): }, 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_is_scriptable_and_serializable(self, tmp_path): @@ -1985,38 +1836,20 @@ class TestSymmetrizedModelWrap: """Test exported-model capabilities, dependencies, and execution.""" @pytest.mark.parametrize("max_o3_lambda_character", [None, 1]) - def test_transfers_metadata_and_declared_capabilities( - self, - max_o3_lambda_character, - ): - """Publish truthful diagnostics without duplicating deprecated aliases.""" - metadata = ModelMetadata( - name="base model", - description="Metadata that should remain unchanged.", - authors=["A. Developer"], - references={"implementation": ["doi:10.0000/example"]}, - extra={"version": "test"}, - ) + def test_wrap_declares_capabilities(self, max_o3_lambda_character): + """Wrapping declares averages and diagnostics with squared units.""" source_outputs = { "energy": ModelOutput( unit="eV", sample_kind="system", explicit_gradients=["positions"], - description="Original energy description.", - ), - "mass": ModelOutput( - unit="u", - sample_kind="atom", - description="Original mass description.", - ), - "mtt::pair": ModelOutput( - sample_kind="atom_pair", - description="Original pair description.", ), + "mass": ModelOutput(unit="u", sample_kind="atom"), + "mtt::pair": ModelOutput(sample_kind="atom_pair"), } base = AtomisticModel( _EmptyModel().eval(), - metadata, + ModelMetadata(), ModelCapabilities( outputs=source_outputs, atomic_types=[1, 6, 8], @@ -2034,64 +1867,31 @@ def test_transfers_metadata_and_declared_capabilities( max_o3_lambda_grid=2, ) - actual_metadata = wrapped.metadata() - assert actual_metadata.name == metadata.name - assert actual_metadata.description == metadata.description - assert actual_metadata.authors == metadata.authors - assert actual_metadata.references == metadata.references - assert actual_metadata.extra == metadata.extra - capabilities = wrapped.capabilities() - assert capabilities.atomic_types == [1, 6, 8] - assert capabilities.interaction_range == 4.5 - assert capabilities.length_unit == "A" assert capabilities.supported_devices == ["cuda", "cpu"] - assert capabilities.dtype == "float32" - declared_names = set(wrapped._model_capabilities_outputs_names) expected_names = set(source_outputs) expected_names.update("o3::variance::" + name for name in source_outputs) if max_o3_lambda_character is not None: expected_names.update( "o3::character_projection::" + name for name in source_outputs ) - assert declared_names == expected_names - - # ``AtomisticModel`` adds this compatibility alias, but it must not become - # another declared source with its own diagnostics. - assert "masses" in capabilities.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 - source_units = {"energy": "eV", "mass": "u", "mtt::pair": ""} - source_sample_kinds = { - "energy": "system", - "mass": "atom", - "mtt::pair": "atom_pair", - } for name, source_output in source_outputs.items(): - average = capabilities.outputs[name] - assert average.unit == source_units[name] - assert average.sample_kind == source_sample_kinds[name] - assert average.explicit_gradients == [] - assert source_output.description in average.description - squared_unit = ( "" if source_output.unit == "" else f"({source_output.unit})^2" ) - variance = capabilities.outputs["o3::variance::" + name] - assert variance.unit == squared_unit - assert variance.sample_kind == source_output.sample_kind - assert variance.explicit_gradients == [] - + assert capabilities.outputs["o3::variance::" + name].unit == squared_unit character_name = "o3::character_projection::" + name if max_o3_lambda_character is None: assert character_name not in capabilities.outputs else: - character = capabilities.outputs[character_name] - assert character.unit == squared_unit - assert character.sample_kind == source_output.sample_kind - assert character.explicit_gradients == [] + assert capabilities.outputs[character_name].unit == squared_unit @pytest.mark.parametrize( "source_name", @@ -2115,7 +1915,11 @@ def test_rejects_reserved_source_names(self, source_name): ), ) - with pytest.raises(ValueError, match="prefix reserved"): + 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_o3_lambda_target=0) def test_rejects_models_without_a_supported_device(self): @@ -2133,7 +1937,11 @@ def test_rejects_models_without_a_supported_device(self): ), ) - with pytest.raises(ValueError, match="supports CPU and CUDA"): + 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_o3_lambda_target=0) def test_preserves_requirements_and_runs_after_save_load(self, tmp_path): @@ -2283,15 +2091,6 @@ def test_saved_wrapper_runs_on_cuda(self, tmp_path): ) cuda_system = cpu_system.to(device=cuda_device) - assert cuda_system.positions.device.type == "cuda" - cuda_neighbors = cuda_system.get_neighbor_list(neighbor_options) - assert cuda_neighbors.values.device.type == "cuda" - assert cuda_neighbors.samples.device.type == "cuda" - cuda_field = cuda_system.get_data("mtt::field") - assert cuda_field.keys.device.type == "cuda" - assert cuda_field.block().values.device.type == "cuda" - assert cuda_field.block().samples.device.type == "cuda" - requested_outputs = { "energy": ModelOutput( unit="meV", @@ -2323,44 +2122,8 @@ def test_saved_wrapper_runs_on_cuda(self, tmp_path): ) assert set(actual) == set(requested_outputs) - assert cuda_model.capabilities().dtype == "float32" - assert cuda_model.capabilities().supported_devices == ["cpu", "cuda"] - assert expected["energy"].block().values.item() == pytest.approx( - 295.2, - rel=2.0e-5, - ) - assert actual["energy"].keys.names == ["_"] - variance = actual["o3::variance::energy"] - assert variance.keys.names == ["o3_lambda", "o3_sigma"] - assert variance.keys.values.cpu().tolist() == [[0, 1]] - assert variance.block().components == [] - projection = actual["o3::character_projection::energy"] - assert projection.keys.names == [ - "o3_lambda", - "o3_sigma", - "chi_lambda", - "chi_sigma", - ] - assert projection.keys.values.cpu().tolist() == [ - [0, 1, 0, 1], - [0, 1, 0, -1], - [0, 1, 1, 1], - [0, 1, 1, -1], - ] - for block in projection.blocks(): - assert len(block.components) == 1 - assert block.components[0].names == ["o3_mu"] - assert len(block.components[0]) == 1 + assert actual["energy"].block().values.device.type == "cuda" for name, tensor in actual.items(): - assert tensor.keys.device.type == "cuda" - for block in tensor.blocks(): - assert block.values.device.type == "cuda" - assert block.values.dtype == torch.float32 - assert block.samples.device.type == "cuda" - assert block.properties.device.type == "cuda" - assert all( - component.device.type == "cuda" for component in block.components - ) mts.allclose_raise( tensor.to(device="cpu"), expected[name], @@ -2380,13 +2143,23 @@ def test_system_column_found_by_name(self): assert rotated.values[:, 1].tolist() == [0, 0, 1, 1] +_SAME_SAMPLE_LABELS_MESSAGE = ( + "SymmetrizedModel expects every rotated copy to produce the same sample " + "labels in the same order." +) + + @pytest.mark.parametrize( ("sample_values", "message"), [ - ([[0, 0], [2, 0]], "out-of-range rotated-copy indices"), - ([[0, 0], [0, 1], [1, 0]], "same sample labels"), - ([[0, 0], [0, 1], [0, 2], [1, 0]], "same sample labels"), - ([[0, 0], [0, 1], [1, 0], [1, 2]], "same sample labels"), + ( + [[0, 0], [2, 0]], + "encountered output samples with out-of-range rotated-copy " + "indices: the system column spans [0, 2], expected [0, 1]", + ), + ([[0, 0], [0, 1], [1, 0]], _SAME_SAMPLE_LABELS_MESSAGE), + ([[0, 0], [0, 1], [0, 2], [1, 0]], _SAME_SAMPLE_LABELS_MESSAGE), + ([[0, 0], [0, 1], [1, 0], [1, 2]], _SAME_SAMPLE_LABELS_MESSAGE), ], ) def test_rotated_copy_layout_rejects_inconsistent_samples(sample_values, message): @@ -2399,7 +2172,7 @@ def test_rotated_copy_layout_rejects_inconsistent_samples(sample_values, message properties=Labels.range("property", 1), ) - with pytest.raises(ValueError, match=message): + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): _group_samples_by_rotated_copy(block, n_rotated_copies=2) @@ -2444,29 +2217,6 @@ def test_group_samples_by_rotated_copy( assert shared_values.tolist() == [[3], [5]] -@pytest.mark.parametrize( - ("sample_names", "sample_values", "expected_names", "expected_values"), - [ - ([], [[]], ["system"], [[7]]), - (["atom"], [[3], [5]], ["system", "atom"], [[7, 3], [7, 5]]), - ], -) -def test_restore_input_system_to_samples( - sample_names, sample_values, expected_names, expected_values -): - """The original system index should be restored without changing samples.""" - samples = _restore_input_system_to_samples( - sample_names, - torch.tensor(sample_values, dtype=torch.int64), - input_system_index=7, - device=torch.device("cpu"), - ) - - assert samples.names == expected_names - assert samples.values.tolist() == expected_values - assert samples.device == torch.device("cpu") - - @pytest.mark.parametrize("component_shape", [(), (2, 3)]) def test_weighted_centered_batch_moments(component_shape): """Compute weighted moments and reuse one fixed reference across batches.""" @@ -2541,16 +2291,8 @@ def test_weighted_centered_batch_moments(component_shape): ["system", "item"], torch.tensor([[4, 5], [4, 7]]), ) - assert first_moment.keys == tensor.keys assert first_moment.block().samples == expected_samples - assert first_moment.block().components == components - assert first_moment.block().properties == tensor.block().properties - assert second.block().samples == expected_samples assert second.block().components == [] - assert second.block().properties == tensor.block().properties - assert absolute_second.block().samples == expected_samples - assert absolute_second.block().components == [] - assert absolute_second.block().properties == tensor.block().properties initial_reference_values = values_by_copy[0].clone() assert torch.equal(reference.block().values, initial_reference_values) @@ -2580,81 +2322,10 @@ def test_weighted_centered_batch_moments(component_shape): assert torch.equal(reference.block().values, initial_reference_values) -def test_join_per_system_tensormaps_with_matching_keys(monkeypatch): - """Systems with identical keys should be joined along samples.""" - tensors = [ - TensorMap( - Labels("kind", torch.tensor([[0]])), - [ - TensorBlock( - values=torch.tensor([[value]], dtype=torch.float64), - samples=Labels("system", torch.tensor([[system_index]])), - components=[], - properties=Labels.range("property", 1), - ) - ], - ) - for system_index, value in enumerate((1.0, 2.0)) - ] - - native_join = mts.join - different_keys_arguments = [] - - def record_join(tensors, axis, different_keys): - different_keys_arguments.append(different_keys) - return native_join(tensors, axis, different_keys=different_keys) - - monkeypatch.setattr(mts, "join", record_join) - joined = _join_per_system_tensormaps(tensors) - - assert different_keys_arguments == ["error"] - assert joined.keys == tensors[0].keys - assert joined.block().samples.values.tolist() == [[0], [1]] - assert joined.block().values.tolist() == [[1.0], [2.0]] - - -def test_join_per_system_tensormaps_with_different_keys(monkeypatch): - """System-dependent keys should be joined through their union.""" - tensors = [ - TensorMap( - Labels("kind", torch.tensor([[key]])), - [ - TensorBlock( - values=torch.tensor([[value]], dtype=torch.float64), - samples=Labels("system", torch.tensor([[system_index]])), - components=[], - properties=Labels.range("property", 1), - ) - ], - ) - for system_index, (key, value) in enumerate(((0, 1.0), (1, 2.0))) - ] - - native_join = mts.join - different_keys_arguments = [] - - def record_join(tensors, axis, different_keys): - different_keys_arguments.append(different_keys) - return native_join(tensors, axis, different_keys=different_keys) - - monkeypatch.setattr(mts, "join", record_join) - joined = _join_per_system_tensormaps(tensors) - - assert different_keys_arguments == ["union"] - assert joined.keys.values.tolist() == [[0], [1]] - assert joined.block(0).samples.values.tolist() == [[0]] - assert joined.block(0).values.tolist() == [[1.0]] - assert joined.block(1).samples.values.tolist() == [[1]] - assert joined.block(1).values.tolist() == [[2.0]] - - -@pytest.mark.parametrize( - ("component_shape", "n_samples"), - [((), 2), ((3,), 2), ((2, 3), 2), ((2, 3), 0)], -) -def test_component_norm_squared(component_shape, n_samples): +@pytest.mark.parametrize("component_shape", [(), (3,), (2, 3)]) +def test_component_norm_squared(component_shape): """All component axes should be contracted without changing metadata.""" - shape = (n_samples, *component_shape, 2) + shape = (2, *component_shape, 2) values = torch.arange(int(np.prod(shape)), dtype=torch.float64).reshape(shape) tensor = _make_single_block_tensor_map(values) @@ -2737,10 +2408,10 @@ def test_centered_variance_is_stable_with_large_offset(): ) -@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) -@pytest.mark.parametrize("scale", [1.0e-12, 1.0e12]) -def test_roundoff_negative_diagnostic_uses_its_scale(dtype, scale): +def test_roundoff_negative_diagnostic_uses_its_scale(): """Only negative values within the summation tolerance should be clamped.""" + dtype = torch.float64 + scale = 1.0e12 n_grid_points = 100 n_epsilon = n_grid_points * torch.finfo(dtype).eps gamma = n_epsilon / (1.0 - n_epsilon) @@ -2758,7 +2429,12 @@ def test_roundoff_negative_diagnostic_uses_its_scale(dtype, scale): assert cleaned.block().values[0, 0].item() == 0.0 assert cleaned.block().values[1, 0].item() == 2.0 - with pytest.raises(ValueError, match="materially negative"): + message = ( + "finite O(3) variance is materially negative; the quadrature does not " + "resolve this response. Increase max_o3_lambda_grid above 3 and check " + "convergence" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): _clamp_roundoff_negative_diagnostic( _make_single_block_tensor_map( torch.tensor([[-2.0 * tolerance]], dtype=dtype) @@ -2770,49 +2446,10 @@ def test_roundoff_negative_diagnostic_uses_its_scale(dtype, scale): ) -@pytest.mark.parametrize( - ("value", "scale"), - [ - (float("nan"), 1.0), - (0.0, float("inf")), - (0.0, -1.0), - ], -) -def test_roundoff_negative_diagnostic_rejects_invalid_input(value, scale): - """Values and their numerical scales should be finite and scales non-negative.""" - with pytest.raises(ValueError, match="round-off scale is invalid"): - _clamp_roundoff_negative_diagnostic( - _make_single_block_tensor_map(torch.tensor([[value]], dtype=torch.float64)), - _make_single_block_tensor_map(torch.tensor([[scale]], dtype=torch.float64)), - n_grid_points=100, - quantity="variance", - max_o3_lambda_grid=3, - ) - - -def test_roundoff_negative_diagnostic_rejects_unsupported_dtype(): - """Diagnostics should use one of the supported floating-point dtypes.""" - with pytest.raises(TypeError, match="float32 or float64"): - _clamp_roundoff_negative_diagnostic( - _make_single_block_tensor_map(torch.tensor([[0.0]], dtype=torch.float16)), - _make_single_block_tensor_map(torch.tensor([[1.0]], dtype=torch.float16)), - n_grid_points=100, - quantity="variance", - max_o3_lambda_grid=3, - ) - - -def test_variance_from_centered_moments_is_scriptable(): - """The complete centered-variance calculation should compile with TorchScript.""" - torch.jit.script(_variance_from_centered_moments) - - -@pytest.mark.parametrize( - ("component_shape", "n_samples"), - [((), 2), ((3,), 2), ((2, 3), 2), ((3,), 0)], -) -def test_mean_variance_over_components(component_shape, n_samples): +@pytest.mark.parametrize("component_shape", [(), (3,), (2, 3)]) +def test_mean_variance_over_components(component_shape): """Divide by component count without aggregating or creating samples.""" + n_samples = 2 variance_values = ( torch.arange(n_samples * 2, dtype=torch.float64).reshape(n_samples, 2) + 1.0 ) @@ -2832,101 +2469,6 @@ def test_mean_variance_over_components(component_shape, n_samples): assert result.block().properties == variance.block().properties -@pytest.mark.parametrize( - ("requested_name", "source_name", "calculation"), - [ - ("energy", "energy", "average"), - ("energy/pbe", "energy/pbe", "average"), - ("mtt::aux::features", "mtt::aux::features", "average"), - ("o3::variance::energy/pbe", "energy/pbe", "variance"), - ( - "o3::variance::mtt::aux::features", - "mtt::aux::features", - "variance", - ), - ( - "o3::character_projection::mtt::feature::layer.0", - "mtt::feature::layer.0", - "character_projection", - ), - ( - "o3::variance_extra::energy", - "o3::variance_extra::energy", - "average", - ), - ], -) -def test_parse_output_request(requested_name, source_name, calculation): - """Recognize only complete prefixes and preserve the remaining name.""" - assert _parse_output_request(requested_name) == (source_name, calculation) - - -@pytest.mark.parametrize( - "requested_name", - ["", "o3::variance::", "o3::character_projection::"], -) -def test_parse_output_request_requires_source_name(requested_name): - """Every request should identify an underlying model output.""" - with pytest.raises(ValueError, match="does not identify"): - _parse_output_request(requested_name) - - -def test_group_output_requests_by_source_and_calculation(): - """Group requests while retaining each requested output name and sample kind.""" - outputs = { - "energy": ModelOutput(sample_kind="system"), - "o3::variance::energy": ModelOutput(sample_kind="system"), - "o3::character_projection::energy": ModelOutput(sample_kind="system"), - "o3::variance::mtt::aux::pairs": ModelOutput(sample_kind="atom_pair"), - } - - ( - source_sample_kinds, - average_names, - variance_names, - character_projection_names, - ) = _group_output_requests(outputs) - - assert source_sample_kinds == { - "energy": "system", - "mtt::aux::pairs": "atom_pair", - } - assert average_names == {"energy": "energy"} - assert variance_names == { - "energy": "o3::variance::energy", - "mtt::aux::pairs": "o3::variance::mtt::aux::pairs", - } - assert character_projection_names == {"energy": "o3::character_projection::energy"} - - -def test_group_output_requests_rejects_mixed_sample_kinds(): - """One source cannot share an evaluation at two sample resolutions.""" - outputs = { - "energy": ModelOutput(sample_kind="system"), - "o3::variance::energy": ModelOutput(sample_kind="atom"), - } - - with pytest.raises(ValueError, match="must use the same sample_kind"): - _group_output_requests(outputs) - - -@pytest.mark.parametrize( - ("o3_lambda", "expected"), - [ - (0, [0]), - (1, [-1, 0, 1]), - (2, [-2, -1, 0, 1, 2]), - ], -) -def test_o3_mu_labels(o3_lambda, expected): - """Spherical components should be ordered from -lambda to +lambda.""" - labels = _o3_mu_labels(o3_lambda, torch.device("cpu")) - - assert labels.names == ["o3_mu"] - assert labels.values[:, 0].tolist() == expected - assert labels.device == torch.device("cpu") - - def test_cartesian_vectors_to_spherical(): """Map Cartesian components to the real spherical l=1 ordering.""" values = torch.tensor( @@ -2984,7 +2526,7 @@ def test_cartesian_vectors_to_spherical_commutes_with_o3(inversion): def test_symmetric_matrices_to_spherical_known_components(): - """Identity, traceless diagonal, and skew matrices should map as expected.""" + """Known matrices map as expected and the symmetric 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 @@ -3001,18 +2543,15 @@ def test_symmetric_matrices_to_spherical_known_components(): assert torch.allclose(l0, expected_l0, rtol=0.0, atol=1.0e-12) assert torch.allclose(l2, expected_l2, rtol=0.0, atol=1.0e-12) - -def test_symmetric_matrices_to_spherical_preserves_norm(): - """The spherical norm should equal the symmetric-part Frobenius norm.""" generator = torch.Generator().manual_seed(1234) - matrices = torch.randn( + random_matrices = torch.randn( (4, 3, 3, 2), dtype=torch.float64, generator=generator, ) - symmetric = 0.5 * (matrices + matrices.transpose(1, 2)) + symmetric = 0.5 * (random_matrices + random_matrices.transpose(1, 2)) - l0, l2 = _symmetric_matrices_to_spherical(matrices) + l0, l2 = _symmetric_matrices_to_spherical(random_matrices) spherical_norm_squared = l0.square().sum(dim=1) + l2.square().sum(dim=1) cartesian_norm_squared = symmetric.square().sum(dim=(1, 2)) @@ -3069,71 +2608,6 @@ def test_symmetric_matrices_to_spherical_commutes_with_o3(inversion): assert torch.allclose(transformed_l2, expected_l2, rtol=0.0, atol=1.0e-12) -@pytest.mark.parametrize( - ( - "names", - "values", - "o3_lambda", - "o3_sigma", - "expected_names", - "expected_values", - ), - [ - (["_"], [[0]], 2, 1, ["o3_lambda", "o3_sigma"], [[2, 1]]), - ( - ["channel"], - [[3], [7]], - 1, - -1, - ["channel", "o3_lambda", "o3_sigma"], - [[3, 1, -1], [7, 1, -1]], - ), - ( - ["channel", "o3_lambda"], - [[3, 1], [7, 1]], - 1, - -1, - ["channel", "o3_lambda", "o3_sigma"], - [[3, 1, -1], [7, 1, -1]], - ), - ], -) -def test_add_o3_irrep_to_keys( - names, - values, - o3_lambda, - o3_sigma, - expected_names, - expected_values, -): - """Preserve semantic keys while assigning one O(3) irrep.""" - result = _add_o3_irrep_to_keys( - Labels(names, torch.tensor(values)), - o3_lambda, - o3_sigma, - ) - - assert result.names == expected_names - assert result.values.tolist() == expected_values - - -@pytest.mark.parametrize( - ("names", "values", "message"), - [ - (["_"], [[1]], "placeholder"), - (["channel", "o3_lambda"], [[3, 1], [7, 2]], "o3_lambda"), - ], -) -def test_add_o3_irrep_to_keys_rejects_conflicting_metadata(names, values, message): - """Reject an invalid ``_`` placeholder or conflicting irrep key values.""" - with pytest.raises(ValueError, match=message): - _add_o3_irrep_to_keys( - Labels(names, torch.tensor(values)), - o3_lambda=1, - o3_sigma=1, - ) - - @pytest.mark.parametrize( "source_name", [ @@ -3228,24 +2702,30 @@ def test_decompose_output_does_not_infer_custom_cartesian_semantics(): result = _decompose_output("mtt::custom", tensor) - assert result is tensor + mts.equal_raise(result, tensor) @pytest.mark.parametrize( ("source_name", "shape", "component_names", "message"), [ - ("energy", (1, 3, 1), ["xyz"], "must not have components"), + ( + "energy", + (1, 3, 1), + ["xyz"], + "energy-like outputs must not have components", + ), ( "non_conservative_force", (1, 3, 1), ["component"], - "one 'xyz' component axis", + "non_conservative_force must have one 'xyz' component axis of size 3", ), ( "non_conservative_stress", (1, 3, 3, 1), ["xyz_1", "component"], - "'xyz_1' and 'xyz_2' component axes", + "non_conservative_stress must have 'xyz_1' and 'xyz_2' component " + "axes of size 3", ), ], ) @@ -3261,37 +2741,55 @@ def test_decompose_output_rejects_invalid_standard_components( component_names, ) - with pytest.raises(ValueError, match=message): + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): _decompose_output(source_name, tensor) -def test_decompose_output_rejects_attached_gradients(): - """Decomposition should not silently discard explicit TensorBlock gradients.""" - properties = Labels.range("property", 1) - block = TensorBlock( - values=torch.ones((1, 1), dtype=torch.float64), - samples=Labels.range("system", 1), - 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, - ), - ) - tensor = TensorMap( - Labels("_", torch.tensor([[0]], dtype=torch.int64)), - [block], - ) - - with pytest.raises(ValueError, match="gradients attached to 'energy'"): - _decompose_output("energy", 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_o3_lambda_target=0, + max_o3_lambda_grid=2, + ) -def test_decompose_output_is_scriptable(): - """The output decomposition should compile with TorchScript.""" - torch.jit.script(_decompose_output) + 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, + ) From 6a3a4bbd8bd1230a95a70bf2cc105096de482b10 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Tue, 28 Jul 2026 15:12:59 +0200 Subject: [PATCH 06/18] Close coverage gaps and use physics terminology in docstrings --- .../torch/symmetrized_model/_model.py | 20 +-- .../torch/symmetrized_model/_projections.py | 5 +- .../tests/symmetrized_model.py | 139 +++++++++++++++++- 3 files changed, 150 insertions(+), 14 deletions(-) diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py index da8e86f8..74346f6b 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py @@ -524,20 +524,20 @@ class SymmetrizedModel(torch.nn.Module): :param model: underlying :py:class:`ModelInterface`. The :py:meth:`wrap` method obtains this module from :py:attr:`AtomisticModel.module`. - :param max_o3_lambda_target: largest spherical rank that can be transformed + :param max_o3_lambda_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_o3_lambda_input: largest spherical rank that can be rotated in + :param max_o3_lambda_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 the spherical ranks that + 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_o3_lambda_character: largest character sector included in character + :param max_o3_lambda_character: maximum angular momentum included in character projections. ``None`` disables character-projection outputs; zero enables the - scalar character sector only. + scalar (``o3_lambda = 0``) contribution only. :param batch_size: positive number of transformed systems evaluated in one call to ``model``. The default is 32. :param max_o3_lambda_grid: quadrature integration degree. If ``None``, use the @@ -666,12 +666,12 @@ def wrap( already been saved. :param model: the :py:class:`AtomisticModel` to wrap - :param max_o3_lambda_target: largest spherical rank accepted in + :param max_o3_lambda_target: maximum angular momentum accepted in already-spherical model outputs requested for averaging or variance - :param max_o3_lambda_input: largest spherical rank accepted in custom System - data - :param max_o3_lambda_character: largest character sector to report, or ``None`` - to disable character projections + :param max_o3_lambda_input: maximum angular momentum accepted in custom + System data + :param max_o3_lambda_character: maximum angular momentum in character + projections, or ``None`` to disable them :param batch_size: number of transformed Systems evaluated in one model call :param max_o3_lambda_grid: quadrature integration degree, selected automatically when ``None`` diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_projections.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_projections.py index d635f1c7..9a6f35b9 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_projections.py +++ b/python/metatomic_torch/metatomic/torch/symmetrized_model/_projections.py @@ -58,7 +58,8 @@ def _character_projection_coefficients_from_batch( inverse_wigner_matrices: List[torch.Tensor], input_system_index: int, ) -> TensorMap: - """Accumulate every character rank for one rotation batch.""" + """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 == ["_"]: @@ -140,7 +141,7 @@ def _character_projection_tensormap_from_cosets( proper_coefficients: TensorMap, improper_coefficients: TensorMap, ) -> TensorMap: - """Combine proper and improper coefficient TensorMaps into O(3) sectors.""" + """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") diff --git a/python/metatomic_torch/tests/symmetrized_model.py b/python/metatomic_torch/tests/symmetrized_model.py index 8c1096d3..6bb8e67f 100644 --- a/python/metatomic_torch/tests/symmetrized_model.py +++ b/python/metatomic_torch/tests/symmetrized_model.py @@ -1328,7 +1328,8 @@ def test_energy_results_match_analytic_values_and_reuse_predictions(self): "o3::character_projection::energy": ModelOutput(sample_kind="system"), } - result = model([system], outputs, None) + with torch.inference_mode(): + result = model([system], outputs, None) assert set(result) == set(outputs) n_rotations = len(model._rotation_matrices) @@ -1660,6 +1661,39 @@ def test_empty_selected_atoms_returns_empty_outputs(self): 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_o3_lambda_target=0, + max_o3_lambda_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]]) @@ -1668,6 +1702,7 @@ def test_equivariant_outputs_preserve_values_metadata_and_zero_variance(self): "non_conservative_force", "non_conservative_stress", "mtt::spherical_vector", + "mtt::spherical_quadrupole", ] outputs = { name: ModelOutput( @@ -1707,12 +1742,23 @@ def test_equivariant_outputs_preserve_values_metadata_and_zero_variance(self): 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 = _symmetric_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], [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] @@ -1756,6 +1802,31 @@ def test_dtype_and_implicit_autograd(self, dtype): 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_o3_lambda_target=0, + max_o3_lambda_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_invalid_requests_before_model_evaluation(self): """Invalid public requests should fail without running the source model.""" base_model = _CountingLinearEnergyModel() @@ -1802,6 +1873,57 @@ def test_rejects_invalid_requests_before_model_evaluation(self): ) assert base_model.call_count == 0 + def test_rejects_downcast_integration_buffers(self): + """Calling .float() on the module must fail loudly at the next forward.""" + model = SymmetrizedModel( + _LinearEnergyModel(), + max_o3_lambda_target=0, + max_o3_lambda_grid=2, + ).float() + + message = ( + "SymmetrizedModel integration buffers must remain float64, got " + "torch.float32; do not call .float() or .half() on the module" + ) + 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_model_that_omits_the_requested_output(self): + """Fail loudly when the underlying model does not return a source.""" + model = SymmetrizedModel( + _EmptyModel(), + max_o3_lambda_target=0, + max_o3_lambda_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_o3_lambda_target=0, + max_o3_lambda_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 = { @@ -1849,7 +1971,7 @@ def test_wrap_declares_capabilities(self, max_o3_lambda_character): } base = AtomisticModel( _EmptyModel().eval(), - ModelMetadata(), + ModelMetadata(name="wrapped source model"), ModelCapabilities( outputs=source_outputs, atomic_types=[1, 6, 8], @@ -1868,6 +1990,10 @@ def test_wrap_declares_capabilities(self, max_o3_lambda_character): ) 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) @@ -2091,6 +2217,15 @@ def test_saved_wrapper_runs_on_cuda(self, tmp_path): ) cuda_system = cpu_system.to(device=cuda_device) + cpu_module = SymmetrizedModel(_LinearEnergyModel(), max_o3_lambda_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", From 5ff2605766aa6d8333e927fcebc9491caee50bd4 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Thu, 30 Jul 2026 09:18:08 +0200 Subject: [PATCH 07/18] Address review --- .../src/torch/reference/symmetrized-model.rst | 23 +- metatomic-torch/CHANGELOG.md | 2 +- .../metatomic/torch/__init__.py | 1 + .../{symmetrized_model => o3}/_decompose.py | 107 +++++-- .../{symmetrized_model => o3}/_projections.py | 16 +- .../{symmetrized_model => o3}/_quadrature.py | 17 +- .../_model.py => o3/_symmetrized.py} | 210 +++++++++---- .../metatomic/torch/o3/_tranformations.py | 34 +-- .../torch/{symmetrized_model => o3}/_utils.py | 14 +- .../metatomic/torch/o3/_wigner.py | 69 +++++ .../torch/symmetrized_model/__init__.py | 14 - .../symmetrized_model/_wigner_storage.py | 59 ---- python/metatomic_torch/tests/o3.py | 2 +- .../tests/symmetrized_model.py | 284 ++++++++++++++---- 14 files changed, 597 insertions(+), 255 deletions(-) rename python/metatomic_torch/metatomic/torch/{symmetrized_model => o3}/_decompose.py (65%) rename python/metatomic_torch/metatomic/torch/{symmetrized_model => o3}/_projections.py (93%) rename python/metatomic_torch/metatomic/torch/{symmetrized_model => o3}/_quadrature.py (89%) rename python/metatomic_torch/metatomic/torch/{symmetrized_model/_model.py => o3/_symmetrized.py} (87%) rename python/metatomic_torch/metatomic/torch/{symmetrized_model => o3}/_utils.py (91%) delete mode 100644 python/metatomic_torch/metatomic/torch/symmetrized_model/__init__.py delete mode 100644 python/metatomic_torch/metatomic/torch/symmetrized_model/_wigner_storage.py diff --git a/docs/src/torch/reference/symmetrized-model.rst b/docs/src/torch/reference/symmetrized-model.rst index 897b9d02..9d4ba5cf 100644 --- a/docs/src/torch/reference/symmetrized-model.rst +++ b/docs/src/torch/reference/symmetrized-model.rst @@ -3,12 +3,13 @@ O(3)-symmetrized models ======================= -The :py:mod:`metatomic.torch.symmetrized_model` module wraps an exported +The :py:class:`metatomic.torch.SymmetrizedModel` class wraps an exported :py:class:`~metatomic.torch.AtomisticModel` with finite-quadrature O(3) -averaging and equivariance diagnostics. Ordinary outputs are averaged over -rotated and inverted copies of each input. Additional output names request an -equivariance variance or squared character-projection contributions of the -model response. +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 --------------- @@ -127,18 +128,18 @@ Quadrature ---------- The deterministic grid combines a Lebedev rule on the sphere, uniformly spaced -in-plane rotations, and both O(3) cosets. 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_o3_lambda_grid`` controls the quadrature +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_o3_lambda_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.symmetrized_model +.. py:currentmodule:: metatomic.torch .. autoclass:: SymmetrizedModel :members: - -.. autofunction:: get_rotation_quadrature diff --git a/metatomic-torch/CHANGELOG.md b/metatomic-torch/CHANGELOG.md index 2da208f3..3be532e1 100644 --- a/metatomic-torch/CHANGELOG.md +++ b/metatomic-torch/CHANGELOG.md @@ -18,7 +18,7 @@ a changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows ### Added -- Added `metatomic.torch.symmetrized_model` for finite-quadrature O(3) +- Added `metatomic.torch.SymmetrizedModel` for finite-quadrature O(3) averaging, equivariance variances, and character projections of exported atomistic models. 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/symmetrized_model/_decompose.py b/python/metatomic_torch/metatomic/torch/o3/_decompose.py similarity index 65% rename from python/metatomic_torch/metatomic/torch/symmetrized_model/_decompose.py rename to python/metatomic_torch/metatomic/torch/o3/_decompose.py index a133b271..7a2bc377 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_decompose.py +++ b/python/metatomic_torch/metatomic/torch/o3/_decompose.py @@ -1,10 +1,61 @@ +""" +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 +from typing import Dict, List import torch from metatensor.torch import Labels, TensorBlock, TensorMap +def _standard_quantity_categories() -> Dict[str, str]: + """Return the Cartesian layout of every decomposable standard quantity. + + 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 this module. + + 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_O3_LAMBDA_PER_CATEGORY: Dict[str, int] = { + "scalar": 0, + "cartesian_vector": 1, + "symmetric_matrix": 2, +} + + def _o3_mu_labels(o3_lambda: int, device: torch.device) -> Labels: """Return ``o3_mu`` labels from ``-o3_lambda`` through ``o3_lambda``.""" return Labels( @@ -31,6 +82,8 @@ def _symmetric_matrices_to_spherical( ) -> tuple[torch.Tensor, torch.Tensor]: """Return orthonormal l=0 and l=2 components of the symmetric matrix part. + ``values`` must have shape ``(n_samples, 3, 3, n_properties)``. + The antisymmetric (l=1) part is silently discarded. """ l0 = (values[:, 0, 0, :] + values[:, 1, 1, :] + values[:, 2, 2, :]).unsqueeze( @@ -53,31 +106,32 @@ def _symmetric_matrices_to_spherical( return l0, l2 -def _decompose_output( +def decompose_output( source_name: str, tensor: TensorMap, ) -> TensorMap: - """Decompose standard outputs for variance and character projection.""" + """Decompose standard outputs for variance and character projection. + + This takes the standard Cartesian or scalar 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 = source_name.split("/", 1)[0] - is_energy = quantity in ( - "energy", - "energy_ensemble", - "energy_uncertainty", - ) - is_force = quantity in ( - "non_conservative_force", - "non_conservative_forces", - ) - is_stress = quantity == "non_conservative_stress" - if not (is_energy or is_force or is_stress): + categories = _standard_quantity_categories() + if quantity not in categories: return tensor + category = categories[quantity] - if is_energy: - energy_blocks: List[TensorBlock] = [] + if category == "scalar": + scalar_blocks: List[TensorBlock] = [] for block in tensor.blocks(): if len(block.components) != 0: - raise ValueError("energy-like outputs must not have components") - energy_blocks.append( + raise ValueError(f"'{quantity}' outputs must not have components") + scalar_blocks.append( TensorBlock( values=block.values.unsqueeze(1), samples=block.samples, @@ -87,11 +141,11 @@ def _decompose_output( ) result = TensorMap( _add_o3_irrep_to_keys(tensor.keys, 0, 1), - energy_blocks, + scalar_blocks, ) - elif is_force: - force_blocks: List[TensorBlock] = [] + elif category == "cartesian_vector": + vector_blocks: List[TensorBlock] = [] for block in tensor.blocks(): if ( len(block.components) != 1 @@ -99,10 +153,9 @@ def _decompose_output( or len(block.components[0]) != 3 ): raise ValueError( - "non_conservative_force must have one 'xyz' component axis " - "of size 3" + f"'{quantity}' must have one 'xyz' component axis of size 3" ) - force_blocks.append( + vector_blocks.append( TensorBlock( values=_cartesian_vectors_to_spherical(block.values, 1), samples=block.samples, @@ -112,7 +165,7 @@ def _decompose_output( ) result = TensorMap( _add_o3_irrep_to_keys(tensor.keys, 1, 1), - force_blocks, + vector_blocks, ) else: @@ -127,8 +180,8 @@ def _decompose_output( or len(block.components[1]) != 3 ): raise ValueError( - "non_conservative_stress must have 'xyz_1' and 'xyz_2' " - "component axes of size 3" + f"'{quantity}' must have 'xyz_1' and 'xyz_2' component axes " + "of size 3" ) values_l0, values_l2 = _symmetric_matrices_to_spherical(block.values) diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_projections.py b/python/metatomic_torch/metatomic/torch/o3/_projections.py similarity index 93% rename from python/metatomic_torch/metatomic/torch/symmetrized_model/_projections.py rename to python/metatomic_torch/metatomic/torch/o3/_projections.py index 9a6f35b9..e4cd8e51 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_projections.py +++ b/python/metatomic_torch/metatomic/torch/o3/_projections.py @@ -1,7 +1,9 @@ """Character-projection helpers. The projected quantity is defined in the :py:class:`SymmetrizedModel` class -docstring. +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 @@ -10,8 +12,8 @@ from metatensor.torch import Labels, TensorBlock, TensorMap from ._utils import ( - _group_samples_by_rotated_copy, - _restore_input_system_to_samples, + group_samples_by_rotated_copy, + restore_input_system_to_samples, ) @@ -52,7 +54,7 @@ def _character_projections_from_proper_and_improper_coefficients( ) -def _character_projection_coefficients_from_batch( +def character_projection_coefficients_from_batch( tensor: TensorMap, weights: torch.Tensor, inverse_wigner_matrices: List[torch.Tensor], @@ -80,11 +82,11 @@ def _character_projection_coefficients_from_batch( 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( + values, sample_names, sample_values = group_samples_by_rotated_copy( block, n_rotated_copies, ) - samples = _restore_input_system_to_samples( + samples = restore_input_system_to_samples( sample_names, sample_values, input_system_index, @@ -137,7 +139,7 @@ def _character_projection_coefficients_from_batch( return TensorMap(Labels(key_names + ["chi_lambda"], values), blocks) -def _character_projection_tensormap_from_cosets( +def character_projection_tensormap_from_cosets( proper_coefficients: TensorMap, improper_coefficients: TensorMap, ) -> TensorMap: diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_quadrature.py b/python/metatomic_torch/metatomic/torch/o3/_quadrature.py similarity index 89% rename from python/metatomic_torch/metatomic/torch/symmetrized_model/_quadrature.py rename to python/metatomic_torch/metatomic/torch/o3/_quadrature.py index ea5e4554..eee58024 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_quadrature.py +++ b/python/metatomic_torch/metatomic/torch/o3/_quadrature.py @@ -1,6 +1,13 @@ +""" +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 +from ._utils import validate_integer _LEBEDEV_ORDERS = ( @@ -52,7 +59,7 @@ def _import_scipy(): return lebedev_rule, Rotation -def _choose_quadrature(L_max: int) -> tuple[int, int]: +def choose_quadrature(L_max: int) -> tuple[int, int]: """ Choose a Lebedev quadrature order and number of in-plane rotations to integrate spherical harmonics up to degree ``L_max``. @@ -60,7 +67,7 @@ def _choose_quadrature(L_max: int) -> tuple[int, int]: :param L_max: maximum spherical harmonic degree :return: (lebedev_order, n_inplane_rotations) """ - L_max = _validate_integer("L_max", L_max, 0) + L_max = validate_integer("L_max", L_max, 0) if L_max > _LEBEDEV_ORDERS[-1]: raise ValueError( f"the requested quadrature degree L_max={L_max} exceeds the largest " @@ -87,8 +94,8 @@ def get_euler_angles_quadrature( 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) + 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 " diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py similarity index 87% rename from python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py rename to python/metatomic_torch/metatomic/torch/o3/_symmetrized.py index 74346f6b..8b19bd29 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_model.py +++ b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py @@ -1,3 +1,10 @@ +""" +: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. +""" + from typing import Dict, List, Optional, Tuple import metatensor.torch as mts @@ -14,28 +21,56 @@ register_autograd_neighbors, ) -from ..o3._tranformations import ( - _max_o3_lambda_in_tensor, - _transform_tensor_with_precomputed_matrices, +from ._decompose import ( + MAX_O3_LAMBDA_PER_CATEGORY, + STANDARD_QUANTITY_CATEGORIES, + decompose_output, ) -from ._decompose import _decompose_output from ._projections import ( - _character_projection_coefficients_from_batch, - _character_projection_tensormap_from_cosets, + character_projection_coefficients_from_batch, + character_projection_tensormap_from_cosets, +) +from ._quadrature import choose_quadrature, get_rotation_quadrature +from ._tranformations import ( + _max_o3_lambda_in_tensor, + _transform_tensor_with_precomputed_matrices, ) -from ._quadrature import _choose_quadrature, get_rotation_quadrature from ._utils import ( - _group_samples_by_rotated_copy, - _map_selected_atoms_to_rotated_copies, - _restore_input_system_to_samples, - _validate_integer, + group_samples_by_rotated_copy, + map_selected_atoms_to_rotated_copies, + restore_input_system_to_samples, + validate_integer, ) -from ._wigner_storage import ( - _build_packed_wigner_matrices, - _wigner_matrices_for_lambda, +from ._wigner import ( + build_packed_wigner_matrices, + wigner_matrices_for_lambda, ) +# deprecated quantity names mapped to their current name, mirroring +# ``AtomisticModel._new_names``. ``AtomisticModel`` advertises both spellings of +# each standard output, so an engine can request either one from the wrapper; +# everything inside the wrapper uses the current names only. +_NEW_NAMES: Dict[str, str] = { + "features": "feature", + "non_conservative_forces": "non_conservative_force", + "positions": "position", + "momenta": "momentum", + "masses": "mass", + "velocities": "velocity", + "charges": "charge", +} + + +def _use_new_quantity_name(name: str, new_names: Dict[str, str]) -> str: + """Replace a deprecated base quantity in ``name`` with its current name.""" + parts = name.split("/") + if parts[0] in new_names: + parts[0] = new_names[parts[0]] + return "/".join(parts) + return name + + def _transform_system_geometry_batch( system: System, matrices: torch.Tensor, @@ -161,15 +196,35 @@ def _parse_output_request(requested_name: str) -> Tuple[str, str]: return source_name, calculation +def _record_output_request( + names: Dict[str, str], + source_name: str, + requested_name: str, +) -> None: + """Register the public name a source output must be returned under.""" + if source_name in names: + raise ValueError( + f"'{requested_name}' and '{names[source_name]}' request the same " + f"'{source_name}' output; only use the new name" + ) + names[source_name] = requested_name + + def _group_output_requests( outputs: Dict[str, ModelOutput], + new_names: Dict[str, str], ) -> Tuple[ Dict[str, str], Dict[str, str], Dict[str, str], Dict[str, str], ]: - """Group public requests by underlying output and calculation.""" + """Group public requests by underlying output and calculation. + + Deprecated quantity names are translated here, so the rest of the wrapper + only ever sees the current names. 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] = {} @@ -177,6 +232,7 @@ def _group_output_requests( for requested_name, output in outputs.items(): source_name, calculation = _parse_output_request(requested_name) + source_name = _use_new_quantity_name(source_name, new_names) sample_kind = output.sample_kind if source_name in source_sample_kinds: previous_sample_kind = source_sample_kinds[source_name] @@ -189,11 +245,15 @@ def _group_output_requests( source_sample_kinds[source_name] = sample_kind if calculation == "average": - average_names[source_name] = requested_name + _record_output_request(average_names, source_name, requested_name) elif calculation == "variance": - variance_names[source_name] = requested_name + _record_output_request(variance_names, source_name, requested_name) else: - character_projection_names[source_name] = requested_name + _record_output_request( + character_projection_names, + source_name, + requested_name, + ) return ( source_sample_kinds, @@ -203,6 +263,30 @@ def _group_output_requests( ) +def _infer_max_o3_lambda( + names: Dict[str, ModelOutput], + kind: str, + argument: str, +) -> int: + """Guess an angular-momentum limit from standard quantity names.""" + max_o3_lambda = 0 + for name in names.keys(): + quantity = _use_new_quantity_name(name, _NEW_NAMES).split("/")[0] + if quantity == "feature": + # features are not an irreducible representation of O(3): they are + # passed through unchanged and never rotated back + continue + if quantity not in STANDARD_QUANTITY_CATEGORIES: + raise ValueError( + f"unable to guess {argument} from the non-standard {kind} " + f"'{name}', please set {argument} explicitly" + ) + category = STANDARD_QUANTITY_CATEGORIES[quantity] + max_o3_lambda = max(max_o3_lambda, MAX_O3_LAMBDA_PER_CATEGORY[category]) + + return max_o3_lambda + + def _reduce_weighted_centered_batch( tensor: TensorMap, weights: torch.Tensor, @@ -224,7 +308,7 @@ def _reduce_weighted_centered_batch( reference_blocks: List[TensorBlock] = [] for key, block in tensor.items(): - values, sample_names, sample_values = _group_samples_by_rotated_copy( + values, sample_names, sample_values = group_samples_by_rotated_copy( block, n_rotated_copies ) if reference is None: @@ -254,7 +338,7 @@ def _reduce_weighted_centered_batch( dim=0, ) - samples = _restore_input_system_to_samples( + samples = restore_input_system_to_samples( sample_names, sample_values, input_system_index, @@ -549,6 +633,7 @@ class SymmetrizedModel(torch.nn.Module): """ max_o3_lambda_character: Optional[int] + _new_names: Dict[str, str] _requested_inputs: Dict[str, ModelOutput] _requested_neighbor_lists: List[NeighborListOptions] @@ -564,20 +649,22 @@ def __init__( super().__init__() self._model = model + # TorchScript cannot read a module-level dictionary from ``forward`` + self._new_names = dict(_NEW_NAMES) self._requested_inputs = {} self._requested_neighbor_lists = [] - self.max_o3_lambda_target = _validate_integer( + self.max_o3_lambda_target = validate_integer( "max_o3_lambda_target", max_o3_lambda_target, 0 ) - self.max_o3_lambda_input = _validate_integer( + self.max_o3_lambda_input = validate_integer( "max_o3_lambda_input", max_o3_lambda_input, 0 ) if max_o3_lambda_character is not None: - max_o3_lambda_character = _validate_integer( + max_o3_lambda_character = validate_integer( "max_o3_lambda_character", max_o3_lambda_character, 0 ) self.max_o3_lambda_character = max_o3_lambda_character - self.batch_size = _validate_integer("batch_size", batch_size, 1) + self.batch_size = validate_integer("batch_size", batch_size, 1) if max_o3_lambda_grid is None: max_o3_lambda_grid = 2 * self.max_o3_lambda_target + 1 @@ -587,7 +674,7 @@ def __init__( 2 * self.max_o3_lambda_character, ) else: - max_o3_lambda_grid = _validate_integer( + max_o3_lambda_grid = validate_integer( "max_o3_lambda_grid", max_o3_lambda_grid, 0 ) if ( @@ -610,7 +697,7 @@ def __init__( if device.type != "cpu" and device.type != "cuda": raise ValueError("SymmetrizedModel supports CPU and CUDA execution") - lebedev_order, n_rotations = _choose_quadrature(self.max_o3_lambda_grid) + lebedev_order, n_rotations = choose_quadrature(self.max_o3_lambda_grid) rotations, weights = get_rotation_quadrature( lebedev_order, n_rotations, @@ -629,7 +716,7 @@ def __init__( self.max_o3_lambda_target, 0 if self.max_o3_lambda_character is None else self.max_o3_lambda_character, ) - packed_wigner_matrices = _build_packed_wigner_matrices( + packed_wigner_matrices = build_packed_wigner_matrices( rotation_matrices, max_o3_lambda_wigner, ) @@ -642,8 +729,8 @@ def __init__( def wrap( model: AtomisticModel, *, - max_o3_lambda_target: int, - max_o3_lambda_input: int = 0, + max_o3_lambda_target: Optional[int] = None, + max_o3_lambda_input: Optional[int] = None, max_o3_lambda_character: Optional[int] = None, batch_size: int = 32, max_o3_lambda_grid: Optional[int] = None, @@ -667,9 +754,14 @@ def wrap( :param model: the :py:class:`AtomisticModel` to wrap :param max_o3_lambda_target: maximum angular momentum accepted in - already-spherical model outputs requested for averaging or variance + already-spherical model outputs requested for averaging or variance. + When ``None``, it is guessed from the standard quantities declared by + ``model``; a non-standard output makes the guess impossible and must + be answered with an explicit value. :param max_o3_lambda_input: maximum angular momentum accepted in custom - System data + System data. When ``None``, it is guessed from the standard + quantities in ``model.requested_inputs()``, with the same + restriction on non-standard inputs. :param max_o3_lambda_character: maximum angular momentum in character projections, or ``None`` to disable them :param batch_size: number of transformed Systems evaluated in one model call @@ -691,6 +783,19 @@ def wrap( "wrapped model declares " + str(capabilities.supported_devices) ) + if max_o3_lambda_target is None: + max_o3_lambda_target = _infer_max_o3_lambda( + capabilities.outputs, + "output", + "max_o3_lambda_target", + ) + if max_o3_lambda_input is None: + max_o3_lambda_input = _infer_max_o3_lambda( + model.requested_inputs(use_new_names=True), + "input", + "max_o3_lambda_input", + ) + outputs: Dict[str, ModelOutput] = {} # private field: the as-declared output names, deliberately without the # deprecation aliases added by the public accessors @@ -812,7 +917,7 @@ def forward( average_names, variance_names, character_projection_names, - ) = _group_output_requests(outputs) + ) = _group_output_requests(outputs, self._new_names) if ( len(character_projection_names) != 0 and self.max_o3_lambda_character is None @@ -906,23 +1011,14 @@ def _evaluate_system( if configured_character_max is not None: character_max = configured_character_max - average_references = torch.jit.annotate(Dict[str, TensorMap], {}) - average_first_moments = torch.jit.annotate(Dict[str, TensorMap], {}) - variance_references = torch.jit.annotate(Dict[str, TensorMap], {}) - variance_first_moments = torch.jit.annotate(Dict[str, TensorMap], {}) - variance_second_moments = torch.jit.annotate(Dict[str, TensorMap], {}) - variance_absolute_second_moments = torch.jit.annotate( - Dict[str, TensorMap], - {}, - ) - proper_character_coefficients = torch.jit.annotate( - Dict[str, TensorMap], - {}, - ) - improper_character_coefficients = torch.jit.annotate( - Dict[str, TensorMap], - {}, - ) + 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 @@ -932,7 +1028,7 @@ def _evaluate_system( 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( + local_selected_atoms = map_selected_atoms_to_rotated_copies( selected_atoms, input_system_index, n_rotated_copies, @@ -941,7 +1037,7 @@ def _evaluate_system( input_wigner_matrices: List[torch.Tensor] = [] for o3_lambda in range(self.max_o3_lambda_input + 1): input_wigner_matrices.append( - _wigner_matrices_for_lambda( + wigner_matrices_for_lambda( self._packed_wigner_matrices, n_rotations, o3_lambda, @@ -955,7 +1051,7 @@ def _evaluate_system( if needs_backrotation: for o3_lambda in range(self.max_o3_lambda_target + 1): inverse_target_wigner_matrices.append( - _wigner_matrices_for_lambda( + wigner_matrices_for_lambda( self._packed_wigner_matrices, n_rotations, o3_lambda, @@ -966,7 +1062,7 @@ def _evaluate_system( if len(character_projection_names) != 0: for chi_lambda in range(character_max + 1): inverse_character_wigner_matrices.append( - _wigner_matrices_for_lambda( + wigner_matrices_for_lambda( self._packed_wigner_matrices, n_rotations, chi_lambda, @@ -1061,7 +1157,7 @@ def _evaluate_system( ) if source_name in variance_names: - diagnostic_tensor = _decompose_output( + diagnostic_tensor = decompose_output( source_name, backrotated, ) @@ -1100,8 +1196,8 @@ def _evaluate_system( ) if source_name in character_projection_names: - direct_tensor = _decompose_output(source_name, tensor) - contribution = _character_projection_coefficients_from_batch( + direct_tensor = decompose_output(source_name, tensor) + contribution = character_projection_coefficients_from_batch( direct_tensor, so3_weights, inverse_character_wigner_matrices, @@ -1150,7 +1246,7 @@ def _evaluate_system( ) for source_name, requested_name in character_projection_names.items(): - projection = _character_projection_tensormap_from_cosets( + projection = character_projection_tensormap_from_cosets( proper_character_coefficients[source_name], improper_character_coefficients[source_name], ) diff --git a/python/metatomic_torch/metatomic/torch/o3/_tranformations.py b/python/metatomic_torch/metatomic/torch/o3/_tranformations.py index c0a4736b..95a7c7f4 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_tranformations.py +++ b/python/metatomic_torch/metatomic/torch/o3/_tranformations.py @@ -9,7 +9,7 @@ from metatensor.torch import Labels, LabelsEntry, TensorBlock, TensorMap from .. import System, register_autograd_neighbors -from ._wigner import build_wigner_D_cache +from ._wigner import build_packed_wigner_matrices, wigner_matrices_for_lambda _INTEGER_DTYPES = ( @@ -185,7 +185,7 @@ def __init__(self, matrix: torch.Tensor, max_angular_momentum: int): 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 + self._packed_wigner_D: torch.Tensor | None = None @classmethod def _create_no_checks( @@ -206,30 +206,28 @@ def _create_no_checks( transformation._matrix = matrix transformation._max_angular_momentum = max_angular_momentum transformation._is_improper = is_improper - transformation._wigner_D_cache = None + transformation._packed_wigner_D = 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( + def _ensure_wigner_D_cache(self) -> torch.Tensor: + """Ensure that the packed Wigner-D cache has been built and return it. + + The packed buffer holds every ``ell`` up to ``max_angular_momentum``; it + inherits the dtype and device of the transformation matrix. + """ + if self._packed_wigner_D is None: + self._packed_wigner_D = build_packed_wigner_matrices( + self._matrix.unsqueeze(0), self._max_angular_momentum, - self._matrix, - device=self._matrix.device, - dtype=self._matrix.dtype, ) - return self._wigner_D_cache + return self._packed_wigner_D 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 + return wigner_matrices_for_lambda(self._ensure_wigner_D_cache(), 1, ell)[0] @property def matrix(self) -> torch.Tensor: @@ -606,7 +604,7 @@ def _validate_component_axis_metadata( def _max_o3_lambda_in_tensor(tensor: TensorMap) -> int: - """Return the largest spherical rank in block values or attached gradients. + """Return the largest angular momentum in block values or attached gradients. A TensorMap containing only scalar or Cartesian component axes returns ``-1``. """ @@ -869,7 +867,7 @@ def _transform_component_values_with_precomputed_matrices( 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") + raise ValueError("angular momentum exceeds the Wigner-D storage") axis_matrices = wigner_matrices[ell] parity *= _spherical_parity_factor(ell, sigma, is_improper) else: diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_utils.py b/python/metatomic_torch/metatomic/torch/o3/_utils.py similarity index 91% rename from python/metatomic_torch/metatomic/torch/symmetrized_model/_utils.py rename to python/metatomic_torch/metatomic/torch/o3/_utils.py index acae2e42..a57f033b 100644 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_utils.py +++ b/python/metatomic_torch/metatomic/torch/o3/_utils.py @@ -1,3 +1,9 @@ +""" +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 @@ -5,7 +11,7 @@ from metatensor.torch import Labels, TensorBlock -def _validate_integer(name: str, value, minimum: int) -> int: +def validate_integer(name: str, value, minimum: int) -> int: """Check that ``value`` is an integer at least ``minimum``. Return it as a Python ``int``. @@ -19,7 +25,7 @@ def _validate_integer(name: str, value, minimum: int) -> int: return integer_value -def _map_selected_atoms_to_rotated_copies( +def map_selected_atoms_to_rotated_copies( selected_atoms: Optional[Labels], input_system_index: int, n_rotated_copies: int, @@ -47,7 +53,7 @@ def _map_selected_atoms_to_rotated_copies( return Labels(list(selected_atoms.names), rotated_values) -def _group_samples_by_rotated_copy( +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.""" @@ -114,7 +120,7 @@ def _group_samples_by_rotated_copy( ) -def _restore_input_system_to_samples( +def restore_input_system_to_samples( sample_names: List[str], sample_values: torch.Tensor, input_system_index: int, diff --git a/python/metatomic_torch/metatomic/torch/o3/_wigner.py b/python/metatomic_torch/metatomic/torch/o3/_wigner.py index 8631a71d..a2a102d5 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_wigner.py +++ b/python/metatomic_torch/metatomic/torch/o3/_wigner.py @@ -119,3 +119,72 @@ def build_wigner_D_cache( cache = _compute_real_wigner_d_matrices(o3_lambda_max, angles, complex_to_real) return {ell: tensor.to(device=device, dtype=dtype) for ell, tensor in cache.items()} + + +def build_packed_wigner_matrices( + matrices: torch.Tensor, + max_o3_lambda: int, +) -> torch.Tensor: + """Build and pack proper Wigner-D matrices through ``max_o3_lambda``. + + :param matrices: ``(n_matrices, 3, 3)`` stack of O(3) matrices + :param max_o3_lambda: 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_o3_lambda + 1) * (2 * max_o3_lambda + 1) * (2 * max_o3_lambda + 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_o3_lambda, + matrix, + device=cpu, + dtype=output_dtype, + ) + for o3_lambda in range(max_o3_lambda + 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/metatomic/torch/symmetrized_model/__init__.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/__init__.py deleted file mode 100644 index 0aa60093..00000000 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -""" -O(3) averaging and equivariance diagnostics for atomistic models. - -See :py:class:`SymmetrizedModel` for the method and public output conventions. -""" - -from ._model import SymmetrizedModel -from ._quadrature import get_rotation_quadrature - - -__all__ = [ - "SymmetrizedModel", - "get_rotation_quadrature", -] diff --git a/python/metatomic_torch/metatomic/torch/symmetrized_model/_wigner_storage.py b/python/metatomic_torch/metatomic/torch/symmetrized_model/_wigner_storage.py deleted file mode 100644 index 1980458a..00000000 --- a/python/metatomic_torch/metatomic/torch/symmetrized_model/_wigner_storage.py +++ /dev/null @@ -1,59 +0,0 @@ -import torch - -from ..o3 import O3Transformation - - -def _build_packed_wigner_matrices( - matrices: torch.Tensor, - max_o3_lambda: int, -) -> torch.Tensor: - """Build and pack proper Wigner-D matrices through ``max_o3_lambda``.""" - output_device = matrices.device - output_dtype = matrices.dtype - calculation_matrices = matrices.detach().to(device="cpu") - n_matrices = matrices.size(0) - n_elements_per_matrix = ( - (max_o3_lambda + 1) * (2 * max_o3_lambda + 1) * (2 * max_o3_lambda + 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)): - transformation = O3Transformation(matrix, max_o3_lambda) - for o3_lambda in range(max_o3_lambda + 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_( - transformation.wigner_D_matrix(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..390bfe2c 100644 --- a/python/metatomic_torch/tests/o3.py +++ b/python/metatomic_torch/tests/o3.py @@ -1427,7 +1427,7 @@ 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("angular momentum exceeds the Wigner-D storage") with pytest.raises(ValueError, match=f"^{message}$"): _transform_tensor_with_precomputed_matrices( unavailable_rank, diff --git a/python/metatomic_torch/tests/symmetrized_model.py b/python/metatomic_torch/tests/symmetrized_model.py index 6bb8e67f..dfba550e 100644 --- a/python/metatomic_torch/tests/symmetrized_model.py +++ b/python/metatomic_torch/tests/symmetrized_model.py @@ -14,21 +14,28 @@ ModelMetadata, ModelOutput, NeighborListOptions, + SymmetrizedModel, System, load_atomistic_model, ) from metatomic.torch.o3 import O3Transformation, transform_system -from metatomic.torch.symmetrized_model import ( - SymmetrizedModel, - get_rotation_quadrature, -) -from metatomic.torch.symmetrized_model._decompose import ( +from metatomic.torch.o3._decompose import ( _cartesian_vectors_to_spherical, - _decompose_output, _o3_mu_labels, _symmetric_matrices_to_spherical, + decompose_output, ) -from metatomic.torch.symmetrized_model._model import ( +from metatomic.torch.o3._projections import ( + _character_projection_coefficients_from_rotation_batch, + _character_projections_from_proper_and_improper_coefficients, +) +from metatomic.torch.o3._quadrature import ( + _rotations_from_euler_angles, + choose_quadrature, + get_euler_angles_quadrature, + get_rotation_quadrature, +) +from metatomic.torch.o3._symmetrized import ( _clamp_roundoff_negative_diagnostic, _component_norm_squared, _mean_variance_over_components, @@ -37,22 +44,13 @@ _transform_system_geometry_batch, _variance_from_centered_moments, ) -from metatomic.torch.symmetrized_model._projections import ( - _character_projection_coefficients_from_rotation_batch, - _character_projections_from_proper_and_improper_coefficients, -) -from metatomic.torch.symmetrized_model._quadrature import ( - _choose_quadrature, - _rotations_from_euler_angles, - get_euler_angles_quadrature, -) -from metatomic.torch.symmetrized_model._utils import ( - _group_samples_by_rotated_copy, - _map_selected_atoms_to_rotated_copies, +from metatomic.torch.o3._utils import ( + group_samples_by_rotated_copy, + map_selected_atoms_to_rotated_copies, ) -from metatomic.torch.symmetrized_model._wigner_storage import ( - _build_packed_wigner_matrices, - _wigner_matrices_for_lambda, +from metatomic.torch.o3._wigner import ( + build_packed_wigner_matrices, + wigner_matrices_for_lambda, ) @@ -367,6 +365,46 @@ def forward( 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.""" @@ -719,12 +757,12 @@ def test_transforms_spherical_custom_data(self, is_improper): dtype=torch.float64, ) matrices = -proper_matrices if is_improper else proper_matrices - packed_wigner = _build_packed_wigner_matrices( + packed_wigner = build_packed_wigner_matrices( proper_matrices, max_o3_lambda=1, ) wigner_matrices = [ - _wigner_matrices_for_lambda( + wigner_matrices_for_lambda( packed_wigner, n_matrices=len(matrices), o3_lambda=o3_lambda, @@ -790,12 +828,12 @@ def test_input_limit_distinguishes_spherical_from_cartesian(self): [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=torch.float64, ).unsqueeze(0) - packed_wigner = _build_packed_wigner_matrices( + packed_wigner = build_packed_wigner_matrices( matrix, max_o3_lambda=0, ) wigner_matrices = [ - _wigner_matrices_for_lambda( + wigner_matrices_for_lambda( packed_wigner, n_matrices=1, o3_lambda=0, @@ -1020,7 +1058,7 @@ def test_packed_matrices_match_o3(self, dtype): ) max_o3_lambda = 2 - packed = _build_packed_wigner_matrices(matrices, max_o3_lambda) + packed = build_packed_wigner_matrices(matrices, max_o3_lambda) assert packed.dim() == 1 assert packed.numel() == len(matrices) * sum( @@ -1033,7 +1071,7 @@ def test_packed_matrices_match_o3(self, dtype): O3Transformation(matrix, max_o3_lambda) for matrix in matrices.unbind(0) ] for o3_lambda in range(max_o3_lambda + 1): - actual = _wigner_matrices_for_lambda( + actual = wigner_matrices_for_lambda( packed, len(matrices), o3_lambda, @@ -1050,7 +1088,7 @@ def test_rank_view_rejects_out_of_range_lambda(self): """Rank views should reject ranks beyond the packed storage.""" message = "o3_lambda exceeds the packed Wigner-D storage" with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - _wigner_matrices_for_lambda(torch.empty(1), 1, 1) + wigner_matrices_for_lambda(torch.empty(1), 1, 1) class TestQuadrature: @@ -1059,7 +1097,7 @@ class TestQuadrature: 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) + 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 @@ -1069,7 +1107,7 @@ def test_weights_sum(self): 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) + 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() @@ -1095,15 +1133,15 @@ def test_quadrature_validation(self): "available Lebedev order (131)" ) with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - _choose_quadrature(132) + choose_quadrature(132) message = "L_max must be non-negative, got -1" with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - _choose_quadrature(-1) + choose_quadrature(-1) message = "L_max must be an integer, got float" with pytest.raises(TypeError, match=f"^{re.escape(message)}$"): - _choose_quadrature(1.5) + choose_quadrature(1.5) message = "n_rotations must be positive, got 0" with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): @@ -1124,7 +1162,7 @@ def test_quadrature_validation(self): get_rotation_quadrature(4, 3) def test_degree_two_grid_resolves_l1_products(self): - order, n_rotations = _choose_quadrature(2) + order, n_rotations = choose_quadrature(2) rotations, weights = get_rotation_quadrature(order, n_rotations) function = rotations[:, 2, 0] @@ -1584,6 +1622,68 @@ def test_preserves_variant_and_custom_output_names(self, source_name): atol=1.0e-12, ) + def test_deprecated_quantity_names_are_normalized(self): + """A deprecated request is decomposed as, and returned under, its own name.""" + model = SymmetrizedModel( + _EquivariantOutputModel(), + max_o3_lambda_target=1, + max_o3_lambda_grid=2, + batch_size=5, + ) + outputs = { + "non_conservative_forces": ModelOutput(sample_kind="atom"), + "o3::variance::non_conservative_forces": ModelOutput(sample_kind="atom"), + } + system = _forward_test_system([[1.0, 2.0, 3.0]]) + + result = model([system], outputs, None) + + assert set(result) == set(outputs) + assert torch.allclose( + result["non_conservative_forces"].block().values.squeeze(-1), + system.positions, + atol=1.0e-12, + ) + # the l=1 keys prove the decomposition recognized the singular quantity + variance = result["o3::variance::non_conservative_forces"] + assert variance.keys.values.tolist() == [[1, 1]] + assert torch.allclose( + variance.block().values, + torch.zeros_like(variance.block().values), + 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_o3_lambda_target=0, + max_o3_lambda_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 = [ @@ -2019,6 +2119,84 @@ def test_wrap_declares_capabilities(self, max_o3_lambda_character): else: assert capabilities.outputs[character_name].unit == squared_unit + @pytest.mark.parametrize( + ("outputs", "expected_max_o3_lambda_target"), + [ + ({"energy": ModelOutput(unit="eV", sample_kind="system")}, 0), + ({"feature": ModelOutput(sample_kind="atom")}, 0), + ( + { + "energy": ModelOutput(unit="eV", sample_kind="system"), + "non_conservative_force": ModelOutput( + unit="eV/A", + sample_kind="atom", + ), + }, + 1, + ), + ( + { + "non_conservative_stress": ModelOutput( + unit="eV/A^3", + sample_kind="system", + ) + }, + 2, + ), + ], + ) + def test_guesses_limits_from_standard_quantities( + self, + outputs, + expected_max_o3_lambda_target, + ): + """Both limits default to what the standard quantities require.""" + + class _VelocityInputModel(_EmptyModel): + def requested_inputs(self) -> Dict[str, ModelOutput]: + return {"velocity": ModelOutput(sample_kind="atom")} + + 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_o3_lambda_grid=2) + + assert wrapped.module.max_o3_lambda_target == expected_max_o3_lambda_target + # velocity is a Cartesian vector + assert wrapped.module.max_o3_lambda_input == 1 + + def test_rejects_guessing_a_limit_from_a_custom_output(self): + """A non-standard output must be answered with an explicit limit.""" + base = AtomisticModel( + _EmptyModel().eval(), + ModelMetadata(), + ModelCapabilities( + outputs={"mtt::custom": ModelOutput(sample_kind="system")}, + atomic_types=[1], + interaction_range=0.0, + length_unit="A", + supported_devices=["cpu"], + dtype="float64", + ), + ) + + message = ( + "unable to guess max_o3_lambda_target from the non-standard output " + "'mtt::custom', please set max_o3_lambda_target explicitly" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + SymmetrizedModel.wrap(base) + @pytest.mark.parametrize( "source_name", [ @@ -2099,6 +2277,9 @@ def test_preserves_requirements_and_runs_after_save_load(self, tmp_path): wrapped = SymmetrizedModel.wrap( loaded_base, max_o3_lambda_target=0, + # 'mtt::linear' and 'mtt::field' are not standard quantities, so + # both limits have to be given explicitly + max_o3_lambda_input=0, max_o3_lambda_character=1, max_o3_lambda_grid=2, batch_size=5, @@ -2272,7 +2453,7 @@ def test_system_column_found_by_name(self): # the rotated-copy index must go into the "system" column wherever it # is, not positionally into column 0 selection = Labels(["atom", "system"], torch.tensor([[3, 0], [5, 0]])) - rotated = _map_selected_atoms_to_rotated_copies(selection, 0, 2) + 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] @@ -2308,7 +2489,7 @@ def test_rotated_copy_layout_rejects_inconsistent_samples(sample_values, message ) with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - _group_samples_by_rotated_copy(block, n_rotated_copies=2) + group_samples_by_rotated_copy(block, n_rotated_copies=2) @pytest.mark.parametrize( @@ -2340,7 +2521,7 @@ def test_group_samples_by_rotated_copy( properties=Labels.range("property", 1), ) - grouped_values, shared_names, shared_values = _group_samples_by_rotated_copy( + grouped_values, shared_names, shared_values = group_samples_by_rotated_copy( block, n_rotated_copies ) @@ -2750,15 +2931,16 @@ def test_symmetric_matrices_to_spherical_commutes_with_o3(inversion): "energy/pbe", "energy_ensemble/member", "energy_uncertainty/direct", + "charge", ], ) -def test_decompose_output_energy_like(source_name): - """Energy-like variants should become one scalar spherical block.""" +def test_decompose_output_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_output(source_name, tensor) + result = decompose_output(source_name, tensor) assert result.keys.names == ["o3_lambda", "o3_sigma"] assert result.keys.values.tolist() == [[0, 1]] @@ -2773,11 +2955,11 @@ def test_decompose_output_energy_like(source_name): "source_name", [ "non_conservative_force/direct", - "non_conservative_forces/direct", + "velocity", ], ) -def test_decompose_output_non_conservative_force_preserves_autograd(source_name): - """Both force spellings should become l=1 and preserve implicit autograd.""" +def test_decompose_output_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, @@ -2785,7 +2967,7 @@ def test_decompose_output_non_conservative_force_preserves_autograd(source_name) ) tensor = _tensor_map_with_components(values, ["xyz"]) - result = _decompose_output(source_name, tensor) + result = decompose_output(source_name, tensor) assert result.keys.names == ["o3_lambda", "o3_sigma"] assert result.keys.values.tolist() == [[1, 1]] @@ -2807,7 +2989,7 @@ def test_decompose_output_non_conservative_stress_combines_irreps(): values[1, 1, 0, 0] = -2.0 tensor = _tensor_map_with_components(values, ["xyz_1", "xyz_2"]) - result = _decompose_output("non_conservative_stress/direct", tensor) + result = decompose_output("non_conservative_stress/direct", tensor) assert result.keys.names == ["o3_lambda", "o3_sigma"] assert result.keys.values.tolist() == [[0, 1], [2, 1]] @@ -2835,7 +3017,7 @@ def test_decompose_output_does_not_infer_custom_cartesian_semantics(): ["xyz_1", "xyz_2"], ) - result = _decompose_output("mtt::custom", tensor) + result = decompose_output("mtt::custom", tensor) mts.equal_raise(result, tensor) @@ -2847,19 +3029,19 @@ def test_decompose_output_does_not_infer_custom_cartesian_semantics(): "energy", (1, 3, 1), ["xyz"], - "energy-like outputs must not have components", + "'energy' outputs must not have components", ), ( "non_conservative_force", (1, 3, 1), ["component"], - "non_conservative_force must have one 'xyz' component axis of size 3", + "'non_conservative_force' must have one 'xyz' component axis of size 3", ), ( "non_conservative_stress", (1, 3, 3, 1), ["xyz_1", "component"], - "non_conservative_stress must have 'xyz_1' and 'xyz_2' component " + "'non_conservative_stress' must have 'xyz_1' and 'xyz_2' component " "axes of size 3", ), ], @@ -2877,7 +3059,7 @@ def test_decompose_output_rejects_invalid_standard_components( ) with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - _decompose_output(source_name, tensor) + decompose_output(source_name, tensor) def test_forward_rejects_outputs_with_attached_gradients(): From 8f87505e8fe17d4858cbba59b79979bb43c332c0 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Thu, 30 Jul 2026 13:24:15 +0200 Subject: [PATCH 08/18] Move quantity metadata to _quantities.py --- .../metatomic/torch/_quantities.py | 79 +++++++++++++++++++ .../metatomic_torch/metatomic/torch/model.py | 21 +---- .../metatomic/torch/o3/_decompose.py | 46 +---------- .../metatomic/torch/o3/_symmetrized.py | 27 ++----- 4 files changed, 93 insertions(+), 80 deletions(-) create mode 100644 python/metatomic_torch/metatomic/torch/_quantities.py diff --git a/python/metatomic_torch/metatomic/torch/_quantities.py b/python/metatomic_torch/metatomic/torch/_quantities.py new file mode 100644 index 00000000..0e2a088b --- /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 of every decomposable standard quantity. + + 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_O3_LAMBDA_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..62f917dd 100644 --- a/python/metatomic_torch/metatomic/torch/model.py +++ b/python/metatomic_torch/metatomic/torch/model.py @@ -25,6 +25,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": @@ -395,26 +396,10 @@ def __init__( 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", - } + self._new_names = dict(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 = dict(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/_decompose.py b/python/metatomic_torch/metatomic/torch/o3/_decompose.py index 7a2bc377..dacb01b0 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_decompose.py +++ b/python/metatomic_torch/metatomic/torch/o3/_decompose.py @@ -8,52 +8,12 @@ """ import math -from typing import Dict, List +from typing import List import torch from metatensor.torch import Labels, TensorBlock, TensorMap - -def _standard_quantity_categories() -> Dict[str, str]: - """Return the Cartesian layout of every decomposable standard quantity. - - 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 this module. - - 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_O3_LAMBDA_PER_CATEGORY: Dict[str, int] = { - "scalar": 0, - "cartesian_vector": 1, - "symmetric_matrix": 2, -} +from .._quantities import standard_quantity_categories def _o3_mu_labels(o3_lambda: int, device: torch.device) -> Labels: @@ -121,7 +81,7 @@ def decompose_output( their variance measures the deviation from invariance. """ quantity = source_name.split("/", 1)[0] - categories = _standard_quantity_categories() + categories = standard_quantity_categories() if quantity not in categories: return tensor category = categories[quantity] diff --git a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py index 8b19bd29..3527e849 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py +++ b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py @@ -21,11 +21,12 @@ register_autograd_neighbors, ) -from ._decompose import ( +from .._quantities import ( MAX_O3_LAMBDA_PER_CATEGORY, + NEW_QUANTITY_NAMES, STANDARD_QUANTITY_CATEGORIES, - decompose_output, ) +from ._decompose import decompose_output from ._projections import ( character_projection_coefficients_from_batch, character_projection_tensormap_from_cosets, @@ -47,21 +48,6 @@ ) -# deprecated quantity names mapped to their current name, mirroring -# ``AtomisticModel._new_names``. ``AtomisticModel`` advertises both spellings of -# each standard output, so an engine can request either one from the wrapper; -# everything inside the wrapper uses the current names only. -_NEW_NAMES: Dict[str, str] = { - "features": "feature", - "non_conservative_forces": "non_conservative_force", - "positions": "position", - "momenta": "momentum", - "masses": "mass", - "velocities": "velocity", - "charges": "charge", -} - - def _use_new_quantity_name(name: str, new_names: Dict[str, str]) -> str: """Replace a deprecated base quantity in ``name`` with its current name.""" parts = name.split("/") @@ -271,7 +257,7 @@ def _infer_max_o3_lambda( """Guess an angular-momentum limit from standard quantity names.""" max_o3_lambda = 0 for name in names.keys(): - quantity = _use_new_quantity_name(name, _NEW_NAMES).split("/")[0] + quantity = _use_new_quantity_name(name, NEW_QUANTITY_NAMES).split("/")[0] if quantity == "feature": # features are not an irreducible representation of O(3): they are # passed through unchanged and never rotated back @@ -649,8 +635,11 @@ def __init__( super().__init__() self._model = model + # ``AtomisticModel`` advertises both spellings of each standard output, so an + # engine can request either one from the wrapper; everything inside the wrapper + # uses the current names only. # TorchScript cannot read a module-level dictionary from ``forward`` - self._new_names = dict(_NEW_NAMES) + self._new_names = dict(NEW_QUANTITY_NAMES) self._requested_inputs = {} self._requested_neighbor_lists = [] self.max_o3_lambda_target = validate_integer( From 17ebabc5871c60868e01cc3013beb1bbc499119f Mon Sep 17 00:00:00 2001 From: ppegolo Date: Thu, 30 Jul 2026 14:52:25 +0200 Subject: [PATCH 09/18] Demote decompose shape validation to asserts --- .../metatomic/torch/o3/_decompose.py | 36 ++++++++----------- 1 file changed, 15 insertions(+), 21 deletions(-) diff --git a/python/metatomic_torch/metatomic/torch/o3/_decompose.py b/python/metatomic_torch/metatomic/torch/o3/_decompose.py index dacb01b0..3d7eb6e0 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_decompose.py +++ b/python/metatomic_torch/metatomic/torch/o3/_decompose.py @@ -89,8 +89,9 @@ def decompose_output( if category == "scalar": scalar_blocks: List[TensorBlock] = [] for block in tensor.blocks(): - if len(block.components) != 0: - raise ValueError(f"'{quantity}' outputs must not have components") + assert len(block.components) == 0, ( + f"'{quantity}' outputs must not have components" + ) scalar_blocks.append( TensorBlock( values=block.values.unsqueeze(1), @@ -107,14 +108,11 @@ def decompose_output( elif category == "cartesian_vector": vector_blocks: List[TensorBlock] = [] for block in tensor.blocks(): - if ( - len(block.components) != 1 - or block.components[0].names != ["xyz"] - or len(block.components[0]) != 3 - ): - raise ValueError( - f"'{quantity}' must have one 'xyz' component axis of size 3" - ) + 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), @@ -132,17 +130,13 @@ def decompose_output( blocks_l0: List[TensorBlock] = [] blocks_l2: List[TensorBlock] = [] for block in tensor.blocks(): - if ( - len(block.components) != 2 - or block.components[0].names != ["xyz_1"] - or block.components[1].names != ["xyz_2"] - or len(block.components[0]) != 3 - or len(block.components[1]) != 3 - ): - raise ValueError( - f"'{quantity}' must have 'xyz_1' and 'xyz_2' component axes " - "of size 3" - ) + 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_l2 = _symmetric_matrices_to_spherical(block.values) blocks_l0.append( From 824277b6b354d8a3b5ee2aafd34cb4bf0517f40e Mon Sep 17 00:00:00 2001 From: ppegolo Date: Thu, 30 Jul 2026 14:52:25 +0200 Subject: [PATCH 10/18] Relax angular momentum inference, rename max_o3_lambda parameters --- .../src/torch/reference/symmetrized-model.rst | 2 +- .../metatomic/torch/_quantities.py | 2 +- .../metatomic/torch/o3/_symmetrized.py | 207 ++++++------ .../metatomic/torch/o3/_wigner.py | 25 +- .../tests/symmetrized_model.py | 307 +++++++++--------- 5 files changed, 277 insertions(+), 266 deletions(-) diff --git a/docs/src/torch/reference/symmetrized-model.rst b/docs/src/torch/reference/symmetrized-model.rst index 9d4ba5cf..fb81a5de 100644 --- a/docs/src/torch/reference/symmetrized-model.rst +++ b/docs/src/torch/reference/symmetrized-model.rst @@ -132,7 +132,7 @@ 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_o3_lambda_grid`` controls the quadrature +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. diff --git a/python/metatomic_torch/metatomic/torch/_quantities.py b/python/metatomic_torch/metatomic/torch/_quantities.py index 0e2a088b..ea9a781f 100644 --- a/python/metatomic_torch/metatomic/torch/_quantities.py +++ b/python/metatomic_torch/metatomic/torch/_quantities.py @@ -46,7 +46,7 @@ def standard_quantity_categories() -> Dict[str, str]: STANDARD_QUANTITY_CATEGORIES: Dict[str, str] = standard_quantity_categories() #: maximum angular momentum carried by each category above -MAX_O3_LAMBDA_PER_CATEGORY: Dict[str, int] = { +MAX_ANGULAR_MOMENTUM_PER_CATEGORY: Dict[str, int] = { "scalar": 0, "cartesian_vector": 1, "symmetric_matrix": 2, diff --git a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py index 3527e849..c9388e17 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py +++ b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py @@ -22,7 +22,7 @@ ) from .._quantities import ( - MAX_O3_LAMBDA_PER_CATEGORY, + MAX_ANGULAR_MOMENTUM_PER_CATEGORY, NEW_QUANTITY_NAMES, STANDARD_QUANTITY_CATEGORIES, ) @@ -115,15 +115,15 @@ def _transform_system_geometry_batch( def _check_o3_lambda_limit( tensor: TensorMap, tensor_description: str, - max_o3_lambda: int, + 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_o3_lambda: + 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_o3_lambda}" + f"exceeding {limit_name}={max_angular_momentum}" ) @@ -249,28 +249,41 @@ def _group_output_requests( ) -def _infer_max_o3_lambda( +def _infer_max_angular_momentum( names: Dict[str, ModelOutput], kind: str, argument: str, ) -> int: """Guess an angular-momentum limit from standard quantity names.""" - max_o3_lambda = 0 + max_angular_momentum = 0 + found_standard = False + custom_names: List[str] = [] for name in names.keys(): quantity = _use_new_quantity_name(name, NEW_QUANTITY_NAMES).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: - raise ValueError( - f"unable to guess {argument} from the non-standard {kind} " - f"'{name}', please set {argument} explicitly" - ) + # 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_o3_lambda = max(max_o3_lambda, MAX_O3_LAMBDA_PER_CATEGORY[category]) + 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_o3_lambda + return max_angular_momentum def _reduce_weighted_centered_batch( @@ -451,7 +464,7 @@ def _clamp_roundoff_negative_diagnostic( *, n_grid_points: int, quantity: str, - max_o3_lambda_grid: int, + max_angular_momentum_grid: int, ) -> TensorMap: """Clamp round-off negatives and reject invalid or materially negative values.""" blocks: List[TensorBlock] = [] @@ -492,8 +505,8 @@ def _clamp_roundoff_negative_diagnostic( 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_o3_lambda_grid " - f"above {max_o3_lambda_grid} and check convergence" + "does not resolve this response. Increase max_angular_momentum_grid " + f"above {max_angular_momentum_grid} and check convergence" ) blocks.append( @@ -513,7 +526,7 @@ def _variance_from_centered_moments( absolute_centered_second_moment: TensorMap, *, n_grid_points: int, - max_o3_lambda_grid: 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) @@ -530,7 +543,7 @@ def _variance_from_centered_moments( roundoff_scale, n_grid_points=n_grid_points, quantity="variance", - max_o3_lambda_grid=max_o3_lambda_grid, + max_angular_momentum_grid=max_angular_momentum_grid, ) @@ -579,7 +592,7 @@ class SymmetrizedModel(torch.nn.Module): 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_o3_lambda_character`` is + 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 @@ -594,31 +607,31 @@ class SymmetrizedModel(torch.nn.Module): :param model: underlying :py:class:`ModelInterface`. The :py:meth:`wrap` method obtains this module from :py:attr:`AtomisticModel.module`. - :param max_o3_lambda_target: maximum angular momentum that can be transformed + :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_o3_lambda_input: maximum angular momentum that can be rotated in + :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_o3_lambda_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_character: maximum angular momentum included in + character projections. ``None`` disables character-projection outputs; zero + enables the scalar (``o3_lambda = 0``) contribution only. :param batch_size: positive number of transformed systems evaluated in one call to ``model``. The default is 32. - :param max_o3_lambda_grid: quadrature integration degree. If ``None``, use the - larger of ``2 * max_o3_lambda_target + 1`` and - ``2 * max_o3_lambda_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_o3_lambda_character`` is + :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. """ - max_o3_lambda_character: Optional[int] + max_angular_momentum_character: Optional[int] _new_names: Dict[str, str] _requested_inputs: Dict[str, ModelOutput] _requested_neighbor_lists: List[NeighborListOptions] @@ -626,11 +639,11 @@ class SymmetrizedModel(torch.nn.Module): def __init__( self, model: ModelInterface, - max_o3_lambda_target: int, - max_o3_lambda_input: int = 0, - max_o3_lambda_character: Optional[int] = None, + max_angular_momentum_target: int, + max_angular_momentum_input: int = 0, + max_angular_momentum_character: Optional[int] = None, batch_size: int = 32, - max_o3_lambda_grid: Optional[int] = None, + max_angular_momentum_grid: Optional[int] = None, ): super().__init__() @@ -642,38 +655,39 @@ def __init__( self._new_names = dict(NEW_QUANTITY_NAMES) self._requested_inputs = {} self._requested_neighbor_lists = [] - self.max_o3_lambda_target = validate_integer( - "max_o3_lambda_target", max_o3_lambda_target, 0 + self.max_angular_momentum_target = validate_integer( + "max_angular_momentum_target", max_angular_momentum_target, 0 ) - self.max_o3_lambda_input = validate_integer( - "max_o3_lambda_input", max_o3_lambda_input, 0 + self.max_angular_momentum_input = validate_integer( + "max_angular_momentum_input", max_angular_momentum_input, 0 ) - if max_o3_lambda_character is not None: - max_o3_lambda_character = validate_integer( - "max_o3_lambda_character", max_o3_lambda_character, 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_o3_lambda_character = max_o3_lambda_character + self.max_angular_momentum_character = max_angular_momentum_character self.batch_size = validate_integer("batch_size", batch_size, 1) - if max_o3_lambda_grid is None: - max_o3_lambda_grid = 2 * self.max_o3_lambda_target + 1 - if self.max_o3_lambda_character is not None: - max_o3_lambda_grid = max( - max_o3_lambda_grid, - 2 * self.max_o3_lambda_character, + 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_o3_lambda_grid = validate_integer( - "max_o3_lambda_grid", max_o3_lambda_grid, 0 + max_angular_momentum_grid = validate_integer( + "max_angular_momentum_grid", max_angular_momentum_grid, 0 ) if ( - self.max_o3_lambda_character is not None - and max_o3_lambda_grid < 2 * self.max_o3_lambda_character + self.max_angular_momentum_character is not None + and max_angular_momentum_grid < 2 * self.max_angular_momentum_character ): raise ValueError( - "max_o3_lambda_grid must be at least twice max_o3_lambda_character" + "max_angular_momentum_grid must be at least twice " + "max_angular_momentum_character" ) - self.max_o3_lambda_grid = max_o3_lambda_grid + self.max_angular_momentum_grid = max_angular_momentum_grid device = torch.device("cpu") for parameter in model.parameters(): @@ -686,7 +700,7 @@ def __init__( if device.type != "cpu" and device.type != "cuda": raise ValueError("SymmetrizedModel supports CPU and CUDA execution") - lebedev_order, n_rotations = choose_quadrature(self.max_o3_lambda_grid) + lebedev_order, n_rotations = choose_quadrature(self.max_angular_momentum_grid) rotations, weights = get_rotation_quadrature( lebedev_order, n_rotations, @@ -700,14 +714,16 @@ def __init__( device=device, ) - max_o3_lambda_wigner = max( - self.max_o3_lambda_input, - self.max_o3_lambda_target, - 0 if self.max_o3_lambda_character is None else self.max_o3_lambda_character, + 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_o3_lambda_wigner, + max_angular_momentum_wigner, ) self.register_buffer("_rotation_matrices", rotation_matrices) @@ -718,11 +734,11 @@ def __init__( def wrap( model: AtomisticModel, *, - max_o3_lambda_target: Optional[int] = None, - max_o3_lambda_input: Optional[int] = None, - max_o3_lambda_character: Optional[int] = None, + max_angular_momentum_target: Optional[int] = None, + max_angular_momentum_input: Optional[int] = None, + max_angular_momentum_character: Optional[int] = None, batch_size: int = 32, - max_o3_lambda_grid: Optional[int] = None, + max_angular_momentum_grid: Optional[int] = None, ) -> AtomisticModel: """ Wrap an exported model with O(3) averaging and diagnostics. @@ -730,7 +746,7 @@ def wrap( 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_o3_lambda_character`` + component-averaged equivariance variance. If ``max_angular_momentum_character`` is set, ``o3::character_projection::`` outputs provide squared character projections through that angular momentum. @@ -742,19 +758,19 @@ def wrap( already been saved. :param model: the :py:class:`AtomisticModel` to wrap - :param max_o3_lambda_target: maximum angular momentum accepted in + :param max_angular_momentum_target: maximum angular momentum accepted in already-spherical model outputs requested for averaging or variance. - When ``None``, it is guessed from the standard quantities declared by - ``model``; a non-standard output makes the guess impossible and must - be answered with an explicit value. - :param max_o3_lambda_input: maximum angular momentum accepted in custom - System data. When ``None``, it is guessed from the standard - quantities in ``model.requested_inputs()``, with the same - restriction on non-standard inputs. - :param max_o3_lambda_character: maximum angular momentum in character + 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 batch_size: number of transformed Systems evaluated in one model call - :param max_o3_lambda_grid: quadrature integration degree, selected + :param max_angular_momentum_grid: quadrature integration degree, selected automatically when ``None`` """ if not isinstance(model, AtomisticModel): @@ -772,17 +788,17 @@ def wrap( "wrapped model declares " + str(capabilities.supported_devices) ) - if max_o3_lambda_target is None: - max_o3_lambda_target = _infer_max_o3_lambda( + if max_angular_momentum_target is None: + max_angular_momentum_target = _infer_max_angular_momentum( capabilities.outputs, "output", - "max_o3_lambda_target", + "max_angular_momentum_target", ) - if max_o3_lambda_input is None: - max_o3_lambda_input = _infer_max_o3_lambda( + if max_angular_momentum_input is None: + max_angular_momentum_input = _infer_max_angular_momentum( model.requested_inputs(use_new_names=True), "input", - "max_o3_lambda_input", + "max_angular_momentum_input", ) outputs: Dict[str, ModelOutput] = {} @@ -822,7 +838,7 @@ def wrap( + "' output for each sample, averaged over components." ), ) - if max_o3_lambda_character is not None: + if max_angular_momentum_character is not None: outputs["o3::character_projection::" + name] = ModelOutput( unit=squared_unit, sample_kind=source_output.sample_kind, @@ -837,11 +853,11 @@ def wrap( wrapper = SymmetrizedModel( model.module, - max_o3_lambda_target=max_o3_lambda_target, - max_o3_lambda_input=max_o3_lambda_input, - max_o3_lambda_character=max_o3_lambda_character, + 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_o3_lambda_grid=max_o3_lambda_grid, + max_angular_momentum_grid=max_angular_momentum_grid, ) # private field: the as-declared inputs, deliberately without deprecation # aliases @@ -909,10 +925,11 @@ def forward( ) = _group_output_requests(outputs, self._new_names) if ( len(character_projection_names) != 0 - and self.max_o3_lambda_character is None + and self.max_angular_momentum_character is None ): raise ValueError( - "max_o3_lambda_character must be set to request character projections" + "max_angular_momentum_character must be set to request " + "character projections" ) source_outputs = torch.jit.annotate(Dict[str, ModelOutput], {}) @@ -991,12 +1008,12 @@ def _evaluate_system( _check_o3_lambda_limit( system.get_data(data_name), f"custom input '{data_name}'", - self.max_o3_lambda_input, - "max_o3_lambda_input", + self.max_angular_momentum_input, + "max_angular_momentum_input", ) character_max = 0 - configured_character_max = self.max_o3_lambda_character + configured_character_max = self.max_angular_momentum_character if configured_character_max is not None: character_max = configured_character_max @@ -1024,7 +1041,7 @@ def _evaluate_system( ) input_wigner_matrices: List[torch.Tensor] = [] - for o3_lambda in range(self.max_o3_lambda_input + 1): + for o3_lambda in range(self.max_angular_momentum_input + 1): input_wigner_matrices.append( wigner_matrices_for_lambda( self._packed_wigner_matrices, @@ -1038,7 +1055,7 @@ def _evaluate_system( inverse_target_wigner_matrices: List[torch.Tensor] = [] if needs_backrotation: - for o3_lambda in range(self.max_o3_lambda_target + 1): + for o3_lambda in range(self.max_angular_momentum_target + 1): inverse_target_wigner_matrices.append( wigner_matrices_for_lambda( self._packed_wigner_matrices, @@ -1106,8 +1123,8 @@ def _evaluate_system( _check_o3_lambda_limit( tensor, f"output '{source_name}'", - self.max_o3_lambda_target, - "max_o3_lambda_target", + self.max_angular_momentum_target, + "max_angular_momentum_target", ) backrotated = _transform_tensor_with_precomputed_matrices( tensor, @@ -1223,7 +1240,7 @@ def _evaluate_system( variance_second_moments[source_name], variance_absolute_second_moments[source_name], n_grid_points=2 * n_rotations, - max_o3_lambda_grid=self.max_o3_lambda_grid, + max_angular_momentum_grid=self.max_angular_momentum_grid, ) variance = _mean_variance_over_components( variance, diff --git a/python/metatomic_torch/metatomic/torch/o3/_wigner.py b/python/metatomic_torch/metatomic/torch/o3/_wigner.py index a2a102d5..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,21 +114,23 @@ 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_o3_lambda: int, + max_angular_momentum: int, ) -> torch.Tensor: - """Build and pack proper Wigner-D matrices through ``max_o3_lambda``. + """Build and pack proper Wigner-D matrices through ``max_angular_momentum``. :param matrices: ``(n_matrices, 3, 3)`` stack of O(3) matrices - :param max_o3_lambda: maximum angular momentum to include + :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`` @@ -139,7 +141,10 @@ def build_packed_wigner_matrices( cpu = torch.device("cpu") n_matrices = matrices.size(0) n_elements_per_matrix = ( - (max_o3_lambda + 1) * (2 * max_o3_lambda + 1) * (2 * max_o3_lambda + 3) // 3 + (max_angular_momentum + 1) + * (2 * max_angular_momentum + 1) + * (2 * max_angular_momentum + 3) + // 3 ) packed = torch.empty( n_matrices * n_elements_per_matrix, @@ -149,12 +154,12 @@ def build_packed_wigner_matrices( for matrix_index, matrix in enumerate(calculation_matrices.unbind(0)): cache = build_wigner_D_cache( - max_o3_lambda, + max_angular_momentum, matrix, device=cpu, dtype=output_dtype, ) - for o3_lambda in range(max_o3_lambda + 1): + 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 diff --git a/python/metatomic_torch/tests/symmetrized_model.py b/python/metatomic_torch/tests/symmetrized_model.py index dfba550e..1c72146a 100644 --- a/python/metatomic_torch/tests/symmetrized_model.py +++ b/python/metatomic_torch/tests/symmetrized_model.py @@ -759,7 +759,7 @@ def test_transforms_spherical_custom_data(self, is_improper): matrices = -proper_matrices if is_improper else proper_matrices packed_wigner = build_packed_wigner_matrices( proper_matrices, - max_o3_lambda=1, + max_angular_momentum=1, ) wigner_matrices = [ wigner_matrices_for_lambda( @@ -830,7 +830,7 @@ def test_input_limit_distinguishes_spherical_from_cartesian(self): ).unsqueeze(0) packed_wigner = build_packed_wigner_matrices( matrix, - max_o3_lambda=0, + max_angular_momentum=0, ) wigner_matrices = [ wigner_matrices_for_lambda( @@ -903,12 +903,12 @@ def test_input_limit_distinguishes_spherical_from_cartesian(self): ) model = SymmetrizedModel( _LinearEnergyModel(), - max_o3_lambda_target=0, - max_o3_lambda_grid=2, + max_angular_momentum_target=0, + max_angular_momentum_grid=2, ) message = ( "custom input 'mtt::field' contains o3_lambda=1, exceeding " - "max_o3_lambda_input=0" + "max_angular_momentum_input=0" ) with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): model( @@ -1056,21 +1056,22 @@ def test_packed_matrices_match_o3(self, dtype): -proper_rotation, ] ) - max_o3_lambda = 2 + max_angular_momentum = 2 - packed = build_packed_wigner_matrices(matrices, max_o3_lambda) + packed = build_packed_wigner_matrices(matrices, max_angular_momentum) assert packed.dim() == 1 assert packed.numel() == len(matrices) * sum( - (2 * o3_lambda + 1) ** 2 for o3_lambda in range(max_o3_lambda + 1) + (2 * o3_lambda + 1) ** 2 for o3_lambda in range(max_angular_momentum + 1) ) assert packed.dtype == matrices.dtype assert packed.device == matrices.device transformations = [ - O3Transformation(matrix, max_o3_lambda) for matrix in matrices.unbind(0) + O3Transformation(matrix, max_angular_momentum) + for matrix in matrices.unbind(0) ] - for o3_lambda in range(max_o3_lambda + 1): + for o3_lambda in range(max_angular_momentum + 1): actual = wigner_matrices_for_lambda( packed, len(matrices), @@ -1189,16 +1190,16 @@ def test_constructs_registered_buffers(self): """Constructor limits should determine the grid and Wigner-D storage.""" model = SymmetrizedModel( _EmptyModel(), - max_o3_lambda_target=1, - max_o3_lambda_input=2, - max_o3_lambda_character=1, + max_angular_momentum_target=1, + max_angular_momentum_input=2, + max_angular_momentum_character=1, batch_size=7, ) - assert model.max_o3_lambda_target == 1 - assert model.max_o3_lambda_input == 2 - assert model.max_o3_lambda_character == 1 - assert model.max_o3_lambda_grid == 3 + assert model.max_angular_momentum_target == 1 + assert model.max_angular_momentum_input == 2 + assert model.max_angular_momentum_character == 1 + assert model.max_angular_momentum_grid == 3 assert model.batch_size == 7 buffers = dict(model.named_buffers()) @@ -1221,56 +1222,59 @@ def test_character_limit_controls_default_grid(self): """Character sectors should raise the default grid degree when necessary.""" model = SymmetrizedModel( _EmptyModel(), - max_o3_lambda_target=0, - max_o3_lambda_character=2, + max_angular_momentum_target=0, + max_angular_momentum_character=2, ) - assert model.max_o3_lambda_grid == 4 + 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_o3_lambda_grid must be at least twice max_o3_lambda_character" + 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_o3_lambda_target=0, - max_o3_lambda_character=2, - max_o3_lambda_grid=3, + max_angular_momentum_target=0, + max_angular_momentum_character=2, + max_angular_momentum_grid=3, ) @pytest.mark.parametrize( ("argument", "value", "error", "message"), [ ( - "max_o3_lambda_target", + "max_angular_momentum_target", -1, ValueError, - "max_o3_lambda_target must be non-negative, got -1", + "max_angular_momentum_target must be non-negative, got -1", ), ( - "max_o3_lambda_target", + "max_angular_momentum_target", True, TypeError, - "max_o3_lambda_target must be an integer, got bool", + "max_angular_momentum_target must be an integer, got bool", ), ( - "max_o3_lambda_input", + "max_angular_momentum_input", 1.5, TypeError, - "max_o3_lambda_input must be an integer, got float", + "max_angular_momentum_input must be an integer, got float", ), ( - "max_o3_lambda_character", + "max_angular_momentum_character", -1, ValueError, - "max_o3_lambda_character must be non-negative, got -1", + "max_angular_momentum_character must be non-negative, got -1", ), ("batch_size", 0, ValueError, "batch_size must be positive, got 0"), ( - "max_o3_lambda_grid", + "max_angular_momentum_grid", -1, ValueError, - "max_o3_lambda_grid must be non-negative, got -1", + "max_angular_momentum_grid must be non-negative, got -1", ), ], ) @@ -1282,7 +1286,7 @@ def test_rejects_invalid_constructor_arguments( message, ): """Every integer constructor argument should enforce its documented range.""" - arguments = {"max_o3_lambda_target": 0, argument: value} + arguments = {"max_angular_momentum_target": 0, argument: value} with pytest.raises(error, match=f"^{re.escape(message)}$"): SymmetrizedModel(_EmptyModel(), **arguments) @@ -1294,7 +1298,7 @@ def test_rejects_a_model_stored_on_an_unsupported_device(self): message = "SymmetrizedModel supports CPU and CUDA execution" with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - SymmetrizedModel(base_model, max_o3_lambda_target=0) + SymmetrizedModel(base_model, max_angular_momentum_target=0) class TestSymmetrizedModelForward: @@ -1309,9 +1313,9 @@ def test_character_projection_separates_sectors_through_lambda_three(self): ] model = SymmetrizedModel( _O3PolynomialSectorModel(), - max_o3_lambda_target=0, - max_o3_lambda_character=3, - max_o3_lambda_grid=6, + 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()) @@ -1354,9 +1358,9 @@ def test_energy_results_match_analytic_values_and_reuse_predictions(self): batch_size = 5 model = SymmetrizedModel( base_model, - max_o3_lambda_target=0, - max_o3_lambda_character=1, - max_o3_lambda_grid=2, + 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]]) @@ -1423,9 +1427,9 @@ def test_stress_character_projection_combines_target_and_character_sectors(self) requested_name = "o3::character_projection::non_conservative_stress" model = SymmetrizedModel( _EquivariantOutputModel(), - max_o3_lambda_target=2, - max_o3_lambda_character=2, - max_o3_lambda_grid=4, + 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]]) @@ -1486,9 +1490,9 @@ def test_source_request_contains_only_the_shared_sample_kind( base_model = _CountingLinearEnergyModel() model = SymmetrizedModel( base_model, - max_o3_lambda_target=0, - max_o3_lambda_character=1, - max_o3_lambda_grid=2, + max_angular_momentum_target=0, + max_angular_momentum_character=1, + max_angular_momentum_grid=2, ) model( @@ -1515,12 +1519,12 @@ 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_o3_lambda_target=1, + max_angular_momentum_target=1, ) message = ( "output 'mtt::spherical_quadrupole' contains o3_lambda=2, " - "exceeding max_o3_lambda_target=1" + "exceeding max_angular_momentum_target=1" ) with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): model( @@ -1548,13 +1552,13 @@ def test_rejects_a_negative_quadrature_error_and_converges(self): underresolved = SymmetrizedModel( _DegreeSevenEnergyModel(), - max_o3_lambda_target=0, - max_o3_lambda_grid=12, + 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_o3_lambda_grid above 12 " + "not resolve this response. Increase max_angular_momentum_grid above 12 " "and check convergence" ) with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): @@ -1562,8 +1566,8 @@ def test_rejects_a_negative_quadrature_error_and_converges(self): resolved = SymmetrizedModel( _DegreeSevenEnergyModel(), - max_o3_lambda_target=0, - max_o3_lambda_grid=14, + max_angular_momentum_target=0, + max_angular_momentum_grid=14, batch_size=64, ) outputs = { @@ -1599,8 +1603,8 @@ def test_preserves_variant_and_custom_output_names(self, source_name): base_model = _CountingLinearEnergyModel() model = SymmetrizedModel( base_model, - max_o3_lambda_target=0, - max_o3_lambda_grid=2, + max_angular_momentum_target=0, + max_angular_momentum_grid=2, ) variance_name = "o3::variance::" + source_name outputs = { @@ -1626,8 +1630,8 @@ def test_deprecated_quantity_names_are_normalized(self): """A deprecated request is decomposed as, and returned under, its own name.""" model = SymmetrizedModel( _EquivariantOutputModel(), - max_o3_lambda_target=1, - max_o3_lambda_grid=2, + max_angular_momentum_target=1, + max_angular_momentum_grid=2, batch_size=5, ) outputs = { @@ -1658,8 +1662,8 @@ def test_component_less_output_averages_and_measures_invariance(self): system = _forward_test_system([[1.0, 2.0, 3.0], [0.0, 1.0, 0.0]]) model = SymmetrizedModel( _AtomFeatureModel(), - max_o3_lambda_target=0, - max_o3_lambda_grid=2, + max_angular_momentum_target=0, + max_angular_momentum_grid=2, batch_size=5, ) outputs = { @@ -1692,8 +1696,8 @@ def test_selected_atoms_excludes_unselected_input_systems(self): ] model = SymmetrizedModel( _EquivariantOutputModel(), - max_o3_lambda_target=1, - max_o3_lambda_grid=2, + max_angular_momentum_target=1, + max_angular_momentum_grid=2, batch_size=5, ) outputs = { @@ -1731,8 +1735,8 @@ def test_empty_selected_atoms_returns_empty_outputs(self): ] model = SymmetrizedModel( _EquivariantOutputModel(), - max_o3_lambda_target=1, - max_o3_lambda_grid=2, + max_angular_momentum_target=1, + max_angular_momentum_grid=2, batch_size=5, ) outputs = { @@ -1769,8 +1773,8 @@ def test_multiple_systems_keep_per_system_rows_in_order(self): ] model = SymmetrizedModel( _EquivariantOutputModel(), - max_o3_lambda_target=0, - max_o3_lambda_grid=2, + max_angular_momentum_target=0, + max_angular_momentum_grid=2, batch_size=5, ) outputs = { @@ -1814,8 +1818,8 @@ def test_equivariant_outputs_preserve_values_metadata_and_zero_variance(self): outputs["o3::variance::" + name] = outputs[name] model = SymmetrizedModel( _EquivariantOutputModel(), - max_o3_lambda_target=2, - max_o3_lambda_grid=2, + max_angular_momentum_target=2, + max_angular_momentum_grid=2, batch_size=7, ) @@ -1882,8 +1886,8 @@ def test_dtype_and_implicit_autograd(self, dtype): ) model = SymmetrizedModel( _LinearEnergyModel(), - max_o3_lambda_target=0, - max_o3_lambda_grid=2, + max_angular_momentum_target=0, + max_angular_momentum_grid=2, ) outputs = { "o3::variance::energy": ModelOutput(sample_kind="system"), @@ -1910,8 +1914,8 @@ def test_average_output_preserves_implicit_autograd(self): ) model = SymmetrizedModel( _EquivariantOutputModel(), - max_o3_lambda_target=0, - max_o3_lambda_grid=2, + max_angular_momentum_target=0, + max_angular_momentum_grid=2, ) result = model([system], {"energy": ModelOutput(sample_kind="system")}, None) @@ -1930,14 +1934,17 @@ def test_average_output_preserves_implicit_autograd(self): 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_o3_lambda_target=0) + 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_o3_lambda_character must be set to request character projections" + message = ( + "max_angular_momentum_character must be set to request " + "character projections" + ) with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): model( [system], @@ -1977,8 +1984,8 @@ def test_rejects_downcast_integration_buffers(self): """Calling .float() on the module must fail loudly at the next forward.""" model = SymmetrizedModel( _LinearEnergyModel(), - max_o3_lambda_target=0, - max_o3_lambda_grid=2, + max_angular_momentum_target=0, + max_angular_momentum_grid=2, ).float() message = ( @@ -1996,8 +2003,8 @@ 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_o3_lambda_target=0, - max_o3_lambda_grid=2, + max_angular_momentum_target=0, + max_angular_momentum_grid=2, ) message = "underlying model did not return requested output 'energy'" @@ -2012,8 +2019,8 @@ def test_rejects_a_non_finite_variance(self): """A NaN model response should fail the variance finiteness check.""" model = SymmetrizedModel( _LinearEnergyModel(), - max_o3_lambda_target=0, - max_o3_lambda_grid=2, + 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))" @@ -2027,9 +2034,9 @@ def test_rejects_a_non_finite_variance(self): def test_is_scriptable_and_serializable(self, tmp_path): """The complete forward path should execute after scripting and reloading.""" constructor_arguments = { - "max_o3_lambda_target": 0, - "max_o3_lambda_character": 1, - "max_o3_lambda_grid": 2, + "max_angular_momentum_target": 0, + "max_angular_momentum_character": 1, + "max_angular_momentum_grid": 2, "batch_size": 5, } eager = SymmetrizedModel(_LinearEnergyModel(), **constructor_arguments) @@ -2057,8 +2064,8 @@ def test_is_scriptable_and_serializable(self, tmp_path): class TestSymmetrizedModelWrap: """Test exported-model capabilities, dependencies, and execution.""" - @pytest.mark.parametrize("max_o3_lambda_character", [None, 1]) - def test_wrap_declares_capabilities(self, max_o3_lambda_character): + @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( @@ -2084,9 +2091,9 @@ def test_wrap_declares_capabilities(self, max_o3_lambda_character): wrapped = SymmetrizedModel.wrap( base, - max_o3_lambda_target=0, - max_o3_lambda_character=max_o3_lambda_character, - max_o3_lambda_grid=2, + max_angular_momentum_target=0, + max_angular_momentum_character=max_angular_momentum_character, + max_angular_momentum_grid=2, ) capabilities = wrapped.capabilities() @@ -2098,7 +2105,7 @@ def test_wrap_declares_capabilities(self, max_o3_lambda_character): expected_names = set(source_outputs) expected_names.update("o3::variance::" + name for name in source_outputs) - if max_o3_lambda_character is not None: + if max_angular_momentum_character is not None: expected_names.update( "o3::character_projection::" + name for name in source_outputs ) @@ -2114,13 +2121,13 @@ def test_wrap_declares_capabilities(self, max_o3_lambda_character): ) assert capabilities.outputs["o3::variance::" + name].unit == squared_unit character_name = "o3::character_projection::" + name - if max_o3_lambda_character is None: + if max_angular_momentum_character is None: assert character_name not in capabilities.outputs else: assert capabilities.outputs[character_name].unit == squared_unit @pytest.mark.parametrize( - ("outputs", "expected_max_o3_lambda_target"), + ("outputs", "expected_max_angular_momentum_target"), [ ({"energy": ModelOutput(unit="eV", sample_kind="system")}, 0), ({"feature": ModelOutput(sample_kind="atom")}, 0), @@ -2143,18 +2150,30 @@ def test_wrap_declares_capabilities(self, max_o3_lambda_character): }, 2, ), + # a custom output is skipped, the standard ones still set the limit + ( + { + "energy": ModelOutput(unit="eV", sample_kind="system"), + "mtt::custom": ModelOutput(sample_kind="system"), + }, + 0, + ), ], ) def test_guesses_limits_from_standard_quantities( self, outputs, - expected_max_o3_lambda_target, + expected_max_angular_momentum_target, ): """Both limits default to what the standard quantities require.""" class _VelocityInputModel(_EmptyModel): def requested_inputs(self) -> Dict[str, ModelOutput]: - return {"velocity": ModelOutput(sample_kind="atom")} + # the custom input is skipped by the guess as well + return { + "velocity": ModelOutput(sample_kind="atom"), + "mtt::field": ModelOutput(sample_kind="atom"), + } base = AtomisticModel( _VelocityInputModel().eval(), @@ -2169,19 +2188,25 @@ def requested_inputs(self) -> Dict[str, ModelOutput]: ), ) - wrapped = SymmetrizedModel.wrap(base, max_o3_lambda_grid=2) + wrapped = SymmetrizedModel.wrap(base, max_angular_momentum_grid=2) - assert wrapped.module.max_o3_lambda_target == expected_max_o3_lambda_target + assert ( + wrapped.module.max_angular_momentum_target + == expected_max_angular_momentum_target + ) # velocity is a Cartesian vector - assert wrapped.module.max_o3_lambda_input == 1 + assert wrapped.module.max_angular_momentum_input == 1 - def test_rejects_guessing_a_limit_from_a_custom_output(self): - """A non-standard output must be answered with an explicit limit.""" + 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")}, + outputs={ + "mtt::custom": ModelOutput(sample_kind="system"), + "mtt::other": ModelOutput(sample_kind="system"), + }, atomic_types=[1], interaction_range=0.0, length_unit="A", @@ -2191,8 +2216,9 @@ def test_rejects_guessing_a_limit_from_a_custom_output(self): ) message = ( - "unable to guess max_o3_lambda_target from the non-standard output " - "'mtt::custom', please set max_o3_lambda_target explicitly" + "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) @@ -2224,7 +2250,7 @@ def test_rejects_reserved_source_names(self, source_name): "by SymmetrizedModel" ) with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - SymmetrizedModel.wrap(base, max_o3_lambda_target=0) + 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.""" @@ -2246,7 +2272,7 @@ def test_rejects_models_without_a_supported_device(self): "wrapped model declares ['mps']" ) with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - SymmetrizedModel.wrap(base, max_o3_lambda_target=0) + 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.""" @@ -2276,12 +2302,12 @@ def test_preserves_requirements_and_runs_after_save_load(self, tmp_path): wrapped = SymmetrizedModel.wrap( loaded_base, - max_o3_lambda_target=0, + max_angular_momentum_target=0, # 'mtt::linear' and 'mtt::field' are not standard quantities, so # both limits have to be given explicitly - max_o3_lambda_input=0, - max_o3_lambda_character=1, - max_o3_lambda_grid=2, + max_angular_momentum_input=0, + max_angular_momentum_character=1, + max_angular_momentum_grid=2, batch_size=5, ) assert ( @@ -2380,9 +2406,10 @@ def test_saved_wrapper_runs_on_cuda(self, tmp_path): ) wrapped = SymmetrizedModel.wrap( base, - max_o3_lambda_target=0, - max_o3_lambda_character=1, - max_o3_lambda_grid=2, + 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" @@ -2398,7 +2425,9 @@ def test_saved_wrapper_runs_on_cuda(self, tmp_path): ) cuda_system = cpu_system.to(device=cuda_device) - cpu_module = SymmetrizedModel(_LinearEnergyModel(), max_o3_lambda_target=0) + 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( @@ -2680,7 +2709,7 @@ def test_variance_from_centered_moments(): centered_second_moment, absolute_centered_second_moment, n_grid_points=12, - max_o3_lambda_grid=3, + max_angular_momentum_grid=3, ) assert torch.allclose(variance.block().values, expected_variance) @@ -2713,7 +2742,7 @@ def test_centered_variance_is_stable_with_large_offset(): second, absolute_second, n_grid_points=4, - max_o3_lambda_grid=3, + max_angular_momentum_grid=3, ) assert torch.allclose( @@ -2740,14 +2769,14 @@ def test_roundoff_negative_diagnostic_uses_its_scale(): _make_single_block_tensor_map(torch.tensor([[scale], [scale]], dtype=dtype)), n_grid_points=n_grid_points, quantity="variance", - max_o3_lambda_grid=3, + max_angular_momentum_grid=3, ) assert cleaned.block().values[0, 0].item() == 0.0 assert cleaned.block().values[1, 0].item() == 2.0 message = ( "finite O(3) variance is materially negative; the quadrature does not " - "resolve this response. Increase max_o3_lambda_grid above 3 and check " + "resolve this response. Increase max_angular_momentum_grid above 3 and check " "convergence" ) with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): @@ -2758,7 +2787,7 @@ def test_roundoff_negative_diagnostic_uses_its_scale(): _make_single_block_tensor_map(torch.tensor([[scale]], dtype=dtype)), n_grid_points=n_grid_points, quantity="variance", - max_o3_lambda_grid=3, + max_angular_momentum_grid=3, ) @@ -3022,46 +3051,6 @@ def test_decompose_output_does_not_infer_custom_cartesian_semantics(): mts.equal_raise(result, tensor) -@pytest.mark.parametrize( - ("source_name", "shape", "component_names", "message"), - [ - ( - "energy", - (1, 3, 1), - ["xyz"], - "'energy' outputs must not have components", - ), - ( - "non_conservative_force", - (1, 3, 1), - ["component"], - "'non_conservative_force' must have one 'xyz' component axis of size 3", - ), - ( - "non_conservative_stress", - (1, 3, 3, 1), - ["xyz_1", "component"], - "'non_conservative_stress' must have 'xyz_1' and 'xyz_2' component " - "axes of size 3", - ), - ], -) -def test_decompose_output_rejects_invalid_standard_components( - source_name, - shape, - component_names, - message, -): - """Standard quantities should use their required Cartesian component axes.""" - tensor = _tensor_map_with_components( - torch.zeros(shape, dtype=torch.float64), - component_names, - ) - - with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - decompose_output(source_name, tensor) - - def test_forward_rejects_outputs_with_attached_gradients(): """The wrapper should not silently discard explicit TensorBlock gradients.""" @@ -3097,8 +3086,8 @@ def forward( model = SymmetrizedModel( _AttachedGradientModel(), - max_o3_lambda_target=0, - max_o3_lambda_grid=2, + max_angular_momentum_target=0, + max_angular_momentum_grid=2, ) message = ( From b97b2e6134c4be4b4a6929d522fa66c31b28397a Mon Sep 17 00:00:00 2001 From: ppegolo Date: Thu, 30 Jul 2026 16:43:50 +0200 Subject: [PATCH 11/18] Polish SymmetrizedModel API: group parameters, keyword-only init --- .../metatomic/torch/o3/_quadrature.py | 21 +++++++++++-------- .../metatomic/torch/o3/_symmetrized.py | 11 +++++----- .../tests/symmetrized_model.py | 8 +++---- 3 files changed, 22 insertions(+), 18 deletions(-) diff --git a/python/metatomic_torch/metatomic/torch/o3/_quadrature.py b/python/metatomic_torch/metatomic/torch/o3/_quadrature.py index eee58024..3daeacdf 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_quadrature.py +++ b/python/metatomic_torch/metatomic/torch/o3/_quadrature.py @@ -59,24 +59,27 @@ def _import_scipy(): return lebedev_rule, Rotation -def choose_quadrature(L_max: int) -> tuple[int, int]: +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 degree ``L_max``. + spherical harmonics up to ``max_angular_momentum``. - :param L_max: maximum spherical harmonic degree + :param max_angular_momentum: maximum spherical harmonic degree :return: (lebedev_order, n_inplane_rotations) """ - L_max = validate_integer("L_max", L_max, 0) - if L_max > _LEBEDEV_ORDERS[-1]: + max_angular_momentum = validate_integer( + "max_angular_momentum", max_angular_momentum, 0 + ) + if max_angular_momentum > _LEBEDEV_ORDERS[-1]: raise ValueError( - f"the requested quadrature degree L_max={L_max} exceeds the largest " + "the requested quadrature degree " + f"max_angular_momentum={max_angular_momentum} exceeds the largest " f"available Lebedev order ({_LEBEDEV_ORDERS[-1]})" ) - # pick smallest order >= L_max - n = min(o for o in _LEBEDEV_ORDERS if o >= L_max) + # pick smallest order >= max_angular_momentum + n = min(o for o in _LEBEDEV_ORDERS if o >= max_angular_momentum) # minimal gamma count - K = L_max + 1 + K = max_angular_momentum + 1 return n, K diff --git a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py index c9388e17..6340d86a 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py +++ b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py @@ -621,14 +621,14 @@ class SymmetrizedModel(torch.nn.Module): :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 batch_size: positive number of transformed systems evaluated in one call to - ``model``. The default is 32. :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] @@ -639,11 +639,12 @@ class SymmetrizedModel(torch.nn.Module): def __init__( self, model: ModelInterface, + *, max_angular_momentum_target: int, max_angular_momentum_input: int = 0, max_angular_momentum_character: Optional[int] = None, - batch_size: int = 32, max_angular_momentum_grid: Optional[int] = None, + batch_size: int = 32, ): super().__init__() @@ -737,8 +738,8 @@ def wrap( max_angular_momentum_target: Optional[int] = None, max_angular_momentum_input: Optional[int] = None, max_angular_momentum_character: Optional[int] = None, - batch_size: int = 32, max_angular_momentum_grid: Optional[int] = None, + batch_size: int = 32, ) -> AtomisticModel: """ Wrap an exported model with O(3) averaging and diagnostics. @@ -769,9 +770,9 @@ def wrap( quantities in ``model.requested_inputs()``. :param max_angular_momentum_character: maximum angular momentum in character projections, or ``None`` to disable them - :param batch_size: number of transformed Systems evaluated in one model call :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") diff --git a/python/metatomic_torch/tests/symmetrized_model.py b/python/metatomic_torch/tests/symmetrized_model.py index 1c72146a..ec160d57 100644 --- a/python/metatomic_torch/tests/symmetrized_model.py +++ b/python/metatomic_torch/tests/symmetrized_model.py @@ -1130,17 +1130,17 @@ def test_euler_angle_rotations_are_in_so3(self): def test_quadrature_validation(self): """Quadrature construction rejects invalid degrees, counts, and orders.""" message = ( - "the requested quadrature degree L_max=132 exceeds the largest " - "available Lebedev order (131)" + "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 = "L_max must be non-negative, got -1" + message = "max_angular_momentum must be non-negative, got -1" with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): choose_quadrature(-1) - message = "L_max must be an integer, got float" + message = "max_angular_momentum must be an integer, got float" with pytest.raises(TypeError, match=f"^{re.escape(message)}$"): choose_quadrature(1.5) From a8e5e112e78631ad1e4024681664d46329962f2a Mon Sep 17 00:00:00 2001 From: ppegolo Date: Wed, 5 Aug 2026 16:15:22 +0200 Subject: [PATCH 12/18] Reuse and adapt O3 stuff + other review items --- metatomic-torch/CHANGELOG.md | 6 +- .../metatomic/torch/_quantities.py | 2 +- .../metatomic_torch/metatomic/torch/model.py | 4 +- .../metatomic/torch/o3/_decompose.py | 22 +- .../metatomic/torch/o3/_symmetrized.py | 286 ++-- .../metatomic/torch/o3/_tranformations.py | 1182 +++++++++------ .../metatomic/torch/o3/_utils.py | 7 +- python/metatomic_torch/tests/o3.py | 262 +++- .../tests/symmetrized_model.py | 1335 ++++------------- 9 files changed, 1361 insertions(+), 1745 deletions(-) diff --git a/metatomic-torch/CHANGELOG.md b/metatomic-torch/CHANGELOG.md index 3be532e1..3163eaac 100644 --- a/metatomic-torch/CHANGELOG.md +++ b/metatomic-torch/CHANGELOG.md @@ -19,11 +19,15 @@ a changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows ### Added - Added `metatomic.torch.SymmetrizedModel` for finite-quadrature O(3) - averaging, equivariance variances, and character projections of exported + 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/_quantities.py b/python/metatomic_torch/metatomic/torch/_quantities.py index ea9a781f..d0800858 100644 --- a/python/metatomic_torch/metatomic/torch/_quantities.py +++ b/python/metatomic_torch/metatomic/torch/_quantities.py @@ -12,7 +12,7 @@ def standard_quantity_categories() -> Dict[str, str]: - """Return the Cartesian layout of every decomposable standard quantity. + """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 diff --git a/python/metatomic_torch/metatomic/torch/model.py b/python/metatomic_torch/metatomic/torch/model.py index 62f917dd..3f50f37c 100644 --- a/python/metatomic_torch/metatomic/torch/model.py +++ b/python/metatomic_torch/metatomic/torch/model.py @@ -395,7 +395,9 @@ def __init__( else: raise ValueError(f"unknown dtype in capabilities: {capabilities.dtype}") - # mapping from deprecated output/input names to their new name + # 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 = dict(NEW_QUANTITY_NAMES) # mapping from new names to the corresponding deprecated name diff --git a/python/metatomic_torch/metatomic/torch/o3/_decompose.py b/python/metatomic_torch/metatomic/torch/o3/_decompose.py index 3d7eb6e0..eb46074d 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_decompose.py +++ b/python/metatomic_torch/metatomic/torch/o3/_decompose.py @@ -42,10 +42,10 @@ def _symmetric_matrices_to_spherical( ) -> tuple[torch.Tensor, torch.Tensor]: """Return orthonormal l=0 and l=2 components of the symmetric matrix part. - ``values`` must have shape ``(n_samples, 3, 3, n_properties)``. - - The antisymmetric (l=1) part is silently discarded. + Standard matrix quantities are symmetric, so their antisymmetric (l=1) part + carries no information and is discarded. """ + 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) @@ -66,21 +66,22 @@ def _symmetric_matrices_to_spherical( return l0, l2 -def decompose_output( - source_name: str, +def decompose_quantity( + name: str, tensor: TensorMap, ) -> TensorMap: - """Decompose standard outputs for variance and character projection. + """Decompose standard quantities for variance and character projection. - This takes the standard Cartesian or scalar 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. + 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 = source_name.split("/", 1)[0] + quantity = name.split("/", 1)[0] categories = standard_quantity_categories() if quantity not in categories: return tensor @@ -127,6 +128,7 @@ def decompose_output( ) else: + assert category == "symmetric_matrix" blocks_l0: List[TensorBlock] = [] blocks_l2: List[TensorBlock] = [] for block in tensor.blocks(): diff --git a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py index 6340d86a..99e0af53 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py +++ b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py @@ -5,6 +5,7 @@ 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 @@ -18,7 +19,6 @@ ModelOutput, NeighborListOptions, System, - register_autograd_neighbors, ) from .._quantities import ( @@ -26,16 +26,13 @@ NEW_QUANTITY_NAMES, STANDARD_QUANTITY_CATEGORIES, ) -from ._decompose import decompose_output +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 ._tranformations import ( - _max_o3_lambda_in_tensor, - _transform_tensor_with_precomputed_matrices, -) +from ._tranformations import O3Transformation, _max_o3_lambda_in_tensor from ._utils import ( group_samples_by_rotated_copy, map_selected_atoms_to_rotated_copies, @@ -49,7 +46,12 @@ def _use_new_quantity_name(name: str, new_names: Dict[str, str]) -> str: - """Replace a deprecated base quantity in ``name`` with its current name.""" + """Replace a deprecated base quantity in ``name`` with its current name. + + Callers may request outputs under deprecated spellings (e.g. ``energies``); + translating them here lets the rest of the wrapper deal with the current + names only, whatever the wrapped model declares. + """ parts = name.split("/") if parts[0] in new_names: parts[0] = new_names[parts[0]] @@ -57,61 +59,6 @@ def _use_new_quantity_name(name: str, new_names: Dict[str, str]) -> str: return name -def _transform_system_geometry_batch( - system: System, - matrices: torch.Tensor, -) -> List[System]: - """Transform System geometry and neighbor lists with internal O(3) matrices.""" - 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 != system.positions.dtype - or matrices.device != system.positions.device - ): - raise ValueError("system and 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( - types=system.types, - positions=positions[index], - cell=cells[index], - pbc=system.pbc, - ) - ) - - for options in system.known_neighbor_lists(): - neighbors = system.get_neighbor_list(options) - 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, - ) - - return transformed_systems - - def _check_o3_lambda_limit( tensor: TensorMap, tensor_description: str, @@ -127,37 +74,6 @@ def _check_o3_lambda_limit( ) -def _transform_system_batch( - system: System, - matrices: torch.Tensor, - wigner_matrices: List[torch.Tensor], - is_improper: bool, -) -> List[System]: - """Transform a System batch, including its custom TensorMap data.""" - data_names = system.known_data() - transformed_systems = _transform_system_geometry_batch(system, matrices) - if len(data_names) == 0: - return transformed_systems - - for index in range(len(transformed_systems)): - wigner_matrices_for_copy: List[torch.Tensor] = [] - for rank_matrices in wigner_matrices: - wigner_matrices_for_copy.append(rank_matrices[index : index + 1]) - - for data_name in data_names: - transformed_systems[index].add_data( - data_name, - _transform_tensor_with_precomputed_matrices( - system.get_data(data_name), - matrices[index : index + 1], - wigner_matrices_for_copy, - is_improper, - ), - ) - - return transformed_systems - - def _parse_output_request(requested_name: str) -> Tuple[str, str]: """Return the underlying output name and requested calculation.""" variance_prefix = "o3::variance::" @@ -170,6 +86,12 @@ def _parse_output_request(requested_name: str) -> Tuple[str, str]: 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" @@ -187,7 +109,12 @@ def _record_output_request( source_name: str, requested_name: str, ) -> None: - """Register the public name a source output must be returned under.""" + """Register the public name a source output must be returned under. + + Because deprecated names were normalized, two requested spellings can map to + the same source output and calculation; this rejects such duplicates, since + only one result per (source, calculation) pair can be returned. + """ if source_name in names: raise ValueError( f"'{requested_name}' and '{names[source_name]}' request the same " @@ -298,8 +225,15 @@ def _reduce_weighted_centered_batch( Optional[TensorMap], TensorMap, ]: - """Accumulate one rotation batch's weighted moments, centered on a reference - value so the variance subtraction stays cancellation-safe.""" + """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] = [] @@ -466,7 +400,15 @@ def _clamp_roundoff_negative_diagnostic( quantity: str, max_angular_momentum_grid: int, ) -> TensorMap: - """Clamp round-off negatives and reject invalid or materially negative values.""" + """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 @@ -699,6 +641,8 @@ def __init__( 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) @@ -726,6 +670,7 @@ def __init__( 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) @@ -806,9 +751,7 @@ def wrap( # 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::variance::") or name.startswith( - "o3::character_projection::" - ): + if name.startswith("o3::"): raise ValueError( "the wrapped model output '" + name @@ -907,8 +850,12 @@ def forward( ) -> Dict[str, TensorMap]: """Evaluate the requested O(3) averages and diagnostics.""" if len(outputs) == 0: - return torch.jit.annotate(Dict[str, TensorMap], {}) + 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(): @@ -933,18 +880,16 @@ def forward( "character projections" ) - source_outputs = torch.jit.annotate(Dict[str, ModelOutput], {}) + 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 = torch.jit.annotate( - Dict[str, List[TensorMap]], - {}, - ) + per_output_results: Dict[str, List[TensorMap]] = {} for requested_name in outputs: - per_output_results[requested_name] = torch.jit.annotate(List[TensorMap], []) + empty_results: List[TensorMap] = [] + per_output_results[requested_name] = empty_results for input_system_index, system in enumerate(systems): system_results = self._evaluate_system( @@ -961,7 +906,7 @@ def forward( system_results[requested_name] ) - results = torch.jit.annotate(Dict[str, TensorMap], {}) + results: Dict[str, TensorMap] = {} for requested_name in outputs: results[requested_name] = mts.join( per_output_results[requested_name], @@ -990,11 +935,17 @@ def _evaluate_system( ) if work_device.type != "cpu" and work_device.type != "cuda": raise ValueError("SymmetrizedModel supports CPU and CUDA execution") - if self._rotation_matrices.dtype != torch.float64: - raise ValueError( - "SymmetrizedModel integration buffers must remain float64, got " - f"{dtype_name(self._rotation_matrices.dtype)}; do not call " - ".float() or .half() on the module" + 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 @@ -1029,6 +980,18 @@ def _evaluate_system( 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 @@ -1041,54 +1004,57 @@ def _evaluate_system( n_rotated_copies, ) - input_wigner_matrices: List[torch.Tensor] = [] + 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_matrices.append( - wigner_matrices_for_lambda( - self._packed_wigner_matrices, - n_rotations, - o3_lambda, - )[batch_start:batch_stop].to( - dtype=work_dtype, - device=work_device, - ) + 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, + ) - inverse_target_wigner_matrices: List[torch.Tensor] = [] + # 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): - inverse_target_wigner_matrices.append( - wigner_matrices_for_lambda( - self._packed_wigner_matrices, - n_rotations, - o3_lambda, - )[batch_start:batch_stop].transpose(1, 2) - ) + 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( - wigner_matrices_for_lambda( - self._packed_wigner_matrices, - n_rotations, - chi_lambda, - )[batch_start:batch_stop].transpose(1, 2) + batch_wigner[chi_lambda].transpose(1, 2) ) for coset_index in range(2): is_improper = coset_index == 1 - sign = -1.0 if is_improper else 1.0 - matrices = (sign * proper_matrices).to( - dtype=work_dtype, - device=work_device, - ) - transformed_systems = _transform_system_batch( - system, - matrices, - input_wigner_matrices, - is_improper, - ) + 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, @@ -1102,7 +1068,6 @@ def _evaluate_system( f"'{source_name}'" ) - inverse_matrices = (sign * proper_matrices).transpose(1, 2) for source_name in source_outputs: raw_tensor = raw_outputs[source_name] for block in raw_tensor.blocks(): @@ -1114,7 +1079,7 @@ def _evaluate_system( ) tensor = raw_tensor.to( - dtype=torch.float64, + dtype=integration_dtype, device=work_device, ) if source_name in average_names or source_name in variance_names: @@ -1127,12 +1092,17 @@ def _evaluate_system( self.max_angular_momentum_target, "max_angular_momentum_target", ) - backrotated = _transform_tensor_with_precomputed_matrices( - tensor, - inverse_matrices, - inverse_target_wigner_matrices, - is_improper, - ) + 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 @@ -1164,7 +1134,7 @@ def _evaluate_system( ) if source_name in variance_names: - diagnostic_tensor = decompose_output( + diagnostic_tensor = decompose_quantity( source_name, backrotated, ) @@ -1203,7 +1173,7 @@ def _evaluate_system( ) if source_name in character_projection_names: - direct_tensor = decompose_output(source_name, tensor) + direct_tensor = decompose_quantity(source_name, tensor) contribution = character_projection_coefficients_from_batch( direct_tensor, so3_weights, @@ -1223,7 +1193,7 @@ def _evaluate_system( contribution, ) - results = torch.jit.annotate(Dict[str, TensorMap], {}) + results: Dict[str, TensorMap] = {} for source_name, requested_name in average_names.items(): mean = mts.add( average_references[source_name], diff --git a/python/metatomic_torch/metatomic/torch/o3/_tranformations.py b/python/metatomic_torch/metatomic/torch/o3/_tranformations.py index 95a7c7f4..ee7955cf 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_tranformations.py +++ b/python/metatomic_torch/metatomic/torch/o3/_tranformations.py @@ -1,9 +1,15 @@ """ 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 @@ -61,201 +67,162 @@ def _spherical_parity_factor( 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``. +def _determinants_3x3(matrices: torch.Tensor) -> torch.Tensor: + """Return the determinants of a ``(N, 3, 3)`` batch of matrices. - ``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``). + Written out explicitly because ``torch.linalg`` is not available in + TorchScript. """ - 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})." - ) + 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 single O(3) transformation, represented by a (3, 3) rotation or improper-rotation - matrix. - - The constructor stores a copy of ``matrix``. + 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): + 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) rotation or improper-rotation matrix + :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 """ - 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)." + if _improper is not None: + matrices = matrix + improper = _improper + else: + max_angular_momentum = _validate_nonnegative_integer( + "max_angular_momentum", max_angular_momentum ) - 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)." - ) + 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 matrix disagree with the cached parity and Wigner-D matrices. - self._matrix = matrix.clone() + # 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._is_improper = bool(torch.det(self._matrix) < 0) + self._improper = improper + self._wigner_D = _wigner_D - self._packed_wigner_D: torch.Tensor | None = None + @property + def matrices(self) -> torch.Tensor: + """The ``(N, 3, 3)`` batch of rotation or improper-rotation matrices.""" + return self._matrices - @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._packed_wigner_D = None - return transformation - - def _ensure_wigner_D_cache(self) -> torch.Tensor: - """Ensure that the packed Wigner-D cache has been built and return it. - - The packed buffer holds every ``ell`` up to ``max_angular_momentum``; it - inherits the dtype and device of the transformation matrix. + @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._packed_wigner_D is None: - self._packed_wigner_D = build_packed_wigner_matrices( - self._matrix.unsqueeze(0), - self._max_angular_momentum, + if self._matrices.size(0) != 1: + raise ValueError( + f"this O3Transformation holds {self._matrices.size(0)} " + "operations; use .matrices" ) + return self._matrices[0] - return self._packed_wigner_D - - 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) + @property + def max_angular_momentum(self) -> int: + """The maximum angular momentum with available Wigner-D matrices.""" + return self._max_angular_momentum - return wigner_matrices_for_lambda(self._ensure_wigner_D_cache(), 1, ell)[0] + @property + def improper(self) -> torch.Tensor: + """Boolean mask marking the improper operations in the batch.""" + return self._improper @property - def matrix(self) -> torch.Tensor: - """The (3, 3) rotation or improper-rotation matrix.""" - return self._matrix + 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 matrix.""" - return self._matrix.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 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 + """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]``.""" @@ -268,42 +235,258 @@ def _validate_ell_range(self, ell: int) -> int: 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 transformation to spherical values. + """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: (..., 2*ell+1) tensor of transformed spherical values + :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) - 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 + # 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" + ) - return transformed + 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, + ) + ) - def wigner_D_matrix(self, ell: int) -> torch.Tensor: - """Return the proper-part Wigner-D matrix for ``ell``. + 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 an improper transformation, :meth:`transform_spherical` applies the - inversion-parity factor separately. + 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, + ), + ) - :param ell: angular momentum in ``[0, max_angular_momentum]`` - :return: (2*ell+1, 2*ell+1) Wigner-D matrix + 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 """ - return self._wigner_D_cache_entry(ell) + 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( @@ -331,7 +514,7 @@ def random_transformations( :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 + :return: list of ``n`` single-operation :class:`O3Transformation` objects """ n = _validate_nonnegative_integer("n", n) max_angular_momentum = _validate_nonnegative_integer( @@ -360,12 +543,12 @@ def random_transformations( dim=1, ).reshape(n, 3, 3) - matrices_are_improper = [False] * n + 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) - matrices_are_improper = (signs < 0).tolist() + improper = signs < 0 identity = torch.eye( 3, @@ -380,76 +563,152 @@ def random_transformations( raise ValueError("Generated transformations are not orthogonal.") return [ - O3Transformation._create_no_checks( - matrix, + O3Transformation( + matrix.unsqueeze(0), max_angular_momentum, - is_improper=is_improper, - ) - for matrix, is_improper in zip( - R.unbind(0), - matrices_are_improper, - strict=True, + _improper=improper[index : index + 1], ) + for index, matrix in enumerate(R.unbind(0)) ] -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)] +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``. - 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] + ``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( - 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." + "Expected one transformation per system, but got " + f"len(systems)={n_systems} and " + f"len(transformations)={n_transformations}." ) - return [ - torch.nonzero(system_labels == system_id, as_tuple=False).reshape(-1) - for system_id in system_ids - ] + if system_ids is None: + return None -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 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 "system" not in parent_block.samples.names: + if len(validated_ids) != n_systems: raise ValueError( - "Rotational augmentation expects the values samples to include a 'system' " - "dimension when transforming gradients of multiple systems." + "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()}." ) - 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 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})." + ) - return [ - torch.nonzero( - gradient_system_labels == system_id, - as_tuple=False, - ).reshape(-1) - for system_id in system_ids - ] + +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: @@ -460,10 +719,15 @@ def transform_system(system: System, transformation: O3Transformation) -> System are preserved. :param system: input system - :param transformation: O(3) transformation to apply, matching + :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 @@ -475,65 +739,7 @@ def transform_system(system: System, transformation: O3Transformation) -> System 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) + return transformation.transform_systems(system)[0] def _component_axis_suffix(axis_name: str, prefix: str) -> tuple[bool, str]: @@ -603,73 +809,41 @@ def _validate_component_axis_metadata( return metadata -def _max_o3_lambda_in_tensor(tensor: TensorMap) -> int: - """Return the largest angular momentum in block values or attached gradients. +def _max_o3_lambda_in_block(key: LabelsEntry, block: TensorBlock) -> int: + """Return the largest angular momentum in one block's values or gradients. - A TensorMap containing only scalar or Cartesian component axes returns ``-1``. + A block 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: + 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 - 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 _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 -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 + return max_o3_lambda def transform_block( @@ -688,8 +862,8 @@ def transform_block( 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 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 @@ -697,7 +871,7 @@ def transform_block( :return: block with transformed values and gradients and unchanged labels; when ``systems`` is empty, the block is unchanged """ - system_ids = _validate_system_ids( + validated_ids = _validate_system_ids( systems, transformations, system_ids, @@ -712,51 +886,19 @@ def transform_block( 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, + 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, ) - 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( @@ -782,8 +924,8 @@ def transform_tensor( :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 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 @@ -798,7 +940,7 @@ def transform_tensor( else: system_ids_device = None - system_ids = _validate_system_ids( + validated_ids = _validate_system_ids( systems, transformations, system_ids, @@ -815,22 +957,23 @@ def transform_tensor( 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 + combined = _combine_transformations( + transformations, + _max_o3_lambda_in_tensor(tensor), + ) + return combined.transform_tensormap(tensor, validated_ids) -def _transformation_indices( +def _transformation_local_indices( samples: Labels, n_transformations: int, + system_ids: Optional[torch.Tensor], ) -> torch.Tensor: - """Map sample rows to local transformation indices.""" + """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: @@ -842,34 +985,62 @@ def _transformation_indices( 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 + 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_with_precomputed_matrices( +def _transform_component_values_batched( values: torch.Tensor, components: list[Labels], key: LabelsEntry, - transformation_indices: torch.Tensor, + local_indices: torch.Tensor, matrices: torch.Tensor, wigner_matrices: list[torch.Tensor], - is_improper: bool, + improper: torch.Tensor, ) -> torch.Tensor: - """Transform component axes with precomputed O(3) matrices.""" + """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 - parity = 1 + spherical_parity = 1 for component_index, (is_spherical, ell, sigma) in enumerate(metadata): if is_spherical: if ell >= len(wigner_matrices): - raise ValueError("angular momentum exceeds the Wigner-D storage") + raise ValueError( + f"ell={ell} exceeds " + f"max_angular_momentum={len(wigner_matrices) - 1}." + ) axis_matrices = wigner_matrices[ell] - parity *= _spherical_parity_factor(ell, sigma, is_improper) + # the factor acquired by improper operations; applied per row below + spherical_parity *= _spherical_parity_factor(ell, sigma, True) else: axis_matrices = matrices @@ -877,39 +1048,102 @@ def _transform_component_values_with_precomputed_matrices( 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), - ) + 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 parity != 1: - transformed = transformed * parity + 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_tensor_with_precomputed_matrices( +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], - is_improper: bool, + improper: torch.Tensor, + system_ids: Optional[torch.Tensor], ) -> 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``. + """Transform a TensorMap with a batch of O(3) operations. - 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. + 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 @@ -930,46 +1164,16 @@ def _transform_tensor_with_precomputed_matrices( 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, + blocks.append( + _transform_block_batched( key, - value_indices, + block, 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, - ), + improper, + system_ids, ) - blocks.append(new_block) + ) transformed = TensorMap(tensor.keys, blocks) for info_name, info_value in tensor.info().items(): diff --git a/python/metatomic_torch/metatomic/torch/o3/_utils.py b/python/metatomic_torch/metatomic/torch/o3/_utils.py index a57f033b..646ba840 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_utils.py +++ b/python/metatomic_torch/metatomic/torch/o3/_utils.py @@ -20,7 +20,12 @@ def validate_integer(name: str, value, minimum: int) -> int: raise TypeError(f"{name} must be an integer, got {type(value).__name__}") integer_value = int(value) if integer_value < minimum: - qualifier = "positive" if minimum == 1 else "non-negative" + 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 diff --git a/python/metatomic_torch/tests/o3.py b/python/metatomic_torch/tests/o3.py index 390bfe2c..018ef84f 100644 --- a/python/metatomic_torch/tests/o3.py +++ b/python/metatomic_torch/tests/o3.py @@ -23,7 +23,7 @@ # entry point yet; their tests below compare them against the public transform_tensor. from metatomic.torch.o3._tranformations 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("angular momentum 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 index ec160d57..ff7ff2bd 100644 --- a/python/metatomic_torch/tests/symmetrized_model.py +++ b/python/metatomic_torch/tests/symmetrized_model.py @@ -18,16 +18,12 @@ System, load_atomistic_model, ) -from metatomic.torch.o3 import O3Transformation, transform_system +from metatomic.torch.o3 import O3Transformation from metatomic.torch.o3._decompose import ( _cartesian_vectors_to_spherical, _o3_mu_labels, _symmetric_matrices_to_spherical, - decompose_output, -) -from metatomic.torch.o3._projections import ( - _character_projection_coefficients_from_rotation_batch, - _character_projections_from_proper_and_improper_coefficients, + decompose_quantity, ) from metatomic.torch.o3._quadrature import ( _rotations_from_euler_angles, @@ -35,23 +31,7 @@ get_euler_angles_quadrature, get_rotation_quadrature, ) -from metatomic.torch.o3._symmetrized import ( - _clamp_roundoff_negative_diagnostic, - _component_norm_squared, - _mean_variance_over_components, - _reduce_weighted_centered_batch, - _transform_system_batch, - _transform_system_geometry_batch, - _variance_from_centered_moments, -) -from metatomic.torch.o3._utils import ( - group_samples_by_rotated_copy, - map_selected_atoms_to_rotated_copies, -) -from metatomic.torch.o3._wigner import ( - build_packed_wigner_matrices, - wigner_matrices_for_lambda, -) +from metatomic.torch.o3._utils import map_selected_atoms_to_rotated_copies def _make_single_block_tensor_map( @@ -208,6 +188,63 @@ def forward( 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.""" @@ -564,534 +601,6 @@ def forward( return result -def _system_with_neighbor_lists(dtype: torch.dtype) -> System: - """Create a test system with populated and empty neighbor lists.""" - positions = torch.tensor( - [[0.2, -0.1, 0.3], [1.1, 0.7, -0.4], [-0.3, 0.6, 1.2]], - dtype=dtype, - ) - cell = torch.tensor( - [[2.5, 0.1, 0.0], [0.0, 2.2, 0.2], [0.1, 0.0, 2.7]], - dtype=dtype, - ) - system = System( - types=torch.tensor([6, 1, 8]), - positions=positions, - cell=cell, - pbc=torch.tensor([True, True, True]), - ) - - samples = Labels( - [ - "first_atom", - "second_atom", - "cell_shift_a", - "cell_shift_b", - "cell_shift_c", - ], - torch.tensor([[0, 1, 0, 0, 0], [1, 2, 1, 0, 0]]), - ) - components = [Labels.range("xyz", 3)] - properties = Labels.range("distance", 1) - system.add_neighbor_list( - NeighborListOptions(3.0, False, True, "populated"), - TensorBlock( - values=torch.stack( - [ - positions[1] - positions[0], - positions[2] - positions[1] + cell[0], - ] - ).unsqueeze(-1), - samples=samples, - components=components, - properties=properties, - ), - ) - system.add_neighbor_list( - NeighborListOptions(1.0, True, False, "empty"), - TensorBlock( - values=torch.empty((0, 3, 1), dtype=dtype), - samples=Labels( - list(samples.names), - torch.empty((0, len(samples.names)), dtype=torch.int64), - ), - components=components, - properties=properties, - ), - ) - return system - - -class TestSystemGeometryBatch: - """Test batched O(3) transformation of System geometry.""" - - @pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) - def test_matches_individual_o3_transformations(self, dtype): - """Batched geometry should match one transformation at a time.""" - proper = torch.tensor( - [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], - dtype=dtype, - ) - matrices = torch.stack([torch.eye(3, dtype=dtype), proper, -proper]) - system = _system_with_neighbor_lists(dtype) - - transformed = _transform_system_geometry_batch(system, matrices) - - assert len(transformed) == len(matrices) - for matrix, actual in zip(matrices, transformed, strict=True): - expected = transform_system( - system, - O3Transformation(matrix, max_angular_momentum=0), - ) - assert torch.equal(actual.positions, expected.positions) - assert torch.equal(actual.cell, expected.cell) - assert torch.equal(actual.types, expected.types) - assert torch.equal(actual.pbc, expected.pbc) - assert actual.known_neighbor_lists() == expected.known_neighbor_lists() - for options in expected.known_neighbor_lists(): - assert torch.equal( - actual.get_neighbor_list(options).values, - expected.get_neighbor_list(options).values, - ) - - def test_preserves_neighbor_autograd(self): - """Rotated neighbor vectors should differentiate through positions and cell.""" - positions = torch.tensor( - [[0.2, -0.1, 0.3], [1.1, 0.7, -0.4]], - dtype=torch.float64, - requires_grad=True, - ) - cell = torch.tensor( - [[2.5, 0.1, 0.0], [0.0, 2.2, 0.2], [0.1, 0.0, 2.7]], - dtype=torch.float64, - requires_grad=True, - ) - system = System( - types=torch.tensor([6, 1]), - positions=positions, - cell=cell, - pbc=torch.tensor([True, True, True]), - ) - cell_shift = torch.tensor([1.0, -1.0, 0.0], dtype=torch.float64) - neighbor_vector = positions[1] - positions[0] + cell_shift @ cell - options = NeighborListOptions(4.0, False, True) - system.add_neighbor_list( - options, - TensorBlock( - values=neighbor_vector.reshape(1, 3, 1), - samples=Labels( - [ - "first_atom", - "second_atom", - "cell_shift_a", - "cell_shift_b", - "cell_shift_c", - ], - torch.tensor([[0, 1, 1, -1, 0]]), - ), - components=[Labels.range("xyz", 3)], - properties=Labels.range("distance", 1), - ), - ) - proper = torch.tensor( - [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], - dtype=torch.float64, - ) - matrices = torch.stack([proper, -proper]) - - transformed = _transform_system_geometry_batch(system, matrices) - loss = sum( - transformed_system.get_neighbor_list(options).values.square().sum() - for transformed_system in transformed - ) - position_gradient, cell_gradient = torch.autograd.grad( - loss, - (positions, cell), - ) - - vector_gradient = 2 * len(matrices) * neighbor_vector.detach() - assert torch.allclose( - position_gradient, - torch.stack([-vector_gradient, vector_gradient]), - ) - assert torch.allclose( - cell_gradient, - torch.outer(cell_shift, vector_gradient), - ) - - def test_rejects_invalid_matrix_batches(self): - """Matrix batches should have a non-empty shape and match the System.""" - system = _system_with_neighbor_lists(torch.float64) - invalid_shapes = [(3, 3), (0, 3, 3), (2, 2, 3), (2, 3, 2)] - message = "matrices must have shape (N, 3, 3) with N > 0" - for shape in invalid_shapes: - with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - _transform_system_geometry_batch( - system, - torch.empty(shape, dtype=torch.float64), - ) - - message = "system and matrices must have the same dtype and device" - with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - _transform_system_geometry_batch( - system, - torch.eye(3, dtype=torch.float32).unsqueeze(0), - ) - - -class TestSystemBatch: - """Test batched O(3) transformation of complete Systems.""" - - @pytest.mark.parametrize("is_improper", [False, True]) - def test_transforms_spherical_custom_data(self, is_improper): - """Every transformed System should contain the corresponding custom data.""" - proper_matrices = torch.tensor( - [ - [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], - [ - [-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, - ) - matrices = -proper_matrices if is_improper else proper_matrices - packed_wigner = build_packed_wigner_matrices( - proper_matrices, - max_angular_momentum=1, - ) - wigner_matrices = [ - wigner_matrices_for_lambda( - packed_wigner, - n_matrices=len(matrices), - o3_lambda=o3_lambda, - ) - for o3_lambda in range(2) - ] - - system = System( - types=torch.tensor([6, 8]), - positions=torch.tensor( - [[0.2, -0.1, 0.3], [1.1, 0.7, -0.4]], - dtype=torch.float64, - ), - cell=torch.eye(3, dtype=torch.float64) * 4.0, - pbc=torch.tensor([True, True, True]), - ) - values = torch.tensor( - [[[1.0], [2.0], [3.0]], [[-0.5], [1.5], [0.25]]], - dtype=torch.float64, - ) - system.add_data( - "mtt::field", - TensorMap( - Labels( - ["o3_lambda", "o3_sigma"], - torch.tensor([[1, 1]]), - ), - [ - TensorBlock( - values=values, - samples=Labels.range("atom", 2), - components=[_o3_mu_labels(1, values.device)], - properties=Labels.range("property", 1), - ) - ], - ), - ) - - transformed = torch.jit.script(_transform_system_batch)( - system, - matrices, - wigner_matrices, - is_improper=is_improper, - ) - - assert len(transformed) == len(matrices) - for matrix, transformed_system in zip(matrices, transformed, strict=True): - expected_system = transform_system( - system, - O3Transformation(matrix, max_angular_momentum=1), - ) - assert "mtt::field" in transformed_system.known_data() - mts.allclose_raise( - transformed_system.get_data("mtt::field"), - expected_system.get_data("mtt::field"), - rtol=0.0, - atol=1.0e-12, - ) - - def test_input_limit_distinguishes_spherical_from_cartesian(self): - """A zero spherical-rank limit should still allow Cartesian custom data.""" - matrix = torch.tensor( - [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], - dtype=torch.float64, - ).unsqueeze(0) - packed_wigner = build_packed_wigner_matrices( - matrix, - max_angular_momentum=0, - ) - wigner_matrices = [ - wigner_matrices_for_lambda( - packed_wigner, - n_matrices=1, - o3_lambda=0, - ) - ] - system = System( - types=torch.tensor([6]), - positions=torch.tensor([[0.2, -0.1, 0.3]], dtype=torch.float64), - cell=torch.eye(3, dtype=torch.float64) * 4.0, - pbc=torch.tensor([True, True, True]), - ) - cartesian = TensorMap( - Labels("_", torch.tensor([[0]])), - [ - TensorBlock( - values=torch.tensor( - [[[1.0], [2.0], [3.0]]], - dtype=torch.float64, - ), - samples=Labels.range("atom", 1), - components=[Labels.range("xyz", 3)], - properties=Labels.range("property", 1), - ) - ], - ) - system.add_data("mtt::field", cartesian) - - transformed = _transform_system_batch( - system, - matrix, - wigner_matrices, - is_improper=False, - ) - expected = transform_system( - system, - O3Transformation(matrix[0], max_angular_momentum=0), - ) - mts.allclose_raise( - transformed[0].get_data("mtt::field"), - expected.get_data("mtt::field"), - rtol=0.0, - atol=1.0e-12, - ) - - spherical_system = System( - types=system.types, - positions=system.positions, - cell=system.cell, - pbc=system.pbc, - ) - 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), - ) - ], - ), - ) - model = SymmetrizedModel( - _LinearEnergyModel(), - max_angular_momentum_target=0, - max_angular_momentum_grid=2, - ) - 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], - {"energy": ModelOutput(sample_kind="system")}, - None, - ) - - -class TestCharacterProjections: - """Test construction of character projections from rotated model responses.""" - - @pytest.mark.parametrize("n_samples", [0, 2]) - def test_batch_coefficients_match_rotation_by_rotation_sum(self, n_samples): - """Batching should match summing the weighted rotations individually.""" - generator = torch.Generator().manual_seed(7) - n_rotations = 4 - dimension = 3 - values = torch.randn( - (n_rotations, n_samples, 2, 3), - dtype=torch.float64, - generator=generator, - ) - weights = torch.tensor( - [0.50, -0.25, 0.30, 0.45], - dtype=torch.float32, - ) - inverse_wigner_matrices = torch.randn( - (n_rotations, dimension, dimension), - dtype=torch.float32, - generator=generator, - ) - - coefficients = _character_projection_coefficients_from_rotation_batch( - values, - weights, - inverse_wigner_matrices, - ) - - expected = torch.zeros( - (n_samples, dimension, dimension, 2, 3), - dtype=torch.float64, - ) - for rotation in range(n_rotations): - expected += ( - weights[rotation].to(torch.float64) - * inverse_wigner_matrices[rotation] - .to(torch.float64) - .reshape(1, dimension, dimension, 1, 1) - * values[rotation].reshape(n_samples, 1, 1, 2, 3) - ) - - assert torch.allclose(coefficients, expected, rtol=0.0, atol=1e-12) - - @pytest.mark.parametrize("chi_lambda", [0, 1, 2]) - def test_factorization_matches_all_rotation_pairs(self, chi_lambda): - """The factorization should match summing every pair of rotations.""" - generator = torch.Generator().manual_seed(11 + chi_lambda) - n_rotations = 4 - dimension = 2 * chi_lambda + 1 - proper_values = torch.randn( - (n_rotations, 2, 2, 1), - dtype=torch.float64, - generator=generator, - ) - improper_values = torch.randn( - (n_rotations, 2, 2, 1), - dtype=torch.float64, - generator=generator, - ) - weights = torch.tensor( - [0.50, -0.25, 0.30, 0.45], - dtype=torch.float64, - ) - inverse_wigner_matrices = torch.randn( - (n_rotations, dimension, dimension), - dtype=torch.float64, - generator=generator, - ) - proper_coefficients = _character_projection_coefficients_from_rotation_batch( - proper_values, - weights, - inverse_wigner_matrices, - ) - improper_coefficients = _character_projection_coefficients_from_rotation_batch( - improper_values, - weights, - inverse_wigner_matrices, - ) - - sigma_plus, sigma_minus = ( - _character_projections_from_proper_and_improper_coefficients( - proper_coefficients, - improper_coefficients, - chi_lambda, - ) - ) - - expected = [] - for chi_sigma in (1, -1): - combined_values = proper_values + ( - chi_sigma * (-1) ** chi_lambda * improper_values - ) - direct_sum = torch.zeros_like(combined_values[0]) - for first_rotation in range(n_rotations): - for second_rotation in range(n_rotations): - character = torch.sum( - inverse_wigner_matrices[first_rotation] - * inverse_wigner_matrices[second_rotation] - ) - direct_sum += ( - float(dimension) - / 4.0 - * weights[first_rotation] - * weights[second_rotation] - * character - * combined_values[first_rotation] - * combined_values[second_rotation] - ) - expected.append(direct_sum) - - assert torch.allclose(sigma_plus, expected[0], rtol=0.0, atol=1e-12) - assert torch.allclose(sigma_minus, expected[1], rtol=0.0, atol=1e-12) - assert torch.all(sigma_plus >= 0) - assert torch.all(sigma_minus >= 0) - - -class TestWignerStorage: - """Test persistent Wigner-D storage for the quadrature grid.""" - - @pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) - def test_packed_matrices_match_o3(self, dtype): - """Packing and rank views should preserve the public O(3) matrices.""" - proper_rotation = torch.tensor( - [ - [0.0, -1.0, 0.0], - [1.0, 0.0, 0.0], - [0.0, 0.0, 1.0], - ], - dtype=dtype, - ) - matrices = torch.stack( - [ - torch.eye(3, dtype=dtype), - -proper_rotation, - ] - ) - max_angular_momentum = 2 - - packed = build_packed_wigner_matrices(matrices, max_angular_momentum) - - assert packed.dim() == 1 - assert packed.numel() == len(matrices) * sum( - (2 * o3_lambda + 1) ** 2 for o3_lambda in range(max_angular_momentum + 1) - ) - assert packed.dtype == matrices.dtype - assert packed.device == matrices.device - - transformations = [ - O3Transformation(matrix, max_angular_momentum) - for matrix in matrices.unbind(0) - ] - for o3_lambda in range(max_angular_momentum + 1): - actual = wigner_matrices_for_lambda( - packed, - len(matrices), - o3_lambda, - ) - expected = torch.stack( - [ - transformation.wigner_D_matrix(o3_lambda) - for transformation in transformations - ] - ) - assert torch.equal(actual, expected) - - def test_rank_view_rejects_out_of_range_lambda(self): - """Rank views should reject ranks beyond the packed storage.""" - message = "o3_lambda exceeds the packed Wigner-D storage" - with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - wigner_matrices_for_lambda(torch.empty(1), 1, 1) - - class TestQuadrature: """Test quadrature weights and grid properties.""" @@ -1186,38 +695,6 @@ def test_rotation_quadrature_matrices(self): class TestSymmetrizedModelConstruction: """Test construction of the quadrature and persistent Wigner-D storage.""" - def test_constructs_registered_buffers(self): - """Constructor limits should determine the grid and Wigner-D storage.""" - model = SymmetrizedModel( - _EmptyModel(), - max_angular_momentum_target=1, - max_angular_momentum_input=2, - max_angular_momentum_character=1, - batch_size=7, - ) - - assert model.max_angular_momentum_target == 1 - assert model.max_angular_momentum_input == 2 - assert model.max_angular_momentum_character == 1 - assert model.max_angular_momentum_grid == 3 - assert model.batch_size == 7 - - buffers = dict(model.named_buffers()) - assert set(buffers) == { - "_rotation_matrices", - "_rotation_weights", - "_packed_wigner_matrices", - } - assert buffers["_rotation_matrices"].dtype == torch.float64 - assert buffers["_rotation_weights"].dtype == torch.float64 - assert buffers["_packed_wigner_matrices"].dtype == torch.float64 - - n_rotations = len(buffers["_rotation_matrices"]) - expected_wigner_elements = n_rotations * sum( - (2 * o3_lambda + 1) ** 2 for o3_lambda in range(3) - ) - assert buffers["_packed_wigner_matrices"].numel() == expected_wigner_elements - def test_character_limit_controls_default_grid(self): """Character sectors should raise the default grid degree when necessary.""" model = SymmetrizedModel( @@ -1238,58 +715,51 @@ def test_rejects_grid_too_small_for_character_sectors(self): SymmetrizedModel( _EmptyModel(), max_angular_momentum_target=0, - max_angular_momentum_character=2, - max_angular_momentum_grid=3, + 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, ) - @pytest.mark.parametrize( - ("argument", "value", "error", "message"), - [ - ( - "max_angular_momentum_target", - -1, - ValueError, - "max_angular_momentum_target must be non-negative, got -1", - ), - ( - "max_angular_momentum_target", - True, - TypeError, - "max_angular_momentum_target must be an integer, got bool", - ), - ( - "max_angular_momentum_input", - 1.5, - TypeError, - "max_angular_momentum_input must be an integer, got float", - ), - ( - "max_angular_momentum_character", - -1, - ValueError, - "max_angular_momentum_character must be non-negative, got -1", - ), - ("batch_size", 0, ValueError, "batch_size must be positive, got 0"), - ( - "max_angular_momentum_grid", - -1, - ValueError, - "max_angular_momentum_grid must be non-negative, got -1", - ), - ], - ) - def test_rejects_invalid_constructor_arguments( - self, - argument, - value, - error, - message, - ): - """Every integer constructor argument should enforce its documented range.""" - arguments = {"max_angular_momentum_target": 0, argument: value} + 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, + ) - with pytest.raises(error, match=f"^{re.escape(message)}$"): - SymmetrizedModel(_EmptyModel(), **arguments) + 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.""" @@ -1591,9 +1061,6 @@ def test_rejects_a_negative_quadrature_error_and_converges(self): [ "energy/pbe", "mtt::feature::node", - # "o3::variance_extra::" is not the reserved prefix: the full name - # is passed through as a source output - "o3::variance_extra::energy", # the reserved prefix is stripped exactly once, keeping "mtt::aux::" "mtt::aux::features", ], @@ -1931,6 +1398,112 @@ def test_average_output_preserves_implicit_autograd(self): 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() @@ -1980,25 +1553,27 @@ def test_rejects_invalid_requests_before_model_evaluation(self): ) assert base_model.call_count == 0 - def test_rejects_downcast_integration_buffers(self): - """Calling .float() on the module must fail loudly at the next forward.""" + 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) - message = ( - "SymmetrizedModel integration buffers must remain float64, got " - "torch.float32; do not call .float() or .half() on the module" - ) - with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - model( - [_forward_test_system([[1.0, 2.0, 3.0]])], - {"energy": ModelOutput(sample_kind="system")}, + 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( @@ -2060,6 +1635,22 @@ def test_is_scriptable_and_serializable(self, tmp_path): 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.""" @@ -2126,45 +1717,7 @@ def test_wrap_declares_capabilities(self, max_angular_momentum_character): else: assert capabilities.outputs[character_name].unit == squared_unit - @pytest.mark.parametrize( - ("outputs", "expected_max_angular_momentum_target"), - [ - ({"energy": ModelOutput(unit="eV", sample_kind="system")}, 0), - ({"feature": ModelOutput(sample_kind="atom")}, 0), - ( - { - "energy": ModelOutput(unit="eV", sample_kind="system"), - "non_conservative_force": ModelOutput( - unit="eV/A", - sample_kind="atom", - ), - }, - 1, - ), - ( - { - "non_conservative_stress": ModelOutput( - unit="eV/A^3", - sample_kind="system", - ) - }, - 2, - ), - # a custom output is skipped, the standard ones still set the limit - ( - { - "energy": ModelOutput(unit="eV", sample_kind="system"), - "mtt::custom": ModelOutput(sample_kind="system"), - }, - 0, - ), - ], - ) - def test_guesses_limits_from_standard_quantities( - self, - outputs, - expected_max_angular_momentum_target, - ): + def test_guesses_limits_from_standard_quantities(self): """Both limits default to what the standard quantities require.""" class _VelocityInputModel(_EmptyModel): @@ -2175,27 +1728,35 @@ def requested_inputs(self) -> Dict[str, ModelOutput]: "mtt::field": ModelOutput(sample_kind="atom"), } - 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) - - assert ( - wrapped.module.max_angular_momentum_target - == expected_max_angular_momentum_target - ) - # velocity is a Cartesian vector - assert wrapped.module.max_angular_momentum_input == 1 + 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.""" @@ -2386,7 +1947,11 @@ def test_preserves_requirements_and_runs_after_save_load(self, tmp_path): @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.""" + """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"), @@ -2477,341 +2042,13 @@ def test_saved_wrapper_runs_on_cuda(self, tmp_path): ) -class TestSelectedAtomsColumnOrder: - def test_system_column_found_by_name(self): - # the rotated-copy index must go into the "system" column wherever it - # is, not positionally into column 0 - 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] - - -_SAME_SAMPLE_LABELS_MESSAGE = ( - "SymmetrizedModel expects every rotated copy to produce the same sample " - "labels in the same order." -) - - -@pytest.mark.parametrize( - ("sample_values", "message"), - [ - ( - [[0, 0], [2, 0]], - "encountered output samples with out-of-range rotated-copy " - "indices: the system column spans [0, 2], expected [0, 1]", - ), - ([[0, 0], [0, 1], [1, 0]], _SAME_SAMPLE_LABELS_MESSAGE), - ([[0, 0], [0, 1], [0, 2], [1, 0]], _SAME_SAMPLE_LABELS_MESSAGE), - ([[0, 0], [0, 1], [1, 0], [1, 2]], _SAME_SAMPLE_LABELS_MESSAGE), - ], -) -def test_rotated_copy_layout_rejects_inconsistent_samples(sample_values, message): - """Samples from different rotated copies must never be mixed.""" - samples = Labels(["system", "atom"], torch.tensor(sample_values)) - block = TensorBlock( - values=torch.zeros((len(samples), 1), dtype=torch.float64), - samples=samples, - components=[], - properties=Labels.range("property", 1), - ) - - with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - group_samples_by_rotated_copy(block, n_rotated_copies=2) - - -@pytest.mark.parametrize( - ("sample_values", "values", "n_rotated_copies", "expected_values"), - [ - ( - [[3, 0], [5, 0]], - [3.0, 5.0], - 1, - [[[3.0], [5.0]]], - ), - ( - [[3, 1], [3, 0], [5, 1], [5, 0]], - [13.0, 3.0, 15.0, 5.0], - 2, - [[[3.0], [5.0]], [[13.0], [15.0]]], - ), - ], -) -def test_group_samples_by_rotated_copy( - sample_values, values, n_rotated_copies, expected_values -): - """Values and shared labels should remain aligned after grouping.""" - samples = Labels(["atom", "system"], torch.tensor(sample_values)) - block = TensorBlock( - values=torch.tensor(values, dtype=torch.float64).reshape(-1, 1), - samples=samples, - components=[], - properties=Labels.range("property", 1), - ) - - grouped_values, shared_names, shared_values = group_samples_by_rotated_copy( - block, n_rotated_copies - ) - - assert torch.equal( - grouped_values, - torch.tensor(expected_values, dtype=torch.float64), - ) - assert shared_names == ["atom"] - assert shared_values.tolist() == [[3], [5]] - - -@pytest.mark.parametrize("component_shape", [(), (2, 3)]) -def test_weighted_centered_batch_moments(component_shape): - """Compute weighted moments and reuse one fixed reference across batches.""" - n_rotated_copies = 3 - n_samples = 2 - n_properties = 2 - values = torch.arange( - n_rotated_copies * n_samples * int(np.prod(component_shape)) * n_properties, - dtype=torch.float64, - ).reshape(n_rotated_copies * n_samples, *component_shape, n_properties) - components = [ - Labels.range(name, size) - for name, size in zip(("a", "b"), component_shape, strict=False) - ] - tensor = TensorMap( - Labels("kind", torch.tensor([[0]])), - [ - TensorBlock( - values=values, - samples=Labels( - ["system", "item"], - torch.tensor( - [ - [copy, item] - for copy in range(n_rotated_copies) - for item in (5, 7) - ] - ), - ), - components=components, - properties=Labels.range("property", n_properties), - ) - ], - ) - weights = torch.tensor([0.2, -0.1, 0.4], dtype=torch.float64) - - moments = _reduce_weighted_centered_batch( - tensor, - weights, - input_system_index=4, - reference=None, - compute_second_moments=True, - ) - first_moment, second, absolute_second, reference = moments - - values_by_copy = values.reshape( - n_rotated_copies, n_samples, *component_shape, n_properties - ) - centered = values_by_copy - values_by_copy[0] - weight_shape = (n_rotated_copies,) + (1,) * (centered.ndim - 1) - assert torch.allclose( - first_moment.block().values, - torch.sum(weights.reshape(weight_shape) * centered, dim=0), - ) - squared_norms = centered**2 - if component_shape: - squared_norms = squared_norms.sum(dim=tuple(range(2, 2 + len(component_shape)))) - assert second is not None - assert absolute_second is not None - assert torch.allclose( - second.block().values, - torch.sum(weights.reshape(n_rotated_copies, 1, 1) * squared_norms, dim=0), - ) - assert torch.allclose( - absolute_second.block().values, - torch.sum( - torch.abs(weights).reshape(n_rotated_copies, 1, 1) * squared_norms, - dim=0, - ), - ) - expected_samples = Labels( - ["system", "item"], - torch.tensor([[4, 5], [4, 7]]), - ) - assert first_moment.block().samples == expected_samples - assert second.block().components == [] - - initial_reference_values = values_by_copy[0].clone() - assert torch.equal(reference.block().values, initial_reference_values) - - # Simulate a later batch with the same layout but different response values. - tensor.block().values.add_(10.0) - later_values_by_copy = tensor.block().values.reshape( - n_rotated_copies, n_samples, *component_shape, n_properties - ) - later_centered = later_values_by_copy - initial_reference_values.unsqueeze(0) - - later_moments = _reduce_weighted_centered_batch( - tensor, - weights, - input_system_index=4, - reference=reference, - compute_second_moments=False, - ) - first_moment, second, absolute_second, reused_reference = later_moments - assert torch.allclose( - first_moment.block().values, - torch.sum(weights.reshape(weight_shape) * later_centered, dim=0), - ) - assert second is None - assert absolute_second is None - assert reused_reference is reference - assert torch.equal(reference.block().values, initial_reference_values) - - -@pytest.mark.parametrize("component_shape", [(), (3,), (2, 3)]) -def test_component_norm_squared(component_shape): - """All component axes should be contracted without changing metadata.""" - shape = (2, *component_shape, 2) - values = torch.arange(int(np.prod(shape)), dtype=torch.float64).reshape(shape) - tensor = _make_single_block_tensor_map(values) - - result = _component_norm_squared(tensor) - - expected = values.square() - if component_shape: - expected = expected.sum(dim=tuple(range(1, 1 + len(component_shape)))) - assert torch.equal(result.block().values, expected) - assert result.keys == tensor.keys - assert result.block().samples == tensor.block().samples - assert result.block().components == [] - assert result.block().properties == tensor.block().properties - - -def test_variance_from_centered_moments(): - """Centered first and second moments should give component-summed variance.""" - component_shape = (2, 3) - shape = (2, *component_shape, 2) - centered_first_moment_values = ( - torch.arange(int(np.prod(shape)), dtype=torch.float64).reshape(shape) / 10 - ) - centered_first_moment = _make_single_block_tensor_map(centered_first_moment_values) - - norm_squared = centered_first_moment_values.square().sum(dim=(1, 2)) - expected_variance = torch.tensor([[0.25, 0.5], [0.75, 1.0]], dtype=torch.float64) - centered_second_moment = _make_single_block_tensor_map( - norm_squared + expected_variance - ) - absolute_centered_second_moment = _make_single_block_tensor_map( - norm_squared + expected_variance + 1.0 - ) - - variance = _variance_from_centered_moments( - centered_first_moment, - centered_second_moment, - absolute_centered_second_moment, - n_grid_points=12, - max_angular_momentum_grid=3, - ) - - assert torch.allclose(variance.block().values, expected_variance) - assert variance.keys == centered_first_moment.keys - assert variance.block().samples == centered_first_moment.block().samples - assert variance.block().components == [] - assert variance.block().properties == centered_first_moment.block().properties - - -def test_centered_variance_is_stable_with_large_offset(): - """A common offset should not cause cancellation in the variance.""" - values = torch.tensor( - [1.0e12, 1.0e12 + 1.0, 1.0e12 + 2.0, 1.0e12 + 3.0], - dtype=torch.float64, - ).reshape(-1, 1) - tensor = _make_single_block_tensor_map(values, sample_name="system") - weights = torch.tensor([0.125, 0.375, 0.375, 0.125], dtype=torch.float64) - - first, second, absolute_second, _ = _reduce_weighted_centered_batch( - tensor, - weights, - input_system_index=7, - reference=None, - compute_second_moments=True, - ) - assert second is not None - assert absolute_second is not None - variance = _variance_from_centered_moments( - first, - second, - absolute_second, - n_grid_points=4, - max_angular_momentum_grid=3, - ) - - assert torch.allclose( - variance.block().values, - torch.tensor([[0.75]], dtype=torch.float64), - rtol=0.0, - atol=1.0e-12, - ) - - -def test_roundoff_negative_diagnostic_uses_its_scale(): - """Only negative values within the summation tolerance should be clamped.""" - dtype = torch.float64 - scale = 1.0e12 - n_grid_points = 100 - n_epsilon = n_grid_points * torch.finfo(dtype).eps - gamma = n_epsilon / (1.0 - n_epsilon) - tolerance = 64.0 * gamma * scale - - cleaned = _clamp_roundoff_negative_diagnostic( - _make_single_block_tensor_map( - torch.tensor([[-0.5 * tolerance], [2.0]], dtype=dtype) - ), - _make_single_block_tensor_map(torch.tensor([[scale], [scale]], dtype=dtype)), - n_grid_points=n_grid_points, - quantity="variance", - max_angular_momentum_grid=3, - ) - assert cleaned.block().values[0, 0].item() == 0.0 - assert cleaned.block().values[1, 0].item() == 2.0 - - message = ( - "finite O(3) variance is materially negative; the quadrature does not " - "resolve this response. Increase max_angular_momentum_grid above 3 and check " - "convergence" - ) - with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - _clamp_roundoff_negative_diagnostic( - _make_single_block_tensor_map( - torch.tensor([[-2.0 * tolerance]], dtype=dtype) - ), - _make_single_block_tensor_map(torch.tensor([[scale]], dtype=dtype)), - n_grid_points=n_grid_points, - quantity="variance", - max_angular_momentum_grid=3, - ) - - -@pytest.mark.parametrize("component_shape", [(), (3,), (2, 3)]) -def test_mean_variance_over_components(component_shape): - """Divide by component count without aggregating or creating samples.""" - n_samples = 2 - variance_values = ( - torch.arange(n_samples * 2, dtype=torch.float64).reshape(n_samples, 2) + 1.0 - ) - variance = _make_single_block_tensor_map(variance_values, sample_name="atom") - component_layout = _make_single_block_tensor_map( - torch.zeros(n_samples, *component_shape, 2, dtype=torch.float64), - sample_name="atom", - ) - - result = _mean_variance_over_components(variance, component_layout) - - n_components = int(np.prod(component_shape)) if component_shape else 1 - assert torch.equal(result.block().values, variance_values / n_components) - assert result.keys == variance.keys - assert result.block().samples == variance.block().samples - assert result.block().components == [] - assert result.block().properties == variance.block().properties +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(): @@ -2963,13 +2200,13 @@ def test_symmetric_matrices_to_spherical_commutes_with_o3(inversion): "charge", ], ) -def test_decompose_output_scalar_quantities(source_name): +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_output(source_name, tensor) + result = decompose_quantity(source_name, tensor) assert result.keys.names == ["o3_lambda", "o3_sigma"] assert result.keys.values.tolist() == [[0, 1]] @@ -2987,7 +2224,7 @@ def test_decompose_output_scalar_quantities(source_name): "velocity", ], ) -def test_decompose_output_cartesian_vectors_preserve_autograd(source_name): +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]]], @@ -2996,7 +2233,7 @@ def test_decompose_output_cartesian_vectors_preserve_autograd(source_name): ) tensor = _tensor_map_with_components(values, ["xyz"]) - result = decompose_output(source_name, tensor) + result = decompose_quantity(source_name, tensor) assert result.keys.names == ["o3_lambda", "o3_sigma"] assert result.keys.values.tolist() == [[1, 1]] @@ -3010,7 +2247,7 @@ def test_decompose_output_cartesian_vectors_preserve_autograd(source_name): assert torch.equal(values.grad, torch.ones_like(values)) -def test_decompose_output_non_conservative_stress_combines_irreps(): +def test_decompose_quantity_non_conservative_stress_combines_irreps(): """Stress should return l=0 and l=2 blocks and silently discard skew.""" values = torch.zeros((2, 3, 3, 1), dtype=torch.float64) values[0, :, :, 0] = torch.eye(3, dtype=torch.float64) @@ -3018,7 +2255,7 @@ def test_decompose_output_non_conservative_stress_combines_irreps(): values[1, 1, 0, 0] = -2.0 tensor = _tensor_map_with_components(values, ["xyz_1", "xyz_2"]) - result = decompose_output("non_conservative_stress/direct", tensor) + result = decompose_quantity("non_conservative_stress/direct", tensor) assert result.keys.names == ["o3_lambda", "o3_sigma"] assert result.keys.values.tolist() == [[0, 1], [2, 1]] @@ -3039,14 +2276,14 @@ def test_decompose_output_non_conservative_stress_combines_irreps(): assert block_l2.properties == tensor.block().properties -def test_decompose_output_does_not_infer_custom_cartesian_semantics(): +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_output("mtt::custom", tensor) + result = decompose_quantity("mtt::custom", tensor) mts.equal_raise(result, tensor) From 41c348ddff07c759ccef20377d024f50dc7c244b Mon Sep 17 00:00:00 2001 From: ppegolo Date: Wed, 5 Aug 2026 16:18:24 +0200 Subject: [PATCH 13/18] Fix typo in _transformations.py file name --- python/metatomic_torch/metatomic/torch/o3/__init__.py | 2 +- python/metatomic_torch/metatomic/torch/o3/_symmetrized.py | 2 +- .../torch/o3/{_tranformations.py => _transformations.py} | 0 python/metatomic_torch/tests/o3.py | 2 +- 4 files changed, 3 insertions(+), 3 deletions(-) rename python/metatomic_torch/metatomic/torch/o3/{_tranformations.py => _transformations.py} (100%) 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/_symmetrized.py b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py index 99e0af53..3b05ac44 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py +++ b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py @@ -32,7 +32,7 @@ character_projection_tensormap_from_cosets, ) from ._quadrature import choose_quadrature, get_rotation_quadrature -from ._tranformations import O3Transformation, _max_o3_lambda_in_tensor +from ._transformations import O3Transformation, _max_o3_lambda_in_tensor from ._utils import ( group_samples_by_rotated_copy, map_selected_atoms_to_rotated_copies, diff --git a/python/metatomic_torch/metatomic/torch/o3/_tranformations.py b/python/metatomic_torch/metatomic/torch/o3/_transformations.py similarity index 100% rename from python/metatomic_torch/metatomic/torch/o3/_tranformations.py rename to python/metatomic_torch/metatomic/torch/o3/_transformations.py diff --git a/python/metatomic_torch/tests/o3.py b/python/metatomic_torch/tests/o3.py index 018ef84f..e66a1bd5 100644 --- a/python/metatomic_torch/tests/o3.py +++ b/python/metatomic_torch/tests/o3.py @@ -21,7 +21,7 @@ # 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_tensormap_batched, ) From 7c41bab0c57b769e88cf785df5fe5c7464ba9962 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Wed, 5 Aug 2026 16:54:02 +0200 Subject: [PATCH 14/18] Address second-round review: drop wrapper-side name aliases, warn on skew stress --- .../metatomic_torch/metatomic/torch/model.py | 5 +- .../metatomic/torch/o3/_decompose.py | 14 ++++- .../metatomic/torch/o3/_symmetrized.py | 59 +++++-------------- .../tests/symmetrized_model.py | 42 +++---------- 4 files changed, 37 insertions(+), 83 deletions(-) diff --git a/python/metatomic_torch/metatomic/torch/model.py b/python/metatomic_torch/metatomic/torch/model.py index 3f50f37c..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 @@ -398,10 +399,10 @@ def __init__( # 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 = dict(NEW_QUANTITY_NAMES) + self._new_names = copy.deepcopy(NEW_QUANTITY_NAMES) # mapping from new names to the corresponding deprecated name - self._deprecated_names = dict(DEPRECATED_QUANTITY_NAMES) + 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/_decompose.py b/python/metatomic_torch/metatomic/torch/o3/_decompose.py index eb46074d..eb3094aa 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_decompose.py +++ b/python/metatomic_torch/metatomic/torch/o3/_decompose.py @@ -8,6 +8,7 @@ """ import math +import warnings from typing import List import torch @@ -43,9 +44,20 @@ def _symmetric_matrices_to_spherical( """Return orthonormal l=0 and l=2 components of the symmetric matrix part. Standard matrix quantities are symmetric, so their antisymmetric (l=1) part - carries no information and is discarded. + carries no information and is discarded. A model output that is materially + non-symmetric triggers a warning, since its antisymmetric part would be + silently excluded from the diagnostics. """ assert values.dim() == 4 and values.size(1) == 3 and values.size(2) == 3 + + antisymmetric_norm = (0.5 * (values - values.permute(0, 2, 1, 3))).norm() + if antisymmetric_norm > 1.0e-6 * values.norm(): + warnings.warn( + "a symmetric-matrix quantity has a materially antisymmetric part, " + "which is discarded by the O(3) diagnostics", + stacklevel=2, + ) + l0 = (values[:, 0, 0, :] + values[:, 1, 1, :] + values[:, 2, 2, :]).unsqueeze( 1 ) / math.sqrt(3.0) diff --git a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py index 3b05ac44..a4816267 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py +++ b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py @@ -45,16 +45,17 @@ ) -def _use_new_quantity_name(name: str, new_names: Dict[str, str]) -> str: +def _use_new_quantity_name(name: str) -> str: """Replace a deprecated base quantity in ``name`` with its current name. - Callers may request outputs under deprecated spellings (e.g. ``energies``); - translating them here lets the rest of the wrapper deal with the current - names only, whatever the wrapped model declares. + 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_names: - parts[0] = new_names[parts[0]] + if parts[0] in NEW_QUANTITY_NAMES: + parts[0] = NEW_QUANTITY_NAMES[parts[0]] return "/".join(parts) return name @@ -104,28 +105,8 @@ def _parse_output_request(requested_name: str) -> Tuple[str, str]: return source_name, calculation -def _record_output_request( - names: Dict[str, str], - source_name: str, - requested_name: str, -) -> None: - """Register the public name a source output must be returned under. - - Because deprecated names were normalized, two requested spellings can map to - the same source output and calculation; this rejects such duplicates, since - only one result per (source, calculation) pair can be returned. - """ - if source_name in names: - raise ValueError( - f"'{requested_name}' and '{names[source_name]}' request the same " - f"'{source_name}' output; only use the new name" - ) - names[source_name] = requested_name - - def _group_output_requests( outputs: Dict[str, ModelOutput], - new_names: Dict[str, str], ) -> Tuple[ Dict[str, str], Dict[str, str], @@ -134,9 +115,8 @@ def _group_output_requests( ]: """Group public requests by underlying output and calculation. - Deprecated quantity names are translated here, so the rest of the wrapper - only ever sees the current names. The returned dictionaries map each source - name to the exact spelling the caller requested it under. + 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] = {} @@ -145,7 +125,6 @@ def _group_output_requests( for requested_name, output in outputs.items(): source_name, calculation = _parse_output_request(requested_name) - source_name = _use_new_quantity_name(source_name, new_names) sample_kind = output.sample_kind if source_name in source_sample_kinds: previous_sample_kind = source_sample_kinds[source_name] @@ -158,15 +137,11 @@ def _group_output_requests( source_sample_kinds[source_name] = sample_kind if calculation == "average": - _record_output_request(average_names, source_name, requested_name) + average_names[source_name] = requested_name elif calculation == "variance": - _record_output_request(variance_names, source_name, requested_name) + variance_names[source_name] = requested_name else: - _record_output_request( - character_projection_names, - source_name, - requested_name, - ) + character_projection_names[source_name] = requested_name return ( source_sample_kinds, @@ -186,7 +161,7 @@ def _infer_max_angular_momentum( found_standard = False custom_names: List[str] = [] for name in names.keys(): - quantity = _use_new_quantity_name(name, NEW_QUANTITY_NAMES).split("/")[0] + 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 @@ -574,7 +549,6 @@ class SymmetrizedModel(torch.nn.Module): """ max_angular_momentum_character: Optional[int] - _new_names: Dict[str, str] _requested_inputs: Dict[str, ModelOutput] _requested_neighbor_lists: List[NeighborListOptions] @@ -591,11 +565,6 @@ def __init__( super().__init__() self._model = model - # ``AtomisticModel`` advertises both spellings of each standard output, so an - # engine can request either one from the wrapper; everything inside the wrapper - # uses the current names only. - # TorchScript cannot read a module-level dictionary from ``forward`` - self._new_names = dict(NEW_QUANTITY_NAMES) self._requested_inputs = {} self._requested_neighbor_lists = [] self.max_angular_momentum_target = validate_integer( @@ -870,7 +839,7 @@ def forward( average_names, variance_names, character_projection_names, - ) = _group_output_requests(outputs, self._new_names) + ) = _group_output_requests(outputs) if ( len(character_projection_names) != 0 and self.max_angular_momentum_character is None diff --git a/python/metatomic_torch/tests/symmetrized_model.py b/python/metatomic_torch/tests/symmetrized_model.py index ff7ff2bd..5ca80338 100644 --- a/python/metatomic_torch/tests/symmetrized_model.py +++ b/python/metatomic_torch/tests/symmetrized_model.py @@ -1093,37 +1093,6 @@ def test_preserves_variant_and_custom_output_names(self, source_name): atol=1.0e-12, ) - def test_deprecated_quantity_names_are_normalized(self): - """A deprecated request is decomposed as, and returned under, its own name.""" - model = SymmetrizedModel( - _EquivariantOutputModel(), - max_angular_momentum_target=1, - max_angular_momentum_grid=2, - batch_size=5, - ) - outputs = { - "non_conservative_forces": ModelOutput(sample_kind="atom"), - "o3::variance::non_conservative_forces": ModelOutput(sample_kind="atom"), - } - system = _forward_test_system([[1.0, 2.0, 3.0]]) - - result = model([system], outputs, None) - - assert set(result) == set(outputs) - assert torch.allclose( - result["non_conservative_forces"].block().values.squeeze(-1), - system.positions, - atol=1.0e-12, - ) - # the l=1 keys prove the decomposition recognized the singular quantity - variance = result["o3::variance::non_conservative_forces"] - assert variance.keys.values.tolist() == [[1, 1]] - assert torch.allclose( - variance.block().values, - torch.zeros_like(variance.block().values), - 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]]) @@ -2116,7 +2085,8 @@ def test_symmetric_matrices_to_spherical_known_components(): matrices[2, 0, 1, 0] = 2.0 matrices[2, 1, 0, 0] = -2.0 - l0, l2 = _symmetric_matrices_to_spherical(matrices) + with pytest.warns(UserWarning, match="materially antisymmetric"): + l0, l2 = _symmetric_matrices_to_spherical(matrices) expected_l0 = torch.zeros((3, 1, 1), dtype=torch.float64) expected_l0[0, 0, 0] = 3.0**0.5 @@ -2133,7 +2103,8 @@ def test_symmetric_matrices_to_spherical_known_components(): ) symmetric = 0.5 * (random_matrices + random_matrices.transpose(1, 2)) - l0, l2 = _symmetric_matrices_to_spherical(random_matrices) + with pytest.warns(UserWarning, match="materially antisymmetric"): + l0, l2 = _symmetric_matrices_to_spherical(random_matrices) spherical_norm_squared = l0.square().sum(dim=1) + l2.square().sum(dim=1) cartesian_norm_squared = symmetric.square().sum(dim=(1, 2)) @@ -2248,14 +2219,15 @@ def test_decompose_quantity_cartesian_vectors_preserve_autograd(source_name): def test_decompose_quantity_non_conservative_stress_combines_irreps(): - """Stress should return l=0 and l=2 blocks and silently discard skew.""" + """Stress should return l=0 and l=2 blocks, warning about discarded skew.""" 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) + with pytest.warns(UserWarning, match="materially antisymmetric"): + result = decompose_quantity("non_conservative_stress/direct", tensor) assert result.keys.names == ["o3_lambda", "o3_sigma"] assert result.keys.values.tolist() == [[0, 1], [2, 1]] From ce188e9f7bbf1a6abd195f575134cd4ea0950d7c Mon Sep 17 00:00:00 2001 From: ppegolo Date: Wed, 5 Aug 2026 17:02:32 +0200 Subject: [PATCH 15/18] Keep the antisymmetric pseudovector sector in stress diagnostics --- .../src/torch/reference/symmetrized-model.rst | 12 +-- .../metatomic/torch/o3/_decompose.py | 59 +++++++++------ .../tests/symmetrized_model.py | 73 +++++++++++-------- 3 files changed, 89 insertions(+), 55 deletions(-) diff --git a/docs/src/torch/reference/symmetrized-model.rst b/docs/src/torch/reference/symmetrized-model.rst index fb81a5de..3555d1b8 100644 --- a/docs/src/torch/reference/symmetrized-model.rst +++ b/docs/src/torch/reference/symmetrized-model.rst @@ -84,16 +84,18 @@ For diagnostics, the standard quantities are represented as follows: * - ``non_conservative_force`` - ``o3_lambda=1``, ``o3_sigma=1`` * - ``non_conservative_stress`` - - ``(o3_lambda, o3_sigma)=(0,1)`` and ``(2,1)`` + - ``(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`. Models should provide -symmetric ``non_conservative_stress`` tensors. Stress diagnostics retain only -the scalar trace and symmetric-traceless sectors, silently discarding any -antisymmetric part. +: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 diff --git a/python/metatomic_torch/metatomic/torch/o3/_decompose.py b/python/metatomic_torch/metatomic/torch/o3/_decompose.py index eb3094aa..a3e3827d 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_decompose.py +++ b/python/metatomic_torch/metatomic/torch/o3/_decompose.py @@ -8,8 +8,7 @@ """ import math -import warnings -from typing import List +from typing import List, Tuple import torch from metatensor.torch import Labels, TensorBlock, TensorMap @@ -38,31 +37,37 @@ def _cartesian_vectors_to_spherical( return values.roll(-1, dims=component_axis) -def _symmetric_matrices_to_spherical( +def _matrices_to_spherical( values: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - """Return orthonormal l=0 and l=2 components of the symmetric matrix part. +) -> 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) part - carries no information and is discarded. A model output that is materially - non-symmetric triggers a warning, since its antisymmetric part would be - silently excluded from the diagnostics. + 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 - antisymmetric_norm = (0.5 * (values - values.permute(0, 2, 1, 3))).norm() - if antisymmetric_norm > 1.0e-6 * values.norm(): - warnings.warn( - "a symmetric-matrix quantity has a materially antisymmetric part, " - "which is discarded by the O(3) diagnostics", - stacklevel=2, - ) - 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, @@ -75,7 +80,7 @@ def _symmetric_matrices_to_spherical( dim=1, ) - return l0, l2 + return l0, l1, l2 def decompose_quantity( @@ -142,6 +147,7 @@ def decompose_quantity( else: assert category == "symmetric_matrix" blocks_l0: List[TensorBlock] = [] + blocks_l1: List[TensorBlock] = [] blocks_l2: List[TensorBlock] = [] for block in tensor.blocks(): assert ( @@ -152,7 +158,7 @@ def decompose_quantity( and len(block.components[1]) == 3 ), f"'{quantity}' must have 'xyz_1' and 'xyz_2' component axes of size 3" - values_l0, values_l2 = _symmetric_matrices_to_spherical(block.values) + values_l0, values_l1, values_l2 = _matrices_to_spherical(block.values) blocks_l0.append( TensorBlock( values=values_l0, @@ -161,6 +167,14 @@ def decompose_quantity( 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, @@ -171,13 +185,16 @@ def decompose_quantity( ) 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_l2.values], dim=0), + torch.cat([keys_l0.values, keys_l1.values, keys_l2.values], dim=0), ), - blocks_l0 + blocks_l2, + blocks_l0 + blocks_l1 + blocks_l2, ) for info_name, info_value in tensor.info().items(): diff --git a/python/metatomic_torch/tests/symmetrized_model.py b/python/metatomic_torch/tests/symmetrized_model.py index 5ca80338..98605114 100644 --- a/python/metatomic_torch/tests/symmetrized_model.py +++ b/python/metatomic_torch/tests/symmetrized_model.py @@ -21,8 +21,8 @@ from metatomic.torch.o3 import O3Transformation from metatomic.torch.o3._decompose import ( _cartesian_vectors_to_spherical, + _matrices_to_spherical, _o3_mu_labels, - _symmetric_matrices_to_spherical, decompose_quantity, ) from metatomic.torch.o3._quadrature import ( @@ -582,7 +582,7 @@ def forward( for system in systems ] ).unsqueeze(-1) - _, spherical = _symmetric_matrices_to_spherical(matrices) + _, _, spherical = _matrices_to_spherical(matrices) result["mtt::spherical_quadrupole"] = TensorMap( Labels( ["o3_lambda", "o3_sigma"], @@ -921,19 +921,22 @@ def test_stress_character_projection_combines_target_and_character_sectors(self) assert { tuple(int(value) for value in key.values) for key in projection.keys } == { - (o3_lambda, 1, chi_lambda, chi_sigma) - for o3_lambda in (0, 2) + (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)] - if chi_lambda == o3_lambda and chi_sigma == 1: + # 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( @@ -1284,7 +1287,7 @@ def test_equivariant_outputs_preserve_values_metadata_and_zero_variance(self): ) quadrupole = result["mtt::spherical_quadrupole"] assert quadrupole.keys.values.tolist() == [[2, 1]] - _, expected_quadrupole = _symmetric_matrices_to_spherical( + _, _, expected_quadrupole = _matrices_to_spherical( torch.outer(system.positions[0], system.positions[0]).reshape(1, 3, 3, 1) ) assert torch.allclose( @@ -1296,7 +1299,7 @@ def test_equivariant_outputs_preserve_values_metadata_and_zero_variance(self): expected_target_keys = { "o3::variance::energy": [[0, 1]], "o3::variance::non_conservative_force": [[1, 1]], - "o3::variance::non_conservative_stress": [[0, 1], [2, 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]], } @@ -2076,8 +2079,8 @@ def test_cartesian_vectors_to_spherical_commutes_with_o3(inversion): ) -def test_symmetric_matrices_to_spherical_known_components(): - """Known matrices map as expected and the symmetric norm is preserved.""" +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 @@ -2085,14 +2088,17 @@ def test_symmetric_matrices_to_spherical_known_components(): matrices[2, 0, 1, 0] = 2.0 matrices[2, 1, 0, 0] = -2.0 - with pytest.warns(UserWarning, match="materially antisymmetric"): - l0, l2 = _symmetric_matrices_to_spherical(matrices) + 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) @@ -2101,13 +2107,13 @@ def test_symmetric_matrices_to_spherical_known_components(): dtype=torch.float64, generator=generator, ) - symmetric = 0.5 * (random_matrices + random_matrices.transpose(1, 2)) - with pytest.warns(UserWarning, match="materially antisymmetric"): - l0, l2 = _symmetric_matrices_to_spherical(random_matrices) + l0, l1, l2 = _matrices_to_spherical(random_matrices) - spherical_norm_squared = l0.square().sum(dim=1) + l2.square().sum(dim=1) - cartesian_norm_squared = symmetric.square().sum(dim=(1, 2)) + 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, @@ -2117,7 +2123,7 @@ def test_symmetric_matrices_to_spherical_known_components(): @pytest.mark.parametrize("inversion", [1.0, -1.0]) -def test_symmetric_matrices_to_spherical_commutes_with_o3(inversion): +def test_matrices_to_spherical_commutes_with_o3(inversion): """Cartesian and spherical transformations should give the same components.""" proper_rotation = torch.tensor( [ @@ -2131,10 +2137,11 @@ def test_symmetric_matrices_to_spherical_commutes_with_o3(inversion): 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.7, 1.1, 0.8], [2.3, 0.8, -0.4]], - [[-0.2, 1.4, 0.5], [1.4, 0.9, -1.1], [0.5, -1.1, 2.0]], + [[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) @@ -2146,18 +2153,22 @@ def test_symmetric_matrices_to_spherical_commutes_with_o3(inversion): matrices, matrix, ) - transformed_l0, transformed_l2 = _symmetric_matrices_to_spherical( + transformed_l0, transformed_l1, transformed_l2 = _matrices_to_spherical( transformed_matrices ) - l0, l2 = _symmetric_matrices_to_spherical(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) @@ -2219,21 +2230,22 @@ def test_decompose_quantity_cartesian_vectors_preserve_autograd(source_name): def test_decompose_quantity_non_conservative_stress_combines_irreps(): - """Stress should return l=0 and l=2 blocks, warning about discarded skew.""" + """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"]) - with pytest.warns(UserWarning, match="materially antisymmetric"): - result = decompose_quantity("non_conservative_stress/direct", tensor) + result = decompose_quantity("non_conservative_stress/direct", tensor) assert result.keys.names == ["o3_lambda", "o3_sigma"] - assert result.keys.values.tolist() == [[0, 1], [2, 1]] + 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, @@ -2241,11 +2253,14 @@ def test_decompose_quantity_non_conservative_stress_combines_irreps(): 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)) - assert block_l0.samples == tensor.block().samples - assert block_l2.samples == tensor.block().samples - assert block_l0.properties == tensor.block().properties - assert block_l2.properties == tensor.block().properties + 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(): From 754246503b91493ad9d21ca4ccf1c16f4d6c89d4 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Thu, 6 Aug 2026 10:19:31 +0200 Subject: [PATCH 16/18] Further simplify code and tests --- .../src/torch/reference/symmetrized-model.rst | 10 + metatomic-torch/CHANGELOG.md | 7 +- .../metatomic/torch/_quantities.py | 43 +- .../metatomic_torch/metatomic/torch/model.py | 4 +- .../metatomic/torch/o3/_decompose.py | 28 +- .../metatomic/torch/o3/_projections.py | 20 +- .../metatomic/torch/o3/_symmetrized.py | 129 ++--- .../metatomic/torch/o3/_transformations.py | 47 +- .../metatomic/torch/o3/_utils.py | 43 +- .../metatomic/torch/o3/_wigner.py | 21 +- python/metatomic_torch/tests/o3.py | 148 ++---- .../tests/symmetrized_model.py | 464 ++++++------------ 12 files changed, 340 insertions(+), 624 deletions(-) diff --git a/docs/src/torch/reference/symmetrized-model.rst b/docs/src/torch/reference/symmetrized-model.rst index 3555d1b8..4cd1056a 100644 --- a/docs/src/torch/reference/symmetrized-model.rst +++ b/docs/src/torch/reference/symmetrized-model.rst @@ -33,6 +33,16 @@ 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. +The meaning of the variance depends on the metadata of the underlying output. +When its blocks carry recognized component labels (``o3_mu``-style spherical +or ``xyz``-style Cartesian axes), each response is rotated back to the input +frame first, and the variance measures the breaking of *equivariance*. Outputs +without such labels cannot be rotated back: their responses are compared as-is +across the quadrature, so their variance measures the deviation from +*invariance* only. An equivariant but unlabelled output — for example an +internal feature vector — reports a large variance even when it transforms +correctly. + Average and variance -------------------- diff --git a/metatomic-torch/CHANGELOG.md b/metatomic-torch/CHANGELOG.md index 3163eaac..2392cd91 100644 --- a/metatomic-torch/CHANGELOG.md +++ b/metatomic-torch/CHANGELOG.md @@ -24,10 +24,9 @@ a changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows ### 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. +- `O3Transformation` now holds a batch of one or more operations, is usable + inside scripted models, and gained `inverse`, `with_inversion`, + `transform_systems`, and `transform_tensormap`. - Renamed `O3Transformation.is_inverted` to `is_improper`. - `wigners >= 0.4.0` is now required. diff --git a/python/metatomic_torch/metatomic/torch/_quantities.py b/python/metatomic_torch/metatomic/torch/_quantities.py index d0800858..a6fa13a1 100644 --- a/python/metatomic_torch/metatomic/torch/_quantities.py +++ b/python/metatomic_torch/metatomic/torch/_quantities.py @@ -1,12 +1,4 @@ -""" -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. -""" +"""Cartesian layout and spherical character for standard quantities.""" from typing import Dict @@ -17,8 +9,8 @@ def standard_quantity_categories() -> Dict[str, str]: 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. + (singular) spellings appear here: deprecated names are not recognized, and + the code using this table treats them as custom quantities. TorchScript cannot read a module-level dictionary from a compiled function, so the table is built by this function and bound to @@ -53,25 +45,16 @@ def standard_quantity_categories() -> Dict[str, str]: } -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 deprecated quantity names to their current name +NEW_QUANTITY_NAMES: Dict[str, str] = { + "features": "feature", + "non_conservative_forces": "non_conservative_force", + "positions": "position", + "momenta": "momentum", + "masses": "mass", + "velocities": "velocity", + "charges": "charge", +} #: mapping from current quantity names to the corresponding deprecated name DEPRECATED_QUANTITY_NAMES: Dict[str, str] = { diff --git a/python/metatomic_torch/metatomic/torch/model.py b/python/metatomic_torch/metatomic/torch/model.py index 15fc498e..73b7282b 100644 --- a/python/metatomic_torch/metatomic/torch/model.py +++ b/python/metatomic_torch/metatomic/torch/model.py @@ -396,9 +396,7 @@ def __init__( else: raise ValueError(f"unknown dtype in capabilities: {capabilities.dtype}") - # mapping from deprecated output/input names to their new name, copied - # onto the instance because TorchScript methods cannot read module-level - # dictionaries + # TorchScript methods cannot read module-level dicts: copy onto the instance self._new_names = copy.deepcopy(NEW_QUANTITY_NAMES) # mapping from new names to the corresponding deprecated name diff --git a/python/metatomic_torch/metatomic/torch/o3/_decompose.py b/python/metatomic_torch/metatomic/torch/o3/_decompose.py index a3e3827d..47112a4a 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_decompose.py +++ b/python/metatomic_torch/metatomic/torch/o3/_decompose.py @@ -14,9 +14,10 @@ from metatensor.torch import Labels, TensorBlock, TensorMap from .._quantities import standard_quantity_categories +from ._utils import copy_tensormap_info, strip_placeholder_key -def _o3_mu_labels(o3_lambda: int, device: torch.device) -> Labels: +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", @@ -114,7 +115,7 @@ def decompose_quantity( TensorBlock( values=block.values.unsqueeze(1), samples=block.samples, - components=[_o3_mu_labels(0, block.values.device)], + components=[o3_mu_labels(0, block.values.device)], properties=block.properties, ) ) @@ -135,7 +136,7 @@ def decompose_quantity( TensorBlock( values=_cartesian_vectors_to_spherical(block.values, 1), samples=block.samples, - components=[_o3_mu_labels(1, block.values.device)], + components=[o3_mu_labels(1, block.values.device)], properties=block.properties, ) ) @@ -163,7 +164,7 @@ def decompose_quantity( TensorBlock( values=values_l0, samples=block.samples, - components=[_o3_mu_labels(0, block.values.device)], + components=[o3_mu_labels(0, block.values.device)], properties=block.properties, ) ) @@ -171,7 +172,7 @@ def decompose_quantity( TensorBlock( values=values_l1, samples=block.samples, - components=[_o3_mu_labels(1, block.values.device)], + components=[o3_mu_labels(1, block.values.device)], properties=block.properties, ) ) @@ -179,7 +180,7 @@ def decompose_quantity( TensorBlock( values=values_l2, samples=block.samples, - components=[_o3_mu_labels(2, block.values.device)], + components=[o3_mu_labels(2, block.values.device)], properties=block.properties, ) ) @@ -197,9 +198,7 @@ def decompose_quantity( blocks_l0 + blocks_l1 + blocks_l2, ) - for info_name, info_value in tensor.info().items(): - result.set_info(info_name, info_value) - return result + return copy_tensormap_info(tensor, result) def _add_o3_irrep_to_keys( @@ -208,16 +207,7 @@ def _add_o3_irrep_to_keys( 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] + names, values = strip_placeholder_key(list(keys.names), keys.values) for name, expected in ( ("o3_lambda", o3_lambda), diff --git a/python/metatomic_torch/metatomic/torch/o3/_projections.py b/python/metatomic_torch/metatomic/torch/o3/_projections.py index e4cd8e51..c1520a11 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_projections.py +++ b/python/metatomic_torch/metatomic/torch/o3/_projections.py @@ -14,6 +14,7 @@ from ._utils import ( group_samples_by_rotated_copy, restore_input_system_to_samples, + strip_placeholder_key, ) @@ -62,15 +63,10 @@ def character_projection_coefficients_from_batch( ) -> 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] + key_names, key_values = strip_placeholder_key( + list(tensor.keys.names), + tensor.keys.values, + ) if "chi_lambda" in key_names or "chi_sigma" in key_names: raise ValueError( @@ -145,10 +141,8 @@ def character_projection_tensormap_from_cosets( ) -> 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'") + # only ever consumes the sibling accumulator's output within one forward + assert "chi_lambda" in key_names and "chi_sigma" not in key_names chi_lambda_column = key_names.index("chi_lambda") blocks: List[TensorBlock] = [] diff --git a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py index a4816267..4c9bf4e5 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py +++ b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py @@ -23,7 +23,6 @@ from .._quantities import ( MAX_ANGULAR_MOMENTUM_PER_CATEGORY, - NEW_QUANTITY_NAMES, STANDARD_QUANTITY_CATEGORIES, ) from ._decompose import decompose_quantity @@ -32,8 +31,9 @@ character_projection_tensormap_from_cosets, ) from ._quadrature import choose_quadrature, get_rotation_quadrature -from ._transformations import O3Transformation, _max_o3_lambda_in_tensor +from ._transformations import O3Transformation, max_o3_lambda_in_tensor from ._utils import ( + copy_tensormap_info, group_samples_by_rotated_copy, map_selected_atoms_to_rotated_copies, restore_input_system_to_samples, @@ -45,21 +45,6 @@ ) -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, @@ -67,7 +52,7 @@ def _check_o3_lambda_limit( limit_name: str, ) -> None: """Check a TensorMap's spherical component ranks against one limit.""" - tensor_max_o3_lambda = _max_o3_lambda_in_tensor(tensor) + 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}, " @@ -161,7 +146,7 @@ def _infer_max_angular_momentum( found_standard = False custom_names: List[str] = [] for name in names.keys(): - quantity = _use_new_quantity_name(name).split("/")[0] + quantity = name.split("/", 1)[0] if quantity == "feature": # features are not an irreducible representation of O(3): they are # passed through unchanged and never rotated back @@ -202,12 +187,9 @@ def _reduce_weighted_centered_batch( ]: """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. + Centering on the first rotated copy keeps both terms of ``E[X^2] - E[X]^2`` + of the order of the variation itself, so their subtraction does not lose + significant digits to cancellation when the mean response is large. """ n_rotated_copies = weights.numel() centered_first_moment_blocks: List[TensorBlock] = [] @@ -342,13 +324,6 @@ def _add_tensormap_contribution( 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] = [] @@ -389,10 +364,10 @@ def _clamp_roundoff_negative_diagnostic( 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()): + if bool(torch.any(~torch.isfinite(scale_values)).item()): raise ValueError( - f"round-off scale of the O(3) {quantity} is negative or not " - f"finite for block ({key.print()})" + f"round-off scale of the O(3) {quantity} is not finite for " + f"block ({key.print()})" ) # TorchScript does not support torch.finfo; use the IEEE-754 values for @@ -469,20 +444,18 @@ def _mean_variance_over_components( 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") + # both maps were built from the same moments earlier in this forward + assert variance.keys == component_layout.keys blocks: List[TensorBlock] = [] for key, block in variance.items(): - if len(block.components) != 0: - raise ValueError("component-summed variance must not have components") + assert len(block.components) == 0 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") + assert ( + layout_block.samples == block.samples + and layout_block.properties == block.properties + ) n_components = 1 for component in layout_block.components: @@ -511,12 +484,14 @@ class SymmetrizedModel(torch.nn.Module): ``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 + squared character-projection contributions. Outputs whose blocks carry no + recognized component labels are not rotated back, so their variance + measures the deviation from invariance only. 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, + 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 @@ -609,11 +584,6 @@ def __init__( 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, @@ -692,16 +662,6 @@ def wrap( 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( @@ -794,7 +754,7 @@ def wrap( atomic_types=capabilities.atomic_types, interaction_range=capabilities.interaction_range, length_unit=capabilities.length_unit, - supported_devices=supported_devices, + supported_devices=capabilities.supported_devices, dtype=capabilities.dtype, ) return AtomisticModel( @@ -860,6 +820,19 @@ def forward( empty_results: List[TensorMap] = [] per_output_results[requested_name] = empty_results + integration_dtype = self._rotation_matrices.dtype + if integration_dtype != torch.float64: + if integration_dtype != torch.float32: + raise TypeError( + "SymmetrizedModel integration buffers must use float32 or " + f"float64, got {dtype_name(integration_dtype)}" + ) + warnings.warn( + "SymmetrizedModel integration buffers were downcast from " + "float64; averages and diagnostics will be less accurate", + stacklevel=2, + ) + for input_system_index, system in enumerate(systems): system_results = self._evaluate_system( system, @@ -897,25 +870,12 @@ def _evaluate_system( """Stream all quadrature batches for one input System.""" work_dtype = system.positions.dtype work_device = system.positions.device + integration_dtype = self._rotation_matrices.dtype 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 @@ -1061,10 +1021,8 @@ def _evaluate_system( self.max_angular_momentum_target, "max_angular_momentum_target", ) - if backrotation_rotations is None: - raise RuntimeError( - "backrotation transformations were not prepared" - ) + # prepared above whenever any output needs back-rotation + assert backrotation_rotations is not None if is_improper: backrotation = ( backrotation_rotations.with_inversion().inverse() @@ -1091,7 +1049,7 @@ def _evaluate_system( compute_second_moments=False, ) if not has_average_reference: - updated_average_reference = _copy_tensormap_info( + updated_average_reference = copy_tensormap_info( backrotated, updated_average_reference, ) @@ -1122,8 +1080,9 @@ def _evaluate_system( variance_reference, compute_second_moments=True, ) - if second_moment is None or absolute_second_moment is None: - raise RuntimeError("variance moments were not computed") + # always computed with compute_second_moments=True + assert second_moment is not None + assert absolute_second_moment is not None variance_references[source_name] = variance_reference _add_tensormap_contribution( variance_first_moments, @@ -1168,7 +1127,7 @@ def _evaluate_system( average_references[source_name], average_first_moments[source_name], ) - mean = _copy_tensormap_info(average_references[source_name], mean) + mean = copy_tensormap_info(average_references[source_name], mean) results[requested_name] = mean.to( dtype=work_dtype, device=work_device, diff --git a/python/metatomic_torch/metatomic/torch/o3/_transformations.py b/python/metatomic_torch/metatomic/torch/o3/_transformations.py index ee7955cf..84fdedac 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_transformations.py +++ b/python/metatomic_torch/metatomic/torch/o3/_transformations.py @@ -15,6 +15,7 @@ from metatensor.torch import Labels, LabelsEntry, TensorBlock, TensorMap from .. import System, register_autograd_neighbors +from ._utils import copy_tensormap_info, validate_integer from ._wigner import build_packed_wigner_matrices, wigner_matrices_for_lambda @@ -30,22 +31,6 @@ ) -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, @@ -122,8 +107,8 @@ def __init__( matrices = matrix improper = _improper else: - max_angular_momentum = _validate_nonnegative_integer( - "max_angular_momentum", max_angular_momentum + max_angular_momentum = validate_integer( + "max_angular_momentum", max_angular_momentum, 0 ) if matrix.dim() == 2: @@ -226,7 +211,7 @@ def device(self) -> torch.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) + ell = validate_integer("ell", ell, 0) if ell > self._max_angular_momentum: raise ValueError( @@ -435,7 +420,7 @@ def transform_systems(self, system: System) -> list[System]: 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: + 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] = [] @@ -478,7 +463,7 @@ def transform_tensormap( information """ wigner_matrices: list[torch.Tensor] = [] - if _max_o3_lambda_in_tensor(tensor) >= 0: + if max_o3_lambda_in_tensor(tensor) >= 0: wigner_matrices = self._wigner_D_matrices() return _transform_tensormap_batched( tensor, @@ -516,9 +501,9 @@ def random_transformations( ``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 + n = validate_integer("n", n, 0) + max_angular_momentum = validate_integer( + "max_angular_momentum", max_angular_momentum, 0 ) if dtype not in (torch.float32, torch.float64): @@ -780,10 +765,7 @@ def _validate_component_axis_metadata( ) metadata.append((False, 0, 1)) elif is_spherical: - ell = _validate_nonnegative_integer( - "ell", - int(key["o3_lambda" + suffix]), - ) + ell = validate_integer("ell", int(key["o3_lambda" + suffix]), 0) sigma = int(key["o3_sigma" + suffix]) _spherical_parity_factor(ell, sigma, is_improper=False) @@ -832,7 +814,7 @@ def _max_o3_lambda_in_block(key: LabelsEntry, block: TensorBlock) -> int: return max_o3_lambda -def _max_o3_lambda_in_tensor(tensor: TensorMap) -> int: +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``. @@ -959,7 +941,7 @@ def transform_tensor( combined = _combine_transformations( transformations, - _max_o3_lambda_in_tensor(tensor), + max_o3_lambda_in_tensor(tensor), ) return combined.transform_tensormap(tensor, validated_ids) @@ -1175,7 +1157,4 @@ def _transform_tensormap_batched( ) ) - transformed = TensorMap(tensor.keys, blocks) - for info_name, info_value in tensor.info().items(): - transformed.set_info(info_name, info_value) - return transformed + return copy_tensormap_info(tensor, TensorMap(tensor.keys, blocks)) diff --git a/python/metatomic_torch/metatomic/torch/o3/_utils.py b/python/metatomic_torch/metatomic/torch/o3/_utils.py index 646ba840..939ff5d3 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_utils.py +++ b/python/metatomic_torch/metatomic/torch/o3/_utils.py @@ -8,17 +8,27 @@ from typing import List, Optional, Tuple import torch -from metatensor.torch import Labels, TensorBlock +from metatensor.torch import Labels, TensorBlock, TensorMap -def validate_integer(name: str, value, minimum: int) -> int: +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 validate_integer(name: str, value: int, 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 torch.jit.is_scripting(): + integer_value = value + else: + 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" @@ -30,6 +40,20 @@ def validate_integer(name: str, value, minimum: int) -> int: return integer_value +def strip_placeholder_key( + names: List[str], + values: torch.Tensor, +) -> Tuple[List[str], torch.Tensor]: + """Drop a ``"_"`` placeholder key dimension after checking its single zero entry.""" + if names == ["_"]: + if values.size(0) != 1 or int(values[0, 0]) != 0: + raise ValueError( + "the '_' placeholder must contain exactly one key with value 0" + ) + return [], values[:, :0] + return names, values + + def map_selected_atoms_to_rotated_copies( selected_atoms: Optional[Labels], input_system_index: int, @@ -72,15 +96,6 @@ def group_samples_by_rotated_copy( ], 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 " diff --git a/python/metatomic_torch/metatomic/torch/o3/_wigner.py b/python/metatomic_torch/metatomic/torch/o3/_wigner.py index 7052adcd..a7e19d12 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_wigner.py +++ b/python/metatomic_torch/metatomic/torch/o3/_wigner.py @@ -123,6 +123,11 @@ def build_wigner_D_cache( return {ell: tensor.to(device=device, dtype=dtype) for ell, tensor in cache.items()} +def _packed_elements_before(o3_lambda: int) -> int: + """Number of packed elements per matrix below ``o3_lambda`` (sum of (2l+1)^2).""" + return o3_lambda * (4 * o3_lambda * o3_lambda - 1) // 3 + + def build_packed_wigner_matrices( matrices: torch.Tensor, max_angular_momentum: int, @@ -140,12 +145,7 @@ def build_packed_wigner_matrices( 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 - ) + n_elements_per_matrix = _packed_elements_before(max_angular_momentum + 1) packed = torch.empty( n_matrices * n_elements_per_matrix, dtype=output_dtype, @@ -161,7 +161,7 @@ def build_packed_wigner_matrices( ) 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 + elements_before = _packed_elements_before(o3_lambda) offset = n_matrices * elements_before + matrix_index * dimension * dimension packed[offset : offset + dimension * dimension].copy_( cache[o3_lambda].reshape(-1) @@ -182,11 +182,10 @@ def wigner_matrices_for_lambda( # 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 + offset = n_matrices * _packed_elements_before(o3_lambda) length = n_matrices * dimension * dimension - if offset + length > packed.numel(): - raise ValueError("o3_lambda exceeds the packed Wigner-D storage") + # callers validate o3_lambda against their own maximum angular momentum + assert offset + length <= packed.numel() return packed[offset : offset + length].view( n_matrices, diff --git a/python/metatomic_torch/tests/o3.py b/python/metatomic_torch/tests/o3.py index e66a1bd5..2563b1b8 100644 --- a/python/metatomic_torch/tests/o3.py +++ b/python/metatomic_torch/tests/o3.py @@ -1,3 +1,4 @@ +import io import re import metatensor.torch as mts @@ -19,13 +20,6 @@ transform_tensor, ) -# 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._transformations import ( - _max_o3_lambda_in_tensor, - _transform_tensormap_batched, -) - # The complex-to-real spherical harmonics conversion is defined only here for now. from metatomic.torch.o3._wigner import _complex_to_real_spherical_harmonics_transform @@ -109,28 +103,9 @@ def _single_block_tensor_map( ) -def _stack_o3_matrices(transformations, max_angular_momentum): - """Stack Cartesian, Wigner, and parity tensors from O3 transformations.""" - matrices = torch.cat( - [transformation.matrices for transformation in transformations] - ) - wigner_matrices = [ - torch.cat( - [ - transformation.wigner_D_matrices(ell) - for transformation in transformations - ] - ) - for ell in range(max_angular_momentum + 1) - ] - improper = torch.cat( - [transformation.improper for transformation in transformations] - ) - return matrices, wigner_matrices, improper - - -def test_max_o3_lambda_in_tensor(): - """Check `_max_o3_lambda_in_tensor`, including in torchscript mode""" +def test_gradient_components_require_wigner_ranks(): + """Angular momenta appearing only in gradient components still require the + matching Wigner-D matrices, while purely Cartesian data needs none.""" properties = Labels("property", torch.tensor([[0]])) block = TensorBlock( values=torch.ones((1, 3, 1), dtype=torch.float64), @@ -160,15 +135,33 @@ def test_max_o3_lambda_in_tensor(): ), [block], ) + + # the values only reach ell=1, but the gradient carries an ell=3 axis + rotation = O3Transformation( + _rotation_90_degrees_around_z(), + max_angular_momentum=1, + ) + message = re.escape("ell=3 exceeds max_angular_momentum=1.") + with pytest.raises(ValueError, match=f"^{message}$"): + rotation.transform_tensormap(spherical) + + # Cartesian data transforms without any Wigner-D matrices cartesian = _single_block_tensor_map( values=torch.ones((1, 3, 1), dtype=torch.float64), samples=Labels("system", torch.tensor([[0]])), components=[Labels("xyz", torch.arange(3).reshape(-1, 1))], ) - - scripted_maximum = torch.jit.script(_max_o3_lambda_in_tensor) - assert scripted_maximum(spherical) == 3 - assert scripted_maximum(cartesian) == -1 + no_wigner = O3Transformation( + _rotation_90_degrees_around_z(), + max_angular_momentum=0, + ) + transformed = no_wigner.transform_tensormap(cartesian) + assert torch.allclose( + transformed.block().values, + torch.tensor([[[-1.0], [1.0], [1.0]]], dtype=torch.float64), + rtol=0.0, + atol=1e-12, + ) @pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE) @@ -332,14 +325,14 @@ def test_transform_system_preserves_neighbor_gradients(): def test_transformation_validation(): """Realistic construction mistakes fail with a clear error.""" # negative counts and angular momentum limits - message = re.escape("max_angular_momentum must be a non-negative integer, got -1.") + message = re.escape("max_angular_momentum must be non-negative, got -1") with pytest.raises(ValueError, match=f"^{message}$"): O3Transformation(torch.eye(3, dtype=torch.float64), -1) with pytest.raises(ValueError, match=f"^{message}$"): random_transformations( 0, max_angular_momentum=-1, device=torch.device("cpu"), dtype=torch.float64 ) - message = re.escape("n must be a non-negative integer, got -1.") + message = re.escape("n must be non-negative, got -1") with pytest.raises(ValueError, match=f"^{message}$"): random_transformations(-1, device=torch.device("cpu"), dtype=torch.float64) @@ -364,7 +357,7 @@ def test_transformation_validation(): # Wigner-D requests need a valid ell transformation = O3Transformation(torch.eye(3, dtype=torch.float64), 1) - message = re.escape("ell must be a non-negative integer, got -1.") + message = re.escape("ell must be non-negative, got -1") with pytest.raises(ValueError, match=f"^{message}$"): transformation.wigner_D_matrix(-1) @@ -1220,8 +1213,8 @@ def test_batched_tensor_transform_matches_transform_tensor( dtype, parities, ): - """The scripted batched kernel matches ``transform_tensor``, including for - batches mixing proper and improper operations.""" + """A batched ``O3Transformation`` 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 proper_matrices = [ @@ -1236,8 +1229,8 @@ def test_batched_tensor_transform_matches_transform_tensor( O3Transformation(sign * matrix, max_angular_momentum=2) for sign, matrix in zip(parities, proper_matrices, strict=True) ] - matrices, wigner_matrices, improper = _stack_o3_matrices( - transformations, + batch = O3Transformation( + torch.cat([transformation.matrices for transformation in transformations]), max_angular_momentum=2, ) @@ -1305,14 +1298,7 @@ def test_batched_tensor_transform_matches_transform_tensor( ], transformations, ) - scripted_transform = torch.jit.script(_transform_tensormap_batched) - result = scripted_transform( - tensor, - matrices, - wigner_matrices, - improper, - None, - ) + result = batch.transform_tensormap(tensor) mts.allclose_raise(result, expected, rtol=0.0, atol=atol) assert result.info() == expected.info() @@ -1336,16 +1322,12 @@ def test_batched_tensor_transform_matches_transform_tensor( ) -def test_batched_tensor_transform_single_transformation_is_scriptable(): - """The scripted singleton path should not require a ``system`` sample label.""" +def test_single_transformation_needs_no_system_label(): + """A single transformation applies to samples without a ``system`` label.""" transformation = O3Transformation( _rotation_90_degrees_around_z(), max_angular_momentum=1, ) - matrices, wigner_matrices, improper = _stack_o3_matrices( - [transformation], - max_angular_momentum=1, - ) tensor = _single_block_tensor_map( keys=Labels( ["o3_lambda", "o3_sigma"], @@ -1361,14 +1343,7 @@ def test_batched_tensor_transform_single_transformation_is_scriptable(): ], ) - scripted_transform = torch.jit.script(_transform_tensormap_batched) - result = scripted_transform( - tensor, - matrices, - wigner_matrices, - improper, - None, - ) + result = transformation.transform_tensormap(tensor) expected = transform_tensor( tensor, [_make_system([1])], @@ -1385,12 +1360,10 @@ def test_batched_tensor_transform_single_transformation_is_scriptable(): 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, improper = _stack_o3_matrices( - transformations, + batch = O3Transformation( + torch.stack( + [torch.eye(3, dtype=torch.float64), _rotation_90_degrees_around_z()] + ), max_angular_momentum=1, ) @@ -1401,13 +1374,7 @@ def test_batched_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_tensormap_batched( - missing_system, - matrices, - wigner_matrices, - improper, - None, - ) + batch.transform_tensormap(missing_system) for system_index in (-1, 2): out_of_range = _single_block_tensor_map( @@ -1417,13 +1384,7 @@ def test_batched_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_tensormap_batched( - out_of_range, - matrices, - wigner_matrices, - improper, - None, - ) + batch.transform_tensormap(out_of_range) unavailable_rank = _single_block_tensor_map( keys=Labels( @@ -1436,15 +1397,13 @@ def test_batched_tensor_transform_rejects_invalid_routing_and_wigner_rank(): Labels("o3_mu", torch.arange(-1, 2).reshape(-1, 1)), ], ) + too_low = O3Transformation( + torch.eye(3, dtype=torch.float64), + max_angular_momentum=0, + ) message = re.escape("ell=1 exceeds max_angular_momentum=0.") with pytest.raises(ValueError, match=f"^{message}$"): - _transform_tensormap_batched( - unavailable_rank, - matrices, - wigner_matrices[:1], - improper, - None, - ) + too_low.transform_tensormap(unavailable_rank) def test_batched_transformation_matches_singles(): @@ -1588,10 +1547,11 @@ def forward(self, tensor: TensorMap) -> TensorMap: dtype=torch.float64, generator=torch.Generator().manual_seed(3), ) - matrices, wigner_matrices, _improper = _stack_o3_matrices( - singles, - max_angular_momentum=2, - ) + matrices = torch.cat([single.matrices for single in singles]) + wigner_matrices = [ + torch.cat([single.wigner_D_matrices(ell) for single in singles]) + for ell in range(3) + ] module = BackRotate(matrices, wigner_matrices) scripted = torch.jit.script(module) @@ -1615,8 +1575,6 @@ def forward(self, tensor: TensorMap) -> TensorMap: atol=1e-12, ) - import io - buffer = io.BytesIO() torch.jit.save(scripted, buffer) buffer.seek(0) diff --git a/python/metatomic_torch/tests/symmetrized_model.py b/python/metatomic_torch/tests/symmetrized_model.py index 98605114..23915efe 100644 --- a/python/metatomic_torch/tests/symmetrized_model.py +++ b/python/metatomic_torch/tests/symmetrized_model.py @@ -19,18 +19,14 @@ 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, -) + +# These helpers back the exported SymmetrizedModel and have no public entry +# point: decompose_quantity is its Cartesian-to-spherical boundary (tested here +# directly for its analytic values, and against the public O3Transformation for +# equivariance), the quadrature functions build its integration buffers, and +# map_selected_atoms_to_rotated_copies routes its selected_atoms argument. +from metatomic.torch.o3._decompose import decompose_quantity, o3_mu_labels +from metatomic.torch.o3._quadrature import choose_quadrature, get_rotation_quadrature from metatomic.torch.o3._utils import map_selected_atoms_to_rotated_copies @@ -61,23 +57,32 @@ def _tensor_map_with_components( component_names, ) -> TensorMap: """Create a one-block TensorMap with the requested component-axis names.""" + device = values.device components = [ - Labels.range(name, values.shape[axis + 1]) + Labels.range(name, values.shape[axis + 1]).to(device=device) for axis, name in enumerate(component_names) ] return TensorMap( - Labels("_", torch.tensor([[0]], dtype=torch.int64)), + Labels("_", torch.tensor([[0]], dtype=torch.int64, device=device)), [ TensorBlock( values=values, - samples=Labels.range("system", values.shape[0]), + samples=Labels.range("system", values.shape[0]).to(device=device), components=components, - properties=Labels.range("property", values.shape[-1]), + properties=Labels.range("property", values.shape[-1]).to(device=device), ) ], ) +def _spherical_quadrupole_values(matrices: torch.Tensor) -> torch.Tensor: + """Return the l=2 spherical components of ``(n, 3, 3, p)`` matrices, + computed through the ``decompose_quantity`` boundary.""" + tensor = _tensor_map_with_components(matrices, ["xyz_1", "xyz_2"]) + decomposed = decompose_quantity("non_conservative_stress", tensor) + return decomposed.block({"o3_lambda": 2}).values + + class _EmptyModel(torch.nn.Module): """Provide the model interface without producing any outputs.""" @@ -167,10 +172,7 @@ 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] = [] + self.requested_outputs: List[ModelOutput] = [] def forward( self, @@ -180,11 +182,7 @@ def forward( ) -> 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) + self.requested_outputs.extend(outputs.values()) return super().forward(systems, outputs, selected_atoms) @@ -569,7 +567,7 @@ def forward( TensorBlock( spherical, system_samples, - [_o3_mu_labels(1, device)], + [o3_mu_labels(1, device)], properties, ) ], @@ -582,7 +580,7 @@ def forward( for system in systems ] ).unsqueeze(-1) - _, _, spherical = _matrices_to_spherical(matrices) + spherical = _spherical_quadrupole_values(matrices) result["mtt::spherical_quadrupole"] = TensorMap( Labels( ["o3_lambda", "o3_sigma"], @@ -592,7 +590,7 @@ def forward( TensorBlock( spherical, system_samples, - [_o3_mu_labels(2, device)], + [o3_mu_labels(2, device)], properties, ) ], @@ -608,19 +606,15 @@ 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 + _, w = get_rotation_quadrature(lebedev_order, n_inplane) 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.""" + """Quadrature 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() + matrices, _ = get_rotation_quadrature(lebedev_order, n_inplane) identity = np.broadcast_to(np.eye(3), matrices.shape) assert np.allclose( @@ -645,22 +639,6 @@ def test_quadrature_validation(self): 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), @@ -672,6 +650,7 @@ def test_quadrature_validation(self): get_rotation_quadrature(4, 3) def test_degree_two_grid_resolves_l1_products(self): + """A degree-2 grid integrates products of l=1 functions exactly.""" order, n_rotations = choose_quadrature(2) rotations, weights = get_rotation_quadrature(order, n_rotations) function = rotations[:, 2, 0] @@ -720,56 +699,19 @@ def test_rejects_grid_too_small_for_character_sectors(self): ) def test_rejects_invalid_constructor_arguments(self): - """Every integer constructor argument should enforce its documented range.""" + """Integer arguments are validated with the argument name in the message.""" 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" + message = "batch_size 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, + batch_size=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.""" @@ -932,7 +874,7 @@ def test_stress_character_projection_combines_target_and_character_sectors(self) 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)] + 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 @@ -946,19 +888,7 @@ def test_stress_character_projection_combines_target_and_character_sectors(self) 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, - ): + def test_source_request_contains_only_the_shared_sample_kind(self): """Do not pass diagnostic metadata to the underlying source output.""" base_model = _CountingLinearEnergyModel() model = SymmetrizedModel( @@ -971,22 +901,26 @@ def test_source_request_contains_only_the_shared_sample_kind( model( [_forward_test_system([[1.0, 2.0, 3.0]])], { - requested_name: ModelOutput( + name: ModelOutput( unit=unit, sample_kind="system", description="Metadata for the public result.", ) + for name, unit in [ + ("energy", "eV"), + ("o3::variance::energy", "(eV)^2"), + ("o3::character_projection::energy", "(eV)^2"), + ] }, 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) == {""} + for output in base_model.requested_outputs: + assert output.sample_kind == "system" + assert output.unit == "" + assert output.explicit_gradients == [] + assert output.description == "" def test_rejects_an_output_above_the_declared_target_rank(self): """Reject a rank-two spherical output when the declared limit is one.""" @@ -1064,7 +998,9 @@ def test_rejects_a_negative_quadrature_error_and_converges(self): [ "energy/pbe", "mtt::feature::node", - # the reserved prefix is stripped exactly once, keeping "mtt::aux::" + # "o3::variance::mtt::aux::features" must evaluate the source output + # "mtt::aux::features": only the "o3::variance::" prefix is removed, + # the "mtt::" namespace inside the source name is left alone "mtt::aux::features", ], ) @@ -1287,7 +1223,7 @@ def test_equivariant_outputs_preserve_values_metadata_and_zero_variance(self): ) quadrupole = result["mtt::spherical_quadrupole"] assert quadrupole.keys.values.tolist() == [[2, 1]] - _, _, expected_quadrupole = _matrices_to_spherical( + expected_quadrupole = _spherical_quadrupole_values( torch.outer(system.positions[0], system.positions[0]).reshape(1, 3, 3, 1) ) assert torch.allclose( @@ -1317,7 +1253,10 @@ def test_equivariant_outputs_preserve_values_metadata_and_zero_variance(self): @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.""" + """Averages and variances keep the model dtype and implicit backward path.""" + tolerance = 2.0e-5 if dtype == torch.float32 else 1.0e-12 + + # variance of a linear model system = _forward_test_system( [[1.0, 2.0, 3.0]], dtype=dtype, @@ -1337,7 +1276,6 @@ def test_dtype_and_implicit_autograd(self, dtype): 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, @@ -1345,10 +1283,10 @@ def test_dtype_and_implicit_autograd(self, dtype): atol=tolerance, ) - def test_average_output_preserves_implicit_autograd(self): - """The averaged output should keep the implicit backward path to positions.""" + # average of an invariant model output system = _forward_test_system( [[1.0, 2.0, 3.0], [-0.5, 0.25, 1.0]], + dtype=dtype, requires_grad=True, ) model = SymmetrizedModel( @@ -1359,15 +1297,14 @@ def test_average_output_preserves_implicit_autograd(self): result = model([system], {"energy": ModelOutput(sample_kind="system")}, None) - gradient = torch.autograd.grad( - result["energy"].block().values.sum(), - system.positions, - )[0] + average = result["energy"].block().values + assert average.dtype == dtype + gradient = torch.autograd.grad(average.sum(), system.positions)[0] assert torch.allclose( gradient, 2.0 * system.positions, rtol=0.0, - atol=1.0e-12, + atol=tolerance, ) def test_rejects_unknown_o3_requests(self): @@ -1424,7 +1361,7 @@ def test_input_limit_distinguishes_spherical_from_cartesian(self): TensorBlock( values=torch.ones((1, 3, 1), dtype=torch.float64), samples=Labels.range("atom", 1), - components=[_o3_mu_labels(1, torch.device("cpu"))], + components=[o3_mu_labels(1, torch.device("cpu"))], properties=Labels.range("property", 1), ) ], @@ -1543,8 +1480,7 @@ def test_downcast_integration_buffers_warn_and_run(self): 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) + assert bool(torch.isfinite(variance).all()) def test_rejects_a_model_that_omits_the_requested_output(self): """Fail loudly when the underlying model does not return a source.""" @@ -1664,7 +1600,7 @@ def test_wrap_declares_capabilities(self, max_angular_momentum_character): assert capabilities.atomic_types == [1, 6, 8] assert capabilities.interaction_range == 4.5 assert capabilities.length_unit == "A" - assert capabilities.supported_devices == ["cuda", "cpu"] + assert capabilities.supported_devices == ["cuda", "mps", "cpu"] expected_names = set(source_outputs) expected_names.update("o3::variance::" + name for name in source_outputs) @@ -1785,28 +1721,6 @@ def test_rejects_reserved_source_names(self, source_name): 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") @@ -1918,12 +1832,18 @@ def test_preserves_requirements_and_runs_after_save_load(self, tmp_path): ) @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. + def test_rejects_systems_on_a_different_device(self): + """A CPU module refuses CUDA Systems instead of silently mixing devices.""" + model = SymmetrizedModel(_LinearEnergyModel(), max_angular_momentum_target=0) + system = _forward_test_system([[1.0, 2.0, 3.0]]).to(device="cuda") - The wrapper's model dtype is float32, while its integration buffers stay - float64 throughout; nothing here downcasts the module itself. - """ + message = "SymmetrizedModel and input Systems must use the same device" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + model([system], {"energy": ModelOutput(sample_kind="system")}, None) + + @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.""" base = AtomisticModel( _LinearModelWithRequirements().eval(), ModelMetadata(name="CUDA source model"), @@ -1962,17 +1882,6 @@ def test_saved_wrapper_runs_on_cuda(self, tmp_path): ) 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", @@ -2023,27 +1932,17 @@ def test_selected_atoms_system_column_found_by_name(): 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.""" +@pytest.mark.parametrize( + ("source_name", "component_names"), + [ + ("velocity", ["xyz"]), + # random matrices are non-symmetric, so the l=1 pseudovector is non-zero + ("non_conservative_stress", ["xyz_1", "xyz_2"]), + ], +) +def test_decompose_quantity_commutes_with_o3(inversion, source_name, component_names): + """Decomposing then transforming matches transforming then decomposing.""" proper_rotation = torch.tensor( [ [-2.0 / 3.0, 2.0 / 15.0, 11.0 / 15.0], @@ -2054,134 +1953,30 @@ def test_cartesian_vectors_to_spherical_commutes_with_o3(inversion): ) 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, + max_angular_momentum=2, ) + generator = torch.Generator().manual_seed(1234) + shape = (2,) + (3,) * len(component_names) + (1,) + values = torch.randn(shape, dtype=torch.float64, generator=generator) + tensor = _tensor_map_with_components(values, component_names) - transformed_cartesian = _cartesian_vectors_to_spherical( - transformation.transform_cartesian(cartesian), - component_axis=1, + transformed_then_decomposed = decompose_quantity( + source_name, + transformation.transform_tensormap(tensor), ) - transformed_spherical = transformation.transform_spherical( - _cartesian_vectors_to_spherical(cartesian, component_axis=1), - ell=1, - sigma=1, + decomposed_then_transformed = transformation.transform_tensormap( + decompose_quantity(source_name, tensor) ) - assert torch.allclose( - transformed_cartesian, - transformed_spherical, + mts.allclose_raise( + transformed_then_decomposed, + decomposed_then_transformed, 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", - ], -) +@pytest.mark.parametrize("source_name", ["energy", "charge/direct"]) 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) @@ -2194,19 +1989,13 @@ def test_decompose_quantity_scalar_quantities(source_name): 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().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): +def test_decompose_quantity_cartesian_vectors_preserve_autograd(): + source_name = "non_conservative_force/direct" """Cartesian vectors should become l=1 and preserve implicit autograd.""" values = torch.tensor( [[[1.0], [2.0], [3.0]]], @@ -2219,7 +2008,7 @@ def test_decompose_quantity_cartesian_vectors_preserve_autograd(source_name): 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 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), @@ -2244,9 +2033,9 @@ def test_decompose_quantity_non_conservative_stress_combines_irreps(): 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 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), @@ -2262,6 +2051,49 @@ def test_decompose_quantity_non_conservative_stress_combines_irreps(): assert block.samples == tensor.block().samples assert block.properties == tensor.block().properties + # the symmetric traceless x^2 - y^2 part lands in the l=2, m=+2 component + traceless = torch.zeros((1, 3, 3, 1), dtype=torch.float64) + traceless[0, 0, 0, 0] = 1.0 + traceless[0, 1, 1, 0] = -1.0 + result = decompose_quantity( + "non_conservative_stress", + _tensor_map_with_components(traceless, ["xyz_1", "xyz_2"]), + ) + expected_l2 = torch.zeros((1, 5, 1), dtype=torch.float64) + expected_l2[0, 4, 0] = 2.0**0.5 + assert torch.allclose( + result.block({"o3_lambda": 2}).values, + expected_l2, + rtol=0.0, + atol=1.0e-12, + ) + assert torch.allclose( + result.block({"o3_lambda": 0}).values, + torch.zeros((1, 1, 1), dtype=torch.float64), + rtol=0.0, + atol=1.0e-12, + ) + + # the decomposition is orthonormal: the Frobenius norm is preserved + generator = torch.Generator().manual_seed(1234) + random_matrices = torch.randn( + (4, 3, 3, 2), dtype=torch.float64, generator=generator + ) + result = decompose_quantity( + "non_conservative_stress", + _tensor_map_with_components(random_matrices, ["xyz_1", "xyz_2"]), + ) + spherical_norm_squared = torch.zeros((4, 2), dtype=torch.float64) + for block in result.blocks(): + spherical_norm_squared += block.values.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, + ) + def test_decompose_quantity_does_not_infer_custom_cartesian_semantics(): """A generic 3x3 output should pass through unchanged.""" From 1092aafc4938d1ed81f691fad7c1e58613fd52d4 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Thu, 6 Aug 2026 13:03:32 +0200 Subject: [PATCH 17/18] Recognize deprecated quantity names when decomposing outputs metatrain does not rename deprecated target names on checkpoint upgrade: models keep answering to their native (plural) spellings, and metatomic's AtomisticModel is the alias bridge for engines. Share that bridge's table and name normalization in _quantities.py and use it when categorizing outputs for decomposition, so wrapped models trained with plural names still get per-irrep diagnostics. --- .../metatomic/torch/_quantities.py | 58 ++++++++++++++----- .../metatomic_torch/metatomic/torch/model.py | 17 +----- .../metatomic/torch/o3/_decompose.py | 9 ++- .../metatomic/torch/o3/_symmetrized.py | 3 +- .../tests/symmetrized_model.py | 16 +++++ 5 files changed, 70 insertions(+), 33 deletions(-) diff --git a/python/metatomic_torch/metatomic/torch/_quantities.py b/python/metatomic_torch/metatomic/torch/_quantities.py index a6fa13a1..dbbac736 100644 --- a/python/metatomic_torch/metatomic/torch/_quantities.py +++ b/python/metatomic_torch/metatomic/torch/_quantities.py @@ -9,8 +9,8 @@ def standard_quantity_categories() -> Dict[str, str]: 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 not recognized, and - the code using this table treats them as custom quantities. + (singular) spellings appear here: normalize deprecated aliases with + :py:func:`current_quantity_name` before looking them up. TorchScript cannot read a module-level dictionary from a compiled function, so the table is built by this function and bound to @@ -45,18 +45,44 @@ def standard_quantity_categories() -> Dict[str, str]: } -#: mapping from deprecated quantity names to their current name -NEW_QUANTITY_NAMES: Dict[str, str] = { - "features": "feature", - "non_conservative_forces": "non_conservative_force", - "positions": "position", - "momenta": "momentum", - "masses": "mass", - "velocities": "velocity", - "charges": "charge", -} +def new_quantity_names() -> Dict[str, str]: + """Return the mapping from deprecated quantity names to their current name. -#: 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() -} + TorchScript cannot read a module-level dictionary from a compiled function, + so the table is built by this function. + """ + return { + "features": "feature", + "non_conservative_forces": "non_conservative_force", + "positions": "position", + "momenta": "momentum", + "masses": "mass", + "velocities": "velocity", + "charges": "charge", + } + + +def deprecated_quantity_names() -> Dict[str, str]: + """Return the mapping from current quantity names to their deprecated name.""" + result: Dict[str, str] = {} + for deprecated, new in new_quantity_names().items(): + result[new] = deprecated + return result + + +def current_quantity_name(name: str) -> str: + """Replace a deprecated base quantity in ``name`` with its current name.""" + base = name.split("/")[0] + names = new_quantity_names() + if base in names: + return name.replace(base, names[base], 1) + return name + + +def deprecated_quantity_name(name: str) -> str: + """Replace a current base quantity in ``name`` with its deprecated name.""" + base = name.split("/")[0] + names = deprecated_quantity_names() + if base in names: + return name.replace(base, names[base], 1) + return name diff --git a/python/metatomic_torch/metatomic/torch/model.py b/python/metatomic_torch/metatomic/torch/model.py index 73b7282b..bac77a77 100644 --- a/python/metatomic_torch/metatomic/torch/model.py +++ b/python/metatomic_torch/metatomic/torch/model.py @@ -1,4 +1,3 @@ -import copy import datetime import json import math @@ -26,7 +25,7 @@ ) from . import __version__ as metatomic_version from ._extensions import _collect_extensions -from ._quantities import DEPRECATED_QUANTITY_NAMES, NEW_QUANTITY_NAMES +from ._quantities import current_quantity_name, deprecated_quantity_name def load_atomistic_model(path, extensions_directory=None) -> "AtomisticModel": @@ -396,12 +395,6 @@ def __init__( else: raise ValueError(f"unknown dtype in capabilities: {capabilities.dtype}") - # TorchScript methods cannot read module-level dicts: copy onto the instance - self._new_names = copy.deepcopy(NEW_QUANTITY_NAMES) - - # mapping from new names to the corresponding deprecated name - self._deprecated_names = copy.deepcopy(DEPRECATED_QUANTITY_NAMES) - # Pretend that the model can output either the new or deprecated names new_outputs = {} for name in self._model_capabilities_outputs_names: @@ -507,14 +500,10 @@ def requested_inputs(self, use_new_names: bool = False) -> Dict[str, ModelOutput return inputs def _get_new_name(self, name: str) -> str: - base = name.split("/")[0] - new_base = self._new_names.get(base, base) - return name.replace(base, new_base, 1) + return current_quantity_name(name) def _get_deprecated_name(self, name: str) -> str: - base = name.split("/")[0] - deprecated_base = self._deprecated_names.get(base, base) - return name.replace(base, deprecated_base, 1) + return deprecated_quantity_name(name) def forward( self, diff --git a/python/metatomic_torch/metatomic/torch/o3/_decompose.py b/python/metatomic_torch/metatomic/torch/o3/_decompose.py index 47112a4a..745f6b49 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_decompose.py +++ b/python/metatomic_torch/metatomic/torch/o3/_decompose.py @@ -13,7 +13,7 @@ import torch from metatensor.torch import Labels, TensorBlock, TensorMap -from .._quantities import standard_quantity_categories +from .._quantities import current_quantity_name, standard_quantity_categories from ._utils import copy_tensormap_info, strip_placeholder_key @@ -98,8 +98,13 @@ def decompose_quantity( ``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. + + Deprecated spellings are normalized before the lookup, so a model that only + answers to e.g. ``non_conservative_forces`` still gets its output decomposed + as the current ``non_conservative_force`` quantity, matching the aliasing + that :py:class:`AtomisticModel` applies for engines. """ - quantity = name.split("/", 1)[0] + quantity = current_quantity_name(name).split("/", 1)[0] categories = standard_quantity_categories() if quantity not in categories: return tensor diff --git a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py index 4c9bf4e5..062d2c78 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py +++ b/python/metatomic_torch/metatomic/torch/o3/_symmetrized.py @@ -24,6 +24,7 @@ from .._quantities import ( MAX_ANGULAR_MOMENTUM_PER_CATEGORY, STANDARD_QUANTITY_CATEGORIES, + current_quantity_name, ) from ._decompose import decompose_quantity from ._projections import ( @@ -146,7 +147,7 @@ def _infer_max_angular_momentum( found_standard = False custom_names: List[str] = [] for name in names.keys(): - quantity = name.split("/", 1)[0] + quantity = current_quantity_name(name).split("/", 1)[0] if quantity == "feature": # features are not an irreducible representation of O(3): they are # passed through unchanged and never rotated back diff --git a/python/metatomic_torch/tests/symmetrized_model.py b/python/metatomic_torch/tests/symmetrized_model.py index 23915efe..56a0ae63 100644 --- a/python/metatomic_torch/tests/symmetrized_model.py +++ b/python/metatomic_torch/tests/symmetrized_model.py @@ -2095,6 +2095,22 @@ def test_decompose_quantity_non_conservative_stress_combines_irreps(): ) +def test_decompose_quantity_recognizes_deprecated_spellings(): + """'non_conservative_forces' decomposes like 'non_conservative_force'. + + Models trained before the singular renaming (e.g. metatrain PET + checkpoints) natively answer to the plural names, and metatrain does not + rename them on checkpoint upgrade. + """ + values = torch.tensor([[[1.0], [2.0], [3.0]]], dtype=torch.float64) + tensor = _tensor_map_with_components(values, ["xyz"]) + + result = decompose_quantity("non_conservative_forces", tensor) + + assert result.keys.names == ["o3_lambda", "o3_sigma"] + assert result.keys.values.tolist() == [[1, 1]] + + def test_decompose_quantity_does_not_infer_custom_cartesian_semantics(): """A generic 3x3 output should pass through unchanged.""" tensor = _tensor_map_with_components( From b3104885fb818df5ba889e3abc7e93b24a46feaf Mon Sep 17 00:00:00 2001 From: ppegolo Date: Thu, 6 Aug 2026 13:03:33 +0200 Subject: [PATCH 18/18] Remove the transform_system/block/tensor free functions O3Transformation.transform_systems and .transform_tensormap cover the same functionality; keeping both meant a second validation layer and a second entry point for every operation. transform_tensormap now accepts integer system_ids of any dtype, converting to long internally. --- docs/src/torch/reference/o3.rst | 13 +- metatomic-torch/CHANGELOG.md | 9 +- .../metatomic/torch/o3/__init__.py | 11 +- .../metatomic/torch/o3/_transformations.py | 304 +----------------- python/metatomic_torch/tests/o3.py | 293 +++++------------ 5 files changed, 105 insertions(+), 525 deletions(-) diff --git a/docs/src/torch/reference/o3.rst b/docs/src/torch/reference/o3.rst index 37a9dd79..c7c8dcd3 100644 --- a/docs/src/torch/reference/o3.rst +++ b/docs/src/torch/reference/o3.rst @@ -11,9 +11,10 @@ augmentation. Conventions ----------- -To transform a :py:class:`~metatensor.torch.TensorBlock`, :py:func:`transform_block` -and :py:func:`transform_tensor` need to know, for each component axis, whether it -carries a Cartesian or a spherical tensor. This is inferred from the axis name: +To transform a :py:class:`~metatensor.torch.TensorMap`, +:py:meth:`O3Transformation.transform_tensormap` needs to know, for each component +axis, whether it carries a Cartesian or a spherical tensor. This is inferred from +the axis name: - Cartesian axes are named ``xyz``, or ``xyz_1``, ``xyz_2``, ... for blocks with several Cartesian axes (e.g. rank-2 Cartesian tensors). These are rotated directly @@ -82,9 +83,3 @@ Reference :members: .. autofunction:: metatomic.torch.o3.random_transformations - -.. autofunction:: metatomic.torch.o3.transform_system - -.. autofunction:: metatomic.torch.o3.transform_tensor - -.. autofunction:: metatomic.torch.o3.transform_block diff --git a/metatomic-torch/CHANGELOG.md b/metatomic-torch/CHANGELOG.md index 2392cd91..122c7efb 100644 --- a/metatomic-torch/CHANGELOG.md +++ b/metatomic-torch/CHANGELOG.md @@ -30,12 +30,19 @@ a changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows - Renamed `O3Transformation.is_inverted` to `is_improper`. - `wigners >= 0.4.0` is now required. +### Removed + +- Removed the `transform_system`, `transform_block`, and `transform_tensor` + free functions from `metatomic.torch.o3`; use the `O3Transformation` + methods `transform_systems` and `transform_tensormap` instead. + ### Fixed - `O3Transformation.transform_spherical` no longer applies the `(-1)^ell` parity factor for proper transformations with `sigma = -1`. - Wigner-D evaluation is now stable near the ZYZ Euler-angle poles. -- `transform_system` now preserves autograd for registered neighbor lists. +- Transforming a `System` now preserves autograd for registered neighbor + lists. ## [Version 0.1.16](https://github.com/metatensor/metatomic/releases/tag/metatomic-torch-v0.1.16) - 2026-07-13 diff --git a/python/metatomic_torch/metatomic/torch/o3/__init__.py b/python/metatomic_torch/metatomic/torch/o3/__init__.py index e27c7fbc..3eee0f5b 100644 --- a/python/metatomic_torch/metatomic/torch/o3/__init__.py +++ b/python/metatomic_torch/metatomic/torch/o3/__init__.py @@ -7,19 +7,10 @@ spherical components in a :py:class:`~metatensor.torch.TensorBlock`. """ -from ._transformations import ( - O3Transformation, - random_transformations, - transform_block, - transform_system, - transform_tensor, -) +from ._transformations import O3Transformation, random_transformations __all__ = [ "O3Transformation", "random_transformations", - "transform_system", - "transform_tensor", - "transform_block", ] diff --git a/python/metatomic_torch/metatomic/torch/o3/_transformations.py b/python/metatomic_torch/metatomic/torch/o3/_transformations.py index 84fdedac..20e1b78b 100644 --- a/python/metatomic_torch/metatomic/torch/o3/_transformations.py +++ b/python/metatomic_torch/metatomic/torch/o3/_transformations.py @@ -5,7 +5,7 @@ :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. +``forward`` and share one implementation with the eager API. """ from numbers import Integral @@ -19,18 +19,6 @@ 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 _spherical_parity_factor( ell: int, sigma: int, @@ -462,6 +450,8 @@ def transform_tensormap( :return: transformed TensorMap with the same metadata and global information """ + if system_ids is not None: + system_ids = system_ids.to(dtype=torch.long) wigner_matrices: list[torch.Tensor] = [] if max_o3_lambda_in_tensor(tensor) >= 0: wigner_matrices = self._wigner_D_matrices() @@ -557,176 +547,6 @@ def random_transformations( ] -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"] @@ -828,124 +648,6 @@ def max_o3_lambda_in_tensor(tensor: TensorMap) -> int: 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, diff --git a/python/metatomic_torch/tests/o3.py b/python/metatomic_torch/tests/o3.py index 2563b1b8..8f7bf888 100644 --- a/python/metatomic_torch/tests/o3.py +++ b/python/metatomic_torch/tests/o3.py @@ -1,7 +1,6 @@ import io import re -import metatensor.torch as mts import numpy as np import pytest import torch @@ -12,13 +11,7 @@ System, register_autograd_neighbors, ) -from metatomic.torch.o3 import ( - O3Transformation, - random_transformations, - transform_block, - transform_system, - transform_tensor, -) +from metatomic.torch.o3 import O3Transformation, random_transformations # The complex-to-real spherical harmonics conversion is defined only here for now. from metatomic.torch.o3._wigner import _complex_to_real_spherical_harmonics_transform @@ -103,6 +96,16 @@ def _single_block_tensor_map( ) +def _batched(transformations): + """Combine single O3Transformations into one batched transformation.""" + return O3Transformation( + torch.cat([transformation.matrices for transformation in transformations]), + max_angular_momentum=max( + transformation.max_angular_momentum for transformation in transformations + ), + ) + + def test_gradient_components_require_wigner_ranks(): """Angular momenta appearing only in gradient components still require the matching Wigner-D matrices, while purely Cartesian data needs none.""" @@ -242,7 +245,7 @@ def test_transform_system(device, dtype): ) transformation = O3Transformation(matrix, max_angular_momentum=0) - rotated = transform_system(system, transformation) + rotated = transformation.transform_systems(system)[0] assert torch.allclose(rotated.positions, system.positions @ matrix.T, atol=atol) assert torch.allclose(rotated.cell, system.cell @ matrix.T, atol=atol) @@ -304,10 +307,8 @@ def test_transform_system_preserves_neighbor_gradients(): system.add_neighbor_list(options, neighbors) rotation = _rotation_90_degrees_around_z() - transformed = transform_system( - system, - O3Transformation(rotation, max_angular_momentum=0), - ) + transformation = O3Transformation(rotation, max_angular_momentum=0) + transformed = transformation.transform_systems(system)[0] loss = torch.sum(transformed.get_neighbor_list(options).values ** 2) gradient = torch.autograd.grad(loss, positions)[0] @@ -365,11 +366,10 @@ def test_transformation_validation(): system = _make_system([1], dtype=torch.float64) transformation = O3Transformation(torch.eye(3, dtype=torch.float32), 0) message = re.escape( - "System has positions with dtype/device (torch.float64, cpu) differing " - "from the transformations (torch.float32, cpu)." + "system and transformation matrices must have the same dtype and device" ) with pytest.raises(ValueError, match=f"^{message}$"): - transform_system(system, transformation) + transformation.transform_systems(system) def test_random_rotations_are_orthogonal(): @@ -622,10 +622,6 @@ def test_wigner_D_roundoff_at_euler_poles(dtype): def test_gradient_rows_follow_parent_system(): """Gradient rows use their parent sample's transformation; inputs unchanged.""" - systems = [ - _make_system([1, 1]), - _make_system([8, 8, 8]), - ] R_92 = torch.tensor(_axis_angle([1.0, 2.0, 3.0], 0.7), dtype=torch.float64) R_38 = torch.tensor(_axis_angle([0.0, 1.0, 1.0], 1.9), dtype=torch.float64) @@ -675,15 +671,13 @@ def test_gradient_rows_follow_parent_system(): pos_grad_before = pos_grad.detach().clone() strain_grad_before = strain_grad.detach().clone() - transformed = transform_tensor( - tensor, - systems, + batch = _batched( [ O3Transformation(R_92, max_angular_momentum=1), O3Transformation(R_38, max_angular_momentum=1), - ], - system_ids=[92, 38], + ] ) + transformed = batch.transform_tensormap(tensor, torch.tensor([92, 38])) transformed_block = transformed.block() assert torch.equal(transformed_block.values, values_before) @@ -729,8 +723,7 @@ def test_gradient_rows_follow_parent_system(): def test_component_metadata_validation(): """Hand-built blocks with wrong component metadata fail with clear errors.""" - systems = [_make_system([1])] - transformations = [O3Transformation(torch.eye(3, dtype=torch.float64), 1)] + transformation = O3Transformation(torch.eye(3, dtype=torch.float64), 1) # unknown component axis name tensor = _single_block_tensor_map( @@ -744,7 +737,7 @@ def test_component_metadata_validation(): "it can not be transformed." ) with pytest.raises(ValueError, match=f"^{message}$"): - transform_tensor(tensor, systems, transformations) + transformation.transform_tensormap(tensor) # o3_sigma outside {-1, +1} tensor = _single_block_tensor_map( @@ -755,7 +748,7 @@ def test_component_metadata_validation(): ) message = re.escape("sigma must be either -1 or +1, got 2.") with pytest.raises(ValueError, match=f"^{message}$"): - transform_tensor(tensor, systems, transformations) + transformation.transform_tensormap(tensor) # misordered Cartesian labels tensor = _single_block_tensor_map( @@ -767,7 +760,7 @@ def test_component_metadata_validation(): "Cartesian component axis 'xyz' must use labels [0, 1, 2] in x, y, z order." ) with pytest.raises(ValueError, match=f"^{message}$"): - transform_tensor(tensor, systems, transformations) + transformation.transform_tensormap(tensor) # misordered spherical labels on an empty block: the validation is # metadata-driven, not row-driven @@ -786,7 +779,7 @@ def test_component_metadata_validation(): "through 1 in ascending order." ) with pytest.raises(ValueError, match=f"^{message}$"): - transform_tensor(tensor, systems, transformations) + transformation.transform_tensormap(tensor) def test_insufficient_max_angular_momentum(): @@ -797,7 +790,6 @@ def test_insufficient_max_angular_momentum(): samples=Labels(["system"], torch.tensor([[0]])), components=[Labels(["o3_mu"], torch.arange(-1, 2).reshape(-1, 1))], ) - systems = [_make_system([1])] transformations = random_transformations( 1, @@ -806,7 +798,7 @@ def test_insufficient_max_angular_momentum(): ) message = re.escape("ell=1 exceeds max_angular_momentum=0.") with pytest.raises(ValueError, match=f"^{message}$"): - transform_tensor(tensor, systems, transformations) + transformations[0].transform_tensormap(tensor) transformations = random_transformations( 1, @@ -814,7 +806,7 @@ def test_insufficient_max_angular_momentum(): device=torch.device("cpu"), dtype=torch.float64, ) - transformed = transform_tensor(tensor, systems, transformations) + transformed = transformations[0].transform_tensormap(tensor) # a rotation (with sigma parity +-1) preserves the norm of ell=1 values torch.testing.assert_close( transformed.block().values.norm(), @@ -841,11 +833,9 @@ def test_transform_tensor_combines_spherical_axis_parities(): ], ) - transformed = transform_tensor( - tensor, - [_make_system([1])], - [O3Transformation(-torch.eye(3, dtype=torch.float64), 1)], - ) + transformed = O3Transformation( + -torch.eye(3, dtype=torch.float64), 1 + ).transform_tensormap(tensor) assert torch.allclose( transformed.block().values, @@ -856,8 +846,7 @@ def test_transform_tensor_combines_spherical_axis_parities(): def test_rows_route_by_system_id(): - """Default, list, and tensor IDs route each row to its own transformation.""" - systems = [_make_system([1]), _make_system([8])] + """Default and explicit IDs route each row to its own transformation.""" transformations = [ O3Transformation( torch.tensor( @@ -907,29 +896,18 @@ def make_block(sample_system_ids): properties=Labels(["p"], torch.tensor([[0]])), ) + batch = _batched(transformations) for system_ids, sample_system_ids in [ (None, [0, 1]), - ([92, 38], [92, 38]), + (torch.tensor([92, 38]), [92, 38]), (torch.tensor([92, -7], dtype=torch.int32), [92, -7]), ]: - kwargs = {} if system_ids is None else {"system_ids": system_ids} - transformed = transform_tensor( + transformed = batch.transform_tensormap( TensorMap(keys, [make_block(sample_system_ids)]), - systems, - transformations, - **kwargs, + system_ids, ).block() assert torch.equal(transformed.values, expected) - # the public transform_block entry point agrees - transformed_block = transform_block( - keys[0], - make_block([0, 1]), - systems, - transformations, - ) - assert torch.equal(transformed_block.values, expected) - def _scalar_two_system_tensor(device="cpu"): """A minimal two-row scalar TensorMap with ``"system"`` labels 92 and 38.""" @@ -940,121 +918,37 @@ def _scalar_two_system_tensor(device="cpu"): ) -def test_system_ids_validation(): - """Bad system assignments fail up front instead of misrouting rows.""" - systems = [_make_system([1]), _make_system([8])] - transformations = [ - O3Transformation(torch.eye(3, dtype=torch.float64), 0), - O3Transformation(torch.eye(3, dtype=torch.float64), 0), - ] - - # one transformation per system, one distinct id per system - message = re.escape( - "Expected one transformation per system, but got len(systems)=2 and " - "len(transformations)=1." - ) - with pytest.raises(ValueError, match=f"^{message}$"): - transform_tensor( - _scalar_two_system_tensor(), systems, transformations[:1], [92, -7] - ) - message = re.escape( - "system_ids must contain exactly one entry per system, but got " - "len(system_ids)=1 and len(systems)=2." - ) - with pytest.raises(ValueError, match=f"^{message}$"): - transform_tensor(_scalar_two_system_tensor(), systems, transformations, [92]) - message = re.escape( - "system_ids must contain one distinct entry per system, but got [92, 92]." - ) - with pytest.raises(ValueError, match=f"^{message}$"): - transform_tensor( - _scalar_two_system_tensor(), systems, transformations, [92, 92] - ) - - # ids must live with the values ("meta" needs no accelerator hardware) - message = re.escape( - "system_ids are on device cpu, but the values to transform are on device meta." +def test_rejects_unassigned_system_labels(): + """Every ``"system"`` label present in a block must appear in ``system_ids``.""" + batch = O3Transformation( + torch.eye(3, dtype=torch.float64).expand(2, 3, 3), + max_angular_momentum=0, ) - with pytest.raises(ValueError, match=f"^{message}$"): - transform_tensor( - _scalar_two_system_tensor(device="meta"), - systems, - transformations, - torch.tensor([92, -7]), - ) - # every "system" label present in a block must appear in system_ids message = re.escape( "Block samples contain system labels [38] that are not in " "system_ids=[92, 99]. Every sample must be assigned to a system in the " "transformation." ) with pytest.raises(ValueError, match=f"^{message}$"): - transform_tensor( - _scalar_two_system_tensor(), systems, transformations, [92, 99] - ) - - # with multiple systems, each row must identify its system - no_system_column = _single_block_tensor_map( - values=torch.zeros((1, 1), dtype=torch.float64), - samples=Labels(["atom"], torch.tensor([[0]])), - components=[], - ) - message = re.escape("multiple transformations require a 'system' sample dimension") - with pytest.raises(ValueError, match=f"^{message}$"): - transform_tensor(no_system_column, systems, transformations) - - # a transformation with no assigned rows is still validated: a silent - # mismatch would only surface once such rows appear in another batch - block = TensorBlock( - values=torch.ones((1, 1), dtype=torch.float64), - samples=Labels(["system"], torch.tensor([[92]])), - components=[], - properties=Labels(["p"], torch.tensor([[0]])), - ) - transformations[1] = O3Transformation(torch.eye(3, dtype=torch.float32), 0) - message = re.escape( - "Transformation at index 1 has dtype/device (torch.float32, cpu), " - "differing from the values to transform (torch.float64, cpu)." - ) - with pytest.raises(ValueError, match=f"^{message}$"): - transform_tensor( - TensorMap(Labels(["_"], torch.tensor([[0]])), [block]), - systems, - transformations, - system_ids=[92, 38], - ) + batch.transform_tensormap(_scalar_two_system_tensor(), torch.tensor([92, 99])) def test_transform_empty_inputs_are_no_ops(): - """Empty blocks, empty TensorMaps, and empty system lists pass through.""" - non_empty = _single_block_tensor_map( - values=torch.tensor([[1.0], [2.0]], dtype=torch.float64), - samples=Labels(["system"], torch.tensor([[0], [1]])), - components=[], + """Blocks with no samples and TensorMaps with no blocks pass through.""" + transformation = O3Transformation( + _rotation_90_degrees_around_z(), + max_angular_momentum=0, ) - passthrough = transform_tensor(non_empty, [], []) - assert torch.equal(passthrough.block().values, non_empty.block().values) - block = TensorBlock( + no_samples = _single_block_tensor_map( values=torch.empty((0, 1), dtype=torch.float64), - samples=Labels( - ["system"], - torch.empty((0, 1), dtype=torch.int64), - ), + samples=Labels(["system"], torch.empty((0, 1), dtype=torch.int64)), components=[], - properties=Labels(["p"], torch.tensor([[0]])), ) - - transformed = transform_block( - Labels(["_"], torch.tensor([[0]]))[0], - block, - [], - [], - ) - - assert torch.equal(transformed.values, block.values) - assert transformed.samples == block.samples + transformed = transformation.transform_tensormap(no_samples) + assert torch.equal(transformed.block().values, no_samples.block().values) + assert transformed.block().samples == no_samples.block().samples empty = TensorMap( Labels( @@ -1063,31 +957,14 @@ def test_transform_empty_inputs_are_no_ops(): ), [], ) - # the mixed transformation dtypes check that no values dtype is imposed - # when there is nothing to transform - transformed_map = transform_tensor( - empty, - [_make_system([1]), _make_system([8])], - [ - O3Transformation( - torch.eye(3, dtype=torch.float64), - max_angular_momentum=0, - ), - O3Transformation( - torch.eye(3, dtype=torch.float32), - max_angular_momentum=0, - ), - ], - ) + transformed_map = transformation.transform_tensormap(empty) assert len(transformed_map) == 0 assert transformed_map.keys == empty.keys def test_single_system_routes_all_rows(): - """With one system, all rows are transformed regardless of ``"system"`` labels.""" - systems = [_make_system([1, 1])] - + """One operation transforms all rows regardless of ``"system"`` labels.""" transformation = O3Transformation(_rotation_90_degrees_around_z(), 1) values = torch.tensor( [ @@ -1113,14 +990,12 @@ def test_single_system_routes_all_rows(): samples=samples, components=[Labels(["xyz"], torch.arange(3).reshape(-1, 1))], ) - transformed = transform_tensor(tensor, systems, [transformation]) + transformed = transformation.transform_tensormap(tensor) assert torch.equal(transformed.block().values, expected) def test_block_with_subset_of_systems(): """Blocks may cover only some systems; ids must survive beyond int32.""" - systems = [_make_system([1]), _make_system([8])] - transformation_92 = O3Transformation(_rotation_90_degrees_around_z(), 1) transformation_38 = O3Transformation( torch.eye(3, dtype=torch.float64), @@ -1139,11 +1014,9 @@ def test_block_with_subset_of_systems(): components=[Labels(["xyz"], torch.arange(3).reshape(-1, 1))], ) - transformed = transform_tensor( + transformed = _batched([transformation_92, transformation_38]).transform_tensormap( tensor, - systems, - [transformation_92, transformation_38], - system_ids=[92, 2**40], + torch.tensor([92, 2**40]), ) expected = torch.tensor( @@ -1156,8 +1029,6 @@ def test_block_with_subset_of_systems(): def test_pair_samples_routing(): """Row routing by the "system" column also works for pair-sampled blocks (e.g. atom-pair targets), which carry extra sample columns beyond "system"/"atom".""" - systems = [_make_system([1, 8]), _make_system([1, 8])] - R0 = O3Transformation(torch.eye(3, dtype=torch.float64), 1) # 90-degree rotation around z: (x,y) -> (-y, x) c, s = np.cos(np.pi / 2), np.sin(np.pi / 2) @@ -1194,7 +1065,7 @@ def test_pair_samples_routing(): ], ) - transformed = transform_tensor(tensor, systems, [R0, R1]) + transformed = _batched([R0, R1]).transform_tensormap(tensor) result_block = transformed.block() expected = torch.tensor( @@ -1208,13 +1079,13 @@ def test_pair_samples_routing(): @pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE) @pytest.mark.parametrize("parities", [(1.0, 1.0), (-1.0, -1.0), (1.0, -1.0)]) -def test_batched_tensor_transform_matches_transform_tensor( +def test_batched_tensor_transform_matches_singleton_reference( device, dtype, parities, ): - """A batched ``O3Transformation`` matches ``transform_tensor``, including - for batches mixing proper and improper operations.""" + """Rows routed to each operation of a batch match the single-operation + result, including for batches mixing proper and improper operations.""" dtype = getattr(torch, dtype) atol = 1.0e-5 if dtype == torch.float32 else 1.0e-12 proper_matrices = [ @@ -1290,18 +1161,32 @@ def test_batched_tensor_transform_matches_transform_tensor( values_before = values.detach().clone() gradient_values_before = gradient_values.detach().clone() - expected = transform_tensor( - tensor, - [ - _make_system([1], device=device, dtype=dtype), - _make_system([8], device=device, dtype=dtype), - ], - transformations, - ) result = batch.transform_tensormap(tensor) - mts.allclose_raise(result, expected, rtol=0.0, atol=atol) - assert result.info() == expected.info() + # reference: a single operation transforms every row through the singleton + # path, so the batch rows routed to it must match row by row + system_column = tensor.block(0).samples.column("system").to(dtype=torch.long) + gradient_parents = ( + tensor.block(0).gradient("parameter").samples.column("sample") + ).to(dtype=torch.long) + for index, single in enumerate(transformations): + reference = single.transform_tensormap(tensor) + rows = system_column == index + assert torch.allclose( + result.block(0).values[rows], + reference.block(0).values[rows], + rtol=0.0, + atol=atol, + ) + gradient_rows = rows[gradient_parents] + assert torch.allclose( + result.block(0).gradient("parameter").values[gradient_rows], + reference.block(0).gradient("parameter").values[gradient_rows], + rtol=0.0, + atol=atol, + ) + + assert result.info() == tensor.info() assert torch.equal(values, values_before) assert torch.equal(gradient_values, gradient_values_before) @@ -1344,15 +1229,15 @@ def test_single_transformation_needs_no_system_label(): ) result = transformation.transform_tensormap(tensor) - expected = transform_tensor( - tensor, - [_make_system([1])], - [transformation], - ) + expected = transformation.transform_spherical( + tensor.block().values.squeeze(-1), + ell=1, + sigma=1, + ).unsqueeze(-1) assert torch.allclose( result.block().values, - expected.block().values, + expected, rtol=0.0, atol=1.0e-12, ) @@ -1453,7 +1338,7 @@ def test_batched_transformation_matches_singles(): batch_systems = batch.transform_systems(system) assert len(batch_systems) == 4 for index, single in enumerate(singles): - expected_system = transform_system(system, single) + expected_system = single.transform_systems(system)[0] assert torch.allclose( batch_systems[index].positions, expected_system.positions,