From 9a4bcdb92f5a4f03fdbac8c93e39701daa0626d9 Mon Sep 17 00:00:00 2001 From: Michelangelo Domina Date: Thu, 23 Jul 2026 00:48:10 +0200 Subject: [PATCH 01/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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)