PyTorch API#

The main class for computing spherical harmonics using a torch-compatible framework follows the same syntax as the Python version sphericart.SphericalHarmonics. Depending on the device the tensor is stored on, and its dtype, the calculations will be performed using 32- or 64- bits floating point arythmetics, and using the CPU or CUDA implementation.

class sphericart.torch.SphericalHarmonics(l_max, normalized=False, backward_second_derivatives=False)#

Spherical harmonics calculator, up to degree l_max.

By default, this class computes a non-normalized form of the real spherical harmonics, i.e. \(r^l Y^l_m\). These scaled spherical harmonics are homogeneous polynomials in the Cartesian coordinates of the input points. normalized=True can be set to compute the normalized spherical harmonics \(Y^l_m\), which are instead homogeneous polynomials of x/r, y/r, z/r.

This class can be used similarly to sphericart.SphericalHarmonics (its Python/NumPy counterpart), and it allows to return explicit forward gradients and/or Hessians. For example:

>>> import torch
>>> import sphericart.torch as sct
>>> sh = sct.SphericalHarmonics(l_max=8, normalized=False)
>>> xyz = torch.rand(size=(10,3))
>>> sh_values, sh_grads = sh.compute_with_gradients(xyz)
>>> sh_grads.shape
torch.Size([10, 3, 81])

Alternatively, if compute() is called, the outputs support single and double backpropagation.

>>> xyz = xyz.detach().clone().requires_grad_()
>>> sh = sct.SphericalHarmonics(l_max=8, normalized=False)
>>> sh_values = sh.compute(xyz)
>>> sh_values.sum().backward()
>>> torch.allclose(xyz.grad, sh_grads.sum(axis=-1))
True

By default, only single backpropagation with respect to xyz is enabled (this includes mixed second derivatives where xyz appears as only one of the differentiation steps). To activate support for double backpropagation with respect to xyz, please set backward_second_derivatives=True at class creation. Warning: if backward_second_derivatives is not set to True and double differentiation with respect to xyz is requested, the results may be incorrect and no warnings will be displayed. This is necessary to provide optimal performance for both use cases.

This class supports TorchScript.

Parameters:

l_max – the maximum degree of the spherical harmonics to be calculated
normalized – whether to normalize the spherical harmonics (default: False)
backward_second_derivatives – if this parameter is set to True, second derivatives of the spherical harmonics are calculated and stored during forward calls to compute (provided that xyz.requires_grad is True), making it possible to perform double reverse-mode differentiation with respect to xyz. If False, only the first derivatives will be computed and only a single reverse-mode differentiation step will be possible with respect to xyz.

Returns:

a calculator, in the form of a SphericalHarmonics object

compute(xyz: Tensor) → Tensor#

Calculates the spherical harmonics for a set of 3D points.

The coordinates should be stored in the xyz array. If xyz has requires_grad = True it stores the forward derivatives which are then used in the backward pass. The type of the entries of xyz determines the precision used, and the device the tensor is stored on determines whether the CPU or CUDA implementation is used for the calculation backend. It always supports single reverse-mode differentiation, as well as double reverse-mode differentiation if backward_second_derivatives was set to True during class creation.

Parameters:: xyz – The Cartesian coordinates of the 3D points, as a torch.Tensor with shape (n_samples, 3).
Returns:: A tensor of shape (n_samples, (l_max+1)**2) containing all the spherical harmonics up to degree l_max in lexicographic order. For example, if l_max = 2, The last axis will correspond to spherical harmonics with (l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, 1), (2, 2), in this order.

compute_with_gradients(xyz: Tensor) → Tuple[Tensor, Tensor]#

Calculates the spherical harmonics for a set of 3D points, and also returns the forward-mode derivatives.

The coordinates should be stored in the xyz array. The type of the entries of xyz determines the precision used, and the device the tensor is stored on determines whether the CPU or CUDA implementation is used for the calculation backend. Reverse-mode differentiation is not supported for this function.

Parameters:

xyz – The Cartesian coordinates of the 3D points, as a torch.Tensor with shape (n_samples, 3).

Returns:

A tuple that contains:

A (n_samples, (l_max+1)**2) tensor containing all the spherical harmonics up to degree l_max in lexicographic order. For example, if l_max = 2, The last axis will correspond to spherical harmonics with (l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, 1), (2, 2), in this order.
A tensor of shape (n_samples, 3, (l_max+1)**2) containing all the spherical harmonics’ derivatives up to degree l_max. The last axis is organized in the same way as in the spherical harmonics return array, while the second-to-last axis refers to derivatives in the the x, y, and z directions, respectively.

compute_with_hessians(xyz: Tensor) → Tuple[Tensor, Tensor, Tensor]#

Calculates the spherical harmonics for a set of 3D points, and also returns the forward derivatives and second derivatives.

Parameters:

xyz – The Cartesian coordinates of the 3D points, as a torch.Tensor with shape (n_samples, 3).

Returns:

A tuple that contains:

A (n_samples, (l_max+1)**2) tensor containing all the spherical harmonics up to degree l_max in lexicographic order. For example, if l_max = 2, The last axis will correspond to spherical harmonics with (l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, 1), (2, 2), in this order.
A tensor of shape (n_samples, 3, (l_max+1)**2) containing all the spherical harmonics’ derivatives up to degree l_max. The last axis is organized in the same way as in the spherical harmonics return array, while the second-to-last axis refers to derivatives in the the x, y, and z directions, respectively.
A tensor of shape (n_samples, 3, 3, (l_max+1)**2) containing all the spherical harmonics’ second derivatives up to degree l_max. The last axis is organized in the same way as in the spherical harmonics return array, while the two intermediate axes represent the hessian dimensions.

omp_num_threads()#: Returns the number of threads available for calculations on the CPU.

l_max()#: Returns the maximum angular momentum setting for this calculator.

normalized()#: Returns normalization setting for this calculator.

The implementation also contains a couple of utility functions to facilitate the integration of sphericart into code using `e3nn.

sphericart.torch.e3nn_spherical_harmonics(l_list: List[int] | int, x: Tensor, normalize: bool | None = False, normalization: str | None = 'integral') → Tensor#

Computes spherical harmonics with an interface similar to the e3nn package.

Provides an interface that is similar to e3nn.o3.spherical_harmonics() but uses SphericalHarmonics for the actual calculation. Uses the same ordering of the [x,y,z] axes, and supports the same options for input and harmonics normalization as e3nn. However, it does not support defining the irreps through a e3nn.o3._irreps.Irreps or a string specification, but just as a single integer or a list of integers.

Parameters:

l_list – Either a single integer or a list of integers specifying which \(Y^l_m\) should be computed. All values up to the maximum l value are computed, so this may be inefficient for use cases requiring a single, or few, angular momentum channels.
x – A torch.Tensor containing the coordinates, in the same format expected by the e3nn function.
normalize – Flag specifying whether the input positions should be normalized, or whether the function should compute scaled :math:` ilde{Y}^l_m`
normalization – String that can be “integral”, “norm”, “component”, that controls a further scaling of the \(Y_m^l\). See the documentation of e3nn.o3.spherical_harmonics() for a detailed explanation of the different conventions.

sphericart.torch.patch_e3nn(e3nn_module: module) → None#

Patches the e3nn module so that sphericart_torch.e3nn_spherical_harmonics() is called in lieu of the built-in function.

Parameters:: e3nn_module – The alias that has been chosen for the e3nn module, usually just e3nn.

sphericart.torch.unpatch_e3nn(e3nn_module: module) → None#: Restore the original spherical_harmonics function in the e3nn module.