Coverage for gamdpy/integrators/NVU_RT.py: 3%
723 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1import h5py
2import math
3import numpy as np
4import numba as nb
5from numba import cuda
6import math
7from typing import Any, Literal, Tuple
8import numpy.typing as npt
9from .integrator import Integrator
12# Just to get rid of certain IDE errors but it is not important
13CudaArray = npt.NDArray[Any]
16class NVU_RT(Integrator):
17 """Potential energy conserving integrator.
18 Calculate the positions by reflecting on the constant potential energy
19 hypersurface and doing Ray Tracing (RT).
21 Uses parabola approximations, newton-rhapson or bisection to perform raytrcing.
23 Parameters
24 ----------
26 target_u : float
27 Target Potential Energy (U_0) to maintain constant along the simulation
29 threshold : float
30 Width of the potential energy "shell" relative to U_0.
31 [Iterative Method] When the potential energy of a certain configuration is within threshold if the potential energy
32 of the fist step, iteration is finished. |U(t) / U_0 - 1| < threshold. It needs to be small enough so that the
33 iterative method is precise enough (it can get very chaotic if it is too high). For example: 1e-6.
35 initial_step : float, default=0.1
36 [Iterative Metod] Initial step in configuration space so that x = positions + d * initial_step
37 is the initial guess for the root algorithm. `d` is the velocity normalized. For method bisection,
38 it needs to be big enough so that it steps away from the initial position but small enough so that it doesn't
39 reach out of the surface. If it is two big the algorithm will try to correct it. For method parabola it needs
40 to be as big as possible but it needs to be a point that is below U_0. For example: 0.01.
41 It is better to be based on the density so a good value is something like 0.5 / rho^(1/3).
43 initial_step_if_high : float, default=0.01
44 [Iterative Metod] Initial step if the potential energy of the initial configuration (time == 0) is higher than the
45 target potential energy. It should be high enough to "enter" the potential energy surface U = U_0. For example: same
46 setting as initial step.
48 step : float, default=1
49 [Iterative Method] (only for bisection method) Step to look for a point with u > u0. For example: 1.
50 As in initial it is better for it to based on the density.
52 max_steps : int, default=20
53 [Iterative Method] Maximum calls to the interactions kernel to find a point outside the hypersurface.
54 Good enough value is maybe 10 for parabola method and 100 for bisection method.
56 max_initial_step_corrections : int, default=20
57 [Iterative Method] If initial_step is too big the algorithm will try to correct it this amount of times.
58 At the nth correction step initial_step_n = initial_step_0 * (1/2)^n. In the first iteration if U > U0
59 then initial step is corrected to make it bigger (s_n = s_0 * 2^n). For example: 10 (max correction will
60 be approximately 1000).
62 max_abs_val : float, default=2
63 [Iterative Metod] (only for bisection) Some potential energy functions increas rapidly at a certain configurations.
64 To prevent numerical errors, if the absolute potential energy of a given configuration, |U|,
65 is above U_0 * max_abs_val, the iterative method will discard that configuration as invalid and reach
66 "less far" into configurational space. For example: 2.
68 eps : float, default=1e-7
69 [Iterative Method] Because of numerical inaccuracies, it could happen that the same value calculated twice
70 once is positive and once is negative. Values of the potential energy relative to the target potential enery
71 with |x| < eps are considered neither positive or negative in the algo. Needs to be smaller than threshold.
73 debug_print : bool, default=False
74 If a root is not found, the 0th thread prints useful debugging information if debug_print is enabled.
76 mode : {"reflection", "no-inertia", "reflection-mass_scaling"}, default = "reflection-mass_scaling"
77 Mode to perform the reflection.
78 `reflection-mass_scaling` applies a correction that takes into account the mass of
79 each particle. If the setup does not include particles with different masses, ``reflection``
80 will be faster (not that much).
81 `no-inertia` is a testing feature: instead of reflecitng velocities in the hyper surface the new velocities
82 follow the direction of the normal vector (the force).
84 save_path_u : optional, default=False
85 Save the potential energy between two consecutive points in the iteration
87 raytracing_method : {"parabola", "parabola-newton", "bisection"}, default = "parabola"
88 Methodology to find the next point in the surface with same potential energy
90 float_type : {"64", "32"}, default = "64"
91 Float type for the potential energy. Higher threshold can work with float32
93 """
95 outputs = ("its", "cos_v_f", "time", "dt", )
97 def __init__(
98 self,
99 target_u: float,
100 threshold: float,
101 initial_step: float = 0.1,
102 initial_step_if_high: float = 0.01,
103 step: float = 1,
104 max_steps: int = 20,
105 max_initial_step_corrections: int = 20,
106 max_abs_val: float = 2,
107 eps: float = 1e-7,
108 debug_print: bool = False,
109 mode: Literal["reflection", "no-inertia", "reflection-mass_scaling"] = "reflection-mass_scaling",
110 save_path_u: bool = False,
111 raytracing_method: Literal["parabola", "parabola-newton", "bisection"] = "parabola",
112 float_type: Literal["32", "64"] = "64"
113 ):
114 if str(float_type) == "32":
115 u_dtype = np.float32
116 elif str(float_type) == "64":
117 u_dtype = np.float64
118 else:
119 raise ValueError(f"Expected \"64\" or \"32\" for parameter `float_type`, but got {float_type}")
121 self.target_u = u_dtype(target_u)
122 self.d_pot_energy = cuda.to_device(np.array([np.nan], dtype=u_dtype))
123 self.max_abs_val = np.float32(max_abs_val)
124 self.threshold = np.float32(threshold)
125 self.step = np.float32(step)
126 self.eps = np.float32(eps)
127 if self.threshold <= 0:
128 raise ValueError(f"`threshold` has to be strictly bigger than 0, got {threshold}")
129 if self.eps >= self.threshold:
130 raise ValueError(f"`eps` has to be smaller bigger than threshold ({threshold}), got {eps}")
131 self.max_steps = np.int32(max_steps)
132 self.max_initial_step_corrections = np.int32(max_initial_step_corrections)
133 self.initial_step = np.float32(initial_step)
134 self.d_initial_step = cuda.to_device(np.array([self.initial_step], dtype=np.float32))
135 self.initial_step_if_high = np.float32(initial_step_if_high)
136 # Simluation requires that integrators have dt
137 self.dt = 1
138 self.d_scalars_shared = cuda.device_array(16, dtype=np.float32) # type: ignore
139 self.output_ids = {name: idx for idx, name in enumerate(self.outputs)}
140 self.d_integrator_output = cuda.device_array(len(self.outputs), dtype=np.float32) # type: ignore
141 self.d_path_u = cuda.to_device(np.zeros((100, 10, 5), dtype=np.float32) + np.nan)
142 self.d_step = cuda.to_device(np.zeros(1, dtype=np.int32))
143 self.debug_print = np.bool_(debug_print)
144 self.mode = mode
145 self.raytracing_method = raytracing_method
146 self.save_path_u = np.bool_(save_path_u)
147 self.d_broken_simulation = cuda.to_device(np.zeros(1, dtype=np.bool_))
148 self.d_last_a = cuda.to_device(np.empty(1, dtype=np.float32) + np.nan)
149 self.d_u_higher_than_target_in_time_0 = cuda.to_device(np.empty(1, dtype=np.bool_))
151 def get_params(self, configuration, interaction_params, verbose = False):
152 # NOTE: for some reason the first param has to be delta time
153 return (
154 self.dt,
155 self.d_integrator_output,
156 self.d_initial_step,
157 interaction_params,
158 self.d_scalars_shared,
159 self.d_pot_energy,
160 self.d_path_u,
161 self.d_step,
162 self.d_broken_simulation,
163 self.d_last_a,
164 self.d_u_higher_than_target_in_time_0,
165 )
167 def update_at_end_of_timeblock(self, storage: str, nblocks: int, block: int):
168 if self.save_path_u:
169 self.d_step[0] = 0
170 if block == 0:
171 with h5py.File(storage, 'a') as f:
172 if "path_u" in f.keys():
173 del f["path_u"]
174 f.create_dataset('path_u', shape=(nblocks, *self.d_path_u.shape),
175 chunks=(1, *self.d_path_u.shape), dtype=np.float32)
176 with h5py.File(storage, "a") as f:
177 f["path_u"][block, :] = self.d_path_u.copy_to_host() # type: ignore
179 def get_kernel(self, configuration, compute_plan, compute_flags, interactions_kernel, verbose=False):
180 # Unpack parameters from configuration and compute_plan
181 num_dim, num_part = configuration.D, configuration.N
182 pb, tp, gridsync = [compute_plan[key] for key in ['pb', 'tp', 'gridsync']]
183 num_blocks = (num_part - 1) // pb + 1
185 if verbose:
186 print(f'Generating NVU kernel for {num_part} particles in {num_dim} dimensions:')
187 print(f'\tpb: {pb}, tp:{tp}, num_blocks:{num_blocks}')
188 print(f'\tNumber (virtual) particles: {num_blocks * pb}')
189 print(f'\tNumber of threads {num_blocks * pb * tp}')
191 # Unpack indices for vectors and scalars to be compiled into kernel
192 r_id, v_id, f_id = [configuration.vectors.indices[key] for key in ['r', 'v', 'f']]
193 m_id = configuration.sid['m']
194 if not compute_flags['U'] or not compute_flags['Fsq']:
195 raise ValueError('NVU_RT requires both U and Fsq in scalars')
196 u_id = configuration.sid['U']
197 fsq_id = configuration.sid['Fsq']
200 (
201 forces_sql_id,
202 vel_sql_id,
203 dot_v_f_id,
204 should_break_id,
205 should_return_id,
206 reached_max_abs_val_id,
207 should_move_r_copy_id,
208 should_flip_velocities_id,
209 second_derivative_id,
210 first_derivative_id,
211 alpha_den_id,
212 *debug_ids
213 ) = range(self.d_scalars_shared.shape[0])
214 debug_ids = tuple(debug_ids)
216 o_its, o_cos_v_f, o_time, o_dt = (self.output_ids[name] for name in ["its", "cos_v_f", "time", "dt", ])
218 # JIT compile functions to be compiled into kernel
219 apply_PBC = nb.jit(configuration.simbox.get_apply_PBC())
221 save_path_u = self.save_path_u
222 save_path_max_saves = self.d_path_u.shape[0]
223 save_path_u_divisions = self.d_path_u.shape[1]
225 if gridsync:
226 @cuda.jit(device=gridsync)
227 def kernel(
228 grid, vectors, scalars, r_im, sim_box,
229 integrator_params, time, ptype,
230 ):
231 (dt, d_integrator_output, d_initial_step, interaction_params, d_scalars_shared, d_pot_energy, d_path_u,
232 d_step, d_broken_simulation, d_last_a, d_u_higher_than_target_in_time_0) = integrator_params
233 if time > 0:
234 d_u_higher_than_target_in_time_0[0] = False
236 if d_broken_simulation[0]:
237 return
239 global_id, my_t = cuda.grid(2) # type: ignore
240 my_m = scalars[global_id, m_id]
241 velocities = vectors[v_id]
242 my_v = velocities[global_id]
243 forces = vectors[f_id]
244 my_f = forces[global_id]
245 positions = vectors[r_id]
246 my_r = positions[global_id]
247 my_r_im = r_im[global_id]
248 its: CudaArray = cuda.local.array(1, dtype=np.int32) # type: ignore
249 its[0] = 0
251 if global_id == 0 and my_t == 0:
252 for k in range(d_scalars_shared.shape[0]):
253 d_scalars_shared[k] = 0
254 grid.sync()
255 if global_id == 0 and my_t == 0:
256 d_integrator_output[o_time] = time
258 if global_id < num_part and my_t == 0:
259 x = np.float32(0)
260 for k in range(num_dim):
261 x += my_f[k] * my_f[k]
262 cuda.atomic.add(d_scalars_shared, forces_sql_id, x) # type: ignore
263 scalars[global_id][fsq_id] = x
264 get_dot_in_conf_space(my_v, my_v, d_scalars_shared, vel_sql_id)
265 get_dot_in_conf_space(my_f, my_v, d_scalars_shared, dot_v_f_id)
266 grid.sync()
267 if global_id == 0 and my_t == 0:
268 if d_scalars_shared[dot_v_f_id] > 0:
269 d_broken_simulation[0] = True
270 vel_l = math.sqrt(d_scalars_shared[vel_sql_id])
272 calculate_next_velocities(grid, my_f, d_scalars_shared, my_v, my_m)
273 grid.sync()
275 if global_id == 0 and my_t == 0:
276 cos_v_f = d_scalars_shared[dot_v_f_id] / vel_l / math.sqrt(d_scalars_shared[forces_sql_id])
277 d_integrator_output[o_cos_v_f] = cos_v_f
279 my_dir: CudaArray = cuda.local.array(num_dim, velocities.dtype) # type: ignore
280 if global_id < num_part and my_t == 0:
281 for k in range(num_dim):
282 my_dir[k] = my_v[k] / vel_l
284 ## R_{i+1} = R_i + t * V_{i+1} | U(R_{i+1}) = U_0
285 #################################################
286 r0: CudaArray = cuda.local.array(num_dim, my_r.dtype) # type: ignore
287 r0_im: CudaArray = cuda.local.array(num_dim, my_r_im.dtype) # type: ignore
288 copy_positions_and_images(my_r, my_r_im, r0, r0_im)
290 step_x = raytracing_kernel(
291 time,
292 interactions_kernel, grid, vectors, scalars, ptype,
293 sim_box, interaction_params,
294 r_im,
295 d_pot_energy,
296 its,
297 d_last_a,
298 d_scalars_shared,
299 vel_l,
300 my_dir,
301 d_u_higher_than_target_in_time_0,
302 d_initial_step,
303 )
305 if save_path_u and d_step[0] < save_path_max_saves:
306 save_my_r: CudaArray = cuda.local.array(num_dim, my_r.dtype) # type: ignore
307 save_my_r_im: CudaArray = cuda.local.array(num_dim, my_r_im.dtype) # type: ignore
308 copy_positions_and_images(my_r, my_r_im, save_my_r, save_my_r_im)
310 save_potential_energy_path(
311 save_path_u_divisions, step_x, d_path_u, d_step, d_scalars_shared,
312 r0, r0_im, my_r, my_r_im, my_dir,
313 interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its
314 )
316 copy_positions_and_images(save_my_r, save_my_r_im, my_r, my_r_im)
317 calculate_potential_energy(interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its)
318 if global_id == 0 and my_t == 0:
319 d_integrator_output[o_dt] = step_x / vel_l
320 d_integrator_output[o_its] = its[0]
321 else:
322 raise NotImplementedError()
324 (
325 get_dot_in_conf_space,
326 copy_positions_and_images,
327 calculate_potential_energy,
328 add_step_in_dir,
329 ) = util_functions(gridsync, num_part, num_dim, apply_PBC, u_id)
330 calculate_next_velocities = self.get_update_velocities_kernel(
331 gridsync=gridsync,
332 num_part=num_part,
333 num_dim=num_dim,
334 dot_v_f_id=dot_v_f_id,
335 forces_sql_id=forces_sql_id,
336 vel_sql_id=vel_sql_id,
337 alpha_den_id=alpha_den_id,
338 )
340 @cuda.jit(device=gridsync)
341 def save_potential_energy_path(
342 save_path_u_divisions, delta_x, d_path_u, d_step, d_scalars_shared,
343 r0, r0_im, my_r, my_r_im, my_dir,
344 interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its
345 ):
346 global_id, my_t = cuda.grid(2) # type: ignore
347 it = np.int32(0)
348 while it < save_path_u_divisions:
349 path_step = it * delta_x / (save_path_u_divisions - 1)
350 add_step_in_dir(r0, r0_im, my_dir, path_step, my_r, my_r_im, sim_box)
351 grid.sync()
352 calculate_potential_energy(interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its)
353 if global_id == 0 and my_t == 0:
354 d_path_u[d_step[0], it, 0] = path_step
355 d_path_u[d_step[0], it, 1] = d_pot_energy[0]
356 d_path_u[d_step[0], it, 2] = 0
357 d_path_u[d_step[0], it, 3] = d_scalars_shared[vel_sql_id]
358 # v · f is of the non-reflected v so we change signs
359 d_path_u[d_step[0], it, 4] = d_scalars_shared[dot_v_f_id]
360 d_path_u[d_step[0], it, 5] = delta_x
361 grid.sync()
362 if global_id < num_part and my_t == 0:
363 cuda.atomic.add(d_path_u, (d_step[0], it, 2), scalars[global_id][lap_id]) # type: ignore
364 grid.sync()
365 it += 1
366 if global_id == 0 and my_t == 0:
367 d_step[0] += 1
369 if self.raytracing_method == "bisection":
370 raytracing_kernel = self.get_kernel_bisection(
371 gridsync=gridsync,
372 num_dim=num_dim,
373 r_id=r_id,
374 should_break_id=should_break_id,
375 should_move_r_copy_id=should_move_r_copy_id,
376 reached_max_abs_val_id=reached_max_abs_val_id,
377 should_return_id=should_return_id,
378 copy_positions_and_images=copy_positions_and_images,
379 add_step_in_dir=add_step_in_dir,
380 calculate_potential_energy=calculate_potential_energy,
381 )
382 elif self.raytracing_method == "parabola":
383 raytracing_kernel = self.get_kernel_parabola(
384 gridsync=gridsync,
385 num_part=num_part,
386 num_dim=num_dim,
387 r_id=r_id,
388 u_id=u_id,
389 dot_v_f_id=dot_v_f_id,
390 should_break_id=should_break_id,
391 copy_positions_and_images=copy_positions_and_images,
392 add_step_in_dir=add_step_in_dir,
393 calculate_potential_energy=calculate_potential_energy,
394 )
395 elif self.raytracing_method == "parabola-newton":
396 raytracing_kernel = self.get_kernel_parabola_newton(
397 gridsync=gridsync,
398 num_part=num_part,
399 num_dim=num_dim,
400 r_id=r_id,
401 f_id=f_id,
402 u_id=u_id,
403 dot_v_f_id=dot_v_f_id,
404 should_break_id=should_break_id,
405 first_derivative_id=first_derivative_id,
406 copy_positions_and_images=copy_positions_and_images,
407 add_step_in_dir=add_step_in_dir,
408 calculate_potential_energy=calculate_potential_energy,
409 )
410 else:
411 assert False, "unreachable"
413 return kernel
415 def get_update_velocities_kernel(
416 self,
417 gridsync: bool,
418 num_part: int,
419 num_dim: int,
420 dot_v_f_id: int,
421 forces_sql_id: int,
422 vel_sql_id: int,
423 alpha_den_id: int,
424 ):
425 if self.mode == "reflection-mass_scaling":
426 @cuda.jit(device=gridsync)
427 def update_velocities_kernel(
428 grid, my_f, d_scalars_shared, my_v, my_m,
429 ):
430 global_id, my_t = cuda.grid(2) # type: ignore
432 if global_id < num_part and my_t == 0:
433 den = 0
434 for k in range(num_dim):
435 den += my_f[k] ** 2
436 den = den / my_m
437 cuda.atomic.add(d_scalars_shared, alpha_den_id, den) # type: ignore
438 grid.sync()
440 alpha = - 2 * d_scalars_shared[dot_v_f_id] / d_scalars_shared[alpha_den_id]
442 if global_id < num_part and my_t == 0:
443 for k in range(num_dim):
444 my_v[k] = my_v[k] + alpha / my_m * my_f[k]
446 if global_id == 0 and my_t == 0:
447 d_scalars_shared[dot_v_f_id] = - d_scalars_shared[dot_v_f_id]
449 elif self.mode == "reflection":
450 @cuda.jit(device=gridsync)
451 def update_velocities_kernel(
452 grid, my_f, d_scalars_shared, my_v, my_m,
453 ):
454 global_id, my_t = cuda.grid(2) # type: ignore
456 ## v_{i+1} = v_i -2 * (F_i * v_i)/|F_i|**2 * F_i
457 ################################################
458 if global_id < num_part and my_t == 0:
459 for k in range(num_dim):
460 my_v[k] = my_v[k] - 2 * d_scalars_shared[dot_v_f_id] / d_scalars_shared[forces_sql_id] * my_f[k]
461 if global_id == 0 and my_t == 0:
462 d_scalars_shared[dot_v_f_id] = - d_scalars_shared[dot_v_f_id]
463 elif self.mode == "no-inertia":
464 @cuda.jit(device=gridsync)
465 def update_velocities_kernel(
466 grid, my_f, d_scalars_shared, my_v, my_m,
467 ):
468 global_id, my_t = cuda.grid(2) # type: ignore
470 vel_l = math.sqrt(d_scalars_shared[vel_sql_id])
471 force_l = math.sqrt(d_scalars_shared[forces_sql_id])
472 if global_id < num_part and my_t == 0:
473 for k in range(num_dim):
474 my_v[k] = vel_l * my_f[k] / force_l
475 if global_id == 0 and my_t == 0:
476 d_scalars_shared[dot_v_f_id] = vel_l * force_l
477 else:
478 assert False, "Unreachable"
479 return update_velocities_kernel
481 def get_setup_kernel(self, configuration, compute_plan, interactions_kernel):
482 num_dim, num_part = configuration.D, configuration.N
483 pb, tp, gridsync = [compute_plan[key] for key in ['pb', 'tp', 'gridsync']]
484 num_blocks = (num_part - 1) // pb + 1
486 # Unpack indices for vectors and scalars to be compiled into kernel
487 r_id, v_id, f_id = [configuration.vectors.indices[key] for key in ['r', 'v', 'f']]
488 u_id, m_id = configuration.sid['U'], configuration.sid['m']
490 apply_PBC = nb.jit(configuration.simbox.apply_PBC)
491 (
492 forces_sql_id,
493 vel_sql_id,
494 dot_v_f_id,
495 should_break_id,
496 should_return_id,
497 reached_max_abs_val_id,
498 should_move_r_copy_id,
499 should_flip_velocities_id,
500 second_derivative_id,
501 first_derivative_id,
502 alpha_den_id,
503 *debug_ids
504 ) = range(self.d_scalars_shared.shape[0])
505 debug_ids = tuple(debug_ids)
507 max_initial_step_corrections = self.max_initial_step_corrections
508 initial_step_if_high = self.initial_step_if_high
509 threshold = self.threshold
510 eps = self.eps
511 target_u = self.target_u
512 debug_print = self.debug_print
514 if gridsync:
515 @cuda.jit(device=gridsync)
516 def kernel_setup(
517 grid, vectors, scalars, r_im, sim_box,
518 integrator_params, ptype,
519 ):
520 (dt, d_integrator_output, d_initial_step, interaction_params, d_scalars_shared, d_pot_energy, d_path_u,
521 d_step, d_broken_simulation, d_last_a, d_u_higher_than_target_in_time_0) = integrator_params
522 global_id, my_t = cuda.grid(2) # type: ignore
524 if global_id == 0 and my_t == 0:
525 for k in range(d_scalars_shared.shape[0]):
526 d_scalars_shared[k] = 0
527 grid.sync()
529 my_m = scalars[global_id, m_id]
530 my_v = vectors[v_id, global_id]
531 forces = vectors[f_id]
532 my_f = forces[global_id]
533 positions = vectors[r_id]
534 my_r = positions[global_id]
535 my_r_im = r_im[global_id]
536 its: CudaArray = cuda.local.array(1, dtype=np.int32) # type: ignore
537 its[0] = 0
539 calculate_potential_energy(interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its)
540 get_dot_in_conf_space(my_f, my_f, d_scalars_shared, forces_sql_id)
541 get_dot_in_conf_space(my_v, my_v, d_scalars_shared, vel_sql_id)
542 get_dot_in_conf_space(my_f, my_v, d_scalars_shared, dot_v_f_id)
543 grid.sync()
544 vel_l = math.sqrt(d_scalars_shared[vel_sql_id])
546 if global_id == 0 and my_t == 0:
547 u = d_pot_energy[0]
548 du_rel = (u - target_u) / abs(target_u)
549 if abs(du_rel) < threshold:
550 pass
551 elif du_rel > threshold:
552 d_u_higher_than_target_in_time_0[0] = True
553 if debug_print:
554 print("Initial configuration has U > U0:", u, target_u, du_rel, math.log10(abs(du_rel)))
555 print("In the past, this is not a `good` point to start the simulation. If it does not work "
556 "try running with a configuration that has U <= U0")
557 else:
558 if debug_print:
559 print("Initial configuration has U < U0:", u, target_u, du_rel, math.log10(abs(du_rel)))
561 if d_u_higher_than_target_in_time_0[0]:
562 if global_id == 0 and my_t == 0:
563 d_initial_step[0] = initial_step_if_high
564 if d_scalars_shared[dot_v_f_id] > 0:
565 d_scalars_shared[should_flip_velocities_id] = True
566 grid.sync()
567 if global_id < num_part and my_t == 0:
568 if d_scalars_shared[should_flip_velocities_id]:
569 for k in range(num_dim):
570 my_v[k] = my_v[k] * np.float32(-1)
571 if global_id == 0 and my_t == 0:
572 d_scalars_shared[dot_v_f_id] = - d_scalars_shared[dot_v_f_id]
573 grid.sync()
575 my_dir: CudaArray = cuda.local.array(num_dim, my_v.dtype) # type: ignore
576 if global_id < num_part and my_t == 0:
577 for k in range(num_dim):
578 my_dir[k] = my_v[k]
580 calculate_next_velocities(grid, my_f, d_scalars_shared, my_dir, my_m)
581 grid.sync()
582 if global_id < num_part and my_t == 0:
583 for k in range(num_dim):
584 my_dir[k] = my_dir[k] / vel_l
586 # Used to reset positions later
587 r0: CudaArray = cuda.local.array(num_dim, my_r.dtype) # type: ignore
588 r0_im: CudaArray = cuda.local.array(num_dim, my_r_im.dtype) # type: ignore
589 copy_positions_and_images(my_r, my_r_im, r0, r0_im)
590 setup_raytracing(
591 my_f, d_broken_simulation, d_last_a,
592 d_scalars_shared, debug_print, 0,
593 r0, r0_im, my_dir, d_initial_step, my_r, my_r_im,
594 interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its,
595 d_u_higher_than_target_in_time_0,
596 )
597 copy_positions_and_images(r0, r0_im, my_r, my_r_im)
598 else:
599 raise NotImplementedError()
601 (
602 get_dot_in_conf_space,
603 copy_positions_and_images,
604 calculate_potential_energy,
605 add_step_in_dir,
606 ) = util_functions(gridsync, num_part, num_dim, apply_PBC, u_id)
607 calculate_next_velocities = self.get_update_velocities_kernel(
608 gridsync=gridsync,
609 num_part=num_part,
610 num_dim=num_dim,
611 dot_v_f_id=dot_v_f_id,
612 forces_sql_id=forces_sql_id,
613 vel_sql_id=vel_sql_id,
614 alpha_den_id=alpha_den_id,
615 )
617 if self.raytracing_method == "bisection":
618 @cuda.jit(device=gridsync)
619 def setup_raytracing(
620 my_f, d_broken_simulation, d_last_a,
621 d_scalars_shared, debug_print, time,
622 r0, r0_im, my_dir, d_initial_step, my_r, my_r_im,
623 interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its,
624 d_u_higher_than_target_in_time_0,
625 ):
626 return
628 elif self.raytracing_method == "parabola-newton" or self.raytracing_method == "parabola":
629 @cuda.jit(device=gridsync)
630 def setup_raytracing(
631 my_f, d_broken_simulation, d_last_a,
632 d_scalars_shared, debug_print, time,
633 r0, r0_im, my_dir, d_initial_step, my_r, my_r_im,
634 interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its,
635 d_u_higher_than_target_in_time_0,
636 ):
637 global_id, my_t = cuda.grid(2) # type: ignore
638 force0: CudaArray = cuda.local.array(num_dim, my_f.dtype) # type: ignore
639 for k in range(num_dim):
640 force0[k] = my_f[k]
642 if global_id == 0 and my_t == 0:
643 d_scalars_shared[should_break_id] = False
644 grid.sync()
646 x_initial_step = d_initial_step[0]
647 corrected_step = np.int32(0)
648 while corrected_step <= max_initial_step_corrections and not d_scalars_shared[should_break_id]:
649 corrected_step += 1
651 # Initial step go into a bit further from the surface into lower energies
652 add_step_in_dir(r0, r0_im, my_dir, x_initial_step, my_r, my_r_im, sim_box)
653 grid.sync()
654 calculate_potential_energy(interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its)
656 if global_id == 0 and my_t == 0:
657 u = d_pot_energy[0]
658 du_rel = (u - target_u) / abs(target_u)
659 if du_rel < -eps:
660 d_scalars_shared[should_break_id] = True
661 grid.sync()
662 if not d_scalars_shared[should_break_id]:
663 if d_u_higher_than_target_in_time_0[0]:
664 x_initial_step = x_initial_step * 2
665 else:
666 x_initial_step = x_initial_step / 2
668 if corrected_step > max_initial_step_corrections and not d_scalars_shared[should_break_id]:
669 if global_id == 0 and my_t == 0:
670 d_broken_simulation[0] = True
671 uf = d_pot_energy[0]
672 copy_positions_and_images(r0, r0_im, my_r, my_r_im)
673 grid.sync()
674 calculate_potential_energy(
675 interactions_kernel, grid, vectors, scalars, ptype,
676 sim_box, interaction_params, d_pot_energy, its)
677 d = d_scalars_shared[dot_v_f_id]
678 u1 = d_pot_energy[0]
679 if global_id == 0 and my_t == 0:
680 d_scalars_shared[dot_v_f_id] = 0
681 grid.sync()
682 get_dot_in_conf_space(my_f, my_dir, d_scalars_shared, dot_v_f_id)
683 grid.sync()
684 b1 = - d_scalars_shared[dot_v_f_id]
686 add_step_in_dir(r0, r0_im, my_dir, d_initial_step[0], my_r, my_r_im, sim_box)
687 grid.sync()
688 calculate_potential_energy(
689 interactions_kernel, grid, vectors, scalars, ptype,
690 sim_box, interaction_params, d_pot_energy, its)
691 copy_positions_and_images(r0, r0_im, my_r, my_r_im)
692 u2 = d_pot_energy[0]
693 if global_id == 0 and my_t == 0:
694 d_scalars_shared[dot_v_f_id] = 0
695 grid.sync()
696 get_dot_in_conf_space(my_f, my_dir, d_scalars_shared, dot_v_f_id)
697 grid.sync()
698 b2 = - d_scalars_shared[dot_v_f_id]
699 a = 0.5 * (b2 - b1) / d_initial_step[0]
700 if global_id == 0 and my_t == 0:
701 if debug_print:
702 print("ERROR:", time, "Reached max_initial_step_corrections")
703 print(" uf =", uf)
704 print(" u1 =", u1)
705 print(" b1 =", b1)
706 print(" u2 =", u2)
707 print(" b2 =", b2)
708 print(" a =", a)
709 print(" initial_step0 = ", d_initial_step[0])
710 print(" initial_stepf = ", x_initial_step)
711 print(" dot_v_f0 = ", d)
712 # return
713 upp = get_second_derivative(
714 grid, my_f, force0, my_dir, x_initial_step, d_scalars_shared
715 )
716 if global_id == 0 and my_t == 0:
717 a = upp / 2
718 d_last_a[0] = a
719 else:
720 assert False, "Unreachable"
722 @cuda.jit(device=gridsync)
723 def get_second_derivative(
724 grid, my_f, force0, my_dir, delta, d_scalars_shared
725 ):
726 global_id, my_t = cuda.grid(2) # type: ignore
727 if global_id < num_part and my_t == 0:
728 acc = np.float32(0)
729 for k in range(num_dim):
730 acc += - (my_f[k] - force0[k]) * my_dir[k] / delta
731 cuda.atomic.add(d_scalars_shared, second_derivative_id, acc) # type: ignore
732 grid.sync()
733 return d_scalars_shared[second_derivative_id]
735 if gridsync:
736 @cuda.jit
737 def integrator_setup(
738 vectors, scalars, r_im, sim_box,
739 integrator_params, ptype,
740 ):
741 grid = cuda.cg.this_grid()
742 kernel_setup(
743 grid, vectors, scalars, r_im, sim_box,
744 integrator_params, ptype,
745 )
746 return integrator_setup[num_blocks, (pb, tp)]
747 else:
748 raise NotImplementedError()
750 def get_kernel_parabola(
751 self,
752 gridsync: bool,
753 num_part: int,
754 num_dim: int,
755 r_id: int,
756 u_id: int,
757 dot_v_f_id: int,
758 should_break_id: int,
760 copy_positions_and_images: Any,
761 add_step_in_dir: Any,
762 calculate_potential_energy: Any,
763 ):
764 max_steps = self.max_steps
765 target_u = self.target_u
766 threshold = self.threshold
768 if gridsync:
769 @cuda.jit(device=gridsync)
770 def parabola_kernel(
771 time,
772 interactions_kernel, grid, vectors, scalars, ptype,
773 sim_box, interaction_params,
774 r_im,
775 d_pot_energy,
776 its,
777 d_last_a,
778 d_scalars_shared,
779 vel_l,
780 my_dir,
781 d_u_higher_than_target_in_time_0,
782 d_initial_step,
783 ):
784 global_id, my_t = cuda.grid(2) # type: ignore
786 positions = vectors[r_id]
787 my_r = positions[global_id]
788 my_r_im = r_im[global_id]
789 r0: CudaArray = cuda.local.array(num_dim, my_r.dtype) # type: ignore
790 r0_im: CudaArray = cuda.local.array(num_dim, my_r_im.dtype) # type: ignore
791 copy_positions_and_images(my_r, my_r_im, r0, r0_im)
793 a = d_last_a[0]
794 b = - d_scalars_shared[dot_v_f_id] / vel_l
796 if global_id == 0 and my_t == 0:
797 d_pot_energy[0] = 0
798 grid.sync()
799 if global_id < num_part and my_t == 0:
800 cuda.atomic.add(d_pot_energy, 0, scalars[global_id][u_id]) # type: ignore
801 grid.sync()
802 c = d_pot_energy[0] - target_u
804 if global_id == 0 and my_t == 0:
805 d_scalars_shared[should_break_id] = False
806 grid.sync()
808 x1 = - b / (2*a) * (math.sqrt(1 - 4 * a * c / b**2) + 1)
809 add_step_in_dir(r0, r0_im, my_dir, x1, my_r, my_r_im, sim_box)
810 grid.sync()
811 calculate_potential_energy(
812 interactions_kernel, grid, vectors, scalars, ptype,
813 sim_box, interaction_params, d_pot_energy, its)
814 u = d_pot_energy[0]
815 if global_id == 0 and my_t == 0:
816 du_rel = (u - target_u) / abs(target_u)
817 if abs(du_rel) < threshold:
818 d_scalars_shared[should_break_id] = True
819 grid.sync()
820 if d_scalars_shared[should_break_id]:
821 return x1
823 u1 = u - target_u
824 x2 = b * x1**2 / (b*x1 - u1)
825 u2 = 0 # Only to solve unbound issues
827 steps_done = 1
828 while steps_done <= max_steps and not d_scalars_shared[should_break_id]:
829 steps_done += 1
830 add_step_in_dir(r0, r0_im, my_dir, x2, my_r, my_r_im, sim_box)
831 grid.sync()
832 calculate_potential_energy(
833 interactions_kernel, grid, vectors, scalars, ptype,
834 sim_box, interaction_params, d_pot_energy, its)
835 u = d_pot_energy[0]
836 if global_id == 0 and my_t == 0:
837 du_rel = (u - target_u) / abs(target_u)
838 if abs(du_rel) < threshold:
839 d_scalars_shared[should_break_id] = True
840 grid.sync()
841 u2 = u - target_u
842 if not d_scalars_shared[should_break_id]:
843 s = (u2 * x1**2 - u1 * x2**2) / (u2 * x1 - u1 * x2)
844 u1 = u2
845 x1 = x2
846 x2 = s
847 if global_id == 0 and my_t == 0:
848 d_last_a[0] = - b / x2
849 return x2
850 else:
851 raise NotImplementedError()
852 return parabola_kernel
854 def get_kernel_parabola_newton(
855 self,
856 gridsync: bool,
857 num_part: int,
858 num_dim: int,
859 r_id: int,
860 f_id: int,
861 u_id: int,
862 dot_v_f_id: int,
863 should_break_id: int,
864 first_derivative_id: int,
866 copy_positions_and_images: Any,
867 add_step_in_dir: Any,
868 calculate_potential_energy: Any,
869 ):
870 max_steps = self.max_steps
871 target_u = self.target_u
872 threshold = self.threshold
874 if gridsync:
875 @cuda.jit(device=gridsync)
876 def parabola_newton_kernel(
877 time,
878 interactions_kernel, grid, vectors, scalars, ptype,
879 sim_box, interaction_params,
880 r_im,
881 d_pot_energy,
882 its,
883 d_last_a,
884 d_scalars_shared,
885 vel_l,
886 my_dir,
887 d_u_higher_than_target_in_time_0,
888 d_initial_step,
889 ):
890 global_id, my_t = cuda.grid(2) # type: ignore
892 forces = vectors[f_id]
893 my_f = forces[global_id]
894 positions = vectors[r_id]
895 my_r = positions[global_id]
896 my_r_im = r_im[global_id]
897 r0: CudaArray = cuda.local.array(num_dim, my_r.dtype) # type: ignore
898 r0_im: CudaArray = cuda.local.array(num_dim, my_r_im.dtype) # type: ignore
899 copy_positions_and_images(my_r, my_r_im, r0, r0_im)
901 a = d_last_a[0]
903 b = - d_scalars_shared[dot_v_f_id] / vel_l
905 if global_id == 0 and my_t == 0:
906 d_pot_energy[0] = 0
907 grid.sync()
908 if global_id < num_part and my_t == 0:
909 cuda.atomic.add(d_pot_energy, 0, scalars[global_id][u_id]) # type: ignore
910 grid.sync()
911 c = d_pot_energy[0] - target_u
913 if global_id == 0 and my_t == 0:
914 d_scalars_shared[should_break_id] = False
915 grid.sync()
917 step = - b / (2*a) * (math.sqrt(1 - 4 * a * c / b**2) + 1)
919 if global_id == 0 and my_t == 0:
920 d_scalars_shared[should_break_id] = False
921 grid.sync()
923 steps_done = np.int32(0)
924 while steps_done <= max_steps and not d_scalars_shared[should_break_id]:
925 steps_done += 1
926 add_step_in_dir(r0, r0_im, my_dir, step, my_r, my_r_im, sim_box)
927 grid.sync()
928 calculate_potential_energy(
929 interactions_kernel, grid, vectors, scalars, ptype,
930 sim_box, interaction_params, d_pot_energy, its)
932 if global_id == 0 and my_t == 0:
933 u = d_pot_energy[0]
934 du_rel = (u - target_u) / abs(target_u)
935 if abs(du_rel) < threshold:
936 d_scalars_shared[should_break_id] = True
937 grid.sync()
938 if not d_scalars_shared[should_break_id]:
939 u_prime = get_first_derivative(
940 grid, my_f, my_dir, d_scalars_shared
941 )
942 # x_i+1 = x_i - f(x_n) / f'(x_n)
943 step = step - (d_pot_energy[0] - target_u) / u_prime
944 return step
945 else:
946 raise NotImplementedError()
948 @cuda.jit(device=gridsync)
949 def get_first_derivative(
950 grid, my_f, my_dir, d_scalars_shared
951 ):
952 global_id, my_t = cuda.grid(2) # type: ignore
953 if global_id == 0 and my_t == 0:
954 d_scalars_shared[first_derivative_id] = 0
955 grid.sync()
956 if global_id < num_part and my_t == 0:
957 acc = np.float32(0)
958 for k in range(num_dim):
959 acc += - my_f[k] * my_dir[k]
960 cuda.atomic.add(d_scalars_shared, first_derivative_id, acc) # type: ignore
961 grid.sync()
962 return d_scalars_shared[first_derivative_id]
964 return parabola_newton_kernel
966 def get_kernel_bisection(
967 self,
968 gridsync: bool,
969 num_dim: int,
970 r_id: int,
971 should_break_id: int,
972 should_move_r_copy_id: int,
973 reached_max_abs_val_id: int,
974 should_return_id: int,
976 copy_positions_and_images: Any,
977 add_step_in_dir: Any,
978 calculate_potential_energy: Any,
979 ):
980 max_steps = self.max_steps
981 target_u = self.target_u
982 threshold = self.threshold
983 eps = self.eps
984 initial_step0 = self.initial_step
985 debug_print = self.debug_print
986 step = self.step
987 max_initial_step_corrections = self.max_initial_step_corrections
988 max_abs_val = self.max_abs_val
990 if gridsync: # construct and return device function
991 @cuda.jit(device=gridsync)
992 def bisection_kernel(
993 time,
994 interactions_kernel, grid, vectors, scalars, ptype,
995 sim_box, interaction_params,
996 r_im,
997 d_pot_energy,
998 its,
999 d_last_a,
1000 d_scalars_shared,
1001 vel_l,
1002 my_dir,
1003 d_u_higher_than_target_in_time_0,
1004 d_initial_step,
1005 ):
1006 global_id, my_t = cuda.grid(2) # type: ignore
1007 positions = vectors[r_id]
1008 my_r = positions[global_id]
1009 my_r_im = r_im[global_id]
1010 # INITIAL STEP
1011 ##############
1012 r_copy: CudaArray = cuda.local.array(num_dim, my_r.dtype) # type: ignore
1013 r_im_copy: CudaArray = cuda.local.array(num_dim, r_im.dtype) # type: ignore
1014 r_copy_u: CudaArray = cuda.local.array(1, d_pot_energy.dtype) # type: ignore
1015 r_copy_u[0] = d_pot_energy[0]
1016 copy_positions_and_images(my_r, my_r_im, r_copy, r_im_copy)
1018 delta_x = 0
1019 if global_id == 0 and my_t == 0:
1020 d_scalars_shared[should_break_id] = False
1021 grid.sync()
1023 x_initial_step = d_initial_step[0]
1024 corrected_step = np.int32(0)
1025 while corrected_step <= max_initial_step_corrections and not d_scalars_shared[should_break_id]:
1026 corrected_step += 1
1028 # Initial step go into a bit further from the surface into lower energies
1029 add_step_in_dir(r_copy, r_im_copy, my_dir, x_initial_step, my_r, my_r_im, sim_box)
1030 grid.sync()
1031 calculate_potential_energy(interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its)
1033 if global_id == 0 and my_t == 0:
1034 u = d_pot_energy[0]
1035 du_rel = (u - target_u) / abs(target_u)
1036 if du_rel < -eps:
1037 d_scalars_shared[should_break_id] = True
1038 grid.sync()
1039 if not d_scalars_shared[should_break_id]:
1040 if d_u_higher_than_target_in_time_0[0]:
1041 x_initial_step = x_initial_step * 2
1042 else:
1043 x_initial_step = x_initial_step / 2
1045 if global_id == 0 and my_t == 0:
1046 if d_u_higher_than_target_in_time_0[0]:
1047 d_initial_step[0] = initial_step0
1048 else:
1049 d_initial_step[0] = x_initial_step
1051 delta_x += x_initial_step
1052 if corrected_step > max_initial_step_corrections and not d_scalars_shared[should_break_id]:
1053 # TODO: deal with this: reached max steps for the initial step correction
1054 # Most likely this is an equilibrium position.
1055 # We should think about setting the reflected ray to a random direction
1056 if global_id == 0 and my_t == 0:
1057 if debug_print:
1058 print("ERROR:", time, "Reached max_initial_step_corrections =>",
1059 (d_pot_energy[0] - target_u) / abs(target_u),
1060 x_initial_step)
1061 return delta_x
1063 r_copy_u[0] = d_pot_energy[0]
1064 copy_positions_and_images(my_r, my_r_im, r_copy, r_im_copy)
1066 # FIND POINT OUTSIDE SURFACE
1067 ############################
1069 should_return_above, x_above, steps_done = perform_find_point_above(
1070 d_scalars_shared, time,
1071 r_copy, r_im_copy, my_dir, step, my_r, my_r_im, r_copy_u,
1072 interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its
1073 )
1074 delta_x += x_above
1075 if should_return_above:
1076 return delta_x
1078 ## PERFORM BISECTION
1079 ####################
1081 ### Root between r_copy(r_copy_u) and my_r(d_pot_energy[0])
1083 if global_id == 0 and my_t == 0 and debug_print:
1084 before_bisec = (r_copy_u[0] - target_u) / abs(target_u), (d_pot_energy[0] - target_u) / abs(target_u)
1085 else:
1086 before_bisec = np.nan, np.nan
1088 should_return_bisection, x_bisection = perform_bisection(
1089 steps_done, time, before_bisec,
1090 d_scalars_shared,
1091 r_copy, r_im_copy, my_dir, step, my_r, my_r_im, r_copy_u,
1092 interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its
1093 )
1094 delta_x += x_bisection
1095 if should_return_bisection:
1096 return delta_x
1097 return delta_x
1098 else: # return python function, which makes kernel-calls
1099 raise ValueError("Currently no gridsync is not supported for NVU_RT")
1101 @cuda.jit(device=gridsync)
1102 def perform_find_point_above(
1103 d_scalars_shared, time,
1104 r_copy, r_im_copy, my_dir, step, my_r, my_r_im, r_copy_u,
1105 interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its
1106 ):
1107 global_id, my_t = cuda.grid(2) # type: ignore
1108 if global_id == 0 and my_t == 0:
1109 d_scalars_shared[should_break_id] = False
1110 # print("Potential Energy before searching:", d_scalars_shared[pot_energy_id])
1111 grid.sync()
1113 delta_x = np.float32(0)
1114 steps_done = np.int32(0)
1115 while steps_done <= max_steps and not d_scalars_shared[should_break_id]:
1116 steps_done += 1
1118 add_step_in_dir(r_copy, r_im_copy, my_dir, step, my_r, my_r_im, sim_box)
1119 grid.sync()
1120 calculate_potential_energy(interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its)
1122 if global_id == 0 and my_t == 0:
1123 du = d_pot_energy[0] - target_u
1124 du_rel = du / abs(target_u)
1125 # Max Abs Val is there to prevent any overflows that can happen if U suddenly rises
1126 d_scalars_shared[should_move_r_copy_id] = False
1127 if abs(du_rel) > max_abs_val:
1128 if debug_print:
1129 print("ERROR:", time, "Reached max_abs_val in", steps_done, "step",
1130 du_rel, "Consider using a lower `step` parameter")
1131 d_scalars_shared[should_break_id] = True
1132 d_scalars_shared[reached_max_abs_val_id] = True
1133 elif abs(du_rel) < threshold:
1134 d_scalars_shared[should_break_id] = True
1135 d_scalars_shared[should_return_id] = True
1136 elif du_rel > eps:
1137 d_scalars_shared[should_break_id] = True
1138 elif du_rel < -eps:
1139 d_scalars_shared[should_move_r_copy_id] = True
1140 else:
1141 # This means that abs(du_rel) < eps, which is a case already covered with theshold
1142 # because we requiere eps < threshold
1143 pass
1144 grid.sync()
1145 if d_scalars_shared[should_move_r_copy_id]:
1146 delta_x += step
1147 copy_positions_and_images(my_r, my_r_im, r_copy, r_im_copy)
1148 r_copy_u[0] = d_pot_energy[0]
1149 grid.sync()
1151 if d_scalars_shared[should_return_id]:
1152 # We found the root sooner
1153 # if global_id == 0 and my_t == 0:
1154 # print("SUCCESS before bisec:", time, " => ", du_rel)
1155 delta_x += step
1156 return True, delta_x, steps_done
1157 if d_scalars_shared[reached_max_abs_val_id]:
1158 copy_positions_and_images(r_copy, r_im_copy, my_r, my_r_im)
1159 if global_id == 0 and my_t == 0:
1160 if debug_print:
1161 du_rel = (d_pot_energy[0] - target_u)/ abs(target_u)
1162 print("ERROR:", time, "Reached max_abs_val before bisection.", du_rel)
1163 return True, delta_x, steps_done
1164 if not d_scalars_shared[should_break_id] and steps_done > max_steps:
1165 # TODO: deal with this: reached max steps
1166 # Use r_copy because my_r is propbably above U0 bevause if not in the wjiole loop r_copy
1167 # should have been moved and so r_copy == my_r
1168 copy_positions_and_images(r_copy, r_im_copy, my_r, my_r_im)
1169 if global_id == 0 and my_t == 0:
1170 if debug_print:
1171 du_rel = (d_pot_energy[0] - target_u)/ abs(target_u)
1172 print("ERROR:", time, "Reached max_steps before bisection.", du_rel, math.log10(abs(du_rel)))
1173 return True, delta_x, steps_done
1174 return False, delta_x, steps_done
1176 @cuda.jit(device=gridsync)
1177 def perform_bisection(
1178 steps_done, time, before_bisec,
1179 d_scalars_shared,
1180 r_copy, r_im_copy, my_dir, step, my_r, my_r_im, r_copy_u,
1181 interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its
1182 ):
1183 global_id, my_t = cuda.grid(2) # type: ignore
1185 if global_id == 0 and my_t == 0:
1186 d_scalars_shared[should_break_id] = False
1187 grid.sync()
1189 delta_x = np.float32(0)
1190 while steps_done <= max_steps and not d_scalars_shared[should_break_id]:
1191 steps_done += 1
1192 step = step / 2
1194 add_step_in_dir(r_copy, r_im_copy, my_dir, step, my_r, my_r_im, sim_box)
1195 grid.sync()
1196 calculate_potential_energy(interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its)
1198 if global_id == 0 and my_t == 0:
1199 d_scalars_shared[should_move_r_copy_id] = False
1200 du = d_pot_energy[0] - target_u
1201 du_rel = du / abs(target_u)
1202 if abs(du_rel) > max_abs_val:
1203 pass
1204 elif abs(du_rel) < threshold:
1205 # elif (abs(u / target_u - 1) < threshold):
1206 d_scalars_shared[should_break_id] = True
1207 elif du_rel < -eps:
1208 d_scalars_shared[should_move_r_copy_id] = True
1209 elif du_rel > eps:
1210 pass
1211 else:
1212 pass
1213 grid.sync()
1215 if d_scalars_shared[should_move_r_copy_id]:
1216 delta_x += step
1217 copy_positions_and_images(my_r, my_r_im, r_copy, r_im_copy)
1218 r_copy_u[0] = d_pot_energy[0]
1219 grid.sync()
1221 if d_scalars_shared[should_break_id] and steps_done > max_steps:
1222 # TODO: deal with this: reached max steps.
1223 # For now use the lower energy positions
1224 copy_positions_and_images(r_copy, r_im_copy, my_r, my_r_im)
1225 grid.sync()
1226 if debug_print:
1227 if global_id == 0 and my_t == 0:
1228 du_rel = (d_pot_energy[0] - target_u) / abs(target_u)
1229 else:
1230 du_rel = np.nan
1231 calculate_potential_energy(interactions_kernel, grid, vectors, scalars, ptype, sim_box, interaction_params, d_pot_energy, its)
1232 if global_id == 0 and my_t == 0:
1233 du_rel2 = (d_pot_energy[0] - target_u) / abs(target_u)
1234 print("ERROR:", time, "Reached max steps within bisection.", steps_done)
1235 print("Before bisec:", before_bisec[0], before_bisec[1], math.log10(abs(before_bisec[0])), math.log10(abs(before_bisec[1])))
1236 print("After bisec:", du_rel2, du_rel, math.log10(abs(du_rel2)), math.log10(abs(du_rel)))
1237 return True, delta_x
1239 delta_x += step
1240 return False, delta_x
1241 return bisection_kernel
1245def util_functions(
1246 gridsync: bool, num_part: int, num_dim: int, apply_PBC: Any,
1247 u_id: int,
1248):
1249 @cuda.jit(device=gridsync)
1250 def get_dot_in_conf_space(my_a, my_b, result_arr, result_id):
1251 global_id, my_t = cuda.grid(2) # type: ignore
1252 if global_id < num_part and my_t == 0:
1253 x = np.float32(0)
1254 for k in range(num_dim):
1255 x += my_a[k] * my_b[k]
1256 cuda.atomic.add(result_arr, result_id, x) # type: ignore
1257 @cuda.jit(device=gridsync)
1258 def copy_positions_and_images(source_r, source_r_im, dest_r, dest_r_im):
1259 global_id, my_t = cuda.grid(2) # type: ignore
1260 if global_id < num_part and my_t == 0:
1261 for k in range(num_dim):
1262 dest_r[k] = source_r[k]
1263 dest_r_im[k] = source_r_im[k]
1265 @cuda.jit(device=gridsync)
1266 def calculate_potential_energy(
1267 interactions_kernel, grid, vectors, scalars, ptype,
1268 sim_box, interaction_params, d_pot_energy, its):
1269 its[0] += 1
1270 global_id, my_t = cuda.grid(2) # type: ignore
1271 interactions_kernel(grid, vectors, scalars, ptype, sim_box, interaction_params)
1272 grid.sync()
1273 if global_id == 0 and my_t == 0:
1274 d_pot_energy[0] = 0
1275 grid.sync()
1276 if global_id < num_part and my_t == 0:
1277 cuda.atomic.add(d_pot_energy, 0, scalars[global_id][u_id]) # type: ignore
1278 grid.sync()
1280 @cuda.jit(device=gridsync)
1281 def add_step_in_dir(source, source_im, dir, step, dest, dest_im, sim_box):
1282 global_id, my_t = cuda.grid(2) # type: ignore
1283 if global_id < num_part and my_t == 0:
1284 for k in range(num_dim):
1285 dest[k] = source[k] + dir[k] * step
1286 dest_im[k] = source_im[k]
1287 apply_PBC(dest, dest_im, sim_box) # type: ignore
1289 return (
1290 get_dot_in_conf_space,
1291 copy_positions_and_images,
1292 calculate_potential_energy,
1293 add_step_in_dir,
1294 )