Add pair_mask parameter to potential calculations in Calculator and CoulombPotential

Author: E-Rum

src/torchpme/calculators/calculator.py

@@ -51,17 +51,21 @@ def _compute_rspace(
         charges: torch.Tensor,
         neighbor_indices: torch.Tensor,
         neighbor_distances: torch.Tensor,
+        node_mask: torch.Tensor | None = None,
+        pair_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
         # Compute the pair potential terms V(r_ij) for each pair of atoms (i,j)
         # contained in the neighbor list
         with profiler.record_function("compute bare potential"):
             if self.potential.smearing is None:
                 if self.potential.exclusion_radius is None:
-                    potentials_bare = self.potential.from_dist(neighbor_distances)
-                else:
-                    potentials_bare = self.potential.from_dist(neighbor_distances) * (
-                        1 - self.potential.f_cutoff(neighbor_distances)
+                    potentials_bare = self.potential.from_dist(
+                        neighbor_distances, pair_mask=pair_mask
                     )
+                else:
+                    potentials_bare = self.potential.from_dist(
+                        neighbor_distances, pair_mask=pair_mask
+                    ) * (1 - self.potential.f_cutoff(neighbor_distances))
             else:
                 potentials_bare = self.potential.sr_from_dist(neighbor_distances)
 
@@ -109,6 +113,8 @@ def forward(
         neighbor_indices: torch.Tensor,
         neighbor_distances: torch.Tensor,
         periodic: Optional[torch.Tensor] = None,
+        node_mask: torch.Tensor | None = None,
+        pair_mask: torch.Tensor | None = None,
     ):
         r"""
         Compute the potential "energy".
@@ -161,6 +167,8 @@ def forward(
             charges=charges,
             neighbor_indices=neighbor_indices,
             neighbor_distances=neighbor_distances,
+            node_mask=node_mask,
+            pair_mask=pair_mask,
         )
 
         if self.potential.smearing is None:

src/torchpme/potentials/coulomb.py

@@ -39,14 +39,21 @@ def __init__(
     ):
         super().__init__(smearing, exclusion_radius, exclusion_degree)
 
-    def from_dist(self, dist: torch.Tensor) -> torch.Tensor:
+    def from_dist(
+        self, dist: torch.Tensor, pair_mask: torch.Tensor | None = None
+    ) -> torch.Tensor:
         """
         Full :math:`1/r` potential as a function of :math:`r`.
 
         :param dist: torch.tensor containing the distances at which the potential is to
             be evaluated.
         """
-        return 1.0 / dist
+        result = 1.0 / dist.clamp(min=1e-12)
+
+        if pair_mask is not None:
+            result = result * pair_mask  # elementwise multiply, keeps shape fixed
+
+        return result
 
     def lr_from_dist(self, dist: torch.Tensor) -> torch.Tensor:
         """