1from numpy import (arange, array, mean, histogram, 

2 sqrt, einsum, exp, float64, ones_like, 

3 var, savez_compressed) 

4from scipy.optimize import curve_fit 

5from scipy.integrate import simps 

6from math import gamma 

7from os import listdir 

8 

9# from numexpr import evaluate 

10 

11# from sys import platform as _platform 

12# from mayavi_xyz import setup as mayavi 

13# from multiprocessing import cpu_count 

14# import subprocess 

15 

16from mango.constants import c 

17 

18from mango.io import read_data, file_write # notfinished 

19import mango.imports as imports 

20 

21from mango.pp.util import (Dataarrays, getvars, onevent, showgraph, loader) 

22from mango.pp.acf import get_suscep, save_suscep 

23from mango.pp.com import CoM_motion_rm 

24from mango.pp.plot import plot_col, plotter, lengthcheckplot 

25from mango.debug import debug 

26 

27 

28@debug(['io']) 

29def readin(fname, run, directory, block, flg, fig_event=None): 

30 """ 

31 Manipulate Data. 

32 

33 This function reads in data from file and then produces visualisation or files as required 

34 

35 Possible options from input flags: 

36 

37 * Extended XYZ file (Positions, Magnetisations, Forces) 

38 * Suceptibility (Data needs to be collected specifically for the purpose) 

39 * Data comparison plots, plot set(s) of data against time etc. 

40 * Kinetic Temperature 

41 * Chain length check 

42 

43 Parameters 

44 ---------- 

45 fname: string 

46 filename (no extension) 

47 run: int 

48 run number 

49 directory: string 

50 file directory 

51 flg: class 

52 Required flags: save_type, column, suscep, files, kinetic, lengthcheck, 

53 align, showg, saveg, save_type, ufile, 

54 

55 """ 

56 if flg.save_type == "txt": 56 ↛ 57line 56 didn't jump to line 57, because the condition on line 56 was never true

57 c.Error("F Postprocessing is not implemented for this filetype and process") 

58 

59 mpl = imports.matplotlib() 

60 

61 if not flg.ufile: 61 ↛ 67line 61 didn't jump to line 67, because the condition on line 61 was never false

62 names = [name for name in listdir(directory) 

63 if name.startswith(fname[len(directory):-2]) and not name.endswith(("svg", "dat"))] 

64 

65 reader = read_data("{}{}".format(directory, names[0]), flg.save_type) 

66 else: 

67 reader = read_data(*flg.ufile[0].rsplit(".", 1), xyzpp=flg.ufile[0].endswith('xyz')) 

68 

69 name = reader.fname 

70 

71 var, incomplete_run, tauN_V = get_ppvars(flg, reader, extras=True) 

72 

73 (timescale, datalen, xyz_data, 

74 energy_data, momenta_data, sus_data, stats) = collectdata(reader, name, flg, var, flg.eq) 

75 

76 if (flg.suscep or flg.column) and flg.showg: 76 ↛ 77line 76 didn't jump to line 77, because the condition on line 76 was never true

77 fig_event = onevent(mpl) 

78 

79 # create datafiles 

80 for i in flg.files: 

81 if flg.files[i]: 81 ↛ 82line 81 didn't jump to line 82, because the condition on line 81 was never true

82 data = locals()[f'{i}_data'] 

83 kind = i[0].upper() 

84 # paralisable? 

85 for no, i in enumerate(stats): 

86 file_write(data[no, ...], kind, timescale, directory, 

87 "{}.{}{}".format(run, i, "nofin" if incomplete_run else ""), flg, mmag=var.ms * var.vol, boxsize=var.boxsize) 

88 

89 # plot raw data graphs 

90 if flg.column: 90 ↛ 91line 90 didn't jump to line 91, because the condition on line 90 was never true

91 location = {"xyz": {"position": xyz_data[..., 0:3], 

92 "momentum": xyz_data[..., 3:6], 

93 "magnetisation": xyz_data[..., 6:9], 

94 "forces": xyz_data[..., 9:12]} if flg.files['xyz'] else None, 

95 "energy": {"Etotal": energy_data[..., 3], 

96 "mag_pot": energy_data[..., 2], 

97 "trans_pot": energy_data[..., 1], 

98 "kinetic": energy_data[..., 0]} if flg.files['energy'] else None, 

99 "momenta": {"Mtotal": momenta_data[..., 0:3], 

100 "total_mag": momenta_data[..., 3:6], 

101 "total_angular": momenta_data[..., 6:9], 

102 "CoM": momenta_data[..., 9:12], 

103 "CoM_vel": momenta_data[..., 12:15]} if flg.files['momenta'] else None, 

104 "time": timescale} 

105 

106 plot_col(mpl, fig_event, flg, location, directory, run, stats, var.temp, var.no_mol) 

107 

108 # compute the atc functions 

109 if flg.suscep and not isinstance(flg.suscep, bool): 109 ↛ 110line 109 didn't jump to line 110, because the condition on line 109 was never true

110 atc_data = get_suscep(flg, var.skip, var.dt, sus_data, block, 

111 stats, var.mass, var.vol, var.ms) 

112 

113 for bl in atc_data.keys(): 

114 for i in atc_data[bl].keys(): 

115 save_suscep(flg, directory, run, var.no_mol, bl, 

116 **atc_data[bl][i]) 

117 

118 # kinetic temperature 

119 if flg.kinetic: 119 ↛ 120line 119 didn't jump to line 120, because the condition on line 119 was never true

120 baseline, histo, err = kinetic_temperature(flg.align, name, var.mass, xyz_data[..., 3:6]) 

121 xy = [[baseline, 'units'], [histo, 'units', err]] 

122 xyname = [array([['x'], ['']]), array([['y'], ['']])] 

123 plotter(mpl, fig_event, flg, xy, xyname, 2, directory, run) 

124 

125 # Distance between pairs 

126 if flg.lengthcheck: 126 ↛ 127line 126 didn't jump to line 127, because the condition on line 126 was never true

127 ld = loader(flg, run, directory) 

128 for bl in block: 

129 try: 

130 out = ld.lengthcheck(bl) 

131 c.Error("M Loaded length data from npz file") 

132 except FileNotFoundError: 

133 out = lengthcheck(flg, run, directory, xyz_data[..., 0:3], bl) 

134 print(out['overallmean'], out['overallerr']) 

135 lengthcheckplot(mpl, fig_event, flg, directory, var.no_mol, run, timescale['X'], var.radius, **out) 

136 

137 showgraph(mpl, flg.showg) 

138 

139 

140def get_ppflags(reader): 

141 """Get post processing flags if stored in save file.""" 

142 return getvars(reader.read("flags"), 

143 ['files', 'kinetic', 'lengthcheck', 'save_type', 

144 'ufile', 'suscep', 'prog']) 

145 

146 

147def get_ppvars(flg, reader, extras=False): 

148 """Get post processing variables.""" 

149 variables = reader.read("vars") 

150 

151 var = getvars(variables, 

152 ['boxsize', 'no_molecules', 'stats', 'dt', 

153 'nmax', 'vol', 'radius', 

154 'ms', 'temp', 'mass', 'skip', 'written', 

155 'skip_iters', 'epsilon', 'sigma', 'limit']) 

156 

157 if 'tauN' in variables: 

158 tauN_V = getvars(variables, ['tauN_0', 'keff', 'temp0']) 

159 flg.neel = True 

160 else: 

161 flg.neel = False 

162 tauN_V = '' 

163 

164 var.skip_iters += 1 

165 

166 incomplete_run = False 

167 

168 if var.written != var.skip_iters: 168 ↛ 169line 168 didn't jump to line 169, because the condition on line 168 was never true

169 if var.written < var.skip_iters: 

170 c.Error(">W Incomplete run only {}/{} steps completed".format(var.written, var.skip_iters)) 

171 incomplete_run = True 

172 var.skip_iters = var.written 

173 var.nmax = var.skip_iters * var.skip 

174 

175 return var if not extras else (var, incomplete_run, tauN_V) 

176 

177 

178def collectdata(reader, name, flg, var, eq=0.1): 

179 """ 

180 Collect data from files as required. 

181 

182 Parameters 

183 ---------- 

184 reader: instance 

185 file reader 

186 name: str 

187 file name 

188 flg: instance 

189 flg storage instance 

190 requires - files kinetic lengthcheck save_type ufile suscep align prog 

191 var: instance 

192 requires - dt, no_mol, stats, skip, skip_iters, epsilon, sigma, limit, mass 

193 eq: float 

194 equilibration fraction 

195 

196 Returns 

197 ------- 

198 timescale 

199 datalen 

200 xyz_data 

201 energy_data 

202 momenta_data 

203 sus_data 

204 stats 

205 

206 """ 

207 da = Dataarrays(flg, reader, var, name, eq) 

208 

209 timescale = {'X': (arange(0, da.var.skip_iters) * var.dt * var.skip)[int(da.var.skip_iters * eq):], 

210 'F': (arange(0, da.var.skip_iters) * var.dt * var.skip)} 

211 

212 for no, (stat, name) in enumerate(da.names.items()): 

213 

214 if flg.prog: 

215 c.progress.sections = da.datalen 

216 c.progress.bar("Reading In Data", (stat + 1 / da.datalen)) 

217 else: 

218 print(name + " ...", end='', flush=True) 

219 

220 da.xyz(no, name) 

221 

222 da.energy(no, name) 

223 

224 da.momenta(no, name) 

225 

226 da.sus(no, name) 

227 

228 if not flg.prog: 

229 print("Done") 

230 

231 print("\n") 

232 CoM = CoM_motion_rm(var.mass, var.no_mol) 

233 

234 if flg.align: 

235 CoM.align(da.xyz_data, da.sus_data) 

236 elif any(anin in flg.suscep if not isinstance(flg.suscep, bool) else False for anin in ['angular', 'inertia']): 236 ↛ 237line 236 didn't jump to line 237, because the condition on line 236 was never true

237 da.sus_data.angular, da.sus_data.inertia = CoM.inertia_angular(da.xyz_data, da.sus_data) 

238 

239 return timescale, da.datalen, da.xyz_data, da.energy_data, da.momenta_data, da.sus_data, da.var.stats 

240 

241 

242def kinetic_temperature(align, basename, mass, mom, equi=0): 

243 """ 

244 Kinetic Temperature. 

245 

246 Calculate the kinetic temperature of the system and write to file. 

247 

248 Parameters 

249 ---------- 

250 TODO 

251 

252 """ 

253 # Mass rescaling 

254 dummy = einsum("ijk...,k->ijk...", mom[equi:], 1 / sqrt(mass)) 

255 

256 # Speed 

257 dummy = sqrt(einsum("ijkl,ijkl->ijk", dummy, dummy)) 

258 

259 # Histogram 

260 histo, edges = histogram(dummy, bins='fd', density=False) # dummy is flatten by default 

261 

262 baseline = 0.5 * (edges[1:] + edges[:-1]) 

263 

264 histo = histo.astype(float64) 

265 

266 err = sqrt(histo) # Poisson's estimate 

267 

268 norm = simps(histo, baseline) 

269 

270 histo /= norm 

271 

272 err /= norm 

273 

274 # The multidimensional Maxwell-Boltzmann distribution 

275 # v^(n-1)*exp(-v^2/2)/2^(n/2-1)/gamma(n/2) 

276 def func(x, amb, dim): 

277 return x**(dim - 1) * exp(-0.5 * (x / amb)**2) / (amb**dim) / (2**(0.5 * dim - 1)) / gamma(0.5 * dim) 

278 

279 dim = 3. 

280 

281 amb = (gamma(0.5 * dim) / gamma(0.5 * (dim + 1.)) / sqrt(2.)) * mean(dummy) # amb = sqrt(kB T) 

282 

283 popt, pcov = curve_fit(func, baseline, histo, p0=[amb, dim]) 

284 

285 amb = popt[0] 

286 

287 dim = popt[1] 

288 

289 temp_mb = amb**2 / c.KB 

290 

291 print("# Kinetic temperature (histogram): %12.5f [K]\n" % (temp_mb)) 

292 

293 print("# Kinetic dimension (histogram): %12.5f []\n" % (dim)) 

294 

295 if align: 

296 name = basename + "_HISTO_ALIGNED.dat" 

297 else: 

298 name = basename + "_HISTO.dat" 

299 

300 with open(name, 'w') as file: 

301 for b, h, e in zip(baseline, histo, err): 

302 file.write(" %12.5e %12.5e %12.5e %12.5e\n" % (b, h, e, func(b, amb, dim))) 

303 

304 return baseline, histo, err 

305 

306 

307def lengthcheck(flg, run, directory, pos, block): 

308 """ 

309 Length of chain check. 

310 

311 Assumes the particles furthest away from each other are particle 0 and N 

312 

313 Improvements: 

314 generalise for systems with multiple chains 

315 generalise for systems without chains eg furthest distance between particles in a cluster 

316 

317 """ 

318 # If attempted manually diff will be 8*EPS2 sometimes due to dx**2+dy**2+dz**2 

319 dummy = ones_like(pos) 

320 dist = einsum("kliz,kljz->klijz", pos, dummy) - einsum("kliz,kljz->kljiz", pos, dummy) 

321 dm = sqrt(einsum('klijz,klijz->klij', dist, dist)) 

322 

323 dmeandm = mean(dm, axis=0) 

324 vardm = var(dm, ddof=1, axis=0) / block 

325 errdm = sqrt(vardm) 

326 

327 overallmean = mean(dmeandm, axis=0) 

328 overallerr = sqrt(mean(vardm, axis=0)) 

329 name = "{}{}_distance_arrays_{}{}".format(directory, "{}Run{}".format( 

330 'S_' if flg.suscep else '', run), block, "_ALIGNED" if flg.align else '') 

331 savez_compressed(name, dmeandm=dmeandm, errdm=errdm, overallmean=overallmean, overallerr=overallerr) 

332 return {"dmeandm": dmeandm, "errdm": errdm, "overallmean": overallmean, "overallerr": overallerr} 

333 

334 

335if __name__ == "__main__": 

336 pass