PyTorch API

The classes for computing spherical harmonics with torch follow a similar syntax to their Python counterparts sphericart.SphericalHarmonics and sphericart.SolidHarmonics, while inheriting from torch.nn.Module. 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: int)

Spherical harmonics calculator, which computes the real spherical harmonics \(Y^m_l\) up to degree l_max. The calculated spherical harmonics are consistent with the definition of real spherical harmonics from Wikipedia.

This class can be used similarly to sphericart.SphericalHarmonics (its Python/NumPy counterpart).

>>> import torch
>>> import sphericart.torch
>>> xyz = torch.rand(size=(10, 3))
>>> xyz = xyz.detach().clone().requires_grad_()
>>> sh = sphericart.torch.SphericalHarmonics(l_max=8)
>>> sh_values = sh(xyz)  # or sh.compute(xyz)
>>> sh_values.sum().backward()
>>> torch.allclose(xyz.grad, sh_grads.sum(axis=-1))
True

Reverse-mode differentiation with respect to xyz supports single and double backpropagation.

Alternatively, the class allows to return explicit forward gradients and/or Hessians of the spherical harmonics. For example:

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

This class fully supports TorchScript and torch.compile. Some limitations are present when using backward differentiation with torch.func.

Parameters:

l_max – the maximum degree of the spherical harmonics to be calculated

Returns:

a calculator, in the form of a SphericalHarmonics object

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(xyz: Tensor) → Tensor

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

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. It supports single and double reverse-mode differentiation with respect to xyz.

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(xyz: Tensor) → Tensor

Equivalent to forward

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.

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.

  • 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.

class sphericart.torch.SolidHarmonics(l_max: int)

Solid harmonics calculator, up to degree l_max.

This class computes the solid harmonics, a non-normalized form of the real spherical harmonics, i.e. \(r^lY^m_l\). These scaled spherical harmonics are polynomials in the Cartesian coordinates of the input points.

The usage of this class is identical to sphericart.SphericalHarmonics.

Parameters:

l_max – the maximum degree of the spherical harmonics to be calculated

Returns:

a calculator, in the form of a SolidHarmonics object

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(xyz: Tensor) → Tensor

See SphericalHarmonics.forward()

compute(xyz: Tensor) → Tensor

Equivalent to forward

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

See SphericalHarmonics.compute_with_gradients()

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

See SphericalHarmonics.compute_with_hessians()

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.

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^m_l\) 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 (resulting in the computation of the spherical harmonics \(Y^m_l\)), or whether the function should compute the solid harmonics \(r^lY^m_l\).

  • 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: ModuleType) → 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: ModuleType) → None

Restore the original spherical_harmonics function in the e3nn module.