feat: support exclude atypes in atomic model

deepmd/dpmodel/atomic_model/make_base_atomic_model.py

@@ -48,6 +48,16 @@
         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/pt/model/atomic_model/base_atomic_model.py

@@ -1,20 +1,99 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
 
+from typing import (
+    Dict,
+    List,
+    Optional,
+    Tuple,
+)
+
 import torch
 
 from deepmd.dpmodel.atomic_model import (
     make_base_atomic_model,
 )
+from deepmd.pt.utils import (
+    AtomExcludeMask,
+    PairExcludeMask,
+)
 
 BaseAtomicModel_ = make_base_atomic_model(torch.Tensor)
 
 
 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)
+
     # export public methods that are not abstract
     get_nsel = torch.jit.export(BaseAtomicModel_.get_nsel)
     get_nnei = torch.jit.export(BaseAtomicModel_.get_nnei)
 
     @torch.jit.export
     def get_model_def_script(self) -> str:
         return self.model_def_script
+
+    def forward_common_atomic(
+        self,
+        extended_coord: torch.Tensor,
+        extended_atype: torch.Tensor,
+        nlist: torch.Tensor,
+        mapping: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+    ) -> Dict[str, torch.Tensor]:
+        nf, nloc, nnei = nlist.shape
+        nf, nall = extended_atype.shape
+        atype = extended_atype[:, :nloc]
+        if self.pair_excl is not None:
+            pair_mask = self.pair_excl(nlist, extended_atype)
+            # exclude neighbors in the nlist
+            nlist = torch.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(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/pt/model/atomic_model/dp_atomic_model.py

@@ -49,8 +49,14 @@
             For example `type_map[1]` gives the name of the type 1.
     """
 
-    def __init__(self, descriptor, fitting, type_map: Optional[List[str]]):
-        super().__init__()
+    def __init__(
+        self,
+        descriptor,
+        fitting,
+        type_map: Optional[List[str]],
+        **kwargs,
+    ):
+        torch.nn.Module.__init__(self)
         self.model_def_script = ""
         ntypes = len(type_map)
         self.type_map = type_map
@@ -59,6 +65,8 @@
         self.rcut = self.descriptor.get_rcut()
         self.sel = self.descriptor.get_sel()
         self.fitting_net = fitting
+        # order matters ntypes and type_map should be initialized first.
+        BaseAtomicModel.__init__(self, **kwargs)
 
     def fitting_output_def(self) -> FittingOutputDef:
         """Get the output def of the fitting net."""
@@ -95,22 +103,29 @@
         return self.descriptor.mixed_types()
 
     def serialize(self) -> dict:
-        return {
-            "@class": "Model",
-            "type": "standard",
-            "@version": 1,
-            "type_map": self.type_map,
-            "descriptor": self.descriptor.serialize(),
-            "fitting": self.fitting_net.serialize(),
-        }
+        dd = BaseAtomicModel.serialize(self)
+        dd.update(
+            {
+                "@class": "Model",
+                "@version": 1,
+                "type": "standard",
+                "type_map": self.type_map,
+                "descriptor": self.descriptor.serialize(),
+                "fitting": self.fitting_net.serialize(),
+            }
+        )
+        return dd
 
     @classmethod
     def deserialize(cls, data) -> "DPAtomicModel":
         data = copy.deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 1, 1)
-        descriptor_obj = BaseDescriptor.deserialize(data["descriptor"])
-        fitting_obj = BaseFitting.deserialize(data["fitting"])
-        obj = cls(descriptor_obj, fitting_obj, type_map=data["type_map"])
+        data.pop("@class", None)
+        data.pop("type", None)
+        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 forward_atomic(

deepmd/pt/model/atomic_model/linear_atomic_model.py

@@ -54,10 +54,11 @@
         models: List[BaseAtomicModel],
         **kwargs,
     ):
-        super().__init__()
+        torch.nn.Module.__init__(self)
         self.models = torch.nn.ModuleList(models)
         self.atomic_bias = None
         self.mixed_types_list = [model.mixed_types() for model in self.models]
+        BaseAtomicModel.__init__(self, **kwargs)
 
     def mixed_types(self) -> bool:
         """If true, the model
@@ -307,32 +308,39 @@
         self.zbl_weight = torch.empty(0, dtype=torch.float64, device=env.DEVICE)
 
     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",
+                "@version": 1,
+                "type": "zbl",
+                "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)
-        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"))
 
+        data.pop("@class", None)
+        data.pop("type", None)
         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/pt/model/atomic_model/pairtab_atomic_model.py

@@ -52,11 +52,12 @@
     def __init__(
         self, tab_file: str, rcut: float, sel: Union[int, List[int]], **kwargs
     ):
-        super().__init__()
+        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)
 
         # handle deserialization with no input file
         if self.tab_file is not None:
@@ -125,23 +126,29 @@
         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",
+                "@version": 1,
+                "type": "pairtab",
+                "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)
-        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"))
+        data.pop("@class", None)
+        data.pop("type", None)
+        tab_model = cls(None, rcut, sel, **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))

deepmd/pt/model/descriptor/se_a.py

@@ -133,6 +133,13 @@
         """Update mean and stddev for descriptor elements."""
         return self.sea.compute_input_stats(merged, path)
 
+    def reinit_exclude(
+        self,
+        exclude_types: List[Tuple[int, int]] = [],
+    ):
+        """Update the type exclusions."""
+        self.sea.reinit_exclude(exclude_types)
+
     def forward(
         self,
         coord_ext: torch.Tensor,
@@ -271,10 +278,10 @@
         self.prec = PRECISION_DICT[self.precision]
         self.resnet_dt = resnet_dt
         self.old_impl = old_impl
-        self.exclude_types = exclude_types
         self.ntypes = len(sel)
-        self.emask = PairExcludeMask(len(sel), exclude_types=exclude_types)
         self.type_one_side = type_one_side
+        # order matters, placed after the assignment of self.ntypes
+        self.reinit_exclude(exclude_types)
 
         self.sel = sel
         self.sec = torch.tensor(
@@ -407,6 +414,13 @@
             )
         return self.stats
 
+    def reinit_exclude(
+        self,
+        exclude_types: List[Tuple[int, int]] = [],
+    ):
+        self.exclude_types = exclude_types
+        self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types)
+
     def forward(
         self,
         nlist: torch.Tensor,

deepmd/pt/model/model/make_model.py

@@ -193,7 +193,7 @@
             nframes, nall = extended_atype.shape[:2]
             extended_coord = extended_coord.view(nframes, -1, 3)
             nlist = self.format_nlist(extended_coord, extended_atype, nlist)
-            atomic_ret = self.forward_atomic(
+            atomic_ret = self.forward_common_atomic(
                 extended_coord,
                 extended_atype,
                 nlist,

deepmd/pt/model/task/fitting.py

@@ -288,9 +288,8 @@
         self.precision = precision
         self.prec = PRECISION_DICT[self.precision]
         self.rcond = rcond
-        self.exclude_types = exclude_types
-
-        self.emask = AtomExcludeMask(self.ntypes, self.exclude_types)
+        # order matters, should be place after the assignment of ntypes
+        self.reinit_exclude(exclude_types)
 
         net_dim_out = self._net_out_dim()
         # init constants
@@ -366,6 +365,13 @@
             log.info("Set seed to %d in fitting net.", seed)
             torch.manual_seed(seed)
 
+    def reinit_exclude(
+        self,
+        exclude_types: List[int] = [],
+    ):
+        self.exclude_types = exclude_types
+        self.emask = AtomExcludeMask(self.ntypes, self.exclude_types)
+
     def serialize(self) -> dict:
         """Serialize the fitting to dict."""
         return {