Fix: LinearModel Stat (#3575)

Refactor stat related methods in LinearAtomicModel.
This should resolve #3585 #3586 #3528

---------

Signed-off-by: Anyang Peng <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -79,27 +79,6 @@ def mixed_types(self) -> bool:
         """
         return self.descriptor.mixed_types()
 
-    def set_out_bias(self, out_bias: np.ndarray, add=False) -> None:
-        """
-        Modify the output bias for the atomic model.
-
-        Parameters
-        ----------
-        out_bias : np.ndarray
-            The new bias to be applied.
-        add : bool, optional
-            Whether to add the new bias to the existing one.
-            If False, the output bias will be directly replaced by the new bias.
-            If True, the new bias will be added to the existing one.
-        """
-        self.fitting["bias_atom_e"] = (
-            out_bias + self.fitting["bias_atom_e"] if add else out_bias
-        )
-
-    def get_out_bias(self) -> np.ndarray:
-        """Return the output bias of the atomic model."""
-        return self.fitting["bias_atom_e"]
-
     def forward_atomic(
         self,
         extended_coord: np.ndarray,

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -162,7 +162,6 @@ def forward_atomic(
             )
         ]
         ener_list = []
-
         for i, model in enumerate(self.models):
             mapping = self.mapping_list[i]
             ener_list.append(
@@ -176,13 +175,10 @@ def forward_atomic(
                 )["energy"]
             )
         self.weights = self._compute_weight(extended_coord, extended_atype, nlists_)
-        self.atomic_bias = None
-        if self.atomic_bias is not None:
-            raise NotImplementedError("Need to add bias in a future PR.")
-        else:
-            fit_ret = {
-                "energy": np.sum(np.stack(ener_list) * np.stack(self.weights), axis=0),
-            }  # (nframes, nloc, 1)
+
+        fit_ret = {
+            "energy": np.sum(np.stack(ener_list) * np.stack(self.weights), axis=0),
+        }  # (nframes, nloc, 1)
         return fit_ret
 
     @staticmethod
@@ -252,7 +248,8 @@ def _compute_weight(
     ) -> List[np.ndarray]:
         """This should be a list of user defined weights that matches the number of models to be combined."""
         nmodels = len(self.models)
-        return [np.ones(1) / nmodels for _ in range(nmodels)]
+        nframes, nloc, _ = nlists_[0].shape
+        return [np.ones((nframes, nloc, 1)) / nmodels for _ in range(nmodels)]
 
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""
@@ -275,27 +272,6 @@ def get_sel_type(self) -> List[int]:
         # join all the selected types
         return list(set().union(*[model.get_sel_type() for model in self.models]))
 
-    def set_out_bias(self, out_bias: np.ndarray, add=False) -> None:
-        """
-        Modify the output bias for all the models in the linear atomic model.
-
-        Parameters
-        ----------
-        out_bias : torch.Tensor
-            The new bias to be applied.
-        add : bool, optional
-            Whether to add the new bias to the existing one.
-            If False, the output bias will be directly replaced by the new bias.
-            If True, the new bias will be added to the existing one.
-        """
-        for model in self.models:
-            model.set_out_bias(out_bias, add=add)
-
-    def get_out_bias(self) -> np.ndarray:
-        """Return the weighted output bias of the linear atomic model."""
-        # TODO add get_out_bias for linear atomic model
-        raise NotImplementedError
-
     def is_aparam_nall(self) -> bool:
         """Check whether the shape of atomic parameters is (nframes, nall, ndim).
 

deepmd/dpmodel/atomic_model/make_base_atomic_model.py

@@ -95,25 +95,6 @@ def get_sel_type(self) -> List[int]:
             If returning an empty list, all atom types are selected.
             """
 
-        @abstractmethod
-        def set_out_bias(self, out_bias: t_tensor, add=False) -> None:
-            """
-            Modify the output bias for the atomic model.
-
-            Parameters
-            ----------
-            out_bias : t_tensor
-                The new bias to be applied.
-            add : bool, optional
-                Whether to add the new bias to the existing one.
-                If False, the output bias will be directly replaced by the new bias.
-                If True, the new bias will be added to the existing one.
-            """
-
-        @abstractmethod
-        def get_out_bias(self) -> t_tensor:
-            """Return the output bias of the atomic model."""
-
         @abstractmethod
         def is_aparam_nall(self) -> bool:
             """Check whether the shape of atomic parameters is (nframes, nall, ndim).

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -130,25 +130,6 @@ def mixed_types(self) -> bool:
         # to match DPA1 and DPA2.
         return True
 
-    def set_out_bias(self, out_bias: np.ndarray, add=False) -> None:
-        """
-        Modify the output bias for the atomic model.
-
-        Parameters
-        ----------
-        out_bias : torch.Tensor
-            The new bias to be applied.
-        add : bool, optional
-            Whether to add the new bias to the existing one.
-            If False, the output bias will be directly replaced by the new bias.
-            If True, the new bias will be added to the existing one.
-        """
-        self.bias_atom_e = out_bias + self.bias_atom_e if add else out_bias
-
-    def get_out_bias(self) -> np.ndarray:
-        """Return the output bias of the atomic model."""
-        return self.bias_atom_e
-
     def serialize(self) -> dict:
         dd = BaseAtomicModel.serialize(self)
         dd.update(

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -366,7 +366,6 @@ def change_out_bias(
                 rcond=self.rcond,
                 preset_bias=self.preset_out_bias,
             )
-            # self.set_out_bias(delta_bias, add=True)
             self._store_out_stat(delta_bias, out_std, add=True)
         elif bias_adjust_mode == "set-by-statistic":
             bias_out, std_out = compute_output_stats(
@@ -377,7 +376,6 @@ def change_out_bias(
                 rcond=self.rcond,
                 preset_bias=self.preset_out_bias,
             )
-            # self.set_out_bias(bias_out)
             self._store_out_stat(bias_out, std_out)
         else:
             raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode)

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -176,6 +176,9 @@ def forward_atomic(
         )
         return fit_ret
 
+    def get_out_bias(self) -> torch.Tensor:
+        return self.out_bias
+
     def compute_or_load_stat(
         self,
         sampled_func,
@@ -217,27 +220,6 @@ def wrapped_sampler():
         self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path)
         self.compute_or_load_out_stat(wrapped_sampler, stat_file_path)
 
-    def set_out_bias(self, out_bias: torch.Tensor, add=False) -> None:
-        """
-        Modify the output bias for the atomic model.
-
-        Parameters
-        ----------
-        out_bias : torch.Tensor
-            The new bias to be applied.
-        add : bool, optional
-            Whether to add the new bias to the existing one.
-            If False, the output bias will be directly replaced by the new bias.
-            If True, the new bias will be added to the existing one.
-        """
-        self.fitting_net["bias_atom_e"] = (
-            out_bias + self.fitting_net["bias_atom_e"] if add else out_bias
-        )
-
-    def get_out_bias(self) -> torch.Tensor:
-        """Return the output bias of the atomic model."""
-        return self.fitting_net["bias_atom_e"]
-
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""
         return self.fitting_net.get_dim_fparam()