From 6c05a71f8c3e9116c139e1ef4505380416d66a72 Mon Sep 17 00:00:00 2001 From: AsymmetryChou <181240085@smail.nju.edu.cn> Date: Wed, 26 Nov 2025 11:21:24 +0800 Subject: [PATCH 1/3] fix: Convert mask to numpy and Ensure arrays are proper numpy --- dptb/data/AtomicData.py | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/dptb/data/AtomicData.py b/dptb/data/AtomicData.py index 46b5dd37..fa689d90 100644 --- a/dptb/data/AtomicData.py +++ b/dptb/data/AtomicData.py @@ -1009,16 +1009,22 @@ def neighbor_list_and_relative_vec( shifts = shifts[mask] # 2. for i == j - + mask = torch.ones(len(first_idex), dtype=torch.bool) mask[first_idex == second_idex] = False # get index bool type ~mask for i == j. - o_first_idex = first_idex[~mask] - o_second_idex = second_idex[~mask] - o_shift = shifts[~mask] + # Convert mask to numpy for consistent indexing behavior + mask_np = mask.numpy() + o_first_idex = first_idex[~mask_np] + o_second_idex = second_idex[~mask_np] + o_shift = shifts[~mask_np] o_mask = mask[~mask] # this is all False, with length being the number all the bonds with i == j. - + # Ensure arrays are proper numpy arrays (not scalars) for isolated systems + o_first_idex = np.atleast_1d(o_first_idex) + o_second_idex = np.atleast_1d(o_second_idex) + o_shift = np.atleast_1d(o_shift) + # using the dict key to remove the duplicate bonds, because it is O(1) to check if a key is in the dict. rev_dict = {} for i in range(len(o_first_idex)): @@ -1042,10 +1048,12 @@ def neighbor_list_and_relative_vec( del o_shift mask[~mask] = o_mask del o_mask - - first_idex = torch.LongTensor(first_idex[mask], device=out_device) - second_idex = torch.LongTensor(second_idex[mask], device=out_device) - shifts = torch.as_tensor(shifts[mask], dtype=out_dtype, device=out_device) + + # Convert mask to numpy for indexing numpy arrays (avoids torch/numpy compatibility issues) + mask_np = mask.numpy() + first_idex = torch.as_tensor(first_idex[mask_np], dtype=torch.long, device=out_device) + second_idex = torch.as_tensor(second_idex[mask_np], dtype=torch.long, device=out_device) + shifts = torch.as_tensor(shifts[mask_np], dtype=out_dtype, device=out_device) if not reduce: assert self_interaction == False, "for self_interaction = True, i i 0 0 0 will be duplicated." From 0eb5c52512dce81a3a9a4803667e8380c5846a7a Mon Sep 17 00:00:00 2001 From: AsymmetryChou <181240085@smail.nju.edu.cn> Date: Wed, 26 Nov 2025 14:02:58 +0800 Subject: [PATCH 2/3] fix: Ensure o_shift is a 2D numpy array for isolated systems --- dptb/data/AtomicData.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dptb/data/AtomicData.py b/dptb/data/AtomicData.py index fa689d90..61915b2a 100644 --- a/dptb/data/AtomicData.py +++ b/dptb/data/AtomicData.py @@ -1023,7 +1023,7 @@ def neighbor_list_and_relative_vec( # Ensure arrays are proper numpy arrays (not scalars) for isolated systems o_first_idex = np.atleast_1d(o_first_idex) o_second_idex = np.atleast_1d(o_second_idex) - o_shift = np.atleast_1d(o_shift) + o_shift = np.atleast_2d(o_shift) # using the dict key to remove the duplicate bonds, because it is O(1) to check if a key is in the dict. rev_dict = {} From 0e074be6c6897fa8061dfdb3a44487495dc0a925 Mon Sep 17 00:00:00 2001 From: AsymmetryChou <181240085@smail.nju.edu.cn> Date: Wed, 26 Nov 2025 14:12:00 +0800 Subject: [PATCH 3/3] fix: Use .cpu().numpy() for mask conversion to ensure compatibility with numpy arrays --- dptb/data/AtomicData.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dptb/data/AtomicData.py b/dptb/data/AtomicData.py index 61915b2a..bc05bfe7 100644 --- a/dptb/data/AtomicData.py +++ b/dptb/data/AtomicData.py @@ -1014,7 +1014,7 @@ def neighbor_list_and_relative_vec( mask[first_idex == second_idex] = False # get index bool type ~mask for i == j. # Convert mask to numpy for consistent indexing behavior - mask_np = mask.numpy() + mask_np = mask.cpu().numpy() o_first_idex = first_idex[~mask_np] o_second_idex = second_idex[~mask_np] o_shift = shifts[~mask_np] @@ -1050,7 +1050,7 @@ def neighbor_list_and_relative_vec( del o_mask # Convert mask to numpy for indexing numpy arrays (avoids torch/numpy compatibility issues) - mask_np = mask.numpy() + mask_np = mask.cpu().numpy() first_idex = torch.as_tensor(first_idex[mask_np], dtype=torch.long, device=out_device) second_idex = torch.as_tensor(second_idex[mask_np], dtype=torch.long, device=out_device) shifts = torch.as_tensor(shifts[mask_np], dtype=out_dtype, device=out_device)