Coverage for gamdpy/integrators/NPT_Atomic.py: 96%
57 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1import numpy as np
2# LC: check https://orbit.dtu.dk/en/publications/cudarray-cuda-based-numpy
3import numba
4from numba import cuda
5from gamdpy import Configuration
6from gamdpy.misc.make_function import make_function_constant
7from .integrator import Integrator
9class NPT_Atomic(Integrator):
10 """ Constant NPT integrator for atomic systems
12 Integrator keeping N (number of particles), P (pressure), and T (temperature) constant,
13 using the leapfrog algorithm and the thermostat-barostat by G Martyna, DJ Tobias and ML Klein in
14 Molecular Physics 87(5), 1117 (1996), `doi:10.1063/1.467468 <https://doi.org/10.1080/00268979600100761>`_
15 Note that the thermostat and barostat states defined here are :math:`p_\\xi/Q` and :math:`p_\\epsilon/W` from Eq. 2.9 in the paper.
17 Parameters
18 ----------
19 temperature : float or function
20 Target temperature as a function of time or constant
22 tau : float
23 Thermostat relaxation time
25 pressure : float or function
26 Target pressure as a function of time or constant
28 tau_p : float
29 Barostat relaxation time
31 dt : float
32 Timestep size
34 """
36 def __init__(self, temperature, tau: float, pressure, tau_p : float, dt: float) -> None:
37 self.temperature = temperature
38 self.tau = tau
39 self.pressure = pressure
40 self.tau_p = tau_p
41 self.dt = dt
42 self.thermostat_state = np.zeros(2, dtype=np.float32)
43 self.barostat_state = np.zeros(3, dtype=np.float32) # NOTE: array is (barostat_state, virial, volume)
45 def get_params(self, configuration: Configuration, interactions_params: tuple, verbose=False) -> tuple:
46 dt = np.float32(self.dt)
47 degrees = configuration.N * configuration.D - configuration.D # number of degrees of freedom
48 factor = np.float32(1./(4*np.pi*np.pi))
49 mass_t = np.float32((degrees-1)*configuration.D * factor * self.tau * self.tau ) # This quantity is the thermostat mass expect for a factor of temperature
50 mass_p = np.float32((degrees-1)*configuration.D * factor * self.tau_p * self.tau_p) # the temperature is missing because at this stage could be a function
51 self.barostat_state[2] = configuration.get_volume() # Copy starting volume (can be avoided)
52 # Copy state variables to device
53 self.d_barostat_state = numba.cuda.to_device(self.barostat_state)
54 self.d_thermostat_state = numba.cuda.to_device(self.thermostat_state)
55 return (dt, mass_t, mass_p, degrees, self.d_thermostat_state, self.d_barostat_state) # Needs to be compatible with unpacking in
56 # step() and update_thermostat_state() below.
58 def get_kernel(self, configuration: Configuration, compute_plan: dict, compute_flags:dict, interactions_kernel, verbose=False):
60 # Unpack parameters from configuration and compute_plan
61 D, num_part = configuration.D, configuration.N
62 pb, tp, gridsync = [compute_plan[key] for key in ['pb', 'tp', 'gridsync']]
63 num_blocks = (num_part - 1) // pb + 1
65 # Convert temperature to a function if isn't already (better be a number then...)
66 if callable(self.temperature):
67 temperature_function = self.temperature
68 else:
69 temperature_function = make_function_constant(value=float(self.temperature))
70 # Convert pressure to a function if isn't already (better be a number then...)
71 if callable(self.pressure):
72 pressure_function = self.pressure
73 else:
74 pressure_function = make_function_constant(value=float(self.pressure))
76 if verbose:
77 print(f'Generating NPT kernel for {num_part} particles in {D} dimensions:')
78 print(f'\tpb: {pb}, tp:{tp}, num_blocks:{num_blocks}')
79 print(f'\tNumber (virtual) particles: {num_blocks * pb}')
80 print(f'\tNumber of threads {num_blocks * pb * tp}')
82 # Unpack indices for vectors and scalars to be compiled into kernel
83 compute_k = compute_flags['K']
84 compute_fsq = compute_flags['Fsq']
85 r_id, v_id, f_id = [configuration.vectors.indices[key] for key in ['r', 'v', 'f']]
86 m_id = configuration.sid['m']
87 if not compute_flags['W']:
88 raise ValueError("NPT_Atomic requires virial")
89 else:
90 w_id = configuration.sid['W']
92 if compute_k:
93 k_id = configuration.sid['K']
94 if compute_fsq:
95 fsq_id = configuration.sid['Fsq']
97 # JIT compile functions to be compiled into kernel
98 temperature_function = numba.njit(temperature_function)
99 pressure_function = numba.njit(pressure_function)
100 apply_PBC = numba.njit(configuration.simbox.get_apply_PBC())
103 def step(grid, vectors, scalars, r_im, sim_box, integrator_params, time): # pragma: no cover
104 """ Make one NPT timestep using Leap-frog
105 Kernel configuration: [num_blocks, (pb, tp)]
106 """
108 # Unpack parameters. MUST be compatible with get_params() above
109 dt, mass_t, mass_p, degrees, thermostat_state, barostat_state = integrator_params
111 # NOTE: thermostat_state and barostat_state have units of inverse time ([t]**-1)
112 factor = np.float32(0.5) * dt * ((1+3./degrees)*barostat_state[0] + thermostat_state[0]) # D=3
113 plus = np.float32(1.) / np.float32(1. + factor)
114 minus = np.float32(1. - factor)
115 rfactor = barostat_state[0]*dt
117 global_id, my_t = cuda.grid(2)
118 if global_id < num_part and my_t == 0:
119 my_r = vectors[r_id][global_id]
120 my_v = vectors[v_id][global_id]
121 my_f = vectors[f_id][global_id]
122 my_w = scalars[global_id][w_id] # Get virial from scalars
123 my_m = scalars[global_id][m_id]
124 my_k = numba.float32(0.0) # Kinetic energy
125 if compute_fsq:
126 my_fsq = numba.float32(0.0) # force squared
128 for k in range(D):
129 if compute_fsq:
130 my_fsq += my_f[k] * my_f[k]
131 my_v[k] = plus * (minus * my_v[k] + my_f[k] / my_m * dt)
132 my_k += numba.float32(0.5) * my_m * my_v[k] * my_v[k] # This is kinetic energy at the half step
133 my_r[k] += my_v[k] * dt + rfactor*my_r[k]
135 apply_PBC(my_r, r_im[global_id], sim_box)
137 cuda.atomic.add(thermostat_state, 1, my_k) # Probably slow? Not really!
138 cuda.atomic.add(barostat_state , 1, my_w)
139 if compute_k:
140 scalars[global_id][k_id] = my_k
141 if compute_fsq:
142 scalars[global_id][fsq_id] = my_fsq
143 return
145 # This function update barostat (P) and thermostat (T) states
146 def update_thermostat_barostat_state(vectors, sim_box, integrator_params, time): # pragma: no cover
147 # Unpack parameters. MUST be compatible with get_params() above
148 dt, mass_t, mass_p, degrees, thermostat_state, barostat_state = integrator_params
150 global_id, my_t = cuda.grid(2)
151 if global_id == 0 and my_t == 0:
152 temperature = 2*thermostat_state[1]/degrees
153 scale_factor_3 = 1 + dt*D*barostat_state[0] # assumes D=3
154 mass_t = mass_t * temperature # fix masses including temperature factor
155 mass_p = mass_p * temperature # thermostat/barostat masses are extensive quantities
156 # Scale volume
157 barostat_state[2] *= scale_factor_3
158 # Scale simbox lenghts
159 scale_factor = scale_factor_3**(1./3)
160 sim_box[0] *= scale_factor
161 sim_box[1] *= scale_factor
162 sim_box[2] *= scale_factor
163 # Update states
164 target_temperature = temperature_function(time)
165 target_pressure = pressure_function(time)
166 instant_pressure = (temperature*num_part + barostat_state[1])/barostat_state[2] # PV = N*T + W
167 # Note that thermostat_state[0] and barostat_state[0] are intensive quantities
168 ke_deviation = (np.float32(2.0) * thermostat_state[1] + barostat_state[0]*barostat_state[0]*mass_p) / ((degrees+1)*target_temperature) - np.float32(1.0)
169 p_deviation = D * (barostat_state[2] * (instant_pressure - target_pressure) + np.float32(0.5)*thermostat_state[1]/degrees) # D=3
170 # Update states
171 thermostat_state[0] += dt * ke_deviation * (degrees+1)*target_temperature / mass_t
172 barostat_state[0] += dt * (p_deviation / mass_p + thermostat_state[0]*barostat_state[0])
173 # Reset tmp variables, these variables are used to pass values between step() and here
174 thermostat_state[1] = np.float32(0.)
175 barostat_state[1] = np.float32(0.)
176 return
178 # Scale the simulation box to the new density
179 def scale_box(vectors, sim_box, integrator_params): # pragma: no cover
180 # Unpack parameters. MUST be compatible with get_params() above
181 dt, mass_t, mass_p, degrees, thermostat_state, barostat_state = integrator_params
183 scale_factor_3 = 1 + dt*D*barostat_state[0] # assumes D=3
184 scale_factor = scale_factor_3**(1./3)
186 global_id, my_t = cuda.grid(2)
187 if global_id < num_part and my_t == 0:
188 for k in range(D):
189 vectors[r_id][global_id][k] = scale_factor*vectors[r_id][global_id][k]
190 return
192 step = numba.cuda.jit(device=gridsync)(step)
193 update_thermostat_barostat_state = numba.cuda.jit(device=gridsync)(update_thermostat_barostat_state)
194 scale_box = numba.cuda.jit(device=gridsync)(scale_box)
196 if gridsync: # pragma: no cover
197 # construct and return device function
198 def kernel(grid, vectors, scalars, r_im, sim_box, integrator_params, time, ptype):
199 step( grid, vectors, scalars, r_im, sim_box, integrator_params, time)
200 grid.sync()
201 scale_box(vectors, sim_box, integrator_params)
202 grid.sync()
203 update_thermostat_barostat_state(vectors, sim_box, integrator_params, time)
204 return
205 return numba.cuda.jit(device=gridsync)(kernel)
206 else: # pragma: no cover
207 # return python function, which makes kernel-calls
208 def kernel(grid, vectors, scalars, r_im, sim_box, integrator_params, time, ptype):
209 step[num_blocks, (pb, 1)](grid, vectors, scalars, r_im, sim_box, integrator_params, time)
210 scale_box[num_blocks, (pb, 1)](vectors, sim_box, integrator_params)
211 update_thermostat_barostat_state[1, (1, 1)](vectors, sim_box, integrator_params, time)
212 return
213 return kernel