pt: add exported methods to BaseAtomicModel (#3258)

Signed-off-by: Jinzhe Zeng <[email protected]>

deepmd/dpmodel/model/dp_atomic_model.py

@@ -143,3 +143,27 @@ def deserialize(cls, data) -> "DPAtomicModel":
         )
         obj = cls(descriptor_obj, fitting_obj, type_map=data["type_map"])
         return obj
+
+    def get_dim_fparam(self) -> int:
+        """Get the number (dimension) of frame parameters of this atomic model."""
+        return 0
+
+    def get_dim_aparam(self) -> int:
+        """Get the number (dimension) of atomic parameters of this atomic model."""
+        return 0
+
+    def get_sel_type(self) -> List[int]:
+        """Get the selected atom types of this model.
+
+        Only atoms with selected atom types have atomic contribution
+        to the result of the model.
+        If returning an empty list, all atom types are selected.
+        """
+        return []
+
+    def is_aparam_nall(self) -> bool:
+        """Check whether the shape of atomic parameters is (nframes, nall, ndim).
+
+        If False, the shape is (nframes, nloc, ndim).
+        """
+        return False

deepmd/dpmodel/model/linear_atomic_model.py

@@ -199,6 +199,34 @@ def _compute_weight(
         """This should be a list of user defined weights that matches the number of models to be combined."""
         raise NotImplementedError
 
+    def get_dim_fparam(self) -> int:
+        """Get the number (dimension) of frame parameters of this atomic model."""
+        # tricky...
+        return max([model.get_dim_fparam() for model in self.models])
+
+    def get_dim_aparam(self) -> int:
+        """Get the number (dimension) of atomic parameters of this atomic model."""
+        return max([model.get_dim_aparam() for model in self.models])
+
+    def get_sel_type(self) -> List[int]:
+        """Get the selected atom types of this model.
+
+        Only atoms with selected atom types have atomic contribution
+        to the result of the model.
+        If returning an empty list, all atom types are selected.
+        """
+        if any(model.get_sel_type() == [] for model in self.models):
+            return []
+        # join all the selected types
+        return list(set().union(*[model.get_sel_type() for model in self.models]))
+
+    def is_aparam_nall(self) -> bool:
+        """Check whether the shape of atomic parameters is (nframes, nall, ndim).
+
+        If False, the shape is (nframes, nloc, ndim).
+        """
+        return False
+
 
 class DPZBLLinearAtomicModel(LinearAtomicModel):
     """Model linearly combine a list of AtomicModels.

deepmd/dpmodel/model/make_base_atomic_model.py

@@ -61,6 +61,30 @@ def get_nnei(self) -> int:
             """Returns the total number of selected neighboring atoms in the cut-off radius."""
             return self.get_nsel()
 
+        @abstractmethod
+        def get_dim_fparam(self) -> int:
+            """Get the number (dimension) of frame parameters of this atomic model."""
+
+        @abstractmethod
+        def get_dim_aparam(self) -> int:
+            """Get the number (dimension) of atomic parameters of this atomic model."""
+
+        @abstractmethod
+        def get_sel_type(self) -> List[int]:
+            """Get the selected atom types of this model.
+
+            Only atoms with selected atom types have atomic contribution
+            to the result of the model.
+            If returning an empty list, all atom types are selected.
+            """
+
+        @abstractmethod
+        def is_aparam_nall(self) -> bool:
+            """Check whether the shape of atomic parameters is (nframes, nall, ndim).
+
+            If False, the shape is (nframes, nloc, ndim).
+            """
+
         @abstractmethod
         def distinguish_types(self) -> bool:
             """Returns if the model requires a neighbor list that distinguish different

deepmd/dpmodel/model/pairtab_atomic_model.py

@@ -312,3 +312,27 @@ def _calculate_ener(coef: np.ndarray, uu: np.ndarray) -> np.ndarray:
         etmp = (a3 * uu + a2) * uu + a1  # this should be elementwise operations.
         ener = etmp * uu + a0  # this energy has the extrapolated value when rcut > rmax
         return ener
+
+    def get_dim_fparam(self) -> int:
+        """Get the number (dimension) of frame parameters of this atomic model."""
+        return 0
+
+    def get_dim_aparam(self) -> int:
+        """Get the number (dimension) of atomic parameters of this atomic model."""
+        return 0
+
+    def get_sel_type(self) -> List[int]:
+        """Get the selected atom types of this model.
+
+        Only atoms with selected atom types have atomic contribution
+        to the result of the model.
+        If returning an empty list, all atom types are selected.
+        """
+        return []
+
+    def is_aparam_nall(self) -> bool:
+        """Check whether the shape of atomic parameters is (nframes, nall, ndim).
+
+        If False, the shape is (nframes, nloc, ndim).
+        """
+        return False

deepmd/pt/infer/deep_eval.py

@@ -136,11 +136,11 @@ def get_type_map(self) -> List[str]:
 
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this DP."""
-        return 0
+        return self.dp.model["Default"].get_dim_fparam()
 
     def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this DP."""
-        return 0
+        return self.dp.model["Default"].get_dim_aparam()
 
     @property
     def model_type(self) -> "DeepEvalWrapper":
@@ -169,7 +169,7 @@ def get_sel_type(self) -> List[int]:
         to the result of the model.
         If returning an empty list, all atom types are selected.
         """
-        return []
+        return self.dp.model["Default"].get_sel_type()
 
     def get_numb_dos(self) -> int:
         """Get the number of DOS."""

deepmd/pt/model/model/base_atomic_model.py

@@ -1,9 +1,16 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
+
 import torch
 
 from deepmd.dpmodel.model import (
     make_base_atomic_model,
 )
 
-BaseAtomicModel = make_base_atomic_model(torch.Tensor)
+BaseAtomicModel_ = make_base_atomic_model(torch.Tensor)
+
+
+class BaseAtomicModel(BaseAtomicModel_):
+    # export public methods that are not abstract
+    get_nsel = torch.jit.export(BaseAtomicModel_.get_nsel)
+    get_nnei = torch.jit.export(BaseAtomicModel_.get_nnei)

deepmd/pt/model/model/dp_atomic_model.py

@@ -202,3 +202,34 @@ def compute_or_load_stat(
             self.fitting_net.compute_or_load_stat(
                 type_map, sampled, stat_file_path_dict["fitting_net"]
             )
+
+    @torch.jit.export
+    def get_dim_fparam(self) -> int:
+        """Get the number (dimension) of frame parameters of this atomic model."""
+        # TODO: self.fitting_net.get_dim_fparam()
+        return 0
+
+    @torch.jit.export
+    def get_dim_aparam(self) -> int:
+        """Get the number (dimension) of atomic parameters of this atomic model."""
+        # TODO: self.fitting_net.get_dim_aparam()
+        return 0
+
+    @torch.jit.export
+    def get_sel_type(self) -> List[int]:
+        """Get the selected atom types of this model.
+
+        Only atoms with selected atom types have atomic contribution
+        to the result of the model.
+        If returning an empty list, all atom types are selected.
+        """
+        # TODO: self.fitting_net.get_sel_type()
+        return []
+
+    @torch.jit.export
+    def is_aparam_nall(self) -> bool:
+        """Check whether the shape of atomic parameters is (nframes, nall, ndim).
+
+        If False, the shape is (nframes, nloc, ndim).
+        """
+        return False

deepmd/pt/model/model/linear_atomic_model.py

@@ -213,6 +213,46 @@ def _compute_weight(
         """This should be a list of user defined weights that matches the number of models to be combined."""
         raise NotImplementedError
 
+    @torch.jit.export
+    def get_dim_fparam(self) -> int:
+        """Get the number (dimension) of frame parameters of this atomic model."""
+        # tricky...
+        return max([model.get_dim_fparam() for model in self.models])
+
+    @torch.jit.export
+    def get_dim_aparam(self) -> int:
+        """Get the number (dimension) of atomic parameters of this atomic model."""
+        return max([model.get_dim_aparam() for model in self.models])
+
+    @torch.jit.export
+    def get_sel_type(self) -> List[int]:
+        """Get the selected atom types of this model.
+
+        Only atoms with selected atom types have atomic contribution
+        to the result of the model.
+        If returning an empty list, all atom types are selected.
+        """
+        if any(model.get_sel_type() == [] for model in self.models):
+            return []
+        # join all the selected types
+        # make torch.jit happy...
+        return torch.unique(
+            torch.cat(
+                [
+                    torch.as_tensor(model.get_sel_type(), dtype=torch.int32)
+                    for model in self.models
+                ]
+            )
+        ).tolist()
+
+    @torch.jit.export
+    def is_aparam_nall(self) -> bool:
+        """Check whether the shape of atomic parameters is (nframes, nall, ndim).
+
+        If False, the shape is (nframes, nloc, ndim).
+        """
+        return False
+
 
 class DPZBLLinearAtomicModel(LinearAtomicModel):
     """Model linearly combine a list of AtomicModels.

deepmd/pt/model/model/pairtab_atomic_model.py

@@ -337,3 +337,31 @@ def _calculate_ener(coef: torch.Tensor, uu: torch.Tensor) -> torch.Tensor:
         etmp = (a3 * uu + a2) * uu + a1  # this should be elementwise operations.
         ener = etmp * uu + a0  # this energy has the extrapolated value when rcut > rmax
         return ener
+
+    @torch.jit.export
+    def get_dim_fparam(self) -> int:
+        """Get the number (dimension) of frame parameters of this atomic model."""
+        return 0
+
+    @torch.jit.export
+    def get_dim_aparam(self) -> int:
+        """Get the number (dimension) of atomic parameters of this atomic model."""
+        return 0
+
+    @torch.jit.export
+    def get_sel_type(self) -> List[int]:
+        """Get the selected atom types of this model.
+
+        Only atoms with selected atom types have atomic contribution
+        to the result of the model.
+        If returning an empty list, all atom types are selected.
+        """
+        return []
+
+    @torch.jit.export
+    def is_aparam_nall(self) -> bool:
+        """Check whether the shape of atomic parameters is (nframes, nall, ndim).
+
+        If False, the shape is (nframes, nloc, ndim).
+        """
+        return False