Coverage for gamdpy/interactions/nblist_linked_lists.py: 39%
191 statements
« prev ^ index » next coverage.py v7.4.4, created at 2025-06-14 15:25 +0200
2import numpy as np
3import numba
4import math
5from numba import cuda
7class NbListLinkedLists():
9 def __init__(self, configuration, exclusions, max_num_nbs):
10 self.configuration = configuration
11 self.N, self.D = configuration.N, configuration.D
12 self.nblist = np.zeros((configuration.N, max_num_nbs+1), dtype=np.int32)
13 self.nbflag = np.zeros(3, dtype=np.int32)
14 self.r_ref = np.zeros_like(configuration['r']) # Inherents also data type
15 self.exclusions = exclusions # Should be able to be a list (eg from bonds, angles, etc), and merge
16 self.d_simbox_last_rebuild = cuda.to_device(np.zeros(configuration.simbox.len_sim_box_data, dtype=np.float32))
18 def copy_to_device(self):
19 self.d_nblist = cuda.to_device(self.nblist)
20 self.d_nbflag = cuda.to_device(self.nbflag)
21 self.d_r_ref = cuda.to_device(self.r_ref)
22 self.d_exclusions = cuda.to_device(self._exclusions)
23 self.d_cells_per_dimension = cuda.to_device(self.cells_per_dimension) # Mapping cells to 1D, so need 'shape'
24 self.d_cells = cuda.to_device(self.cells)
25 self.d_my_cell = cuda.to_device(self.my_cell)
26 self.d_next_particle_in_cell = cuda.to_device(self.next_particle_in_cell)
28 def get_params(self, max_cut, compute_plan, verbose=True):
29 self.max_cut = max_cut
30 self.skin, = [compute_plan[key] for key in ['skin']]
32 if type(self.exclusions) == np.ndarray: # Don't change user-set properties
33 self._exclusions = self.exclusions.copy()
34 else:
35 self._exclusions = np.zeros((self.r_ref.shape[0], 2), dtype=np.int32)
37 self.cells_per_dimension = np.zeros((self.D), dtype=np.int32)
38 min_cells_size = (self.max_cut + self.skin)/2
39 simbox_lengths = self.configuration.simbox.get_lengths()
40 for i in range(self.D):
41 self.cells_per_dimension[i] = int(simbox_lengths[i]/min_cells_size)
42 assert self.cells_per_dimension[i] > 4
43 if i == 0:
44 assert self.cells_per_dimension[i] > 6
45 # TODO: Take care of
46 # - changing simbox size during simulation (NPT, compression)
49 #print(self.cells_per_dimension)
50 self.my_cell = np.zeros((self.N, self.D), np.int32) # my_cell[:, -1] 1d index
51 self.cells = -np.ones(self.cells_per_dimension, dtype=np.int32)
52 #self.cells = -np.ones(np.prod(self.cells_per_dimension), dtype=np.int32) # Map to 1D array
53 #print(self.cells.shape)
55 self.next_particle_in_cell = -np.ones(self.N, dtype=np.int32) # -1 = no further particles in list
56 self.copy_to_device()
57 return (np.float32(self.max_cut), np.float32(self.skin), self.d_nbflag, self.d_r_ref, self.d_exclusions,
58 self.d_cells_per_dimension, self.d_cells, self.d_my_cell, self.d_next_particle_in_cell, self.d_simbox_last_rebuild)
60 def get_kernel(self, configuration, compute_plan, compute_flags, verbose=False, force_update=False):
62 # Unpack parameters from configuration and compute_plan
63 D, num_part = configuration.D, configuration.N
64 pb, tp, gridsync, UtilizeNIII = [compute_plan[key] for key in ['pb', 'tp', 'gridsync', 'UtilizeNIII']]
65 num_blocks = (num_part - 1) // pb + 1
66 compute_stresses = compute_flags['stresses']
68 loop_x_addition = configuration.simbox.get_loop_x_addition()
70 # Unpack indices for vectors and scalars to be compiled into kernel
71 r_id, f_id = [configuration.vectors.indices[key] for key in ['r', 'f']]
72 if compute_stresses:
73 sx_id = configuration.vectors.indices['sx']
74 if D > 1:
75 sy_id = configuration.vectors.indices['sy']
76 if D > 2:
77 sz_id = configuration.vectors.indices['sz']
78 if D > 3:
79 sw_id = configuration.vectors.indices['sw']
82 # JIT compile functions to be compiled into kernel
83 dist_sq_function = numba.njit(configuration.simbox.get_dist_sq_function())
84 dist_moved_exceeds_limit_function = numba.njit(configuration.simbox.get_dist_moved_exceeds_limit_function())
85 loop_x_shift_function = numba.njit(configuration.simbox.get_loop_x_shift_function())
87 @cuda.jit( device=gridsync )
88 def nblist_check(vectors, sim_box, skin, r_ref, nbflag, simbox_last_rebuild, cut):
89 """ Check validity of nblist, i.e. did any particle mode more than skin/2 since last nblist update?
90 Each tread-block checks the assigned particles (global_id)
91 nbflag[0] = 0 : No update needed
92 nbflag[0] = num_blocks : Update needed
93 Kernel configuration: [num_blocks, (pb, tp)]
94 """
96 global_id, my_t = cuda.grid(2)
97 if force_update: # nblist update forced (for benchmark or similar)
98 if global_id==0 and my_t==0:
99 nbflag[0]=num_blocks
101 if global_id < num_part and my_t==0:
102 if dist_moved_exceeds_limit_function(vectors[r_id][global_id], r_ref[global_id], sim_box, simbox_last_rebuild, skin, cut):
103 nbflag[0] = num_blocks
105 #if global_id < num_part and my_t==0:
106 # dist_sq = dist_sq_function(vectors[r_id][global_id], r_ref[global_id], sim_box)
107 # if dist_sq > skin*skin*numba.float32(0.25):
108 # nbflag[0]=num_blocks
111 return
113 @cuda.jit(device=gridsync)
114 def put_particles_in_cells(vectors, sim_box, nbflag, cells_per_dimension, cells, my_cell, next_particle_in_cell):
115 """ Each particle computers its cells coordinates, and inserts itself in cell-list
116 Kernel configuration: [num_blocks, (pb, tp)]
117 """
119 global_id, my_t = cuda.grid(2)
120 if global_id < num_part and my_t == 0:
121 for k in range(D):
122 #my_cell[global_id,k] = int(math.floor(vectors[r_id][global_id,k]*cells_per_dimension[k]/sim_box[k]))%cells_per_dimension[k]
123 my_cell[global_id,k] = int(math.floor((vectors[r_id][global_id,k]+0.5*sim_box[k])*cells_per_dimension[k]/sim_box[k]))%cells_per_dimension[k]
124 #if my_cell[global_id,k]<0:
125 # print(global_id,k, my_cell[global_id,k],vectors[r_id][global_id,k])
126 index = (my_cell[global_id,0], my_cell[global_id,1], my_cell[global_id,2]) # 3D
127 next_particle_in_cell[global_id] = cuda.atomic.exch(cells, index, global_id) # index needs to be tuple when multidim
129 @cuda.jit(device=gridsync)
130 def nblist_update_from_linked_lists(vectors, sim_box, cut_plus_skin, nbflag, cells_per_dimension, cells, my_cell, next_particle_in_cell, nblist, r_ref, exclusions):
131 """ Order N neighbor-list update from linked lists
132 Kernel configuration: [num_blocks, (pb, tp)] USING ONLY MY_T == 0 for now...
133 """
135 global_id, my_t = cuda.grid(2)
136 local_id = cuda.threadIdx.x
138 cell_length_x = sim_box[0] / cells_per_dimension[0]
139 loop_x_shift = loop_x_shift_function(sim_box, cell_length_x)
141 max_nbs = nblist.shape[1]-1 # Last index is used for storing number of neighbors
143 if global_id < num_part and my_t==0:
144 my_num_nbs = 0
145 my_num_exclusions = exclusions[global_id, -1]
147 for ix in range(-2-loop_x_addition,3+loop_x_addition,1):
148 for iy in range(-2,3,1):
149 # Correct handling of LEBC requires modifyng the loop over neighbor cells to take the box shift into account.
150 other_cell_y_unwrapped = my_cell[global_id, 1]+iy
151 y_wrap_cell = 1 if other_cell_y_unwrapped >= cells_per_dimension[1] else -1 if other_cell_y_unwrapped < 0 else 0
152 shifted_ix = ix + y_wrap_cell * loop_x_shift
154 for iz in range(-2,3,1):
155 other_index = (
156 (my_cell[global_id, 0]+shifted_ix)%cells_per_dimension[0],
157 (my_cell[global_id, 1]+iy)%cells_per_dimension[1],
158 (my_cell[global_id, 2]+iz)%cells_per_dimension[2])
159 other_global_id = cells[other_index]
160 while other_global_id >= 0: # To use tp>1: read tp particles ahead, and pick yours
161 if UtilizeNIII: # Could be done per cell basis...
162 #flag = other_global_id < global_id
163 TwodN = 2*(other_global_id - global_id)
164 flag = other_global_id < num_part and (0 < TwodN <= num_part or TwodN < -num_part)
165 else:
166 flag = other_global_id != global_id
167 if flag:
168 dist_sq = dist_sq_function(vectors[r_id][other_global_id], vectors[r_id][global_id], sim_box)
169 if dist_sq < cut_plus_skin*cut_plus_skin:
170 not_excluded = True # Check exclusion list. Do later ???
171 for k in range(my_num_exclusions):
172 if exclusions[global_id, k] == other_global_id:
173 not_excluded = False
174 if not_excluded:
175 my_num_nbs += 1
176 if my_num_nbs < max_nbs:
177 nblist[global_id, my_num_nbs-1] = other_global_id # Last entry is number of neighbors
178 other_global_id = next_particle_in_cell[other_global_id]
179 nblist[global_id, -1] = my_num_nbs
181 # Various house-keeping
182 for k in range(D):
183 r_ref[global_id, k] = vectors[r_id][global_id, k] # Store positions for wich nblist was updated ( used in nblist_check() )
184 #if local_id == 0 and my_t==0:
185 # cuda.atomic.add(nbflag, 0, -1) # nbflag[0] = 0 by when all blocks are done
186 if global_id == 0 and my_t==0:
187 cuda.atomic.add(nbflag, 2, 1) # Count how many updates are done in nbflag[2]
188 if my_num_nbs >= max_nbs: # Overflow detected, nbflag[1] should be checked later, and then
189 cuda.atomic.max(nbflag, 1, my_num_nbs) # re-allocate larger nb-list, and redo computations from last safe state
192 @cuda.jit(device=gridsync)
193 def clear_cells(vectors, sim_box, nbflag, cells_per_dimension, cells, my_cell, next_particle_in_cell):
194 """ Particles clears cell-list (only one (e.g. tail) actually needs to do this)
195 Kernel configuration: [num_blocks, (pb, tp)]
196 """
198 global_id, my_t = cuda.grid(2)
199 local_id = cuda.threadIdx.x
201 if global_id < num_part and my_t == 0:
202 for k in range(D):
203 if my_cell[global_id,k]<0:
204 print(global_id,k, my_cell[global_id,k],vectors[r_id][global_id,k])
207 if global_id < num_part and my_t == 0: # Change to simple Nullify kernel?
208 #cells[my_cell[global_id,:]] = -1 # Every body writes, but thats OK
209 cells[my_cell[global_id,0], my_cell[global_id,1], my_cell[global_id,2]] = -1 # Every body writes, but thats OK
210 next_particle_in_cell[global_id] = -1
213 if cuda.threadIdx.x == 0 and my_t==0: # One thread per threadblock decreases from numblocks
214 cuda.atomic.add(nbflag, 0, -1) # i.e., nbflag[0] = 0 by when all blocks are done
215 # (Not necessarry, after check moved to final kernel?)
217 @cuda.jit(device=gridsync)
218 def nblist_update(vectors, sim_box, cut_plus_skin, nbflag, nblist, r_ref, exclusions):
219 """ N^2 Update neighbor-list using numba.cuda
220 Kernel configuration: [num_blocks, (pb, tp)]
221 """
223 global_id, my_t = cuda.grid(2)
224 local_id = cuda.threadIdx.x
226 max_nbs = nblist.shape[1]-1 # Last index is used for storing number of neighbors
228 if global_id < num_part and my_t==0:
229 nblist[global_id, max_nbs] = 0 # Set number of neighbors (stored at index max_nbs) to zero
231 cuda.syncthreads() # wait for nblist[global_id, max_nbs] to be ready
233 if global_id < num_part:
234 my_num_exclusions = exclusions[global_id,-1]
235 for i in range(0, num_part, pb*tp): # Loop over blocks
236 for j in range(pb): # Loop over particles the pb in block
237 other_global_id = j + i + my_t*pb
238 if UtilizeNIII:
239 TwodN = 2*(other_global_id - global_id)
240 flag = other_global_id < num_part and (0 < TwodN <= num_part or TwodN < -num_part)
241 else:
242 flag = other_global_id != global_id and other_global_id < num_part
243 if flag:
244 dist_sq = dist_sq_function(vectors[r_id][other_global_id], vectors[r_id][global_id], sim_box)
245 if dist_sq < cut_plus_skin*cut_plus_skin:
246 not_excluded = True # Check exclusion list
247 for k in range(my_num_exclusions):
248 if exclusions[global_id, k] == other_global_id:
249 not_excluded = False
250 if not_excluded:
251 my_num_nbs = cuda.atomic.add(nblist, (global_id, max_nbs), 1) # Find next free index in nblist
252 if my_num_nbs < max_nbs:
253 nblist[global_id, my_num_nbs] = other_global_id # Last entry is number of neighbors
255 # Various house-keeping
256 if global_id < num_part and my_t==0:
257 for k in range(D):
258 r_ref[global_id, k] = vectors[r_id][global_id, k] # Store positions for wich nblist was updated ( used in nblist_check() )
259 #if local_id == 0 and my_t==0:
260 # cuda.atomic.add(nbflag, 0, -1) # nbflag[0] = 0 by when all blocks are done
261 if global_id == 0 and my_t==0:
262 cuda.atomic.add(nbflag, 2, 1) # Count how many updates are done in nbflag[2]
263 if my_num_nbs >= max_nbs: # Overflow detected, nbflag[1] should be checked later, and then
264 cuda.atomic.max(nbflag, 1, my_num_nbs) # re-allocate larger nb-list, and redo computations from last safe state
266 return
268 if gridsync==True:
269 # A device function, calling a number of device functions, using gridsync to syncronize
270 @cuda.jit( device=gridsync )
271 def check_and_update(grid, vectors, scalars, ptype, sim_box, nblist, nblist_parameters):
272 max_cut, skin, nbflag, r_ref, exclusions, cells_per_dimension, cells, my_cell, next_particle_in_cell, simbox_last_rebuild = nblist_parameters
273 nblist_check(vectors, sim_box, skin, r_ref, nbflag, simbox_last_rebuild, max_cut)
274 grid.sync()
275 if nbflag[0] > 0:
276 put_particles_in_cells(vectors, sim_box, nbflag, cells_per_dimension, cells, my_cell, next_particle_in_cell)
277 grid.sync()
278 nblist_update_from_linked_lists(vectors, sim_box, max_cut+skin, nbflag, cells_per_dimension, cells, my_cell, next_particle_in_cell, nblist, r_ref, exclusions)
279 #nblist_update(vectors, sim_box, max_cut+skin, nbflag, nblist, r_ref, exclusions)
280 grid.sync()
281 clear_cells(vectors, sim_box, nbflag, cells_per_dimension, cells, my_cell, next_particle_in_cell)
282 return
283 return check_and_update
285 else:
286 # A python function, making several kernel calls to syncronize
287 def check_and_update(grid, vectors, scalars, ptype, sim_box, nblist, nblist_parameters):
288 max_cut, skin, nbflag, r_ref, exclusions, cells_per_dimension, cells, my_cell, next_particle_in_cell, simbox_last_rebuild = nblist_parameters
289 nblist_check[num_blocks, (pb, 1)](vectors, sim_box, skin, r_ref, nbflag, simbox_last_rebuild, max_cut)
290 if nbflag[0] > 0:
291 put_particles_in_cells[num_blocks, (pb, 1)](vectors, sim_box, nbflag, cells_per_dimension, cells, my_cell, next_particle_in_cell)
292 nblist_update_from_linked_lists[num_blocks, (pb, 1)](vectors, sim_box, max_cut+skin, nbflag, cells_per_dimension, cells, my_cell, next_particle_in_cell, nblist, r_ref, exclusions)
293 #nblist_update[num_blocks, (pb, tp)](vectors, sim_box, max_cut+skin, nbflag, nblist, r_ref, exclusions) # Maybe tp = 1 ???
294 clear_cells[num_blocks, (pb, 1)](vectors, sim_box, nbflag, cells_per_dimension, cells, my_cell, next_particle_in_cell)
295 return
296 return check_and_update