Coverage for gamdpy/interactions/pair_potential.py: 45%
204 statements
« prev ^ index » next coverage.py v7.4.4, created at 2025-06-14 15:25 +0200
1import numpy as np
2import numba
3import math
4from numba import cuda
5import matplotlib.pyplot as plt
6import gamdpy as gp
7from .interaction import Interaction
9class PairPotential(Interaction):
10 """ Pair potential """
12 def __init__(self, pairpotential_function, params, max_num_nbs, exclusions=None):
13 def params_function(i_type, j_type, params):
14 result = params[i_type, j_type] # default: read from params array
15 return result
17 self.pairpotential_function = pairpotential_function
18 self.params_function = params_function
19 self.params_user = params
20 self.exclusions = exclusions
21 self.max_num_nbs = max_num_nbs
23 def convert_user_params(self):
24 # Upgrade any scalar parameters to 1x1 numpy array
25 num_params = len(self.params_user)
26 params_list = []
27 for parameter in self.params_user:
28 if np.isscalar(parameter):
29 params_list.append(np.ones((1,1))*parameter)
30 else:
31 params_list.append(np.array(parameter, dtype=np.float32))
33 # Ensure all parameters are the right format (num_types x num_types) numpy arrays
34 num_types = params_list[0].shape[0]
35 for parameter in params_list:
36 assert len(parameter.shape) == 2
37 assert parameter.shape[0] == num_types
38 assert parameter.shape[1] == num_types
40 # Convert params to the format required by kernels (num_types x num_types) array of tuples (p0, p1, ..., cutoff)
41 params = np.zeros((num_types, num_types), dtype="f,"*num_params)
42 for i in range(num_types):
43 for j in range(num_types):
44 plist = []
45 for parameter in params_list:
46 plist.append(parameter[i,j])
47 params[i,j] = tuple(plist)
49 max_cut = np.float32(np.max(params_list[-1]))
51 return params, max_cut
53 def plot(self, xlim=None, ylim=(-3,6), figsize=(8,4), names=None):
54 params, max_cut = self.convert_user_params()
55 num_types = len(params[0])
56 if names==None:
57 names = np.arange(num_types)
58 plt.figure(figsize=figsize)
59 for i in range(num_types):
60 for j in range(num_types):
61 r = np.linspace(0, params[i,j][-1], 1000)
62 u, s, lap = self.pairpotential_function(r, params[i,j])
63 plt.plot(r, u, label=f'{names[i]} - {names[j]}')
64 if xlim!=None:
65 plt.xlim(xlim)
66 plt.ylim(ylim)
67 plt.xlabel('Pair distance')
68 plt.ylabel('Pair potential')
69 plt.legend()
70 plt.show()
72 def evaluate_potential_function(self, r, types):
73 params, max_cut = self.convert_user_params()
74 u, s, lap = self.pairpotential_function(r, params[types[0], types[1]])
75 return u
78 def check_datastructure_validity(self) -> bool:
79 nbflag = self.nblist.d_nbflag.copy_to_host()
80 if nbflag[0] != 0 or nbflag[1] != 0:
81 raise RuntimeError(f'Neighbor-list is invalid. Try allocating space for more neighbors (max_num_nbs in PairPot object). Allocated size: {self.max_num_nbs}, but {nbflag[1]+1} neighbours found. {nbflag=}.')
82 return True
84 def get_params(self, configuration: gp.Configuration, compute_plan: dict, verbose=False) -> tuple:
86 self.params, max_cut = self.convert_user_params()
87 self.d_params = cuda.to_device(self.params)
89 if compute_plan['nblist'] == 'N squared':
90 self.nblist = gp.NbList2(configuration, self.exclusions, self.max_num_nbs)
91 elif compute_plan['nblist'] == 'linked lists':
92 self.nblist = gp.NbListLinkedLists(configuration, self.exclusions, self.max_num_nbs)
93 else:
94 raise ValueError(f"No lblist called: {compute_plan['nblist']}. Use either 'N squared' or 'linked lists'")
95 nblist_params = self.nblist.get_params(max_cut, compute_plan, verbose)
97 return (self.d_params, self.nblist.d_nblist, nblist_params)
99 def get_kernel(self, configuration: gp.Configuration, compute_plan: dict, compute_flags: dict[str,bool], verbose=False):
100 num_cscalars = configuration.num_cscalars
102 compute_u = compute_flags['U']
103 compute_w = compute_flags['W']
104 compute_lap = compute_flags['lapU']
105 compute_stresses = compute_flags['stresses']
107 # Unpack parameters from configuration and compute_plan
108 D, num_part = configuration.D, configuration.N
109 pb, tp, gridsync, UtilizeNIII = [compute_plan[key] for key in ['pb', 'tp', 'gridsync', 'UtilizeNIII']]
110 num_blocks = (num_part - 1) // pb + 1
112 if verbose:
113 print(f'\tpb: {pb}, tp:{tp}, num_blocks:{num_blocks}')
114 print(f'\tNumber (virtual) particles: {num_blocks*pb}')
115 print(f'\tNumber of threads {num_blocks*pb*tp}')
116 if compute_stresses:
117 print('\tIncluding computation of stress tensor in pair potential')
118 # Unpack indices for vectors and scalars to be compiled into kernel
119 r_id, f_id = [configuration.vectors.indices[key] for key in ['r', 'f']]
121 if compute_u:
122 u_id = configuration.sid['U']
123 if compute_w:
124 w_id = configuration.sid['W']
125 if compute_lap:
126 lap_id = configuration.sid['lapU']
128 if compute_stresses:
129 sx_id = configuration.vectors.indices['sx']
130 if D > 1:
131 sy_id = configuration.vectors.indices['sy']
132 if D > 2:
133 sz_id = configuration.vectors.indices['sz']
134 if D > 3:
135 sw_id = configuration.vectors.indices['sw']
139 pairpotential_function = self.pairpotential_function
141 if UtilizeNIII:
142 virial_factor_NIII = numba.float32( 1.0/configuration.D)
143 def pairpotential_calculator(ij_dist, ij_params, dr, my_f, cscalars, my_stress, f, other_id):
144 u, s, umm = pairpotential_function(ij_dist, ij_params)
145 for k in range(D):
146 cuda.atomic.add(f, (other_id, k), dr[k]*s)
147 my_f[k] = my_f[k] - dr[k]*s # Force
148 if compute_w:
149 cscalars[w_id] += dr[k]*dr[k]*s*virial_factor_NIII # Virial
150 if compute_stresses:
151 for k2 in range(D):
152 my_stress[k,k2] -= dr[k]*dr[k2]*s
154 if compute_u:
155 cscalars[u_id] += u # Potential energy
156 if compute_lap:
157 cscalars[lap_id] += (numba.float32(1-D)*s + umm)*numba.float32( 2.0 ) # Laplacian
160 return
162 else:
163 virial_factor = numba.float32( 0.5/configuration.D )
164 def pairpotential_calculator(ij_dist, ij_params, dr, my_f, cscalars, my_stress, f, other_id):
165 u, s, umm = pairpotential_function(ij_dist, ij_params)
166 half = numba.float32(0.5)
167 for k in range(D):
168 my_f[k] = my_f[k] - dr[k]*s # Force
169 if compute_w:
170 cscalars[w_id] += dr[k]*dr[k]*s*virial_factor # Virial
171 if compute_stresses:
172 for k2 in range(D):
173 my_stress[k,k2] -= half*dr[k]*dr[k2]*s # stress tensor
174 if compute_u:
175 cscalars[u_id] += half*u # Potential energy
176 if compute_lap:
177 cscalars[lap_id] += numba.float32(1-D)*s + umm # Laplacian
178 return
180 ptype_function = numba.njit(configuration.ptype_function)
181 params_function = numba.njit(self.params_function)
182 pairpotential_calculator = numba.njit(pairpotential_calculator)
183 dist_sq_dr_function = numba.njit(configuration.simbox.get_dist_sq_dr_function())
185 @cuda.jit( device=gridsync )
186 def calc_forces(vectors, cscalars, ptype, sim_box, nblist, params):
187 """ Calculate forces as given by pairpotential_calculator() (needs to exist in outer-scope) using nblist
188 Kernel configuration: [num_blocks, (pb, tp)]
189 """
191 my_block = cuda.blockIdx.x
192 local_id = cuda.threadIdx.x
193 global_id = my_block*pb + local_id
194 my_t = cuda.threadIdx.y
196 max_nbs = nblist.shape[1]-1
198 #if global_id < num_part and my_t==0: # Initialize global variables. Should be controlled by flag if more pair-potentials
199 # for k in range(num_cscalars):
200 # cscalars[global_id, k] = numba.float32(0.0)
202 my_f = cuda.local.array(shape=D,dtype=numba.float32)
203 my_dr = cuda.local.array(shape=D,dtype=numba.float32)
204 my_cscalars = cuda.local.array(shape=num_cscalars, dtype=numba.float32)
205 if compute_stresses:
206 my_stress = cuda.local.array(shape=(D,D), dtype=numba.float32)
207 else:
208 my_stress = cuda.local.array(shape=(1,1), dtype=numba.float32)
210 if global_id < num_part:
211 for k in range(D):
212 #my_r[k] = r[global_id, k]
213 my_f[k] = numba.float32(0.0)
214 if compute_stresses:
215 for k2 in range(D):
216 my_stress[k,k2] = numba.float32(0.0)
217 for k in range(num_cscalars):
218 my_cscalars[k] = numba.float32(0.0)
219 my_type = ptype_function(global_id, ptype)
221 cuda.syncthreads() # Make sure initializing global variables to zero is done
223 if global_id < num_part:
224 for i in range(my_t, nblist[global_id, max_nbs], tp):
225 other_id = nblist[global_id, i]
226 other_type = ptype_function(other_id, ptype)
227 dist_sq = dist_sq_dr_function(vectors[r_id][other_id], vectors[r_id][global_id], sim_box, my_dr)
228 ij_params = params_function(my_type, other_type, params)
229 cut = ij_params[-1]
230 if dist_sq < cut*cut:
231 pairpotential_calculator(math.sqrt(dist_sq), ij_params, my_dr, my_f, my_cscalars, my_stress, vectors[f_id], other_id)
232 for k in range(D):
233 cuda.atomic.add(vectors[f_id], (global_id, k), my_f[k])
234 if compute_stresses:
235 cuda.atomic.add(vectors[sx_id], (global_id, k), my_stress[0,k])
236 if D > 1:
237 cuda.atomic.add(vectors[sy_id], (global_id, k), my_stress[1,k])
238 if D > 2:
239 cuda.atomic.add(vectors[sz_id], (global_id, k), my_stress[2,k])
240 if D > 3:
241 cuda.atomic.add(vectors[sw_id], (global_id, k), my_stress[3,k])
243 for k in range(num_cscalars):
244 cuda.atomic.add(cscalars, (global_id, k), my_cscalars[k])
246 return
248 nblist_check_and_update = self.nblist.get_kernel(configuration, compute_plan, compute_flags, verbose)
250 if gridsync:
251 # A device function, calling a number of device functions, using gridsync to syncronize
252 @cuda.jit( device=gridsync )
253 def compute_interactions(grid, vectors, scalars, ptype, sim_box, interaction_parameters):
254 params, nblist, nblist_parameters = interaction_parameters
255 nblist_check_and_update(grid, vectors, scalars, ptype, sim_box, nblist, nblist_parameters)
256 grid.sync()
257 calc_forces(vectors, scalars, ptype, sim_box, nblist, params)
258 return
259 return compute_interactions
261 else:
262 # A python function, making several kernel calls to syncronize
263 def compute_interactions(grid, vectors, scalars, ptype, sim_box, interaction_parameters):
264 params, nblist, nblist_parameters = interaction_parameters
265 nblist_check_and_update(grid, vectors, scalars, ptype, sim_box, nblist, nblist_parameters)
266 calc_forces[num_blocks, (pb, tp)](vectors, scalars, ptype, sim_box, nblist, params)
267 return
268 return compute_interactions