Add more examples,shape tests, example and code blocks in docs

.gitignore

@@ -6,9 +6,10 @@ __pycache__
 *.swo
 *DS_Store
 *coverage*
+sg_execution_times.rst
+*.gz
 
 .tox/
 build/
 dist/
 docs/src/examples
-playground/

docs/requirements.txt

@@ -4,3 +4,4 @@ sphinx-gallery
 sphinx-toggleprompt
 pydata-sphinx-theme
 tomli
+chemiscope

examples/basic-tutorial.py

@@ -0,0 +1,200 @@
+"""
+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.
+
+"""
+
+import ase
+import chemiscope
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+import meshlode
+
+
+torch.set_default_dtype(torch.float64)
+
+
+# plot a 3D mesh with a stack of 2D plots
+def sliceplot(mesh, sz=12, cmap="viridis", vmin=None, vmax=None):
+    mesh = mesh.detach().numpy()
+    if vmin is None:
+        vmin = mesh.min()
+    if vmax is None:
+        vmax = mesh.max()
+    fig, ax = plt.subplots(
+        1,
+        mesh.shape[-1],
+        figsize=(sz, sz / mesh.shape[-1]),
+        sharey=True,
+        constrained_layout=True,
+    )
+    for i in range(mesh.shape[-1]):
+        ax[i].matshow(mesh[:, :, i], vmin=vmin, vmax=vmax, cmap=cmap)
+        ax[i].set_xticklabels([])
+        ax[i].set_yticklabels([])
+
+
+# %%
+# 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!
+#
+
+positions = torch.tensor([[0, 0, 0], [0.5, 0.5, 0.5]]) * 4
+atomic_numbers = torch.tensor([55, 17])  # Cs and Cl
+cell = torch.eye(3) * 4
+ase_frame = ase.Atoms(positions=positions, cell=cell, numbers=atomic_numbers).repeat(
+    [2, 2, 2]
+)
+ase_frame.positions[:] += np.random.normal(size=ase_frame.positions.shape) * 0.1
+charges = torch.tensor([1.0, -1.0] * 8)
+frame = meshlode.System(
+    species=torch.tensor(ase_frame.numbers),
+    positions=torch.tensor(np.array(ase_frame.positions)),
+    cell=torch.tensor(ase_frame.cell),
+)
+
+cs = chemiscope.show(
+    frames=[ase_frame],
+    mode="structure",
+    settings={"structure": [{"unitCell": True, "axes": "xyz"}]},
+)
+
+if chemiscope.jupyter._is_running_in_notebook():
+    from IPython.display import display
+
+    display(cs)
+else:
+    cs.save("cscl.json.gz")
+
+
+# %%
+# MeshInterpolator
+# ----------------
+#
+# ``MeshInterpolator`` serves as a utility class to compute a mesh
+# representation of points, and/or to project a function defined on the
+# mesh on a set of points. Computing the mesh representation is a two-step
+# procedure. First, the weights associated with the interpolation of the
+# point positions are evaluated, then they are combined with one or more
+# list of atom weights to yield the mesh values.
+#
+
+interpol = meshlode.mesh_interpolator.MeshInterpolator(
+    frame.cell, torch.tensor([16, 16, 16]), interpolation_order=3
+)
+
+interpol.compute_interpolation_weights(frame.positions)
+
+
+# %%
+# We use two sets of weights: ones (giving the atom density irrespective
+# of the species) and charges (giving a smooth representation of the point
+# charges).
+#
+
+atom_weights = torch.ones((len(charges), 2))
+atom_weights[:, 1] = charges
+mesh = interpol.points_to_mesh(atom_weights)
+
+# there are two densities
+mesh.shape
+
+# %%
+# the first corresponds to plain density
+sliceplot(mesh[0, :, :, :5])
+
+# %%
+# the second to the charge-weighted one
+sliceplot(mesh[1, :, :, :5], cmap="seismic", vmax=1, vmin=-1)
+
+
+# %%
+# Fourier filter
+# --------------
+#
+# This module computes a Fourier-domain filter, that can be used e.g. to
+# smear the density and/or compute a 1/r^p potential field. This can also
+# be easily extended to compute an arbitrary filter
+#
+
+fsc = meshlode.fourier_convolution.FourierSpaceConvolution(frame.cell)
+
+# %%
+# plain smearing
+rho_mesh = fsc.compute(mesh, potential_exponent=0, smearing=1)
+
+sliceplot(rho_mesh[0, :, :, :5])
+
+# %%
+# coulomb-like potential, no smearing
+coulomb_mesh = fsc.compute(mesh, potential_exponent=1, smearing=0)
+
+sliceplot(coulomb_mesh[1, :, :, :5], cmap="seismic")
+
+
+# %%
+# Back-interpolation (on the same points)
+# ---------------------------------------
+#
+# The same ``MeshInterpolator`` object can be used to compute a field on
+# the same points used initially to generate the atom density
+#
+
+potentials = interpol.mesh_to_points(coulomb_mesh)
+
+potentials
+
+
+# %%
+# Back-interpolation (on different points)
+# ----------------------------------------
+#
+# In order to compute the field on a different set of points, it is
+# sufficient to build another ``MeshInterpolator`` object and to compute
+# it with the desired field. One can also use a different
+# ``interpolation_order``, if wanted.
+#
+
+interpol_slice = meshlode.mesh_interpolator.MeshInterpolator(
+    frame.cell, torch.tensor([16, 16, 16]), interpolation_order=4
+)
+
+# Compute a denser grid on a 2D slice
+n_points = 50
+x = torch.linspace(0, frame.cell[0, 0], n_points + 1)[:n_points]
+y = torch.linspace(0, frame.cell[1, 1], n_points + 1)[:n_points]
+xx, yy = torch.meshgrid(x, y, indexing="ij")
+
+# Flatten xx and yy, and concatenate with a zero column for the z-coordinate
+slice_points = torch.cat(
+    (xx.reshape(-1, 1), yy.reshape(-1, 1), 0.5 * torch.ones(n_points**2, 1)), dim=1
+)
+
+
+# %%
+interpol_slice.compute_interpolation_weights(slice_points)
+
+coulomb_slice = interpol_slice.mesh_to_points(coulomb_mesh)
+
+plt.contourf(
+    xx,
+    yy,
+    coulomb_slice[:, 1].reshape(n_points, n_points).T,
+    cmap="seismic",
+    vmin=-1,
+    vmax=1,
+)

examples/madelung.py

@@ -7,13 +7,15 @@
 """
 
 import torch
+
 from meshlode import MeshPotential, System
 
+
 # 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
+positions = torch.tensor([[0, 0, 0], [0.5, 0.5, 0.5]])
+atomic_numbers = torch.tensor([55, 17])  # Cs and Cl
 cell = torch.eye(3)
-charges = torch.tensor([1.,-1.])
+charges = torch.tensor([1.0, -1.0])
 frame = System(species=atomic_numbers, positions=positions, cell=torch.eye(3))
 
 # Define parameters in MeshLODE
@@ -23,10 +25,12 @@
 interpolation_order = 2
 
 # Compute features
-MP = MeshPotential(atomic_smearing=atomic_smearing,
-                   mesh_spacing=mesh_spacing,
-                   interpolation_order=interpolation_order,
-                   subtract_self=True)
+MP = MeshPotential(
+    atomic_smearing=atomic_smearing,
+    mesh_spacing=mesh_spacing,
+    interpolation_order=interpolation_order,
+    subtract_self=True,
+)
 potentials_mesh = MP.compute(frame)
 
 # The ``potentials'' that have been computed so far are not the actual electrostatic
@@ -55,6 +59,6 @@
 # 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)
+print("Computed energies on each atom = \n", atomic_energies)
+print("Reference Madelung constant = \n", madelung)
+print("Total energy = \n", total_energy)