1from numpy import (empty, sqrt, mean, array, einsum, 

2 exp, floor, log10, absolute, load) 

3from math import pi 

4from os import path 

5from contextlib import suppress 

6from contextlib import contextmanager 

7from io import StringIO 

8import sys 

9 

10from mango.constants import c, _variables 

11from mango.io import _input 

12from mango.imports import inquirer 

13from mango.arguments import column_manip 

14from mango.pp.energy import energy_calc 

15 

16 

17@contextmanager 

18def captured_output(): 

19 nout, nerr = StringIO(), StringIO() 

20 oout, oerr = sys.stdout, sys.stderr 

21 try: 

22 sys.stdout, sys.stderr = nout, nerr 

23 yield sys.stdout, sys.stderr 

24 finally: 

25 sys.stdout, sys.stderr = oout, oerr 

26 

27 

28def getvars(varib, lst): 

29 rlst = [] 

30 for x in lst: 

31 try: 

32 if x in ['nmax', 'written', 'skip_iters']: 

33 rlst.append(int(varib[x])) 

34 else: 

35 rlst.append(varib[x]) 

36 except KeyError: 

37 rlst.append(None) 

38 with suppress(ValueError): 

39 lst[lst.index('no_molecules')] = 'no_mol' 

40 return _variables(**dict(zip(lst, rlst))) 

41 

42 

43def get_expon(number): 

44 return (floor(log10(abs(number)))).astype(int).astype(float) 

45 

46 

47def at_least2d_end(*args): 

48 lst = [] 

49 for i in args: 

50 if len(i.shape) < 2: 

51 lst.append(i[:, None]) 

52 else: 

53 lst.append(i) 

54 if len(lst) == 1: 

55 lst = lst[0] 

56 return lst 

57 

58 

59def interactive_getfiles(flg, check_name, files): 

60 """ 

61 Get postprocessing jobs if none selected, Check file exists. 

62 

63 Parameters 

64 ---------- 

65 Error: func 

66 flg: instance 

67 instance of required flags 

68 check_name: str 

69 name of file to read in 

70 files: list 

71 list of files in selected directory 

72 

73 """ 

74 if (not (flg.column or flg.lengthcheck or 

75 flg.files['xyz'] or flg.files['energy'] or flg.files['momenta'] or flg.kinetic) and 

76 ((flg.suscep and isinstance(flg.suscep, bool) or not flg.suscep))): 

77 

78 inq, (Checkbox, ConsoleRender) = inquirer() 

79 

80 if inq and c.tinfo['otty']: 

81 

82 c.Error(">G Nothing to do please select:\n ") 

83 

84 choices = {'XYZ File': "xyz", 'Momenta File': "momenta", 'Energy File': "energy", 

85 'Raw Data plotting': "column", 'Susceptibility': "suscep"} 

86 

87 question = Checkbox('pp', message="{}{}".format("Which postprocessing tool do you want?", 

88 " (right arrow to select)"), 

89 choices=choices.keys()) 

90 

91 selections = set(choices.keys()).intersection(ConsoleRender().render(question, {})) 

92 

93 if selections == set(): 

94 print('') 

95 c.Error("F No postprocessing tool selected" + 30 * ' ') 

96 

97 for i in selections: 

98 if i == "Raw Data plotting": 

99 flg.column, flg.average, flg.files = column_manip([_input("Enter x string: "), 

100 _input("Enter y string: ")], 

101 flg.average, flg.files) 

102 elif i == "Susceptibility": 

103 flg.suscep = _input("Graphs [chi mag vel]?: ").split() 

104 else: 

105 # av = _input("Averaged {} [y/n]?: ".format(i)) 

106 # flg.average[choices[i][0]] = (av == "y") 

107 flg.files[choices[i]] = True 

108 

109 

110 else: 

111 c.Error("F No postprocessing tool selected") 

112 

113 while not any(file.startswith(check_name) for file in files) and c.tinfo['otty']: 

114 

115 c.Error(">G Cannot find {}run".format("Susceptibility " if flg.suscep else "")) 

116 

117 run_number = _input("\nRun number or rerun [y/n]?: ") 

118 

119 if run_number == "y" or run_number == "yes": 

120 flg.run = False 

121 break 

122 elif run_number == "n" or run_number == "no": 

123 c.Error(">ME Exiting\n\n") 

124 try: 

125 flg.run = int(run_number) 

126 check_name = '{}Run{:g}_mol-'.format("S_" if flg.suscep else "", flg.run) 

127 except ValueError: 

128 c.Error("F Not a number, Exiting") 

129 

130 if not c.tinfo['otty'] and not any(file.startswith(check_name) for file in files): 

131 c.Error("F Run not found") 

132 

133 return flg, check_name 

134 

135 

136def getlocation(keys, dictionary, stats, kd={}): 

137 """ 

138 Apply any basic arithmetic to keys with newly found location. 

139 

140 Arrays can be split eg array[:,1] etc 

141 """ 

142 for no, xy in enumerate(keys): 

143 ustr = string = '' 

144 op = {} 

145 units = {} 

146 std_var = {} 

147 sumkey = {} 

148 ar = {} 

149 

150 for k, key in enumerate(xy): 

151 if key[1] == 'False': 

152 op[k] = key[0] 

153 string += " {} ".format(op[k]) 

154 else: 

155 sumkey[k], units[k] = dictfind(key[0], dictionary) 

156 string += "sumkey[{}]".format(k) 

157 if key[1] != 'True': 

158 ar[k] = key[1] 

159 string += ar[k] 

160 

161 add = {o: p for o, p in op.items() if p in ["*", "/", "-", "+"]} 

162 

163 if len(units.keys()) > 1: 

164 for num in range(k + 1): 

165 if num in units: 

166 ustr += units[num] 

167 if num in add: 

168 ustr += add[num] 

169 else: 

170 ustr += units[0] 

171 

172 # If numexpr becomes a req this is safer and quicker 

173 # 

174 # string = string.split() 

175 # new_str = "" 

176 # for i, j in enumerate(string): 

177 # if j in ["*", "/", "-", "+"]: 

178 # new_str += j 

179 # else: 

180 # exec('var_' + str(i) + ' = ' + j) 

181 # new_str += "var_{}".format(i) 

182 # evaled = evaluate(new_str) 

183 evaled = eval(string) 

184 kd[no] = [evaled, ustr] 

185 

186 return kd[0], kd[1] 

187 

188 

189def dictfind(key, dictionary): 

190 """ 

191 Find key in nested dictionary, Returns as soon as found. 

192 

193 Parameters 

194 ---------- 

195 key: str 

196 key to find 

197 dictionary: dict 

198 2 level dictionary of all data 

199 

200 Returns 

201 ------- 

202 d: ndarray 

203 data array 

204 u: str 

205 units 

206 std_var: ndarray/bool 

207 standard deviation array, False if data not averaged 

208 

209 """ 

210 av = ["x", "e", "m"] 

211 

212 if key == "time": 

213 return dictionary["time"], c.units["time"] 

214 

215 for no, i in enumerate(["xyz", "energy", "momenta"]): 215 ↛ exitline 215 didn't return from function 'dictfind', because the loop on line 215 didn't complete

216 with suppress(TypeError, KeyError): 

217 d = dictionary[i][key] 

218 u = c.units[i][key] 

219 return d, u 

220 

221 

222class loader(): 

223 """Load postprocessed data files.""" 

224 

225 def __init__(self, flg, run, directory): 

226 self.flg = flg 

227 self.run = run 

228 self.directory = directory 

229 

230 def lengthcheck(self, block): 

231 """Load numpy length check array.""" 

232 name = "{}{}_distance_arrays_{}{}.npz".format(self.directory, "{}Run{}".format( 

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

234 return load(name) 

235 

236 

237class Dataarrays(): 

238 """Data array collection.""" 

239 

240 def __init__(self, flg, reader, var, name, eq=0.1): 

241 """ 

242 flg - files kinetic lengthcheck save_type ufile suscep 

243 var - stats no_mol skip_iters epsilon sigma limit mass 

244 """ 

245 var.stats.sort() 

246 # stats = [0, 1] 

247 

248 self.flg = flg 

249 self.var = var 

250 self.reader = reader 

251 self.name = name 

252 

253 print("Checking Files...", end='', flush=True) 

254 self.datalen, self.names = self._showmissing() 

255 self.length() 

256 print("Done") 

257 

258 self.xyz_data = self.momenta_data = self.energy_data = None 

259 self.angular = self.pos_com = None 

260 

261 self.eq = int(self.var.skip_iters * eq) 

262 skip_iters = int(self.var.skip_iters - self.eq) 

263 

264 if flg.files['xyz'] or flg.files['energy'] or flg.kinetic or flg.lengthcheck: 

265 self.xyz_data = empty((self.datalen, skip_iters, self.var.no_mol, 12)) 

266 if flg.files['energy']: 

267 self.energy_data = empty((self.datalen, self.var.skip_iters, 4)) 

268 if flg.files['momenta']: 

269 self.momenta_data = empty((self.datalen, self.var.skip_iters, 15)) 

270 self.sus_data, self.sus_list, self.slocs = self._sus_em(skip_iters, self.var.no_mol) 

271 

272 def _showmissing(self, fmiss='', names={}): 

273 """Show missing statistics from list.""" 

274 stat_copy = [] 

275 for stat in self.var.stats: 

276 

277 self.name = "{}.{:g}.{}".format(self.name.rsplit(".", maxsplit=2)[0], stat + 1, 

278 self.flg.save_type if not self.flg.ufile else self.flg.ufile[0].split('.')[-1]) 

279 

280 if path.isfile(self.name): 

281 names[stat] = self.name 

282 stat_copy.append(stat) 

283 

284 self.var.stats = stat_copy 

285 

286 rsts = list(set(self.var.stats).difference(names.keys())) 

287 

288 if len(rsts) > 0: 

289 

290 splt = [0] + [i for i in range(1, len(rsts)) if rsts[i] - rsts[i - 1] > 1] + [None] 

291 missing = [rsts[b:e] for (b, e) in [(splt[i - 1], splt[i]) for i in range(1, len(splt))]] 

292 for m in missing: 

293 if len(m) > 1: 

294 fmiss += "{} to {}, ".format(m[0] + 1, m[-1] + 1) 

295 else: 

296 fmiss += "{}, ".format(m[0] + 1) 

297 

298 c.Error(">W Can't find the following repetitions: {}".format(fmiss[:-2])) 

299 

300 return len(list(names.keys())), names 

301 

302 def length(self): 

303 lstore = [] 

304 self.pstore = {} 

305 for n in self.names.values(): 

306 p, ls = self.reader.read(fname=n, lengthcheck=True) 

307 self.pstore[n] = p 

308 lstore.append(ls) 

309 

310 small = min(lstore) + 1 

311 if small < self.var.skip_iters: 

312 c.Error(f">W All runs not the same length shortest simulation is {small}") 

313 self.var.skip_iters = small 

314 

315 def _sus_em(self, skip_iters, no_mol, sus_list=[], sus_data=None, slocs=None): 

316 """Create empty arrays as needed for susceptibility data.""" 

317 if self.flg.suscep and not isinstance(self.flg.suscep, bool): 

318 sus_data = {} 

319 slocs = {"pos": [(0, 3), 'position'], 

320 "mom": [(3, 6), 'momentum'], 

321 "mag": [(6, 9), 'magnetisation']} 

322 

323 sus_list = list(slocs.keys()) 

324 sus_list += ['angular', 'inertia'] 

325 

326 for i in sus_list: 

327 sus_data[i] = empty((self.datalen, skip_iters, no_mol, 3)) 

328 

329 return sus_data, sus_list, slocs 

330 

331 def _em_vars(self, epsilon=None, sigma=None, limit=None, mass=None): 

332 for i, j in zip(['epsilon', 'sigma', 'limit', 'mass'], [epsilon, sigma, limit, mass]): 

333 if j is not None: 

334 setattr(self.var, i, j) 

335 

336 def energy(self, stat, name, epsilon=None, sigma=None, limit=None, mass=None): 

337 """Get total energy data.""" 

338 # TODO expose more energy data to user 

339 self._em_vars(epsilon, sigma, limit, mass) 

340 if self.flg.files['energy']: 

341 self.energy_sv = energy_calc(self.pstore[name][:self.var.skip_iters], 

342 self.reader.read("momentum", name)[:self.var.skip_iters], 

343 self.reader.read("magnetisation", name)[:self.var.skip_iters], 

344 self.var.epsilon, self.var.sigma, self.var.limit, self.var.mass) 

345 self.energy_data[stat, :, 0] = self.energy_sv['kinetic'] 

346 self.energy_data[stat, :, 1] = self.energy_sv['trans_pot'] 

347 self.energy_data[stat, :, 2] = self.energy_sv['mag_pot'] 

348 self.energy_data[stat, :, 3] = self.energy_sv['total_E'] 

349 

350 def momenta(self, stat, name, epsilon=None, sigma=None, limit=None, mass=None): 

351 """Get momenta data.""" 

352 self._em_vars(epsilon, sigma, limit, mass) 

353 

354 if self.flg.files['momenta']: 

355 if hasattr(self, 'energy_sv'): 

356 momenta = self.energy_sv 

357 else: 

358 momenta = energy_calc(self.pstore[name][:self.var.skip_iters], 

359 self.reader.read("momentum", name)[:self.var.skip_iters], 

360 self.reader.read("magnetisation", name)[:self.var.skip_iters], 

361 self.var.epsilon, self.var.sigma, self.var.limit, self.var.mass) 

362 self.momenta_data[stat, :, 0:3] = momenta['total_M'] 

363 self.momenta_data[stat, :, 3:6] = momenta['total_mag'] 

364 self.momenta_data[stat, :, 6:9] = momenta['total_angular'] 

365 self.momenta_data[stat, :, 9:12] = momenta['CoM'] 

366 self.momenta_data[stat, :, 12:15] = momenta['CoM_vel'] 

367 

368 def sus(self, stat, name): 

369 """Get susceptibility data.""" 

370 if self.flg.suscep and not isinstance(self.flg.suscep, bool): 

371 for i in self.sus_list: 

372 if i in ['angular', 'inertia']: 

373 pass 

374 elif not (self.flg.files['xyz'] or self.flg.files['energy'] or self.flg.kinetic or self.flg.lengthcheck): 

375 if self.slocs[i][1] == "position": 

376 self.sus_data[i][stat, ...] = self.pstore[name][self.eq:self.var.skip_iters] 

377 else: 

378 self.sus_data[i][stat, ...] = self.reader.read(self.slocs[i][1], name)[self.eq:self.var.skip_iters] 

379 else: 

380 self.sus_data[i][stat, ...] = self.xyz_data[stat, :, :, self.slocs[i][0][0]:self.slocs[i][0][1]] 

381 

382 def xyz(self, stat, name): 

383 """Get xyz data.""" 

384 if self.flg.files['xyz'] or self.flg.kinetic or self.flg.lengthcheck: 

385 self.xyz_data[stat, :, :, 0:3] = self.pstore[name][self.eq:self.var.skip_iters] 

386 self.xyz_data[stat, :, :, 3:6] = self.reader.read("momentum", name)[self.eq:self.var.skip_iters] 

387 self.xyz_data[stat, :, :, 6:9] = self.reader.read("magnetisation", name)[self.eq:self.var.skip_iters] 

388 self.xyz_data[stat, :, :, 9:12] = self.reader.read("forces", name)[self.eq:self.var.skip_iters] 

389 

390 

391def showgraph(mpl, showg=False): 

392 """Show graphs onscreen at end of run.""" 

393 if showg: 393 ↛ exitline 393 didn't return from function 'showgraph', because the condition on line 393 was never false

394 show, get_fignums = mpl.showg() 

395 if get_fignums(): 

396 if not c.using_IP: 396 ↛ 400line 396 didn't jump to line 400, because the condition on line 396 was never false

397 c.Error("{}{}{}".format("M Press Enter to close all graphs,\n", 

398 "Space to close current graph\n", 

399 "Click on legend to show/hide lines")) 

400 with suppress(AttributeError): 

401 show() 

402 else: 

403 c.Error("M >No Graph to Show") 

404 

405 

406class onevent(): 

407 

408 def __init__(self, mpl): 

409 self.close = mpl.close() 

410 self.figcall = {} 

411 from matplotlib.lines import Line2D 

412 self.l2d = Line2D 

413 

414 def connect(self, fig, leg=None, line=None): 

415 """ 

416 Connect figures to events. 

417 

418 Parameters 

419 ---------- 

420 fig: figure 

421 leg: legend 

422 line: line 

423 

424 """ 

425 if fig.number not in self.figcall: 

426 self.figcall[fig.number] = {"key": fig.canvas.mpl_connect('key_press_event', self.onkey), 

427 "pick": fig.canvas.mpl_connect('pick_event', self.onpick), 

428 "fig": fig} 

429 

430 self.figcall[fig.number]['scales'] = self.fig_d(fig) 

431 

432 if leg is not None: 

433 if "ld" in self.figcall[fig.number]: 

434 self.figcall[fig.number]['ld'].update(self.line_d(leg.legendHandles, line)) 

435 else: 

436 self.figcall[fig.number]['ld'] = self.line_d(leg.legendHandles, line) 

437 

438 @staticmethod 

439 def fig_d(fig): 

440 """ 

441 Set picker for yaxis. 

442 

443 Parameters 

444 ---------- 

445 fig: fig instance 

446 

447 """ 

448 scales = {} 

449 for ax in fig.axes: 

450 ax.get_yaxis().set_picker(100) 

451 scales[ax] = {ax.get_yscale(): ax.get_ylim()} 

452 return scales 

453 

454 @staticmethod 

455 def line_d(legh, line): 

456 """ 

457 Set picker for legend. 

458 

459 Parameters 

460 ---------- 

461 legh: legendHandles 

462 

463 line: line object 

464 

465 Returns 

466 ------- 

467 dict: 

468 a dictionary linking legend entry and line 

469 

470 """ 

471 lst = {} 

472 for l, o in zip(legh, line): 

473 l.set_picker(5) 

474 lst[l] = o 

475 return lst 

476 

477 # def setlim(self, xmin, xmax, ymin, ymax): 

478 # self.figcall[self.cn]['xm'] = xmin 

479 # self.figcall[self.cn]['ym'] = ymin 

480 # self.figcall[self.cn]['xM'] = xmax 

481 # self.figcall[self.cn]['yM'] = ymax 

482 

483 def onpick(self, event): 

484 """Switch lines visibility.""" 

485 fig_dict = self.figcall[event.canvas.figure.number] 

486 lline = event.artist 

487 if lline in fig_dict['ld']: 

488 oline = fig_dict['ld'][lline] 

489 if isinstance(oline, self.l2d): 

490 vis = not oline.get_visible() 

491 oline.set_visible(vis) 

492 else: 

493 for o in oline: 

494 vis = not o.get_visible() 

495 o.set_visible(vis) 

496 

497 lline.set_alpha(1.0) if vis else lline.set_alpha(0.2) 

498 

499 else: 

500 ax = event.artist.axes 

501 scale = ax.get_yscale() 

502 if scale not in fig_dict['scales'][ax]: 

503 fig_dict['scales'][ax][scale] = ax.get_ylim() 

504 ax.set_yscale("linear" if scale == 'log' else 'log') 

505 with suppress(KeyError): 

506 ax.set_ylim(fig_dict['scales'][ax]["linear" if scale == 'log' else 'log']) 

507 fig_dict['fig'].canvas.draw() 

508 

509 def onkey(self, event): 

510 """On key events.""" 

511 if event.key == "enter": 

512 self.close('all') 

513 elif event.key == " ": 

514 self.close(event.canvas.figure.number)