diff --git a/docs/src/how-to/index.rst b/docs/src/how-to/index.rst index f54490aed..5d963acb8 100644 --- a/docs/src/how-to/index.rst +++ b/docs/src/how-to/index.rst @@ -14,6 +14,5 @@ are a total beginner, you can go to :ref:`userdoc-get-started` section. sample-selection property-selection keys-selection - le-basis splined-radial-integral long-range diff --git a/docs/src/how-to/le-basis.rst b/docs/src/how-to/le-basis.rst deleted file mode 100644 index b7a1e7b53..000000000 --- a/docs/src/how-to/le-basis.rst +++ /dev/null @@ -1,27 +0,0 @@ -.. _userdoc-how-to-le-basis: - -Laplacian eigenstate basis -========================== - -Temporarily deactivated example - -.. This examples shows how to calculate a spherical expansion using a Laplacian -.. eigenstate basis. - -.. .. tabs:: - -.. .. group-tab:: Python - -.. .. container:: sphx-glr-footer sphx-glr-footer-example - -.. .. container:: sphx-glr-download sphx-glr-download-python - -.. :download:`Download Python source code for this example: le-basis.py <../examples/le-basis.py>` - -.. .. container:: sphx-glr-download sphx-glr-download-jupyter - -.. :download:`Download Jupyter notebook for this example: le-basis.ipynb <../examples/le-basis.ipynb>` - -.. .. include:: ../examples/le-basis.rst -.. :start-after: start-body -.. :end-before: end-body diff --git a/docs/src/references/api/python/basis.rst b/docs/src/references/api/python/basis.rst index 2a1fef155..0e25bb8ce 100644 --- a/docs/src/references/api/python/basis.rst +++ b/docs/src/references/api/python/basis.rst @@ -7,6 +7,9 @@ Basis functions .. autoclass:: rascaline.basis.Explicit :show-inheritance: +.. autoclass:: rascaline.basis.LaplacianEigenstate + :show-inheritance: + .. autoclass:: rascaline.basis.ExpansionBasis :members: diff --git a/python/rascaline/examples/le-basis.py b/python/rascaline/examples/le-basis.py deleted file mode 100644 index 23882fc88..000000000 --- a/python/rascaline/examples/le-basis.py +++ /dev/null @@ -1,286 +0,0 @@ -""" -.. _userdoc-tutorials-le-basis: - -LE basis -======== - -Temporarily deactivated example - -""" - -# .. start-body - -# This example illustrates how to generate a spherical expansion using the Laplacian -# eigenstate (LE) basis (https://doi.org/10.1063/5.0124363), using two different basis -# truncations approaches. The basis can be truncated in the "traditional" way, using -# all values below a limit in the angular and radial direction; or using a "ragged -# truncation", where basis functions are selected according to an eigenvalue threshold. - -# The main ideas behind the LE basis are: - -# 1. use a basis of controllable *smoothness* (intended in the same sense as the -# smoothness of a low-pass-truncated Fourier expansion) -# 2. apply a "ragged truncation" strategy in which different angular channels are -# truncated at a different number of radial channels, so as to obtain more balanced -# smoothness level in the radial and angular direction, for a given number of basis -# functions. - -# Here we use :class:`rascaline.utils.SphericalBesselBasis` to create a spline of the -# radial integral corresponding to the LE basis. An detailed how-to guide how to -# construct radial integrals is given in :ref:`userdoc-how-to-splined-radial-integral`. -# """ - -# import ase.io -# import matplotlib.pyplot as plt -# import numpy as np -# from metatensor import Labels, TensorBlock, TensorMap - -# import rascaline - - -# # %% -# # -# # Let's start by using a traditional/square basis truncation. Here we will select all -# # basis functions with ``l <= max_angular`` and ``n < max_radial``. The basis -# # functions are the solution of a radial Laplacian eigenvalue problem (spherical -# # Bessel functions). - -# cutoff = 4.4 -# max_angular = 6 -# max_radial = 8 - -# # create a spliner for the SOAP radial integral, using delta functions for the atomic -# # density and spherical Bessel functions for the basis -# spliner = rascaline.utils.SoapSpliner( -# cutoff=cutoff, -# max_radial=max_radial, -# max_angular=max_angular, -# basis=rascaline.utils.SphericalBesselBasis( -# cutoff=cutoff, max_radial=max_radial, max_angular=max_angular -# ), -# density=rascaline.utils.DeltaDensity(), -# accuracy=1e-8, -# ) - -# # %% -# # -# # We can now plot the radial integral splines for a couple of functions. This gives an -# # idea of the smoothness of the different components - -# splined_basis = spliner.compute() -# grid = [p["position"] for p in splined_basis["TabulatedRadialIntegral"]["points"]] -# values = np.array( -# [ -# np.array(p["values"]["data"]).reshape(p["values"]["dim"]) -# for p in splined_basis["TabulatedRadialIntegral"]["points"] -# ] -# ) - -# plt.plot(grid, values[:, 1, 1], "b-", label="l=1, n=1") -# plt.plot(grid, values[:, 4, 1], "r-", label="l=4, n=1") -# plt.plot(grid, values[:, 1, 4], "g-", label="l=1, n=4") -# plt.plot(grid, values[:, 4, 4], "y-", label="l=4, n=4") -# plt.xlabel("$r$") -# plt.ylabel(r"$R_{nl}$") -# plt.legend() -# plt.show() - -# # %% -# # -# # We can use this spline basis in a :py:class:`SphericalExpansion` calculator to -# # evaluate spherical expansion coefficients. - -# calculator = rascaline.SphericalExpansion( -# cutoff=cutoff, -# max_radial=max_radial, -# max_angular=max_angular, -# center_atom_weight=1.0, -# radial_basis=splined_basis, -# atomic_gaussian_width=-1.0, # will not be used due to the delta density above -# cutoff_function={"ShiftedCosine": {"width": 0.5}}, -# ) - -# # %% -# # -# # This calculator defaults to the "traditional" basis function selection, so we have -# # the same maximal ``n`` value for all ``l``. - -# systems = ase.io.read("dataset.xyz", ":10") - -# descriptor = calculator.compute(systems) -# descriptor = descriptor.keys_to_properties("neighbor_type") -# descriptor = descriptor.keys_to_samples("center_type") - -# for key, block in descriptor.items(): -# n_max = np.max(block.properties["n"]) + 1 -# print(f"l = {key['o3_lambda']}, n_max = {n_max}") - -# # %% -# # -# # **Selecting basis with an eigenvalue threshold** -# # -# # Now we will calculate the same basis with an eigenvalue threshold. The idea is to -# # treat on the same footings the radial and angular dimension, and select all -# # functions with a mean Laplacian below a certain threshold. This is similar to the -# # common practice in plane-wave electronic-structure methods to use a kinetic energy -# # cutoff where :math:`k_x^2 + k_y^2 + k_z^2 < k_\text{max}^2` - -# eigenvalue_threshold = 20 - -# # %% -# # -# # Let's start by computing a lot of Laplacian eigenvalues, which are related to the -# # squares of the zeros of spherical Bessel functions. - -# l_max_large = 49 # just used to get the eigenvalues -# n_max_large = 50 # just used to get the eigenvalues - -# # compute the zeros of the spherical Bessel functions -# zeros_ln = rascaline.utils.SphericalBesselBasis.compute_zeros(l_max_large,n_max_large) - -# # %% -# # -# # We have a 50x50 array containing the position of the zero of the different spherical -# # Bessel functions, indexed by ``l`` and ``n``. - -# print("zeros_ln.shape =", zeros_ln.shape) -# print("zeros_ln =", zeros_ln[:3, :3]) - -# # calculate the Laplacian eigenvalues -# eigenvalues_ln = zeros_ln**2 / cutoff**2 - -# # %% -# # -# # We can now determine the set of ``l, n`` pairs to include all eigenvalues below the -# # threshold. - -# max_radial_by_angular = [] -# for ell in range(l_max_large + 1): -# # for each l, calculate how many radial basis functions we want to include -# max_radial = len(np.where(eigenvalues_ln[ell] < eigenvalue_threshold)[0]) -# max_radial_by_angular.append(max_radial) -# if max_radial_by_angular[-1] == 0: -# # all eigenvalues for this `l` are over the threshold -# max_radial_by_angular.pop() -# max_angular = ell - 1 -# break - -# # %% -# # -# # Comparing this eigenvalues threshold with the one based on a square selection, we -# # see that the eigenvalues threshold leads to a gradual decrease of ``max_radial`` -# # for high ``l`` values - -# square_max_angular = 10 -# square_max_radial = 4 -# plt.fill_between( -# [0, square_max_angular], -# [square_max_radial, square_max_radial], -# label=r"$l_\mathrm{max}$, $n_\mathrm{max}$ threshold " -# + f"({(square_max_angular + 1) * square_max_radial} functions)", -# color="gray", -# ) -# plt.fill_between( -# np.arange(max_angular + 1), -# max_radial_by_angular, -# label=f"Eigenvalues threshold ({sum(max_radial_by_angular)} functions)", -# alpha=0.5, -# ) -# plt.xlabel(r"$\ell$") -# plt.ylabel("n radial basis functions") -# plt.ylim(-0.5, max_radial_by_angular[0] + 0.5) -# plt.legend() -# plt.show() - -# # %% -# # -# # **Using a subset of basis functions with rascaline** -# # -# # We can tweak the default basis selection of rascaline by specifying a larger total -# # basis; and then only asking for a subset of properties to be computed. See -# # :ref:`userdoc-how-to-property-selection` for more details on properties selection. - -# # extract all the atomic types from our dataset -# all_atomic_types = list( -# np.unique(np.concatenate([system.numbers for system in systems])) -# ) - -# keys = [] -# blocks = [] -# for center_type in all_atomic_types: -# for neighbor_type in all_atomic_types: -# for ell in range(max_angular + 1): -# max_radial = max_radial_by_angular[ell] - -# keys.append([ell, 1, center_type, neighbor_type]) -# blocks.append( -# TensorBlock( -# values=np.zeros((0, max_radial)), -# samples=Labels.empty("_"), -# components=[], -# properties=Labels("n", np.arange(max_radial).reshape(-1, 1)), -# ) -# ) - -# selected_properties = TensorMap( -# keys=Labels( -# names=["o3_lambda", "o3_sigma", "center_type", "neighbor_type"], -# values=np.array(keys), -# ), -# blocks=blocks, -# ) - -# # %% -# # -# # With this, we can build a calculator and calculate the spherical expansion -# # coefficients - -# # the biggest max_radial will be for l=0 -# max_radial = max_radial_by_angular[0] - - -# # set up a spliner object for the spherical Bessel functions this radial basis will be -# # used to compute the spherical expansion -# spliner = rascaline.utils.SoapSpliner( -# cutoff=cutoff, -# max_radial=max_radial, -# max_angular=max_angular, -# basis=rascaline.utils.SphericalBesselBasis( -# cutoff=cutoff, max_radial=max_radial, max_angular=max_angular -# ), -# density=rascaline.utils.DeltaDensity(), -# accuracy=1e-8, -# ) - -# calculator = rascaline.SphericalExpansion( -# cutoff=cutoff, -# max_radial=max_radial, -# max_angular=max_angular, -# center_atom_weight=1.0, -# radial_basis=spliner.compute(), -# atomic_gaussian_width=-1.0, # will not be used due to the delta density above -# cutoff_function={"ShiftedCosine": {"width": 0.5}}, -# ) - -# # %% -# # -# # And check that we do get the expected Eigenvalues truncation for the calculated -# # features! - -# descriptor = calculator.compute( -# systems, -# # we tell the calculator to only compute the selected properties -# # (the desired set of (l,n) expansion coefficients -# selected_properties=selected_properties, -# ) - -# descriptor = descriptor.keys_to_properties("neighbor_type") -# descriptor = descriptor.keys_to_samples("center_type") - -# for key, block in descriptor.items(): -# n_max = np.max(block.properties["n"]) + 1 -# print(f"l = {key['o3_lambda']}, n_max = {n_max}") - -# # %% -# # -# # .. end-body diff --git a/python/rascaline/rascaline/basis.py b/python/rascaline/rascaline/basis.py index 809966d6c..d6409e874 100644 --- a/python/rascaline/rascaline/basis.py +++ b/python/rascaline/rascaline/basis.py @@ -54,8 +54,9 @@ def get_hypers(self): Return the native hyper parameters corresponding to this set of basis functions """ raise NotImplementedError( - f"this radial basis function ({self.__class__.__name__}) does not have " - "matching hyper parameters in the native calculators" + f"This radial basis function ({self.__class__.__name__}) does not have " + "matching hyper parameters in the native calculators. It should be used " + "through one of the spliner class instead of directly." ) @property @@ -406,8 +407,9 @@ def get_hypers(self): Return the native hyper parameters corresponding to this set of basis functions """ raise NotImplementedError( - f"this basis functions set ({self.__class__.__name__}) does not have " - "matching hyper parameters in the native calculators" + f"This basis functions set ({self.__class__.__name__}) does not have " + "matching hyper parameters in the native calculators. It should be used " + "through one of the spliner class instead of directly." ) @abc.abstractmethod @@ -439,6 +441,12 @@ def __init__( radial: RadialBasis, spline_accuracy: Optional[float] = 1e-8, ): + """ + :param max_angular: Largest angular channel to include in the basis + :param radial: radial basis to use for all angular channels + :param spline_accuracy: requested accuracy of the splined radial integrals, + defaults to 1e-8 + """ self.max_angular = int(max_angular) self.radial = radial @@ -467,7 +475,7 @@ def radial_basis(self, angular: int) -> RadialBasis: class Explicit(ExpansionBasis): - r""" + """ An expansion basis where combinations of radial and angular functions is picked explicitly. @@ -483,6 +491,12 @@ def __init__( by_angular: Dict[int, RadialBasis], spline_accuracy: Optional[float] = 1e-8, ): + """ + :param by_angular: definition of the radial basis for each angular channel to + include. + :param spline_accuracy: requested accuracy of the splined radial integrals, + defaults to 1e-8 + """ self.by_angular = by_angular if spline_accuracy is None: @@ -507,3 +521,101 @@ def angular_channels(self) -> List[int]: def radial_basis(self, angular: int) -> RadialBasis: return self.by_angular[angular] + + +class LaplacianEigenstate(ExpansionBasis): + """ + The Laplacian eigenstate basis, introduced in https://doi.org/10.1063/5.0124363, is + a set of basis functions for spherical expansion which is both *smooth* (in the same + sense as the smoothness of a low-pass-truncated Fourier expansion) and *ragged*, + using a different number of radial function for each angular channel. This is + intended to obtain a more balanced smoothness level in the radial and angular + direction for a given total number of basis functions. + + This expansion basis is not directly implemented in the native calculators, but is + intended to be used with the :py:class:`rascaline.splines.SoapSpliner` to create + splines of the radial integrals. + """ + + def __init__( + self, + *, + radius: float, + max_radial: int, + max_angular: Optional[int] = None, + spline_accuracy: Optional[float] = 1e-8, + ): + """ + :param radius: radius of the basis functions + :param max_radial: number of radial basis function for the ``L=0`` angular + channel. All other angular channels will have fewer radial basis functions. + :param max_angular: Truncate the set of radial functions at this angular + channel. If ``None``, this will be set to a high enough value to include all + basis functions with an Laplacian eigenvalue below the one for ``l=0, + n=max_radial``. + :param spline_accuracy: requested accuracy of the splined radial integrals, + defaults to 1e-8 + """ + self.radius = float(radius) + assert self.radius >= 0 + + self.max_radial = int(max_radial) + assert self.max_radial >= 0 + + if max_angular is None: + self.max_angular = self.max_radial + else: + self.max_angular = int(max_angular) + assert self.max_angular >= 0 + + if spline_accuracy is None: + self.spline_accuracy = None + else: + self.spline_accuracy = float(spline_accuracy) + assert self.spline_accuracy > 0 + + # compute the zeros of the spherical Bessel functions + zeros_ln = SphericalBessel._compute_zeros( + self.max_angular + 1, self.max_radial + 1 + ) + + # determine the eigenvalue cutoff + eigenvalues = zeros_ln**2 / self.radius**2 + max_eigenvalue = eigenvalues[0, max_radial] + + # find the actual `max_angular` if the user did not specify one by repeatedly + # increasing the size of `eigenvalues`, until we find an angular channel where + # all eigenvalues are above the cutoff. + if max_angular is None: + while eigenvalues[-1, 0] < max_eigenvalue: + self.max_angular += self.max_radial + zeros_ln = SphericalBessel._compute_zeros( + self.max_angular + 1, self.max_radial + 1 + ) + eigenvalues = zeros_ln**2 / self.radius**2 + + self.max_angular = len(np.where(eigenvalues[:, 0] <= max_eigenvalue)[0]) - 1 + assert self.max_angular >= 0 + + by_angular = {} + for angular in range(self.max_angular + 1): + max_radial = len(np.where(eigenvalues[angular, :] <= max_eigenvalue)[0]) - 1 + by_angular[angular] = SphericalBessel( + angular_channel=angular, + max_radial=max_radial, + radius=self.radius, + ) + + self._explicit = Explicit( + by_angular=by_angular, + spline_accuracy=self.spline_accuracy, + ) + + def get_hypers(self): + return self._explicit.get_hypers() + + def angular_channels(self) -> List[int]: + return self._explicit.angular_channels() + + def radial_basis(self, angular: int) -> RadialBasis: + return self._explicit.radial_basis(angular) diff --git a/python/rascaline/rascaline/density.py b/python/rascaline/rascaline/density.py index 3e09a4c9d..c4660d711 100644 --- a/python/rascaline/rascaline/density.py +++ b/python/rascaline/rascaline/density.py @@ -166,8 +166,9 @@ def get_hypers(self): Return the native hyper parameters corresponding to this atomic density """ raise NotImplementedError( - f"this density ({self.__class__.__name__}) does not have matching " - "hyper parameters in the native calculators" + f"This density ({self.__class__.__name__}) does not have matching " + "hyper parameters in the native calculators. It should be used " + "through one of the spliner class instead of directly." ) def _rascaline_hypers(self): diff --git a/python/rascaline/tests/basis.py b/python/rascaline/tests/basis.py index 7e462fd13..b899338b5 100644 --- a/python/rascaline/tests/basis.py +++ b/python/rascaline/tests/basis.py @@ -1,7 +1,14 @@ -import numpy as np +import numpy as np # noqa: I001 import pytest -from rascaline.basis import Gto, Monomials, RadialBasis, SphericalBessel +from rascaline.utils import hypers_to_json +from rascaline.basis import ( + Gto, + LaplacianEigenstate, + Monomials, + RadialBasis, + SphericalBessel, +) pytest.importorskip("scipy") @@ -86,3 +93,25 @@ def test_derivative(analytical_basis): analytical_basis.compute_primitive(positions, n, derivative=True), atol=1e-6, ) + + +def test_le_basis(): + basis = LaplacianEigenstate(max_radial=4, radius=3.4) + assert basis.max_angular == 10 + assert basis.angular_channels() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + + for angular in basis.angular_channels(): + radial = basis.radial_basis(angular) + assert radial.max_radial <= 4 + + basis = LaplacianEigenstate(max_radial=4, max_angular=4, radius=3.4) + assert basis.max_angular == 4 + assert basis.angular_channels() == [0, 1, 2, 3, 4] + + message = ( + "This radial basis function \\(SphericalBessel\\) does not have matching " + "hyper parameters in the native calculators. It should be used through one of " + "the spliner class instead of directly" + ) + with pytest.raises(NotImplementedError, match=message): + hypers_to_json(basis.get_hypers()) diff --git a/python/rascaline/tests/calculators/hypers.py b/python/rascaline/tests/calculators/hypers.py index dcd42be51..ef6705750 100644 --- a/python/rascaline/tests/calculators/hypers.py +++ b/python/rascaline/tests/calculators/hypers.py @@ -176,8 +176,9 @@ def compute(self, positions, derivative): pass message = ( - "this density \\(MyCustomDensity\\) does not have matching hyper " - "parameters in the native calculators" + "This density \\(MyCustomDensity\\) does not have matching hyper " + "parameters in the native calculators. It should be used " + "through one of the spliner class instead of directly" ) with pytest.raises(NotImplementedError, match=message): SphericalExpansion( @@ -197,8 +198,9 @@ def radial_basis(self, angular) -> rascaline.basis.RadialBasis: return rascaline.basis.Gto(width=0.3) message = ( - "this basis functions set \\(MyCustomExpansionBasis\\) does not have " - "matching hyper parameters in the native calculators" + "This basis functions set \\(MyCustomExpansionBasis\\) does not have " + "matching hyper parameters in the native calculators. It should be used " + "through one of the spliner class instead of directly" ) with pytest.raises(NotImplementedError, match=message): SphericalExpansion( @@ -215,8 +217,9 @@ def compute_primitive(self, positions, n, derivative): pass message = ( - "this radial basis function \\(MyCustomRadialBasis\\) does not have " - "matching hyper parameters in the native calculators" + "This radial basis function \\(MyCustomRadialBasis\\) does not have " + "matching hyper parameters in the native calculators. It should be used " + "through one of the spliner class instead of directly" ) with pytest.raises(NotImplementedError, match=message): SphericalExpansion(