import numpy as np
[docs]def compute_sed(traj, ideal, prim, k_points):
""" Computes spectral energy density for a trajectory.
Parameters
----------
traj : list
trajectory with atoms objects with velocities
ideal : ASE atoms object
ideal atoms object
prim : ASE atoms object
compatible primitive cell. Must be aligned correctly
k_points : list
list of k points in cart coord (2pi must be included)
"""
velocities = []
for atom in traj:
velocities.append(atom.get_velocities())
velocities = np.array(velocities)
velocities = velocities.transpose(1, 2, 0).copy()
velocities = np.fft.fft(velocities, axis=2)
masses = prim.get_masses()
indices, offsets = _index_offset(ideal, prim)
pos = np.dot(k_points, np.dot(offsets, prim.cell).T)
exppos = np.exp(1.0j * pos)
density = np.zeros((len(k_points), velocities.shape[2]))
for alpha in range(3):
for b in range(len(masses)):
tmp = np.zeros(density.shape, dtype=np.complex)
for i in range(len(indices)):
index = indices[i]
if index != b:
continue
tmp += np.outer(exppos[:, i], velocities[i, alpha])
density += masses[b] * np.abs(tmp)**2
return density
def _index_offset(atoms, prim, atol=1e-3, rtol=0.0):
index, offset = [], []
for pos in atoms.positions:
spos = np.linalg.solve(prim.cell.T, pos)
for i, spos2 in enumerate(prim.get_scaled_positions()):
off = spos - spos2
off_round = np.round(off)
if not np.allclose(off, off_round, atol=atol, rtol=rtol):
continue
index.append(i)
off = off_round.astype(int)
assert np.allclose(off, off_round)
offset.append(off)
break
else:
raise ValueError('prim not compatible with atoms')
index, offset = np.array(index), np.array(offset)
return index, offset