Coverage for gamdpy/integrators/NVT_Langevin.py: 54%

1 

2import numpy as np 

3import numba 

4from numba import cuda 

5import math 

6from numba.cuda.random import create_xoroshiro128p_states 

7from numba.cuda.random import xoroshiro128p_normal_float32 

8import gamdpy as gp 

9from .integrator import Integrator 

10 

11 

12class NVT_Langevin(Integrator): 

13 """ NVT Langevin Leap-frog integrator 

14 

15 The langevin thermostat is a stochastic thermostat that keeps the system at a constant temperature: 

16 

17 .. math:: 

18 m a = f - \\alpha v + \\beta 

19 

20 where :math:`a` is the acceleration, :math:`f` is the force, :math:`v` is the velocity, :math:`m` is the mass, 

21 :math:`\\alpha` is the friction coefficient, and :math:`\\beta` is a random number drawn from a normal distribution. 

22 The implementation uses the leap-frog algorithm described in reference https://arxiv.org/pdf/1303.7011.pdf 

23 

24 Parameters 

25 ---------- 

26 

27 temperature : float or function 

28 Temperature of the thermostat. If a function, it must take a single argument, time, and return a float. 

29 

30 alpha : float 

31 Friction coefficient of the thermostat. 

32 

33 dt : float 

34 Time step for the integration. 

35 

36 """ 

37 

38 def __init__(self, temperature, alpha: float, dt: float, seed: int) -> None: 

39 self.temperature = temperature 

40 self.alpha = alpha 

41 self.dt = dt 

42 self.seed = seed 

43 

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

45 dt = np.float32(self.dt) 

46 alpha = np.float32(self.alpha) 

47 rng_states = create_xoroshiro128p_states(configuration.N, seed=self.seed) 

48 old_beta = np.zeros((configuration.N, configuration.D), dtype=np.float32) 

49 d_old_beta = cuda.to_device(old_beta) 

50 return (dt, alpha, rng_states, d_old_beta) # Needs to be compatible with unpacking in 

51 # step() below 

52 

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

54 

55 # Unpack parameters from configuration and compute_plan 

56 D, num_part = configuration.D, configuration.N 

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

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

59 

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

61 if callable(self.temperature): 

62 temperature_function = self.temperature 

63 else: 

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

65 

66 if verbose: 

67 print(f'Generating NVT langevin integrator for {num_part} particles in {D} dimensions:') 

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

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

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

71 

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

73 compute_k = compute_flags['K'] 

74 compute_fsq = compute_flags['Fsq'] 

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

76 m_id = configuration.sid['m'] 

77 if compute_k: 

78 k_id = configuration.sid['K'] 

79 if compute_fsq: 

80 fsq_id = configuration.sid['Fsq'] 

81 

82 

83 # JIT compile functions to be compiled into kernel 

84 temperature_function = numba.njit(temperature_function) 

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

86 

87 

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

89 """ Make one NVT Langevin timestep using Leap-frog 

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

91 REF: https://arxiv.org/pdf/1303.7011.pdf 

92 """ 

93 

94 dt, alpha, rng_states, old_beta = integrator_params 

95 temperature = temperature_function(time) 

96 

97 global_id, my_t = cuda.grid(2) 

98 if global_id < num_part and my_t == 0: 

99 my_r = vectors[r_id][global_id] 

100 my_v = vectors[v_id][global_id] 

101 my_f = vectors[f_id][global_id] 

102 my_m = scalars[global_id][m_id] 

103 if compute_k: 

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

105 if compute_fsq: 

106 my_fsq = numba.float32(0.0) # force squared energy 

107 

108 

109 for k in range(D): 

110 # REF: https://arxiv.org/pdf/1303.7011.pdf sec. 2.C. 

111 random_number = xoroshiro128p_normal_float32(rng_states, global_id) 

112 beta = math.sqrt(numba.float32(2.0) * alpha * temperature * dt) * random_number 

113 # Eq. (16) in https://arxiv.org/pdf/1303.7011.pdf 

114 numerator = numba.float32(2.0)*my_m - alpha * dt 

115 denominator = numba.float32(2.0)*my_m + alpha * dt 

116 a = numerator / denominator 

117 b_over_m = numba.float32(2.0) / denominator 

118 if compute_k: 

119 my_k += numba.float32(0.5) * my_m * my_v[k] * my_v[k] # Half step kinetic energy 

120 if compute_fsq: 

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

122 my_v[k] = a * my_v[k] + b_over_m * my_f[k] * dt + b_over_m * np.float32(0.5)*(beta+old_beta[global_id,k]) 

123 old_beta[global_id,k] = beta # Store beta for next step 

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

125 

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

127 if compute_k: 

128 scalars[global_id][k_id] = my_k 

129 if compute_fsq: 

130 scalars[global_id][fsq_id] = my_fsq 

131 return 

132 

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

134 

135 if gridsync: 

136 return step # return device function 

137 else: 

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