pt: Add support for dipole and polar training

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -78,7 +78,12 @@
         if self.atom_excl is not None:
             atom_mask = self.atom_excl.build_type_exclude_mask(atype)
             for kk in ret_dict.keys():
-                ret_dict[kk] = ret_dict[kk] * atom_mask[:, :, None]
+                out_shape = ret_dict[kk].shape
+                ret_dict[kk] = (
+                    ret_dict[kk].reshape([out_shape[0], out_shape[1], -1])
+                    * atom_mask[:, :, None]
+                ).reshape(out_shape)
+            ret_dict["mask"] = atom_mask
 
         return ret_dict
 

deepmd/dpmodel/model/make_model.py

@@ -159,6 +159,9 @@
                 aparam=ap,
                 do_atomic_virial=do_atomic_virial,
             )
+            nf, nloc = nlist.shape[:2]
+            if "mask" in model_predict_lower:
+                model_predict_lower["mask"] = model_predict_lower["mask"][:, :nloc]
             model_predict = communicate_extended_output(
                 model_predict_lower,
                 self.model_output_def(),

deepmd/dpmodel/model/transform_output.py

@@ -28,6 +28,8 @@
     """
     model_ret = dict(fit_ret.items())
     for kk, vv in fit_ret.items():
+        if kk in ["mask"]:
+            continue
         vdef = fit_output_def[kk]
         shap = vdef.shape
         atom_axis = -(len(shap) + 1)
@@ -59,6 +61,8 @@
 
     """
     new_ret = {}
+    if "mask" in model_ret:
+        new_ret["mask"] = model_ret["mask"]
     for kk in model_output_def.keys_outp():
         vv = model_ret[kk]
         vdef = model_output_def[kk]

deepmd/pt/loss/__init__.py

@@ -8,9 +8,13 @@
 from .loss import (
     TaskLoss,
 )
+from .tensor import (
+    TensorLoss,
+)
 
 __all__ = [
     "DenoiseLoss",
     "EnergyStdLoss",
+    "TensorLoss",
     "TaskLoss",
 ]

deepmd/pt/loss/tensor.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    List,
+)
+
+import torch
+
+from deepmd.pt.loss.loss import (
+    TaskLoss,
+)
+from deepmd.pt.utils import (
+    env,
+)
+from deepmd.utils.data import (
+    DataRequirementItem,
+)
+
+
+class TensorLoss(TaskLoss):
+    def __init__(
+        self,
+        tensor_name: str,
+        tensor_size: int,
+        label_name: str,
+        pref_atomic: float = 0.0,
+        pref: float = 0.0,
+        inference=False,
+        **kwargs,
+    ):
+        r"""Construct a loss for local and global tensors.
+
+        Parameters
+        ----------
+        tensor_name : str
+            The name of the tensor in the model predictions to compute the loss.
+        tensor_size : int
+            The size (dimension) of the tensor.
+        label_name : str
+            The name of the tensor in the labels to compute the loss.
+        pref_atomic : float
+            The prefactor of the weight of atomic loss. It should be larger than or equal to 0.
+        pref : float
+            The prefactor of the weight of global loss. It should be larger than or equal to 0.
+        inference : bool
+            If true, it will output all losses found in output, ignoring the pre-factors.
+        **kwargs
+            Other keyword arguments.
+        """
+        super().__init__()
+        self.tensor_name = tensor_name
+        self.tensor_size = tensor_size
+        self.label_name = label_name
+        self.local_weight = pref_atomic
+        self.global_weight = pref
+        self.inference = inference
+
+        assert (
+            self.local_weight >= 0.0 and self.global_weight >= 0.0
+        ), "Can not assign negative weight to `pref` and `pref_atomic`"
+        self.has_local_weight = self.local_weight > 0.0 or inference
+        self.has_global_weight = self.global_weight > 0.0 or inference
+        assert self.has_local_weight or self.has_global_weight, AssertionError(
+            "Can not assian zero weight both to `pref` and `pref_atomic`"
+        )
+
+    def forward(self, model_pred, label, natoms, learning_rate=0.0, mae=False):
+        """Return loss on local and global tensors.
+
+        Parameters
+        ----------
+        model_pred : dict[str, torch.Tensor]
+            Model predictions.
+        label : dict[str, torch.Tensor]
+            Labels.
+        natoms : int
+            The local atom number.
+
+        Returns
+        -------
+        loss: torch.Tensor
+            Loss for model to minimize.
+        more_loss: dict[str, torch.Tensor]
+            Other losses for display.
+        """
+        del learning_rate, mae
+        loss = torch.tensor(0.0, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)
+        more_loss = {}
+        if (
+            self.has_local_weight
+            and self.tensor_name in model_pred
+            and "atomic_" + self.label_name in label
+        ):
+            local_tensor_pred = model_pred[self.tensor_name].reshape(
+                [-1, natoms, self.tensor_size]
+            )
+            local_tensor_label = label["atomic_" + self.label_name].reshape(
+                [-1, natoms, self.tensor_size]
+            )
+            diff = (local_tensor_pred - local_tensor_label).reshape(
+                [-1, self.tensor_size]
+            )
+            if "mask" in model_pred:
+                diff = diff[model_pred["mask"].reshape([-1]).bool()]
+            l2_local_loss = torch.mean(torch.square(diff))
+            if not self.inference:
+                more_loss[f"l2_local_{self.tensor_name}_loss"] = l2_local_loss.detach()
+            loss += self.local_weight * l2_local_loss
+            rmse_local = l2_local_loss.sqrt()
+            more_loss[f"rmse_local_{self.tensor_name}"] = rmse_local.detach()
+        if (
+            self.has_global_weight
+            and "global_" + self.tensor_name in model_pred
+            and self.label_name in label
+        ):
+            global_tensor_pred = model_pred["global_" + self.tensor_name].reshape(
+                [-1, self.tensor_size]
+            )
+            global_tensor_label = label[self.label_name].reshape([-1, self.tensor_size])
+            diff = global_tensor_pred - global_tensor_label
+            if "mask" in model_pred:
+                atom_num = model_pred["mask"].sum(-1, keepdim=True)
+                l2_global_loss = torch.mean(
+                    torch.sum(torch.square(diff) * atom_num, dim=0) / atom_num.sum()
+                )
+                atom_num = torch.mean(atom_num.float())
+            else:
+                atom_num = natoms
+                l2_global_loss = torch.mean(torch.square(diff))
+            if not self.inference:
+                more_loss[
+                    f"l2_global_{self.tensor_name}_loss"
+                ] = l2_global_loss.detach()
+            loss += self.global_weight * l2_global_loss
+            rmse_global = l2_global_loss.sqrt() / atom_num
+            more_loss[f"rmse_global_{self.tensor_name}"] = rmse_global.detach()
+        return loss, more_loss
+
+    @property
+    def label_requirement(self) -> List[DataRequirementItem]:
+        """Return data label requirements needed for this loss calculation."""
+        label_requirement = []
+        if self.has_local_weight:
+            label_requirement.append(
+                DataRequirementItem(
+                    "atomic_" + self.label_name,
+                    ndof=self.tensor_size,
+                    atomic=True,
+                    must=False,
+                    high_prec=False,
+                )
+            )
+        if self.has_global_weight:
+            label_requirement.append(
+                DataRequirementItem(
+                    self.label_name,
+                    ndof=self.tensor_size,
+                    atomic=False,
+                    must=False,
+                    high_prec=False,
+                )
+            )
+        return label_requirement

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -89,7 +89,12 @@
         if self.atom_excl is not None:
             atom_mask = self.atom_excl(atype)
             for kk in ret_dict.keys():
-                ret_dict[kk] = ret_dict[kk] * atom_mask[:, :, None]
+                out_shape = ret_dict[kk].shape
+                ret_dict[kk] = (
+                    ret_dict[kk].reshape([out_shape[0], out_shape[1], -1])
+                    * atom_mask[:, :, None]
+                ).reshape(out_shape)
+            ret_dict["mask"] = atom_mask
 
         return ret_dict
 

deepmd/pt/model/model/__init__.py

@@ -94,7 +94,7 @@
     fitting_net["type"] = fitting_net.get("type", "ener")
     fitting_net["ntypes"] = descriptor.get_ntypes()
     fitting_net["mixed_types"] = descriptor.mixed_types()
-    fitting_net["embedding_width"] = descriptor.get_dim_out()
+    fitting_net["embedding_width"] = descriptor.get_dim_emb()
     fitting_net["dim_descrpt"] = descriptor.get_dim_out()
     grad_force = "direct" not in fitting_net["type"]
     if not grad_force:

deepmd/pt/model/model/dipole_model.py

@@ -50,6 +50,8 @@
                     model_predict["atom_virial"] = model_ret["dipole_derv_c"].squeeze(
                         -3
                     )
+            if "mask" in model_ret:
+                model_predict["mask"] = model_ret["mask"]
         else:
             model_predict = model_ret
             model_predict["updated_coord"] += coord

deepmd/pt/model/model/dp_zbl_model.py

@@ -63,6 +63,8 @@
                 model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3)
         else:
             model_predict["force"] = model_ret["dforce"]
+        if "mask" in model_ret:
+            model_predict["mask"] = model_ret["mask"]
         return model_predict
 
     @torch.jit.export

deepmd/pt/model/model/ener_model.py

@@ -52,6 +52,8 @@
                     )
             else:
                 model_predict["force"] = model_ret["dforce"]
+            if "mask" in model_ret:
+                model_predict["mask"] = model_ret["mask"]
         else:
             model_predict = model_ret
             model_predict["updated_coord"] += coord

deepmd/pt/model/model/make_model.py

@@ -154,6 +154,9 @@
                 fparam=fp,
                 aparam=ap,
             )
+            nf, nloc = nlist.shape[:2]
+            if "mask" in model_predict_lower:
+                model_predict_lower["mask"] = model_predict_lower["mask"][:, :nloc]
             model_predict = communicate_extended_output(
                 model_predict_lower,
                 self.model_output_def(),

deepmd/pt/model/model/polar_model.py

@@ -42,6 +42,8 @@
             model_predict = {}
             model_predict["polar"] = model_ret["polar"]
             model_predict["global_polar"] = model_ret["polar_redu"]
+            if "mask" in model_ret:
+                model_predict["mask"] = model_ret["mask"]
         else:
             model_predict = model_ret
             model_predict["updated_coord"] += coord

deepmd/pt/model/model/transform_output.py

@@ -154,6 +154,8 @@
     redu_prec = env.GLOBAL_PT_ENER_FLOAT_PRECISION
     model_ret = dict(fit_ret.items())
     for kk, vv in fit_ret.items():
+        if kk in ["mask"]:
+            continue
         vdef = fit_output_def[kk]
         shap = vdef.shape
         atom_axis = -(len(shap) + 1)
@@ -192,6 +194,8 @@
     """
     redu_prec = env.GLOBAL_PT_ENER_FLOAT_PRECISION
     new_ret = {}
+    if "mask" in model_ret:
+        new_ret["mask"] = model_ret["mask"]
     for kk in model_output_def.keys_outp():
         vv = model_ret[kk]
         vdef = model_output_def[kk]

deepmd/pt/model/task/dipole.py

@@ -157,7 +157,7 @@
             The path to the stat file.
 
         """
-        raise NotImplementedError
+        pass
 
     def forward(
         self,

deepmd/pt/model/task/polarizability.py

@@ -184,7 +184,7 @@
             The path to the stat file.
 
         """
-        raise NotImplementedError
+        pass
 
     def forward(
         self,