1""" 

2Magnetic motion setup. 

3 

4Set up environment and start calculation loop 

5""" 

6# External Dependencies 

7from numpy import (array, zeros, zeros_like, 

8 sin, cos, ones, maximum, 

9 sqrt, errstate, where, einsum) 

10 

11from numpy.core.multiarray import c_einsum 

12 

13from copy import deepcopy 

14from scipy import linalg 

15from math import pi 

16from contextlib import suppress 

17 

18# Internal Dependencies 

19from mango.constants import c, nestedDict 

20from mango.position import position 

21from mango.io import write_data 

22from mango.initpositions import initpos 

23from mango.time import _time, end, grace 

24from mango.multiproc import mp_handle 

25from mango.debug import debug, profile 

26from mango.managers import addpid 

27 

28 

29def calc_wrap(calc, stat): 

30 """ 

31 Wrap calculate to allow decoration in multiprocessing. 

32 

33 Could be used in future for any function 

34 """ 

35 with suppress(KeyboardInterrupt): 

36 return profile(calc.run, stat, file="./wprofiler.prof") if c.profile and calc.flg.parallel else calc.run(stat) 

37 

38 

39def field(time, H0=167.0, nu=267.0e3): 

40 """Calculate the phase of the external magnetic field.""" 

41 if nu > c.EPS: 

42 H = H0 * sin(2.0 * pi * nu * time) 

43 

44 else: 

45 H = H0 

46 

47 return H 

48 

49 

50def num_iters(iters, skip): 

51 """Get number of iterations including first and last iterations.""" 

52 return 1 + ((iters - 1) // skip) 

53 

54 

55def sanity_checks(var, pos, mag, xyz): 

56 """Sanity check input.""" 

57 new_nm = pos.shape[0] 

58 

59 new_size = zeros(new_nm) 

60 

61 if "no_molecules" not in var.defaults and var.no_molecules != new_nm: 61 ↛ 62line 61 didn't jump to line 62, because the condition on line 61 was never true

62 c.Error("F Number of molecules is not consistant with input file") 

63 elif var.no_molecules != new_nm: 63 ↛ 64line 63 didn't jump to line 64, because the condition on line 63 was never true

64 var.no_molecules = new_nm 

65 name = list(var.name) 

66 name[-2] = str(var.no_molecules) 

67 var.name = "".join(name) 

68 

69 if "radius_1" not in var.defaults and var.radius.shape[0] != new_nm: 69 ↛ 70line 69 didn't jump to line 70, because the condition on line 69 was never true

70 c.Error("F Number of molecule radii is not consistant with input file") 

71 elif "radius_1" in var.defaults: 71 ↛ 74line 71 didn't jump to line 74, because the condition on line 71 was never false

72 var.radius = var.radius[0] + new_size 

73 

74 if "dens_1" not in var.defaults and var.dens.shape[0] != new_nm: 74 ↛ 75line 74 didn't jump to line 75, because the condition on line 74 was never true

75 c.Error("F Number of molecule densities is not consistant with input file") 

76 elif "dens_1" in var.defaults: 76 ↛ 80line 76 didn't jump to line 80, because the condition on line 76 was never false

77 var.dens = var.dens[0] + new_size 

78 var.ms = var.ms[0] + new_size 

79 

80 if mag is not None: 80 ↛ 81line 80 didn't jump to line 81, because the condition on line 80 was never true

81 abs_mag = linalg.norm(mag, axis=1) 

82 if (abs_mag > 1 + c.EPS).any() ^ (not xyz and (abs_mag > (var.ms * var.vol) + c.EPS).any()): 

83 c.Error("F The magnetisation should be less than the saturation magnetisation") 

84 

85 

86def new_vars(var): 

87 """Calculate New variables from input.""" 

88 var.hradius = var.radius 

89 var.vol = (4.0 / 3.0) * pi * (var.radius**3) # [nm^3] 

90 var.hvol = (4.0 / 3.0) * pi * (var.hradius**3) # [nm^3] 

91 

92 var.phi = c.random.normal(pi, (2 * pi), size=var.no_molecules) 

93 var.theta = c.random.normal(pi / 2, pi, size=var.no_molecules) 

94 var.mass = var.vol * var.dens 

95 var.alpha = 6.0 * var.eta * c.GBARE / var.ms # dimensionless damping parameter 

96 var.geff = c.GBARE / (1.0 + var.alpha**2) # effective gyromagnetic ratio 

97 # size of the magnetic fluctuations 

98 # Note that c2m is divided here by sqrt(dt), but finally multiplied by dt. 

99 # The stochastic part correctly scales with sqrt(dt) 

100 var.c2 = sqrt(2.0 * c.KB * var.temp * var.alpha / (var.ms * c.GBARE * var.hvol * var.dt)) 

101 var.c2 = where(var.c2 > c.EPS2, var.c2, 0) 

102 

103 with errstate(all='raise'): 

104 try: 

105 var.chi0 = (var.ms**2) * var.vol / (3.0 * c.KB * var.temp) 

106 var.tauB = 3.0 * var.eta * var.vol / (c.KB * var.temp) 

107 except FloatingPointError: 

108 c.Error(">W ZeroDivisionError\n{} and {} have been set to infinity".format(c.ustr.chi, c.ustr.taub)) 

109 var.chi0 = float("inf") * ones(var.no_molecules) 

110 var.tauB = float("inf") * ones(var.no_molecules) 

111 

112 

113def get_particles(var, flg): 

114 

115 if var.location is not None: 115 ↛ 116line 115 didn't jump to line 116, because the condition on line 115 was never true

116 fpos = {} 

117 fmom = {} 

118 fmag = {} 

119 

120 # TODO probably could improve splitting filenames 

121 # currently splits on all spaces 

122 # move splitting to arguments.py etc 

123 location = var.location.split() 

124 

125 for stat, loc in enumerate(location): 

126 pos, mom, mag = _get_particles(var, flg, loc) 

127 fpos[stat] = pos 

128 fmom[stat] = mom 

129 fmag[stat] = mag 

130 

131 if len(var.location.split()) == 1: 

132 return pos, mom, mag, pos.shape 

133 pos = fpos 

134 mom = fmom 

135 mag = fmag 

136 

137 return pos, mom, mag, pos[0].shape 

138 else: 

139 pos, mom, mag = _get_particles(var, flg, var.location) 

140 return pos, mom, mag, pos.shape 

141 

142 

143def _get_particles(var, flg, location): 

144 (pos, mom, mag), xyz = initpos(var.no_molecules, var.radius * 2, location, var.boxsize) 

145 

146 sanity_checks(var, pos, mag, xyz) 

147 

148 new_vars(var) 

149 

150 if mag is None: 150 ↛ 160line 150 didn't jump to line 160, because the condition on line 150 was never false

151 mag = zeros_like(pos) 

152 # # Transposed (ji not ij) due to theta and phi being row vectors 

153 einsum_str = "i, ji -> ij" if var.no_molecules > 1 else "i, ji -> j" 

154 mag[:] = einsum(einsum_str, 

155 var.ms, array([sin(var.theta) * cos(var.phi), 

156 sin(var.theta) * sin(var.phi), 

157 cos(var.theta)])) 

158 # mag[:, 2] = 1 

159 

160 mag = einsum("...iz,...i->...iz", mag, var.ms * var.vol / maximum(sqrt(einsum("...iz,...iz->...i", mag, mag)), c.EPS2)) 

161 

162 if mom is None: 162 ↛ 164line 162 didn't jump to line 164, because the condition on line 162 was never false

163 mom = zeros_like(pos) 

164 elif flg.labview: 

165 mom -= einsum("...iz->...z", mom) / mom.shape[0] # unit: 1.e-12 g*cm/s 

166 

167 return pos, mom, mag 

168 

169 

170# Main calculation 

171@debug(['core']) 

172def integrate(var={}, flg={}): 

173 """ 

174 Set up of Calculation. 

175 

176 Setting up of initial positions magnetisations and momentums 

177 as well as neel relaxations for each molecule 

178 

179 This is all passed to the parallelisation of the statistical separations 

180 """ 

181 c.reseed(54321) 

182 

183 if flg.restart: 183 ↛ 184line 183 didn't jump to line 184, because the condition on line 183 was never true

184 pos = var.pos 

185 mag = var.mag 

186 mom = var.mom 

187 del var.pos 

188 del var.mom 

189 del var.mag 

190 

191 movement = {"var": var, "flg": flg, "mom": mom, "mag": mag, "pos": pos} 

192 else: 

193 

194 pos, mom, mag, shape = get_particles(var, flg) 

195 

196 # Particle movement variables 

197 movement = {"var": var, "flg": flg, "mom": mom, "mag": mag, "pos": pos, 

198 "op_noise": c.random.standard_normal(size=shape) if flg.op_thermal else None} 

199 

200 print(var.name[:-1]) 

201 

202 # Time recording 

203 timer = end(var.finishtime) 

204 

205 # Initial conditions 

206 calc = calculate(**{"Error": c.Error, "posit": position(**movement), 

207 "var": var, "flg": flg, "timer": timer}) 

208 

209 if hasattr(integrate, "DBG"): 209 ↛ 210line 209 didn't jump to line 210, because the condition on line 209 was never true

210 var.print_contents() 

211 flg.print_contents() 

212 

213 # Parallelisation of calculations 

214 mp_handle()(calc_wrap, calc, var.stats, flg.parallel) 

215 

216 return timer, var.name 

217 

218 

219class calculate(): 

220 """ 

221 1st parallel section. 

222 

223 This function does the boilerplate around the brunt of the calculation 

224 it is run once for each statistic 

225 

226 All the data is then saved in one of three user specified forms 

227 

228 1. hdf5 

229 2. pickle 

230 3. plain text 

231 

232 """ 

233 

234 @debug(["calc"]) 

235 def __init__(self, **argd): 

236 """ 

237 Initialise constant and variables. 

238 

239 dictionary keys: var, flg, mom, mag, pos, timer, posit (initialised position class) 

240 

241 Parameters 

242 ---------- 

243 argd: dict 

244 dictionary of data 

245 

246 """ 

247 self.__dict__.update(argd) 

248 

249 self.var.skip_iters = num_iters(self.var.nmax, self.var.skip) 

250 self.savechunk() 

251 

252 self.dictcreate = dictionary_creation(self.var, self.flg) 

253 self.save = save_m(self.flg.neel, write_data(self.var, self.flg)) 

254 self.grtime = grace(self.var.stats).time 

255 self.end = False 

256 self.nsplit = 0 

257 

258 self.h_axis = array([[0.0, 0.0, 1.0]]) * ones((self.var.no_molecules, 3)) 

259 

260 if self.flg.noise: 

261 self.getnoise = self._returnnoise 

262 else: 

263 self.getnoise = self._noop 

264 

265 if self.flg.suscep or self.flg.field is False: 

266 self.h_axis *= 0 

267 self.getfield = self._returnnoextfield 

268 else: 

269 self.getfield = self._returnextfield 

270 

271 if self.flg.prog: 271 ↛ 272line 271 didn't jump to line 272, because the condition on line 271 was never true

272 self.progress = self._prog_verb 

273 self.progresstime = 1 # max(600 // self.flg.nout, 1) 

274 elif 3 > self.flg.nout > 1: 

275 self.progress = self._prog_nverb 

276 self.progresstime = max(600 // self.flg.nout, 1) 

277 else: 

278 self.progress = self._noop 

279 self.progresstime = 300 

280 

281 if hasattr(self, 'DBG'): 281 ↛ 282line 281 didn't jump to line 282, because the condition on line 281 was never true

282 self._refresh_dict = self._refresh_dict_debug 

283 

284 @staticmethod 

285 def _noop(): 

286 pass 

287 

288 def _prog_verb(self): 

289 c.progress(f"bar {self.moldata_dict['name'][len(self.var.directory):]} {self.count / self.var.nmax} {self.stat}") 

290 

291 def _prog_nverb(self): 

292 print("{} of {} statistic {}".format(self.count, self.var.nmax, self.stat)) 

293 

294 def _returnnoise(self): 

295 c.random.standard_normal(out=self.noise_setup) 

296 self.Hext += c_einsum("ij, i -> ij", self.noise_setup[0], self.var.c2) 

297 

298 def _returnnoextfield(self): 

299 self.Hext[:] = self.h_axis 

300 

301 def _returnextfield(self): 

302 self.Hext[:] = field(self.time.time, H0=self.var.H_0, nu=self.var.nu) * self.h_axis 

303 

304 def progress_report(self): 

305 """Progress reporter.""" 

306 self.count += self.var.skip 

307 

308 self.countdown = self.timer.gettimegone() 

309 if (self.countdown - self.oldcount) > self.progresstime: 309 ↛ 310line 309 didn't jump to line 310, because the condition on line 309 was never true

310 self.oldcount = self.countdown 

311 self.progress() 

312 if self.timer.gettimeleft() <= self.grtime: 

313 self.moldata_dict = self.save.stop(self.moldata_dict, self.posit.count) 

314 c.Error("W Hit walltime, trying to exit cleanly") 

315 

316 def savechunk(self): 

317 """Adjust savechunk.""" 

318 self.bksavechunk = self.var.savechunk 

319 if self.var.savechunk > self.var.skip_iters: 

320 # split can't be larger than no_iterations 

321 self.var.savechunk = self.var.skip_iters 

322 

323 def iterationadjust(self): 

324 """Adjust iterations for writing and storing.""" 

325 if self.flg.restart: 325 ↛ 326line 325 didn't jump to line 326, because the condition on line 325 was never true

326 self.rstep = ((self.var.extra_iter[self.stname]) * self.var.skip) 

327 self.cstep = self.var.nmax - self.rstep 

328 # Doesn't include post optimise initial positions step 

329 # Does include post optimise initial positions step 

330 self.cwritten = self.var.skip_iters - self.var.extra_iter[self.stname] + 1 

331 print("Stat {} Written: {} / {}".format(self.stat, self.cwritten, self.var.skip_iters + 1)) 

332 if self.cwritten == self.var.skip_iters + 1: 

333 self.finishdata(0, self.cstep) 

334 exit(0) 

335 elif self.cwritten > self.var.skip_iters + 1: 

336 print("Bigger?", self.cwritten, self.var.skip_iters, num_iters(self.rstep, self.var.skip)) 

337 exit(1) 

338 

339 self.var.skip_iters = num_iters(self.rstep, self.var.skip) 

340 nmax = self.rstep 

341 else: 

342 nmax = self.var.nmax 

343 

344 if self.flg.restart or isinstance(self.posit.pos, dict): 344 ↛ 345line 344 didn't jump to line 345, because the condition on line 344 was never true

345 self.posit.pos = self.posit.pos[self.stat] 

346 self.posit.mag = self.posit.mag[self.stat] 

347 self.posit.mom = self.posit.mom[self.stat] 

348 

349 self.savechunk() 

350 

351 # savechunk as 0 == save at the end 

352 self.writes = nmax // (self.var.savechunk * self.var.skip) if self.var.savechunk != 0 else 0 

353 

354 self.state_int = max(self.writes // 10, 1) 

355 self.finalmemoryloop = (nmax % (self.var.savechunk * self.var.skip)) / self.var.skip if self.var.savechunk != 0 else 0 

356 

357 self.finalskiploop = nmax % self.var.skip 

358 if self.finalskiploop == 0 and self.finalmemoryloop % 1 != 0: 358 ↛ 360line 358 didn't jump to line 360, because the condition on line 358 was never true

359 # finalmemory loop not whole number and skip fits exactly into nmax 

360 self.finalskiploop = self.var.skip 

361 

362 # whole number 

363 self.finalmemoryloop = int(self.finalmemoryloop) 

364 

365 # last dictionary size: Number of memory loops + last skip loop 

366 self.remainder = self.finalmemoryloop + 1 if self.finalskiploop > 0 else 0 

367 

368 if self.finalmemoryloop == 0 and self.finalskiploop == 0: 368 ↛ 370line 368 didn't jump to line 370, because the condition on line 368 was never true

369 # Always run final loop at least once 

370 self.finalmemoryloop = self.var.savechunk 

371 self.writes -= 1 if self.writes > 0 else 0 

372 

373 if hasattr(self, 'DBG'): 373 ↛ 374line 373 didn't jump to line 374, because the condition on line 373 was never true

374 print(f"Stat {self.stat}", f"Write_loop: {self.writes}", 

375 f"Savechunk: {self.var.savechunk}", f"Skip: {self.var.skip}", 

376 f"Final_M_loop: {self.finalmemoryloop}", 

377 f"Final_S_loop: {self.finalskiploop}", 

378 f"Remainder: {self.remainder}", 

379 f"Total_tosave: {self.var.skip_iters}") 

380 

381 def randomnumbers(self): 

382 """Set up random number starting point.""" 

383 if self.flg.restart: 383 ↛ 384line 383 didn't jump to line 384, because the condition on line 383 was never true

384 c.rgen.state = self.var.RandNoState[self.stname] 

385 else: 

386 c.reseed(54321) 

387 c.jump(self.stat) 

388 

389 def dictionarys(self): 

390 """Create dictionaries for data storage.""" 

391 self.moldata_dict_end = dictionary_creation(self.var, self.flg, self.remainder) 

392 if self.bksavechunk != self.var.savechunk: 392 ↛ 393line 392 didn't jump to line 393, because the condition on line 392 was never true

393 self.dictcreate = dictionary_creation(self.var, self.flg) 

394 self.moldata_dict = self.dictcreate.copy() 

395 # File naming 

396 self.moldata_dict['name'] += "{:g}.{}".format(self.stat + 1, self.flg.save_type) 

397 self.name = self.moldata_dict['name'] 

398 

399 def _refresh_dict(self, dictionary): 

400 self.moldata_dict = dictionary.copy() 

401 self.moldata_dict['name'] = self.name 

402 self.save.remove_state(dictionary) 

403 

404 def _refresh_dict_debug(self, dictionary): 

405 self.moldata_dict = deepcopy(dictionary) 

406 self.moldata_dict['name'] = self.name 

407 self.save.remove_state(dictionary) 

408 

409 @addpid("Error") 

410 def setup(self): 

411 """Set up instance for each parallel calculation.""" 

412 self.noise_setup = zeros((3, self.var.no_molecules, 3)) 

413 self.Hext = zeros((self.var.no_molecules, 3)) 

414 

415 self.stname = self.save.stname = f"stat{self.stat}" 

416 

417 self.iterationadjust() 

418 self.randomnumbers() 

419 self.dictionarys() 

420 

421 self.flg.save_type = self.save.setup(self.name) 

422 

423 self.starttime() 

424 self.firststep() 

425 self.prerange() 

426 

427 # Progress report varables 

428 self.oldcount = self.timer.gettimegone() 

429 

430 def prerange(self): 

431 """Preallocate ranges.""" 

432 self.dlrange = range(self.writes) 

433 self.mlrange = range(self.var.savechunk) 

434 self.slrange = range(self.var.skip) 

435 

436 self.mlrangefinal = range(self.finalmemoryloop) 

437 self.slrangefinal = range(self.finalskiploop) 

438 

439 def run(self, stat): 

440 """ 

441 Run calculation. 

442 

443 Parameters 

444 ---------- 

445 stat: int 

446 statistic number 

447 

448 """ 

449 self.stat = stat 

450 self.setup() 

451 self.disksaveloop() 

452 

453 def starttime(self): 

454 """Start timestep recording.""" 

455 self.time = _time(self.var.dt, self.var.time) 

456 

457 self.posit.time = self.time 

458 

459 def firststep(self): 

460 """Set up initial conditions and calculate first step.""" 

461 self.posit.initialconditions() 

462 if self.flg.restart is None: 462 ↛ 468line 462 didn't jump to line 468, because the condition on line 462 was never false

463 self.count = 0 

464 first_step = dictionary_creation(self.var, self.flg, 1) 

465 first_step["name"] = self.moldata_dict['name'] 

466 self.save.file(self.save.memory(first_step, self.posit.prop, self.time)) 

467 else: 

468 self.count = self.cstep 

469 self.save.restart(self.moldata_dict['name'], self.var.extra_iter, self.var.SCFcount, self.cwritten) 

470 

471 def disksaveloop(self): 

472 """Save to disk loop.""" 

473 for dl in self.dlrange: 

474 self.memorysaveloop(self.mlrange) 

475 if self.save.end: 475 ↛ 476line 475 didn't jump to line 476, because the condition on line 475 was never true

476 return 

477 if dl % self.state_int == 0 and dl != 0: 

478 self.save.state(self.moldata_dict, self.posit.count) 

479 self.save.file(self.moldata_dict, self.posit.count) 

480 self._refresh_dict(self.dictcreate) 

481 

482 self.finaliteration() 

483 

484 def memorysaveloop(self, loop): 

485 """ 

486 Save to memory loop. 

487 

488 Parameters 

489 ---------- 

490 loop: int 

491 loop size 

492 

493 """ 

494 for _ in loop: 

495 self.skiploop(self.slrange) 

496 

497 # Collect data 

498 self.moldata_dict = self.save.memory(self.moldata_dict, self.posit.prop, self.time) 

499 

500 self.progress_report() 

501 

502 def skiploop(self, loop): 

503 """ 

504 Skip loop, no saving required. 

505 

506 Parameters 

507 ---------- 

508 loop: int 

509 loop size 

510 

511 """ 

512 for _ in loop: 

513 # External field 

514 self.getfield() 

515 

516 # System noise 

517 self.getnoise() 

518 

519 # Iterate 

520 self.posit.propagate(self.Hext, self.noise_setup) 

521 

522 # time update 

523 self.time.time_update() 

524 

525 def finaliteration(self): 

526 """ 

527 Last iteration loop. 

528 

529 Setup for last step saving useful things for restart 

530 """ 

531 if self.moldata_dict_end is not None: 531 ↛ 534line 531 didn't jump to line 534, because the condition on line 531 was never false

532 self._refresh_dict(self.moldata_dict_end) 

533 else: 

534 self._refresh_dict(self.dictcreate) 

535 

536 self.memorysaveloop(self.mlrangefinal) 

537 

538 if self.finalskiploop > 0: 538 ↛ 543line 538 didn't jump to line 543, because the condition on line 538 was never false

539 self.count += self.finalskiploop 

540 self.skiploop(self.slrangefinal) 

541 self.moldata_dict = self.save.memory(self.moldata_dict, self.posit.prop, self.time) 

542 

543 self.save.end = True 

544 self.save.file(self.moldata_dict, self.posit.count) 

545 

546 self.progress() 

547 self.finishdata(self.posit.count, self.count) 

548 

549 def finishdata(self, SCFcount, count): 

550 SCFcount += self.var.SCFcount if self.flg.restart else 0 

551 print(f"\nStat {self.stat}\n", 

552 f"SCF cycles: {SCFcount}\n", 

553 f"Total Iterations: {self.var.nmax}\n", 

554 f"Completed Iterations: {count}\n", 

555 f"SCF per Iteration: {(SCFcount / count) if count > 0 else 0}") 

556 

557 

558class save_m(): 

559 """ 

560 Storing function. 

561 

562 Stores data in memory until asked to save to file where it calls the writer 

563 

564 Parameters 

565 ---------- 

566 neel: bool 

567 Is neel relaxation required? 

568 wd: instance 

569 instance of writing class 

570 

571 """ 

572 

573 @debug(["save"]) 

574 def __init__(self, neel, wd): 

575 """Initialise save routine.""" 

576 self.wd = wd 

577 

578 self.end = False 

579 self.neel = neel 

580 if self.neel: 

581 self.neel_save = self._neel_save 

582 else: 

583 self.neel_save = self._no_neel_save 

584 self.written = 0 

585 self.SCFcount = 0 

586 self._reset_ind() 

587 

588 def _neel_save(self, data_dict, data): 

589 data_dict['neel_relaxation'][self.splitcount, :] = data["neel_count"].count 

590 return data_dict 

591 

592 @staticmethod 

593 def _no_neel_save(data_dict, data): 

594 return data_dict 

595 

596 def _reset_ind(self): 

597 self.splitcount = 0 

598 

599 def _incr_splitcount(self): 

600 self.splitcount += 1 

601 self.written += 1 

602 

603 def setup(self, name): 

604 """Set up writing class.""" 

605 return self.wd.setup(name) 

606 

607 def restart(self, name, extra_iter, SCFcount, written): 

608 """Save extra_iters, SCFcount and current written.""" 

609 self.written = written 

610 self.SCFcount = SCFcount 

611 self.wd.write({'name': name, 'vars': {'extra_iter': extra_iter}}) 

612 

613 def file(self, data_dict, SCFcount=0): 

614 """ 

615 Save current data block to file. 

616 

617 Parameters 

618 ---------- 

619 data_dict: dict 

620 Storage dictionary 

621 

622 """ 

623 if self.end: 

624 self.state(data_dict, SCFcount) 

625 self.wd.write(data_dict) 

626 self._reset_ind() 

627 

628 def state(self, data_dict, SCFcount): 

629 """Save current state.""" 

630 data_dict['vars']['RandNoState'][self.stname] = c.rgen.state 

631 data_dict['vars']['written'] = self.written 

632 data_dict['vars']['SCFcount'] = self.SCFcount + SCFcount 

633 

634 def remove_state(self, data_dict): 

635 """Remove current state.""" 

636 for val in ['RandNoState', 'written', 'SCFcount']: 

637 data_dict['vars'][val] = {} 

638 

639 def stop(self, data_dict, SCFcount=0): 

640 """ 

641 Emergency Stop saving. 

642 

643 Save current data block to file. 

644 

645 Parameters 

646 ---------- 

647 data_dict: dict 

648 Storage dictionary 

649 

650 """ 

651 self.end = True 

652 sc = self.splitcount # save to avoid reset 

653 

654 self.file(self.cutter(data_dict, deepcopy(data_dict), None, sc), SCFcount) 

655 

656 return self.cutter(data_dict, data_dict, sc, None) 

657 

658 def cutter(self, orig, copy, start, end): 

659 """Return a section of input dictionary.""" 

660 copy['position'] = orig['position'][start:end, :, :] 

661 copy['iter_time'] = orig['iter_time'][start:end, :] 

662 copy['magnetisation'] = orig['magnetisation'][start:end, :, :] 

663 copy['forces'] = orig['forces'][start:end, :, :] 

664 copy['momentum'] = orig['momentum'][start:end, :, :] 

665 

666 if self.neel: 

667 copy['neel_relaxation'][start:end, :] = orig['neel_relaxation'][start:end, :] 

668 

669 return copy 

670 

671 def memory(self, data_dict, data, time): 

672 """ 

673 Save to memory current iteration data. 

674 

675 Parameters 

676 ---------- 

677 data_dict: dict 

678 Storage dictionary 

679 data: dict 

680 data to be stored 

681 time: instance 

682 timestep instance 

683 

684 Returns 

685 ------- 

686 data_dict: dict 

687 Storage dictionary 

688 

689 """ 

690 data_dict['position'][self.splitcount, :, :] = data["pos"] 

691 data_dict['iter_time'][self.splitcount, :] = time.iter, time.time 

692 data_dict['magnetisation'][self.splitcount, :, :] = data["mag"] 

693 data_dict['forces'][self.splitcount, :, :] = data["forces"] 

694 data_dict['momentum'][self.splitcount, :, :] = data["mom"] 

695 

696 data_dict = self.neel_save(data_dict, data) 

697 

698 self._incr_splitcount() 

699 

700 return data_dict 

701 

702 

703@debug(['core']) 

704def dictionary_creation(var, flg, extra=False): 

705 """ 

706 Creation of initial dictionary. 

707 

708 - all data is the same at this point 

709 

710 - Uses nested dictionaries for ease of storage 

711 """ 

712 if extra is False: 

713 size = var.savechunk 

714 elif extra > 0: 

715 size = extra 

716 else: 

717 return None 

718 

719 moldata_dict = _getdicts(size, var.no_molecules, flg.neel) 

720 

721 moldata_dict['mango_version'] = c.version 

722 moldata_dict['name'] = var.name 

723 

724 moldata_dict['flags'] = flg.__dict__.copy() 

725 

726 moldata_dict['vars'] = var.__dict__.copy() 

727 

728 del moldata_dict['vars']['finishtime'] 

729 with suppress(KeyError): 

730 del moldata_dict['vars']['defaults'] 

731 del moldata_dict['vars']['time'] 

732 del moldata_dict['vars']['name'] 

733 

734 for val in ['RandNoState', 'extra_iter', 'SCFcount', 'written']: 

735 moldata_dict['vars'][val] = {} 

736 

737 return moldata_dict 

738 

739 

740def _getdicts(x, y, neel): 

741 new_dict = nestedDict() 

742 new_dict['iter_time'] = zeros((x, 2)) 

743 new_dict['position'] = zeros((x, y, 3)) 

744 new_dict['magnetisation'] = zeros_like(new_dict['position']) 

745 new_dict['forces'] = zeros_like(new_dict['position']) 

746 new_dict['momentum'] = zeros_like(new_dict['position']) 

747 

748 # new_dict['CoM'] = zeros((x, 3)) 

749 # new_dict['CoM_vel'] = zeros_like(new_dict['CoM']) 

750 # new_dict['momenta']['total'] = zeros_like(new_dict['CoM']) 

751 # new_dict['momenta']['total_mag'] = zeros_like(new_dict['CoM']) 

752 # new_dict['momenta']['total_angular'] = zeros_like(new_dict['CoM']) 

753 # new_dict['energy']['kinetic'] = zeros(x) 

754 # new_dict['energy']['trans_pot'] = zeros_like(new_dict['energy']['kinetic']) 

755 # new_dict['energy']['mag_pot'] = zeros_like(new_dict['energy']['kinetic']) 

756 # new_dict['energy']['total'] = zeros_like(new_dict['energy']['kinetic']) 

757 # new_dict['energy']['kineticm'] = zeros((x, y)) 

758 

759 if neel: 

760 new_dict['neel_relaxation'] = zeros((x, y)) 

761 

762 return new_dict 

763 

764 

765if __name__ == "__main__": 

766 pass