Merge branch 'devel' into D0516_dynamic_sel

Author: iProzd

.pre-commit-config.yaml

@@ -29,7 +29,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.11.9
+    rev: v0.11.10
     hooks:
       - id: ruff
         args: ["--fix"]

deepmd/pd/model/model/transform_output.py

@@ -82,7 +82,9 @@ def task_deriv_one(
     assert extended_force is not None
     extended_force = -extended_force
     if do_virial:
-        extended_virial = extended_force.unsqueeze(-1) @ extended_coord.unsqueeze(-2)
+        extended_virial = paddle.einsum(
+            "...ik,...ij->...ikj", extended_force, extended_coord
+        )
         # the correction sums to zero, which does not contribute to global virial
         if do_atomic_virial:
             extended_virial_corr = atomic_virial_corr(extended_coord, atom_energy)

deepmd/pt/model/model/transform_output.py

@@ -85,7 +85,9 @@ def task_deriv_one(
     assert extended_force is not None
     extended_force = -extended_force
     if do_virial:
-        extended_virial = extended_force.unsqueeze(-1) @ extended_coord.unsqueeze(-2)
+        extended_virial = torch.einsum(
+            "...ik,...ij->...ikj", extended_force, extended_coord
+        )
         # the correction sums to zero, which does not contribute to global virial
         if do_atomic_virial:
             extended_virial_corr = atomic_virial_corr(extended_coord, atom_energy)

deepmd/tf/infer/deep_eval.py

@@ -259,7 +259,9 @@ def _init_attr(self) -> None:
             self.numb_dos = 0
         self.tmap = tmap.decode("utf-8").split()
         if self.tensors["modifier_type"] is not None:
-            self.modifier_type = run_sess(self.sess, [self.tensors["modifier_type"]])[0]
+            self.modifier_type = run_sess(self.sess, [self.tensors["modifier_type"]])[
+                0
+            ].decode()
         else:
             self.modifier_type = None
 
@@ -761,15 +763,17 @@ def eval(
             odef.name: oo for oo, odef in zip(output, self.output_def.var_defs.values())
         }
         # ugly!!
-        if self.modifier_type is not None and isinstance(self.model_type, DeepPot):
+        if self.modifier_type is not None and issubclass(self.model_type, DeepPot):
             if atomic:
                 raise RuntimeError("modifier does not support atomic modification")
             me, mf, mv = self.dm.eval(coords, cells, atom_types)
-            output = list(output)  # tuple to list
-            e, f, v = output[:3]
-            output_dict["energy_redu"] += me.reshape(e.shape)
-            output_dict["energy_deri_r"] += mf.reshape(f.shape)
-            output_dict["energy_deri_c_redu"] += mv.reshape(v.shape)
+            output_dict["energy_redu"] += me.reshape(output_dict["energy_redu"].shape)
+            output_dict["energy_derv_r"] += mf.reshape(
+                output_dict["energy_derv_r"].shape
+            )
+            output_dict["energy_derv_c_redu"] += mv.reshape(
+                output_dict["energy_derv_c_redu"].shape
+            )
         return output_dict
 
     def _prepare_feed_dict(
@@ -1348,6 +1352,8 @@ def sort_input(
             natoms = atom_type[0].size
             idx_map = np.arange(natoms)  # pylint: disable=no-explicit-dtype
             return coord, atom_type, idx_map
+        if atom_type.ndim > 1:
+            atom_type = atom_type[0]
         if sel_atoms is not None:
             selection = [False] * np.size(atom_type)
             for ii in sel_atoms: