import concurrent.futures
import functools
from itertools import combinations_with_replacement
from typing import Optional
import numba
import numpy as np
from ase import Atoms
from ase.units import fs
from numpy.typing import NDArray
from dynasor.logging_tools import logger
from dynasor.trajectory import Trajectory, WindowIterator
from dynasor.sample import DynamicSample, StaticSample
from dynasor.tools.acfs import psd_from_acf_2d
from dynasor.core.time_averager import TimeAverager, truncate_nans
from dynasor.core.reciprocal import calc_rho_q, calc_rho_j_q
from dynasor.qpoints.tools import is_qpoint_commensurate
from dynasor.tools.structures import get_offset_index
from dynasor.units import radians_per_fs_to_meV
def _validate_qpoints_commensurate(q_points: NDArray[float], cell: NDArray[float]) -> None:
"""Raise if any q-point is not commensurate with the simulation cell."""
n_bad = sum(not is_qpoint_commensurate(q, cell) for q in q_points)
if n_bad > 0:
raise ValueError(
f'{n_bad} of {len(q_points)} q-points are not commensurate with the cell; use '
'get_spherical_qpoints or get_supercell_qpoints_along_path to generate '
'commensurate q-points.')
[docs]
def compute_dynamic_structure_factors(
traj: Trajectory,
q_points: NDArray[float],
dt: float,
window_size: int,
window_step: Optional[int] = 1,
calculate_currents: Optional[bool] = False,
calculate_incoherent: Optional[bool] = False,
logging_interval: Optional[int] = 1000,
) -> DynamicSample:
r"""Compute the dynamic structure factors. The results are returned in the
form of a :class:`DynamicSample <dynasor.sample.DynamicSample>`
object.
Parameters
----------
traj
Input trajectory.
q_points
Array of q-points in units of rad/Å with shape ``(N_qpoints, 3)`` in Cartesian coordinates.
dt
Time difference in femtoseconds between two consecutive snapshots
in the trajectory. Note that you should *not* change :attr:`dt` if you change
:attr:`frame_step <dynasor.trajectory.Trajectory.frame_step>` in :attr:`traj`.
window_size
Maximum time lag, expressed as a number of frame intervals, for which to calculate
correlations. Each calculation window therefore contains ``window_size + 1`` frames,
including the frame at time lag zero. This parameter determines the smallest frequency
resolved.
window_step
Window step (or stride) given as the number of frames between consecutive trajectory
windows. This parameter does *not* affect the time between consecutive frames in the
calculation. If :attr:`window_step` > :attr:`window_size` + 1, some frames will not be
used.
calculate_currents
Calculate the current correlations. Requires velocities to be available in :attr:`traj`.
calculate_incoherent
Calculate the incoherent part (self-part) of :math:`F_\text{incoh}`.
logging_interval
Log progress at ``INFO`` level every this many windows. Set to ``0`` to disable
progress logging.
"""
# sanity check input args
if q_points.ndim != 2 or q_points.shape[1] != 3:
raise ValueError('q-points array has the wrong shape.')
if dt <= 0:
raise ValueError(f'dt must be positive: dt= {dt}')
if window_size <= 2:
raise ValueError(f'window_size must be larger than 2: window_size= {window_size}')
if window_step <= 0:
raise ValueError(f'window_step must be positive: window_step= {window_step}')
if calculate_currents and not traj.has_velocities:
raise ValueError('calculate_currents=True requires velocities to be available in the '
'trajectory, but traj does not provide velocities.')
_validate_qpoints_commensurate(q_points, traj.cell)
# define internal parameters
n_qpoints = q_points.shape[0]
delta_t = traj.frame_step * dt
N_tc = window_size + 1
# log all setup information
n_fft = 2 * window_size + 1
dw = 2 * np.pi / (n_fft * delta_t)
w_max = window_size * dw
w_N = np.pi / delta_t
dw_mev = dw * radians_per_fs_to_meV
w_max_mev = w_max * radians_per_fs_to_meV
w_N_mev = w_N * radians_per_fs_to_meV
logger.info(f'Spacing between samples (frame_step): {traj.frame_step}')
logger.info(f'Time between consecutive frames in input trajectory (dt): {dt} fs')
logger.info(f'Time between consecutive frames used (dt * frame_step): {delta_t} fs')
logger.info(f'Time window size (dt * frame_step * window_size): {delta_t * window_size:.1f} fs')
logger.info(f'Angular frequency resolution: dw = {dw:.6f} rad/fs = {dw_mev:.3f} meV')
logger.info(f'Maximum angular frequency: {w_max:.6f} rad/fs = {w_max_mev:.3f} meV '
f'(Nyquist limit: {w_N:.6f} rad/fs = {w_N_mev:.3f} meV)')
if calculate_currents:
logger.info('Calculating current (velocity) correlations')
if calculate_incoherent:
logger.info('Calculating incoherent part (self-part) of correlations')
# log some info regarding q-points
logger.info(f'Number of q-points: {n_qpoints}')
q_directions = q_points.astype(float, copy=True)
q_distances = np.linalg.norm(q_points, axis=1)
nonzero = q_distances > 0
q_directions[nonzero] /= q_distances[nonzero].reshape(-1, 1)
# setup functions to process frames
def f2_rho(frame):
rho_qs_dict = dict()
for atom_type in frame.positions_by_type.keys():
x = frame.positions_by_type[atom_type]
rho_qs_dict[atom_type] = calc_rho_q(x, q_points)
frame.rho_qs_dict = rho_qs_dict
return frame
def f2_rho_and_j(frame):
rho_qs_dict = dict()
jz_qs_dict = dict()
jper_qs_dict = dict()
for atom_type in frame.positions_by_type.keys():
x = frame.positions_by_type[atom_type]
v = frame.velocities_by_type[atom_type]
rho_qs, j_qs = calc_rho_j_q(x, v, q_points)
jz_qs = np.sum(j_qs * q_directions, axis=1)
jper_qs = j_qs - (jz_qs[:, None] * q_directions)
rho_qs_dict[atom_type] = rho_qs
jz_qs_dict[atom_type] = jz_qs
jper_qs_dict[atom_type] = jper_qs
frame.rho_qs_dict = rho_qs_dict
frame.jz_qs_dict = jz_qs_dict
frame.jper_qs_dict = jper_qs_dict
return frame
if calculate_currents:
element_processor = f2_rho_and_j
else:
element_processor = f2_rho
# setup window iterator
window_iterator = WindowIterator(traj, width=N_tc, window_step=window_step,
element_processor=element_processor)
# define all atom types and pairs
atom_types = traj.atom_types
pairs = list(combinations_with_replacement(atom_types, r=2))
particle_counts = {key: len(val) for key, val in traj.atomic_indices.items()}
logger.debug('Considering pairs:')
for pair in pairs:
logger.debug(f' {pair}')
# set up all time averager instances
F_q_t_averager = dict()
for pair in pairs:
F_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
if calculate_currents:
Cl_q_t_averager = dict()
Ct_q_t_averager = dict()
for pair in pairs:
Cl_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
Ct_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
if calculate_incoherent:
F_s_q_t_averager = dict()
for pair in atom_types:
F_s_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
# define correlation function
#
# Note: calc_corr is dispatched concurrently across OS threads via
# ThreadPoolExecutor.map() below, so it must never call into a numba
# parallel=True/prange kernel (e.g. calc_rho_q, calc_rho_j_q). Numba's
# default threading layer is not safe against concurrent entry from
# multiple Python threads and aborts the process if this happens. Such
# calls belong in calc_incoherent instead, which is always run strictly
# sequentially in the main thread.
def calc_corr(window, time_i):
# Calculate correlations between two frames in the window without normalization 1/N
f0 = window[0]
fi = window[time_i]
for s1, s2 in pairs:
Fqt = np.real(f0.rho_qs_dict[s1] * fi.rho_qs_dict[s2].conjugate())
if s1 != s2:
Fqt += np.real(f0.rho_qs_dict[s2] * fi.rho_qs_dict[s1].conjugate())
F_q_t_averager[(s1, s2)].add_sample(time_i, Fqt)
if calculate_currents:
for s1, s2 in pairs:
Clqt = np.real(f0.jz_qs_dict[s1] * fi.jz_qs_dict[s2].conjugate())
Ctqt = 0.5 * np.real(np.sum(f0.jper_qs_dict[s1] *
fi.jper_qs_dict[s2].conjugate(), axis=1))
if s1 != s2:
Clqt += np.real(f0.jz_qs_dict[s2] * fi.jz_qs_dict[s1].conjugate())
Ctqt += 0.5 * np.real(np.sum(f0.jper_qs_dict[s2] *
fi.jper_qs_dict[s1].conjugate(), axis=1))
Cl_q_t_averager[(s1, s2)].add_sample(time_i, Clqt)
Ct_q_t_averager[(s1, s2)].add_sample(time_i, Ctqt)
def calc_incoherent(window, time_i):
# Calculate the incoherent (self) part between two frames in the window.
#
# calc_rho_q is backed by a numba parallel=True/prange kernel, which must
# not be entered concurrently from multiple Python threads (see the note
# on calc_corr above). This function is therefore always called from a
# plain sequential loop in the main thread, never from within the
# ThreadPoolExecutor used for calc_corr.
f0 = window[0]
fi = window[time_i]
for atom_type in atom_types:
xi = fi.positions_by_type[atom_type]
x0 = f0.positions_by_type[atom_type]
Fsqt = np.real(calc_rho_q(xi - x0, q_points))
F_s_q_t_averager[atom_type].add_sample(time_i, Fsqt)
# run calculation
with concurrent.futures.ThreadPoolExecutor() as tpe:
# This is the "main loop" over the trajectory
for window in window_iterator:
if logging_interval and window[0].frame_index % logging_interval == 0:
logger.info(f'Processing window {window[0].frame_index} to {window[-1].frame_index}') # noqa
else:
logger.debug(f'Processing window {window[0].frame_index} to {window[-1].frame_index}') # noqa
# The map conveniently applies calc_corr to all time-lags. However,
# as everything is done in place nothing gets returned so in order
# to start and wait for the processes to finish we must iterate
# over the None values returned
for _ in tpe.map(functools.partial(calc_corr, window), range(len(window))):
pass
# Run the incoherent part strictly sequentially in the main thread
# (see calc_incoherent for why it must not run inside the thread pool).
if calculate_incoherent:
for time_i in range(len(window)):
calc_incoherent(window, time_i)
# collect results into dict with numpy arrays (n_qpoints, N_tc)
data_dict_corr = dict()
data_dict_corr['q_points'] = q_points
time = None
for pair in pairs:
key = '_'.join(pair)
F_q_t = 1 / traj.n_atoms * truncate_nans(F_q_t_averager[pair].get_average_all())
w, S_q_w = psd_from_acf_2d(F_q_t, delta_t)
S_q_w = np.array(S_q_w)
if time is None:
# Determine the actual length of time signal after truncation (if needed)
time = delta_t * np.arange(F_q_t.shape[1], dtype=float)
N_tc_actual = len(time)
F_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
S_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
else:
assert F_q_t.shape[1] == len(time)
data_dict_corr['omega'] = w
data_dict_corr[f'Fqt_coh_{key}'] = F_q_t
data_dict_corr[f'Sqw_coh_{key}'] = S_q_w
# sum all partials to the total
F_q_t_tot += F_q_t
S_q_w_tot += S_q_w
if N_tc_actual < N_tc:
logger.warning('Truncating ACF due to NaNs, likely time_window is longer than length of Trajectory') # noqa
dw = float(w[1] - w[0])
w_max = float(w[-1])
data_dict_corr['time'] = time
data_dict_corr['Fqt_coh'] = F_q_t_tot
data_dict_corr['Sqw_coh'] = S_q_w_tot
if calculate_currents:
Cl_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
Ct_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
Cl_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
Ct_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
for pair in pairs:
key = '_'.join(pair)
Cl_q_t = 1 / traj.n_atoms * truncate_nans(Cl_q_t_averager[pair].get_average_all())
Ct_q_t = 1 / traj.n_atoms * truncate_nans(Ct_q_t_averager[pair].get_average_all())
_, Cl_q_w = psd_from_acf_2d(Cl_q_t, delta_t)
_, Ct_q_w = psd_from_acf_2d(Ct_q_t, delta_t)
data_dict_corr[f'Clqt_{key}'] = Cl_q_t
data_dict_corr[f'Ctqt_{key}'] = Ct_q_t
data_dict_corr[f'Clqw_{key}'] = Cl_q_w
data_dict_corr[f'Ctqw_{key}'] = Ct_q_w
# sum all partials to the total
Cl_q_t_tot += Cl_q_t
Ct_q_t_tot += Ct_q_t
Cl_q_w_tot += Cl_q_w
Ct_q_w_tot += Ct_q_w
data_dict_corr['Clqt'] = Cl_q_t_tot
data_dict_corr['Ctqt'] = Ct_q_t_tot
data_dict_corr['Clqw'] = Cl_q_w_tot
data_dict_corr['Ctqw'] = Ct_q_w_tot
if calculate_incoherent:
Fs_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
Ss_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
for atom_type in atom_types:
Fs_q_t = 1 / traj.n_atoms * truncate_nans(F_s_q_t_averager[atom_type].get_average_all())
_, Ss_q_w = psd_from_acf_2d(Fs_q_t, delta_t)
data_dict_corr[f'Fqt_incoh_{atom_type}'] = Fs_q_t
data_dict_corr[f'Sqw_incoh_{atom_type}'] = Ss_q_w
# sum all partials to the total
Fs_q_t_tot += Fs_q_t
Ss_q_w_tot += Ss_q_w
data_dict_corr['Fqt_incoh'] = Fs_q_t_tot
data_dict_corr['Sqw_incoh'] = Ss_q_w_tot
# finalize results with additional metadata
new_sample = DynamicSample(
data_dict_corr,
simulation_data=dict(
atom_types=atom_types, pairs=pairs,
particle_counts=particle_counts,
cell=traj.cell,
time_between_frames=delta_t,
maximum_time_lag=float(time[-1]),
angular_frequency_resolution=dw,
maximum_angular_frequency=w_max,
number_of_frames=traj.number_of_frames_read,
))
new_sample._append_history(
'compute_dynamic_structure_factors',
dict(
dt=dt,
window_size=window_size,
window_step=window_step,
calculate_currents=calculate_currents,
calculate_incoherent=calculate_incoherent,
))
return new_sample
[docs]
def compute_static_structure_factors(
traj: Trajectory,
q_points: NDArray[float],
logging_interval: Optional[int] = 1000,
) -> StaticSample:
r"""Compute the static structure factors. The results are returned in the
form of a :class:`StaticSample <dynasor.sample.StaticSample>`
object.
Parameters
----------
traj
Input trajectory.
q_points
Array of q-points in units of rad/Å with shape ``(N_qpoints, 3)`` in Cartesian coordinates.
logging_interval
Log progress at ``INFO`` level every this many frames. Set to ``0`` to disable
progress logging.
"""
# sanity check input args
if q_points.ndim != 2 or q_points.shape[1] != 3:
raise ValueError('q-points array has the wrong shape.')
_validate_qpoints_commensurate(q_points, traj.cell)
n_qpoints = q_points.shape[0]
logger.info(f'Number of q-points: {n_qpoints}')
# define all pairs
pairs = list(combinations_with_replacement(traj.atom_types, r=2))
particle_counts = {key: len(val) for key, val in traj.atomic_indices.items()}
logger.debug('Considering pairs:')
for pair in pairs:
logger.debug(f' {pair}')
# processing function
def f2_rho(frame):
rho_qs_dict = dict()
for atom_type in frame.positions_by_type.keys():
x = frame.positions_by_type[atom_type]
rho_qs_dict[atom_type] = calc_rho_q(x, q_points)
frame.rho_qs_dict = rho_qs_dict
return frame
# setup averager
Sq_averager = dict()
for pair in pairs:
Sq_averager[pair] = TimeAverager(1, n_qpoints) # time average with only timelag=0
# main loop
for frame in traj:
# process_frame
f2_rho(frame)
if logging_interval and frame.frame_index % logging_interval == 0:
logger.info(f'Processing frame {frame.frame_index}')
else:
logger.debug(f'Processing frame {frame.frame_index}')
for s1, s2 in pairs:
# compute correlation
Sq_pair = np.real(frame.rho_qs_dict[s1] * frame.rho_qs_dict[s2].conjugate())
if s1 != s2:
Sq_pair += np.real(frame.rho_qs_dict[s2] * frame.rho_qs_dict[s1].conjugate())
Sq_averager[(s1, s2)].add_sample(0, Sq_pair)
# collect results
data_dict = dict()
data_dict['q_points'] = q_points
S_q_tot = np.zeros((n_qpoints, 1))
for s1, s2 in pairs:
Sq = 1 / traj.n_atoms * Sq_averager[(s1, s2)].get_average_at_timelag(0).reshape(-1, 1)
data_dict[f'Sq_{s1}_{s2}'] = Sq
S_q_tot += Sq
data_dict['Sq'] = S_q_tot
# finalize results
new_sample = StaticSample(
data_dict,
simulation_data=dict(
atom_types=traj.atom_types,
pairs=pairs,
particle_counts=particle_counts,
cell=traj.cell,
number_of_frames=traj.number_of_frames_read,
))
new_sample._append_history('compute_static_structure_factors')
return new_sample
[docs]
def compute_spectral_energy_density(
traj: Trajectory,
ideal_supercell: Atoms,
primitive_cell: Atoms,
q_points: NDArray[float],
dt: float,
partial: Optional[bool] = False,
logging_interval: Optional[int] = 1000,
) -> tuple[NDArray[float], NDArray[float]]:
r"""
Compute the spectral energy density (SED) at specific q-points. The results
are returned in the form of a tuple, which comprises the angular
frequencies in an array of length ``N_times`` in units of rad/fs and the
SED in units of eV/(rad/fs) as an array of shape ``(N_qpoints, N_times)``.
The normalization is chosen such that integrating the SED of a q-point
together with the supplied angular frequencies omega (rad/fs) yields
`1/2 kB T` * number of bands (where number of bands = `len(prim) * 3`)
More details can be found in Thomas *et al.*, Physical Review B **81**, 081411 (2010),
which should be cited when using this function along with the dynasor reference.
**Note 1:**
SED analysis is only suitable for crystalline materials without diffusion as
atoms are assumed to move around fixed reference positions throughout the entire trajectory.
**Note 2:**
This implementation reads the full trajectory and can thus consume a lot of memory.
Parameters
----------
traj
Input trajectory.
ideal_supercell
Ideal structure defining the reference positions. Do not change the masses
in the ASE :class:`Atoms` objects to dynasor internal units, this will be
done internally. Its atom count is checked against :attr:`traj`; a mismatched
cell only triggers a warning (some thermal expansion relative to the reference
structure is normal), and the atom ordering is not checked at all, so it must
match the trajectory's.
primitive_cell
Underlying primitive structure. Must be aligned correctly with :attr:`ideal_supercell`.
q_points
Array of q-points in units of rad/Å with shape ``(N_qpoints, 3)`` in Cartesian coordinates.
dt
Time difference in femtoseconds between two consecutive snapshots in
the trajectory. Note that you should not change :attr:`dt` if you change
:attr:`frame_step <dynasor.trajectory.Trajectory.frame_step>` in :attr:`traj`.
partial
If True the SED will be returned decomposed per basis and Cartesian direction.
The shape is ``(N_qpoints, N_frequencies, len(primitive_cell), 3)``.
logging_interval
Log progress at ``INFO`` level every this many frames. Set to ``0`` to disable
progress logging.
"""
if q_points.ndim != 2 or q_points.shape[1] != 3:
raise ValueError('q-points array has the wrong shape.')
if dt <= 0:
raise ValueError(f'dt must be positive: dt= {dt}')
delta_t = traj.frame_step * dt
# logger
logger.info('Running SED')
logger.info(f'Time between consecutive frames (dt * frame_step): {delta_t} fs')
logger.info(f'Number of atoms in primitive_cell: {len(primitive_cell)}')
logger.info(f'Number of atoms in ideal_supercell: {len(ideal_supercell)}')
logger.info(f'Number of q-points: {q_points.shape[0]}')
# check that the ideal supercell agrees with traj
if traj.n_atoms != len(ideal_supercell):
raise ValueError('ideal_supercell must contain the same number of atoms as the trajectory.')
if not np.allclose(traj.cell, ideal_supercell.cell, atol=1e-5, rtol=0.0):
logger.warning('ideal_supercell cell does not match the trajectory cell.')
if len(primitive_cell) > len(ideal_supercell):
raise ValueError('primitive_cell contains more atoms than ideal_supercell.')
if not traj.has_velocities:
raise ValueError('compute_spectral_energy_density requires velocities to be available '
'in the trajectory, but traj does not provide velocities.')
# q-points must be commensurate with the ideal supercell (which defines the
# phase factors below), not with the possibly thermally-expanded traj cell
_validate_qpoints_commensurate(q_points, ideal_supercell.cell)
# collect all velocities, and scale with sqrt(masses)
masses = ideal_supercell.get_masses().reshape(-1, 1) / fs**2 # From Dalton to dmu
velocities = []
for it, frame in enumerate(traj):
if logging_interval and it % logging_interval == 0:
logger.info(f'Reading frame {it}')
else:
logger.debug(f'Reading frame {it}')
v = frame.get_velocities_as_array(traj.atomic_indices) # in Å/fs
velocities.append(np.sqrt(masses) * v)
logger.info(f'Number of snapshots: {len(velocities)}')
# Perform the FFT on the last axis for extra speed (maybe not needed)
N_samples = len(velocities)
velocities = np.array(velocities)
# places time index last and makes a copy for continuity
velocities = velocities.transpose(1, 2, 0).copy()
# #atoms in supercell x 3 directions x #frequencies
velocities = np.fft.rfft(velocities, axis=2)
# Calculate indices and offsets needed for the SED method
offsets, indices = get_offset_index(primitive_cell, ideal_supercell)
# Phase factor for use in FT. #qpoints x #atoms in supercell
cell_positions = np.dot(offsets, primitive_cell.cell)
phase = np.dot(q_points, cell_positions.T) # #qpoints x #unit cells
phase_factors = np.exp(1.0j * phase)
# This dict maps the offsets to an index so ndarrays can be over
# offset,index instead of atoms in supercell
offset_dict = {off: n for n, off in enumerate(set(tuple(offset) for offset in offsets))}
# Pick out some shapes
n_super, _, n_w = velocities.shape
n_qpts = len(q_points)
n_prim = len(primitive_cell)
n_offsets = len(offset_dict)
# This new array will be indexed by index and offset instead (and also transposed)
new_velocities = np.zeros((n_w, 3, n_prim, n_offsets), dtype=velocities.dtype)
for i in range(n_super):
j = indices[i] # atom with index i in the supercell is of basis type j ...
n = offset_dict[tuple(offsets[i])] # and its offset has index n
new_velocities[:, :, j, n] = velocities[i].T
velocities = new_velocities
# Same story with the spatial phase factors
new_phase_factors = np.zeros((n_qpts, n_prim, n_offsets), dtype=phase_factors.dtype)
for i in range(n_super):
j = indices[i]
n = offset_dict[tuple(offsets[i])]
new_phase_factors[:, j, n] = phase_factors[:, i]
phase_factors = new_phase_factors
# calculate the density in a numba function
density = _sed_inner_loop(phase_factors, velocities)
if not partial:
density = np.sum(density, axis=(2, 3))
# units
# make so that the velocities were originally in Angstrom / fs to be compatible with eV and Da
# the time delta in the fourier transform
density = density * delta_t**2
# Divide by the length of the time signal
density = density / (N_samples * delta_t)
# Divide by the number of primitive cells
density = density / (n_super / n_prim)
# Factor so the sed can be integrated together with the returned omega
# numpy fft works with ordinary/linear frequencies and not angular freqs
density = density / (2*np.pi)
# angular frequencies
w = 2 * np.pi * np.fft.rfftfreq(N_samples, delta_t) # rad/fs
return w, density
@numba.njit(parallel=True, fastmath=True)
def _sed_inner_loop(phase_factors, velocities):
"""This numba function calculates the spatial
Fourier transform using precomputed phase factors.
As the use case can be one or many q-points the parallelization is over the
temporal frequency components instead.
"""
n_qpts = phase_factors.shape[0] # q-point index
n_prim = phase_factors.shape[1] # basis atom index
n_super = phase_factors.shape[2] # unit cell index
n_freqs = velocities.shape[0] # frequency, direction, basis atom, unit cell
density = np.zeros((n_qpts, n_freqs, n_prim, 3), dtype=np.float64)
for w in numba.prange(n_freqs):
for k in range(n_qpts):
for a in range(3):
for b in range(n_prim):
tmp = 0.0j
for n in range(n_super):
tmp += phase_factors[k, b, n] * velocities[w, a, b, n]
density[k, w, b, a] += np.abs(tmp)**2
return density
