From 9ccb4d8dffe6f8378b106a778572109b4ffe5433 Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Thu, 7 Nov 2024 10:59:06 +0100 Subject: [PATCH] Update all examples to use the new hypers --- docs/src/how-to/long-range.rst | 58 +- docs/src/how-to/splined-radial-integral.rst | 28 +- docs/src/references/api/python/misc.rst | 2 + python/rascaline/examples/.gitignore | 1 + .../examples/long-range-descriptor.py | 614 ++++++++---------- .../examples/splined-radial-integral.py | 357 +++++----- python/rascaline/rascaline/basis.py | 1 + python/rascaline/rascaline/calculators.py | 36 +- python/rascaline/rascaline/utils/__init__.py | 6 +- .../rascaline/utils/hypers/__init__.py | 20 + 10 files changed, 556 insertions(+), 567 deletions(-) create mode 100644 python/rascaline/examples/.gitignore diff --git a/docs/src/how-to/long-range.rst b/docs/src/how-to/long-range.rst index 5d4bda09d..622baea77 100644 --- a/docs/src/how-to/long-range.rst +++ b/docs/src/how-to/long-range.rst @@ -3,44 +3,44 @@ Long-range only LODE descriptor =============================== -Temporarily deactivated example +The :py:class:`LodeSphericalExpansion ` allows +the calculation of a descriptor that includes all atoms within the system and +projects them onto a spherical expansion/ fingerprint within a given ``cutoff``. +This is very useful if long-range interactions between atoms are important to +describe the physics and chemistry of a collection of atoms. However, as stated +the descriptor contains **ALL** atoms of the system and sometimes it can be +desired to only have a long-range/exterior only descriptor that only includes +the atoms outside a given cutoff. Sometimes there descriptors are also denoted +by far-field descriptors. -.. The :py:class:`LodeSphericalExpansion ` allows the -.. calculation of a descriptor that includes all atoms within the system and projects them -.. onto a spherical expansion/ fingerprint within a given ``cutoff``. This is very useful -.. if long-range interactions between atoms are important to describe the physics and -.. chemistry of a collection of atoms. However, as stated the descriptor contains **ALL** -.. atoms of the system and sometimes it can be desired to only have a long-range/exterior -.. only descriptor that only includes the atoms outside a given cutoff. Sometimes there -.. descriptors are also denoted by far-field descriptors. +A long range only descriptor can be particular useful when one already has a +good descriptor for the short-range density like (SOAP) and the long-range +descriptor (far field) should contain different information from what the +short-range descriptor already offers. -.. A long range only descriptor can be particular useful when one already has a good -.. descriptor for the short-range density like (SOAP) and the long-range descriptor (far -.. field) should contain different information from what the short-range descriptor already -.. offers. +Such descriptor can be constructed within `rascaline` as sketched by the image +below. -.. Such descriptor can be constructed within `rascaline` as sketched by the image below. +.. figure:: ../../static/images/long-range-descriptor.* + :align: center -.. .. figure:: ../../static/images/long-range-descriptor.* -.. :align: center +In this example will construct such a descriptor using the :ref:`radial integral +splining ` tools of `rascaline`. -.. In this example will construct such a descriptor using the :ref:`radial integral -.. splining ` tools of `rascaline`. +.. tabs:: -.. .. tabs:: + .. group-tab:: Python -.. .. group-tab:: Python + .. container:: sphx-glr-footer sphx-glr-footer-example -.. .. container:: sphx-glr-footer sphx-glr-footer-example + .. container:: sphx-glr-download sphx-glr-download-python -.. .. container:: sphx-glr-download sphx-glr-download-python + :download:`Download Python source code for this example: long-range-descriptor.py <../examples/long-range-descriptor.py>` -.. :download:`Download Python source code for this example: long-range-descriptor.py <../examples/long-range-descriptor.py>` + .. container:: sphx-glr-download sphx-glr-download-jupyter -.. .. container:: sphx-glr-download sphx-glr-download-jupyter + :download:`Download Jupyter notebook for this example: long-range-descriptor.ipynb <../examples/long-range-descriptor.ipynb>` -.. :download:`Download Jupyter notebook for this example: long-range-descriptor.ipynb <../examples/long-range-descriptor.ipynb>` - -.. .. include:: ../examples/long-range-descriptor.rst -.. :start-after: start-body -.. :end-before: end-body + .. include:: ../examples/long-range-descriptor.rst + :start-after: start-body + :end-before: end-body diff --git a/docs/src/how-to/splined-radial-integral.rst b/docs/src/how-to/splined-radial-integral.rst index 201dfa243..bce9c40cd 100644 --- a/docs/src/how-to/splined-radial-integral.rst +++ b/docs/src/how-to/splined-radial-integral.rst @@ -3,26 +3,24 @@ Splined radial integral ======================= -Temporarily deactivated example +This examples shows how to feed custom radial integrals (as splines) to the Rust +calculators that use radial integrals: the SOAP and LODE spherical expansions, +and any other calculator based on these. -.. This examples shows how to feed custom radial integrals (as splines) to the Rust -.. calculators that use radial integrals: the SOAP and LODE spherical expansions, -.. and any other calculator based on these. +.. tabs:: -.. .. tabs:: + .. group-tab:: Python -.. .. group-tab:: Python + .. container:: sphx-glr-footer sphx-glr-footer-example -.. .. container:: sphx-glr-footer sphx-glr-footer-example + .. container:: sphx-glr-download sphx-glr-download-python -.. .. container:: sphx-glr-download sphx-glr-download-python + :download:`Download Python source code for this example: tabulated.py <../examples/splined-radial-integral.py>` -.. :download:`Download Python source code for this example: tabulated.py <../examples/splined-radial-integral.py>` + .. container:: sphx-glr-download sphx-glr-download-jupyter -.. .. container:: sphx-glr-download sphx-glr-download-jupyter + :download:`Download Jupyter notebook for this example: tabulated.ipynb <../examples/splined-radial-integral.ipynb>` -.. :download:`Download Jupyter notebook for this example: tabulated.ipynb <../examples/splined-radial-integral.ipynb>` - -.. .. include:: ../examples/splined-radial-integral.rst -.. :start-after: start-body -.. :end-before: end-body + .. include:: ../examples/splined-radial-integral.rst + :start-after: start-body + :end-before: end-body diff --git a/docs/src/references/api/python/misc.rst b/docs/src/references/api/python/misc.rst index 6706b5509..81e750346 100644 --- a/docs/src/references/api/python/misc.rst +++ b/docs/src/references/api/python/misc.rst @@ -15,3 +15,5 @@ Miscellaneous :undoc-members: .. autofunction:: rascaline.convert_hypers + +.. autofunction:: rascaline.utils.hypers_to_json diff --git a/python/rascaline/examples/.gitignore b/python/rascaline/examples/.gitignore new file mode 100644 index 000000000..f879f70bf --- /dev/null +++ b/python/rascaline/examples/.gitignore @@ -0,0 +1 @@ +splined-hypers.json diff --git a/python/rascaline/examples/long-range-descriptor.py b/python/rascaline/examples/long-range-descriptor.py index 5855fd83a..8723cc769 100644 --- a/python/rascaline/examples/long-range-descriptor.py +++ b/python/rascaline/examples/long-range-descriptor.py @@ -4,438 +4,388 @@ Long-range only LODE descriptor =============================== -Temporarily deactivated example -""" - -# .. start-body - -# We start the example by loading the required packages -# """ - -# # %% - -# import ase -# import ase.visualize.plot -# import matplotlib.pyplot as plt -# import numpy as np -# from ase.build import molecule -# from metatensor import LabelsEntry, TensorMap - -# from rascaline import LodeSphericalExpansion, SphericalExpansion -# from rascaline.utils import ( -# GaussianDensity, -# LodeDensity, -# LodeSpliner, -# MonomialBasis, -# SoapSpliner, -# ) - - -# # %% -# # -# # **Single water molecule (short range) system** -# # -# # Our first test system is a single water molecule with a :math:`15\,\mathrm{Å}` -# # vacuum layer around it. - +.. start-body -# atoms = molecule("H2O", vacuum=15, pbc=True) - -# # %% -# # We choose a ``cutoff`` for the projection of the spherical expansion and the -# # neighbor search of the real space spherical expansion. - -# cutoff = 3 - -# # %% -# # We can use ase's visualization tools to plot the system and draw a gray circle to -# # indicate the ``cutoff`` radius. - -# fig, ax = plt.subplots() - -# ase.visualize.plot.plot_atoms(atoms) - -# cutoff_circle = plt.Circle( -# xy=atoms[0].position[:2], -# radius=cutoff, -# color="gray", -# ls="dashed", -# fill=False, -# ) -# ax.add_patch(cutoff_circle) - -# ax.set_xlabel("Å") -# ax.set_ylabel("Å") - -# fig.show() +We start the example by loading the required packages +""" +# %% -# # %% -# # -# # As you can see, for a single water molecule, the ``cutoff`` includes all atoms of -# # the system. The combination of the test system and the ``cutoff`` aims to -# # demonstrate that the full atomic fingerprint is contained within the ``cutoff``. -# # By later subtracting the short-range density from the LODE density, we will observe -# # that the difference between them is almost zero, indicating that a single water -# # molecule is a short-range system. -# # -# # To start this construction we choose a high potential exponent to emulate the -# # rapidly decaying LODE density and mimic the polar-polar interactions of water. +import ase +import ase.visualize.plot +import matplotlib.pyplot as plt +import numpy as np +from ase.build import molecule +from metatensor import LabelsEntry, TensorMap +import rascaline +from rascaline import LodeSphericalExpansion, SphericalExpansion -# potential_exponent = 3 +# %% +# +# **Single water molecule (short range) system** +# +# Our first test system is a single water molecule with a :math:`15\,\mathrm{Å}` +# vacuum layer around it. -# # %% -# # We now define some typical hyperparameters to compute the spherical expansions. -# max_radial = 5 -# max_angular = 1 -# atomic_gaussian_width = 1.2 -# center_atom_weight = 1.0 +atoms = molecule("H2O", vacuum=15, pbc=True) +# %% +# We choose a ``cutoff`` for the projection of the spherical expansion and the +# neighbor search of the real space spherical expansion. -# # %% -# # We choose a relatively low spline accuracy (default is ``1e-8``) to achieve quick -# # computation of the spline points. You can increase the spline accuracy if required, -# # but be aware that the time to compute these points will increase significantly! +cutoff = 3 +# %% +# We can use ase's visualization tools to plot the system and draw a gray circle to +# indicate the ``cutoff`` radius. -# spline_accuracy = 1e-2 +fig, ax = plt.subplots() +ase.visualize.plot.plot_atoms(atoms) -# # %% -# # As a projection basis, we don't use the usual :py:class:`GtoBasis -# # ` which is commonly used for short range descriptors. -# # Instead, we select the :py:class:`MonomialBasis ` -# # which is the optimal radial basis for the LODE descriptor as discussed in -# # `Huguenin-Dumittan et al. `_ +cutoff_circle = plt.Circle( + xy=atoms[0].position[:2], + radius=cutoff, + color="gray", + ls="dashed", + fill=False, +) +ax.add_patch(cutoff_circle) +ax.set_xlabel("Å") +ax.set_ylabel("Å") -# basis = MonomialBasis(cutoff=cutoff) +fig.show() -# # %% -# # For the density, we choose a smeared power law as used in LODE, which does not decay -# # exponentially like a :py:class:`Gaussian density ` -# # and is therefore suited to describe long-range interactions between atoms. +# %% +# +# As you can see, for a single water molecule, the ``cutoff`` includes all atoms of +# the system. The combination of the test system and the ``cutoff`` aims to +# demonstrate that the full atomic fingerprint is contained within the ``cutoff``. +# By later subtracting the short-range density from the LODE density, we will observe +# that the difference between them is almost zero, indicating that a single water +# molecule is a short-range system. -# density = LodeDensity( -# atomic_gaussian_width=atomic_gaussian_width, -# potential_exponent=potential_exponent, -# ) +# %% +# +# For the density, we choose a smeared power law as used in LODE, which does not decay +# exponentially like a :py:class:`Gaussian ` density and is +# therefore suited to describe long-range interactions between atoms. -# # %% -# # To visualize this we plot ``density`` together with a Gaussian density -# # (``gaussian_density``) with the same ``atomic_gaussian_width`` in a log-log plot. +density = rascaline.density.SmearedPowerLaw(smearing=1.2, exponent=3) -# radial_positions = np.geomspace(1e-5, 10, num=1000) -# gaussian_density = GaussianDensity(atomic_gaussian_width=atomic_gaussian_width) -# plt.plot(radial_positions, density.compute(radial_positions), label="LodeDensity") -# plt.plot( -# radial_positions, -# gaussian_density.compute(radial_positions), -# label="GaussianDensity", -# ) +# %% +# +# To visualize this we plot ``density`` together with a Gaussian density +# (``gaussian_density``) with the same ``width`` in a log-log plot. +radial_positions = np.geomspace(1e-5, 10, num=1000) +gaussian_density = rascaline.density.Gaussian(width=density.smearing) -# positions_indicator = np.array([3.0, 8.0]) -# plt.plot( -# positions_indicator, -# 2 * positions_indicator**-potential_exponent, -# c="k", -# label=f"p={potential_exponent}", -# ) +plt.plot( + radial_positions, + density.compute(radial_positions, derivative=False), + label="SmearedPowerLaw", +) +plt.plot( + radial_positions, + gaussian_density.compute(radial_positions, derivative=False), + label="Gaussian", +) -# plt.legend() -# plt.xlim(1e-1, 10) -# plt.ylim(1e-3, 5e-1) +positions_indicator = np.array([3.0, 8.0]) +plt.plot( + positions_indicator, + 2 * positions_indicator ** (-density.exponent), + c="k", + label=f"exponent={density.exponent}", +) -# plt.xlabel("radial positions / Å") -# plt.ylabel("atomic density") +plt.legend() -# plt.xscale("log") -# plt.yscale("log") +plt.xlim(1e-1, 10) +plt.ylim(1e-3, 5e-1) -# # %% -# # We see that the ``LodeDensity`` decays with a power law of 3, which is the potential -# # exponent we picked above, wile the :py:class:`Gaussian density -# # ` decays exponentially and is therefore not suited -# # for long-range descriptors. -# # -# # We now have all building blocks to construct the spline points for the real and -# # Fourier space spherical expansions. +plt.xlabel("radial positions / Å") +plt.ylabel("atomic density") +plt.xscale("log") +plt.yscale("log") -# real_space_splines = SoapSpliner( -# cutoff=cutoff, -# max_radial=max_radial, -# max_angular=max_angular, -# basis=basis, -# density=density, -# accuracy=spline_accuracy, -# ).compute() +# %% +# +# We see that the ``SmearedPowerLaw`` decays with a power law of 3, which is the +# potential exponent we picked above, wile the :py:class:`Gaussian +# ` density decays exponentially and is therefore not suited +# for long-range descriptors. +# %% +# +# As a projection basis, we don't use the usual :py:class:`Gto ` +# which is commonly used for short range descriptors. Instead, we select the +# :py:class:`Monomials ` which is the optimal radial basis +# for the LODE descriptor as discussed in `Huguenin-Dumittan et al. +# `_ -# # This value gives good convergences for the Fourier space version -# k_cutoff = 1.2 * np.pi / atomic_gaussian_width +by_angular = {} +for angular in range(2): + by_angular[angular] = rascaline.basis.Monomials( + radius=cutoff, angular_channel=angular, max_radial=4 + ) -# fourier_space_splines = LodeSpliner( -# k_cutoff=k_cutoff, -# max_radial=max_radial, -# max_angular=max_angular, -# basis=basis, -# density=density, -# accuracy=spline_accuracy, -# ).compute() +basis = rascaline.basis.Explicit( + by_angular=by_angular, + # We choose a relatively low spline accuracy (default is ``1e-8``) to achieve quick + # computation of the spline points. You can increase the spline accuracy if + # required, but be aware that the time to compute these points will increase + # significantly! + spline_accuracy=1e-2, +) -# # %% -# # .. note:: -# # You might want to save the spline points using :py:func:`json.dump` to a file and -# # load them with :py:func:`json.load` later without recalculating them. Saving them -# # is especially useful if the spline calculations are expensive, i.e., if you -# # increase the ``spline_accuracy``. -# # -# # With the spline points ready, we now compute the real space spherical expansion +# %% +# +# We now have all building blocks to construct the spline for the real and +# Fourier space spherical expansions. -# real_space_calculator = SphericalExpansion( -# cutoff=cutoff, -# max_radial=max_radial, -# max_angular=max_angular, -# atomic_gaussian_width=atomic_gaussian_width, -# radial_basis=real_space_splines, -# center_atom_weight=center_atom_weight, -# cutoff_function={"Step": {}}, -# radial_scaling=None, -# ) +real_space_spliner = rascaline.splines.SoapSpliner( + # We don't use a ``smoothing`` function in the cutoff or a ``radial_scaling`` in the + # density to ensure the correct construction of the long-range only descriptor + cutoff=rascaline.cutoff.Cutoff(radius=cutoff, smoothing=None), + basis=basis, + density=density, +) +real_space_hypers = real_space_spliner.get_hypers() -# real_space_expansion = real_space_calculator.compute(atoms) +fourier_space_spliner = rascaline.splines.LodeSpliner( + basis=basis, + density=density, +) +fourier_space_hypers = fourier_space_spliner.get_hypers() +# %% +# +# With the splines ready, we now compute the two spherical expansions -# # %% -# # where we don't use a smoothing ``cutoff_function`` or a ``radial_scaling`` to ensure -# # the correct construction of the long-range only descriptor. Next, we compute the -# # Fourier Space / LODE spherical expansion +real_space_calculator = SphericalExpansion(**real_space_hypers) +real_space_expansion = real_space_calculator.compute(atoms) -# fourier_space_calculator = LodeSphericalExpansion( -# cutoff=cutoff, -# max_radial=max_radial, -# max_angular=max_angular, -# atomic_gaussian_width=atomic_gaussian_width, -# center_atom_weight=center_atom_weight, -# potential_exponent=potential_exponent, -# radial_basis=fourier_space_splines, -# k_cutoff=k_cutoff, -# ) -# fourier_space_expansion = fourier_space_calculator.compute(atoms) +fourier_space_calculator = LodeSphericalExpansion(**fourier_space_hypers) +fourier_space_expansion = fourier_space_calculator.compute(atoms) -# # %% -# # As described in the beginning, we now subtract the real space LODE contributions -# # from Fourier space to obtain a descriptor that only contains the contributions from -# # atoms outside of the ``cutoff``. +# %% +# As described in the beginning, we now subtract the real space LODE contributions +# from Fourier space to obtain a descriptor that only contains the contributions from +# atoms outside of the ``cutoff``. -# subtracted_expansion = fourier_space_expansion - real_space_expansion +delta_expansion = fourier_space_expansion - real_space_expansion -# # %% You can now use the ``subtracted_expansion`` as a long-range descriptor in -# # combination with a short-range descriptor like -# # :py:class:`rascaline.SphericalExpansion` for your machine learning models. We now -# # verify that for our test ``atoms`` the LODE spherical expansion only contains -# # short-range contributions. To demonstrate this, we densify the -# # :py:class:`metatensor.TensorMap` to have only one block per ``"center_type"`` and -# # visualize our result. Since we have to perform the densify operation several times -# # in this how-to, we define a helper function ``densify_tensormap``. +# %% +# +# You can now use the ``delta_expansion`` as a purely long-range descriptor in +# combination with a short-range descriptor like +# :py:class:`rascaline.SphericalExpansion` for your machine learning models. +# +# We now verify that for our test ``atoms`` the LODE spherical expansion only contains +# short-range contributions. To demonstrate this, we densify the +# :py:class:`metatensor.TensorMap` to have only one block per ``"center_type"`` and +# visualize our result. Since we have to perform the densify operation several times in +# this how-to, we define a helper function ``densify_tensormap``. -# def densify_tensormap(tensor: TensorMap) -> TensorMap: -# dense_tensor = tensor.components_to_properties("o3_mu") -# dense_tensor = dense_tensor.keys_to_samples("neighbor_type") -# dense_tensor = dense_tensor.keys_to_properties(["o3_lambda", "o3_sigma"]) +def densify_tensormap(tensor: TensorMap) -> TensorMap: + dense_tensor = tensor.components_to_properties("o3_mu") + dense_tensor = dense_tensor.keys_to_samples("neighbor_type") + dense_tensor = dense_tensor.keys_to_properties(["o3_lambda", "o3_sigma"]) -# return dense_tensor + return dense_tensor -# # %% -# # We apply the function to the Fourier space spherical expansion -# # ``fourier_space_expansion`` and ``subtracted_expansion``. +# %% +# We apply the function to the Fourier space spherical expansion +# ``fourier_space_expansion`` and ``subtracted_expansion``. -# fourier_space_expansion = densify_tensormap(fourier_space_expansion) -# subtracted_expansion = densify_tensormap(subtracted_expansion) +fourier_space_expansion = densify_tensormap(fourier_space_expansion) +delta_expansion = densify_tensormap(delta_expansion) -# # %% -# # Finally, we plot the values of each block for the Fourier Space spherical expansion -# # in the upper panel and the difference between the Fourier Space and the real space -# # in the lower panel. And since we will do this plot several times we again define a -# # small plot function to help us +# %% +# Finally, we plot the values of each block for the Fourier Space spherical expansion +# in the upper panel and the difference between the Fourier Space and the real space +# in the lower panel. And since we will do this plot several times we again define a +# small plot function to help us -# def plot_value_comparison( -# key: LabelsEntry, -# fourier_space_expansion: TensorMap, -# subtracted_expansion: TensorMap, -# ): -# fig, ax = plt.subplots(2, layout="tight") +def plot_value_comparison( + key: LabelsEntry, + fourier_space_expansion: TensorMap, + subtracted_expansion: TensorMap, +): + fig, ax = plt.subplots(2, layout="tight") -# values_subtracted = subtracted_expansion[key].values -# values_fourier_space = fourier_space_expansion[key].values + values_subtracted = subtracted_expansion[key].values + values_fourier_space = fourier_space_expansion[key].values -# ax[0].set_title(f"center_type={key.values[0]}\n Fourier space sph. expansion") -# im = ax[0].matshow(values_fourier_space, vmin=-0.25, vmax=0.5) -# ax[0].set_ylabel("sample index") + ax[0].set_title(f"center_type={key.values[0]}\n Fourier space sph. expansion") + im = ax[0].matshow(values_fourier_space, vmin=-0.25, vmax=0.5) + ax[0].set_ylabel("sample index") -# ax[1].set_title("Difference between Fourier and real space sph. expansion") -# ax[1].matshow(values_subtracted, vmin=-0.25, vmax=0.5) -# ax[1].set_ylabel("sample index") -# ax[1].set_xlabel("property index") + ax[1].set_title("Difference between Fourier and real space sph. expansion") + ax[1].matshow(values_subtracted, vmin=-0.25, vmax=0.5) + ax[1].set_ylabel("sample index") + ax[1].set_xlabel("property index") -# fig.colorbar(im, ax=ax[0], orientation="horizontal", fraction=0.1, label="values") + fig.colorbar(im, ax=ax[0], orientation="horizontal", fraction=0.1, label="values") -# # %% -# # We first plot the values of the TensorMaps for center_type=1 (hydrogen) +# %% +# We first plot the values of the TensorMaps for center_type=1 (hydrogen) -# plot_value_comparison( -# fourier_space_expansion.keys[0], fourier_space_expansion, subtracted_expansion -# ) +plot_value_comparison( + fourier_space_expansion.keys[0], fourier_space_expansion, delta_expansion +) -# # %% -# # and for center_type=8 (oxygen) +# %% +# and for center_type=8 (oxygen) -# plot_value_comparison( -# fourier_space_expansion.keys[1], fourier_space_expansion, subtracted_expansion -# ) +plot_value_comparison( + fourier_space_expansion.keys[1], fourier_space_expansion, delta_expansion +) -# # %% -# # The plot shows that the spherical expansion for the Fourier space is non-zero while -# # the difference between the two expansions is very small. -# # -# # .. warning:: -# # Small residual values may stems from the contribution of the periodic images. You -# # can verify and reduce those contributions by either increasing the cell and/or -# # increase the ``potential_exponent``. -# # -# # **Two water molecule (long range) system** -# # -# # We now add a second water molecule shifted by :math:`3\,\mathrm{Å}` in each -# # direction from our first water molecule to show that such a system has non -# # negligible long range effects. +# %% +# The plot shows that the spherical expansion for the Fourier space is non-zero while +# the difference between the two expansions is very small. +# +# .. warning:: +# Small residual values may stems from the contribution of the periodic images. You +# can verify and reduce those contributions by either increasing the cell and/or +# increase the ``potential_exponent``. +# +# **Two water molecule (long range) system** +# +# We now add a second water molecule shifted by :math:`3\,\mathrm{Å}` in each +# direction from our first water molecule to show that such a system has non +# negligible long range effects. -# atoms_shifted = molecule("H2O", vacuum=10, pbc=True) -# atoms_shifted.positions = atoms.positions + 3 +atoms_shifted = molecule("H2O", vacuum=10, pbc=True) +atoms_shifted.positions = atoms.positions + 3 -# atoms_long_range = atoms + atoms_shifted +atoms_long_range = atoms + atoms_shifted -# fig, ax = plt.subplots() +fig, ax = plt.subplots() -# ase.visualize.plot.plot_atoms(atoms_long_range, ax=ax) +ase.visualize.plot.plot_atoms(atoms_long_range, ax=ax) -# cutoff_circle = plt.Circle( -# xy=atoms[0].position[1:], -# radius=cutoff, -# color="gray", -# ls="dashed", -# fill=False, -# ) +cutoff_circle = plt.Circle( + xy=atoms[0].position[1:], + radius=cutoff, + color="gray", + ls="dashed", + fill=False, +) -# cutoff_circle_shifted = plt.Circle( -# xy=atoms_shifted[0].position[1:], -# radius=cutoff, -# color="gray", -# ls="dashed", -# fill=False, -# ) +cutoff_circle_shifted = plt.Circle( + xy=atoms_shifted[0].position[1:], + radius=cutoff, + color="gray", + ls="dashed", + fill=False, +) -# ax.add_patch(cutoff_circle) -# ax.add_patch(cutoff_circle_shifted) +ax.add_patch(cutoff_circle) +ax.add_patch(cutoff_circle_shifted) -# ax.set_xlabel("Å") -# ax.set_ylabel("Å") +ax.set_xlabel("Å") +ax.set_ylabel("Å") -# fig.show() +fig.show() -# # %% -# # As you can see, the ``cutoff`` radii of the two molecules are completely disjoint. -# # Therefore, a short-range model will not able to describe the intermolecular -# # interactions between our two molecules. To verify we now again create a long-range -# # only descriptor for this system. We use the already defined -# # ``real_space_expansion_long_range`` and ``fourier_space_expansion_long_range`` +# %% +# As you can see, the ``cutoff`` radii of the two molecules are completely disjoint. +# Therefore, a short-range model will not able to describe the intermolecular +# interactions between our two molecules. To verify we now again create a long-range +# only descriptor for this system. We use the already defined +# ``real_space_expansion_long_range`` and ``fourier_space_expansion_long_range`` -# real_space_expansion_long_range = real_space_calculator.compute(atoms_long_range) -# fourier_space_expansion_long_range=fourier_space_calculator.compute(atoms_long_range) +real_space_expansion_long_range = real_space_calculator.compute(atoms_long_range) +fourier_space_expansion_long_range = fourier_space_calculator.compute(atoms_long_range) -# # %% -# # We now first verify that the contribution from the short-range descriptors is the -# # same as for a single water molecule. Exemplarily, we compare only the first -# # (Hydrogen) block of each tensor. +# %% +# We now first verify that the contribution from the short-range descriptors is the +# same as for a single water molecule. Exemplarily, we compare only the first +# (Hydrogen) block of each tensor. -# print("Single water real space spherical expansion") -# print(np.round(real_space_expansion[1].values, 3)) +print("Single water real space spherical expansion") +print(np.round(real_space_expansion[1].values, 3)) -# print("\nTwo water real space spherical expansion") -# print(np.round(real_space_expansion_long_range[1].values, 3)) +print("\nTwo water real space spherical expansion") +print(np.round(real_space_expansion_long_range[1].values, 3)) -# # %% -# # Since the values of the block are the same, we can conclude that there is no -# # information shared between the two molecules and that the short-range descriptor is -# # not able to distinguish the system with only one or two water molecules. Note that -# # the different number of `samples` in ``real_space_expansion_long_range`` reflects -# # the fact that the second system has more atoms then the first. -# # -# # As above, we construct a long-range only descriptor and densify the result for -# # plotting the values. +# %% +# Since the values of the block are the same, we can conclude that there is no +# information shared between the two molecules and that the short-range descriptor is +# not able to distinguish the system with only one or two water molecules. Note that +# the different number of `samples` in ``real_space_expansion_long_range`` reflects +# the fact that the second system has more atoms then the first. +# +# As above, we construct a long-range only descriptor and densify the result for +# plotting the values. -# subtracted_expansion_long_range = ( -# fourier_space_expansion_long_range - real_space_expansion_long_range -# ) +delta_expansion_long_range = ( + fourier_space_expansion_long_range - real_space_expansion_long_range +) -# fourier_space_expansion_long_range = densify_tensormap( -# fourier_space_expansion_long_range -# ) -# subtracted_expansion_long_range = densify_tensormap(subtracted_expansion_long_range) +fourier_space_expansion_long_range = densify_tensormap( + fourier_space_expansion_long_range +) +delta_expansion_long_range = densify_tensormap(delta_expansion_long_range) -# # %% -# # As above, we plot the values of the spherical expansions for the Fourier and the -# # subtracted (long range only) spherical expansion. First for hydrogen -# # (``center_species=1``) +# %% +# As above, we plot the values of the spherical expansions for the Fourier and the +# subtracted (long range only) spherical expansion. First for hydrogen +# (``center_species=1``) -# plot_value_comparison( -# fourier_space_expansion_long_range.keys[0], -# fourier_space_expansion_long_range, -# subtracted_expansion_long_range, -# ) +plot_value_comparison( + fourier_space_expansion_long_range.keys[0], + fourier_space_expansion_long_range, + delta_expansion_long_range, +) -# # %% -# # amd second for oxygen (``center_species=8``) +# %% +# amd second for oxygen (``center_species=8``) -# plot_value_comparison( -# fourier_space_expansion_long_range.keys[1], -# fourier_space_expansion_long_range, -# subtracted_expansion_long_range, -# ) +plot_value_comparison( + fourier_space_expansion_long_range.keys[1], + fourier_space_expansion_long_range, + delta_expansion_long_range, +) -# # %% -# # We clearly see that the values of the subtracted spherical are much larger compared -# # to the system with only a single water molecule, thus confirming the presence of -# # long-range contributions in the descriptor for a system with two water molecules. -# # -# # .. end-body +# %% +# We clearly see that the values of the subtracted spherical are much larger compared +# to the system with only a single water molecule, thus confirming the presence of +# long-range contributions in the descriptor for a system with two water molecules. +# +# .. end-body diff --git a/python/rascaline/examples/splined-radial-integral.py b/python/rascaline/examples/splined-radial-integral.py index 8388501e1..87529a9bc 100644 --- a/python/rascaline/examples/splined-radial-integral.py +++ b/python/rascaline/examples/splined-radial-integral.py @@ -1,168 +1,201 @@ """ Splined radial integrals ======================== -""" - -# .. start-body - -# This example illustrates how to generate splined radial basis functions/integrals, -# using a "rectangular" Laplacian eigenstate (LE) basis -# (https://doi.org/10.1063/5.0124363) as the example, i.e, a LE basis truncated -# with ``l_max``, ``n_max`` hyper-parameters. - -# Note that the same basis is also directly available through -# :class:`rascaline.utils.SphericalBesselBasis` with an how-to guide given in -# :ref:`userdoc-how-to-le-basis`. -# """ - -# # %% - -# import ase -# import numpy as np -# import scipy as sp -# from scipy.special import spherical_jn as j_l - -# from rascaline import SphericalExpansion -# from rascaline.utils import RadialIntegralFromFunction, SphericalBesselBasis - - -# # %% -# # Set some hyper-parameters - -# max_angular = 6 -# max_radial = 8 -# cutoff = 5.0 - -# # %% -# # -# # where ``cutoff`` is also the radius of the LE sphere. Now we compute the zeros of -# # the spherical bessel functions. - -# z_ln = SphericalBesselBasis.compute_zeros(max_angular, max_radial) -# z_nl = z_ln.T - -# # %% -# # and define the radial basis functions - - -# def R_nl(n, el, r): -# # Un-normalized LE radial basis functions -# return j_l(el, z_nl[n, el] * r / cutoff) - - -# def N_nl(n, el): -# # Normalization factor for LE basis functions, excluding the a**(-1.5) factor -# def function_to_integrate_to_get_normalization_factor(x): -# return j_l(el, x) ** 2 * x**2 -# integral, _ = sp.integrate.quadrature( -# function_to_integrate_to_get_normalization_factor, 0.0, z_nl[n, el] -# ) -# return (1.0 / z_nl[n, el] ** 3 * integral) ** (-0.5) - - -# def laplacian_eigenstate_basis(n, el, r): -# R = np.zeros_like(r) -# for i in range(r.shape[0]): -# R[i] = R_nl(n, el, r[i]) -# return N_nl(n, el) * R * cutoff ** (-1.5) - - -# # %% -# # Quick normalization check: - -# normalization_check_integral, _ = sp.integrate.quadrature( -# lambda x: laplacian_eigenstate_basis(1, 1, x) ** 2 * x**2, -# 0.0, -# cutoff, -# ) -# print(f"Normalization check (needs to be close to 1): {normalization_check_integral}") - - -# # %% -# # Now the derivatives (by finite differences): - - -# def laplacian_eigenstate_basis_derivative(n, el, r): -# delta = 1e-6 -# all_derivatives_except_at_zero = ( -# laplacian_eigenstate_basis(n, el, r[1:] + delta) -# - laplacian_eigenstate_basis(n, el, r[1:] - delta) -# ) / (2.0 * delta) -# derivative_at_zero = ( -# laplacian_eigenstate_basis(n, el, np.array([delta / 10.0])) -# - laplacian_eigenstate_basis(n, el, np.array([0.0])) -# ) / (delta / 10.0) -# return np.concatenate([derivative_at_zero, all_derivatives_except_at_zero]) - - -# # %% -# # The radial basis functions and their derivatives can be input into a spline -# # generator class. This will output the positions of the spline points, the values -# # of the basis functions evaluated at the spline points, and the corresponding -# # derivatives. - -# spliner = RadialIntegralFromFunction( -# radial_integral=laplacian_eigenstate_basis, -# radial_integral_derivative=laplacian_eigenstate_basis_derivative, -# spline_cutoff=cutoff, -# max_radial=max_radial, -# max_angular=max_angular, -# accuracy=1e-5, -# ) - -# # %% -# # The, we feed the splines to the Rust calculator: Note that the -# # ``atomic_gaussian_width`` will be ignored since we are not uisng a Gaussian basis. - -# hypers_spherical_expansion = { -# "cutoff": cutoff, -# "max_radial": max_radial, -# "max_angular": max_angular, -# "center_atom_weight": 0.0, -# "radial_basis": spliner.compute(), -# "atomic_gaussian_width": 1.0, # ignored -# "cutoff_function": {"Step": {}}, -# } -# calculator = SphericalExpansion(**hypers_spherical_expansion) - -# # %% -# # -# # Create dummy systems to test if the calculator outputs correct radial functions: - - -# def get_dummy_systems(r_array): -# dummy_systems = [] -# for r in r_array: -# dummy_systems.append(ase.Atoms("CH", positions=[(0, 0, 0), (0, 0, r)])) -# return dummy_systems +.. start-body +This example illustrates how to generate splines and use custom basis function and +density when computing density-based representations, such as SOAP or LODE. +""" -# r = np.linspace(0.1, 4.9, 20) -# systems = get_dummy_systems(r) -# spherical_expansion_coefficients = calculator.compute(systems) - -# # %% -# # Extract ``l = 0`` features and check that the ``n = 2`` predictions are the same: - -# block_C_l0 = spherical_expansion_coefficients.block( -# center_type=6, o3_lambda=0, neighbor_type=1 -# ) -# block_C_l0_n2 = block_C_l0.values[:, :, 2].flatten() -# spherical_harmonics_0 = 1.0 / np.sqrt(4.0 * np.pi) - -# # %% -# # radial function = feature / spherical harmonics function -# rascaline_output_radial_function = block_C_l0_n2 / spherical_harmonics_0 - -# assert np.allclose( -# rascaline_output_radial_function, -# laplacian_eigenstate_basis(2, 0, r), -# atol=1e-5, -# ) -# print("Assertion passed successfully!") - - -# # %% -# # -# # .. end-body +# %% + +import json + +import ase.build +import matplotlib.pyplot as plt +import numpy as np +import scipy + +import rascaline +from rascaline import SphericalExpansion +from rascaline.basis import RadialBasis +from rascaline.splines import SoapSpliner + + +# %% +# +# For this example, we will define a new custom radial basis for the SOAP spherical +# expansion, based on Chebyshev polynomials of the first kind. This basis will then be +# used in combination with spherical harmonics to expand the density of neighboring +# atoms around a central atom. +# +# In rascaline, defining custom radial basis is done by creating a class inheriting from +# :py:class:`rascaline.basis.RadialBasis`, and implementing the required method. The +# main one is ``compute_primitive``, which evaluates the radial basis on a set of +# points. This function should also be able to evaluate the derivative of the radial +# basis. If needed :py:meth:`rascaline.basis.RadialBasis.finite_differences_derivative` +# can be used to compute the derivative with finite differences. + + +class Chebyshev(RadialBasis): + def __init__(self, max_radial, radius): + # initialize `RadialBasis` + super().__init__(max_radial=max_radial, radius=radius) + + def compute_primitive(self, positions, n, *, derivative=False): + # map argument from [0, cutoff] to [-1, 1] + z = 2 * positions / self.radius - 1 + if derivative: + return -2 * n / self.radius * scipy.special.chebyu(n)(z) + else: + return scipy.special.chebyt(n + 1)(z) + + @property + def integration_radius(self): + return self.radius + + +# %% +# +# We can now look at the basis functions and their derivatives +radius = 4.5 +basis = Chebyshev(max_radial=4, radius=radius) + +r = np.linspace(0, radius) +for n in range(basis.size): + plt.plot(r, basis.compute_primitive(r, n, derivative=False)) + +plt.title("Chebyshev radial basis functions") +plt.show() + +# %% +for n in range(basis.size): + plt.plot(r, basis.compute_primitive(r, n, derivative=True)) +plt.title("Chebyshev radial basis functions' derivatives") +plt.show() + +# %% +# +# Before being used by rascaline, the basis functions we implemented will be +# orthogonalized and normalized, to improve conditioning of the produced features. This +# is done automatically, and one can access the orthonormalized basis functions with the +# :py:meth:`rascaline.basis.RadialBasis.compute` method. + +basis_orthonormal = basis.compute(r, derivative=False) +for n in range(basis.size): + plt.plot(r, basis_orthonormal[:, n]) + +plt.title("Orthonormalized Chebyshev radial basis functions") +plt.show() + + +# %% +# +# With this, our new radial basis definition is ready to be used with +# :py:class:`rascaline.splines.SoapSpliner`. This class will take the whole set of hyper +# parameters, use them to compute a spline of the radial integral, and give us back new +# hypers that can be used with the native calculators to compute the expansion with our +# custom basis. +# + +spliner = SoapSpliner( + cutoff=rascaline.cutoff.Cutoff( + radius=radius, + smoothing=rascaline.cutoff.ShiftedCosine(width=0.3), + ), + density=rascaline.density.Gaussian(width=0.5), + basis=rascaline.basis.TensorProduct( + max_angular=4, + radial=Chebyshev(max_radial=4, radius=radius), + spline_accuracy=1e-4, + ), +) + +hypers = spliner.get_hypers() + +# %% +# +# The hyper parameters have been transformed from what we gave to the +# :py:class:`rascaline.splines.SoapSpliner`: + +print("hypers['basis'] is", type(hypers["basis"])) +print("hypers['density'] is", type(hypers["density"])) + +# %% +# +# And the new hypers can be used directly with the calculators: + +calculator_splined = SphericalExpansion(**hypers) + +# %% +# +# As a comparison, let's look at the expansion coefficient for formic acid, using both +# our splined radial basis and the classic GTO radial basis: + +atoms = ase.build.molecule("HCOOH", vacuum=4, pbc=True) + +calculator_gto = SphericalExpansion( + # same parameters, only the radial basis changed + cutoff=rascaline.cutoff.Cutoff( + radius=radius, + smoothing=rascaline.cutoff.ShiftedCosine(width=0.3), + ), + density=rascaline.density.Gaussian(width=0.5), + basis=rascaline.basis.TensorProduct( + max_angular=4, + radial=rascaline.basis.Gto(max_radial=4, radius=radius), + spline_accuracy=1e-4, + ), +) + +expansion_splined = calculator_splined.compute(atoms) +expansion_gto = calculator_gto.compute(atoms) + +# %% +# +# As you can see, the coefficients ends up different, with values assigned to different +# basis functions. In practice, which basis function will be the best will depend on the +# use case and exact dataset, so you should try a couple and check how they performe for +# you! + +selection = dict(o3_lambda=0, center_type=8, neighbor_type=1) + +plt.matshow(expansion_splined.block(selection).values.reshape(2, 5)) +plt.matshow(expansion_gto.block(selection).values.reshape(2, 5)) + + +# %% +# +# Since the calculation of the splines requires computing some integral numerically, the +# creation of the splines might take a while. After an initial calculation, you can save +# the splines data in JSON files; and then reload them later to re-use: + +# convert the hypers from classes to a pure JSON-compatible dictionary +json_hypers = rascaline.utils.hypers_to_json(hypers) + +# save the data to a file +with open("splined-hypers.json", "w") as fp: + json.dump(json_hypers, fp) + + +# load the data from the file +with open("splined-hypers.json", "r") as fp: + json_hypers = json.load(fp) + +# the hypers can be used directly with the calculators +calculator = rascaline.SphericalExpansion(**json_hypers) + + +# %% +# +# Finally, you can use the same method to define custom +# :py:class:`rascaline.basis.ExpansionBasis` and custom +# :py:class:`rascaline.density.AtomicDensity`; by creating a new class inheriting from +# the corresponding base class and implementing the corresponding methods. This allow +# you to create a fully custom spherical expansion, and evaluate them efficiently +# through the splines. + +# %% +# +# .. end-body diff --git a/python/rascaline/rascaline/basis.py b/python/rascaline/rascaline/basis.py index 61ad9d36b..809966d6c 100644 --- a/python/rascaline/rascaline/basis.py +++ b/python/rascaline/rascaline/basis.py @@ -44,6 +44,7 @@ def __init__(self, *, max_radial: int, radius: float): assert self.max_radial >= 0 assert self.radius > 0.0 + self._orthonormalization_matrix = None def _rascaline_hypers(self): return self.get_hypers() diff --git a/python/rascaline/rascaline/calculators.py b/python/rascaline/rascaline/calculators.py index b04b3f2dc..a9d18efba 100644 --- a/python/rascaline/rascaline/calculators.py +++ b/python/rascaline/rascaline/calculators.py @@ -1,6 +1,6 @@ import json -from .utils import BadHyperParameters, convert_hypers +from .utils import BadHyperParameters, convert_hypers, hypers_to_json try: @@ -10,22 +10,6 @@ from .calculator_base import CalculatorBase -def _expand_hypers(hypers): - """ - Starting from a Dict[str, Any], recursively expand all values in the dict by - calling their ``_rascaline_hypers`` method if such method exists. - """ - cleaned = {} - for key, value in hypers.items(): - if hasattr(value, "_rascaline_hypers"): - value = value._rascaline_hypers() - - if isinstance(value, dict): - value = _expand_hypers(value) - cleaned[key] = value - return cleaned - - class AtomicComposition(CalculatorBase): """An atomic composition calculator for obtaining the stoichiometric information. @@ -37,7 +21,7 @@ class AtomicComposition(CalculatorBase): """ def __init__(self, *, per_system): - parameters = _expand_hypers( + parameters = hypers_to_json( { "per_system": per_system, } @@ -47,7 +31,7 @@ def __init__(self, *, per_system): class DummyCalculator(CalculatorBase): def __init__(self, *, cutoff, delta, name): - parameters = _expand_hypers( + parameters = hypers_to_json( { "cutoff": cutoff, "delta": delta, @@ -80,7 +64,7 @@ class NeighborList(CalculatorBase): """ def __init__(self, *, cutoff, full_neighbor_list, self_pairs=False): - parameters = _expand_hypers( + parameters = hypers_to_json( { "cutoff": cutoff, "full_neighbor_list": full_neighbor_list, @@ -106,7 +90,7 @@ class SortedDistances(CalculatorBase): """ def __init__(self, *, cutoff, max_neighbors, separate_neighbor_types): - parameters = _expand_hypers( + parameters = hypers_to_json( { "cutoff": cutoff, "max_neighbors": max_neighbors, @@ -160,7 +144,7 @@ def __init__(self, *, cutoff=None, density=None, basis=None, **kwargs): if len(kwargs) != 0 or density is None or basis is None: _check_for_old_hypers("SphericalExpansion", {"cutoff": cutoff, **kwargs}) - parameters = _expand_hypers( + parameters = hypers_to_json( { "cutoff": cutoff, "density": density, @@ -210,7 +194,7 @@ def __init__(self, *, cutoff=None, density=None, basis=None, **kwargs): "SphericalExpansionByPair", {"cutoff": cutoff, **kwargs} ) - parameters = _expand_hypers( + parameters = hypers_to_json( { "cutoff": cutoff, "density": density, @@ -242,7 +226,7 @@ def __init__(self, *, cutoff=None, density=None, basis=None, **kwargs): if len(kwargs) != 0 or density is None or basis is None: _check_for_old_hypers("SoapRadialSpectrum", {"cutoff": cutoff, **kwargs}) - parameters = _expand_hypers( + parameters = hypers_to_json( { "cutoff": cutoff, "density": density, @@ -278,7 +262,7 @@ def __init__(self, *, cutoff=None, density=None, basis=None, **kwargs): if len(kwargs) != 0 or density is None or basis is None: _check_for_old_hypers("SoapPowerSpectrum", {"cutoff": cutoff, **kwargs}) - parameters = _expand_hypers( + parameters = hypers_to_json( { "cutoff": cutoff, "density": density, @@ -312,7 +296,7 @@ def __init__(self, *, density=None, basis=None, k_cutoff=None, **kwargs): "LodeSphericalExpansion", {"k_cutoff": k_cutoff, **kwargs} ) - parameters = _expand_hypers( + parameters = hypers_to_json( { "k_cutoff": k_cutoff, "density": density, diff --git a/python/rascaline/rascaline/utils/__init__.py b/python/rascaline/rascaline/utils/__init__.py index 5d6cecfbd..5be9b5144 100644 --- a/python/rascaline/rascaline/utils/__init__.py +++ b/python/rascaline/rascaline/utils/__init__.py @@ -1,13 +1,13 @@ import os -from .clebsch_gordan import ( # noqa +from .clebsch_gordan import ( # noqa: F401 ClebschGordanProduct, DensityCorrelations, calculate_cg_coefficients, cartesian_to_spherical, ) -from .hypers import BadHyperParameters, convert_hypers # noqa -from .power_spectrum import PowerSpectrum # noqa +from .hypers import BadHyperParameters, convert_hypers, hypers_to_json # noqa: F401 +from .power_spectrum import PowerSpectrum # noqa: F401 _HERE = os.path.dirname(__file__) diff --git a/python/rascaline/rascaline/utils/hypers/__init__.py b/python/rascaline/rascaline/utils/hypers/__init__.py index fa42977a4..54979940e 100644 --- a/python/rascaline/rascaline/utils/hypers/__init__.py +++ b/python/rascaline/rascaline/utils/hypers/__init__.py @@ -1,5 +1,6 @@ import json import re +from typing import Any, Dict from . import _rascaline @@ -47,3 +48,22 @@ def convert_hypers(origin, representation=None, hypers=None): return f"{representation}(**{hypers_dict})" else: raise ValueError(f"no hyper conversion exists for {origin} software") + + +def hypers_to_json(hypers_dict: Dict[str, Any]): + """ + Convert from class version of rascaline hyper-parameters to the JSON version. + + The class version would contain something like ``{"cutoff": Cutoff(radius=3.4)}``, + which this function transforms into ``{"cutoff": {"radius": 3.4", "smoothing": + {"type": "Step"}}}``. + """ + json = {} + for key, value in hypers_dict.items(): + if hasattr(value, "_rascaline_hypers"): + value = value._rascaline_hypers() + + if isinstance(value, dict): + value = hypers_to_json(value) + json[key] = value + return json