Revamped tuning

docs/src/references/changelog.rst

@@ -29,12 +29,16 @@ Added
 
 * Added a PyTorch implementation of the exponential integral function
 * Added ``dtype`` and ``device`` for ``Calculator`` classses
+* Added an example on the tuning scheme and usage, and how to optimize the ``cutoff``
 
 Changed
 #######
 
 * Removed ``utils`` module. ``utils.tuning`` and ``utils.prefactor`` are now in the root
   of the package. ``utils.splines`` is now in the ``lib`` module.
+* The tuning now uses a grid-search based scheme, instead of a gradient based scheme.
+* The tuning functions no longer takes the ``cutoff`` parameter, and thus does not
+  support a built-in NL calculation.
 
 Fixed
 #####

docs/src/references/index.rst

@@ -22,7 +22,7 @@ refer to the :ref:`userdoc-how-to` section.
 
     potentials/index
     calculators/index
-    tuning
+    tuning/index
     prefactors
     metatensor
     lib/index

docs/src/references/tuning/base_classes.rst

@@ -0,0 +1,52 @@
+Base Classes
+############
+
+Current scheme behind all tuning functions is grid-searching based, focusing on the Fourier
+space parameters like ``lr_wavelength``, ``mesh_spacing`` and ``interpolation_nodes``.
+For real space parameter ``cutoff``, it is treated as a hyperparameter here, which
+should be manually specified by the user. The parameter ``smearing`` is determined by
+the real space error formula and is set to achieve a real space error of
+``desired_accuracy / 4``.
+
+The Fourier space parameters are all discrete, so it's convenient to do the grid-search.
+Default searching-ranges are provided for those parameters. For ``lr_wavelength``, the
+values are chosen to be with a minimum of 1 and a maximum of 13 mesh points in each
+spatial direction ``(x, y, z)``. For ``mesh_spacing``, the values are set to have
+minimally 2 and maximally 7 mesh points in each spatial direction, for both the P3M and
+PME method. The values of ``interpolation_nodes`` are the same as those supported in
+:class:`torchpme.lib.MeshInterpolator`.
+
+In the grid-searching, all possible parameter combinations are evaluated. The error
+associated with the parameter is estimated by the error formulas implemented in the
+subclasses of :class:`torchpme.tuning.tuner.TuningErrorBounds`. Parameter with
+the error within the desired accuracy are benchmarked for computational time by
+:class:`torchpme.tuning.tuner.TuningTimings` The timing of the other parameters are
+not tested and set to infinity.
+
+The return of these tuning functions contains the ``smearing`` and a dictionary, in
+which there is parameter for the Fourier space. The parameter is that of the desired
+accuracy and the shortest timing. The parameter of the smallest error will be returned
+in the case that no parameter can fulfill the accuracy requirement.
+
+
+.. autoclass:: torchpme.tuning.tuner.TunerBase
+    :members:
+
+.. autoclass:: torchpme.tuning.tuner.GridSearchTuner
+    :members:
+
+.. autoclass:: torchpme.tuning.tuner.TuningTimings
+    :members:
+
+.. autoclass:: torchpme.tuning.tuner.TuningErrorBounds
+    :members:
+
+Examples using Tuning Classes
+-----------------------------
+
+.. minigallery::
+
+    torchpme.tuning.tuner.TunerBase
+    torchpme.tuning.tuner.GridSearchTuner
+    torchpme.tuning.tuner.TuningTimings
+    torchpme.tuning.tuner.TuningErrorBounds

docs/src/references/tuning/index.rst

@@ -0,0 +1,24 @@
+Tuning
+######
+
+The choice of parameters like the neighborlist ``cutoff``, the ``smearing`` or the
+``lr_wavelength``/``mesh_spacing`` has a large influence one the accuracy of the
+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 :class:`Coulomb potential
+   <torchpme.CoulombPotential>`.
+
+.. toctree::
+   :maxdepth: 1
+   :glob:
+
+   ./*

docs/src/references/tuning/tune_ewald.rst

@@ -0,0 +1,24 @@
+Tune Ewald
+##########
+
+The tuning is based on the following error formulas:
+
+.. math::
+      \Delta F_\mathrm{real}
+        \approx \frac{Q^2}{\sqrt{N}}
+                \frac{2}{\sqrt{r_{\text{cutoff}} V}}
+                e^{-r_{\text{cutoff}}^2 / 2 \sigma^2}
+
+.. math::
+    \Delta F_\mathrm{Fourier}^\mathrm{Ewald}
+        \approx \frac{Q^2}{\sqrt{N}}
+                \frac{\sqrt{2} / \sigma}{\pi\sqrt{2 V / h}} e^{-2\pi^2 \sigma^2 / h ^ 2}
+
+where :math:`N` is the number of charges, :math:`Q^2 = \sum_{i = 1}^N q_i^2`, is the sum of squared
+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.tune_ewald
+
+.. autoclass:: torchpme.tuning.ewald.EwaldErrorBounds
+    :members:

docs/src/references/tuning/tune_p3m.rst

@@ -0,0 +1,27 @@
+Tune P3M
+#########
+
+The tuning is based on the following error formulas:
+
+.. math::
+      \Delta F_\mathrm{real}
+        \approx \frac{Q^2}{\sqrt{N}}
+                \frac{2}{\sqrt{r_{\text{cutoff}} V}}
+                e^{-r_{\text{cutoff}}^2 / 2 \sigma^2}
+
+.. math::
+    \Delta F_\mathrm{Fourier}^\mathrm{P3M}
+        \approx \frac{Q^2}{L^2}(\frac{\sqrt{2}H}{\sigma})^p
+                    \sqrt{\frac{\sqrt{2}L}{N\sigma}
+                    \sqrt{2\pi}\sum_{m=0}^{p-1}a_m^{(p)}(\frac{\sqrt{2}H}{\sigma})^{2m}}
+
+where :math:`N` is the number of charges, :math:`Q^2 = \sum_{i = 1}^N q_i^2`, is the sum of squared
+charges, :math:`r_{\text{cutoff}}` is the short-range cutoff, :math:`V` is the volume of the
+simulation box, :math:`p` is the order of the interpolation scheme, :math:`H` is the spacing of mesh
+points and :math:`a_m^{(p)}` is an expansion coefficient.
+
+
+.. autofunction:: torchpme.tuning.tune_p3m
+
+.. autoclass:: torchpme.tuning.p3m.P3MErrorBounds
+    :members:

docs/src/references/tuning/tune_pme.rst

@@ -0,0 +1,27 @@
+Tune PME
+#########
+
+The tuning is based on the following error formulas:
+
+.. math::
+      \Delta F_\mathrm{real}
+        \approx \frac{Q^2}{\sqrt{N}}
+                \frac{2}{\sqrt{r_{\text{cutoff}} V}}
+                e^{-r_{\text{cutoff}}^2 / 2 \sigma^2}
+
+.. math::
+    \Delta F_\mathrm{Fourier}^\mathrm{PME}
+        \approx 2\pi^{1/4}\sqrt{\frac{3\sqrt{2} / \sigma}{N(2p+3)}}
+                \frac{Q^2}{L^2}\frac{(\sqrt{2}H/\sigma)^{p+1}}{(p+1)!} \times
+                \exp{\frac{(p+1)[\log{(p+1)} - \log 2 - 1]}{2}} \left< \phi_p^2 \right> ^{1/2}
+
+where :math:`N` is the number of charges, :math:`Q^2 = \sum_{i = 1}^N q_i^2`, is the sum of squared
+charges, :math:`r_{\text{cutoff}}` is the short-range cutoff, :math:`V` is the volume of the
+simulation box, :math:`p` is the order of the interpolation scheme, :math:`H` is the spacing of mesh
+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.tune_pme
+
+.. autoclass:: torchpme.tuning.pme.PMEErrorBounds
+    :members:

examples/1-charges-example.py → examples/01-charges-example.py

@@ -37,13 +37,15 @@
 from metatensor.torch.atomistic import NeighborListOptions, System
 
 import torchpme
+from torchpme.tuning import tune_pme
 
 # %%
 #
 # Create the properties CsCl unit cell
 
 symbols = ("Cs", "Cl")
 types = torch.tensor([55, 17])
+charges = torch.tensor([[1.0], [-1.0]], dtype=torch.float64)
 positions = torch.tensor([(0, 0, 0), (0.5, 0.5, 0.5)], dtype=torch.float64)
 cell = torch.eye(3, dtype=torch.float64)
 pbc = torch.tensor([True, True, True])
@@ -55,8 +57,21 @@
 # The ``sum_squared_charges`` is equal to ``2.0`` becaue each atom either has a charge
 # of 1 or -1 in units of elementary charges.
 
-smearing, pme_params, cutoff = torchpme.tuning.tune_pme(
-    sum_squared_charges=2.0, cell=cell, positions=positions
+cutoff = 4.4
+nl = vesin.torch.NeighborList(cutoff=cutoff, full_list=False)
+neighbor_indices, neighbor_distances = nl.compute(
+    points=positions.to(dtype=torch.float64, device="cpu"),
+    box=cell.to(dtype=torch.float64, device="cpu"),
+    periodic=True,
+    quantities="Pd",
+)
+smearing, pme_params, _ = tune_pme(
+    charges=charges,
+    cell=cell,
+    positions=positions,
+    cutoff=cutoff,
+    neighbor_indices=neighbor_indices,
+    neighbor_distances=neighbor_distances,
 )
 
 # %%

examples/2-neighbor-lists-usage.py → examples/02-neighbor-lists-usage.py

@@ -46,6 +46,7 @@
 import vesin.torch
 
 import torchpme
+from torchpme.tuning import tune_pme
 
 # %%
 #
@@ -92,9 +93,22 @@
 cell = torch.from_numpy(atoms.cell.array)
 
 sum_squared_charges = float(torch.sum(charges**2))
+cutoff = 4.4
+nl = vesin.torch.NeighborList(cutoff=cutoff, full_list=False)
+neighbor_indices, neighbor_distances = nl.compute(
+    points=positions.to(dtype=torch.float64, device="cpu"),
+    box=cell.to(dtype=torch.float64, device="cpu"),
+    periodic=True,
+    quantities="Pd",
+)
 
-smearing, pme_params, cutoff = torchpme.tuning.tune_pme(
-    sum_squared_charges=sum_squared_charges, cell=cell, positions=positions
+smearing, pme_params, _ = tune_pme(
+    charges=charges,
+    cell=cell,
+    positions=positions,
+    cutoff=cutoff,
+    neighbor_indices=neighbor_indices,
+    neighbor_distances=neighbor_distances,
 )
 
 # %%

examples/5-autograd-demo.py → examples/05-autograd-demo.py

@@ -17,6 +17,8 @@
 exercise to the reader.
 """
 
+# %%
+
 from time import time
 
 import ase
@@ -477,10 +479,11 @@ def forward(self, positions, cell, charges):
 )
 
 # %%
-# We can also time the difference in execution
+# We can also evaluate the difference in execution
 # time between the Pytorch and scripted versions of the
 # module (depending on the system, the relative efficiency
-# of the two evaluations could go either way!)
+# of the two evaluations could go either way, as this is
+# a too small system to make a difference!)
 
 duration = 0.0
 for _i in range(20):
@@ -513,5 +516,3 @@ def forward(self, positions, cell, charges):
 
 # %%
 print(f"Evaluation time:\nPytorch: {time_python}ms\nJitted:  {time_jit}ms")
-
-# %%