Split torch.Tensor and metatensor bindings (#4)

Co-authored-by: Kevin Kazuki Huguenin-Dumittan <kvhuguenin@gmail.com>
lab-cosmo · Dec 6, 2023 · 49e9497 · 49e9497
1 parent c57b6e9
commit 49e9497
Show file tree

Hide file tree

Showing 32 changed files with 543 additions and 204 deletions.
diff --git a/docs/requirements.txt b/docs/requirements.txt
@@ -2,6 +2,5 @@ furo
 sphinx > 7.0
 sphinx-gallery
 sphinx-toggleprompt
-pydata-sphinx-theme
 tomli
 chemiscope
diff --git a/docs/src/conf.py b/docs/src/conf.py
@@ -72,7 +72,7 @@
     "footer_icons": [
         {
             "name": "GitHub",
-            "url": "https://github.com/ceriottm/MeshLODE",
+            "url": project_dict["urls"]["repository"],
             "html": "",
             "class": "fa-brands fa-github fa-2x",
         },

diff --git a/docs/src/index.rst b/docs/src/index.rst
@@ -1,4 +1,7 @@
-.. automodule:: meshlode
+MeshLODE
+########
+
+Particle-mesh based calculation of Long Distance Equivariants.
 
 .. toctree::
    :hidden:

diff --git a/docs/src/references/calculators/index.rst b/docs/src/references/calculators/index.rst
@@ -1,3 +1,19 @@
+.. _calculators:
+
+Calculators
+###########
+
+Below is a list of all calculators available. Calculators are the core of MeshLODE and
+are algorithms for transforming Cartesian coordinates into representations suitable for
+machine learning.
+
+Our calculator API follows the `rascaline <https://luthaf.fr/rascaline>`_ API and coding
+guidelines to promote usability and interoperability with existing workflows.
+
+Calculators return the representations as a :py:obj:`List` of :py:class:`torch.Tensor`.
+We also provide a return values as a :py:class:`metatensor.TensorMap` in
+:ref:`metatensor`.
+
 .. automodule:: meshlode.calculators
 
 .. toctree::

diff --git a/docs/src/references/fourier_convolution.rst b/docs/src/references/fourier_convolution.rst
diff --git a/docs/src/references/index.rst b/docs/src/references/index.rst
@@ -12,6 +12,5 @@ functions toegther you might take a look at the :ref:`userdoc-how-to` section.
     :maxdepth: 1
 
     calculators/index
-    fourier_convolution
-    mesh_interpolator
-    system
+    metatensor/index
+    lib/index
diff --git a/docs/src/references/lib/fourier_convolution.rst b/docs/src/references/lib/fourier_convolution.rst
@@ -0,0 +1,7 @@
+Fourier Convolution
+===================
+
+.. autoclass:: meshlode.lib.FourierSpaceConvolution
+    :members:
+    :undoc-members:
+    :show-inheritance:
diff --git a/docs/src/references/lib/index.rst b/docs/src/references/lib/index.rst
@@ -0,0 +1,12 @@
+General Library functions
+#########################
+
+General Library functions and classes of meshLODE mathematical helper functions, which
+are used for the meshLODE calculators.
+
+.. toctree::
+   :maxdepth: 1
+
+   fourier_convolution
+   mesh_interpolator
+   system
diff --git a/docs/src/references/lib/mesh_interpolator.rst b/docs/src/references/lib/mesh_interpolator.rst
@@ -0,0 +1,7 @@
+Mesh Interpolator
+=================
+
+.. autoclass:: meshlode.lib.MeshInterpolator
+    :members:
+    :undoc-members:
+    :show-inheritance:
diff --git a/docs/src/references/system.rst → docs/src/references/lib/system.rst b/docs/src/references/system.rst → docs/src/references/lib/system.rst
@@ -1,4 +1,7 @@
-.. automodule:: meshlode.system
+System
+======
+
+.. autoclass:: meshlode.System
     :members:
     :undoc-members:
     :show-inheritance:
diff --git a/docs/src/references/mesh_interpolator.rst b/docs/src/references/mesh_interpolator.rst
diff --git a/docs/src/references/metatensor/index.rst b/docs/src/references/metatensor/index.rst
@@ -0,0 +1,19 @@
+.. _metatensor:
+
+Metatensor Bindings
+###################
+
+MeshLODE calculators returning representations as :py:class:`metatensor.TensorMap`.
+For using these bindings you need to install the ``metatensor.torch`` optional
+dependencies.
+
+.. code-block:: bash
+
+   pip install .[metatensor]
+
+For a plain :py:class:`torch.Tensor` refer to :ref:`calculators`.
+
+.. toctree::
+   :maxdepth: 1
+
+   meshpotential
diff --git a/docs/src/references/metatensor/meshpotential.rst b/docs/src/references/metatensor/meshpotential.rst
@@ -0,0 +1,7 @@
+MeshPotential
+#############
+
+.. autoclass:: meshlode.metatensor.MeshPotential
+    :members:
+    :undoc-members:
+    :show-inheritance:
diff --git a/examples/README.rst b/examples/README.rst
@@ -6,3 +6,10 @@ Examples
 This section list introductory examples and recipes to the various classes and functions
 of ``MeshLODE``. For details on the API specification of the functions take a look at
 the :ref:`userdoc-reference` section.
+
+For running the all examples install with the ``examples`` and ``metatensor`` optional
+dependencies.
+
+.. code-block:: bash
+
+    pip install .[examples,metatensor]
diff --git a/examples/basic-tutorial.py → examples/library-tutorial.py b/examples/basic-tutorial.py → examples/library-tutorial.py
@@ -1,17 +1,16 @@
 """
-Basic Tutorial
-==============
-
-This examples provides an illustration of the functioning of
-``meshlode`` and the construction LODE descriptors (`Grisafi
-2019 <https://doi.org/10.1063/1.5128375>`__, `Grisafi
-2021 <https://doi.org/10.1039/D0SC04934D>`__, `Huguenin
-2023 <10.1021/acs.jpclett.3c02375>`__). It builds the (simple and
-weighted) atom density for a CsCl-type structure, computes a smeared
-version and the Coulomb potential, and projects on a separate set of
-points.
-
+Basic Tutorial for Library functions
+====================================
+
+This examples provides an illustration of the functioning of the underlaying library
+functions of ``meshlode`` and the construction LODE descriptors (`Grisafi 2019
+<https://doi.org/10.1063/1.5128375>`__, `Grisafi 2021
+<https://doi.org/10.1039/D0SC04934D>`__, `Huguenin 2023
+<10.1021/acs.jpclett.3c02375>`__). It builds the (simple and weighted) atom density for
+a CsCl-type structure, computes a smeared version and the Coulomb potential, and
+projects on a separate set of points.
 """
+# %%
 
 import ase
 import chemiscope
@@ -32,7 +31,7 @@ def sliceplot(mesh, sz=12, cmap="viridis", vmin=None, vmax=None):
         vmin = mesh.min()
     if vmax is None:
         vmax = mesh.max()
-    fig, ax = plt.subplots(
+    _, ax = plt.subplots(
         1,
         mesh.shape[-1],
         figsize=(sz, sz / mesh.shape[-1]),
@@ -49,8 +48,9 @@ def sliceplot(mesh, sz=12, cmap="viridis", vmin=None, vmax=None):
 # Builds the structure
 # --------------------
 #
-# Nothing special to see here. Builds a CsCl structure by replicating the
-# primitive cell. Add a bit of noise to make it less boring!
+# Builds a CsCl structure by replicating the primitive cell using ase and convert it to
+# a :py:class:`List` of :py:class:`meshlode.System`. We add a bit of noise to make
+# it less boring!
 #
 
 positions = torch.tensor([[0, 0, 0], [0.5, 0.5, 0.5]]) * 4
@@ -93,7 +93,7 @@ def sliceplot(mesh, sz=12, cmap="viridis", vmin=None, vmax=None):
 # list of atom weights to yield the mesh values.
 #
 
-interpol = meshlode.mesh_interpolator.MeshInterpolator(
+interpol = meshlode.lib.mesh_interpolator.MeshInterpolator(
     frame.cell, torch.tensor([16, 16, 16]), interpolation_order=3
 )
 
@@ -131,7 +131,7 @@ def sliceplot(mesh, sz=12, cmap="viridis", vmin=None, vmax=None):
 # be easily extended to compute an arbitrary filter
 #
 
-fsc = meshlode.fourier_convolution.FourierSpaceConvolution(frame.cell)
+fsc = meshlode.lib.fourier_convolution.FourierSpaceConvolution(frame.cell)
 
 # %%
 # plain smearing
@@ -169,7 +169,7 @@ def sliceplot(mesh, sz=12, cmap="viridis", vmin=None, vmax=None):
 # ``interpolation_order``, if wanted.
 #
 
-interpol_slice = meshlode.mesh_interpolator.MeshInterpolator(
+interpol_slice = meshlode.lib.mesh_interpolator.MeshInterpolator(
     frame.cell, torch.tensor([16, 16, 16]), interpolation_order=4
 )
 

diff --git a/examples/madelung.py b/examples/madelung.py
@@ -3,62 +3,131 @@
 ==========================
 
 In this tutorial we show how to calculate the Madelung constants and total electrostatic
-energy of atomic structures using MeshLODE.
+energy of atomic structures using the :py:class:`meshlode.MeshPotential` and
+:py:class:`meshlode.metatensor.MeshPotential` calculator.
 """
 
+# %%
+import math
+
 import torch
 
-from meshlode import MeshPotential, System
+import meshlode
 
 
-# Define simple example structure having the CsCl structure
-positions = torch.tensor([[0, 0, 0], [0.5, 0.5, 0.5]])
-atomic_numbers = torch.tensor([55, 17])  # Cs and Cl
+# %%
+# Define simple example structure having the CsCl structure and compute the reference
+# values. MeshPotential by default outputs the species sorted according to the atomic
+# number. Thus, we input the compound "CsCl" and "ClCs" since Cl and Cs have atomic
+# numbers 17 and 55, respectively.
+atomic_numbers = torch.tensor([17, 55])  # Cl and Cs
+charges = torch.tensor([-1.0, 1.0])
 cell = torch.eye(3)
-charges = torch.tensor([1.0, -1.0])
-frame = System(species=atomic_numbers, positions=positions, cell=torch.eye(3))
+positions = torch.tensor([[0, 0, 0], [0.5, 0.5, 0.5]])
+frame = meshlode.System(species=atomic_numbers, positions=positions, cell=torch.eye(3))
+
+# %%
+# Define the expected values of the energy
+n_atoms = len(positions)
+madelung = 2 * 1.7626 / math.sqrt(3)
+energies_ref = -madelung * torch.ones((n_atoms, 1))
+
+# %%
+# We first define general parameters for our calculation MeshLODE
 
-# Define parameters in MeshLODE
 atomic_smearing = 0.1
 cell = torch.eye(3)
 mesh_spacing = atomic_smearing / 4
 interpolation_order = 2
 
-# Compute features
-MP = MeshPotential(
+# %%
+# Computation using ``meshlode``
+# ------------------------------
+#
+# Compute features using
+
+MP = meshlode.MeshPotential(
     atomic_smearing=atomic_smearing,
     mesh_spacing=mesh_spacing,
     interpolation_order=interpolation_order,
     subtract_self=True,
 )
-potentials_mesh = MP.compute(frame)
+potentials_torch = MP.compute(frame)
 
-# The ``potentials'' that have been computed so far are not the actual electrostatic
+# %%
+# The "potentials" that have been computed so far are not the actual electrostatic
 # potentials. For instance, for the Cs atom, we are separately storing the contributions
 # to the potential (at the location of the Cs atom) from the Cs atoms and Cl atoms
 # separately. Thus, to get the Madelung constant, we need to take a linear combination
 # of these "potentials" weighted by the charges of the atoms.
-n_atoms = len(positions)
-atomic_energies = torch.zeros((n_atoms, 1))
+
+atomic_energies_torch = torch.zeros((n_atoms, 1))
+for idx_c in range(n_atoms):
+    for idx_n in range(n_atoms):
+        # The coulomb potential between atoms i and j is charge_i * charge_j / d_ij
+        # The features are simply computing a pure 1/r potential with no prefactors.
+        # Thus, to compute the energy between atoms of species i and j, we need to
+        # multiply by the charges of i and j.
+        print(charges[idx_c] * charges[idx_n], potentials_torch[idx_n, idx_c])
+        atomic_energies_torch[idx_c] += (
+            charges[idx_c] * charges[idx_n] * potentials_torch[idx_c, idx_n]
+        )
+
+# %%
+# The total energy is just the sum of all atomic energies
+total_energy_torch = torch.sum(atomic_energies_torch)
+
+# %%
+# Compare against reference Madelung constant and reference energy:
+print("Using the torch version")
+print(f"Computed energies on each atom = {atomic_energies_torch.tolist()}")
+print(f"Reference Madelung constant = {madelung:.3f}")
+print(f"Total energy = {total_energy_torch:.3f}\n")
+
+
+# %%
+# Computation using ``meshlode.metatensor``
+# -----------------------------------------
+#
+# We now compute the same constants using the metatensor based calculator
+
+MP = meshlode.metatensor.MeshPotential(
+    atomic_smearing=atomic_smearing,
+    mesh_spacing=mesh_spacing,
+    interpolation_order=interpolation_order,
+    subtract_self=True,
+)
+potential_metatensor = MP.compute(frame)
+
+
+# %%
+# To get the Madelung constant, we again need to take a linear combination
+# of the "potentials" weighted by the charges of the atoms.
+
+atomic_energies_metatensor = torch.zeros((n_atoms, 1))
 for idx_c, c in enumerate(atomic_numbers):
     for idx_n, n in enumerate(atomic_numbers):
         # Take the coefficients with the correct center atom and neighbor atom species
-        block = potentials_mesh.block(
+        block = potential_metatensor.block(
             {"species_center": int(c), "species_neighbor": int(n)}
         )
 
         # The coulomb potential between atoms i and j is charge_i * charge_j / d_ij
         # The features are simply computing a pure 1/r potential with no prefactors.
         # Thus, to compute the energy between atoms of species i and j, we need to
         # multiply by the charges of i and j.
-        atomic_energies[idx_c] += charges[idx_c] * charges[idx_n] * block.values[0, 0]
+        print(c, n, charges[idx_c] * charges[idx_n], block.values[0, 0])
+        atomic_energies_metatensor[idx_c] += (
+            charges[idx_c] * charges[idx_n] * block.values[0, 0]
+        )
 
+# %%
 # The total energy is just the sum of all atomic energies
-total_energy = torch.sum(atomic_energies)
+total_energy_metatensor = torch.sum(atomic_energies_metatensor)
 
-# Compare against reference Madelung constant:
-madelung = 2 * 1.7626 / torch.sqrt(torch.tensor(3))
-energies_ref = -madelung * torch.ones((n_atoms, 1))
-print("Computed energies on each atom = \n", atomic_energies)
-print("Reference Madelung constant = \n", madelung)
-print("Total energy = \n", total_energy)
+# %%
+# Compare against reference Madelung constant and reference energy:
+print("Using the metatensor version")
+print(f"Computed energies on each atom = {atomic_energies_metatensor.tolist()}")
+print(f"Reference Madelung constant = {madelung:.3f}")
+print(f"Total energy = {total_energy_metatensor:.3f}")
diff --git a/pyproject.toml b/pyproject.toml
@@ -35,7 +35,6 @@ keywords = [
 ]
 dependencies = [
     "torch >= 1.11",
-    "metatensor[torch]",
 ]
 dynamic = ["version"]
 
@@ -44,6 +43,9 @@ examples = [
     "ase",
     "matplotlib",
 ]
+metatensor = [
+    "metatensor[torch]",
+]
 
 [project.urls]
 homepage = "http://meshlode.readthedocs.io"