dp model format for dp models.

deepmd/model_format/__init__.py

@@ -7,6 +7,9 @@
 from .dp_atomic_model import (
     DPAtomicModel,
 )
+from .dp_model import (
+    DPModel,
+)
 from .env_mat import (
     EnvMat,
 )
@@ -40,6 +43,7 @@
 )
 
 __all__ = [
+    "DPModel",
     "DPAtomicModel",
     "InvarFitting",
     "DescrptSeA",

deepmd/model_format/atomic_model.py

@@ -10,7 +10,7 @@
     Optional,
 )
 
-from deepmd.model_format import (
+from .output_def import (
     FittingOutputDef,
 )
 

deepmd/model_format/se_e2_a.py

@@ -200,6 +200,14 @@ def get_dim_out(self):
         """Returns the output dimension of this descriptor."""
         return self.neuron[-1] * self.axis_neuron
 
+    def get_rcut(self):
+        """Returns cutoff radius."""
+        return self.rcut
+
+    def get_sel(self):
+        """Returns cutoff radius."""
+        return self.sel
+
     def distinguish_types(self):
         """Returns if the descriptor uses different nets for
         different atomic types.

source/tests/common/test_model_format_utils.py

@@ -11,6 +11,7 @@
 from deepmd.model_format import (
     DescrptSeA,
     DPAtomicModel,
+    DPModel,
     EmbeddingNet,
     EnvMat,
     FittingNet,
@@ -277,7 +278,7 @@ def test_zero_dim(self):
         )
 
 
-class TestDPModel(unittest.TestCase):
+class TestSaveLoadDPModel(unittest.TestCase):
     def setUp(self) -> None:
         self.w = np.full((3, 2), 3.0)
         self.b = np.full((3,), 4.0)
@@ -534,6 +535,38 @@ def test_self_consistency(
         np.testing.assert_allclose(ret0["energy"], ret1["energy"])
 
 
+class TestDPModel(unittest.TestCase, TestCaseSingleFrameWithNlist):
+    def setUp(self):
+        TestCaseSingleFrameWithNlist.setUp(self)
+
+    def test_self_consistency(
+        self,
+    ):
+        rng = np.random.default_rng()
+        nf, nloc, nnei = self.nlist.shape
+        ds = DescrptSeA(
+            self.rcut,
+            self.rcut_smth,
+            self.sel,
+        )
+        ft = InvarFitting(
+            "energy",
+            self.nt,
+            ds.get_dim_out(),
+            1,
+            distinguish_types=ds.distinguish_types(),
+        )
+        type_map = ["foo", "bar"]
+        md0 = DPModel(ds, ft, type_map=type_map)
+        md1 = DPModel.deserialize(md0.serialize())
+
+        ret0 = md0.call_lower(self.coord_ext, self.atype_ext, self.nlist)
+        ret1 = md1.call_lower(self.coord_ext, self.atype_ext, self.nlist)
+
+        np.testing.assert_allclose(ret0["energy"], ret1["energy"])
+        np.testing.assert_allclose(ret0["energy_redu"], ret1["energy_redu"])
+
+
 class TestRegion(unittest.TestCase):
     def setUp(self):
         self.cell = np.array(