Coverage for dynasor/correlation_functions.py: 98%
304 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-02 07: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 traj._log_progress(window[-1].frame_index, window_start=window[0].frame_index)
245 for time_i in range(len(window)):
246 calc_corr(window, time_i)
248 if calculate_incoherent:
249 calc_incoherent(window)
251 # collect results into dict with numpy arrays (n_qpoints, N_tc)
252 data_dict_corr = dict()
253 data_dict_corr['q_points'] = q_points
255 time = None
256 for pair in pairs:
257 key = '_'.join(pair)
258 F_q_t = 1 / traj.n_atoms * truncate_nans(F_q_t_averager[pair].get_average_all())
259 w, S_q_w = psd_from_acf_2d(F_q_t, delta_t)
260 S_q_w = np.array(S_q_w)
262 if time is None:
263 # Determine the actual length of time signal after truncation (if needed)
264 time = delta_t * np.arange(F_q_t.shape[1], dtype=float)
265 N_tc_actual = len(time)
266 F_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
267 S_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
268 else:
269 assert F_q_t.shape[1] == len(time)
271 data_dict_corr['omega'] = w
272 data_dict_corr[f'Fqt_coh_{key}'] = F_q_t
273 data_dict_corr[f'Sqw_coh_{key}'] = S_q_w
275 # sum all partials to the total
276 F_q_t_tot += F_q_t
277 S_q_w_tot += S_q_w
279 if N_tc_actual < N_tc:
280 logger.warning('Truncating ACF due to NaNs, likely time_window is longer than length of Trajectory') # noqa
282 dw = float(w[1] - w[0])
283 w_max = float(w[-1])
285 data_dict_corr['time'] = time
286 data_dict_corr['Fqt_coh'] = F_q_t_tot
287 data_dict_corr['Sqw_coh'] = S_q_w_tot
289 if calculate_currents:
290 Cl_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
291 Ct_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
292 Cl_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
293 Ct_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
294 for pair in pairs:
295 key = '_'.join(pair)
296 Cl_q_t = 1 / traj.n_atoms * truncate_nans(Cl_q_t_averager[pair].get_average_all())
297 Ct_q_t = 1 / traj.n_atoms * truncate_nans(Ct_q_t_averager[pair].get_average_all())
298 _, Cl_q_w = psd_from_acf_2d(Cl_q_t, delta_t)
299 _, Ct_q_w = psd_from_acf_2d(Ct_q_t, delta_t)
300 data_dict_corr[f'Clqt_{key}'] = Cl_q_t
301 data_dict_corr[f'Ctqt_{key}'] = Ct_q_t
302 data_dict_corr[f'Clqw_{key}'] = Cl_q_w
303 data_dict_corr[f'Ctqw_{key}'] = Ct_q_w
305 # sum all partials to the total
306 Cl_q_t_tot += Cl_q_t
307 Ct_q_t_tot += Ct_q_t
308 Cl_q_w_tot += Cl_q_w
309 Ct_q_w_tot += Ct_q_w
310 data_dict_corr['Clqt'] = Cl_q_t_tot
311 data_dict_corr['Ctqt'] = Ct_q_t_tot
312 data_dict_corr['Clqw'] = Cl_q_w_tot
313 data_dict_corr['Ctqw'] = Ct_q_w_tot
315 if calculate_incoherent:
316 Fs_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
317 Ss_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
318 for atom_type in atom_types:
319 Fs_q_t = 1 / traj.n_atoms * truncate_nans(F_s_q_t_averager[atom_type].get_average_all())
320 _, Ss_q_w = psd_from_acf_2d(Fs_q_t, delta_t)
321 data_dict_corr[f'Fqt_incoh_{atom_type}'] = Fs_q_t
322 data_dict_corr[f'Sqw_incoh_{atom_type}'] = Ss_q_w
324 # sum all partials to the total
325 Fs_q_t_tot += Fs_q_t
326 Ss_q_w_tot += Ss_q_w
328 data_dict_corr['Fqt_incoh'] = Fs_q_t_tot
329 data_dict_corr['Sqw_incoh'] = Ss_q_w_tot
331 # finalize results with additional metadata
332 new_sample = DynamicSample(
333 data_dict_corr,
334 simulation_data=dict(
335 atom_types=atom_types, pairs=pairs,
336 particle_counts=particle_counts,
337 cell=traj.cell,
338 time_between_frames=delta_t,
339 maximum_time_lag=float(time[-1]),
340 angular_frequency_resolution=dw,
341 maximum_angular_frequency=w_max,
342 number_of_frames=traj.number_of_frames_read,
343 ))
344 new_sample._append_history(
345 'compute_dynamic_structure_factors',
346 dict(
347 dt=dt,
348 window_size=window_size,
349 window_step=window_step,
350 calculate_currents=calculate_currents,
351 calculate_incoherent=calculate_incoherent,
352 ))
354 return new_sample
357def compute_static_structure_factors(
358 traj: Trajectory,
359 q_points: NDArray[float],
360 logging_interval: Optional[int] = 1000,
361) -> StaticSample:
362 r"""Compute the static structure factors. The results are returned in the
363 form of a :class:`StaticSample <dynasor.sample.StaticSample>`
364 object.
366 The computational backend and precision are set globally via
367 :func:`set_compute_config <dynasor.set_compute_config>`, not as parameters of this
368 function; see :ref:`the backends reference page <backends>` for details.
370 Parameters
371 ----------
372 traj
373 Input trajectory.
374 q_points
375 Array of q-points in units of rad/Å with shape ``(N_qpoints, 3)`` in Cartesian coordinates.
376 logging_interval
377 Log progress at ``INFO`` level every this many frames. Set to ``0`` to disable
378 progress logging.
379 """
380 # sanity check input args
381 if q_points.ndim != 2 or q_points.shape[1] != 3:
382 raise ValueError('q-points array has the wrong shape.')
383 if not traj.has_positions:
384 raise ValueError('compute_static_structure_factors requires positions to be available '
385 'in the trajectory, but traj does not provide positions.')
386 _validate_qpoints_commensurate(q_points, traj.cell)
388 n_qpoints = q_points.shape[0]
389 logger.info(f'Number of q-points: {n_qpoints}')
391 # define all pairs
392 pairs = list(combinations_with_replacement(traj.atom_types, r=2))
393 particle_counts = {key: len(val) for key, val in traj.atomic_indices.items()}
394 logger.debug('Considering pairs:')
395 for pair in pairs:
396 logger.debug(f' {pair}')
398 # processing function / batch processor
399 compute_config = get_compute_config()
400 backend = compute_config.backend
401 precision = compute_config.precision
402 gpu_batch_size = compute_config.batch_size
404 is_batched, processor = get_frame_processor(q_points, backend, precision, gpu_batch_size)
406 # setup averager
407 Sq_averager = dict()
408 for pair in pairs:
409 Sq_averager[pair] = TimeAverager(1, n_qpoints) # time average with only timelag=0
411 # main loop: GPU uses batched prefetch, numba uses per-frame map
412 _frame_source = (_prefetch_batched_iter(traj, processor, batch_size=gpu_batch_size)
413 if is_batched
414 else map(processor, traj))
416 for frame in _frame_source:
418 if logging_interval and frame.frame_index % logging_interval == 0:
419 traj._log_progress(frame.frame_index)
421 for s1, s2 in pairs:
422 # compute correlation
423 Sq_pair = np.real(frame.rho_qs_dict[s1] * frame.rho_qs_dict[s2].conjugate())
424 if s1 != s2:
425 Sq_pair += np.real(frame.rho_qs_dict[s2] * frame.rho_qs_dict[s1].conjugate())
426 Sq_averager[(s1, s2)].add_sample(0, Sq_pair)
428 # collect results
429 data_dict = dict()
430 data_dict['q_points'] = q_points
431 S_q_tot = np.zeros((n_qpoints, 1))
432 for s1, s2 in pairs:
433 Sq = 1 / traj.n_atoms * Sq_averager[(s1, s2)].get_average_at_timelag(0).reshape(-1, 1)
434 data_dict[f'Sq_{s1}_{s2}'] = Sq
435 S_q_tot += Sq
436 data_dict['Sq'] = S_q_tot
438 # finalize results
439 new_sample = StaticSample(
440 data_dict,
441 simulation_data=dict(
442 atom_types=traj.atom_types,
443 pairs=pairs,
444 particle_counts=particle_counts,
445 cell=traj.cell,
446 number_of_frames=traj.number_of_frames_read,
447 ))
448 new_sample._append_history('compute_static_structure_factors')
450 return new_sample
453def compute_spectral_energy_density(
454 traj: Trajectory,
455 ideal_supercell: Atoms,
456 primitive_cell: Atoms,
457 q_points: NDArray[float],
458 dt: float,
459 partial: Optional[bool] = False,
460 logging_interval: Optional[int] = 1000,
461) -> tuple[NDArray[float], NDArray[float]]:
462 r"""
463 Compute the spectral energy density (SED) at specific q-points. The results
464 are returned in the form of a tuple, which comprises the angular
465 frequencies in an array of length ``N_times`` in units of rad/fs and the
466 SED in units of eV/(rad/fs) as an array of shape ``(N_qpoints, N_times)``.
467 The normalization is chosen such that integrating the SED of a q-point
468 together with the supplied angular frequencies omega (rad/fs) yields
469 `1/2 kB T` * number of bands (where number of bands = `len(prim) * 3`)
471 More details can be found in Thomas *et al.*, Physical Review B **81**, 081411 (2010),
472 which should be cited when using this function along with the dynasor reference.
474 **Note 1:**
475 SED analysis is only suitable for crystalline materials without diffusion as
476 atoms are assumed to move around fixed reference positions throughout the entire trajectory.
478 **Note 2:**
479 This implementation reads the full trajectory and can thus consume a lot of memory.
481 Parameters
482 ----------
483 traj
484 Input trajectory.
485 ideal_supercell
486 Ideal structure defining the reference positions. Do not change the masses
487 in the ASE :class:`~ase.Atoms` objects to dynasor internal units, this will be
488 done internally. Its atom count is checked against :attr:`traj`; a mismatched
489 cell only triggers a warning (some thermal expansion relative to the reference
490 structure is normal), and the atom ordering is not checked at all, so it must
491 match the trajectory's.
492 primitive_cell
493 Underlying primitive structure. Must be aligned correctly with :attr:`ideal_supercell`.
494 q_points
495 Array of q-points in units of rad/Å with shape ``(N_qpoints, 3)`` in Cartesian coordinates.
496 dt
497 Time difference in femtoseconds between two consecutive snapshots in
498 the trajectory. Note that you should not change :attr:`dt` if you change
499 :attr:`frame_step <dynasor.Trajectory.frame_step>` in :attr:`traj`.
500 partial
501 If True the SED will be returned decomposed per basis and Cartesian direction.
502 The shape is ``(N_qpoints, N_frequencies, len(primitive_cell), 3)``.
503 logging_interval
504 Log progress at ``INFO`` level every this many frames. Set to ``0`` to disable
505 progress logging.
506 """
508 if q_points.ndim != 2 or q_points.shape[1] != 3:
509 raise ValueError('q-points array has the wrong shape.')
510 if dt <= 0:
511 raise ValueError(f'dt must be positive: dt= {dt}')
513 delta_t = traj.frame_step * dt
515 # logger
516 logger.info('Running SED')
517 logger.info(f'Time between consecutive frames (dt * frame_step): {delta_t} fs')
518 logger.info(f'Number of atoms in primitive_cell: {len(primitive_cell)}')
519 logger.info(f'Number of atoms in ideal_supercell: {len(ideal_supercell)}')
520 logger.info(f'Number of q-points: {q_points.shape[0]}')
522 # check that the ideal supercell agrees with traj
523 if traj.n_atoms != len(ideal_supercell):
524 raise ValueError('ideal_supercell must contain the same number of atoms as the trajectory.')
525 if not np.allclose(traj.cell, ideal_supercell.cell, atol=1e-5, rtol=0.0):
526 logger.warning('ideal_supercell cell does not match the trajectory cell.')
528 if len(primitive_cell) > len(ideal_supercell):
529 raise ValueError('primitive_cell contains more atoms than ideal_supercell.')
531 if not traj.has_velocities:
532 raise ValueError('compute_spectral_energy_density requires velocities to be available '
533 'in the trajectory, but traj does not provide velocities.')
535 # q-points must be commensurate with the ideal supercell (which defines the
536 # phase factors below), not with the possibly thermally-expanded traj cell
537 _validate_qpoints_commensurate(q_points, ideal_supercell.cell)
539 # collect all velocities, and scale with sqrt(masses)
540 masses = ideal_supercell.get_masses().reshape(-1, 1) / fs**2 # From Dalton to dmu
541 velocities = []
542 for it, frame in enumerate(traj):
543 if logging_interval and it % logging_interval == 0:
544 traj._log_progress(frame.frame_index)
545 v = frame.get_velocities_as_array(traj.atomic_indices) # in Å/fs
546 velocities.append(np.sqrt(masses) * v)
547 logger.info(f'Number of snapshots: {len(velocities)}')
549 # Perform the FFT on the last axis for extra speed (maybe not needed)
550 N_samples = len(velocities)
551 velocities = np.array(velocities)
552 # places time index last and makes a copy for continuity
553 velocities = velocities.transpose(1, 2, 0).copy()
554 # #atoms in supercell x 3 directions x #frequencies
555 velocities = np.fft.rfft(velocities, axis=2)
557 # Calculate indices and offsets needed for the SED method
558 offsets, indices = get_offset_index(primitive_cell, ideal_supercell)
560 # Phase factor for use in FT. #qpoints x #atoms in supercell
561 cell_positions = np.dot(offsets, primitive_cell.cell)
562 phase = np.dot(q_points, cell_positions.T) # #qpoints x #unit cells
563 phase_factors = np.exp(1.0j * phase)
565 # This dict maps the offsets to an index so ndarrays can be over
566 # offset,index instead of atoms in supercell
567 offset_dict = {off: n for n, off in enumerate(set(tuple(offset) for offset in offsets))}
569 # Pick out some shapes
570 n_super, _, n_w = velocities.shape
571 n_qpts = len(q_points)
572 n_prim = len(primitive_cell)
573 n_offsets = len(offset_dict)
575 # This new array will be indexed by index and offset instead (and also transposed)
576 new_velocities = np.zeros((n_w, 3, n_prim, n_offsets), dtype=velocities.dtype)
578 for i in range(n_super):
579 j = indices[i] # atom with index i in the supercell is of basis type j ...
580 n = offset_dict[tuple(offsets[i])] # and its offset has index n
581 new_velocities[:, :, j, n] = velocities[i].T
583 velocities = new_velocities
585 # Same story with the spatial phase factors
586 new_phase_factors = np.zeros((n_qpts, n_prim, n_offsets), dtype=phase_factors.dtype)
588 for i in range(n_super):
589 j = indices[i]
590 n = offset_dict[tuple(offsets[i])]
591 new_phase_factors[:, j, n] = phase_factors[:, i]
593 phase_factors = new_phase_factors
595 # calculate the density in a numba function
596 density = _sed_inner_loop(phase_factors, velocities)
597 warn_if_numba_threading_layer_is_slow()
599 if not partial:
600 density = np.sum(density, axis=(2, 3))
602 # units
603 # make so that the velocities were originally in Angstrom / fs to be compatible with eV and Da
605 # the time delta in the fourier transform
606 density = density * delta_t**2
608 # Divide by the length of the time signal
609 density = density / (N_samples * delta_t)
611 # Divide by the number of primitive cells
612 density = density / (n_super / n_prim)
614 # Factor so the sed can be integrated together with the returned omega
615 # numpy fft works with ordinary/linear frequencies and not angular freqs
616 density = density / (2*np.pi)
618 # angular frequencies
619 w = 2 * np.pi * np.fft.rfftfreq(N_samples, delta_t) # rad/fs
621 return w, density
624@numba.njit(parallel=True, fastmath=True)
625def _sed_inner_loop(phase_factors, velocities):
626 """This numba function calculates the spatial
627 Fourier transform using precomputed phase factors.
629 As the use case can be one or many q-points the parallelization is over the
630 temporal frequency components instead.
631 """
633 n_qpts = phase_factors.shape[0] # q-point index
634 n_prim = phase_factors.shape[1] # basis atom index
635 n_super = phase_factors.shape[2] # unit cell index
637 n_freqs = velocities.shape[0] # frequency, direction, basis atom, unit cell
639 density = np.zeros((n_qpts, n_freqs, n_prim, 3), dtype=np.float64)
641 for w in numba.prange(n_freqs):
642 for k in range(n_qpts):
643 for a in range(3):
644 for b in range(n_prim):
645 tmp = 0.0j
646 for n in range(n_super):
647 tmp += phase_factors[k, b, n] * velocities[w, a, b, n]
648 density[k, w, b, a] += np.abs(tmp)**2
649 return density