Feat: add DipoleModel and PolarModel

deepmd/dpmodel/fitting/dipole_fitting.py

@@ -68,9 +68,14 @@ class DipoleFitting(GeneralFitting):
     mixed_types
             If true, use a uniform fitting net for all atom types, otherwise use
             different fitting nets for different atom types.
-    exclude_types: List[int]
+    exclude_types
             Atomic contributions of the excluded atom types are set zero.
-
+    r_differentiable
+            If the variable is differentiated with respect to coordinates of atoms.
+            Only reduciable variable are differentiable.
+    c_differentiable
+            If the variable is differentiated with respect to the cell tensor (pbc case).
+            Only reduciable variable are differentiable.
     """
 
     def __init__(
@@ -94,6 +99,8 @@ def __init__(
         spin: Any = None,
         mixed_types: bool = False,
         exclude_types: List[int] = [],
+        r_differentiable: bool = True,
+        c_differentiable: bool = True,
         old_impl=False,
     ):
         # seed, uniform_seed are not included
@@ -109,6 +116,8 @@ def __init__(
             raise NotImplementedError("atom_ener is not implemented")
 
         self.embedding_width = embedding_width
+        self.r_differentiable = r_differentiable
+        self.c_differentiable = c_differentiable
         super().__init__(
             var_name=var_name,
             ntypes=ntypes,
@@ -139,6 +148,8 @@ def serialize(self) -> dict:
         data = super().serialize()
         data["embedding_width"] = self.embedding_width
         data["old_impl"] = self.old_impl
+        data["r_differentiable"] = self.r_differentiable
+        data["c_differentiable"] = self.c_differentiable
         return data
 
     def output_def(self):
@@ -148,8 +159,8 @@ def output_def(self):
                     self.var_name,
                     [3],
                     reduciable=True,
-                    r_differentiable=True,
-                    c_differentiable=True,
+                    r_differentiable=self.r_differentiable,
+                    c_differentiable=self.c_differentiable,
                 ),
             ]
         )

deepmd/dpmodel/fitting/polarizability_fitting.py

@@ -176,8 +176,8 @@ def output_def(self):
                     self.var_name,
                     [3, 3],
                     reduciable=True,
-                    r_differentiable=True,
-                    c_differentiable=True,
+                    r_differentiable=False,
+                    c_differentiable=False,
                 ),
             ]
         )

deepmd/infer/deep_tensor.py

@@ -105,6 +105,8 @@ def eval(
             **kwargs,
         )
         sel_natoms = self._get_sel_natoms(atom_types[0])
+        if sel_natoms == 0:
+            sel_natoms = atom_types.shape[-1]  # set to natoms
         if atomic:
             return results[self.output_tensor_name].reshape(nframes, sel_natoms, -1)
         else:
@@ -184,22 +186,24 @@ def eval_full(
             aparam=aparam,
             **kwargs,
         )
+
         sel_natoms = self._get_sel_natoms(atom_types[0])
+        if sel_natoms == 0:
+            sel_natoms = atom_types.shape[-1]  # set to natoms
         energy = results[f"{self.output_tensor_name}_redu"].reshape(nframes, -1)
         force = results[f"{self.output_tensor_name}_derv_r"].reshape(
             nframes, -1, natoms, 3
         )
         virial = results[f"{self.output_tensor_name}_derv_c_redu"].reshape(
             nframes, -1, 9
         )
-        atomic_energy = results[self.output_tensor_name].reshape(
-            nframes, sel_natoms, -1
-        )
-        atomic_virial = results[f"{self.output_tensor_name}_derv_c"].reshape(
-            nframes, -1, natoms, 9
-        )
-
         if atomic:
+            atomic_energy = results[self.output_tensor_name].reshape(
+                nframes, sel_natoms, -1
+            )
+            atomic_virial = results[f"{self.output_tensor_name}_derv_c"].reshape(
+                nframes, -1, natoms, 9
+            )
             return (
                 energy,
                 force,

deepmd/pt/model/model/dipole_model.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Dict,
+    Optional,
+)
+
+import torch
+
+from .dp_model import (
+    DPModel,
+)
+
+
+class DipoleModel(DPModel):
+    model_type = "dipole"
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+    def forward(
+        self,
+        coord,
+        atype,
+        box: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+        do_atomic_virial: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        model_ret = self.forward_common(
+            coord,
+            atype,
+            box,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+        )
+        if self.fitting_net is not None:
+            model_predict = {}
+            model_predict["dipole"] = model_ret["dipole"]
+            model_predict["global_dipole"] = model_ret["dipole_redu"]
+            if self.do_grad("dipole"):
+                model_predict["force"] = model_ret["dipole_derv_r"].squeeze(-2)
+                if do_atomic_virial:
+                    model_predict["atom_virial"] = model_ret["dipole_derv_c"].squeeze(
+                        -3
+                    )
+                model_predict["virial"] = model_ret["dipole_derv_c_redu"].squeeze(-2)
+        else:
+            model_predict = model_ret
+            model_predict["updated_coord"] += coord
+        return model_predict
+
+    def forward_lower(
+        self,
+        extended_coord,
+        extended_atype,
+        nlist,
+        mapping: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+        do_atomic_virial: bool = False,
+    ):
+        model_ret = self.forward_common_lower(
+            extended_coord,
+            extended_atype,
+            nlist,
+            mapping,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+        )
+        if self.fitting_net is not None:
+            model_predict = {}
+            model_predict["dipole"] = model_ret["dipole"]
+            model_predict["global_dipole"] = model_ret["dipole_redu"]
+            if self.do_grad("dipole"):
+                model_predict["force"] = model_ret["dipole_derv_r"].squeeze(-2)
+                if do_atomic_virial:
+                    model_predict["atom_virial"] = model_ret["dipole_derv_c"].squeeze(
+                        -3
+                    )
+                model_predict["virial"] = model_ret["dipole_derv_c_redu"].squeeze(-2)
+        else:
+            model_predict = model_ret
+        return model_predict

deepmd/pt/model/model/polar_model.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Dict,
+    Optional,
+)
+
+import torch
+
+from .dp_model import (
+    DPModel,
+)
+
+
+class PolarModel(DPModel):
+    model_type = "polar"
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+    def forward(
+        self,
+        coord,
+        atype,
+        box: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+        do_atomic_virial: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        model_ret = self.forward_common(
+            coord,
+            atype,
+            box,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+        )
+        if self.fitting_net is not None:
+            model_predict = {}
+            model_predict["polar"] = model_ret["polar"]
+            model_predict["global_polar"] = model_ret["polar_redu"]
+        else:
+            model_predict = model_ret
+            model_predict["updated_coord"] += coord
+        return model_predict
+
+    def forward_lower(
+        self,
+        extended_coord,
+        extended_atype,
+        nlist,
+        mapping: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+        do_atomic_virial: bool = False,
+    ):
+        model_ret = self.forward_common_lower(
+            extended_coord,
+            extended_atype,
+            nlist,
+            mapping,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+        )
+        if self.fitting_net is not None:
+            model_predict = {}
+            model_predict["polar"] = model_ret["polar"]
+            model_predict["global_polar"] = model_ret["polar_redu"]
+        else:
+            model_predict = model_ret
+        return model_predict

deepmd/pt/model/task/dipole.py

@@ -56,6 +56,12 @@ class DipoleFittingNet(GeneralFitting):
         The condition number for the regression of atomic energy.
     seed : int, optional
         Random seed.
+    r_differentiable
+        If the variable is differentiated with respect to coordinates of atoms.
+        Only reduciable variable are differentiable.
+    c_differentiable
+        If the variable is differentiated with respect to the cell tensor (pbc case).
+        Only reduciable variable are differentiable.
     """
 
     def __init__(
@@ -74,9 +80,13 @@ def __init__(
         rcond: Optional[float] = None,
         seed: Optional[int] = None,
         exclude_types: List[int] = [],
+        r_differentiable: bool = True,
+        c_differentiable: bool = True,
         **kwargs,
     ):
         self.embedding_width = embedding_width
+        self.r_differentiable = r_differentiable
+        self.c_differentiable = c_differentiable
         super().__init__(
             var_name=var_name,
             ntypes=ntypes,
@@ -103,6 +113,8 @@ def serialize(self) -> dict:
         data = super().serialize()
         data["embedding_width"] = self.embedding_width
         data["old_impl"] = self.old_impl
+        data["r_differentiable"] = self.r_differentiable
+        data["c_differentiable"] = self.c_differentiable
         return data
 
     def output_def(self) -> FittingOutputDef:
@@ -112,8 +124,8 @@ def output_def(self) -> FittingOutputDef:
                     self.var_name,
                     [3],
                     reduciable=True,
-                    r_differentiable=True,
-                    c_differentiable=True,
+                    r_differentiable=self.r_differentiable,
+                    c_differentiable=self.c_differentiable,
                 ),
             ]
         )

deepmd/pt/model/task/polarizability.py

@@ -144,8 +144,8 @@ def output_def(self) -> FittingOutputDef:
                     self.var_name,
                     [3, 3],
                     reduciable=True,
-                    r_differentiable=True,
-                    c_differentiable=True,
+                    r_differentiable=False,
+                    c_differentiable=False,
                 ),
             ]
         )