Feat: add zbl training

deepmd/pt/model/atomic_model/linear_atomic_model.py

@@ -25,6 +25,9 @@
     get_multiple_nlist_key,
     nlist_distinguish_types,
 )
+from deepmd.utils.path import (
+    DPPath,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -184,14 +187,17 @@
 
         weights = self._compute_weight(extended_coord, extended_atype, nlists_)
 
-        if self.atomic_bias is not None:
-            raise NotImplementedError("Need to add bias in a future PR.")
-        else:
-            fit_ret = {
-                "energy": torch.sum(
-                    torch.stack(ener_list) * torch.stack(weights), dim=0
-                ),
-            }  # (nframes, nloc, 1)
+        atype = extended_atype[:, :nloc]
+        for idx, m in enumerate(self.models):
+            if isinstance(m, DPAtomicModel) and m.fitting_net is not None:
+                bias_atom_e = m.fitting_net.bias_atom_e
+            elif isinstance(m, PairTabAtomicModel):
+                bias_atom_e = m.bias_atom_e
+            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 fitting_output_def(self) -> FittingOutputDef:
@@ -307,6 +313,29 @@
         # this is a placeholder being updated in _compute_weight, to handle Jit attribute init error.
         self.zbl_weight = torch.empty(0, dtype=torch.float64, device=env.DEVICE)
 
+    def compute_or_load_stat(
+        self,
+        sampled_func,
+        stat_file_path: Optional[DPPath] = None,
+    ):
+        """
+        Compute or load the statistics parameters of the model,
+        such as mean and standard deviation of descriptors or the energy bias of the fitting net.
+        When `sampled` is provided, all the statistics parameters will be calculated (or re-calculated for update),
+        and saved in the `stat_file_path`(s).
+        When `sampled` is not provided, it will check the existence of `stat_file_path`(s)
+        and load the calculated statistics parameters.
+
+        Parameters
+        ----------
+        sampled_func
+            The lazy sampled function to get data frames from different data systems.
+        stat_file_path
+            The dictionary of paths to the statistics files.
+        """
+        self.dp_model.compute_or_load_stat(sampled_func, stat_file_path)
+        self.zbl_model.compute_output_stats(sampled_func, stat_file_path)
+
     def serialize(self) -> dict:
         dd = BaseAtomicModel.serialize(self)
         dd.update(

deepmd/pt/model/atomic_model/pairtab_atomic_model.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import copy
 from typing import (
+    Callable,
     Dict,
     List,
     Optional,
@@ -13,9 +14,21 @@
     FittingOutputDef,
     OutputVariableDef,
 )
+from deepmd.pt.utils import (
+    env,
+)
+from deepmd.pt.utils.utils import (
+    to_numpy_array,
+)
+from deepmd.utils.out_stat import (
+    compute_stats_from_redu,
+)
 from deepmd.utils.pair_tab import (
     PairTab,
 )
+from deepmd.utils.path import (
+    DPPath,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -57,6 +70,7 @@
         self.tab_file = tab_file
         self.rcut = rcut
         self.tab = self._set_pairtab(tab_file, rcut)
+
         BaseAtomicModel.__init__(self, **kwargs)
 
         # handle deserialization with no input file
@@ -70,6 +84,7 @@
         else:
             self.register_buffer("tab_info", None)
             self.register_buffer("tab_data", None)
+        self.bias_atom_e = None
 
         # self.model_type = "ener"
         # self.model_version = MODEL_VERSION ## this shoud be in the parent class
@@ -154,6 +169,65 @@
         tab_model.register_buffer("tab_data", torch.from_numpy(tab_model.tab.tab_data))
         return tab_model
 
+    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.
+
+        """
+        if stat_file_path is not None:
+            stat_file_path = stat_file_path / "bias_atom_e"
+        if stat_file_path is not None and stat_file_path.is_file():
+            bias_atom_e = stat_file_path.load_numpy()
+        else:
+            if callable(merged):
+                # only get data for once
+                sampled = merged()
+            else:
+                sampled = merged
+            energy = [item["energy"] for item in sampled]
+            data_mixed_type = "real_natoms_vec" in sampled[0]
+            if data_mixed_type:
+                input_natoms = [item["real_natoms_vec"] for item in sampled]
+            else:
+                input_natoms = [item["natoms"] for item in sampled]
+            # shape: (nframes, ndim)
+            merged_energy = to_numpy_array(torch.cat(energy))
+            # shape: (nframes, ntypes)
+            merged_natoms = to_numpy_array(torch.cat(input_natoms)[:, 2:])
+
+            bias_atom_e, _ = compute_stats_from_redu(
+                merged_energy,
+                merged_natoms,
+                assigned_bias=None,
+                rcond=None,
+            )
+            if stat_file_path is not None:
+                stat_file_path.save_numpy(bias_atom_e)
+        assert all(x is not None for x in [bias_atom_e])
+        ntypes = merged_natoms.shape[1]
+        self.bias_atom_e = torch.empty(
+            [ntypes, 1], dtype=torch.float64, device=env.DEVICE
+        )
+        self.bias_atom_e.copy_(
+            torch.tensor(bias_atom_e, device=env.DEVICE).view([ntypes, 1])
+        )
+
     def forward_atomic(
         self,
         extended_coord: torch.Tensor,

deepmd/pt/train/training.py

@@ -26,6 +26,7 @@
 )
 from deepmd.pt.model.model import (
     get_model,
+    get_zbl_model,
 )
 from deepmd.pt.optimizer import (
     KFOptimizerWrapper,
@@ -243,7 +244,10 @@
         def get_single_model(
             _model_params,
         ):
-            model = get_model(deepcopy(_model_params)).to(DEVICE)
+            if "use_srtab" in _model_params:
+                model = get_zbl_model(deepcopy(_model_params)).to(DEVICE)
+            else:
+                model = get_model(deepcopy(_model_params)).to(DEVICE)
             return model
 
         def get_lr(lr_params):