test(pt/dp): add universal uts for all models

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -196,11 +196,11 @@ def forward_atomic(
         ]
         ener_list = []
         for i, model in enumerate(self.models):
-            mapping = self.mapping_list[i]
+            type_map_model = self.mapping_list[i]
             ener_list.append(
                 model.forward_atomic(
                     extended_coord,
-                    mapping[extended_atype],
+                    type_map_model[extended_atype],
                     nlists_[i],
                     mapping,
                     fparam,
@@ -414,7 +414,12 @@ def _compute_weight(
         )
 
         numerator = np.sum(
-            pairwise_rr * np.exp(-pairwise_rr / self.smin_alpha), axis=-1
+            np.where(
+                nlist_larger != -1,
+                pairwise_rr * np.exp(-pairwise_rr / self.smin_alpha),
+                np.zeros_like(nlist_larger),
+            ),
+            axis=-1,
         )  # masked nnei will be zero, no need to handle
         denominator = np.sum(
             np.where(
@@ -436,5 +441,7 @@ def _compute_weight(
             smooth = -6 * u**5 + 15 * u**4 - 10 * u**3 + 1
         coef[mid_mask] = smooth[mid_mask]
         coef[right_mask] = 0
+        # to handle masked atoms
+        coef = np.where(sigma != 0, coef, np.zeros_like(coef))
         self.zbl_weight = coef
         return [1 - np.expand_dims(coef, -1), np.expand_dims(coef, -1)]

deepmd/dpmodel/descriptor/se_r.py

@@ -193,7 +193,7 @@ def get_dim_out(self):
         return self.neuron[-1]
 
     def get_dim_emb(self):
-        """Returns the embedding (g2) dimension of this descriptor."""
+        """Returns the output dimension."""
         raise NotImplementedError
 
     def get_rcut(self):

deepmd/dpmodel/fitting/dipole_fitting.py

@@ -175,6 +175,7 @@ def output_def(self):
                     reduciable=True,
                     r_differentiable=self.r_differentiable,
                     c_differentiable=self.c_differentiable,
+                    rot_invariant=False,
                 ),
             ]
         )

deepmd/dpmodel/fitting/polarizability_fitting.py

@@ -218,6 +218,7 @@ def output_def(self):
                     reduciable=True,
                     r_differentiable=False,
                     c_differentiable=False,
+                    rot_invariant=False,
                 ),
             ]
         )

deepmd/dpmodel/model/__init__.py

@@ -15,6 +15,9 @@
 from .dp_model import (
     DPModelCommon,
 )
+from .ener_model import (
+    EnergyModel,
+)
 from .make_model import (
     make_model,
 )
@@ -23,6 +26,7 @@
 )
 
 __all__ = [
+    "EnergyModel",
     "DPModelCommon",
     "SpinModel",
     "make_model",

deepmd/dpmodel/output_def.py

@@ -186,6 +186,8 @@ class OutputVariableDef:
           If hessian is requred
     magnetic : bool
           If the derivatives of variable have magnetic parts.
+    rot_invariant : bool
+          If the variable is rotationally invariant.
     """
 
     def __init__(
@@ -199,6 +201,7 @@ def __init__(
         category: int = OutputVariableCategory.OUT.value,
         r_hessian: bool = False,
         magnetic: bool = False,
+        rot_invariant: bool = True,
     ):
         self.name = name
         self.shape = list(shape)
@@ -218,6 +221,7 @@ def __init__(
         self.category = category
         self.r_hessian = r_hessian
         self.magnetic = magnetic
+        self.rot_invariant = rot_invariant
         if self.r_hessian:
             if not self.reduciable:
                 raise ValueError("only reduciable variable can calculate hessian")
@@ -228,6 +232,11 @@ def __init__(
     def size(self):
         return self.output_size
 
+    def squeeze(self, dim):
+        # squeeze the shape on given dimension
+        if -len(self.shape) <= dim < len(self.shape) != 1 and self.shape[dim] == 1:
+            self.shape.pop(dim)
+
 
 class FittingOutputDef:
     """Defines the shapes and other properties of the fitting network outputs.
@@ -306,7 +315,6 @@ def __getitem__(
 
     def get_data(
         self,
-        key: str,
     ) -> Dict[str, OutputVariableDef]:
         return self.var_defs
 
@@ -417,6 +425,7 @@ def do_reduce(
                 c_differentiable=False,
                 atomic=False,
                 category=apply_operation(vv, OutputVariableOperation.REDU),
+                rot_invariant=vv.rot_invariant,
             )
     return def_redu
 
@@ -465,6 +474,7 @@ def do_derivative(
                 c_differentiable=False,
                 atomic=True,
                 category=apply_operation(vv, OutputVariableOperation.DERV_R),
+                rot_invariant=False,
             )
             if vv.magnetic:
                 def_derv_r[rkrm] = OutputVariableDef(
@@ -478,6 +488,7 @@ def do_derivative(
                     atomic=True,
                     category=apply_operation(vv, OutputVariableOperation.DERV_R),
                     magnetic=True,
+                    rot_invariant=False,
                 )
 
         if vv.c_differentiable:
@@ -490,6 +501,7 @@ def do_derivative(
                 c_differentiable=False,
                 atomic=True,
                 category=apply_operation(vv, OutputVariableOperation.DERV_C),
+                rot_invariant=False,
             )
             if vv.magnetic:
                 def_derv_r[rkcm] = OutputVariableDef(
@@ -501,5 +513,6 @@ def do_derivative(
                     atomic=True,
                     category=apply_operation(vv, OutputVariableOperation.DERV_C),
                     magnetic=True,
+                    rot_invariant=False,
                 )
     return def_derv_r, def_derv_c

deepmd/pt/model/atomic_model/linear_atomic_model.py

@@ -224,12 +224,12 @@ def forward_atomic(
         ener_list = []
 
         for i, model in enumerate(self.models):
-            mapping = self.mapping_list[i]
+            type_map_model = self.mapping_list[i]
             # apply bias to each individual model
             ener_list.append(
                 model.forward_common_atomic(
                     extended_coord,
-                    mapping[extended_atype],
+                    type_map_model[extended_atype],
                     nlists_[i],
                     mapping,
                     fparam,

deepmd/pt/model/descriptor/se_r.py

@@ -305,6 +305,7 @@ def forward(
         atype_ext: torch.Tensor,
         nlist: torch.Tensor,
         mapping: Optional[torch.Tensor] = None,
+        comm_dict: Optional[Dict[str, torch.Tensor]] = None,
     ):
         """Compute the descriptor.
 
@@ -318,6 +319,8 @@ def forward(
             The neighbor list. shape: nf x nloc x nnei
         mapping
             The index mapping, not required by this descriptor.
+        comm_dict
+            The data needed for communication for parallel inference.
 
         Returns
         -------

deepmd/pt/model/model/__init__.py

@@ -210,6 +210,9 @@ def get_model(model_params):
     "get_model",
     "DPModelCommon",
     "EnergyModel",
+    "DipoleModel",
+    "PolarModel",
+    "DOSModel",
     "FrozenModel",
     "SpinModel",
     "SpinEnergyModel",

deepmd/pt/model/model/dipole_model.py

@@ -1,4 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from copy import (
+    deepcopy,
+)
 from typing import (
     Dict,
     Optional,
@@ -35,6 +38,24 @@ def __init__(
         DPModelCommon.__init__(self)
         DPDOSModel_.__init__(self, *args, **kwargs)
 
+    def translated_output_def(self):
+        out_def_data = self.model_output_def().get_data()
+        output_def = {
+            "dipole": deepcopy(out_def_data["dipole"]),
+            "global_dipole": deepcopy(out_def_data["dipole_redu"]),
+        }
+        if self.do_grad_r("dipole"):
+            output_def["force"] = deepcopy(out_def_data["dipole_derv_r"])
+            output_def["force"].squeeze(-2)
+        if self.do_grad_c("dipole"):
+            output_def["virial"] = deepcopy(out_def_data["dipole_derv_c_redu"])
+            output_def["virial"].squeeze(-2)
+            output_def["atom_virial"] = deepcopy(out_def_data["dipole_derv_c"])
+            output_def["atom_virial"].squeeze(-3)
+        if "mask" in out_def_data:
+            output_def["mask"] = deepcopy(out_def_data["mask"])
+        return output_def
+
     def forward(
         self,
         coord,
@@ -96,13 +117,13 @@ def forward_lower(
             model_predict["dipole"] = model_ret["dipole"]
             model_predict["global_dipole"] = model_ret["dipole_redu"]
             if self.do_grad_r("dipole"):
-                model_predict["force"] = model_ret["dipole_derv_r"].squeeze(-2)
+                model_predict["extended_force"] = model_ret["dipole_derv_r"].squeeze(-2)
             if self.do_grad_c("dipole"):
                 model_predict["virial"] = model_ret["dipole_derv_c_redu"].squeeze(-2)
                 if do_atomic_virial:
-                    model_predict["atom_virial"] = model_ret["dipole_derv_c"].squeeze(
-                        -3
-                    )
+                    model_predict["extended_virial"] = model_ret[
+                        "dipole_derv_c"
+                    ].squeeze(-3)
         else:
             model_predict = model_ret
         return model_predict

deepmd/pt/model/model/dos_model.py

@@ -1,4 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from copy import (
+    deepcopy,
+)
 from typing import (
     Dict,
     Optional,
@@ -35,6 +38,16 @@ def __init__(
         DPModelCommon.__init__(self)
         DPDOSModel_.__init__(self, *args, **kwargs)
 
+    def translated_output_def(self):
+        out_def_data = self.model_output_def().get_data()
+        output_def = {
+            "atom_dos": deepcopy(out_def_data["dos"]),
+            "dos": deepcopy(out_def_data["dos_redu"]),
+        }
+        if "mask" in out_def_data:
+            output_def["mask"] = deepcopy(out_def_data["mask"])
+        return output_def
+
     def forward(
         self,
         coord,

deepmd/pt/model/model/dp_zbl_model.py

@@ -1,4 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from copy import (
+    deepcopy,
+)
 from typing import (
     Dict,
     List,
@@ -39,6 +42,24 @@ def __init__(
     ):
         super().__init__(*args, **kwargs)
 
+    def translated_output_def(self):
+        out_def_data = self.model_output_def().get_data()
+        output_def = {
+            "atom_energy": deepcopy(out_def_data["energy"]),
+            "energy": deepcopy(out_def_data["energy_redu"]),
+        }
+        if self.do_grad_r("energy"):
+            output_def["force"] = deepcopy(out_def_data["energy_derv_r"])
+            output_def["force"].squeeze(-2)
+        if self.do_grad_c("energy"):
+            output_def["virial"] = deepcopy(out_def_data["energy_derv_c_redu"])
+            output_def["virial"].squeeze(-2)
+            output_def["atom_virial"] = deepcopy(out_def_data["energy_derv_c"])
+            output_def["atom_virial"].squeeze(-3)
+        if "mask" in out_def_data:
+            output_def["mask"] = deepcopy(out_def_data["mask"])
+        return output_def
+
     def forward(
         self,
         coord,

deepmd/pt/model/model/ener_model.py

@@ -1,4 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from copy import (
+    deepcopy,
+)
 from typing import (
     Dict,
     Optional,
@@ -35,6 +38,24 @@ def __init__(
         DPModelCommon.__init__(self)
         DPEnergyModel_.__init__(self, *args, **kwargs)
 
+    def translated_output_def(self):
+        out_def_data = self.model_output_def().get_data()
+        output_def = {
+            "atom_energy": deepcopy(out_def_data["energy"]),
+            "energy": deepcopy(out_def_data["energy_redu"]),
+        }
+        if self.do_grad_r("energy"):
+            output_def["force"] = deepcopy(out_def_data["energy_derv_r"])
+            output_def["force"].squeeze(-2)
+        if self.do_grad_c("energy"):
+            output_def["virial"] = deepcopy(out_def_data["energy_derv_c_redu"])
+            output_def["virial"].squeeze(-2)
+            output_def["atom_virial"] = deepcopy(out_def_data["energy_derv_c"])
+            output_def["atom_virial"].squeeze(-3)
+        if "mask" in out_def_data:
+            output_def["mask"] = deepcopy(out_def_data["mask"])
+        return output_def
+
     def forward(
         self,
         coord,

deepmd/pt/model/model/polar_model.py

@@ -1,4 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from copy import (
+    deepcopy,
+)
 from typing import (
     Dict,
     Optional,
@@ -35,6 +38,16 @@ def __init__(
         DPModelCommon.__init__(self)
         DPDOSModel_.__init__(self, *args, **kwargs)
 
+    def translated_output_def(self):
+        out_def_data = self.model_output_def().get_data()
+        output_def = {
+            "polar": deepcopy(out_def_data["polarizability"]),
+            "global_polar": deepcopy(out_def_data["polarizability_redu"]),
+        }
+        if "mask" in out_def_data:
+            output_def["mask"] = deepcopy(out_def_data["mask"])
+        return output_def
+
     def forward(
         self,
         coord,

deepmd/pt/model/task/dipole.py

@@ -145,6 +145,7 @@ def output_def(self) -> FittingOutputDef:
                     reduciable=True,
                     r_differentiable=self.r_differentiable,
                     c_differentiable=self.c_differentiable,
+                    rot_invariant=False,
                 ),
             ]
         )

deepmd/pt/model/task/polarizability.py

@@ -219,6 +219,7 @@ def output_def(self) -> FittingOutputDef:
                     reduciable=True,
                     r_differentiable=False,
                     c_differentiable=False,
+                    rot_invariant=False,
                 ),
             ]
         )