Feat: add zbl training (#3398)

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/__init__.py

@@ -22,8 +22,8 @@
     DPAtomicModel,
 )
 from .linear_atomic_model import (
-    DPZBLLinearAtomicModel,
-    LinearAtomicModel,
+    DPZBLLinearEnergyAtomicModel,
+    LinearEnergyAtomicModel,
 )
 from .make_base_atomic_model import (
     make_base_atomic_model,
@@ -37,6 +37,6 @@
     "BaseAtomicModel",
     "DPAtomicModel",
     "PairTabAtomicModel",
-    "LinearAtomicModel",
-    "DPZBLLinearAtomicModel",
+    "LinearEnergyAtomicModel",
+    "DPZBLLinearEnergyAtomicModel",
 ]

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -38,7 +38,7 @@
 )
 
 
-class LinearAtomicModel(BaseAtomicModel):
+class LinearEnergyAtomicModel(BaseAtomicModel):
     """Linear model make linear combinations of several existing models.
 
     Parameters
@@ -59,14 +59,16 @@ def __init__(
         self.models = models
         sub_model_type_maps = [md.get_type_map() for md in models]
         err_msg = []
+        self.mapping_list = []
         common_type_map = set(type_map)
+        self.type_map = type_map
         for tpmp in sub_model_type_maps:
             if not common_type_map.issubset(set(tpmp)):
                 err_msg.append(
                     f"type_map {tpmp} is not a subset of type_map {type_map}"
                 )
+            self.mapping_list.append(self.remap_atype(tpmp, self.type_map))
         assert len(err_msg) == 0, "\n".join(err_msg)
-        self.type_map = type_map
         self.mixed_types_list = [model.mixed_types() for model in self.models]
         super().__init__(**kwargs)
 
@@ -163,17 +165,20 @@ def forward_atomic(
                 self.mixed_types_list, raw_nlists, self.get_model_sels()
             )
         ]
-        ener_list = [
-            model.forward_atomic(
-                extended_coord,
-                extended_atype,
-                nl,
-                mapping,
-                fparam,
-                aparam,
-            )["energy"]
-            for model, nl in zip(self.models, nlists_)
-        ]
+        ener_list = []
+
+        for i, model in enumerate(self.models):
+            mapping = self.mapping_list[i]
+            ener_list.append(
+                model.forward_atomic(
+                    extended_coord,
+                    mapping[extended_atype],
+                    nlists_[i],
+                    mapping,
+                    fparam,
+                    aparam,
+                )["energy"]
+            )
         self.weights = self._compute_weight(extended_coord, extended_atype, nlists_)
         self.atomic_bias = None
         if self.atomic_bias is not None:
@@ -184,6 +189,29 @@ def forward_atomic(
             }  # (nframes, nloc, 1)
         return fit_ret
 
+    @staticmethod
+    def remap_atype(ori_map: List[str], new_map: List[str]) -> np.ndarray:
+        """
+        This method is used to map the atype from the common type_map to the original type_map of
+        indivial AtomicModels.
+
+        Parameters
+        ----------
+        ori_map : List[str]
+            The original type map of an AtomicModel.
+        new_map : List[str]
+            The common type map of the DPZBLLinearEnergyAtomicModel, created by the `get_type_map` method,
+            must be a subset of the ori_map.
+
+        Returns
+        -------
+        np.ndarray
+        """
+        type_2_idx = {atp: idx for idx, atp in enumerate(ori_map)}
+        # this maps the atype in the new map to the original map
+        mapping = np.array([type_2_idx[new_map[idx]] for idx in range(len(new_map))])
+        return mapping
+
     def fitting_output_def(self) -> FittingOutputDef:
         return FittingOutputDef(
             [
@@ -261,7 +289,7 @@ def is_aparam_nall(self) -> bool:
         return False
 
 
-class DPZBLLinearAtomicModel(LinearAtomicModel):
+class DPZBLLinearEnergyAtomicModel(LinearEnergyAtomicModel):
     """Model linearly combine a list of AtomicModels.
 
     Parameters
@@ -308,7 +336,7 @@ def serialize(self) -> dict:
                 "@class": "Model",
                 "type": "zbl",
                 "@version": 1,
-                "models": LinearAtomicModel.serialize(
+                "models": LinearEnergyAtomicModel.serialize(
                     [self.dp_model, self.zbl_model], self.type_map
                 ),
                 "sw_rmin": self.sw_rmin,
@@ -319,7 +347,7 @@ def serialize(self) -> dict:
         return dd
 
     @classmethod
-    def deserialize(cls, data) -> "DPZBLLinearAtomicModel":
+    def deserialize(cls, data) -> "DPZBLLinearEnergyAtomicModel":
         data = copy.deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class")
@@ -328,7 +356,7 @@ def deserialize(cls, data) -> "DPZBLLinearAtomicModel":
         sw_rmax = data.pop("sw_rmax")
         smin_alpha = data.pop("smin_alpha")
 
-        ([dp_model, zbl_model], type_map) = LinearAtomicModel.deserialize(
+        ([dp_model, zbl_model], type_map) = LinearEnergyAtomicModel.deserialize(
             data.pop("models")
         )
 

deepmd/dpmodel/atomic_model/make_base_atomic_model.py

@@ -54,18 +54,13 @@ def get_rcut(self) -> float:
             pass
 
         @abstractmethod
-        def get_type_map(self) -> Optional[List[str]]:
+        def get_type_map(self) -> List[str]:
             """Get the type map."""
+            pass
 
         def get_ntypes(self) -> int:
             """Get the number of atom types."""
-            tmap = self.get_type_map()
-            if tmap is not None:
-                return len(tmap)
-            else:
-                raise ValueError(
-                    "cannot infer the number of types from a None type map"
-                )
+            return len(self.get_type_map())
 
         @abstractmethod
         def get_sel(self) -> List[int]:

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -66,9 +66,17 @@ def __init__(
         self.type_map = type_map
 
         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
 
@@ -145,7 +153,8 @@ def deserialize(cls, data) -> "PairTabAtomicModel":
         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/__init__.py

@@ -21,8 +21,8 @@
     DPAtomicModel,
 )
 from .linear_atomic_model import (
-    DPZBLLinearAtomicModel,
-    LinearAtomicModel,
+    DPZBLLinearEnergyAtomicModel,
+    LinearEnergyAtomicModel,
 )
 from .pairtab_atomic_model import (
     PairTabAtomicModel,
@@ -32,6 +32,6 @@
     "BaseAtomicModel",
     "DPAtomicModel",
     "PairTabAtomicModel",
-    "LinearAtomicModel",
-    "DPZBLLinearAtomicModel",
+    "LinearEnergyAtomicModel",
+    "DPZBLLinearEnergyAtomicModel",
 ]