O(3) symmetries¶
The metatomic.torch.o3 module provides tools to deal with O(3)
symmetries. First it offers functions to rotate and invert
System and TensorMap
objects, for example to generate randomly rotated copies of a structure for data
augmentation. Second, it offers tools to convert
TensorMap between Cartesian and spherical forms.
Conventions for O(3) representations¶
To transform a TensorMap (or a
TensorBlock), the
O3Transformations.transform_tensormap() (resp.
O3Transformations.transform_block()) need to know, for each component
axis, whether it carries a Cartesian or a spherical tensor. This is inferred
from the axis name: axis, whether it carries a Cartesian or a spherical tensor.
This is inferred from the axis name:
Cartesian axes are named
xyz, orxyz_1,xyz_2, … for blocks with several Cartesian axes (e.g. rank-2 Cartesian tensors). These are rotated directly with the (3, 3) transformation matrix \(R\) passed toO3Transformations, following the usual column-vector convention \(v' = R v\) (equivalently, since values are stored as rows,values_transformed = values @ R.T).Spherical axes are named
o3_mu, oro3_mu_1,o3_mu_2, … They are rotated with the real Wigner-D matrix for the angular momentumo3_lambda(respectivelyo3_lambda_1,o3_lambda_2, …) found in the block’s key. Real, rather than complex, spherical harmonics are used throughout, matching the convention used elsewhere in metatomic and metatensor for spherical targets.
Wigner-D matrices¶
Complex Wigner-D matrices are computed by the wigners package, using the convention
with ZYZ Euler angles \((\alpha, \beta, \gamma)\) extracted from the proper part of \(R\), i.e. from \(R\) itself when \(\det R = 1\), or from \(-R\) when \(\det R = -1\), such that this proper part equals \(R_z(\alpha) R_y(\beta) R_z(\gamma)\).
These are then converted to real Wigner-D matrices through the unitary change of basis \(T\) mapping complex spherical harmonics \(Y_{\ell}^{m}\) to real ones:
so that the real Wigner-D matrix is \(D^{\ell}_{\text{real}} = T^{*}
D^{\ell} T^{T}\). Spherical components follow the same column-vector convention
as Cartesian ones, \(f' = D^{\ell}_{\text{real}} f\) (equivalently, since
values are stored as rows, values_transformed = values @ D.T).
For an improper rotation (a rotation composed with an inversion), Cartesian axes
are flipped as part of the transformation matrix itself, while spherical axes
pick up an extra parity factor \((-1)^{\ell} \sigma\), where \(\ell\) is
o3_lambda and \(\sigma\) is the block’s o3_sigma (its behavior under
inversion, +1 or -1), also read from the key.
A TensorBlock in general contains values
associated with several systems. This information is contained in the
"system" samples dimension; each row is rotated with the transformation of
the system it belongs to. Gradient blocks are routed the same way, via their
parent block’s "system" column. When only one system is being transformed,
the "system" column is optional and, if present, is ignored: every row is
rotated with the (single) given transformation.
- class metatomic.torch.o3.O3Transformations(matrices: Tensor, max_angular_momentum: int)[source]¶
A batch of one or more O(3) transformations, represented by
(N, 3, 3)rotation or improper-rotation matrices.The module stores the matrices, an improper-operations mask, and a packed buffer of Wigner-D matrices. It is constructed eagerly in
__init__(building the Wigner-D matrices requires thewignerspackage); the transform methods are TorchScript compatible and can be called from a scriptedforward.Transform methods accept an
indicesargument to select a sub-batch of operations, and aadd_inversionflag to compose every selected operation with the inversion. Inverse transformations are available through theinverse_transform_*methods, which share their kernel with the forward ones and use transposed matrices / Wigner-D matrices.- Parameters:
- wigner_D_matrices(ell: int) Tensor[source]¶
Return the proper-part Wigner-D matrices for
ell.For improper operations, the inversion-parity factor
sigma * (-1) ** ellis applied separately when transforming spherical values.
- inverse_wigner_D_matrices(ell: int) Tensor[source]¶
Return the transposed (inverse) proper-part Wigner-D matrices.
The inverse of a real Wigner-D matrix is its transpose. This is the accessor used by projection integrals that contract values against the inverse representation directly.
- transform_cartesian(vectors: Tensor, add_inversion: bool = False) Tensor[source]¶
Apply the transformations to Cartesian vectors.
- inverse_transform_cartesian(vectors: Tensor, add_inversion: bool = False) Tensor[source]¶
Apply the inverse transformations to Cartesian vectors.
The inverse of an orthogonal matrix is its transpose.
- Parameters:
- Returns:
transformed vectors, with the input shape for a single transformation or a leading batch axis for a batch of more than one
- Return type:
- transform_spherical(values: Tensor, ell: int, sigma: int, add_inversion: bool = False) Tensor[source]¶
Apply the transformations to spherical values.
- Parameters:
values (Tensor) – (…, 2*ell+1) tensor of spherical values
ell (int) – angular momentum in
[0, max_angular_momentum]sigma (int) –
+1for a proper spherical representation or-1for a pseudo one. Under an improper transformation, the representation acquires the factorsigma * (-1) ** ell.add_inversion (bool) – compose every operation with the inversion
- Returns:
transformed values, with the input shape for a single transformation or a leading batch axis for a batch of more than one
- Return type:
- inverse_transform_spherical(values: Tensor, ell: int, sigma: int, add_inversion: bool = False) Tensor[source]¶
Apply the inverse transformations to spherical values.
The inverse Wigner-D matrix is its transpose; the inversion-parity factor for an improper operation is unchanged by the transposition.
- Parameters:
- Returns:
transformed values, with the input shape for a single transformation or a leading batch axis for a batch of more than one
- Return type:
- transform_systems(systems: List[System], add_inversion: bool = False, start: int = 0, stop: int | None = None) List[System][source]¶
Apply transformations to a list of systems.
systems[i]is transformed by operationi. To apply a batch of transformations to the same input system, pass that system repeatedly in the list.Positions, cell vectors, neighbor-list displacements, and custom data following Conventions for O(3) representations are transformed. Atomic types and periodic-boundary flags are preserved.
- Parameters:
systems (List[System]) – input systems, one per operation, matching the transformation matrices in dtype and device
add_inversion (bool) – compose every operation with the inversion
start (int) – first operation to apply
stop (int | None) – end of the range of operations to apply, defaulting to all of them
- Returns:
one transformed
Systemper input system- Return type:
- inverse_transform_systems(systems: List[System], add_inversion: bool = False, start: int = 0, stop: int | None = None) List[System][source]¶
Apply the inverse transformations to a list of systems.
See
transform_systems()for the per-system routing; the inverse uses transposed matrices and Wigner-D matrices.- Parameters:
- Returns:
one transformed
Systemper input system- Return type:
- transform_tensormap(tensor: TensorMap, system_ids: Tensor | None = None, add_inversion: bool = False, start: int = 0, stop: int | None = None) TensorMap[source]¶
Apply the transformations to a TensorMap and its gradients.
Scalar, Cartesian, and spherical data are identified by their component-axis names, following Conventions for O(3) representations. With a batch of more than one operation, the
"system"sample label assigns each value row to an operation: whensystem_idsisNone, the labels index the batch directly, and otherwise rows labelledsystem_ids[i]use operationi. Gradient rows use the operation of the value row referenced by their"sample"label. With a single operation, the"system"label is optional and ignored.- Parameters:
tensor (TensorMap) – TensorMap to transform, matching the transformation matrices in dtype and device
system_ids (Tensor | None) – optional one-dimensional tensor with one distinct
"system"sample label per operationadd_inversion (bool) – compose every operation with the inversion
start (int) – first operation to apply
stop (int | None) – end of the range of operations to apply, defaulting to all of them
- Returns:
transformed TensorMap with the same metadata and global information
- Return type:
- inverse_transform_tensormap(tensor: TensorMap, system_ids: Tensor | None = None, add_inversion: bool = False, start: int = 0, stop: int | None = None) TensorMap[source]¶
Apply the inverse transformations to a TensorMap and its gradients.
The inverse of an orthogonal matrix is its transpose, and the inverse of a real Wigner-D matrix is its transpose, so this shares its kernel with
transform_tensormap()using transposed matrices and Wigner-D matrices. Seetransform_tensormap()for row routing.- Parameters:
tensor (TensorMap) – TensorMap to transform, matching the transformation matrices in dtype and device
system_ids (Tensor | None) – optional one-dimensional tensor with one distinct
"system"sample label per operationadd_inversion (bool) – compose every operation with the inversion before inverting
start (int) – first operation to invert
stop (int | None) – end of the range of operations to invert, defaulting to all of them
- Returns:
transformed TensorMap with the same metadata and global information
- Return type:
- transform_block(key: LabelsEntry, block: TensorBlock, system_ids: Tensor | None = None, add_inversion: bool = False) TensorBlock[source]¶
Apply the transformations to a TensorBlock and its gradients.
See
transform_tensormap()for the row-routing conventions. Thekeymust be the block’s key, carrying theo3_lambda/o3_sigmametadata needed for spherical components.- Parameters:
key (LabelsEntry) – the block’s key
block (TensorBlock) – TensorBlock to transform, matching the transformation matrices in dtype and device
system_ids (Tensor | None) – optional one-dimensional tensor with one distinct
"system"sample label per operationadd_inversion (bool) – compose every operation with an inversion
- Returns:
transformed TensorBlock with the same metadata
- Return type:
- inverse_transform_block(key: LabelsEntry, block: TensorBlock, system_ids: Tensor | None = None, add_inversion: bool = False) TensorBlock[source]¶
Apply the inverse transformations to a TensorBlock and its gradients.
See
transform_block()for the parameters; the inverse uses transposed matrices and Wigner-D matrices.- Parameters:
key (LabelsEntry) – the block’s key
block (TensorBlock) – TensorBlock to transform, matching the transformation matrices in dtype and device
system_ids (Tensor | None) – optional one-dimensional tensor with one distinct
"system"sample label per operationadd_inversion (bool) – compose every operation with an inversion before inverting
- Returns:
transformed TensorBlock with the same metadata
- Return type:
- metatomic.torch.o3.random_transformations(n: int, max_angular_momentum: int = 0, *, device: device, dtype: dtype, add_inversions: bool = False, generator: Generator | None = None) O3Transformations[source]¶
Sample
ntransformations uniformly from SO(3), or from O(3) when inversions are included.Rotations are sampled from the Haar measure on SO(3) via random unit quaternions. When
add_inversionsisTrue, each matrix is independently negated with probability 0.5, giving a uniform distribution over the full O(3) group.- Parameters:
n (int) – positive number of transformations to generate
max_angular_momentum (int) – non-negative maximum angular momentum for Wigner-D matrices
device (device) – target device for the output tensors
dtype (dtype) – target dtype for the output tensors; must be
torch.float32ortorch.float64add_inversions (bool) – if
True, sample from O(3) instead of SO(3)generator (Generator | None) – optional
torch.Generatorfor reproducible sampling; whenNonethe global RNG is used
- Returns:
a single
O3Transformationsholdingnoperations- Return type:
Cartesian and spherical tensors¶
cartesian_to_spherical() converts a Cartesian
tensor of arbitrary rank, stored on a product basis with one xyz_1,
xyz_2, … component axis per Cartesian index, into a set of blocks on a
basis of real spherical harmonics (with a single o3_mu component axis, and
o3_lambda / o3_sigma keys). This is done by coupling the Cartesian
components two by two with Clebsch-Gordan coefficients, starting from the last
ones. For tensors of rank 3 and above, there are multiple ways to get the same
o3_lambda, and the coupling history is stored in additional l_{i} and
k_{i} keys. spherical_to_cartesian() performs
the inverse transformation. These functions can be used from TorchScript, in
which case the Clebsch-Gordan coefficients have to be computed in advance with
calculate_cg_coefficients().
- metatomic.torch.o3.cartesian_to_spherical(tensor: TensorMap, components: List[str], keep_l_in_keys: bool | None = None, remove_blocks_threshold: float | None = 1e-09, cg_coefficients: TensorMap | None = None) TensorMap[source]¶
Transform a
tensorof arbitrary rank from Cartesian form to a spherical form.Starting from a tensor on a basis of product of Cartesian coordinates, this function computes the same tensor using a basis of real spherical harmonics
Y^M_L. For example, a rank 1 tensor with a singlexyzcomponent would be represented as a single L=1 spherical harmonic; while a rank 5 tensor using a product basisℝ^3 ⊗ ℝ^3 ⊗ ℝ^3 ⊗ ℝ^3 ⊗ ℝ^3would require multiple blocks up to L=5 spherical harmonics.A single
TensorBlockin the input might correspond to multipleTensorBlockin the output. The output keys will contain all the dimensions of the input keys, pluso3_lambda(indicating the spherical harmonics degree) ando3_sigma(indicating that this block is a proper or improper tensor with+1and-1respectively). Ifkeep_l_in_keysisTrueor if the input tensor is a tensor of rank 3 or more, the keys will also contain multiplel_{i}andk_{i}dimensions, which indicate which angular momenta have been coupled together in which order to get this block.componentsspecifies which ones of the components of the inputTensorMapshould be transformed from Cartesian to spherical. All these components will be replaced in the output by a singleo3_mucomponent, corresponding to the spherical harmonicsM. Ifcomponentsis empty, the tensor is treated as a rank 0 (scalar) tensor, and the output contains a singleo3_lambda=0,o3_sigma=1block for each input block, with an additionalo3_mucomponent placed after all the existing components.If the input keys only contain the
_placeholder dimension (as created bymetatensor.torch.Labels.single()), it is not included in the output keys.By default, symmetric tensors will only contain blocks corresponding to
o3_sigma=1. This is achieved by checking the norm of the blocks after the full calculation; and dropping any block with a norm belowremove_blocks_threshold. To keep all blocks regardless of their norm, you can setremove_blocks_threshold=None.Gradients are not supported yet, and this function raises an error if any of the input blocks has gradients.
- Parameters:
tensor (TensorMap) – input tensor, using a Cartesian product basis
components (List[str]) – components of the input tensor to transform into spherical components, or an empty list for scalar tensors
keep_l_in_keys (bool | None) –
should the output contains the values of angular momenta that were combined together? This defaults to
Falsefor rank 1 and 2 tensors, andTruefor all other tensors.Keys named
l_{i}correspond to the inputcomponents, withl_1being the last entry incomponentsandl_Nthe first one. Keys namedk_{i}correspond to intermediary spherical components created during the calculation, i.e. ak_{i}used to beo3_lambda.remove_blocks_threshold (float | None) – Numerical tolerance to use when determining if a block’s norm is zero or not. Blocks with zero norm will be excluded from the output. Set this to
Noneto keep all blocks in the output.cg_coefficients (TensorMap | None) – Clebsch-Gordan coefficients computed with
calculate_cg_coefficients(), withlambda_maxat least equal to the number ofcomponents. This is optional except when using this function from TorchScript.
- Returns:
TensorMapcontaining spherical components instead of Cartesian components.- Return type:
- metatomic.torch.o3.spherical_to_cartesian(tensor: TensorMap, new_component_names: List[str] | None = None, rank: int | None = None, cg_coefficients: TensorMap | None = None) TensorMap[source]¶
Transform a
tensorfrom spherical form back to Cartesian form.This is the exact inverse of
cartesian_to_spherical(). Starting from a tensor expressed on a basis of real spherical harmonicsY^M_L(as produced bycartesian_to_spherical()), this function reconstructs the corresponding Cartesian product-basis tensor. Blocks that were removed from the spherical tensor (because their norm was zero, see theremove_blocks_thresholdparameter ofcartesian_to_spherical()) are treated as blocks full of zeros.The input keys must contain
o3_lambdaando3_sigma, and the singleo3_mucomponent will be replaced by a number of Cartesian components (according to the rank of the tensor, each with values[0, 1, 2]). For tensors of rank 3 and above the keys must also contain thek_{i}dimensions describing the coupling history (these are added automatically bycartesian_to_spherical()).The rank of the tensor is given by the number of
l_{i}dimensions in the keys. If these dimensions are missing (for example when usingkeep_l_in_keys=Falseincartesian_to_spherical()), therankmust be given explicitly. It can not be determined from theo3_lambdain the keys, since some of the blocks might be missing: for example, an antisymmetric rank 2 tensor only contains ao3_lambda=1block, like a rank 1 tensor.Gradients are not supported yet, and this function raises an error if any of the input blocks has gradients.
- Parameters:
tensor (TensorMap) – input tensor, using a spherical harmonics basis
new_component_names (List[str] | None) – names to give to the reconstructed Cartesian components. There should be as many of them as is the rank of the tensor. If
None, the components are namedxyz_1,xyz_2, …rank (int | None) – rank of the Cartesian tensor. This is required if the keys do not contain the
l_{i}dimensions, and must match the number ofl_{i}dimensions otherwise. Userank=0to convert scalar tensors, removing theo3_mucomponent.cg_coefficients (TensorMap | None) – Clebsch-Gordan coefficients computed with
calculate_cg_coefficients(), withlambda_maxat least equal to the rank of the tensor. This is optional except when using this function from TorchScript to convert tensors of rank 2 and above.
- Returns:
TensorMapcontaining Cartesian components instead of spherical components.- Return type:
- metatomic.torch.o3.calculate_cg_coefficients(lambda_max: int, *, dtype: dtype, device: device) TensorMap[source]¶
Calculate the Clebsch-Gordan coefficients coupling two real spherical harmonics of degrees
l1andl2into a single real spherical harmonic of degreelambda, for alll1,l2andlambdaup tolambda_max.These coefficients are used by
cartesian_to_spherical()andspherical_to_cartesian(), and must be computed in advance and given to these functions explicitly when calling them from TorchScript. Converting a tensor of rankNrequireslambda_max >= N.The coefficients are stored in a
TensorMapwithl1,l2andlambdakeys. Each block contains a single sample (_), them1,m2andmucomponents (with values from-ltolfor the corresponding degree), and a single property (cg_coefficient).- Parameters:
- Returns:
TensorMapcontaining the Clebsch-Gordan coefficients- Return type:
Coupled and uncoupled spherical tensors¶
A block with several spherical component axes o3_mu_1, o3_mu_2, …
(with the corresponding o3_lambda_1 / o3_sigma_1, o3_lambda_2 /
o3_sigma_2, … keys) transforms like a product of spherical harmonics. This
is the uncoupled form, used for example for the blocks of a Hamiltonian
expressed on a basis of atomic orbitals.
uncoupled_to_coupled() uses Clebsch-Gordan
coefficients to transform such a tensor into the coupled form, where each block
has a single o3_mu component axis and transforms like a single spherical
harmonic, with degree and parity given by the o3_lambda and o3_sigma
keys. The degrees and parities of the uncoupled components are kept in the
l_1 / s_1, l_2 / s_2, … keys of the coupled tensor, and
coupled_to_uncoupled() uses them to perform the
inverse transformation. As for the conversion between Cartesian and spherical
tensors, the Clebsch-Gordan coefficients have to be computed in advance with
calculate_cg_coefficients() to use these functions
from TorchScript.
- metatomic.torch.o3.uncoupled_to_coupled(tensor: TensorMap, components: List[str] | None = None, remove_blocks_threshold: float | None = None, cg_coefficients: TensorMap | None = None) TensorMap[source]¶
Transform a
tensorfrom an uncoupled spherical form to a coupled spherical form.In the uncoupled form, the blocks contain multiple spherical components (named
o3_mu_1,o3_mu_2, …), each transforming like a real spherical harmonic with degree and parity given by the correspondingo3_lambda_1/o3_sigma_1,o3_lambda_2/o3_sigma_2, … keys. Such a block transforms like a product of spherical harmonicsY^{m_1}_{l_1} ⊗ Y^{m_2}_{l_2} ⊗ .... In the coupled form, these components are replaced by a singleo3_mucomponent and each block transforms like a single spherical harmonicY^M_L, using Clebsch-Gordan coefficients to combine the spherical harmonics. For example, coupling two components withl_1 = 1andl_2 = 2gives three blocks, withL = 1, 2, 3.The output keys contain
o3_lambdaando3_sigma(the degree and parity of the coupled spherical harmonic), followed by the keys of the input. Theo3_sigmaof the output is given by the product of theo3_sigmaof the coupled components, times(-1)^(l_1 + l_2 + ... + L). The degree and parity of the coupled components are needed to go back to the uncoupled form withcoupled_to_uncoupled(), so theo3_lambda_{i}ando3_sigma_{i}keys are kept in the output, renamed tol_{i}ands_{i}. When coupling three or more components, the components are coupled two by two, starting from the last ones, and the degrees of the intermediate spherical harmonics are stored ink_{i}keys.Depending on the symmetries of the input (for example when coupling a component with itself), some of the blocks will be zero everywhere. Any block with a squared norm below
remove_blocks_thresholdis not included in the output. To keep all blocks regardless of their norm, you can setremove_blocks_threshold=None.Gradients are not supported yet, and this function raises an error if any of the input blocks has gradients.
- Parameters:
tensor (TensorMap) – input tensor, in uncoupled form
components (List[str] | None) – names of the spherical components to couple. These must be consecutive components named
o3_mu_{suffix}, with the correspondingo3_lambda_{suffix}ando3_sigma_{suffix}keys. If this isNone, all the components namedo3_mu_{suffix}are coupled, in the order they appear in the blocks.remove_blocks_threshold (float | None) – Numerical tolerance to use when determining if a block’s norm is zero or not. Blocks with zero norm will be excluded from the output. Set this to
Noneto keep all blocks in the output (the default).cg_coefficients (TensorMap | None) – Clebsch-Gordan coefficients computed with
calculate_cg_coefficients(), withlambda_maxat least equal to the sum of the degrees of the coupled components. This is optional except when using this function from TorchScript.
- Returns:
TensorMapin coupled form- Return type:
- metatomic.torch.o3.coupled_to_uncoupled(tensor: TensorMap, components: List[str] | None = None, cg_coefficients: TensorMap | None = None) TensorMap[source]¶
Transform a
tensorfrom a coupled spherical form back to an uncoupled spherical form.This is the inverse of
uncoupled_to_coupled(). The singleo3_mucomponent of the input is replaced by the sphericalcomponents, using the degrees stored in the correspondingl_{suffix}keys. Thel_{suffix}ands_{suffix}keys are renamed back too3_lambda_{suffix}ando3_sigma_{suffix}in the output, to match the components. All the blocks that only differ by theiro3_lambdaando3_sigma(andk_{i}when uncoupling three or more components) are combined into a single block, and any block missing from the input (for example because its norm was zero, see theremove_blocks_thresholdparameter ofuncoupled_to_coupled()) is treated as a block full of zeros.Gradients are not supported yet, and this function raises an error if any of the input blocks has gradients.
- Parameters:
tensor (TensorMap) – input tensor, in coupled form. The keys must contain
o3_lambda,o3_sigma, thel_{suffix}ands_{suffix}for allcomponents, and thek_{i}describing the intermediate couplings when uncoupling three or more components.components (List[str] | None) – names of the spherical components to create, in the same order as they were given to
uncoupled_to_coupled(). These must be namedo3_mu_{suffix}. If this isNone, one component is created for eachl_{suffix}key, in the order they appear in the keys.cg_coefficients (TensorMap | None) – Clebsch-Gordan coefficients computed with
calculate_cg_coefficients(), withlambda_maxat least equal to the sum of the degrees of the uncoupled components. This is optional except when using this function from TorchScript.
- Returns:
TensorMapin uncoupled form- Return type: