Expose Slab Correction in Electrostatics Wrappers (#103)

atulcthakur · web-flow · commit 63247ba4560e · 2026-06-09T21:23:20.000Z
* Add slab correction support to electrostatics wrappers

* Add slab correction wrapper tests

* Address comments

Signed-off-by: atulcthakur &lt;atthakur@nvidia.com&gt;

---------

Signed-off-by: atulcthakur &lt;atthakur@nvidia.com&gt;
diff --git a/nvalchemi/models/ewald.py b/nvalchemi/models/ewald.py
@@ -94,7 +94,13 @@ class EwaldModelWrapper(nn.Module, BaseModelMixin):
     coulomb_constant : float, optional
         Coulomb prefactor :math:`k_e` in eV·Å/e².
         Defaults to ``14.3996`` (standard value for Å/e/eV unit system).
-
+    slab_correction : bool, optional
+        Whether to enable the two-dimensional slab correction. Defaults to
+        ``False``. When enabled, the input batch must provide ``data.pbc`` as
+        a boolean tensor with shape ``(B, 3)``. Rows with exactly one
+        ``False`` entry mark slab systems, for example ``[True, True, False]``
+        for a non-periodic z axis. Fully periodic rows are no-ops, so mixed
+        slab and three-dimensional periodic batches are supported.
 
     Attributes
     ----------
@@ -116,12 +122,14 @@ def __init__(
         accuracy: float = 1e-6,
         coulomb_constant: float = 14.3996,
         hybrid_forces: bool = True,
+        slab_correction: bool = False,
     ) -> None:
         super().__init__()
         self.cutoff = cutoff
         self.accuracy = accuracy
         self.coulomb_constant = coulomb_constant
         self.hybrid_forces = hybrid_forces
+        self.slab_correction = slab_correction
         self.model_config = ModelConfig(
             outputs=frozenset({"energy", "forces", "stress"}),
             active_outputs={"energy", "forces"},
@@ -184,7 +192,10 @@ def direct_derivative_keys(self) -> set[str]:
 
     def input_data(self) -> set[str]:
         """Return required input keys (override to drop ``atomic_numbers``)."""
-        return {"positions", "charges", "neighbor_matrix", "num_neighbors"}
+        keys = {"positions", "charges", "neighbor_matrix", "num_neighbors"}
+        if self.slab_correction:
+            keys.add("pbc")
+        return keys
 
     # ------------------------------------------------------------------
     # Cache management
@@ -251,6 +262,12 @@ def adapt_input(self, data: AtomicData | Batch, **kwargs: Any) -> dict[str, Any]
         for key in self.input_data():
             value = getattr(data, key, None)
             if value is None:
+                if key == "pbc" and self.slab_correction:
+                    raise ValueError(
+                        "EwaldModelWrapper with slab_correction=True requires "
+                        "periodic boundary condition flags "
+                        "(data.pbc must be present)."
+                    )
                 raise KeyError(f"'{key}' required but not found in input data.")
             input_dict[key] = value
 
@@ -268,6 +285,15 @@ def adapt_input(self, data: AtomicData | Batch, **kwargs: Any) -> dict[str, Any]
                 "(data.cell must be present)."
             )
 
+        if self.slab_correction:
+            pbc = getattr(data, "pbc", None)
+            if pbc is None:
+                raise ValueError(
+                    "EwaldModelWrapper with slab_correction=True requires "
+                    "periodic boundary condition flags (data.pbc must be present)."
+                )
+            input_dict["pbc"] = pbc  # (B, 3)
+
         # neighbor_matrix and num_neighbors are already collected by the
         # input_data() loop above.  In a pipeline, the pipeline adapts them
         # to this model's cutoff/format before calling forward().
@@ -327,8 +353,7 @@ def forward(self, data: AtomicData | Batch, **kwargs: Any) -> ModelOutputs:
             ``-W/V``).
         """
         from nvalchemiops.torch.interactions.electrostatics.ewald import (  # lazy
-            ewald_real_space,
-            ewald_reciprocal_space,
+            ewald_summation,
         )
 
         inp = self.adapt_input(data, **kwargs)
@@ -343,6 +368,7 @@ def forward(self, data: AtomicData | Batch, **kwargs: Any) -> ModelOutputs:
         B: int = inp["num_graphs"]
         neighbor_matrix = inp["neighbor_matrix"].contiguous()
         neighbor_matrix_shifts = inp.get("neighbor_matrix_shifts")
+        pbc = inp.get("pbc")
 
         compute_forces = "forces" in self.model_config.active_outputs
         compute_stresses = "stress" in self.model_config.active_outputs
@@ -385,36 +411,25 @@ def forward(self, data: AtomicData | Batch, **kwargs: Any) -> ModelOutputs:
                 self._null_shifts_shape = (N, K)
             neighbor_matrix_shifts = self._null_shifts
 
-        # --- Real-space contribution ---
-        real_result = ewald_real_space(
+        result = ewald_summation(
             positions=positions,
             charges=charges,
             cell=cell,
             alpha=alpha,
+            k_vectors=k_vectors,
             neighbor_matrix=neighbor_matrix,
             neighbor_matrix_shifts=neighbor_matrix_shifts.contiguous(),
             mask_value=fill_value,
             batch_idx=batch_idx,
             compute_forces=compute_forces,
             compute_virial=compute_stresses,
             hybrid_forces=self.hybrid_forces,
+            pbc=pbc,
+            slab_correction=self.slab_correction,
         )
 
-        # --- Reciprocal-space contribution ---
-        recip_result = ewald_reciprocal_space(
-            positions=positions,
-            charges=charges,
-            cell=cell,
-            k_vectors=k_vectors,
-            alpha=alpha,
-            batch_idx=batch_idx,
-            compute_forces=compute_forces,
-            compute_virial=compute_stresses,
-            hybrid_forces=self.hybrid_forces,
-        )
-
-        # Unpack results (energies always first; forces and virial follow
-        # in the order they were requested).
+        # Unpack results (energies always first; forces and virial follow in
+        # the order they were requested).
         def _unpack(result, compute_f: bool, compute_v: bool):
             """Extract (energies, forces_or_None, virial_or_None) from result."""
             if isinstance(result, torch.Tensor):
@@ -431,23 +446,10 @@ def _unpack(result, compute_f: bool, compute_v: bool):
                 v = result_list[idx]
             return e, f, v
 
-        e_real, f_real, v_real = _unpack(real_result, compute_forces, compute_stresses)
-        e_recip, f_recip, v_recip = _unpack(
-            recip_result, compute_forces, compute_stresses
+        per_atom_energies, forces, virial = _unpack(
+            result, compute_forces, compute_stresses
         )
-
-        # Sum real + reciprocal contributions.
-        per_atom_energies = (e_real + e_recip).to(
-            positions.dtype
-        )  # (N,) float64->dtype
-
-        forces: torch.Tensor | None = None
-        if compute_forces and f_real is not None and f_recip is not None:
-            forces = f_real + f_recip  # (N, 3)
-
-        virial: torch.Tensor | None = None
-        if compute_stresses and v_real is not None and v_recip is not None:
-            virial = v_real + v_recip  # (B, 3, 3)
+        per_atom_energies = per_atom_energies.to(positions.dtype)
 
         # Scale by Coulomb constant.
         per_atom_energies = per_atom_energies * self.coulomb_constant
diff --git a/nvalchemi/models/pme.py b/nvalchemi/models/pme.py
@@ -109,6 +109,13 @@ class PMEModelWrapper(nn.Module, BaseModelMixin):
     coulomb_constant : float, optional
         Coulomb prefactor :math:`k_e` in eV·Å/e².
         Defaults to ``14.3996`` (standard value for Å/e/eV unit system).
+    slab_correction : bool, optional
+        Whether to enable the two-dimensional slab correction. Defaults to
+        ``False``. When enabled, the input batch must provide ``data.pbc`` as
+        a boolean tensor with shape ``(B, 3)``. Rows with exactly one
+        ``False`` entry mark slab systems, for example ``[True, True, False]``
+        for a non-periodic z axis. Fully periodic rows are no-ops, so mixed
+        slab and three-dimensional periodic batches are supported.
 
     Attributes
     ----------
@@ -134,6 +141,7 @@ def __init__(
         accuracy: float = 1e-6,
         coulomb_constant: float = 14.3996,
         hybrid_forces: bool = True,
+        slab_correction: bool = False,
     ) -> None:
         super().__init__()
         self.cutoff = cutoff
@@ -144,6 +152,7 @@ def __init__(
         self.accuracy = accuracy
         self.coulomb_constant = coulomb_constant
         self.hybrid_forces = hybrid_forces
+        self.slab_correction = slab_correction
         self.model_config = ModelConfig(
             outputs=frozenset({"energy", "forces", "stress"}),
             active_outputs={"energy", "forces"},
@@ -207,7 +216,10 @@ def direct_derivative_keys(self) -> set[str]:
 
     def input_data(self) -> set[str]:
         """Return required input keys (override to drop ``atomic_numbers``)."""
-        return {"positions", "charges", "neighbor_matrix", "num_neighbors"}
+        keys = {"positions", "charges", "neighbor_matrix", "num_neighbors"}
+        if self.slab_correction:
+            keys.add("pbc")
+        return keys
 
     # ------------------------------------------------------------------
     # Cache management
@@ -287,6 +299,11 @@ def adapt_input(self, data: AtomicData | Batch, **kwargs: Any) -> dict[str, Any]
         for key in self.input_data():
             value = getattr(data, key, None)
             if value is None:
+                if key == "pbc" and self.slab_correction:
+                    raise ValueError(
+                        "PMEModelWrapper with slab_correction=True requires periodic "
+                        "boundary condition flags (data.pbc must be present)."
+                    )
                 raise KeyError(f"'{key}' required but not found in input data.")
             input_dict[key] = value
 
@@ -304,6 +321,15 @@ def adapt_input(self, data: AtomicData | Batch, **kwargs: Any) -> dict[str, Any]
                 "(data.cell must be present)."
             )
 
+        if self.slab_correction:
+            pbc = getattr(data, "pbc", None)
+            if pbc is None:
+                raise ValueError(
+                    "PMEModelWrapper with slab_correction=True requires periodic "
+                    "boundary condition flags (data.pbc must be present)."
+                )
+            input_dict["pbc"] = pbc  # (B, 3)
+
         # neighbor_matrix and num_neighbors are already collected by the
         # input_data() loop above.  In a pipeline, the pipeline adapts them
         # to this model's cutoff/format before calling forward().
@@ -377,6 +403,7 @@ def forward(self, data: AtomicData | Batch, **kwargs: Any) -> ModelOutputs:
         B: int = inp["num_graphs"]
         neighbor_matrix = inp["neighbor_matrix"].contiguous()
         neighbor_matrix_shifts = inp.get("neighbor_matrix_shifts")
+        pbc = inp.get("pbc")
 
         compute_forces = "forces" in self.model_config.active_outputs
         compute_stresses = "stress" in self.model_config.active_outputs
@@ -450,6 +477,8 @@ def forward(self, data: AtomicData | Batch, **kwargs: Any) -> ModelOutputs:
             compute_virial=compute_stresses,
             accuracy=self.accuracy,
             hybrid_forces=self.hybrid_forces,
+            pbc=pbc,
+            slab_correction=self.slab_correction,
         )
 
         # Unpack tuple: (energies, [forces], [virial]).
diff --git a/test/models/test_ewald.py b/test/models/test_ewald.py
@@ -208,6 +208,11 @@ def test_input_data_override(self):
         assert "neighbor_matrix" in keys
         assert "num_neighbors" in keys
 
+    def test_input_data_includes_pbc_for_slab_correction(self):
+        """Slab-enabled Ewald declares pbc as a required input."""
+        w = _make_ewald(slab_correction=True)
+        assert "pbc" in w.input_data()
+
     def test_output_data_with_forces(self):
         w = _make_ewald()
         out = w.output_data()
@@ -514,31 +519,14 @@ def test_forward_stress_when_requested(self):
 
     def test_forward_stress_is_negative_virial_over_volume(self):
         """ASE-style stress == -virial / volume (eV/A^3)."""
-        w = _make_ewald()
+        w = _make_ewald(slab_correction=True)
         w.model_config.active_outputs = {"energy", "forces", "stress"}
         batch = _make_charged_batch(box_size=10.0)
         self._build_nl(batch, w)
 
-        real_virial_value = 2.0
-        recip_virial_value = 3.0
-
-        def fake_real_space(**kw):
-            positions = kw["positions"]
-            cell = kw["cell"]
-            return (
-                torch.zeros(
-                    positions.shape[0], dtype=positions.dtype, device=positions.device
-                ),
-                torch.zeros_like(positions),
-                torch.full(
-                    (cell.shape[0], 3, 3),
-                    real_virial_value,
-                    dtype=positions.dtype,
-                    device=positions.device,
-                ),
-            )
+        virial_value = 5.0
 
-        def fake_reciprocal_space(**kw):
+        def fake_ewald_summation(**kw):
             positions = kw["positions"]
             cell = kw["cell"]
             return (
@@ -548,33 +536,30 @@ def fake_reciprocal_space(**kw):
                 torch.zeros_like(positions),
                 torch.full(
                     (cell.shape[0], 3, 3),
-                    recip_virial_value,
+                    virial_value,
                     dtype=positions.dtype,
                     device=positions.device,
                 ),
             )
 
-        with (
-            patch(
-                "nvalchemiops.torch.interactions.electrostatics.ewald.ewald_real_space",
-                side_effect=fake_real_space,
-            ),
-            patch(
-                "nvalchemiops.torch.interactions.electrostatics.ewald.ewald_reciprocal_space",
-                side_effect=fake_reciprocal_space,
-            ),
-        ):
+        with patch(
+            "nvalchemiops.torch.interactions.electrostatics.ewald.ewald_summation",
+            side_effect=fake_ewald_summation,
+        ) as mock_ewald_summation:
             out = w.forward(batch)
 
+        call_kwargs = mock_ewald_summation.call_args.kwargs
+        torch.testing.assert_close(call_kwargs["pbc"], batch.pbc)
+        assert call_kwargs["slab_correction"] is True
+        assert call_kwargs["compute_virial"] is True
+
         volume = torch.det(batch.cell).abs().view(-1, 1, 1)
-        expected = (
-            -(real_virial_value + recip_virial_value) * w.coulomb_constant / volume
-        )
+        expected = -virial_value * w.coulomb_constant / volume
         torch.testing.assert_close(out["stress"], expected.expand_as(out["stress"]))
 
     def test_forward_raises_when_virial_none(self):
         """RuntimeError when stress is requested but kernels return no virial."""
-        w = _make_ewald()
+        w = _make_ewald(slab_correction=True)
         w.model_config.active_outputs = {"energy", "forces", "stress"}
         batch = _make_charged_batch()
         self._build_nl(batch, w)
@@ -586,19 +571,18 @@ def _fake_kernel(**kw):
             forces = torch.zeros(N, 3, dtype=torch.float64)
             return energies, forces
 
-        with (
-            patch(
-                "nvalchemiops.torch.interactions.electrostatics.ewald.ewald_real_space",
-                side_effect=_fake_kernel,
-            ),
-            patch(
-                "nvalchemiops.torch.interactions.electrostatics.ewald.ewald_reciprocal_space",
-                side_effect=_fake_kernel,
-            ),
-        ):
+        with patch(
+            "nvalchemiops.torch.interactions.electrostatics.ewald.ewald_summation",
+            side_effect=_fake_kernel,
+        ) as mock_ewald_summation:
             with pytest.raises(RuntimeError, match="kernel did not return a virial"):
                 w.forward(batch)
 
+        call_kwargs = mock_ewald_summation.call_args.kwargs
+        torch.testing.assert_close(call_kwargs["pbc"], batch.pbc)
+        assert call_kwargs["slab_correction"] is True
+        assert call_kwargs["compute_virial"] is True
+
     def test_cache_populated_after_forward(self):
         w = _make_ewald()
         batch = _make_charged_batch()
diff --git a/test/models/test_pme.py b/test/models/test_pme.py
@@ -242,6 +242,11 @@ def test_input_data_override(self):
         assert "neighbor_matrix" in keys
         assert "num_neighbors" in keys
 
+    def test_input_data_includes_pbc_for_slab_correction(self):
+        """Slab-enabled PME declares pbc as a required input."""
+        w = _make_pme(slab_correction=True)
+        assert "pbc" in w.input_data()
+
     def test_output_data_with_forces(self):
         w = _make_pme()
         out = w.output_data()
diff --git a/test/models/test_slab.py b/test/models/test_slab.py
