Merge branch 'devel' into train_rf

Signed-off-by: Duo <[email protected]>

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -1,8 +1,89 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Dict,
+    List,
+    Optional,
+    Tuple,
+)
+
 import numpy as np
 
+from deepmd.dpmodel.utils import (
+    AtomExcludeMask,
+    PairExcludeMask,
+)
+
 from .make_base_atomic_model import (
     make_base_atomic_model,
 )
 
-BaseAtomicModel = make_base_atomic_model(np.ndarray)
+BaseAtomicModel_ = make_base_atomic_model(np.ndarray)
+
+
+class BaseAtomicModel(BaseAtomicModel_):
+    def __init__(
+        self,
+        atom_exclude_types: List[int] = [],
+        pair_exclude_types: List[Tuple[int, int]] = [],
+    ):
+        super().__init__()
+        self.reinit_atom_exclude(atom_exclude_types)
+        self.reinit_pair_exclude(pair_exclude_types)
+
+    def reinit_atom_exclude(
+        self,
+        exclude_types: List[int] = [],
+    ):
+        self.atom_exclude_types = exclude_types
+        if exclude_types == []:
+            self.atom_excl = None
+        else:
+            self.atom_excl = AtomExcludeMask(self.get_ntypes(), self.atom_exclude_types)
+
+    def reinit_pair_exclude(
+        self,
+        exclude_types: List[Tuple[int, int]] = [],
+    ):
+        self.pair_exclude_types = exclude_types
+        if exclude_types == []:
+            self.pair_excl = None
+        else:
+            self.pair_excl = PairExcludeMask(self.get_ntypes(), self.pair_exclude_types)
+
+    def forward_common_atomic(
+        self,
+        extended_coord: np.ndarray,
+        extended_atype: np.ndarray,
+        nlist: np.ndarray,
+        mapping: Optional[np.ndarray] = None,
+        fparam: Optional[np.ndarray] = None,
+        aparam: Optional[np.ndarray] = None,
+    ) -> Dict[str, np.ndarray]:
+        _, nloc, _ = nlist.shape
+        atype = extended_atype[:, :nloc]
+        if self.pair_excl is not None:
+            pair_mask = self.pair_excl.build_type_exclude_mask(nlist, extended_atype)
+            # exclude neighbors in the nlist
+            nlist = np.where(pair_mask == 1, nlist, -1)
+
+        ret_dict = self.forward_atomic(
+            extended_coord,
+            extended_atype,
+            nlist,
+            mapping=mapping,
+            fparam=fparam,
+            aparam=aparam,
+        )
+
+        if self.atom_excl is not None:
+            atom_mask = self.atom_excl.build_type_exclude_mask(atype)
+            for kk in ret_dict.keys():
+                ret_dict[kk] = ret_dict[kk] * atom_mask[:, :, None]
+
+        return ret_dict
+
+    def serialize(self) -> dict:
+        return {
+            "atom_exclude_types": self.atom_exclude_types,
+            "pair_exclude_types": self.pair_exclude_types,
+        }

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -46,11 +46,12 @@ def __init__(
         descriptor,
         fitting,
         type_map: Optional[List[str]] = None,
+        **kwargs,
     ):
-        super().__init__()
         self.type_map = type_map
         self.descriptor = descriptor
         self.fitting = fitting
+        super().__init__(**kwargs)
 
     def fitting_output_def(self) -> FittingOutputDef:
         """Get the output def of the fitting net."""
@@ -132,24 +133,29 @@ def forward_atomic(
         return ret
 
     def serialize(self) -> dict:
-        return {
-            "@class": "Model",
-            "type": "standard",
-            "@version": 1,
-            "type_map": self.type_map,
-            "descriptor": self.descriptor.serialize(),
-            "fitting": self.fitting.serialize(),
-        }
+        dd = super().serialize()
+        dd.update(
+            {
+                "@class": "Model",
+                "type": "standard",
+                "@version": 1,
+                "type_map": self.type_map,
+                "descriptor": self.descriptor.serialize(),
+                "fitting": self.fitting.serialize(),
+            }
+        )
+        return dd
 
     @classmethod
     def deserialize(cls, data) -> "DPAtomicModel":
         data = copy.deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class")
         data.pop("type")
-        descriptor_obj = BaseDescriptor.deserialize(data["descriptor"])
-        fitting_obj = BaseFitting.deserialize(data["fitting"])
-        obj = cls(descriptor_obj, fitting_obj, type_map=data["type_map"])
+        descriptor_obj = BaseDescriptor.deserialize(data.pop("descriptor"))
+        fitting_obj = BaseFitting.deserialize(data.pop("fitting"))
+        type_map = data.pop("type_map", None)
+        obj = cls(descriptor_obj, fitting_obj, type_map=type_map, **data)
         return obj
 
     def get_dim_fparam(self) -> int:

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -52,9 +52,9 @@ def __init__(
         models: List[BaseAtomicModel],
         **kwargs,
     ):
-        super().__init__()
         self.models = models
         self.mixed_types_list = [model.mixed_types() for model in self.models]
+        super().__init__(**kwargs)
 
     def mixed_types(self) -> bool:
         """If true, the model
@@ -273,34 +273,39 @@ def __init__(
         self.smin_alpha = smin_alpha
 
     def serialize(self) -> dict:
-        return {
-            "@class": "Model",
-            "type": "zbl",
-            "@version": 1,
-            "models": LinearAtomicModel.serialize([self.dp_model, self.zbl_model]),
-            "sw_rmin": self.sw_rmin,
-            "sw_rmax": self.sw_rmax,
-            "smin_alpha": self.smin_alpha,
-        }
+        dd = BaseAtomicModel.serialize(self)
+        dd.update(
+            {
+                "@class": "Model",
+                "type": "zbl",
+                "@version": 1,
+                "models": LinearAtomicModel.serialize([self.dp_model, self.zbl_model]),
+                "sw_rmin": self.sw_rmin,
+                "sw_rmax": self.sw_rmax,
+                "smin_alpha": self.smin_alpha,
+            }
+        )
+        return dd
 
     @classmethod
     def deserialize(cls, data) -> "DPZBLLinearAtomicModel":
         data = copy.deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class")
         data.pop("type")
-        sw_rmin = data["sw_rmin"]
-        sw_rmax = data["sw_rmax"]
-        smin_alpha = data["smin_alpha"]
+        sw_rmin = data.pop("sw_rmin")
+        sw_rmax = data.pop("sw_rmax")
+        smin_alpha = data.pop("smin_alpha")
 
-        dp_model, zbl_model = LinearAtomicModel.deserialize(data["models"])
+        dp_model, zbl_model = LinearAtomicModel.deserialize(data.pop("models"))
 
         return cls(
             dp_model=dp_model,
             zbl_model=zbl_model,
             sw_rmin=sw_rmin,
             sw_rmax=sw_rmax,
             smin_alpha=smin_alpha,
+            **data,
         )
 
     def _compute_weight(

deepmd/dpmodel/atomic_model/make_base_atomic_model.py

@@ -48,6 +48,16 @@ def get_rcut(self) -> float:
         def get_type_map(self) -> Optional[List[str]]:
             """Get the type map."""
 
+        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"
+                )
+
         @abstractmethod
         def get_sel(self) -> List[int]:
             """Returns the number of selected atoms for each type."""

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -109,25 +109,29 @@ def mixed_types(self) -> bool:
         return True
 
     def serialize(self) -> dict:
-        return {
-            "@class": "Model",
-            "type": "pairtab",
-            "@version": 1,
-            "tab": self.tab.serialize(),
-            "rcut": self.rcut,
-            "sel": self.sel,
-        }
+        dd = BaseAtomicModel.serialize(self)
+        dd.update(
+            {
+                "@class": "Model",
+                "type": "pairtab",
+                "@version": 1,
+                "tab": self.tab.serialize(),
+                "rcut": self.rcut,
+                "sel": self.sel,
+            }
+        )
+        return dd
 
     @classmethod
     def deserialize(cls, data) -> "PairTabAtomicModel":
         data = copy.deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 1, 1)
         data.pop("@class")
         data.pop("type")
-        rcut = data["rcut"]
-        sel = data["sel"]
-        tab = PairTab.deserialize(data["tab"])
-        tab_model = cls(None, rcut, sel)
+        rcut = data.pop("rcut")
+        sel = data.pop("sel")
+        tab = PairTab.deserialize(data.pop("tab"))
+        tab_model = cls(None, rcut, sel, **data)
         tab_model.tab = tab
         tab_model.tab_info = tab_model.tab.tab_info
         tab_model.tab_data = tab_model.tab.tab_data

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -133,6 +133,22 @@ def deserialize(cls, data: dict) -> "BD":
                 return BD.get_class_by_type(data["type"]).deserialize(data)
             raise NotImplementedError("Not implemented in class %s" % cls.__name__)
 
+        @classmethod
+        @abstractmethod
+        def update_sel(cls, global_jdata: dict, local_jdata: dict):
+            """Update the selection and perform neighbor statistics.
+
+            Parameters
+            ----------
+            global_jdata : dict
+                The global data, containing the training section
+            local_jdata : dict
+                The local data refer to the current class
+            """
+            # call subprocess
+            cls = cls.get_class_by_type(j_get_type(local_jdata, cls.__name__))
+            return cls.update_sel(global_jdata, local_jdata)
+
     setattr(BD, fwd_method_name, BD.fwd)
     delattr(BD, "fwd")
 

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -3,6 +3,9 @@
 
 import numpy as np
 
+from deepmd.dpmodel.utils.update_sel import (
+    UpdateSel,
+)
 from deepmd.env import (
     GLOBAL_NP_FLOAT_PRECISION,
 )
@@ -392,3 +395,17 @@ def deserialize(cls, data: dict) -> "DescrptSeA":
         obj.embeddings = NetworkCollection.deserialize(embeddings)
         obj.env_mat = EnvMat.deserialize(env_mat)
         return obj
+
+    @classmethod
+    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+        """Update the selection and perform neighbor statistics.
+
+        Parameters
+        ----------
+        global_jdata : dict
+            The global data, containing the training section
+        local_jdata : dict
+            The local data refer to the current class
+        """
+        local_jdata_cpy = local_jdata.copy()
+        return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False)

deepmd/dpmodel/descriptor/se_r.py

@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import numpy as np
 
+from deepmd.dpmodel.utils.update_sel import (
+    UpdateSel,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -328,3 +331,17 @@ def deserialize(cls, data: dict) -> "DescrptSeR":
         obj.embeddings = NetworkCollection.deserialize(embeddings)
         obj.env_mat = EnvMat.deserialize(env_mat)
         return obj
+
+    @classmethod
+    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+        """Update the selection and perform neighbor statistics.
+
+        Parameters
+        ----------
+        global_jdata : dict
+            The global data, containing the training section
+        local_jdata : dict
+            The local data refer to the current class
+        """
+        local_jdata_cpy = local_jdata.copy()
+        return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False)

deepmd/dpmodel/model/base_model.py

@@ -146,6 +146,21 @@ def get_nsel(self) -> int:
             """Returns the total number of selected neighboring atoms in the cut-off radius."""
             pass
 
+        @classmethod
+        @abstractmethod
+        def update_sel(cls, global_jdata: dict, local_jdata: dict):
+            """Update the selection and perform neighbor statistics.
+
+            Parameters
+            ----------
+            global_jdata : dict
+                The global data, containing the training section
+            local_jdata : dict
+                The local data refer to the current class
+            """
+            cls = cls.get_class_by_type(local_jdata.get("type", "standard"))
+            return cls.update_sel(global_jdata, local_jdata)
+
     return BaseBaseModel