From 20b51b2cd8e77750929d40f995636976f5ca08bf Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Wed, 28 Aug 2024 17:12:36 +0200 Subject: [PATCH 01/12] Index, installation, APIs up to C --- docs/src/api.rst | 15 +- docs/src/index.rst | 4 +- docs/src/installation.rst | 9 +- sphericart/include/sphericart.h | 227 +++++++++--------------------- sphericart/include/sphericart.hpp | 27 +++- 5 files changed, 101 insertions(+), 181 deletions(-) diff --git a/docs/src/api.rst b/docs/src/api.rst index ba2c2e961..cd1c40a05 100644 --- a/docs/src/api.rst +++ b/docs/src/api.rst @@ -1,19 +1,10 @@ API documentation ================= -The core implementation of ``sphericart`` is written in C++. It relies on templates and -C++17 features such as ``if constexpr`` to reduce the runtime overhead of implementing -different normalization styles, and providing both a version with and without derivatives. - -Spherical harmonics and their derivatives are computed with optimized hard-coded expressions -for low values of the principal angular momentum number :math:`l`, then switch to an efficient -recursive evaluation. The API involves initializing a calculator that allocates buffer space -and computes some constant factors, and then using it to compute :math:`Y_l^m` (and possibly its -first and/or second derivatives) for one or more points in 3D space. - -This core C++ library is then made available to different environments through a C API. +The core implementation of ``sphericart`` is written in C++ and CUDA. This core library is +then also made available to different environments (C, Python, PyTorch, JAX). This section contains a description of the interface of the ``sphericart`` library for the -different languages it supports. +different languages and frameworks it supports. .. toctree:: :maxdepth: 1 diff --git a/docs/src/index.rst b/docs/src/index.rst index e56c03ef3..fad66e1a7 100644 --- a/docs/src/index.rst +++ b/docs/src/index.rst @@ -15,8 +15,8 @@ The theory behind this efficient implementation is detailed in this `paper `_. The core library is implemented in C++ (with OpenMP parallelism) and CUDA. -It provides APIs for C, Python (NumPy), PyTorch and JAX. The torch and JAX -implementations provide fast spherical harmonics evaluations on GPUs. +It also provides APIs for C, Python (NumPy), PyTorch and JAX. The torch and JAX +implementations provide fast spherical harmonics on GPUs. A native Julia package is also available. diff --git a/docs/src/installation.rst b/docs/src/installation.rst index b36c84317..8df3f4d6b 100644 --- a/docs/src/installation.rst +++ b/docs/src/installation.rst @@ -19,8 +19,8 @@ This basic package makes use of NumPy. A PyTorch-based implementation can be ins This pre-built version available on PyPI sacrifices some performance to ensure it can run on all systems, and it does not include GPU support. -If you need an extra 5-10% of performance or you want to evaluate the spherical harmonics on GPUs, -you should build the code from source: +If you need an extra 5-10% of performance, you want to evaluate the spherical harmonics on GPUs, +and/or you want to use it in JAX, you should build the code from source: .. code-block:: bash @@ -35,8 +35,9 @@ you should build the code from source: # torch bindings (CPU-only) pip install --extra-index-url https://download.pytorch.org/whl/cpu .[torch] -Before installing the JAX version of ``sphericart``, you should already have the JAX -library installed according to the official JAX installation instructions. +Before installing the JAX version of ``sphericart``, make sure you already have the JAX +library installed according to the official JAX installation instructions at +. Julia package diff --git a/sphericart/include/sphericart.h b/sphericart/include/sphericart.h index f7c4c257e..149759907 100644 --- a/sphericart/include/sphericart.h +++ b/sphericart/include/sphericart.h @@ -1,5 +1,16 @@ /** \file sphericart.h - * Defines the C API for `sphericart`. + * Defines the C API for `sphericart`. Similar types and functions are + * available whether one is using the `double` or `float` data type, and whether + * one is using spherical or solid harmonics calculators. + * Types and functions for the `float` data type contain `_f` in their name, + * while those for the `double` data type do not. + * Similarly, types and functions for spherical harmonics calculations contain + * `_spherical_harmonics` in their name, and they calculate the + * real spherical harmonics :math:`Y^m_l: as defined on Wikipedia, + * which are homogeneous polynomials of (x/r, y/r, z/r). In contrast, types + * and functions for solid harmonics calculations contain `_solid_harmonics` in + * their name, and thay calculate the same polynomials but as a function of the + * Cartesian coordinates (x, y, z), or, equivalently, :math:`r^l Y^m_l`. */ #ifndef SPHERICART_H @@ -109,30 +120,29 @@ SPHERICART_EXPORT void sphericart_spherical_harmonics_delete_f( ); /** - * This function calculates the spherical harmonics and, optionally, their - * derivatives for an array of 3D points. + * This function calculates the spherical harmonics for an array of 3D points. * * @param calculator A pointer to a `sphericart_spherical_harmonics_calculator_t` struct * that holds prefactors and options to compute the spherical - * harmonics. + * harmonics. * @param xyz An array of size `n_samples x 3`. It contains the Cartesian * coordinates of the 3D points for which the spherical harmonics are - * to be computed, organized along two dimensions. The outer dimension is + * to be computed, organized along two dimensions. The outer dimension is * `n_samples` long, accounting for different samples, while the inner * dimension has size 3 and it represents the x, y, and z coordinates * respectively. * @param xyz_length size of the xyz allocation, i.e, `3 * n_samples` * @param sph pointer to the first element of an array containing `n_samples - * * (l_max + 1) * (l_max + 1)` elements. On exit, this array will contain + * * (l_max + 1) * (l_max + 1)` elements. On exit, this array will contain * the spherical harmonics organized along two dimensions. The leading * dimension is `n_samples` long and it represents the different - * samples, while the inner dimension size is `(l_max + 1) * (l_max + 1)` - * long and it contains the spherical harmonics. These are laid out in + * samples, while the inner dimension size is `(l_max + 1) * (l_max + 1)` + * long and it contains the spherical harmonics. These are laid out in * lexicographic order. For example, if `l_max=2`, it will contain - * `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, + * `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, * 1), (2, 2)`, in this order. * @param sph_length size of the sph allocation, should be `n_samples * - * (l_max + 1) * (l_max + 1)` + * (l_max + 1) * (l_max + 1)` */ SPHERICART_EXPORT void sphericart_spherical_harmonics_compute_array( sphericart_spherical_harmonics_calculator_t* calculator, @@ -148,37 +158,37 @@ SPHERICART_EXPORT void sphericart_spherical_harmonics_compute_array( * * @param calculator A pointer to a `sphericart_spherical_harmonics_calculator_t` struct * that holds prefactors and options to compute the spherical - * harmonics. + * harmonics. * @param xyz An array of size `n_samples x 3`. It contains the Cartesian * coordinates of the 3D points for which the spherical harmonics are - * to be computed, organized along two dimensions. The outer dimension is + * to be computed, organized along two dimensions. The outer dimension is * `n_samples` long, accounting for different samples, while the inner * dimension has size 3 and it represents the x, y, and z coordinates * respectively. * @param xyz_length size of the xyz allocation, i.e, `3 * n_samples` * @param sph pointer to the first element of an array containing `n_samples - * * (l_max + 1) * (l_max + 1)` elements. On exit, this array will contain + * * (l_max + 1) * (l_max + 1)` elements. On exit, this array will contain * the spherical harmonics organized along two dimensions. The leading * dimension is `n_samples` long and it represents the different - * samples, while the inner dimension size is `(l_max + 1) * (l_max + 1)` - * long and it contains the spherical harmonics. These are laid out in + * samples, while the inner dimension size is `(l_max + 1) * (l_max + 1)` + * long and it contains the spherical harmonics. These are laid out in * lexicographic order. For example, if `l_max=2`, it will contain - * `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, + * `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, * 1), (2, 2)`, in this order. * @param sph_length size of the sph allocation, should be `n_samples * - * (l_max + 1) * (l_max + 1)` + * (l_max + 1) * (l_max + 1)` * @param dsph pointer to the first element of an array containing `n_samples - * * `n_samples * 3 * (l_max + 1) * (l_max + 1)` elements. On exit, this + * * `n_samples * 3 * (l_max + 1) * (l_max + 1)` elements. On exit, this * array will contain the spherical harmonics' derivatives organized * along three dimensions. As for the `sph` parameter, the leading * dimension represents the different samples, while the inner-most * dimension size is `(l_max + 1) * (l_max + 1)`, and it represents - * the degree and order of the spherical harmonics (again, organized in + * the degree and order of the spherical harmonics (again, organized in * lexicographic order). The intermediate dimension corresponds to * different spatial derivatives of the spherical harmonics: x, y, - * and z, respectively. + * and z, respectively. * @param dsph_length size of the dsph allocation, which should be `n_samples - * * 3 * (l_max + 1) * (l_max + 1)` + * * 3 * (l_max + 1) * (l_max + 1)` */ SPHERICART_EXPORT void sphericart_spherical_harmonics_compute_array_with_gradients( sphericart_spherical_harmonics_calculator_t* calculator, @@ -196,49 +206,49 @@ SPHERICART_EXPORT void sphericart_spherical_harmonics_compute_array_with_gradien * * @param calculator A pointer to a `sphericart_spherical_harmonics_calculator_t` struct * that holds prefactors and options to compute the spherical - * harmonics. + * harmonics. * @param xyz An array of size `n_samples x 3`. It contains the Cartesian * coordinates of the 3D points for which the spherical harmonics are - * to be computed, organized along two dimensions. The outer dimension is + * to be computed, organized along two dimensions. The outer dimension is * `n_samples` long, accounting for different samples, while the inner * dimension has size 3 and it represents the x, y, and z coordinates * respectively. * @param xyz_length size of the xyz allocation, i.e, `3 * n_samples` * @param sph pointer to the first element of an array containing `n_samples - * * (l_max + 1) * (l_max + 1)` elements. On exit, this array will contain + * * (l_max + 1) * (l_max + 1)` elements. On exit, this array will contain * the spherical harmonics organized along two dimensions. The leading * dimension is `n_samples` long and it represents the different - * samples, while the inner dimension size is `(l_max + 1) * (l_max + 1)` - * long and it contains the spherical harmonics. These are laid out in + * samples, while the inner dimension size is `(l_max + 1) * (l_max + 1)` + * long and it contains the spherical harmonics. These are laid out in * lexicographic order. For example, if `l_max=2`, it will contain - * `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, + * `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, * 1), (2, 2)`, in this order. * @param sph_length size of the sph allocation, should be `n_samples * - * (l_max + 1) * (l_max + 1)` + * (l_max + 1) * (l_max + 1)` * @param dsph pointer to the first element of an array containing * `n_samples * 3 * (l_max + 1) * (l_max + 1)` elements. On exit, - * this array will contain the spherical harmonics' derivatives organized + * this array will contain the spherical harmonics' derivatives organized * along three dimensions. As for the `sph` parameter, the leading * dimension represents the different samples, while the inner-most * dimension size is `(l_max + 1) * (l_max + 1)`, and it represents - * the degree and order of the spherical harmonics (again, organized in + * the degree and order of the spherical harmonics (again, organized in * lexicographic order). The intermediate dimension corresponds to * different spatial derivatives of the spherical harmonics: x, y, - * and z, respectively. + * and z, respectively. * @param dsph_length size of the dsph allocation, which should be `n_samples - * * 3 * (l_max + 1) * (l_max + 1)` + * * 3 * (l_max + 1) * (l_max + 1)` * @param ddsph pointer to the first element of an array containing * `n_samples * 3 * 3 * (l_max + 1) * (l_max + 1)` elements. On exit, - * this array will contain the spherical harmonics' second derivatives - * organized along four dimensions. As for the `sph` parameter, the leading - * dimension represents the different samples, while the inner-most dimension - * size is `(l_max + 1) * (l_max + 1)`, and it represents the degree and - * order of the spherical harmonics (again, organized in lexicographic - * order). The intermediate dimensions correspond to the different spatial - * second derivatives of the spherical harmonics, i.e., to the dimensions of - * the hessian matrix. + * this array will contain the spherical harmonics' second derivatives + * organized along four dimensions. As for the `sph` parameter, the leading + * dimension represents the different samples, while the inner-most dimension + * size is `(l_max + 1) * (l_max + 1)`, and it represents the degree and + * order of the spherical harmonics (again, organized in lexicographic + * order). The intermediate dimensions correspond to the different spatial + * second derivatives of the spherical harmonics, i.e., to the dimensions of + * the hessian matrix. * @param ddsph_length size of the dsph allocation, which should be - * `n_samples * 3 * 3* (l_max + 1) * (l_max + 1)` + * `n_samples * 3 * 3* (l_max + 1) * (l_max + 1)` */ SPHERICART_EXPORT void sphericart_spherical_harmonics_compute_array_with_hessians( sphericart_spherical_harmonics_calculator_t* calculator, @@ -344,6 +354,10 @@ SPHERICART_EXPORT int sphericart_spherical_harmonics_omp_num_threads( sphericart_spherical_harmonics_calculator_t* calculator ); +/** + * Similar to :func:`sphericart_spherical_harmonics_omp_num_threads`, but for + * a `float` calculator. + */ SPHERICART_EXPORT int sphericart_spherical_harmonics_omp_num_threads_f( sphericart_spherical_harmonics_calculator_f_t* calculator ); @@ -391,15 +405,8 @@ SPHERICART_EXPORT void sphericart_spherical_harmonics_compute_sample_with_hessia ); /** - * Initializes a solid harmonics calculator and returns a pointer that - * can then be used by functions that evaluate solid harmonics over - * arrays or individual samples. - * - * @param l_max The maximum degree of the solid harmonics to be - * calculated. - * - * @return A pointer to a `sphericart_solid_harmonics_calculator_t` object - * + * Similar to :func:`sphericart_spherical_harmonics_new`, but it returns a + * `sphericart_solid_harmonics_calculator_t`, which perform solid harmonics calculations. */ SPHERICART_EXPORT sphericart_solid_harmonics_calculator_t* sphericart_solid_harmonics_new(size_t l_max ); @@ -426,30 +433,8 @@ SPHERICART_EXPORT void sphericart_solid_harmonics_delete_f( ); /** - * This function calculates the solid harmonics and, optionally, their - * derivatives for an array of 3D points. - * - * @param calculator A pointer to a `sphericart_solid_harmonics_calculator_t` struct - * that holds prefactors and options to compute the solid + * Similar to :func:`sphericart_spherical_harmonics_compute_array`, but it computes the solid * harmonics. - * @param xyz An array of size `n_samples x 3`. It contains the Cartesian - * coordinates of the 3D points for which the solid harmonics are - * to be computed, organized along two dimensions. The outer dimension is - * `n_samples` long, accounting for different samples, while the inner - * dimension has size 3 and it represents the x, y, and z coordinates - * respectively. - * @param xyz_length size of the xyz allocation, i.e, `3 * n_samples` - * @param sph pointer to the first element of an array containing `n_samples - * * (l_max + 1) * (l_max + 1)` elements. On exit, this array will contain - * the solid harmonics organized along two dimensions. The leading - * dimension is `n_samples` long and it represents the different - * samples, while the inner dimension size is `(l_max + 1) * (l_max + 1)` - * long and it contains the solid harmonics. These are laid out in - * lexicographic order. For example, if `l_max=2`, it will contain - * `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, - * 1), (2, 2)`, in this order. - * @param sph_length size of the sph allocation, should be `n_samples * - * (l_max + 1) * (l_max + 1)` */ SPHERICART_EXPORT void sphericart_solid_harmonics_compute_array( sphericart_solid_harmonics_calculator_t* calculator, @@ -460,42 +445,8 @@ SPHERICART_EXPORT void sphericart_solid_harmonics_compute_array( ); /** - * This function calculates the solid harmonics and their - * derivatives for an array of 3D points. - * - * @param calculator A pointer to a `sphericart_solid_harmonics_calculator_t` struct - * that holds prefactors and options to compute the solid - * harmonics. - * @param xyz An array of size `n_samples x 3`. It contains the Cartesian - * coordinates of the 3D points for which the solid harmonics are - * to be computed, organized along two dimensions. The outer dimension is - * `n_samples` long, accounting for different samples, while the inner - * dimension has size 3 and it represents the x, y, and z coordinates - * respectively. - * @param xyz_length size of the xyz allocation, i.e, `3 * n_samples` - * @param sph pointer to the first element of an array containing `n_samples - * * (l_max + 1) * (l_max + 1)` elements. On exit, this array will contain - * the solid harmonics organized along two dimensions. The leading - * dimension is `n_samples` long and it represents the different - * samples, while the inner dimension size is `(l_max + 1) * (l_max + 1)` - * long and it contains the solid harmonics. These are laid out in - * lexicographic order. For example, if `l_max=2`, it will contain - * `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, - * 1), (2, 2)`, in this order. - * @param sph_length size of the sph allocation, should be `n_samples * - * (l_max + 1) * (l_max + 1)` - * @param dsph pointer to the first element of an array containing `n_samples - * * `n_samples * 3 * (l_max + 1) * (l_max + 1)` elements. On exit, this - * array will contain the solid harmonics' derivatives organized - * along three dimensions. As for the `sph` parameter, the leading - * dimension represents the different samples, while the inner-most - * dimension size is `(l_max + 1) * (l_max + 1)`, and it represents - * the degree and order of the solid harmonics (again, organized in - * lexicographic order). The intermediate dimension corresponds to - * different spatial derivatives of the solid harmonics: x, y, - * and z, respectively. - * @param dsph_length size of the dsph allocation, which should be `n_samples - * * 3 * (l_max + 1) * (l_max + 1)` + * Similar to :func:`sphericart_solid_harmonics_compute_array_with_gradients`, but it computes the + * solid harmonics and their derivatives. */ SPHERICART_EXPORT void sphericart_solid_harmonics_compute_array_with_gradients( sphericart_solid_harmonics_calculator_t* calculator, @@ -508,54 +459,8 @@ SPHERICART_EXPORT void sphericart_solid_harmonics_compute_array_with_gradients( ); /** - * This function calculates the solid harmonics, their - * derivatives and second derivatives for an array of 3D points. - * - * @param calculator A pointer to a `sphericart_solid_harmonics_calculator_t` struct - * that holds prefactors and options to compute the solid - * harmonics. - * @param xyz An array of size `n_samples x 3`. It contains the Cartesian - * coordinates of the 3D points for which the solid harmonics are - * to be computed, organized along two dimensions. The outer dimension is - * `n_samples` long, accounting for different samples, while the inner - * dimension has size 3 and it represents the x, y, and z coordinates - * respectively. - * @param xyz_length size of the xyz allocation, i.e, `3 * n_samples` - * @param sph pointer to the first element of an array containing `n_samples - * * (l_max + 1) * (l_max + 1)` elements. On exit, this array will contain - * the solid harmonics organized along two dimensions. The leading - * dimension is `n_samples` long and it represents the different - * samples, while the inner dimension size is `(l_max + 1) * (l_max + 1)` - * long and it contains the solid harmonics. These are laid out in - * lexicographic order. For example, if `l_max=2`, it will contain - * `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, - * 1), (2, 2)`, in this order. - * @param sph_length size of the sph allocation, should be `n_samples * - * (l_max + 1) * (l_max + 1)` - * @param dsph pointer to the first element of an array containing - * `n_samples * 3 * (l_max + 1) * (l_max + 1)` elements. On exit, - * this array will contain the solid harmonics' derivatives organized - * along three dimensions. As for the `sph` parameter, the leading - * dimension represents the different samples, while the inner-most - * dimension size is `(l_max + 1) * (l_max + 1)`, and it represents - * the degree and order of the solid harmonics (again, organized in - * lexicographic order). The intermediate dimension corresponds to - * different spatial derivatives of the solid harmonics: x, y, - * and z, respectively. - * @param dsph_length size of the dsph allocation, which should be `n_samples - * * 3 * (l_max + 1) * (l_max + 1)` - * @param ddsph pointer to the first element of an array containing - * `n_samples * 3 * 3 * (l_max + 1) * (l_max + 1)` elements. On exit, - * this array will contain the solid harmonics' second derivatives - * organized along four dimensions. As for the `sph` parameter, the leading - * dimension represents the different samples, while the inner-most dimension - * size is `(l_max + 1) * (l_max + 1)`, and it represents the degree and - * order of the solid harmonics (again, organized in lexicographic - * order). The intermediate dimensions correspond to the different spatial - * second derivatives of the solid harmonics, i.e., to the dimensions of - * the hessian matrix. - * @param ddsph_length size of the dsph allocation, which should be - * `n_samples * 3 * 3* (l_max + 1) * (l_max + 1)` + * Similar to :func:`sphericart_solid_harmonics_compute_array_with_hessians`, but it computes the + * solid harmonics and their derivatives. */ SPHERICART_EXPORT void sphericart_solid_harmonics_compute_array_with_hessians( sphericart_solid_harmonics_calculator_t* calculator, @@ -696,13 +601,17 @@ SPHERICART_EXPORT void sphericart_solid_harmonics_compute_sample_with_hessians_f ); /** - * Get the number of OpenMP threads used by a calculator. - * If `sphericart` is computed without OpenMP support returns 1. + * Similar to :func:`sphericart_spherical_harmonics_omp_num_threads`, but for a solid harmonics + * calculator. */ SPHERICART_EXPORT int sphericart_solid_harmonics_omp_num_threads( sphericart_solid_harmonics_calculator_t* calculator ); +/** + * Similar to :func:`sphericart_spherical_harmonics_omp_num_threads`, but for a `float` solid + * harmonics calculator. + */ SPHERICART_EXPORT int sphericart_solid_harmonics_omp_num_threads_f( sphericart_solid_harmonics_calculator_f_t* calculator ); diff --git a/sphericart/include/sphericart.hpp b/sphericart/include/sphericart.hpp index 73251552e..01654f5fb 100644 --- a/sphericart/include/sphericart.hpp +++ b/sphericart/include/sphericart.hpp @@ -1,5 +1,10 @@ /** \file sphericart.hpp - * Defines the C++ API for `sphericart`. + * Defines the C++ API for `sphericart`. Two classes are available: + * `SphericalHarmonics` and `SolidHarmonics`. The former calculates the + * real spherical harmonics :math:`Y^m_l: as defined on Wikipedia, + * which are homogeneous polynomials of (x/r, y/r, z/r). The latter + * calculates the same polynomials but as a function of the Cartesian coordinates + * (x, y, z), or, equivalently, :math:`r^l Y^m_l`. */ #ifndef SPHERICART_HPP @@ -352,14 +357,21 @@ template class SphericalHarmonics { size_t ddsph_length ); + /** + * Returns the maximum degree of the spherical harmonics computed by this + * calculator. + */ + size_t get_l_max() { return this->l_max; } + /** Returns the number of threads used in the calculation */ int get_omp_num_threads() { return this->omp_num_threads; } - template friend class SolidHarmonics; /* @cond */ private: + template friend class SolidHarmonics; + size_t l_max; // maximum l value computed by this class size_t size_y; // size of the Ylm rows (l_max+1)**2 size_t size_q; // size of the prefactor-like arrays (l_max+1)*(l_max+2)/2 @@ -381,13 +393,20 @@ template class SphericalHarmonics { /* @endcond */ }; +/** + * A solid harmonics calculator. + * + * Its interface is the same as that of the `SphericalHarmonics` class, but it + * calculates the solid harmonics :math:`r^l Y^m_l` instead of the real spherical + * harmonics :math:`Y^m_l`, allowing for faster computations. + */ template class SolidHarmonics : public SphericalHarmonics { public: - /** Initialize the SphericalHarmonics class setting maximum degree and + /** Initialize the SolidHarmonics class setting maximum degree and * normalization * * @param l_max - * The maximum degree of the spherical harmonics to be calculated. + * The maximum degree of the solid harmonics to be calculated. */ SolidHarmonics(size_t l_max); }; From 3d0ac03955e0f6f8864e032c55a92c0c74626544 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Wed, 28 Aug 2024 22:23:06 +0200 Subject: [PATCH 02/12] Python and JAX APIs --- python/src/sphericart/spherical_harmonics.py | 54 ++++++------------- .../sphericart/jax/spherical_harmonics.py | 38 +++++++------ 2 files changed, 39 insertions(+), 53 deletions(-) diff --git a/python/src/sphericart/spherical_harmonics.py b/python/src/sphericart/spherical_harmonics.py index bd21a8281..c8b4874d4 100644 --- a/python/src/sphericart/spherical_harmonics.py +++ b/python/src/sphericart/spherical_harmonics.py @@ -8,12 +8,12 @@ class SphericalHarmonics: """ - Spherical harmonics calculator, up to degree ``l_max``. + Spherical harmonics calculator, which computes the real spherical harmonics + :math:`Y^l_m` up to degree ``l_max``. The calculated spherical harmonics + are consistent with the definition of real spherical harmonics from Wikipedia. - This class computes the real spherical harmonics :math:`Y^l_m`. - - In order to minimize the cost of each call, the `SphericalHarmonics` object - computes prefactors and initializes buffers upon creation + The `SphericalHarmonics` object computes prefactors and initializes buffers + upon creation >>> import numpy as np >>> import sphericart as sc @@ -38,7 +38,7 @@ class SphericalHarmonics: :param l_max: the maximum degree of the spherical harmonics to be calculated - :return: a calculator, in the form of a `SphericalHarmonics` object + :return: a calculator, in the form of a ``SphericalHarmonics`` object """ def __init__(self, l_max: int): @@ -64,7 +64,7 @@ def __del__(self): def compute(self, xyz: np.ndarray) -> np.ndarray: """ Calculates the spherical harmonics for a set of 3D points, whose - coordinates are in the ``xyz`` array. + coordinates are given by the ``xyz`` array. >>> import numpy as np >>> import sphericart as sc @@ -141,12 +141,13 @@ def compute_with_gradients(self, xyz: np.ndarray) -> Tuple[np.ndarray, np.ndarra :return: A tuple containing: - * an array of shape ``(n_samples, (l_max+1)**2)`` containing all the + + - an array 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. - * An array of shape ``(n_samples, 3, (l_max+1)**2)`` containing all + - an array 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 @@ -231,17 +232,17 @@ def compute_with_hessians( :return: A tuple containing: - * an array of shape ``(n_samples, (l_max+1)**2)`` containing all the + - an array 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. - * An array of shape ``(n_samples, 3, (l_max+1)**2)`` containing all + - an array 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. - * An array of shape ``(n_samples, 3, 3, (l_max+1)**2)`` containing all + - an array 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 @@ -315,33 +316,10 @@ class SolidHarmonics: This class computes the solid harmonics, a non-normalized form of the real spherical harmonics, i.e. :math:`r^l Y^l_m`. These scaled spherical harmonics - are polynomials in the Cartesian coordinates of the input points. - - In order to minimize the cost of each call, the `SolidHarmonics` object - computes prefactors and initializes buffers upon creation - - >>> import numpy as np - >>> import sphericart as sc - >>> sh = sc.SolidHarmonics(l_max=8) + are polynomials in the Cartesian coordinates of the input points, and they + are therefore faster to compute. - Then, the :py:func:`compute` method can be called on an array of 3D - Cartesian points to compute the solid harmonics - - >>> xyz = np.random.normal(size=(10,3)) - >>> sh_values = sh.compute(xyz) - >>> sh_values.shape - (10, 81) - - In order to also compute derivatives, you can use - - >>> sh_values, sh_grads = sh.compute_with_gradients(xyz) - >>> sh_grads.shape - (10, 3, 81) - - which returns the gradient as a tensor with size - `(n_samples, 3, (l_max+1)**2)`. - - :param l_max: the maximum degree of the spherical harmonics to be calculated + :param l_max: the maximum degree of the solid harmonics to be calculated :return: a calculator, in the form of a `SolidHarmonics` object """ diff --git a/sphericart-jax/python/sphericart/jax/spherical_harmonics.py b/sphericart-jax/python/sphericart/jax/spherical_harmonics.py index d3acb3185..a60a36065 100644 --- a/sphericart-jax/python/sphericart/jax/spherical_harmonics.py +++ b/sphericart-jax/python/sphericart/jax/spherical_harmonics.py @@ -1,10 +1,15 @@ +import jax + from .sph import sph -def spherical_harmonics(xyz, l_max): - """Computes the Spherical harmonics and their derivatives within +def spherical_harmonics(xyz: jax.Array, l_max: int): + """Computes the spherical harmonics and their derivatives within the JAX framework. + The definition of the real spherical harmonics is consistent with the + Wikipedia spherical harmonics page. + Note that the ``l_max`` argument (position 1 in the signature) should be tagged as static when jit-ing the function: @@ -12,18 +17,16 @@ def spherical_harmonics(xyz, l_max): >>> import sphericart.jax >>> jitted_sph_fn = jax.jit(sphericart.jax.spherical_harmonics, static_argnums=1) - Parameters - ---------- - xyz : jax array [..., 3] - single vector or set of vectors in 3D. All dimensions are optional except for - the last - l_max : int - maximum order of the spherical harmonics (included) - - Returns - ------- - jax array [..., (l_max+1)**2] - Spherical harmonics expansion of `xyz` + :param xyz: single vector or set of vectors in 3D. All dimensions are optional + except for the last. Shape ``[..., 3]``. + :param l_max: the maximum degree of the spherical harmonics to be calculated + (included) + + :return: Spherical harmonics expansion of ``xyz``. Shape ``[..., (l_max+1)**2]``. + The last dimension is organized 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. """ if xyz.shape[-1] != 3: raise ValueError("the last axis of xyz must have size 3") @@ -32,9 +35,14 @@ def spherical_harmonics(xyz, l_max): return output -def solid_harmonics(xyz, l_max, normalized=False): +def solid_harmonics(xyz: jax.Array, l_max: int): """ Same as `spherical_harmonics`, but computes the solid harmonics instead. + + These are a non-normalized form of the real + spherical harmonics, i.e. :math:`r^l Y^l_m`. These scaled spherical harmonics + are polynomials in the Cartesian coordinates of the input points, and they + are therefore less expoensive to compute. """ if xyz.shape[-1] != 3: raise ValueError("the last axis of xyz must have size 3") From 8aaccfbab79d33eed5fcb815965436e8d7063c68 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Wed, 28 Aug 2024 22:36:28 +0200 Subject: [PATCH 03/12] Update math page --- docs/src/maths.rst | 27 +++++----- python/src/sphericart/spherical_harmonics.py | 53 +++++++++----------- 2 files changed, 40 insertions(+), 40 deletions(-) diff --git a/docs/src/maths.rst b/docs/src/maths.rst index f26e5b271..eec92c352 100644 --- a/docs/src/maths.rst +++ b/docs/src/maths.rst @@ -12,11 +12,21 @@ There are multiple conventions for choosing normalization and phases, and it is possible to reformulate the spherical harmonics in a real-valued form, which leads to even further ambiguity in the definitions. -Within `sphericart` we take an opinionated stance: we compute only real-valued -harmonics, we express them as a function of the full Cartesian coordinates of a -point in three dimensions :math:`(x,y,z)` and compute by default "scaled" -versions :math:`\tilde{Y}^m_l(x, y, z)` which correspond to homogeneous polynomials -of the Cartesian coordinates: +Within `sphericart`, we compute only real-valued spherical harmonics and we express +them as a function of the full Cartesian coordinates of a point in three dimensions. +These correspond to the real spherical harmonics as defined in the corresponding +`Wikipedia article `_, which we +refer to as :math:`Y^m_l`. + +We also offer the possibility to compute "solid" harmonics, which are given by +:math:`\tilde{Y}^m_l = r^l\,{Y}_l^m`. Since these can be expressed as homogeneous +polynomials of the Cartesian coordinates :math:`(x,y,z)`, as opposed to +:math:`(x/r,y/r,z/r)`, they are less computationally expensive to evaluate. +Besides being slightly faster, they can also +provide a more natural scaling if used together with a radial expansion. + +The formulas used to compute the solid harmonics (and, with few modifications, +also for the spherical harmonics) are: .. math :: \tilde{Y}_l^m(x, y, z) = r^l\,{Y}_l^m(x, y, z) = F_l^{|m|} Q_l^{|m|}(z, r) \times @@ -41,13 +51,6 @@ If we neglect some constant normalization factors, these correspond to the See also the `reference paper `_ for further implementation details. -The radially normalized version of the spherical harmonics can also be computed by providing -the appropriate flag when creating the `sphericart` calculators. These correspond to -the real spherical harmonics as defined in the corresponding -`Wikipedia article `_. -However, we recommend using the scaled versions, which are slightly faster and -provide a more natural scaling if used together with a radial expansion. - The :math:`\tilde{Y}^m_l(x)` are stored contiguously in memory, e.g. as :math:`\{ (l,m)=(0,0), (1,-1), (1,0), (1,1), (2,-2), \ldots \}`. With zero-based indexing of the arrays, the ``(l,m)`` term is stored at diff --git a/python/src/sphericart/spherical_harmonics.py b/python/src/sphericart/spherical_harmonics.py index c8b4874d4..a8038cfa0 100644 --- a/python/src/sphericart/spherical_harmonics.py +++ b/python/src/sphericart/spherical_harmonics.py @@ -135,23 +135,21 @@ def compute_with_gradients(self, xyz: np.ndarray) -> Tuple[np.ndarray, np.ndarra >>> sh_grads.shape (10, 3, 81) - :param xyz: - The Cartesian coordinates of the 3D points, as an array with + :param xyz: The Cartesian coordinates of the 3D points, as an array with shape ``(n_samples, 3)``. - :return: - A tuple containing: + :return: A tuple containing: - an array 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. + 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. - an array 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. + 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. """ @@ -226,27 +224,26 @@ def compute_with_hessians( >>> sh_hessians.shape (10, 3, 3, 81) - :param xyz: - The Cartesian coordinates of the 3D points, as an array with + :param xyz: The Cartesian coordinates of the 3D points, as an array with shape ``(n_samples, 3)``. - :return: - A tuple containing: + :return: A tuple containing: + - an array 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. + 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. - an array 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. + 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. - an array 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. + 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. """ From df496272ab5eb50048dda647c46ad8a17c8ceeb6 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Thu, 29 Aug 2024 08:25:40 +0200 Subject: [PATCH 04/12] Fix C API rendering issue --- sphericart/include/sphericart.h | 54 ++++++++++++++++----------------- 1 file changed, 27 insertions(+), 27 deletions(-) diff --git a/sphericart/include/sphericart.h b/sphericart/include/sphericart.h index 149759907..8b60f66a7 100644 --- a/sphericart/include/sphericart.h +++ b/sphericart/include/sphericart.h @@ -131,18 +131,18 @@ SPHERICART_EXPORT void sphericart_spherical_harmonics_delete_f( * `n_samples` long, accounting for different samples, while the inner * dimension has size 3 and it represents the x, y, and z coordinates * respectively. - * @param xyz_length size of the xyz allocation, i.e, `3 * n_samples` + * @param xyz_length size of the xyz allocation, i.e, `3 x n_samples` * @param sph pointer to the first element of an array containing `n_samples - * * (l_max + 1) * (l_max + 1)` elements. On exit, this array will contain + * x (l_max + 1) x (l_max + 1)` elements. On exit, this array will contain * the spherical harmonics organized along two dimensions. The leading * dimension is `n_samples` long and it represents the different - * samples, while the inner dimension size is `(l_max + 1) * (l_max + 1)` + * samples, while the inner dimension size is `(l_max + 1) x (l_max + 1)` * long and it contains the spherical harmonics. These are laid out in * lexicographic order. For example, if `l_max=2`, it will contain * `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, * 1), (2, 2)`, in this order. - * @param sph_length size of the sph allocation, should be `n_samples * - * (l_max + 1) * (l_max + 1)` + * @param sph_length size of the sph allocation, should be `n_samples x + * (l_max + 1) x (l_max + 1)` */ SPHERICART_EXPORT void sphericart_spherical_harmonics_compute_array( sphericart_spherical_harmonics_calculator_t* calculator, @@ -165,30 +165,30 @@ SPHERICART_EXPORT void sphericart_spherical_harmonics_compute_array( * `n_samples` long, accounting for different samples, while the inner * dimension has size 3 and it represents the x, y, and z coordinates * respectively. - * @param xyz_length size of the xyz allocation, i.e, `3 * n_samples` + * @param xyz_length size of the xyz allocation, i.e, `3 x n_samples`` * @param sph pointer to the first element of an array containing `n_samples - * * (l_max + 1) * (l_max + 1)` elements. On exit, this array will contain + * x (l_max + 1) x (l_max + 1)` elements. On exit, this array will contain * the spherical harmonics organized along two dimensions. The leading * dimension is `n_samples` long and it represents the different - * samples, while the inner dimension size is `(l_max + 1) * (l_max + 1)` + * samples, while the inner dimension size is ``(l_max + 1) x (l_max + 1)`` * long and it contains the spherical harmonics. These are laid out in * lexicographic order. For example, if `l_max=2`, it will contain * `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, * 1), (2, 2)`, in this order. * @param sph_length size of the sph allocation, should be `n_samples * - * (l_max + 1) * (l_max + 1)` + * (l_max + 1) x (l_max + 1)` * @param dsph pointer to the first element of an array containing `n_samples - * * `n_samples * 3 * (l_max + 1) * (l_max + 1)` elements. On exit, this + * x `n_samples x 3 x (l_max + 1) x (l_max + 1)` elements. On exit, this * array will contain the spherical harmonics' derivatives organized * along three dimensions. As for the `sph` parameter, the leading * dimension represents the different samples, while the inner-most - * dimension size is `(l_max + 1) * (l_max + 1)`, and it represents + * dimension size is `(l_max + 1) x (l_max + 1)`, and it represents * the degree and order of the spherical harmonics (again, organized in * lexicographic order). The intermediate dimension corresponds to * different spatial derivatives of the spherical harmonics: x, y, * and z, respectively. * @param dsph_length size of the dsph allocation, which should be `n_samples - * * 3 * (l_max + 1) * (l_max + 1)` + * x 3 x (l_max + 1) x (l_max + 1)` */ SPHERICART_EXPORT void sphericart_spherical_harmonics_compute_array_with_gradients( sphericart_spherical_harmonics_calculator_t* calculator, @@ -213,42 +213,42 @@ SPHERICART_EXPORT void sphericart_spherical_harmonics_compute_array_with_gradien * `n_samples` long, accounting for different samples, while the inner * dimension has size 3 and it represents the x, y, and z coordinates * respectively. - * @param xyz_length size of the xyz allocation, i.e, `3 * n_samples` - * @param sph pointer to the first element of an array containing `n_samples - * * (l_max + 1) * (l_max + 1)` elements. On exit, this array will contain - * the spherical harmonics organized along two dimensions. The leading - * dimension is `n_samples` long and it represents the different - * samples, while the inner dimension size is `(l_max + 1) * (l_max + 1)` + * @param xyz_length size of the xyz allocation, i.e, ``3 x n_samples`` + * @param sph pointer to the first element of an array containing + * `n_samples x (l_max + 1) x (l_max + 1)` elements. On exit, this array + * will contain the spherical harmonics organized along two dimensions. + * The leading dimension is `n_samples` long and it represents the different + * samples, while the inner dimension size is `(l_max + 1) x (l_max + 1)` * long and it contains the spherical harmonics. These are laid out in * lexicographic order. For example, if `l_max=2`, it will contain - * `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, - * 1), (2, 2)`, in this order. + * ``(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, + * 1), (2, 2)``, in this order. * @param sph_length size of the sph allocation, should be `n_samples * - * (l_max + 1) * (l_max + 1)` + * (l_max + 1) x (l_max + 1)` * @param dsph pointer to the first element of an array containing - * `n_samples * 3 * (l_max + 1) * (l_max + 1)` elements. On exit, + * `n_samples x 3 x (l_max + 1) x (l_max + 1)` elements. On exit, * this array will contain the spherical harmonics' derivatives organized * along three dimensions. As for the `sph` parameter, the leading * dimension represents the different samples, while the inner-most - * dimension size is `(l_max + 1) * (l_max + 1)`, and it represents + * dimension size is `(l_max + 1) x (l_max + 1)`, and it represents * the degree and order of the spherical harmonics (again, organized in * lexicographic order). The intermediate dimension corresponds to * different spatial derivatives of the spherical harmonics: x, y, * and z, respectively. * @param dsph_length size of the dsph allocation, which should be `n_samples - * * 3 * (l_max + 1) * (l_max + 1)` + * x 3 x (l_max + 1) x (l_max + 1)` * @param ddsph pointer to the first element of an array containing - * `n_samples * 3 * 3 * (l_max + 1) * (l_max + 1)` elements. On exit, + * `n_samples x 3 x 3 x (l_max + 1) x (l_max + 1)` elements. On exit, * this array will contain the spherical harmonics' second derivatives * organized along four dimensions. As for the `sph` parameter, the leading * dimension represents the different samples, while the inner-most dimension - * size is `(l_max + 1) * (l_max + 1)`, and it represents the degree and + * size is `(l_max + 1) x (l_max + 1)`, and it represents the degree and * order of the spherical harmonics (again, organized in lexicographic * order). The intermediate dimensions correspond to the different spatial * second derivatives of the spherical harmonics, i.e., to the dimensions of * the hessian matrix. * @param ddsph_length size of the dsph allocation, which should be - * `n_samples * 3 * 3* (l_max + 1) * (l_max + 1)` + * `n_samples x 3 x 3* (l_max + 1) x (l_max + 1)` */ SPHERICART_EXPORT void sphericart_spherical_harmonics_compute_array_with_hessians( sphericart_spherical_harmonics_calculator_t* calculator, From 9c35c69b7bf2455c3d84da4d9b8e90a0606fb1f3 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Thu, 29 Aug 2024 09:29:27 +0200 Subject: [PATCH 05/12] Fix a few math rendering issues --- docs/src/jax-api.rst | 2 ++ docs/src/python-api.rst | 3 +++ python/src/sphericart/spherical_harmonics.py | 4 ++-- .../python/sphericart/jax/spherical_harmonics.py | 2 +- sphericart-torch/python/sphericart/torch/__init__.py | 12 ++++++------ sphericart/include/sphericart.h | 4 ++-- sphericart/include/sphericart.hpp | 8 ++++---- 7 files changed, 20 insertions(+), 15 deletions(-) diff --git a/docs/src/jax-api.rst b/docs/src/jax-api.rst index ff9776948..c7876e86a 100644 --- a/docs/src/jax-api.rst +++ b/docs/src/jax-api.rst @@ -14,3 +14,5 @@ using the CPU or CUDA implementation. .. autofunction:: sphericart.jax.spherical_harmonics +.. autofunction:: sphericart.jax.solid_harmonics + diff --git a/docs/src/python-api.rst b/docs/src/python-api.rst index eef7fd2b4..0cd529b62 100644 --- a/docs/src/python-api.rst +++ b/docs/src/python-api.rst @@ -3,3 +3,6 @@ Python API .. autoclass:: sphericart.SphericalHarmonics :members: + +.. autoclass:: sphericart.SolidHarmonics + :members: diff --git a/python/src/sphericart/spherical_harmonics.py b/python/src/sphericart/spherical_harmonics.py index a8038cfa0..ad549017a 100644 --- a/python/src/sphericart/spherical_harmonics.py +++ b/python/src/sphericart/spherical_harmonics.py @@ -9,7 +9,7 @@ class SphericalHarmonics: """ Spherical harmonics calculator, which computes the real spherical harmonics - :math:`Y^l_m` up to degree ``l_max``. The calculated spherical harmonics + :math:`Y^m_l` up to degree ``l_max``. The calculated spherical harmonics are consistent with the definition of real spherical harmonics from Wikipedia. The `SphericalHarmonics` object computes prefactors and initializes buffers @@ -312,7 +312,7 @@ class SolidHarmonics: 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. :math:`r^l Y^l_m`. These scaled spherical harmonics + spherical harmonics, i.e. :math:`r^l Y^m_l`. These scaled spherical harmonics are polynomials in the Cartesian coordinates of the input points, and they are therefore faster to compute. diff --git a/sphericart-jax/python/sphericart/jax/spherical_harmonics.py b/sphericart-jax/python/sphericart/jax/spherical_harmonics.py index a60a36065..d9274ed92 100644 --- a/sphericart-jax/python/sphericart/jax/spherical_harmonics.py +++ b/sphericart-jax/python/sphericart/jax/spherical_harmonics.py @@ -40,7 +40,7 @@ def solid_harmonics(xyz: jax.Array, l_max: int): Same as `spherical_harmonics`, but computes the solid harmonics instead. These are a non-normalized form of the real - spherical harmonics, i.e. :math:`r^l Y^l_m`. These scaled spherical harmonics + spherical harmonics, i.e. :math:`r^lY^m_l`. These scaled spherical harmonics are polynomials in the Cartesian coordinates of the input points, and they are therefore less expoensive to compute. """ diff --git a/sphericart-torch/python/sphericart/torch/__init__.py b/sphericart-torch/python/sphericart/torch/__init__.py index 06aa3b7ab..d3a9e64d2 100644 --- a/sphericart-torch/python/sphericart/torch/__init__.py +++ b/sphericart-torch/python/sphericart/torch/__init__.py @@ -74,7 +74,7 @@ class SphericalHarmonics(torch.nn.Module): """ Spherical harmonics calculator, up to degree ``l_max``. - This class computes :math:`Y^l_m`, which are instead homogeneous polynomials + This class computes :math:`Y^m_l`, which are instead homogeneous polynomials of x/r, y/r, z/r. This class can be used similarly to :py:class:`sphericart.SphericalHarmonics` @@ -248,7 +248,7 @@ class SolidHarmonics(torch.nn.Module): 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. :math:`r^l Y^l_m`. These scaled spherical harmonics + spherical harmonics, i.e. :math:`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 :py:class:`sphericart.SphericalHarmonics`. @@ -319,7 +319,7 @@ def e3nn_spherical_harmonics( :param l_list: Either a single integer or a list of integers specifying which - :math:`Y^l_m` should be computed. All values up to the maximum + :math:`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. :param x: @@ -327,12 +327,12 @@ def e3nn_spherical_harmonics( expected by the `e3nn` function. :param normalize: Flag specifying whether the input positions should be normalized - (resulting in the computation of the spherical harmonics :math:`Y^l_m`), + (resulting in the computation of the spherical harmonics :math:`Y^m_l`), or whether the function should compute the solid harmonics - :math:`\tilde{Y}^l_m`. + :math:`\tilde{Y}^m_l=r^lY^m_l`. :param normalization: String that can be "integral", "norm", "component", that controls - a further scaling of the :math:`Y_m^l`. See the + a further scaling of the :math:`Y^m_l`. See the documentation of :py:func:`e3nn.o3.spherical_harmonics()` for a detailed explanation of the different conventions. """ diff --git a/sphericart/include/sphericart.h b/sphericart/include/sphericart.h index 8b60f66a7..1013cd942 100644 --- a/sphericart/include/sphericart.h +++ b/sphericart/include/sphericart.h @@ -6,11 +6,11 @@ * while those for the `double` data type do not. * Similarly, types and functions for spherical harmonics calculations contain * `_spherical_harmonics` in their name, and they calculate the - * real spherical harmonics :math:`Y^m_l: as defined on Wikipedia, + * real spherical harmonics \f$ Y^m_l \f$ as defined on Wikipedia, * which are homogeneous polynomials of (x/r, y/r, z/r). In contrast, types * and functions for solid harmonics calculations contain `_solid_harmonics` in * their name, and thay calculate the same polynomials but as a function of the - * Cartesian coordinates (x, y, z), or, equivalently, :math:`r^l Y^m_l`. + * Cartesian coordinates (x, y, z), or, equivalently, \f$ r^l\,Y^m_l \f$. */ #ifndef SPHERICART_H diff --git a/sphericart/include/sphericart.hpp b/sphericart/include/sphericart.hpp index 01654f5fb..c93f3951a 100644 --- a/sphericart/include/sphericart.hpp +++ b/sphericart/include/sphericart.hpp @@ -1,10 +1,10 @@ /** \file sphericart.hpp * Defines the C++ API for `sphericart`. Two classes are available: * `SphericalHarmonics` and `SolidHarmonics`. The former calculates the - * real spherical harmonics :math:`Y^m_l: as defined on Wikipedia, + * real spherical harmonics \f$ Y^m_l \f$ as defined on Wikipedia, * which are homogeneous polynomials of (x/r, y/r, z/r). The latter * calculates the same polynomials but as a function of the Cartesian coordinates - * (x, y, z), or, equivalently, :math:`r^l Y^m_l`. + * (x, y, z), or, equivalently, \f$ r^l\,Y^m_l \f$. */ #ifndef SPHERICART_HPP @@ -397,8 +397,8 @@ template class SphericalHarmonics { * A solid harmonics calculator. * * Its interface is the same as that of the `SphericalHarmonics` class, but it - * calculates the solid harmonics :math:`r^l Y^m_l` instead of the real spherical - * harmonics :math:`Y^m_l`, allowing for faster computations. + * calculates the solid harmonics \f$ r^l\,Y^m_l \f$ instead of the real spherical + * harmonics \f$ Y^m_l \f$, allowing for faster computations. */ template class SolidHarmonics : public SphericalHarmonics { public: From d2c3c9d411f23fe3ad7afb0154795f6ec7c7cb70 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Thu, 29 Aug 2024 14:08:33 +0200 Subject: [PATCH 06/12] Update CUDA C++ documentation --- sphericart/include/sphericart.hpp | 3 +- sphericart/include/sphericart_cuda.hpp | 95 ++++++++++++++++++++----- sphericart/src/sphericart_cuda.cu | 20 +++--- sphericart/src/sphericart_cuda_stub.cpp | 8 +-- 4 files changed, 93 insertions(+), 33 deletions(-) diff --git a/sphericart/include/sphericart.hpp b/sphericart/include/sphericart.hpp index c93f3951a..74a9c24aa 100644 --- a/sphericart/include/sphericart.hpp +++ b/sphericart/include/sphericart.hpp @@ -402,8 +402,7 @@ template class SphericalHarmonics { */ template class SolidHarmonics : public SphericalHarmonics { public: - /** Initialize the SolidHarmonics class setting maximum degree and - * normalization + /** Initialize the SolidHarmonics class setting its maximum degree * * @param l_max * The maximum degree of the solid harmonics to be calculated. diff --git a/sphericart/include/sphericart_cuda.hpp b/sphericart/include/sphericart_cuda.hpp index a7f54c54c..3df183845 100644 --- a/sphericart/include/sphericart_cuda.hpp +++ b/sphericart/include/sphericart_cuda.hpp @@ -1,6 +1,12 @@ /** \file sphericart_cuda.hpp - * Defines the CUDA API for `sphericart`. + * Defines the CUDA C++ API for `sphericart`. Two classes are available: + * `SphericalHarmonics` and `SolidHarmonics`. The former calculates the + * real spherical harmonics \f$ Y^m_l \f$ as defined on Wikipedia, + * which are homogeneous polynomials of (x/r, y/r, z/r). The latter + * calculates the same polynomials but as a function of the Cartesian coordinates + * (x, y, z), or, equivalently, \f$ r^l\,Y^m_l \f$. */ + #ifndef SPHERICART_CUDA_HPP #define SPHERICART_CUDA_HPP @@ -13,7 +19,7 @@ namespace sphericart { /* @endcond */ /** - * The ``sphericart::cuda`` namespace contains the CUDA API for `sphericart`. + * The `sphericart::cuda` namespace contains the CUDA API for `sphericart`. */ namespace cuda { @@ -52,40 +58,87 @@ template class SphericalHarmonics { * The outer dimension is `n_samples` long, accounting for different * samples, while the inner dimension has size 3 and it represents * the x, y, and z coordinates respectively. - * @param nsamples Number of samples contained within `xyz`. + * @param n_samples Number of samples contained within `xyz`. + * @param sph On entry, a preallocated device-side array of size `n_samples + * x (l_max + 1) x (l_max + 1)`. On exit, this array will contain the + * spherical harmonics organized along two dimensions. The leading dimension + * is `n_samples` long and it represents the different samples, while the + * inner dimension is + * `(l_max + 1) x (l_max + 1)` long and it contains the spherical harmonics. + * These are laid out in lexicographic order. For example, if `l_max=2`, it + * will contain `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), + * (2, 0), (2, 1), (2, 2)`, in this order. + * @param cuda_stream Pointer to a cudaStream_t or nullptr. If this is + * nullptr, the kernel launch will be performed on the default stream. + */ + void compute(const T* xyz, const size_t n_samples, T* sph, void* cuda_stream = nullptr); + + /** Computes the spherical harmonics and their first derivatives for one or + * more 3D points, using pre-allocated device-side pointers + * + * @param xyz A pre-allocated device-side array of size `n_samples x 3`. It + * contains the Cartesian coordinates of the 3D points for which the + * spherical harmonics are to be computed, organized along two dimensions. + * The outer dimension is `n_samples` long, accounting for different + * samples, while the inner dimension has size 3 and it represents + * the x, y, and z coordinates respectively. + * @param n_samples Number of samples contained within `xyz`. + * @param sph On entry, a preallocated device-side array of size `n_samples + * * (l_max + 1) x (l_max + 1)`. On exit, this array will contain the + * spherical harmonics organized along two dimensions. The leading dimension + * is `n_samples` long and it represents the different samples, while the + * inner dimension is + * `(l_max + 1) x (l_max + 1)` long and it contains the spherical harmonics. + * These are laid out in lexicographic order. For example, if `l_max=2`, it + * will contain `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), + * (2, 0), (2, 1), (2, 2)`, in this order. + * @param dsph On entry, nullptr or a preallocated device-side array of size + * `n_samples x 3 x (l_max + 1) x (l_max + 1)`. If the pointer is not + * nullptr, then compute_with_gradients must also be true in order for + * gradients to be computed. + * @param cuda_stream Pointer to a cudaStream_t or nullptr. If this is + * nullptr, the kernel launch will be performed on the default stream. + */ + void compute_with_gradients( + const T* xyz, const size_t n_samples, T* sph, T* dsph, void* cuda_stream = nullptr + ); + + /** Computes the spherical harmonics and their first and second derivatives + * for one or more 3D points, using pre-allocated device-side pointers + * + * @param xyz A pre-allocated device-side array of size `n_samples x 3`. It + * contains the Cartesian coordinates of the 3D points for which the + * spherical harmonics are to be computed, organized along two dimensions. + * The outer dimension is `n_samples` long, accounting for different + * samples, while the inner dimension has size 3 and it represents + * the x, y, and z coordinates respectively. + * @param n_samples Number of samples contained within `xyz`. * @param compute_with_gradients Whether we should compute dsph. If true, * the pointer dsph must also be allocated on device. * @param compute_with_hessians Whether we should compute ddsph. If true, * the pointer ddsph must also be allocated on device. * @param sph On entry, a preallocated device-side array of size `n_samples - * * (l_max + 1) * (l_max + 1)`. On exit, this array will contain the + * * (l_max + 1) x (l_max + 1)`. On exit, this array will contain the * spherical harmonics organized along two dimensions. The leading dimension * is `n_samples` long and it represents the different samples, while the * inner dimension is - * `(l_max + 1) * (l_max + 1)` long and it contains the spherical harmonics. + * `(l_max + 1) x (l_max + 1)` long and it contains the spherical harmonics. * These are laid out in lexicographic order. For example, if `l_max=2`, it * will contain `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), * (2, 0), (2, 1), (2, 2)`, in this order. * @param dsph On entry, nullptr or a preallocated device-side array of size - * `n_samples * 3 * (l_max + 1) * (l_max + 1)`. If the pointer is not + * `n_samples x 3 x (l_max + 1) x (l_max + 1)`. If the pointer is not * nullptr, then compute_with_gradients must also be true in order for * gradients to be computed. * @param ddsph On entry, nullptr or a preallocated device-side array of - * size `n_samples * 3 * 3 * (l_max + 1) * (l_max + 1)`. If the pointer is + * size `n_samples x 3 x 3 x (l_max + 1) x (l_max + 1)`. If the pointer is * not nullptr, then compute_with_hessians must also be true in order for * gradients to be computed. * @param cuda_stream Pointer to a cudaStream_t or nullptr. If this is * nullptr, the kernel launch will be performed on the default stream. */ - - void compute(const T* xyz, const size_t nsamples, T* sph, void* cuda_stream = nullptr); - - void compute_with_gradients( - const T* xyz, const size_t nsamples, T* sph, T* dsph, void* cuda_stream = nullptr - ); - void compute_with_hessians( - const T* xyz, const size_t nsamples, T* sph, T* dsph, T* ddsph, void* cuda_stream = nullptr + const T* xyz, const size_t n_samples, T* sph, T* dsph, T* ddsph, void* cuda_stream = nullptr ); template friend class SolidHarmonics; @@ -105,7 +158,7 @@ template class SphericalHarmonics { void compute_internal( const T* xyz, - const size_t nsamples, + const size_t n_samples, bool compute_with_gradients, bool compute_with_hessian, T* sph, @@ -116,10 +169,16 @@ template class SphericalHarmonics { /* @endcond */ }; +/** + * A solid harmonics calculator. + * + * Its interface is the same as that of the `SphericalHarmonics` class, but it + * calculates the solid harmonics \f$ r^l\,Y^m_l \f$ instead of the real spherical + * harmonics \f$ Y^m_l \f$, allowing for faster computations. + */ template class SolidHarmonics : public SphericalHarmonics { public: - /** Initialize the SolidHarmonics class setting maximum degree and - * normalization + /** Initialize the SolidHarmonics class setting its maximum degree * * @param l_max * The maximum degree of the spherical harmonics to be calculated. diff --git a/sphericart/src/sphericart_cuda.cu b/sphericart/src/sphericart_cuda.cu index d751ffe39..b6f4f1241 100644 --- a/sphericart/src/sphericart_cuda.cu +++ b/sphericart/src/sphericart_cuda.cu @@ -107,7 +107,7 @@ template SphericalHarmonics::~SphericalHarmonics() { template void SphericalHarmonics::compute_internal( const T* xyz, - const size_t nsamples, + const size_t n_samples, bool compute_with_gradients, bool compute_with_hessian, T* sph, @@ -115,7 +115,7 @@ void SphericalHarmonics::compute_internal( T* ddsph, void* cuda_stream ) { - if (nsamples == 0) { + if (n_samples == 0) { // nothing to compute; we return here because some libraries (e.g. torch) // seem to use nullptrs for tensors with 0 elements return; @@ -193,7 +193,7 @@ void SphericalHarmonics::compute_internal( sphericart::cuda::spherical_harmonics_cuda_base( xyz, - nsamples, + n_samples, this->prefactors_cuda[attributes.device], this->nprefactors, this->l_max, @@ -211,26 +211,28 @@ void SphericalHarmonics::compute_internal( CUDA_CHECK(cudaSetDevice(current_device)); } template -void SphericalHarmonics::compute(const T* xyz, const size_t nsamples, T* sph, void* cuda_stream) { +void SphericalHarmonics::compute(const T* xyz, const size_t n_samples, T* sph, void* cuda_stream) { SphericalHarmonics::compute_internal( - xyz, nsamples, false, false, sph, nullptr, nullptr, cuda_stream + xyz, n_samples, false, false, sph, nullptr, nullptr, cuda_stream ); } template void SphericalHarmonics::compute_with_gradients( - const T* xyz, const size_t nsamples, T* sph, T* dsph, void* cuda_stream + const T* xyz, const size_t n_samples, T* sph, T* dsph, void* cuda_stream ) { SphericalHarmonics::compute_internal( - xyz, nsamples, true, false, sph, dsph, nullptr, cuda_stream + xyz, n_samples, true, false, sph, dsph, nullptr, cuda_stream ); } template void SphericalHarmonics::compute_with_hessians( - const T* xyz, const size_t nsamples, T* sph, T* dsph, T* ddsph, void* cuda_stream + const T* xyz, const size_t n_samples, T* sph, T* dsph, T* ddsph, void* cuda_stream ) { - SphericalHarmonics::compute_internal(xyz, nsamples, true, true, sph, dsph, ddsph, cuda_stream); + SphericalHarmonics::compute_internal( + xyz, n_samples, true, true, sph, dsph, ddsph, cuda_stream + ); } template diff --git a/sphericart/src/sphericart_cuda_stub.cpp b/sphericart/src/sphericart_cuda_stub.cpp index 2b4a595c0..f6f1e2cee 100644 --- a/sphericart/src/sphericart_cuda_stub.cpp +++ b/sphericart/src/sphericart_cuda_stub.cpp @@ -11,7 +11,7 @@ template SphericalHarmonics::~SphericalHarmonics() {} template void SphericalHarmonics::compute_internal( const T* /*xyz*/, - const size_t /*nsamples*/, + const size_t /*n_samples*/, bool /*compute_with_gradients*/, bool /*compute_with_hessian*/, T* /*sph*/, @@ -24,21 +24,21 @@ void SphericalHarmonics::compute_internal( template void SphericalHarmonics::compute( - const T* /*xyz*/, const size_t /*nsamples*/, T* /*sph*/, void* /*cuda_stream*/ + const T* /*xyz*/, const size_t /*n_samples*/, T* /*sph*/, void* /*cuda_stream*/ ) { throw std::runtime_error("sphericart was not compiled with CUDA support"); } template void SphericalHarmonics::compute_with_gradients( - const T* /*xyz*/, const size_t /*nsamples*/, T* /*sph*/, T* /*dsph*/, void* /*cuda_stream*/ + const T* /*xyz*/, const size_t /*n_samples*/, T* /*sph*/, T* /*dsph*/, void* /*cuda_stream*/ ) { throw std::runtime_error("sphericart was not compiled with CUDA support"); } template void SphericalHarmonics::compute_with_hessians( - const T* /*xyz*/, const size_t /*nsamples*/, T* /*sph*/, T* /*dsph*/, T* /*ddsph*/, void* /*cuda_stream*/ + const T* /*xyz*/, const size_t /*n_samples*/, T* /*sph*/, T* /*dsph*/, T* /*ddsph*/, void* /*cuda_stream*/ ) { throw std::runtime_error("sphericart was not compiled with CUDA support"); } From 1b9f52502a9fe1b559fae240956f0846616c614f Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Thu, 29 Aug 2024 14:49:35 +0200 Subject: [PATCH 07/12] Fix more breathe rendering issues --- sphericart/include/sphericart.h | 36 +++--- sphericart/include/sphericart.hpp | 160 ++++++++++++------------- sphericart/include/sphericart_cuda.hpp | 18 +-- 3 files changed, 102 insertions(+), 112 deletions(-) diff --git a/sphericart/include/sphericart.h b/sphericart/include/sphericart.h index 1013cd942..03f1d9203 100644 --- a/sphericart/include/sphericart.h +++ b/sphericart/include/sphericart.h @@ -133,16 +133,16 @@ SPHERICART_EXPORT void sphericart_spherical_harmonics_delete_f( * respectively. * @param xyz_length size of the xyz allocation, i.e, `3 x n_samples` * @param sph pointer to the first element of an array containing `n_samples - * x (l_max + 1) x (l_max + 1)` elements. On exit, this array will contain + * x (l_max + 1)^2` elements. On exit, this array will contain * the spherical harmonics organized along two dimensions. The leading * dimension is `n_samples` long and it represents the different - * samples, while the inner dimension size is `(l_max + 1) x (l_max + 1)` + * samples, while the inner dimension size is `(l_max + 1)^2` * long and it contains the spherical harmonics. These are laid out in * lexicographic order. For example, if `l_max=2`, it will contain * `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, * 1), (2, 2)`, in this order. * @param sph_length size of the sph allocation, should be `n_samples x - * (l_max + 1) x (l_max + 1)` + * (l_max + 1)^2` */ SPHERICART_EXPORT void sphericart_spherical_harmonics_compute_array( sphericart_spherical_harmonics_calculator_t* calculator, @@ -167,28 +167,28 @@ SPHERICART_EXPORT void sphericart_spherical_harmonics_compute_array( * respectively. * @param xyz_length size of the xyz allocation, i.e, `3 x n_samples`` * @param sph pointer to the first element of an array containing `n_samples - * x (l_max + 1) x (l_max + 1)` elements. On exit, this array will contain + * x (l_max + 1)^2` elements. On exit, this array will contain * the spherical harmonics organized along two dimensions. The leading * dimension is `n_samples` long and it represents the different - * samples, while the inner dimension size is ``(l_max + 1) x (l_max + 1)`` + * samples, while the inner dimension size is ``(l_max + 1)^2`` * long and it contains the spherical harmonics. These are laid out in * lexicographic order. For example, if `l_max=2`, it will contain * `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, * 1), (2, 2)`, in this order. * @param sph_length size of the sph allocation, should be `n_samples * - * (l_max + 1) x (l_max + 1)` + * (l_max + 1)^2` * @param dsph pointer to the first element of an array containing `n_samples - * x `n_samples x 3 x (l_max + 1) x (l_max + 1)` elements. On exit, this + * x `n_samples x 3 x (l_max + 1)^2` elements. On exit, this * array will contain the spherical harmonics' derivatives organized * along three dimensions. As for the `sph` parameter, the leading * dimension represents the different samples, while the inner-most - * dimension size is `(l_max + 1) x (l_max + 1)`, and it represents + * dimension size is `(l_max + 1)^2`, and it represents * the degree and order of the spherical harmonics (again, organized in * lexicographic order). The intermediate dimension corresponds to * different spatial derivatives of the spherical harmonics: x, y, * and z, respectively. * @param dsph_length size of the dsph allocation, which should be `n_samples - * x 3 x (l_max + 1) x (l_max + 1)` + * x 3 x (l_max + 1)^2` */ SPHERICART_EXPORT void sphericart_spherical_harmonics_compute_array_with_gradients( sphericart_spherical_harmonics_calculator_t* calculator, @@ -215,40 +215,40 @@ SPHERICART_EXPORT void sphericart_spherical_harmonics_compute_array_with_gradien * respectively. * @param xyz_length size of the xyz allocation, i.e, ``3 x n_samples`` * @param sph pointer to the first element of an array containing - * `n_samples x (l_max + 1) x (l_max + 1)` elements. On exit, this array + * `n_samples x (l_max + 1)^2` elements. On exit, this array * will contain the spherical harmonics organized along two dimensions. * The leading dimension is `n_samples` long and it represents the different - * samples, while the inner dimension size is `(l_max + 1) x (l_max + 1)` + * samples, while the inner dimension size is `(l_max + 1)^2` * long and it contains the spherical harmonics. These are laid out in * lexicographic order. For example, if `l_max=2`, it will contain * ``(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), (2, 0), (2, * 1), (2, 2)``, in this order. * @param sph_length size of the sph allocation, should be `n_samples * - * (l_max + 1) x (l_max + 1)` + * (l_max + 1)^2` * @param dsph pointer to the first element of an array containing - * `n_samples x 3 x (l_max + 1) x (l_max + 1)` elements. On exit, + * `n_samples x 3 x (l_max + 1)^2` elements. On exit, * this array will contain the spherical harmonics' derivatives organized * along three dimensions. As for the `sph` parameter, the leading * dimension represents the different samples, while the inner-most - * dimension size is `(l_max + 1) x (l_max + 1)`, and it represents + * dimension size is `(l_max + 1)^2`, and it represents * the degree and order of the spherical harmonics (again, organized in * lexicographic order). The intermediate dimension corresponds to * different spatial derivatives of the spherical harmonics: x, y, * and z, respectively. * @param dsph_length size of the dsph allocation, which should be `n_samples - * x 3 x (l_max + 1) x (l_max + 1)` + * x 3 x (l_max + 1)^2` * @param ddsph pointer to the first element of an array containing - * `n_samples x 3 x 3 x (l_max + 1) x (l_max + 1)` elements. On exit, + * `n_samples x 3 x 3 x (l_max + 1)^2` elements. On exit, * this array will contain the spherical harmonics' second derivatives * organized along four dimensions. As for the `sph` parameter, the leading * dimension represents the different samples, while the inner-most dimension - * size is `(l_max + 1) x (l_max + 1)`, and it represents the degree and + * size is `(l_max + 1)^2`, and it represents the degree and * order of the spherical harmonics (again, organized in lexicographic * order). The intermediate dimensions correspond to the different spatial * second derivatives of the spherical harmonics, i.e., to the dimensions of * the hessian matrix. * @param ddsph_length size of the dsph allocation, which should be - * `n_samples x 3 x 3* (l_max + 1) x (l_max + 1)` + * `n_samples x 3 x 3* (l_max + 1)^2` */ SPHERICART_EXPORT void sphericart_spherical_harmonics_compute_array_with_hessians( sphericart_spherical_harmonics_calculator_t* calculator, diff --git a/sphericart/include/sphericart.hpp b/sphericart/include/sphericart.hpp index 74a9c24aa..089f435f7 100644 --- a/sphericart/include/sphericart.hpp +++ b/sphericart/include/sphericart.hpp @@ -54,14 +54,14 @@ template class SphericalHarmonics { * single point, the class will call a simpler function that * directly evaluates the point, without a loop. * @param sph On entry, a (possibly uninitialized) `std::vector`, which, if - * needed, will be resized to `n_samples * (l_max + 1) * (l_max + - * 1)`. On exit, this array will contain the spherical harmonics organized + * needed, will be resized to `n_samples x (l_max + 1)^2. + * On exit, this array will contain the spherical harmonics organized * along two dimensions. The leading dimension is `n_samples` long - * and it represents the different samples, while the inner dimension is - * `(l_max + 1) * (l_max + 1)` long and it contains the spherical harmonics. - * These are laid out in lexicographic order. For example, if `l_max=2`, it - * will contain `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), - * (2, 0), (2, 1), (2, 2)`, in this order. + * and it represents the different samples, while the inner dimension is + * `(l_max + 1)^2` long and it contains the spherical harmonics. + * These are laid out in lexicographic order. For example, if `l_max=2`, it + * will contain `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), + * (2, 0), (2, 1), (2, 2)`, in this order. */ void compute(const std::vector& xyz, std::vector& sph); @@ -78,21 +78,21 @@ template class SphericalHarmonics { * single point, the class will call a simpler functions that * directly evaluates the point, without a loop. * @param sph On entry, a (possibly uninitialized) `std::vector`, which, if - * needed, will be resized to `n_samples * (l_max + 1) * (l_max + - * 1)`. On exit, this array will contain the spherical harmonics organized + * needed, will be resized to `n_samples x (l_max + 1)^2`. + * On exit, this array will contain the spherical harmonics organized * along two dimensions. The leading dimension is `n_samples` long - * and it represents the different samples, while the inner dimension size - * is `(l_max + 1) * (l_max + 1)` long and it contains the spherical + * and it represents the different samples, while the inner dimension size + * is `(l_max + 1)^2` long and it contains the spherical * harmonics. These are laid out in lexicographic order. For example, * if `l_max=2`, it will contain `(l, m) = (0, 0), (1, -1), (1, 0), * (1, 1), (2, -2), (2, -1), (2, 0), (2, 1), (2, 2)`, in this order. * @param dsph `std::vector` for spherical harmonics derivatives. * It is a (possibly uninitialized) `std::vector`, which, if needed, - * will be resized to `n_samples * 3 * (l_max + 1) * (l_max + 1)`. On + * will be resized to `n_samples x 3 x (l_max + 1)^2`. On * exit, this array will contain the derivatives of the spherical - * harmonics organized along three dimensions. As for the `sph` parameter, - * the leading dimension represents the different samples, while the - * inner-most dimension size is `(l_max + 1) * (l_max + 1)`, and it + * harmonics organized along three dimensions. As for the `sph` parameter, + * the leading dimension represents the different samples, while the + * inner-most dimension size is `(l_max + 1)^2`, and it * represents the degree and order of the spherical harmonics (again, * organized in lexicographic order). The intermediate dimension * corresponds to different spatial derivatives of the spherical @@ -113,36 +113,36 @@ template class SphericalHarmonics { * single point, the class will call a simpler functions that * directly evaluates the point, without a loop. * @param sph On entry, a (possibly uninitialized) `std::vector`, which, if - * needed, will be resized to `n_samples * (l_max + 1) * (l_max + - * 1)`. On exit, this array will contain the spherical harmonics organized + * needed, will be resized to `n_samples x (l_max + 1) ^ 2`. + * On exit, this array will contain the spherical harmonics organized * along two dimensions. The leading dimension is `n_samples` long * and it represents the different samples, while the inner dimension size - * is `(l_max + 1) * (l_max + 1)` long and it contains the spherical + * is `(l_max + 1)^2` long and it contains the spherical * harmonics. These are laid out in lexicographic order. For example, * if `l_max=2`, it will contain `(l, m) = (0, 0), (1, -1), (1, 0), * (1, 1), (2, -2), (2, -1), (2, 0), (2, 1), (2, 2)`, in this order. * @param dsph `std::vector` for spherical harmonics derivatives. * It is a (possibly uninitialized) `std::vector`, which, if needed, - * will be resized to `n_samples * 3 * (l_max + 1) * (l_max + 1)`. On + * will be resized to `n_samples x 3 x (l_max + 1)^2`. On * exit, this array will contain the derivatives of the spherical - * harmonics organized along three dimensions. As for the `sph` parameter, - * the leading dimension represents the different samples, while the - * inner-most dimension size is `(l_max + 1) * (l_max + 1)`, and it + * harmonics organized along three dimensions. As for the `sph` parameter, + * the leading dimension represents the different samples, while the + * inner-most dimension size is `(l_max + 1)^2`, and it * represents the degree and order of the spherical harmonics (again, * organized in lexicographic order). The intermediate dimension * corresponds to different spatial derivatives of the spherical * harmonics: x, y, and z, respectively. * @param ddsph `std::vector` for spherical harmonics second derivatives. * It is a (possibly uninitialized) `std::vector`, which, if needed, - * will be resized to `n_samples * 3 * 3 * (l_max + 1) * (l_max + - * 1)`. On exit, this array will contain the second derivatives of the - * spherical harmonics organized along four dimensions. As for the `sph` - * parameter, the leading dimension represents the different samples, while - * the inner-most dimension size is `(l_max + 1) * (l_max + 1)`, and it + * will be resized to `n_samples x 3 x 3 x (l_max + 1)^2`. + * On exit, this array will contain the second derivatives of the + * spherical harmonics organized along four dimensions. As for the `sph` + * parameter, the leading dimension represents the different samples, while + * the inner-most dimension size is `(l_max + 1)^2`, and it * represents the degree and order of the spherical harmonics (again, * organized in lexicographic order). The intermediate dimensions * correspond to the different spatial second derivatives of the - * spherical harmonics, i.e., to the dimensions of the Hessian matrix. + * spherical harmonics, i.e., to the dimensions of the Hessian matrix. */ void compute_with_hessians( const std::vector& xyz, std::vector& sph, std::vector& dsph, std::vector& ddsph @@ -157,7 +157,7 @@ template class SphericalHarmonics { * outer dimension is `n_samples` long, accounting for different * samples, while the inner dimension has size 3 and it represents * the x, y, and z coordinates respectively. - * @param xyz_length Total length of the `xyz` array: `n_samples * 3`. + * @param xyz_length Total length of the `xyz` array: `n_samples x 3`. */ void compute_array(const T* xyz, size_t xyz_length, T* sph, size_t sph_length); @@ -170,30 +170,29 @@ template class SphericalHarmonics { * outer dimension is `n_samples` long, accounting for different * samples, while the inner dimension has size 3 and it represents * the x, y, and z coordinates respectively. - * @param xyz_length Total length of the `xyz` array: `n_samples * 3`. - * @param sph On entry, an array of size `n_samples * (l_max + 1) * (l_max + - * 1)`. On exit, this array will contain the spherical harmonics organized - * along two dimensions. The leading dimension is `n_samples` long and it + * @param xyz_length Total length of the `xyz` array: `n_samples x 3`. + * @param sph On entry, an array of size `n_samples x (l_max + 1)^2`. + * On exit, this array will contain the spherical harmonics organized + * along two dimensions. The leading dimension is `n_samples` long and it * represents the different samples, while the inner dimension size - * is `(l_max + 1) * (l_max + 1)` long and it contains the spherical + * is `(l_max + 1)^2` long and it contains the spherical * harmonics. These are laid out in lexicographic order. For example, * if `l_max=2`, it will contain `(l, m) = (0, 0), (1, -1), (1, 0), * (1, 1), (2, -2), (2, -1), (2, 0), (2, 1), (2, 2)`, in this order. - * @param sph_length Total length of the `sph` array: `n_samples * (l_max + - * 1) - * * (l_max + 1)`. + * @param sph_length Total length of the `sph` array: `n_samples x (l_max + + * 1)^2`. * @param dsph Array for spherical harmonics derivatives. - * It is an array of size `n_samples * 3 * (l_max + 1) * (l_max + - * 1)`. On exit, this array will contain the derivatives of the spherical - * harmonics organized along three dimensions. As for the `sph` parameter, - * the leading dimension represents the different samples, while the - * inner-most dimension size is `(l_max + 1) * (l_max + 1)`, and it + * It is an array of size `n_samples x 3 x (l_max + 1)^2`. + * On exit, this array will contain the derivatives of the spherical + * harmonics organized along three dimensions. As for the `sph` parameter, + * the leading dimension represents the different samples, while the + * inner-most dimension size is `(l_max + 1)^2`, and it * represents the degree and order of the spherical harmonics (again, * organized in lexicographic order). The intermediate dimension * corresponds to the different spatial derivatives of the spherical * harmonics: x, y, and z, respectively. - * @param dsph_length Total length of the `dsph` array: `n_samples * 3 * - * (l_max + 1) * (l_max + 1)`. + * @param dsph_length Total length of the `dsph` array: `n_samples x 3 x + * (l_max + 1)^2`. */ void compute_array_with_gradients( const T* xyz, size_t xyz_length, T* sph, size_t sph_length, T* dsph, size_t dsph_length @@ -208,42 +207,41 @@ template class SphericalHarmonics { * outer dimension is `n_samples` long, accounting for different * samples, while the inner dimension has size 3 and it represents * the x, y, and z coordinates respectively. - * @param xyz_length Total length of the `xyz` array: `n_samples * 3`. - * @param sph On entry, an array of size `n_samples * (l_max + 1) * (l_max + - * 1)`. On exit, this array will contain the spherical harmonics organized - * along two dimensions. The leading dimension is `n_samples` long and it + * @param xyz_length Total length of the `xyz` array: `n_samples x 3`. + * @param sph On entry, an array of size `n_samples x (l_max + 1)^2`. + * On exit, this array will contain the spherical harmonics organized + * along two dimensions. The leading dimension is `n_samples` long and it * represents the different samples, while the inner dimension size - * is `(l_max + 1) * (l_max + 1)` long and it contains the spherical + * is `(l_max + 1)^2` long and it contains the spherical * harmonics. These are laid out in lexicographic order. For example, * if `l_max=2`, it will contain `(l, m) = (0, 0), (1, -1), (1, 0), * (1, 1), (2, -2), (2, -1), (2, 0), (2, 1), (2, 2)`, in this order. - * @param sph_length Total length of the `sph` array: `n_samples * (l_max + - * 1) - * * (l_max + 1)`. + * @param sph_length Total length of the `sph` array: `n_samples x (l_max + + * 1)^2`. * @param dsph Array for spherical harmonics derivatives. - * It is an array of size `n_samples * 3 * (l_max + 1) * (l_max + - * 1)`. On exit, this array will contain the derivatives of the spherical + * It is an array of size `n_samples x 3 x (l_max + 1)^2`. + * On exit, this array will contain the derivatives of the spherical * harmonics organized along three dimensions. As for the `sph` parameter, * the leading dimension represents the different samples, while the - * inner-most dimension size is `(l_max + 1) * (l_max + 1)`, and it + * inner-most dimension size is `(l_max + 1)^2`, and it * represents the degree and order of the spherical harmonics (again, * organized in lexicographic order). The intermediate dimension * corresponds to the different spatial derivatives of the spherical * harmonics: x, y, and z, respectively. - * @param dsph_length Total length of the `dsph` array: `n_samples * 3 * - * (l_max + 1) * (l_max + 1)`. + * @param dsph_length Total length of the `dsph` array: `n_samples x 3 x + * (l_max + 1)^2`. * @param ddsph Array for spherical harmonics second derivatives. - * It is an array of size `n_samples * 3 * 3 * (l_max + 1) * (l_max + - * 1)`. On exit, this array will contain the second derivatives of the + * It is an array of size `n_samples x 3 x 3 x (l_max + 1)^2`. + * On exit, this array will contain the second derivatives of the * spherical harmonics organized along four dimensions. As for the `sph` * parameter, the leading dimension represents the different samples, while - * the inner-most dimension size is `(l_max + 1) * (l_max + 1)`, and it + * the inner-most dimension size is `(l_max + 1)^2`, and it * represents the degree and order of the spherical harmonics (again, * organized in lexicographic order). The intermediate dimensions * correspond to the different spatial second derivatives of the * spherical harmonics, i.e., to the dimensions of the Hessian matrix. - * @param ddsph_length Total length of the `ddsph` array: `n_samples * 9 * - * (l_max + 1) * (l_max + 1)`. + * @param ddsph_length Total length of the `ddsph` array: `n_samples x 9 x + * (l_max + 1)^2`. */ void compute_array_with_hessians( const T* xyz, @@ -264,14 +262,13 @@ template class SphericalHarmonics { * harmonics are to be computed. x, y, and z coordinates * respectively. * @param xyz_length Length of the `xyz` array: 3. - * @param sph On entry, an array of size `(l_max + 1) * (l_max + 1)`. + * @param sph On entry, an array of size `(l_max + 1)^2`. * On exit, this array will contain the spherical harmonics laid out * in lexicographic order. For example, * if `l_max=2`, it will contain the spherical harmonics in the * following order: `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, * -1), (2, 0), (2, 1), (2, 2)`. - * @param sph_length Total length of the `sph` array: `(l_max + 1) * (l_max - * + 1)`. + * @param sph_length Total length of the `sph` array: `(l_max + 1)^2`. */ void compute_sample(const T* xyz, size_t xyz_length, T* sph, size_t sph_length); @@ -283,25 +280,22 @@ template class SphericalHarmonics { * harmonics are to be computed. x, y, and z coordinates * respectively. * @param xyz_length Length of the `xyz` array: 3. - * @param sph On entry, an array of size `(l_max + 1) * (l_max + 1)`. + * @param sph On entry, an array of size `(l_max + 1)^2`. * On exit, this array will contain the spherical harmonics laid out * in lexicographic order. For example, * if `l_max=2`, it will contain the spherical harmonics in the * following order: `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, * -1), (2, 0), (2, 1), (2, 2)`. - * @param sph_length Total length of the `sph` array: `(l_max + 1) * (l_max - * + 1)`. + * @param sph_length Total length of the `sph` array: `(l_max + 1)^2`. * @param dsph Array for spherical harmonics derivatives. - * It is an array of size `3 * (l_max + 1) * (l_max + 1)`. + * It is an array of size `3 x (l_max + 1)^2`. * On exit, this array will contain the spherical harmonics' * derivatives organized along two dimensions. The second dimension's size - * is - * `(l_max + 1) * (l_max + 1)`, and it represents the degree and order of + * is `(l_max + 1)^2`, and it represents the degree and order of * the spherical harmonics (again, organized in lexicographic order). The * first dimension corresponds to the different spatial derivatives of the * spherical harmonics: x, y, and z, respectively. - * @param dsph_length Total length of the `dsph` array: `3 * (l_max + 1) * - * (l_max + 1)`. + * @param dsph_length Total length of the `dsph` array: `3 x (l_max + 1)^2`. */ void compute_sample_with_gradients( const T* xyz, size_t xyz_length, T* sph, size_t sph_length, T* dsph, size_t dsph_length @@ -315,36 +309,32 @@ template class SphericalHarmonics { * harmonics are to be computed. x, y, and z coordinates * respectively. * @param xyz_length Length of the `xyz` array: 3. - * @param sph On entry, an array of size `(l_max + 1) * (l_max + 1)`. + * @param sph On entry, an array of size `(l_max + 1)^2`. * On exit, this array will contain the spherical harmonics laid out * in lexicographic order. For example, * if `l_max=2`, it will contain the spherical harmonics in the * following order: `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, * -1), (2, 0), (2, 1), (2, 2)`. - * @param sph_length Total length of the `sph` array: `(l_max + 1) * (l_max - * + 1)`. + * @param sph_length Total length of the `sph` array: `(l_max + 1)^2`. * @param dsph Array for spherical harmonics derivatives. - * It is an array of size `3 * (l_max + 1) * (l_max + 1)`. + * It is an array of size `3 x (l_max + 1)^2`. * On exit, this array will contain the spherical harmonics' * derivatives organized along two dimensions. The second dimension's size - * is - * `(l_max + 1) * (l_max + 1)`, and it represents the degree and order of + * is `(l_max + 1)^2`, and it represents the degree and order of * the spherical harmonics (again, organized in lexicographic order). The * first dimension corresponds to the different spatial derivatives of the * spherical harmonics: x, y, and z, respectively. - * @param dsph_length Total length of the `dsph` array: `3 * (l_max + 1) * - * (l_max + 1)`. + * @param dsph_length Total length of the `dsph` array: `3 x (l_max + 1)^2`. * @param ddsph Array for spherical harmonics second derivatives. - * It is an array of size `3 * 3 * (l_max + 1) * (l_max + 1)`. + * It is an array of size `3 x 3 x (l_max + 1)^2`. * On exit, this array will contain the second derivatives of the * spherical harmonics organized along three dimensions. As for the `sph` - * parameter, the inner-most dimension size is `(l_max + 1) * (l_max + 1)`, + * parameter, the inner-most dimension size is `(l_max + 1)^2`, * and it represents the degree and order of the spherical harmonics (again, * organized in lexicographic order). The first two dimensions * correspond to the different spatial second derivatives of the * spherical harmonics, i.e., to the dimensions of the Hessian matrix. - * @param ddsph_length Total length of the `ddsph` array: `9 * (l_max + 1) * - * (l_max + 1)`. + * @param ddsph_length Total length of the `ddsph` array: `9 x (l_max + 1)^2`. */ void compute_sample_with_hessians( const T* xyz, diff --git a/sphericart/include/sphericart_cuda.hpp b/sphericart/include/sphericart_cuda.hpp index 3df183845..04e1bd8b7 100644 --- a/sphericart/include/sphericart_cuda.hpp +++ b/sphericart/include/sphericart_cuda.hpp @@ -60,11 +60,11 @@ template class SphericalHarmonics { * the x, y, and z coordinates respectively. * @param n_samples Number of samples contained within `xyz`. * @param sph On entry, a preallocated device-side array of size `n_samples - * x (l_max + 1) x (l_max + 1)`. On exit, this array will contain the + * x (l_max + 1)^2`. On exit, this array will contain the * spherical harmonics organized along two dimensions. The leading dimension * is `n_samples` long and it represents the different samples, while the * inner dimension is - * `(l_max + 1) x (l_max + 1)` long and it contains the spherical harmonics. + * `(l_max + 1)^2` long and it contains the spherical harmonics. * These are laid out in lexicographic order. For example, if `l_max=2`, it * will contain `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), * (2, 0), (2, 1), (2, 2)`, in this order. @@ -84,16 +84,16 @@ template class SphericalHarmonics { * the x, y, and z coordinates respectively. * @param n_samples Number of samples contained within `xyz`. * @param sph On entry, a preallocated device-side array of size `n_samples - * * (l_max + 1) x (l_max + 1)`. On exit, this array will contain the + * x (l_max + 1)^2`. On exit, this array will contain the * spherical harmonics organized along two dimensions. The leading dimension * is `n_samples` long and it represents the different samples, while the * inner dimension is - * `(l_max + 1) x (l_max + 1)` long and it contains the spherical harmonics. + * `(l_max + 1)^2` long and it contains the spherical harmonics. * These are laid out in lexicographic order. For example, if `l_max=2`, it * will contain `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), * (2, 0), (2, 1), (2, 2)`, in this order. * @param dsph On entry, nullptr or a preallocated device-side array of size - * `n_samples x 3 x (l_max + 1) x (l_max + 1)`. If the pointer is not + * `n_samples x 3 x (l_max + 1)^2`. If the pointer is not * nullptr, then compute_with_gradients must also be true in order for * gradients to be computed. * @param cuda_stream Pointer to a cudaStream_t or nullptr. If this is @@ -118,20 +118,20 @@ template class SphericalHarmonics { * @param compute_with_hessians Whether we should compute ddsph. If true, * the pointer ddsph must also be allocated on device. * @param sph On entry, a preallocated device-side array of size `n_samples - * * (l_max + 1) x (l_max + 1)`. On exit, this array will contain the + * x (l_max + 1)^2`. On exit, this array will contain the * spherical harmonics organized along two dimensions. The leading dimension * is `n_samples` long and it represents the different samples, while the * inner dimension is - * `(l_max + 1) x (l_max + 1)` long and it contains the spherical harmonics. + * `(l_max + 1)^2` long and it contains the spherical harmonics. * These are laid out in lexicographic order. For example, if `l_max=2`, it * will contain `(l, m) = (0, 0), (1, -1), (1, 0), (1, 1), (2, -2), (2, -1), * (2, 0), (2, 1), (2, 2)`, in this order. * @param dsph On entry, nullptr or a preallocated device-side array of size - * `n_samples x 3 x (l_max + 1) x (l_max + 1)`. If the pointer is not + * `n_samples x 3 x (l_max + 1)^2`. If the pointer is not * nullptr, then compute_with_gradients must also be true in order for * gradients to be computed. * @param ddsph On entry, nullptr or a preallocated device-side array of - * size `n_samples x 3 x 3 x (l_max + 1) x (l_max + 1)`. If the pointer is + * size `n_samples x 3 x 3 x (l_max + 1)^2`. If the pointer is * not nullptr, then compute_with_hessians must also be true in order for * gradients to be computed. * @param cuda_stream Pointer to a cudaStream_t or nullptr. If this is From 81252e145a228557e62933ea61cf7cb85716fbb2 Mon Sep 17 00:00:00 2001 From: Michele Ceriotti Date: Fri, 30 Aug 2024 16:06:35 +0200 Subject: [PATCH 08/12] Fixed a few typos --- docs/src/_static/sphericart_icon.svg | 222 +++++++++--------- docs/src/index.rst | 7 +- docs/src/installation.rst | 7 +- docs/src/maths.rst | 11 +- .../sphericart/jax/spherical_harmonics.py | 2 +- 5 files changed, 130 insertions(+), 119 deletions(-) diff --git a/docs/src/_static/sphericart_icon.svg b/docs/src/_static/sphericart_icon.svg index b29b4a481..ffdc5e991 100644 --- a/docs/src/_static/sphericart_icon.svg +++ b/docs/src/_static/sphericart_icon.svg @@ -2,13 +2,13 @@ `_. +The theory behind this efficient implementation is detailed in +`this paper `_. The core library is implemented in C++ (with OpenMP parallelism) and CUDA. It also provides APIs for C, Python (NumPy), PyTorch and JAX. The torch and JAX implementations provide fast spherical harmonics on GPUs. -A native Julia package is also available. +A native Julia package is also available, contributed by +`Christoph Ortner `_. This documentation contains an installation guide, an API overview, some examples of how to use the library, and a brief explanation of the mathematics involved. diff --git a/docs/src/installation.rst b/docs/src/installation.rst index 8df3f4d6b..0bc54f0a4 100644 --- a/docs/src/installation.rst +++ b/docs/src/installation.rst @@ -36,15 +36,16 @@ and/or you want to use it in JAX, you should build the code from source: pip install --extra-index-url https://download.pytorch.org/whl/cpu .[torch] Before installing the JAX version of ``sphericart``, make sure you already have the JAX -library installed according to the official JAX installation instructions at -. +library installed according to the `official JAX installation instructions +`_. Julia package ------------- The native Julia package can be installed by opening a REPL, -switching to the package manager by typing ``]`` and then ``add SpheriCart``. +switching to the package manager by typing ``]`` and then executing +the command ``add SpheriCart``. C/C++/CUDA library diff --git a/docs/src/maths.rst b/docs/src/maths.rst index eec92c352..f234f4179 100644 --- a/docs/src/maths.rst +++ b/docs/src/maths.rst @@ -16,14 +16,19 @@ Within `sphericart`, we compute only real-valued spherical harmonics and we expr them as a function of the full Cartesian coordinates of a point in three dimensions. These correspond to the real spherical harmonics as defined in the corresponding `Wikipedia article `_, which we -refer to as :math:`Y^m_l`. +refer to as :math:`Y^m_l`. +If you need complex spherical harmonics, or use a different convention for normalization +and storage order it is usually simple - if tedious and inefficient - to perform the +conversion manually, see :ref:`spherical-complex` for a simple example. + We also offer the possibility to compute "solid" harmonics, which are given by :math:`\tilde{Y}^m_l = r^l\,{Y}_l^m`. Since these can be expressed as homogeneous polynomials of the Cartesian coordinates :math:`(x,y,z)`, as opposed to :math:`(x/r,y/r,z/r)`, they are less computationally expensive to evaluate. -Besides being slightly faster, they can also -provide a more natural scaling if used together with a radial expansion. +Besides being slightly faster, they can also provide a more natural scaling if +used together with a radial expansion, and we recommend using them unless you +need the normalized version. The formulas used to compute the solid harmonics (and, with few modifications, also for the spherical harmonics) are: diff --git a/sphericart-jax/python/sphericart/jax/spherical_harmonics.py b/sphericart-jax/python/sphericart/jax/spherical_harmonics.py index d9274ed92..4149e8e59 100644 --- a/sphericart-jax/python/sphericart/jax/spherical_harmonics.py +++ b/sphericart-jax/python/sphericart/jax/spherical_harmonics.py @@ -42,7 +42,7 @@ def solid_harmonics(xyz: jax.Array, l_max: int): These are a non-normalized form of the real spherical harmonics, i.e. :math:`r^lY^m_l`. These scaled spherical harmonics are polynomials in the Cartesian coordinates of the input points, and they - are therefore less expoensive to compute. + are therefore less expensive to compute. """ if xyz.shape[-1] != 3: raise ValueError("the last axis of xyz must have size 3") From c341a710d6692832d2b13ed4a625894042224cf7 Mon Sep 17 00:00:00 2001 From: Michele Ceriotti Date: Fri, 30 Aug 2024 16:23:38 +0200 Subject: [PATCH 09/12] Ooops --- docs/src/maths.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/maths.rst b/docs/src/maths.rst index f234f4179..19988e610 100644 --- a/docs/src/maths.rst +++ b/docs/src/maths.rst @@ -19,7 +19,7 @@ These correspond to the real spherical harmonics as defined in the corresponding refer to as :math:`Y^m_l`. If you need complex spherical harmonics, or use a different convention for normalization and storage order it is usually simple - if tedious and inefficient - to perform the -conversion manually, see :ref:`spherical-complex` for a simple example. +conversion manually, see :doc:`spherical-complex` for a simple example. We also offer the possibility to compute "solid" harmonics, which are given by From 9e06dd8dc2fa4e90af8005a7d33fac6bb441e256 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Fri, 30 Aug 2024 17:30:22 +0200 Subject: [PATCH 10/12] Update installation instructions for CUDA torch/JAX --- docs/src/installation.rst | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/src/installation.rst b/docs/src/installation.rst index 0bc54f0a4..66ea42efa 100644 --- a/docs/src/installation.rst +++ b/docs/src/installation.rst @@ -39,6 +39,10 @@ Before installing the JAX version of ``sphericart``, make sure you already have library installed according to the `official JAX installation instructions `_. +In addition, if you want to use the CUDA functionalities of sphericart (either with torch +or JAX), make sure you have installed the CUDA toolkit and set up the environment variables +``CUDA_HOME``, ``LD_LIBRARY_FLAGS``, and ``PATH`` accordingly. + Julia package ------------- From c2a8b04cdf7a45f66e3dafda9486102d97899be7 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Fri, 30 Aug 2024 17:44:23 +0200 Subject: [PATCH 11/12] Fix double backticks --- python/src/sphericart/spherical_harmonics.py | 8 +-- .../python/sphericart/torch/__init__.py | 56 +++++++++---------- 2 files changed, 32 insertions(+), 32 deletions(-) diff --git a/python/src/sphericart/spherical_harmonics.py b/python/src/sphericart/spherical_harmonics.py index ad549017a..d1b859ada 100644 --- a/python/src/sphericart/spherical_harmonics.py +++ b/python/src/sphericart/spherical_harmonics.py @@ -34,7 +34,7 @@ class SphericalHarmonics: (10, 3, 81) which returns the gradient as a tensor with size - `(n_samples, 3, (l_max+1)**2)`. + ``(n_samples, 3, (l_max+1)**2)``. :param l_max: the maximum degree of the spherical harmonics to be calculated @@ -141,7 +141,7 @@ def compute_with_gradients(self, xyz: np.ndarray) -> Tuple[np.ndarray, np.ndarra :return: A tuple containing: - an array of shape ``(n_samples, (l_max+1)**2)`` containing all the - spherical harmonics up to degree `l_max` in lexicographic order. + 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. @@ -230,7 +230,7 @@ def compute_with_hessians( :return: A tuple containing: - an array of shape ``(n_samples, (l_max+1)**2)`` containing all the - spherical harmonics up to degree `l_max` in lexicographic order. + 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. @@ -318,7 +318,7 @@ class SolidHarmonics: :param l_max: the maximum degree of the solid harmonics to be calculated - :return: a calculator, in the form of a `SolidHarmonics` object + :return: a calculator, in the form of a ``SolidHarmonics`` object """ def __init__(self, l_max: int): diff --git a/sphericart-torch/python/sphericart/torch/__init__.py b/sphericart-torch/python/sphericart/torch/__init__.py index d3a9e64d2..a23a67ee3 100644 --- a/sphericart-torch/python/sphericart/torch/__init__.py +++ b/sphericart-torch/python/sphericart/torch/__init__.py @@ -89,7 +89,7 @@ class SphericalHarmonics(torch.nn.Module): >>> sh_grads.shape torch.Size([10, 3, 81]) - Alternatively, if `compute()` is used, or if the class is called directly, + Alternatively, if ``compute()`` is used, or if the class is called directly, the outputs support single and double backpropagation. >>> xyz = xyz.detach().clone().requires_grad_() @@ -100,12 +100,12 @@ class SphericalHarmonics(torch.nn.Module): True By default, only single backpropagation with respect to `xyz` is - enabled (this includes mixed second derivatives where `xyz` appears + 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 + 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. @@ -114,12 +114,12 @@ class SphericalHarmonics(torch.nn.Module): :param l_max: the maximum degree of the spherical harmonics to be calculated :param 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 + 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`. + differentiation step will be possible with respect to ``xyz``. :return: a calculator, in the form of a SphericalHarmonics object """ @@ -139,14 +139,14 @@ def forward(self, 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 + 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, + 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. + double reverse-mode differentiation if ``backward_second_derivatives`` + was set to ``True`` during class creation. :param xyz: The Cartesian coordinates of the 3D points, as a `torch.Tensor` with @@ -171,7 +171,7 @@ def compute_with_gradients(self, xyz: Tensor) -> Tuple[Tensor, Tensor]: 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, + 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. @@ -184,7 +184,7 @@ def compute_with_gradients(self, xyz: Tensor) -> Tuple[Tensor, Tensor]: 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. + 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. @@ -203,20 +203,20 @@ def compute_with_hessians(self, xyz: Tensor) -> Tuple[Tensor, Tensor, Tensor]: 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, + 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. :param xyz: - The Cartesian coordinates of the 3D points, as a `torch.Tensor` with + The Cartesian coordinates of the 3D points, as a ``torch.Tensor`` with shape ``(n_samples, 3)``. :return: 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. + 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. @@ -256,12 +256,12 @@ class SolidHarmonics(torch.nn.Module): :param l_max: the maximum degree of the spherical harmonics to be calculated :param 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 + 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`. + differentiation step will be possible with respect to ``xyz``. :return: a calculator, in the form of a SolidHarmonics object """ @@ -323,13 +323,13 @@ def e3nn_spherical_harmonics( l value are computed, so this may be inefficient for use cases requiring a single, or few, angular momentum channels. :param x: - A `torch.Tensor` containing the coordinates, in the same format - expected by the `e3nn` function. + A ``torch.Tensor`` containing the coordinates, in the same format + expected by the ``e3nn`` function. :param normalize: Flag specifying whether the input positions should be normalized (resulting in the computation of the spherical harmonics :math:`Y^m_l`), or whether the function should compute the solid harmonics - :math:`\tilde{Y}^m_l=r^lY^m_l`. + :math:`r^lY^m_l`. :param normalization: String that can be "integral", "norm", "component", that controls a further scaling of the :math:`Y^m_l`. See the From cde80263077a171455a983ab4219544dd6e3f561 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Sat, 31 Aug 2024 09:56:23 +0200 Subject: [PATCH 12/12] Add gradient and Hessian calculations to CUDA C++ examples --- examples/cuda/example.cu | 37 +++++++++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/examples/cuda/example.cu b/examples/cuda/example.cu index cb26f277d..569756584 100644 --- a/examples/cuda/example.cu +++ b/examples/cuda/example.cu @@ -58,6 +58,7 @@ int main() { // internal buffers and numerical factors are initalized at construction sphericart::cuda::SphericalHarmonics calculator_cuda(l_max); + // allcate device memory double* xyz_cuda; CUDA_CHECK(cudaMalloc(&xyz_cuda, n_samples * 3 * sizeof(double))); CUDA_CHECK( @@ -65,8 +66,21 @@ int main() { ); double* sph_cuda; CUDA_CHECK(cudaMalloc(&sph_cuda, n_samples * (l_max + 1) * (l_max + 1) * sizeof(double))); + double* dsph_cuda; + CUDA_CHECK(cudaMalloc(&dsph_cuda, n_samples * 3 * (l_max + 1) * (l_max + 1) * sizeof(double))); + double* ddsph_cuda; + CUDA_CHECK( + cudaMalloc(&ddsph_cuda, n_samples * 3 * 3 * (l_max + 1) * (l_max + 1) * sizeof(double)) + ); + // calculation examples calculator_cuda.compute(xyz_cuda, n_samples, sph_cuda); // no gradients + calculator_cuda.compute_with_gradients( + xyz_cuda, n_samples, sph_cuda, dsph_cuda + ); // with gradients + calculator_cuda.compute_with_hessians( + xyz_cuda, n_samples, sph_cuda, dsph_cuda, ddsph_cuda + ); // with gradients and hessians CUDA_CHECK(cudaMemcpy( sph.data(), sph_cuda, n_samples * (l_max + 1) * (l_max + 1) * sizeof(double), cudaMemcpyDeviceToHost @@ -82,8 +96,21 @@ int main() { ); float* sph_cuda_f; CUDA_CHECK(cudaMalloc(&sph_cuda_f, n_samples * (l_max + 1) * (l_max + 1) * sizeof(float))); + float* dsph_cuda_f; + CUDA_CHECK(cudaMalloc(&dsph_cuda_f, n_samples * 3 * (l_max + 1) * (l_max + 1) * sizeof(float))); + float* ddsph_cuda_f; + CUDA_CHECK( + cudaMalloc(&ddsph_cuda_f, n_samples * 3 * 3 * (l_max + 1) * (l_max + 1) * sizeof(float)) + ); + // calculation examples (float) calculator_cuda_f.compute(xyz_cuda_f, n_samples, sph_cuda_f); // no gradients + calculator_cuda_f.compute_with_gradients( + xyz_cuda_f, n_samples, sph_cuda_f, dsph_cuda_f + ); // with gradients + calculator_cuda_f.compute_with_hessians( + xyz_cuda_f, n_samples, sph_cuda_f, dsph_cuda_f, ddsph_cuda_f + ); // with gradients and hessians CUDA_CHECK(cudaMemcpy( sph_f.data(), @@ -101,5 +128,15 @@ int main() { } printf("Float vs double relative error: %12.8e\n", sqrt(sph_error / sph_norm)); + /* ===== free device memory ===== */ + CUDA_CHECK(cudaFree(xyz_cuda)); + CUDA_CHECK(cudaFree(sph_cuda)); + CUDA_CHECK(cudaFree(dsph_cuda)); + CUDA_CHECK(cudaFree(ddsph_cuda)); + CUDA_CHECK(cudaFree(xyz_cuda_f)); + CUDA_CHECK(cudaFree(sph_cuda_f)); + CUDA_CHECK(cudaFree(dsph_cuda_f)); + CUDA_CHECK(cudaFree(ddsph_cuda_f)); + return 0; }