Fix: LinearModel Stat

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 @@
             )
         ]
         ener_list = []
-
         for i, model in enumerate(self.models):
             mapping = self.mapping_list[i]
             ener_list.append(
@@ -176,13 +175,10 @@
                 )["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
@@ -275,27 +271,6 @@
         # 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 @@
         )
         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 @@
         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()

deepmd/pt/model/atomic_model/linear_atomic_model.py

@@ -73,7 +73,6 @@
             self.mapping_list.append(self.remap_atype(tpmp, self.type_map))
         assert len(err_msg) == 0, "\n".join(err_msg)
 
-        self.atomic_bias = None
         self.mixed_types_list = [model.mixed_types() for model in self.models]
         self.rcuts = torch.tensor(
             self.get_model_rcuts(), dtype=torch.float64, device=env.DEVICE
@@ -92,6 +91,9 @@
         """
         return True
 
+    def get_out_bias(self) -> torch.Tensor:
+        return self.out_bias
+
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
         return max(self.get_model_rcuts())
@@ -188,36 +190,46 @@
 
         for i, model in enumerate(self.models):
             mapping = self.mapping_list[i]
+            raw_ret = model.forward_atomic(
+                extended_coord,
+                mapping[extended_atype],
+                nlists_[i],
+                mapping,
+                fparam,
+                aparam,
+            )
+            # apply bias to each individual model
             ener_list.append(
-                model.forward_atomic(
-                    extended_coord,
-                    mapping[extended_atype],
-                    nlists_[i],
-                    mapping,
-                    fparam,
-                    aparam,
-                )["energy"]
+                model.apply_out_stat(raw_ret, mapping[extended_atype][:, :nloc])[
+                    "energy"
+                ]
             )
-
         weights = self._compute_weight(extended_coord, extended_atype, nlists_)
 
-        atype = extended_atype[:, :nloc]
-        for idx, model in enumerate(self.models):
-            # TODO: provide interfaces for atomic models to access bias_atom_e
-            if isinstance(model, DPAtomicModel):
-                bias_atom_e = model.fitting_net.bias_atom_e
-            elif isinstance(model, PairTabAtomicModel):
-                bias_atom_e = model.bias_atom_e
-            else:
-                bias_atom_e = None
-            if bias_atom_e is not None:
-                ener_list[idx] += bias_atom_e[atype]
-
         fit_ret = {
             "energy": torch.sum(torch.stack(ener_list) * torch.stack(weights), dim=0),
         }  # (nframes, nloc, 1)
         return fit_ret
 
+    def apply_out_stat(
+        self,
+        ret: Dict[str, torch.Tensor],
+        atype: torch.Tensor,
+    ):
+        """Apply the stat to each atomic output.
+        The developer may override the method to define how the bias is applied
+        to the atomic output of the model.
+
+        Parameters
+        ----------
+        ret
+            The returned dict by the forward_atomic method
+        atype
+            The atom types. nf x nloc
+
+        """
+        return ret
+
     @staticmethod
     def remap_atype(ori_map: List[str], new_map: List[str]) -> torch.Tensor:
         """
@@ -289,27 +301,6 @@
             for _ in range(nmodels)
         ]
 
-    def set_out_bias(self, out_bias: torch.Tensor, 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) -> torch.Tensor:
-        """Return the weighted output bias of the linear atomic model."""
-        # TODO add get_out_bias for linear atomic model
-        raise NotImplementedError
-
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""
         # tricky...

deepmd/pt/model/atomic_model/pairtab_atomic_model.py

@@ -136,6 +136,9 @@
             ]
         )
 
+    def get_out_bias(self) -> torch.Tensor:
+        return self.out_bias
+
     def get_rcut(self) -> float:
         return self.rcut
 
@@ -226,25 +229,6 @@
         """
         self.compute_or_load_out_stat(merged, 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.bias_atom_e = out_bias + self.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.bias_atom_e
-
     def forward_atomic(
         self,
         extended_coord: torch.Tensor,

deepmd/pt/train/training.py

@@ -1168,7 +1168,7 @@ def _model_change_out_bias(
     idx_type_map = sorter[np.searchsorted(model_type_map, new_type_map, sorter=sorter)]
     log.info(
         f"Change output bias of {new_type_map!s} "
-        f"from {to_numpy_array(old_bias[idx_type_map]).reshape(-1)!s} "
-        f"to {to_numpy_array(new_bias[idx_type_map]).reshape(-1)!s}."
+        f"from {to_numpy_array(old_bias[:,idx_type_map]).reshape(-1)!s} "
+        f"to {to_numpy_array(new_bias[:,idx_type_map]).reshape(-1)!s}."
     )
     return _model