Clebsch-Gordan products

class featomic.utils.DensityCorrelations(*, n_correlations: int, max_angular: int, skip_redundant: bool = True, cg_backend: str | None = None, arrays_backend: str | None = None, dtype: dtype | dtype | None = None, device: device | str | None = None)

Takes n_correlations of iterative CG tensor products of a density with itself to produce a density auto-correlation tensor of higher correlation order.

The constructor computes and stores the CG coefficients. The compute() method computes the auto-correlation by CG tensor product of the input density.

Parameters:
  • n_correlationsint, the number of iterative CG tensor products to perform.

  • max_angularint, the maximum angular momentum to compute CG coefficients for. This must be large enough to perform the desired number of correlations on the density passed to the compute() method.

  • skip_redundantbool, whether to skip redundant CG combination steps without losing information in the output. Setting to True can save computation time, but does not preserve the norm of the output tensors.

  • cg_backendstr, the backend to use for the CG tensor product. If None, the backend is automatically selected based on the arrays backend.

  • arrays_backendstr, the backend to use for array operations. If None, the backend is automatically selected based on the environment. Possible values are “numpy” and “torch”.

  • dtype – the scalar type to use to store coefficients

  • device – the computational device to use for calculations. This must be "cpu" if array_backend="numpy".

compute(density: TensorMap, angular_cutoff: int | None = None, selected_keys: Labels | None = None) TensorMap

Takes n_correlations of iterative CG tensor products of a density with itself, to generate density auto-correlations of arbitrary correlation order.

\[\rho^{\nu=n_{corr} + 1} = \rho^{\nu=1} \otimes \rho^{\nu=1} \ldots \otimes \rho^{\nu=1}\]

where rho^{nu=1} is the input density of correlation order 1 (body order 2), and rho^{nu=n_{corr} + 1} is the output density of correlation order n_correlations + 1

Before performing any correlations, the properties dimensions of density are modified to carry a “_1” suffix. At each iteration, the dimension names of the copy of the density being correlated are incremented by one each time.

selected_keys can be passed to select the keys to compute on the final iteration. If None, all keys are computed. To limit the maximum angular momenta to compute on intermediate iterations, pass angular_cutoff.

If angular_cutoff and selected_keys are both passed, angular_cutoff is ignored on the final iteration.

Parameters:
  • densityTensorMap, the input density tensor of correlation order 1.

  • angular_cutoffint, the maximum angular momentum to compute output blocks for at each iteration. If selected_keys is passed, this parameter is applied to all intermediate (i.e. prior to the final) iterations. If selected_keys is not passed, this parameter is applied globally to all iterations. If None, all angular momenta are computed.

  • selected_keysLabels, the keys to compute on the final iteration. If None, all keys are computed. Subsets of key dimensions can be passed to compute output blocks that match in these dimensions.

Returns:

TensorMap, the output density auto-correlation tensor of correlation order n_correlations + 1.

forward(density: TensorMap, angular_cutoff: int | None = None, selected_keys: Labels | None = None) TensorMap

Calls the compute() method.

This is intended for torch.nn.Module compatibility, and should be ignored in pure Python mode.

See compute() for a full description of the parameters.

compute_metadata(density: TensorMap, angular_cutoff: int | None = None, selected_keys: Labels | None = None) TensorMap

Returns the metadata-only TensorMap that would be output by the function compute() for the same calculator under the same settings, without performing the actual Clebsch-Gordan tensor products.

See compute() for a full description of the parameters.

featomic.utils.cartesian_to_spherical(tensor: TensorMap, components: List[str], keep_l_in_keys: bool | None = None, remove_blocks_threshold: float | None = 1e-09, cg_backend: str | None = None, cg_coefficients: TensorMap | None = None) TensorMap

Transform a tensor of 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 spherical harmonics Y^M_L. For example, a rank 1 tensor with a single “xyz” component would be represented as a single L=1 spherical harmonic; while a rank 5 tensor using a product basis ℝ^3 ℝ^3 ℝ^3 ℝ^3 ℝ^3 would require multiple blocks up to L=5 spherical harmonics.

A single TensorBlock in the input might correspond to multiple TensorBlock in the output. The output keys will contain all the dimensions of the input keys, plus o3_lambda (indicating the spherical harmonics degree) and o3_sigma (indicating that this block is a proper- or improper tensor with +1 and -1 respectively). If keep_l_in_keys is True or if the input tensor is a tensor of rank 3 or more, the keys will also contain multiple l_{i} and k_{i} dimensions, which indicate which angular momenta have been coupled together in which order to get this block.

components specifies which ones of the components of the input TensorMap should be transformed from cartesian to spherical. All these components will be replaced in the output by a single o3_mu component, corresponding to the spherical harmonics M.

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 below remove_blocks_epsilon. To keep all blocks regardless of their norm, you can set remove_blocks_epsilon=None.

Parameters:
  • tensor – input tensor, using a cartesian product basis

  • components – components of the input tensor to transform into spherical components

  • keep_l_in_keys

    should the output contains the values of angular momenta that were combined together? This defaults to False for rank 1 and 2 tensors, and True for all other tensors.

    Keys named l_{i} correspond to the input components, with l_1 being the last entry in components and l_N the first one. Keys named k_{i} correspond to intermediary spherical components created during the calculation, i.e. a k_{i} used to be o3_lambda.

  • remove_blocks_threshold – 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 None to keep all blocks in the output.

  • cg_backend – Backend to use for Clebsch-Gordan calculations. This can be "python-dense" or "python-sparse" for dense or sparse operations respectively. If None, this is automatically determined.

  • cg_coefficients – Cache containing Clebsch-Gordan coefficients. This is optional except when using this function from TorchScript. The coefficients should be computed with calculate_cg_coefficients(), using the same cg_backend as this function.

Returns:

TensorMap containing spherical components instead of cartesian components.

Low-level functionalities

class featomic.utils.ClebschGordanProduct(*, max_angular: int, cg_backend: str | None = None, keys_filter: Callable[[Labels], List[int]] | None = None, arrays_backend: str | None = None, dtype: dtype | dtype | None = None, device: device | str | None = None)

A general class for computing the Clebsch-Gordan (CG) tensor product between two arbitrary TensorMap.

The constructor computes and stores the CG coefficients. The compute() method computes the CG tensor product between two tensors.

Parameters:
  • max_angularint, the maximum angular momentum to compute CG coefficients for.

  • cg_backendstr, the backend to use for the CG tensor product. If None, the backend is automatically selected based on the arrays backend (“numpy” when importing this class from featomic.utils, and “torch” when importing from featomic.torch.utils).

  • keys_filter – A function to remove more keys from the output. This is applied after any user-provided key_selection in compute(). This function should take one argument keys: Labels, and return the indices of keys to keep.

  • arrays_backendstr, the backend to use for array operations. If None, the backend is automatically selected based on the environment. Possible values are “numpy” and “torch”.

  • dtype – the scalar type to use to store coefficients

  • device – the computational device to use for calculations. This must be "cpu" if array_backend="numpy".

compute(tensor_1: TensorMap, tensor_2: TensorMap, o3_lambda_1_new_name: str, o3_lambda_2_new_name: str, selected_keys: Labels | None = None) TensorMap

Computes the Clebsch-Gordan (CG) tensor product between tensor_1 and tensor_2.

This function assumes the metadata of tensor_1 and tensor_2 has been modified to be compatible for the CG tensor product, according to the following rules:

  • both tensor_1 and tensor_2 must have key dimensions "o3_lambda" and "o3_sigma" as these are used to determine the symmetry of output blocks upon combination;

  • o3_lambda_1_new_name and o3_lambda_2_new_name define the output key dimension names that store the "o3_lambda" values of the blocks combined from tensor_1 and tensor_2 respectively;

  • any other key dimensions that have equivalent names in both tensor_1 and tensor_2 will be matched, such that only blocks that have equal values in these dimensions will be combined;

  • any other named key dimensions that are present in tensor_1 but not in

    tensor_2, and vice versa, will have the full product computed;

  • tensor_1 and tensor_2 must have a single component axis with a single key dimension named "o3_mu". tensor_1 and tensor_2 may have different samples, but all samples in the tensor with fewer samples dimensions must be present in the samples of the other tensor.

A Labels object can be passed in selected_keys to select specific keys to compute. The full set of output keys are computed, then are filtered to match the key dimensions passed in selected_keys.

This parameter can be used to match the output tensor to a given target basis set definition, and/or enhance performance by limiting the combinations computed.

For instance, passing just the "o3_lambda" key dimension with a range of values 0, ..., max_angular can be used to perform angular selection, limiting the maximum angular momentum of the output blocks.

Note that using selected_keys to perform any kind of selection will reduce the information content of the output tensor. This may be important if the returned tensor is used in further CG tensor products.

Parameters:
  • tensor_1 – The first TensorMap, containing data with SO(3) character.

  • tensor_2 – The first TensorMap, containing data with SO(3) character.

  • o3_lambda_1_new_namestr, the name of the output key dimension that stores the "o3_lambda" values of the blocks combined from tensor_1.

  • o3_lambda_2_new_namestr, the name of the output key dimension that stores the "o3_lambda" values of the blocks combined from tensor_2.

  • selected_keysLabels, the keys to compute on the final iteration. If None, all keys are computed. Subsets of key dimensions can be passed to compute output blocks that match in these dimensions.

Returns:

A TensorMap containing the Clebsch-Gordan tensor product of tensor_1 and tensor_2.

forward(tensor_1: TensorMap, tensor_2: TensorMap, o3_lambda_1_new_name: str, o3_lambda_2_new_name: str, selected_keys: Labels | None = None) TensorMap

Calls the ClebschGordanProduct.compute() function.

This is intended for torch.nn.Module compatibility, and should be ignored in pure Python mode.

See compute() for a full description of the parameters.

compute_metadata(tensor_1: TensorMap, tensor_2: TensorMap, o3_lambda_1_new_name: str, o3_lambda_2_new_name: str, selected_keys: Labels | None = None) TensorMap

Returns the metadata-only TensorMap that would be output by the function compute() for the same calculator under the same settings, without performing the actual Clebsch-Gordan tensor products.

See compute() for a full description of the parameters.

featomic.utils.calculate_cg_coefficients(lambda_max: int, cg_backend: str, arrays_backend: str, dtype: dtype | dtype, device: str | device) TensorMap

Calculates the Clebsch-Gordan coefficients for all possible combination of angular momenta up to lambda_max.

The structure of the returned TensorMap depends on whether the backend used to perform CG tensor products uses sparse or dense operations.

Dense:
  • samples: _, i.e. a dummy sample.

  • components: [m1, m2, mu] on separate components axes, where m1 and m2 are the m component values for the two arrays being combined and mu is the m component value for the resulting array.

  • properties: cg_coefficient = [0]

Sparse:
  • samples: (m1, m2, mu), where m1 and m2 are the m component values for the two arrays being combined and mu is the m component value for the resulting array.

  • components: [], i.e. no components axis.

  • properties: cg_coefficient = [0]

Parameters:
  • lambda_max – maximum angular momentum value to compute CG coefficients for.

  • cg_backend – whether to use "python-spare" or "python-dense" format for storing the CG coefficients.

  • arrays_backend – whether to use "numpy" or "torch" arrays to store the coefficients.

  • dtype – the scalar type to use to store coefficients

  • device – the computational device to use for calculations. This must be "cpu" if array_backend="numpy".

Returns:

TensorMap of the Clebsch-Gordan coefficients.