Coverage for dynasor/correlation_functions.py: 100%
302 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-10 08:27 +0000
1import concurrent.futures
2import functools
3from itertools import combinations_with_replacement
4from typing import Optional
6import numba
7import numpy as np
8from ase import Atoms
9from ase.units import fs
10from numpy.typing import NDArray
12from dynasor.logging_tools import logger
13from dynasor.trajectory import Trajectory, WindowIterator
14from dynasor.sample import DynamicSample, StaticSample
16from dynasor.tools.acfs import psd_from_acf_2d
17from dynasor.core.time_averager import TimeAverager, truncate_nans
18from dynasor.core.reciprocal import calc_rho_q, calc_rho_j_q
19from dynasor.tools.structures import get_offset_index
20from dynasor.units import radians_per_fs_to_meV
23def compute_dynamic_structure_factors(
24 traj: Trajectory,
25 q_points: NDArray[float],
26 dt: float,
27 window_size: int,
28 window_step: Optional[int] = 1,
29 calculate_currents: Optional[bool] = False,
30 calculate_incoherent: Optional[bool] = False,
31 logging_interval: Optional[int] = 1000,
32) -> DynamicSample:
33 r"""Compute the dynamic structure factors. The results are returned in the
34 form of a :class:`DynamicSample <dynasor.sample.DynamicSample>`
35 object.
37 Parameters
38 ----------
39 traj
40 Input trajectory.
41 q_points
42 Array of q-points in units of rad/Å with shape ``(N_qpoints, 3)`` in Cartesian coordinates.
43 dt
44 Time difference in femtoseconds between two consecutive snapshots
45 in the trajectory. Note that you should *not* change :attr:`dt` if you change
46 :attr:`frame_step <dynasor.trajectory.Trajectory.frame_step>` in :attr:`traj`.
47 window_size
48 Maximum time lag, expressed as a number of frame intervals, for which to calculate
49 correlations. Each calculation window therefore contains ``window_size + 1`` frames,
50 including the frame at time lag zero. This parameter determines the smallest frequency
51 resolved.
52 window_step
53 Window step (or stride) given as the number of frames between consecutive trajectory
54 windows. This parameter does *not* affect the time between consecutive frames in the
55 calculation. If :attr:`window_step` > :attr:`window_size` + 1, some frames will not be
56 used.
57 calculate_currents
58 Calculate the current correlations. Requires velocities to be available in :attr:`traj`.
59 calculate_incoherent
60 Calculate the incoherent part (self-part) of :math:`F_\text{incoh}`.
61 logging_interval
62 Log progress at ``INFO`` level every this many windows. Set to ``0`` to disable
63 progress logging.
64 """
65 # sanity check input args
66 if q_points.ndim != 2 or q_points.shape[1] != 3:
67 raise ValueError('q-points array has the wrong shape.')
68 if dt <= 0:
69 raise ValueError(f'dt must be positive: dt= {dt}')
70 if window_size <= 2:
71 raise ValueError(f'window_size must be larger than 2: window_size= {window_size}')
72 if window_step <= 0:
73 raise ValueError(f'window_step must be positive: window_step= {window_step}')
74 if calculate_currents and not traj.has_velocities:
75 raise ValueError('calculate_currents=True requires velocities to be available in the '
76 'trajectory, but traj does not provide velocities.')
78 # define internal parameters
79 n_qpoints = q_points.shape[0]
80 delta_t = traj.frame_step * dt
81 N_tc = window_size + 1
83 # log all setup information
84 n_fft = 2 * window_size + 1
85 dw = 2 * np.pi / (n_fft * delta_t)
86 w_max = window_size * dw
87 w_N = np.pi / delta_t
88 dw_mev = dw * radians_per_fs_to_meV
89 w_max_mev = w_max * radians_per_fs_to_meV
90 w_N_mev = w_N * radians_per_fs_to_meV
91 logger.info(f'Spacing between samples (frame_step): {traj.frame_step}')
92 logger.info(f'Time between consecutive frames in input trajectory (dt): {dt} fs')
93 logger.info(f'Time between consecutive frames used (dt * frame_step): {delta_t} fs')
94 logger.info(f'Time window size (dt * frame_step * window_size): {delta_t * window_size:.1f} fs')
95 logger.info(f'Angular frequency resolution: dw = {dw:.6f} rad/fs = {dw_mev:.3f} meV')
96 logger.info(f'Maximum angular frequency: {w_max:.6f} rad/fs = {w_max_mev:.3f} meV '
97 f'(Nyquist limit: {w_N:.6f} rad/fs = {w_N_mev:.3f} meV)')
99 if calculate_currents:
100 logger.info('Calculating current (velocity) correlations')
101 if calculate_incoherent:
102 logger.info('Calculating incoherent part (self-part) of correlations')
104 # log some info regarding q-points
105 logger.info(f'Number of q-points: {n_qpoints}')
107 q_directions = q_points.astype(float, copy=True)
108 q_distances = np.linalg.norm(q_points, axis=1)
109 nonzero = q_distances > 0
110 q_directions[nonzero] /= q_distances[nonzero].reshape(-1, 1)
112 # setup functions to process frames
113 def f2_rho(frame):
114 rho_qs_dict = dict()
115 for atom_type in frame.positions_by_type.keys():
116 x = frame.positions_by_type[atom_type]
117 rho_qs_dict[atom_type] = calc_rho_q(x, q_points)
118 frame.rho_qs_dict = rho_qs_dict
119 return frame
121 def f2_rho_and_j(frame):
122 rho_qs_dict = dict()
123 jz_qs_dict = dict()
124 jper_qs_dict = dict()
126 for atom_type in frame.positions_by_type.keys():
127 x = frame.positions_by_type[atom_type]
128 v = frame.velocities_by_type[atom_type]
129 rho_qs, j_qs = calc_rho_j_q(x, v, q_points)
130 jz_qs = np.sum(j_qs * q_directions, axis=1)
131 jper_qs = j_qs - (jz_qs[:, None] * q_directions)
133 rho_qs_dict[atom_type] = rho_qs
134 jz_qs_dict[atom_type] = jz_qs
135 jper_qs_dict[atom_type] = jper_qs
137 frame.rho_qs_dict = rho_qs_dict
138 frame.jz_qs_dict = jz_qs_dict
139 frame.jper_qs_dict = jper_qs_dict
140 return frame
142 if calculate_currents:
143 element_processor = f2_rho_and_j
144 else:
145 element_processor = f2_rho
147 # setup window iterator
148 window_iterator = WindowIterator(traj, width=N_tc, window_step=window_step,
149 element_processor=element_processor)
151 # define all atom types and pairs
152 atom_types = traj.atom_types
153 pairs = list(combinations_with_replacement(atom_types, r=2))
154 particle_counts = {key: len(val) for key, val in traj.atomic_indices.items()}
155 logger.debug('Considering pairs:')
156 for pair in pairs:
157 logger.debug(f' {pair}')
159 # set up all time averager instances
160 F_q_t_averager = dict()
161 for pair in pairs:
162 F_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
163 if calculate_currents:
164 Cl_q_t_averager = dict()
165 Ct_q_t_averager = dict()
166 for pair in pairs:
167 Cl_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
168 Ct_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
169 if calculate_incoherent:
170 F_s_q_t_averager = dict()
171 for pair in atom_types:
172 F_s_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
174 # define correlation function
175 #
176 # Note: calc_corr is dispatched concurrently across OS threads via
177 # ThreadPoolExecutor.map() below, so it must never call into a numba
178 # parallel=True/prange kernel (e.g. calc_rho_q, calc_rho_j_q). Numba's
179 # default threading layer is not safe against concurrent entry from
180 # multiple Python threads and aborts the process if this happens. Such
181 # calls belong in calc_incoherent instead, which is always run strictly
182 # sequentially in the main thread.
183 def calc_corr(window, time_i):
184 # Calculate correlations between two frames in the window without normalization 1/N
185 f0 = window[0]
186 fi = window[time_i]
187 for s1, s2 in pairs:
188 Fqt = np.real(f0.rho_qs_dict[s1] * fi.rho_qs_dict[s2].conjugate())
189 if s1 != s2:
190 Fqt += np.real(f0.rho_qs_dict[s2] * fi.rho_qs_dict[s1].conjugate())
191 F_q_t_averager[(s1, s2)].add_sample(time_i, Fqt)
193 if calculate_currents:
194 for s1, s2 in pairs:
195 Clqt = np.real(f0.jz_qs_dict[s1] * fi.jz_qs_dict[s2].conjugate())
196 Ctqt = 0.5 * np.real(np.sum(f0.jper_qs_dict[s1] *
197 fi.jper_qs_dict[s2].conjugate(), axis=1))
198 if s1 != s2:
199 Clqt += np.real(f0.jz_qs_dict[s2] * fi.jz_qs_dict[s1].conjugate())
200 Ctqt += 0.5 * np.real(np.sum(f0.jper_qs_dict[s2] *
201 fi.jper_qs_dict[s1].conjugate(), axis=1))
203 Cl_q_t_averager[(s1, s2)].add_sample(time_i, Clqt)
204 Ct_q_t_averager[(s1, s2)].add_sample(time_i, Ctqt)
206 def calc_incoherent(window, time_i):
207 # Calculate the incoherent (self) part between two frames in the window.
208 #
209 # calc_rho_q is backed by a numba parallel=True/prange kernel, which must
210 # not be entered concurrently from multiple Python threads (see the note
211 # on calc_corr above). This function is therefore always called from a
212 # plain sequential loop in the main thread, never from within the
213 # ThreadPoolExecutor used for calc_corr.
214 f0 = window[0]
215 fi = window[time_i]
216 for atom_type in atom_types:
217 xi = fi.positions_by_type[atom_type]
218 x0 = f0.positions_by_type[atom_type]
219 Fsqt = np.real(calc_rho_q(xi - x0, q_points))
220 F_s_q_t_averager[atom_type].add_sample(time_i, Fsqt)
222 # run calculation
223 with concurrent.futures.ThreadPoolExecutor() as tpe:
224 # This is the "main loop" over the trajectory
225 for window in window_iterator:
226 if logging_interval and window[0].frame_index % logging_interval == 0:
227 logger.info(f'Processing window {window[0].frame_index} to {window[-1].frame_index}') # noqa
228 else:
229 logger.debug(f'Processing window {window[0].frame_index} to {window[-1].frame_index}') # noqa
231 # The map conveniently applies calc_corr to all time-lags. However,
232 # as everything is done in place nothing gets returned so in order
233 # to start and wait for the processes to finish we must iterate
234 # over the None values returned
235 for _ in tpe.map(functools.partial(calc_corr, window), range(len(window))):
236 pass
238 # Run the incoherent part strictly sequentially in the main thread
239 # (see calc_incoherent for why it must not run inside the thread pool).
240 if calculate_incoherent:
241 for time_i in range(len(window)):
242 calc_incoherent(window, time_i)
244 # collect results into dict with numpy arrays (n_qpoints, N_tc)
245 data_dict_corr = dict()
246 data_dict_corr['q_points'] = q_points
248 time = None
249 for pair in pairs:
250 key = '_'.join(pair)
251 F_q_t = 1 / traj.n_atoms * truncate_nans(F_q_t_averager[pair].get_average_all())
252 w, S_q_w = psd_from_acf_2d(F_q_t, delta_t)
253 S_q_w = np.array(S_q_w)
255 if time is None:
256 # Determine the actual length of time signal after truncation (if needed)
257 time = delta_t * np.arange(F_q_t.shape[1], dtype=float)
258 N_tc_actual = len(time)
259 F_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
260 S_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
261 else:
262 assert F_q_t.shape[1] == len(time)
264 data_dict_corr['omega'] = w
265 data_dict_corr[f'Fqt_coh_{key}'] = F_q_t
266 data_dict_corr[f'Sqw_coh_{key}'] = S_q_w
268 # sum all partials to the total
269 F_q_t_tot += F_q_t
270 S_q_w_tot += S_q_w
272 if N_tc_actual < N_tc:
273 logger.warning('Truncating ACF due to NaNs, likely time_window is longer than length of Trajectory') # noqa
275 dw = float(w[1] - w[0])
276 w_max = float(w[-1])
278 data_dict_corr['time'] = time
279 data_dict_corr['Fqt_coh'] = F_q_t_tot
280 data_dict_corr['Sqw_coh'] = S_q_w_tot
282 if calculate_currents:
283 Cl_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
284 Ct_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
285 Cl_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
286 Ct_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
287 for pair in pairs:
288 key = '_'.join(pair)
289 Cl_q_t = 1 / traj.n_atoms * truncate_nans(Cl_q_t_averager[pair].get_average_all())
290 Ct_q_t = 1 / traj.n_atoms * truncate_nans(Ct_q_t_averager[pair].get_average_all())
291 _, Cl_q_w = psd_from_acf_2d(Cl_q_t, delta_t)
292 _, Ct_q_w = psd_from_acf_2d(Ct_q_t, delta_t)
293 data_dict_corr[f'Clqt_{key}'] = Cl_q_t
294 data_dict_corr[f'Ctqt_{key}'] = Ct_q_t
295 data_dict_corr[f'Clqw_{key}'] = Cl_q_w
296 data_dict_corr[f'Ctqw_{key}'] = Ct_q_w
298 # sum all partials to the total
299 Cl_q_t_tot += Cl_q_t
300 Ct_q_t_tot += Ct_q_t
301 Cl_q_w_tot += Cl_q_w
302 Ct_q_w_tot += Ct_q_w
303 data_dict_corr['Clqt'] = Cl_q_t_tot
304 data_dict_corr['Ctqt'] = Ct_q_t_tot
305 data_dict_corr['Clqw'] = Cl_q_w_tot
306 data_dict_corr['Ctqw'] = Ct_q_w_tot
308 if calculate_incoherent:
309 Fs_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
310 Ss_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
311 for atom_type in atom_types:
312 Fs_q_t = 1 / traj.n_atoms * truncate_nans(F_s_q_t_averager[atom_type].get_average_all())
313 _, Ss_q_w = psd_from_acf_2d(Fs_q_t, delta_t)
314 data_dict_corr[f'Fqt_incoh_{atom_type}'] = Fs_q_t
315 data_dict_corr[f'Sqw_incoh_{atom_type}'] = Ss_q_w
317 # sum all partials to the total
318 Fs_q_t_tot += Fs_q_t
319 Ss_q_w_tot += Ss_q_w
321 data_dict_corr['Fqt_incoh'] = Fs_q_t_tot
322 data_dict_corr['Sqw_incoh'] = Ss_q_w_tot
324 # finalize results with additional metadata
325 new_sample = DynamicSample(
326 data_dict_corr,
327 simulation_data=dict(
328 atom_types=atom_types, pairs=pairs,
329 particle_counts=particle_counts,
330 cell=traj.cell,
331 time_between_frames=delta_t,
332 maximum_time_lag=float(time[-1]),
333 angular_frequency_resolution=dw,
334 maximum_angular_frequency=w_max,
335 number_of_frames=traj.number_of_frames_read,
336 ))
337 new_sample._append_history(
338 'compute_dynamic_structure_factors',
339 dict(
340 dt=dt,
341 window_size=window_size,
342 window_step=window_step,
343 calculate_currents=calculate_currents,
344 calculate_incoherent=calculate_incoherent,
345 ))
347 return new_sample
350def compute_static_structure_factors(
351 traj: Trajectory,
352 q_points: NDArray[float],
353 logging_interval: Optional[int] = 1000,
354) -> StaticSample:
355 r"""Compute the static structure factors. The results are returned in the
356 form of a :class:`StaticSample <dynasor.sample.StaticSample>`
357 object.
359 Parameters
360 ----------
361 traj
362 Input trajectory.
363 q_points
364 Array of q-points in units of rad/Å with shape ``(N_qpoints, 3)`` in Cartesian coordinates.
365 logging_interval
366 Log progress at ``INFO`` level every this many frames. Set to ``0`` to disable
367 progress logging.
368 """
369 # sanity check input args
370 if q_points.ndim != 2 or q_points.shape[1] != 3:
371 raise ValueError('q-points array has the wrong shape.')
373 n_qpoints = q_points.shape[0]
374 logger.info(f'Number of q-points: {n_qpoints}')
376 # define all pairs
377 pairs = list(combinations_with_replacement(traj.atom_types, r=2))
378 particle_counts = {key: len(val) for key, val in traj.atomic_indices.items()}
379 logger.debug('Considering pairs:')
380 for pair in pairs:
381 logger.debug(f' {pair}')
383 # processing function
384 def f2_rho(frame):
385 rho_qs_dict = dict()
386 for atom_type in frame.positions_by_type.keys():
387 x = frame.positions_by_type[atom_type]
388 rho_qs_dict[atom_type] = calc_rho_q(x, q_points)
389 frame.rho_qs_dict = rho_qs_dict
390 return frame
392 # setup averager
393 Sq_averager = dict()
394 for pair in pairs:
395 Sq_averager[pair] = TimeAverager(1, n_qpoints) # time average with only timelag=0
397 # main loop
398 for frame in traj:
400 # process_frame
401 f2_rho(frame)
402 if logging_interval and frame.frame_index % logging_interval == 0:
403 logger.info(f'Processing frame {frame.frame_index}')
404 else:
405 logger.debug(f'Processing frame {frame.frame_index}')
407 for s1, s2 in pairs:
408 # compute correlation
409 Sq_pair = np.real(frame.rho_qs_dict[s1] * frame.rho_qs_dict[s2].conjugate())
410 if s1 != s2:
411 Sq_pair += np.real(frame.rho_qs_dict[s2] * frame.rho_qs_dict[s1].conjugate())
412 Sq_averager[(s1, s2)].add_sample(0, Sq_pair)
414 # collect results
415 data_dict = dict()
416 data_dict['q_points'] = q_points
417 S_q_tot = np.zeros((n_qpoints, 1))
418 for s1, s2 in pairs:
419 Sq = 1 / traj.n_atoms * Sq_averager[(s1, s2)].get_average_at_timelag(0).reshape(-1, 1)
420 data_dict[f'Sq_{s1}_{s2}'] = Sq
421 S_q_tot += Sq
422 data_dict['Sq'] = S_q_tot
424 # finalize results
425 new_sample = StaticSample(
426 data_dict,
427 simulation_data=dict(
428 atom_types=traj.atom_types,
429 pairs=pairs,
430 particle_counts=particle_counts,
431 cell=traj.cell,
432 number_of_frames=traj.number_of_frames_read,
433 ))
434 new_sample._append_history('compute_static_structure_factors')
436 return new_sample
439def compute_spectral_energy_density(
440 traj: Trajectory,
441 ideal_supercell: Atoms,
442 primitive_cell: Atoms,
443 q_points: NDArray[float],
444 dt: float,
445 partial: Optional[bool] = False,
446 logging_interval: Optional[int] = 1000,
447) -> tuple[NDArray[float], NDArray[float]]:
448 r"""
449 Compute the spectral energy density (SED) at specific q-points. The results
450 are returned in the form of a tuple, which comprises the angular
451 frequencies in an array of length ``N_times`` in units of rad/fs and the
452 SED in units of eV/(rad/fs) as an array of shape ``(N_qpoints, N_times)``.
453 The normalization is chosen such that integrating the SED of a q-point
454 together with the supplied angular frequencies omega (rad/fs) yields
455 `1/2 kB T` * number of bands (where number of bands = `len(prim) * 3`)
457 More details can be found in Thomas *et al.*, Physical Review B **81**, 081411 (2010),
458 which should be cited when using this function along with the dynasor reference.
460 **Note 1:**
461 SED analysis is only suitable for crystalline materials without diffusion as
462 atoms are assumed to move around fixed reference positions throughout the entire trajectory.
464 **Note 2:**
465 This implementation reads the full trajectory and can thus consume a lot of memory.
467 Parameters
468 ----------
469 traj
470 Input trajectory.
471 ideal_supercell
472 Ideal structure defining the reference positions. Do not change the masses
473 in the ASE :class:`Atoms` objects to dynasor internal units, this will be
474 done internally. Its atom count is checked against :attr:`traj`; a mismatched
475 cell only triggers a warning (some thermal expansion relative to the reference
476 structure is normal), and the atom ordering is not checked at all, so it must
477 match the trajectory's.
478 primitive_cell
479 Underlying primitive structure. Must be aligned correctly with :attr:`ideal_supercell`.
480 q_points
481 Array of q-points in units of rad/Å with shape ``(N_qpoints, 3)`` in Cartesian coordinates.
482 dt
483 Time difference in femtoseconds between two consecutive snapshots in
484 the trajectory. Note that you should not change :attr:`dt` if you change
485 :attr:`frame_step <dynasor.trajectory.Trajectory.frame_step>` in :attr:`traj`.
486 partial
487 If True the SED will be returned decomposed per basis and Cartesian direction.
488 The shape is ``(N_qpoints, N_frequencies, len(primitive_cell), 3)``.
489 logging_interval
490 Log progress at ``INFO`` level every this many frames. Set to ``0`` to disable
491 progress logging.
492 """
494 if q_points.ndim != 2 or q_points.shape[1] != 3:
495 raise ValueError('q-points array has the wrong shape.')
496 if dt <= 0:
497 raise ValueError(f'dt must be positive: dt= {dt}')
499 delta_t = traj.frame_step * dt
501 # logger
502 logger.info('Running SED')
503 logger.info(f'Time between consecutive frames (dt * frame_step): {delta_t} fs')
504 logger.info(f'Number of atoms in primitive_cell: {len(primitive_cell)}')
505 logger.info(f'Number of atoms in ideal_supercell: {len(ideal_supercell)}')
506 logger.info(f'Number of q-points: {q_points.shape[0]}')
508 # check that the ideal supercell agrees with traj
509 if traj.n_atoms != len(ideal_supercell):
510 raise ValueError('ideal_supercell must contain the same number of atoms as the trajectory.')
511 if not np.allclose(traj.cell, ideal_supercell.cell, atol=1e-5, rtol=0.0):
512 logger.warning('ideal_supercell cell does not match the trajectory cell.')
514 if len(primitive_cell) > len(ideal_supercell):
515 raise ValueError('primitive_cell contains more atoms than ideal_supercell.')
517 if not traj.has_velocities:
518 raise ValueError('compute_spectral_energy_density requires velocities to be available '
519 'in the trajectory, but traj does not provide velocities.')
521 # collect all velocities, and scale with sqrt(masses)
522 masses = ideal_supercell.get_masses().reshape(-1, 1) / fs**2 # From Dalton to dmu
523 velocities = []
524 for it, frame in enumerate(traj):
525 if logging_interval and it % logging_interval == 0:
526 logger.info(f'Reading frame {it}')
527 else:
528 logger.debug(f'Reading frame {it}')
529 v = frame.get_velocities_as_array(traj.atomic_indices) # in Å/fs
530 velocities.append(np.sqrt(masses) * v)
531 logger.info(f'Number of snapshots: {len(velocities)}')
533 # Perform the FFT on the last axis for extra speed (maybe not needed)
534 N_samples = len(velocities)
535 velocities = np.array(velocities)
536 # places time index last and makes a copy for continuity
537 velocities = velocities.transpose(1, 2, 0).copy()
538 # #atoms in supercell x 3 directions x #frequencies
539 velocities = np.fft.rfft(velocities, axis=2)
541 # Calculate indices and offsets needed for the SED method
542 offsets, indices = get_offset_index(primitive_cell, ideal_supercell)
544 # Phase factor for use in FT. #qpoints x #atoms in supercell
545 cell_positions = np.dot(offsets, primitive_cell.cell)
546 phase = np.dot(q_points, cell_positions.T) # #qpoints x #unit cells
547 phase_factors = np.exp(1.0j * phase)
549 # This dict maps the offsets to an index so ndarrays can be over
550 # offset,index instead of atoms in supercell
551 offset_dict = {off: n for n, off in enumerate(set(tuple(offset) for offset in offsets))}
553 # Pick out some shapes
554 n_super, _, n_w = velocities.shape
555 n_qpts = len(q_points)
556 n_prim = len(primitive_cell)
557 n_offsets = len(offset_dict)
559 # This new array will be indexed by index and offset instead (and also transposed)
560 new_velocities = np.zeros((n_w, 3, n_prim, n_offsets), dtype=velocities.dtype)
562 for i in range(n_super):
563 j = indices[i] # atom with index i in the supercell is of basis type j ...
564 n = offset_dict[tuple(offsets[i])] # and its offset has index n
565 new_velocities[:, :, j, n] = velocities[i].T
567 velocities = new_velocities
569 # Same story with the spatial phase factors
570 new_phase_factors = np.zeros((n_qpts, n_prim, n_offsets), dtype=phase_factors.dtype)
572 for i in range(n_super):
573 j = indices[i]
574 n = offset_dict[tuple(offsets[i])]
575 new_phase_factors[:, j, n] = phase_factors[:, i]
577 phase_factors = new_phase_factors
579 # calculate the density in a numba function
580 density = _sed_inner_loop(phase_factors, velocities)
582 if not partial:
583 density = np.sum(density, axis=(2, 3))
585 # units
586 # make so that the velocities were originally in Angstrom / fs to be compatible with eV and Da
588 # the time delta in the fourier transform
589 density = density * delta_t**2
591 # Divide by the length of the time signal
592 density = density / (N_samples * delta_t)
594 # Divide by the number of primitive cells
595 density = density / (n_super / n_prim)
597 # Factor so the sed can be integrated together with the returned omega
598 # numpy fft works with ordinary/linear frequencies and not angular freqs
599 density = density / (2*np.pi)
601 # angular frequencies
602 w = 2 * np.pi * np.fft.rfftfreq(N_samples, delta_t) # rad/fs
604 return w, density
607@numba.njit(parallel=True, fastmath=True)
608def _sed_inner_loop(phase_factors, velocities):
609 """This numba function calculates the spatial
610 Fourier transform using precomputed phase factors.
612 As the use case can be one or many q-points the parallelization is over the
613 temporal frequency components instead.
614 """
616 n_qpts = phase_factors.shape[0] # q-point index
617 n_prim = phase_factors.shape[1] # basis atom index
618 n_super = phase_factors.shape[2] # unit cell index
620 n_freqs = velocities.shape[0] # frequency, direction, basis atom, unit cell
622 density = np.zeros((n_qpts, n_freqs, n_prim, 3), dtype=np.float64)
624 for w in numba.prange(n_freqs):
625 for k in range(n_qpts):
626 for a in range(3):
627 for b in range(n_prim):
628 tmp = 0.0j
629 for n in range(n_super):
630 tmp += phase_factors[k, b, n] * velocities[w, a, b, n]
631 density[k, w, b, a] += np.abs(tmp)**2
632 return density