[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

deepmd/pt/model/descriptor/repflow_layer.py

@@ -210,8 +210,12 @@ 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(
@@ -327,7 +331,10 @@ 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
@@ -375,9 +382,16 @@ 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
@@ -530,7 +544,9 @@ 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)
@@ -541,7 +557,11 @@ 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
 
@@ -565,15 +585,19 @@ 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)
@@ -582,7 +606,13 @@ 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(
@@ -615,7 +645,9 @@ 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(
@@ -643,7 +675,9 @@ 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)
@@ -742,7 +776,9 @@ 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]
@@ -815,7 +851,9 @@ 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,
                     ],
@@ -828,7 +866,9 @@ 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(
@@ -864,7 +904,9 @@ 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:
@@ -920,7 +962,9 @@ 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 = (
@@ -938,18 +982,24 @@ 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)
@@ -987,9 +1037,15 @@ 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(
                     [
@@ -1017,7 +1073,9 @@ 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,
@@ -1029,8 +1087,12 @@ 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")
 
@@ -1088,7 +1150,9 @@ 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
@@ -1106,7 +1170,9 @@ def list_update_res_residual(self, update_list: list[torch.Tensor], update_name:
         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/network/utils.py

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