Coverage for gamdpy/integrators/NVT.py: 52%

1import numpy as np 

2import numba 

3from numba import cuda 

4import gamdpy as gp 

5from .integrator import Integrator 

6 

7class NVT(Integrator): 

8 """ The leapfrog algorithm with a Nose-Hoover thermostat 

9 

10 Integrator keeping N (number of particles), V (volume), and T (temperature) constant, 

11 using the leapfrog algorithm and a Nose-Hoover thermostat. 

12 

13 Parameters 

14 ---------- 

15 temperature : float or function 

16 The temperature to keep the system at. If a function, it should take time as argument. 

17 

18 tau : float 

19 The relaxation time of the thermostat. 

20 

21 dt : float 

22 The time step of the integration. 

23 

24 """ 

25 

26 def __init__(self, temperature, tau: float, dt: float) -> None: 

27 self.temperature = temperature 

28 self.tau = tau 

29 self.dt = dt 

30 self.thermostat_state = np.zeros(2, dtype=np.float32) # Right time to allocate and copy to device? 

31 self.d_thermostat_state = cuda.to_device(self.thermostat_state) # - or in get_params 

32 

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

34 dt = np.float32(self.dt) 

35 omega2 = np.float32(4.0 * np.pi * np.pi / self.tau / self.tau) 

36 degrees = configuration.N * configuration.D - configuration.D 

37 return (dt, omega2, degrees, self.d_thermostat_state) # Needs to be compatible with unpacking in 

38 # step() and update_thermostat_state() below. 

39 

40 def get_kernel(self, configuration: gp.Configuration, compute_plan: dict, compute_flags: dict, interactions_kernel, verbose=False): 

41 

42 # Unpack parameters from configuration and compute_plan 

43 D, num_part = configuration.D, configuration.N 

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

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

46 

47 # Convert temperature to a function if isn't allready (better be a number then...) 

48 if callable(self.temperature): 

49 temperature_function = self.temperature 

50 else: 

51 temperature_function = gp.make_function_constant(value=float(self.temperature)) 

52 

53 if verbose: 

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

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

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

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

58 

59 # Unpack indices for vectors and scalars to be compiled into kernel 

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

61 m_id = configuration.sid['m'] 

62 compute_k = compute_flags['K'] 

63 compute_fsq = compute_flags['Fsq'] 

64 if compute_k: 

65 k_id = configuration.sid['K'] 

66 if compute_fsq: 

67 fsq_id = configuration.sid['Fsq'] 

68 

69 # JIT compile functions to be compiled into kernel 

70 temperature_function = numba.njit(temperature_function) 

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

72 

73 

74 

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

76 """ Make one NVT timestep using Leap-frog 

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

78 """ 

79 

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

81 dt, omega2, degrees, thermostat_state = integrator_params 

82 

83 factor = np.float32(0.5) * thermostat_state[0] * dt 

84 plus = np.float32(1.) / (np.float32(1.) + factor) # Possibly change to exp(...) 

85 minus = np.float32(1.) - factor # Possibly change to exp(...) 

86 

87 global_id, my_t = cuda.grid(2) 

88 if global_id < num_part and my_t == 0: 

89 my_r = vectors[r_id][global_id] 

90 my_v = vectors[v_id][global_id] 

91 my_f = vectors[f_id][global_id] 

92 my_m = scalars[global_id][m_id] 

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

94 if compute_fsq: 

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

96 

97 for k in range(D): 

98 if compute_fsq: 

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

100 my_v[k] = plus * (minus * my_v[k] + my_f[k] / my_m * dt) 

101 my_k += numba.float32(0.5) * my_m * my_v[k] * my_v[k] 

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

103 

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

105 

106 cuda.atomic.add(thermostat_state, 1, my_k) # Probably slow? Not really! 

107 if compute_k: 

108 scalars[global_id][k_id] = my_k 

109 if compute_fsq: 

110 scalars[global_id][fsq_id] = my_fsq 

111 return 

112 

113 def update_thermostat_state(integrator_params, time): 

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

115 dt, omega2, degrees, thermostat_state = integrator_params 

116 

117 global_id, my_t = cuda.grid(2) 

118 if global_id == 0 and my_t == 0: 

119 target_temperature = temperature_function(time) 

120 

121 ke_deviation = np.float32(2.0) * thermostat_state[1] / (degrees * target_temperature) - np.float32(1.0) 

122 thermostat_state[0] += dt * omega2 * ke_deviation 

123 thermostat_state[1] = np.float32(0.) 

124 return 

125 

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

127 update_thermostat_state = cuda.jit(device=gridsync)(update_thermostat_state) 

128 

129 if gridsync: # construct and return device function 

130 def kernel(grid, vectors, scalars, r_im, sim_box, integrator_params, time, ptype): 

131 step( grid, vectors, scalars, r_im, sim_box, integrator_params, time) 

132 grid.sync() 

133 update_thermostat_state(integrator_params, time) 

134 return 

135 return cuda.jit(device=gridsync)(kernel) 

136 else: # return python function, which makes kernel-calls 

137 def kernel(grid, vectors, scalars, r_im, sim_box, integrator_params, time, ptype): 

138 step[num_blocks, (pb, 1)](grid, vectors, scalars, r_im, sim_box, integrator_params, time) 

139 update_thermostat_state[1, (1, 1)](integrator_params, time) 

140 return 

141 return kernel