Coverage for mango/pp/acf.py : 18%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from numpy import (arange, array, conjugate, where, block,
2 mean, amax, zeros_like, zeros, var, sum as nsum, linalg,
3 sqrt, savetxt, einsum, exp, log, log2, diag)
4from scipy.optimize import fminbound, curve_fit
5from scipy.interpolate import interp1d
6from math import pi
7from os import path
8from contextlib import suppress
10from mango.constants import c, nestedDict, asciistring
11from mango.pp.util import at_least2d_end # debye,
12import mango.imports as imports
14from mango.debug import debug
16# from numpy import seterr
18# seterr(all='raise')
21def get_suscep(flg,
22 skip, dt, sus_data, blocks,
23 stats, mass, vol, ms):
24 """
25 Calculate the ACF of data for different properties.
27 Parameters
28 ----------
30 Returns
31 -------
32 atc_data: dict
33 A dictionary of all the calculated ACF
35 """
36 autocorr_labels = {"mag": {"x": r'$\omega$ [$10^{{{}}}$]',
37 "y": r"MDOS [$10^{{{}}}$]",
38 "type": "MDOS"},
39 "vel": {"x": r'$\omega$ [$10^{{{}}}$]',
40 "y": r"VDOS [$10^{{{}}}$]",
41 "type": "VDOS"},
42 "angular": {"x": r'$\omega$ [$10^{{{}}}$]',
43 "y": r"ADOS [$10^{{{}}}$]",
44 "type": "ADOS"},
45 "inertia": {"x": r'$\omega$ [$10^{{{}}}$]',
46 "y": r"IDOS [$10^{{{}}}$]",
47 "type": "IDOS"},
48 "rotation": {"x": r'$\omega$ [$10^{{{}}}$]',
49 "y": r"RDOS [$10^{{{}}}$]",
50 "type": "RDOS"},
51 "mag_ifft": {"x": r"Time [$10^{{{}}}$s]",
52 "y": r"MACF [$10^{{{}}}$]",
53 "type": "MACF"},
54 "vel_ifft": {"x": r"Time [$10^{{{}}}$s]",
55 "y": r"VACF [$10^{{{}}}$]",
56 "type": "VACF"},
57 "angular_ifft": {"x": r"Time [$10^{{{}}}$s]",
58 "y": r"AACF [$10^{{{}}}$]",
59 "type": "AACF"},
60 "inertia_ifft": {"x": r"Time [$10^{{{}}}$s]",
61 "y": r"IACF [$10^{{{}}}$]",
62 "type": "IACF"},
63 "rotation_ifft": {"x": r"Time [$10^{{{}}}$s]",
64 "y": r"RACF [$10^{{{}}}$]",
65 "type": "RACF"}}
67 if 'vel' in flg.suscep:
68 sus_data["vel"] = einsum("hijk, j -> hijk", sus_data["mom"], 1 / mass)
70 atc_data = nestedDict()
72 autocorr = Autocorr(flg, sus_data, dt, skip, ms * vol)
74 autocorr.conv(len(stats))
76 for bl in blocks:
78 autocorr.analysis(blocks=bl)
80 for i in flg.suscep:
82 atc_data[bl][i] = {"xaxis": autocorr.freq, "yaxis": autocorr.ret[i]['data'], "stddev_y": autocorr.ret[i]['error'],
83 # "popt": tau, "popt_fitted": ad_tau.popt_fit,
84 "autocorr_labels": autocorr_labels[i], 'scale': autocorr.ret[i]['scale']}
86 iatc = i + '_ifft'
87 atc_data[bl][iatc] = {"xaxis": autocorr.time, "yaxis": autocorr.ret[i]['ifft'], "stddev_y": autocorr.ret[i]['ifft_error'],
88 # "popt": tau, "popt_fitted": ad_tau.popt_fit,
89 "autocorr_labels": autocorr_labels[iatc], 'scale': autocorr.ret[i]['scale']}
91 return atc_data
94class Autocorr():
95 """Calculates the ACF and MDOS of provided data."""
97 @debug(['acf'])
98 def __init__(self, flg, sus_data, dt, skip, mmag):
99 """
100 Store data to be processed.
102 Parameters
103 ----------
104 flg: instance
105 flags class instance
106 sus_data: dict
107 data dictionary
108 dt: float
109 timestep
110 skip: int
111 number of iterations skipped
112 mmag: float
113 Mmag
115 """
116 self.flg = flg
117 self.dt = dt
118 self.skip = skip
119 self.mmag = mmag
121 self.fft, self.ffts, self.ifft, self.byte_align = imports.fftw()
122 self.shape = sus_data[list(sus_data.keys())[0]].shape
124 self.ret = {}
126 self.sus_collect = {}
128 self._shape_data(sus_data)
130 if hasattr(self, 'DBG'):
131 self.sus_data = sus_data.copy()
133 def _calc(self, func, *args):
134 # Paralisable loop c.processors/numkeys
135 for autocorr_type in self.sus_collect.keys():
136 func(autocorr_type, *args)
138 if hasattr(self, 'DBG'):
139 from mango.pp.acf_checker import checker, autoold
140 ret2, freq2, time2 = autoold(c.Error, self.sus_data, self.dt, self.skip, self.mmag, blocks=self.blocks)
141 checker(self.ret, ret2, self.freq, freq2, self.time, time2)
143 def conv(self, no_stats, RMS=True):
144 """Carry out blocking convergence tests."""
145 self._calc(self._blockingconvergence, no_stats, RMS)
147 def analysis(self, blocks):
148 """
149 Carry out DOS and ACF calculation for provided data and blocks.
151 Parameters
152 ----------
153 blocks: int
154 number of blocks
156 """
157 self.sample = self.shape[0] * self.shape[1] // blocks
158 self.sample_pad = self.ffts.next_fast_len(2 * self.sample - 1)
159 self.freq = 2.0 * pi * self.ffts.fftshift(self.ffts.fftfreq(self.sample_pad, d=(self.dt * self.skip)))
160 self.norm_val = self.freq[1] - self.freq[0]
161 # arange(-self.sample + 1, self.sample) * self.dt * self.skip
162 self.time = 2. * pi * self.ffts.fftshift(self.ffts.fftfreq(self.sample_pad, d=self.norm_val))
164 self.blocks = blocks
165 self._calc(self._blockinganalysis)
167 def _shape_data(self, sus_data):
168 """
169 Shape arrays dependent on type.
171 Parameters
172 ----------
173 sus_data: dict
174 dict of data arrays
176 """
177 reshape = []
179 for autocorr_type in self.flg.suscep:
180 # if autocorr_type in ['chi']:
181 # continue
182 # else:
183 value = sus_data[autocorr_type]
184 if autocorr_type in ['mag']:
185 # Sum over the number of atoms
186 # Flatten and refold -> Blocking will only trim from
187 # first trajetory for correct dimensions
188 # data =
189 # stddev = self._get_av_err(data, av=False) # - (stddev if self.flg.align else 0)
190 # - (mean(einsum("il, il ->il", data[:, 0], data[:, 0]), axis=1) if self.flg.align else 0
191 data = self._remove_av(einsum('ijkl->ijl', value)).reshape(-1, 3)
192 elif autocorr_type in ["rotation"]:
193 if self.flg.align:
194 data = self._remove_av(einsum('ijkl, il -> ijk', value,
195 mean(einsum("ijkz ->ijz", sus_data["mag"]), axis=1))).reshape(-1, 3)
196 else:
197 c.Error("W rotation autocorrelation only works with align=True.")
198 else:
199 data = self._remove_av(value).reshape([-1, *value.shape[2:]])
201 if autocorr_type in ['inertia', "vel"]:
202 reshape = [-1, 3]
203 else:
204 reshape = [3]
206 print(autocorr_type, data.shape)
207 e_string = "{0}, {0} -> ".format(''.join(asciistring(data.ndim)))
209 self.sus_collect[autocorr_type] = {'data': data, 'scale': 2 * pi * einsum(e_string, data, data) / data.size,
210 'reshapeto': reshape}
212 @staticmethod
213 def _remove_av(value):
214 """
215 Remove the average over the iterations before any modifcation.
217 Trajectory may not be fully relaxed
218 Removing the mean after blocking leads to spurious errors
220 """
221 return value - mean(value, axis=1)[:, None]
223 def _fft_av_err(self, dummy, scale, blocks):
224 """Scale fft errors."""
225 ave, err = self._get_av_err(dummy, blocks)
226 print(ave[0])
227 return scale * self.ffts.fftshift(ave), scale * self.ffts.fftshift(err)
229 @staticmethod
230 def _get_av_err(dummy, blocks=1, axis=0, av=True):
231 """
232 Average data over blocks and calculate errors.
234 Parameters
235 ----------
236 dummy: array
237 data array
238 blocks: int
239 number of blocks
241 Returns
242 -------
243 ave: array
244 averaged data array
245 err: array
246 standard deviation of averaged data
248 """
249 if dummy.shape[0] == 1:
250 err = zeros(dummy.shape[1:])
251 else:
252 err = sqrt(var(dummy, ddof=1, axis=axis) / blocks)
254 return (mean(dummy, axis=axis), err) if av else err
256 def _blockingconvergence(self, autocorr_type, no_stats, RMS=True):
257 """
258 Convergence test for blocking.
260 Parameters
261 ----------
262 autocorr_type: str
263 autocorr name
264 RMS: bool
265 Root mean squared or mean
267 """
268 print('# Block averages {}'.format('magnetic moment [emu]' if autocorr_type == 'mag' else autocorr_type))
269 print('# Block, sample, average, error')
270 data = self.sus_collect[autocorr_type]['data']
272 if RMS:
273 def RMSwrap(data, **kw):
274 e_string = "{0}, {0} -> ab".format(''.join(asciistring(data.ndim)))
275 return sqrt(mean(einsum(e_string, data, data), **kw))
276 else:
277 def RMSwrap(data, **kw):
278 e_string = "{} -> ab".format(''.join(asciistring(data.ndim)))
279 return mean(einsum(e_string, data), **kw)
281 for local_sample in (no_stats * 2**x for x in range(int(log2(self.shape[0] * self.shape[1]) - log2(self.shape[0])))):
282 no_blocks = self.shape[0] * self.shape[1] // local_sample
283 blocked = data[:local_sample * no_blocks].reshape([no_blocks, local_sample, *self.sus_collect[autocorr_type]['reshapeto']])
284 reblock = RMSwrap(blocked, axis=1)
285 ave, err = self._get_av_err(reblock, no_blocks)
286 print(local_sample, no_blocks, ave, err) # if local_sample > 1 else 0)
288 @staticmethod
289 def _get_estring(ndim):
290 """
291 Return the 3 character strings for einsum of DOS.
293 eg 3 dimensions would return:
295 dos = 'abc'
296 cdos = 'abd'
297 fstr = 'abcd'
299 Parameters
300 ----------
301 ndim: int
302 number of dimensions of array
304 Returns
305 -------
306 dos_str: str
307 string of char length ndim
308 cdos_str: str
309 string of char length ndim
310 fstr: str
311 string of char length of ndim + 1
313 """
314 chr_list = asciistring(ndim + 1)
315 dos_str = ''.join(chr_list[:-1])
316 cdos_str = ''.join(chr_list[:-2] + [chr_list[-1]])
317 return dos_str, cdos_str, ''.join(chr_list)
319 @staticmethod
320 def _pmsum(data):
321 # TODO per particle correleation
322 # if shape[1] == no_mol and shape[0] == iter (?)
323 # if data.ndim == 3:
324 # return einsum("ijk -> ik", data)
325 return data
327 def _blockinganalysis(self, autocorr_type):
328 """
329 DOS and ACF calculation.
331 creates storage dictionary 'ret'
333 Parameters
334 ----------
335 autocorr_type: str
336 autocorr name
338 """
339 scale = self.sus_collect[autocorr_type]['scale']
340 reshapeto = [self.blocks, self.sample, *self.sus_collect[autocorr_type]['reshapeto']]
341 # Scale
342 # if autocorr_type is 'mag': # 2*pi*<M*M>/3
343 # # scale = 2. * pi * einsum("i, i -> ", self.mmag, self.mmag) / (3. * self.mmag.size)
344 # scale = einsum("ijk, ijk -> ", mxyz, mxyz) / mxyz.size # This is equal to <M*M>/3 for independent particles
345 # else:
346 # 2*pi*<v*v>/3, 2*pi*<M*M>/3
348 # Blocking
349 print(autocorr_type, self.sus_collect[autocorr_type]['data'].shape)
351 mxyz = self.sus_collect[autocorr_type]['data'][:self.sample * self.blocks].reshape(reshapeto)
353 self.ret[autocorr_type] = {'scale': scale}
355 print('DOS')
356 # DOS
357 dos = self.fft(self.byte_align(mxyz), n=self.sample_pad, axis=1, threads=c.processors)
359 print("DOne")
360 # modulus of the data gives separate real and imaginary parts
361 dos = self._pmsum(dos)
363 dos_str, cdos_str, f_str = self._get_estring(dos.ndim)
365 mod_dos = einsum("{}, {} -> {}".format(dos_str, cdos_str, f_str), dos, conjugate(dos))
366 # Calculate eigenvalues
367 eig_dos = linalg.eigvalsh(mod_dos)
369 # Normalisation: Integral of DOS equal 1
370 norm = self.norm_val * einsum('{} -> '.format(dos_str), eig_dos).real / self.blocks
371 eig_dos /= norm
373 print('IDOS')
374 # Normalisation: ACF(0) = 1
375 ifftx = 2 * pi * self.ifft(self.byte_align(eig_dos), axis=1, threads=c.processors) / (self.dt * self.skip)
376 print("DOne")
378 # Error and averaging
379 (self.ret[autocorr_type]['data'],
380 self.ret[autocorr_type]['error']) = self._fft_av_err(eig_dos, scale, self.blocks)
381 (self.ret[autocorr_type]['ifft'],
382 self.ret[autocorr_type]['ifft_error']) = self._fft_av_err(ifftx, scale, self.blocks)
385@debug(['save'])
386def save_suscep(flg, directory, run, no_mol, cblock, xaxis, yaxis, stddev_y,
387 autocorr_labels, scale, line_d={}, **kwargs):
388 """
389 Save the susceptibility calculation.
391 Saves autocorrelation data in raw text,
392 The columns are shown below with yaxis and std_y will always be 3 columns:
394 xaxis yaxis... std_y...
396 Parameters
397 ----------
398 directory: str
399 save location
400 run: int
401 run number
402 no_mol: int
403 number of particles
404 cblock: int
405 current block length
406 xaxis: ndarray
407 x data
408 yaxis: ndarray
409 y data
410 stddev_y: ndarray
411 y error data
412 autocorr_labels: dict
413 dictionary of x and y labels
414 scale: float
415 scale value for data
417 """
419 # if (path.isfile(figname) or path.isfile("{}{}.dat".format(figname, "ALIGN" if flg.align else ''))) and flg.saveg:
420 # c.Error('W File {} exists, skipping'.format(figname.rsplit('/', 1)[-1]))
421 # return None
423 xaxis, yaxis, stddev_y = at_least2d_end(xaxis, yaxis, stddev_y)
425 if ('DOS' in autocorr_labels["type"]) or ('ACF' in autocorr_labels["type"]):
426 header = f'scale factor: {scale}' # Remove 2 pi which is from Fourier transform definintion
428 data = block([xaxis, yaxis, stddev_y])
429 print(data.shape, xaxis.shape, yaxis.shape, stddev_y.shape)
430 name = "{}S_Run{}_mol-{:g}_autocorr_{}_blk{}".format(directory, run, no_mol, autocorr_labels["type"], cblock)
431 savetxt("{}{}.dat".format(name, "ALIGN" if flg.align else ''), data, header=header)