Add mixed periodic boundary condition support to EwaldCalculator (#200)

* Add mixed periodic boundary condition support to EwaldCalculator

* Fix typos

* Another bug fix

* bug fixes + corrections for charged cells

* Yet another bug correction

* small cleaning

* Move periodic param from init to forward

* Add support for single non-periodic axis in Ewald summation

* Fix existing tests and lint

* Fix metatensor

* Add tests for slabs and fix docs

* Update changelog

* Move slab correction to the potential definiton

* Fix tests

* Rename _2d_correction to pbc_correction and update documentation

* typo fix

Author: E-Rum

docs/src/references/changelog.rst

@@ -27,6 +27,8 @@ changelog <https://keepachangelog.com/en/1.1.0/>`_ format. This project follows
 Added
 #####
 
+* Add support for slab geometries in Ewald and PME calculators
+
 Fixed
 #####
 

examples/01-charges-example.py

@@ -31,8 +31,8 @@
 # %%
 
 import torch
+import vesin.metatomic
 import vesin.torch
-import vesin.torch.metatensor
 from metatensor.torch import Labels, TensorBlock, TensorMap
 from metatomic.torch import NeighborListOptions, System
 
@@ -231,7 +231,7 @@
 # the cutoff and the type of list.
 
 options = NeighborListOptions(cutoff=4.0, full_list=True, strict=False)
-nl_mts = vesin.torch.metatensor.NeighborList(options, length_unit="Angstrom")
+nl_mts = vesin.metatomic.NeighborList(options, length_unit="Angstrom")
 neighbors = nl_mts.compute(system)
 
 # %%

examples/07-lode-demo.py

@@ -442,6 +442,7 @@ def forward(
         positions: torch.Tensor,
         neighbor_indices: Optional[torch.Tensor] = None,
         neighbor_distances: Optional[torch.Tensor] = None,
+        periodic: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         # Update meshes
         assert self.potential.smearing is not None  # otherwise mypy complains

src/torchpme/_utils.py

@@ -10,6 +10,7 @@ def _validate_parameters(
     neighbor_indices: torch.Tensor,
     neighbor_distances: torch.Tensor,
     smearing: Union[float, None],
+    periodic: Union[torch.Tensor, None] = None,
 ) -> None:
     dtype = positions.dtype
     device = positions.device
@@ -106,3 +107,16 @@ def _validate_parameters(
             f"type of `neighbor_distances` ({neighbor_distances.dtype}) must be same "
             f"as that of the `positions` class ({dtype})"
         )
+
+    if periodic is not None:
+        if periodic.shape != (3,):
+            raise ValueError(
+                "`periodic` must be a tensor of shape (3,), got "
+                f"tensor with shape {list(periodic.shape)}"
+            )
+
+        if periodic.device != device:
+            raise ValueError(
+                f"device of `periodic` ({periodic.device}) must be same as that of "
+                f"the `positions` class ({device})"
+            )

src/torchpme/calculators/calculator.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 from torch import profiler
 
@@ -106,6 +108,7 @@ def forward(
         positions: torch.Tensor,
         neighbor_indices: torch.Tensor,
         neighbor_distances: torch.Tensor,
+        periodic: Optional[torch.Tensor] = None,
     ):
         r"""
         Compute the potential "energy".
@@ -139,6 +142,9 @@ def forward(
             which the potential should be computed in real space.
         :param neighbor_distances: torch.tensor with the pair distances of the neighbors
             for which the potential should be computed in real space.
+        :param periodic: optional torch.tensor of shape ``(3,)`` indicating which
+            directions are periodic (True) and which are not (False). If not
+            provided, full periodicity is assumed.
         """
         _validate_parameters(
             charges=charges,
@@ -147,6 +153,7 @@ def forward(
             neighbor_indices=neighbor_indices,
             neighbor_distances=neighbor_distances,
             smearing=self.potential.smearing,
+            periodic=periodic,
         )
 
         # Compute short-range (SR) part using a real space sum
@@ -163,6 +170,7 @@ def forward(
             charges=charges,
             cell=cell,
             positions=positions,
+            periodic=periodic,
         )
 
         return self.prefactor * (potential_sr + potential_lr)

src/torchpme/calculators/ewald.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 
 from ..lib import generate_kvectors_for_ewald
@@ -12,7 +14,7 @@ class EwaldCalculator(Calculator):
     Scaling as :math:`\mathcal{O}(N^2)` with respect to the number of particles
     :math:`N`.
 
-    For getting reasonable values for the ``smaring`` of the potential class and  the
+    For getting reasonable values for the ``smearing`` of the potential class and  the
     ``lr_wavelength`` based on a given accuracy for a specific structure you should use
     :func:`torchpme.tuning.tune_ewald`. This function will also find the optimal
     ``cutoff`` for the  **neighborlist**.
@@ -78,6 +80,7 @@ def _compute_kspace(
         charges: torch.Tensor,
         cell: torch.Tensor,
         positions: torch.Tensor,
+        periodic: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         # Define k-space cutoff from required real-space resolution
         k_cutoff = 2 * torch.pi / self.lr_wavelength
@@ -131,4 +134,5 @@ def _compute_kspace(
         prefac = self.potential.background_correction()
         energy -= 2 * prefac * charge_tot * ivolume
         # Compensate for double counting of pairs (i,j) and (j,i)
+        energy += self.potential.pbc_correction(periodic, positions, cell, charges)
         return energy / 2

src/torchpme/calculators/pme.py

@@ -1,5 +1,6 @@
+from typing import Optional
+
 import torch
-from torch import profiler
 
 from ..lib.kspace_filter import KSpaceFilter
 from ..lib.kvectors import get_ns_mesh
@@ -97,51 +98,48 @@ def _compute_kspace(
         charges: torch.Tensor,
         cell: torch.Tensor,
         positions: torch.Tensor,
+        periodic: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         # TODO: Kernel function `G` and initialization of `MeshInterpolator` only depend
         # on `cell`. Caching may save up to 15% but issues with AD need to be resolved.
 
-        with profiler.record_function("init 0: preparation"):
-            # Compute number of times each basis vector of the reciprocal space can be
-            # scaled until the cutoff is reached
-            ns = get_ns_mesh(cell, self.mesh_spacing)
+        # Compute number of times each basis vector of the reciprocal space can be
+        # scaled until the cutoff is reached
+        ns = get_ns_mesh(cell, self.mesh_spacing)
 
-        with profiler.record_function("init 1: update mesh interpolator"):
-            self.mesh_interpolator.update(cell, ns)
+        self.mesh_interpolator.update(cell, ns)
 
-        with profiler.record_function("update the mesh for the k-space filter"):
-            self.kspace_filter.update(cell, ns)
+        self.kspace_filter.update(cell, ns)
 
-        with profiler.record_function("step 1: compute density interpolation"):
-            self.mesh_interpolator.compute_weights(positions)
-            rho_mesh = self.mesh_interpolator.points_to_mesh(particle_weights=charges)
+        self.mesh_interpolator.compute_weights(positions)
+        rho_mesh = self.mesh_interpolator.points_to_mesh(particle_weights=charges)
 
-        with profiler.record_function("step 2: perform actual convolution using FFT"):
-            potential_mesh = self.kspace_filter.forward(rho_mesh)
+        potential_mesh = self.kspace_filter.forward(rho_mesh)
 
-        with profiler.record_function("step 3: back interpolation + volume scaling"):
-            ivolume = torch.abs(cell.det()).pow(-1)
-            interpolated_potential = (
-                self.mesh_interpolator.mesh_to_points(potential_mesh) * ivolume
-            )
+        ivolume = torch.abs(cell.det()).pow(-1)
+        interpolated_potential = (
+            self.mesh_interpolator.mesh_to_points(potential_mesh) * ivolume
+        )
 
-        with profiler.record_function("step 4: remove the self-contribution"):
-            # Using the Coulomb potential as an example, this is the potential generated
-            # at the origin by the fictituous Gaussian charge density in order to split
-            # the potential into a SR and LR part. This contribution always should be
-            # subtracted since it depends on the smearing parameter, which is purely a
-            # convergence parameter.
-            interpolated_potential -= charges * self.potential.self_contribution()
-
-        with profiler.record_function("step 5: charge neutralization"):
-            # If the cell has a net charge (i.e. if sum(charges) != 0), the method
-            # implicitly assumes that a homogeneous background charge of the opposite
-            # sign is present to make the cell neutral. In this case, the potential has
-            # to be adjusted to compensate for this. An extra factor of 2 is added to
-            # compensate for the division by 2 later on
-            charge_tot = torch.sum(charges, dim=0)
-            prefac = self.potential.background_correction()
-            interpolated_potential -= 2 * prefac * charge_tot * ivolume
+        # Using the Coulomb potential as an example, this is the potential generated
+        # at the origin by the fictituous Gaussian charge density in order to split
+        # the potential into a SR and LR part. This contribution always should be
+        # subtracted since it depends on the smearing parameter, which is purely a
+        # convergence parameter.
+        interpolated_potential -= charges * self.potential.self_contribution()
+
+        # If the cell has a net charge (i.e. if sum(charges) != 0), the method
+        # implicitly assumes that a homogeneous background charge of the opposite
+        # sign is present to make the cell neutral. In this case, the potential has
+        # to be adjusted to compensate for this. An extra factor of 2 is added to
+        # compensate for the division by 2 later on
+        charge_tot = torch.sum(charges, dim=0)
+        prefac = self.potential.background_correction()
+        interpolated_potential -= 2 * prefac * charge_tot * ivolume
+
+        interpolated_potential += self.potential.pbc_correction(
+            periodic, positions, cell, charges
+        )
 
         # Compensate for double counting of pairs (i,j) and (j,i)
         return interpolated_potential / 2

src/torchpme/potentials/coulomb.py

@@ -103,5 +103,56 @@ def background_correction(self) -> torch.Tensor:
             )
         return torch.pi * self.smearing**2
 
+    @staticmethod
+    def pbc_correction(
+        periodic: Optional[torch.Tensor],
+        positions: torch.Tensor,
+        cell: torch.Tensor,
+        charges: torch.Tensor,
+    ) -> torch.Tensor:
+        # "2D periodicity" correction for 1/r potential
+        if periodic is None:
+            periodic = torch.tensor([True, True, True], device=cell.device)
+
+        n_periodic = torch.sum(periodic).item()
+        if n_periodic == 3:
+            periodicity = 3
+            nonperiodic_axis = None
+        elif n_periodic == 2:
+            periodicity = 2
+            nonperiodic_axis = torch.where(~periodic)[0]
+            max_distance = torch.max(positions[:, nonperiodic_axis]) - torch.min(
+                positions[:, nonperiodic_axis]
+            )
+            cell_size = torch.linalg.norm(cell[nonperiodic_axis])
+            if max_distance > cell_size / 3:
+                raise ValueError(
+                    f"Maximum distance along non-periodic axis ({max_distance}) "
+                    f"exceeds one third of cell size ({cell_size})."
+                )
+        else:
+            raise ValueError(
+                "K-space summation is not implemented for 1D or non-periodic systems."
+            )
+
+        if periodicity == 2:
+            charge_tot = torch.sum(charges, dim=0)
+            axis = nonperiodic_axis
+            z_i = positions[:, axis].view(-1, 1)
+            basis_len = torch.linalg.norm(cell[axis])
+            M_axis = torch.sum(charges * z_i, dim=0)
+            M_axis_sq = torch.sum(charges * z_i**2, dim=0)
+            V = torch.abs(torch.linalg.det(cell))
+            E_slab = (4.0 * torch.pi / V) * (
+                z_i * M_axis
+                - 0.5 * (M_axis_sq + charge_tot * z_i**2)
+                - charge_tot / 12.0 * basis_len**2
+            )
+        else:
+            E_slab = torch.zeros_like(charges)
+
+        return E_slab
+
     self_contribution.__doc__ = Potential.self_contribution.__doc__
     background_correction.__doc__ = Potential.background_correction.__doc__
+    pbc_correction.__doc__ = Potential.pbc_correction.__doc__

src/torchpme/potentials/inversepowerlaw.py

@@ -5,6 +5,7 @@
 
 from torchpme.lib import gamma, gammaincc_over_powerlaw
 
+from .coulomb import CoulombPotential
 from .potential import Potential
 
 
@@ -141,5 +142,11 @@ def background_correction(self) -> torch.Tensor:
         prefac /= (3 - self.exponent) * gamma(self.exponent / 2)
         return prefac
 
+    def pbc_correction(self, periodic, positions, cell, charges):
+        if self.exponent == 1:
+            return CoulombPotential.pbc_correction(periodic, positions, cell, charges)
+        return super().pbc_correction(periodic, positions, cell, charges)
+
     self_contribution.__doc__ = Potential.self_contribution.__doc__
     background_correction.__doc__ = Potential.background_correction.__doc__
+    pbc_correction.__doc__ = Potential.pbc_correction.__doc__

src/torchpme/potentials/potential.py

@@ -171,3 +171,17 @@ def background_correction(self) -> torch.Tensor:
         raise NotImplementedError(
             f"background_correction is not implemented for {self.__class__.__name__}"
         )
+
+    @torch.jit.export
+    def pbc_correction(
+        self,
+        periodic: Optional[torch.Tensor],
+        positions: torch.Tensor,
+        cell: torch.Tensor,
+        charges: torch.Tensor,
+    ) -> torch.Tensor:
+        """A correction term that is only relevant for systems with 2D periodicity."""
+        if periodic is None or torch.all(periodic):
+            return torch.zeros_like(charges)
+
+        raise NotImplementedError(f"pbc_correction is not implemented for {self}")