style: require explicit device and dtype

.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/linear_atomic_model.py

@@ -285,7 +285,7 @@
         """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/pairtab_atomic_model.py

@@ -204,7 +204,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 +301,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/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 @@
                 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/utils/nlist.py

@@ -123,7 +123,8 @@ def build_neighbor_list(
         nlist = nlist[:, :, :nsel]
     else:
         rr = np.concatenate(
-            [rr, np.ones([batch_size, nloc, nsel - nnei]) + rcut], axis=-1
+            [rr, np.ones([batch_size, nloc, nsel - nnei]) + rcut],  # pylint: disable=no-explicit-dtype
+            axis=-1,
         )
         nlist = np.concatenate(
             [nlist, np.ones([batch_size, nloc, nsel - nnei], dtype=nlist.dtype)],
@@ -262,7 +263,7 @@ def extend_coord_with_ghosts(
 
     """
     nf, nloc = atype.shape
-    aidx = np.tile(np.arange(nloc)[np.newaxis, :], (nf, 1))
+    aidx = np.tile(np.arange(nloc)[np.newaxis, :], (nf, 1))  # pylint: disable=no-explicit-dtype
     if cell is None:
         nall = nloc
         extend_coord = coord.copy()
@@ -274,9 +275,9 @@ def extend_coord_with_ghosts(
         to_face = to_face_distance(cell)
         nbuff = np.ceil(rcut / to_face).astype(int)
         nbuff = np.max(nbuff, axis=0)
-        xi = np.arange(-nbuff[0], nbuff[0] + 1, 1)
-        yi = np.arange(-nbuff[1], nbuff[1] + 1, 1)
-        zi = np.arange(-nbuff[2], nbuff[2] + 1, 1)
+        xi = np.arange(-nbuff[0], nbuff[0] + 1, 1)  # pylint: disable=no-explicit-dtype
+        yi = np.arange(-nbuff[1], nbuff[1] + 1, 1)  # pylint: disable=no-explicit-dtype
+        zi = np.arange(-nbuff[2], nbuff[2] + 1, 1)  # pylint: disable=no-explicit-dtype
         xyz = np.outer(xi, np.array([1, 0, 0]))[:, np.newaxis, np.newaxis, :]
         xyz = xyz + np.outer(yi, np.array([0, 1, 0]))[np.newaxis, :, np.newaxis, :]
         xyz = xyz + np.outer(zi, np.array([0, 0, 1]))[np.newaxis, np.newaxis, :, :]

deepmd/infer/model_devi.py

@@ -329,7 +329,7 @@ def calc_model_devi(
     forces = np.array(forces)
     virials = np.array(virials)
 
-    devi = [np.arange(coord.shape[0]) * frequency]
+    devi = [np.arange(coord.shape[0]) * frequency]  # pylint: disable=no-explicit-dtype
     if real_data is None:
         devi += list(calc_model_devi_v(virials, relative=relative_v))
         devi_f = list(calc_model_devi_f(forces, relative=relative, atomic=atomic))
@@ -503,7 +503,7 @@ def make_model_devi(
             nframes_tot += coord.shape[0]
             devis.append(devi)
         devis = np.vstack(devis)
-        devis[:, 0] = np.arange(nframes_tot) * frequency
+        devis[:, 0] = np.arange(nframes_tot) * frequency  # pylint: disable=no-explicit-dtype
         write_model_devi_out(devis, output, header=system, atomic=atomic)
         devis_coll.append(devis)
     return devis_coll

deepmd/pt/infer/deep_eval.py

@@ -426,7 +426,9 @@ def _eval_model(
                 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 _eval_model_spin(
@@ -502,7 +504,9 @@ def _eval_model_spin(
                 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):
@@ -666,28 +670,28 @@ def eval_model(
                 logits_out.append(batch_output["logits"])
     if not return_tensor:
         energy_out = (
-            np.concatenate(energy_out) if energy_out else np.zeros([nframes, 1])
+            np.concatenate(energy_out) if energy_out else np.zeros([nframes, 1])  # pylint: disable=no-explicit-dtype
         )
         atomic_energy_out = (
             np.concatenate(atomic_energy_out)
             if atomic_energy_out
-            else np.zeros([nframes, natoms, 1])
+            else np.zeros([nframes, natoms, 1])  # pylint: disable=no-explicit-dtype
         )
         force_out = (
-            np.concatenate(force_out) if force_out else np.zeros([nframes, natoms, 3])
+            np.concatenate(force_out) if force_out else np.zeros([nframes, natoms, 3])  # pylint: disable=no-explicit-dtype
         )
         force_mag_out = (
             np.concatenate(force_mag_out)
             if force_mag_out
-            else np.zeros([nframes, natoms, 3])
+            else np.zeros([nframes, natoms, 3])  # pylint: disable=no-explicit-dtype
         )
         virial_out = (
-            np.concatenate(virial_out) if virial_out else np.zeros([nframes, 3, 3])
+            np.concatenate(virial_out) if virial_out else np.zeros([nframes, 3, 3])  # pylint: disable=no-explicit-dtype
         )
         atomic_virial_out = (
             np.concatenate(atomic_virial_out)
             if atomic_virial_out
-            else np.zeros([nframes, natoms, 3, 3])
+            else np.zeros([nframes, natoms, 3, 3])  # pylint: disable=no-explicit-dtype
         )
         updated_coord_out = (
             np.concatenate(updated_coord_out) if updated_coord_out else None

deepmd/pt/model/atomic_model/linear_atomic_model.py

@@ -90,7 +90,7 @@ def __init__(
         self.rcuts = torch.tensor(
             self.get_model_rcuts(), dtype=torch.float64, device=env.DEVICE
         )
-        self.nsels = torch.tensor(self.get_model_nsels(), device=env.DEVICE)
+        self.nsels = torch.tensor(self.get_model_nsels(), device=env.DEVICE)  # pylint: disable=no-explicit-dtype
 
     def mixed_types(self) -> bool:
         """If true, the model
@@ -285,7 +285,7 @@ def remap_atype(ori_map: List[str], new_map: List[str]) -> torch.Tensor:
         """
         type_2_idx = {atp: idx for idx, atp in enumerate(ori_map)}
         # this maps the atype in the new map to the original map
-        mapping = torch.tensor(
+        mapping = torch.tensor(  # pylint: disable=no-explicit-dtype
             [type_2_idx[new_map[idx]] for idx in range(len(new_map))], device=env.DEVICE
         )
         return mapping

deepmd/pt/model/atomic_model/pairtab_atomic_model.py

@@ -267,7 +267,7 @@ def forward_atomic(
         # i_type : (nframes, nloc), this is atype.
         # j_type : (nframes, nloc, nnei)
         j_type = extended_atype[
-            torch.arange(extended_atype.size(0), device=extended_coord.device)[
+            torch.arange(extended_atype.size(0), device=extended_coord.device)[  # pylint: disable=no-explicit-dtype
                 :, None, None
             ],
             masked_nlist,

deepmd/pt/model/descriptor/descriptor.py

@@ -147,8 +147,8 @@ def share_params(self, base_class, shared_level, resume=False):
                         base_env.stats[kk] += self.get_stats()[kk]
                     mean, stddev = base_env()
                     if not base_class.set_davg_zero:
-                        base_class.mean.copy_(torch.tensor(mean, device=env.DEVICE))
-                    base_class.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
+                        base_class.mean.copy_(torch.tensor(mean, device=env.DEVICE))  # pylint: disable=no-explicit-dtype
+                    base_class.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))  # pylint: disable=no-explicit-dtype
                 # must share, even if not do stat
                 self.mean = base_class.mean
                 self.stddev = base_class.stddev

deepmd/pt/model/descriptor/gaussian_lcc.py

@@ -107,7 +107,7 @@
             sel = [sel]
 
         self.ntypes = ntypes
-        self.sec = torch.tensor(sel)
+        self.sec = torch.tensor(sel)  # pylint: disable=no-explicit-dtype,no-explicit-device
         self.nnei = sum(sel)
 
         if self.do_tag_embedding:
@@ -199,7 +199,7 @@
         nall = extended_coord.shape[1]
         nlist2 = torch.cat(
             [
-                torch.arange(0, nloc, device=nlist.device)
+                torch.arange(0, nloc, device=nlist.device)  # pylint: disable=no-explicit-dtype
                 .reshape(1, nloc, 1)
                 .expand(nframes, -1, -1),
                 nlist,
@@ -208,7 +208,7 @@
         )
         nlist_loc2 = torch.cat(
             [
-                torch.arange(0, nloc, device=nlist_loc.device)
+                torch.arange(0, nloc, device=nlist_loc.device)  # pylint: disable=no-explicit-dtype
                 .reshape(1, nloc, 1)
                 .expand(nframes, -1, -1),
                 nlist_loc,

deepmd/pt/model/descriptor/repformer_layer_old_impl.py

@@ -320,7 +320,7 @@ def __init__(
         sel = [sel] if isinstance(sel, int) else sel
         self.nnei = sum(sel)
         assert len(sel) == 1
-        self.sel = torch.tensor(sel, device=env.DEVICE)
+        self.sel = torch.tensor(sel, device=env.DEVICE)  # pylint: disable=no-explicit-dtype
         self.sec = self.sel
         self.axis_neuron = axis_neuron
         self.set_davg_zero = set_davg_zero

deepmd/pt/model/descriptor/repformers.py

@@ -485,8 +485,8 @@
                     comm_dict["recv_num"],
                     g1,
                     comm_dict["communicator"],
-                    torch.tensor(nloc),
-                    torch.tensor(nall - nloc),
+                    torch.tensor(nloc),  # pylint: disable=no-explicit-dtype,no-explicit-device
+                    torch.tensor(nall - nloc),  # pylint: disable=no-explicit-dtype,no-explicit-device
                 )
                 g1_ext = ret[0].unsqueeze(0)
             g1, g2, h2 = ll.forward(
@@ -543,8 +543,8 @@
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:
-            self.mean.copy_(torch.tensor(mean, device=env.DEVICE))
-        self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
+            self.mean.copy_(torch.tensor(mean, device=env.DEVICE))  # pylint: disable=no-explicit-dtype
+        self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))  # pylint: disable=no-explicit-dtype
 
     def get_stats(self) -> Dict[str, StatItem]:
         """Get the statistics of the descriptor."""

deepmd/pt/model/descriptor/se_a.py

@@ -578,8 +578,8 @@
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:
-            self.mean.copy_(torch.tensor(mean, device=env.DEVICE))
-        self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
+            self.mean.copy_(torch.tensor(mean, device=env.DEVICE))  # pylint: disable=no-explicit-dtype
+        self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))  # pylint: disable=no-explicit-dtype
 
     def get_stats(self) -> Dict[str, StatItem]:
         """Get the statistics of the descriptor."""
@@ -635,7 +635,7 @@
             dmatrix = dmatrix.view(
                 -1, self.ndescrpt
             )  # shape is [nframes*nall, self.ndescrpt]
-            xyz_scatter = torch.empty(
+            xyz_scatter = torch.empty(  # pylint: disable=no-explicit-dtype
                 1,
                 device=env.DEVICE,
             )

deepmd/pt/model/descriptor/se_atten.py

@@ -407,8 +407,8 @@
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:
-            self.mean.copy_(torch.tensor(mean, device=env.DEVICE))
-        self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
+            self.mean.copy_(torch.tensor(mean, device=env.DEVICE))  # pylint: disable=no-explicit-dtype
+        self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))  # pylint: disable=no-explicit-dtype
 
     def get_stats(self) -> Dict[str, StatItem]:
         """Get the statistics of the descriptor."""

deepmd/pt/model/descriptor/se_r.py

@@ -208,8 +208,8 @@
                     base_env.stats[kk] += self.get_stats()[kk]
                 mean, stddev = base_env()
                 if not base_class.set_davg_zero:
-                    base_class.mean.copy_(torch.tensor(mean, device=env.DEVICE))
-                base_class.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
+                    base_class.mean.copy_(torch.tensor(mean, device=env.DEVICE))  # pylint: disable=no-explicit-dtype
+                base_class.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))  # pylint: disable=no-explicit-dtype
                 self.mean = base_class.mean
                 self.stddev = base_class.stddev
             # self.load_state_dict(base_class.state_dict()) # this does not work, because it only inits the model
@@ -268,8 +268,8 @@
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:
-            self.mean.copy_(torch.tensor(mean, device=env.DEVICE))
-        self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
+            self.mean.copy_(torch.tensor(mean, device=env.DEVICE))  # pylint: disable=no-explicit-dtype
+        self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))  # pylint: disable=no-explicit-dtype
 
     def get_stats(self) -> Dict[str, StatItem]:
         """Get the statistics of the descriptor."""

deepmd/pt/model/descriptor/se_t.py

@@ -605,8 +605,8 @@
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:
-            self.mean.copy_(torch.tensor(mean, device=env.DEVICE))
-        self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
+            self.mean.copy_(torch.tensor(mean, device=env.DEVICE))  # pylint: disable=no-explicit-dtype
+        self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))  # pylint: disable=no-explicit-dtype
 
     def get_stats(self) -> Dict[str, StatItem]:
         """Get the statistics of the descriptor."""

deepmd/pt/model/model/__init__.py

@@ -71,7 +71,7 @@ def get_spin_model(model_params):
     if not model_params["spin"]["use_spin"] or isinstance(
         model_params["spin"]["use_spin"][0], int
     ):
-        use_spin = np.full(len(model_params["type_map"]), False)
+        use_spin = np.full(len(model_params["type_map"]), False)  # pylint: disable=no-explicit-dtype
         use_spin[model_params["spin"]["use_spin"]] = True
         model_params["spin"]["use_spin"] = use_spin.tolist()
     # include virtual spin and placeholder types