test(pt): add common test case for model/atomic model

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -9,6 +9,9 @@
 
 import numpy as np
 
+from deepmd.dpmodel.common import (
+    NativeOP,
+)
 from deepmd.dpmodel.output_def import (
     FittingOutputDef,
     OutputVariableDef,
@@ -25,7 +28,7 @@
 BaseAtomicModel_ = make_base_atomic_model(np.ndarray)
 
 
-class BaseAtomicModel(BaseAtomicModel_):
+class BaseAtomicModel(BaseAtomicModel_, NativeOP):
     def __init__(
         self,
         type_map: List[str],
@@ -183,6 +186,24 @@ def forward_common_atomic(
 
         return ret_dict
 
+    def call(
+        self,
+        extended_coord: np.ndarray,
+        extended_atype: np.ndarray,
+        nlist: np.ndarray,
+        mapping: Optional[np.ndarray] = None,
+        fparam: Optional[np.ndarray] = None,
+        aparam: Optional[np.ndarray] = None,
+    ) -> Dict[str, np.ndarray]:
+        return self.forward_common_atomic(
+            extended_coord,
+            extended_atype,
+            nlist,
+            mapping=mapping,
+            fparam=fparam,
+            aparam=aparam,
+        )
+
     def serialize(self) -> dict:
         return {
             "type_map": self.type_map,

deepmd/dpmodel/model/make_model.py

@@ -473,4 +473,8 @@ def atomic_output_def(self) -> FittingOutputDef:
             """Get the output def of the atomic model."""
             return self.atomic_model.atomic_output_def()
 
+        def get_ntypes(self) -> int:
+            """Get the number of types."""
+            return len(self.get_type_map())
+
     return CM

deepmd/dpmodel/utils/nlist.py

@@ -9,10 +9,42 @@
 import numpy as np
 
 from .region import (
+    normalize_coord,
     to_face_distance,
 )
 
 
+def extend_input_and_build_neighbor_list(
+    coord,
+    atype,
+    rcut: float,
+    sel: List[int],
+    mixed_types: bool = False,
+    box: Optional[np.ndarray] = None,
+):
+    nframes, nloc = atype.shape[:2]
+    if box is not None:
+        coord_normalized = normalize_coord(
+            coord.reshape(nframes, nloc, 3),
+            box.reshape(nframes, 3, 3),
+        )
+    else:
+        coord_normalized = coord
+    extended_coord, extended_atype, mapping = extend_coord_with_ghosts(
+        coord_normalized, atype, box, rcut
+    )
+    nlist = build_neighbor_list(
+        extended_coord,
+        extended_atype,
+        nloc,
+        rcut,
+        sel,
+        distinguish_types=(not mixed_types),
+    )
+    extended_coord = extended_coord.reshape(nframes, -1, 3)
+    return extended_coord, extended_atype, mapping, nlist
+
+
 ## translated from torch implemantation by chatgpt
 def build_neighbor_list(
     coord: np.ndarray,

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -256,6 +256,26 @@ def forward_common_atomic(
 
         return ret_dict
 
+    def forward(
+        self,
+        extended_coord: torch.Tensor,
+        extended_atype: torch.Tensor,
+        nlist: torch.Tensor,
+        mapping: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+        comm_dict: Optional[Dict[str, torch.Tensor]] = None,
+    ) -> Dict[str, torch.Tensor]:
+        return self.forward_common_atomic(
+            extended_coord,
+            extended_atype,
+            nlist,
+            mapping=mapping,
+            fparam=fparam,
+            aparam=aparam,
+            comm_dict=comm_dict,
+        )
+
     def serialize(self) -> dict:
         return {
             "type_map": self.type_map,

source/tests/universal/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Universal tests for the project."""

source/tests/universal/atomic_model/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later

source/tests/universal/atomic_model/test_ener_model.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+
+from deepmd.dpmodel.atomic_model.dp_atomic_model import DPAtomicModel as DPAtomicModelDP
+from deepmd.dpmodel.descriptor.se_e2_a import DescrptSeA as DescrptSeADP
+from deepmd.dpmodel.fitting.ener_fitting import EnergyFittingNet as EnergyFittingNetDP
+
+from ..utils import (
+    INSTALLED_PT,
+)
+
+if INSTALLED_PT:
+    from deepmd.pt.model.atomic_model.energy_atomic_model import (
+        DPEnergyAtomicModel as DPEnergyAtomicModelPT,
+    )
+
+from .utils import (
+    AtomicModelTestCase,
+)
+
+
+class TestEnerAtomicModel(unittest.TestCase, AtomicModelTestCase):
+    def setUp(self) -> None:
+        self.expected_rcut = 5.0
+        self.expected_type_map = ["foo", "bar"]
+        self.expected_dim_fparam = 0
+        self.expected_dim_aparam = 0
+        self.expected_sel_type = [0, 1]
+        self.expected_aparam_nall = False
+        self.expected_model_output_type = ["energy", "mask"]
+        self.expected_sel = [8, 12]
+        ds = DescrptSeADP(
+            rcut=self.expected_rcut,
+            rcut_smth=self.expected_rcut / 2,
+            sel=self.expected_sel,
+        )
+        ft = EnergyFittingNetDP(
+            ntypes=len(self.expected_type_map),
+            dim_descrpt=ds.get_dim_out(),
+            mixed_types=ds.mixed_types(),
+        )
+        self.dp_module = DPAtomicModelDP(
+            ds,
+            ft,
+            type_map=self.expected_type_map,
+        )
+
+        if INSTALLED_PT:
+            self.pt_module = DPEnergyAtomicModelPT.deserialize(
+                self.dp_module.serialize()
+            )

source/tests/universal/atomic_model/utils.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    List,
+)
+
+import numpy as np
+
+from deepmd.dpmodel.utils.nlist import (
+    extend_input_and_build_neighbor_list,
+)
+
+from ..utils import (
+    CommonTestCase,
+    forward_wrapper,
+)
+
+
+class AtomicModelTestCase(CommonTestCase):
+    """Common test case for atomic model."""
+
+    expected_type_map: List[str]
+    """Expected type map."""
+    expected_rcut: float
+    """Expected cut-off radius."""
+    expected_dim_fparam: int
+    """Expected number (dimension) of frame parameters."""
+    expected_dim_aparam: int
+    """Expected number (dimension) of atomic parameters."""
+    expected_sel_type: List[int]
+    """Expected selected atom types."""
+    expected_aparam_nall: bool
+    """Expected shape of atomic parameters."""
+    expected_model_output_type: List[str]
+    """Expected output type for the model."""
+    expected_sel: List[int]
+    """Expected number of neighbors."""
+
+    def test_get_type_map(self):
+        """Test get_type_map."""
+        for module in self.modules_to_test:
+            self.assertEqual(module.get_type_map(), self.expected_type_map)
+
+    def test_get_rcut(self):
+        """Test get_rcut."""
+        for module in self.modules_to_test:
+            self.assertAlmostEqual(module.get_rcut(), self.expected_rcut)
+
+    def test_get_dim_fparam(self):
+        """Test get_dim_fparam."""
+        for module in self.modules_to_test:
+            self.assertEqual(module.get_dim_fparam(), self.expected_dim_fparam)
+
+    def test_get_dim_aparam(self):
+        """Test get_dim_aparam."""
+        for module in self.modules_to_test:
+            self.assertEqual(module.get_dim_aparam(), self.expected_dim_aparam)
+
+    def test_get_sel_type(self):
+        """Test get_sel_type."""
+        for module in self.modules_to_test:
+            self.assertEqual(module.get_sel_type(), self.expected_sel_type)
+
+    def test_is_aparam_nall(self):
+        """Test is_aparam_nall."""
+        for module in self.modules_to_test:
+            self.assertEqual(module.is_aparam_nall(), self.expected_aparam_nall)
+
+    def test_get_nnei(self):
+        """Test get_nnei."""
+        expected_nnei = sum(self.expected_sel)
+        for module in self.modules_to_test:
+            self.assertEqual(module.get_nnei(), expected_nnei)
+
+    def test_get_ntypes(self):
+        """Test get_ntypes."""
+        for module in self.modules_to_test:
+            self.assertEqual(module.get_ntypes(), len(self.expected_type_map))
+
+    def test_forward(self):
+        """Test forward."""
+        nf = 1
+        nloc = 3
+        coord = np.array(
+            [
+                [0, 0, 0],
+                [0, 1, 0],
+                [0, 0, 1],
+            ],
+            dtype=np.float64,
+        ).reshape([nf, -1])
+        atype = np.array([0, 0, 1], dtype=int).reshape([nf, -1])
+        cell = 6.0 * np.eye(3).reshape([nf, 9])
+        coord_ext, atype_ext, mapping, nlist = extend_input_and_build_neighbor_list(
+            coord,
+            atype,
+            self.expected_rcut,
+            self.expected_sel,
+            mixed_types=True,
+            box=cell,
+        )
+        ret_lower = []
+        for module in self.modules_to_test:
+            module = forward_wrapper(module)
+
+            ret_lower.append(module(coord_ext, atype_ext, nlist))
+        for kk in ret_lower[0].keys():
+            subret = []
+            for rr in ret_lower:
+                if rr is not None:
+                    subret.append(rr[kk])
+            if len(subret):
+                for ii, rr in enumerate(subret[1:]):
+                    np.testing.assert_allclose(
+                        subret[0], rr, err_msg=f"compare {kk} between 0 and {ii}"
+                    )

source/tests/universal/model/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later

source/tests/universal/model/test_ener_model.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+
+from deepmd.dpmodel.descriptor.se_e2_a import DescrptSeA as DescrptSeADP
+from deepmd.dpmodel.fitting.ener_fitting import EnergyFittingNet as EnergyFittingNetDP
+from deepmd.dpmodel.model.ener_model import EnergyModel as EnergyModelDP
+
+from ..utils import (
+    INSTALLED_PT,
+)
+
+if INSTALLED_PT:
+    from deepmd.pt.model.model.ener_model import (
+        EnergyModel as EnergyModelPT,
+    )
+
+from .utils import (
+    ModelTestCase,
+)
+
+
+class TestEnerModel(unittest.TestCase, ModelTestCase):
+    def setUp(self) -> None:
+        self.expected_rcut = 5.0
+        self.expected_type_map = ["foo", "bar"]
+        self.expected_dim_fparam = 0
+        self.expected_dim_aparam = 0
+        self.expected_sel_type = [0, 1]
+        self.expected_aparam_nall = False
+        self.expected_model_output_type = ["energy", "mask"]
+        self.expected_sel = [8, 12]
+        ds = DescrptSeADP(
+            rcut=self.expected_rcut,
+            rcut_smth=self.expected_rcut / 2,
+            sel=self.expected_sel,
+        )
+        ft = EnergyFittingNetDP(
+            ntypes=len(self.expected_type_map),
+            dim_descrpt=ds.get_dim_out(),
+            mixed_types=ds.mixed_types(),
+        )
+        self.dp_module = EnergyModelDP(
+            ds,
+            ft,
+            type_map=self.expected_type_map,
+        )
+
+        if INSTALLED_PT:
+            self.pt_module = EnergyModelPT.deserialize(self.dp_module.serialize())