Merge branch 'devel' into virial_test

.pre-commit-config.yaml

@@ -144,5 +144,12 @@ repos:
         entry: DeepMD|DeepMd|Pytorch|Tensorflow|Numpy|Github|Lammps|I-Pi|I-PI|i-Pi
         # unclear why PairDeepMD is used instead of PairDeePMD
         exclude: .pre-commit-config.yaml|source/lmp
+  # customized pylint rules
+  - repo: https://github.com/pylint-dev/pylint/
+    rev: v3.2.6
+    hooks:
+      - id: pylint
+        entry: env PYTHONPATH=source/checker pylint
+        files: ^deepmd/
 ci:
   autoupdate_branch: devel

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -57,8 +57,8 @@ def init_out_stat(self):
             [self.atomic_output_def()[kk].size for kk in self.bias_keys]
         )
         self.n_out = len(self.bias_keys)
-        out_bias_data = np.zeros([self.n_out, ntypes, self.max_out_size])
-        out_std_data = np.ones([self.n_out, ntypes, self.max_out_size])
+        out_bias_data = np.zeros([self.n_out, ntypes, self.max_out_size])  # pylint: disable=no-explicit-dtype
+        out_std_data = np.ones([self.n_out, ntypes, self.max_out_size])  # pylint: disable=no-explicit-dtype
         self.out_bias = out_bias_data
         self.out_std = out_std_data
 

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -84,6 +84,10 @@ def has_message_passing(self) -> bool:
         """Returns whether the atomic model has message passing."""
         return self.descriptor.has_message_passing()
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the atomic model needs sorted nlist when using `forward_lower`."""
+        return self.descriptor.need_sorted_nlist_for_lower()
+
     def forward_atomic(
         self,
         extended_coord: np.ndarray,

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -96,6 +96,10 @@ def has_message_passing(self) -> bool:
         """Returns whether the atomic model has message passing."""
         return any(model.has_message_passing() for model in self.models)
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the atomic model needs sorted nlist when using `forward_lower`."""
+        return True
+
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
         return max(self.get_model_rcuts())
@@ -285,7 +289,7 @@ def _compute_weight(
         """This should be a list of user defined weights that matches the number of models to be combined."""
         nmodels = len(self.models)
         nframes, nloc, _ = nlists_[0].shape
-        return [np.ones((nframes, nloc, 1)) / nmodels for _ in range(nmodels)]
+        return [np.ones((nframes, nloc, 1)) / nmodels for _ in range(nmodels)]  # pylint: disable=no-explicit-dtype
 
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""

deepmd/dpmodel/atomic_model/make_base_atomic_model.py

@@ -119,6 +119,10 @@ def mixed_types(self) -> bool:
         def has_message_passing(self) -> bool:
             """Returns whether the descriptor has message passing."""
 
+        @abstractmethod
+        def need_sorted_nlist_for_lower(self) -> bool:
+            """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+
         @abstractmethod
         def fwd(
             self,

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -135,6 +135,10 @@ def has_message_passing(self) -> bool:
         """Returns whether the atomic model has message passing."""
         return False
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the atomic model needs sorted nlist when using `forward_lower`."""
+        return False
+
     def change_type_map(
         self, type_map: List[str], model_with_new_type_stat=None
     ) -> None:
@@ -204,7 +208,7 @@ def forward_atomic(
 
         # (nframes, nloc, nnei)
         j_type = extended_atype[
-            np.arange(extended_atype.shape[0])[:, None, None], masked_nlist
+            np.arange(extended_atype.shape[0])[:, None, None], masked_nlist  # pylint: disable=no-explicit-dtype
         ]
 
         raw_atomic_energy = self._pair_tabulated_inter(
@@ -301,7 +305,7 @@ def _get_pairwise_dist(coords: np.ndarray, nlist: np.ndarray) -> np.ndarray:
         np.ndarray
             The pairwise distance between the atoms (nframes, nloc, nnei).
         """
-        batch_indices = np.arange(nlist.shape[0])[:, None, None]
+        batch_indices = np.arange(nlist.shape[0])[:, None, None]  # pylint: disable=no-explicit-dtype
         neighbor_atoms = coords[batch_indices, nlist]
         loc_atoms = coords[:, : nlist.shape[1], :]
         pairwise_dr = loc_atoms[:, :, None, :] - neighbor_atoms

deepmd/dpmodel/descriptor/descriptor.py

@@ -132,6 +132,10 @@ def call(
     def has_message_passing(self) -> bool:
         """Returns whether the descriptor block has message passing."""
 
+    @abstractmethod
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor block needs sorted nlist when using `forward_lower`."""
+
 
 def extend_descrpt_stat(des, type_map, des_with_stat=None):
     r"""

deepmd/dpmodel/descriptor/dpa1.py

@@ -371,6 +371,10 @@ def has_message_passing(self) -> bool:
         """Returns whether the descriptor has message passing."""
         return self.se_atten.has_message_passing()
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+        return self.se_atten.need_sorted_nlist_for_lower()
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.se_atten.get_env_protection()
@@ -956,6 +960,10 @@ def has_message_passing(self) -> bool:
         """Returns whether the descriptor block has message passing."""
         return False
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor block needs sorted nlist when using `forward_lower`."""
+        return False
+
 
 class NeighborGatedAttention(NativeOP):
     def __init__(

deepmd/dpmodel/descriptor/dpa2.py

@@ -553,6 +553,10 @@ def has_message_passing(self) -> bool:
             [self.repinit.has_message_passing(), self.repformers.has_message_passing()]
         )
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+        return True
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.env_protection

deepmd/dpmodel/descriptor/hybrid.py

@@ -146,6 +146,10 @@ def has_message_passing(self) -> bool:
         """Returns whether the descriptor has message passing."""
         return any(descrpt.has_message_passing() for descrpt in self.descrpt_list)
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+        return True
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix. All descriptors should be the same."""
         all_protection = [descrpt.get_env_protection() for descrpt in self.descrpt_list]

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -104,6 +104,10 @@ def mixed_types(self) -> bool:
         def has_message_passing(self) -> bool:
             """Returns whether the descriptor has message passing."""
 
+        @abstractmethod
+        def need_sorted_nlist_for_lower(self) -> bool:
+            """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+
         @abstractmethod
         def get_env_protection(self) -> float:
             """Returns the protection of building environment matrix."""

deepmd/dpmodel/descriptor/repformers.py

@@ -401,6 +401,10 @@ def has_message_passing(self) -> bool:
         """Returns whether the descriptor block has message passing."""
         return True
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor block needs sorted nlist when using `forward_lower`."""
+        return False
+
 
 # translated by GPT and modified
 def get_residual(

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -265,6 +265,10 @@ def has_message_passing(self) -> bool:
         """Returns whether the descriptor has message passing."""
         return False
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+        return False
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.env_protection

deepmd/dpmodel/descriptor/se_r.py

@@ -223,6 +223,10 @@ def has_message_passing(self) -> bool:
         """Returns whether the descriptor has message passing."""
         return False
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+        return False
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.env_protection

deepmd/dpmodel/descriptor/se_t.py

@@ -215,6 +215,10 @@ def has_message_passing(self) -> bool:
         """Returns whether the descriptor has message passing."""
         return False
 
+    def need_sorted_nlist_for_lower(self) -> bool:
+        """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""
+        return False
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.env_protection

deepmd/dpmodel/fitting/general_fitting.py

@@ -144,18 +144,18 @@ def __init__(
         net_dim_out = self._net_out_dim()
         # init constants
         if bias_atom_e is None:
-            self.bias_atom_e = np.zeros([self.ntypes, net_dim_out])
+            self.bias_atom_e = np.zeros([self.ntypes, net_dim_out])  # pylint: disable=no-explicit-dtype
         else:
             assert bias_atom_e.shape == (self.ntypes, net_dim_out)
             self.bias_atom_e = bias_atom_e
         if self.numb_fparam > 0:
-            self.fparam_avg = np.zeros(self.numb_fparam)
-            self.fparam_inv_std = np.ones(self.numb_fparam)
+            self.fparam_avg = np.zeros(self.numb_fparam)  # pylint: disable=no-explicit-dtype
+            self.fparam_inv_std = np.ones(self.numb_fparam)  # pylint: disable=no-explicit-dtype
         else:
             self.fparam_avg, self.fparam_inv_std = None, None
         if self.numb_aparam > 0:
-            self.aparam_avg = np.zeros(self.numb_aparam)
-            self.aparam_inv_std = np.ones(self.numb_aparam)
+            self.aparam_avg = np.zeros(self.numb_aparam)  # pylint: disable=no-explicit-dtype
+            self.aparam_inv_std = np.ones(self.numb_aparam)  # pylint: disable=no-explicit-dtype
         else:
             self.aparam_avg, self.aparam_inv_std = None, None
         # init networks
@@ -405,7 +405,7 @@ def _call_common(
 
         # calcualte the prediction
         if not self.mixed_types:
-            outs = np.zeros([nf, nloc, net_dim_out])
+            outs = np.zeros([nf, nloc, net_dim_out])  # pylint: disable=no-explicit-dtype
             for type_i in range(self.ntypes):
                 mask = np.tile(
                     (atype == type_i).reshape([nf, nloc, 1]), [1, 1, net_dim_out]

deepmd/dpmodel/fitting/polarizability_fitting.py

@@ -308,7 +308,7 @@ def call(
             bias = self.constant_matrix[atype]
             # (nframes, nloc, 1)
             bias = np.expand_dims(bias, axis=-1) * self.scale[atype]
-            eye = np.eye(3)
+            eye = np.eye(3)  # pylint: disable=no-explicit-dtype
             eye = np.tile(eye, (nframes, nloc, 1, 1))
             # (nframes, nloc, 3, 3)
             bias = np.expand_dims(bias, axis=-1) * eye

deepmd/dpmodel/infer/deep_eval.py

@@ -343,11 +343,13 @@ def _eval_model(
                 if batch_output[dp_name] is not None:
                     out = batch_output[dp_name].reshape(shape)
                 else:
-                    out = np.full(shape, np.nan)
+                    out = np.full(shape, np.nan)  # pylint: disable=no-explicit-dtype
                 results.append(out)
             else:
                 shape = self._get_output_shape(odef, nframes, natoms)
-                results.append(np.full(np.abs(shape), np.nan))  # this is kinda hacky
+                results.append(
+                    np.full(np.abs(shape), np.nan)  # pylint: disable=no-explicit-dtype
+                )  # this is kinda hacky
         return tuple(results)
 
     def _get_output_shape(self, odef, nframes, natoms):

deepmd/dpmodel/model/base_model.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import inspect
+import json
 from abc import (
     ABC,
     abstractmethod,
@@ -193,6 +194,30 @@ def update_sel(
             cls = cls.get_class_by_type(model_type)
             return cls.update_sel(train_data, type_map, local_jdata)
 
+        @classmethod
+        def get_model(cls, model_params: dict) -> "BaseBaseModel":
+            """Get the model by the parameters.
+
+            By default, all the parameters are directly passed to the constructor.
+            If not, override this method.
+
+            Parameters
+            ----------
+            model_params : dict
+                The model parameters
+
+            Returns
+            -------
+            BaseBaseModel
+                The model
+            """
+            model_params_old = model_params.copy()
+            model_params = model_params.copy()
+            model_params.pop("type", None)
+            model = cls(**model_params)
+            model.model_def_script = json.dumps(model_params_old)
+            return model
+
     return BaseBaseModel
 
 

deepmd/dpmodel/model/make_model.py

@@ -214,7 +214,12 @@ def call_lower(
             """
             nframes, nall = extended_atype.shape[:2]
             extended_coord = extended_coord.reshape(nframes, -1, 3)
-            nlist = self.format_nlist(extended_coord, extended_atype, nlist)
+            nlist = self.format_nlist(
+                extended_coord,
+                extended_atype,
+                nlist,
+                extra_nlist_sort=self.need_sorted_nlist_for_lower(),
+            )
             cc_ext, _, fp, ap, input_prec = self.input_type_cast(
                 extended_coord, fparam=fparam, aparam=aparam
             )
@@ -311,6 +316,7 @@ def format_nlist(
             extended_coord: np.ndarray,
             extended_atype: np.ndarray,
             nlist: np.ndarray,
+            extra_nlist_sort: bool = False,
         ):
             """Format the neighbor list.
 
@@ -336,6 +342,8 @@ def format_nlist(
                 atomic type in extended region. nf x nall
             nlist
                 neighbor list. nf x nloc x nsel
+            extra_nlist_sort
+                whether to forcibly sort the nlist.
 
             Returns
             -------
@@ -345,7 +353,12 @@ def format_nlist(
             """
             n_nf, n_nloc, n_nnei = nlist.shape
             mixed_types = self.mixed_types()
-            ret = self._format_nlist(extended_coord, nlist, sum(self.get_sel()))
+            ret = self._format_nlist(
+                extended_coord,
+                nlist,
+                sum(self.get_sel()),
+                extra_nlist_sort=extra_nlist_sort,
+            )
             if not mixed_types:
                 ret = nlist_distinguish_types(ret, extended_atype, self.get_sel())
             return ret
@@ -355,6 +368,7 @@ def _format_nlist(
             extended_coord: np.ndarray,
             nlist: np.ndarray,
             nnei: int,
+            extra_nlist_sort: bool = False,
         ):
             n_nf, n_nloc, n_nnei = nlist.shape
             extended_coord = extended_coord.reshape([n_nf, -1, 3])
@@ -370,7 +384,9 @@ def _format_nlist(
                     ],
                     axis=-1,
                 )
-            elif n_nnei > nnei:
+
+            if n_nnei > nnei or extra_nlist_sort:
+                n_nf, n_nloc, n_nnei = nlist.shape
                 # make a copy before revise
                 m_real_nei = nlist >= 0
                 ret = np.where(m_real_nei, nlist, 0)
@@ -384,9 +400,11 @@ def _format_nlist(
                 ret = np.take_along_axis(ret, ret_mapping, axis=2)
                 ret = np.where(rr > rcut, -1, ret)
                 ret = ret[..., :nnei]
-            else:  # n_nnei == nnei:
-                # copy anyway...
+            # not extra_nlist_sort and n_nnei <= nnei:
+            elif n_nnei == nnei:
                 ret = nlist
+            else:
+                pass
             assert ret.shape[-1] == nnei
             return ret
 
@@ -483,6 +501,10 @@ def has_message_passing(self) -> bool:
             """Returns whether the model has message passing."""
             return self.atomic_model.has_message_passing()
 
+        def need_sorted_nlist_for_lower(self) -> bool:
+            """Returns whether the model needs sorted nlist when using `forward_lower`."""
+            return self.atomic_model.need_sorted_nlist_for_lower()
+
         def atomic_output_def(self) -> FittingOutputDef:
             """Get the output def of the atomic model."""
             return self.atomic_model.atomic_output_def()