Feat: Add consistency test for ZBL between dp and pt

deepmd/dpmodel/model/dp_zbl_model.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Optional,
+)
+
+from deepmd.dpmodel.atomic_model.linear_atomic_model import (
+    DPZBLLinearEnergyAtomicModel,
+)
+from deepmd.dpmodel.model.base_model import (
+    BaseModel,
+)
+from deepmd.dpmodel.model.dp_model import (
+    DPModelCommon,
+)
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
+
+from .make_model import (
+    make_model,
+)
+
+DPEnergyModel_ = make_model(DPZBLLinearEnergyAtomicModel)
+
+
+@BaseModel.register("zbl")
+class DPZBLModel(DPEnergyModel_):
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        DPEnergyModel_.__init__(self, *args, **kwargs)
+
+
+@classmethod
+def update_sel(
+    cls,
+    train_data: DeepmdDataSystem,
+    type_map: Optional[list[str]],
+    local_jdata: dict,
+) -> tuple[dict, Optional[float]]:
+    """Update the selection and perform neighbor statistics.
+
+    Parameters
+    ----------
+    train_data : DeepmdDataSystem
+        data used to do neighbor statistics
+    type_map : list[str], optional
+        The name of each type of atoms
+    local_jdata : dict
+        The local data refer to the current class
+
+    Returns
+    -------
+    dict
+        The updated local data
+    float
+        The minimum distance between two atoms
+    """
+    local_jdata_cpy = local_jdata.copy()
+    local_jdata_cpy["dpmodel"], min_nbor_dist = DPModelCommon.update_sel(
+        train_data, type_map, local_jdata["dpmodel"]
+    )
+    return local_jdata_cpy, min_nbor_dist

deepmd/dpmodel/model/model.py

@@ -1,4 +1,13 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.dpmodel.atomic_model.dp_atomic_model import (
+    DPAtomicModel,
+)
+from deepmd.dpmodel.atomic_model.pairtab_atomic_model import (
+    PairTabAtomicModel,
+)
+from deepmd.dpmodel.descriptor.base_descriptor import (
+    BaseDescriptor,
+)
 from deepmd.dpmodel.descriptor.se_e2_a import (
     DescrptSeA,
 )
@@ -8,6 +17,9 @@
 from deepmd.dpmodel.model.base_model import (
     BaseModel,
 )
+from deepmd.dpmodel.model.dp_zbl_model import (
+    DPZBLModel,
+)
 from deepmd.dpmodel.model.ener_model import (
     EnergyModel,
 )
@@ -55,6 +67,45 @@ def get_standard_model(data: dict) -> EnergyModel:
     )
 
 
+def get_zbl_model(data: dict):
+    data["descriptor"]["ntypes"] = len(data["type_map"])
+    descriptor = BaseDescriptor(**data["descriptor"])
+    fitting_type = data["fitting_net"].pop("type")
+    if fitting_type == "ener":
+        fitting = EnergyFittingNet(
+            ntypes=descriptor.get_ntypes(),
+            dim_descrpt=descriptor.get_dim_out(),
+            mixed_types=descriptor.mixed_types(),
+            **data["fitting_net"],
+        )
+    else:
+        raise ValueError(f"Unknown fitting type {fitting_type}")
+
+    dp_model = DPAtomicModel(descriptor, fitting, type_map=data["type_map"])
+    # pairtab
+    filepath = data["use_srtab"]
+    pt_model = PairTabAtomicModel(
+        filepath,
+        data["descriptor"]["rcut"],
+        data["descriptor"]["sel"],
+        type_map=data["type_map"],
+    )
+
+    rmin = data["sw_rmin"]
+    rmax = data["sw_rmax"]
+    atom_exclude_types = data.get("atom_exclude_types", [])
+    pair_exclude_types = data.get("pair_exclude_types", [])
+    return DPZBLModel(
+        dp_model,
+        pt_model,
+        rmin,
+        rmax,
+        type_map=data["type_map"],
+        atom_exclude_types=atom_exclude_types,
+        pair_exclude_types=pair_exclude_types,
+    )
+
+
 def get_spin_model(data: dict) -> SpinModel:
     """Get a spin model from a dictionary.
 
@@ -100,6 +151,8 @@ def get_model(data: dict):
     if model_type == "standard":
         if "spin" in data:
             return get_spin_model(data)
+        elif "use_srtab" in data:
+            return get_zbl_model(data)
         else:
             return get_standard_model(data)
     else:

source/tests/consistent/common.py

@@ -75,7 +75,7 @@
 class CommonTest(ABC):
     data: ClassVar[dict]
     """Arguments data."""
-    addtional_data: ClassVar[dict] = {}
+    additional_data: ClassVar[dict] = {}
     """Additional data that will not be checked."""
     tf_class: ClassVar[Optional[type]]
     """TensorFlow model class."""
@@ -128,7 +128,7 @@ def init_backend_cls(self, cls) -> Any:
 
     def pass_data_to_cls(self, cls, data) -> Any:
         """Pass data to the class."""
-        return cls(**data, **self.addtional_data)
+        return cls(**data, **self.additional_data)
 
     @abstractmethod
     def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]:

source/tests/consistent/fitting/test_dipole.py

@@ -83,7 +83,7 @@ def setUp(self):
         self.atype.sort()
 
     @property
-    def addtional_data(self) -> dict:
+    def additional_data(self) -> dict:
         (
             resnet_dt,
             precision,

source/tests/consistent/fitting/test_dos.py

@@ -124,7 +124,7 @@ def setUp(self):
         ).reshape(-1, 1)
 
     @property
-    def addtional_data(self) -> dict:
+    def additional_data(self) -> dict:
         (
             resnet_dt,
             precision,

source/tests/consistent/fitting/test_ener.py

@@ -134,7 +134,7 @@ def setUp(self):
         ).reshape(-1, 1)
 
     @property
-    def addtional_data(self) -> dict:
+    def additional_data(self) -> dict:
         (
             resnet_dt,
             precision,

source/tests/consistent/fitting/test_polar.py

@@ -83,7 +83,7 @@ def setUp(self):
         self.atype.sort()
 
     @property
-    def addtional_data(self) -> dict:
+    def additional_data(self) -> dict:
         (
             resnet_dt,
             precision,

source/tests/consistent/fitting/test_property.py

@@ -104,7 +104,7 @@ def setUp(self):
         ).reshape(-1, 1)
 
     @property
-    def addtional_data(self) -> dict:
+    def additional_data(self) -> dict:
         (
             resnet_dt,
             precision,

source/tests/consistent/model/test_ener.py

@@ -130,7 +130,7 @@ def pass_data_to_cls(self, cls, data) -> Any:
             return get_model_pt(data)
         elif cls is EnergyModelJAX:
             return get_model_jax(data)
-        return cls(**data, **self.addtional_data)
+        return cls(**data, **self.additional_data)
 
     def setUp(self):
         CommonTest.setUp(self)