Coverage for gamdpy/integrators/NPT_Langevin.py: 44%
154 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1import numpy as np
2import numba
3from numba import cuda
4import math
5from numba.cuda.random import create_xoroshiro128p_states
6from numba.cuda.random import xoroshiro128p_normal_float32
7import gamdpy as gp
8from .integrator import Integrator
10class NPT_Langevin(Integrator):
11 """ Constant NPT Langevin integrator
12 NPT Langevin Leap-frog integrator based on
13 N. Grønbech-Jensen and Oded Farago, J. Chem. Phys. 141, 194108 (2014),
14 `doi:10.1063/1.4901303 <https://doi.org/10.1063/1.4901303>`_
15 """
17 def __init__(self, temperature, pressure, alpha:float, alpha_baro, mass_baro,
18 volume_velocity, barostatModeISO, boxFlucCoord, dt:float, seed:int) -> None:
19 self.temperature = temperature
20 self.pressure = pressure
21 self.alpha = alpha
22 self.alpha_baro = alpha_baro
23 self.mass_baro = mass_baro
24 self.volume_velocity = volume_velocity
25 self.barostatModeISO = barostatModeISO
26 self.boxFlucCoord = boxFlucCoord
27 self.dt = dt
28 self.seed = seed
30 def get_params(self, configuration: gp.Configuration, interactions_params: tuple, verbose=False) -> tuple:
31 dt = np.float32(self.dt)
32 alpha = np.float32(self.alpha)
33 alpha_baro = np.float32(self.alpha_baro)
34 mass_baro = np.float32(self.mass_baro)
35 rng_states = create_xoroshiro128p_states(configuration.N+1, seed=self.seed) # +1 for barostat dynamics
36 barostat_state = np.array([1.0, self.volume_velocity], dtype=np.float64) # [0] = new_vol / old_vol , [1] = vol velocity
37 d_barostat_state = cuda.to_device(barostat_state)
38 barostatVirial = np.array([0.0], dtype=np.float32)
39 d_barostatVirial = cuda.to_device(barostatVirial)
40 d_length_ratio = cuda.to_device(np.ones(3, dtype=np.float32))
41 return (dt, alpha, alpha_baro, mass_baro, # Needs to be compatible with unpacking in step() below
42 self.barostatModeISO, np.int32(self.boxFlucCoord),
43 rng_states, d_barostat_state, d_barostatVirial, d_length_ratio)
45 def get_kernel(self, configuration: gp.Configuration, compute_plan: dict, compute_flags:dict[str,bool], interactions_kernel, verbose=False):
47 # Unpack parameters from configuration and compute_plan
48 D, num_part = configuration.D, configuration.N
49 pb, tp, gridsync = [compute_plan[key] for key in ['pb', 'tp', 'gridsync']]
50 num_blocks = (num_part - 1) // pb + 1
52 # Convert temperature and pressure to a functions id needed
53 if callable(self.temperature):
54 temperature_function = self.temperature
55 else:
56 temperature_function = gp.make_function_constant(value=float(self.temperature))
57 if callable(self.pressure):
58 pressure_function = self.pressure
59 else:
60 pressure_function = gp.make_function_constant(value=float(self.pressure))
62 if verbose:
63 print(f'Generating NPT langevin integrator for {num_part} particles in {D} dimensions:')
64 print(f'\tpb: {pb}, tp:{tp}, num_blocks:{num_blocks}')
65 print(f'\tNumber (virtual) particles: {num_blocks * pb}')
66 print(f'\tNumber of threads {num_blocks * pb * tp}')
68 # Unpack indices for vectors and scalars to be compiled into kernel
69 compute_k = compute_flags['K']
70 compute_fsq = compute_flags['Fsq']
71 r_id, v_id, f_id = [configuration.vectors.indices[key] for key in ['r', 'v', 'f']]
72 m_id = configuration.sid['m']
73 if not compute_flags['W']:
74 raise ValueError("NPT_Langevin requires virial")
75 else:
76 w_id = configuration.sid['W']
78 if compute_k:
79 k_id = configuration.sid['K']
80 if compute_fsq:
81 fsq_id = configuration.sid['Fsq']
84 # JIT compile functions to be compiled into kernel
85 temperature_function = numba.njit(temperature_function)
86 pressure_function = numba.njit(pressure_function)
87 apply_PBC = numba.njit(configuration.simbox.get_apply_PBC())
90 def copyParticleVirial(scalars, integrator_params):
91 dt, alpha, alpha_baro, mass_baro, barostatModeISO, boxFlucCoord, rng_states, barostat_state, barostatVirial, length_ratio = integrator_params
92 global_id, my_t = cuda.grid(2)
93 my_w = scalars[global_id][w_id]
95 #reset the barostatVirial to zero
96 #if global_id == 0 and my_t == 0:
97 # barostatVirial[0] = numba.float32(0.0)
98 # # Not safe to set to zero like this! (moved to end of update_barostat_state)
99 #cuda.syncthreads() # - particles in other blocks might allready have added their w
101 if global_id < num_part and my_t == 0:
102 cuda.atomic.add(barostatVirial, 0, my_w) # factor of 6 already accounted for using virial_factor and virial_factor_NIII
104 return
106 def update_barostat_state(sim_box, integrator_params, time):
107 dt, alpha, alpha_baro, mass_baro, barostatModeISO, boxFlucCoord, rng_states, barostat_state, barostatVirial, length_ratio = integrator_params
108 temperature = temperature_function(time)
109 pressure = pressure_function(time)
111 global_id, my_t = cuda.grid(2)
112 if global_id == 0 and my_t == 0:
114 #Copy barostat_state into current_barostat_state using a local aaray
115 current_barostat_state = cuda.local.array(2, numba.float64)
116 current_barostat_state[0] = barostat_state[0]
117 current_barostat_state[1] = barostat_state[1]
119 volume = sim_box[0]*sim_box[1]*sim_box[2]
120 targetConfPressure = pressure - temperature * num_part / volume
121 barostatForce = barostatVirial[0] / volume - targetConfPressure
123 random_number = xoroshiro128p_normal_float32(rng_states, 0) # 0th random number state is reserved for barostat
124 barostatRandomForce = math.sqrt(numba.float32(2.0) * alpha_baro * temperature * dt) * random_number
126 current_volume_velocity = current_barostat_state[1]
127 inv_baro_mass = numba.float32(1.0) / mass_baro
128 scaled_dt = numba.float32(0.5) * dt * alpha_baro * inv_baro_mass
129 b_tilde = numba.float64(1.0) / (numba.float64(1.0) + scaled_dt)
130 a_tilde = b_tilde * (numba.float64(1.0) - scaled_dt)
132 new_volume_vel = a_tilde * current_volume_velocity + b_tilde * inv_baro_mass * (barostatForce * dt + barostatRandomForce)
133 new_volume = volume + dt * new_volume_vel
135 # Update the barostat state
136 barostat_state[0] = new_volume / volume
137 barostat_state[1] = new_volume_vel
139 # reset length_ratio to 1.0. WHY?
140 for i in range(3):
141 length_ratio[i] = numba.float64(1.0)
143 vol_scale_factor = barostat_state[0]
144 lr_iso = math.pow(vol_scale_factor, numba.float64(1.0) / numba.float64(3.0))
146 #update box length
147 sim_box[0] += sim_box[0] * (lr_iso - 1.) # Better: sim_box[0] = sim_box[0] * lr_iso ?
148 sim_box[1] += sim_box[1] * (lr_iso - 1.)
149 sim_box[2] += sim_box[2] * (lr_iso - 1.)
151 # update length_ratio using cuda loop
152 for i in range(3):
153 length_ratio[i] = lr_iso
155 barostatVirial[0] = numba.float32(0.0)
157 return
160 def step(grid, vectors, scalars, r_im, sim_box, integrator_params, time):
161 """ Make one NPT timestep using Leap-frog
162 Kernel configuration: [num_blocks, (pb, 1)]
163 REF: https://arxiv.org/pdf/1303.7011.pdf
164 """
166 dt, alpha, alpha_baro, mass_baro, barostatModeISO, boxFlucCoord, rng_states, barostat_state, barostatVirial, length_ratio = integrator_params
167 temperature = temperature_function(time)
169 global_id, my_t = cuda.grid(2)
170 if global_id < num_part and my_t == 0:
171 my_r = vectors[r_id][global_id]
172 my_v = vectors[v_id][global_id]
173 my_f = vectors[f_id][global_id]
174 my_m = scalars[global_id][m_id]
175 if compute_k:
176 my_k = numba.float32(0.0) # Kinetic energy
177 if compute_fsq:
178 my_fsq = numba.float32(0.0) # force squared energy
180 for k in range(D):
181 random_number = xoroshiro128p_normal_float32(rng_states, global_id + 1) # +1 to avoid using the same random number state as the barostat
182 beta = math.sqrt(numba.float32(2.0) * alpha * temperature * dt) * random_number
184 scaled_dt = numba.float32(0.5) * dt * alpha * numba.float32(1.0)/my_m
185 prm_b = numba.float64(1.0) / (numba.float64(1.0) + scaled_dt)
186 prm_a = prm_b * (numba.float64(1.0) - scaled_dt)
188 if compute_k:
189 my_k += numba.float32(0.5) * my_m * my_v[k] * my_v[k] # ke before
190 if compute_fsq:
191 my_fsq += my_f[k] * my_f[k]
193 my_v[k] = prm_a * my_v[k] + prm_b * (numba.float32(1.0)/my_m) * (my_f[k] * dt + beta)
194 if compute_k:
195 my_k += numba.float32(0.5) * my_m * my_v[k] * my_v[k] # ke after
197 L_factor = 2.*length_ratio[k] / (1. + length_ratio[k])
198 my_r[k] = length_ratio[k] * my_r[k] + L_factor * my_v[k] * dt
200 apply_PBC(my_r, r_im[global_id], sim_box)
202 if compute_k:
203 scalars[global_id][k_id] = numba.float32(0.5) * my_k
204 if compute_fsq:
205 scalars[global_id][fsq_id] = my_fsq
207 return
209 copyParticleVirial = cuda.jit(device=gridsync)(copyParticleVirial)
210 update_barostat_state = cuda.jit(device=gridsync)(update_barostat_state)
211 step = cuda.jit(device=gridsync)(step)
213 if gridsync:
215 def kernel(grid, vectors, scalars, r_im, sim_box, integrator_params, time, ptype):
216 copyParticleVirial(scalars, integrator_params)
217 grid.sync()
218 update_barostat_state(sim_box, integrator_params, time)
219 grid.sync()
220 step(grid, vectors, scalars, r_im, sim_box, integrator_params, time)
221 return
223 return cuda.jit(device=gridsync)(kernel)
225 else:
227 def kernel(grid, vectors, scalars, r_im, sim_box, integrator_params, time, ptype):
228 copyParticleVirial[num_blocks, (pb, 1)](scalars, integrator_params)
229 update_barostat_state[1, (1, 1)](sim_box, integrator_params, time)
230 step[num_blocks, (pb, 1)](grid, vectors, scalars, r_im, sim_box, integrator_params, time)
231 return
233 return kernel