add torchsim integration and tests

Author: cw-tan

README.md

@@ -60,6 +60,31 @@ calculator = NequIPPolCalculator.from_compiled_model(
 )
 ```
 
+### Compile for TorchSim
+
+Use the `allegro-pol` batched target to generate a compiled model for
+`allegro_pol.integrations.torchsim.NequIPPolTorchSimCalc`:
+
+```bash
+nequip-compile \
+  path/to/model.ckpt \
+  path/to/compiled_model.nequip.pt2 \
+  --device cuda \
+  --mode aotinductor \
+  --target batch_pol_bc
+```
+
+Then load it with the `allegro-pol` TorchSim calculator:
+
+```python
+from allegro_pol.integrations.torchsim import NequIPPolTorchSimCalc
+
+calculator = NequIPPolTorchSimCalc.from_compiled_model(
+    compile_path="path/to/compiled_model.nequip.pt2",
+    device="cuda",  # or "cpu"
+)
+```
+
 ### Compile for LAMMPS pair styles
 
 Use an `allegro-pol` LAMMPS target with `nequip-compile` for pair-style integrations:

allegro_pol/_compile.py

@@ -2,8 +2,10 @@
 from nequip.data import AtomicDataDict
 from nequip.scripts._compile_utils import (
     ASE_OUTPUTS,
+    BATCH_INPUTS,
     LMP_OUTPUTS,
     PAIR_NEQUIP_INPUTS,
+    batched_data_settings,
     single_frame_data_settings,
     single_frame_batch_map_settings,
     register_compile_targets,
@@ -16,6 +18,7 @@
 AOTI_PAIR_ALLEGRO_POL_TARGET = "pair_allegro_pol"
 AOTI_PAIR_ALLEGRO_POL_BC_TARGET = "pair_allegro_pol_bc"
 AOTI_ASE_POL_BC_TARGET = "ase_pol_bc"
+AOTI_BATCH_POL_BC_TARGET = "batch_pol_bc"
 
 
 PAIR_ALLEGRO_POL_OUTPUTS = [*LMP_OUTPUTS, AtomicDataDict.POLARIZATION_KEY]
@@ -53,10 +56,18 @@
     "data_settings": single_frame_data_settings,
 }
 
+BATCH_POL_BC_TARGET = {
+    "input": BATCH_INPUTS,
+    "output": ASE_POL_BC_OUTPUTS,
+    "batch_map_settings": lambda batch_map: batch_map,  # no static shapes
+    "data_settings": batched_data_settings,
+}
+
 register_compile_targets(
     {
         AOTI_PAIR_ALLEGRO_POL_TARGET: PAIR_ALLEGRO_POL_TARGET,
         AOTI_PAIR_ALLEGRO_POL_BC_TARGET: PAIR_ALLEGRO_POL_BC_TARGET,
         AOTI_ASE_POL_BC_TARGET: ASE_POL_BC_TARGET,
+        AOTI_BATCH_POL_BC_TARGET: BATCH_POL_BC_TARGET,
     }
 )

allegro_pol/integrations/torchsim.py

@@ -0,0 +1,33 @@
+# This file is a part of the `allegro-pol` package. Please see LICENSE and README at the root for information on using it.
+from typing import Dict
+
+import torch
+
+from nequip.data import AtomicDataDict
+from nequip.integrations.torchsim import NequIPTorchSimCalc
+from nequip.scripts._compile_utils import COMPILE_TARGET_DICT
+
+from .._compile import AOTI_BATCH_POL_BC_TARGET
+from .._keys import POLARIZABILITY_KEY
+
+
+class NequIPPolTorchSimCalc(NequIPTorchSimCalc):
+    @classmethod
+    def _get_aoti_compile_target(cls) -> Dict:
+        return COMPILE_TARGET_DICT[AOTI_BATCH_POL_BC_TARGET]
+
+    def save_extra_outputs(
+        self, out: dict[str, torch.Tensor], results: dict[str, torch.Tensor]
+    ) -> None:
+        if AtomicDataDict.POLARIZATION_KEY in out:
+            results[AtomicDataDict.POLARIZATION_KEY] = out[
+                AtomicDataDict.POLARIZATION_KEY
+            ].detach()
+
+        if AtomicDataDict.BORN_CHARGE_KEY in out:
+            results[AtomicDataDict.BORN_CHARGE_KEY] = out[
+                AtomicDataDict.BORN_CHARGE_KEY
+            ].detach()
+
+        if POLARIZABILITY_KEY in out:
+            results[POLARIZABILITY_KEY] = out[POLARIZABILITY_KEY].detach()

tests/model/test_allegro_pol.py

@@ -1,8 +1,13 @@
 # This file is a part of the `allegro-pol` package. Please see LICENSE and README at the root for information on using it.
 import pytest
-from nequip.utils.unittests.model_tests_compilation import CompilationTestsMixin
+from nequip.data import AtomicDataDict
+from nequip.utils.unittests.model_tests_ase_integration import ASEIntegrationMixin
+from nequip.utils.unittests.model_tests_train_time_compile import TrainTimeCompileMixin
+from nequip.utils.unittests.model_tests_torchsim import TorchSimIntegrationMixin
 from nequip.utils.versions import _TORCH_GE_2_6
 
+from allegro_pol._keys import POLARIZABILITY_KEY
+
 _CUEQ_INSTALLED = False
 
 if _TORCH_GE_2_6:
@@ -51,15 +56,17 @@
 )
 
 
-class TestAllegroPol(CompilationTestsMixin):
+class TestAllegroPol(
+    TrainTimeCompileMixin, ASEIntegrationMixin, TorchSimIntegrationMixin
+):
     """Test suite for Allegro Polarization models"""
 
     @pytest.fixture
     def strict_locality(self):
         return True
 
     @pytest.fixture(scope="class")
-    def nequip_compile_tol(self, model_dtype):
+    def ase_integration_tol(self, model_dtype):
         return {"float32": 5e-5, "float64": 1e-10}[model_dtype]
 
     @pytest.fixture(scope="class")
@@ -69,11 +76,49 @@ def ase_calculator_cls(self):
         return NequIPPolCalculator
 
     @pytest.fixture(scope="class")
-    def ase_aoti_compile_target(self):
+    def ase_aoti_target(self):
         from allegro_pol._compile import AOTI_ASE_POL_BC_TARGET
 
         return AOTI_ASE_POL_BC_TARGET
 
+    @pytest.fixture(scope="class")
+    def ase_properties_to_compare(self):
+        return [
+            "energy",
+            "forces",
+            AtomicDataDict.POLARIZATION_KEY,
+            AtomicDataDict.BORN_CHARGE_KEY,
+        ]
+
+    @pytest.fixture(scope="class")
+    def torchsim_calculator_cls(self):
+        from allegro_pol.integrations.torchsim import NequIPPolTorchSimCalc
+
+        return NequIPPolTorchSimCalc
+
+    @pytest.fixture(scope="class")
+    def torchsim_reference_ase_calculator_cls(self):
+        from allegro_pol.integrations.ase import NequIPPolCalculator
+
+        return NequIPPolCalculator
+
+    @pytest.fixture(scope="class")
+    def torchsim_aoti_target(self):
+        from allegro_pol._compile import AOTI_BATCH_POL_BC_TARGET
+
+        return AOTI_BATCH_POL_BC_TARGET
+
+    @pytest.fixture(scope="class")
+    def torchsim_properties_to_compare(self):
+        return [
+            "energy",
+            "forces",
+            "stress",
+            AtomicDataDict.POLARIZATION_KEY,
+            AtomicDataDict.BORN_CHARGE_KEY,
+            POLARIZABILITY_KEY,
+        ]
+
     @pytest.fixture(
         scope="class",
         params=[None]
@@ -83,7 +128,7 @@ def ase_aoti_compile_target(self):
             else []
         ),
     )
-    def nequip_compile_acceleration_modifiers(self, request):
+    def ase_compile_modifiers(self, request):
         if request.param is None:
             return None
 
@@ -103,6 +148,10 @@ def modifier_handler(mode, device, model_dtype):
 
         return modifier_handler
 
+    @pytest.fixture(scope="class")
+    def torchsim_compile_modifiers(self, ase_compile_modifiers):
+        return ase_compile_modifiers
+
     @pytest.fixture(
         scope="class",
         params=[None]
@@ -112,11 +161,11 @@ def modifier_handler(mode, device, model_dtype):
             else []
         ),
     )
-    def train_time_compile_acceleration_modifiers(self, request):
+    def train_time_compile_modifiers(self, request):
         if request.param is None:
             return None
 
-        def modifier_handler(device):
+        def modifier_handler(device, model_dtype):
             if request.param == "enable_CuEquivarianceContracter":
                 if device == "cpu":
                     pytest.skip("CuEquivarianceContracter tests skipped for CPU")