diff --git a/.codecov.yml b/.codecov.yml index 3b8bba5e6..94bab7aef 100644 --- a/.codecov.yml +++ b/.codecov.yml @@ -1,15 +1,11 @@ coverage: - ignore: - - tests/.* status: project: default: target: 80% patch: off + ignore: + - "tests/.*" + - "examples/.*" -comment: off - -fixes: - # map coverage collected inside tox virtual environments - # to the source dir in git - - ".tox/all-deps/lib/*/site-packages/::python/" +comment: false diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 6949a6c0d..8ef085c1a 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -38,7 +38,7 @@ jobs: - name: install python dependencies run: | python -m pip install --upgrade pip - python -m pip install tox + python -m pip install tox coverage - name: install lcov run: sudo apt install -y lcov @@ -47,10 +47,12 @@ jobs: uses: actions/cache@v3 with: path: .tox - key: tox-${{ matrix.os }}-${{ hashFiles('pyproject.toml', 'setup.cfg', 'tox.ini') }} + key: tox-${{ hashFiles('pyproject.toml', 'setup.cfg', 'tox.ini') }} - name: Setup sccache uses: mozilla-actions/sccache-action@v0.0.3 + with: + version: "v0.5.4" - name: Setup sccache environnement variables run: | @@ -69,10 +71,22 @@ jobs: lcov --remove coverage.info '/usr/*' "$(pwd)/rascaline-c-api/tests/*" "$(pwd)/rascaline-c-api/examples/*" --output-file coverage.info - name: collect Python coverage - run: tox -e all-deps + run: | + tox -e all-deps + tox -e torch-tests + env: + # Use the CPU only version of torch when building/running the code + PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu + + - name: combine Python coverage files + run: | + coverage combine --append \ + ./.coverage \ + ./python/rascaline-torch/.coverage + coverage xml - name: upload to codecov.io uses: codecov/codecov-action@v3 with: fail_ci_if_error: true - files: target/tarpaulin/cobertura.xml,.tox/coverage.xml,coverage.info + files: target/tarpaulin/cobertura.xml,coverage.xml,coverage.info diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index 7627b9508..889b38b42 100644 --- a/.github/workflows/python-tests.yml +++ b/.github/workflows/python-tests.yml @@ -19,7 +19,7 @@ jobs: matrix: include: - os: ubuntu-20.04 - python-version: "3.7" + python-version: "3.8" - os: ubuntu-20.04 python-version: "3.11" - os: macos-11 @@ -47,6 +47,8 @@ jobs: - name: Setup sccache uses: mozilla-actions/sccache-action@v0.0.3 + with: + version: "v0.5.4" - name: Setup sccache environnement variables run: | @@ -73,7 +75,7 @@ jobs: name: Python ${{ matrix.python-version }} / check build strategy: matrix: - python-version: ['3.7', '3.11'] + python-version: ['3.8', '3.11'] os: [ubuntu-20.04] steps: - uses: actions/checkout@v3 diff --git a/.github/workflows/rust-tests.yml b/.github/workflows/rust-tests.yml index 20d7f83a0..999542e2e 100644 --- a/.github/workflows/rust-tests.yml +++ b/.github/workflows/rust-tests.yml @@ -100,6 +100,8 @@ jobs: - name: Setup sccache uses: mozilla-actions/sccache-action@v0.0.3 + with: + version: "v0.5.4" - name: Setup sccache environnement variables run: | @@ -140,6 +142,8 @@ jobs: - name: Setup sccache uses: mozilla-actions/sccache-action@v0.0.3 + with: + version: "v0.5.4" - name: Setup sccache environnement variables run: | diff --git a/.github/workflows/torch-tests.yml b/.github/workflows/torch-tests.yml index 79ec93e2d..c9557e22e 100644 --- a/.github/workflows/torch-tests.yml +++ b/.github/workflows/torch-tests.yml @@ -19,7 +19,7 @@ jobs: include: - os: ubuntu-20.04 torch-version: 1.11.* - python-version: "3.7" + python-version: "3.8" cargo-test-flags: --release - os: ubuntu-20.04 @@ -60,6 +60,8 @@ jobs: - name: Setup sccache uses: mozilla-actions/sccache-action@v0.0.3 + with: + version: "v0.5.4" - name: Setup sccache environnement variables run: | diff --git a/.gitignore b/.gitignore index 7bcc9343a..4f077544b 100644 --- a/.gitignore +++ b/.gitignore @@ -3,10 +3,12 @@ target/ Cargo.lock .tox/ -.coverage build/ dist/ *.egg-info __pycache__/ .vscode *.DS_Store +htmlcov/ +.coverage +coverage.xml diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst index 7b61cb772..ead7509a3 100644 --- a/CONTRIBUTING.rst +++ b/CONTRIBUTING.rst @@ -176,6 +176,24 @@ of just testing it. .. _`cargo` : https://doc.rust-lang.org/cargo/ .. _valgrind: https://valgrind.org/ +Inspecting Python code coverage +------------------------------- + +The code coverage is reported at `codecov`_. You can also inspect the coverage locally. +To get the full coverage first combine all reports and open produced html file in a +browser + +.. code-block:: bash + + tox + coverage combine --append \ + ./.coverage \ + ./python/rascaline-torch/.coverage + coverage html + firefox htmlcov/index.html + +.. _codecov: https://codecov.io/gh/lab-cosmo/metatensor + Writing your own calculator --------------------------- diff --git a/docs/src/conf.py b/docs/src/conf.py index 3b45ea3f6..a8d8a21da 100644 --- a/docs/src/conf.py +++ b/docs/src/conf.py @@ -140,6 +140,7 @@ def setup(app): "metatensor": ("https://lab-cosmo.github.io/metatensor/latest/", None), "matplotlib": ("https://matplotlib.org/stable/", None), "numpy": ("https://numpy.org/doc/stable/", None), + "scipy": ("https://docs.scipy.org/doc/scipy/", None), "skmatter": ("https://scikit-matter.readthedocs.io/en/latest/", None), "torch": ("https://pytorch.org/docs/stable/", None), "python": ("https://docs.python.org/3", None), diff --git a/docs/src/devdoc/explanations/index.rst b/docs/src/devdoc/explanations/index.rst index 8d054cc20..9ccd23922 100644 --- a/docs/src/devdoc/explanations/index.rst +++ b/docs/src/devdoc/explanations/index.rst @@ -12,3 +12,4 @@ to give you more clarity and understanding of what rascaline is all about. architecture interfaces + radial-integral diff --git a/docs/src/devdoc/explanations/radial-integral.rst b/docs/src/devdoc/explanations/radial-integral.rst new file mode 100644 index 000000000..c17f5a901 --- /dev/null +++ b/docs/src/devdoc/explanations/radial-integral.rst @@ -0,0 +1,275 @@ +.. _radial-integral: + +SOAP and LODE radial integrals +=================================== + +On this page, we describe the exact mathematical expression that are implemented in the +radial integral and the splined radial integral classes i.e. +:ref:`python-splined-radial-integral`. Note that this page assumes knowledge of +spherical expansion & friends and currently serves as a reference page for +the developers to support the implementation. + +Preliminaries +------------- + +In this subsection, we briefly provide all the preliminary knowledge that is needed to +understand what the radial integral class is doing. The actual explanation for what is +computed in the radial integral class can be found in the next subsection (1.2). The +spherical expansion coefficients :math:`\langle anlm | \rho_i \rangle` are completely +determined by specifying two ingredients: + +- the atomic density function :math:`g(r)` as implemented in + :ref:`python-atomic-density`, often chosen to be a Gaussian or Delta function, that + defined the type of density under consideration. For a given center atom :math:`i` in + the structure, the total density function :math:`\rho_i(\boldsymbol{r})` around is + then defined as :math:`\rho_i(\boldsymbol{r}) = \sum_{j} g(\boldsymbol{r} - + \boldsymbol{r}_{ij})`. + +- the radial basis functions :math:`R_{nl}(r)` as implementated + :ref:`python-radial-basis`, on which the density :math:`\rho_i` is projected. To be + more precise, the actual basis functions are of the form + + .. math:: + + B_{nlm}(\boldsymbol{r}) = R_{nl}(r)Y_{lm}(\hat{r}), + + where :math:`Y_{lm}(\hat{r})` are the real spherical harmonics evaluated at the point + :math:`\hat{r}`, i.e. at the spherical angles :math:`(\theta, \phi)` that determine + the orientation of the unit vector :math:`\hat{r} = \boldsymbol{r}/r`. + +The spherical expansion coefficient :math:`\langle nlm | \rho_i \rangle` (we ommit the +chemical species index :math:`a` for simplicity) is then defined as + +.. math:: + + \begin{aligned} + \langle nlm | \rho_i \rangle & = \int \mathrm{d}^3\boldsymbol{r} + B_{nlm}(\boldsymbol{r}) \rho_i(\boldsymbol{r}) \\ \label{expansion_coeff_def} & = + \int \mathrm{d}^3\boldsymbol{r} R_{nl}(r)Y_{lm}(\hat{r})\rho_i(\boldsymbol{r}). + \end{aligned} + +In practice, the atom centered density :math:`\rho_i` is a superposition of the neighbor +contributions, namely :math:`\rho_i(\boldsymbol{r}) = \sum_{j} g(\boldsymbol{r} - +\boldsymbol{r}_{ij})`. Due to linearity of integration, evaluating the integral can then +be simplified to + +.. math:: + + \begin{aligned} + \langle nlm | \rho_i \rangle & = \int \mathrm{d}^3\boldsymbol{r} + R_{nl}(r)Y_{lm}(\hat{r})\rho_i(\boldsymbol{r}) \\ & = \int + \mathrm{d}^3\boldsymbol{r} R_{nl}(r)Y_{lm}(\hat{r})\left( \sum_{j} + g(\boldsymbol{r} - \boldsymbol{r}_{ij})\right) \\ & = \sum_{j} \int + \mathrm{d}^3\boldsymbol{r} R_{nl}(r)Y_{lm}(\hat{r}) g(\boldsymbol{r} - + \boldsymbol{r}_{ij}) \\ & = \sum_j \langle nlm | g;\boldsymbol{r}_{ij} \rangle. + \end{aligned} + +Thus, instead of having to compute integrals for arbitrary densities :math:`\rho_i`, we +have reduced our problem to the evaluation of integrals of the form + +.. math:: + + \begin{aligned} + \langle nlm | g;\boldsymbol{r}_{ij} \rangle & = \int \mathrm{d}^3\boldsymbol{r} + R_{nl}(r)Y_{lm}(\hat{r})g(\boldsymbol{r} - \boldsymbol{r}_{ij}), + \end{aligned} + +which are completely specified by + +- the density function :math:`g(\boldsymbol{r})` + +- the radial basis :math:`R_{nl}(r)` + +- the position of the neighbor atom :math:`\boldsymbol{r}_{ij}` relative to the center + atom + +The radial integral class +------------------------- + +In the previous subsection, we have explained how the computation of the spherical +expansion coefficients can be reduced to integrals of the form + +.. math:: + + \begin{aligned} + \langle nlm | g;\boldsymbol{r}_{ij} \rangle & = \int \mathrm{d}^3\boldsymbol{r} + R_{nl}(r)Y_{lm}(\hat{r})g(\boldsymbol{r} - \boldsymbol{r}_{ij}). + \end{aligned} + +If the atomic density is spherically symmetric, i.e. if :math:`g(\boldsymbol{r}) = g(r)` +this integral can always be written in the following form: + +.. math:: + + \begin{aligned} \label{expansion_coeff_spherical_symmetric} + \langle nlm | g;\boldsymbol{r}_{ij} \rangle & = + Y_{lm}(\hat{r}_{ij})I_{nl}(r_{ij}). + \end{aligned} + +The key point is that the dependence on the vectorial position +:math:`\boldsymbol{r}_{ij}` is split into a factor that contains information about the +orientation of this vector, namely :math:`Y_{lm}(\hat{r}_{ij})`, which is just the +spherical harmonic evaluated at :math:`\hat{r}_{ij}`, and a remaining part that captures +the dependence on the distance of atom :math:`j` from the center atom :math:`i`, namely +:math:`I_{nl}(r_{ij})`, which we shall call the radial integral. The radial integral +class computes and outputs this radial part :math:`I_{nl}(r_{ij})`. Since the angular +part is just the usual spherical harmonic, this is the part that also depends on the +choice of atomic density :math:`g(r)`, as well as the radial basis :math:`R_{nl}(r)`. In +the following, for users only interested in a specific type of density, we provide the +explicit expressions of :math:`I_{nl}(r)` for the Delta and Gaussian densities, followed +by the general expression. + +Delta Densities +~~~~~~~~~~~~~~~ + +Here, we consider the especially simple special case where the atomic density function +:math:`g(\boldsymbol{r}) = \delta(\boldsymbol{r})`. Then: + +.. math:: + + \begin{aligned} + \langle nlm | g;\boldsymbol{r}_{ij} \rangle & = \int \mathrm{d}^3\boldsymbol{r} + R_{nl}(r)Y_{lm}(\hat{r})g(\boldsymbol{r} - \boldsymbol{r}_{ij}) \\ & = \int + \mathrm{d}^3\boldsymbol{r} R_{nl}(r)Y_{lm}(\hat{r})\delta(\boldsymbol{r} - + \boldsymbol{r}_{ij}) \\ & = R_{nl}(r) Y_{lm}(\hat{r}_{ij}) = + B_{nlm}(\boldsymbol{r}_{ij}). + \end{aligned} + +Thus, in this particularly simple case, the radial integral is simply the radial basis +function evaluated at the pair distance :math:`r_{ij}`, and we see that the integrals +have indeed the form presented above. + +Gaussian Densities +~~~~~~~~~~~~~~~~~~ + +Here, we consider another popular use case, where the atomic density function is a +Gaussian. In rascaline, we use the convention + +.. math:: + + g(r) = \frac{1}{(\pi \sigma^2)^{3/4}}e^{-\frac{r^2}{2\sigma^2}}. + +The prefactor was chosen such that the “L2-norm” of the Gaussian + +.. math:: + + \begin{aligned} + \|g\|^2 = \int \mathrm{d}^3\boldsymbol{r} |g(r)|^2 = 1, + \end{aligned} + +but does not affect the following calculations in any way. With these conventions, it +can be shown that the integral has the desired form + +.. math:: + + \begin{aligned} + \langle nlm | g;\boldsymbol{r}_{ij} \rangle & = \int \mathrm{d}^3\boldsymbol{r} + R_{nl}(r)Y_{lm}(\hat{r})g(\boldsymbol{r} - \boldsymbol{r}_{ij}) \\ & = + Y_{lm}(\hat{r}_{ij}) \cdot I_{nl}(r_{ij}) + \end{aligned} + +with + +.. math:: + + I_{nl}(r_{ij}) = \frac{1}{(\pi \sigma^2)^{3/4}}4\pi e^{-\frac{r_{ij}^2}{2\sigma^2}} + \int_0^\infty \mathrm{d}r r^2 R_{nl}(r) e^{-\frac{r^2}{2\sigma^2}} + i_l\left(\frac{rr_{ij}}{\sigma^2}\right), + +where :math:`i_l` is a modified spherical Bessel function. The first factor, of course, +is just the normalization factor of the Gaussian density. See the next two subsections +for a derivation of this formula. + +Derivation of the General Case +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +We now derive an explicit formula for radial integral that works for any density. Let +:math:`g(r)` be a generic spherically symmetric density function. Our goal will be to +show that + +.. math:: + + \langle nlm | g;\boldsymbol{r}_{ij} \rangle = Y_{lm}(\hat{r}_{ij}) \left[2\pi + \int_0^\infty \mathrm{d}r r^2 R_{nl}(r) \int_{-1}^1 \mathrm{d}(\cos\theta) + P_l(\cos\theta) g(\sqrt{r^2+r_{ij}^2-2rr_{ij}\cos\theta}) \right] + +and thus we have the desired form :math:`\langle nlm | g;\boldsymbol{r}_{ij} \rangle = +Y_{lm}(\hat{r}_{ij}) I_{nl}(r_{ij})` with + +.. math:: + + \begin{aligned} + I_{nl}(r_{ij}) = 2\pi \int_0^\infty \mathrm{d}r r^2 R_{nl}(r) \int_{-1}^1 + \mathrm{d}u P_l(u) g(\sqrt{r^2+r_{ij}^2-2rr_{ij}u}), + \end{aligned} + +where :math:`P_l(x)` is the :math:`l`-th Legendre polynomial. + +Derivation of the explicit radial integral for Gaussian densities +----------------------------------------------------------------- + +Denoting by :math:`\theta(\boldsymbol{r},\boldsymbol{r}_{ij})` the angle between a +generic position vector :math:`\boldsymbol{r}` and the vector +:math:`\boldsymbol{r}_{ij}`, we can write + +.. math:: + + \begin{aligned} + g(\boldsymbol{r}- \boldsymbol{r}_{ij}) & = \frac{1}{(\pi + \sigma^2)^{3/4}}e^{-\frac{(\boldsymbol{r}- \boldsymbol{r}_{ij})^2}{2\sigma^2}} \\ + & = \frac{1}{(\pi + \sigma^2)^{3/4}}e^{-\frac{(r_{ij})^2}{2\sigma^2}}e^{-\frac{(\boldsymbol{r}^2- + 2\boldsymbol{r}\boldsymbol{r}_{ij})}{2\sigma^2}}, + \end{aligned} + +where the first factor no longer depends on the integration variable :math:`r`. + +Analytical Expressions for the GTO Basis +---------------------------------------- + +While the above integrals are hard to compute in general, the GTO basis is one of the +few sets of basis functions for which many of the integrals can be evaluated +analytically. This is also useful to test the correctness of more numerical +implementations. + +The primitive basis functions are defined as + +.. math:: + + \begin{aligned} + R_{nl}(r) = R_n(r) = r^n e^{-\frac{r^2}{2\sigma_n^2}} + \end{aligned} + +In this form, the basis functions are not yet orthonormal, which requires an extra +linear transformation. Since this transformation can also be applied after computing the +integrals, we simply evaluate the radial integral with respect to these primitive basis +functions. + +Real Space Integral for Gaussian Densities +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +We now evaluate + +.. math:: + + \begin{aligned} + I_{nl}(r_{ij}) & = \frac{1}{(\pi \sigma^2)^{3/4}}4\pi + e^{-\frac{r_{ij}^2}{2\sigma^2}} \int_0^\infty \mathrm{d}r r^2 R_{nl}(r) + e^{-\frac{r^2}{2\sigma^2}} i_l\left(\frac{rr_{ij}}{\sigma^2}\right) \\ & = + \frac{1}{(\pi \sigma^2)^{3/4}}4\pi e^{-\frac{r_{ij}^2}{2\sigma^2}} \int_0^\infty + \mathrm{d}r r^2 r^n e^{-\frac{r^2}{2\sigma_n^2}} e^{-\frac{r^2}{2\sigma^2}} + i_l\left(\frac{rr_{ij}}{\sigma^2}\right), + \end{aligned} + +the result of which can be conveniently expressed using :math:`a=\frac{1}{2\sigma^2}`, +:math:`b_n = \frac{1}{2\sigma_n^2}`, :math:`n_\mathrm{eff}=\frac{n+l+3}{2}` and +:math:`l_\mathrm{eff}=l+\frac{3}{2}` as + +.. math:: + + \begin{aligned} + I_{nl}(r_{ij}) = \frac{1}{(\pi \sigma^2)^{3/4}} \cdot + \pi^{\frac{3}{2}}\frac{\Gamma\left(n_\mathrm{eff}\right)}{\Gamma\left(l_\mathrm{eff}\right)}\frac{(ar_{ij})^l}{(a+b)^{n_\mathrm{eff}}}M\left(n_\mathrm{eff},l_\mathrm{eff},\frac{a^2r_{ij}^2}{a^2+b^2}\right), + \end{aligned} + +where :math:`M(a,b,z)` is the confluent hypergeometric function (hyp1f1). diff --git a/docs/src/explanations/index.rst b/docs/src/explanations/index.rst index 8643a7235..77eb9ffa0 100644 --- a/docs/src/explanations/index.rst +++ b/docs/src/explanations/index.rst @@ -13,3 +13,4 @@ all about. concepts soap + rotation_adapted \ No newline at end of file diff --git a/docs/src/explanations/rotation_adapted.rst b/docs/src/explanations/rotation_adapted.rst new file mode 100644 index 000000000..e98bb332d --- /dev/null +++ b/docs/src/explanations/rotation_adapted.rst @@ -0,0 +1,88 @@ +Rotation-Adapted Features +========================= + +Equivariance +------------ + +Descriptors like SOAP are translation, rotation, and permutation invariant. +Indeed, such invariances are extremely useful if one wants to learn an invariant target (e.g., the energy). +Being already encoded in the descriptor, the learning algorithm does not have to learn such a physical requirement. + +The situation is different if the target is not invariant. For example, one may want to learn a dipole. The dipole rotates with a rotation of the molecule, and as such, invariant descriptors do not have the required symmetries for this task. + +Instead, one would need a rotation equivariant descriptor. +Rotation equivariance means that, if I first rotate the structure and compute the descriptor, I obtain the same result as first computing the descriptor and then applying the rotation, i.e., the descriptor behaves correctly upon rotation operations. +Denoting a structure as :math:`A`, the function computing the descriptor as :math:`f(\cdot)`, and the rotation operator as :math:`\hat{R}`, rotation equivariance can be expressed as: + +.. math:: + :name: eq:equivariance + + f(\hat{R} A) = \hat{R} f(A) + +Of course, invariance is a special case of equivariance. + + +Rotation Equivariance of the Spherical Expansion +------------------------------------------------ + +The spherical expansion is a rotation equivariant descriptor. +Let's consider the expansion coefficients of :math:`\rho_i(\mathbf{r})`. +We have: + +.. math:: + + \hat{R} \rho_i(\mathbf{r}) &= \sum_{nlm} c_{nlm}^{i} R_n(r) \hat{R} Y_l^m(\hat{\mathbf{r}}) \nonumber \\ + &= \sum_{nlmm'} c_{nlm}^{i} R_n(r) D_{m,m'}^{l}(\hat{R}) Y_l^{m'}(\hat{\mathbf{r}}) \nonumber \\ + &= \sum_{nlm} \left( \sum_{m'} D_{m',m}^l(\hat{R}) c_{nlm'}^{i}\right) B_{nlm}(\mathbf{r}) \nonumber + +and noting that :math:`Y_l^m(\hat{R} \hat{\mathbf{r}}) = \hat{R} Y_l^m(\hat{\mathbf{r}})` and :math:`\hat{R}r = r`, equation :ref:`(1) ` is satisfied and we conclude that the expansion coefficients :math:`c_{nlm}^{i}` are rotation equivariant. +Indeed, each :math:`c_{nlm}^{i}` transforms under rotation as the spherical harmonics :math:`Y_l^m(\hat{\mathbf{r}})`. + +Using the Dirac notation, the coefficient :math:`c_{nlm}^{i}` can be expressed as :math:`\braket{nlm\vert\rho_i}`. +Equivalently, and to stress the fact that this coefficient describes something that transforms under rotation as a spherical harmonics :math:`Y_l^m(\hat{\mathbf{r}})`, it is sometimes written as :math:`\braket{n\vert\rho_i;lm}`, i.e., the atomic density is "tagged" with a label that tells how it transforms under rotations. + + +Completeness Relations of Spherical Harmonics +--------------------------------------------- + +Spherical harmonics can be combined together using rules coming from standard theory of angular momentum: + +.. math:: + :name: eq:cg_coupling + + \ket{lm} \propto \ket{l_1 l_2 l m} = \sum_{m_1 m_2} C_{m_1 m_2 m}^{l_1 l_2 l} \ket{l_1 m_1} \ket{l_2 m_2} + +where :math:`C_{m_1 m_2 m}^{l_1 l_2 l}` is a Clebsch-Gordan (CG) coefficient. + +Thanks to the one-to-one correspondence (under rotation) between :math:`c_{nlm}^{i}` and :math:`Y_l^m`, +:ref:`(2) ` means that one can take products of two spherical expansion coefficients (which amounts to considering density correlations), and combine them with CG coefficients to get new coefficients that transform as a single spherical harmonics. +This process is known as coupling, from the uncoupled basis of angular momentum (formed by the product of rotation eigenstates) to a coupled basis (a single rotation eigenstate). + +One can also write the inverse of :ref:`(2) `: + +.. math:: + :name: eq:cg_decoupling + + \ket{l_1 m_1} \ket{l_2 m_2} = \sum_{l m} C_{m_1 m_2 m}^{l_1 l_2 l m} \ket{l_1 l_2 l m} + +that express the product of two rotation eigenstates in terms of one. This process is known as decoupling. + +Example: :math:`\lambda`-SOAP +----------------------------- + +A straightforward application of :ref:`(2) ` is the construction of :math:`\lambda`-SOAP features. +Indeed, :math:`\lambda`-SOAP was created in order to have a rotation and inversion equivariant version of the 3-body density correlations. +The :math:`\lambda` represents the degree of a spherical harmonics, :math:`Y_{\lambda}^{\mu}(\hat{\mathbf{r}})`, +and it indicates that this descriptor can transform under rotations as a spherical harmonics, i.e., it is rotation equivariant. + +It is then obtained by considering two expansion coefficients of the atomic density, and combining them with a CG iteration to a coupled basis, +as in :ref:`(2) `. +The :math:`\lambda`-SOAP descriptor is then: + +.. math:: + + \braket{n_1 l_1 n_2 l_2\vert\overline{\rho_i^{\otimes 2}, \sigma, \lambda \mu}} = + \frac{\delta_{\sigma, (-1)^{l_1 + l_2 + \lambda}}}{\sqrt{2 \lambda + 1}} + \sum_{m} C_{m (\mu-m) \mu}^{l_1 l_2 \lambda} c_{n_1 l_1 m}^{i} c_{n_2 l_2 (\mu - m)}^{i} + +where we have assumed real spherical harmonics coefficients. \ No newline at end of file diff --git a/docs/src/references/api/python/utils/atomic-density.rst b/docs/src/references/api/python/utils/atomic-density.rst new file mode 100644 index 000000000..4c0fc5591 --- /dev/null +++ b/docs/src/references/api/python/utils/atomic-density.rst @@ -0,0 +1 @@ +.. automodule:: rascaline.utils.splines.atomic_density diff --git a/docs/src/references/api/python/utils/index.rst b/docs/src/references/api/python/utils/index.rst index abe7e11a0..3d88c4da9 100644 --- a/docs/src/references/api/python/utils/index.rst +++ b/docs/src/references/api/python/utils/index.rst @@ -7,5 +7,7 @@ Utility functions and classes that extend the core usage of rascaline. .. toctree:: :maxdepth: 1 + atomic-density + radial-basis power-spectrum splines diff --git a/docs/src/references/api/python/utils/radial-basis.rst b/docs/src/references/api/python/utils/radial-basis.rst new file mode 100644 index 000000000..fa8bacf62 --- /dev/null +++ b/docs/src/references/api/python/utils/radial-basis.rst @@ -0,0 +1 @@ +.. automodule:: rascaline.utils.splines.radial_basis diff --git a/docs/src/references/api/python/utils/splines.rst b/docs/src/references/api/python/utils/splines.rst index 7e93ba9db..4a3b0a912 100644 --- a/docs/src/references/api/python/utils/splines.rst +++ b/docs/src/references/api/python/utils/splines.rst @@ -1 +1 @@ -.. automodule:: rascaline.utils.splines +.. automodule:: rascaline.utils.splines.splines diff --git a/pyproject.toml b/pyproject.toml index 924103358..2cce98f8f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "rascaline" -dynamic = ["version", "authors"] -requires-python = ">=3.7" +dynamic = ["version", "authors", "optional-dependencies"] +requires-python = ">=3.8" readme = "README.rst" license = {text = "BSD-3-Clause"} @@ -27,6 +27,7 @@ classifiers = [ dependencies = [ "metatensor-core >=0.1.0,<0.2.0", + "wigners", ] [project.urls] diff --git a/python/rascaline-torch/build-backend/backend.py b/python/rascaline-torch/build-backend/backend.py index 057d0157c..84f667978 100644 --- a/python/rascaline-torch/build-backend/backend.py +++ b/python/rascaline-torch/build-backend/backend.py @@ -18,7 +18,7 @@ RASCALINE_DEP = f"rascaline @ file://{RASCALINE}?{uuid}" else: # we are building from a sdist - RASCALINE_DEP = "rascaline == 0.1.0" + RASCALINE_DEP = "rascaline >=0.1.0.dev0,<0.2.0" prepare_metadata_for_build_wheel = build_meta.prepare_metadata_for_build_wheel diff --git a/python/rascaline-torch/pyproject.toml b/python/rascaline-torch/pyproject.toml index 4c4daacf4..098a8ba90 100644 --- a/python/rascaline-torch/pyproject.toml +++ b/python/rascaline-torch/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "rascaline-torch" dynamic = ["version", "authors", "dependencies"] -requires-python = ">=3.7" +requires-python = ">=3.8" readme = "README.rst" license = {text = "BSD-3-Clause"} diff --git a/python/rascaline-torch/rascaline/torch/__init__.py b/python/rascaline-torch/rascaline/torch/__init__.py index 177d0c688..971e0a3a1 100644 --- a/python/rascaline-torch/rascaline/torch/__init__.py +++ b/python/rascaline-torch/rascaline/torch/__init__.py @@ -1,15 +1,7 @@ -import sys +import importlib.metadata -if (sys.version_info.major >= 3) and (sys.version_info.minor >= 8): - import importlib.metadata - - __version__ = importlib.metadata.version("rascaline-torch") - -else: - from pkg_resources import get_distribution - - __version__ = get_distribution("rascaline-torch").version +__version__ = importlib.metadata.version("rascaline-torch") from ._c_lib import _load_library diff --git a/python/rascaline-torch/rascaline/torch/utils.py b/python/rascaline-torch/rascaline/torch/utils.py deleted file mode 100644 index 96bce84d6..000000000 --- a/python/rascaline-torch/rascaline/torch/utils.py +++ /dev/null @@ -1,21 +0,0 @@ -import os - -import metatensor.torch - -import rascaline - - -_HERE = os.path.realpath(os.path.dirname(__file__)) - -_rascaline_torch_cmake_prefix = os.path.join(os.path.dirname(__file__), "lib", "cmake") - -cmake_prefix_path = ";".join( - [ - _rascaline_torch_cmake_prefix, - rascaline.utils.cmake_prefix_path, - metatensor.torch.utils.cmake_prefix_path, - ] -) -""" -Path containing the CMake configuration files for the underlying C library -""" diff --git a/python/rascaline-torch/rascaline/torch/utils/__init__.py b/python/rascaline-torch/rascaline/torch/utils/__init__.py index bc31628b9..3a73c2d20 100644 --- a/python/rascaline-torch/rascaline/torch/utils/__init__.py +++ b/python/rascaline-torch/rascaline/torch/utils/__init__.py @@ -1,4 +1,13 @@ +import os + from .power_spectrum import PowerSpectrum +_HERE = os.path.dirname(__file__) + +cmake_prefix_path = os.path.realpath(os.path.join(_HERE, "..", "lib", "cmake")) +""" +Path containing the CMake configuration files for the underlying C library +""" + __all__ = ["PowerSpectrum"] diff --git a/python/rascaline-torch/setup.py b/python/rascaline-torch/setup.py index dea776aee..0d357482e 100644 --- a/python/rascaline-torch/setup.py +++ b/python/rascaline-torch/setup.py @@ -38,6 +38,7 @@ class cmake_ext(build_ext): """Build the native library using cmake""" def run(self): + import metatensor import metatensor.torch import torch @@ -52,6 +53,7 @@ def run(self): # Tell CMake where to find rascaline & torch cmake_prefix_path = [ rascaline.utils.cmake_prefix_path, + metatensor.utils.cmake_prefix_path, metatensor.torch.utils.cmake_prefix_path, torch.utils.cmake_prefix_path, ] @@ -238,10 +240,13 @@ def git_extra_version(): with open(os.path.join(ROOT, authors[0])) as fd: authors = fd.read().splitlines() - install_requires = ["torch >= 1.11"] + install_requires = [ + "torch >= 1.11", + "metatensor-torch", + ] if os.path.exists(RASCALINE_C_API): # we are building from a git checkout - rascaline_path = os.path.join(ROOT, "..", "..") + rascaline_path = os.path.realpath(os.path.join(ROOT, "..", "..")) # add a random uuid to the file url to prevent pip from using a cached # wheel for rascaline, and force it to re-build from scratch diff --git a/python/rascaline-torch/tests/calculator.py b/python/rascaline-torch/tests/calculator.py index 4f6d4017a..1c26be1bb 100644 --- a/python/rascaline-torch/tests/calculator.py +++ b/python/rascaline-torch/tests/calculator.py @@ -62,7 +62,6 @@ def test_compute(system): assert torch.all(gradient.values[i, 2, :] == torch.tensor([0, 1])) assert len(gradient.samples) == 8 - print(gradient.samples.values) assert gradient.samples.names == ["sample", "structure", "atom"] assert tuple(gradient.samples[0]) == (0, 0, 0) assert tuple(gradient.samples[1]) == (0, 0, 1) @@ -191,3 +190,4 @@ def forward( with tmpdir.as_cwd(): torch.jit.save(module, "test-save.torch") + module = torch.jit.load("test-save.torch") diff --git a/python/rascaline-torch/tests/utils/cmake_prefix.py b/python/rascaline-torch/tests/utils/cmake_prefix.py new file mode 100644 index 000000000..d6172a3bb --- /dev/null +++ b/python/rascaline-torch/tests/utils/cmake_prefix.py @@ -0,0 +1,7 @@ +import os + +import rascaline.torch + + +def test_cmake_prefix(): + assert os.path.exists(rascaline.torch.utils.cmake_prefix_path) diff --git a/python/rascaline/rascaline/systems/pyscf.py b/python/rascaline/rascaline/systems/pyscf.py index 2f0ac6d85..868a2b314 100644 --- a/python/rascaline/rascaline/systems/pyscf.py +++ b/python/rascaline/rascaline/systems/pyscf.py @@ -21,7 +21,7 @@ def _std_symbol_without_ghost(symb_or_chg): rawsymb = pyscf.data.elements._rm_digit(symb_or_chg) if rawsymb in pyscf.data.elements._ELEMENTS_UPPER: return pyscf.data.elements._ELEMENTS_UPPER[rawsymb] - elif len(rawsymb) > 1 and symb_or_chg[0] == "X" and symb_or_chg[:2] != "XE": + elif len(rawsymb) > 1 and (symb_or_chg[0] == "X" and symb_or_chg[:2] != "XE"): rawsymb = rawsymb[1:] # Remove the prefix X return pyscf.data.elements._ELEMENTS_UPPER[rawsymb] elif len(rawsymb) > 5 and rawsymb[:5] == "GHOST": @@ -34,24 +34,35 @@ def _std_symbol_without_ghost(symb_or_chg): class PyscfSystem(SystemBase): - """Implements :py:class:`rascaline.SystemBase` wrapping a `pyscf.gto.mole.Mole`_ - or `pyscf.pbc.gto.cell.Cell`_. + """Implements :py:class:`rascaline.SystemBase` wrapping a + `pyscf.gto.mole.Mole`_ or `pyscf.pbc.gto.cell.Cell`_. Since pyscf does not offer a neighbors list, this - implementation of system can only be used with ``use_native_system=True`` in + implementation of system can only be used with + ``use_native_system=True`` in :py:func:`rascaline.calculators.CalculatorBase.compute`. - Atomic species are assigned as the atomic number if the atom ``type`` is one - of the periodic table elements; or their opposite if they are ghost atoms. + Atomic species are assigned as the atomic number if the atom ``type`` is + one of the periodic table elements; or their opposite if they are + ghost atoms. (Pyscf does not seem to support anything else) + Please note that while pyscf uses Bohrs as length units internally, + we convert those back into Angströms for rascaline. + A pyscf object's "unit" attribute determines the units of the coordinates + given *to pyscf*, which are by default angströms. + .. _pyscf.gto.mole.Mole: https://pyscf.org/user/gto.html .. _pyscf.pbc.gto.cell.Cell: https://pyscf.org/user/pbc/gto.html """ @staticmethod def can_wrap(o): - return isinstance(o, (pyscf.gto.mole.Mole, pyscf.pbc.gto.cell.Cell)) + # assumption: if we have a periodic system, then pyscf.pbc is defined + if hasattr(pyscf, "pbc"): + return isinstance(o, (pyscf.gto.mole.Mole, pyscf.pbc.gto.cell.Cell)) + else: + return isinstance(o, pyscf.gto.mole.Mole) def __init__(self, frame): """ @@ -59,7 +70,7 @@ def __init__(self, frame): in this ``ChemfilesSystem`` """ super().__init__() - if not isinstance(frame, (pyscf.gto.mole.Mole, pyscf.pbc.gto.cell.Cell)): + if not self.can_wrap(frame): raise Exception( "this class expects pyscf.gto.mole.Mole" + "or pyscf.pbc.gto.cell.Cell objects" @@ -81,11 +92,17 @@ def species(self): return self._species def positions(self): - return self._frame.atom_coords() + return pyscf.data.nist.BOHR * self._frame.atom_coords() def cell(self): if self.is_periodic: - return self._frame.a + cell = self._frame.a + if self._frame.unit[0].lower() == "a": + # assume angströms, we are good to go + return cell + else: + # assume bohrs, correct this + return pyscf.data.nist.BOHR * np.asarray(cell) else: return np.zeros((3, 3), float) diff --git a/python/rascaline/rascaline/utils/__init__.py b/python/rascaline/rascaline/utils/__init__.py index 88d7e9e90..8309ac44c 100644 --- a/python/rascaline/rascaline/utils/__init__.py +++ b/python/rascaline/rascaline/utils/__init__.py @@ -1,16 +1,26 @@ import os -import metatensor - +from .clebsch_gordan import * # noqa from .power_spectrum import PowerSpectrum # noqa -from .splines import RadialIntegralFromFunction, RadialIntegralSplinerBase # noqa +from .splines import ( # noqa + AtomicDensityBase, + DeltaDensity, + GaussianDensity, + GtoBasis, + LodeDensity, + LodeSpliner, + MonomialBasis, + RadialBasisBase, + RadialIntegralFromFunction, + RadialIntegralSplinerBase, + SoapSpliner, + SphericalBesselBasis, +) -# path that can be used with cmake to access the rascaline library and headers -_HERE = os.path.realpath(os.path.dirname(__file__)) +_HERE = os.path.dirname(__file__) -_rascaline_cmake_prefix = os.path.realpath(os.path.join(_HERE, "..", "lib", "cmake")) -cmake_prefix_path = f"{_rascaline_cmake_prefix};{metatensor.utils.cmake_prefix_path}" +cmake_prefix_path = os.path.realpath(os.path.join(_HERE, "..", "lib", "cmake")) """ Path containing the CMake configuration files for the underlying C library """ diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/__init__.py b/python/rascaline/rascaline/utils/clebsch_gordan/__init__.py new file mode 100644 index 000000000..38777d18e --- /dev/null +++ b/python/rascaline/rascaline/utils/clebsch_gordan/__init__.py @@ -0,0 +1,7 @@ +from .clebsch_gordan import correlate_density, correlate_density_metadata # noqa + + +__all__ = [ + "correlate_density", + "correlate_density_metadata", +] diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/_cg_cache.py b/python/rascaline/rascaline/utils/clebsch_gordan/_cg_cache.py new file mode 100644 index 000000000..8ad082fb8 --- /dev/null +++ b/python/rascaline/rascaline/utils/clebsch_gordan/_cg_cache.py @@ -0,0 +1,551 @@ +""" +Module that stores the ClebschGordanReal class for computing and caching Clebsch +Gordan coefficients for use in CG combinations. +""" +from typing import Union + +import numpy as np +import wigners + +from . import _dispatch + + +try: + from mops import sparse_accumulation_of_products as sap # noqa F401 + + HAS_MOPS = True +except ImportError: + HAS_MOPS = False + +try: + from torch import Tensor as TorchTensor +except ImportError: + + class TorchTensor: + pass + + +UNKNOWN_ARRAY_TYPE = ( + "unknown array type, only numpy arrays and torch tensors are supported" +) + + +# ================================= +# ===== ClebschGordanReal class +# ================================= + + +class ClebschGordanReal: + """ + Class for computing Clebsch-Gordan coefficients for real spherical + harmonics. + + Stores the coefficients in a dictionary in the `self.coeffs` attribute, + which is built at initialization. There are 3 current use cases for the + format of these coefficients. By default, sparse accumulation of products is + performed, whether or not Mops is installed. + + Case 1: standard sparse format. + + Each dictionary entry is a dictionary with entries for each (m1, m2, mu) + combination. + + { + (l1, l2, lambda): { + (m1, m2, mu) : cg_{m1, m2, mu}^{l1, l2, lambda} + for m1 in range(-l1, l1 + 1), + for m2 in range(-l2, l2 + 1), + }, + ... + for l1 in range(0, l1_list) + for l2 in range(0, l2_list) + for lambda in range(0, range(|l1 - l2|, ..., |l1 + l2|)) + } + + Case 2: standard dense format. + + Each dictionary entry is a dense array with shape (2 * l1 + 1, 2 * l2 + 1, 2 + * lambda + 1). + + { + (l1, l2, lambda): + array( + cg_{m1, m2, mu}^{l1, l2, lambda} + ... + for m1 in range(-l1, l1 + 1), + for m2 in range(-l2, l2 + 1), + for mu in range(-lambda, lambda + 1), + + shape=(2 * l1 + 1, 2 * l2 + 1, 2 * lambda + 1), + ) + ... + for l1 in range(0, l1_list) + for l2 in range(0, l2_list) + for lambda in range(0, range(|l1 - l2|, ..., |l1 + l2|)) + } + + Case 3: MOPS sparse format. + + Each dictionary entry contains a tuple with four 1D arrays, corresponding to + the CG coeffs and m1, m2, mu indices respectively. All of these arrays are + sorted according to the mu index. This format is used for Sparse + Accumulation of Products (SAP) as implemented in MOPS. See + https://github.com/lab-cosmo/mops . + + { + (l1, l2, lambda): + ( + [ + cg_{m1, m2, mu}^{l1, l2, lambda} + ... + for m1 in range(-l1, l1 + 1), + for m2 in range(-l2, l2 + 1), + for mu in range(-lambda, lambda + 1) + ], + [ + m1 for m1 in range(-l1, l1 + 1), + ], + [ + m2 for m2 in range(-l2, l2 + 1), + ], + [ + mu for mu in range(-lambda, lambda + 1), + ], + ) + + + } + + where `cg_{m1, m2, mu}^{l1, l2, lambda}` is the Clebsch-Gordan coefficient + that describes the combination of the `m1` irreducible component of the `l1` + angular channel and the `m2` irreducible component of the `l2` angular + channel into the irreducible tensor of order `lambda`. In all cases, these + correspond to the non-zero CG coefficients, i.e. those in the range |-l, + ..., +l| for each angular order l in {l1, l2, lambda}. + + :param lambda_max: maximum lambda value to compute CG coefficients for. + :param sparse: whether to store the CG coefficients in sparse format. + :param use_mops: whether to store the CG coefficients in MOPS sparse format. + This is recommended as the default for sparse accumulation, but can only + be used if Mops is installed. + """ + + def __init__(self, lambda_max: int, sparse: bool = True, use_mops: bool = HAS_MOPS): + self._lambda_max = lambda_max + self._sparse = sparse + + if sparse: + if not HAS_MOPS: + # TODO: provide a warning once Mops is fully ready + # import warnings + # warnings.warn( + # "It is recommended to use MOPS for sparse accumulation. " + # " This can be installed with ``pip install" + # " git+https://github.com/lab-cosmo/mops`." + # " Falling back to numpy for now." + # ) + self._use_mops = False + else: + self._use_mops = True + + else: + # TODO: provide a warning once Mops is fully ready + # if HAS_MOPS: + # import warnings + # warnings.warn( + # "Mops is installed, but not being used" + # " as dense operations chosen." + # ) + self._use_mops = False + + self._coeffs = ClebschGordanReal.build_coeff_dict( + self._lambda_max, + self._sparse, + self._use_mops, + ) + + @property + def lambda_max(self): + return self._lambda_max + + @property + def sparse(self): + return self._sparse + + @property + def use_mops(self): + return self._use_mops + + @property + def coeffs(self): + return self._coeffs + + @staticmethod + def build_coeff_dict(lambda_max: int, sparse: bool, use_mops: bool): + """ + Builds a dictionary of Clebsch-Gordan coefficients for all possible + combination of l1 and l2, up to lambda_max. + """ + # real-to-complex and complex-to-real transformations as matrices + r2c = {} + c2r = {} + coeff_dict = {} + for lambda_ in range(0, lambda_max + 1): + c2r[lambda_] = _complex2real(lambda_) + r2c[lambda_] = _real2complex(lambda_) + + for l1 in range(lambda_max + 1): + for l2 in range(lambda_max + 1): + for lambda_ in range( + max(l1, l2) - min(l1, l2), min(lambda_max, (l1 + l2)) + 1 + ): + complex_cg = _complex_clebsch_gordan_matrix(l1, l2, lambda_) + + real_cg = (r2c[l1].T @ complex_cg.reshape(2 * l1 + 1, -1)).reshape( + complex_cg.shape + ) + + real_cg = real_cg.swapaxes(0, 1) + real_cg = (r2c[l2].T @ real_cg.reshape(2 * l2 + 1, -1)).reshape( + real_cg.shape + ) + real_cg = real_cg.swapaxes(0, 1) + + real_cg = real_cg @ c2r[lambda_].T + + if (l1 + l2 + lambda_) % 2 == 0: + cg_l1l2lam = np.real(real_cg) + else: + cg_l1l2lam = np.imag(real_cg) + + if sparse: + # Find the m1, m2, mu idxs of the nonzero CG coeffs + nonzeros_cg_coeffs_idx = np.where(np.abs(cg_l1l2lam) > 1e-15) + if use_mops: + # Store CG coeffs in a specific format for use in + # MOPS. Here we need the m1, m2, mu, and CG coeffs + # to be stored as separate 1D arrays. + m1_arr, m2_arr, mu_arr, C_arr = [], [], [], [] + for m1, m2, mu in zip(*nonzeros_cg_coeffs_idx): + m1_arr.append(m1) + m2_arr.append(m2) + mu_arr.append(mu) + C_arr.append(cg_l1l2lam[m1, m2, mu]) + + # Reorder the arrays based on sorted mu values + mu_idxs = np.argsort(mu_arr) + m1_arr = np.array(m1_arr)[mu_idxs] + m2_arr = np.array(m2_arr)[mu_idxs] + mu_arr = np.array(mu_arr)[mu_idxs] + C_arr = np.array(C_arr)[mu_idxs] + cg_l1l2lam = (C_arr, m1_arr, m2_arr, mu_arr) + else: + # Otherwise fall back to torch/numpy and store as + # sparse dicts. + cg_l1l2lam = { + (m1, m2, mu): cg_l1l2lam[m1, m2, mu] + for m1, m2, mu in zip(*nonzeros_cg_coeffs_idx) + } + + # Store + coeff_dict[(l1, l2, lambda_)] = cg_l1l2lam + + return coeff_dict + + +# ============================ +# ===== Helper functions +# ============================ + + +def _real2complex(lambda_: int) -> np.ndarray: + """ + Computes a matrix that can be used to convert from real to complex-valued + spherical harmonics(coefficients) of order ``lambda_``. + + This is meant to be applied to the left: ``real2complex @ [-lambda_, ..., + +lambda_]``. + + See https://en.wikipedia.org/wiki/Spherical_harmonics#Real_form for details + on the convention for how these tranformations are defined. + """ + result = np.zeros((2 * lambda_ + 1, 2 * lambda_ + 1), dtype=np.complex128) + inv_sqrt_2 = 1.0 / np.sqrt(2) + i_sqrt_2 = 1j / np.sqrt(2) + for m in range(-lambda_, lambda_ + 1): + if m < 0: + # Positve part + result[lambda_ + m, lambda_ + m] = +i_sqrt_2 + # Negative part + result[lambda_ - m, lambda_ + m] = -i_sqrt_2 * ((-1) ** m) + + if m == 0: + result[lambda_, lambda_] = +1.0 + + if m > 0: + # Negative part + result[lambda_ - m, lambda_ + m] = +inv_sqrt_2 + # Positive part + result[lambda_ + m, lambda_ + m] = +inv_sqrt_2 * ((-1) ** m) + + return result + + +def _complex2real(lambda_: int) -> np.ndarray: + """ + Converts from complex to real spherical harmonics. This is just given by the + conjugate tranpose of the real->complex transformation matrices. + """ + return np.conjugate(_real2complex(lambda_)).T + + +def _complex_clebsch_gordan_matrix(l1, l2, lambda_): + r"""clebsch-gordan matrix + Computes the Clebsch-Gordan (CG) matrix for + transforming complex-valued spherical harmonics. + The CG matrix is computed as a 3D array of elements + < l1 m1 l2 m2 | lambda_ mu > + where the first axis loops over m1, the second loops over m2, + and the third one loops over mu. The matrix is real. + For example, using the relation: + | l1 l2 lambda_ mu > = + \sum_{m1, m2} + | l1 m1 > | l2 m2 > + (https://en.wikipedia.org/wiki/Clebsch–Gordan_coefficients, section + "Formal definition of Clebsch-Gordan coefficients", eq 2) + one can obtain the spherical harmonics lambda_ from two sets of + spherical harmonics with l1 and l2 (up to a normalization factor). + E.g.: + Args: + l1: l number for the first set of spherical harmonics + l2: l number for the second set of spherical harmonics + lambda_: l number For the third set of spherical harmonics + Returns: + cg: CG matrix for transforming complex-valued spherical harmonics + >>> from scipy.special import sph_harm + >>> import numpy as np + >>> import wigners + >>> C_112 = _complex_clebsch_gordan_matrix(1, 1, 2) + >>> comp_sph_1 = np.array([sph_harm(m, 1, 0.2, 0.2) for m in range(-1, 1 + 1)]) + >>> comp_sph_2 = np.array([sph_harm(m, 1, 0.2, 0.2) for m in range(-1, 1 + 1)]) + >>> # obtain the (unnormalized) spherical harmonics + >>> # with l = 2 by contraction over m1 and m2 + >>> comp_sph_2_u = np.einsum("ijk,i,j->k", C_112, comp_sph_1, comp_sph_2) + >>> # we can check that they differ from the spherical harmonics + >>> # by a constant factor + >>> comp_sph_2 = np.array([sph_harm(m, 2, 0.2, 0.2) for m in range(-2, 2 + 1)]) + >>> ratio = comp_sph_2 / comp_sph_2_u + >>> np.allclose(ratio[0], ratio) + True + """ + if np.abs(l1 - l2) > lambda_ or np.abs(l1 + l2) < lambda_: + return np.zeros((2 * l1 + 1, 2 * l2 + 1, 2 * lambda_ + 1), dtype=np.double) + else: + return wigners.clebsch_gordan_array(l1, l2, lambda_) + + +# ================================================= +# ===== Functions for performing CG combinations +# ================================================= + + +def combine_arrays( + arr_1: Union[np.ndarray, TorchTensor], + arr_2: Union[np.ndarray, TorchTensor], + lambda_: int, + cg_cache, + return_empty_array: bool = False, +) -> Union[np.ndarray, TorchTensor]: + """ + Couples arrays `arr_1` and `arr_2` corresponding to the irreducible + spherical components of 2 angular channels l1 and l2 using the appropriate + Clebsch-Gordan coefficients. As l1 and l2 can be combined to form multiple + lambda channels, this function returns the coupling to a single specified + channel `lambda`. The angular channels l1 and l2 are inferred from the size + of the components axis (axis 1) of the input arrays. + + `arr_1` has shape (n_i, 2 * l1 + 1, n_p) and `arr_2` has shape (n_i, 2 * l2 + + 1, n_q). n_i is the number of samples, n_p and n_q are the number of + properties in each array. The number of samples in each array must be the + same. + + The ouput array has shape (n_i, 2 * lambda + 1, n_p * n_q), where lambda is + the input parameter `lambda_`. + + The Clebsch-Gordan coefficients are cached in `cg_cache`. Currently, these + must be produced by the ClebschGordanReal class in this module. These + coefficients can be stored in either sparse dictionaries or dense arrays. + + The combination operation is dispatched such that numpy arrays or torch + tensors are automatically handled. + + `return_empty_array` can be used to return an empty array of the correct + shape, without performing the CG combination step. This can be useful for + probing the outputs of CG iterations in terms of metadata without the + computational cost of performing the CG combinations - i.e. using the + function :py:func:`combine_single_center_to_body_order_metadata_only`. + + :param arr_1: array with the m values for l1 with shape [n_samples, 2 * l1 + + 1, n_q_properties] + :param arr_2: array with the m values for l2 with shape [n_samples, 2 * l2 + + 1, n_p_properties] + :param lambda_: int value of the resulting coupled channel + :param cg_cache: either a sparse dictionary with keys (m1, m2, mu) and array + values being sparse blocks of shape , or a dense array + of shape [(2 * l1 +1) * (2 * l2 +1), (2 * lambda_ + 1)]. + + :returns: array of shape [n_samples, (2*lambda_+1), q_properties * p_properties] + """ + # If just precomputing metadata, return an empty array + if return_empty_array: + return sparse_combine(arr_1, arr_2, lambda_, cg_cache, return_empty_array=True) + + # Otherwise, perform the CG combination + # Spare CG cache + if cg_cache.sparse: + return sparse_combine(arr_1, arr_2, lambda_, cg_cache, return_empty_array=False) + + # Dense CG cache + return dense_combine(arr_1, arr_2, lambda_, cg_cache) + + +def sparse_combine( + arr_1: Union[np.ndarray, TorchTensor], + arr_2: Union[np.ndarray, TorchTensor], + lambda_: int, + cg_cache, + return_empty_array: bool = False, +) -> Union[np.ndarray, TorchTensor]: + """ + Performs a Clebsch-Gordan combination step on 2 arrays using sparse + operations. The angular channel of each block is inferred from the size of + its component axis, and the blocks are combined to the desired output + angular channel `lambda_` using the appropriate Clebsch-Gordan coefficients. + + :param arr_1: array with the m values for l1 with shape [n_samples, 2 * l1 + + 1, n_q_properties] + :param arr_2: array with the m values for l2 with shape [n_samples, 2 * l2 + + 1, n_p_properties] + :param lambda_: int value of the resulting coupled channel + :param cg_cache: sparse dictionary with keys (m1, m2, mu) and array values + being sparse blocks of shape + + :returns: array of shape [n_samples, (2*lambda_+1), q_properties * p_properties] + """ + # Samples dimensions must be the same + assert arr_1.shape[0] == arr_2.shape[0] + + # Infer l1 and l2 from the len of the length of axis 1 of each tensor + l1 = (arr_1.shape[1] - 1) // 2 + l2 = (arr_2.shape[1] - 1) // 2 + + # Define other useful dimensions + n_i = arr_1.shape[0] # number of samples + n_p = arr_1.shape[2] # number of properties in arr_1 + n_q = arr_2.shape[2] # number of properties in arr_2 + + if return_empty_array: # used when only computing metadata + return _dispatch.zeros_like((n_i, 2 * lambda_ + 1, n_p * n_q), like=arr_1) + + if isinstance(arr_1, np.ndarray) and HAS_MOPS: + # Reshape + arr_1 = np.repeat(arr_1[:, :, :, None], n_q, axis=3).reshape( + n_i, 2 * l1 + 1, n_p * n_q + ) + arr_2 = np.repeat(arr_2[:, :, None, :], n_p, axis=2).reshape( + n_i, 2 * l2 + 1, n_p * n_q + ) + + arr_1 = _dispatch.swapaxes(arr_1, 1, 2).reshape(n_i * n_p * n_q, 2 * l1 + 1) + arr_2 = _dispatch.swapaxes(arr_2, 1, 2).reshape(n_i * n_p * n_q, 2 * l2 + 1) + + # Do SAP + arr_out = sap( + arr_1, + arr_2, + *cg_cache._coeffs[(l1, l2, lambda_)], + output_size=2 * lambda_ + 1, + ) + assert arr_out.shape == (n_i * n_p * n_q, 2 * lambda_ + 1) + + # Reshape back + arr_out = arr_out.reshape(n_i, n_p * n_q, 2 * lambda_ + 1) + arr_out = _dispatch.swapaxes(arr_out, 1, 2) + + return arr_out + + if isinstance(arr_1, np.ndarray) or isinstance(arr_1, TorchTensor): + # Initialise output array + arr_out = _dispatch.zeros_like((n_i, 2 * lambda_ + 1, n_p * n_q), like=arr_1) + + # Get the corresponding Clebsch-Gordan coefficients + cg_coeffs = cg_cache.coeffs[(l1, l2, lambda_)] + + # Fill in each mu component of the output array in turn + for m1, m2, mu in cg_coeffs.keys(): + # Broadcast arrays, multiply together and with CG coeff + arr_out[:, mu, :] += ( + arr_1[:, m1, :, None] * arr_2[:, m2, None, :] * cg_coeffs[(m1, m2, mu)] + ).reshape(n_i, n_p * n_q) + + return arr_out + + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) + + +def dense_combine( + arr_1: Union[np.ndarray, TorchTensor], + arr_2: Union[np.ndarray, TorchTensor], + lambda_: int, + cg_cache, +) -> Union[np.ndarray, TorchTensor]: + """ + Performs a Clebsch-Gordan combination step on 2 arrays using a dense + operation. The angular channel of each block is inferred from the size of + its component axis, and the blocks are combined to the desired output + angular channel `lambda_` using the appropriate Clebsch-Gordan coefficients. + + :param arr_1: array with the m values for l1 with shape [n_samples, 2 * l1 + + 1, n_q_properties] + :param arr_2: array with the m values for l2 with shape [n_samples, 2 * l2 + + 1, n_p_properties] + :param lambda_: int value of the resulting coupled channel + :param cg_cache: dense array of shape [(2 * l1 +1) * (2 * l2 +1), (2 * lambda_ + + 1)] + + :returns: array of shape [n_samples, (2*lambda_+1), q_properties * p_properties] + """ + if isinstance(arr_1, np.ndarray) or isinstance(arr_1, TorchTensor): + # Infer l1 and l2 from the len of the length of axis 1 of each tensor + l1 = (arr_1.shape[1] - 1) // 2 + l2 = (arr_2.shape[1] - 1) // 2 + cg_coeffs = cg_cache.coeffs[(l1, l2, lambda_)] + + # (samples None None l1_mu q) * (samples l2_mu p None None) + # -> (samples l2_mu p l1_mu q) we broadcast it in this way + # so we only need to do one swapaxes in the next step + arr_out = arr_1[:, None, None, :, :] * arr_2[:, :, :, None, None] + + # (samples l2_mu p l1_mu q) -> (samples q p l1_mu l2_mu) + arr_out = _dispatch.swapaxes(arr_out, 1, 4) + + # samples (q p l1_mu l2_mu) -> (samples (q p) (l1_mu l2_mu)) + arr_out = arr_out.reshape( + -1, + arr_1.shape[2] * arr_2.shape[2], + arr_1.shape[1] * arr_2.shape[1], + ) + + # (l1_mu l2_mu lam_mu) -> ((l1_mu l2_mu) lam_mu) + cg_coeffs = cg_coeffs.reshape(-1, 2 * lambda_ + 1) + + # (samples (q p) (l1_mu l2_mu)) @ ((l1_mu l2_mu) lam_mu) + # -> samples (q p) lam_mu + arr_out = arr_out @ cg_coeffs + + # (samples (q p) lam_mu) -> (samples lam_mu (q p)) + return _dispatch.swapaxes(arr_out, 1, 2) + + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/_dispatch.py b/python/rascaline/rascaline/utils/clebsch_gordan/_dispatch.py new file mode 100644 index 000000000..f44aa3713 --- /dev/null +++ b/python/rascaline/rascaline/utils/clebsch_gordan/_dispatch.py @@ -0,0 +1,133 @@ +""" +Module containing dispatch functions for numpy/torch CG combination operations. +""" +from typing import List, Optional + +import numpy as np + + +try: + import torch + from torch import Tensor as TorchTensor +except ImportError: + + class TorchTensor: + pass + + +UNKNOWN_ARRAY_TYPE = ( + "unknown array type, only numpy arrays and torch tensors are supported" +) + + +def unique(array, axis: Optional[int] = None): + """Find the unique elements of an array.""" + if isinstance(array, TorchTensor): + return torch.unique(array, dim=axis) + elif isinstance(array, np.ndarray): + return np.unique(array, axis=axis) + + +def int_range_like(min_val, max_val, like): + """Returns an array of integers from min to max, non-inclusive, based on the + type of `like`""" + if isinstance(like, TorchTensor): + return torch.arange(min_val, max_val, dtype=torch.int64, device=like.device) + elif isinstance(like, np.ndarray): + return np.arange(min_val, max_val).astype(np.int64) + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) + + +def int_array_like(int_list: List[int], like): + """ + Converts the input list of int to a numpy array or torch tensor + based on the type of `like`. + """ + if isinstance(like, TorchTensor): + return torch.tensor(int_list, dtype=torch.int64, device=like.device) + elif isinstance(like, np.ndarray): + return np.array(int_list).astype(np.int64) + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) + + +def concatenate(arrays, axis: Optional[int] = 0): + """Concatenate arrays along an axis.""" + if isinstance(arrays[0], TorchTensor): + return torch.cat(arrays, dim=axis) + elif isinstance(arrays[0], np.ndarray): + return np.concatenate(arrays, axis=axis) + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) + + +def all(array, axis: Optional[int] = None): + """Test whether all array elements along a given axis evaluate to True. + + This function has the same behavior as + ``np.all(array,axis=axis)``. + """ + if isinstance(array, bool): + array = np.array(array) + if isinstance(array, list): + array = np.array(array) + + if isinstance(array, TorchTensor): + # torch.all has two implementation, and picks one depending if more than one + # parameter is given. The second one does not supports setting dim to `None` + if axis is None: + return torch.all(input=array) + else: + return torch.all(input=array, dim=axis) + elif isinstance(array, np.ndarray): + return np.all(a=array, axis=axis) + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) + + +def any(array): + """Test whether any array elements along a given axis evaluate to True. + + This function has the same behavior as + ``np.any(array)``. + """ + if isinstance(array, bool): + array = np.array(array) + if isinstance(array, list): + array = np.array(array) + if isinstance(array, TorchTensor): + return torch.any(array) + elif isinstance(array, np.ndarray): + return np.any(array) + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) + + +def zeros_like(shape, like): + """Return an array of zeros with the same shape and type as a given array. + + This function has the same behavior as + ``np.zeros_like(array)``. + """ + if isinstance(like, TorchTensor): + return torch.zeros( + shape, + requires_grad=like.requires_grad, + dtype=like.dtype, + device=like.device, + ) + elif isinstance(like, np.ndarray): + return np.zeros(shape, dtype=like.dtype) + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) + + +def swapaxes(array, axis0: int, axis1: int): + """Swaps axes of an array.""" + if isinstance(array, TorchTensor): + return torch.swapaxes(array, axis0, axis1) + elif isinstance(array, np.ndarray): + return np.swapaxes(array, axis0, axis1) + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/clebsch_gordan.py b/python/rascaline/rascaline/utils/clebsch_gordan/clebsch_gordan.py new file mode 100644 index 000000000..5572c7598 --- /dev/null +++ b/python/rascaline/rascaline/utils/clebsch_gordan/clebsch_gordan.py @@ -0,0 +1,772 @@ +""" +Module for computing Clebsch-gordan iterations with metatensor TensorMaps. +""" +import itertools +from typing import List, Optional, Tuple, Union + +from metatensor import Labels, TensorBlock, TensorMap + +from . import _cg_cache, _dispatch + + +# ====================================================================== +# ===== Public API functions +# ====================================================================== + + +def correlate_density( + density: TensorMap, + correlation_order: int, + angular_cutoff: Optional[int] = None, + selected_keys: Optional[Union[Labels, List[Labels]]] = None, + skip_redundant: Optional[Union[bool, List[bool]]] = False, + output_selection: Optional[Union[bool, List[bool]]] = None, +) -> Union[TensorMap, List[TensorMap]]: + """ + Takes iterative Clebsch-Gordan (CG) tensor products of a density descriptor + with itself up to the desired correlation order. Returns + :py:class:`TensorMap`(s) corresponding to the density correlations output + from the specified iteration(s). + + A density descriptor necessarily is body order 2 (i.e. correlation order 1), + but can be single- or multi-center. The output is a :py:class:`list` of + density correlations for each iteration specified in `output_selection`, up + to the target order passed in `correlation_order`. By default only the last + correlation (i.e. the correlation of order ``correlation_order``) is + returned. + + This function is an iterative special case of the more general + :py:func:`correlate_tensors`. As a density is being correlated with itself, + some redundant CG tensor products can be skipped with the `skip_redundant` + keyword. + + Selections on the angular and parity channels at each iteration can also be + controlled with arguments `angular_cutoff`, `angular_selection` and + `parity_selection`. + + :param density: A density descriptor of body order 2 (correlation order 1), + in :py:class:`TensorMap` format. This may be, for example, a rascaline + :py:class:`SphericalExpansion` or :py:class:`LodeSphericalExpansion`. + Alternatively, this could be multi-center descriptor, such as a pair + density. + :param correlation_order: The desired correlation order of the output + descriptor. Must be >= 1. + :param angular_cutoff: The maximum angular channel to compute at any given + CG iteration, applied globally to all iterations until the target + correlation order is reached. + :param selected_keys: :py:class:`Labels` or `List[:py:class:`Labels`]` + specifying the angular and/or parity channels to output at each + iteration. All :py:class:`Labels` objects passed here must only contain + key names "spherical_harmonics_l" and "inversion_sigma". If a single + :py:class:`Labels` object is passed, this is applied to the final + iteration only. If a :py:class:`list` of :py:class:`Labels` objects is + passed, each is applied to its corresponding iteration. If None is + passed, all angular and parity channels are output at each iteration, + with the global `angular_cutoff` applied if specified. + :param skip_redundant: Whether to skip redundant CG combinations. Defaults + to False, which means all combinations are performed. If a + :py:class:`list` of :py:class:`bool` is passed, this is applied to each + iteration. If a single :py:class:`bool` is passed, this is applied to + all iterations. + :param output_selection: A :py:class:`list` of :py:class:`bool` specifying + whether to output a :py:class:`TensorMap` for each iteration. If a + single :py:class:`bool` is passed as True, outputs from all iterations + will be returned. If a :py:class:`list` of :py:class:`bool` is passed, + this controls the output at each corresponding iteration. If None is + passed, only the final iteration is output. + + :return: A :py:class:`list` of :py:class:`TensorMap` corresponding to the + density correlations output from the specified iterations. If the output + from a single iteration is requested, a :py:class:`TensorMap` is + returned instead. + """ + return _correlate_density( + density, + correlation_order, + angular_cutoff, + selected_keys, + skip_redundant, + output_selection, + compute_metadata_only=False, + sparse=True, # sparse CG cache by default + ) + + +def correlate_density_metadata( + density: TensorMap, + correlation_order: int, + angular_cutoff: Optional[int] = None, + selected_keys: Optional[Union[Labels, List[Labels]]] = None, + skip_redundant: Optional[Union[bool, List[bool]]] = False, + output_selection: Optional[Union[bool, List[bool]]] = None, +) -> Union[TensorMap, List[TensorMap]]: + """ + Returns the metadata-only :py:class:`TensorMap`(s) that would be output by + the function :py:func:`correlate_density` under the same settings, without + perfoming the actual Clebsch-Gordan tensor products. See this function for + full documentation. + """ + + return _correlate_density( + density, + correlation_order, + angular_cutoff, + selected_keys, + skip_redundant, + output_selection, + compute_metadata_only=True, + ) + + +# ==================================================================== +# ===== Private functions that do the work on the TensorMap level +# ==================================================================== + + +def _correlate_density( + density: TensorMap, + correlation_order: int, + angular_cutoff: Optional[int] = None, + selected_keys: Optional[Union[Labels, List[Labels]]] = None, + skip_redundant: Optional[Union[bool, List[bool]]] = False, + output_selection: Optional[Union[bool, List[bool]]] = None, + compute_metadata_only: bool = False, + sparse: bool = True, +) -> Union[TensorMap, List[TensorMap]]: + """ + Performs the density correlations for public functions + :py:func:`correlate_density` and :py:func:`correlate_density_metadata`. + """ + # Check inputs + if correlation_order <= 1: + raise ValueError("`correlation_order` must be > 1") + # TODO: implement combinations of gradients too + if _dispatch.any([len(list(block.gradients())) > 0 for block in density]): + raise NotImplementedError( + "Clebsch Gordan combinations with gradients not yet implemented." + " Use metatensor.remove_gradients to remove gradients from the input." + ) + # Check metadata + if not ( + _dispatch.all(density.keys.names == ["spherical_harmonics_l", "species_center"]) + or _dispatch.all( + density.keys.names + == ["spherical_harmonics_l", "species_center", "species_neighbor"] + ) + ): + raise ValueError( + "input `density` must have key names" + ' ["spherical_harmonics_l", "species_center"] or' + ' ["spherical_harmonics_l", "species_center", "species_neighbor"]' + ) + if not _dispatch.all(density.component_names == ["spherical_harmonics_m"]): + raise ValueError( + "input `density` must have a single component" + " axis with name `spherical_harmonics_m`" + ) + n_iterations = correlation_order - 1 # num iterations + density = _standardize_keys(density) # standardize metadata + density_correlation = density # create a copy to combine with itself + + # Parse the selected keys + selected_keys = _parse_selected_keys( + n_iterations=n_iterations, + angular_cutoff=angular_cutoff, + selected_keys=selected_keys, + like=density.keys.values, + ) + # Parse the bool flags that control skipping of redundant CG combinations + # and TensorMap output from each iteration + skip_redundant, output_selection = _parse_bool_iteration_filters( + n_iterations, + skip_redundant=skip_redundant, + output_selection=output_selection, + ) + + # Pre-compute the keys needed to perform each CG iteration + key_metadata = _precompute_keys( + density.keys, + density.keys, + n_iterations=n_iterations, + selected_keys=selected_keys, + skip_redundant=skip_redundant, + ) + # Compute CG coefficient cache + if compute_metadata_only: + cg_cache = None + else: + angular_max = max( + _dispatch.concatenate( + [density.keys.column("spherical_harmonics_l")] + + [mdata[2].column("spherical_harmonics_l") for mdata in key_metadata] + ) + ) + # TODO: keys have been precomputed, so perhaps we don't need to + # compute all CG coefficients up to angular_max here. + # TODO: use sparse cache by default until we understand under which + # circumstances (and if) dense is faster. + cg_cache = _cg_cache.ClebschGordanReal(angular_max, sparse=sparse) + + # Perform iterative CG tensor products + density_correlations = [] + for iteration in range(n_iterations): + # Define the correlation order of the current iteration + correlation_order_it = iteration + 2 + + # Combine block pairs + blocks_out = [] + for key_1, key_2, key_out in zip(*key_metadata[iteration]): + block_out = _combine_blocks_same_samples( + density_correlation[key_1], + density[key_2], + key_out["spherical_harmonics_l"], + cg_cache, + compute_metadata_only=compute_metadata_only, + ) + blocks_out.append(block_out) + keys_out = key_metadata[iteration][2] + density_correlation = TensorMap(keys=keys_out, blocks=blocks_out) + + # If this tensor is to be included in the output, move the [l1, l2, ...] + # keys to properties and store + if output_selection[iteration]: + density_correlations.append( + density_correlation.keys_to_properties( + [f"l{i}" for i in range(1, correlation_order_it + 1)] + + [f"k{i}" for i in range(2, correlation_order_it)] + ) + ) + + # Drop redundant key names. TODO: these should be part of the global + # matadata associated with the TensorMap. Awaiting this functionality in + # metatensor. + for i, tensor in enumerate(density_correlations): + keys = tensor.keys + if len(_dispatch.unique(tensor.keys.column("order_nu"))) == 1: + keys = keys.remove(name="order_nu") + if len(_dispatch.unique(tensor.keys.column("inversion_sigma"))) == 1: + keys = keys.remove(name="inversion_sigma") + density_correlations[i] = TensorMap( + keys=keys, blocks=[b.copy() for b in tensor.blocks()] + ) + + # Return a single TensorMap in the simple case + if len(density_correlations) == 1: + return density_correlations[0] + + # Otherwise return a list of TensorMaps + return density_correlations + + +# ================================================================== +# ===== Functions to handle metadata +# ================================================================== + + +def _standardize_keys(tensor: TensorMap) -> TensorMap: + """ + Takes a nu=1 tensor and standardizes its metadata. This involves: 1) moving + the "species_neighbor" key to properties, if present as a dimension in the + keys, and 2) adding dimensions in the keys for tracking the body order + ("order_nu") and parity ("inversion_sigma") of the blocks. + + Checking for the presence of the "species_neighbor" key in the keys allows + the option of the user pre-moving this key to the properties before calling + `n_body_iteration_single_center`, allowing sparsity in a set of global + neighbors to be created if desired. + + Assumes that the input `tensor` is nu=1, and has only even parity blocks. + """ + if "species_neighbor" in tensor.keys.names: + tensor = tensor.keys_to_properties(keys_to_move="species_neighbor") + keys = tensor.keys.insert( + name="order_nu", + values=_dispatch.int_array_like([1], like=tensor.keys.values), + index=0, + ) + keys = keys.insert( + name="inversion_sigma", + values=_dispatch.int_array_like([1], like=tensor.keys.values), + index=1, + ) + return TensorMap(keys=keys, blocks=[b.copy() for b in tensor.blocks()]) + + +def _parse_selected_keys( + n_iterations: int, + angular_cutoff: Optional[int] = None, + selected_keys: Optional[Union[Labels, List[Labels]]] = None, + like=None, +) -> List[Union[None, Labels]]: + """ + Parses the `selected_keys` argument passed to public functions. Checks the + values and returns a :py:class:`list` of :py:class:`Labels` objects, one for + each iteration of CG combination. + + `like` is required if a new :py:class:`Labels` object is to be created by + :py:mod:`_dispatch`. + """ + # Check angular_cutoff arg + if angular_cutoff is not None: + if not isinstance(angular_cutoff, int): + raise TypeError("`angular_cutoff` must be passed as an int") + if angular_cutoff < 1: + raise ValueError("`angular_cutoff` must be >= 1") + + if selected_keys is None: + if angular_cutoff is None: # no selections at all + selected_keys = [None] * n_iterations + else: + # Create a key selection with all angular channels <= the specified + # angular cutoff + selected_keys = [ + Labels( + names=["spherical_harmonics_l"], + values=_dispatch.int_range_like( + 0, angular_cutoff, like=like + ).reshape(-1, 1), + ) + ] * n_iterations + + if isinstance(selected_keys, Labels): + # Create a list, but only apply a key selection at the final iteration + selected_keys = [None] * (n_iterations - 1) + [selected_keys] + + # Check the selected_keys + if not isinstance(selected_keys, List): + raise TypeError( + "`selected_keys` must be a `Labels` or List[Union[None, `Labels`]]" + ) + if not len(selected_keys) == n_iterations: + raise ValueError( + "`selected_keys` must be a List[Union[None, Labels]] of length" + " `correlation_order` - 1" + ) + if not _dispatch.all( + [isinstance(val, (Labels, type(None))) for val in selected_keys] + ): + raise TypeError("`selected_keys` must be a Labels or List[Union[None, Labels]]") + + # Now iterate over each of the Labels (or None) in the list and check + for slct in selected_keys: + if slct is None: + continue + assert isinstance(slct, Labels) + if not _dispatch.all( + [ + name in ["spherical_harmonics_l", "inversion_sigma"] + for name in slct.names + ] + ): + raise ValueError( + "specified key names in `selected_keys` must be either" + " 'spherical_harmonics_l' or 'inversion_sigma'" + ) + if "spherical_harmonics_l" in slct.names: + if angular_cutoff is not None: + if not _dispatch.all( + slct.column("spherical_harmonics_l") <= angular_cutoff + ): + raise ValueError( + "specified angular channels in `selected_keys` must be <= the" + " specified `angular_cutoff`" + ) + if not _dispatch.all( + [angular_l >= 0 for angular_l in slct.column("spherical_harmonics_l")] + ): + raise ValueError( + "specified angular channels in `selected_keys` must be >= 0" + ) + if "inversion_sigma" in slct.names: + if not _dispatch.all( + [parity_s in [-1, +1] for parity_s in slct.column("inversion_sigma")] + ): + raise ValueError( + "specified parities in `selected_keys` must be -1 or +1" + ) + + return selected_keys + + +def _parse_bool_iteration_filters( + n_iterations: int, + skip_redundant: Optional[Union[bool, List[bool]]] = False, + output_selection: Optional[Union[bool, List[bool]]] = None, +) -> List[List[bool]]: + """ + Parses the `skip_redundant` and `output_selection` arguments passed to + public functions. + """ + if isinstance(skip_redundant, bool): + skip_redundant = [skip_redundant] * n_iterations + if not _dispatch.all([isinstance(val, bool) for val in skip_redundant]): + raise TypeError("`skip_redundant` must be a `bool` or `list` of `bool`") + if not len(skip_redundant) == n_iterations: + raise ValueError( + "`skip_redundant` must be a bool or `list` of `bool` of length" + " `correlation_order` - 1" + ) + if output_selection is None: + output_selection = [False] * (n_iterations - 1) + [True] + else: + if isinstance(output_selection, bool): + output_selection = [output_selection] * n_iterations + if not isinstance(output_selection, List): + raise TypeError("`output_selection` must be passed as `list` of `bool`") + + if not len(output_selection) == n_iterations: + raise ValueError( + "`output_selection` must be a ``list`` of ``bool`` of length" + " corresponding to the number of CG iterations" + ) + if not _dispatch.all([isinstance(v, bool) for v in output_selection]): + raise TypeError("`output_selection` must be passed as a `list` of `bool`") + if not _dispatch.all([isinstance(v, bool) for v in output_selection]): + raise TypeError("`output_selection` must be passed as a `list` of `bool`") + + return skip_redundant, output_selection + + +def _precompute_keys( + keys_1: Labels, + keys_2: Labels, + n_iterations: int, + selected_keys: List[Union[None, Labels]], + skip_redundant: List[bool], +) -> List[Tuple[Labels, List[List[int]]]]: + """ + Computes all the keys metadata needed to perform `n_iterations` of CG + combination steps. + + At each iteration, a full product of the keys of two tensors, i.e. `keys_1` + and `keys_2` is computed. Then, key selections are applied according to the + user-defined settings: the maximum angular channel cutoff + (`angular_cutoff`), and angular and/or parity selections specified in + `selected_keys`. + + If `skip_redundant` is True, then keys that represent redundant CG + operations are not included in the output keys at each step. + """ + keys_metadata = [] + keys_out = keys_1 + for iteration in range(n_iterations): + # Get the keys metadata for the combination of the 2 tensors + keys_1_entries, keys_2_entries, keys_out = _precompute_keys_full_product( + keys_1=keys_out, + keys_2=keys_2, + ) + if selected_keys[iteration] is not None: + keys_1_entries, keys_2_entries, keys_out = _apply_key_selection( + keys_1_entries, + keys_2_entries, + keys_out, + selected_keys=selected_keys[iteration], + ) + + if skip_redundant[iteration]: + keys_1_entries, keys_2_entries, keys_out = _remove_redundant_keys( + keys_1_entries, keys_2_entries, keys_out + ) + + # Check that some keys are produced as a result of the combination + if len(keys_out) == 0: + raise ValueError( + f"invalid selections: iteration {iteration + 1} produces no" + " valid combinations. Check the `angular_cutoff` and" + " `selected_keys` args and try again." + ) + + keys_metadata.append((keys_1_entries, keys_2_entries, keys_out)) + + return keys_metadata + + +def _precompute_keys_full_product( + keys_1: Labels, keys_2: Labels +) -> Tuple[List, List, Labels]: + """ + Given the keys of 2 TensorMaps, returns the keys that would be present after + a full CG product of these TensorMaps. + + Assumes that `keys_1` corresponds to a TensorMap with arbitrary body order, + while `keys_2` corresponds to a TensorMap with body order 1. `keys_1` must + follow the key name convention: + + ["order_nu", "inversion_sigma", "spherical_harmonics_l", "species_center", + "l1", "l2", ..., f"l{`nu`}", "k2", ..., f"k{`nu`-1}"]. The "lx" columns + track the l values of the nu=1 blocks that were previously combined. The + "kx" columns tracks the intermediate lambda values of nu > 1 blocks that + have been combined. + + For instance, a TensorMap of body order nu=4 will have key names + ["order_nu", "inversion_sigma", "spherical_harmonics_l", "species_center", + "l1", "l2", "l3", "l4", "k2", "k3"]. Two nu=1 TensorMaps with blocks of + order "l1" and "l2" were combined to form a nu=2 TensorMap with blocks of + order "k2". This was combined with a nu=1 TensorMap with blocks of order + "l3" to form a nu=3 TensorMap with blocks of order "k3". Finally, this was + combined with a nu=1 TensorMap with blocks of order "l4" to form a nu=4. + + .. math :: + + \\bra{ + n_1 l_1 ; n_2 l_2 k_2 ; ... ; + n_{\nu-1} l_{\\nu-1} k_{\\nu-1} ; + n_{\\nu} l_{\\nu} k_{\\nu}; \\lambda + } + \\ket{ \\rho^{\\otimes \\nu}; \\lambda M } + + `keys_2` must follow the key name convention: ["order_nu", + "inversion_sigma", "spherical_harmonics_l", "species_center"] + + Returned is Tuple[List, List, Labels]. The first two lists correspond to the + LabelsEntry objects of the keys being combined. The third element is a + Labels object corresponding to the keys of the output TensorMap. Each entry + in this Labels object corresponds to the keys is formed by combination of + the pair of blocks indexed by correspoding key pairs in the first two lists. + """ + # Get the correlation order of the first TensorMap. + unique_nu = _dispatch.unique(keys_1.column("order_nu")) + if len(unique_nu) > 1: + raise ValueError( + "keys_1 must correspond to a tensor of a single correlation order." + f" Found {len(unique_nu)} body orders: {unique_nu}" + ) + nu1 = unique_nu[0] + + # Define new correlation order of output TensorMap + nu = nu1 + 1 + + # The correlation order of the second TensorMap should be nu = 1. + assert _dispatch.all(keys_2.column("order_nu") == 1) + + # If nu1 = 1, the key names don't yet have any "lx" columns + if nu1 == 1: + l_list_names = [] + new_l_list_names = ["l1", "l2"] + else: + l_list_names = [f"l{angular_l}" for angular_l in range(1, nu1 + 1)] + new_l_list_names = l_list_names + [f"l{nu}"] + + # Check key names + assert _dispatch.all( + keys_1.names + == ["order_nu", "inversion_sigma", "spherical_harmonics_l", "species_center"] + + l_list_names + + [f"k{k}" for k in range(2, nu1)] + ) + assert _dispatch.all( + keys_2.names + == ["order_nu", "inversion_sigma", "spherical_harmonics_l", "species_center"] + ) + + # Define key names of output Labels (i.e. for combined TensorMap) + new_names = ( + ["order_nu", "inversion_sigma", "spherical_harmonics_l", "species_center"] + + new_l_list_names + + [f"k{k}" for k in range(2, nu)] + ) + + new_key_values = [] + keys_1_entries = [] + keys_2_entries = [] + for key_1, key_2 in itertools.product(keys_1, keys_2): + # Unpack relevant key values + sig1, lam1, a = key_1.values[1:4] + sig2, lam2, a2 = key_2.values[1:4] + + # Only combine blocks of the same chemical species + if a != a2: + continue + + # First calculate the possible non-zero angular channels that can be + # formed from combination of blocks of order `lam1` and `lam2`. This + # corresponds to values in the inclusive range { |lam1 - lam2|, ..., + # |lam1 + lam2| } + nonzero_lams = _dispatch.int_range_like( + abs(lam1 - lam2), abs(lam1 + lam2) + 1, like=key_1.values + ) + + # Now iterate over the non-zero angular channels and apply the custom + # selections + for lambda_ in nonzero_lams: + # Calculate new sigma + sig = sig1 * sig2 * (-1) ** (lam1 + lam2 + lambda_) + + # Extract the l and k lists from keys_1 + l_list = key_1.values[4 : 4 + nu1].tolist() + k_list = key_1.values[4 + nu1 :].tolist() + + # Build the new keys values. l{nu} is `lam2`` (i.e. + # "spherical_harmonics_l" of the key from `keys_2`. k{nu-1} is + # `lam1` (i.e. "spherical_harmonics_l" of the key from `keys_1`). + new_vals = [nu, sig, lambda_, a] + l_list + [lam2] + k_list + [lam1] + new_key_values.append(new_vals) + keys_1_entries.append(key_1) + keys_2_entries.append(key_2) + + # Define new keys as the full product of keys_1 and keys_2 + keys_out = Labels( + names=new_names, + values=_dispatch.int_array_like(new_key_values, like=keys_1.values), + ) + + return keys_1_entries, keys_2_entries, keys_out + + +def _apply_key_selection( + keys_1_entries: List, keys_2_entries: List, keys_out: Labels, selected_keys: Labels +) -> Tuple[List, List, Labels]: + """ + Applies a selection according to `selected_keys` to the keys of an output + TensorMap `keys_out` produced by combination of blocks indexed by keys + entries in `keys_1_entries` and `keys_2_entries` lists. + + After application of the selections, returned is a reduced set of keys and + set of corresponding parents key entries. + + If a selection in `selected_keys` is not valid based on the keys in + `keys_out`, an error is raised. + """ + # Extract the relevant columns from `selected_keys` that the selection will + # be performed on + keys_out_vals = [[k[name] for name in selected_keys.names] for k in keys_out] + + # First check that all of the selected keys exist in the output keys + for slct in selected_keys.values: + if not _dispatch.any([_dispatch.all(slct == k) for k in keys_out_vals]): + raise ValueError( + f"selected key {selected_keys.names} = {slct} not found" + " in the output keys. Check the `selected_keys` argument." + ) + + # Build a mask of the selected keys + mask = [ + _dispatch.any([_dispatch.all(i == j) for j in selected_keys.values]) + for i in keys_out_vals + ] + + # Apply the mask to key entries and keys and return + keys_1_entries = [k for k, isin in zip(keys_1_entries, mask) if isin] + keys_2_entries = [k for k, isin in zip(keys_2_entries, mask) if isin] + keys_out = Labels(names=keys_out.names, values=keys_out.values[mask]) + + return keys_1_entries, keys_2_entries, keys_out + + +def _remove_redundant_keys( + keys_1_entries: List, keys_2_entries: List, keys_out: Labels +) -> Tuple[List, List, Labels]: + """ + For a Labels object `keys_out` that corresponds to the keys of a TensorMap + formed by combined of the blocks described by the entries in the lists + `keys_1_entries` and `keys_2_entries`, removes redundant keys. + + These are the keys that correspond to blocks that have the same sorted l + list. The block where the l values are already sorted (i.e. l1 <= l2 <= ... + <= ln) is kept. + """ + # Get and check the correlation order of the input keys + nu1 = keys_1_entries[0]["order_nu"] + nu2 = keys_2_entries[0]["order_nu"] + assert nu2 == 1 + + # Get the correlation order of the output TensorMap + nu = nu1 + 1 + + # Identify keys of redundant blocks and remove them + key_idxs_to_keep = [] + for key_idx, key in enumerate(keys_out): + # Get the important key values. This is all of the keys, excpet the k + # list + key_vals_slice = key.values[: 4 + (nu + 1)].tolist() + first_part, l_list = key_vals_slice[:4], key_vals_slice[4:] + + # Sort the l list + l_list_sorted = sorted(l_list) + + # Compare the sliced key with the one recreated when the l list is + # sorted. If they are identical, this is the key of the block that we + # want to compute a CG combination for. + key_slice_tuple = tuple(first_part + l_list) + key_slice_sorted_tuple = tuple(first_part + l_list_sorted) + if _dispatch.all(key_slice_tuple == key_slice_sorted_tuple): + key_idxs_to_keep.append(key_idx) + + # Build a reduced Labels object for the combined keys, with redundancies removed + keys_out_red = Labels( + names=keys_out.names, + values=_dispatch.int_array_like( + [keys_out[idx].values for idx in key_idxs_to_keep], + like=keys_1_entries[0].values, + ), + ) + + # Store the list of reduced entries that combine to form the reduced output keys + keys_1_entries_red = [keys_1_entries[idx] for idx in key_idxs_to_keep] + keys_2_entries_red = [keys_2_entries[idx] for idx in key_idxs_to_keep] + + return keys_1_entries_red, keys_2_entries_red, keys_out_red + + +# ================================================================== +# ===== Functions to perform the CG combinations of blocks +# ================================================================== + + +def _combine_blocks_same_samples( + block_1: TensorBlock, + block_2: TensorBlock, + lambda_: int, + cg_cache, + compute_metadata_only: bool = False, +) -> TensorBlock: + """ + For a given pair of TensorBlocks and desired angular channel, combines the + values arrays and returns a new TensorBlock. + """ + + # Do the CG combination - single center so no shape pre-processing required + if compute_metadata_only: + combined_values = _cg_cache.combine_arrays( + block_1.values, block_2.values, lambda_, cg_cache, return_empty_array=True + ) + else: + combined_values = _cg_cache.combine_arrays( + block_1.values, block_2.values, lambda_, cg_cache, return_empty_array=False + ) + + # Infer the new nu value: block 1's properties are nu pairs of + # "species_neighbor_x" and "nx". + combined_nu = int((len(block_1.properties.names) / 2) + 1) + + # Define the new property names for "nx" and "species_neighbor_x" + n_names = [f"n{i}" for i in range(1, combined_nu + 1)] + neighbor_names = [f"species_neighbor_{i}" for i in range(1, combined_nu + 1)] + prop_names = [item for i in zip(neighbor_names, n_names) for item in i] + + # Create a TensorBlock + combined_block = TensorBlock( + values=combined_values, + samples=block_1.samples, + components=[ + Labels( + names=["spherical_harmonics_m"], + values=_dispatch.int_range_like( + min_val=-lambda_, max_val=lambda_ + 1, like=block_1.values + ).reshape(-1, 1), + ), + ], + properties=Labels( + names=prop_names, + values=_dispatch.int_array_like( + [ + _dispatch.concatenate((b2, b1)) + for b2 in block_2.properties.values + for b1 in block_1.properties.values + ], + like=block_1.values, + ), + ), + ) + + return combined_block diff --git a/python/rascaline/rascaline/utils/power_spectrum/calculator.py b/python/rascaline/rascaline/utils/power_spectrum/calculator.py index 6815f4b16..139a3e091 100644 --- a/python/rascaline/rascaline/utils/power_spectrum/calculator.py +++ b/python/rascaline/rascaline/utils/power_spectrum/calculator.py @@ -217,7 +217,10 @@ def compute( ell = key[0] species_center = key[1] - factor = 1 / sqrt(2 * ell + 1) + # For consistency with a full Clebsch-Gordan product we need to add + # a `-1^l / sqrt(2 l + 1)` factor to the power spectrum invariants + factor = (-1) ** ell / sqrt(2 * ell + 1) + # Find that block indices that have the same spherical_harmonics_l and # species_center selection = Labels( diff --git a/python/rascaline/rascaline/utils/splines.py b/python/rascaline/rascaline/utils/splines.py deleted file mode 100644 index af8534b5e..000000000 --- a/python/rascaline/rascaline/utils/splines.py +++ /dev/null @@ -1,403 +0,0 @@ -""" -Splined radial integrals -======================== - -Classes for generating splines which can be used as tabulated radial integrals in the -various SOAP and LODE calculators. For an complete example of how to use these classes -see :ref:`example-splines`. - -.. autoclass:: rascaline.utils.RadialIntegralSplinerBase - :members: - :show-inheritance: - -.. autoclass:: rascaline.utils.RadialIntegralFromFunction - :members: - :show-inheritance: -""" - -from abc import ABC, abstractmethod -from typing import Callable, Dict, Optional, Union - -import numpy as np - - -class RadialIntegralSplinerBase(ABC): - """Base class for splining arbitrary radial integrals. - - If ``_radial_integral_derivative`` is not implemented in a child class it will - computed based on finite differences. - - :parameter max_angular: number of radial components - :parameter max_radial: number of angular components - :parameter spline_cutoff: cutoff radius for the spline interpolation. This is also - the maximal value that can be interpolated. - :parameter accuracy: accuracy of the numerical integration and the splining. - Accuracy is reached when either the mean absolute error or the mean relative - error gets below the ``accuracy`` threshold. - """ - - def __init__( - self, - max_radial: int, - max_angular: int, - spline_cutoff: float, - accuracy: float, - ): - self.max_radial = max_radial - self.max_angular = max_angular - self.spline_cutoff = spline_cutoff - self.accuracy = accuracy - - @abstractmethod - def _radial_integral(self, n: int, ell: int, positions: np.ndarray) -> np.ndarray: - """Method calculating the radial integral.""" - ... - - @property - def _center_contribution(self) -> Union[None, np.ndarray]: - r"""Contribution of the central atom required for LODE calculations.""" - - return None - - def _radial_integral_derivative( - self, n: int, ell: int, positions: np.ndarray - ) -> np.ndarray: - """Method calculating the derivatice of the radial integral.""" - displacement = 1e-6 - mean_abs_positions = np.abs(positions).mean() - - if mean_abs_positions <= 1.0: - raise ValueError( - "Numerically derivative of the radial integral can not be performed " - "since positions are too small. Mean of the absolute positions is " - f"{mean_abs_positions:.1e} but should be at least 1." - ) - - radial_integral_pos = self._radial_integral( - n, ell, positions + displacement / 2 - ) - radial_integral_neg = self._radial_integral( - n, ell, positions - displacement / 2 - ) - - return (radial_integral_pos - radial_integral_neg) / displacement - - def _value_evaluator_3D( - self, - positions: np.ndarray, - derivative: bool, - ): - values = np.zeros([len(positions), self.max_angular + 1, self.max_radial]) - for ell in range(self.max_angular + 1): - for n in range(self.max_radial): - if derivative: - values[:, ell, n] = self._radial_integral_derivative( - n, ell, positions - ) - else: - values[:, ell, n] = self._radial_integral(n, ell, positions) - - return values - - def compute( - self, - n_spline_points: Optional[int] = None, - ) -> Dict: - """Compute the spline for rascaline's tabulated radial integrals. - - :parameter n_spline_points: Use fixed number of spline points instead of find - the number based on the provided ``accuracy``. - :returns dict: dictionary for the input as the ``radial_basis`` parameter of a - rascaline calculator. - """ - - def value_evaluator_3D(positions): - return self._value_evaluator_3D(positions, derivative=False) - - def derivative_evaluator_3D(positions): - return self._value_evaluator_3D(positions, derivative=True) - - if n_spline_points is not None: - positions = np.linspace(0, self.spline_cutoff, n_spline_points) - values = value_evaluator_3D(positions) - derivatives = derivative_evaluator_3D(positions) - else: - dynamic_spliner = DynamicSpliner( - 0, - self.spline_cutoff, - value_evaluator_3D, - derivative_evaluator_3D, - self.accuracy, - ) - positions, values, derivatives = dynamic_spliner.spline() - - # Convert positions, values, derivatives into the appropriate json formats: - spline_points = [] - for position, value, derivative in zip(positions, values, derivatives): - spline_points.append( - { - "position": position, - "values": { - "v": 1, - "dim": value.shape, - "data": value.flatten().tolist(), - }, - "derivatives": { - "v": 1, - "dim": derivative.shape, - "data": derivative.flatten().tolist(), - }, - } - ) - - parameters = {"points": spline_points} - - center_contribution = self._center_contribution - if center_contribution is not None: - parameters["center_contribution"] = center_contribution - - return {"TabulatedRadialIntegral": parameters} - - -class DynamicSpliner: - def __init__( - self, - start: float, - stop: float, - values_fn: Callable[[np.ndarray], np.ndarray], - derivatives_fn: Callable[[np.ndarray], np.ndarray], - accuracy: float = 1e-8, - ) -> None: - """Dynamic spline generator. - - This class can be used to spline any set of functions defined within the - start-stop interval. Cubic Hermite splines - (https://en.wikipedia.org/wiki/Cubic_Hermite_spline) are used. The same spline - points will be used for all functions, and more will be added until either the - relative error or the absolute error fall below the requested accuracy on - average across all functions. The functions are specified via values_fn and - derivatives_fn. These must be able to take a numpy 1D array of positions as - their input, and they must output a numpy array where the first dimension - corresponds to the input positions, while other dimensions are arbitrary and can - correspond to any way in which the target functions can be classified. The - splines can be obtained via the spline method. - """ - - self.start = start - self.stop = stop - self.values_fn = values_fn - self.derivatives_fn = derivatives_fn - self.requested_accuracy = accuracy - - # initialize spline with 11 points - positions = np.linspace(start, stop, 11) - self.spline_positions = positions - self.spline_values = values_fn(positions) - self.spline_derivatives = derivatives_fn(positions) - - self.number_of_custom_axes = len(self.spline_values.shape) - 1 - - def spline(self): - """Calculates and outputs the splines. - - The outputs of this function are, respectively: - A numpy 1D array containing - the spline positions. These are equally spaced in the start-stop interval. - - A numpy ndarray containing the values of the splined functions at the spline - positions. The first dimension corresponds to the spline positions, while all - subsequent dimensions are consistent with the values_fn and - `get_function_derivative` provided during initialization of the class. - - A numpy ndarray containing the derivatives of the splined functions at the - spline positions, with the same structure as that of the ndarray of values. - """ - - while True: - n_intermediate_positions = len(self.spline_positions) - 1 - - if n_intermediate_positions >= 50000: - raise ValueError( - "Maximum number of spline points reached. \ - There might be a problem with the functions to be splined" - ) - - half_step = (self.spline_positions[1] - self.spline_positions[0]) / 2 - intermediate_positions = np.linspace( - self.start + half_step, self.stop - half_step, n_intermediate_positions - ) - - estimated_values = self._compute_from_spline(intermediate_positions) - new_values = self.values_fn(intermediate_positions) - - mean_absolute_error = np.mean(np.abs(estimated_values - new_values)) - with np.errstate(divide="ignore"): # Ignore divide-by-zero warnings - mean_relative_error = np.mean( - np.abs((estimated_values - new_values) / new_values) - ) - - if ( - mean_absolute_error < self.requested_accuracy - or mean_relative_error < self.requested_accuracy - ): - break - - new_derivatives = self.derivatives_fn(intermediate_positions) - - concatenated_positions = np.concatenate( - [self.spline_positions, intermediate_positions], axis=0 - ) - concatenated_values = np.concatenate( - [self.spline_values, new_values], axis=0 - ) - concatenated_derivatives = np.concatenate( - [self.spline_derivatives, new_derivatives], axis=0 - ) - - sort_indices = np.argsort(concatenated_positions, axis=0) - - self.spline_positions = concatenated_positions[sort_indices] - self.spline_values = concatenated_values[sort_indices] - self.spline_derivatives = concatenated_derivatives[sort_indices] - - return self.spline_positions, self.spline_values, self.spline_derivatives - - def _compute_from_spline(self, positions): - x = positions - delta_x = self.spline_positions[1] - self.spline_positions[0] - n = (np.floor(x / delta_x)).astype(np.int32) - - t = (x - n * delta_x) / delta_x - t_2 = t**2 - t_3 = t**3 - - h00 = 2.0 * t_3 - 3.0 * t_2 + 1.0 - h10 = t_3 - 2.0 * t_2 + t - h01 = -2.0 * t_3 + 3.0 * t_2 - h11 = t_3 - t_2 - - p_k = self.spline_values[n] - p_k_1 = self.spline_values[n + 1] - - m_k = self.spline_derivatives[n] - m_k_1 = self.spline_derivatives[n + 1] - - new_shape = (-1,) + (1,) * self.number_of_custom_axes - h00 = h00.reshape(new_shape) - h10 = h10.reshape(new_shape) - h01 = h01.reshape(new_shape) - h11 = h11.reshape(new_shape) - - interpolated_values = ( - h00 * p_k + h10 * delta_x * m_k + h01 * p_k_1 + h11 * delta_x * m_k_1 - ) - - return interpolated_values - - -class RadialIntegralFromFunction(RadialIntegralSplinerBase): - r"""Compute the radial integral spline points based on a provided function. - - :parameter radial_integral: Function to compute the radial integral. Function must - take ``n``, ``l``, and ``positions`` as inputs, where ``n`` and ``l`` are - integers and ``positions`` is a numpy 1-D array that contains the spline points - at which the radial integral will be evaluated. The function must return a numpy - 1-D array containing the values of the radial integral. - :parameter spline_cutoff: cutoff radius for the spline interpolation. This is also - the maximal value that can be interpolated. - :parameter max_radial: number of angular componentss - :parameter max_angular: number of radial components - :parameter radial_integral_derivative: The derivative of the radial integral taking - the same paramaters as ``radial_integral``. If it is :py:obj:`None` (default), - finite differences are used to calculate the derivative of the radial integral. - It is recommended to provide this parameter if possible. Derivatives from finite - differences can cause problems when evaluating at the edges of the domain (i.e., - at ``0`` and ``spline_cutoff``) because the function might not be defined - outside of the domain. - :parameter accuracy: accuracy of the numerical integration and the splining. - Accuracy is reached when either the mean absolute error or the mean relative - error gets below the ``accuracy`` threshold. - :parameter center_contribution: Contribution of the central atom required for LODE - calculations. The ``center_contribution`` is defined as - - .. math:: - c_n = \sqrt{4π}\int_0^\infty dr r^2 R_n(r) g(r) - - where :math:`g(r)` is the radially symmetric density function, `R_n(r)` the - radial basis function and :math:`n` the current radial channel. This should be - pre-computed and provided as a separate parameter. - - Example - ------- - First define a ``radial_integral`` function - - >>> def radial_integral(n, ell, r): - ... return np.sin(r) - ... - - and provide this as input to the spline generator - - >>> spliner = RadialIntegralFromFunction( - ... radial_integral=radial_integral, - ... max_radial=12, - ... max_angular=9, - ... spline_cutoff=8.0, - ... ) - - Finally, we can use the ``spliner`` directly in the ``radial_integral`` section of a - calculator - - >>> from rascaline import SoapPowerSpectrum - >>> calculator = SoapPowerSpectrum( - ... cutoff=8.0, - ... max_radial=12, - ... max_angular=9, - ... center_atom_weight=1.0, - ... radial_basis=spliner.compute(), - ... atomic_gaussian_width=1.0, # ignored - ... cutoff_function={"Step": {}}, - ... ) - - The ``atomic_gaussian_width`` paramater is required by the calculator but will be - will be ignored during the feature computation. - - A more in depth example using a "rectangular" Laplacian eigenstate basis - is provided in the :ref:`example section`. - """ - - def __init__( - self, - radial_integral: Callable[[int, int, np.ndarray], np.ndarray], - spline_cutoff: float, - max_radial: int, - max_angular: int, - radial_integral_derivative: Optional[ - Callable[[int, int, np.ndarray], np.ndarray] - ] = None, - center_contribution: Optional[np.ndarray] = None, - accuracy: float = 1e-8, - ): - self.radial_integral_function = radial_integral - self.radial_integral_derivative_funcion = radial_integral_derivative - self.center_contribution = center_contribution - - super().__init__( - max_radial=max_radial, - max_angular=max_angular, - spline_cutoff=spline_cutoff, - accuracy=accuracy, - ) - - def _radial_integral(self, n: int, ell: int, positions: np.ndarray) -> np.ndarray: - return self.radial_integral_function(n, ell, positions) - - @property - def _center_contribution(self) -> Union[None, np.ndarray]: - # Test that ``len(self.center_contribution) == max_radial`` is performed by the - # calculator. - return self.center_contribution - - def _radial_integral_derivative( - self, n: int, ell: int, positions: np.ndarray - ) -> np.ndarray: - if self.radial_integral_derivative_funcion is None: - return super()._radial_integral_derivative(n, ell, positions) - else: - return self.radial_integral_derivative_funcion(n, ell, positions) diff --git a/python/rascaline/rascaline/utils/splines/__init__.py b/python/rascaline/rascaline/utils/splines/__init__.py new file mode 100644 index 000000000..247aaa83e --- /dev/null +++ b/python/rascaline/rascaline/utils/splines/__init__.py @@ -0,0 +1,18 @@ +from .atomic_density import ( # noqa + AtomicDensityBase, + DeltaDensity, + GaussianDensity, + LodeDensity, +) +from .radial_basis import ( # noqa + GtoBasis, + MonomialBasis, + RadialBasisBase, + SphericalBesselBasis, +) +from .splines import ( # noqa + LodeSpliner, + RadialIntegralFromFunction, + RadialIntegralSplinerBase, + SoapSpliner, +) diff --git a/python/rascaline/rascaline/utils/splines/atomic_density.py b/python/rascaline/rascaline/utils/splines/atomic_density.py new file mode 100644 index 000000000..b1e074391 --- /dev/null +++ b/python/rascaline/rascaline/utils/splines/atomic_density.py @@ -0,0 +1,166 @@ +r""" +.. _python-atomic-density: + +Atomic Density +============== + +the atomic density function :math:`g(r)`, often chosen to be a Gaussian or Delta +function, that defined the type of density under consideration. For a given center atom +:math:`i` in the structure, the total density function :math:`\rho_i(\boldsymbol{r})` +around is then defined as :math:`\rho_i(\boldsymbol{r}) = \sum_{j} g(\boldsymbol{r} - +\boldsymbol{r}_{ij})`. + +Atomic densities are represented as different child class of +:py:class:`rascaline.utils.AtomicDensityBase`: :py:class:`rascaline.utils.DeltaDensity`, +:py:class:`rascaline.utils.GaussianDensity`, and :py:class:`rascaline.utils.LodeDensity` +are provided, and you can implement your own by defining a new class. + +.. autoclass:: rascaline.utils.AtomicDensityBase + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.DeltaDensity + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.GaussianDensity + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.LodeDensity + :members: + :show-inheritance: + +""" +from abc import ABC, abstractmethod +from typing import Union + +import numpy as np + + +try: + from scipy.special import gamma, gammainc + + HAS_SCIPY = True +except ImportError: + HAS_SCIPY = False + + +class AtomicDensityBase(ABC): + """Base class representing atomic densities.""" + + @abstractmethod + def compute(self, positions: Union[float, np.ndarray]) -> Union[float, np.ndarray]: + """Compute the atomic density arising from atoms at ``positions`` + + :param positions: positions to evaluate the atomic densities + :returns: evaluated atomic density + """ + + +class DeltaDensity(AtomicDensityBase): + r"""Delta atomic densities of the form :math:`g(r)=\delta(r)`.""" + + def compute(self, positions: Union[float, np.ndarray]) -> Union[float, np.ndarray]: + raise ValueError( + "Compute function of the delta density should never called directly." + ) + + +class GaussianDensity(AtomicDensityBase): + r"""Gaussian atomic density function. + + In rascaline, we use the convention + + .. math:: + + g(r) = \frac{1}{(\pi \sigma^2)^{3/4}}e^{-\frac{r^2}{2\sigma^2}} \,. + + The prefactor was chosen such that the "L2-norm" of the Gaussian + + .. math:: + + \|g\|^2 = \int \mathrm{d}^3\boldsymbol{r} |g(r)|^2 = 1\,, + + :param atomic_gaussian_width: Width of the atom-centered gaussian used to create the + atomic density + """ + + def __init__(self, atomic_gaussian_width: float): + self.atomic_gaussian_width = atomic_gaussian_width + + def compute(self, positions: Union[float, np.ndarray]) -> Union[float, np.ndarray]: + atomic_gaussian_width_sq = self.atomic_gaussian_width**2 + return np.exp(-0.5 * positions**2 / atomic_gaussian_width_sq) / ( + np.pi * atomic_gaussian_width_sq + ) ** (3 / 4) + + +class LodeDensity(AtomicDensityBase): + r"""Smeared power law density, as used in LODE. + + It is defined as + + .. math:: + + g(r) = \frac{1}{\Gamma\left(\frac{p}{2}\right)} + \frac{\gamma\left( \frac{p}{2}, \frac{r^2}{2\sigma^2} \right)} + {r^p}, + + where :math:`\Gamma(z)` is the Gamma function and :math:`\gamma(a, x)` is the + incomplete lower Gamma function. However its evaluation at :math:`r=0` is + problematic because :math:`g(r)` is of the form :math:`0/0`. For practical + implementations, it is thus more convenient to rewrite the density as + + .. math:: + + g(r) = \frac{1}{\Gamma(a)}\frac{1}{\left(2 \sigma^2\right)^a} + \begin{cases} + \frac{1}{a} - \frac{x}{a+1} + \frac{x^2}{2(a+2)} + \mathcal{O}(x^3) + & x < 10^{-5} \\ + \frac{\gamma(a,x)}{x^a} + & x \geq 10^{-5} + \end{cases} + + It is convenient to use the expression for sufficiently small :math:`x` since the + relative weight of the first neglected term is on the order of :math:`1/6x^3`. + Therefore, the threshold :math:`x = 10^{-5}` leads to relative errors on the order + of the machine epsilon. + + :param atomic_gaussian_width: Width of the atom-centered gaussian used to create the + atomic density + :param potential_exponent: Potential exponent of the decorated atom density. + Currently only implemented for potential_exponent < 10. Some exponents can be + connected to SOAP or physics-based quantities: p=0 uses Gaussian densities as in + SOAP, p=1 uses 1/r Coulomb like densities, p=6 uses 1/r^6 dispersion like + densities. + """ + + def __init__(self, atomic_gaussian_width: float, potential_exponent: int): + if not HAS_SCIPY: + raise ValueError("LodeDensity requires scipy to be installed") + + self.potential_exponent = potential_exponent + self.atomic_gaussian_width = atomic_gaussian_width + + def _short_range(self, a, x): + return 1 / a - x / (a + 1) + x**2 / (2 * (a + 2)) + + def _long_range(self, a, x): + return gamma(a) * gammainc(a, x) / x**a + + def compute(self, positions: Union[float, np.ndarray]) -> Union[float, np.ndarray]: + if self.potential_exponent == 0: + return GaussianDensity.compute(self, positions=positions) + else: + atomic_gaussian_width_sq = self.atomic_gaussian_width**2 + a = self.potential_exponent / 2 + x = positions**2 / (2 * atomic_gaussian_width_sq) + + prefac = 1 / gamma(a) / (2 * atomic_gaussian_width_sq) ** a + + return prefac * np.where( + x < 1e-5, + self._short_range(a, x), + self._long_range(a, x), + ) diff --git a/python/rascaline/rascaline/utils/splines/radial_basis.py b/python/rascaline/rascaline/utils/splines/radial_basis.py new file mode 100644 index 000000000..1b68396b3 --- /dev/null +++ b/python/rascaline/rascaline/utils/splines/radial_basis.py @@ -0,0 +1,331 @@ +r""" +.. _python-radial-basis: + +Radial Basis +============ + +Radial basis functions :math:`R_{nl}(\boldsymbol{r})` are besides :ref:`atomic densities +` :math:`\rho_i` the central ingredients to compute spherical +expansion coefficients :math:`\langle anlm\vert\rho_i\rangle`. Radial basis functions, +define how which the atomic density is projected. To be more precise, the actual basis +functions are of + +.. math:: + + B_{nlm}(\boldsymbol{r}) = R_{nl}(r)Y_{lm}(\hat{r}) \,, + +where :math:`Y_{lm}(\hat{r})` are the real spherical harmonics evaluated at the point +:math:`\hat{r}`, i.e. at the spherical angles :math:`(\theta, \phi)` that determine the +orientation of the unit vector :math:`\hat{r} = \boldsymbol{r}/r`. + +Radial basis are represented as different child class of +:py:class:`rascaline.utils.RadialBasisBase`: :py:class:`rascaline.utils.GtoBasis`, +:py:class:`rascaline.utils.MonomialBasis`, and +:py:class:`rascaline.utils.SphericalBesselBasis` are provided, and you can implement +your own by defining a new class. + +.. autoclass:: rascaline.utils.RadialBasisBase + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.GtoBasis + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.MonomialBasis + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.SphericalBesselBasis + :members: + :show-inheritance: +""" + +from abc import ABC, abstractmethod +from typing import Union + +import numpy as np + + +try: + import scipy.integrate + import scipy.optimize + import scipy.special + + HAS_SCIPY = True +except ImportError: + HAS_SCIPY = False + + +class RadialBasisBase(ABC): + r""" + Base class to define radial basis and their evaluation. + + The class provides methods to evaluate the radial basis :math:`R_{nl}(r)` as well as + its (numerical) derivative with respect to positions :math:`r`. + + :parameter integration_radius: Value up to which the radial integral should be + performed. The usual value is :math:`\infty`. + """ + + def __init__(self, integration_radius: float): + self.integration_radius = integration_radius + + @abstractmethod + def compute( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + """Compute the ``n``/``l`` radial basis at all given ``integrand_positions`` + + :param n: radial channel + :param ell: angular channel + :param integrand_positions: positions to evaluate the radial basis + :returns: evaluated radial basis + """ + + def compute_derivative( + self, n: float, ell: float, integrand_positions: np.ndarray + ) -> np.ndarray: + """Compute the derivative of the ``n``/``l`` radial basis at all given + ``integrand_positions`` + + This is used for radial integrals with delta-like atomic densities. If not + defined in a child class, a numerical derivative based on finite differences of + ``integrand_positions`` will be used instead. + + :param n: radial channel + :param ell: angular channel + :param integrand_positions: positions to evaluate the radial basis + :returns: evaluated derivative of the radial basis + """ + displacement = 1e-6 + mean_abs_positions = np.abs(integrand_positions).mean() + + if mean_abs_positions < 1.0: + raise ValueError( + "Numerically derivative of the radial integral can not be performed " + "since positions are too small. Mean of the absolute positions is " + f"{mean_abs_positions:.1e} but should be at least 1." + ) + + radial_basis_pos = self.compute(n, ell, integrand_positions + displacement / 2) + radial_basis_neg = self.compute(n, ell, integrand_positions - displacement / 2) + + return (radial_basis_pos - radial_basis_neg) / displacement + + def compute_gram_matrix( + self, + max_radial: int, + max_angular: int, + ) -> np.ndarray: + """Gram matrix of the current basis. + + :parameter max_radial: number of angular components + :parameter max_angular: number of radial components + :returns: orthonormalization matrix of shape + ``(max_angular + 1, max_radial, max_radial)`` + """ + + if not HAS_SCIPY: + raise ValueError("Orthonormalization requires scipy!") + + # Gram matrix (also called overlap matrix or inner product matrix) + gram_matrix = np.zeros((max_angular + 1, max_radial, max_radial)) + + def integrand( + integrand_positions: np.ndarray, + n1: int, + n2: int, + ell: int, + ) -> np.ndarray: + return ( + integrand_positions**2 + * self.compute(n1, ell, integrand_positions) + * self.compute(n2, ell, integrand_positions) + ) + + for ell in range(max_angular + 1): + for n1 in range(max_radial): + for n2 in range(max_radial): + gram_matrix[ell, n1, n2] = scipy.integrate.quad( + func=integrand, + a=0, + b=self.integration_radius, + args=(n1, n2, ell), + )[0] + + return gram_matrix + + def compute_orthonormalization_matrix( + self, + max_radial: int, + max_angular: int, + ) -> np.ndarray: + """Compute orthonormalization matrix + + :parameter max_radial: number of angular components + :parameter max_angular: number of radial components + :returns: orthonormalization matrix of shape (max_angular + 1, max_radial, + max_radial) + """ + + gram_matrix = self.compute_gram_matrix(max_radial, max_angular) + + # Get the normalization constants from the diagonal entries + normalizations = np.zeros((max_angular + 1, max_radial)) + + for ell in range(max_angular + 1): + for n in range(max_radial): + normalizations[ell, n] = 1 / np.sqrt(gram_matrix[ell, n, n]) + + # Rescale orthonormalization matrix to be defined + # in terms of the normalized (but not yet orthonormalized) + # basis functions + gram_matrix[ell, n, :] *= normalizations[ell, n] + gram_matrix[ell, :, n] *= normalizations[ell, n] + + orthonormalization_matrix = np.zeros_like(gram_matrix) + for ell in range(max_angular + 1): + eigvals, eigvecs = np.linalg.eigh(gram_matrix[ell]) + orthonormalization_matrix[ell] = ( + eigvecs @ np.diag(np.sqrt(1.0 / eigvals)) @ eigvecs.T + ) + + # Rescale the orthonormalization matrix so that it + # works with respect to the primitive (not yet normalized) + # radial basis functions + for ell in range(max_angular + 1): + for n in range(max_radial): + orthonormalization_matrix[ell, :, n] *= normalizations[ell, n] + + return orthonormalization_matrix + + +class GtoBasis(RadialBasisBase): + r"""Primitive (not normalized nor orthonormalized) GTO radial basis. + + It is defined as + + .. math:: + + R_{nl}(r) = R_n(r) = r^n e^{-\frac{r^2}{2\sigma_n^2}}, + + where :math:`\sigma_n = \sqrt{n} r_\mathrm{cut}/n_\mathrm{max}` with + :math:`r_\mathrm{cut}` being the ``cutoff`` and :math:`n_\mathrm{max}` the maximal + number of radial components. + + :parameter cutoff: spherical cutoff for the radial basis + :parameter max_radial: number of radial components + """ + + def __init__(self, cutoff, max_radial): + # choosing infinity leads to problems when calculating the radial integral with + # `quad`! + super().__init__(integration_radius=5 * cutoff) + self.max_radial = max_radial + self.cutoff = cutoff + self.sigmas = np.ones(self.max_radial, dtype=float) + + for n in range(1, self.max_radial): + self.sigmas[n] = np.sqrt(n) + self.sigmas *= self.cutoff / self.max_radial + + def compute( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + return integrand_positions**n * np.exp( + -0.5 * (integrand_positions / self.sigmas[n]) ** 2 + ) + + def compute_derivative( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + return n / integrand_positions * self.compute( + n, ell, integrand_positions + ) - integrand_positions / self.sigmas[n] ** 2 * self.compute( + n, ell, integrand_positions + ) + + +class MonomialBasis(RadialBasisBase): + r"""Monomial basis. + + Basis is consisting of functions + + .. math:: + R_{nl}(r) = r^{l+2n}, + + where :math:`n` runs from :math:`0,1,...,n_\mathrm{max}-1`. These capture precisely + the radial dependence if we compute the Taylor expansion of a generic function + defined in 3D space. + + :parameter cutoff: spherical cutoff for the radial basis + """ + + def __init__(self, cutoff): + super().__init__(integration_radius=cutoff) + + def compute( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + return integrand_positions ** (ell + 2 * n) + + def compute_derivative( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + return (ell + 2 * n) * integrand_positions ** (ell + 2 * n - 1) + + +def _spherical_jn_swapped(z, n): + """spherical_jn with swapped arguments for usage in `fsolve`.""" + return scipy.special.spherical_jn(n=n, z=z) + + +class SphericalBesselBasis(RadialBasisBase): + """Spherical Bessel functions used in the Laplacian eigenstate (LE) basis. + + :parameter cutoff: spherical cutoff for the radial basis + :parameter max_radial: number of angular components + :parameter max_angular: number of radial components + """ + + def __init__(self, cutoff, max_radial, max_angular): + if not HAS_SCIPY: + raise ValueError("SphericalBesselBasis requires scipy!") + + super().__init__(integration_radius=cutoff) + + self.max_radial = max_radial + self.max_angular = max_angular + self.roots = np.zeros([max_angular + 1, self.max_radial]) + + # Define target function and the estimated location of roots obtained from the + # asymptotic expansion of the spherical Bessel functions for large arguments x + for ell in range(self.max_angular + 1): + roots_guesses = np.pi * (np.arange(1, self.max_radial + 1) + ell / 2) + # Compute roots from initial guess using Newton method + for n, root_guess in enumerate(roots_guesses): + self.roots[ell, n] = scipy.optimize.fsolve( + func=_spherical_jn_swapped, x0=root_guess, args=(ell,) + )[0] + + def compute( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + return scipy.special.spherical_jn( + ell, + integrand_positions * self.roots[ell, n] / self.integration_radius, + ) + + def compute_derivative( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + return ( + self.roots[ell, n] + / self.integration_radius + * scipy.special.spherical_jn( + ell, + integrand_positions * self.roots[ell, n] / self.integration_radius, + derivative=True, + ) + ) diff --git a/python/rascaline/rascaline/utils/splines/splines.py b/python/rascaline/rascaline/utils/splines/splines.py new file mode 100644 index 000000000..42ad3c968 --- /dev/null +++ b/python/rascaline/rascaline/utils/splines/splines.py @@ -0,0 +1,849 @@ +""" +.. _python-splined-radial-integral: + +Splined radial integrals +======================== + +Classes for generating splines which can be used as tabulated radial integrals in the +various SOAP and LODE calculators. + +All classes are based on :py:class:`rascaline.utils.RadialIntegralSplinerBase`. We +provides several ways to compute a radial integral: you may chose and initialize a pre +defined atomic density and radial basis and provide them to +:py:class:`rascaline.utils.SoapSpliner` or :py:class:`rascaline.utils.LodeSpliner`. Both +classes require `scipy`_ to be installed in order to perform the numerical integrals. + +Alternatively, you can also explicitly provide functions for the radial integral and its +derivative and passing them to :py:class:`rascaline.utils.RadialIntegralFromFunction`. + +.. autoclass:: rascaline.utils.RadialIntegralSplinerBase + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.SoapSpliner + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.LodeSpliner + :members: + :show-inheritance: + +.. autoclass:: rascaline.utils.RadialIntegralFromFunction + :members: + :show-inheritance: + + +.. _`scipy`: https://scipy.org +""" + +from abc import ABC, abstractmethod +from typing import Callable, Dict, Optional, Union + +import numpy as np + + +try: + from scipy.integrate import quad, quad_vec + from scipy.special import spherical_in, spherical_jn + + HAS_SCIPY = True +except ImportError: + HAS_SCIPY = False + +from .atomic_density import AtomicDensityBase, DeltaDensity, GaussianDensity +from .radial_basis import RadialBasisBase + + +class RadialIntegralSplinerBase(ABC): + """Base class for splining arbitrary radial integrals. + + If :py:meth:`RadialIntegralSplinerBase.radial_integral_derivative` is not + implemented in a child class it will computed based on finite differences. + + :parameter max_angular: number of radial components + :parameter max_radial: number of angular components + :parameter spline_cutoff: cutoff radius for the spline interpolation. This is also + the maximal value that can be interpolated. + :parameter basis: Provide a :class:`RadialBasisBase` instance to orthonormalize the + radial integral. + :parameter accuracy: accuracy of the numerical integration and the splining. + Accuracy is reached when either the mean absolute error or the mean relative + error gets below the ``accuracy`` threshold. + """ + + def __init__( + self, + max_radial: int, + max_angular: int, + spline_cutoff: float, + basis: Optional[RadialBasisBase], + accuracy: float, + ): + self.max_radial = max_radial + self.max_angular = max_angular + self.spline_cutoff = spline_cutoff + self.basis = basis + self.accuracy = accuracy + + def compute( + self, + n_spline_points: Optional[int] = None, + ) -> Dict: + """Compute the spline for rascaline's tabulated radial integrals. + + :parameter n_spline_points: Use fixed number of spline points instead of find + the number based on the provided ``accuracy``. + :returns dict: dictionary for the input as the ``radial_basis`` parameter of a + rascaline calculator. + """ + + if self.basis is not None: + orthonormalization_matrix = self.basis.compute_orthonormalization_matrix( + self.max_radial, self.max_angular + ) + else: + orthonormalization_matrix = None + + def value_evaluator_3D(positions): + return self._value_evaluator_3D( + positions, orthonormalization_matrix, derivative=False + ) + + def derivative_evaluator_3D(positions): + return self._value_evaluator_3D( + positions, orthonormalization_matrix, derivative=True + ) + + if n_spline_points is not None: + positions = np.linspace(0, self.spline_cutoff, n_spline_points) + values = value_evaluator_3D(positions) + derivatives = derivative_evaluator_3D(positions) + else: + dynamic_spliner = DynamicSpliner( + 0, + self.spline_cutoff, + value_evaluator_3D, + derivative_evaluator_3D, + self.accuracy, + ) + positions, values, derivatives = dynamic_spliner.spline() + + # Convert positions, values, derivatives into the appropriate json formats: + spline_points = [] + for position, value, derivative in zip(positions, values, derivatives): + spline_points.append( + { + "position": position, + "values": { + "v": 1, + "dim": value.shape, + "data": value.flatten().tolist(), + }, + "derivatives": { + "v": 1, + "dim": derivative.shape, + "data": derivative.flatten().tolist(), + }, + } + ) + + parameters = {"points": spline_points} + + center_contribution = self.center_contribution + if center_contribution is not None: + if self.basis is not None: + # consider only `l=0` component of the `orthonormalization_matrix` + parameters["center_contribution"] = list( + orthonormalization_matrix[0] @ center_contribution + ) + else: + parameters["center_contribution"] = center_contribution + + return {"TabulatedRadialIntegral": parameters} + + @abstractmethod + def radial_integral(self, n: int, ell: int, positions: np.ndarray) -> np.ndarray: + """evaluate the radial integral""" + ... + + @property + def center_contribution(self) -> Union[None, np.ndarray]: + r"""Pre-computed value for the contribution of the central atom. + + Required for LODE calculations. The central atom contribution will be + orthonormalized in the same way as the radial integral. + """ + + return None + + def radial_integral_derivative( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + """evaluate the derivative of the radial integral""" + displacement = 1e-6 + mean_abs_positions = np.mean(np.abs(positions)) + + if mean_abs_positions < 1.0: + raise ValueError( + "Numerically derivative of the radial integral can not be performed " + "since positions are too small. Mean of the absolute positions is " + f"{mean_abs_positions:.1e} but should be at least 1." + ) + + radial_integral_pos = self.radial_integral(n, ell, positions + displacement / 2) + radial_integral_neg = self.radial_integral(n, ell, positions - displacement / 2) + + return (radial_integral_pos - radial_integral_neg) / displacement + + def _value_evaluator_3D( + self, + positions: np.ndarray, + orthonormalization_matrix: Optional[np.ndarray], + derivative: bool, + ): + values = np.zeros([len(positions), self.max_angular + 1, self.max_radial]) + for ell in range(self.max_angular + 1): + for n in range(self.max_radial): + if derivative: + values[:, ell, n] = self.radial_integral_derivative( + n, ell, positions + ) + else: + values[:, ell, n] = self.radial_integral(n, ell, positions) + + if orthonormalization_matrix is not None: + # For each l channel we do a dot product of the orthonormalization_matrix of + # shape (n, n) with the values which should have the shape (n, n_positions). + # To achieve the correct broadcasting we have to transpose twice. + for ell in range(self.max_angular + 1): + values[:, ell, :] = ( + orthonormalization_matrix[ell] @ values[:, ell, :].T + ).T + + return values + + +class DynamicSpliner: + def __init__( + self, + start: float, + stop: float, + values_fn: Callable[[np.ndarray], np.ndarray], + derivatives_fn: Callable[[np.ndarray], np.ndarray], + accuracy: float = 1e-8, + ) -> None: + """Dynamic spline generator. + + This class can be used to spline any set of functions defined within the + start-stop interval. Cubic Hermite splines + (https://en.wikipedia.org/wiki/Cubic_Hermite_spline) are used. The same spline + points will be used for all functions, and more will be added until either the + relative error or the absolute error fall below the requested accuracy on + average across all functions. The functions are specified via values_fn and + derivatives_fn. These must be able to take a numpy 1D array of positions as + their input, and they must output a numpy array where the first dimension + corresponds to the input positions, while other dimensions are arbitrary and can + correspond to any way in which the target functions can be classified. The + splines can be obtained via the spline method. + """ + + self.start = start + self.stop = stop + self.values_fn = values_fn + self.derivatives_fn = derivatives_fn + self.requested_accuracy = accuracy + + # initialize spline with 11 points + positions = np.linspace(start, stop, 11) + self.spline_positions = positions + self.spline_values = values_fn(positions) + self.spline_derivatives = derivatives_fn(positions) + + self.number_of_custom_axes = len(self.spline_values.shape) - 1 + + def spline(self): + """Calculates and outputs the splines. + + The outputs of this function are, respectively: + + - A numpy 1D array containing the spline positions. These are equally spaced in + the start-stop interval. + - A numpy ndarray containing the values of the splined functions at the spline + positions. The first dimension corresponds to the spline positions, while all + subsequent dimensions are consistent with the values_fn and + `get_function_derivative` provided during initialization of the class. + - A numpy ndarray containing the derivatives of the splined functions at the + spline positions, with the same structure as that of the ndarray of values. + """ + + while True: + n_intermediate_positions = len(self.spline_positions) - 1 + + if n_intermediate_positions >= 50000: + raise ValueError( + "Maximum number of spline points reached. \ + There might be a problem with the functions to be splined" + ) + + half_step = (self.spline_positions[1] - self.spline_positions[0]) / 2 + intermediate_positions = np.linspace( + self.start + half_step, self.stop - half_step, n_intermediate_positions + ) + + estimated_values = self._compute_from_spline(intermediate_positions) + new_values = self.values_fn(intermediate_positions) + + mean_absolute_error = np.mean(np.abs(estimated_values - new_values)) + with np.errstate(divide="ignore"): # Ignore divide-by-zero warnings + mean_relative_error = np.mean( + np.abs((estimated_values - new_values) / new_values) + ) + + if ( + mean_absolute_error < self.requested_accuracy + or mean_relative_error < self.requested_accuracy + ): + break + + new_derivatives = self.derivatives_fn(intermediate_positions) + + concatenated_positions = np.concatenate( + [self.spline_positions, intermediate_positions], axis=0 + ) + concatenated_values = np.concatenate( + [self.spline_values, new_values], axis=0 + ) + concatenated_derivatives = np.concatenate( + [self.spline_derivatives, new_derivatives], axis=0 + ) + + sort_indices = np.argsort(concatenated_positions, axis=0) + + self.spline_positions = concatenated_positions[sort_indices] + self.spline_values = concatenated_values[sort_indices] + self.spline_derivatives = concatenated_derivatives[sort_indices] + + return self.spline_positions, self.spline_values, self.spline_derivatives + + def _compute_from_spline(self, positions): + x = positions + delta_x = self.spline_positions[1] - self.spline_positions[0] + n = (np.floor(x / delta_x)).astype(np.int32) + + t = (x - n * delta_x) / delta_x + t_2 = t**2 + t_3 = t**3 + + h00 = 2.0 * t_3 - 3.0 * t_2 + 1.0 + h10 = t_3 - 2.0 * t_2 + t + h01 = -2.0 * t_3 + 3.0 * t_2 + h11 = t_3 - t_2 + + p_k = self.spline_values[n] + p_k_1 = self.spline_values[n + 1] + + m_k = self.spline_derivatives[n] + m_k_1 = self.spline_derivatives[n + 1] + + new_shape = (-1,) + (1,) * self.number_of_custom_axes + h00 = h00.reshape(new_shape) + h10 = h10.reshape(new_shape) + h01 = h01.reshape(new_shape) + h11 = h11.reshape(new_shape) + + interpolated_values = ( + h00 * p_k + h10 * delta_x * m_k + h01 * p_k_1 + h11 * delta_x * m_k_1 + ) + + return interpolated_values + + +class RadialIntegralFromFunction(RadialIntegralSplinerBase): + r"""Compute radial integral spline points based on a provided function. + + :parameter radial_integral: Function to compute the radial integral. Function must + take ``n``, ``l``, and ``positions`` as inputs, where ``n`` and ``l`` are + integers and ``positions`` is a numpy 1-D array that contains the spline points + at which the radial integral will be evaluated. The function must return a numpy + 1-D array containing the values of the radial integral. + :parameter spline_cutoff: cutoff radius for the spline interpolation. This is also + the maximal value that can be interpolated. + :parameter max_radial: number of angular components + :parameter max_angular: number of radial components + :parameter radial_integral_derivative: The derivative of the radial integral taking + the same parameters as ``radial_integral``. If it is :py:obj:`None` (default), + finite differences are used to calculate the derivative of the radial integral. + It is recommended to provide this parameter if possible. Derivatives from finite + differences can cause problems when evaluating at the edges of the domain (i.e., + at ``0`` and ``spline_cutoff``) because the function might not be defined + outside of the domain. + :parameter accuracy: accuracy of the numerical integration and the splining. + Accuracy is reached when either the mean absolute error or the mean relative + error gets below the ``accuracy`` threshold. + :parameter center_contribution: Contribution of the central atom required for LODE + calculations. The ``center_contribution`` is defined as + + .. math:: + c_n = \sqrt{4π}\int_0^\infty dr r^2 R_n(r) g(r) + + where :math:`g(r)` is the radially symmetric density function, `R_n(r)` the + radial basis function and :math:`n` the current radial channel. This should be + pre-computed and provided as a separate parameter. + + Example + ------- + First define a ``radial_integral`` function + + >>> def radial_integral(n, ell, r): + ... return np.sin(r) + ... + + and provide this as input to the spline generator + + >>> spliner = RadialIntegralFromFunction( + ... radial_integral=radial_integral, + ... max_radial=12, + ... max_angular=9, + ... spline_cutoff=8.0, + ... ) + + Finally, we can use the ``spliner`` directly in the ``radial_integral`` section of a + calculator + + >>> from rascaline import SoapPowerSpectrum + >>> calculator = SoapPowerSpectrum( + ... cutoff=8.0, + ... max_radial=12, + ... max_angular=9, + ... center_atom_weight=1.0, + ... radial_basis=spliner.compute(), + ... atomic_gaussian_width=1.0, # ignored + ... cutoff_function={"Step": {}}, + ... ) + + The ``atomic_gaussian_width`` parameter is required by the calculator but will be + will be ignored during the feature computation. + + A more in depth example using a "rectangular" Laplacian eigenstate basis is provided + in the :ref:`example section`. + """ + + def __init__( + self, + radial_integral: Callable[[int, int, np.ndarray], np.ndarray], + spline_cutoff: float, + max_radial: int, + max_angular: int, + radial_integral_derivative: Optional[ + Callable[[int, int, np.ndarray], np.ndarray] + ] = None, + center_contribution: Optional[np.ndarray] = None, + accuracy: float = 1e-8, + ): + self.radial_integral_function = radial_integral + self.radial_integral_derivative_function = radial_integral_derivative + + if center_contribution is not None and len(center_contribution) != max_radial: + raise ValueError( + f"center contribution has {len(center_contribution)} entries but " + f"should be the same as max_radial ({max_radial})" + ) + + self._center_contribution = center_contribution + + super().__init__( + max_radial=max_radial, + max_angular=max_angular, + spline_cutoff=spline_cutoff, + basis=None, # do no orthonormalize the radial integral + accuracy=accuracy, + ) + + def radial_integral(self, n: int, ell: int, positions: np.ndarray) -> np.ndarray: + return self.radial_integral_function(n, ell, positions) + + @property + def center_contribution(self) -> Union[None, np.ndarray]: + return self._center_contribution + + def radial_integral_derivative( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + if self.radial_integral_derivative_function is None: + return super().radial_integral_derivative(n, ell, positions) + else: + return self.radial_integral_derivative_function(n, ell, positions) + + +class SoapSpliner(RadialIntegralSplinerBase): + """Compute radial integral spline points for real space calculators. + + Use only in combination with a real space calculators like + :class:`rascaline.SphericalExpansion` or :class:`rascaline.SoapPowerSpectrum`. For + k-space spherical expansions use :class:`LodeSpliner`. + + If ``density`` is either :class:`rascaline.utils.DeltaDensity` or + :class:`rascaline.utils.GaussianDensity` the radial integral will be partly solved + analytical. These simpler expressions result in a faster and more stable evaluation. + + :parameter cutoff: spherical cutoff for the radial basis + :parameter max_radial: number of angular components + :parameter max_angular: number of radial components + :parameter basis: definition of the radial basis + :parameter density: definition of the atomic density + :parameter accuracy: accuracy of the numerical integration and the splining. + Accuracy is reached when either the mean absolute error or the mean relative + error gets below the ``accuracy`` threshold. + :raise ValueError: if `scipy`_ is not available + + Example + ------- + + First import the necessary classed and define hyper parameters for the spherical + expansions. + + >>> from rascaline import SphericalExpansion + >>> from rascaline.utils import GaussianDensity, GtoBasis + + >>> cutoff = 2 + >>> max_radial = 6 + >>> max_angular = 4 + >>> atomic_gaussian_width = 1.0 + + Next we initialize our radial basis and the density + + >>> basis = GtoBasis(cutoff=cutoff, max_radial=max_radial) + >>> density = GaussianDensity(atomic_gaussian_width=atomic_gaussian_width) + + And finally the actual spliner instance + + >>> spliner = SoapSpliner( + ... cutoff=cutoff, + ... max_radial=max_radial, + ... max_angular=max_angular, + ... basis=basis, + ... density=density, + ... accuracy=1e-3, + ... ) + + Above we reduced ``accuracy`` from the default value of ``1e-8`` to ``1e-3`` to + speed up calculations. + + As for all spliner classes you can use the output + :meth:`RadialIntegralSplinerBase.compute` method directly as the ``radial_basis`` + parameter. + + >>> calculator = SphericalExpansion( + ... cutoff=cutoff, + ... max_radial=max_radial, + ... max_angular=max_angular, + ... center_atom_weight=1.0, + ... atomic_gaussian_width=atomic_gaussian_width, + ... radial_basis=spliner.compute(), + ... cutoff_function={"Step": {}}, + ... ) + + You can now use ``calculator`` to obtain the spherical expansion coefficients of + your systems. Note that the the spliner based used here will produce the same + coefficients as if ``radial_basis={"Gto": {}}`` would be used. + + .. seealso:: + :class:`LodeSpliner` for a spliner class that works with + :class:`rascaline.LodeSphericalExpansion` + """ + + def __init__( + self, + cutoff: float, + max_radial: int, + max_angular: int, + basis: RadialBasisBase, + density: AtomicDensityBase, + accuracy: float = 1e-8, + ): + if not HAS_SCIPY: + raise ValueError("Spliner class requires scipy!") + + self.density = density + + super().__init__( + max_radial=max_radial, + max_angular=max_angular, + spline_cutoff=cutoff, + basis=basis, + accuracy=accuracy, + ) + + def radial_integral(self, n: int, ell: int, positions: np.ndarray) -> np.ndarray: + if type(self.density) is DeltaDensity: + return self._radial_integral_delta(n, ell, positions) + elif type(self.density) is GaussianDensity: + return self._radial_integral_gaussian(n, ell, positions) + else: + return self._radial_integral_custom(n, ell, positions) + + def radial_integral_derivative( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + if type(self.density) is DeltaDensity: + return self._radial_integral_delta_derivative(n, ell, positions) + elif type(self.density) is GaussianDensity: + return self._radial_integral_gaussian_derivative(n, ell, positions) + else: + return self._radial_integral_custom_derivative(n, ell, positions) + + def _radial_integral_delta( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + return self.basis.compute(n, ell, positions) + + def _radial_integral_delta_derivative( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + return self.basis.compute_derivative(n, ell, positions) + + def _radial_integral_gaussian( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + atomic_gaussian_width_sq = self.density.atomic_gaussian_width**2 + + prefactor = ( + (4 * np.pi) + / (np.pi * atomic_gaussian_width_sq) ** (3 / 4) + * np.exp(-0.5 * positions**2 / atomic_gaussian_width_sq) + ) + + def integrand( + integrand_position: float, n: int, ell: int, positions: np.array + ) -> np.ndarray: + return ( + integrand_position**2 + * self.basis.compute(n, ell, integrand_position) + * np.exp(-0.5 * integrand_position**2 / atomic_gaussian_width_sq) + * spherical_in( + ell, + integrand_position * positions / atomic_gaussian_width_sq, + ) + ) + + return ( + prefactor + * quad_vec( + f=lambda x: integrand(x, n, ell, positions), + a=0, + b=self.basis.integration_radius, + )[0] + ) + + def _radial_integral_gaussian_derivative( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + atomic_gaussian_width_sq = self.density.atomic_gaussian_width**2 + + prefactor = ( + (4 * np.pi) + / (np.pi * atomic_gaussian_width_sq) ** (3 / 4) + * np.exp(-0.5 * positions**2 / atomic_gaussian_width_sq) + ) + + def integrand( + integrand_position: float, n: int, ell: int, positions: np.array + ) -> np.ndarray: + return ( + integrand_position**3 + * self.basis.compute(n, ell, integrand_position) + * np.exp(-(integrand_position**2) / (2 * atomic_gaussian_width_sq)) + * spherical_in( + ell, + integrand_position * positions / atomic_gaussian_width_sq, + derivative=True, + ) + ) + + return atomic_gaussian_width_sq**-1 * ( + prefactor + * quad_vec( + f=lambda x: integrand(x, n, ell, positions), + a=0, + b=self.basis.integration_radius, + )[0] + - positions * self._radial_integral_gaussian(n, ell, positions) + ) + + def _radial_integral_custom( + self, n: int, ell: int, positions: np.ndarray, derivative: bool + ) -> np.ndarray: + raise NotImplementedError( + "Radial integral with custom atomic densities is not implemented yet!" + ) + + def _radial_integral_custom_derivative( + self, n: int, ell: int, positions: np.ndarray, derivative: bool + ) -> np.ndarray: + raise NotImplementedError( + "Radial integral with custom atomic densities is not implemented yet!" + ) + + +class LodeSpliner(RadialIntegralSplinerBase): + r"""Compute radial integral spline points for k-space calculators. + + Use only in combination with a k-space/Fourier-space calculators like + :class:`rascaline.LodeSphericalExpansion`. For real space spherical expansions use + :class:`SoapSpliner`. + + :parameter k_cutoff: spherical reciprocal cutoff + :parameter max_radial: number of angular components + :parameter max_angular: number of radial components + :parameter basis: definition of the radial basis + :parameter density: definition of the atomic density + :parameter accuracy: accuracy of the numerical integration and the splining. + Accuracy is reached when either the mean absolute error or the mean relative + error gets below the ``accuracy`` threshold. + :raise ValueError: if `scipy`_ is not available + + Example + ------- + + First import the necessary classed and define hyper parameters for the spherical + expansions. + + >>> from rascaline import LodeSphericalExpansion + >>> from rascaline.utils import GaussianDensity, GtoBasis + + Note that ``cutoff`` defined below denotes the maximal distance for the projection + of the density. In contrast to SOAP, LODE also takes atoms outside of this + ``cutoff`` into account for the density. + + >>> cutoff = 2 + >>> max_radial = 6 + >>> max_angular = 4 + >>> atomic_gaussian_width = 1.0 + + :math:`1.2 \, \pi \, \sigma` where :math:`\sigma` is the ``atomic_gaussian_width`` + which is a reasonable value for most systems. + + >>> k_cutoff = 1.2 * np.pi / atomic_gaussian_width + + Next we initialize our radial basis and the density + + >>> basis = GtoBasis(cutoff=cutoff, max_radial=max_radial) + >>> density = GaussianDensity(atomic_gaussian_width=atomic_gaussian_width) + + And finally the actual spliner instance + + >>> spliner = LodeSpliner( + ... k_cutoff=k_cutoff, + ... max_radial=max_radial, + ... max_angular=max_angular, + ... basis=basis, + ... density=density, + ... ) + + As for all spliner classes you can use the output + :meth:`RadialIntegralSplinerBase.compute` method directly as the ``radial_basis`` + parameter. + + >>> calculator = LodeSphericalExpansion( + ... cutoff=cutoff, + ... max_radial=max_radial, + ... max_angular=max_angular, + ... center_atom_weight=1.0, + ... atomic_gaussian_width=atomic_gaussian_width, + ... potential_exponent=1, + ... radial_basis=spliner.compute(), + ... ) + + You can now use ``calculator`` to obtain the spherical expansion coefficients of + your systems. Note that the the spliner based used here will produce the same + coefficients as if ``radial_basis={"Gto": {}}`` would be used. + + .. seealso:: + :class:`SoapSpliner` for a spliner class that works with + :class:`rascaline.SphericalExpansion` + """ + + def __init__( + self, + k_cutoff: float, + max_radial: int, + max_angular: int, + basis: RadialBasisBase, + density: AtomicDensityBase, + accuracy: float = 1e-8, + ): + if not HAS_SCIPY: + raise ValueError("Spliner class requires scipy!") + + self.density = density + + super().__init__( + max_radial=max_radial, + max_angular=max_angular, + basis=basis, + spline_cutoff=k_cutoff, # use k_cutoff here because we spline in k-space + accuracy=accuracy, + ) + + def radial_integral(self, n: int, ell: int, positions: np.ndarray) -> np.ndarray: + def integrand( + integrand_position: float, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + return ( + integrand_position**2 + * self.basis.compute(n, ell, integrand_position) + * spherical_jn(ell, integrand_position * positions) + ) + + return quad_vec( + f=lambda x: integrand(x, n, ell, positions), + a=0, + b=self.basis.integration_radius, + )[0] + + def radial_integral_derivative( + self, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + def integrand( + integrand_position: float, n: int, ell: int, positions: np.ndarray + ) -> np.ndarray: + return ( + integrand_position**3 + * self.basis.compute(n, ell, integrand_position) + * spherical_jn(ell, integrand_position * positions, derivative=True) + ) + + return quad_vec( + f=lambda x: integrand(x, n, ell, positions), + a=0, + b=self.basis.integration_radius, + )[0] + + @property + def center_contribution(self) -> np.ndarray: + if type(self.density) is DeltaDensity: + center_contrib = self._center_contribution_delta + else: + center_contrib = self._center_contribution_custom + + return [np.sqrt(4 * np.pi) * center_contrib(n) for n in range(self.max_radial)] + + def _center_contribution_delta(self, n: int): + raise NotImplementedError( + "center contribution for delta distributions is not implemented yet." + ) + + def _center_contribution_custom(self, n: int): + def integrand(integrand_position: float, n: int) -> np.ndarray: + return ( + integrand_position**2 + * self.basis.compute(n, 0, integrand_position) + * self.density.compute(integrand_position) + ) + + return quad( + func=integrand, + a=0, + b=self.basis.integration_radius, + args=(n), + )[0] diff --git a/python/rascaline/rascaline/version.py b/python/rascaline/rascaline/version.py index a7248880b..e69652e02 100644 --- a/python/rascaline/rascaline/version.py +++ b/python/rascaline/rascaline/version.py @@ -1,12 +1,4 @@ -import sys +import importlib.metadata -if (sys.version_info.major >= 3) and (sys.version_info.minor >= 8): - import importlib.metadata - - __version__ = importlib.metadata.version("rascaline") - -else: - from pkg_resources import get_distribution - - __version__ = get_distribution("rascaline").version +__version__ = importlib.metadata.version("rascaline") diff --git a/python/rascaline/tests/systems/pyscf.py b/python/rascaline/tests/systems/pyscf.py new file mode 100644 index 000000000..35d39d99e --- /dev/null +++ b/python/rascaline/tests/systems/pyscf.py @@ -0,0 +1,71 @@ +import numpy as np +import pytest + +from rascaline.systems import PyscfSystem + + +pyscf = pytest.importorskip("pyscf") + + +@pytest.fixture +def system(): + atoms = pyscf.M( + atom="C 0 0 0; O 0 0 1.4; O 0 0 -1.6", + ) + # atoms.pbc = [False, False, False] + return PyscfSystem(atoms) + + +def test_system_implementation(system): + assert system.size() == 3 + assert np.all(system.species() == [6, 8, 8]) + + positions = np.array( + [ + (0, 0, 0), + (0, 0, 1.4), + (0, 0, -1.6), + ] + ) + print(system.positions(), positions) + assert np.allclose(system.positions(), positions, rtol=1e-14) + assert np.all(system.cell() == [[0, 0, 0], [0, 0, 0], [0, 0, 0]]) + + +def test_pbc_data(): + import pyscf.pbc + + atoms = pyscf.pbc.gto.Cell( + atom="H 0 0 0; H 1 1 1", + a=np.array([[2, 0, 0], [0, 2, 1], [0, 0, 2]], dtype=float), + ).build() + ras_sys = PyscfSystem(atoms) + assert np.allclose( + ras_sys.positions(), + np.array([[0, 0, 0], [1, 1, 1]], dtype=float), + rtol=1e-14, + ) + + +def test_explicit_units(): + import pyscf.pbc + + cell = np.array([[2, 0, 0], [0, 2, 1], [0, 0, 2]], dtype=float) + + at1 = pyscf.pbc.gto.Cell( + atom="H 0 0 0; H 1 1 1", + a=cell, + unit="Angstrom", + ).build() + at2 = pyscf.pbc.gto.Cell( + atom=[("H", at1.atom_coord(0)), ("H", at1.atom_coord(1))], + a=cell / pyscf.data.nist.BOHR, + unit="Bohr", + ).build() + at1 = PyscfSystem(at1) + at2 = PyscfSystem(at2) + + assert np.allclose(at1.positions(), at2.positions()) + assert np.allclose(at1.positions(), np.array([[0, 0, 0], [1, 1, 1]], dtype=float)) + assert np.allclose(at1.cell(), at2.cell()) + assert np.allclose(at1.cell(), cell) diff --git a/python/rascaline/tests/utils/clebsch_gordan.py b/python/rascaline/tests/utils/clebsch_gordan.py new file mode 100644 index 000000000..38b7fab32 --- /dev/null +++ b/python/rascaline/tests/utils/clebsch_gordan.py @@ -0,0 +1,540 @@ +# -*- coding: utf-8 -*- +import os +from typing import List + +import metatensor +import numpy as np +import pytest +from metatensor import Labels, TensorBlock, TensorMap + +import rascaline +from rascaline.utils import PowerSpectrum +from rascaline.utils.clebsch_gordan._cg_cache import ClebschGordanReal +from rascaline.utils.clebsch_gordan.clebsch_gordan import ( + _correlate_density, + _standardize_keys, + correlate_density, + correlate_density_metadata, +) + + +# Try to import some modules +ase = pytest.importorskip("ase") +import ase.io # noqa: E402 + + +try: + import metatensor.operations + + HAS_METATENSOR_OPERATIONS = True +except ImportError: + HAS_METATENSOR_OPERATIONS = False +try: + import sympy # noqa: F401 + + HAS_SYMPY = True +except ImportError: + HAS_SYMPY = False + +if HAS_SYMPY: + from .rotations import WignerDReal, transform_frame_o3, transform_frame_so3 + + +DATA_ROOT = os.path.join(os.path.dirname(__file__), "data") + +SPHEX_HYPERS = { + "cutoff": 3.0, # Angstrom + "max_radial": 6, # Exclusive + "max_angular": 4, # Inclusive + "atomic_gaussian_width": 0.2, + "radial_basis": {"Gto": {}}, + "cutoff_function": {"ShiftedCosine": {"width": 0.5}}, + "center_atom_weight": 1.0, +} + +SPHEX_HYPERS_SMALL = { + "cutoff": 3.0, # Angstrom + "max_radial": 1, # Exclusive + "max_angular": 2, # Inclusive + "atomic_gaussian_width": 0.2, + "radial_basis": {"Gto": {}}, + "cutoff_function": {"ShiftedCosine": {"width": 0.5}}, + "center_atom_weight": 1.0, +} + + +# ============ Pytest fixtures ============ + + +@pytest.fixture() +def cg_cache_sparse(): + return ClebschGordanReal(lambda_max=5, sparse=True) + + +@pytest.fixture() +def cg_cache_dense(): + return ClebschGordanReal(lambda_max=5, sparse=False) + + +# ============ Helper functions ============ + + +def h2_isolated(): + return ase.io.read(os.path.join(DATA_ROOT, "h2_isolated.xyz"), ":") + + +def h2o_isolated(): + return ase.io.read(os.path.join(DATA_ROOT, "h2o_isolated.xyz"), ":") + + +def h2o_periodic(): + return ase.io.read(os.path.join(DATA_ROOT, "h2o_periodic.xyz"), ":") + + +def wigner_d_matrices(lmax: int): + return WignerDReal(lmax=lmax) + + +def spherical_expansion(frames: List[ase.Atoms]): + """Returns a rascaline SphericalExpansion""" + calculator = rascaline.SphericalExpansion(**SPHEX_HYPERS) + return calculator.compute(frames) + + +def spherical_expansion_small(frames: List[ase.Atoms]): + """Returns a rascaline SphericalExpansion""" + calculator = rascaline.SphericalExpansion(**SPHEX_HYPERS_SMALL) + return calculator.compute(frames) + + +def power_spectrum(frames: List[ase.Atoms]): + """Returns a rascaline PowerSpectrum constructed from a + SphericalExpansion""" + return PowerSpectrum(rascaline.SphericalExpansion(**SPHEX_HYPERS)).compute(frames) + + +def power_spectrum_small(frames: List[ase.Atoms]): + """Returns a rascaline PowerSpectrum constructed from a + SphericalExpansion""" + return PowerSpectrum(rascaline.SphericalExpansion(**SPHEX_HYPERS_SMALL)).compute( + frames + ) + + +def get_norm(tensor: TensorMap): + """ + Calculates the norm used in CG iteration tests. Assumes that the TensorMap + is sliced to a single sample. + + For a given atomic sample, the norm is calculated for each feature vector, + as a sum over lambda, sigma, and m. + """ + # Check that there is only one sample + assert ( + len( + metatensor.unique_metadata( + tensor, "samples", ["structure", "center", "species_center"] + ).values + ) + == 1 + ) + norm = 0.0 + for key, block in tensor.items(): # Sum over lambda and sigma + angular_l = key["spherical_harmonics_l"] + norm += np.sum( + [ + np.linalg.norm(block.values[0, m, :]) ** 2 + for m in range(-angular_l, angular_l + 1) + ] + ) + + return norm + + +# ============ Test equivariance ============ + + +@pytest.mark.skipif( + not HAS_SYMPY or not HAS_METATENSOR_OPERATIONS, + reason="SymPy or metatensor-operations are not installed", +) +@pytest.mark.parametrize( + "frames, nu_target, angular_cutoff, selected_keys", + [ + ( + h2_isolated(), + 3, + None, + Labels( + names=["spherical_harmonics_l", "inversion_sigma"], + values=np.array([[0, 1], [4, 1], [5, 1]]), + ), + ), + (h2o_isolated(), 2, 5, None), + (h2o_periodic(), 2, 5, None), + ], +) +def test_so3_equivariance(frames, nu_target, angular_cutoff, selected_keys): + """ + Tests that the output of :py:func:`correlate_density` is equivariant under + SO(3) transformations. + """ + wig = wigner_d_matrices(nu_target * SPHEX_HYPERS["max_angular"]) + frames_so3 = [transform_frame_so3(frame, wig.angles) for frame in frames] + + nu_1 = spherical_expansion(frames) + nu_1_so3 = spherical_expansion(frames_so3) + + nu_3 = correlate_density( + density=nu_1, + correlation_order=nu_target, + angular_cutoff=angular_cutoff, + selected_keys=selected_keys, + ) + nu_3_so3 = correlate_density( + density=nu_1_so3, + correlation_order=nu_target, + angular_cutoff=angular_cutoff, + selected_keys=selected_keys, + ) + + nu_3_transf = wig.transform_tensormap_so3(nu_3) + assert metatensor.allclose(nu_3_transf, nu_3_so3) + + +@pytest.mark.skipif( + not HAS_SYMPY or not HAS_METATENSOR_OPERATIONS, + reason="SymPy or metatensor-operations are not installed", +) +@pytest.mark.parametrize( + "frames, nu_target, angular_cutoff, selected_keys", + [ + ( + h2_isolated(), + 3, + None, + Labels( + names=["spherical_harmonics_l"], + values=np.array([0, 4, 5]).reshape(-1, 1), + ), + ), + (h2o_isolated(), 2, 5, None), + (h2o_periodic(), 2, 5, None), + ], +) +def test_o3_equivariance(frames, nu_target, angular_cutoff, selected_keys): + """ + Tests that the output of :py:func:`correlate_density` is equivariant under + O(3) transformations. + """ + wig = wigner_d_matrices(nu_target * SPHEX_HYPERS["max_angular"]) + frames_o3 = [transform_frame_o3(frame, wig.angles) for frame in frames] + + nu_1 = spherical_expansion(frames) + nu_1_o3 = spherical_expansion(frames_o3) + + nu_3 = correlate_density( + density=nu_1, + correlation_order=nu_target, + angular_cutoff=angular_cutoff, + selected_keys=selected_keys, + ) + nu_3_o3 = correlate_density( + density=nu_1_o3, + correlation_order=nu_target, + angular_cutoff=angular_cutoff, + selected_keys=selected_keys, + ) + + nu_3_transf = wig.transform_tensormap_o3(nu_3) + assert metatensor.allclose(nu_3_transf, nu_3_o3) + + +# ============ Test lambda-SOAP vs PowerSpectrum ============ + + +@pytest.mark.skipif( + not HAS_METATENSOR_OPERATIONS, reason="metatensor-operations is not installed" +) +@pytest.mark.parametrize("frames", [h2_isolated()]) +@pytest.mark.parametrize( + "sphex_powspec", + [ + (spherical_expansion, power_spectrum), + (spherical_expansion_small, power_spectrum_small), + ], +) +def test_lambda_soap_vs_powerspectrum(frames, sphex_powspec): + """ + Tests for exact equivalence between the invariant block of a generated + lambda-SOAP equivariant and the Python implementation of PowerSpectrum in + rascaline utils. + """ + # Build a PowerSpectrum + ps = sphex_powspec[1](frames) + + # Build a lambda-SOAP + density = sphex_powspec[0](frames) + lsoap = correlate_density( + density=density, + correlation_order=2, + selected_keys=Labels( + names=["spherical_harmonics_l"], values=np.array([0]).reshape(-1, 1) + ), + ) + keys = lsoap.keys.remove(name="spherical_harmonics_l") + lsoap = TensorMap(keys=keys, blocks=[b.copy() for b in lsoap.blocks()]) + + # Manipulate metadata to match that of PowerSpectrum: + # 1) remove components axis + # 2) change "l1" and "l2" properties dimensions to just "l" (as l1 == l2) + blocks = [] + for block in lsoap.blocks(): + n_samples, n_props = block.values.shape[0], block.values.shape[2] + new_props = block.properties + new_props = new_props.remove(name="l1") + new_props = new_props.rename(old="l2", new="l") + blocks.append( + TensorBlock( + values=block.values.reshape((n_samples, n_props)), + samples=block.samples, + components=[], + properties=new_props, + ) + ) + lsoap = TensorMap(keys=lsoap.keys, blocks=blocks) + + # Compare metadata + assert metatensor.equal_metadata(lsoap, ps) + + # allclose on values + assert metatensor.allclose(lsoap, ps) + + +# ============ Test norm preservation ============ + + +@pytest.mark.skipif( + not HAS_METATENSOR_OPERATIONS, reason="metatensor-operations is not installed" +) +@pytest.mark.parametrize("frames", [h2_isolated(), h2o_periodic()]) +@pytest.mark.parametrize("correlation_order", [2, 3, 4]) +def test_correlate_density_norm(frames, correlation_order): + """ + Checks \\|ρ^\\nu\\| = \\|ρ\\|^\\nu in the case where l lists are not + sorted. If l lists are sorted, thus saving computation of redundant block + combinations, the norm check will not hold for target body order greater + than 2. + """ + + # Build nu=1 SphericalExpansion + nu1 = spherical_expansion_small(frames) + + # Build higher body order tensor without sorting the l lists + nux = correlate_density( + nu1, + correlation_order=correlation_order, + angular_cutoff=None, + selected_keys=None, + skip_redundant=False, + ) + # Build higher body order tensor *with* sorting the l lists + nux_sorted_l = correlate_density( + nu1, + correlation_order=correlation_order, + angular_cutoff=None, + selected_keys=None, + skip_redundant=True, + ) + + # Standardize the features by passing through the CG combination code but with + # no iterations (i.e. body order 1 -> 1) + nu1 = _standardize_keys(nu1) + + # Make only lambda and sigma part of keys + nu1 = nu1.keys_to_samples(["species_center"]) + nux = nux.keys_to_samples(["species_center"]) + nux_sorted_l = nux_sorted_l.keys_to_samples(["species_center"]) + + # The norm shoudl be calculated for each sample. First find the unqiue + # samples + uniq_samples = metatensor.unique_metadata( + nux, "samples", names=["structure", "center", "species_center"] + ) + grouped_labels = [ + Labels(names=nux.sample_names, values=uniq_samples.values[i].reshape(1, 3)) + for i in range(len(uniq_samples)) + ] + + # Calculate norms + norm_nu1 = 0.0 + norm_nux = 0.0 + norm_nux_sorted_l = 0.0 + for sample in grouped_labels: + # Slice the TensorMaps + nu1_sliced = metatensor.slice(nu1, "samples", labels=sample) + nux_sliced = metatensor.slice(nux, "samples", labels=sample) + nux_sorted_sliced = metatensor.slice(nux_sorted_l, "samples", labels=sample) + + # Calculate norms + norm_nu1 += get_norm(nu1_sliced) ** correlation_order + norm_nux += get_norm(nux_sliced) + norm_nux_sorted_l += get_norm(nux_sorted_sliced) + + # Without sorting the l list we should get the same norm + assert np.allclose(norm_nu1, norm_nux) + + # But with sorting the l list we should get a different norm + assert not np.allclose(norm_nu1, norm_nux_sorted_l) + + +# ============ Test CG cache ============ + + +@pytest.mark.parametrize("l1, l2", [(1, 2), (2, 3), (0, 5)]) +def test_clebsch_gordan_orthogonality(cg_cache_dense, l1, l2): + """ + Test orthogonality relationships of cached dense CG coefficients. + + See + https://en.wikipedia.org/wiki/Clebsch%E2%80%93Gordan_coefficients#Orthogonality_relations + for details. + """ + lam_min = abs(l1 - l2) + lam_max = l1 + l2 + + # We test lam dimension + # \sum_{-m1 \leq l1 \leq m1, -m2 \leq l2 \leq m2} + # <λμ|l1m1,l2m2> = δ_μμ' + for lam in range(lam_min, lam_max): + cg_mat = cg_cache_dense.coeffs[(l1, l2, lam)].reshape(-1, 2 * lam + 1) + dot_product = cg_mat.T @ cg_mat + diag_mask = np.zeros(dot_product.shape, dtype=np.bool_) + diag_mask[np.diag_indices(len(dot_product))] = True + assert np.allclose( + dot_product[~diag_mask], np.zeros(dot_product.shape)[~diag_mask] + ) + assert np.allclose(dot_product[diag_mask], dot_product[diag_mask][0]) + + # We test l1 l2 dimension + # \sum_{|l1-l2| \leq λ \leq l1+l2} \sum_{-μ \leq λ \leq μ} + # <λμ|l1m1,l2m2> = δ_m1m1' δ_m2m2' + l1l2_dim = (2 * l1 + 1) * (2 * l2 + 1) + dot_product = np.zeros((l1l2_dim, l1l2_dim)) + for lam in range(lam_min, lam_max + 1): + cg_mat = cg_cache_dense.coeffs[(l1, l2, lam)].reshape(-1, 2 * lam + 1) + dot_product += cg_mat @ cg_mat.T + diag_mask = np.zeros(dot_product.shape, dtype=np.bool_) + diag_mask[np.diag_indices(len(dot_product))] = True + + assert np.allclose(dot_product[~diag_mask], np.zeros(dot_product.shape)[~diag_mask]) + assert np.allclose(dot_product[diag_mask], dot_product[diag_mask][0]) + + +@pytest.mark.skipif( + not HAS_METATENSOR_OPERATIONS, reason="metatensor-operations is not installed" +) +@pytest.mark.parametrize("frames", [h2_isolated(), h2o_isolated()]) +def test_correlate_density_dense_sparse_agree(frames): + """ + Tests for agreement between nu=3 tensors built using both sparse and dense + CG coefficient caches. + """ + density = spherical_expansion_small(frames) + + # NOTE: testing the private function here so we can control the use of + # sparse v dense CG cache + n_body_sparse = _correlate_density( + density, + correlation_order=3, + compute_metadata_only=False, + sparse=True, + ) + n_body_dense = _correlate_density( + density, + correlation_order=3, + compute_metadata_only=False, + sparse=False, + ) + + assert metatensor.allclose(n_body_sparse, n_body_dense, atol=1e-8, rtol=1e-8) + + +# ============ Test metadata ============ + + +@pytest.mark.skipif( + not HAS_METATENSOR_OPERATIONS, reason="metatensor-operations is not installed" +) +@pytest.mark.parametrize("frames", [h2o_isolated()]) +@pytest.mark.parametrize("correlation_order", [2, 3]) +@pytest.mark.parametrize("skip_redundant", [True, False]) +def test_correlate_density_metadata_agree(frames, correlation_order, skip_redundant): + """ + Tests that the metadata of outputs from :py:func:`correlate_density` and + :py:func:`correlate_density_metadata` agree. + """ + for nu1 in [spherical_expansion_small(frames), spherical_expansion(frames)]: + # Build higher body order tensor with CG computation + nux = correlate_density( + nu1, + correlation_order=correlation_order, + angular_cutoff=None, + selected_keys=None, + skip_redundant=skip_redundant, + ) + # Build higher body order tensor without CG computation - i.e. metadata + # only + nux_metadata_only = correlate_density_metadata( + nu1, + correlation_order=correlation_order, + angular_cutoff=None, + selected_keys=None, + skip_redundant=skip_redundant, + ) + assert metatensor.equal_metadata(nux, nux_metadata_only) + + +@pytest.mark.parametrize("frames", [h2o_isolated()]) +@pytest.mark.parametrize( + "selected_keys", + [ + None, + Labels( + names=["spherical_harmonics_l"], values=np.array([1, 2, 4]).reshape(-1, 1) + ), + ], +) +@pytest.mark.parametrize("skip_redundant", [True, False]) +def test_correlate_density_angular_selection( + frames: List[ase.Atoms], + selected_keys: Labels, + skip_redundant: bool, +): + """ + Tests that the correct angular channels are outputted based on the + specified ``selected_keys``. + """ + nu_1 = spherical_expansion(frames) + + nu_2 = correlate_density( + density=nu_1, + correlation_order=2, + angular_cutoff=None, + selected_keys=selected_keys, + skip_redundant=skip_redundant, + ) + + if selected_keys is None: + assert np.all( + [ + angular in np.arange(SPHEX_HYPERS["max_angular"] * 2 + 1) + for angular in np.unique(nu_2.keys.column("spherical_harmonics_l")) + ] + ) + + else: + assert np.all( + np.sort(np.unique(nu_2.keys.column("spherical_harmonics_l"))) + == np.sort(selected_keys.column("spherical_harmonics_l")) + ) diff --git a/python/rascaline/tests/utils/data/h2_isolated.xyz b/python/rascaline/tests/utils/data/h2_isolated.xyz new file mode 100644 index 000000000..ec5f59680 --- /dev/null +++ b/python/rascaline/tests/utils/data/h2_isolated.xyz @@ -0,0 +1,4 @@ +2 +pbc="F F F" +H 1.97361700 1.73067300 2.47063400 +H 1.97361700 3.26932700 2.47063400 diff --git a/python/rascaline/tests/utils/data/h2o_isolated.xyz b/python/rascaline/tests/utils/data/h2o_isolated.xyz new file mode 100644 index 000000000..fc876d2ba --- /dev/null +++ b/python/rascaline/tests/utils/data/h2o_isolated.xyz @@ -0,0 +1,5 @@ +3 +pbc="F F F" +O 2.56633400 2.50000000 2.50370100 +H 1.97361700 1.73067300 2.47063400 +H 1.97361700 3.26932700 2.47063400 diff --git a/python/rascaline/tests/utils/data/h2o_periodic.xyz b/python/rascaline/tests/utils/data/h2o_periodic.xyz new file mode 100644 index 000000000..3374566e6 --- /dev/null +++ b/python/rascaline/tests/utils/data/h2o_periodic.xyz @@ -0,0 +1,5 @@ +3 +Lattice="5.0 0.0 0.0 0.0 5.0 0.0 0.0 0.0 5.0" pbc="T T T" +O 2.56633400 2.50000000 2.50370100 +H 1.97361700 1.73067300 2.47063400 +H 1.97361700 3.26932700 2.47063400 diff --git a/python/rascaline/tests/utils/radial_basis.py b/python/rascaline/tests/utils/radial_basis.py new file mode 100644 index 000000000..986717fd0 --- /dev/null +++ b/python/rascaline/tests/utils/radial_basis.py @@ -0,0 +1,91 @@ +from typing import Union + +import numpy as np +import pytest +from numpy.testing import assert_allclose + +from rascaline.utils import ( + GtoBasis, + MonomialBasis, + RadialBasisBase, + SphericalBesselBasis, +) + + +pytest.importorskip("scipy") + + +class RtoNRadialBasis(RadialBasisBase): + def compute( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + return integrand_positions**n + + +def test_radial_basis_gram(): + """Test that quad integration of the gram matrix is the same as an analytical.""" + + integration_radius = 1 + max_radial = 4 + max_angular = 2 + + test_basis = RtoNRadialBasis(integration_radius=integration_radius) + + numerical_gram = test_basis.compute_gram_matrix(max_radial, max_angular) + analytical_gram = np.zeros_like(numerical_gram) + + for ell in range(max_angular + 1): + for n1 in range(max_radial): + for n2 in range(max_radial): + exp = 3 + n1 + n2 + analytical_gram[ell, n1, n2] = integration_radius**exp / exp + + assert_allclose(numerical_gram, analytical_gram) + + +def test_radial_basis_orthornormalization(): + integration_radius = 1 + max_radial = 4 + max_angular = 2 + + test_basis = RtoNRadialBasis(integration_radius=integration_radius) + + gram = test_basis.compute_gram_matrix(max_radial, max_angular) + ortho = test_basis.compute_orthonormalization_matrix(max_radial, max_angular) + + for ell in range(max_angular): + eye = ortho[ell] @ gram[ell] @ ortho[ell].T + assert_allclose(eye, np.eye(max_radial, max_radial), atol=1e-11) + + +@pytest.mark.parametrize( + "analytical_basis", + [ + GtoBasis(cutoff=4, max_radial=6), + MonomialBasis(cutoff=4), + SphericalBesselBasis(cutoff=4, max_radial=6, max_angular=4), + ], +) +def test_derivative(analytical_basis: RadialBasisBase): + """Finite difference test for testing the derivative of a radial basis""" + + class NumericalRadialBasis(RadialBasisBase): + def compute( + self, n: float, ell: float, integrand_positions: Union[float, np.ndarray] + ) -> Union[float, np.ndarray]: + return analytical_basis.compute(n, ell, integrand_positions) + + numerical_basis = NumericalRadialBasis(integration_radius=np.inf) + + cutoff = 4 + max_radial = 6 + max_angular = 4 + positions = np.linspace(2, cutoff) + + for n in range(max_radial): + for ell in range(max_angular): + assert_allclose( + numerical_basis.compute_derivative(n, ell, positions), + analytical_basis.compute_derivative(n, ell, positions), + atol=1e-9, + ) diff --git a/python/rascaline/tests/utils/rotations.py b/python/rascaline/tests/utils/rotations.py new file mode 100644 index 000000000..2abe22714 --- /dev/null +++ b/python/rascaline/tests/utils/rotations.py @@ -0,0 +1,330 @@ +""" +Class for generating real Wigner-D matrices, and using them to rotate ASE frames +and TensorMaps of density coefficients in the spherical basis. +""" +from typing import Sequence + +import pytest + + +ase = pytest.importorskip("ase") + +import numpy as np # noqa: E402 +from metatensor import TensorBlock, TensorMap # noqa: E402 +from scipy.spatial.transform import Rotation # noqa: E402 + + +try: + import torch # noqa: E402 + from torch import Tensor as TorchTensor # noqa: E402 +except ImportError: + + class TorchTensor: + pass + + +# ===== Functions for transformations in the Cartesian basis ===== + + +def cartesian_rotation(angles: Sequence[float]): + """ + Returns a Cartesian rotation matrix in the appropriate convention (ZYZ, + implicit rotations) to be consistent with the common Wigner D definition. + + `angles` correspond to the alpha, beta, gamma Euler angles in the ZYZ + convention, in radians. + """ + return Rotation.from_euler("ZYZ", angles).as_matrix() + + +def transform_frame_so3(frame: ase.Atoms, angles: Sequence[float]) -> ase.Atoms: + """ + Transforms the positions and cell coordinates of an ASE frame by a SO(3) + rigid rotation. + """ + new_frame = frame.copy() + + # Build cartesian rotation matrix + R = cartesian_rotation(angles) + + # Rotate its positions and cell + new_frame.positions = new_frame.positions @ R.T + new_frame.cell = new_frame.cell @ R.T + + return new_frame + + +def transform_frame_o3(frame: ase.Atoms, angles: Sequence[float]) -> ase.Atoms: + """ + Transforms the positions and cell coordinates of an ASE frame by an O(3) + rotation. This involves a rigid SO(3) rotation of the positions and cell + according to the Euler `angles`, then an inversion by multiplying just the + positions by -1. + """ + new_frame = frame.copy() + + # Build cartesian rotation matrix + R = cartesian_rotation(angles) + + # Rotate its positions and cell + new_frame.positions = new_frame.positions @ R.T + new_frame.cell = new_frame.cell @ R.T + + # Invert the atom positions + new_frame.positions *= -1 + + return new_frame + + +# ===== WignerDReal for transformations in the spherical basis ===== + + +class WignerDReal: + """ + A helper class to compute Wigner D matrices given the Euler angles of a rotation, + and apply them to spherical harmonics (or coefficients). Built to function with + real-valued coefficients. + """ + + def __init__(self, lmax: int, angles: Sequence[float] = None): + """ + Initialize the WignerDReal class. + + :param lmax: int, the maximum angular momentum channel for which the + Wigner D matrices are computed + :param angles: Sequence[float], the alpha, beta, gamma Euler angles, in + radians. + """ + self.lmax = lmax + # Randomly generate Euler angles between 0 and 2 pi if none are provided + if angles is None: + angles = np.random.uniform(size=(3)) * 2 * np.pi + self.angles = angles + self.rotation = cartesian_rotation(angles) + + r2c_mats = {} + c2r_mats = {} + for L in range(0, self.lmax + 1): + r2c_mats[L] = np.hstack( + [_r2c(np.eye(2 * L + 1)[i])[:, np.newaxis] for i in range(2 * L + 1)] + ) + c2r_mats[L] = np.conjugate(r2c_mats[L]).T + self.matrices = {} + for L in range(0, self.lmax + 1): + wig = _wigner_d(L, self.angles) + self.matrices[L] = np.real(c2r_mats[L] @ np.conjugate(wig) @ r2c_mats[L]) + + def rotate_coeff_vector( + self, + frame: ase.Atoms, + coeffs: np.ndarray, + lmax: dict, + nmax: dict, + ) -> np.ndarray: + """ + Rotates the irreducible spherical components (ISCs) of basis set + coefficients in the spherical basis passed in as a flat vector. + + Required is the basis set definition specified by ``lmax`` and ``nmax``. + This are dicts of the form: + + lmax = {symbol: lmax_value, ...} + nmax = {(symbol, l): nmax_value, ...} + + where ``symbol`` is the chemical symbol of the atom, ``lmax_value`` is + its corresponding max l channel value. For each combination of species + symbol and lmax, there exists a max radial channel value ``nmax_value``. + + Then, the assumed ordering of basis function coefficients follows a + hierarchy, which can be read as nested loops over the various indices. + Be mindful that some indices range are from 0 to x (exclusive) and + others from 0 to x + 1 (exclusive). The ranges reported below are + ordered. + + 1. Loop over atoms (index ``i``, of chemical species ``a``) in the + structure. ``i`` takes values 0 to N (** exclusive **), where N is the + number of atoms in the structure. + + 2. Loop over spherical harmonics channel (index ``l``) for each atom. + ``l`` takes values from 0 to ``lmax[a] + 1`` (** exclusive **), where + ``a`` is the chemical species of atom ``i``, given by the chemical + symbol at the ``i``th position of ``symbol_list``. + + 3. Loop over radial channel (index ``n``) for each atom ``i`` and + spherical harmonics channel ``l`` combination. ``n`` takes values from 0 + to ``nmax[(a, l)]`` (** exclusive **). + + 4. Loop over spherical harmonics component (index ``m``) for each atom. + ``m`` takes values from ``-l`` to ``l`` (** inclusive **). + + :param frame: the atomic structure in ASE format for which the + coefficients are defined. + :param coeffs: the coefficients in the spherical basis, as a flat + vector. + :param lmax: dict containing the maximum spherical harmonics (l) value + for each atom type. + :param nmax: dict containing the maximum radial channel (n) value for + each combination of atom type and l. + + :return: the rotated coefficients in the spherical basis, as a flat + vector with the same order as the input vector. + """ + # Initialize empty vector for storing the rotated ISCs + rot_vect = np.empty_like(coeffs) + + # Iterate over atomic species of the atoms in the frame + curr_idx = 0 + for symbol in frame.get_chemical_symbols(): + # Get the basis set lmax value for this species + sym_lmax = lmax[symbol] + for angular_l in range(sym_lmax + 1): + # Get the number of radial functions for this species and l value + sym_l_nmax = nmax[(symbol, angular_l)] + # Get the Wigner D Matrix for this l value + wig_mat = self.matrices[angular_l].T + for _n in range(sym_l_nmax): + # Retrieve the irreducible spherical component + isc = coeffs[curr_idx : curr_idx + (2 * angular_l + 1)] + # Rotate the ISC and store + rot_isc = isc @ wig_mat + rot_vect[curr_idx : curr_idx + (2 * angular_l + 1)][:] = rot_isc[:] + # Update the start index for the next ISC + curr_idx += 2 * angular_l + 1 + + return rot_vect + + def rotate_tensorblock(self, angular_l: int, block: TensorBlock) -> TensorBlock: + """ + Rotates a TensorBlock ``block``, represented in the spherical basis, + according to the Wigner D Real matrices for the given ``l`` value. + Assumes the components of the block are [("spherical_harmonics_m",),]. + """ + # Get the Wigner matrix for this l value + wig = self.matrices[angular_l].T + + # Copy the block + block_rotated = block.copy() + vals = block_rotated.values + + # Perform the rotation, either with numpy or torch, by taking the + # tensordot product of the irreducible spherical components. Modify + # in-place the values of the copied TensorBlock. + if isinstance(vals, TorchTensor): + wig = torch.tensor(wig) + block_rotated.values[:] = torch.tensordot( + vals.swapaxes(1, 2), wig, dims=1 + ).swapaxes(1, 2) + elif isinstance(block.values, np.ndarray): + block_rotated.values[:] = np.tensordot( + vals.swapaxes(1, 2), wig, axes=1 + ).swapaxes(1, 2) + else: + raise TypeError("TensorBlock values must be a numpy array or torch tensor.") + + return block_rotated + + def transform_tensormap_so3(self, tensor: TensorMap) -> TensorMap: + """ + Transforms a TensorMap by a by an SO(3) rigid rotation using Wigner-D + matrices. + + Assumes the input tensor follows the metadata structure consistent with + those produce by rascaline. + """ + # Retrieve the key and the position of the l value in the key names + keys = tensor.keys + idx_l_value = keys.names.index("spherical_harmonics_l") + + # Iterate over the blocks and rotate + rotated_blocks = [] + for key in keys: + # Retrieve the l value + angular_l = key[idx_l_value] + + # Rotate the block and store + rotated_blocks.append(self.rotate_tensorblock(angular_l, tensor[key])) + + return TensorMap(keys, rotated_blocks) + + def transform_tensormap_o3(self, tensor: TensorMap) -> TensorMap: + """ + Transforms a TensorMap by a by an O(3) transformation: this involves an + SO(3) rigid rotation using Wigner-D Matrices followed by an inversion. + + Assumes the input tensor follows the metadata structure consistent with + those produce by rascaline. + """ + # Retrieve the key and the position of the l value in the key names + keys = tensor.keys + idx_l_value = keys.names.index("spherical_harmonics_l") + + # Iterate over the blocks and rotate + new_blocks = [] + for key in keys: + # Retrieve the l value + angular_l = key[idx_l_value] + + # Rotate the block + new_block = self.rotate_tensorblock(angular_l, tensor[key]) + + # Work out the inversion multiplier according to the convention + inversion_multiplier = 1 + if key["spherical_harmonics_l"] % 2 == 1: + inversion_multiplier *= -1 + + # "inversion_sigma" may not be present if CG iterations haven't been + # performed (i.e. nu=1 rascaline SphericalExpansion) + if "inversion_sigma" in keys.names: + if key["inversion_sigma"] == -1: + inversion_multiplier *= -1 + + # Invert the block by applying the inversion multiplier + new_block = TensorBlock( + values=new_block.values * inversion_multiplier, + samples=new_block.samples, + components=new_block.components, + properties=new_block.properties, + ) + new_blocks.append(new_block) + + return TensorMap(keys, new_blocks) + + +# ===== Helper functions for WignerDReal + + +def _wigner_d(angular_l: int, angles: Sequence[float]) -> np.ndarray: + """ + Computes the Wigner D matrix: + D^l_{mm'}(alpha, beta, gamma) + from sympy and converts it to numerical values. + + `angles` are the alpha, beta, gamma Euler angles (radians, ZYZ convention) + and l the irrep. + """ + try: + from sympy.physics.wigner import wigner_d + except ModuleNotFoundError: + raise ModuleNotFoundError( + "Calculation of Wigner D matrices requires a sympy installation" + ) + return np.complex128(wigner_d(angular_l, *angles)) + + +def _r2c(sp): + """ + Real to complex SPH. Assumes a block with 2l+1 reals corresponding + to real SPH with m indices from -l to +l + """ + + i_sqrt_2 = 1.0 / np.sqrt(2) + + angular_l = (len(sp) - 1) // 2 # infers l from the vector size + rc = np.zeros(len(sp), dtype=np.complex128) + rc[angular_l] = sp[angular_l] + for m in range(1, angular_l + 1): + rc[angular_l + m] = ( + (sp[angular_l + m] + 1j * sp[angular_l - m]) * i_sqrt_2 * (-1) ** m + ) + rc[angular_l - m] = (sp[angular_l + m] - 1j * sp[angular_l - m]) * i_sqrt_2 + return rc diff --git a/python/rascaline/tests/utils/splines.py b/python/rascaline/tests/utils/splines.py index a08cc54e3..9ff918e99 100644 --- a/python/rascaline/tests/utils/splines.py +++ b/python/rascaline/tests/utils/splines.py @@ -2,7 +2,22 @@ import pytest from numpy.testing import assert_allclose, assert_equal -from rascaline.utils import RadialIntegralFromFunction +from rascaline import LodeSphericalExpansion, SphericalExpansion +from rascaline.utils import ( + DeltaDensity, + GaussianDensity, + GtoBasis, + LodeDensity, + LodeSpliner, + RadialIntegralFromFunction, + SoapSpliner, +) + +from ..test_systems import SystemForTests + + +pytest.importorskip("scipy") +from scipy.special import gamma, hyp1f1 # noqa def sine(n: int, ell: int, positions: np.ndarray) -> np.ndarray: @@ -102,3 +117,214 @@ def test_splines_numerical_derivative_error(): match = "Numerically derivative of the radial integral can not be performed" with pytest.raises(ValueError, match=match): RadialIntegralFromFunction(**kwargs).compute() + + +def test_kspace_radial_integral(): + """Test against analytical integral with Gaussian densities and GTOs""" + + cutoff = 2 + max_radial = 6 + max_angular = 3 + atomic_gaussian_width = 1.0 + k_cutoff = 1.2 * np.pi / atomic_gaussian_width + + basis = GtoBasis(cutoff=cutoff, max_radial=max_radial) + + spliner = LodeSpliner( + max_radial=max_radial, + max_angular=max_angular, + k_cutoff=k_cutoff, + basis=basis, + density=DeltaDensity(), # density does not enter in a Kspace radial integral + accuracy=1e-8, + ) + + Neval = 100 + kk = np.linspace(0, k_cutoff, Neval) + + sigma = np.ones(max_radial, dtype=float) + for i in range(1, max_radial): + sigma[i] = np.sqrt(i) + sigma *= cutoff / max_radial + + factors = np.sqrt(np.pi) * np.ones((max_radial, max_angular + 1)) + + coeffs_num = np.zeros([max_radial, max_angular + 1, Neval]) + coeffs_exact = np.zeros_like(coeffs_num) + + for ell in range(max_angular + 1): + for n in range(max_radial): + i1 = 0.5 * (3 + n + ell) + i2 = 1.5 + ell + factors[n, ell] *= ( + 2 ** (0.5 * (n - ell - 1)) + * gamma(i1) + / gamma(i2) + * sigma[n] ** (2 * i1) + ) + coeffs_exact[n, ell] = ( + factors[n, ell] + * kk**ell + * hyp1f1(i1, i2, -0.5 * (kk * sigma[n]) ** 2) + ) + + coeffs_num[n, ell] = spliner.radial_integral(n, ell, kk) + + assert_allclose(coeffs_num, coeffs_exact) + + +def test_rspace_delta(): + cutoff = 2 + max_radial = 6 + max_angular = 3 + + basis = GtoBasis(cutoff=cutoff, max_radial=max_radial) + density = DeltaDensity() + + spliner = SoapSpliner( + max_radial=max_radial, + max_angular=max_angular, + cutoff=cutoff, + basis=basis, + density=density, + accuracy=1e-8, + ) + + positions = np.linspace(0, cutoff) + + for ell in range(max_angular + 1): + for n in range(max_radial): + assert_equal( + spliner.radial_integral(n, ell, positions), + basis.compute(n, ell, positions), + ) + assert_equal( + spliner.radial_integral_derivative(n, ell, positions), + basis.compute_derivative(n, ell, positions), + ) + + +def test_real_space_spliner(): + """Compare splined spherical expansion with GTOs and a Gaussian density to + analytical implementation.""" + cutoff = 8.0 + max_radial = 12 + max_angular = 9 + atomic_gaussian_width = 1.2 + + # We choose an accuracy that is larger then the default one (1e-8) to limit the time + # consumption of the test. + accuracy = 1e-4 + + spliner = SoapSpliner( + cutoff=cutoff, + max_radial=max_radial, + max_angular=max_angular, + basis=GtoBasis(cutoff=cutoff, max_radial=max_radial), + density=GaussianDensity(atomic_gaussian_width=atomic_gaussian_width), + accuracy=accuracy, + ) + + hypers_spherical_expansion = { + "cutoff": cutoff, + "max_radial": max_radial, + "max_angular": max_angular, + "center_atom_weight": 1.0, + "atomic_gaussian_width": atomic_gaussian_width, + "cutoff_function": {"Step": {}}, + } + + analytic = SphericalExpansion( + radial_basis={"Gto": {}}, **hypers_spherical_expansion + ).compute(SystemForTests()) + splined = SphericalExpansion( + radial_basis=spliner.compute(), **hypers_spherical_expansion + ).compute(SystemForTests()) + + for key, block_analytic in analytic.items(): + block_splined = splined.block(key) + assert_allclose( + block_splined.values, block_analytic.values, rtol=5e-4, atol=2e-5 + ) + + +@pytest.mark.parametrize("center_atom_weight", [1.0, 0.0]) +@pytest.mark.parametrize("potential_exponent", [0, 1]) +def test_fourier_space_spliner(center_atom_weight, potential_exponent): + """Compare splined LODE spherical expansion with GTOs and a Gaussian density to + analytical implementation.""" + + cutoff = 2 + max_radial = 6 + max_angular = 4 + atomic_gaussian_width = 0.8 + k_cutoff = 1.2 * np.pi / atomic_gaussian_width + + spliner = LodeSpliner( + k_cutoff=k_cutoff, + max_radial=max_radial, + max_angular=max_angular, + basis=GtoBasis(cutoff=cutoff, max_radial=max_radial), + density=LodeDensity( + atomic_gaussian_width=atomic_gaussian_width, + potential_exponent=potential_exponent, + ), + ) + + hypers_spherical_expansion = { + "cutoff": cutoff, + "max_radial": max_radial, + "max_angular": max_angular, + "center_atom_weight": center_atom_weight, + "atomic_gaussian_width": atomic_gaussian_width, + "potential_exponent": potential_exponent, + } + + analytic = LodeSphericalExpansion( + radial_basis={"Gto": {}}, **hypers_spherical_expansion + ).compute(SystemForTests()) + splined = LodeSphericalExpansion( + radial_basis=spliner.compute(), **hypers_spherical_expansion + ).compute(SystemForTests()) + + for key, block_analytic in analytic.items(): + block_splined = splined.block(key) + assert_allclose(block_splined.values, block_analytic.values, atol=1e-14) + + +def test_center_contribution_gto_gaussian(): + cutoff = 2.0 + max_radial = 6 + max_angular = 4 + atomic_gaussian_width = 0.8 + k_cutoff = 1.2 * np.pi / atomic_gaussian_width + + # Numerical evaluation of center contributions + spliner = LodeSpliner( + k_cutoff=k_cutoff, + max_radial=max_radial, + max_angular=max_angular, + basis=GtoBasis(cutoff=cutoff, max_radial=max_radial), + density=GaussianDensity(atomic_gaussian_width=atomic_gaussian_width), + ) + + # Analytical evaluation of center contributions + center_contr_analytical = np.zeros((max_radial)) + + normalization = 1.0 / (np.pi * atomic_gaussian_width**2) ** (3 / 4) + sigma_radial = np.ones(max_radial, dtype=float) + + for n in range(1, max_radial): + sigma_radial[n] = np.sqrt(n) + sigma_radial *= cutoff / max_radial + + for n in range(max_radial): + sigmatemp_sq = 1.0 / ( + 1.0 / atomic_gaussian_width**2 + 1.0 / sigma_radial[n] ** 2 + ) + neff = 0.5 * (3 + n) + center_contr_analytical[n] = (2 * sigmatemp_sq) ** neff * gamma(neff) + + center_contr_analytical *= normalization * 2 * np.pi / np.sqrt(4 * np.pi) + + assert_allclose(spliner.center_contribution, center_contr_analytical, rtol=1e-14) diff --git a/rascaline-torch/include/rascaline/torch/calculator.hpp b/rascaline-torch/include/rascaline/torch/calculator.hpp index bdc09bb1f..ee54929de 100644 --- a/rascaline-torch/include/rascaline/torch/calculator.hpp +++ b/rascaline-torch/include/rascaline/torch/calculator.hpp @@ -75,7 +75,8 @@ class RASCALINE_TORCH_EXPORT CalculatorHolder: public torch::CustomClassHolder { public: /// Create a new calculator with the given `name` and JSON `parameters` CalculatorHolder(std::string name, std::string parameters): - calculator_(std::move(name), std::move(parameters)) + c_name_(std::move(name)), + calculator_(c_name_, std::move(parameters)) {} /// Get the name of this calculator @@ -83,6 +84,11 @@ class RASCALINE_TORCH_EXPORT CalculatorHolder: public torch::CustomClassHolder { return calculator_.name(); } + /// Get the name used to register this calculator + std::string c_name() const { + return c_name_; + } + /// Get the parameters of this calculator std::string parameters() const { return calculator_.parameters(); @@ -100,6 +106,7 @@ class RASCALINE_TORCH_EXPORT CalculatorHolder: public torch::CustomClassHolder { ); private: + std::string c_name_; rascaline::Calculator calculator_; }; diff --git a/rascaline-torch/src/register.cpp b/rascaline-torch/src/register.cpp index 965b10865..75cf845d4 100644 --- a/rascaline-torch/src/register.cpp +++ b/rascaline-torch/src/register.cpp @@ -52,13 +52,13 @@ TORCH_LIBRARY(rascaline, module) { }) .def_pickle( // __getstate__ - [](const TorchCalculator& self) -> std::vector { - return {self->name(), self->parameters()}; + [](const TorchCalculator& self) -> std::tuple { + return {self->c_name(), self->parameters()}; }, // __setstate__ - [](std::vector state) -> TorchCalculator { + [](std::tuple state) -> TorchCalculator { return c10::make_intrusive( - state[0], state[1] + std::get<0>(state), std::get<1>(state) ); }) ; diff --git a/rascaline/src/calculators/lode/radial_integral/mod.rs b/rascaline/src/calculators/lode/radial_integral/mod.rs index 5b98cadb3..c76b626df 100644 --- a/rascaline/src/calculators/lode/radial_integral/mod.rs +++ b/rascaline/src/calculators/lode/radial_integral/mod.rs @@ -81,7 +81,7 @@ impl LodeRadialIntegralCache { // the largest value the spline should interpolate is // the k-space cutoff, not the real-space cutoff // associated with the GTO basis - cutoff: parameters.k_cutoff, + k_cutoff: parameters.k_cutoff, }; Box::new(LodeRadialIntegralSpline::with_accuracy( @@ -95,7 +95,7 @@ impl LodeRadialIntegralCache { let parameters = LodeRadialIntegralSplineParameters { max_radial: parameters.max_radial, max_angular: parameters.max_angular, - cutoff: parameters.cutoff, + k_cutoff: parameters.k_cutoff, }; let center_contribution = center_contribution.ok_or(Error::InvalidParameter( diff --git a/rascaline/src/calculators/lode/radial_integral/spline.rs b/rascaline/src/calculators/lode/radial_integral/spline.rs index eabeb00e5..7597de15f 100644 --- a/rascaline/src/calculators/lode/radial_integral/spline.rs +++ b/rascaline/src/calculators/lode/radial_integral/spline.rs @@ -24,8 +24,8 @@ pub struct LodeRadialIntegralSplineParameters { pub max_radial: usize, /// Number of angular components pub max_angular: usize, - /// cutoff radius, this is also the maximal value that can be interpolated - pub cutoff: f64, + /// k-space cutoff radius, this is also the maximal value that can be interpolated + pub k_cutoff: f64, } impl LodeRadialIntegralSpline { @@ -44,7 +44,7 @@ impl LodeRadialIntegralSpline { let parameters = SplineParameters { start: 0.0, - stop: parameters.cutoff, + stop: parameters.k_cutoff, shape: vec![parameters.max_angular + 1, parameters.max_radial], }; @@ -74,7 +74,7 @@ impl LodeRadialIntegralSpline { let spline_parameters = SplineParameters { start: 0.0, - stop: parameters.cutoff, + stop: parameters.k_cutoff, shape: vec![parameters.max_angular + 1, parameters.max_radial], }; @@ -123,18 +123,17 @@ mod tests { #[test] fn high_accuracy() { - // Check that even with high accuracy and large domain MAX_SPLINE_SIZE - // is enough + // Check that even with high accuracy and large domain MAX_SPLINE_SIZE is enough let parameters = LodeRadialIntegralSplineParameters { max_radial: 15, max_angular: 10, - cutoff: 12.0, + k_cutoff: 10.0, }; let gto = LodeRadialIntegralGto::new(LodeRadialIntegralGtoParameters { max_radial: parameters.max_radial, max_angular: parameters.max_angular, - cutoff: parameters.cutoff, + cutoff: 5.0, atomic_gaussian_width: 0.5, potential_exponent: 1, }).unwrap(); @@ -150,20 +149,19 @@ mod tests { let parameters = LodeRadialIntegralSplineParameters { max_radial: max_radial, max_angular: max_angular, - cutoff: 5.0, + k_cutoff: 10.0, }; let gto = LodeRadialIntegralGto::new(LodeRadialIntegralGtoParameters { max_radial: parameters.max_radial, max_angular: parameters.max_angular, - cutoff: parameters.cutoff, + cutoff: 5.0, atomic_gaussian_width: 0.5, potential_exponent: 1, }).unwrap(); - // even with very bad accuracy, we want the gradients of the spline to - // match the values produces by the spline, and not necessarily the - // actual GTO gradients. + // even with very bad accuracy, we want the gradients of the spline to match the + // values produces by the spline, and not necessarily the actual GTO gradients. let spline = LodeRadialIntegralSpline::with_accuracy(parameters, 1e-2, gto).unwrap(); let rij = 3.4; @@ -204,13 +202,13 @@ mod tests { let parameters = LodeRadialIntegralSplineParameters { max_radial: max_radial, max_angular: max_angular, - cutoff: 5.0, + k_cutoff: 10.0, }; let gto = LodeRadialIntegralGto::new(LodeRadialIntegralGtoParameters { max_radial: parameters.max_radial, max_angular: parameters.max_angular, - cutoff: parameters.cutoff, + cutoff: 5.0, atomic_gaussian_width: 0.5, potential_exponent: 1, }).unwrap(); diff --git a/rascaline/src/calculators/soap/power_spectrum.rs b/rascaline/src/calculators/soap/power_spectrum.rs index 045b94f00..d3c581a81 100644 --- a/rascaline/src/calculators/soap/power_spectrum.rs +++ b/rascaline/src/calculators/soap/power_spectrum.rs @@ -358,8 +358,19 @@ impl SoapPowerSpectrum { let property_1 = block_1.properties.position(&[n1]).expect("missing n1"); let property_2 = block_2.properties.position(&[n2]).expect("missing n2"); + let spherical_harmonics_l = l.usize(); + + // For consistency with a full Clebsch-Gordan product we need to add + // a `-1^l / sqrt(2 l + 1)` factor to the power spectrum invariants + let normalization = if spherical_harmonics_l % 2 == 0 { + f64::sqrt((2 * spherical_harmonics_l + 1) as f64) + } else { + -f64::sqrt((2 * spherical_harmonics_l + 1) as f64) + }; + SpxPropertiesToCombine { - spherical_harmonics_l: l.usize(), + spherical_harmonics_l, + normalization, property_1, property_2, spx_1: block_1.clone(), @@ -375,6 +386,8 @@ impl SoapPowerSpectrum { struct SpxPropertiesToCombine<'a> { /// value of l spherical_harmonics_l: usize, + /// normalization factor $-1^l * \sqrt{2 l + 1}$ + normalization: f64, /// position of n1 in the first spherical expansion properties property_1: usize, /// position of n2 in the second spherical expansion properties @@ -599,7 +612,7 @@ impl CalculatorBase for SoapPowerSpectrum { } unsafe { - *values.uget_mut(property_i) = sum / f64::sqrt((2 * spx.spherical_harmonics_l + 1) as f64); + *values.uget_mut(property_i) = sum / spx.normalization; } } }); @@ -655,10 +668,9 @@ impl CalculatorBase for SoapPowerSpectrum { } } - let normalization = f64::sqrt((2 * spx.spherical_harmonics_l + 1) as f64); for d in 0..3 { unsafe { - *values.uget_mut([d, property_i]) = sum[d] / normalization; + *values.uget_mut([d, property_i]) = sum[d] / spx.normalization; } } } @@ -694,7 +706,6 @@ impl CalculatorBase for SoapPowerSpectrum { let value_2 = spx_2.values.uget([spx_sample_2, m, spx.property_2]); for d1 in 0..3 { for d2 in 0..3 { - // TODO: ensure that gradient samples are 0..nsamples sum[d1][d2] += value_2 * spx_1_gradient.uget([spx_sample_1, d1, d2, m, spx.property_1]); } } @@ -707,7 +718,6 @@ impl CalculatorBase for SoapPowerSpectrum { let value_1 = spx_1.values.uget([spx_sample_1, m, spx.property_1]); for d1 in 0..3 { for d2 in 0..3 { - // TODO: ensure that gradient samples are 0..nsamples sum[d1][d2] += value_1 * spx_2_gradient.uget([spx_sample_2, d1, d2, m, spx.property_2]); } } @@ -723,12 +733,10 @@ impl CalculatorBase for SoapPowerSpectrum { } } - let normalization = f64::sqrt((2 * spx.spherical_harmonics_l + 1) as f64); - for d1 in 0..3 { for d2 in 0..3 { unsafe { - *values.uget_mut([d1, d2, property_i]) = sum[d1][d2] / normalization; + *values.uget_mut([d1, d2, property_i]) = sum[d1][d2] / spx.normalization; } } } diff --git a/rascaline/tests/soap-power-spectrum.rs b/rascaline/tests/soap-power-spectrum.rs index 42bfdb1da..632b02e6b 100644 --- a/rascaline/tests/soap-power-spectrum.rs +++ b/rascaline/tests/soap-power-spectrum.rs @@ -24,8 +24,10 @@ fn values() { let block = &descriptor.block_by_id(0); let array = block.values().to_array(); - let expected = &data::load_expected_values("soap-power-spectrum-values.npy.gz"); - assert_relative_eq!(array, expected, max_relative=1e-5); + let mut expected = data::load_expected_values("soap-power-spectrum-values.npy.gz"); + correct_factor(&mut expected, block.properties()); + + assert_relative_eq!(array, &expected, max_relative=1e-5); } #[test] @@ -51,13 +53,15 @@ fn gradients() { let gradients = block.gradient("positions").unwrap(); let array = sum_gradients(n_atoms, gradients); - let expected = &data::load_expected_values("soap-power-spectrum-positions-gradient.npy.gz"); - assert_relative_eq!(array, expected, max_relative=1e-6); + let mut expected = data::load_expected_values("soap-power-spectrum-positions-gradient.npy.gz"); + correct_factor(&mut expected, block.properties()); + assert_relative_eq!(array, &expected, max_relative=1e-6); let gradient = block.gradient("cell").unwrap(); let array = gradient.values().to_array(); - let expected = &data::load_expected_values("soap-power-spectrum-cell-gradient.npy.gz"); - assert_relative_eq!(array, expected, max_relative=1e-6); + let mut expected = data::load_expected_values("soap-power-spectrum-cell-gradient.npy.gz"); + correct_factor(&mut expected, block.properties()); + assert_relative_eq!(array, &expected, max_relative=1e-6); } fn sum_gradients(n_atoms: usize, gradients: TensorBlockRef<'_>) -> ArrayD { @@ -72,3 +76,15 @@ fn sum_gradients(n_atoms: usize, gradients: TensorBlockRef<'_>) -> ArrayD { sum } + + +// Add back the missing (-1)^l factor to the reference data +fn correct_factor(reference: &mut ndarray::ArrayD, properties: Labels) { + assert!(properties.names() == [ "species_neighbor_1", "species_neighbor_2", "l", "n1", "n2"]); + + let last_axis = reference.shape().len() - 1; + + for (&[_, _, l, _, _], mut column) in properties.iter_fixed_size().zip(reference.axis_iter_mut(Axis(last_axis))) { + column *= (-1.0_f64).powi(l.i32()); + } +} diff --git a/scripts/clean-python.sh b/scripts/clean-python.sh index 650c09747..992aa6c06 100755 --- a/scripts/clean-python.sh +++ b/scripts/clean-python.sh @@ -10,6 +10,7 @@ cd "$ROOT_DIR" rm -rf dist rm -rf build +rm -rf .coverage rm -rf python/rascaline-torch/dist rm -rf python/rascaline-torch/build diff --git a/setup.py b/setup.py index d26c8a651..4d7fca5f9 100644 --- a/setup.py +++ b/setup.py @@ -3,6 +3,7 @@ import shutil import subprocess import sys +import uuid from setuptools import Extension, setup from setuptools.command.bdist_egg import bdist_egg @@ -12,6 +13,7 @@ ROOT = os.path.realpath(os.path.dirname(__file__)) +RASCALINE_TORCH = os.path.join(ROOT, "python", "rascaline-torch") RASCALINE_BUILD_TYPE = os.environ.get("RASCALINE_BUILD_TYPE", "release") if RASCALINE_BUILD_TYPE not in ["debug", "release"]: @@ -250,9 +252,22 @@ def git_extra_version(): with open(os.path.join(ROOT, "AUTHORS")) as fd: authors = fd.read().splitlines() + extras_require = {} + if os.path.exists(RASCALINE_TORCH): + # we are building from a git checkout + + # add a random uuid to the file url to prevent pip from using a cached + # wheel for rascaline-torch, and force it to re-build from scratch + uuid = uuid.uuid4() + extras_require["torch"] = f"rascaline-torch @ file://{RASCALINE_TORCH}?{uuid}" + else: + # we are building from a sdist/installing from a wheel + extras_require["torch"] = "rascaline-torch >=0.1.0.dev0,<0.2.0" + setup( version=version, author=", ".join(authors), + extras_require=extras_require, ext_modules=[ # only declare the extension, it is built & copied as required by cmake # in the build_ext command diff --git a/tox.ini b/tox.ini index c4df209aa..0eec76844 100644 --- a/tox.ini +++ b/tox.ini @@ -1,11 +1,11 @@ [tox] -min_version = 4.0 +min_version = 4.9.0 # these are the default environments, i.e. the list of tests running when you # execute `tox` in the command-line without anything else envlist = lint - all-deps min-deps + all-deps docs-tests torch-tests @@ -25,15 +25,17 @@ metatensor-torch-requirement = metatensor-torch >=0.1.0,<0.2.0 build-single-wheel = --no-deps --no-build-isolation --check-build-dependencies - -commands = - # error if the user gives a wrong testenv name in `tox -e` - python -c "import sys; print('environement {env_name} does not exist'); sys.exit(1)" +test_options = + --cov={env_site_packages_dir}/rascaline \ + --cov-append \ + --cov-report= \ + --import-mode=append [testenv:build-rascaline] -# This environment is only used to build the wheels which are then re-used by -# all other environments requiring rascaline to be installed +description = + This environment is only used to build the wheels which are then re-used by + all other environments requiring rascaline to be installed passenv = * deps = setuptools @@ -45,34 +47,45 @@ commands = [testenv:all-deps] -# Run Python unit tests with all dependencies installed (ase & chemfiles are -# optional dependencies) +# note: platform_system can be "Linux","Darwin", or "Windows". +description = + Run Python unit tests with all dependencies installed (ase, pyscf, + and chemfiles are optional dependencies) deps = {[testenv]metatensor-core-requirement} ase chemfiles + metatensor-operations pytest pytest-cov - + scipy + sympy + pyscf;platform_system!="Windows" + wigners + # TODO: add mops once it becomes stable enough (and potentially supports windows) + #mops@git+https://github.com/lab-cosmo/mops ; platform_system!="Windows" commands = - pytest --cov={env_site_packages_dir}/rascaline --cov-report xml:.tox/coverage.xml --import-mode=append {posargs} + pytest {[testenv]test_options} {posargs} [testenv:min-deps] -# Run Python unit tests with the minimal dependencies installed +description = Run Python unit tests with the minimal dependencies installed deps = {[testenv]metatensor-core-requirement} pytest + pytest-cov commands = - pytest --import-mode=append {posargs} + pytest {[testenv]test_options} {posargs} [testenv:torch-tests] +description = Run Python unit tests using torch deps = {[testenv]metatensor-torch-requirement} pytest + pytest-cov numpy torch ase @@ -85,9 +98,10 @@ commands = # install rascaline-torch pip install . {[testenv]build-single-wheel} --force-reinstall # run the unit tests - pytest --import-mode=append --assert=plain {posargs} + pytest {[testenv]test_options} --assert=plain {posargs} [testenv:docs] +description = Build the package documentation. deps = -r docs/requirements.txt cmake @@ -105,10 +119,11 @@ commands = [testenv:docs-tests] -# this environement runs the doctests defined in any metatensor package +description = Runs the doctests defined in any metatensor package deps = {[testenv]metatensor-core-requirement} ase + pyscf;platform_system!="Windows" pytest commands = @@ -116,8 +131,9 @@ commands = [testenv:lint] -# lint the Python code with flake8 (code linter), black (code formatter), and -# isort (sorting of imports) +description = + lint the Python code with flake8 (code linter), black (code formatter), and isort + (sorting of imports) package = skip deps = black @@ -135,9 +151,8 @@ commands = [testenv:format] +description = Abuse tox to do actual formatting on all files. package = skip -# Abuse tox to do actual formatting. Users can call `tox -e format` to run -# formatting on all files deps = black blackdoc @@ -183,3 +198,19 @@ commands = [flake8] max_line_length = 88 extend-ignore = E203 + +[coverage:report] +skip_covered = True +show_missing = True +omit = + tests/.* + examples/.* + +[coverage:paths] +rascaline = + python/rascaline/rascaline + .tox/*/lib/python*/site-packages/rascaline + +rascaline_torch = + python/rascaline-torch/rascaline/torch + .tox/*/lib/python*/site-packages/rascaline/torch