Chore: make type_map complusory model attribute

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -47,19 +47,35 @@
         The cutoff radius.
     sel : int or list[int]
         The maxmum number of atoms in the cut-off radius.
+    type_map: List[str]
+        Mapping atom type to the name (str) of the type.
+        For example `type_map[1]` gives the name of the type 1.
     """
 
     def __init__(
-        self, tab_file: str, rcut: float, sel: Union[int, List[int]], **kwargs
+        self,
+        tab_file: str,
+        rcut: float,
+        sel: Union[int, List[int]],
+        type_map: List[str],
+        **kwargs,
     ):
         super().__init__()
         self.tab_file = tab_file
         self.rcut = rcut
 
         self.tab = PairTab(self.tab_file, rcut=rcut)
+        self.type_map = type_map
+        self.ntypes = len(type_map)
 
         if self.tab_file is not None:
             self.tab_info, self.tab_data = self.tab.get()
+            nspline, ntypes_tab = self.tab_info[-2:].astype(int)
+            self.tab_data = self.tab_data.reshape(ntypes_tab, ntypes_tab, nspline, 4)
+            if self.ntypes != ntypes_tab:
+                raise ValueError(
+                    "The `type_map` provided does not match the number of columns in the table."
+                )
         else:
             self.tab_info, self.tab_data = None, None
 
@@ -118,6 +134,7 @@
                 "tab": self.tab.serialize(),
                 "rcut": self.rcut,
                 "sel": self.sel,
+                "type_map": self.type_map,
             }
         )
         return dd
@@ -130,11 +147,13 @@
         data.pop("type")
         rcut = data.pop("rcut")
         sel = data.pop("sel")
+        type_map = data.pop("type_map")
         tab = PairTab.deserialize(data.pop("tab"))
-        tab_model = cls(None, rcut, sel, **data)
+        tab_model = cls(None, rcut, sel, type_map, **data)
         tab_model.tab = tab
         tab_model.tab_info = tab_model.tab.tab_info
-        tab_model.tab_data = tab_model.tab.tab_data
+        nspline, ntypes = tab_model.tab_info[-2:].astype(int)
+        tab_model.tab_data = tab_model.tab.tab_data.reshape(ntypes, ntypes, nspline, 4)
         return tab_model
 
     def forward_atomic(

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -50,7 +50,7 @@ def __init__(
         self,
         descriptor,
         fitting,
-        type_map: Optional[List[str]],
+        type_map: List[str],
         **kwargs,
     ):
         torch.nn.Module.__init__(self)

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,
 )
@@ -77,10 +80,9 @@
         """Get the cut-off radius."""
         return max(self.get_model_rcuts())
 
-    @torch.jit.export
+    @abstractmethod
     def get_type_map(self) -> List[str]:
         """Get the type map."""
-        raise NotImplementedError("TODO: implement this method")
 
     def get_model_rcuts(self) -> List[float]:
         """Get the cut-off radius for each individual models."""
@@ -103,8 +105,8 @@
         nsels = torch.tensor(self.get_model_nsels(), device=device)
         zipped = torch.stack(
             [
-                torch.tensor(rcuts, device=device),
-                torch.tensor(nsels, device=device),
+                rcuts,
+                nsels,
             ],
             dim=0,
         ).T
@@ -184,14 +186,20 @@
 
         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, model in enumerate(self.models):
+            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 fitting_output_def(self) -> FittingOutputDef:
@@ -307,6 +315,39 @@
         # 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)
 
+    @torch.jit.export
+    def get_type_map(self) -> List[str]:
+        dp_map = self.dp_model.get_type_map()
+        zbl_map = self.zbl_model.get_type_map()
+        return self.dp_model.get_type_map()
+
+    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_or_load_stat(sampled_func, stat_file_path)
+
+    def change_energy_bias(self):
+        # need to implement
+        pass
+
     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,18 @@
     FittingOutputDef,
     OutputVariableDef,
 )
+from deepmd.pt.utils import (
+    env,
+)
+from deepmd.pt.utils.stat import (
+    compute_output_stats,
+)
 from deepmd.utils.pair_tab import (
     PairTab,
 )
+from deepmd.utils.path import (
+    DPPath,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -47,29 +57,59 @@
         The cutoff radius.
     sel : int or list[int]
         The maxmum number of atoms in the cut-off radius.
+    type_map: List[str]
+        Mapping atom type to the name (str) of the type.
+        For example `type_map[1]` gives the name of the type 1.
+    rcond : float, optional
+        The condition number for the regression of atomic energy.
+    atom_ener
+        Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set.
     """
 
     def __init__(
-        self, tab_file: str, rcut: float, sel: Union[int, List[int]], **kwargs
+        self,
+        tab_file: str,
+        rcut: float,
+        sel: Union[int, List[int]],
+        type_map: List[str],
+        rcond: Optional[float] = None,
+        atom_ener: Optional[List[float]] = None,
+        **kwargs,
     ):
         torch.nn.Module.__init__(self)
         self.model_def_script = ""
         self.tab_file = tab_file
         self.rcut = rcut
         self.tab = self._set_pairtab(tab_file, rcut)
+
         BaseAtomicModel.__init__(self, **kwargs)
+        self.rcond = rcond
+        self.atom_ener = atom_ener
+        self.type_map = type_map
+        self.ntypes = len(type_map)
 
         # handle deserialization with no input file
         if self.tab_file is not None:
             (
                 tab_info,
                 tab_data,
             ) = self.tab.get()  # this returns -> Tuple[np.array, np.array]
+            nspline, ntypes_tab = tab_info[-2:].astype(int)
             self.register_buffer("tab_info", torch.from_numpy(tab_info))
-            self.register_buffer("tab_data", torch.from_numpy(tab_data))
+            self.register_buffer(
+                "tab_data",
+                torch.from_numpy(tab_data).reshape(ntypes_tab, ntypes_tab, nspline, 4),
+            )
+            if self.ntypes != ntypes_tab:
+                raise ValueError(
+                    "The `type_map` provided does not match the number of columns in the table."
+                )
         else:
             self.register_buffer("tab_info", None)
             self.register_buffer("tab_data", None)
+        self.bias_atom_e = torch.zeros(
+            self.ntypes, 1, dtype=env.GLOBAL_PT_ENER_FLOAT_PRECISION, device=env.DEVICE
+        )
 
         # self.model_type = "ener"
         # self.model_version = MODEL_VERSION ## this shoud be in the parent class
@@ -103,8 +143,8 @@
         return self.rcut
 
     @torch.jit.export
-    def get_type_map(self) -> Optional[List[str]]:
-        raise NotImplementedError("TODO: implement this method")
+    def get_type_map(self) -> List[str]:
+        return self.type_map
 
     def get_sel(self) -> List[int]:
         return [self.sel]
@@ -135,6 +175,9 @@
                 "tab": self.tab.serialize(),
                 "rcut": self.rcut,
                 "sel": self.sel,
+                "type_map": self.type_map,
+                "rcond": self.rcond,
+                "atom_ener": self.atom_ener,
             }
         )
         return dd
@@ -145,15 +188,56 @@
         check_version_compatibility(data.pop("@version", 1), 1, 1)
         rcut = data.pop("rcut")
         sel = data.pop("sel")
+        type_map = data.pop("type_map")
+        rcond = data.pop("rcond")
+        atom_ener = data.pop("atom_ener")
         tab = PairTab.deserialize(data.pop("tab"))
         data.pop("@class", None)
         data.pop("type", None)
-        tab_model = cls(None, rcut, sel, **data)
+        tab_model = cls(None, rcut, sel, type_map, rcond, atom_ener, **data)
         tab_model.tab = tab
         tab_model.register_buffer("tab_info", torch.from_numpy(tab_model.tab.tab_info))
-        tab_model.register_buffer("tab_data", torch.from_numpy(tab_model.tab.tab_data))
+        nspline, ntypes = tab_model.tab.tab_info[-2:].astype(int)
+        tab_model.register_buffer(
+            "tab_data",
+            torch.from_numpy(tab_model.tab.tab_data).reshape(
+                ntypes, ntypes, nspline, 4
+            ),
+        )
         return tab_model
 
+    def compute_or_load_stat(
+        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.
+
+        """
+        bias_atom_e = compute_output_stats(
+            merged, stat_file_path, self.rcond, self.atom_ener
+        )
+        self.bias_atom_e.copy_(
+            torch.tensor(bias_atom_e, device=env.DEVICE).view([self.ntypes, 1])
+        )
+
+    def change_energy_bias(self) -> None:
+        # need to implement
+        pass
+
     def forward_atomic(
         self,
         extended_coord: torch.Tensor,

deepmd/pt/model/model/__init__.py

@@ -66,7 +66,10 @@ def get_zbl_model(model_params):
     # pairtab
     filepath = model_params["use_srtab"]
     pt_model = PairTabAtomicModel(
-        filepath, model_params["descriptor"]["rcut"], model_params["descriptor"]["sel"]
+        filepath,
+        model_params["descriptor"]["rcut"],
+        model_params["descriptor"]["sel"],
+        type_map=model_params["type_map"],
     )
 
     rmin = model_params["sw_rmin"]