Coverage for /builds/kinetik161/ase/ase/md/analysis.py: 63.48%
115 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-10 11:04 +0000
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[0]))]
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 = [1]
48 else:
49 self.types_of_atoms = sorted(
50 set(traj[0].symbols[self.atom_indices]))
51 self.no_of_atoms = [traj[0].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[0].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[0].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[0] += 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)[0]
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: {} ({})'.format(
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('---')