Coverage for gamdpy/calculators/calculator_widom_insertion.py: 13%

1import math 

2 

3import numpy as np 

4import numba 

5from numba import cuda 

6 

7from ..simulation.get_default_compute_plan import get_default_compute_plan 

8 

9class CalculatorWidomInsertion: 

10 """ Calculator class for Widom's particle insertion method for computing chemical potentials 

11 

12 This class is used to calculate the chemical potential of a 

13 not to dense fluid using Widom's particle insertion method. 

14 The excess chemical potential, μᵉˣ, fluid can be computed  

15 as μᵉˣ = -kT ln 〈exp(-ΔU/kT)〉 where ΔU is the energy difference  

16 between the system with and without a ghost particle. 

17 Here, 〈...〉 is an average for all possible positions  

18 of the ghost particle (sometimes writte as an integral over space). 

19 

20 The method should be used on a NVT trajectory in equlibrium. 

21 

22 Parameters 

23 ---------- 

24 

25 configuration : gamdpy.Configuration 

26 The configuration object representing the system. 

27 

28 pair_potential : gamdpy.PairPotential 

29 The pair potential object representing the interaction between particles. 

30 

31 temperature : float 

32 The temperature of the system. 

33 

34 ghost_positions : ndarray 

35 The positions of the ghost particles for which the chemical potential is to be calculated. 

36 

37 ptype_ghost : int, optional 

38 The particle type of the ghost particles. Default is 0. 

39 

40 compute_plan : dict, optional 

41 The compute plan for the system. Default is None (then a default plan is used). 

42 

43 backend : str, optional 

44 The backend to use for the calculations. Default is 'GPU'.  

45 Supported backends are 'CPU' (testing) and 'GPU' (recomended). 

46  

47 Example 

48 ------- 

49 >>> import numpy as np 

50 >>> import gamdpy as gp 

51 >>> sim = gp.get_default_sim() # Replace with your own equbriliated simulation 

52 >>> pair_pot = sim.interactions[0] 

53 >>> num_ghost_particles = 500_000 

54 >>> ghost_positions = np.random.rand(num_ghost_particles, sim.configuration.D) * sim.configuration.simbox.get_lengths() 

55 >>> calc_widom = gp.CalculatorWidomInsertion(sim.configuration, pair_pot, sim.integrator.temperature, ghost_positions) 

56 >>> for block in sim.run_timeblocks(): 

57 ... calc_widom.update() 

58 >>> calc_widom_data = calc_widom.read() # Dictionary with chemical potential and more 

59 >>> calc_widom_data.keys() 

60 dict_keys(['chemical_potential', 'boltzmann_factors', 'chemical_potentials']) 

61 >>> mu_ex = calc_widom_data['chemical_potential'] # Estimated excess chemical potential 

62 """ 

63 

64 KNOWN_BACKENDS = ['CPU', 'GPU'] 

65 

66 def __init__(self, configuration, pair_potential, temperature, ghost_positions, ptype_ghost=0, compute_plan=None, backend='GPU') -> None: 

67 if backend not in self.KNOWN_BACKENDS: 

68 raise ValueError(f'Unknown backend. Try one of {self.KNOWN_BACKENDS=}') 

69 

70 self.configuration = configuration 

71 self.pair_potential = pair_potential 

72 self.temperature = np.float64(temperature) 

73 self.ghost_positions = np.array(ghost_positions, dtype=np.float32) 

74 self.ptype_ghost = ptype_ghost # The ghost particle of this type 

75 self.num_updates = 0 # How many times have statistics been added to? 

76 self.chemical_potential = 0 # Current best estimate of the chemical potential 

77 self.chemical_potentials = [] # All estimates of the chemical potential 

78 

79 # Expect that positions are in the same dimension as the configuration 

80 if len(self.ghost_positions.shape) != 2: 

81 raise ValueError('The ghost positions must be a 2D array (like [[x1, y1, z1], [x2, y2, z2], ...])') 

82 if self.ghost_positions.shape[1] != self.configuration.D: 

83 raise ValueError('The ghost positions must have the same spatial dimension as the configuration') 

84 

85 self.compute_plan = compute_plan 

86 if self.compute_plan is None: 

87 self.compute_plan = get_default_compute_plan(configuration=configuration) 

88 

89 self.backend = backend 

90 self.boltzmann_factors = np.zeros(len(self.ghost_positions), dtype=np.float64) 

91 self.boltzmann_factors_sum = np.zeros_like(self.boltzmann_factors) 

92 self.boltzmann_factors_timeblock = [] 

93 

94 # Make kernel 

95 if self.backend == 'GPU': 

96 self.update_kernel = self._make_updater_kernel() 

97 

98 def _update_CPU(self): 

99 """ Return the boltzmann factors for the ghost particles in the current configuration (CPU backend). """ 

100 pair_pot_func = self.pair_potential.pairpotential_function 

101 this_boltzmann_factors = np.zeros(len(self.ghost_positions), dtype=np.float32) 

102 for idx_ghost in range(len(self.ghost_positions)): 

103 ghost_pos = self.ghost_positions[idx_ghost] 

104 this_u = 0.0 

105 for n in range(self.configuration.N): 

106 ptype = self.configuration.ptype[n] 

107 r = self.configuration.vectors['r'][n] 

108 dr2 = self.configuration.simbox.dist_sq_function(r, ghost_pos, self.configuration.simbox.get_lengths()) 

109 params = self.pair_potential.params[ptype, self.ptype_ghost] 

110 r_cut = params[-1] 

111 if dr2 < r_cut*r_cut: 

112 dr = np.sqrt(dr2) 

113 this_u += pair_pot_func(dr, params)[0] 

114 this_boltzmann_factors[idx_ghost] = math.exp(-this_u / self.temperature) 

115 return this_boltzmann_factors 

116 

117 def _make_updater_kernel(self): 

118 """ Make the kernel for updating the radial distribution function (GPU backend). """ 

119 # Unpack parameters from configuration and compute_plan 

120 D, num_part = self.configuration.D, self.configuration.N 

121 num_ghost_particles = len(self.ghost_positions) 

122 pb, tp, gridsync, UtilizeNIII = [self.compute_plan[key] for key in ['pb', 'tp', 'gridsync', 'UtilizeNIII']] 

123 num_blocks = (num_ghost_particles - 1) // pb + 1 

124 

125 # Allocate to device 

126 self.d_ghost_positions = cuda.to_device(self.ghost_positions) 

127 self.d_ptype_ghost = np.int32(self.ptype_ghost) 

128 self.d_boltzmann_factors = cuda.to_device(self.boltzmann_factors_sum) 

129 self.d_temperature = cuda.to_device(self.temperature) 

130 

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

132 r_id, f_id = [self.configuration.vectors.indices[key] for key in ['r', 'f']] 

133 

134 # Prepare user-specified functions for inclusion in kernel(s) 

135 ptype_function = numba.njit(self.configuration.ptype_function) 

136 params_function = numba.njit(self.pair_potential.params_function) 

137 # pairpotential_function = numba.njit(self.pair_potential.pairpotential_function) 

138 pairpotential_function = self.pair_potential.pairpotential_function 

139 dist_sq_function = numba.njit(self.configuration.simbox.get_dist_sq_function()) 

140 

141 def update_kernel(vectors, sim_box, ptype, params, ghost_positions, ptype_ghost, temperature, boltzmann_factors): 

142 global_id = cuda.grid(1) 

143 if global_id < len(ghost_positions): 

144 boltzmann_factors[global_id] = np.float64(0.0) 

145 

146 this_u = np.float64(0.0) 

147 for other_global_id in range(0, num_part, 1): 

148 dist_sq = dist_sq_function(vectors[r_id][other_global_id], 

149 ghost_positions[global_id], sim_box) 

150 ij_params = params_function(ptype[other_global_id], ptype_ghost, params) 

151 cut = ij_params[-1] 

152 if dist_sq < cut*cut: 

153 dist = math.sqrt(dist_sq) 

154 pair_energy = pairpotential_function(dist, ij_params)[0] 

155 this_u += np.float64(pair_energy) 

156 boltzmann_factors[global_id] = math.exp(-this_u/temperature) 

157 

158 return 

159 

160 return cuda.jit(device=False)(update_kernel)[num_blocks, (pb,)] 

161 

162 def update(self): 

163 """ Update state the current configuration. """ 

164 self.num_updates += 1 

165 this_boltzmann_factors = None 

166 if self.backend == 'CPU': 

167 this_boltzmann_factors = self._update_CPU() 

168 elif self.backend == 'GPU': 

169 self.update_kernel(self.configuration.d_vectors, 

170 self.configuration.simbox.d_data, 

171 self.configuration.d_ptype, 

172 self.pair_potential.d_params, 

173 self.d_ghost_positions, 

174 self.d_ptype_ghost, 

175 self.temperature, 

176 self.d_boltzmann_factors) 

177 this_boltzmann_factors = self.d_boltzmann_factors.copy_to_host() 

178 else: 

179 raise ValueError(f'Unknown backend. Try one of {self.KNOWN_BACKENDS=}') 

180 self.boltzmann_factors_sum += this_boltzmann_factors 

181 self.boltzmann_factors = self.boltzmann_factors_sum/self.num_updates 

182 self.chemical_potential = -self.temperature*math.log(np.mean(self.boltzmann_factors)) 

183 this_chemical_potential = -self.temperature*math.log(np.mean(this_boltzmann_factors)) 

184 self.chemical_potentials.append(this_chemical_potential) 

185 

186 def read(self): 

187 """ Read data 

188  

189 Return the current chemical potential,  

190 average Boltzmann factors, and timeblock-specific chemical potentials for the ghost particles. 

191 

192 Returns 

193 ------- 

194 dict 

195 A dictionary containing the following keys: 

196  

197 - 'chemical_potential': float 

198 The best estimate of the overall chemical potential for the system. 

199  

200 - 'boltzmann_factors': ndarray 

201 The average Boltzmann factors for each ghost particle, representing the statistical weights based on their interaction energies. 

202  

203 - 'chemical_potentials': ndarray 

204 An array of estimated chemical potentials for each timeblock, providing insight into the variability of the chemical potential over time. 

205 """ 

206 

207 return { 

208 'chemical_potential': self.chemical_potential, # Best estimate of the chemical potential 

209 'boltzmann_factors': self.boltzmann_factors, # Average Boltzmann factors for the ghost particles 

210 'chemical_potentials': self.chemical_potentials # Estimates of the chemical potentials for each timeblock 

211 } 

212