.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/2-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_2-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-50 .. 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 .. GENERATED FROM PYTHON SOURCE LINES 51-55 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 56-68 .. code-block:: Python 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 69-71 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 72-83 .. code-block:: Python rng = np.random.default_rng(42) displacement = rng.normal(loc=0.0, scale=0.1, size=(len(atoms), 3)) atoms.positions += displacement chemiscope.show( frames=[atoms], mode="structure", settings=chemiscope.quick_settings(structure_settings={"unitCell": True}), ) .. chemiscope:: _datasets/fig_2-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-102 .. 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)) smearing, pme_params, cutoff = torchpme.utils.tune_pme( sum_squared_charges=sum_squared_charges, cell=cell, positions=positions ) .. GENERATED FROM PYTHON SOURCE LINES 103-104 The tuning found the following best values for our system. .. GENERATED FROM PYTHON SOURCE LINES 105-110 .. 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.2083905187289359 PME parameters: {'mesh_spacing': 0.7851426755551525, 'interpolation_nodes': 4} cutoff: 5.534948153534669 .. GENERATED FROM PYTHON SOURCE LINES 111-116 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 117-123 .. code-block:: Python nl = vesin.NeighborList(cutoff=cutoff, full_list=False) i, j, S = nl.compute( points=atoms.positions, box=atoms.cell.array, periodic=True, quantities="ijS" ) .. GENERATED FROM PYTHON SOURCE LINES 124-126 We now define a function that (re-)computes the distances in a way that torch can track these operations. .. GENERATED FROM PYTHON SOURCE LINES 127-155 .. 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 156-158 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 159-169 .. code-block:: Python positions.requires_grad = True i = torch.from_numpy(i.astype(int)) j = torch.from_numpy(j.astype(int)) neighbor_indices = torch.stack([i, j], dim=1) neighbor_shifts = torch.from_numpy(S) .. GENERATED FROM PYTHON SOURCE LINES 170-171 Now, we start to re-compute the distances .. GENERATED FROM PYTHON SOURCE LINES 172-180 .. code-block:: Python neighbor_distances = distances( positions=positions, neighbor_indices=neighbor_indices, cell=cell, neighbor_shifts=neighbor_shifts, ) .. GENERATED FROM PYTHON SOURCE LINES 181-183 and initialize a :class:`PMECalculator` instance using a :class:`CoulombPotential` to compute the potential. .. GENERATED FROM PYTHON SOURCE LINES 184-198 .. 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([[-1.2624], [ 1.0095], [-1.0479], [ 0.9498], [-1.0533], [ 1.1807], [-0.8604], [ 1.0190], [-1.2350], [ 0.9824], [-1.0189], [ 0.9397], [-1.0748], [ 1.1874], [-0.7674], [ 1.0140]], dtype=torch.float64, grad_fn=) .. GENERATED FROM PYTHON SOURCE LINES 199-200 The energy is given by the scalar product of the potential with the charges. .. GENERATED FROM PYTHON SOURCE LINES 201-204 .. code-block:: Python energy = charges.T @ potential .. GENERATED FROM PYTHON SOURCE LINES 205-206 Finally, we can compute and print the forces in CGS units as erg/Å. .. GENERATED FROM PYTHON SOURCE LINES 207-212 .. 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.9093, 0.1146, 0.1891], [ 0.9286, -0.4931, -0.1914], [ 0.5059, -0.5171, 0.4208], [ 0.3930, 1.4463, -0.7444], [ 0.0229, -0.1376, -0.2500], [-0.9241, -0.7591, -0.4926], [-0.7989, -0.0829, 0.4588], [-0.1867, -0.4855, -0.4608], [-1.2066, 1.0370, 0.4063], [ 0.1030, 0.0138, 0.9360], [ 0.5667, -0.1388, -0.2156], [-1.2830, 0.9800, 0.0463], [ 0.5705, -1.5602, -0.5621], [-0.0867, 0.7624, 0.6774], [ 0.7111, -0.2498, 0.1696], [-0.2251, 0.0702, -0.3875]], dtype=torch.float64) .. GENERATED FROM PYTHON SOURCE LINES 213-222 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 223-227 .. code-block:: Python positions_new = positions.detach().clone() positions_new.requires_grad = True .. GENERATED FROM PYTHON SOURCE LINES 228-229 and create new distances in a similar manner as above. .. GENERATED FROM PYTHON SOURCE LINES 230-236 .. code-block:: Python nl = vesin.torch.NeighborList(cutoff=1.0, full_list=False) i, j, d = nl.compute(points=positions_new, box=cell, periodic=True, quantities="ijd") neighbor_indices_new = torch.stack([i, j], dim=1) .. GENERATED FROM PYTHON SOURCE LINES 237-238 Following the same steps as above, we compute the forces. .. GENERATED FROM PYTHON SOURCE LINES 239-254 .. 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([[ 5.8156e-01, 4.0017e-01, -1.1581e-01], [ 8.6043e-01, -1.5703e+00, -4.7596e-01], [ 1.5938e+00, -2.1923e-01, 5.6314e-01], [ 9.9873e-01, 2.2980e+00, -9.5430e-01], [-4.0018e-01, -2.5289e-01, -1.3555e+00], [-8.3079e-01, -1.7730e+00, -2.3673e-01], [-1.6350e+00, 3.8185e-01, 1.1528e+00], [ 7.7320e-02, -1.7556e+00, -3.5352e-01], [-5.9445e-01, 3.5112e-01, -2.1645e-01], [ 9.9394e-01, 3.9357e-01, 1.4560e+00], [ 1.1421e+00, 9.5488e-01, 5.5840e-02], [-3.5063e+00, 1.9421e+00, -1.5203e-01], [ 1.9623e-04, -6.7670e-01, 7.3247e-02], [-2.7440e-01, 4.9361e-01, 3.1006e-02], [ 1.7108e+00, -8.3318e-02, -7.3458e-02], [-7.1777e-01, -8.8422e-01, 6.0166e-01]], dtype=torch.float64) .. GENERATED FROM PYTHON SOURCE LINES 255-257 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 258-267 .. 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_2-neighbor-lists-usage_001.png :alt: 2 neighbor lists usage :srcset: /examples/images/sphx_glr_2-neighbor-lists-usage_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 7.177 seconds) .. _sphx_glr_download_examples_2-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: 2-neighbor-lists-usage.ipynb <2-neighbor-lists-usage.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 2-neighbor-lists-usage.py <2-neighbor-lists-usage.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 2-neighbor-lists-usage.zip <2-neighbor-lists-usage.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_