Coverage for dynasor/correlation_functions.py: 98%
307 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-03 19:46 +0000
1from itertools import combinations_with_replacement
2from typing import Optional
4import numba
5import numpy as np
6from ase import Atoms
7from ase.units import fs
8from numpy.typing import NDArray
10from dynasor.compute_config import get_compute_config
11from dynasor.logging_tools import logger, warn_if_numba_threading_layer_is_slow
12from dynasor.trajectory import Trajectory, WindowIterator
13from dynasor.trajectory.prefetch import (_drop_frames_outside_windows,
14 _prefetch_batched_iter)
15from dynasor.sample import DynamicSample, StaticSample
17from dynasor.tools.acfs import psd_from_acf_2d
18from dynasor.core.time_averager import TimeAverager, truncate_nans
19from dynasor.core.reciprocal import _calc_rho_q_incoherent
20from dynasor.core.frame_processors import get_frame_processor
21from dynasor.qpoints.tools import is_qpoint_commensurate
22from dynasor.tools.structures import get_offset_index
23from dynasor.units import radians_per_fs_to_meV
26def _validate_qpoints_commensurate(q_points: NDArray[float], cell: NDArray[float]) -> None:
27 """Raise if any q-point is not commensurate with the simulation cell."""
28 n_bad = sum(not is_qpoint_commensurate(q, cell) for q in q_points)
29 if n_bad > 0:
30 raise ValueError(
31 f'{n_bad} of {len(q_points)} q-points are not commensurate with the cell; use '
32 'get_spherical_qpoints or get_supercell_qpoints_along_path to generate '
33 'commensurate q-points.')
36def compute_dynamic_structure_factors(
37 traj: Trajectory,
38 q_points: NDArray[float],
39 dt: float,
40 window_size: int,
41 window_step: Optional[int] = 1,
42 calculate_currents: Optional[bool] = False,
43 calculate_incoherent: Optional[bool] = False,
44 logging_interval: Optional[int] = 1000,
45) -> DynamicSample:
46 r"""Compute the dynamic structure factors. The results are returned in the
47 form of a :class:`DynamicSample <dynasor.sample.DynamicSample>`
48 object.
50 The computational backend and precision are set globally via
51 :func:`set_compute_config <dynasor.set_compute_config>`, not as parameters of this
52 function; see :ref:`the backends reference page <backends>` for details, including
53 which parts of this calculation the GPU backends do (and do not) accelerate.
55 Parameters
56 ----------
57 traj
58 Input trajectory.
59 q_points
60 Array of q-points in units of rad/Å with shape ``(N_qpoints, 3)`` in Cartesian coordinates.
61 dt
62 Time difference in femtoseconds between two consecutive snapshots
63 in the trajectory. Note that you should *not* change :attr:`dt` if you change
64 :attr:`frame_step <dynasor.Trajectory.frame_step>` in :attr:`traj`.
65 window_size
66 Maximum time lag, expressed as a number of frame intervals, for which to calculate
67 correlations. Each calculation window therefore contains ``window_size + 1`` frames,
68 including the frame at time lag zero. This parameter determines the smallest frequency
69 resolved.
70 window_step
71 Window step (or stride) given as the number of frames between consecutive trajectory
72 windows. This parameter does *not* affect the time between consecutive frames in the
73 calculation. If :attr:`window_step` > :attr:`window_size` + 1, some frames will not be
74 used.
75 calculate_currents
76 Calculate the current correlations. Requires velocities to be available in :attr:`traj`.
77 calculate_incoherent
78 Calculate the incoherent part (self-part) of :math:`F_\text{incoh}`. Always evaluated on
79 CPU via the Numba kernel, regardless of the configured backend (see
80 :func:`set_compute_config <dynasor.set_compute_config>`).
81 logging_interval
82 Log progress at ``INFO`` level every this many windows. Set to ``0`` to disable
83 progress logging.
84 """
85 # sanity check input args
86 if q_points.ndim != 2 or q_points.shape[1] != 3:
87 raise ValueError('q-points array has the wrong shape.')
88 if dt <= 0:
89 raise ValueError(f'dt must be positive: dt= {dt}')
90 if window_size <= 2:
91 raise ValueError(f'window_size must be larger than 2: window_size= {window_size}')
92 if window_step <= 0:
93 raise ValueError(f'window_step must be positive: window_step= {window_step}')
94 if not traj.has_positions:
95 raise ValueError('compute_dynamic_structure_factors requires positions to be available '
96 'in the trajectory, but traj does not provide positions.')
97 if calculate_currents and not traj.has_velocities:
98 raise ValueError('calculate_currents=True requires velocities to be available in the '
99 'trajectory, but traj does not provide velocities.')
100 _validate_qpoints_commensurate(q_points, traj.cell)
102 # define internal parameters
103 n_qpoints = q_points.shape[0]
104 delta_t = traj.frame_step * dt
105 N_tc = window_size + 1
107 # log all setup information
108 n_fft = 2 * window_size + 1
109 dw = 2 * np.pi / (n_fft * delta_t)
110 w_max = window_size * dw
111 w_N = np.pi / delta_t
112 dw_mev = dw * radians_per_fs_to_meV
113 w_max_mev = w_max * radians_per_fs_to_meV
114 w_N_mev = w_N * radians_per_fs_to_meV
115 logger.info(f'Spacing between samples (frame_step): {traj.frame_step}')
116 logger.info(f'Time between consecutive frames in input trajectory (dt): {dt} fs')
117 logger.info(f'Time between consecutive frames used (dt * frame_step): {delta_t} fs')
118 logger.info(f'Time window size (dt * frame_step * window_size): {delta_t * window_size:.1f} fs')
119 logger.info(f'Angular frequency resolution: dw = {dw:.6f} rad/fs = {dw_mev:.3f} meV')
120 logger.info(f'Maximum angular frequency: {w_max:.6f} rad/fs = {w_max_mev:.3f} meV '
121 f'(Nyquist limit: {w_N:.6f} rad/fs = {w_N_mev:.3f} meV)')
123 if calculate_currents:
124 logger.info('Calculating current (velocity) correlations')
125 if calculate_incoherent:
126 logger.info('Calculating incoherent part (self-part) of correlations')
128 # log some info regarding q-points
129 logger.info(f'Number of q-points: {n_qpoints}')
131 q_directions = q_points.astype(float, copy=True)
132 q_distances = np.linalg.norm(q_points, axis=1)
133 nonzero = q_distances > 0
134 q_directions[nonzero] /= q_distances[nonzero].reshape(-1, 1)
136 # setup functions to process frames
137 compute_config = get_compute_config()
138 backend = compute_config.backend
139 precision = compute_config.precision
140 gpu_batch_size = compute_config.batch_size
142 is_batched, processor = get_frame_processor(
143 q_points, backend, precision, gpu_batch_size,
144 calculate_currents=calculate_currents, q_directions=q_directions)
146 # Build the frame source and window iterator.
147 # GPU: batching + background prefetch thread (processor reads B frames,
148 # GPU-processes them, and queues results while the main thread runs
149 # calc_corr on the previous window).
150 # Numba: standard per-frame element_processor inside WindowIterator.
151 if is_batched: 151 ↛ 156line 151 didn't jump to line 156 because the condition on line 151 was never true
152 # The batched processor runs ahead of the window iterator, so frames
153 # that fall between two windows have to be dropped here; skipping them
154 # in the window iterator would leave them already transformed. What
155 # remains is one group of N_tc frames per window, hence the step.
156 if window_step >= N_tc:
157 raw_frames = _drop_frames_outside_windows(traj, N_tc, window_step)
158 iterator_step = N_tc
159 else:
160 raw_frames = traj
161 iterator_step = window_step
162 _frame_source = _prefetch_batched_iter(raw_frames, processor, batch_size=gpu_batch_size)
163 window_iterator = WindowIterator(_frame_source, width=N_tc, window_step=iterator_step)
164 else:
165 window_iterator = WindowIterator(traj, width=N_tc, window_step=window_step,
166 element_processor=processor)
168 # define all atom types and pairs
169 atom_types = traj.atom_types
170 pairs = list(combinations_with_replacement(atom_types, r=2))
171 particle_counts = {key: len(val) for key, val in traj.atomic_indices.items()}
172 logger.debug('Considering pairs:')
173 for pair in pairs:
174 logger.debug(f' {pair}')
176 # set up all time averager instances
177 F_q_t_averager = dict()
178 for pair in pairs:
179 F_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
180 if calculate_currents:
181 Cl_q_t_averager = dict()
182 Ct_q_t_averager = dict()
183 for pair in pairs:
184 Cl_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
185 Ct_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
186 if calculate_incoherent:
187 F_s_q_t_averager = dict()
188 for pair in atom_types:
189 F_s_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
190 incoherent_displacements = dict()
191 incoherent_rho = dict()
193 # define correlation function
194 def calc_corr(window, time_i):
195 # Calculate correlations between two frames in the window without normalization 1/N
196 f0 = window[0]
197 fi = window[time_i]
198 for s1, s2 in pairs:
199 Fqt = np.real(f0.rho_qs_dict[s1] * fi.rho_qs_dict[s2].conjugate())
200 if s1 != s2:
201 Fqt += np.real(f0.rho_qs_dict[s2] * fi.rho_qs_dict[s1].conjugate())
202 F_q_t_averager[(s1, s2)].add_sample(time_i, Fqt)
204 if calculate_currents:
205 for s1, s2 in pairs:
206 Clqt = np.real(f0.jz_qs_dict[s1] * fi.jz_qs_dict[s2].conjugate())
207 Ctqt = 0.5 * np.real(np.sum(f0.jper_qs_dict[s1] *
208 fi.jper_qs_dict[s2].conjugate(), axis=1))
209 if s1 != s2:
210 Clqt += np.real(f0.jz_qs_dict[s2] * fi.jz_qs_dict[s1].conjugate())
211 Ctqt += 0.5 * np.real(np.sum(f0.jper_qs_dict[s2] *
212 fi.jper_qs_dict[s1].conjugate(), axis=1))
214 Cl_q_t_averager[(s1, s2)].add_sample(time_i, Clqt)
215 Ct_q_t_averager[(s1, s2)].add_sample(time_i, Ctqt)
217 def calc_incoherent(window):
218 # Calculate the incoherent (self) part for all lags in the window at once.
219 n_window = len(window)
220 f0 = window[0]
221 for atom_type in atom_types:
222 x0 = f0.positions_by_type[atom_type]
223 if atom_type not in incoherent_displacements:
224 incoherent_displacements[atom_type] = np.empty(
225 (n_window, 3, x0.shape[0]), dtype=x0.dtype)
226 incoherent_rho[atom_type] = np.empty(
227 (n_window, n_qpoints), dtype=np.float64)
229 displacements = incoherent_displacements[atom_type][:n_window]
230 for time_i, frame in enumerate(window):
231 np.subtract(frame.positions_by_type[atom_type].T, x0.T,
232 out=displacements[time_i])
234 rho = incoherent_rho[atom_type][:n_window]
235 _calc_rho_q_incoherent(displacements, q_points, out=rho)
236 for time_i, sample in enumerate(rho):
237 F_s_q_t_averager[atom_type].add_sample(time_i, sample)
239 # run calculation
240 # This is the "main loop" over the trajectory
241 for window in window_iterator:
242 if logging_interval and window[0].frame_index % logging_interval == 0:
243 logger.info(f'Processing window {window[0].frame_index} to {window[-1].frame_index}') # noqa
244 else:
245 logger.debug(f'Processing window {window[0].frame_index} to {window[-1].frame_index}') # noqa
247 for time_i in range(len(window)):
248 calc_corr(window, time_i)
250 if calculate_incoherent:
251 calc_incoherent(window)
253 # collect results into dict with numpy arrays (n_qpoints, N_tc)
254 data_dict_corr = dict()
255 data_dict_corr['q_points'] = q_points
257 time = None
258 for pair in pairs:
259 key = '_'.join(pair)
260 F_q_t = 1 / traj.n_atoms * truncate_nans(F_q_t_averager[pair].get_average_all())
261 w, S_q_w = psd_from_acf_2d(F_q_t, delta_t)
262 S_q_w = np.array(S_q_w)
264 if time is None:
265 # Determine the actual length of time signal after truncation (if needed)
266 time = delta_t * np.arange(F_q_t.shape[1], dtype=float)
267 N_tc_actual = len(time)
268 F_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
269 S_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
270 else:
271 assert F_q_t.shape[1] == len(time)
273 data_dict_corr['omega'] = w
274 data_dict_corr[f'Fqt_coh_{key}'] = F_q_t
275 data_dict_corr[f'Sqw_coh_{key}'] = S_q_w
277 # sum all partials to the total
278 F_q_t_tot += F_q_t
279 S_q_w_tot += S_q_w
281 if N_tc_actual < N_tc:
282 logger.warning('Truncating ACF due to NaNs, likely time_window is longer than length of Trajectory') # noqa
284 dw = float(w[1] - w[0])
285 w_max = float(w[-1])
287 data_dict_corr['time'] = time
288 data_dict_corr['Fqt_coh'] = F_q_t_tot
289 data_dict_corr['Sqw_coh'] = S_q_w_tot
291 if calculate_currents:
292 Cl_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
293 Ct_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
294 Cl_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
295 Ct_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
296 for pair in pairs:
297 key = '_'.join(pair)
298 Cl_q_t = 1 / traj.n_atoms * truncate_nans(Cl_q_t_averager[pair].get_average_all())
299 Ct_q_t = 1 / traj.n_atoms * truncate_nans(Ct_q_t_averager[pair].get_average_all())
300 _, Cl_q_w = psd_from_acf_2d(Cl_q_t, delta_t)
301 _, Ct_q_w = psd_from_acf_2d(Ct_q_t, delta_t)
302 data_dict_corr[f'Clqt_{key}'] = Cl_q_t
303 data_dict_corr[f'Ctqt_{key}'] = Ct_q_t
304 data_dict_corr[f'Clqw_{key}'] = Cl_q_w
305 data_dict_corr[f'Ctqw_{key}'] = Ct_q_w
307 # sum all partials to the total
308 Cl_q_t_tot += Cl_q_t
309 Ct_q_t_tot += Ct_q_t
310 Cl_q_w_tot += Cl_q_w
311 Ct_q_w_tot += Ct_q_w
312 data_dict_corr['Clqt'] = Cl_q_t_tot
313 data_dict_corr['Ctqt'] = Ct_q_t_tot
314 data_dict_corr['Clqw'] = Cl_q_w_tot
315 data_dict_corr['Ctqw'] = Ct_q_w_tot
317 if calculate_incoherent:
318 Fs_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
319 Ss_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
320 for atom_type in atom_types:
321 Fs_q_t = 1 / traj.n_atoms * truncate_nans(F_s_q_t_averager[atom_type].get_average_all())
322 _, Ss_q_w = psd_from_acf_2d(Fs_q_t, delta_t)
323 data_dict_corr[f'Fqt_incoh_{atom_type}'] = Fs_q_t
324 data_dict_corr[f'Sqw_incoh_{atom_type}'] = Ss_q_w
326 # sum all partials to the total
327 Fs_q_t_tot += Fs_q_t
328 Ss_q_w_tot += Ss_q_w
330 data_dict_corr['Fqt_incoh'] = Fs_q_t_tot
331 data_dict_corr['Sqw_incoh'] = Ss_q_w_tot
333 # finalize results with additional metadata
334 new_sample = DynamicSample(
335 data_dict_corr,
336 simulation_data=dict(
337 atom_types=atom_types, pairs=pairs,
338 particle_counts=particle_counts,
339 cell=traj.cell,
340 time_between_frames=delta_t,
341 maximum_time_lag=float(time[-1]),
342 angular_frequency_resolution=dw,
343 maximum_angular_frequency=w_max,
344 number_of_frames=traj.number_of_frames_read,
345 ))
346 new_sample._append_history(
347 'compute_dynamic_structure_factors',
348 dict(
349 dt=dt,
350 window_size=window_size,
351 window_step=window_step,
352 calculate_currents=calculate_currents,
353 calculate_incoherent=calculate_incoherent,
354 ))
356 return new_sample
359def compute_static_structure_factors(
360 traj: Trajectory,
361 q_points: NDArray[float],
362 logging_interval: Optional[int] = 1000,
363) -> StaticSample:
364 r"""Compute the static structure factors. The results are returned in the
365 form of a :class:`StaticSample <dynasor.sample.StaticSample>`
366 object.
368 The computational backend and precision are set globally via
369 :func:`set_compute_config <dynasor.set_compute_config>`, not as parameters of this
370 function; see :ref:`the backends reference page <backends>` for details.
372 Parameters
373 ----------
374 traj
375 Input trajectory.
376 q_points
377 Array of q-points in units of rad/Å with shape ``(N_qpoints, 3)`` in Cartesian coordinates.
378 logging_interval
379 Log progress at ``INFO`` level every this many frames. Set to ``0`` to disable
380 progress logging.
381 """
382 # sanity check input args
383 if q_points.ndim != 2 or q_points.shape[1] != 3:
384 raise ValueError('q-points array has the wrong shape.')
385 if not traj.has_positions:
386 raise ValueError('compute_static_structure_factors requires positions to be available '
387 'in the trajectory, but traj does not provide positions.')
388 _validate_qpoints_commensurate(q_points, traj.cell)
390 n_qpoints = q_points.shape[0]
391 logger.info(f'Number of q-points: {n_qpoints}')
393 # define all pairs
394 pairs = list(combinations_with_replacement(traj.atom_types, r=2))
395 particle_counts = {key: len(val) for key, val in traj.atomic_indices.items()}
396 logger.debug('Considering pairs:')
397 for pair in pairs:
398 logger.debug(f' {pair}')
400 # processing function / batch processor
401 compute_config = get_compute_config()
402 backend = compute_config.backend
403 precision = compute_config.precision
404 gpu_batch_size = compute_config.batch_size
406 is_batched, processor = get_frame_processor(q_points, backend, precision, gpu_batch_size)
408 # setup averager
409 Sq_averager = dict()
410 for pair in pairs:
411 Sq_averager[pair] = TimeAverager(1, n_qpoints) # time average with only timelag=0
413 # main loop: GPU uses batched prefetch, numba uses per-frame map
414 _frame_source = (_prefetch_batched_iter(traj, processor, batch_size=gpu_batch_size)
415 if is_batched
416 else map(processor, traj))
418 for frame in _frame_source:
420 if logging_interval and frame.frame_index % logging_interval == 0:
421 logger.info(f'Processing frame {frame.frame_index}')
422 else:
423 logger.debug(f'Processing frame {frame.frame_index}')
425 for s1, s2 in pairs:
426 # compute correlation
427 Sq_pair = np.real(frame.rho_qs_dict[s1] * frame.rho_qs_dict[s2].conjugate())
428 if s1 != s2:
429 Sq_pair += np.real(frame.rho_qs_dict[s2] * frame.rho_qs_dict[s1].conjugate())
430 Sq_averager[(s1, s2)].add_sample(0, Sq_pair)
432 # collect results
433 data_dict = dict()
434 data_dict['q_points'] = q_points
435 S_q_tot = np.zeros((n_qpoints, 1))
436 for s1, s2 in pairs:
437 Sq = 1 / traj.n_atoms * Sq_averager[(s1, s2)].get_average_at_timelag(0).reshape(-1, 1)
438 data_dict[f'Sq_{s1}_{s2}'] = Sq
439 S_q_tot += Sq
440 data_dict['Sq'] = S_q_tot
442 # finalize results
443 new_sample = StaticSample(
444 data_dict,
445 simulation_data=dict(
446 atom_types=traj.atom_types,
447 pairs=pairs,
448 particle_counts=particle_counts,
449 cell=traj.cell,
450 number_of_frames=traj.number_of_frames_read,
451 ))
452 new_sample._append_history('compute_static_structure_factors')
454 return new_sample
457def compute_spectral_energy_density(
458 traj: Trajectory,
459 ideal_supercell: Atoms,
460 primitive_cell: Atoms,
461 q_points: NDArray[float],
462 dt: float,
463 partial: Optional[bool] = False,
464 logging_interval: Optional[int] = 1000,
465) -> tuple[NDArray[float], NDArray[float]]:
466 r"""
467 Compute the spectral energy density (SED) at specific q-points. The results
468 are returned in the form of a tuple, which comprises the angular
469 frequencies in an array of length ``N_times`` in units of rad/fs and the
470 SED in units of eV/(rad/fs) as an array of shape ``(N_qpoints, N_times)``.
471 The normalization is chosen such that integrating the SED of a q-point
472 together with the supplied angular frequencies omega (rad/fs) yields
473 `1/2 kB T` * number of bands (where number of bands = `len(prim) * 3`)
475 More details can be found in Thomas *et al.*, Physical Review B **81**, 081411 (2010),
476 which should be cited when using this function along with the dynasor reference.
478 **Note 1:**
479 SED analysis is only suitable for crystalline materials without diffusion as
480 atoms are assumed to move around fixed reference positions throughout the entire trajectory.
482 **Note 2:**
483 This implementation reads the full trajectory and can thus consume a lot of memory.
485 Parameters
486 ----------
487 traj
488 Input trajectory.
489 ideal_supercell
490 Ideal structure defining the reference positions. Do not change the masses
491 in the ASE :class:`~ase.Atoms` objects to dynasor internal units, this will be
492 done internally. Its atom count is checked against :attr:`traj`; a mismatched
493 cell only triggers a warning (some thermal expansion relative to the reference
494 structure is normal), and the atom ordering is not checked at all, so it must
495 match the trajectory's.
496 primitive_cell
497 Underlying primitive structure. Must be aligned correctly with :attr:`ideal_supercell`.
498 q_points
499 Array of q-points in units of rad/Å with shape ``(N_qpoints, 3)`` in Cartesian coordinates.
500 dt
501 Time difference in femtoseconds between two consecutive snapshots in
502 the trajectory. Note that you should not change :attr:`dt` if you change
503 :attr:`frame_step <dynasor.Trajectory.frame_step>` in :attr:`traj`.
504 partial
505 If True the SED will be returned decomposed per basis and Cartesian direction.
506 The shape is ``(N_qpoints, N_frequencies, len(primitive_cell), 3)``.
507 logging_interval
508 Log progress at ``INFO`` level every this many frames. Set to ``0`` to disable
509 progress logging.
510 """
512 if q_points.ndim != 2 or q_points.shape[1] != 3:
513 raise ValueError('q-points array has the wrong shape.')
514 if dt <= 0:
515 raise ValueError(f'dt must be positive: dt= {dt}')
517 delta_t = traj.frame_step * dt
519 # logger
520 logger.info('Running SED')
521 logger.info(f'Time between consecutive frames (dt * frame_step): {delta_t} fs')
522 logger.info(f'Number of atoms in primitive_cell: {len(primitive_cell)}')
523 logger.info(f'Number of atoms in ideal_supercell: {len(ideal_supercell)}')
524 logger.info(f'Number of q-points: {q_points.shape[0]}')
526 # check that the ideal supercell agrees with traj
527 if traj.n_atoms != len(ideal_supercell):
528 raise ValueError('ideal_supercell must contain the same number of atoms as the trajectory.')
529 if not np.allclose(traj.cell, ideal_supercell.cell, atol=1e-5, rtol=0.0):
530 logger.warning('ideal_supercell cell does not match the trajectory cell.')
532 if len(primitive_cell) > len(ideal_supercell):
533 raise ValueError('primitive_cell contains more atoms than ideal_supercell.')
535 if not traj.has_velocities:
536 raise ValueError('compute_spectral_energy_density requires velocities to be available '
537 'in the trajectory, but traj does not provide velocities.')
539 # q-points must be commensurate with the ideal supercell (which defines the
540 # phase factors below), not with the possibly thermally-expanded traj cell
541 _validate_qpoints_commensurate(q_points, ideal_supercell.cell)
543 # collect all velocities, and scale with sqrt(masses)
544 masses = ideal_supercell.get_masses().reshape(-1, 1) / fs**2 # From Dalton to dmu
545 velocities = []
546 for it, frame in enumerate(traj):
547 if logging_interval and it % logging_interval == 0:
548 logger.info(f'Reading frame {it}')
549 else:
550 logger.debug(f'Reading frame {it}')
551 v = frame.get_velocities_as_array(traj.atomic_indices) # in Å/fs
552 velocities.append(np.sqrt(masses) * v)
553 logger.info(f'Number of snapshots: {len(velocities)}')
555 # Perform the FFT on the last axis for extra speed (maybe not needed)
556 N_samples = len(velocities)
557 velocities = np.array(velocities)
558 # places time index last and makes a copy for continuity
559 velocities = velocities.transpose(1, 2, 0).copy()
560 # #atoms in supercell x 3 directions x #frequencies
561 velocities = np.fft.rfft(velocities, axis=2)
563 # Calculate indices and offsets needed for the SED method
564 offsets, indices = get_offset_index(primitive_cell, ideal_supercell)
566 # Phase factor for use in FT. #qpoints x #atoms in supercell
567 cell_positions = np.dot(offsets, primitive_cell.cell)
568 phase = np.dot(q_points, cell_positions.T) # #qpoints x #unit cells
569 phase_factors = np.exp(1.0j * phase)
571 # This dict maps the offsets to an index so ndarrays can be over
572 # offset,index instead of atoms in supercell
573 offset_dict = {off: n for n, off in enumerate(set(tuple(offset) for offset in offsets))}
575 # Pick out some shapes
576 n_super, _, n_w = velocities.shape
577 n_qpts = len(q_points)
578 n_prim = len(primitive_cell)
579 n_offsets = len(offset_dict)
581 # This new array will be indexed by index and offset instead (and also transposed)
582 new_velocities = np.zeros((n_w, 3, n_prim, n_offsets), dtype=velocities.dtype)
584 for i in range(n_super):
585 j = indices[i] # atom with index i in the supercell is of basis type j ...
586 n = offset_dict[tuple(offsets[i])] # and its offset has index n
587 new_velocities[:, :, j, n] = velocities[i].T
589 velocities = new_velocities
591 # Same story with the spatial phase factors
592 new_phase_factors = np.zeros((n_qpts, n_prim, n_offsets), dtype=phase_factors.dtype)
594 for i in range(n_super):
595 j = indices[i]
596 n = offset_dict[tuple(offsets[i])]
597 new_phase_factors[:, j, n] = phase_factors[:, i]
599 phase_factors = new_phase_factors
601 # calculate the density in a numba function
602 density = _sed_inner_loop(phase_factors, velocities)
603 warn_if_numba_threading_layer_is_slow()
605 if not partial:
606 density = np.sum(density, axis=(2, 3))
608 # units
609 # make so that the velocities were originally in Angstrom / fs to be compatible with eV and Da
611 # the time delta in the fourier transform
612 density = density * delta_t**2
614 # Divide by the length of the time signal
615 density = density / (N_samples * delta_t)
617 # Divide by the number of primitive cells
618 density = density / (n_super / n_prim)
620 # Factor so the sed can be integrated together with the returned omega
621 # numpy fft works with ordinary/linear frequencies and not angular freqs
622 density = density / (2*np.pi)
624 # angular frequencies
625 w = 2 * np.pi * np.fft.rfftfreq(N_samples, delta_t) # rad/fs
627 return w, density
630@numba.njit(parallel=True, fastmath=True)
631def _sed_inner_loop(phase_factors, velocities):
632 """This numba function calculates the spatial
633 Fourier transform using precomputed phase factors.
635 As the use case can be one or many q-points the parallelization is over the
636 temporal frequency components instead.
637 """
639 n_qpts = phase_factors.shape[0] # q-point index
640 n_prim = phase_factors.shape[1] # basis atom index
641 n_super = phase_factors.shape[2] # unit cell index
643 n_freqs = velocities.shape[0] # frequency, direction, basis atom, unit cell
645 density = np.zeros((n_qpts, n_freqs, n_prim, 3), dtype=np.float64)
647 for w in numba.prange(n_freqs):
648 for k in range(n_qpts):
649 for a in range(3):
650 for b in range(n_prim):
651 tmp = 0.0j
652 for n in range(n_super):
653 tmp += phase_factors[k, b, n] * velocities[w, a, b, n]
654 density[k, w, b, a] += np.abs(tmp)**2
655 return density