Clebsch-Gordan products

class featomic.clebsch_gordan.EquivariantPowerSpectrum(calculator_1: CalculatorBase, calculator_2: CalculatorBase | None = None, neighbor_types: List[int] | None = None, *, dtype: dtype | dtype | None = None, device: device | str | None = None)

Computes a general equivariant power spectrum descriptor of two calculators.

If only calculator_1 is provided, the power spectrum is computed as the density auto-correlation of the density produced by the first calculator. If calculator_2 is also provided, the power spectrum is computed as the density cross-correlation of the densities produced by the two calculators.

Example

As an example we calculate the equivariant power spectrum for a short range (sr) spherical expansion and a long-range (lr) LODE spherical expansion for a NaCl crystal.

>>> import featomic
>>> import ase

Construct the NaCl crystal

>>> atoms = ase.Atoms(
...     symbols="NaCl",
...     positions=[[0, 0, 0], [0.5, 0.5, 0.5]],
...     pbc=True,
...     cell=[1, 1, 1],
... )

Define the hyper parameters for the short-range spherical expansion

>>> sr_hypers = {
...     "cutoff": {
...         "radius": 1.0,
...         "smoothing": {"type": "ShiftedCosine", "width": 0.5},
...     },
...     "density": {
...         "type": "Gaussian",
...         "width": 0.3,
...     },
...     "basis": {
...         "type": "TensorProduct",
...         "max_angular": 2,
...         "radial": {"type": "Gto", "max_radial": 5},
...     },
... }

Define the hyper parameters for the long-range LODE spherical expansion from the hyper parameters of the short-range spherical expansion

>>> lr_hypers = {
...     "density": {
...         "type": "SmearedPowerLaw",
...         "smearing": 0.3,
...         "exponent": 1,
...     },
...     "basis": {
...         "type": "TensorProduct",
...         "max_angular": 2,
...         "radial": {"type": "Gto", "max_radial": 3, "radius": 1.0},
...     },
... }

Construct the calculators

>>> sr_calculator = featomic.SphericalExpansion(**sr_hypers)
>>> lr_calculator = featomic.LodeSphericalExpansion(**lr_hypers)

Construct the power spectrum calculators and compute the spherical expansion

>>> calculator = featomic.clebsch_gordan.EquivariantPowerSpectrum(
...     sr_calculator, lr_calculator
... )
>>> power_spectrum = calculator.compute(atoms, neighbors_to_properties=True)

The resulting equivariants are stored as metatensor.TensorMap as for any other calculator. The keys contain the symmetry information:

>>> power_spectrum.keys
Labels(
    o3_lambda  o3_sigma  center_type
        0         1          11
        1         1          11
        2         1          11
        1         -1         11
        2         -1         11
        3         1          11
        3         -1         11
        4         1          11
        0         1          17
        1         1          17
        2         1          17
        1         -1         17
        2         -1         17
        3         1          17
        3         -1         17
        4         1          17
)

The block properties contain the angular order of the combined blocks (“l_1”, “l_2”), along with the neighbor types (“neighbor_1_type”, “neighbor_2_type”) and radial channel indices.

>>> power_spectrum[0].properties.names
['l_1', 'l_2', 'neighbor_1_type', 'n_1', 'neighbor_2_type', 'n_2']

See also

Faster power spectrum calculator specifically for invariant descriptors can be found at featomic.SoapPowerSpectrum and featomic.clebsch_gordan.PowerSpectrum.

Constructs the equivariant power spectrum calculator.

param calculator_1:

first calculator that computes a density descriptor, either a featomic.SphericalExpansion or featomic.LodeSphericalExpansion.

param calculator_2:

optional second calculator that computes a density descriptor, either a featomic.SphericalExpansion or featomic.LodeSphericalExpansion. If None, the equivariant power spectrum is computed as the auto-correlation of the first calculator. Defaults to None.

param neighbor_types:

List of "neighbor_type" to use in the properties of the output. This option might be useful when running the calculation on subset of a whole dataset and trying to join along the sample dimension after the calculation. If None, blocks are filled with "neighbor_type" found in the systems. This parameter is only used if neighbors_to_properties=True is passed to the compute() method.

param dtype:

the scalar type to use to store coefficients

param device:

the computational device to use for calculations.

property name

Name of this calculator.

compute(systems: IntoSystem | List[IntoSystem], selected_keys: Labels | None = None, neighbors_to_properties: bool = False) TensorMap

Computes an equivariant power spectrum, also called “Lambda-SOAP” when doing a self-correlation of the SOAP density.

First computes a SphericalExpansion density descriptor of body order 2.

Before performing the Clebsch-Gordan tensor product, the spherical expansion density can be densified by moving the key dimension “neighbor_type” to the block properties. This is controlled by the neighbors_to_properties parameter. Depending on the specific systems descriptors are being computed for, the sparsity or density of the density can affect the computational cost of the Clebsch-Gordan tensor product.

If neighbors_to_properties=True and neighbor_types have been passed to the constructor, property dimensions are created for all of these global atom types when moving the key dimension to properties. This ensures that the output properties dimension is of consistent size across all systems passed in systems.

Finally a single Clebsch-Gordan tensor product is taken to produce a body order 3 equivariant power spectrum.

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

  • neighbors_to_propertiesbool, if true, densifies the spherical expansion by moving key dimension “neighbor_type” to properties prior to performing the Clebsch Gordan product step. Defaults to false.

Returns:

TensorMap, the output equivariant power spectrum.

forward(systems: IntoSystem | List[IntoSystem], selected_keys: Labels | None = None, neighbors_to_properties: bool = False) 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(systems: IntoSystem | List[IntoSystem], selected_keys: Labels | None = None, neighbors_to_properties: bool = False) 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 in the second step.

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

class featomic.clebsch_gordan.PowerSpectrum(calculator_1: CalculatorBase, calculator_2: CalculatorBase | None = None, neighbor_types: List[int] | None = None)

General power spectrum of one or of two calculators.

If calculator_2 is provided, the invariants \(p_{nl}\) are generated by taking quadratic combinations of calculator_1’s spherical expansion \(\rho_{nlm}\) and calculator_2’s spherical expansion \(\nu_{nlm}\) according to Bartók et. al.

\[p_{nl} = \rho_{nlm}^\dagger \cdot \nu_{nlm}\]

where we use the Einstein summation convention. If gradients are present the invariants of those are constructed as

\[\nabla p_{nl} = \nabla \rho_{nlm}^\dagger \cdot \nu_{nlm} + \rho_{nlm}^\dagger \cdot \nabla \nu_{nlm}\]

Note

Currently only supports gradients with respect to positions.

If calculator_2=None invariants are generated by combining the coefficients of the spherical expansion of calculator_1. The spherical expansions given as input can only be featomic.SphericalExpansion or featomic.LodeSphericalExpansion.

Example

As an example we calculate the power spectrum for a short range (sr) spherical expansion and a long-range (lr) LODE spherical expansion for a NaCl crystal.

>>> import featomic
>>> import ase

Construct the NaCl crystal

>>> atoms = ase.Atoms(
...     symbols="NaCl",
...     positions=[[0, 0, 0], [0.5, 0.5, 0.5]],
...     pbc=True,
...     cell=[1, 1, 1],
... )

Define the hyper parameters for the short-range spherical expansion

>>> sr_hypers = {
...     "cutoff": {
...         "radius": 1.0,
...         "smoothing": {"type": "ShiftedCosine", "width": 0.5},
...     },
...     "density": {
...         "type": "Gaussian",
...         "width": 0.3,
...     },
...     "basis": {
...         "type": "TensorProduct",
...         "max_angular": 2,
...         "radial": {"type": "Gto", "max_radial": 5},
...     },
... }

Define the hyper parameters for the long-range LODE spherical expansion from the hyper parameters of the short-range spherical expansion

>>> lr_hypers = {
...     "density": {
...         "type": "SmearedPowerLaw",
...         "smearing": 0.3,
...         "exponent": 1,
...     },
...     "basis": {
...         "type": "TensorProduct",
...         "max_angular": 2,
...         "radial": {"type": "Gto", "max_radial": 5, "radius": 1.0},
...     },
... }

Construct the calculators

>>> sr_calculator = featomic.SphericalExpansion(**sr_hypers)
>>> lr_calculator = featomic.LodeSphericalExpansion(**lr_hypers)

Construct the power spectrum calculators and compute the spherical expansion

>>> calculator = featomic.clebsch_gordan.PowerSpectrum(sr_calculator, lr_calculator)
>>> power_spectrum = calculator.compute(atoms)

The resulting invariants are stored as metatensor.TensorMap as for any other calculator

>>> power_spectrum.keys
Labels(
    center_type
        11
        17
)
>>> power_spectrum[0]
TensorBlock
    samples (1): ['system', 'atom']
    components (): []
    properties (432): ['l', 'neighbor_1_type', 'n_1', 'neighbor_2_type', 'n_2']
    gradients: None

See also

If you are interested in the SOAP power spectrum you can the use the faster featomic.SoapPowerSpectrum. For an equivariant version of this calculator, computing the power spectrum for covariant as well as invariant blocks, see featomic.EquivariantPowerSpectrum.

Constructs the power spectrum calculator.

param calculator_1:

first calculator that computes a density descriptor, either a featomic.SphericalExpansion or featomic.LodeSphericalExpansion.

param calculator_2:

optional second calculator that computes a density descriptor, either a featomic.SphericalExpansion or featomic.LodeSphericalExpansion. If None, the power spectrum is computed as the auto-correlation of the first calculator. Defaults to None.

param neighbor_types:

List of "neighbor_type" to use in the properties of the output. This option might be useful when running the calculation on subset of a whole dataset and trying to join along the sample dimension after the calculation. If None, blocks are filled with "neighbor_type" found in the systems.

property name

Name of this calculator.

compute(systems: IntoSystem | List[IntoSystem], gradients: List[str] | None = None, use_native_system: bool = True) TensorMap

Runs a calculation with this calculator on the given systems.

See featomic.calculators.CalculatorBase.compute() for details on the parameters.

Raises:

NotImplementedError – If a spherical expansions contains a gradient with respect to an unknwon parameter.

forward(systems: IntoSystem | List[IntoSystem], gradients: List[str] | None = None, use_native_system: bool = True) TensorMap

Calls the PowerSpectrum.compute() function.

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

class featomic.clebsch_gordan.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.clebsch_gordan.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.clebsch_gordan.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.clebsch_gordan.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.