Coverage for mango/pp/energy.py : 97%
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 einsum, triu_indices, zeros
2from scipy.special import comb
4from mango.constants import c, asciistring
5from mango.position import Communal, energy
6from mango.boundaries import periodic
7from mango.debug import debug
10@debug(['energy'])
11def conserv(e, pos, mom, mag, dist_matrix, dist, mass):
12 """Calculate the energy of the system and its CoM and CoM momentum."""
13 # TODO
14 # fix for proper array sizes (stats, multiple iterations)
15 invmass = 1 / mass
17 CoM = einsum("aij, i, i -> aj", pos, mass, invmass) / pos.shape[0]
19 CoM_vel = einsum("aij, i -> aj", mom, invmass)
21 einsum('xyz, aix, aiy -> az', c.eijk, pos - CoM[:, None], mom, out=e.angular) # May fail due to CoM dimensions differ from pos
23 einsum("aij -> aj", mag, out=e.tmag) # unit: 1.e-12 emu
25 e.tot[:] = e.angular + e.tmag / c.GBARE
27 ekin_m = einsum("aij -> ai", 0.5 * einsum("aiz, i -> aiz", einsum("aiz, aiz -> aiz", mom, mom), invmass))
29 ekin = einsum("ai -> a", ekin_m)
31 e.energy_vars(dist_matrix)
32 e.vv_trans()
34 epot1_m = e.epot_tr
35 epot2_m = e.vv_mag(mag, dist)
37 dims = ''.join(asciistring(epot1_m.ndim))
38 estr = '{} -> {}'.format(dims, dims[:-1])
40 epot1 = einsum(estr, epot1_m)
42 epot2 = einsum(estr, epot2_m)
44 etot = ekin + epot1 + epot2
46 return {"total_M": e.tot, "total_mag": e.tmag, "total_angular": e.angular,
47 "kinetic": ekin, "trans_pot": epot1, "mag_pot": epot2,
48 "total_E": etot, "CoM": CoM, "CoM_vel": CoM_vel,
49 "ekin_pm": ekin_m, 'trans_pot_pm': epot1_m, 'mag_pot_pm': epot2_m}
52def energy_calc(pos, mom, mag, epsilon, sigma, limit, mass):
53 """Direct energy calculation."""
54 nm = pos.shape[-2]
55 tri = triu_indices(n=nm, k=1)
56 cmb = (pos.shape[0], comb(nm, 2, exact=True))
57 dist, dist_matrix = periodic().setup_full(pos)
58 e = energy(Communal(epsilon, sigma, tri, cmb, limit), singleframe=False)
59 e.set_attr(['tot', 'tmag', 'angular'], zeros((pos.shape[0], 3)))
61 return conserv(e, pos, mom, mag, dist_matrix, dist, mass)
64def get_eandm(reader, var, name=None):
65 """Help function for getting energy and momenta data."""
66 return energy_calc(reader.read('position', name)[:var.skip_iters],
67 reader.read("momentum", name)[:var.skip_iters],
68 reader.read("magnetisation", name)[:var.skip_iters],
69 var.epsilon, var.sigma, var.limit, var.mass)