 r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# flake8: noqa

2import numpy as np

5class DiffusionCoefficient:

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

8 """

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

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

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

16 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

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

20 Parameters:

21 traj (Trajectory):

22 Trajectory of atoms objects (images)

23 timestep (Float):

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

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

26 atom_indices (List of Int):

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

28 molecule (Boolean)

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

31 """

33 self.traj = traj

34 self.timestep = timestep

36 # Condition used if user wants to calculate diffusion coefficients for

37 # specific atoms or all atoms

38 self.atom_indices = atom_indices

39 if self.atom_indices is None:

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

42 # Condition if we are working with the mobility of a molecule, need to

43 # manage arrays slightly differently

44 self.is_molecule = molecule

45 if self.is_molecule:

46 self.types_of_atoms = ["molecule"]

47 self.no_of_atoms = 

48 else:

49 self.types_of_atoms = sorted(

50 set(traj.symbols[self.atom_indices]))

51 self.no_of_atoms = [traj.get_chemical_symbols().count(

52 symbol) for symbol in self.types_of_atoms]

54 # Dummy initialisation for important results data object

55 self._slopes = []

57 @property

58 def no_of_types_of_atoms(self):

59 """

61 Dynamically returns the number of different atoms in the system

63 """

64 return len(self.types_of_atoms)

66 @property

67 def slopes(self):

68 """

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

72 """

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

74 self.calculate()

75 return self._slopes

77 @slopes.setter

78 def slopes(self, values):

79 """

81 Method to set slopes as fitted to datapoints

83 """

84 self._slopes = values

86 def _initialise_arrays(self, ignore_n_images, number_of_segments):

87 """

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

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

92 Parameters:

93 ignore_n_images (Int):

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

95 number_of_segments (Int):

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

98 """

100 total_images = len(self.traj) - ignore_n_images

101 self.no_of_segments = number_of_segments

102 self.len_segments = total_images // self.no_of_segments

104 # These are the data objects we need when plotting information. First

105 # the x-axis, timesteps

106 self.timesteps = np.linspace(

107 0, total_images*self.timestep, total_images+1)

108 # This holds all the data points for the diffusion coefficients,

109 # averaged over atoms

110 self.xyz_segment_ensemble_average = np.zeros(

111 (self.no_of_segments, self.no_of_types_of_atoms, 3, self.len_segments))

112 # This holds all the information on linear fits, from which we get the

113 # diffusion coefficients

114 self.slopes = np.zeros(

115 (self.no_of_types_of_atoms, self.no_of_segments, 3))

116 self.intercepts = np.zeros(

117 (self.no_of_types_of_atoms, self.no_of_segments, 3))

119 self.cont_xyz_segment_ensemble_average = 0

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

122 """

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

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

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

128 Parameters:

129 ignore_n_images (Int):

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

131 number_of_segments (Int):

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

134 """

136 # Setup all the arrays we need to store information

137 self._initialise_arrays(ignore_n_images, number_of_segments)

139 for segment_no in range(self.no_of_segments):

140 start = segment_no*self.len_segments

141 end = start + self.len_segments

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

144 # If we are considering a molecular system, work out the COM for the

145 # starting structure

146 if self.is_molecule:

147 com_orig = np.zeros(3)

148 for atom_no in self.atom_indices:

149 com_orig += seg.positions[atom_no] / \

150 len(self.atom_indices)

152 # For each image, calculate displacement.

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

154 # t = 0, but this is a data point that needs fitting too and so

155 # should be included

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

157 # This object collects the xyz displacements for all atom

158 # species in the image

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

161 # Calculating for each atom individually, grouping by species

162 # type (e.g. solid state)

163 if not self.is_molecule:

164 # For each atom, work out displacement from start coordinate

165 # and collect information with like atoms

166 for atom_no in self.atom_indices:

167 sym_index = self.types_of_atoms.index(

168 seg[image_no].symbols[atom_no])

169 xyz_disp[sym_index] += np.square(

170 seg[image_no].positions[atom_no] - seg.positions[atom_no])

172 # Calculating for group of atoms (molecule) and work out squared

173 # displacement

174 else:

175 com_disp = np.zeros(3)

176 for atom_no in self.atom_indices:

177 com_disp += seg[image_no].positions[atom_no] / \

178 len(self.atom_indices)

179 xyz_disp += np.square(com_disp - com_orig)

181 # For each atom species or molecule, use xyz_disp to calculate

182 # the average data

183 for sym_index in range(self.no_of_types_of_atoms):

184 # Normalise by degrees of freedom and average overall atoms

185 # for each axes over entire segment

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

187 for xyz in range(3):

188 self.xyz_segment_ensemble_average[segment_no][sym_index][xyz][image_no] = (

189 xyz_disp[sym_index][xyz]/denominator)

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

192 # fit the data.

193 for sym_index in range(self.no_of_types_of_atoms):

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

195 self.xyz_segment_ensemble_average[segment_no][sym_index])

197 def _fit_data(self, x, y):

198 """

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

202 Parameters:

203 x (Array of floats):

204 Linear list of timesteps in the calculation

205 y (Array of floats):

206 Mean square displacement as a function of time.

208 """

210 # Simpler implementation but disabled as fails Conda tests.

211 # from scipy.stats import linregress

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

214 # Initialise objects

215 slopes = np.zeros(3)

216 intercepts = np.zeros(3)

218 # Convert into suitable format for lstsq

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

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

221 for xyz in range(3):

222 slopes[xyz], intercepts[xyz] = np.linalg.lstsq(

223 x_edited, y[xyz], rcond=-1)

225 return slopes, intercepts

227 def get_diffusion_coefficients(self):

228 """

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

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

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

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

236 """

238 slopes = [np.mean(self.slopes[sym_index])

239 for sym_index in range(self.no_of_types_of_atoms)]

240 std = [np.std(self.slopes[sym_index])

241 for sym_index in range(self.no_of_types_of_atoms)]

243 return slopes, std

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

246 """

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

250 Parameters:

251 ax (Matplotlib.axes.Axes)

252 Axes object on to which plot can be created

253 show (Boolean)

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

256 """

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

259 import matplotlib.pyplot as plt

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

262 from ase.units import fs as fs_conversion

264 if ax is None:

265 ax = plt.gca()

267 # Define some aesthetic variables

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

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

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

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

273 graph_timesteps = self.timesteps / fs_conversion

275 for segment_no in range(self.no_of_segments):

276 start = segment_no*self.len_segments

277 end = start + self.len_segments

278 label = None

280 for sym_index in range(self.no_of_types_of_atoms):

281 for xyz in range(3):

282 if segment_no == 0:

283 label = 'Species: %s (%s)' % (

284 self.types_of_atoms[sym_index], xyz_labels[xyz])

285 # Add scatter graph for the mean square displacement data

286 # in this segment

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

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

290 # Print the line of best fit for segment

291 line = np.mean(self.slopes[sym_index][segment_no])*fs_conversion * \

292 graph_timesteps[start:end] + \

293 np.mean(self.intercepts[sym_index][segment_no])

294 if segment_no == 0:

295 label = 'Segment Mean : %s' % (

296 self.types_of_atoms[sym_index])

297 ax.plot(graph_timesteps[start:end], line, color='C%d' % (

298 sym_index), label=label, linestyle='--')

300 # Plot separator at end of segment

301 x_coord = graph_timesteps[end-1]

302 ax.plot([x_coord,

303 x_coord],

304 [-0.001,

305 1.05*np.amax(self.xyz_segment_ensemble_average)],

306 color='grey',

307 linestyle=":")

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

310 # 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

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

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

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

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

316 # Aesthetic parts of the plot

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

318 ax.legend(loc='best')

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

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

322 if show:

323 plt.show()

325 def print_data(self):

326 """

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

330 """

332 from ase.units import fs as fs_conversion

334 # Collect statistical data for diffusion coefficient over all segments

335 slopes, std = self.get_diffusion_coefficients()

337 # Useful notes for any consideration of conversion.

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

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

340 #

341 # Note currently in ASE internal time units

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

343 # (fs_conversion)

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

346 for sym_index in range(self.no_of_types_of_atoms):

347 print('---')

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

349 print('---')

350 for segment_no in range(self.no_of_segments):

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

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

354 # Print average overall data.

355 print('---')

356 for sym_index in range(self.no_of_types_of_atoms):

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

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

359 print('---')