1""" 

2The Main engine. 

3 

4Iteration loop module for updating properties of particles. 

5""" 

6 

7from numpy import (zeros, zeros_like, ravel, ones, errstate, 

8 exp, sqrt, array, 

9 where, append, maximum, allclose, 

10 add, subtract, multiply, einsum) 

11 

12from numpy.core.multiarray import c_einsum 

13 

14 

15from math import pi 

16from scipy.special import comb 

17from scipy.optimize import minimize 

18from functools import partial 

19 

20from mango.constants import c, bin_count 

21from mango.boundaries import periodic 

22from mango.imports import getpotentials 

23from mango.debug import debug 

24 

25 

26class Communal(): 

27 """Reused variables and methods for propogation.""" 

28 

29 def __init__(self, epsilon, sigma, tri, comb, limit): 

30 """Initialise variables.""" 

31 self.epsilon = epsilon 

32 self.sigma = sigma 

33 self.triag_ind = tri 

34 self.ia = tri[0] 

35 self.ja = tri[1] 

36 self.limit = limit 

37 self.comb = comb 

38 self.rij = zeros(self.comb) 

39 self.set_attr(['lj', 'lj6', 'rij3', 'rij4', 'magi', 'magj', 'dmag'], zeros_like(self.rij)) 

40 self.lind = zeros_like(self.rij, dtype=bool) 

41 self.hind = zeros_like(self.rij, dtype=bool) # self.hind[:] = ~lind 

42 

43 def set_attr(self, var_list, val): 

44 """Initialise multiple arrays with same shape.""" 

45 for i in var_list: 

46 setattr(self, i, val.copy()) 

47 

48 

49class countneel(object): 

50 """ 

51 Counting and storing neel relaxations. 

52 

53 initialise with timestep(dt) and Neel relaxation time (tauN) 

54 

55 run with each magnetisation step 

56 

57 """ 

58 

59 def __init__(self, var, dt): 

60 """Initialise variables.""" 

61 with errstate(all='raise'): 

62 try: 

63 sigmaN = exp((var.keff * var.vol) / (c.KB * (var.temp - var.temp0))) 

64 var.tauN = var.tauN_0 * sigmaN 

65 except FloatingPointError: 

66 c.Error(">W OverflowError {} is infinite".format(c.ustr.taun)) 

67 var.tauN = float("inf") * ones(var.no_molecules) 

68 self.count = zeros(var.no_molecules) 

69 self.tauN = var.tauN 

70 self.dt = dt 

71 self._modifydt(self.dt) 

72 self.no_mol = var.no_molecules 

73 

74 def _modifydt(self, dt): 

75 self.A = exp(-dt / self.tauN) 

76 self.dt = dt 

77 

78 def run(self, dt, mag): 

79 """Counter.""" 

80 if dt != self.dt: 

81 self._modifydt(dt) 

82 

83 self.mag = mag 

84 self.x = c.random.random(size=self.no_mol) 

85 

86 ind = where(self.x > self.A) 

87 

88 self.mag[ind] = -self.mag[ind] 

89 self.count[ind] += 1 

90 

91 

92class position(Communal): 

93 """ 

94 Main calculation class. 

95 

96 Calling propogate will calculate 1 iteration 

97 """ 

98 

99 boundary = periodic() 

100 

101 @debug(['init']) 

102 def __init__(self, **argd): 

103 """ 

104 Initialise propgation routine. 

105 

106 Sets up potentials and variables required. 

107 Optimises the positional configuration of the particles if required. 

108 

109 dictionary keys: var, flg, mom, mag, pos, op_noise 

110 

111 Parameters 

112 ---------- 

113 argd: dict 

114 dictionary of data 

115 

116 """ 

117 self.__dict__.update(argd) 

118 self.count = 0 

119 self.mnorm = c.EPS_val + 1 

120 self.halfdt = self.var.dt * 0.5 

121 self.maxscf = 1 + self.var.scfiter * 2**5 

122 self.Mmag = self.var.ms * self.var.vol 

123 

124 self.inv_Mmag = 1 / self.Mmag 

125 self.inv_Mmag2 = 1 / c_einsum('i,i -> i', self.Mmag, self.Mmag) 

126 self.invmass = 1 / self.var.mass 

127 self.invnatom = 1 / self.var.no_molecules 

128 self.inv_Mmag2_atm = self.inv_Mmag2 * self.invnatom 

129 self.sqEPS = c.EPS_val ** 2 

130 self.sqmnorm = zeros((0, 1)) 

131 

132 self.abslimit = c.EPS_val * 1.2 

133 self.prerange = range(0, 6) 

134 self.scfrange = range(0, self.var.scfiter) 

135 self.new_norm = [0, 0] 

136 self.sqmnorm = zeros(1) 

137 self.stoc_constant = 6 * pi * self.var.eta * self.var.hradius * self.invmass # gammat 

138 

139 self.estrings() 

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

141 self.set_attr(['c1t', 'c2t'], zeros(self.var.no_molecules)) 

142 self.set_attr(['tot', 'tmag', 'angular'], zeros(3)) 

143 self.set_attr(['force', 'Hext', 'dm_mag_old', 'dm_umag', 

144 'dm_diff', 'dm_hfield', 'dm_Heff', 

145 'pos_tmp', 'noise_tmp'], zeros((self.var.no_molecules, 3))) 

146 

147 self.tri = c.tri 

148 self.comb = comb(self.var.no_molecules, 2, exact=True) 

149 

150 self.set_pot() 

151 self.init_pot() 

152 

153 if self.flg.opt: 153 ↛ 154line 153 didn't jump to line 154, because the condition on line 153 was never true

154 self.dist, self.dist_matrix = self.boundary.setup(self.var.boxsize, self.flg.pbc, self.pos, self.var.sigma, self.tri) 

155 self.optimise(toll=c.EPS, iprint=2 if self.flg.nout >= 5 else 0) 

156 

157 if self.flg.neel: 157 ↛ 158line 157 didn't jump to line 158, because the condition on line 157 was never true

158 self.neel_count = countneel(self.var, self.halfdt) 

159 self.magnetisation_step = partial(self.neel, self.magnetisation_step) 

160 else: 

161 self.neel_count = None 

162 

163 def _noise(self, noise_row): 

164 sqrt(c.KB * self.var.temp * self.var.mass * (1. - self.c1t**2), out=self.c2t) # [A/fs] 

165 c_einsum("i, ij -> ij", self.c2t, self.noise_setup[noise_row, ...], out=self.noise_tmp) 

166 add(self.mom, self.noise_tmp, out=self.mom) 

167 

168 def estrings(self): 

169 self.dm1 = "iz,i->iz" 

170 self.dm2 = "iz,iz->i" 

171 self.dm3 = 'i,xyz,ix,iy->iz' 

172 self.scf1 = "iz,iz,i->" 

173 self.p_s = "ij, i -> ij" 

174 

175 @staticmethod 

176 def _noop(*args): 

177 pass 

178 

179 def set_pot(self): 

180 """Set new potentials if required.""" 

181 if self.var.potential: 

182 new_pot = getpotentials(self.var.potential) 

183 functionlist = [f for f in dir(self) if not f.startswith('_')] 

184 

185 for i in new_pot.__all__: 

186 if ("vv" in i or i == "energy_vars") and not hasattr(energy, "_" + i): 

187 setattr(energy, "_" + i, getattr(energy, i)) 

188 setattr(energy, i, getattr(new_pot, i)) 

189 if ("ff" in i or "hh" in i or i == "force_vars") and not hasattr(force, "_" + i): 

190 setattr(force, "_" + i, getattr(force, i)) 

191 setattr(force, i, getattr(new_pot, i)) 

192 if i in functionlist: 192 ↛ 193line 192 didn't jump to line 193, because the condition on line 192 was never true

193 setattr(self, "_" + i, getattr(self, i)) 

194 setattr(self, i, getattr(new_pot, i)) 

195 

196 def init_pot(self): 

197 """Initialise potentials.""" 

198 com = Communal(self.var.epsilon, self.var.sigma, self.tri, self.comb, self.var.limit) 

199 self.e = energy(com) 

200 self.f = force(com, self.var.no_molecules) 

201 

202 def initialconditions(self): 

203 """ 

204 Set initial conditions of particles and boundary conditions. 

205 

206 If MPI is used reinitialise potentials. 

207 """ 

208 if c.comm.Get_size() > 1 and c.comm.Get_rank() != 0: 208 ↛ 209line 208 didn't jump to line 209, because the condition on line 208 was never true

209 self.set_pot() 

210 self.init_pot() 

211 if not hasattr(self.boundary, 'wrap'): 

212 self.boundary.setup(self.var.boxsize, self.flg.pbc, self.pos, self.var.sigma, self.tri) 

213 else: 

214 self.boundary.reset_wrap() 

215 if not self.flg.mag_sw: 

216 self.magnetisation_step = self._noop 

217 self.noise = self._noise if self.flg.noise else self._noop 

218 self.pos[:] = self.boundary.wrapping(self.pos) 

219 self.dist, self.dist_matrix = self.boundary.distance(self.pos) 

220 self.f.force_vars(self.dist, self.dist_matrix) 

221 self.force_step() 

222 

223 self.prop = {"pos": self.pos, "mom": self.mom, "mag": self.mag, 

224 "neel_count": self.neel_count, "forces": self.force} 

225 

226 def neel(self, magstep, dt): 

227 """Wrap magnetisation if neel relaxation is used.""" 

228 self.neel_count.run(self.halfdt, self.mag) 

229 self.mag[:] = self.neel_count.mag 

230 magstep(dt) 

231 self.neel_count.run(self.halfdt, self.mag) 

232 self.mag[:] = self.neel_count.mag 

233 

234 @debug(['prop']) 

235 def propagate(self, Hext, noise_setup): 

236 """Calculate propagation loop.""" 

237 self.Hext = Hext 

238 self.noise_setup = noise_setup 

239 

240 # half position_step 

241 self.position_step(self.halfdt) 

242 # half momentum_step 

243 self.f.force_vars(self.dist, self.dist_matrix) 

244 self.momentum_step(self.halfdt) 

245 # half stochastic_step 

246 self.stochastic_step(self.halfdt, 1) 

247 

248 # full magnetisation_step 

249 self.magnetisation_step(self.var.dt) 

250 

251 # half stochastic_step 

252 self.stochastic_step(self.halfdt, 2) 

253 # half momentum_step 

254 self.momentum_step(self.halfdt) 

255 # half position_step 

256 self.position_step(self.halfdt) 

257 

258 @debug(['pos']) 

259 @boundary 

260 def position_step(self, dt): 

261 """Calculate the new position of the particle.""" 

262 c_einsum(self.p_s, self.mom, self.invmass, out=self.pos_tmp) 

263 add(self.pos, self.pos_tmp * dt, out=self.pos) 

264 

265 def force_step(self): 

266 """Calculate the forces between atoms.""" 

267 add(self.f._ff_trans_return(), self.f.ff_mag(self.mag), out=self.force) 

268 

269 @debug(["mom"]) 

270 def momentum_step(self, dt): 

271 """Calculate the momentum of all the particles.""" 

272 self.force_step() 

273 

274 add(self.mom, self.force * dt, out=self.mom) 

275 

276 def stochastic_step(self, dt, noise_row): 

277 """Apply noise to system.""" 

278 exp(- self.stoc_constant * dt, out=self.c1t) # dimensionless 

279 

280 c_einsum(self.p_s, self.mom, self.c1t, out=self.mom) 

281 

282 self.noise(noise_row) 

283 

284 def optimise(self, toll=1.e-6, iprint=2): 

285 """Optimise particle position based on magnetisation.""" 

286 def func(x): 

287 pos = x[:self.pos.size].reshape(self.pos.shape) 

288 

289 mag = x[self.pos.size:self.pos.size + self.mag.size].reshape(self.mag.shape) 

290 

291 d, dm = self.boundary.distance(pos) 

292 

293 self.e.energy_vars(dm) 

294 self.e.vv_trans() 

295 # Possible bug with 1 np, vv_mag returns [] 

296 return einsum('i ->', self.e.epot_tr + self.e.vv_mag(mag, d)) 

297 

298 def dfunc(x): 

299 pos = x[:self.pos.size].reshape(self.pos.shape) 

300 

301 mag = x[self.pos.size:self.pos.size + self.mag.size].reshape(self.mag.shape) 

302 

303 d, dm = self.boundary.distance(pos) 

304 

305 # Note that the gradient with respect to the magnetisation 

306 # gives the magnetic field. This cannot go to zero because 

307 # of the constrained normalisation (see below). 

308 self.f.force_vars(d, dm) 

309 return -append(self.f._ff_trans_return(), self.f.hh_mag(mag)) 

310 

311 def normalisation(x): 

312 mag = x[self.pos.size:self.pos.size + self.mag.size].reshape(self.mag.shape) 

313 

314 return c_einsum("iz,iz->i", mag, mag) - c_einsum("iz,iz->i", self.mag, self.mag) 

315 

316 x0 = append(self.pos, self.mag) 

317 

318 # x = fmin_slsqp(func, x0, fprime=dfunc, f_eqcons=normalisation, acc=toll, iter=2000, iprint=iprint) 

319 res = minimize(func, x0, jac=dfunc, method='slsqp', 

320 constraints=dict(type='eq', fun=normalisation), 

321 options=dict(maxiter=2000, iprint=iprint, ftol=toll)) 

322 x = res.x 

323 

324 self.pos[:] = x[:self.pos.size].reshape(self.pos.shape) # unit: 1.e-6 cm 

325 

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

327 c_einsum("i, ij -> ij", sqrt(self.var.mass * c.KB * self.var.temp), self.op_noise, out=self.mom) 

328 

329 if all(self.Mmag > c.EPS_val): 329 ↛ exitline 329 didn't return from function 'optimise', because the condition on line 329 was never false

330 self.mag = x[self.pos.size:self.pos.size + self.mag.size].reshape(self.mag.shape) # unit: 1.e-12 emu 

331 

332 scale = self.Mmag / maximum(sqrt(c_einsum("iz,iz->i", self.mag, self.mag)), c.EPS_val) 

333 

334 self.mag = c_einsum("iz,i->iz", self.mag, scale) # unit: 1.e-12 emu 

335 

336 def deltamag(self, mag, dt, umag, hfield, _Heff): 

337 """Update magnetisation.""" 

338 add(mag, self.mag, out=mag) 

339 

340 c_einsum(self.dm1, mag, 1 / sqrt(c_einsum(self.dm2, mag, mag)), out=umag) # dimensionless 

341 

342 c_einsum(self.dm1, umag, self.Mmag, out=mag) # unit: 1.e-12 emu 

343 

344 # WARNING: The new mag is used just to compute hfield! 

345 add(self.f.hh_mag(mag), self.Hext, out=hfield) # unit: 1.e6 oersted 

346 

347 # sLLG is remaining lines 

348 # WARNING: alpha is negative! Heff = hfield - _Heff 

349 c_einsum(self.dm3, self.var.alpha, c.eijk, umag, hfield, out=_Heff) # unit: 1.e6 oersted 

350 

351 subtract(hfield, _Heff, out=hfield) 

352 

353 c_einsum(self.dm3, self.var.geff, c.eijk, mag, dt * hfield, out=hfield) # unit: 1.e-12 emu 

354 

355 def scfloop(self, mag, dt, umag, hfield, Heff, mag_old, diff, new_norm, count=0): 

356 """ 

357 SCF loop to update magnetisation and check the magnetisation norm isn't changed. 

358 

359 Most (more than 50%) of the simulation is spent here 

360 improvements to the for loop have the largest impact 

361 """ 

362 for count in self.scfrange: 362 ↛ 373line 362 didn't jump to line 373, because the loop on line 362 didn't complete

363 self.deltamag(mag, dt, umag, hfield, Heff) 

364 add(self.mag, hfield, out=mag) # unit: 1.e-12 emu 

365 subtract(mag, mag_old, out=diff) 

366 

367 if c_einsum(self.scf1, diff, diff, self.inv_Mmag2_atm) <= self.sqEPS: 

368 return count 

369 

370 mag_old[:] = mag # unit: 1.e-12 emu 

371 

372 # From here assumption function is slower, hopefully rarely used 

373 self.var.scfiter *= 2 

374 self.scfrange = range(count, self.var.scfiter) 

375 

376 mnorm = sqrt(c_einsum(self.scf1, diff, diff, self.inv_Mmag2_atm)) 

377 

378 if allclose(new_norm[0], mnorm): 

379 new_norm[1] += 1 

380 else: 

381 new_norm[:] = mnorm, 0 

382 

383 if self.var.scfiter >= self.maxscf: 

384 if new_norm[1] >= 2 and mnorm < 1.3 * self.sqEPS: 

385 c.Error(f"W Norm seems stuck {mnorm}, continuing") 

386 self.var.scfiter //= 32 

387 new_norm = [0, 0] 

388 return count 

389 else: 

390 c.Error(f"F Exceeded self consistant iteration hard limit {new_norm}, {mnorm}") 

391 

392 c.Error( 

393 "W Maximum number of iterations exceeded doubling to {} [MD step {}, mnorm={} ]".format( 

394 self.var.scfiter, self.time.iter, mnorm)) 

395 

396 return self.scfloop(mag, dt, umag, hfield, Heff, mag_old, diff, new_norm, count) 

397 

398 def magnetisation_step(self, dt): 

399 """Update magnetisation.""" 

400 mag = self.mag.copy() # unit: 1.e-12 emu 

401 umag = self.dm_umag 

402 Heff = self.dm_Heff 

403 hfield = self.dm_hfield 

404 

405 # iteration here 

406 self.count += self.scfloop(mag, dt, umag, hfield, Heff, self.dm_mag_old, self.dm_diff, self.new_norm.copy()) 

407 

408 # whole step 

409 self.deltamag(mag, dt, umag, hfield, Heff) 

410 add(self.mag, hfield, out=self.mag) # unit 1.e-12 emu 

411 self.count += 1 

412 

413 

414class energy(Communal): 

415 """Energy potentials.""" 

416 

417 def __init__(self, communal, singleframe=True): 

418 """Initialise energy constants and arrays.""" 

419 Communal.__init__(self, communal.epsilon, communal.sigma, communal.triag_ind, communal.comb, communal.limit) 

420 self.tr_estr = ['{0}, {0}, {0}, {0}, {0}, {0} -> {0}', '{0}, {0} -> {0}'] 

421 self.mag_estr = ['{0}, {0}, {0} -> {0}', '{0}, {1} -> {0}', '{0}, {0} -> {1}'] 

422 if singleframe: 

423 estr = ['i', 'ij'] 

424 self.unitv = zeros((self.comb, 3)) 

425 self.epot_tr = zeros(self.comb if self.comb > 1 else 1) 

426 else: 

427 estr = ['ai', 'aij'] 

428 self.triag_ind = (slice(None), *self.triag_ind) 

429 self.ia = (slice(None), self.ia) 

430 self.ja = (slice(None), self.ja) 

431 self.unitv = zeros((*self.comb, 3)) 

432 self.epot_tr = zeros(self.comb) 

433 self.tr_estr[0] = self.tr_estr[0].format(estr[0]) 

434 self.tr_estr[1] = self.tr_estr[1].format(estr[0]) 

435 self.mag_estr[0] = self.mag_estr[0].format(estr[0]) 

436 self.mag_estr[1] = self.mag_estr[1].format(estr[1], estr[0]) 

437 self.mag_estr[2] = self.mag_estr[2].format(estr[1], estr[0]) 

438 

439 def energy_vars(self, dm): 

440 """Set reused energy variables.""" 

441 self.rij[:] = dm[self.triag_ind] # unit: cm 

442 

443 def vv_trans(self): 

444 """Calculate translation potential energy.""" 

445 # unit: erg 

446 self.lind[:] = where(self.rij > self.limit, True, False) 

447 

448 if True in self.lind: 

449 self.lj[:] = self.sigma * self.rij 

450 c_einsum(self.tr_estr[0], self.lj, self.lj, self.lj, self.lj, self.lj, self.lj, out=self.lj6) 

451 c_einsum(self.tr_estr[1], self.lind, self.epsilon * (1. + 4. * self.lj6 * (self.lj6 - 1.)), out=self.epot_tr) # unit: erg 

452 

453 def vv_mag(self, mag, d): 

454 """Calculate magnetic potential energy.""" 

455 magia = mag[self.ia] 

456 magja = mag[self.ja] 

457 

458 c_einsum(self.mag_estr[0], self.rij, self.rij, self.rij, out=self.rij3) # = self.rij * self.rij * self.rij 

459 

460 c_einsum(self.mag_estr[1], d[self.triag_ind], self.rij, out=self.unitv) # dimensionless 

461 

462 c_einsum(self.mag_estr[2], magia, self.unitv, out=self.magi) 

463 c_einsum(self.mag_estr[2], magja, self.unitv, out=self.magj) 

464 c_einsum(self.mag_estr[2], magia, magja, out=self.dmag) 

465 

466 return -(3. * self.magi * self.magj - self.dmag) * self.rij3 # potential energy unit: erg 

467 

468 

469class force(Communal, bin_count): 

470 """Force potentials.""" 

471 

472 __slots__ = ['epsilon', 'sigma', 'triag_ind', 'comb', 'limit', 

473 'zeros', 'epsilon24', 'rijz', 'ftrans', 'fmag', 'unitv', 

474 'hfia', 'hfja', 'hfia_temp', 'hfja_temp', 'fmag_temp', 

475 'fstr1', 'fstr2', 'fstr3', 'fstr4', 'fstr5', 'fstr6' 

476 'ia', 'ja', 'rij', 

477 'lj', 'lj6', 'rij3', 'rij4', 'magi', 'magj', 'dmag', 

478 'lind', 'hind', 

479 'long_ia', 'long_ja', 

480 'force_trans', 'force_transR', 

481 'magi', 'magj', 'dmag'] 

482 

483 def __init__(self, communal, nm): 

484 """Initialise force constants and arrays.""" 

485 Communal.__init__(self, communal.epsilon, communal.sigma, communal.triag_ind, communal.comb, communal.limit) 

486 self.zeros = zeros((nm, 3)) 

487 self.set_attr(['rijz', 'ftrans', 'fmag', 

488 'unitv', 'hfia', 'hfja', 

489 'hfia_temp', 'hfja_temp', 'fmag_temp'], zeros((self.comb, 3))) 

490 

491 self.epsilon24 = -24 * self.epsilon # avoid unneed operation repetition 

492 self.estrings() 

493 self.bin_count_setup() 

494 

495 def estrings(self): 

496 self.fstr1 = "i, i, i -> i" 

497 self.fstr2 = "ij, i -> ij" 

498 self.fstr3 = "i, i, i, i, i, i -> i" 

499 self.fstr4 = "iz,iz -> i" 

500 self.fstr5 = "i, ij -> ij" 

501 self.fstr6 = "i, i, ij, i-> ij" 

502 

503 def bin_count_setup(self): 

504 """ 

505 bin_count setup. 

506 

507 create variables needed for bincount 

508 """ 

509 self.long_ia = ravel(array([self.ia * 3, self.ia * 3 + 1, self.ia * 3 + 2]), 'F') 

510 self.long_ja = ravel(array([self.ja * 3, self.ja * 3 + 1, self.ja * 3 + 2]), 'F') 

511 self.leng = self.zeros.size 

512 

513 def force_vars(self, d, dm): 

514 """Set reused force variables.""" 

515 self.rij = dm[self.triag_ind] # unit: cm 

516 self.rijz = d[self.triag_ind] 

517 

518 c_einsum(self.fstr1, self.rij, self.rij, self.rij, out=self.rij3) 

519 self.rij4[:] = self.rij3 * self.rij 

520 

521 c_einsum(self.fstr2, self.rijz, self.rij, out=self.unitv) # dimensionles 

522 

523 self.ff_trans() 

524 

525 def ff_trans(self): 

526 """ 

527 Create variables for all forces. 

528 

529 forces are updated twice without positional changes, therefore no change in translational force 

530 """ 

531 # translational force 

532 self.force_trans = self.zeros.copy() 

533 self.force_transR = self.force_trans.ravel() 

534 

535 self.lind[:] = where(self.rij > self.limit, True, False) 

536 

537 if True in self.lind: 

538 self.lj[:] = self.sigma * self.rij 

539 c_einsum(self.fstr3, self.lj, self.lj, self.lj, self.lj, self.lj, self.lj, out=self.lj6) 

540 c_einsum(self.fstr6, self.lind, self.epsilon24 * (2. * self.lj6 - 1.) * self.lj6, 

541 self.rijz, self.rij * self.rij, out=self.ftrans) 

542 

543 ftransR = self.ftrans.ravel() 

544 

545 self.subtractat(self.force_transR, self.long_ia, ftransR) 

546 self.addat(self.force_transR, self.long_ja, ftransR) 

547 

548 def _ff_trans_return(self): 

549 """Calculate Translational Force.""" 

550 return self.force_trans 

551 

552 def ff_mag(self, mag): 

553 """Calculate Magnetic force.""" 

554 force = self.zeros.copy() # unit: dine 

555 forceR = force.ravel() 

556 

557 magia = mag[self.ia] 

558 magja = mag[self.ja] 

559 

560 c_einsum(self.fstr4, magia, self.unitv, out=self.magi) 

561 c_einsum(self.fstr4, magja, self.unitv, out=self.magj) 

562 c_einsum(self.fstr4, magia, magja, out=self.dmag) 

563 

564 subtract(5. * self.magi * self.magj, self.dmag, out=self.dmag) 

565 

566 subtract(c_einsum(self.fstr5, self.dmag, self.unitv), 

567 c_einsum(self.fstr5, self.magi, magja), out=self.fmag_temp) 

568 

569 subtract(self.fmag_temp, c_einsum(self.fstr5, self.magj, magia), out=self.fmag_temp) 

570 

571 c_einsum(self.fstr2, self.fmag_temp, self.rij4, out=self.fmag) # unit: dine 

572 

573 self.fmag *= 3. 

574 

575 fmagR = self.fmag.ravel() 

576 

577 self.subtractat(forceR, self.long_ia, fmagR) 

578 self.addat(forceR, self.long_ja, fmagR) 

579 

580 return force 

581 

582 def hh_mag(self, mag): 

583 """Calculate Applied field.""" 

584 # Possible to split this function into two concurrent streams 

585 # Most calculation time is spent here (up to 50%) 

586 hfield = self.zeros.copy() # unit: oersted 

587 hfieldR = hfield.ravel() 

588 

589 magia = mag[self.ia] 

590 magja = mag[self.ja] 

591 

592 c_einsum(self.fstr4, magia, self.unitv, out=self.magi) 

593 c_einsum(self.fstr4, magja, self.unitv, out=self.magj) 

594 

595 subtract(3. * c_einsum(self.fstr5, self.magj, self.unitv), magja, out=self.hfia_temp) 

596 subtract(3. * c_einsum(self.fstr5, self.magi, self.unitv), magia, out=self.hfja_temp) 

597 

598 c_einsum(self.fstr2, self.hfia_temp, self.rij3, out=self.hfia) # unit: oersted 

599 c_einsum(self.fstr2, self.hfja_temp, self.rij3, out=self.hfja) # unit: oersted 

600 

601 self.addat(hfieldR, self.long_ia, self.hfia.ravel()) 

602 self.addat(hfieldR, self.long_ja, self.hfja.ravel()) 

603 

604 return hfield 

605 

606 

607if __name__ == "__main__": 

608 pass