Coverage for gamdpy/integrators/NVT.py: 53%
89 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 gamdpy as gp
5from .integrator import Integrator
7class NVT(Integrator):
8 """ The leapfrog algorithm with a Nose-Hoover thermostat
10 Integrator keeping N (number of particles), V (volume), and T (temperature) constant,
11 using the leapfrog algorithm and a Nose-Hoover thermostat.
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.
18 tau : float
19 The relaxation time of the thermostat.
21 dt : float
22 The time step of the integration.
24 """
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
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.
40 def get_kernel(self, configuration: gp.Configuration, compute_plan: dict, compute_flags: dict, interactions_kernel, verbose=False):
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
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))
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}')
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']
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())
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 """
80 # Unpack parameters. MUST be compatible with get_params() above
81 dt, omega2, degrees, thermostat_state = integrator_params
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(...)
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
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
104 apply_PBC(my_r, r_im[global_id], sim_box)
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
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
117 global_id, my_t = cuda.grid(2)
118 if global_id == 0 and my_t == 0:
119 target_temperature = temperature_function(time)
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
126 step = cuda.jit(device=gridsync)(step)
127 update_thermostat_state = cuda.jit(device=gridsync)(update_thermostat_state)
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