Chore: migrate and refactor polar and dos bias

deepmd/dpmodel/fitting/polarizability_fitting.py

@@ -206,7 +206,7 @@
         return FittingOutputDef(
             [
                 OutputVariableDef(
-                    self.var_name,
+                    "polarizability",
                     [3, 3],
                     reduciable=True,
                     r_differentiable=False,
@@ -280,4 +280,4 @@
             # (nframes, nloc, 3, 3)
             bias = np.expand_dims(bias, axis=-1) * eye
             out = out + bias
-        return {self.var_name: out}
+        return {"polarizability": out}

deepmd/entrypoints/test.py

@@ -897,7 +897,7 @@
         polar = polar.reshape((polar.shape[0], -1, 9))[:, sel_mask, :].reshape(
             (polar.shape[0], -1)
         )
-        rmse_f = rmse(polar - test_data["atomic_polarizability"][:numb_test])
+        rmse_f = rmse(polar - test_data["atom_polarizability"][:numb_test])
 
     log.info(f"# number of test data : {numb_test:d} ")
     log.info(f"Polarizability  RMSE       : {rmse_f:e}")
@@ -926,7 +926,7 @@
             pe = np.concatenate(
                 (
                     np.reshape(
-                        test_data["atomic_polarizability"][:numb_test],
+                        test_data["atom_polarizability"][:numb_test],
                         [-1, 9 * sel_natoms],
                     ),
                     np.reshape(polar, [-1, 9 * sel_natoms]),
@@ -1037,7 +1037,7 @@
         dipole = dipole.reshape((dipole.shape[0], -1, 3))[:, sel_mask, :].reshape(
             (dipole.shape[0], -1)
         )
-        rmse_f = rmse(dipole - test_data["atomic_dipole"][:numb_test])
+        rmse_f = rmse(dipole - test_data["atom_dipole"][:numb_test])
 
     log.info(f"# number of test data : {numb_test:d}")
     log.info(f"Dipole  RMSE       : {rmse_f:e}")
@@ -1061,7 +1061,7 @@
             pe = np.concatenate(
                 (
                     np.reshape(
-                        test_data["atomic_dipole"][:numb_test], [-1, 3 * sel_natoms]
+                        test_data["atom_dipole"][:numb_test], [-1, 3 * sel_natoms]
                     ),
                     np.reshape(dipole, [-1, 3 * sel_natoms]),
                 ),

deepmd/pt/loss/tensor.py

@@ -93,14 +93,14 @@
         if (
             self.has_local_weight
             and self.tensor_name in model_pred
-            and "atomic_" + self.label_name in label
+            and "atom_" + self.label_name in label
         ):
-            find_local = label.get("find_" + "atomic_" + self.label_name, 0.0)
+            find_local = label.get("find_" + "atom_" + self.label_name, 0.0)
             local_weight = self.local_weight * find_local
             local_tensor_pred = model_pred[self.tensor_name].reshape(
                 [-1, natoms, self.tensor_size]
             )
-            local_tensor_label = label["atomic_" + self.label_name].reshape(
+            local_tensor_label = label["atom_" + self.label_name].reshape(
                 [-1, natoms, self.tensor_size]
             )
             diff = (local_tensor_pred - local_tensor_label).reshape(

deepmd/pt/model/model/polar_model.py

@@ -54,8 +54,8 @@
         )
         if self.get_fitting_net() is not None:
             model_predict = {}
-            model_predict["polar"] = model_ret["polar"]
-            model_predict["global_polar"] = model_ret["polar_redu"]
+            model_predict["polar"] = model_ret["polarizability"]
+            model_predict["global_polar"] = model_ret["polarizability_redu"]
             if "mask" in model_ret:
                 model_predict["mask"] = model_ret["mask"]
         else:
@@ -85,8 +85,8 @@
         )
         if self.get_fitting_net() is not None:
             model_predict = {}
-            model_predict["polar"] = model_ret["polar"]
-            model_predict["global_polar"] = model_ret["polar_redu"]
+            model_predict["polar"] = model_ret["polarizability"]
+            model_predict["global_polar"] = model_ret["polarizability_redu"]
         else:
             model_predict = model_ret
         return model_predict

deepmd/pt/model/task/dos.py

@@ -2,13 +2,11 @@
 import copy
 import logging
 from typing import (
-    Callable,
     List,
     Optional,
     Union,
 )
 
-import numpy as np
 import torch
 
 from deepmd.dpmodel import (
@@ -30,13 +28,6 @@
 from deepmd.pt.utils.utils import (
     to_numpy_array,
 )
-from deepmd.utils.out_stat import (
-    compute_stats_from_atomic,
-    compute_stats_from_redu,
-)
-from deepmd.utils.path import (
-    DPPath,
-)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -105,63 +96,6 @@ def output_def(self) -> FittingOutputDef:
             ]
         )
 
-    def compute_output_stats(
-        self,
-        merged: Union[Callable[[], List[dict]], List[dict]],
-        stat_file_path: Optional[DPPath] = None,
-    ) -> None:
-        """
-        Compute the output statistics (e.g. dos bias) for the fitting net from packed data.
-
-        Parameters
-        ----------
-        merged : Union[Callable[[], List[dict]], List[dict]]
-            - List[dict]: A list of data samples from various data systems.
-                Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor`
-                originating from the `i`-th data system.
-            - Callable[[], List[dict]]: A lazy function that returns data samples in the above format
-                only when needed. Since the sampling process can be slow and memory-intensive,
-                the lazy function helps by only sampling once.
-        stat_file_path : Optional[DPPath]
-            The path to the stat file.
-
-        """
-        if stat_file_path is not None:
-            stat_file_path = stat_file_path / "bias_dos"
-        if stat_file_path is not None and stat_file_path.is_file():
-            bias_dos = stat_file_path.load_numpy()
-        else:
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-            for sys in range(len(sampled)):
-                nframs = sampled[sys]["atype"].shape[0]
-
-                if "atom_dos" in sampled[sys]:
-                    bias_dos = compute_stats_from_atomic(
-                        sampled[sys]["atom_dos"].numpy(force=True),
-                        sampled[sys]["atype"].numpy(force=True),
-                    )[0]
-                else:
-                    sys_type_count = np.zeros(
-                        (nframs, self.ntypes), dtype=env.GLOBAL_NP_FLOAT_PRECISION
-                    )
-                    for itype in range(self.ntypes):
-                        type_mask = sampled[sys]["atype"] == itype
-                        sys_type_count[:, itype] = type_mask.sum(dim=1).numpy(
-                            force=True
-                        )
-                    sys_bias_redu = sampled[sys]["dos"].numpy(force=True)
-
-                    bias_dos = compute_stats_from_redu(
-                        sys_bias_redu, sys_type_count, rcond=self.rcond
-                    )[0]
-                if stat_file_path is not None:
-                    stat_file_path.save_numpy(bias_dos)
-        self.bias_dos = torch.tensor(bias_dos, device=env.DEVICE)
-
     @classmethod
     def deserialize(cls, data: dict) -> "DOSFittingNet":
         data = copy.deepcopy(data)

deepmd/pt/model/task/invar_fitting.py

@@ -2,10 +2,8 @@
 import copy
 import logging
 from typing import (
-    Callable,
     List,
     Optional,
-    Union,
 )
 
 import torch
@@ -24,12 +22,6 @@
 from deepmd.pt.utils.env import (
     DEFAULT_PRECISION,
 )
-from deepmd.pt.utils.stat import (
-    compute_output_stats,
-)
-from deepmd.utils.path import (
-    DPPath,
-)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -146,40 +138,6 @@ def deserialize(cls, data: dict) -> "GeneralFitting":
         check_version_compatibility(data.pop("@version", 1), 1, 1)
         return super().deserialize(data)
 
-    def compute_output_stats(
-        self,
-        merged: Union[Callable[[], List[dict]], List[dict]],
-        stat_file_path: Optional[DPPath] = None,
-    ):
-        """
-        Compute the output statistics (e.g. energy bias) for the fitting net from packed data.
-
-        Parameters
-        ----------
-        merged : Union[Callable[[], List[dict]], List[dict]]
-            - List[dict]: A list of data samples from various data systems.
-                Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor`
-                originating from the `i`-th data system.
-            - Callable[[], List[dict]]: A lazy function that returns data samples in the above format
-                only when needed. Since the sampling process can be slow and memory-intensive,
-                the lazy function helps by only sampling once.
-        stat_file_path : Optional[DPPath]
-            The path to the stat file.
-
-        """
-        # [0] to get the mean (bias)
-        bias_atom_e = compute_output_stats(
-            merged,
-            self.ntypes,
-            keys=[self.var_name],
-            stat_file_path=stat_file_path,
-            rcond=self.rcond,
-            preset_bias={self.var_name: self.atom_ener}
-            if self.atom_ener is not None
-            else None,
-        )[0][self.var_name]
-        self.bias_atom_e.copy_(bias_atom_e.view([self.ntypes, self.dim_out]))
-
     def output_def(self) -> FittingOutputDef:
         return FittingOutputDef(
             [

deepmd/pt/model/task/polarizability.py

@@ -2,13 +2,11 @@
 import copy
 import logging
 from typing import (
-    Callable,
     List,
     Optional,
     Union,
 )
 
-import numpy as np
 import torch
 
 from deepmd.dpmodel import (
@@ -27,13 +25,6 @@
 from deepmd.pt.utils.utils import (
     to_numpy_array,
 )
-from deepmd.utils.out_stat import (
-    compute_stats_from_atomic,
-    compute_stats_from_redu,
-)
-from deepmd.utils.path import (
-    DPPath,
-)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -185,7 +176,7 @@
         return FittingOutputDef(
             [
                 OutputVariableDef(
-                    self.var_name,
+                    "polarizability",
                     [3, 3],
                     reduciable=True,
                     r_differentiable=False,
@@ -194,82 +185,6 @@
             ]
         )
 
-    def compute_output_stats(
-        self,
-        merged: Union[Callable[[], List[dict]], List[dict]],
-        stat_file_path: Optional[DPPath] = None,
-    ) -> None:
-        """
-        Compute the output statistics (e.g. energy bias) for the fitting net from packed data.
-
-        Parameters
-        ----------
-        merged : Union[Callable[[], List[dict]], List[dict]]
-            - List[dict]: A list of data samples from various data systems.
-                Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor`
-                originating from the `i`-th data system.
-            - Callable[[], List[dict]]: A lazy function that returns data samples in the above format
-                only when needed. Since the sampling process can be slow and memory-intensive,
-                the lazy function helps by only sampling once.
-        stat_file_path : Optional[DPPath]
-            The path to the stat file.
-
-        """
-        if self.shift_diag:
-            if stat_file_path is not None:
-                stat_file_path = stat_file_path / "constant_matrix"
-            if stat_file_path is not None and stat_file_path.is_file():
-                constant_matrix = stat_file_path.load_numpy()
-            else:
-                if callable(merged):
-                    # only get data for once
-                    sampled = merged()
-                else:
-                    sampled = merged
-
-                sys_constant_matrix = []
-                for sys in range(len(sampled)):
-                    nframs = sampled[sys]["atype"].shape[0]
-
-                    if sampled[sys]["find_atomic_polarizability"] > 0.0:
-                        sys_atom_polar = compute_stats_from_atomic(
-                            sampled[sys]["atomic_polarizability"].numpy(force=True),
-                            sampled[sys]["atype"].numpy(force=True),
-                        )[0]
-                    else:
-                        if not sampled[sys]["find_polarizability"] > 0.0:
-                            continue
-                        sys_type_count = np.zeros(
-                            (nframs, self.ntypes), dtype=env.GLOBAL_NP_FLOAT_PRECISION
-                        )
-                        for itype in range(self.ntypes):
-                            type_mask = sampled[sys]["atype"] == itype
-                            sys_type_count[:, itype] = type_mask.sum(dim=1).numpy(
-                                force=True
-                            )
-
-                        sys_bias_redu = sampled[sys]["polarizability"].numpy(force=True)
-
-                        sys_atom_polar = compute_stats_from_redu(
-                            sys_bias_redu, sys_type_count, rcond=self.rcond
-                        )[0]
-                    cur_constant_matrix = np.zeros(
-                        self.ntypes, dtype=env.GLOBAL_NP_FLOAT_PRECISION
-                    )
-
-                    for itype in range(self.ntypes):
-                        cur_constant_matrix[itype] = np.mean(
-                            np.diagonal(sys_atom_polar[itype].reshape(3, 3))
-                        )
-                    sys_constant_matrix.append(cur_constant_matrix)
-                constant_matrix = np.stack(sys_constant_matrix).mean(axis=0)
-
-                # handle nan values.
-                constant_matrix = np.nan_to_num(constant_matrix)
-            if stat_file_path is not None:
-                stat_file_path.save_numpy(constant_matrix)
-            self.constant_matrix = torch.tensor(constant_matrix, device=env.DEVICE)
-
     def forward(
         self,
         descriptor: torch.Tensor,
@@ -314,7 +229,7 @@
             bias = bias.unsqueeze(-1) * eye
             out = out + bias
 
-        return {self.var_name: out.to(env.GLOBAL_PT_FLOAT_PRECISION)}
+        return {"polarizability": out.to(env.GLOBAL_PT_FLOAT_PRECISION)}
 
     # make jit happy with torch 2.0.0
     exclude_types: List[int]

deepmd/pt/train/training.py

@@ -303,9 +303,10 @@
                     tensor_name
                 ].output_size
                 label_name = tensor_name
-                if label_name == "polar":
-                    label_name = "polarizability"
+                if label_name == "polarizability":
+                    label_name = "polar"
                 loss_params["label_name"] = label_name
+                loss_params["tensor_name"] = label_name
                 return TensorLoss(**loss_params)
             else:
                 raise NotImplementedError