Coverage for gamdpy/integrators/SLLOD.py: 98%
45 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1import numpy as np
2import numba
3import gamdpy as gp
4from numba import cuda
5import math
6from .integrator import Integrator
9## TO DO LIST FOR SLLOD (including LEBCs)
10# 1. implement gridsync=False case and check that it runs DONE 20/6
11# 2. Check conservation of KE DONE 25/6
12# 3. Figure out how to run the initialization kernel separately DONE 24/6
13# 4. Correct check of whether nb list needs to be built DONE 12/8
14# 5. Update images when box shift gets wrapped
15# 6. Save box-shift
17class SLLOD(Integrator):
18 """ The SLLOD integrator
20 Shear an atomic system in the xy-plane using the SLLOD equations.
22 Parameters
23 ----------
25 shear_rate : float
26 The shear rate of the system.
28 dt : float
29 The time step of the simulation.
31 """
32 def __init__(self, shear_rate, dt):
33 self.shear_rate = shear_rate
34 self.dt = dt
36 def get_params(self, configuration, interactions_params, verbose=False):
37 dt = np.float32(self.dt)
38 sr = np.float32(self.shear_rate)
40 # three 'groups' of three sum variables
41 self.thermostat_sums = np.zeros(9, dtype=np.float32)
43 # before first time-step, group 0, ie elements 0, 1, 2 needs to be initialized with sum_pxpy, sum_pypy, sum_p2
44 D, num_part = configuration.D, configuration.N
46 v = configuration['v']
47 m = configuration['m']
49 self.thermostat_sums[0] = np.sum(v[:,0] * v[:,1] * m)
50 self.thermostat_sums[1] = np.sum(v[:,1] * v[:,1] * m)
51 self.thermostat_sums[2] = np.sum(np.sum(v**2, axis=1)*m)
53 self.d_thermostat_sums = cuda.to_device(self.thermostat_sums)
55 return (dt,sr, self.d_thermostat_sums)
57 def get_kernel(self, configuration, compute_plan, compute_flags, interactions_kernel, verbose=False):
59 # Unpack parameters from configuration and compute_plan
60 D, num_part = configuration.D, configuration.N
61 pb, tp, gridsync = [compute_plan[key] for key in ['pb', 'tp', 'gridsync']]
62 num_blocks = (num_part - 1) // pb + 1
64 if verbose:
65 print(f'Generating SLLOD kernel for {num_part} particles in {D} dimensions:')
66 print(f'\tpb: {pb}, tp:{tp}, num_blocks:{num_blocks}')
67 print(f'\tNumber (virtual) particles: {num_blocks * pb}')
68 print(f'\tNumber of threads {num_blocks * pb * tp}')
70 # Unpack indices for vectors and scalars
72 r_id, v_id, f_id = [configuration.vectors.indices[key] for key in ['r', 'v', 'f']]
73 compute_k = compute_flags['K']
74 compute_fsq = compute_flags['Fsq']
75 m_id = configuration.sid['m']
76 if compute_k:
77 k_id = configuration.sid['K']
78 if compute_fsq:
79 fsq_id = configuration.sid['Fsq']
80 # was thinking that using a function could avoid synchronization
81 # issues for updating the boxshift. But now I'm not sure if it really
82 # makes sense to use a function (the same way that NVT
83 # does for temperature). There the temperature isn't stored anywhere.
84 # Here I'm pretty sure the box_shift has to be stored together with the
85 # other box details so interactions can always access it. So any
86 # function has to update that one location and then we have to worry
87 # about synchronization anyway
88 #def strain_function(time):
89 # strain = self.shear_rate*time
91 # JIT compile functions to be compiled into kernel
92 apply_PBC = numba.njit(configuration.simbox.get_apply_PBC())
93 update_box_shift = numba.njit(configuration.simbox.get_update_box_shift())
96 def call_update_box_shift(sim_box, integrator_params): # pragma: no cover
97 dt, sr, thermostat_sums = integrator_params
98 global_id, my_t = cuda.grid(2)
99 if global_id == 0 and my_t == 0:
100 delta_shift = sim_box[1] * sr * dt
101 update_box_shift(sim_box, delta_shift)
104 def integrate_sllod_b1(grid, vectors, scalars, integrator_params, time): # pragma: no cover
105 dt, sr, thermostat_sums = integrator_params
107 global_id, my_t = cuda.grid(2)
108 if global_id < num_part and my_t == 0:
109 my_v = vectors[v_id][global_id]
110 my_f = vectors[f_id][global_id]
111 my_m = scalars[global_id][m_id]
113 # read sums from group 0
114 sum_pxpy = thermostat_sums[0]
115 sum_pypy = thermostat_sums[1]
116 sum_p2 = thermostat_sums[2]
117 # compute coefficients
118 c1 = sr*sum_pxpy / sum_p2
119 c2 = sr*sr * sum_pypy / sum_p2
120 # g-factor for a half time-step
121 g_factor = 1./math.sqrt(1. - (c1*dt - 0.25 * c2 * dt**2)) # double precision
123 # update velocity - multiply by g_factor in double precision
124 my_v[0] = numba.float32(g_factor * (my_v[0] - 0.5*sr*dt*my_v[1]))
125 for k in range(1, D):
126 my_v[k] = numba.float32(g_factor * my_v[k])
128 # add to sums in group 1 needed for step B2
129 my_p2 = numba.float32(0.)
130 my_fp = numba.float32(0.)
131 my_f2 = numba.float32(0.)
132 for k in range(D):
133 my_p2 += my_v[k] * my_v[k] * my_m
134 my_fp += my_f[k] * my_v[k]
135 my_f2 += my_f[k] * my_f[k] / my_m
136 cuda.atomic.add(thermostat_sums, 3, my_p2)
137 cuda.atomic.add(thermostat_sums, 4, my_fp)
138 cuda.atomic.add(thermostat_sums, 5, my_f2)
140 # and reset group 2 sums to zero
141 if global_id == 0 and my_t == 0:
142 thermostat_sums[6] = 0.
143 thermostat_sums[7] = 0.
144 thermostat_sums[8] = 0.
147 def integrate_sllod_b2(grid, vectors, scalars, integrator_params, time): # pragma: no cover
148 dt,sr, thermostat_sums = integrator_params
150 global_id, my_t = cuda.grid(2)
151 if global_id < num_part and my_t == 0:
152 my_v = vectors[v_id][global_id]
153 my_f = vectors[f_id][global_id]
154 my_m = scalars[global_id][m_id]
155 if compute_fsq:
156 my_fsq = numba.float32(0.0) # force squared
157 # read sums from group 1
158 sum_p2 = thermostat_sums[3]
159 sum_fp = thermostat_sums[4]
160 sum_f2 = thermostat_sums[5]
162 # compute coefficients
163 alpha = sum_fp / sum_p2
164 beta = math.sqrt(sum_f2 / sum_p2)
165 h = (alpha + beta) / (alpha - beta)
166 e = math.exp(-beta * dt)
167 one = numba.float32(1.)
168 integrate_coefficient1 = (one - h) / (e - h/e)
169 integrate_coefficient2 = (one + h - e - h/e)/((one-h)*beta)
170 # update velocity
172 for k in range(D):
173 if compute_fsq:
174 my_fsq += my_f[k] * my_f[k]
175 my_v[k] = integrate_coefficient1 * (my_v[k] + integrate_coefficient2 * my_f[k] / my_m)
177 # add to sums in group 2
178 my_pxpy = my_v[0] * my_v[1] * my_m
179 my_pypy = my_v[1] * my_v[1] * my_m
180 my_p2 = numba.float32(0.)
181 for k in range(D):
182 my_p2 += my_v[k]**2
183 my_p2 *= my_m
185 cuda.atomic.add(thermostat_sums, 6, my_pxpy)
186 cuda.atomic.add(thermostat_sums, 7, my_pypy)
187 cuda.atomic.add(thermostat_sums, 8, my_p2)
189 if compute_fsq:
190 scalars[global_id][fsq_id] = my_fsq
192 # and reset group 0 sums to zero
193 if global_id == 0 and my_t == 0:
194 thermostat_sums[0] = numba.float32(0.)
195 thermostat_sums[1] = numba.float32(0.)
196 thermostat_sums[2] = numba.float32(0.)
199 def integrate_sllod_a_b1(grid, vectors, scalars, r_im, sim_box, integrator_params, time): # pragma: no cover
200 dt,sr, thermostat_sums = integrator_params
201 global_id, my_t = cuda.grid(2)
202 if global_id < num_part and my_t == 0:
203 my_r = vectors[r_id][global_id]
204 my_v = vectors[v_id][global_id]
205 my_f = vectors[f_id][global_id]
206 my_m = scalars[global_id][m_id]
208 # read sums from group 2
209 sum_pxpy = thermostat_sums[6]
210 sum_pypy = thermostat_sums[7]
211 sum_p2 = thermostat_sums[8]
212 # compute coefficients
213 c1 = sr*sum_pxpy / sum_p2
214 c2 = sr*sr * sum_pypy / sum_p2
215 # g-factor for a half time-step
216 g_factor = 1./math.sqrt(1. - (c1*dt - 0.25 * c2 * dt**2)) # double precision
218 # update velocity - multiply by g_factor in double precision
219 my_v[0] = numba.float32(g_factor * (my_v[0] - 0.5*sr*dt*my_v[1]))
220 for k in range(1, D):
221 my_v[k] = numba.float32(g_factor * my_v[k])
223 # update position and apply boundary conditions
224 my_r[0] += sr*dt*my_r[1] # rumd-3 has another term which seems to be incorrect (!)
225 # Here is the alternative version (DEBUG)
226 #my_r[0] += (my_v[0] + numba.float32(0.5) * sr*dt*my_v[1]) * dt
227 #for k in range(1, D): # DEBUG, was range(D)
228 for k in range(D):
229 my_r[k] += my_v[k] * dt
231 apply_PBC(my_r, r_im[global_id], sim_box)
233 # add to sums in group 0 for next time integrate_B1 is called (at the next time step)
234 my_pxpy = my_v[0] * my_v[1] * my_m
235 my_pypy = my_v[1] * my_v[1] * my_m
236 my_p2 = numba.float32(0.)
237 for k in range(D):
238 my_p2 += my_v[k]**2
239 my_p2 *= my_m
241 cuda.atomic.add(thermostat_sums, 0, my_pxpy)
242 cuda.atomic.add(thermostat_sums, 1, my_pypy)
243 cuda.atomic.add(thermostat_sums, 2, my_p2)
244 # store ke of this particle
245 if compute_k:
246 scalars[global_id][k_id] = numba.float32(0.5) * my_p2
248 # and reset group 1 sums to zero
249 if global_id == 0 and my_t == 0:
250 thermostat_sums[3] = numba.float32(0.)
251 thermostat_sums[4] = numba.float32(0.)
252 thermostat_sums[5] = numba.float32(0.)
256 call_update_box_shift = cuda.jit(call_update_box_shift)
257 integrate_sllod_b1 = cuda.jit(device=gridsync)(integrate_sllod_b1)
258 integrate_sllod_b2 = cuda.jit(device=gridsync)(integrate_sllod_b2)
259 integrate_sllod_a_b1 = cuda.jit(device=gridsync)(integrate_sllod_a_b1)
262 if gridsync: # pragma: no cover
263 def kernel(grid, vectors, scalars, r_im, sim_box, integrator_params, time, ptype):
264 integrate_sllod_b1(grid, vectors, scalars, integrator_params, time)
265 grid.sync()
266 integrate_sllod_b2(grid, vectors, scalars, integrator_params, time)
267 grid.sync()
268 call_update_box_shift(sim_box, integrator_params)
269 # need to apply wrap to images!
270 # (alternatively store an extra integer with the box to count
271 # how many times it's been wrapped)
272 grid.sync()
273 integrate_sllod_a_b1(grid, vectors, scalars, r_im, sim_box, integrator_params, time)
274 return
275 return cuda.jit(device=gridsync)(kernel)
276 else: # pragma: no cover
277 def kernel(grid, vectors, scalars, r_im, sim_box, integrator_params, time, ptype):
278 integrate_sllod_b1[num_blocks, (pb, 1)](grid, vectors, scalars, integrator_params, time)
279 integrate_sllod_b2[num_blocks, (pb, 1)](grid, vectors, scalars, integrator_params, time)
280 call_update_box_shift[1, (1, 1)](sim_box, integrator_params)
281 # need to apply wrap to images!
282 # (alternatively store an extra integer with the box to count
283 # how many times it's been wrapped)
284 integrate_sllod_a_b1[num_blocks, (pb, 1)](grid, vectors, scalars, r_im, sim_box, integrator_params, time)
285 return
287 return kernel