Revamped tuning

docs/src/references/utils/tuning.rst

@@ -12,8 +12,11 @@ than the given accuracy. Because these methods are gradient-based, be sure to pa
 attention to the ``learning_rate`` and ``max_steps`` parameter. A good choice of these
 two parameters can enhance the optimization speed and performance.
 
-.. autoclass:: torchpme.utils.tune_ewald
+.. autoclass:: torchpme.utils.tuning.ewald.EwaldTuner
     :members:
 
-.. autoclass:: torchpme.utils.tune_pme
+.. autoclass:: torchpme.utils.tuning.pme.PMETuner
+    :members:
+
+.. autoclass:: torchpme.utils.tuning.p3m.P3MTuner
     :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.utils.tuning.pme import PMETuner
 
 # %%
 #
 # 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,9 +57,9 @@
 # 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.utils.tune_pme(
-    sum_squared_charges=2.0, cell=cell, positions=positions
-)
+smearing, pme_params, cutoff = PMETuner(
+    charges=charges, cell=cell, positions=positions, cutoff=4.4
+).tune()
 
 # %%
 #

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

@@ -46,6 +46,7 @@
 import vesin.torch
 
 import torchpme
+from torchpme.utils.tuning.pme import PMETuner
 
 # %%
 #
@@ -93,9 +94,9 @@
 
 sum_squared_charges = float(torch.sum(charges**2))
 
-smearing, pme_params, cutoff = torchpme.utils.tune_pme(
-    sum_squared_charges=sum_squared_charges, cell=cell, positions=positions
-)
+smearing, pme_params, cutoff = PMETuner(
+    charges=charges, cell=cell, positions=positions, cutoff=4.4
+).tune()
 
 # %%
 #

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):
@@ -515,3 +518,82 @@ 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}
+"""
+)