fea: bring in serial sum model

Author: CompRhys

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")

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