diff --git a/docs/src/references/tuning/index.rst b/docs/src/references/tuning/index.rst index 73f3ad2e..85610845 100644 --- a/docs/src/references/tuning/index.rst +++ b/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 diff --git a/docs/src/references/tuning/tune_ewald.rst b/docs/src/references/tuning/tune_ewald.rst index f50e7cee..120a0d25 100644 --- a/docs/src/references/tuning/tune_ewald.rst +++ b/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: diff --git a/docs/src/references/tuning/tune_p3m.rst b/docs/src/references/tuning/tune_p3m.rst index 8b958cbe..93e7c05b 100644 --- a/docs/src/references/tuning/tune_p3m.rst +++ b/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: diff --git a/docs/src/references/tuning/tune_pme.rst b/docs/src/references/tuning/tune_pme.rst index e9efca3d..aa4a1ea2 100644 --- a/docs/src/references/tuning/tune_pme.rst +++ b/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: diff --git a/examples/05-autograd-demo.py b/examples/05-autograd-demo.py index c6e594f4..b44416be 100644 --- a/examples/05-autograd-demo.py +++ b/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) `_) -# 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} -""" -) diff --git a/examples/11-4-site-water.py b/examples/11-4-site-water.py new file mode 100644 index 00000000..58e19b3e --- /dev/null +++ b/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) `_) +# 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} +""" +) diff --git a/src/torchpme/calculators/calculator.py b/src/torchpme/calculators/calculator.py index 4c8798a9..627283e5 100644 --- a/src/torchpme/calculators/calculator.py +++ b/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