Coverage for gamdpy/interactions/nblist_linked_lists.py: 98%
61 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 18:46 +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)
53 self.next_particle_in_cell = -np.ones(self.N, dtype=np.int32) # -1 = no further particles in list
54 self.copy_to_device()
55 return (np.float32(self.max_cut), np.float32(self.skin), self.d_nbflag, self.d_r_ref, self.d_exclusions,
56 self.d_cells_per_dimension, self.d_cells, self.d_my_cell, self.d_next_particle_in_cell, self.d_simbox_last_rebuild)
58 def get_kernel(self, configuration, compute_plan, compute_flags, verbose=False, force_update=False):
60 # Unpack parameters from configuration and compute_plan
61 D, num_part = configuration.D, configuration.N
62 pb, tp, gridsync, UtilizeNIII = [compute_plan[key] for key in ['pb', 'tp', 'gridsync', 'UtilizeNIII']]
63 num_blocks = (num_part - 1) // pb + 1
64 loop_x_addition = configuration.simbox.get_loop_x_addition()
66 # Unpack indices for vectors and scalars to be compiled into kernel
67 r_id, f_id = [configuration.vectors.indices[key] for key in ['r', 'f']]
69 # JIT compile functions to be compiled into kernel
70 dist_sq_function = numba.njit(configuration.simbox.get_dist_sq_function())
71 dist_moved_exceeds_limit_function = numba.njit(configuration.simbox.get_dist_moved_exceeds_limit_function())
72 loop_x_shift_function = numba.njit(configuration.simbox.get_loop_x_shift_function())
74 @cuda.jit( device=gridsync )
75 def nblist_check(vectors, sim_box, skin, r_ref, nbflag, simbox_last_rebuild, cut): # pragma: no cover
76 """ Check validity of nblist, i.e. did any particle mode more than skin/2 since last nblist update?
77 Each tread-block checks the assigned particles (global_id)
78 nbflag[0] = 0 : No update needed
79 nbflag[0] = num_blocks : Update needed
80 Kernel configuration: [num_blocks, (pb, tp)]
81 """
83 global_id, my_t = cuda.grid(2)
84 if force_update: # nblist update forced (for benchmark or similar)
85 if global_id==0 and my_t==0:
86 nbflag[0]=num_blocks
88 if global_id < num_part and my_t==0:
89 if dist_moved_exceeds_limit_function(vectors[r_id][global_id], r_ref[global_id], sim_box, simbox_last_rebuild, skin, cut):
90 nbflag[0] = num_blocks
92 return
94 @cuda.jit(device=gridsync)
95 def put_particles_in_cells(vectors, sim_box, nbflag, cells_per_dimension, cells, my_cell, next_particle_in_cell): # pragma: no cover
96 """ Each particle computers its cells coordinates, and inserts itself in cell-list
97 Kernel configuration: [num_blocks, (pb, tp)]
98 """
100 global_id, my_t = cuda.grid(2)
101 if global_id < num_part and my_t == 0:
102 for k in range(D):
103 #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]
104 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]
105 #if my_cell[global_id,k]<0:
106 # print(global_id,k, my_cell[global_id,k],vectors[r_id][global_id,k])
107 index = (my_cell[global_id,0], my_cell[global_id,1], my_cell[global_id,2]) # 3D
108 next_particle_in_cell[global_id] = cuda.atomic.exch(cells, index, global_id) # index needs to be tuple when multidim
109 return
111 @cuda.jit(device=gridsync)
112 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): # pragma: no cover
113 """ Order N neighbor-list update from linked lists
114 Kernel configuration: [num_blocks, (pb, tp)] USING ONLY MY_T == 0 for now...
115 """
117 global_id, my_t = cuda.grid(2)
119 cell_length_x = sim_box[0] / cells_per_dimension[0]
120 loop_x_shift = loop_x_shift_function(sim_box, cell_length_x)
122 max_nbs = nblist.shape[1]-1 # Last index is used for storing number of neighbors
124 if global_id < num_part and my_t==0:
125 my_num_nbs = 0
126 my_num_exclusions = exclusions[global_id, -1]
128 for ix in range(-2-loop_x_addition,3+loop_x_addition,1):
129 for iy in range(-2,3,1):
130 # Correct handling of LEBC requires modifyng the loop over neighbor cells to take the box shift into account.
131 other_cell_y_unwrapped = my_cell[global_id, 1]+iy
132 y_wrap_cell = 1 if other_cell_y_unwrapped >= cells_per_dimension[1] else -1 if other_cell_y_unwrapped < 0 else 0
133 shifted_ix = ix + y_wrap_cell * loop_x_shift
135 for iz in range(-2,3,1):
136 other_index = (
137 (my_cell[global_id, 0]+shifted_ix)%cells_per_dimension[0],
138 (my_cell[global_id, 1]+iy)%cells_per_dimension[1],
139 (my_cell[global_id, 2]+iz)%cells_per_dimension[2])
140 other_global_id = cells[other_index]
141 while other_global_id >= 0: # To use tp>1: read tp particles ahead, and pick yours
142 if UtilizeNIII: # Could be done per cell basis...
143 #flag = other_global_id < global_id
144 TwodN = 2*(other_global_id - global_id)
145 flag = other_global_id < num_part and (0 < TwodN <= num_part or TwodN < -num_part)
146 else:
147 flag = other_global_id != global_id
148 if flag:
149 dist_sq = dist_sq_function(vectors[r_id][other_global_id], vectors[r_id][global_id], sim_box)
150 if dist_sq < cut_plus_skin*cut_plus_skin:
151 not_excluded = True # Check exclusion list. Do later ???
152 for k in range(my_num_exclusions):
153 if exclusions[global_id, k] == other_global_id:
154 not_excluded = False
155 if not_excluded:
156 my_num_nbs += 1
157 if my_num_nbs < max_nbs:
158 nblist[global_id, my_num_nbs-1] = other_global_id # Last entry is number of neighbors
159 other_global_id = next_particle_in_cell[other_global_id]
160 nblist[global_id, -1] = my_num_nbs
162 # Various house-keeping
163 for k in range(D):
164 r_ref[global_id, k] = vectors[r_id][global_id, k] # Store positions for wich nblist was updated ( used in nblist_check() )
165 #if local_id == 0 and my_t==0:
166 # cuda.atomic.add(nbflag, 0, -1) # nbflag[0] = 0 by when all blocks are done. Moved to clear_cells
167 if global_id == 0 and my_t==0:
168 cuda.atomic.add(nbflag, 2, 1) # Count how many updates are done in nbflag[2]
169 if my_num_nbs >= max_nbs: # Overflow detected, nbflag[1] should be checked later, and then
170 cuda.atomic.max(nbflag, 1, my_num_nbs) # re-allocate larger nb-list, and redo computations from last safe state
172 return
175 @cuda.jit(device=gridsync)
176 def clear_cells(vectors, sim_box, nbflag, cells_per_dimension, cells, my_cell, next_particle_in_cell): # pragma: no cover
177 """ Particles clears cell-list (only one (e.g. tail) actually needs to do this)
178 Kernel configuration: [num_blocks, (pb, tp)]
179 """
181 global_id, my_t = cuda.grid(2)
183 if global_id < num_part and my_t == 0: # Change to simple Nullify kernel?
184 cells[my_cell[global_id,0], my_cell[global_id,1], my_cell[global_id,2]] = -1 # Every body writes, but thats OK
185 next_particle_in_cell[global_id] = -1
188 if cuda.threadIdx.x == 0 and my_t==0: # One thread per threadblock decreases from numblocks
189 cuda.atomic.add(nbflag, 0, -1) # i.e., nbflag[0] = 0 by when all blocks are done
191 return
193 if gridsync:
194 # A device function, calling a number of device functions, using gridsync to syncronize
195 @cuda.jit( device=gridsync )
196 def check_and_update(grid, vectors, scalars, ptype, sim_box, nblist, nblist_parameters): # pragma: no cover
197 max_cut, skin, nbflag, r_ref, exclusions, cells_per_dimension, cells, my_cell, next_particle_in_cell, simbox_last_rebuild = nblist_parameters
198 nblist_check(vectors, sim_box, skin, r_ref, nbflag, simbox_last_rebuild, max_cut)
199 grid.sync()
200 if nbflag[0] > 0:
201 put_particles_in_cells(vectors, sim_box, nbflag, cells_per_dimension, cells, my_cell, next_particle_in_cell)
202 grid.sync()
203 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)
204 grid.sync()
205 clear_cells(vectors, sim_box, nbflag, cells_per_dimension, cells, my_cell, next_particle_in_cell)
206 return
207 return check_and_update
209 else:
210 # A python function, making several kernel calls to syncronize
211 def check_and_update(grid, vectors, scalars, ptype, sim_box, nblist, nblist_parameters):
212 max_cut, skin, nbflag, r_ref, exclusions, cells_per_dimension, cells, my_cell, next_particle_in_cell, simbox_last_rebuild = nblist_parameters
213 nblist_check[num_blocks, (pb, 1)](vectors, sim_box, skin, r_ref, nbflag, simbox_last_rebuild, max_cut)
214 if nbflag[0] > 0:
215 put_particles_in_cells[num_blocks, (pb, 1)](vectors, sim_box, nbflag, cells_per_dimension, cells, my_cell, next_particle_in_cell)
216 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)
217 clear_cells[num_blocks, (pb, 1)](vectors, sim_box, nbflag, cells_per_dimension, cells, my_cell, next_particle_in_cell)
218 return
219 return check_and_update