update some docs

docs/src/references/tuning/index.rst

@@ -6,6 +6,15 @@ The choice of parameters like the neighborlist ``cutoff``, the ``smearing`` or t
 calculation. To help find the parameters that meet the accuracy requirements, this
 module offers tuning methods for the calculators.
 
+For usual tuning procedures we provide simple functions like
+:func:`torchpme.tuning.tune_ewald` that returns for a given system the optimal
+parameters for the Ewald summation. For more complex tuning procedures, we provide
+classes like :class:`torchpme.tuning.ewald.EwaldErrorBounds` that can be used to
+implement custom tuning procedures.
+
+.. important::
+
+   Current tuning methods are only implemented for the Coulomb potential.
 
 .. toctree::
    :maxdepth: 1

docs/src/references/tuning/tune_ewald.rst

@@ -18,7 +18,7 @@ where :math:`N` is the number of charges, :math:`Q^2 = \sum_{i = 1}^N q_i^2`, is
 charges, :math:`r_{\text{cutoff}}` is the short-range cutoff, :math:`V` is the volume of the
 simulation box and :math:`h^2` is the long range wavelength.
 
-.. autofunction:: torchpme.tuning.ewald.tune_ewald
+.. autofunction:: torchpme.tuning.tune_ewald
 
 .. autoclass:: torchpme.tuning.ewald.EwaldErrorBounds
     :members:

docs/src/references/tuning/tune_p3m.rst

@@ -21,7 +21,7 @@ simulation box, :math:`p` is the order of the interpolation scheme, :math:`H` is
 points and :math:`a_m^{(p)}` is an expansion coefficient.
 
 
-.. autofunction:: torchpme.tuning.p3m.tune_p3m
+.. autofunction:: torchpme.tuning.tune_p3m
 
 .. autoclass:: torchpme.tuning.p3m.P3MErrorBounds
     :members:

docs/src/references/tuning/tune_pme.rst

@@ -21,7 +21,7 @@ simulation box, :math:`p` is the order of the interpolation scheme, :math:`H` is
 points, and :math:`\phi_p^2 = H^{-(p+1)}\prod_{s\in S_H^{(p)}}(x - s)`, in which :math:`S_H^{(p)}` is
 the :math:`p+1` mesh points closest to the point :math:`x`.
 
-.. autofunction:: torchpme.tuning.pme.tune_pme
+.. autofunction:: torchpme.tuning.tune_pme
 
 .. autoclass:: torchpme.tuning.pme.PMEErrorBounds
     :members:

examples/05-autograd-demo.py

@@ -516,84 +516,3 @@ def forward(self, positions, cell, charges):
 
 # %%
 print(f"Evaluation time:\nPytorch: {time_python}ms\nJitted:  {time_jit}ms")
-
-# %%
-# Other auto-differentiation ideas
-# --------------------------------
-#
-# There are many other ways the auto-differentiation engine of
-# ``torch`` can be used to facilitate the evaluation of atomistic
-# models.
-
-# %%
-# 4-site water models
-# ~~~~~~~~~~~~~~~~~~~
-#
-# Several water models (starting from the venerable TIP4P model of
-# `Abascal and C. Vega, JCP (2005) <http://doi.org/10.1063/1.2121687>`_)
-# use a center of negative charge that is displaced from the O position.
-# This is easily implemented, yielding the forces on the O and H positions
-# generated by the displaced charge.
-
-structure = ase.Atoms(
-    positions=[
-        [0, 0, 0],
-        [0, 1, 0],
-        [1, -0.2, 0],
-    ],
-    cell=[6, 6, 6],
-    symbols="OHH",
-)
-
-cell = torch.from_numpy(structure.cell.array).to(device=device, dtype=dtype)
-positions = torch.from_numpy(structure.positions).to(device=device, dtype=dtype)
-
-# %%
-# The key step is to create a "fourth site" based on the O positions
-# and use it in the ``interpolate`` step.
-
-charges = torch.tensor(
-    [[-1.0], [0.5], [0.5]],
-    dtype=dtype,
-    device=device,
-)
-
-positions.requires_grad_(True)
-charges.requires_grad_(True)
-cell.requires_grad_(True)
-
-positions_4site = torch.vstack(
-    [
-        ((positions[1::3] + positions[2::3]) * 0.5 + positions[0::3] * 3) / 4,
-        positions[1::3],
-        positions[2::3],
-    ]
-)
-
-ns = torch.tensor([5, 5, 5])
-interpolator = torchpme.lib.MeshInterpolator(
-    cell=cell, ns_mesh=ns, interpolation_nodes=3, method="Lagrange"
-)
-interpolator.compute_weights(positions_4site)
-mesh = interpolator.points_to_mesh(charges)
-
-value = (mesh**2).sum()
-
-# %%
-# The gradients can be computed by just running `backward` on the
-# end result. Gradients are computed on the H and O positions.
-
-value.backward()
-
-print(
-    f"""
-Position gradients:
-{positions.grad.T}
-
-Cell gradients:
-{cell.grad}
-
-Charges gradients:
-{charges.grad.T}
-"""
-)

examples/11-4-site-water.py

@@ -0,0 +1,84 @@
+"""
+.. _example-tip4p-water:
+
+4-site water models
+===================
+
+.. currentmodule:: torchpme
+
+# Several water models (starting from the venerable TIP4P model of
+# `Abascal and C. Vega, JCP (2005) <http://doi.org/10.1063/1.2121687>`_)
+# use a center of negative charge that is displaced from the O position.
+# This is easily implemented, yielding the forces on the O and H positions
+# generated by the displaced charge.
+"""
+
+import ase
+import torch
+
+import torchpme
+
+structure = ase.Atoms(
+    positions=[
+        [0, 0, 0],
+        [0, 1, 0],
+        [1, -0.2, 0],
+    ],
+    cell=[6, 6, 6],
+    symbols="OHH",
+)
+
+cell = torch.from_numpy(structure.cell.array)
+positions = torch.from_numpy(structure.positions)
+
+# %%
+# The key step is to create a "fourth site" based on the oxygen positions and use it in
+# the ``interpolate`` step.
+
+charges = torch.tensor([[-1.0], [0.5], [0.5]])
+
+positions.requires_grad_(True)
+charges.requires_grad_(True)
+cell.requires_grad_(True)
+
+positions_4site = torch.vstack(
+    [
+        ((positions[1::3] + positions[2::3]) * 0.5 + positions[0::3] * 3) / 4,
+        positions[1::3],
+        positions[2::3],
+    ]
+)
+
+# %%
+# .. important::
+#
+#   For the automatic differentiation to work it is important to make a new tensor as
+#   ``positions_4site`` and do not "overwrite" the original tensor.
+
+ns = torch.tensor([5, 5, 5])
+interpolator = torchpme.lib.MeshInterpolator(
+    cell=cell, ns_mesh=ns, interpolation_nodes=3, method="Lagrange"
+)
+interpolator.compute_weights(positions_4site)
+mesh = interpolator.points_to_mesh(charges)
+
+value = (mesh**2).sum()
+
+# %%
+# The gradients can be computed by just running `backward` on the
+# end result. Gradients are computed on the H and O positions.
+
+value.backward()
+
+print(
+    f"""
+Position gradients:
+{positions.grad.T}
+
+Cell gradients:
+{cell.grad}
+
+Charges gradients:
+{charges.grad.T}
+"""
+)

src/torchpme/calculators/calculator.py

@@ -40,9 +40,9 @@ def __init__(
     ):
         super().__init__()
 
-        assert isinstance(
-            potential, Potential
-        ), f"Potential must be an instance of Potential, got {type(potential)}"
+        assert isinstance(potential, Potential), (
+            f"Potential must be an instance of Potential, got {type(potential)}"
+        )
 
         self.device = "cpu" if device is None else device
         self.dtype = torch.get_default_dtype() if dtype is None else dtype