.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/02-neighbor-lists-usage.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_examples_02-neighbor-lists-usage.py: Advanced neighbor list usage ============================ .. currentmodule:: torchpme Accurately calculating forces as derivatives from energy is crucial for predicting system dynamics as well as in training machine learning models. In systems where forces are derived from the gradients of the potential energy, it is essential that the distance calculations between particles are included in the computational graph. This ensures that the force computations respect the dependencies between particle positions and distances, allowing for accurate gradients during backpropagation. .. figure:: ../../static/images/backprop-path.* Visualization of the data flow to compute the energy from the ``cell``, ``positions`` and ``charges`` through a neighborlist calculator and the potential calculator. All operations on the red line have to be tracked to obtain the correct computation of derivatives on the ``positions``. In this tutorial, we demonstrate two methods for maintaining differentiability when computing distances between particles. The **first method** manually recomputes ``distances`` within the computational graph using ``positions``, ``cell`` information, and neighbor shifts, making it suitable for any neighbor list code. The **second method** uses a backpropagable neighbor list from the `vesin-torch `_ library, which automatically ensures that the distance calculations remain differentiable. .. note:: While both approaches yield the same result, a backpropagable neighbor list is generally preferred because it eliminates the need to manually recompute distances. This not only simplifies your code but also improves performance. .. GENERATED FROM PYTHON SOURCE LINES 38-51 .. code-block:: Python from typing import Optional import ase import chemiscope import matplotlib.pyplot as plt import numpy as np import torch import vesin import vesin.torch import torchpme from torchpme.tuning import tune_pme .. GENERATED FROM PYTHON SOURCE LINES 52-56 The test system --------------- As a test system, we use a 2x2x2 supercell of an CsCl crystal in a cubic cell. .. GENERATED FROM PYTHON SOURCE LINES 57-70 .. code-block:: Python dtype = torch.float64 atoms_unitcell = ase.Atoms( symbols=["Cs", "Cl"], positions=np.array([(0, 0, 0), (0.5, 0.5, 0.5)]), cell=np.eye(3), pbc=torch.tensor([True, True, True]), ) charges_unitcell = np.array([1.0, -1.0]) atoms = atoms_unitcell.repeat([2, 2, 2]) charges = np.tile(charges_unitcell, 2 * 2 * 2) .. GENERATED FROM PYTHON SOURCE LINES 71-73 We now slightly displace the atoms from their initial positions randomly based on a Gaussian distribution with a width of 0.1 Å to create non-zero forces. .. GENERATED FROM PYTHON SOURCE LINES 74-83 .. code-block:: Python atoms.rattle(stdev=0.1) chemiscope.show( frames=[atoms], mode="structure", settings=chemiscope.quick_settings(structure_settings={"unitCell": True}), ) .. chemiscope:: _datasets/fig_02-neighbor-lists-usage_001.json.gz :mode: structure .. raw:: html


.. GENERATED FROM PYTHON SOURCE LINES 84-90 Tune paramaters --------------- Based on our system we will first *tune* the PME parameters for an accurate computation. We first convert the ``positions``, ``charges`` and the ``cell`` from NumPy arrays into torch tensors and compute the summed squared charges. .. GENERATED FROM PYTHON SOURCE LINES 91-115 .. code-block:: Python positions = torch.from_numpy(atoms.positions) charges = torch.from_numpy(charges).unsqueeze(1) 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=dtype, 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, ) .. GENERATED FROM PYTHON SOURCE LINES 116-117 The tuning found the following best values for our system. .. GENERATED FROM PYTHON SOURCE LINES 118-123 .. code-block:: Python print("smearing:", smearing) print("PME parameters:", pme_params) print("cutoff:", cutoff) .. rst-class:: sphx-glr-script-out .. code-block:: none smearing: 1.1069526756106463 PME parameters: {'interpolation_nodes': 4, 'mesh_spacing': 0.5714285714285714} cutoff: 4.4 .. GENERATED FROM PYTHON SOURCE LINES 124-129 Generic Neighborlist -------------------- One usual workflow is to compute the distance vectors using default tools like the the default (NumPy) version of the vesin neighbor list. .. GENERATED FROM PYTHON SOURCE LINES 130-136 .. code-block:: Python nl = vesin.NeighborList(cutoff=cutoff, full_list=False) neighbor_indices, S = nl.compute( points=atoms.positions, box=atoms.cell.array, periodic=True, quantities="PS" ) .. GENERATED FROM PYTHON SOURCE LINES 137-139 We now define a function that (re-)computes the distances in a way that torch can track these operations. .. GENERATED FROM PYTHON SOURCE LINES 140-168 .. code-block:: Python def distances( positions: torch.Tensor, neighbor_indices: torch.Tensor, cell: Optional[torch.Tensor] = None, neighbor_shifts: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Compute pairwise distances.""" atom_is = neighbor_indices[:, 0] atom_js = neighbor_indices[:, 1] pos_is = positions[atom_is] pos_js = positions[atom_js] distance_vectors = pos_js - pos_is if cell is not None and neighbor_shifts is not None: shifts = neighbor_shifts.type(cell.dtype) distance_vectors += shifts @ cell elif cell is not None and neighbor_shifts is None: raise ValueError("Provided `cell` but no `neighbor_shifts`.") elif cell is None and neighbor_shifts is not None: raise ValueError("Provided `neighbor_shifts` but no `cell`.") return torch.linalg.norm(distance_vectors, dim=1) .. GENERATED FROM PYTHON SOURCE LINES 169-171 To use this function we now the tracking of operations by setting the :attr:`requires_grad ` property to :obj:`True`. .. GENERATED FROM PYTHON SOURCE LINES 172-180 .. code-block:: Python positions.requires_grad = True neighbor_indices = torch.from_numpy(neighbor_indices.astype(int)) neighbor_shifts = torch.from_numpy(S) .. GENERATED FROM PYTHON SOURCE LINES 181-182 Now, we start to re-compute the distances .. GENERATED FROM PYTHON SOURCE LINES 183-191 .. code-block:: Python neighbor_distances = distances( positions=positions, neighbor_indices=neighbor_indices, cell=cell, neighbor_shifts=neighbor_shifts, ) .. GENERATED FROM PYTHON SOURCE LINES 192-194 and initialize a :class:`PMECalculator` instance using a :class:`CoulombPotential` to compute the potential. .. GENERATED FROM PYTHON SOURCE LINES 195-210 .. code-block:: Python pme = torchpme.PMECalculator( potential=torchpme.CoulombPotential(smearing=smearing), **pme_params, ) potential = pme( charges=charges, cell=cell, positions=positions, neighbor_indices=neighbor_indices, neighbor_distances=neighbor_distances, ) print(potential) .. rst-class:: sphx-glr-script-out .. code-block:: none tensor([[-0.8874], [ 1.2554], [-0.9274], [ 1.4907], [-0.7261], [ 0.7836], [-0.8839], [ 0.8997], [-0.9603], [ 1.0151], [-1.0800], [ 1.3655], [-1.0228], [ 0.6690], [-1.1581], [ 0.8331]], dtype=torch.float64, grad_fn=) .. GENERATED FROM PYTHON SOURCE LINES 211-212 The energy is given by the scalar product of the potential with the charges. .. GENERATED FROM PYTHON SOURCE LINES 213-216 .. code-block:: Python energy = charges.T @ potential .. GENERATED FROM PYTHON SOURCE LINES 217-218 Finally, we can compute and print the forces in CGS units as erg/Å. .. GENERATED FROM PYTHON SOURCE LINES 219-224 .. code-block:: Python forces = torch.autograd.grad(-1.0 * energy, positions)[0] print(forces) .. rst-class:: sphx-glr-script-out .. code-block:: none tensor([[-0.9484, -1.4864, 0.9129], [ 0.1350, -0.6500, 0.4235], [-0.0230, -1.4964, -0.1962], [-0.3949, -0.3494, 0.1579], [-0.0038, 0.9212, 1.0006], [ 0.2718, 0.2613, -0.4679], [-1.2057, 0.9323, -1.0406], [-0.3698, 0.5807, 0.5612], [ 0.5696, -1.1072, 0.0541], [ 0.6476, 0.4007, -0.4058], [ 0.5322, 0.1812, -0.3541], [ 0.5632, 0.4293, 0.3243], [ 0.4923, 0.5976, -0.6239], [ 0.0734, 0.1513, 0.0368], [-0.3666, 0.8548, -0.5614], [ 0.0273, -0.2203, 0.1787]], dtype=torch.float64) .. GENERATED FROM PYTHON SOURCE LINES 225-234 Backpropagable Neighborlist --------------------------- We now repeat the computation of the forces, but instead of using a generic neighbor list and our custom ``distances`` function, we directly use a neighbor list function that tracks the operations, as implemented by the ``vesin-torch`` library. We first ``detach`` and ``clone`` the position tensor to create a new computational graph .. GENERATED FROM PYTHON SOURCE LINES 235-239 .. code-block:: Python positions_new = positions.detach().clone() positions_new.requires_grad = True .. GENERATED FROM PYTHON SOURCE LINES 240-241 and create new distances in a similar manner as above. .. GENERATED FROM PYTHON SOURCE LINES 242-248 .. code-block:: Python nl = vesin.torch.NeighborList(cutoff=cutoff, full_list=False) neighbor_indices_new, d = nl.compute( points=positions_new, box=cell, periodic=True, quantities="Pd" ) .. GENERATED FROM PYTHON SOURCE LINES 249-250 Following the same steps as above, we compute the forces. .. GENERATED FROM PYTHON SOURCE LINES 251-266 .. code-block:: Python potential_new = pme( charges=charges, cell=cell, positions=positions_new, neighbor_indices=neighbor_indices_new, neighbor_distances=d, ) energy_new = charges.T @ potential_new forces_new = torch.autograd.grad(-1.0 * energy_new, positions_new)[0] print(forces_new) .. rst-class:: sphx-glr-script-out .. code-block:: none tensor([[-0.9484, -1.4864, 0.9129], [ 0.1350, -0.6500, 0.4235], [-0.0230, -1.4964, -0.1962], [-0.3949, -0.3494, 0.1579], [-0.0038, 0.9212, 1.0006], [ 0.2718, 0.2613, -0.4679], [-1.2057, 0.9323, -1.0406], [-0.3698, 0.5807, 0.5612], [ 0.5696, -1.1072, 0.0541], [ 0.6476, 0.4007, -0.4058], [ 0.5322, 0.1812, -0.3541], [ 0.5632, 0.4293, 0.3243], [ 0.4923, 0.5976, -0.6239], [ 0.0734, 0.1513, 0.0368], [-0.3666, 0.8548, -0.5614], [ 0.0273, -0.2203, 0.1787]], dtype=torch.float64) .. GENERATED FROM PYTHON SOURCE LINES 267-269 The forces are the same as those we printed above. For better comparison, we can also plot the scalar force for each method. .. GENERATED FROM PYTHON SOURCE LINES 270-279 .. code-block:: Python plt.plot(torch.linalg.norm(forces, dim=1), "o-", label="normal Neighborlist") plt.plot(torch.linalg.norm(forces_new, dim=1), ".-", label="torch Neighborlist") plt.legend() plt.xlabel("atom index") plt.ylabel(r"$|F|~/~\mathrm{erg\,Å^{-1}}$") plt.show() .. image-sg:: /examples/images/sphx_glr_02-neighbor-lists-usage_001.png :alt: 02 neighbor lists usage :srcset: /examples/images/sphx_glr_02-neighbor-lists-usage_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 3.759 seconds) .. _sphx_glr_download_examples_02-neighbor-lists-usage.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 02-neighbor-lists-usage.ipynb <02-neighbor-lists-usage.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 02-neighbor-lists-usage.py <02-neighbor-lists-usage.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 02-neighbor-lists-usage.zip <02-neighbor-lists-usage.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_