Fast evaluation of spherical harmonics with sphericart

The spherical harmonics $Ylm$ are basis functions for the irreducible representations of the SO(3) group,¹ which makes them a key tool for understanding physical phenomena that exhibit rotational symmetries and designing practical algorithms to model them on a computer. Examples include the distribution of charge in atoms,² the behavior of gravitational³ and magnetic^4,5 fields, and the propagation of light⁶ and sound⁷ in the atmosphere. Moreover, they provide a complete set of smooth, orthogonal functions defined on the surface of a sphere, and in this sense, they are widely used in many fields, including computer graphics,^8,9 quantum^10–13 and physical^14–16 chemistry, and signal processing.^17,18 More recently, they have become an essential tool in the context of geometric deep learning,^19,20 as a structural descriptor needed in SO(3)-, O(3)-, and E(3)-equivariant machine learning models^21–24 and more specifically in the construction of symmetry-adapted descriptors of atomic structures in chemical machine learning.^25–29 Derivatives of spherical harmonics are also very often used in these applications. An example is the calculation of forces on a configuration of atoms by an E(3)-invariant machine learning model.³⁰ Likewise, spherical harmonics gradients have been used in computer graphics for mid-range illumination, irradiance volumes, and optimization methods.⁹ 

In contrast, our algorithm aims at calculating $Ỹlm=rlYlm$⁠, where $r=x2+y2+z2$⁠. These are the so-called solid harmonics, which consist of simple homogeneous polynomials of the Cartesian coordinates. This choice avoids the need to normalize r and leads to simple and stable iterations to compute $Ỹlm$ as well as very compact expressions for their derivatives with respect to the Cartesian coordinates, which re-use the same factors needed to evaluate $Ỹlm$⁠. In most applications that require spherical harmonics in Cartesian coordinates, the radial direction is dealt with by a separate expansion [cf. Eq. (1)], and the r^l factor that is included in the scaled $Ỹlm$ can be compensated for at little to no additional cost by a corresponding r^−l factor in the radial term. If, instead, spherical harmonics are needed in their conventional version, they can be recovered easily from the radially scaled (or solid) kind.

Most computational applications require the evaluation of all the spherical harmonics, and possibly their derivatives, for the values of l up to a maximum degree l_max. In many cases, the spherical harmonics have to be computed for many points simultaneously, e.g., for the interatomic separation vectors of all neighbors of a selected atom.

Even though we recommend using the scaled spherical harmonics in applications, accounting for the fact that accompanying radial terms have to compensate for the scaling, we also provide an implementation of “normalized” spherical harmonics by simply evaluating $Ỹlm(x/r,y/r,z/r)$ and applying the chain rule to the derivatives. These additional operations typically result in an overhead of 5%–10% for the calculation of the spherical harmonics and their derivatives.

Our general routines take a list of Cartesian coordinates and evaluate the spherical harmonics and, optionally, their derivatives and/or second derivatives with respect to the inputs. We use C++ for the implementation using templates to exploit compile-time knowledge of the maximum angular momentum, the need to compute derivatives, and so on. We then define a pure C API covering the typical use cases on top of this C++ API. Since almost all programming languages have a way to call C functions, this enables using our code from most languages used in a scientific context. We also provide a Python package with a high-level interface to the C library.

Furthermore, we use the C++ API to provide an implementation compatible with PyTorch,³⁶ using a custom backward function to compute gradients using the derivatives evaluated in the forward pass and making sure the code is compatible with TorchScript, allowing models to be used without a Python interpreter. Finally, we implement a GPU-accelerated version of the PyTorch code for NVIDIA GPUs, using the CUDA language. Table I compares the features of sphericart to those of other spherical harmonics implementations.^41–50 

We apply a number of trivial (and a few less obvious) optimizations. For example, we pre-compute the factors $Flm$ [Eq. (3)] to minimize the number of operations in the inner loop of the iterative algorithm discussed above. In order to further accelerate the evaluation of low-l Cartesian harmonics, we use a computer algebra system to derive expressions that evaluate the full $Ỹlm$ and their derivatives with hard-coded expressions using the smallest possible number of multiplications. As presented in Table II and as noted in previous implementations,^24,33 there is considerable scope for optimization by using ad hoc expressions. However, the speedup is less remarkable when including the calculation of the derivatives through Eq. (10), which re-uses quantities that have already been computed when evaluating $Ỹlm$⁠.

As a compromise between the convenience of a function that works for arbitrary l_max and the efficiency of optimized expressions, we provide an interface for the hard-coded implementation up to l_max = 6. By default, our implementation applies the hard-coded version for small l and then switches to the general expression, making use of the alternative recursion (8) to avoid the computation of the Legendre polynomials $Qlm$⁠. for small l. All our functions can be applied to many 3D points at once, and they are trivially parallelized over the sample index using OpenMP.³⁷ Despite the simplicity of our design, we achieve good parallel scaling, particularly when using a relatively large l_max so that there is a substantial amount of computation for each thread, as illustrated in Fig. 1. The curves show some degree of irregular behavior, suggesting that there might be room for further optimization. As presented in Table III, for large l_max and numbers of data points, there is a substantial performance gain by using single-precision floating-point arithmetics. However, for smaller amounts of computation, there is a very small advantage, and in some corner cases (l_max = 8, n_samples = 10 000, and n_threads = 16), the single-precision version can be slower than that using 64-bit floating-point values, which further suggests that it may be possible to optimize memory-access patterns to improve parallel performance.

We adapt the implementation to a custom CUDA kernel, which allows efficient execution on GPUs. Similarly to the CPU implementation, we use hard-coded spherical harmonics and derivatives up to a reduced hard-coded l_max = 3 to lower shared memory requirements, and we use the general expression for the remaining terms. We parallelize the computation with a two-dimensional thread-block of 16 × 8 threads using a grid dimension of (n_samples/16) blocks. The first dimension in the thread-block parallelizes over samples, while the first thread in the second dimension is responsible for performing the computation. The remaining threads perform coalesced writing of the temporary buffers to global memory. We store the intermediary work buffers for the spherical harmonics and derivatives in fast shared memory, which defaults to 48 KB for most NVIDIA cards. Memory accesses are designed such that each sample-dimension thread is mapped to a shared memory bank contiguously, eliminating possible bank conflicts. The CUDA wrapper automatically re-allocates the necessary amount of shared memory beyond the default 48 KB on the first compute call or attempts to reduce the number of threads launched in the CUDA kernel if the required allocation is too large. As a consequence, for accurate timings, we call the forward step once to perform the initialization and then measure timings thereafter. We note that using the recursion (7) would require storing temporary values for all l ≤ l_max, so that the default 48 KB shared memory allocation would be filled for l_max ≈ 8 when computing derivatives in 64-bit floating-point format. Using the alternative recursion (8) allows us to evaluate one l channel at a time, so values up to l_max ≈ 30 can be computed without adjusting the shared memory allocation or reducing the number of sample-dimension threads. For GPUs that support more than 48 KB shared memory per block, for example, the A100 (164 KB), the recursion (8) allows us to evaluate values up to l_max ≈ 82.

Table IV presents timings for GPU-accelerated computations on an A100 SXM4 (80 GB) card. In general, using single-precision floating-point arithmetic results in half the computation time than double-precision arithmetic. We note that the GPU per-sample timings reduce significantly upon an increasing number of samples, as the GPU is not fully saturated for n_samples = 10 000. For example, for n_samples = 100 000, the per-sample timings reduce by a factor of 1.5–4. Although the parallel CPU implementation is faster than the GPU for l_max ≤ 4 when using n_samples = 10 000, this trend reverses for higher values of l_max. For example, the 64-bit floating-point computations with n_samples = 10 000 and l_max = 16 show that the GPU outperforms the CPU by a factor of 8.

We note in closing that, even though precise timings depend on the hardware and the maximum angular momentum considered, sphericart is between 10 and 40 times faster than the spherical harmonics implementation in e3nn,³⁸ a widely used library for equivariant neural networks. We provide wrappers that use the same conventions as e3nn, to simplify integration of our implementation in existing models and to accelerate computational frameworks that are limited by the evaluation of $Ylm$⁠.

Spherical harmonics are ubiquitous in the computational sciences, and their efficient calculation has become even more important in light of the widespread adoption of equivariant machine-learning models in chemistry. Reformulating the calculation of spherical harmonics in a scaled form that corresponds to real-valued polynomials of the Cartesian coordinates provides simple expressions for their recursive evaluation. Derivatives can be obtained with little overhead, re-using the same factors that enter the definition of the scaled $Ỹlm$⁠, and without numerical instabilities. The conventional form of the spherical harmonics can be recovered easily by computing $Ỹlm$ at (x/r, y/r, z/r) and applying the corresponding correction to the derivatives. In many applications, this normalization may be unnecessary, as the scaling can be incorporated (explicitly or through regression) into additional radial terms. We provide an efficient implementation as a C++ library, complete with a C API, Python, and PyTorch bindings, that tackles the most common use case in scientific computing and geometric machine learning, i.e., the evaluation of real-valued spherical harmonics for all angular momentum channels up to a given cutoff l_max and possibly for many 3D points at once. The function call is parallelized over the sample direction, and it uses more efficient hard-coded expressions for low-l terms. Future efforts will focus on the extension of the software library to different programming languages and frameworks as well as on further optimization of new hardware platforms and improved parallelism.

The Authors would like to thank Kevin Kazuki Huguenin-Dumittan for useful discussions, Philip Loche for help with the sphericart library, and Prashanth Kanduri for help with the continuous integration framework. M.C. acknowledges funding from the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation program (Grant Agreement No. 101001890-FIAMMA). GF acknowledges support from the Platform for Advanced Scientific Computing (PASC). M.C. and F.B. acknowledge support from the NCCR MARVEL, funded by the Swiss National Science Foundation (Grant No. 182892).

The authors have no conflicts to disclose.

Filippo Bigi and Michele Ceriotti conceived the project and performed analytical derivations. All authors contributed to software development and the writing of the paper.

Filippo Bigi: Conceptualization (equal); Formal analysis (equal); Software (equal); Writing – original draft (equal); Writing – review & editing (equal). Guillaume Fraux: Software (equal); Writing – review & editing (supporting). Nicholas J. Browning: Software (equal); Writing – review & editing (supporting). Michele Ceriotti: Conceptualization (equal); Formal analysis (equal); Methodology (equal); Resources (lead); Software (equal); Supervision (lead); Writing – original draft (equal); Writing – review & editing (equal).

The sphericart library can be freely downloaded under the Apache License version 2.0 from the public git repository https://github.com/lab-cosmo/sphericart. A Python package that provides a convenient class to compute spherical harmonics and derivatives from a NumPy³⁹ or a PyTorch³⁶ array of 3D coordinates is available on PyPI at https://pypi.org/project/sphericart/. The timing data can be generated using the benchmarking code that is included in the software distribution.