From cff3a51fa3e4f732ac5247998c1eaa438bf253bd Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Wed, 21 May 2025 22:41:59 +0800
Subject: [PATCH 1/5] feat(pt/dp): add distance init for DPA3 edge feat

---
 deepmd/dpmodel/descriptor/dpa3.py                |  7 +++++++
 deepmd/dpmodel/descriptor/repflows.py            | 16 ++++++++++++++--
 deepmd/pt/model/descriptor/dpa3.py               |  1 +
 deepmd/pt/model/descriptor/repflows.py           | 15 +++++++++++++--
 deepmd/utils/argcheck.py                         | 12 ++++++++++++
 source/tests/consistent/descriptor/test_dpa3.py  | 14 +++++++++++++-
 .../dpmodel/descriptor/test_descriptor.py        |  5 ++++-
 7 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/dpa3.py b/deepmd/dpmodel/descriptor/dpa3.py
index 7d3292dff2..52740948e1 100644
--- a/deepmd/dpmodel/descriptor/dpa3.py
+++ b/deepmd/dpmodel/descriptor/dpa3.py
@@ -123,6 +123,9 @@ class RepFlowArgs:
     smooth_edge_update : bool, optional
         Whether to make edge update smooth.
         If True, the edge update from angle message will not use self as padding.
+    edge_init_use_dist : bool, optional
+        Whether to use direct distance r to initialize the edge features instead of 1/r.
+        Note that when using this option, the activation function will not be used when initializing edge features.
     """
 
     def __init__(
@@ -150,6 +153,7 @@ def __init__(
         skip_stat: bool = False,
         optim_update: bool = True,
         smooth_edge_update: bool = False,
+        edge_init_use_dist: bool = False,
     ) -> None:
         self.n_dim = n_dim
         self.e_dim = e_dim
@@ -176,6 +180,7 @@ def __init__(
         self.a_compress_use_split = a_compress_use_split
         self.optim_update = optim_update
         self.smooth_edge_update = smooth_edge_update
+        self.edge_init_use_dist = edge_init_use_dist
 
     def __getitem__(self, key):
         if hasattr(self, key):
@@ -207,6 +212,7 @@ def serialize(self) -> dict:
             "fix_stat_std": self.fix_stat_std,
             "optim_update": self.optim_update,
             "smooth_edge_update": self.smooth_edge_update,
+            "edge_init_use_dist": self.edge_init_use_dist,
         }
 
     @classmethod
@@ -303,6 +309,7 @@ def init_subclass_params(sub_data, sub_class):
             fix_stat_std=self.repflow_args.fix_stat_std,
             optim_update=self.repflow_args.optim_update,
             smooth_edge_update=self.repflow_args.smooth_edge_update,
+            edge_init_use_dist=self.repflow_args.edge_init_use_dist,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,
diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py
index 7273ba8ebe..ce27483a56 100644
--- a/deepmd/dpmodel/descriptor/repflows.py
+++ b/deepmd/dpmodel/descriptor/repflows.py
@@ -122,6 +122,9 @@ class DescrptBlockRepflows(NativeOP, DescriptorBlock):
     smooth_edge_update : bool, optional
         Whether to make edge update smooth.
         If True, the edge update from angle message will not use self as padding.
+    edge_init_use_dist : bool, optional
+        Whether to use direct distance r to initialize the edge features instead of 1/r.
+        Note that when using this option, the activation function will not be used when initializing edge features.
     ntypes : int
         Number of element types
     activation_function : str, optional
@@ -170,6 +173,7 @@ def __init__(
         fix_stat_std: float = 0.3,
         optim_update: bool = True,
         smooth_edge_update: bool = False,
+        edge_init_use_dist: bool = False,
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
         super().__init__()
@@ -201,6 +205,7 @@ def __init__(
         self.a_compress_use_split = a_compress_use_split
         self.optim_update = optim_update
         self.smooth_edge_update = smooth_edge_update
+        self.edge_init_use_dist = edge_init_use_dist
 
         self.n_dim = n_dim
         self.e_dim = e_dim
@@ -438,8 +443,14 @@ def call(
         # edge_input, h2 = xp.split(dmatrix, [1], axis=-1)
         edge_input = dmatrix[:, :, :, :1]
         h2 = dmatrix[:, :, :, 1:]
-        # nb x nloc x nnei x e_dim
-        edge_ebd = self.act(self.edge_embd(edge_input))
+        if self.edge_init_use_dist:
+            # nb x nloc x nnei x 1
+            edge_input = xp.linalg.vector_norm(diff, axis=-1, keepdims=True)
+            # nb x nloc x nnei x e_dim
+            edge_ebd = self.edge_embd(edge_input)
+        else:
+            # nb x nloc x nnei x e_dim
+            edge_ebd = self.act(self.edge_embd(edge_input))
 
         # get angle nlist (maybe smaller)
         a_dist_mask = (xp.linalg.vector_norm(diff, axis=-1) < self.a_rcut)[
@@ -577,6 +588,7 @@ def serialize(self):
             "precision": self.precision,
             "fix_stat_std": self.fix_stat_std,
             "optim_update": self.optim_update,
+            "edge_init_use_dist": self.edge_init_use_dist,
             # variables
             "edge_embd": self.edge_embd.serialize(),
             "angle_embd": self.angle_embd.serialize(),
diff --git a/deepmd/pt/model/descriptor/dpa3.py b/deepmd/pt/model/descriptor/dpa3.py
index 545da962e7..c4fd4cdac9 100644
--- a/deepmd/pt/model/descriptor/dpa3.py
+++ b/deepmd/pt/model/descriptor/dpa3.py
@@ -150,6 +150,7 @@ def init_subclass_params(sub_data, sub_class):
             fix_stat_std=self.repflow_args.fix_stat_std,
             optim_update=self.repflow_args.optim_update,
             smooth_edge_update=self.repflow_args.smooth_edge_update,
+            edge_init_use_dist=self.repflow_args.edge_init_use_dist,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,
diff --git a/deepmd/pt/model/descriptor/repflows.py b/deepmd/pt/model/descriptor/repflows.py
index 330336b1de..5cf1daa231 100644
--- a/deepmd/pt/model/descriptor/repflows.py
+++ b/deepmd/pt/model/descriptor/repflows.py
@@ -133,6 +133,9 @@ class DescrptBlockRepflows(DescriptorBlock):
     smooth_edge_update : bool, optional
         Whether to make edge update smooth.
         If True, the edge update from angle message will not use self as padding.
+    edge_init_use_dist : bool, optional
+        Whether to use direct distance r to initialize the edge features instead of 1/r.
+        Note that when using this option, the activation function will not be used when initializing edge features.
     optim_update : bool, optional
         Whether to enable the optimized update method.
         Uses a more efficient process when enabled. Defaults to True
@@ -183,6 +186,7 @@ def __init__(
         precision: str = "float64",
         fix_stat_std: float = 0.3,
         smooth_edge_update: bool = False,
+        edge_init_use_dist: bool = False,
         optim_update: bool = True,
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
@@ -215,6 +219,7 @@ def __init__(
         self.a_compress_use_split = a_compress_use_split
         self.optim_update = optim_update
         self.smooth_edge_update = smooth_edge_update
+        self.edge_init_use_dist = edge_init_use_dist
 
         self.n_dim = n_dim
         self.e_dim = e_dim
@@ -414,8 +419,14 @@ def forward(
         n_dim = node_ebd.shape[-1]
         # nb x nloc x nnei x 1,  nb x nloc x nnei x 3
         edge_input, h2 = torch.split(dmatrix, [1, 3], dim=-1)
-        # nb x nloc x nnei x e_dim
-        edge_ebd = self.act(self.edge_embd(edge_input))
+        if self.edge_init_use_dist:
+            # nb x nloc x nnei x 1
+            edge_input = torch.linalg.norm(diff, dim=-1, keepdim=True)
+            # nb x nloc x nnei x e_dim
+            edge_ebd = self.edge_embd(edge_input)
+        else:
+            # nb x nloc x nnei x e_dim
+            edge_ebd = self.act(self.edge_embd(edge_input))
 
         # get angle nlist (maybe smaller)
         a_dist_mask = (torch.linalg.norm(diff, dim=-1) < self.a_rcut)[
diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
index 0260700165..a19af92dad 100644
--- a/deepmd/utils/argcheck.py
+++ b/deepmd/utils/argcheck.py
@@ -1497,6 +1497,10 @@ def dpa3_repflow_args():
         "Whether to make edge update smooth. "
         "If True, the edge update from angle message will not use self as padding."
     )
+    doc_edge_init_use_dist = (
+        "Whether to use direct distance r to initialize the edge features instead of 1/r. "
+        "Note that when using this option, the activation function will not be used when initializing edge features."
+    )
 
     return [
         # repflow args
@@ -1597,6 +1601,14 @@ def dpa3_repflow_args():
             default=False,  # For compatability. This will be True in the future
             doc=doc_smooth_edge_update,
         ),
+        Argument(
+            "edge_init_use_dist",
+            bool,
+            optional=True,
+            default=False,
+            alias=["edge_use_dist"],
+            doc=doc_edge_init_use_dist,
+        ),
     ]
 
 
diff --git a/source/tests/consistent/descriptor/test_dpa3.py b/source/tests/consistent/descriptor/test_dpa3.py
index 56df87f931..d93dc2db89 100644
--- a/source/tests/consistent/descriptor/test_dpa3.py
+++ b/source/tests/consistent/descriptor/test_dpa3.py
@@ -65,6 +65,7 @@
     (1, 2),  # a_compress_e_rate
     (True, False),  # a_compress_use_split
     (True, False),  # optim_update
+    (True, False),  # edge_init_use_dist
     (0.3, 0.0),  # fix_stat_std
     (1, 2),  # n_multi_edge_message
     ("float64",),  # precision
@@ -80,6 +81,7 @@ def data(self) -> dict:
             a_compress_e_rate,
             a_compress_use_split,
             optim_update,
+            edge_init_use_dist,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -103,6 +105,7 @@ def data(self) -> dict:
                     "a_compress_e_rate": a_compress_e_rate,
                     "a_compress_use_split": a_compress_use_split,
                     "optim_update": optim_update,
+                    "edge_init_use_dist": edge_init_use_dist,
                     "fix_stat_std": fix_stat_std,
                     "n_multi_edge_message": n_multi_edge_message,
                     "axis_neuron": 4,
@@ -130,6 +133,7 @@ def skip_pt(self) -> bool:
             a_compress_e_rate,
             a_compress_use_split,
             optim_update,
+            edge_init_use_dist,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -146,11 +150,14 @@ def skip_pd(self) -> bool:
             a_compress_e_rate,
             a_compress_use_split,
             optim_update,
+            edge_init_use_dist,
             fix_stat_std,
             n_multi_edge_message,
             precision,
         ) = self.param
-        return not INSTALLED_PD or precision == "bfloat16"
+        return (
+            not INSTALLED_PD or precision == "bfloat16" or edge_init_use_dist
+        )  # not supported yet
 
     @property
     def skip_dp(self) -> bool:
@@ -162,6 +169,7 @@ def skip_dp(self) -> bool:
             a_compress_e_rate,
             a_compress_use_split,
             optim_update,
+            edge_init_use_dist,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -178,6 +186,7 @@ def skip_tf(self) -> bool:
             a_compress_e_rate,
             a_compress_use_split,
             optim_update,
+            edge_init_use_dist,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -236,6 +245,7 @@ def setUp(self) -> None:
             a_compress_e_rate,
             a_compress_use_split,
             optim_update,
+            edge_init_use_dist,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -315,6 +325,7 @@ def rtol(self) -> float:
             a_compress_e_rate,
             a_compress_use_split,
             optim_update,
+            edge_init_use_dist,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -337,6 +348,7 @@ def atol(self) -> float:
             a_compress_e_rate,
             a_compress_use_split,
             optim_update,
+            edge_init_use_dist,
             fix_stat_std,
             n_multi_edge_message,
             precision,
diff --git a/source/tests/universal/dpmodel/descriptor/test_descriptor.py b/source/tests/universal/dpmodel/descriptor/test_descriptor.py
index 4fa0593419..081285ff4c 100644
--- a/source/tests/universal/dpmodel/descriptor/test_descriptor.py
+++ b/source/tests/universal/dpmodel/descriptor/test_descriptor.py
@@ -482,6 +482,7 @@ def DescriptorParamDPA3(
     a_compress_use_split=False,
     optim_update=True,
     smooth_edge_update=False,
+    edge_init_use_dist=False,
     fix_stat_std=0.3,
     precision="float64",
 ):
@@ -504,6 +505,7 @@ def DescriptorParamDPA3(
                 "a_compress_use_split": a_compress_use_split,
                 "optim_update": optim_update,
                 "smooth_edge_update": smooth_edge_update,
+                "edge_init_use_dist": edge_init_use_dist,
                 "fix_stat_std": fix_stat_std,
                 "n_multi_edge_message": n_multi_edge_message,
                 "axis_neuron": 2,
@@ -540,8 +542,9 @@ def DescriptorParamDPA3(
             "a_compress_use_split": (True, False),
             "optim_update": (True, False),
             "smooth_edge_update": (True,),
+            "edge_init_use_dist": (True, False),
             "fix_stat_std": (0.3,),
-            "n_multi_edge_message": (1, 2),
+            "n_multi_edge_message": (1,),
             "env_protection": (0.0, 1e-8),
             "precision": ("float64",),
         }

From 5d7e15ae5ef7642bae2adcdaa1e136026ec6cd0c Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 26 May 2025 05:17:34 +0000
Subject: [PATCH 2/5] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 source/tests/consistent/descriptor/test_dpa3.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/source/tests/consistent/descriptor/test_dpa3.py b/source/tests/consistent/descriptor/test_dpa3.py
index 3b2dd75cae..c61a6c4fd3 100644
--- a/source/tests/consistent/descriptor/test_dpa3.py
+++ b/source/tests/consistent/descriptor/test_dpa3.py
@@ -162,7 +162,10 @@ def skip_pd(self) -> bool:
             precision,
         ) = self.param
         return (
-            not INSTALLED_PD or precision == "bfloat16" or edge_init_use_dist or use_dynamic_sel
+            not INSTALLED_PD
+            or precision == "bfloat16"
+            or edge_init_use_dist
+            or use_dynamic_sel
         )  # not supported yet
 
     @property

From 37361707afaa8ca50455c2030a174ce52fcf0e81 Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Mon, 26 May 2025 13:26:28 +0800
Subject: [PATCH 3/5] resolve conversations

---
 deepmd/dpmodel/descriptor/repflows.py  | 29 ++++++++-----------------
 deepmd/pt/model/descriptor/repflows.py | 30 +++++++-------------------
 2 files changed, 17 insertions(+), 42 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py
index 355b268ef1..c3f144384c 100644
--- a/deepmd/dpmodel/descriptor/repflows.py
+++ b/deepmd/dpmodel/descriptor/repflows.py
@@ -464,24 +464,6 @@ def call(
         # beyond the cutoff sw should be 0.0
         sw = xp.where(nlist_mask, sw, xp.zeros_like(sw))
 
-        # nb x nloc x tebd_dim
-        atype_embd = atype_embd_ext[:, :nloc, :]
-        assert list(atype_embd.shape) == [nframes, nloc, self.n_dim]
-
-        node_ebd = self.act(atype_embd)
-        # nb x nloc x nnei x 1,  nb x nloc x nnei x 3
-        # edge_input, h2 = xp.split(dmatrix, [1], axis=-1)
-        edge_input = dmatrix[:, :, :, :1]
-        h2 = dmatrix[:, :, :, 1:]
-        if self.edge_init_use_dist:
-            # nb x nloc x nnei x 1
-            edge_input = xp.linalg.vector_norm(diff, axis=-1, keepdims=True)
-            # nb x nloc x nnei x e_dim
-            edge_ebd = self.edge_embd(edge_input)
-        else:
-            # nb x nloc x nnei x e_dim
-            edge_ebd = self.act(self.edge_embd(edge_input))
-
         # get angle nlist (maybe smaller)
         a_dist_mask = (xp.linalg.vector_norm(diff, axis=-1) < self.a_rcut)[
             :, :, : self.a_sel
@@ -517,7 +499,11 @@ def call(
         # get edge and angle embedding input
         # nb x nloc x nnei x 1,  nb x nloc x nnei x 3
         # edge_input, h2 = xp.split(dmatrix, [1], axis=-1)
-        edge_input = dmatrix[:, :, :, :1]
+        if self.edge_init_use_dist:
+            # nb x nloc x nnei x 1
+            edge_input = xp.linalg.vector_norm(diff, axis=-1, keepdims=True)
+        else:
+            edge_input = dmatrix[:, :, :, :1]
         h2 = dmatrix[:, :, :, 1:]
 
         # nf x nloc x a_nnei x 3
@@ -559,7 +545,10 @@ def call(
 
         # get edge and angle embedding
         # nb x nloc x nnei x e_dim [OR] n_edge x e_dim
-        edge_ebd = self.act(self.edge_embd(edge_input))
+        if not self.edge_init_use_dist:
+            edge_ebd = self.act(self.edge_embd(edge_input))
+        else:
+            edge_ebd = self.edge_embd(edge_input)
         # nf x nloc x a_nnei x a_nnei x a_dim [OR] n_angle x a_dim
         angle_ebd = self.angle_embd(angle_input)
 
diff --git a/deepmd/pt/model/descriptor/repflows.py b/deepmd/pt/model/descriptor/repflows.py
index c793db1b86..ea993e1838 100644
--- a/deepmd/pt/model/descriptor/repflows.py
+++ b/deepmd/pt/model/descriptor/repflows.py
@@ -436,27 +436,6 @@ def forward(
         # beyond the cutoff sw should be 0.0
         sw = sw.masked_fill(~nlist_mask, 0.0)
 
-        # [nframes, nloc, tebd_dim]
-        if comm_dict is None:
-            assert isinstance(extended_atype_embd, torch.Tensor)  # for jit
-            atype_embd = extended_atype_embd[:, :nloc, :]
-            assert list(atype_embd.shape) == [nframes, nloc, self.n_dim]
-        else:
-            atype_embd = extended_atype_embd
-        assert isinstance(atype_embd, torch.Tensor)  # for jit
-        node_ebd = self.act(atype_embd)
-        n_dim = node_ebd.shape[-1]
-        # nb x nloc x nnei x 1,  nb x nloc x nnei x 3
-        edge_input, h2 = torch.split(dmatrix, [1, 3], dim=-1)
-        if self.edge_init_use_dist:
-            # nb x nloc x nnei x 1
-            edge_input = torch.linalg.norm(diff, dim=-1, keepdim=True)
-            # nb x nloc x nnei x e_dim
-            edge_ebd = self.edge_embd(edge_input)
-        else:
-            # nb x nloc x nnei x e_dim
-            edge_ebd = self.act(self.edge_embd(edge_input))
-
         # get angle nlist (maybe smaller)
         a_dist_mask = (torch.linalg.norm(diff, dim=-1) < self.a_rcut)[
             :, :, : self.a_sel
@@ -497,6 +476,10 @@ def forward(
         # get edge and angle embedding input
         # nb x nloc x nnei x 1,  nb x nloc x nnei x 3
         edge_input, h2 = torch.split(dmatrix, [1, 3], dim=-1)
+        if self.edge_init_use_dist:
+            # nb x nloc x nnei x 1
+            edge_input = torch.linalg.norm(diff, dim=-1, keepdim=True)
+
         # nf x nloc x a_nnei x 3
         normalized_diff_i = a_diff / (
             torch.linalg.norm(a_diff, dim=-1, keepdim=True) + 1e-6
@@ -533,7 +516,10 @@ def forward(
             )
         # get edge and angle embedding
         # nb x nloc x nnei x e_dim [OR] n_edge x e_dim
-        edge_ebd = self.act(self.edge_embd(edge_input))
+        if not self.edge_init_use_dist:
+            edge_ebd = self.act(self.edge_embd(edge_input))
+        else:
+            edge_ebd = self.edge_embd(edge_input)
         # nf x nloc x a_nnei x a_nnei x a_dim [OR] n_angle x a_dim
         angle_ebd = self.angle_embd(angle_input)
 

From 84aceae0150399befcb829cd687c62312fb873a0 Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Mon, 26 May 2025 13:28:01 +0800
Subject: [PATCH 4/5] Update repflows.py

---
 deepmd/dpmodel/descriptor/repflows.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py
index c3f144384c..284dd58fcf 100644
--- a/deepmd/dpmodel/descriptor/repflows.py
+++ b/deepmd/dpmodel/descriptor/repflows.py
@@ -499,8 +499,8 @@ def call(
         # get edge and angle embedding input
         # nb x nloc x nnei x 1,  nb x nloc x nnei x 3
         # edge_input, h2 = xp.split(dmatrix, [1], axis=-1)
+        # nb x nloc x nnei x 1
         if self.edge_init_use_dist:
-            # nb x nloc x nnei x 1
             edge_input = xp.linalg.vector_norm(diff, axis=-1, keepdims=True)
         else:
             edge_input = dmatrix[:, :, :, :1]

From 9e6027c36fa5d19de1932b9c197a404565ef5537 Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Mon, 26 May 2025 13:37:28 +0800
Subject: [PATCH 5/5] Update argcheck.py

---
 deepmd/utils/argcheck.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
index 901265e358..1b320bcb7c 100644
--- a/deepmd/utils/argcheck.py
+++ b/deepmd/utils/argcheck.py
@@ -1500,6 +1500,7 @@ def dpa3_repflow_args():
     doc_edge_init_use_dist = (
         "Whether to use direct distance r to initialize the edge features instead of 1/r. "
         "Note that when using this option, the activation function will not be used when initializing edge features."
+    )
     doc_use_dynamic_sel = (
         "Whether to dynamically select neighbors within the cutoff radius. "
         "If True, the exact number of neighbors within the cutoff radius is used "