import concurrent
from functools import partial
from itertools import combinations_with_replacement
from typing import Tuple
import numpy as np
from ase import Atoms
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.post_processing import fourier_cos
from dynasor.core.time_averager import TimeAverager
from dynasor.core.reciprocal import calc_rho_q, calc_rho_j_q
from dynasor.qpoints.tools import get_index_offset
[docs]def compute_dynamic_structure_factors(
traj: Trajectory,
q_points: NDArray[float],
dt: float,
window_size: int,
window_step: int = 1,
calculate_currents: bool = False,
calculate_incoherent: bool = False,
) -> DynamicSample:
"""Compute 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 2π/Å 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
Length of the trajectory frame window to use for time correlation calculation.
It is expressed in terms of the number of time lags to consider
and thus 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, e.g., :attr:`window_step` > :attr:`window_size`, 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_incoh`.
"""
# sanity check input args
if 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_size % 2 != 0:
raise ValueError(f'window_size must be even: window_size= {window_size}')
if window_step <= 0:
raise ValueError(f'window_step must be positive: window_step= {window_step}')
# define internal parameters
n_qpoints = q_points.shape[0]
delta_t = traj.frame_step * dt
N_tc = window_size + 1
# log all setup information
dw = 2 * np.pi / (window_size * delta_t)
w_max = dw * window_size
f_N = 1 / (2 * delta_t) # Nyquist frequency
conv = 658.2119 # conversion from 2pi/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} 2pi/fs = {dw * conv:.3f} meV')
logger.info(f'Maximum angular frequency (dw * window_size):'
f' {w_max:.6f} 2pi/fs = {w_max * conv:.3f} meV')
logger.info(f'Nyquist angular frequency (2pi / frame_step / dt / 2):'
f' {f_N * 2 * np.pi:.3f} 2pi/fs = {f_N * 2 * np.pi * conv:.1f} 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.copy()
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 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}')
# setup 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 traj.atom_types:
F_s_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
# define correlation function
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)
if calculate_incoherent:
for atom_type in traj.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:
logger.debug(f'processing window {window[0].frame_index} to {window[-1].frame_index}')
# The map conviniently 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(partial(calc_corr, window), range(len(window))):
pass
# collect results into dict with numpy arrays (n_qpoints, N_tc)
data_dict_corr = dict()
time = delta_t * np.arange(N_tc, dtype=float)
data_dict_corr['q_points'] = q_points
data_dict_corr['time'] = time
F_q_t_tot = np.zeros((n_qpoints, N_tc))
S_q_w_tot = np.zeros((n_qpoints, N_tc))
for pair in pairs:
key = '_'.join(pair)
F_q_t = 1 / traj.n_atoms * F_q_t_averager[pair].get_average_all()
w, S_q_w = zip(*[fourier_cos(F, delta_t) for F in F_q_t])
w = w[0]
S_q_w = np.array(S_q_w)
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
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))
Ct_q_t_tot = np.zeros((n_qpoints, N_tc))
Cl_q_w_tot = np.zeros((n_qpoints, N_tc))
Ct_q_w_tot = np.zeros((n_qpoints, N_tc))
for pair in pairs:
key = '_'.join(pair)
Cl_q_t = 1 / traj.n_atoms * Cl_q_t_averager[pair].get_average_all()
Ct_q_t = 1 / traj.n_atoms * Ct_q_t_averager[pair].get_average_all()
Cl_q_w = np.array([fourier_cos(C, delta_t)[1] for C in Cl_q_t])
Ct_q_w = np.array([fourier_cos(C, delta_t)[1] for C in Ct_q_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))
Ss_q_w_tot = np.zeros((n_qpoints, N_tc))
for atom_type in traj.atom_types:
Fs_q_t = 1 / traj.n_atoms * F_s_q_t_averager[atom_type].get_average_all()
Ss_q_w = np.array([fourier_cos(F, delta_t)[1] for F in Fs_q_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
data_dict_corr['Fqt'] = data_dict_corr['Fqt_coh'] + data_dict_corr['Fqt_incoh']
data_dict_corr['Sqw'] = data_dict_corr['Sqw_coh'] + data_dict_corr['Sqw_incoh']
else:
data_dict_corr['Fqt'] = data_dict_corr['Fqt_coh'].copy()
data_dict_corr['Sqw'] = data_dict_corr['Sqw_coh'].copy()
# finalize results with additional meta dat
result = DynamicSample(data_dict_corr, atom_types=traj.atom_types, pairs=pairs,
particle_counts=particle_counts, cell=traj.cell,
nyquist_frequency=f_N, max_frequency=w_max, frequency_resolution=dw)
return result
[docs]def compute_static_structure_factors(
traj: Trajectory,
q_points: NDArray[float],
) -> StaticSample:
r"""Compute 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 2π/Å with shape ``(N_qpoints, 3)`` in Cartesian coordinates
"""
# sanity check input args
if q_points.shape[1] != 3:
raise ValueError('q-points array has the wrong shape.')
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)
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
result = StaticSample(data_dict, atom_types=traj.atom_types, pairs=pairs,
particle_counts=particle_counts, cell=traj.cell)
return result
[docs]def compute_spectral_energy_density(
traj: Trajectory,
ideal_supercell: Atoms,
primitive_cell: Atoms,
q_points: NDArray[float],
dt: float,
) -> 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 2π/fs
and the SED in units of Da*(Å/fs)² as an array of shape ``(N_qpoints, N_times)``.
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
primitive_cell
Underlying primitive structure. Must be aligned correctly with :attr:`ideal_supercell`.
q_points
Array of q-points in units of 2π/Å 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`.
"""
delta_t = traj.frame_step * dt
# logger
logger.info('Running SED')
logger.info(f'Time between consecutive frames (dt * frame_step): {delta_t} fs')
# 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.')
# colllect all velocities
velocities = []
for it, frame in enumerate(traj):
logger.debug(f'Reading frame {it}')
if frame.velocities_by_type is None:
raise ValueError(f'Could not read velocities from frame {it}')
v = frame.get_velocities_as_array(traj.atomic_indices)
velocities.append(v)
velocities = np.array(velocities)
velocities = velocities.transpose(1, 2, 0).copy()
velocities = np.fft.fft(velocities, axis=2)
# calculate SED
masses = primitive_cell.get_masses()
indices, offsets = get_index_offset(ideal_supercell, primitive_cell)
pos = np.dot(q_points, np.dot(offsets, primitive_cell.cell).T)
exppos = np.exp(1.0j * pos)
density = np.zeros((len(q_points), velocities.shape[2]))
for alpha in range(3):
for b in range(len(masses)):
tmp = np.zeros(density.shape, dtype=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
# frequencies
w = np.linspace(0.0, 2 * np.pi / delta_t, density.shape[1]) # units of 2pi/fs
return w, density