1# flake8: noqa 

2import numpy as np 

3 

4class DiffusionCoefficient: 

5 

6 def __init__(self, traj, timestep, atom_indices=None, molecule=False): 

7 """ 

8 

9 This class calculates the Diffusion Coefficient for the given Trajectory using the Einstein Equation: 

10  

11 ..math:: \\left \\langle \\left | r(t) - r(0) \\right | ^{2} \\right \\rangle = 2nDt 

12  

13 where r(t) is the position of atom at time t, n is the degrees of freedom and D is the Diffusion Coefficient 

14  

15 Solved herein by fitting with :math:`y = mx + c`, i.e. :math:`\\frac{1}{2n} \\left \\langle \\left | r(t) - r(0) \\right | ^{2} \\right \\rangle = Dt`, with m = D and c = 0 

16 

17 wiki : https://en.wikibooks.org/wiki/Molecular_Simulation/Diffusion_Coefficients 

18 

19 Parameters: 

20 traj (Trajectory):  

21 Trajectory of atoms objects (images) 

22 timestep (Float):  

23 Timestep between *each image in the trajectory*, in ASE timestep units 

24 (For an MD simulation with timestep of N, and images written every M iterations, our timestep here is N * M)  

25 atom_indices (List of Int):  

26 The indices of atoms whose Diffusion Coefficient is to be calculated explicitly 

27 molecule (Boolean) 

28 Indicate if we are studying a molecule instead of atoms, therefore use centre of mass in calculations 

29  

30 """ 

31 

32 self.traj = traj 

33 self.timestep = timestep 

34 

35 # Condition used if user wants to calculate diffusion coefficients for specific atoms or all atoms 

36 self.atom_indices = atom_indices 

37 if self.atom_indices == None: 

38 self.atom_indices = [i for i in range(len(traj[0]))] 

39 

40 # Condition if we are working with the mobility of a molecule, need to manage arrays slightly differently 

41 self.is_molecule = molecule 

42 if self.is_molecule: 

43 self.types_of_atoms = ["molecule"] 

44 self.no_of_atoms = [1] 

45 else: 

46 self.types_of_atoms = sorted(set(traj[0].symbols[self.atom_indices])) 

47 self.no_of_atoms = [traj[0].get_chemical_symbols().count(symbol) for symbol in self.types_of_atoms] 

48 

49 # Dummy initialisation for important results data object 

50 self._slopes = [] 

51 

52 @property 

53 def no_of_types_of_atoms(self): 

54 """ 

55 

56 Dynamically returns the number of different atoms in the system 

57  

58 """ 

59 return len(self.types_of_atoms) 

60 

61 @property 

62 def slopes(self): 

63 """ 

64 

65 Method to return slopes fitted to datapoints. If undefined, calculate slopes 

66 

67 """ 

68 if len(self._slopes) == 0: 

69 self.calculate() 

70 return self._slopes 

71 

72 @slopes.setter 

73 def slopes(self, values): 

74 """ 

75  

76 Method to set slopes as fitted to datapoints 

77 

78 """ 

79 self._slopes = values 

80 

81 def _initialise_arrays(self, ignore_n_images, number_of_segments): 

82 """ 

83 

84 Private function to initialise data storage objects. This includes objects to store the total timesteps 

85 sampled, the average diffusivity for species in any given segment, and objects to store gradient and intercept from fitting. 

86 

87 Parameters: 

88 ignore_n_images (Int):  

89 Number of images you want to ignore from the start of the trajectory, e.g. during equilibration 

90 number_of_segments (Int):  

91 Divides the given trajectory in to segments to allow statistical analysis 

92  

93 """ 

94 

95 total_images = len(self.traj) - ignore_n_images 

96 self.no_of_segments = number_of_segments 

97 self.len_segments = total_images // self.no_of_segments 

98 

99 # These are the data objects we need when plotting information. First the x-axis, timesteps 

100 self.timesteps = np.linspace(0,total_images*self.timestep,total_images+1) 

101 # This holds all the data points for the diffusion coefficients, averaged over atoms 

102 self.xyz_segment_ensemble_average = np.zeros((self.no_of_segments,self.no_of_types_of_atoms,3,self.len_segments)) 

103 # This holds all the information on linear fits, from which we get the diffusion coefficients 

104 self.slopes = np.zeros((self.no_of_types_of_atoms,self.no_of_segments,3)) 

105 self.intercepts = np.zeros((self.no_of_types_of_atoms,self.no_of_segments,3)) 

106 

107 self.cont_xyz_segment_ensemble_average = 0 

108 

109 def calculate(self, ignore_n_images=0, number_of_segments=1): 

110 """ 

111  

112 Calculate the diffusion coefficients, using the previously supplied data. The user can break the data into segments and  

113 take the average over these trajectories, therefore allowing statistical analysis and derivation of standard deviations. 

114 Option is also provided to ignore initial images if data is perhaps unequilibrated initially. 

115 

116 Parameters: 

117 ignore_n_images (Int):  

118 Number of images you want to ignore from the start of the trajectory, e.g. during equilibration 

119 number_of_segments (Int):  

120 Divides the given trajectory in to segments to allow statistical analysis 

121 

122 """ 

123 

124 # Setup all the arrays we need to store information 

125 self._initialise_arrays(ignore_n_images, number_of_segments) 

126 

127 for segment_no in range(self.no_of_segments): 

128 start = segment_no*self.len_segments 

129 end = start + self.len_segments 

130 seg = self.traj[ignore_n_images+start:ignore_n_images+end] 

131 

132 # If we are considering a molecular system, work out the COM for the starting structure 

133 if self.is_molecule: 

134 com_orig = np.zeros(3) 

135 for atom_no in self.atom_indices: 

136 com_orig += seg[0].positions[atom_no] / len(self.atom_indices) 

137 

138 # For each image, calculate displacement. 

139 # I spent some time deciding if this should run from 0 or 1, as the displacement will be zero for  

140 # t = 0, but this is a data point that needs fitting too and so should be included 

141 for image_no in range(0,len(seg)): 

142 # This object collects the xyz displacements for all atom species in the image 

143 xyz_disp = np.zeros((self.no_of_types_of_atoms,3)) 

144 

145 # Calculating for each atom individually, grouping by species type (e.g. solid state) 

146 if not self.is_molecule: 

147 # For each atom, work out displacement from start coordinate and collect information with like atoms 

148 for atom_no in self.atom_indices: 

149 sym_index = self.types_of_atoms.index(seg[image_no].symbols[atom_no]) 

150 xyz_disp[sym_index] += np.square(seg[image_no].positions[atom_no] - seg[0].positions[atom_no]) 

151 

152 else: # Calculating for group of atoms (molecule) and work out squared displacement 

153 com_disp = np.zeros(3) 

154 for atom_no in self.atom_indices: 

155 com_disp += seg[image_no].positions[atom_no] / len(self.atom_indices) 

156 xyz_disp[0] += np.square(com_disp - com_orig) 

157 

158 # For each atom species or molecule, use xyz_disp to calculate the average data  

159 for sym_index in range(self.no_of_types_of_atoms): 

160 # Normalise by degrees of freedom and average overall atoms for each axes over entire segment  

161 denominator = (2*self.no_of_atoms[sym_index]) 

162 for xyz in range(3): 

163 self.xyz_segment_ensemble_average[segment_no][sym_index][xyz][image_no] = (xyz_disp[sym_index][xyz]/denominator) 

164 

165 # We've collected all the data for this entire segment, so now to fit the data. 

166 for sym_index in range(self.no_of_types_of_atoms): 

167 self.slopes[sym_index][segment_no], self.intercepts[sym_index][segment_no] = self._fit_data(self.timesteps[start:end], 

168 self.xyz_segment_ensemble_average[segment_no][sym_index]) 

169 

170 def _fit_data(self, x, y): 

171 """ 

172 Private function that returns slope and intercept for linear fit to mean square diffusion data 

173 

174 

175 Parameters: 

176 x (Array of floats):  

177 Linear list of timesteps in the calculation 

178 y (Array of floats):  

179 Mean square displacement as a function of time. 

180  

181 """ 

182 

183 # Simpler implementation but disabled as fails Conda tests. 

184 # from scipy.stats import linregress 

185 # slope, intercept, r_value, p_value, std_err = linregress(x,y) 

186 

187 # Initialise objects 

188 slopes = np.zeros(3) 

189 intercepts = np.zeros(3) 

190 

191 # Convert into suitable format for lstsq 

192 x_edited = np.vstack([np.array(x), np.ones(len(x))]).T 

193 # Calculate slopes for x, y and z-axes 

194 for xyz in range(3): 

195 slopes[xyz], intercepts[xyz] = np.linalg.lstsq(x_edited, y[xyz], rcond=-1)[0] 

196 

197 return slopes, intercepts 

198 

199 def get_diffusion_coefficients(self): 

200 """ 

201  

202 Returns diffusion coefficients for atoms (in alphabetical order) along with standard deviation. 

203  

204 All data is currently passed out in units of Å^2/<ASE time units> 

205 To convert into Å^2/fs => multiply by ase.units.fs 

206 To convert from Å^2/fs to cm^2/s => multiply by (10^-8)^2 / 10^-15 = 10^-1 

207  

208 """ 

209 

210 slopes = [np.mean(self.slopes[sym_index]) for sym_index in range(self.no_of_types_of_atoms)] 

211 std = [np.std(self.slopes[sym_index]) for sym_index in range(self.no_of_types_of_atoms)] 

212 

213 return slopes, std 

214 

215 def plot(self, ax=None, show=False): 

216 """ 

217  

218 Auto-plot of Diffusion Coefficient data. Provides basic framework for visualising analysis. 

219  

220 Parameters: 

221 ax (Matplotlib.axes.Axes) 

222 Axes object on to which plot can be created 

223 show (Boolean) 

224 Whether or not to show the created plot. Default: False 

225  

226 """ 

227 

228 # Necessary if user hasn't supplied an axis.  

229 import matplotlib.pyplot as plt 

230 

231 # Convert from ASE time units to fs (aesthetic) 

232 from ase.units import fs as fs_conversion 

233 

234 if ax is None: 

235 ax = plt.gca() 

236 

237 # Define some aesthetic variables 

238 color_list = plt.cm.Set3(np.linspace(0, 1, self.no_of_types_of_atoms)) 

239 xyz_labels=['X','Y','Z'] 

240 xyz_markers = ['o','s','^'] 

241 

242 # Create an x-axis that is in a more intuitive format for the view 

243 graph_timesteps = self.timesteps / fs_conversion 

244 

245 for segment_no in range(self.no_of_segments): 

246 start = segment_no*self.len_segments 

247 end = start + self.len_segments 

248 label = None 

249 

250 for sym_index in range(self.no_of_types_of_atoms): 

251 for xyz in range(3): 

252 if segment_no == 0: 

253 label = 'Species: %s (%s)'%(self.types_of_atoms[sym_index], xyz_labels[xyz]) 

254 # Add scatter graph for the mean square displacement data in this segment 

255 ax.scatter(graph_timesteps[start:end], self.xyz_segment_ensemble_average[segment_no][sym_index][xyz], 

256 color=color_list[sym_index], marker=xyz_markers[xyz], label=label, linewidth=1, edgecolor='grey') 

257 

258 # Print the line of best fit for segment  

259 line = np.mean(self.slopes[sym_index][segment_no])*fs_conversion*graph_timesteps[start:end]+np.mean(self.intercepts[sym_index][segment_no]) 

260 if segment_no == 0: 

261 label = 'Segment Mean : %s'%(self.types_of_atoms[sym_index]) 

262 ax.plot(graph_timesteps[start:end], line, color='C%d'%(sym_index), label=label, linestyle='--') 

263 

264 # Plot separator at end of segment 

265 x_coord = graph_timesteps[end-1] 

266 ax.plot([x_coord, x_coord],[-0.001, 1.05*np.amax(self.xyz_segment_ensemble_average)], color='grey', linestyle=":") 

267 

268 # Plot the overall mean (average of slopes) for each atom species 

269 # This only makes sense if the data is all plotted on the same x-axis timeframe, which currently we are not - everything is plotted sequentially 

270 #for sym_index in range(self.no_of_types_of_atoms): 

271 # line = np.mean(self.slopes[sym_index])*graph_timesteps+np.mean(self.intercepts[sym_index]) 

272 # label ='Mean, Total : %s'%(self.types_of_atoms[sym_index]) 

273 # ax.plot(graph_timesteps, line, color='C%d'%(sym_index), label=label, linestyle="-") 

274 

275 # Aesthetic parts of the plot 

276 ax.set_ylim(-0.001, 1.05*np.amax(self.xyz_segment_ensemble_average)) 

277 ax.legend(loc='best') 

278 ax.set_xlabel('Time (fs)') 

279 ax.set_ylabel(r'Mean Square Displacement ($\AA^2$)') 

280 

281 if show: 

282 plt.show() 

283 

284 def print_data(self): 

285 """ 

286  

287 Output of statistical analysis for Diffusion Coefficient data. Provides basic framework for understanding calculation. 

288  

289 """ 

290 

291 from ase.units import fs as fs_conversion 

292 

293 # Collect statistical data for diffusion coefficient over all segments 

294 slopes, std = self.get_diffusion_coefficients() 

295 

296 # Useful notes for any consideration of conversion.  

297 # Converting gradient from Å^2/fs to more common units of cm^2/s => multiplying by (10^-8)^2 / 10^-15 = 10^-1 

298 # Converting intercept from Å^2 to more common units of cm^2 => multiply by (10^-8)^2 = 10^-16 

299 # 

300 # Note currently in ASE internal time units 

301 # Converting into fs => divide by 1/(fs_conversion) => multiply by (fs_conversion) 

302 

303 # Print data for each atom, in each segment. 

304 for sym_index in range(self.no_of_types_of_atoms): 

305 print('---') 

306 print(r'Species: %4s' % self.types_of_atoms[sym_index]) 

307 print('---') 

308 for segment_no in range(self.no_of_segments): 

309 print(r'Segment %3d: Diffusion Coefficient = %.10f Å^2/fs; Intercept = %.10f Å^2;' % 

310 (segment_no, np.mean(self.slopes[sym_index][segment_no])*fs_conversion, np.mean(self.intercepts[sym_index][segment_no]))) 

311 

312 # Print average overall data. 

313 print('---') 

314 for sym_index in range(self.no_of_types_of_atoms): 

315 print('Mean Diffusion Coefficient (X, Y and Z) : %s = %.10f Å^2/fs; Std. Dev. = %.10f Å^2/fs' % 

316 (self.types_of_atoms[sym_index], slopes[sym_index]*fs_conversion, std[sym_index]*fs_conversion)) 

317 print('---')