update dynamic sel

Author: iProzd

deepmd/pt/model/descriptor/repflow_layer.py

@@ -210,12 +210,8 @@ def __init__(
                 self.n_a_compress_dim = n_dim
             else:
                 # angle + a_dim/c + a_dim/2c * 2 * e_rate
-                self.angle_dim += (1 + self.a_compress_e_rate) * (
-                    self.a_dim // self.a_compress_rate
-                )
-                self.e_a_compress_dim = (
-                    self.a_dim // (2 * self.a_compress_rate) * self.a_compress_e_rate
-                )
+                self.angle_dim += (1 + self.a_compress_e_rate) * (self.a_dim // self.a_compress_rate)
+                self.e_a_compress_dim = self.a_dim // (2 * self.a_compress_rate) * self.a_compress_e_rate
                 self.n_a_compress_dim = self.a_dim // self.a_compress_rate
                 if not self.a_compress_use_split:
                     self.a_compress_n_linear = MLPLayer(
@@ -331,10 +327,7 @@ def _cal_hg(
         # nb x nloc x nnei x e_dim
         edge_ebd = _apply_nlist_mask(edge_ebd, nlist_mask)
         edge_ebd = _apply_switch(edge_ebd, sw)
-        invnnei = torch.rsqrt(
-            float(nnei)
-            * torch.ones((nb, nloc, 1, 1), dtype=edge_ebd.dtype, device=edge_ebd.device)
-        )
+        invnnei = torch.rsqrt(float(nnei) * torch.ones((nb, nloc, 1, 1), dtype=edge_ebd.dtype, device=edge_ebd.device))
         # nb x nloc x 3 x e_dim
         h2g2 = torch.matmul(torch.transpose(h2, -1, -2), edge_ebd) * invnnei
         return h2g2
@@ -382,16 +375,9 @@ def _cal_hg_dynamic(
         # n_edge x e_dim
         flat_edge_ebd = flat_edge_ebd * flat_sw.unsqueeze(-1)
         # n_edge x 3 x e_dim
-        flat_h2g2 = (flat_h2.unsqueeze(-1) * flat_edge_ebd.unsqueeze(-2)).reshape(
-            -1, 3 * e_dim
-        )
+        flat_h2g2 = (flat_h2.unsqueeze(-1) * flat_edge_ebd.unsqueeze(-2)).reshape(-1, 3 * e_dim)
         # nf x nloc x 3 x e_dim
-        h2g2 = (
-            aggregate(flat_h2g2, owner, average=False, num_owner=num_owner).reshape(
-                nb, nloc, 3, e_dim
-            )
-            * scale_factor
-        )
+        h2g2 = aggregate(flat_h2g2, owner, average=False, num_owner=num_owner).reshape(nb, nloc, 3, e_dim) * scale_factor
         return h2g2
 
     @staticmethod
@@ -544,9 +530,7 @@ def optim_angle_update(
         node_dim = node_ebd.shape[-1]
         edge_dim = edge_ebd.shape[-1]
         # angle_dim, node_dim, edge_dim, edge_dim
-        sub_angle, sub_node, sub_edge_ik, sub_edge_ij = torch.split(
-            matrix, [angle_dim, node_dim, edge_dim, edge_dim]
-        )
+        sub_angle, sub_node, sub_edge_ik, sub_edge_ij = torch.split(matrix, [angle_dim, node_dim, edge_dim, edge_dim])
 
         # nf * nloc * a_sel * a_sel * angle_dim
         sub_angle_update = torch.matmul(angle_ebd, sub_angle)
@@ -557,11 +541,7 @@ def optim_angle_update(
         sub_edge_update_ij = torch.matmul(edge_ebd, sub_edge_ij)
 
         result_update = (
-            bias
-            + sub_node_update.unsqueeze(2).unsqueeze(3)
-            + sub_edge_update_ik.unsqueeze(2)
-            + sub_edge_update_ij.unsqueeze(3)
-            + sub_angle_update
+            bias + sub_node_update.unsqueeze(2).unsqueeze(3) + sub_edge_update_ik.unsqueeze(2) + sub_edge_update_ij.unsqueeze(3) + sub_angle_update
         )
         return result_update
 
@@ -585,19 +565,15 @@ def optim_angle_update_dynamic(
         edge_dim = flat_edge_ebd.shape[-1]
         angle_dim = flat_angle_ebd.shape[-1]
         # angle_dim, node_dim, edge_dim, edge_dim
-        sub_angle, sub_node, sub_edge_ik, sub_edge_ij = torch.split(
-            matrix, [angle_dim, node_dim, edge_dim, edge_dim]
-        )
+        sub_angle, sub_node, sub_edge_ik, sub_edge_ij = torch.split(matrix, [angle_dim, node_dim, edge_dim, edge_dim])
 
         # n_angle * angle_dim
         sub_angle_update = torch.matmul(flat_angle_ebd, sub_angle)
 
         # nf * nloc * angle_dim
         sub_node_update = torch.matmul(node_ebd, sub_node)
         # n_angle * angle_dim
-        sub_node_update = torch.index_select(
-            sub_node_update.reshape(nf * nloc, sub_node_update.shape[-1]), 0, n2a_index
-        )
+        sub_node_update = torch.index_select(sub_node_update.reshape(nf * nloc, sub_node_update.shape[-1]), 0, n2a_index)
 
         # n_edge * angle_dim
         sub_edge_update_ik = torch.matmul(flat_edge_ebd, sub_edge_ik)
@@ -606,13 +582,7 @@ def optim_angle_update_dynamic(
         sub_edge_update_ik = torch.index_select(sub_edge_update_ik, 0, eik2a_index)
         sub_edge_update_ij = torch.index_select(sub_edge_update_ij, 0, eij2a_index)
 
-        result_update = (
-            bias
-            + sub_node_update
-            + sub_edge_update_ik
-            + sub_edge_update_ij
-            + sub_angle_update
-        )
+        result_update = bias + sub_node_update + sub_edge_update_ik + sub_edge_update_ij + sub_angle_update
         return result_update
 
     def optim_edge_update(
@@ -645,9 +615,7 @@ def optim_edge_update(
         # nf * nloc * nnei * node/edge_dim
         sub_edge_update = torch.matmul(edge_ebd, edge)
 
-        result_update = (
-            bias + sub_node_update.unsqueeze(2) + sub_edge_update + sub_node_ext_update
-        )
+        result_update = bias + sub_node_update.unsqueeze(2) + sub_edge_update + sub_node_ext_update
         return result_update
 
     def optim_edge_update_dynamic(
@@ -675,9 +643,7 @@ def optim_edge_update_dynamic(
         # nf * nloc * node/edge_dim
         sub_node_update = torch.matmul(node_ebd, node)
         # n_edge * node/edge_dim
-        sub_node_update = torch.index_select(
-            sub_node_update.reshape(nf * nloc, sub_node_update.shape[-1]), 0, n2e_index
-        )
+        sub_node_update = torch.index_select(sub_node_update.reshape(nf * nloc, sub_node_update.shape[-1]), 0, n2e_index)
 
         # nf * nall * node/edge_dim
         sub_node_ext_update = torch.matmul(node_ebd_ext, node_ext)
@@ -757,14 +723,12 @@ def forward(
         nb, nloc, nnei = nlist.shape
         nall = node_ebd_ext.shape[1]
         node_ebd = node_ebd_ext[:, :nloc, :]
+        n_edge = int(nlist_mask.sum().item()) if self.use_dynamic_sel else 0
         assert (nb, nloc) == node_ebd.shape[:2]
         if not self.use_dynamic_sel:
             assert (nb, nloc, nnei, 3) == h2.shape
-            n_edge = None
         else:
-            # n_edge = int(nlist_mask.sum().item())
-            # assert (n_edge, 3) == h2.shape
-            n_edge = h2.shape[0]
+            assert (n_edge, 3) == h2.shape
         del a_nlist  # may be used in the future
 
         n2e_index, n_ext2e_index = edge_index[0], edge_index[1]
@@ -778,9 +742,7 @@ def forward(
         nei_node_ebd = (
             _make_nei_g1(node_ebd_ext, nlist)
             if not self.use_dynamic_sel
-            else torch.index_select(
-                node_ebd_ext.reshape(-1, self.n_dim), 0, n_ext2e_index
-            )
+            else torch.index_select(node_ebd_ext.reshape(-1, self.n_dim), 0, n_ext2e_index)
         )
 
         n_update_list: list[torch.Tensor] = [node_ebd]
@@ -853,9 +815,7 @@ def forward(
                 # n_edge x (n_dim * 2 + e_dim)
                 edge_info = torch.cat(
                     [
-                        torch.index_select(
-                            node_ebd.reshape(-1, self.n_dim), 0, n2e_index
-                        ),
+                        torch.index_select(node_ebd.reshape(-1, self.n_dim), 0, n2e_index),
                         nei_node_ebd,
                         edge_ebd,
                     ],
@@ -868,9 +828,7 @@ def forward(
         # nb x nloc x nnei x (h * n_dim)
         if not self.optim_update:
             assert edge_info is not None
-            node_edge_update = self.act(
-                self.node_edge_linear(edge_info)
-            ) * sw.unsqueeze(-1)
+            node_edge_update = self.act(self.node_edge_linear(edge_info)) * sw.unsqueeze(-1)
         else:
             node_edge_update = self.act(
                 self.optim_edge_update(
@@ -906,9 +864,7 @@ def forward(
 
         if self.n_multi_edge_message > 1:
             # nb x nloc x h x n_dim
-            node_edge_update_mul_head = node_edge_update.view(
-                nb, nloc, self.n_multi_edge_message, self.n_dim
-            )
+            node_edge_update_mul_head = node_edge_update.view(nb, nloc, self.n_multi_edge_message, self.n_dim)
             for head_index in range(self.n_multi_edge_message):
                 n_update_list.append(node_edge_update_mul_head[..., head_index, :])
         else:
@@ -964,9 +920,7 @@ def forward(
                 # nb x nloc x a_nnei x e_dim
                 edge_ebd_for_angle = edge_ebd_for_angle[..., : self.a_sel, :]
                 # nb x nloc x a_nnei x e_dim
-                edge_ebd_for_angle = torch.where(
-                    a_nlist_mask.unsqueeze(-1), edge_ebd_for_angle, 0.0
-                )
+                edge_ebd_for_angle = torch.where(a_nlist_mask.unsqueeze(-1), edge_ebd_for_angle, 0.0)
             if not self.optim_update:
                 # nb x nloc x a_nnei x a_nnei x n_dim [OR] n_angle x n_dim
                 node_for_angle_info = (
@@ -984,24 +938,18 @@ def forward(
 
                 # nb x nloc x (a_nnei) x a_nnei x e_dim [OR] n_angle x e_dim
                 edge_for_angle_k = (
-                    torch.tile(
-                        edge_ebd_for_angle.unsqueeze(2), (1, 1, self.a_sel, 1, 1)
-                    )
+                    torch.tile(edge_ebd_for_angle.unsqueeze(2), (1, 1, self.a_sel, 1, 1))
                     if not self.use_dynamic_sel
                     else torch.index_select(edge_ebd_for_angle, 0, eik2a_index)
                 )
                 # nb x nloc x a_nnei x (a_nnei) x e_dim [OR] n_angle x e_dim
                 edge_for_angle_j = (
-                    torch.tile(
-                        edge_ebd_for_angle.unsqueeze(3), (1, 1, 1, self.a_sel, 1)
-                    )
+                    torch.tile(edge_ebd_for_angle.unsqueeze(3), (1, 1, 1, self.a_sel, 1))
                     if not self.use_dynamic_sel
                     else torch.index_select(edge_ebd_for_angle, 0, eij2a_index)
                 )
                 # nb x nloc x a_nnei x a_nnei x (e_dim + e_dim) [OR] n_angle x (e_dim + e_dim)
-                edge_for_angle_info = torch.cat(
-                    [edge_for_angle_k, edge_for_angle_j], dim=-1
-                )
+                edge_for_angle_info = torch.cat([edge_for_angle_k, edge_for_angle_j], dim=-1)
                 angle_info_list = [angle_ebd]
                 angle_info_list.append(node_for_angle_info)
                 angle_info_list.append(edge_for_angle_info)
@@ -1039,15 +987,9 @@ def forward(
 
             if not self.use_dynamic_sel:
                 # nb x nloc x a_nnei x a_nnei x e_dim
-                weighted_edge_angle_update = (
-                    a_sw.unsqueeze(-1).unsqueeze(-1)
-                    * a_sw.unsqueeze(-2).unsqueeze(-1)
-                    * edge_angle_update
-                )
+                weighted_edge_angle_update = a_sw.unsqueeze(-1).unsqueeze(-1) * a_sw.unsqueeze(-2).unsqueeze(-1) * edge_angle_update
                 # nb x nloc x a_nnei x e_dim
-                reduced_edge_angle_update = torch.sum(
-                    weighted_edge_angle_update, dim=-2
-                ) / (self.a_sel**0.5)
+                reduced_edge_angle_update = torch.sum(weighted_edge_angle_update, dim=-2) / (self.a_sel**0.5)
                 # nb x nloc x nnei x e_dim
                 padding_edge_angle_update = torch.concat(
                     [
@@ -1075,9 +1017,7 @@ def forward(
                 # will be deprecated in the future
                 # not support dynamic index, will pass anyway
                 if self.use_dynamic_sel:
-                    raise NotImplementedError(
-                        "smooth_edge_update must be True when use_dynamic_sel is True!"
-                    )
+                    raise NotImplementedError("smooth_edge_update must be True when use_dynamic_sel is True!")
                 full_mask = torch.concat(
                     [
                         a_nlist_mask,
@@ -1089,12 +1029,8 @@ def forward(
                     ],
                     dim=-1,
                 )
-                padding_edge_angle_update = torch.where(
-                    full_mask.unsqueeze(-1), padding_edge_angle_update, edge_ebd
-                )
-            e_update_list.append(
-                self.act(self.edge_angle_linear2(padding_edge_angle_update))
-            )
+                padding_edge_angle_update = torch.where(full_mask.unsqueeze(-1), padding_edge_angle_update, edge_ebd)
+            e_update_list.append(self.act(self.edge_angle_linear2(padding_edge_angle_update)))
             # update edge_ebd
             e_updated = self.list_update(e_update_list, "edge")
 
@@ -1152,9 +1088,7 @@ def list_update_res_incr(self, update_list: list[torch.Tensor]) -> torch.Tensor:
         return uu
 
     @torch.jit.export
-    def list_update_res_residual(
-        self, update_list: list[torch.Tensor], update_name: str = "node"
-    ) -> torch.Tensor:
+    def list_update_res_residual(self, update_list: list[torch.Tensor], update_name: str = "node") -> torch.Tensor:
         nitem = len(update_list)
         uu = update_list[0]
         # make jit happy
@@ -1172,9 +1106,7 @@ def list_update_res_residual(
         return uu
 
     @torch.jit.export
-    def list_update(
-        self, update_list: list[torch.Tensor], update_name: str = "node"
-    ) -> torch.Tensor:
+    def list_update(self, update_list: list[torch.Tensor], update_name: str = "node") -> torch.Tensor:
         if self.update_style == "res_avg":
             return self.list_update_res_avg(update_list)
         elif self.update_style == "res_incr":

deepmd/pt/model/descriptor/repflows.py

@@ -472,17 +472,6 @@ def forward(
         # beyond the cutoff sw should be 0.0
         sw = sw.masked_fill(~nlist_mask, 0.0)
 
-        # [nframes, nloc, tebd_dim]
-        assert extended_atype_embd is not None
-        atype_embd = extended_atype_embd[:, :nloc, :]
-        assert list(atype_embd.shape) == [nframes, nloc, self.n_dim]
-        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)
-        # 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

deepmd/pt/model/network/utils.py

@@ -1,4 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Optional,
+)
 
 import torch