From f181a3084a04f8e797f5474e14015d46e597e98a Mon Sep 17 00:00:00 2001
From: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com>
Date: Mon, 8 Jan 2024 13:45:39 +0800
Subject: [PATCH] model format for the embedding net (#3113)

- refact the dir structure for the model_format
- change the input of `NativeNet` from `List[NativeLayer]` to the
`List[dict]`

---------

Signed-off-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com>
Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 deepmd_utils/model_format/__init__.py         | 18 +++++
 .../network.py}                               | 66 ++++++++++++++++++-
 source/tests/test_model_format_utils.py       | 12 ++++
 3 files changed, 93 insertions(+), 3 deletions(-)
 create mode 100644 deepmd_utils/model_format/__init__.py
 rename deepmd_utils/{model_format.py => model_format/network.py} (81%)

diff --git a/deepmd_utils/model_format/__init__.py b/deepmd_utils/model_format/__init__.py
new file mode 100644
index 0000000000..4b33aa0151
--- /dev/null
+++ b/deepmd_utils/model_format/__init__.py
@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from .network import (
+    EmbeddingNet,
+    NativeLayer,
+    NativeNet,
+    load_dp_model,
+    save_dp_model,
+    traverse_model_dict,
+)
+
+__all__ = [
+    "EmbeddingNet",
+    "NativeLayer",
+    "NativeNet",
+    "load_dp_model",
+    "save_dp_model",
+    "traverse_model_dict",
+]
diff --git a/deepmd_utils/model_format.py b/deepmd_utils/model_format/network.py
similarity index 81%
rename from deepmd_utils/model_format.py
rename to deepmd_utils/model_format/network.py
index 131be93121..04aaa75534 100644
--- a/deepmd_utils/model_format.py
+++ b/deepmd_utils/model_format/network.py
@@ -275,10 +275,10 @@ class NativeNet(NativeOP):
         The layers of the network.
     """
 
-    def __init__(self, layers: Optional[List[NativeLayer]] = None) -> None:
+    def __init__(self, layers: Optional[List[dict]] = None) -> None:
         if layers is None:
             layers = []
-        self.layers = layers
+        self.layers = [NativeLayer.deserialize(layer) for layer in layers]
 
     def serialize(self) -> dict:
         """Serialize the network to a dict.
@@ -299,7 +299,7 @@ def deserialize(cls, data: dict) -> "NativeNet":
         data : dict
             The dict to deserialize from.
         """
-        return cls([NativeLayer.deserialize(layer) for layer in data["layers"]])
+        return cls(data["layers"])
 
     def __getitem__(self, key):
         assert isinstance(key, int)
@@ -329,3 +329,63 @@ def call(self, x: np.ndarray) -> np.ndarray:
         for layer in self.layers:
             x = layer.call(x)
         return x
+
+
+class EmbeddingNet(NativeNet):
+    def __init__(
+        self,
+        in_dim,
+        neuron: List[int] = [24, 48, 96],
+        activation_function: str = "tanh",
+        resnet_dt: bool = False,
+    ):
+        layers = []
+        i_in = in_dim
+        rng = np.random.default_rng()
+        for idx, ii in enumerate(neuron):
+            i_ot = ii
+            layers.append(
+                NativeLayer(
+                    rng.normal(size=(i_in, i_ot)),
+                    b=rng.normal(size=(ii)),
+                    idt=rng.normal(size=(ii)) if resnet_dt else None,
+                    activation_function=activation_function,
+                    resnet=True,
+                ).serialize()
+            )
+            i_in = i_ot
+        super().__init__(layers)
+        self.in_dim = in_dim
+        self.neuron = neuron
+        self.activation_function = activation_function
+        self.resnet_dt = resnet_dt
+
+    def serialize(self) -> dict:
+        """Serialize the network to a dict.
+
+        Returns
+        -------
+        dict
+            The serialized network.
+        """
+        return {
+            "in_dim": self.in_dim,
+            "neuron": self.neuron.copy(),
+            "activation_function": self.activation_function,
+            "resnet_dt": self.resnet_dt,
+            "layers": [layer.serialize() for layer in self.layers],
+        }
+
+    @classmethod
+    def deserialize(cls, data: dict) -> "EmbeddingNet":
+        """Deserialize the network from a dict.
+
+        Parameters
+        ----------
+        data : dict
+            The dict to deserialize from.
+        """
+        layers = data.pop("layers")
+        obj = cls(**data)
+        super(EmbeddingNet, obj).__init__(layers)
+        return obj
diff --git a/source/tests/test_model_format_utils.py b/source/tests/test_model_format_utils.py
index af8c4361c8..f26ebbaa8d 100644
--- a/source/tests/test_model_format_utils.py
+++ b/source/tests/test_model_format_utils.py
@@ -9,6 +9,7 @@
 import numpy as np
 
 from deepmd_utils.model_format import (
+    EmbeddingNet,
     NativeLayer,
     NativeNet,
     load_dp_model,
@@ -90,6 +91,17 @@ def test_deserialize(self):
         np.testing.assert_array_equal(network[0]["resnet"], True)
         np.testing.assert_array_equal(network[1]["resnet"], True)
 
+    def test_embedding_net(self):
+        for ni, idt, act in itertools.product(
+            [1, 10],
+            [True, False],
+            ["tanh", "none"],
+        ):
+            en0 = EmbeddingNet(ni)
+            en1 = EmbeddingNet.deserialize(en0.serialize())
+            inp = np.ones([ni])
+            np.testing.assert_allclose(en0.call(inp), en1.call(inp))
+
 
 class TestDPModel(unittest.TestCase):
     def setUp(self) -> None: