Coverage for gamdpy/integrators/NVE.py: 48%

1import numpy as np 

2import numba 

3import gamdpy as gp 

4from numba import cuda 

5from .integrator import Integrator 

6 

7 

8class NVE(Integrator): 

9 """Total energy conserving integrator 

10 

11 Use the Leap-frog algorithm to integrate the equations of motion: 

12 

13 .. math:: 

14 v(t+dt/2) &= v(t-dt/2) + f(t) dt / m 

15 

16 r(t+dt) &= r(t) + v(t+dt/2) dt 

17 

18 

19 Parameters 

20 ---------- 

21 

22 dt : float 

23 Time step for the integration. 

24 

25 """ 

26 def __init__(self, dt): 

27 self.dt = dt 

28 

29 def get_params(self, configuration: gp.Configuration, interactions_params: tuple, verbose=False) -> tuple: 

30 dt = np.float32(self.dt) 

31 return (dt,) 

32 

33 def get_kernel(self, configuration: gp.Configuration, compute_plan: dict, compute_flags: dict[str,bool], interactions_kernel, verbose=False): 

34 

35 # Unpack parameters from configuration and compute_plan 

36 D, num_part = configuration.D, configuration.N 

37 pb, tp, gridsync = [compute_plan[key] for key in ['pb', 'tp', 'gridsync']] 

38 num_blocks = (num_part - 1) // pb + 1 

39 

40 if verbose: 

41 print(f'Generating NVE kernel for {num_part} particles in {D} dimensions:') 

42 print(f'\tpb: {pb}, tp:{tp}, num_blocks:{num_blocks}') 

43 print(f'\tNumber (virtual) particles: {num_blocks * pb}') 

44 print(f'\tNumber of threads {num_blocks * pb * tp}') 

45 

46 # Unpack indices for vectors and scalars 

47 compute_k = compute_flags['K'] 

48 compute_fsq = compute_flags['Fsq'] 

49 r_id, v_id, f_id = [configuration.vectors.indices[key] for key in ['r', 'v', 'f']] 

50 m_id = configuration.sid['m'] 

51 if compute_k: 

52 k_id = configuration.sid['K'] 

53 if compute_fsq: 

54 fsq_id = configuration.sid['Fsq'] 

55 

56 # JIT compile functions to be compiled into kernel 

57 apply_PBC = numba.njit(configuration.simbox.get_apply_PBC()) 

58 

59 def step(grid, vectors, scalars, r_im, sim_box, integrator_params, time, ptype): 

60 """ Make one NVE timestep using Leap-frog 

61 Kernel configuration: [num_blocks, (pb, tp)] 

62 """ 

63 

64 # Unpack parameters. MUST be compatible with get_params() above 

65 dt, = integrator_params 

66 

67 global_id, my_t = cuda.grid(2) 

68 if global_id < num_part and my_t == 0: 

69 my_r = vectors[r_id][global_id] 

70 my_v = vectors[v_id][global_id] 

71 my_f = vectors[f_id][global_id] 

72 my_m = scalars[global_id][m_id] 

73 if compute_k: 

74 my_k = numba.float32(0.0) # Kinetic energy 

75 if compute_fsq: 

76 my_fsq = numba.float32(0.0) # force squared 

77 

78 for k in range(D): 

79 if compute_fsq: 

80 my_fsq += my_f[k] * my_f[k] 

81 v_mean = numba.float32(0.0) 

82 

83 # square before mean to get KE, part 1 

84 #my_k += numba.float32(0.25) * my_m * my_v[k] * my_v[k] 

85 

86 v_mean += my_v[k] # v(t-dt/2) 

87 my_v[k] += my_f[k] / my_m * dt 

88 v_mean += my_v[k] # v(t+dt/2) 

89 v_mean /= numba.float32(2.0) # v(t) = (v(t-dt/2) + v(t+dt/2))/2 

90 # square before mean,part 2 

91 #my_k += numba.float32(0.25) * my_m * my_v[k] * my_v[k] 

92 

93 # Basic: square the mean velocity 

94 if compute_k: 

95 my_k += numba.float32(0.5) * my_m * v_mean * v_mean 

96 

97 my_r[k] += my_v[k] * dt 

98 

99 apply_PBC(my_r, r_im[global_id], sim_box) 

100 

101 if compute_k: 

102 scalars[global_id][k_id] = my_k 

103 if compute_fsq: 

104 scalars[global_id][fsq_id] = my_fsq 

105 return 

106 

107 step = cuda.jit(device=gridsync)(step) 

108 

109 if gridsync: 

110 return step # return device function 

111 else: 

112 return step[num_blocks, (pb, 1)] # return kernel, incl. launch parameters 

113