Coverage for /builds/ase/ase/ase/md/analysis.py: 63.48%
115 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
3# flake8: noqa
4import numpy as np
7class DiffusionCoefficient:
9 def __init__(self, traj, timestep, atom_indices=None, molecule=False):
10 """
12 This class calculates the Diffusion Coefficient for the given Trajectory using the Einstein Equation:
14 ..math:: \\left \\langle \\left | r(t) - r(0) \\right | ^{2} \\right \\rangle = 2nDt
16 where r(t) is the position of atom at time t, n is the degrees of freedom and D is the Diffusion Coefficient
18 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
20 wiki : https://en.wikibooks.org/wiki/Molecular_Simulation/Diffusion_Coefficients
22 Parameters:
23 traj (Trajectory):
24 Trajectory of atoms objects (images)
25 timestep (Float):
26 Timestep between *each image in the trajectory*, in ASE timestep units
27 (For an MD simulation with timestep of N, and images written every M iterations, our timestep here is N * M)
28 atom_indices (List of Int):
29 The indices of atoms whose Diffusion Coefficient is to be calculated explicitly
30 molecule (Boolean)
31 Indicate if we are studying a molecule instead of atoms, therefore use centre of mass in calculations
33 """
35 self.traj = traj
36 self.timestep = timestep
38 # Condition used if user wants to calculate diffusion coefficients for
39 # specific atoms or all atoms
40 self.atom_indices = atom_indices
41 if self.atom_indices is None:
42 self.atom_indices = [i for i in range(len(traj[0]))]
44 # Condition if we are working with the mobility of a molecule, need to
45 # manage arrays slightly differently
46 self.is_molecule = molecule
47 if self.is_molecule:
48 self.types_of_atoms = ["molecule"]
49 self.no_of_atoms = [1]
50 else:
51 self.types_of_atoms = sorted(
52 set(traj[0].symbols[self.atom_indices]))
53 self.no_of_atoms = [traj[0].get_chemical_symbols().count(
54 symbol) for symbol in self.types_of_atoms]
56 # Dummy initialisation for important results data object
57 self._slopes = []
59 @property
60 def no_of_types_of_atoms(self):
61 """
63 Dynamically returns the number of different atoms in the system
65 """
66 return len(self.types_of_atoms)
68 @property
69 def slopes(self):
70 """
72 Method to return slopes fitted to datapoints. If undefined, calculate slopes
74 """
75 if len(self._slopes) == 0:
76 self.calculate()
77 return self._slopes
79 @slopes.setter
80 def slopes(self, values):
81 """
83 Method to set slopes as fitted to datapoints
85 """
86 self._slopes = values
88 def _initialise_arrays(self, ignore_n_images, number_of_segments):
89 """
91 Private function to initialise data storage objects. This includes objects to store the total timesteps
92 sampled, the average diffusivity for species in any given segment, and objects to store gradient and intercept from fitting.
94 Parameters:
95 ignore_n_images (Int):
96 Number of images you want to ignore from the start of the trajectory, e.g. during equilibration
97 number_of_segments (Int):
98 Divides the given trajectory in to segments to allow statistical analysis
100 """
102 total_images = len(self.traj) - ignore_n_images
103 self.no_of_segments = number_of_segments
104 self.len_segments = total_images // self.no_of_segments
106 # These are the data objects we need when plotting information. First
107 # the x-axis, timesteps
108 self.timesteps = np.linspace(
109 0, total_images * self.timestep, total_images + 1)
110 # This holds all the data points for the diffusion coefficients,
111 # averaged over atoms
112 self.xyz_segment_ensemble_average = np.zeros(
113 (self.no_of_segments, self.no_of_types_of_atoms, 3, self.len_segments))
114 # This holds all the information on linear fits, from which we get the
115 # diffusion coefficients
116 self.slopes = np.zeros(
117 (self.no_of_types_of_atoms, self.no_of_segments, 3))
118 self.intercepts = np.zeros(
119 (self.no_of_types_of_atoms, self.no_of_segments, 3))
121 self.cont_xyz_segment_ensemble_average = 0
123 def calculate(self, ignore_n_images=0, number_of_segments=1):
124 """
126 Calculate the diffusion coefficients, using the previously supplied data. The user can break the data into segments and
127 take the average over these trajectories, therefore allowing statistical analysis and derivation of standard deviations.
128 Option is also provided to ignore initial images if data is perhaps unequilibrated initially.
130 Parameters:
131 ignore_n_images (Int):
132 Number of images you want to ignore from the start of the trajectory, e.g. during equilibration
133 number_of_segments (Int):
134 Divides the given trajectory in to segments to allow statistical analysis
136 """
138 # Setup all the arrays we need to store information
139 self._initialise_arrays(ignore_n_images, number_of_segments)
141 for segment_no in range(self.no_of_segments):
142 start = segment_no * self.len_segments
143 end = start + self.len_segments
144 seg = self.traj[ignore_n_images + start:ignore_n_images + end]
146 # If we are considering a molecular system, work out the COM for the
147 # starting structure
148 if self.is_molecule:
149 com_orig = np.zeros(3)
150 for atom_no in self.atom_indices:
151 com_orig += seg[0].positions[atom_no] / \
152 len(self.atom_indices)
154 # For each image, calculate displacement.
155 # I spent some time deciding if this should run from 0 or 1, as the displacement will be zero for
156 # t = 0, but this is a data point that needs fitting too and so
157 # should be included
158 for image_no in range(0, len(seg)):
159 # This object collects the xyz displacements for all atom
160 # species in the image
161 xyz_disp = np.zeros((self.no_of_types_of_atoms, 3))
163 # Calculating for each atom individually, grouping by species
164 # type (e.g. solid state)
165 if not self.is_molecule:
166 # For each atom, work out displacement from start coordinate
167 # and collect information with like atoms
168 for atom_no in self.atom_indices:
169 sym_index = self.types_of_atoms.index(
170 seg[image_no].symbols[atom_no])
171 xyz_disp[sym_index] += np.square(
172 seg[image_no].positions[atom_no] - seg[0].positions[atom_no])
174 # Calculating for group of atoms (molecule) and work out squared
175 # displacement
176 else:
177 com_disp = np.zeros(3)
178 for atom_no in self.atom_indices:
179 com_disp += seg[image_no].positions[atom_no] / \
180 len(self.atom_indices)
181 xyz_disp[0] += np.square(com_disp - com_orig)
183 # For each atom species or molecule, use xyz_disp to calculate
184 # the average data
185 for sym_index in range(self.no_of_types_of_atoms):
186 # Normalise by degrees of freedom and average overall atoms
187 # for each axes over entire segment
188 denominator = (2 * self.no_of_atoms[sym_index])
189 for xyz in range(3):
190 self.xyz_segment_ensemble_average[segment_no][sym_index][xyz][image_no] = (
191 xyz_disp[sym_index][xyz] / denominator)
193 # We've collected all the data for this entire segment, so now to
194 # fit the data.
195 for sym_index in range(self.no_of_types_of_atoms):
196 self.slopes[sym_index][segment_no], self.intercepts[sym_index][segment_no] = self._fit_data(self.timesteps[start:end],
197 self.xyz_segment_ensemble_average[segment_no][sym_index])
199 def _fit_data(self, x, y):
200 """
201 Private function that returns slope and intercept for linear fit to mean square diffusion data
204 Parameters:
205 x (Array of floats):
206 Linear list of timesteps in the calculation
207 y (Array of floats):
208 Mean square displacement as a function of time.
210 """
212 # Simpler implementation but disabled as fails Conda tests.
213 # from scipy.stats import linregress
214 # slope, intercept, r_value, p_value, std_err = linregress(x,y)
216 # Initialise objects
217 slopes = np.zeros(3)
218 intercepts = np.zeros(3)
220 # Convert into suitable format for lstsq
221 x_edited = np.vstack([np.array(x), np.ones(len(x))]).T
222 # Calculate slopes for x, y and z-axes
223 for xyz in range(3):
224 slopes[xyz], intercepts[xyz] = np.linalg.lstsq(
225 x_edited, y[xyz], rcond=-1)[0]
227 return slopes, intercepts
229 def get_diffusion_coefficients(self):
230 """
232 Returns diffusion coefficients for atoms (in alphabetical order) along with standard deviation.
234 All data is currently passed out in units of Å^2/<ASE time units>
235 To convert into Å^2/fs => multiply by ase.units.fs
236 To convert from Å^2/fs to cm^2/s => multiply by (10^-8)^2 / 10^-15 = 10^-1
238 """
240 slopes = [np.mean(self.slopes[sym_index])
241 for sym_index in range(self.no_of_types_of_atoms)]
242 std = [np.std(self.slopes[sym_index])
243 for sym_index in range(self.no_of_types_of_atoms)]
245 return slopes, std
247 def plot(self, ax=None, show=False):
248 """
250 Auto-plot of Diffusion Coefficient data. Provides basic framework for visualising analysis.
252 Parameters:
253 ax (Matplotlib.axes.Axes)
254 Axes object on to which plot can be created
255 show (Boolean)
256 Whether or not to show the created plot. Default: False
258 """
260 # Necessary if user hasn't supplied an axis.
261 import matplotlib.pyplot as plt
263 # Convert from ASE time units to fs (aesthetic)
264 from ase.units import fs as fs_conversion
266 if ax is None:
267 ax = plt.gca()
269 # Define some aesthetic variables
270 color_list = plt.cm.Set3(np.linspace(0, 1, self.no_of_types_of_atoms))
271 xyz_labels = ['X', 'Y', 'Z']
272 xyz_markers = ['o', 's', '^']
274 # Create an x-axis that is in a more intuitive format for the view
275 graph_timesteps = self.timesteps / fs_conversion
277 for segment_no in range(self.no_of_segments):
278 start = segment_no * self.len_segments
279 end = start + self.len_segments
280 label = None
282 for sym_index in range(self.no_of_types_of_atoms):
283 for xyz in range(3):
284 if segment_no == 0:
285 label = 'Species: {} ({})'.format(
286 self.types_of_atoms[sym_index], xyz_labels[xyz])
287 # Add scatter graph for the mean square displacement data
288 # in this segment
289 ax.scatter(graph_timesteps[start:end], self.xyz_segment_ensemble_average[segment_no][sym_index][xyz],
290 color=color_list[sym_index], marker=xyz_markers[xyz], label=label, linewidth=1, edgecolor='grey')
292 # Print the line of best fit for segment
293 line = np.mean(self.slopes[sym_index][segment_no]) * fs_conversion * \
294 graph_timesteps[start:end] + \
295 np.mean(self.intercepts[sym_index][segment_no])
296 if segment_no == 0:
297 label = 'Segment Mean : %s' % (
298 self.types_of_atoms[sym_index])
299 ax.plot(graph_timesteps[start:end], line, color='C%d' % (
300 sym_index), label=label, linestyle='--')
302 # Plot separator at end of segment
303 x_coord = graph_timesteps[end - 1]
304 ax.plot([x_coord,
305 x_coord],
306 [-0.001,
307 1.05 * np.amax(self.xyz_segment_ensemble_average)],
308 color='grey',
309 linestyle=":")
311 # Plot the overall mean (average of slopes) for each atom species
312 # 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
313 # for sym_index in range(self.no_of_types_of_atoms):
314 # line = np.mean(self.slopes[sym_index])*graph_timesteps+np.mean(self.intercepts[sym_index])
315 # label ='Mean, Total : %s'%(self.types_of_atoms[sym_index])
316 # ax.plot(graph_timesteps, line, color='C%d'%(sym_index), label=label, linestyle="-")
318 # Aesthetic parts of the plot
319 ax.set_ylim(-0.001, 1.05 * np.amax(self.xyz_segment_ensemble_average))
320 ax.legend(loc='best')
321 ax.set_xlabel('Time (fs)')
322 ax.set_ylabel(r'Mean Square Displacement ($\AA^2$)')
324 if show:
325 plt.show()
327 def print_data(self):
328 """
330 Output of statistical analysis for Diffusion Coefficient data. Provides basic framework for understanding calculation.
332 """
334 from ase.units import fs as fs_conversion
336 # Collect statistical data for diffusion coefficient over all segments
337 slopes, std = self.get_diffusion_coefficients()
339 # Useful notes for any consideration of conversion.
340 # Converting gradient from Å^2/fs to more common units of cm^2/s => multiplying by (10^-8)^2 / 10^-15 = 10^-1
341 # Converting intercept from Å^2 to more common units of cm^2 => multiply by (10^-8)^2 = 10^-16
342 #
343 # Note currently in ASE internal time units
344 # Converting into fs => divide by 1/(fs_conversion) => multiply by
345 # (fs_conversion)
347 # Print data for each atom, in each segment.
348 for sym_index in range(self.no_of_types_of_atoms):
349 print('---')
350 print(r'Species: %4s' % self.types_of_atoms[sym_index])
351 print('---')
352 for segment_no in range(self.no_of_segments):
353 print(r'Segment %3d: Diffusion Coefficient = %.10f Å^2/fs; Intercept = %.10f Å^2;' %
354 (segment_no, np.mean(self.slopes[sym_index][segment_no]) * fs_conversion, np.mean(self.intercepts[sym_index][segment_no])))
356 # Print average overall data.
357 print('---')
358 for sym_index in range(self.no_of_types_of_atoms):
359 print('Mean Diffusion Coefficient (X, Y and Z) : %s = %.10f Å^2/fs; Std. Dev. = %.10f Å^2/fs' %
360 (self.types_of_atoms[sym_index], slopes[sym_index] * fs_conversion, std[sym_index] * fs_conversion))
361 print('---')