From 0c7d0050d2596d641a064c9bbefef9730faaa0a3 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 21 Oct 2024 22:52:19 +0800 Subject: [PATCH 01/19] fix fitting --- deepmd/dpmodel/fitting/general_fitting.py | 14 ++++++++------ deepmd/dpmodel/fitting/polarizability_fitting.py | 2 +- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index fd80ccb4aa..75e622b714 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -14,6 +14,7 @@ from deepmd.dpmodel import ( DEFAULT_PRECISION, + PRECISION_DICT, NativeOP, ) from deepmd.dpmodel.common import ( @@ -133,6 +134,7 @@ def __init__( self.trainable = [self.trainable] * (len(self.neuron) + 1) self.activation_function = activation_function self.precision = precision + prec = PRECISION_DICT[self.precision.lower()] self.layer_name = layer_name self.use_aparam_as_mask = use_aparam_as_mask self.spin = spin @@ -146,18 +148,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]) # pylint: disable=no-explicit-dtype + self.bias_atom_e = np.zeros([self.ntypes, net_dim_out], dtype=prec) 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) # pylint: disable=no-explicit-dtype - self.fparam_inv_std = np.ones(self.numb_fparam) # pylint: disable=no-explicit-dtype + self.fparam_avg = np.zeros(self.numb_fparam, dtype=prec) + self.fparam_inv_std = np.ones(self.numb_fparam, dtype=prec) else: self.fparam_avg, self.fparam_inv_std = None, None if self.numb_aparam > 0: - 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 + self.aparam_avg = np.zeros(self.numb_aparam, dtype=prec) + self.aparam_inv_std = np.ones(self.numb_aparam, dtype=prec) else: self.aparam_avg, self.aparam_inv_std = None, None # init networks @@ -410,7 +412,7 @@ def _call_common( # calcualte the prediction if not self.mixed_types: - outs = xp.zeros([nf, nloc, net_dim_out]) # pylint: disable=no-explicit-dtype + outs = xp.zeros([nf, nloc, net_dim_out], dtype=descriptor.dtype) for type_i in range(self.ntypes): mask = xp.tile( xp.reshape((atype == type_i), [nf, nloc, 1]), (1, 1, net_dim_out) diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py index 73a691f482..2d96eec580 100644 --- a/deepmd/dpmodel/fitting/polarizability_fitting.py +++ b/deepmd/dpmodel/fitting/polarizability_fitting.py @@ -303,7 +303,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) # pylint: disable=no-explicit-dtype + eye = np.eye(3, dtype=descriptor.dtype) eye = np.tile(eye, (nframes, nloc, 1, 1)) # (nframes, nloc, 3, 3) bias = np.expand_dims(bias, axis=-1) * eye From 0b50629aaf4077efaaba0b2a973611bcce2bddf0 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 22 Oct 2024 22:58:04 +0800 Subject: [PATCH 02/19] fix nlist --- deepmd/dpmodel/utils/nlist.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index c56f1bc061..1801c155cd 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -131,7 +131,7 @@ def build_neighbor_list( nlist = nlist[:, :, :nsel] else: rr = xp.concatenate( - [rr, xp.ones([batch_size, nloc, nsel - nnei]) + rcut], # pylint: disable=no-explicit-dtype + [rr, xp.ones([batch_size, nloc, nsel - nnei], dtype=rr.dtype) + rcut], axis=-1, ) nlist = xp.concatenate( @@ -277,7 +277,7 @@ def extend_coord_with_ghosts( """ xp = array_api_compat.array_namespace(coord, atype) nf, nloc = atype.shape - aidx = xp.tile(xp.arange(nloc)[xp.newaxis, :], (nf, 1)) # pylint: disable=no-explicit-dtype + aidx = xp.tile(xp.arange(nloc, dtype=xp.int64)[xp.newaxis, :], (nf, 1)) if cell is None: nall = nloc extend_coord = coord @@ -289,9 +289,9 @@ def extend_coord_with_ghosts( to_face = to_face_distance(cell) nbuff = xp.astype(xp.ceil(rcut / to_face), xp.int64) nbuff = xp.max(nbuff, axis=0) - xi = xp.arange(-int(nbuff[0]), int(nbuff[0]) + 1, 1) # pylint: disable=no-explicit-dtype - yi = xp.arange(-int(nbuff[1]), int(nbuff[1]) + 1, 1) # pylint: disable=no-explicit-dtype - zi = xp.arange(-int(nbuff[2]), int(nbuff[2]) + 1, 1) # pylint: disable=no-explicit-dtype + xi = xp.arange(-int(nbuff[0]), int(nbuff[0]) + 1, 1, dtype=nbuff.dtype) + yi = xp.arange(-int(nbuff[1]), int(nbuff[1]) + 1, 1, dtype=nbuff.dtype) + zi = xp.arange(-int(nbuff[2]), int(nbuff[2]) + 1, 1, dtype=nbuff.dtype) xyz = xp.linalg.outer(xi, xp.asarray([1, 0, 0]))[:, xp.newaxis, xp.newaxis, :] xyz = ( xyz From 9c64ec902c9d43b557c133527daa24af6b827617 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 22 Oct 2024 23:09:21 +0800 Subject: [PATCH 03/19] fix unittest.mock for python3 --- source/tests/common/test_auto_batch_size.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/source/tests/common/test_auto_batch_size.py b/source/tests/common/test_auto_batch_size.py index cc1e6bf25a..9cb86d4923 100644 --- a/source/tests/common/test_auto_batch_size.py +++ b/source/tests/common/test_auto_batch_size.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import os import unittest +from unittest.mock import ( + patch, +) import array_api_strict as xp @@ -77,7 +80,7 @@ def test_execute_oom_cpu(self): self.assertEqual(nb, 128) self.assertEqual(result.shape, (128, 2)) - @unittest.mock.patch.dict(os.environ, {"DP_INFER_BATCH_SIZE": "256"}, clear=True) + @patch.dict(os.environ, {"DP_INFER_BATCH_SIZE": "256"}, clear=True) def test_execute_oom_environment_variables(self): # DP_INFER_BATCH_SIZE = 256 = 128 * 2, nb is always 128 auto_batch_size = CustomizedAutoBatchSizeGPU(999, 2.0) From 29ac75426b97f4667b0aa5b12eecb825c88153ce Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 22 Oct 2024 23:10:35 +0800 Subject: [PATCH 04/19] fix dpmodel.atomic_models --- deepmd/dpmodel/atomic_model/base_atomic_model.py | 11 +++++++++-- deepmd/dpmodel/atomic_model/linear_atomic_model.py | 9 ++++++++- deepmd/dpmodel/atomic_model/pairtab_atomic_model.py | 10 +++++++--- 3 files changed, 24 insertions(+), 6 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index c29a76b3f1..9aa7ef9bcf 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -17,6 +17,9 @@ AtomExcludeMask, PairExcludeMask, ) +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) from deepmd.utils.finetune import ( get_index_between_two_maps, map_atom_exclude_types, @@ -54,8 +57,12 @@ 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]) # pylint: disable=no-explicit-dtype - out_std_data = np.ones([self.n_out, ntypes, self.max_out_size]) # pylint: disable=no-explicit-dtype + out_bias_data = np.zeros( + [self.n_out, ntypes, self.max_out_size], dtype=GLOBAL_NP_FLOAT_PRECISION + ) + out_std_data = np.ones( + [self.n_out, ntypes, self.max_out_size], dtype=GLOBAL_NP_FLOAT_PRECISION + ) self.out_bias = out_bias_data self.out_std = out_std_data diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 79a51635d2..880c92f504 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -12,6 +12,9 @@ get_multiple_nlist_key, nlist_distinguish_types, ) +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -286,7 +289,11 @@ 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)] # pylint: disable=no-explicit-dtype + # the dtype of weights is the interface data type. + return [ + np.ones((nframes, nloc, 1), dtype=GLOBAL_NP_FLOAT_PRECISION) / nmodels + for _ in range(nmodels) + ] def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index 22471d3f32..ee9c7a9a76 100644 --- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py @@ -204,9 +204,12 @@ def forward_atomic( self.tab.ntypes, self.tab.ntypes, self.tab.nspline, 4 ) - # (nframes, nloc, nnei) + # (nframes, nloc, nnei), dtype is the same as atype. j_type = extended_atype[ - np.arange(extended_atype.shape[0])[:, None, None], masked_nlist # pylint: disable=no-explicit-dtype + np.arange(extended_atype.shape[0], dtype=extended_atype.dtype)[ + :, None, None + ], + masked_nlist, ] raw_atomic_energy = self._pair_tabulated_inter( @@ -303,7 +306,8 @@ 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] # pylint: disable=no-explicit-dtype + # index type is int64 + batch_indices = np.arange(nlist.shape[0], dtype=np.int64)[:, None, None] neighbor_atoms = coords[batch_indices, nlist] loc_atoms = coords[:, : nlist.shape[1], :] pairwise_dr = loc_atoms[:, :, None, :] - neighbor_atoms From de9d865a7bbafd6aaf22a54ce9077a75a5ed68a7 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 22 Oct 2024 23:15:04 +0800 Subject: [PATCH 05/19] fit pt train and model.__init__ --- deepmd/pt/model/model/__init__.py | 2 +- deepmd/pt/train/training.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 613baf440e..f3896ec853 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -100,7 +100,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) # pylint: disable=no-explicit-dtype + use_spin = np.full(len(model_params["type_map"]), False, dtype=bool) use_spin[model_params["spin"]["use_spin"]] = True model_params["spin"]["use_spin"] = use_spin.tolist() # include virtual spin and placeholder types diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 10e841682a..d59ea82d86 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -941,7 +941,9 @@ def log_loss_valid(_task_key="Default"): continue if self.multi_task: chosen_index_list = dp_random.choice( - np.arange(self.num_model), # pylint: disable=no-explicit-dtype + np.arange( + self.num_model, dtype=np.int32 + ), # int32 should be enough... p=np.array(self.model_prob), size=self.world_size, replace=True, From c957eaf873cf6509e85f8e7507567d6a330e80fa Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 22 Oct 2024 23:15:42 +0800 Subject: [PATCH 06/19] add comment --- deepmd/pt/train/training.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index d59ea82d86..756fdfebca 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -943,7 +943,7 @@ def log_loss_valid(_task_key="Default"): chosen_index_list = dp_random.choice( np.arange( self.num_model, dtype=np.int32 - ), # int32 should be enough... + ), # int32 should be enough for # models... p=np.array(self.model_prob), size=self.world_size, replace=True, From 0dcebda255fb410611845b37dc5fa1a8c3f57b0d Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 22 Oct 2024 23:23:37 +0800 Subject: [PATCH 07/19] fix: pt atomic models --- deepmd/pt/model/atomic_model/linear_atomic_model.py | 11 ++++++++--- deepmd/pt/model/atomic_model/pairtab_atomic_model.py | 8 +++++--- deepmd/pt/model/task/ener.py | 7 +++++-- 3 files changed, 18 insertions(+), 8 deletions(-) diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 8d27fbcac4..0aa5afc67f 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -90,7 +90,9 @@ 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) # pylint: disable=no-explicit-dtype + self.nsels = torch.tensor( + self.get_model_nsels(), device=env.DEVICE, dtype=torch.int32 + ) if isinstance(weights, str): assert weights in ["sum", "mean"] @@ -299,8 +301,11 @@ 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( # pylint: disable=no-explicit-dtype - [type_2_idx[new_map[idx]] for idx in range(len(new_map))], device=env.DEVICE + # int32 should be enough for number of atom types. + mapping = torch.tensor( + [type_2_idx[new_map[idx]] for idx in range(len(new_map))], + device=env.DEVICE, + dtype=torch.int32, ) return mapping diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py index 2918bba947..ae8c2d5cb0 100644 --- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py +++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py @@ -269,9 +269,11 @@ 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)[ # pylint: disable=no-explicit-dtype - :, None, None - ], + torch.arange( + extended_atype.size(0), + device=extended_coord.device, + dtype=extended_atype.dtype, + )[:, None, None], masked_nlist, ] diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 2048c05ba9..e0c5b0951e 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -126,10 +126,13 @@ def __init__( self.use_tebd = use_tebd self.out_dim = out_dim if bias_atom_e is None: - bias_atom_e = np.zeros([self.ntypes]) # pylint: disable=no-explicit-dtype + # place holder, dtype does not matter + bias_atom_e = np.zeros([self.ntypes], dtype=np.float64) if not use_tebd: assert self.ntypes == len(bias_atom_e), "Element count mismatches!" - bias_atom_e = torch.tensor(bias_atom_e, device=env.DEVICE) # pylint: disable=no-explicit-dtype + bias_atom_e = torch.tensor( + bias_atom_e, device=env.DEVICE, dtype=env.GLOBAL_PT_FLOAT_PRECISION + ) self.register_buffer("bias_atom_e", bias_atom_e) filter_layers_dipole = [] From f186457c5c0f1d9ffadccb4f3457b57790618380 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 22 Oct 2024 23:43:39 +0800 Subject: [PATCH 08/19] fix pt descriptors --- deepmd/pt/model/descriptor/descriptor.py | 12 ++++++++++-- deepmd/pt/model/descriptor/repformers.py | 20 ++++++++++++++++---- deepmd/pt/model/descriptor/se_a.py | 8 ++++++-- deepmd/pt/model/descriptor/se_atten.py | 8 ++++++-- deepmd/pt/model/descriptor/se_r.py | 20 ++++++++++++++++---- deepmd/pt/model/descriptor/se_t.py | 8 ++++++-- 6 files changed, 60 insertions(+), 16 deletions(-) diff --git a/deepmd/pt/model/descriptor/descriptor.py b/deepmd/pt/model/descriptor/descriptor.py index 78a4608108..03173a7693 100644 --- a/deepmd/pt/model/descriptor/descriptor.py +++ b/deepmd/pt/model/descriptor/descriptor.py @@ -145,8 +145,16 @@ 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)) # pylint: disable=no-explicit-dtype - base_class.stddev.copy_(torch.tensor(stddev, device=env.DEVICE)) # pylint: disable=no-explicit-dtype + base_class.mean.copy_( + torch.tensor( + mean, device=env.DEVICE, dtype=base_class.mean.dtype + ) + ) + base_class.stddev.copy_( + torch.tensor( + stddev, device=env.DEVICE, dtype=base_class.stddev.dtype + ) + ) # must share, even if not do stat self.mean = base_class.mean self.stddev = base_class.stddev diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py index ad4ead4d74..14be0ef742 100644 --- a/deepmd/pt/model/descriptor/repformers.py +++ b/deepmd/pt/model/descriptor/repformers.py @@ -466,8 +466,16 @@ def forward( comm_dict["recv_num"], g1, comm_dict["communicator"], - torch.tensor(nloc), # pylint: disable=no-explicit-dtype,no-explicit-device - torch.tensor(nall - nloc), # pylint: disable=no-explicit-dtype,no-explicit-device + torch.tensor( + nloc, + dtype=torch.int32, + device=env.DEVICE, + ), # should be int of c++ + torch.tensor( + nall - nloc, + dtype=torch.int32, + device=env.DEVICE, + ), # should be int of c++ ) g1_ext = ret[0].unsqueeze(0) g1, g2, h2 = ll.forward( @@ -530,8 +538,12 @@ def compute_input_stats( 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)) # pylint: disable=no-explicit-dtype - self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE)) # pylint: disable=no-explicit-dtype + self.mean.copy_( + torch.tensor(mean, device=env.DEVICE, dtype=self.mean.dtype) + ) + self.stddev.copy_( + torch.tensor(stddev, device=env.DEVICE, dtype=self.stddev.dtype) + ) def get_stats(self) -> dict[str, StatItem]: """Get the statistics of the descriptor.""" diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index e939a2541b..15be572778 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -556,8 +556,12 @@ def compute_input_stats( 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)) # pylint: disable=no-explicit-dtype - self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE)) # pylint: disable=no-explicit-dtype + self.mean.copy_( + torch.tensor(mean, device=env.DEVICE, dtype=self.mean.dtype) + ) + self.stddev.copy_( + torch.tensor(stddev, device=env.DEVICE, dtype=self.stddev.dtype) + ) def get_stats(self) -> dict[str, StatItem]: """Get the statistics of the descriptor.""" diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index c028230e9b..2c346e1ae1 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -364,8 +364,12 @@ def compute_input_stats( 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)) # pylint: disable=no-explicit-dtype - self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE)) # pylint: disable=no-explicit-dtype + self.mean.copy_( + torch.tensor(mean, device=env.DEVICE, dtype=self.mean.dtype) + ) + self.stddev.copy_( + torch.tensor(stddev, device=env.DEVICE, dtype=self.stddev.dtype) + ) def get_stats(self) -> dict[str, StatItem]: """Get the statistics of the descriptor.""" diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index e82bb23dac..0c73bb30d6 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -207,8 +207,16 @@ 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)) # pylint: disable=no-explicit-dtype - base_class.stddev.copy_(torch.tensor(stddev, device=env.DEVICE)) # pylint: disable=no-explicit-dtype + base_class.mean.copy_( + torch.tensor( + mean, device=env.DEVICE, dtype=base_class.mean.dtype + ) + ) + base_class.stddev.copy_( + torch.tensor( + stddev, device=env.DEVICE, dtype=base_class.stddev.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 @@ -267,8 +275,12 @@ def compute_input_stats( 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)) # pylint: disable=no-explicit-dtype - self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE)) # pylint: disable=no-explicit-dtype + self.mean.copy_( + torch.tensor(mean, device=env.DEVICE, dtype=self.mean.dtype) + ) + self.stddev.copy_( + torch.tensor(stddev, device=env.DEVICE, dtype=self.stddev.dtype) + ) def get_stats(self) -> dict[str, StatItem]: """Get the statistics of the descriptor.""" diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py index 072457b48f..5d73b2b46f 100644 --- a/deepmd/pt/model/descriptor/se_t.py +++ b/deepmd/pt/model/descriptor/se_t.py @@ -606,8 +606,12 @@ def compute_input_stats( 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)) # pylint: disable=no-explicit-dtype - self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE)) # pylint: disable=no-explicit-dtype + self.mean.copy_( + torch.tensor(mean, device=env.DEVICE, dtype=self.mean.dtype) + ) + self.stddev.copy_( + torch.tensor(stddev, device=env.DEVICE, dtype=self.stddev.dtype) + ) def get_stats(self) -> dict[str, StatItem]: """Get the statistics of the descriptor.""" From 3a56a5b57751cfe4643c674d98f32d0f80618fc1 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 23 Oct 2024 00:08:24 +0800 Subject: [PATCH 09/19] fix utils --- deepmd/dpmodel/utils/nlist.py | 1 + deepmd/pt/model/descriptor/se_t_tebd.py | 8 ++++++-- deepmd/pt/utils/nlist.py | 25 ++++++++++++++++++++----- deepmd/pt/utils/stat.py | 8 +++++++- deepmd/utils/data.py | 12 ++++++------ deepmd/utils/data_system.py | 10 +++++++--- deepmd/utils/out_stat.py | 8 ++++++-- deepmd/utils/pair_tab.py | 23 ++++++++++++++++------- 8 files changed, 69 insertions(+), 26 deletions(-) diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index 1801c155cd..22952349e1 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -277,6 +277,7 @@ def extend_coord_with_ghosts( """ xp = array_api_compat.array_namespace(coord, atype) nf, nloc = atype.shape + # int64 for index aidx = xp.tile(xp.arange(nloc, dtype=xp.int64)[xp.newaxis, :], (nf, 1)) if cell is None: nall = nloc diff --git a/deepmd/pt/model/descriptor/se_t_tebd.py b/deepmd/pt/model/descriptor/se_t_tebd.py index 437a464709..f5163793e1 100644 --- a/deepmd/pt/model/descriptor/se_t_tebd.py +++ b/deepmd/pt/model/descriptor/se_t_tebd.py @@ -698,8 +698,12 @@ def compute_input_stats( 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)) # pylint: disable=no-explicit-dtype - self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE)) # pylint: disable=no-explicit-dtype + self.mean.copy_( + torch.tensor(mean, device=env.DEVICE, dtype=self.mean.dtype) + ) + self.stddev.copy_( + torch.tensor(stddev, device=env.DEVICE, dtype=self.stddev.dtype) + ) def get_stats(self) -> dict[str, StatItem]: """Get the statistics of the descriptor.""" diff --git a/deepmd/pt/utils/nlist.py b/deepmd/pt/utils/nlist.py index 1060b40ce1..748a02b184 100644 --- a/deepmd/pt/utils/nlist.py +++ b/deepmd/pt/utils/nlist.py @@ -148,7 +148,13 @@ def _trim_mask_distinguish_nlist( nlist = nlist[:, :, :nsel] else: rr = torch.cat( - [rr, torch.ones([batch_size, nloc, nsel - nnei], device=rr.device) + rcut], # pylint: disable=no-explicit-dtype + [ + rr, + torch.ones( + [batch_size, nloc, nsel - nnei], device=rr.device, dtype=rr.dtype + ) + + rcut, + ], dim=-1, ) nlist = torch.cat( @@ -428,7 +434,10 @@ def extend_coord_with_ghosts( """ device = coord.device nf, nloc = atype.shape - aidx = torch.tile(torch.arange(nloc, device=device).unsqueeze(0), [nf, 1]) # pylint: disable=no-explicit-dtype + # int64 for index + aidx = torch.tile( + torch.arange(nloc, device=device, dtype=torch.int64).unsqueeze(0), [nf, 1] + ) if cell is None: nall = nloc extend_coord = coord.clone() @@ -447,9 +456,15 @@ def extend_coord_with_ghosts( # 3 nbuff = torch.amax(nbuff, dim=0) # faster than torch.max nbuff_cpu = nbuff.cpu() - xi = torch.arange(-nbuff_cpu[0], nbuff_cpu[0] + 1, 1, device="cpu") # pylint: disable=no-explicit-dtype - yi = torch.arange(-nbuff_cpu[1], nbuff_cpu[1] + 1, 1, device="cpu") # pylint: disable=no-explicit-dtype - zi = torch.arange(-nbuff_cpu[2], nbuff_cpu[2] + 1, 1, device="cpu") # pylint: disable=no-explicit-dtype + xi = torch.arange( + -nbuff_cpu[0], nbuff_cpu[0] + 1, 1, device="cpu", dtype=nbuff_cpu.dtype + ) + yi = torch.arange( + -nbuff_cpu[1], nbuff_cpu[1] + 1, 1, device="cpu", dtype=nbuff_cpu.dtype + ) + zi = torch.arange( + -nbuff_cpu[2], nbuff_cpu[2] + 1, 1, device="cpu", dtype=nbuff_cpu.dtype + ) eye_3 = torch.eye(3, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device="cpu") xyz = xi.view(-1, 1, 1, 1) * eye_3[0] xyz = xyz + yi.view(1, -1, 1, 1) * eye_3[1] diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 23fb12f2a4..831d2bef76 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -583,7 +583,13 @@ def compute_output_stats_atomic( # correction for missing types missing_types = ntypes - merged_natoms[kk].max() - 1 if missing_types > 0: - nan_padding = np.empty((missing_types, bias_atom_e[kk].shape[1])) # pylint: disable=no-explicit-dtype + assert ( + bias_atom_e[kk].dtype is std_atom_e[kk].dtype + ), "bias and std should be of the same dtypes" + nan_padding = np.empty( + (missing_types, bias_atom_e[kk].shape[1]), + dtype=bias_atom_e[kk].dtype, + ) nan_padding.fill(np.nan) bias_atom_e[kk] = np.concatenate([bias_atom_e[kk], nan_padding], axis=0) std_atom_e[kk] = np.concatenate([std_atom_e[kk], nan_padding], axis=0) diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index 4c77bcf59a..72e3d58660 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -266,7 +266,7 @@ def get_batch(self, batch_size: int) -> dict: iterator_1 = self.iterator + batch_size if iterator_1 >= set_size: iterator_1 = set_size - idx = np.arange(self.iterator, iterator_1) # pylint: disable=no-explicit-dtype + idx = np.arange(self.iterator, iterator_1, dtype=np.int64) self.iterator += batch_size ret = self._get_subdata(self.batch_set, idx) return ret @@ -290,7 +290,7 @@ def get_test(self, ntests: int = -1) -> dict: else self.test_set["type"].shape[0] ) # print('ntest', self.test_set['type'].shape[0], ntests, ntests_) - idx = np.arange(ntests_) # pylint: disable=no-explicit-dtype + idx = np.arange(ntests_, dtype=np.int64) ret = self._get_subdata(self.test_set, idx=idx) if self.modifier is not None: self.modifier.modify_data(ret, self) @@ -378,14 +378,14 @@ def _idx_map_sel(self, atom_type, type_sel): new_types.append(ii) new_types = np.array(new_types, dtype=int) natoms = new_types.shape[0] - idx = np.arange(natoms) # pylint: disable=no-explicit-dtype + idx = np.arange(natoms, dtype=np.int64) idx_map = np.lexsort((idx, new_types)) return idx_map def _get_natoms_2(self, ntypes): sample_type = self.atom_type natoms = len(sample_type) - natoms_vec = np.zeros(ntypes).astype(int) # pylint: disable=no-explicit-dtype + natoms_vec = np.zeros(ntypes, dtype=np.int64) for ii in range(ntypes): natoms_vec[ii] = np.count_nonzero(sample_type == ii) return natoms, natoms_vec @@ -435,7 +435,7 @@ def _load_test_set(self, shuffle_test: bool): def _shuffle_data(self, data): ret = {} nframes = data["coord"].shape[0] - idx = np.arange(nframes) # pylint: disable=no-explicit-dtype + idx = np.arange(nframes, dtype=np.int64) # the training times of each frame idx = np.repeat(idx, np.reshape(data["numb_copy"], (nframes,))) dp_random.shuffle(idx) @@ -676,7 +676,7 @@ def _load_type_mix(self, set_name: DPPath): def _make_idx_map(self, atom_type): natoms = atom_type.shape[0] - idx = np.arange(natoms) # pylint: disable=no-explicit-dtype + idx = np.arange(natoms, dtype=np.int64) if self.sort_atoms: idx_map = np.lexsort((idx, atom_type)) else: diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 7bec0b16f4..2b5fb6e6db 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -437,7 +437,9 @@ def get_batch_standard(self, sys_idx: Optional[int] = None) -> dict: self.pick_idx = sys_idx else: # prob = self._get_sys_probs(sys_probs, auto_prob_style) - self.pick_idx = dp_random.choice(np.arange(self.nsystems), p=self.sys_probs) # pylint: disable=no-explicit-dtype + self.pick_idx = dp_random.choice( + np.arange(self.nsystems, dtype=np.int32), p=self.sys_probs + ) b_data = self.data_systems[self.pick_idx].get_batch( self.batch_size[self.pick_idx] ) @@ -457,7 +459,9 @@ def get_batch_mixed(self) -> dict: batch_size = self.batch_size[0] batch_data = [] for _ in range(batch_size): - self.pick_idx = dp_random.choice(np.arange(self.nsystems), p=self.sys_probs) # pylint: disable=no-explicit-dtype + self.pick_idx = dp_random.choice( + np.arange(self.nsystems, dtype=np.int32), p=self.sys_probs + ) bb_data = self.data_systems[self.pick_idx].get_batch(1) bb_data["natoms_vec"] = self.natoms_vec[self.pick_idx] bb_data["default_mesh"] = self.default_mesh[self.pick_idx] @@ -721,7 +725,7 @@ def prob_sys_size_ext(keywords, nsystems, nbatch): block_weights.append(weight) nblocks = len(block_str) block_probs = np.array(block_weights) / np.sum(block_weights) - sys_probs = np.zeros([nsystems]) # pylint: disable=no-explicit-dtype + sys_probs = np.zeros([nsystems], dtype=np.float64) for ii in range(nblocks): nbatch_block = nbatch[block_stt[ii] : block_end[ii]] tmp_prob = [float(i) for i in nbatch_block] / np.sum(nbatch_block) diff --git a/deepmd/utils/out_stat.py b/deepmd/utils/out_stat.py index 43af191e62..bc765645dc 100644 --- a/deepmd/utils/out_stat.py +++ b/deepmd/utils/out_stat.py @@ -7,6 +7,10 @@ import numpy as np +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) + def compute_stats_from_redu( output_redu: np.ndarray, @@ -115,8 +119,8 @@ def compute_stats_from_atomic( # compute output bias nframes, nloc, ndim = output.shape ntypes = atype.max() + 1 - output_bias = np.zeros((ntypes, ndim)) # pylint: disable=no-explicit-dtype - output_std = np.zeros((ntypes, ndim)) # pylint: disable=no-explicit-dtype + output_bias = np.zeros((ntypes, ndim), dtype=GLOBAL_NP_FLOAT_PRECISION) + output_std = np.zeros((ntypes, ndim), dtype=GLOBAL_NP_FLOAT_PRECISION) for type_i in range(ntypes): mask = atype == type_i output_bias[type_i] = ( diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index cddc358f27..89f66cc994 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -36,6 +36,7 @@ class PairTab: def __init__(self, filename: str, rcut: Optional[float] = None) -> None: """Constructor.""" + self.data_type = np.float64 self.reinit(filename, rcut) def reinit(self, filename: str, rcut: Optional[float] = None) -> None: @@ -56,7 +57,7 @@ def reinit(self, filename: str, rcut: Optional[float] = None) -> None: if filename is None: self.tab_info, self.tab_data = None, None return - self.vdata = np.loadtxt(filename) + self.vdata = np.loadtxt(filename, dtype=self.data_type) self.rmin = self.vdata[0][0] self.rmax = self.vdata[-1][0] self.hh = self.vdata[1][0] - self.vdata[0][0] @@ -168,11 +169,14 @@ def _check_table_upper_boundary(self) -> None: # if table values decay to `0` before rcut, pad table with `0`s. elif self.rcut > self.rmax: - pad_zero = np.zeros((rcut_idx - upper_idx, self.ncol)) # pylint: disable=no-explicit-dtype - pad_zero[:, 0] = np.linspace( # pylint: disable=no-explicit-dtype + pad_zero = np.zeros( + (rcut_idx - upper_idx, self.ncol), dtype=self.vdata.dtype + ) + pad_zero[:, 0] = np.linspace( self.rmax + self.hh, self.rmax + self.hh * (rcut_idx - upper_idx), rcut_idx - upper_idx, + dtype=self.vdata.dtype, ) self.vdata = np.concatenate((self.vdata, pad_zero), axis=0) else: @@ -186,12 +190,15 @@ def _check_table_upper_boundary(self) -> None: log.warning( "The rcut goes beyond table upper boundary, performing extrapolation." ) - pad_extrapolation = np.zeros((rcut_idx - upper_idx, self.ncol)) # pylint: disable=no-explicit-dtype + pad_extrapolation = np.zeros( + (rcut_idx - upper_idx, self.ncol), dtype=self.vdata.dtype + ) - pad_extrapolation[:, 0] = np.linspace( # pylint: disable=no-explicit-dtype + pad_extrapolation[:, 0] = np.linspace( self.rmax + self.hh, self.rmax + self.hh * (rcut_idx - upper_idx), rcut_idx - upper_idx, + dtype=self.vdata.dtype, ) # need to calculate table values to fill in with cubic spline pad_extrapolation = self._extrapolate_table(pad_extrapolation) @@ -252,7 +259,9 @@ def _extrapolate_table(self, pad_extrapolation: np.array) -> np.array: return pad_extrapolation def _make_data(self): - data = np.zeros([self.ntypes * self.ntypes * 4 * self.nspline]) # pylint: disable=no-explicit-dtype + data = np.zeros( + [self.ntypes * self.ntypes * 4 * self.nspline], dtype=self.data_type + ) stride = 4 * self.nspline idx_iter = 0 xx = self.vdata[:, 0] @@ -262,7 +271,7 @@ def _make_data(self): cs = CubicSpline(xx, vv, bc_type="clamped") dd = cs(xx, 1) dd *= self.hh - dtmp = np.zeros(stride) # pylint: disable=no-explicit-dtype + dtmp = np.zeros(stride, dtype=self.data_type) for ii in range(self.nspline): dtmp[ii * 4 + 0] = 2 * vv[ii] - 2 * vv[ii + 1] + dd[ii] + dd[ii + 1] dtmp[ii * 4 + 1] = ( From 1a1f1454dc88e2cf5b8c447691d6515344f63c8f Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 23 Oct 2024 09:40:53 +0800 Subject: [PATCH 10/19] fix infer, model_devi and spin --- deepmd/dpmodel/infer/deep_eval.py | 7 +++++-- deepmd/infer/model_devi.py | 4 ++-- deepmd/pt/infer/deep_eval.py | 19 +++++++++++-------- deepmd/utils/spin.py | 22 ++++++++++++++++------ 4 files changed, 34 insertions(+), 18 deletions(-) diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index 695edb29d2..2b1e74c8de 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/deepmd/dpmodel/infer/deep_eval.py @@ -24,6 +24,9 @@ from deepmd.dpmodel.utils.serialization import ( load_dp_model, ) +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) from deepmd.infer.deep_dipole import ( DeepDipole, ) @@ -340,12 +343,12 @@ 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) # pylint: disable=no-explicit-dtype + out = np.full(shape, np.nan, dtype=GLOBAL_NP_FLOAT_PRECISION) results.append(out) else: shape = self._get_output_shape(odef, nframes, natoms) results.append( - np.full(np.abs(shape), np.nan) # pylint: disable=no-explicit-dtype + np.full(np.abs(shape), np.nan, dtype=GLOBAL_NP_FLOAT_PRECISION) ) # this is kinda hacky return tuple(results) diff --git a/deepmd/infer/model_devi.py b/deepmd/infer/model_devi.py index 29e1eec741..68100ba739 100644 --- a/deepmd/infer/model_devi.py +++ b/deepmd/infer/model_devi.py @@ -328,7 +328,7 @@ def calc_model_devi( forces = np.array(forces) virials = np.array(virials) - devi = [np.arange(coord.shape[0]) * frequency] # pylint: disable=no-explicit-dtype + devi = [np.arange(coord.shape[0], dtype=np.int64) * frequency] 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)) @@ -502,7 +502,7 @@ def make_model_devi( nframes_tot += coord.shape[0] devis.append(devi) devis = np.vstack(devis) - devis[:, 0] = np.arange(nframes_tot) * frequency # pylint: disable=no-explicit-dtype + devis[:, 0] = np.arange(nframes_tot, dtype=np.int64) * frequency write_model_devi_out(devis, output, header=system, atomic=atomic) devis_coll.append(devis) return devis_coll diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 0a77a38135..ee5a7e39c8 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -394,6 +394,7 @@ def _eval_model( request_defs: list[OutputVariableDef], ): model = self.dp.to(DEVICE) + prec = NP_PRECISION_DICT[RESERVED_PRECISON_DICT[GLOBAL_PT_FLOAT_PRECISION]] nframes = coords.shape[0] if len(atom_types.shape) == 1: @@ -403,9 +404,7 @@ def _eval_model( natoms = len(atom_types[0]) coord_input = torch.tensor( - coords.reshape([nframes, natoms, 3]).astype( - NP_PRECISION_DICT[RESERVED_PRECISON_DICT[GLOBAL_PT_FLOAT_PRECISION]] - ), + coords.reshape([nframes, natoms, 3]).astype(prec), dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE, ) @@ -416,9 +415,7 @@ def _eval_model( ) if cells is not None: box_input = torch.tensor( - cells.reshape([nframes, 3, 3]).astype( - NP_PRECISION_DICT[RESERVED_PRECISON_DICT[GLOBAL_PT_FLOAT_PRECISION]] - ), + cells.reshape([nframes, 3, 3]).astype(prec), dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE, ) @@ -460,7 +457,7 @@ def _eval_model( else: shape = self._get_output_shape(odef, nframes, natoms) results.append( - np.full(np.abs(shape), np.nan) # pylint: disable=no-explicit-dtype + np.full(np.abs(shape), np.nan, dtype=prec) ) # this is kinda hacky return tuple(results) @@ -540,7 +537,13 @@ def _eval_model_spin( else: shape = self._get_output_shape(odef, nframes, natoms) results.append( - np.full(np.abs(shape), np.nan) # pylint: disable=no-explicit-dtype + np.full( + np.abs(shape), + np.nan, + dtype=NP_PRECISION_DICT[ + RESERVED_PRECISON_DICT[GLOBAL_PT_FLOAT_PRECISION] + ], + ) ) # this is kinda hacky return tuple(results) diff --git a/deepmd/utils/spin.py b/deepmd/utils/spin.py index 41ea52df88..9ea5fc0713 100644 --- a/deepmd/utils/spin.py +++ b/deepmd/utils/spin.py @@ -6,6 +6,10 @@ import numpy as np +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) + class Spin: """Class for spin, mainly processes the spin type-related information. @@ -36,6 +40,7 @@ def __init__( use_spin: list[bool], virtual_scale: Union[list[float], float], ) -> None: + type_dtype = np.int32 self.ntypes_real = len(use_spin) self.ntypes_spin = use_spin.count(True) self.use_spin = np.array(use_spin) @@ -43,19 +48,24 @@ def __init__( self.ntypes_real_and_spin = self.ntypes_real + self.ntypes_spin self.ntypes_placeholder = self.ntypes_real - self.ntypes_spin self.ntypes_input = 2 * self.ntypes_real # with placeholder for input types - self.real_type = np.arange(self.ntypes_real) # pylint: disable=no-explicit-dtype - self.spin_type = np.arange(self.ntypes_real)[self.use_spin] + self.ntypes_real # pylint: disable=no-explicit-dtype + self.real_type = np.arange(self.ntypes_real, dtype=type_dtype) + self.spin_type = self.real_type[self.use_spin] + self.ntypes_real self.real_and_spin_type = np.concatenate([self.real_type, self.spin_type]) self.placeholder_type = ( - np.arange(self.ntypes_real)[~self.use_spin] + self.ntypes_real # pylint: disable=no-explicit-dtype + np.arange(self.ntypes_real, dtype=type_dtype)[~self.use_spin] + + self.ntypes_real ) - self.spin_placeholder_type = np.arange(self.ntypes_real) + self.ntypes_real # pylint: disable=no-explicit-dtype - self.input_type = np.arange(self.ntypes_real * 2) # pylint: disable=no-explicit-dtype + self.spin_placeholder_type = ( + np.arange(self.ntypes_real, dtype=type_dtype) + self.ntypes_real + ) + self.input_type = np.arange(self.ntypes_real * 2, dtype=type_dtype) if isinstance(virtual_scale, list): if len(virtual_scale) == self.ntypes_real: self.virtual_scale = virtual_scale elif len(virtual_scale) == self.ntypes_spin: - self.virtual_scale = np.zeros(self.ntypes_real) # pylint: disable=no-explicit-dtype + self.virtual_scale = np.zeros( + self.ntypes_real, dtype=GLOBAL_NP_FLOAT_PRECISION + ) self.virtual_scale[self.use_spin] = virtual_scale else: raise ValueError( From a7427c9c1cb7a1aa588e454deccfe42bedb3ba9d Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 23 Oct 2024 09:59:04 +0800 Subject: [PATCH 11/19] fix index type --- deepmd/dpmodel/atomic_model/pairtab_atomic_model.py | 6 ++---- deepmd/pt/model/atomic_model/pairtab_atomic_model.py | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index ee9c7a9a76..2899f106bc 100644 --- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py @@ -204,11 +204,9 @@ def forward_atomic( self.tab.ntypes, self.tab.ntypes, self.tab.nspline, 4 ) - # (nframes, nloc, nnei), dtype is the same as atype. + # (nframes, nloc, nnei), index type is int64. j_type = extended_atype[ - np.arange(extended_atype.shape[0], dtype=extended_atype.dtype)[ - :, None, None - ], + np.arange(extended_atype.shape[0], dtype=np.int64)[:, None, None], masked_nlist, ] diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py index ae8c2d5cb0..2bedccbd43 100644 --- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py +++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py @@ -272,7 +272,7 @@ def forward_atomic( torch.arange( extended_atype.size(0), device=extended_coord.device, - dtype=extended_atype.dtype, + dtype=torch.int64, )[:, None, None], masked_nlist, ] From fa7e103bfb703f5ee9589bedea0e787830aebfda Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 23 Oct 2024 10:06:11 +0800 Subject: [PATCH 12/19] fix some fitting precs --- deepmd/dpmodel/fitting/general_fitting.py | 23 ++++++++++++++--------- deepmd/pt/model/task/fitting.py | 4 +++- 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 75e622b714..8e9e228787 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -28,6 +28,9 @@ from deepmd.dpmodel.utils.seed import ( child_seed, ) +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) from deepmd.utils.finetune import ( get_index_between_two_maps, map_atom_exclude_types, @@ -134,7 +137,7 @@ def __init__( self.trainable = [self.trainable] * (len(self.neuron) + 1) self.activation_function = activation_function self.precision = precision - prec = PRECISION_DICT[self.precision.lower()] + self.prec = PRECISION_DICT[self.precision.lower()] self.layer_name = layer_name self.use_aparam_as_mask = use_aparam_as_mask self.spin = spin @@ -148,18 +151,20 @@ 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], dtype=prec) + self.bias_atom_e = np.zeros( + [self.ntypes, net_dim_out], dtype=GLOBAL_NP_FLOAT_PRECISION + ) else: assert bias_atom_e.shape == (self.ntypes, net_dim_out) - self.bias_atom_e = bias_atom_e + self.bias_atom_e = bias_atom_e.astype(GLOBAL_NP_FLOAT_PRECISION) if self.numb_fparam > 0: - self.fparam_avg = np.zeros(self.numb_fparam, dtype=prec) - self.fparam_inv_std = np.ones(self.numb_fparam, dtype=prec) + self.fparam_avg = np.zeros(self.numb_fparam, dtype=self.prec) + self.fparam_inv_std = np.ones(self.numb_fparam, dtype=self.prec) else: self.fparam_avg, self.fparam_inv_std = None, None if self.numb_aparam > 0: - self.aparam_avg = np.zeros(self.numb_aparam, dtype=prec) - self.aparam_inv_std = np.ones(self.numb_aparam, dtype=prec) + self.aparam_avg = np.zeros(self.numb_aparam, dtype=self.prec) + self.aparam_inv_std = np.ones(self.numb_aparam, dtype=self.prec) else: self.aparam_avg, self.aparam_inv_std = None, None # init networks @@ -412,7 +417,7 @@ def _call_common( # calcualte the prediction if not self.mixed_types: - outs = xp.zeros([nf, nloc, net_dim_out], dtype=descriptor.dtype) + outs = xp.zeros([nf, nloc, net_dim_out], dtype=self.prec) for type_i in range(self.ntypes): mask = xp.tile( xp.reshape((atype == type_i), [nf, nloc, 1]), (1, 1, net_dim_out) @@ -438,4 +443,4 @@ def _call_common( exclude_mask = self.emask.build_type_exclude_mask(atype) # nf x nloc x nod outs = outs * xp.astype(exclude_mask[:, :, None], outs.dtype) - return {self.var_name: outs} + return {self.var_name: outs.astype(GLOBAL_NP_FLOAT_PRECISION)} diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 10f88519e1..6e9829e4b6 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -177,7 +177,9 @@ def __init__( # init constants if bias_atom_e is None: bias_atom_e = np.zeros([self.ntypes, net_dim_out], dtype=np.float64) - bias_atom_e = torch.tensor(bias_atom_e, dtype=self.prec, device=device) + bias_atom_e = torch.tensor( + bias_atom_e, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=device + ) bias_atom_e = bias_atom_e.view([self.ntypes, net_dim_out]) if not self.mixed_types: assert self.ntypes == bias_atom_e.shape[0], "Element count mismatches!" From 11b044052003128cf744525cbde90023bf4ac7e4 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 23 Oct 2024 10:18:30 +0800 Subject: [PATCH 13/19] fix nlist --- deepmd/dpmodel/utils/nlist.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index 22952349e1..3ef17fc6b9 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -290,9 +290,9 @@ def extend_coord_with_ghosts( to_face = to_face_distance(cell) nbuff = xp.astype(xp.ceil(rcut / to_face), xp.int64) nbuff = xp.max(nbuff, axis=0) - xi = xp.arange(-int(nbuff[0]), int(nbuff[0]) + 1, 1, dtype=nbuff.dtype) - yi = xp.arange(-int(nbuff[1]), int(nbuff[1]) + 1, 1, dtype=nbuff.dtype) - zi = xp.arange(-int(nbuff[2]), int(nbuff[2]) + 1, 1, dtype=nbuff.dtype) + xi = xp.arange(-int(nbuff[0]), int(nbuff[0]) + 1, 1, dtype=xp.int64) + yi = xp.arange(-int(nbuff[1]), int(nbuff[1]) + 1, 1, dtype=xp.int64) + zi = xp.arange(-int(nbuff[2]), int(nbuff[2]) + 1, 1, dtype=xp.int64) xyz = xp.linalg.outer(xi, xp.asarray([1, 0, 0]))[:, xp.newaxis, xp.newaxis, :] xyz = ( xyz From dafa435a253699c10c129b52df4db7861cb09dba Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 23 Oct 2024 10:22:51 +0800 Subject: [PATCH 14/19] fix nlist --- deepmd/pt/utils/nlist.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/pt/utils/nlist.py b/deepmd/pt/utils/nlist.py index 748a02b184..82e97c29e6 100644 --- a/deepmd/pt/utils/nlist.py +++ b/deepmd/pt/utils/nlist.py @@ -457,13 +457,13 @@ def extend_coord_with_ghosts( nbuff = torch.amax(nbuff, dim=0) # faster than torch.max nbuff_cpu = nbuff.cpu() xi = torch.arange( - -nbuff_cpu[0], nbuff_cpu[0] + 1, 1, device="cpu", dtype=nbuff_cpu.dtype + -nbuff_cpu[0], nbuff_cpu[0] + 1, 1, device="cpu", dtype=torch.int64 ) yi = torch.arange( - -nbuff_cpu[1], nbuff_cpu[1] + 1, 1, device="cpu", dtype=nbuff_cpu.dtype + -nbuff_cpu[1], nbuff_cpu[1] + 1, 1, device="cpu", dtype=torch.int64 ) zi = torch.arange( - -nbuff_cpu[2], nbuff_cpu[2] + 1, 1, device="cpu", dtype=nbuff_cpu.dtype + -nbuff_cpu[2], nbuff_cpu[2] + 1, 1, device="cpu", dtype=torch.int64 ) eye_3 = torch.eye(3, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device="cpu") xyz = xi.view(-1, 1, 1, 1) * eye_3[0] From a261d6740c59532851eb5474f0cfdefc1951166a Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 23 Oct 2024 10:23:54 +0800 Subject: [PATCH 15/19] torch.long -> torch.int64 --- deepmd/pt/utils/nlist.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/pt/utils/nlist.py b/deepmd/pt/utils/nlist.py index 82e97c29e6..c30ec6dd02 100644 --- a/deepmd/pt/utils/nlist.py +++ b/deepmd/pt/utils/nlist.py @@ -452,7 +452,7 @@ def extend_coord_with_ghosts( # nf x 3 # *2: ghost copies on + and - directions # +1: central cell - nbuff = torch.ceil(rcut / to_face).to(torch.long) + nbuff = torch.ceil(rcut / to_face).to(torch.int64) # 3 nbuff = torch.amax(nbuff, dim=0) # faster than torch.max nbuff_cpu = nbuff.cpu() From b213475be23755ae9c5dd1fdc51e2181e18facde Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 23 Oct 2024 18:47:05 +0800 Subject: [PATCH 16/19] fix general_fitting, add xp type dict --- deepmd/dpmodel/common.py | 27 +++++++++++++++++++++++ deepmd/dpmodel/fitting/general_fitting.py | 7 ++++-- source/tests/common/test_common.py | 27 +++++++++++++++++++++++ 3 files changed, 59 insertions(+), 2 deletions(-) create mode 100644 source/tests/common/test_common.py diff --git a/deepmd/dpmodel/common.py b/deepmd/dpmodel/common.py index d9d57d2d6c..5c75229e49 100644 --- a/deepmd/dpmodel/common.py +++ b/deepmd/dpmodel/common.py @@ -50,6 +50,33 @@ DEFAULT_PRECISION = "float64" +def get_xp_precision( + xp: Any, + precision: str, +): + """Get the precision from the API compatible namespace.""" + if precision == "float16" or precision == "half": + return xp.float16 + elif precision == "float32" or precision == "single": + return xp.float32 + elif precision == "float64" or precision == "double": + return xp.float64 + elif precision == "int32": + return xp.int32 + elif precision == "int64": + return xp.int64 + elif precision == "bool": + return bool + elif precision == "default": + return get_xp_precision(xp, RESERVED_PRECISON_DICT[PRECISION_DICT[precision]]) + elif precision == "global": + return get_xp_precision(xp, RESERVED_PRECISON_DICT[GLOBAL_NP_FLOAT_PRECISION]) + elif precision == "bfloat16": + return ml_dtypes.bfloat16 + else: + raise ValueError(f"unsupported precision {precision} for {xp}") + + class NativeOP(ABC): """The unit operation of a native model.""" diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 8e9e228787..016f947206 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -18,6 +18,7 @@ NativeOP, ) from deepmd.dpmodel.common import ( + get_xp_precision, to_numpy_array, ) from deepmd.dpmodel.utils import ( @@ -417,7 +418,9 @@ def _call_common( # calcualte the prediction if not self.mixed_types: - outs = xp.zeros([nf, nloc, net_dim_out], dtype=self.prec) + outs = xp.zeros( + [nf, nloc, net_dim_out], dtype=get_xp_precision(xp, self.precision) + ) for type_i in range(self.ntypes): mask = xp.tile( xp.reshape((atype == type_i), [nf, nloc, 1]), (1, 1, net_dim_out) @@ -443,4 +446,4 @@ def _call_common( exclude_mask = self.emask.build_type_exclude_mask(atype) # nf x nloc x nod outs = outs * xp.astype(exclude_mask[:, :, None], outs.dtype) - return {self.var_name: outs.astype(GLOBAL_NP_FLOAT_PRECISION)} + return {self.var_name: xp.astype(outs, get_xp_precision(xp, "global"))} diff --git a/source/tests/common/test_common.py b/source/tests/common/test_common.py new file mode 100644 index 0000000000..478c512fed --- /dev/null +++ b/source/tests/common/test_common.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import array_api_compat +import ml_dtypes +import numpy as np + +from deepmd.dpmodel.common import ( + get_xp_precision, +) +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) + + +class TestGetXPPrecision(unittest.TestCase): + def test(self): + aa = np.zeros(3) + xp = array_api_compat.array_namespace(aa) + self.assertTrue(get_xp_precision(xp, "float16"), xp.float16) + self.assertTrue(get_xp_precision(xp, "float32"), xp.float32) + self.assertTrue(get_xp_precision(xp, "float64"), xp.float64) + self.assertTrue(get_xp_precision(xp, "single"), xp.float32) + self.assertTrue(get_xp_precision(xp, "double"), xp.float64) + self.assertTrue(get_xp_precision(xp, "global"), GLOBAL_NP_FLOAT_PRECISION) + self.assertTrue(get_xp_precision(xp, "default"), GLOBAL_NP_FLOAT_PRECISION) + self.assertTrue(get_xp_precision(xp, "bfloat16"), ml_dtypes.bfloat16) From dc652fc69a19ffc1ae4143f692aeeeb6ddce1b88 Mon Sep 17 00:00:00 2001 From: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> Date: Wed, 23 Oct 2024 19:58:42 +0800 Subject: [PATCH 17/19] Update source/tests/common/test_common.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> --- source/tests/common/test_common.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/source/tests/common/test_common.py b/source/tests/common/test_common.py index 478c512fed..5b5854a26a 100644 --- a/source/tests/common/test_common.py +++ b/source/tests/common/test_common.py @@ -17,11 +17,15 @@ class TestGetXPPrecision(unittest.TestCase): def test(self): aa = np.zeros(3) xp = array_api_compat.array_namespace(aa) - self.assertTrue(get_xp_precision(xp, "float16"), xp.float16) - self.assertTrue(get_xp_precision(xp, "float32"), xp.float32) - self.assertTrue(get_xp_precision(xp, "float64"), xp.float64) - self.assertTrue(get_xp_precision(xp, "single"), xp.float32) - self.assertTrue(get_xp_precision(xp, "double"), xp.float64) - self.assertTrue(get_xp_precision(xp, "global"), GLOBAL_NP_FLOAT_PRECISION) - self.assertTrue(get_xp_precision(xp, "default"), GLOBAL_NP_FLOAT_PRECISION) - self.assertTrue(get_xp_precision(xp, "bfloat16"), ml_dtypes.bfloat16) + self.assertEqual(get_xp_precision(xp, "float16"), xp.float16) + self.assertEqual(get_xp_precision(xp, "float32"), xp.float32) + self.assertEqual(get_xp_precision(xp, "float64"), xp.float64) + self.assertEqual(get_xp_precision(xp, "single"), xp.float32) + self.assertEqual(get_xp_precision(xp, "double"), xp.float64) + self.assertEqual(get_xp_precision(xp, "global"), GLOBAL_NP_FLOAT_PRECISION) + self.assertEqual(get_xp_precision(xp, "default"), GLOBAL_NP_FLOAT_PRECISION) + self.assertEqual(get_xp_precision(xp, "bfloat16"), ml_dtypes.bfloat16) + + # Test invalid input + with self.assertRaises(ValueError): + get_xp_precision(xp, "invalid_precision") From aca1b1e036704778684cf155ba8907bf360236a1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 23 Oct 2024 11:59:21 +0000 Subject: [PATCH 18/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/common/test_common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/tests/common/test_common.py b/source/tests/common/test_common.py index 5b5854a26a..fe9054d6ad 100644 --- a/source/tests/common/test_common.py +++ b/source/tests/common/test_common.py @@ -25,7 +25,7 @@ def test(self): self.assertEqual(get_xp_precision(xp, "global"), GLOBAL_NP_FLOAT_PRECISION) self.assertEqual(get_xp_precision(xp, "default"), GLOBAL_NP_FLOAT_PRECISION) self.assertEqual(get_xp_precision(xp, "bfloat16"), ml_dtypes.bfloat16) - + # Test invalid input with self.assertRaises(ValueError): get_xp_precision(xp, "invalid_precision") From 0c6b42a51c809141ed757639ebd82c781c9eda0e Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 23 Oct 2024 20:01:25 +0800 Subject: [PATCH 19/19] raise exception --- deepmd/dpmodel/fitting/general_fitting.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 016f947206..62aafc6207 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -138,6 +138,10 @@ def __init__( self.trainable = [self.trainable] * (len(self.neuron) + 1) self.activation_function = activation_function self.precision = precision + if self.precision.lower() not in PRECISION_DICT: + raise ValueError( + f"Unsupported precision '{self.precision}'. Supported options are: {list(PRECISION_DICT.keys())}" + ) self.prec = PRECISION_DICT[self.precision.lower()] self.layer_name = layer_name self.use_aparam_as_mask = use_aparam_as_mask