fea: add polarization model

CompRhys · CompRhys · commit bcc58c25259c · 2026-04-08T16:30:35.000-04:00
diff --git a/tests/models/test_polarization.py b/tests/models/test_polarization.py
@@ -0,0 +1,194 @@
+"""Tests for the polarization electric-field correction model."""
+
+import pytest
+import torch
+
+import torch_sim as ts
+from tests.conftest import DEVICE, DTYPE
+from torch_sim.models.interface import ModelInterface, SerialSumModel
+from torch_sim.models.polarization import UniformPolarizationModel
+
+
+class DummyPolarResponseModel(ModelInterface):
+    def __init__(
+        self,
+        *,
+        polarization_key: str = "polarization",
+        include_born_effective_charges: bool = True,
+        include_polarizability: bool = True,
+        device: torch.device = DEVICE,
+        dtype: torch.dtype = DTYPE,
+    ) -> None:
+        super().__init__()
+        self.polarization_key = polarization_key
+        self.include_born_effective_charges = include_born_effective_charges
+        self.include_polarizability = include_polarizability
+        self._device = device
+        self._dtype = dtype
+        self._compute_forces = True
+        self._compute_stress = True
+
+    def forward(self, state: ts.SimState, **kwargs: object) -> dict[str, torch.Tensor]:
+        del kwargs
+        energy = torch.arange(
+            1, state.n_systems + 1, device=state.device, dtype=state.dtype
+        )
+        forces = (
+            torch.arange(state.n_atoms * 3, device=state.device, dtype=state.dtype)
+            .reshape(state.n_atoms, 3)
+            .div(10.0)
+        )
+        stress = (
+            torch.arange(state.n_systems * 9, device=state.device, dtype=state.dtype)
+            .reshape(state.n_systems, 3, 3)
+            .div(100.0)
+        )
+        polarization = (
+            torch.arange(state.n_systems * 3, device=state.device, dtype=state.dtype)
+            .reshape(state.n_systems, 3)
+            .add(0.5)
+        )
+        output: dict[str, torch.Tensor] = {
+            "energy": energy,
+            "forces": forces,
+            "stress": stress,
+            self.polarization_key: polarization,
+        }
+        if self.include_polarizability:
+            diag = torch.tensor([1.0, 2.0, 3.0], device=state.device, dtype=state.dtype)
+            output["polarizability"] = torch.diag_embed(diag.repeat(state.n_systems, 1))
+        if self.include_born_effective_charges:
+            born_effective_charges = torch.zeros(
+                state.n_atoms, 3, 3, device=state.device, dtype=state.dtype
+            )
+            born_effective_charges[:, 0, 0] = 1.0
+            born_effective_charges[:, 1, 1] = 2.0
+            born_effective_charges[:, 2, 2] = 3.0
+            output["born_effective_charges"] = born_effective_charges
+        return output
+
+
+def test_polarization_model_normalizes_raw_key_without_field(
+    si_double_sim_state: ts.SimState,
+) -> None:
+    base_model = DummyPolarResponseModel()
+    combined_model = SerialSumModel(
+        base_model,
+        UniformPolarizationModel(device=DEVICE, dtype=DTYPE),
+    )
+
+    base_output = base_model(si_double_sim_state)
+    combined_output = combined_model(si_double_sim_state)
+
+    torch.testing.assert_close(combined_output["energy"], base_output["energy"])
+    torch.testing.assert_close(combined_output["forces"], base_output["forces"])
+    torch.testing.assert_close(combined_output["stress"], base_output["stress"])
+    torch.testing.assert_close(
+        combined_output["total_polarization"], base_output["polarization"]
+    )
+    torch.testing.assert_close(
+        combined_output["polarization"], base_output["polarization"]
+    )
+
+
+def test_polarization_model_applies_linear_response_corrections(
+    si_double_sim_state: ts.SimState,
+) -> None:
+    base_model = DummyPolarResponseModel()
+    combined_model = SerialSumModel(
+        base_model,
+        UniformPolarizationModel(device=DEVICE, dtype=DTYPE),
+    )
+    field = torch.tensor(
+        [[0.2, -0.1, 0.05], [-0.3, 0.4, 0.1]],
+        device=DEVICE,
+        dtype=DTYPE,
+    )
+    state = ts.SimState.from_state(si_double_sim_state, external_E_field=field)
+
+    base_output = base_model(state)
+    combined_output = combined_model(state)
+    expected_polarization = base_output["polarization"] + torch.einsum(
+        "sij,sj->si", base_output["polarizability"], field
+    )
+    expected_energy = base_output["energy"] - torch.einsum(
+        "si,si->s", field, base_output["polarization"]
+    )
+    expected_energy = expected_energy - 0.5 * torch.einsum(
+        "si,sij,sj->s", field, base_output["polarizability"], field
+    )
+    expected_forces = base_output["forces"] + torch.einsum(
+        "imn,im->in",
+        base_output["born_effective_charges"],
+        field[state.system_idx],
+    )
+
+    torch.testing.assert_close(combined_output["energy"], expected_energy)
+    torch.testing.assert_close(combined_output["forces"], expected_forces)
+    torch.testing.assert_close(
+        combined_output["total_polarization"], expected_polarization
+    )
+    torch.testing.assert_close(
+        combined_output["polarization"], base_output["polarization"]
+    )
+    torch.testing.assert_close(combined_output["stress"], base_output["stress"])
+
+
+def test_polarization_model_adds_only_delta_for_blessed_name(
+    si_double_sim_state: ts.SimState,
+) -> None:
+    base_model = DummyPolarResponseModel(polarization_key="total_polarization")
+    combined_model = SerialSumModel(
+        base_model,
+        UniformPolarizationModel(device=DEVICE, dtype=DTYPE),
+    )
+    field = torch.tensor([[0.1, 0.0, 0.0], [0.0, -0.2, 0.3]], device=DEVICE, dtype=DTYPE)
+    state = ts.SimState.from_state(si_double_sim_state, external_E_field=field)
+
+    base_output = base_model(state)
+    combined_output = combined_model(state)
+    expected_total_polarization = base_output["total_polarization"] + torch.einsum(
+        "sij,sj->si", base_output["polarizability"], field
+    )
+
+    torch.testing.assert_close(
+        combined_output["total_polarization"], expected_total_polarization
+    )
+    assert "polarization" not in combined_output
+
+
+def test_polarization_model_requires_born_effective_charges_for_force_correction(
+    si_double_sim_state: ts.SimState,
+) -> None:
+    base_model = DummyPolarResponseModel(include_born_effective_charges=False)
+    combined_model = SerialSumModel(
+        base_model,
+        UniformPolarizationModel(device=DEVICE, dtype=DTYPE),
+    )
+    state = ts.SimState.from_state(
+        si_double_sim_state,
+        external_E_field=torch.ones(
+            si_double_sim_state.n_systems, 3, device=DEVICE, dtype=DTYPE
+        ),
+    )
+
+    with pytest.raises(ValueError, match="born_effective_charges"):
+        combined_model(state)
+
+
+def test_polarization_model_rejects_non_uniform_field_shape(
+    si_double_sim_state: ts.SimState,
+) -> None:
+    state = ts.SimState.from_state(
+        si_double_sim_state,
+        external_E_field=torch.zeros(
+            si_double_sim_state.n_systems, 3, device=DEVICE, dtype=DTYPE
+        ),
+    )
+    state._system_extras["external_E_field"] = torch.zeros(  # noqa: SLF001
+        state.n_atoms, 3, device=DEVICE, dtype=DTYPE
+    )
+    model = UniformPolarizationModel(device=DEVICE, dtype=DTYPE)
+
+    with pytest.raises(ValueError, match="shape \\(n_systems, 3\\)"):
+        model(state)
diff --git a/tests/models/test_sum_model.py b/tests/models/test_sum_model.py
@@ -6,11 +6,47 @@
 import torch_sim as ts
 from tests.conftest import DEVICE, DTYPE
 from tests.models.conftest import make_validate_model_outputs_test
-from torch_sim.models.interface import SumModel
+from torch_sim.models.interface import ModelInterface, SerialSumModel, SumModel
 from torch_sim.models.lennard_jones import LennardJonesModel
 from torch_sim.models.morse import MorseModel
 
 
+class ExtraProducerModel(ModelInterface):
+    def __init__(self, device: torch.device = DEVICE, dtype: torch.dtype = DTYPE) -> None:
+        super().__init__()
+        self._device = device
+        self._dtype = dtype
+        self._compute_stress = False
+        self._compute_forces = False
+
+    def forward(self, state: ts.SimState, **kwargs: object) -> dict[str, torch.Tensor]:
+        del kwargs
+        latent = state.positions[:, 0] + 2.0
+        return {
+            "energy": torch.ones(state.n_systems, device=state.device, dtype=state.dtype),
+            "latent": latent,
+        }
+
+
+class ExtraConsumerModel(ModelInterface):
+    seen_latent: torch.Tensor | None
+
+    def __init__(self, device: torch.device = DEVICE, dtype: torch.dtype = DTYPE) -> None:
+        super().__init__()
+        self._device = device
+        self._dtype = dtype
+        self._compute_stress = False
+        self._compute_forces = False
+        self.seen_latent = None
+
+    def forward(self, state: ts.SimState, **kwargs: object) -> dict[str, torch.Tensor]:
+        del kwargs
+        self.seen_latent = state.latent.clone()
+        energy = torch.zeros(state.n_systems, device=state.device, dtype=state.dtype)
+        energy.scatter_add_(0, state.system_idx, state.latent)
+        return {"energy": energy}
+
+
 @pytest.fixture
 def lj_model_a() -> LennardJonesModel:
     return LennardJonesModel(
@@ -43,9 +79,19 @@ def sum_model(lj_model_a: LennardJonesModel, morse_model: MorseModel) -> SumMode
     return SumModel(lj_model_a, morse_model)
 
 
+@pytest.fixture
+def serial_sum_model(
+    lj_model_a: LennardJonesModel, morse_model: MorseModel
+) -> SerialSumModel:
+    return SerialSumModel(lj_model_a, morse_model)
+
+
 test_sum_model_outputs = make_validate_model_outputs_test(
     model_fixture_name="sum_model", device=DEVICE, dtype=DTYPE
 )
+test_serial_sum_model_outputs = make_validate_model_outputs_test(
+    model_fixture_name="serial_sum_model", device=DEVICE, dtype=DTYPE
+)
 
 
 def test_sum_model_requires_two_models(lj_model_a: LennardJonesModel) -> None:
@@ -102,3 +148,39 @@ def test_sum_model_retain_graph(
     assert lj_model_a.retain_graph is True
     assert morse_model.retain_graph is True
     assert sm.retain_graph is True
+
+
+def test_serial_sum_model_matches_parallel_sum_for_independent_models(
+    lj_model_a: LennardJonesModel,
+    morse_model: MorseModel,
+    si_sim_state: ts.SimState,
+) -> None:
+    sum_out = SumModel(lj_model_a, morse_model)(si_sim_state)
+    serial_out = SerialSumModel(lj_model_a, morse_model)(si_sim_state)
+    torch.testing.assert_close(serial_out["energy"], sum_out["energy"])
+    torch.testing.assert_close(serial_out["forces"], sum_out["forces"])
+    torch.testing.assert_close(serial_out["stress"], sum_out["stress"])
+
+
+def test_serial_sum_model_exposes_extras_to_later_models(
+    si_double_sim_state: ts.SimState,
+) -> None:
+    producer = ExtraProducerModel()
+    consumer = ExtraConsumerModel()
+    serial_model = SerialSumModel(producer, consumer)
+    state = si_double_sim_state.clone()
+    expected_latent = state.positions[:, 0] + 2.0
+    expected_energy = torch.ones(state.n_systems, device=state.device, dtype=state.dtype)
+    expected_energy = expected_energy.scatter_add(
+        0,
+        state.system_idx,
+        expected_latent,
+    )
+
+    output = serial_model(state)
+
+    assert consumer.seen_latent is not None
+    torch.testing.assert_close(consumer.seen_latent, expected_latent)
+    torch.testing.assert_close(output["latent"], expected_latent)
+    torch.testing.assert_close(output["energy"], expected_energy)
+    assert not state.has_extras("latent")
diff --git a/torch_sim/models/interface.py b/torch_sim/models/interface.py
@@ -52,6 +52,17 @@ def forward(self, positions, cell, batch, atomic_numbers=None, **kwargs):
 }
 
 
+def _accumulate_model_output(
+    combined: dict[str, torch.Tensor], output: dict[str, torch.Tensor]
+) -> None:
+    """Accumulate one model output into a combined output dict."""
+    for key, tensor in output.items():
+        if key in combined:
+            combined[key] = combined[key] + tensor
+        else:
+            combined[key] = tensor
+
+
 class ModelInterface(torch.nn.Module, ABC):
     """Abstract base class for all simulation models in TorchSim.
 
@@ -300,11 +311,36 @@ def forward(self, state: SimState, **kwargs) -> dict[str, torch.Tensor]:
         combined: dict[str, torch.Tensor] = {}
         for model in self._children():
             output = model(state, **kwargs)
-            for key, tensor in output.items():
-                if key in combined:
-                    combined[key] = combined[key] + tensor
-                else:
-                    combined[key] = tensor
+            _accumulate_model_output(combined, output)
+        return combined
+
+
+class SerialSumModel(SumModel):
+    """Serial additive composition of multiple :class:`ModelInterface` models.
+
+    Unlike :class:`SumModel`, child models do not all see the same input state.
+    Instead, each child runs after the previous child's non-canonical outputs have
+    been stored into a cloned :class:`~torch_sim.state.SimState` via
+    :meth:`torch_sim.state.SimState.store_model_extras`. This lets earlier models
+    expose per-atom or per-system features that later models can consume, while
+    energies, forces, stresses, and any repeated auxiliary outputs are still summed
+    key-by-key.
+
+    Examples:
+        ```py
+        serial_model = SerialSumModel(polarization_model, dispersion_model)
+        output = serial_model(sim_state)
+        ```
+    """
+
+    def forward(self, state: SimState, **kwargs) -> dict[str, torch.Tensor]:
+        """Run child models serially, exposing extras from earlier models."""
+        combined: dict[str, torch.Tensor] = {}
+        serial_state = state.clone()
+        for model in self._children():
+            output = model(serial_state, **kwargs)
+            _accumulate_model_output(combined, output)
+            serial_state.store_model_extras(output)
         return combined
 
 
diff --git a/torch_sim/models/polarization.py b/torch_sim/models/polarization.py
