From 27f14cd6f72e5453580552aca53a6703ab2ad901 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Wed, 31 Jan 2024 14:47:22 +0800
Subject: [PATCH] dp model format for dp models.

---
 deepmd/model_format/__init__.py               |  4 +++
 deepmd/model_format/atomic_model.py           |  2 +-
 deepmd/model_format/se_e2_a.py                |  8 +++++
 .../tests/common/test_model_format_utils.py   | 35 ++++++++++++++++++-
 4 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/deepmd/model_format/__init__.py b/deepmd/model_format/__init__.py
index 71a9ceee21..38873469e8 100644
--- a/deepmd/model_format/__init__.py
+++ b/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",
diff --git a/deepmd/model_format/atomic_model.py b/deepmd/model_format/atomic_model.py
index 808c9a82a9..da3868e8f1 100644
--- a/deepmd/model_format/atomic_model.py
+++ b/deepmd/model_format/atomic_model.py
@@ -10,7 +10,7 @@
     Optional,
 )
 
-from deepmd.model_format import (
+from .output_def import (
     FittingOutputDef,
 )
 
diff --git a/deepmd/model_format/se_e2_a.py b/deepmd/model_format/se_e2_a.py
index 5f773694f8..9e61f2b115 100644
--- a/deepmd/model_format/se_e2_a.py
+++ b/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.
diff --git a/source/tests/common/test_model_format_utils.py b/source/tests/common/test_model_format_utils.py
index a374da103b..0a457e0802 100644
--- a/source/tests/common/test_model_format_utils.py
+++ b/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(