Feat: Add polar stat constant matrix calculation to PT

deepmd/pt/model/task/polarizability.py

@@ -7,6 +7,7 @@
     Union,
 )
 
+import numpy as np
 import torch
 
 from deepmd.dpmodel import (
@@ -25,6 +26,10 @@
 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,
 )
@@ -114,6 +119,7 @@
             self.scale, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE
         ).view(ntypes, 1)
         self.shift_diag = shift_diag
+        self.constant_matrix = torch.zeros(ntypes)
         super().__init__(
             var_name=kwargs.pop("var_name", "polar"),
             ntypes=ntypes,
@@ -184,7 +190,50 @@
             The path to the stat file.
 
         """
-        pass
+        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]["type"].shape[0]
+
+                    if sampled[sys]["find_atomic_polarizability"] > 0.0:
+                        sys_atom_polar = compute_stats_from_atomic(
+                            sampled[sys]["atomic_polarizability"], sampled[sys]["type"]
+                        )[0]
+                    else:
+                        if not sampled[sys]["find_polarizability"] > 0.0:
+                            continue
+                        sys_type_count = np.zeros((nframs, self.ntypes))
+                        for itype in range(self.ntypes):
+                            type_mask = sampled[sys]["type"] == itype
+                            sys_type_count[:, itype] = type_mask.sum(dim=1)
+
+                        sys_bias_redu = sampled[sys]["polarizability"]
+
+                        sys_atom_polar = compute_stats_from_redu(
+                            sys_bias_redu, sys_type_count
+                        )[0]
+                    cur_constant_matrix = np.zeros(self.ntypes)
+                    for itype in range(self.ntypes):
+                        cur_constant_matrix[itype] = torch.mean(
+                            torch.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)
+
+            self.constant_matrix = torch.tensor(constant_matrix, device=env.DEVICE)
+            if stat_file_path is not None:
+                stat_file_path.save_numpy(self.constant_matrix.detach().cpu().numpy())
 
     def forward(
         self,
@@ -218,5 +267,12 @@
             "bim,bmj->bij", gr.transpose(1, 2), out
         )  # (nframes * nloc, 3, 3)
         out = out.view(nframes, nloc, 3, 3)
+        if self.shift_diag:
+            out = (
+                out
+                + self.constant_matrix[atype]
+                * torch.eye(3, device=env.DEVICE)
+                * self.scale[atype]
+            )
 
         return {self.var_name: out.to(env.GLOBAL_PT_FLOAT_PRECISION)}

deepmd/tf/fit/polar.py

@@ -193,18 +193,16 @@ def compute_output_stats(self, all_stat):
                         index_lis = [
                             index
                             for index, w in enumerate(atom_has_polar)
-                            if atom_has_polar[index] == self.sel_type[itype]
+                            if w == self.sel_type[itype]
                         ]  # select index in this type
 
                         sys_matrix.append(np.zeros((1, len(self.sel_type))))
                         sys_matrix[-1][0, itype] = len(index_lis)
 
                         polar_bias.append(
                             np.sum(
-                                all_stat["atomic_polarizability"][ss].reshape((-1, 9))[
-                                    index_lis
-                                ],
-                                axis=0,
+                                all_stat["atomic_polarizability"][ss][:, index_lis, :],
+                                axis=(0, 1),
                             ).reshape((1, 9))
                         )
                 else:  # No atomic polar in this system, so it should have global polar