Coverage for dynasor/correlation_functions.py: 100%
285 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-20 20:02 +0000
1import concurrent
2from functools import partial
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 Length of the trajectory frame window to use for time correlation calculation.
49 It is expressed in terms of the number of time lags to consider
50 and thus determines the smallest frequency resolved.
51 window_step
52 Window step (or stride) given as the number of frames between consecutive trajectory
53 windows. This parameter does *not* affect the time between consecutive frames in the
54 calculation. If, e.g., :attr:`window_step` > :attr:`window_size`, some frames will not
55 be used.
56 calculate_currents
57 Calculate the current correlations. Requires velocities to be available in :attr:`traj`.
58 calculate_incoherent
59 Calculate the incoherent part (self-part) of :math:`F_\text{incoh}`.
60 logging_interval
61 Log progress at ``INFO`` level every this many windows. Set to ``0`` to disable
62 progress logging.
63 """
64 # sanity check input args
65 if q_points.shape[1] != 3:
66 raise ValueError('q-points array has the wrong shape.')
67 if dt <= 0:
68 raise ValueError(f'dt must be positive: dt= {dt}')
69 if window_size <= 2:
70 raise ValueError(f'window_size must be larger than 2: window_size= {window_size}')
71 if window_step <= 0:
72 raise ValueError(f'window_step must be positive: window_step= {window_step}')
73 if calculate_currents and not traj.has_velocities:
74 raise ValueError('calculate_currents=True requires velocities to be available in the '
75 'trajectory, but traj does not provide velocities.')
77 # define internal parameters
78 n_qpoints = q_points.shape[0]
79 delta_t = traj.frame_step * dt
80 N_tc = window_size + 1
82 # log all setup information
83 dw = np.pi / (window_size * delta_t)
84 w_max = dw * window_size
85 w_N = 2 * np.pi / (2 * delta_t) # Nyquist angular frequency
87 logger.info(f'Spacing between samples (frame_step): {traj.frame_step}')
88 logger.info(f'Time between consecutive frames in input trajectory (dt): {dt} fs')
89 logger.info(f'Time between consecutive frames used (dt * frame_step): {delta_t} fs')
90 logger.info(f'Time window size (dt * frame_step * window_size): {delta_t * window_size:.1f} fs')
91 logger.info(f'Angular frequency resolution: dw = {dw:.6f} rad/fs = '
92 f'{dw * radians_per_fs_to_meV:.3f} meV')
93 logger.info(f'Maximum angular frequency (dw * window_size):'
94 f' {w_max:.6f} rad/fs = {w_max * radians_per_fs_to_meV:.3f} meV')
95 logger.info(f'Nyquist angular frequency (2pi / frame_step / dt / 2):'
96 f' {w_N:.6f} rad/fs = {w_N * radians_per_fs_to_meV:.3f} meV')
98 if calculate_currents:
99 logger.info('Calculating current (velocity) correlations')
100 if calculate_incoherent:
101 logger.info('Calculating incoherent part (self-part) of correlations')
103 # log some info regarding q-points
104 logger.info(f'Number of q-points: {n_qpoints}')
106 q_directions = q_points.astype(float, copy=True)
107 q_distances = np.linalg.norm(q_points, axis=1)
108 nonzero = q_distances > 0
109 q_directions[nonzero] /= q_distances[nonzero].reshape(-1, 1)
111 # setup functions to process frames
112 def f2_rho(frame):
113 rho_qs_dict = dict()
114 for atom_type in frame.positions_by_type.keys():
115 x = frame.positions_by_type[atom_type]
116 rho_qs_dict[atom_type] = calc_rho_q(x, q_points)
117 frame.rho_qs_dict = rho_qs_dict
118 return frame
120 def f2_rho_and_j(frame):
121 rho_qs_dict = dict()
122 jz_qs_dict = dict()
123 jper_qs_dict = dict()
125 for atom_type in frame.positions_by_type.keys():
126 x = frame.positions_by_type[atom_type]
127 v = frame.velocities_by_type[atom_type]
128 rho_qs, j_qs = calc_rho_j_q(x, v, q_points)
129 jz_qs = np.sum(j_qs * q_directions, axis=1)
130 jper_qs = j_qs - (jz_qs[:, None] * q_directions)
132 rho_qs_dict[atom_type] = rho_qs
133 jz_qs_dict[atom_type] = jz_qs
134 jper_qs_dict[atom_type] = jper_qs
136 frame.rho_qs_dict = rho_qs_dict
137 frame.jz_qs_dict = jz_qs_dict
138 frame.jper_qs_dict = jper_qs_dict
139 return frame
141 if calculate_currents:
142 element_processor = f2_rho_and_j
143 else:
144 element_processor = f2_rho
146 # setup window iterator
147 window_iterator = WindowIterator(traj, width=N_tc, window_step=window_step,
148 element_processor=element_processor)
150 # define all pairs
151 pairs = list(combinations_with_replacement(traj.atom_types, r=2))
152 particle_counts = {key: len(val) for key, val in traj.atomic_indices.items()}
153 logger.debug('Considering pairs:')
154 for pair in pairs:
155 logger.debug(f' {pair}')
157 # set up all time averager instances
158 F_q_t_averager = dict()
159 for pair in pairs:
160 F_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
161 if calculate_currents:
162 Cl_q_t_averager = dict()
163 Ct_q_t_averager = dict()
164 for pair in pairs:
165 Cl_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
166 Ct_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
167 if calculate_incoherent:
168 F_s_q_t_averager = dict()
169 for pair in traj.atom_types:
170 F_s_q_t_averager[pair] = TimeAverager(N_tc, n_qpoints)
172 # define correlation function
173 def calc_corr(window, time_i):
174 # Calculate correlations between two frames in the window without normalization 1/N
175 f0 = window[0]
176 fi = window[time_i]
177 for s1, s2 in pairs:
178 Fqt = np.real(f0.rho_qs_dict[s1] * fi.rho_qs_dict[s2].conjugate())
179 if s1 != s2:
180 Fqt += np.real(f0.rho_qs_dict[s2] * fi.rho_qs_dict[s1].conjugate())
181 F_q_t_averager[(s1, s2)].add_sample(time_i, Fqt)
183 if calculate_currents:
184 for s1, s2 in pairs:
185 Clqt = np.real(f0.jz_qs_dict[s1] * fi.jz_qs_dict[s2].conjugate())
186 Ctqt = 0.5 * np.real(np.sum(f0.jper_qs_dict[s1] *
187 fi.jper_qs_dict[s2].conjugate(), axis=1))
188 if s1 != s2:
189 Clqt += np.real(f0.jz_qs_dict[s2] * fi.jz_qs_dict[s1].conjugate())
190 Ctqt += 0.5 * np.real(np.sum(f0.jper_qs_dict[s2] *
191 fi.jper_qs_dict[s1].conjugate(), axis=1))
193 Cl_q_t_averager[(s1, s2)].add_sample(time_i, Clqt)
194 Ct_q_t_averager[(s1, s2)].add_sample(time_i, Ctqt)
196 if calculate_incoherent:
197 for atom_type in traj.atom_types:
198 xi = fi.positions_by_type[atom_type]
199 x0 = f0.positions_by_type[atom_type]
200 Fsqt = np.real(calc_rho_q(xi - x0, q_points))
201 F_s_q_t_averager[atom_type].add_sample(time_i, Fsqt)
203 # run calculation
204 with concurrent.futures.ThreadPoolExecutor() as tpe:
205 # This is the "main loop" over the trajectory
206 for window in window_iterator:
207 if logging_interval and window[0].frame_index % logging_interval == 0:
208 logger.info(f'Processing window {window[0].frame_index} to {window[-1].frame_index}') # noqa
209 else:
210 logger.debug(f'Processing window {window[0].frame_index} to {window[-1].frame_index}') # noqa
212 # The map conveniently applies calc_corr to all time-lags. However,
213 # as everything is done in place nothing gets returned so in order
214 # to start and wait for the processes to finish we must iterate
215 # over the None values returned
216 for _ in tpe.map(partial(calc_corr, window), range(len(window))):
217 pass
219 # collect results into dict with numpy arrays (n_qpoints, N_tc)
220 data_dict_corr = dict()
221 data_dict_corr['q_points'] = q_points
223 time = None
224 for pair in pairs:
225 key = '_'.join(pair)
226 F_q_t = 1 / traj.n_atoms * truncate_nans(F_q_t_averager[pair].get_average_all())
227 w, S_q_w = psd_from_acf_2d(F_q_t, delta_t)
228 S_q_w = np.array(S_q_w)
230 if time is None:
231 # Determine the actual length of time signal after truncation (if needed)
232 time = delta_t * np.arange(F_q_t.shape[1], dtype=float)
233 N_tc_actual = len(time)
234 F_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
235 S_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
236 else:
237 assert F_q_t.shape[1] == len(time)
239 data_dict_corr['omega'] = w
240 data_dict_corr[f'Fqt_coh_{key}'] = F_q_t
241 data_dict_corr[f'Sqw_coh_{key}'] = S_q_w
243 # sum all partials to the total
244 F_q_t_tot += F_q_t
245 S_q_w_tot += S_q_w
247 if N_tc_actual < N_tc:
248 logger.warning('Truncating ACF due to NaNs, likely time_window is longer than length of Trajectory') # noqa
249 data_dict_corr['time'] = time
250 data_dict_corr['Fqt_coh'] = F_q_t_tot
251 data_dict_corr['Sqw_coh'] = S_q_w_tot
253 if calculate_currents:
254 Cl_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
255 Ct_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
256 Cl_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
257 Ct_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
258 for pair in pairs:
259 key = '_'.join(pair)
260 Cl_q_t = 1 / traj.n_atoms * truncate_nans(Cl_q_t_averager[pair].get_average_all())
261 Ct_q_t = 1 / traj.n_atoms * truncate_nans(Ct_q_t_averager[pair].get_average_all())
262 _, Cl_q_w = psd_from_acf_2d(Cl_q_t, delta_t)
263 _, Ct_q_w = psd_from_acf_2d(Ct_q_t, delta_t)
264 data_dict_corr[f'Clqt_{key}'] = Cl_q_t
265 data_dict_corr[f'Ctqt_{key}'] = Ct_q_t
266 data_dict_corr[f'Clqw_{key}'] = Cl_q_w
267 data_dict_corr[f'Ctqw_{key}'] = Ct_q_w
269 # sum all partials to the total
270 Cl_q_t_tot += Cl_q_t
271 Ct_q_t_tot += Ct_q_t
272 Cl_q_w_tot += Cl_q_w
273 Ct_q_w_tot += Ct_q_w
274 data_dict_corr['Clqt'] = Cl_q_t_tot
275 data_dict_corr['Ctqt'] = Ct_q_t_tot
276 data_dict_corr['Clqw'] = Cl_q_w_tot
277 data_dict_corr['Ctqw'] = Ct_q_w_tot
279 if calculate_incoherent:
280 Fs_q_t_tot = np.zeros((n_qpoints, N_tc_actual))
281 Ss_q_w_tot = np.zeros((n_qpoints, N_tc_actual))
282 for atom_type in traj.atom_types:
283 Fs_q_t = 1 / traj.n_atoms * truncate_nans(F_s_q_t_averager[atom_type].get_average_all())
284 _, Ss_q_w = psd_from_acf_2d(Fs_q_t, delta_t)
285 data_dict_corr[f'Fqt_incoh_{atom_type}'] = Fs_q_t
286 data_dict_corr[f'Sqw_incoh_{atom_type}'] = Ss_q_w
288 # sum all partials to the total
289 Fs_q_t_tot += Fs_q_t
290 Ss_q_w_tot += Ss_q_w
292 data_dict_corr['Fqt_incoh'] = Fs_q_t_tot
293 data_dict_corr['Sqw_incoh'] = Ss_q_w_tot
295 # finalize results with additional metadata
296 new_sample = DynamicSample(
297 data_dict_corr,
298 simulation_data=dict(
299 atom_types=traj.atom_types, pairs=pairs,
300 particle_counts=particle_counts,
301 cell=traj.cell,
302 time_between_frames=delta_t,
303 maximum_time_lag=delta_t * window_size,
304 angular_frequency_resolution=dw,
305 maximum_angular_frequency=w_max,
306 number_of_frames=traj.number_of_frames_read,
307 ))
308 new_sample._append_history(
309 'compute_dynamic_structure_factors',
310 dict(
311 dt=dt,
312 window_size=window_size,
313 window_step=window_step,
314 calculate_currents=calculate_currents,
315 calculate_incoherent=calculate_incoherent,
316 ))
318 return new_sample
321def compute_static_structure_factors(
322 traj: Trajectory,
323 q_points: NDArray[float],
324 logging_interval: Optional[int] = 1000,
325) -> StaticSample:
326 r"""Compute the static structure factors. The results are returned in the
327 form of a :class:`StaticSample <dynasor.sample.StaticSample>`
328 object.
330 Parameters
331 ----------
332 traj
333 Input trajectory.
334 q_points
335 Array of q-points in units of rad/Å with shape ``(N_qpoints, 3)`` in Cartesian coordinates.
336 logging_interval
337 Log progress at ``INFO`` level every this many frames. Set to ``0`` to disable
338 progress logging.
339 """
340 # sanity check input args
341 if q_points.shape[1] != 3:
342 raise ValueError('q-points array has the wrong shape.')
344 n_qpoints = q_points.shape[0]
345 logger.info(f'Number of q-points: {n_qpoints}')
347 # define all pairs
348 pairs = list(combinations_with_replacement(traj.atom_types, r=2))
349 particle_counts = {key: len(val) for key, val in traj.atomic_indices.items()}
350 logger.debug('Considering pairs:')
351 for pair in pairs:
352 logger.debug(f' {pair}')
354 # processing function
355 def f2_rho(frame):
356 rho_qs_dict = dict()
357 for atom_type in frame.positions_by_type.keys():
358 x = frame.positions_by_type[atom_type]
359 rho_qs_dict[atom_type] = calc_rho_q(x, q_points)
360 frame.rho_qs_dict = rho_qs_dict
361 return frame
363 # setup averager
364 Sq_averager = dict()
365 for pair in pairs:
366 Sq_averager[pair] = TimeAverager(1, n_qpoints) # time average with only timelag=0
368 # main loop
369 for frame in traj:
371 # process_frame
372 f2_rho(frame)
373 if logging_interval and frame.frame_index % logging_interval == 0:
374 logger.info(f'Processing frame {frame.frame_index}')
375 else:
376 logger.debug(f'Processing frame {frame.frame_index}')
378 for s1, s2 in pairs:
379 # compute correlation
380 Sq_pair = np.real(frame.rho_qs_dict[s1] * frame.rho_qs_dict[s2].conjugate())
381 if s1 != s2:
382 Sq_pair += np.real(frame.rho_qs_dict[s2] * frame.rho_qs_dict[s1].conjugate())
383 Sq_averager[(s1, s2)].add_sample(0, Sq_pair)
385 # collect results
386 data_dict = dict()
387 data_dict['q_points'] = q_points
388 S_q_tot = np.zeros((n_qpoints, 1))
389 for s1, s2 in pairs:
390 Sq = 1 / traj.n_atoms * Sq_averager[(s1, s2)].get_average_at_timelag(0).reshape(-1, 1)
391 data_dict[f'Sq_{s1}_{s2}'] = Sq
392 S_q_tot += Sq
393 data_dict['Sq'] = S_q_tot
395 # finalize results
396 new_sample = StaticSample(
397 data_dict,
398 simulation_data=dict(
399 atom_types=traj.atom_types,
400 pairs=pairs,
401 particle_counts=particle_counts,
402 cell=traj.cell,
403 number_of_frames=traj.number_of_frames_read,
404 ))
405 new_sample._append_history('compute_static_structure_factors')
407 return new_sample
410def compute_spectral_energy_density(
411 traj: Trajectory,
412 ideal_supercell: Atoms,
413 primitive_cell: Atoms,
414 q_points: NDArray[float],
415 dt: float,
416 partial: Optional[bool] = False,
417 logging_interval: Optional[int] = 1000,
418) -> tuple[NDArray[float], NDArray[float]]:
419 r"""
420 Compute the spectral energy density (SED) at specific q-points. The results
421 are returned in the form of a tuple, which comprises the angular
422 frequencies in an array of length ``N_times`` in units of rad/fs and the
423 SED in units of eV/(rad/fs) as an array of shape ``(N_qpoints, N_times)``.
424 The normalization is chosen such that integrating the SED of a q-point
425 together with the supplied angular frequencies omega (rad/fs) yields
426 `1/2 kB T` * number of bands (where number of bands = `len(prim) * 3`)
428 More details can be found in Thomas *et al.*, Physical Review B **81**, 081411 (2010),
429 which should be cited when using this function along with the dynasor reference.
431 **Note 1:**
432 SED analysis is only suitable for crystalline materials without diffusion as
433 atoms are assumed to move around fixed reference positions throughout the entire trajectory.
435 **Note 2:**
436 This implementation reads the full trajectory and can thus consume a lot of memory.
438 Parameters
439 ----------
440 traj
441 Input trajectory.
442 ideal_supercell
443 Ideal structure defining the reference positions. Do not change the masses
444 in the ASE :class:`Atoms` objects to dynasor internal units, this will be
445 done internally.
446 primitive_cell
447 Underlying primitive structure. Must be aligned correctly with :attr:`ideal_supercell`.
448 q_points
449 Array of q-points in units of rad/Å with shape ``(N_qpoints, 3)`` in Cartesian coordinates.
450 dt
451 Time difference in femtoseconds between two consecutive snapshots in
452 the trajectory. Note that you should not change :attr:`dt` if you change
453 :attr:`frame_step <dynasor.trajectory.Trajectory.frame_step>` in :attr:`traj`.
454 partial
455 If True the SED will be returned decomposed per basis and Cartesian direction.
456 The shape is ``(N_qpoints, N_frequencies, len(primitive_cell), 3)``.
457 logging_interval
458 Log progress at ``INFO`` level every this many frames. Set to ``0`` to disable
459 progress logging.
460 """
462 delta_t = traj.frame_step * dt
464 # logger
465 logger.info('Running SED')
466 logger.info(f'Time between consecutive frames (dt * frame_step): {delta_t} fs')
467 logger.info(f'Number of atoms in primitive_cell: {len(primitive_cell)}')
468 logger.info(f'Number of atoms in ideal_supercell: {len(ideal_supercell)}')
469 logger.info(f'Number of q-points: {q_points.shape[0]}')
471 # check that the ideal supercell agrees with traj
472 if traj.n_atoms != len(ideal_supercell):
473 raise ValueError('ideal_supercell must contain the same number of atoms as the trajectory.')
475 if len(primitive_cell) >= len(ideal_supercell):
476 raise ValueError('primitive_cell contains more atoms than ideal_supercell.')
478 if not traj.has_velocities:
479 raise ValueError('compute_spectral_energy_density requires velocities to be available '
480 'in the trajectory, but traj does not provide velocities.')
482 # colllect all velocities, and scale with sqrt(masses)
483 masses = ideal_supercell.get_masses().reshape(-1, 1) / fs**2 # From Dalton to dmu
484 velocities = []
485 for it, frame in enumerate(traj):
486 if logging_interval and it % logging_interval == 0:
487 logger.info(f'Reading frame {it}')
488 else:
489 logger.debug(f'Reading frame {it}')
490 v = frame.get_velocities_as_array(traj.atomic_indices) # in Å/fs
491 velocities.append(np.sqrt(masses) * v)
492 logger.info(f'Number of snapshots: {len(velocities)}')
494 # Perform the FFT on the last axis for extra speed (maybe not needed)
495 N_samples = len(velocities)
496 velocities = np.array(velocities)
497 # places time index last and makes a copy for continuity
498 velocities = velocities.transpose(1, 2, 0).copy()
499 # #atoms in supercell x 3 directions x #frequencies
500 velocities = np.fft.rfft(velocities, axis=2)
502 # Calculate indices and offsets needed for the SED method
503 offsets, indices = get_offset_index(primitive_cell, ideal_supercell)
505 # Phase factor for use in FT. #qpoints x #atoms in supercell
506 cell_positions = np.dot(offsets, primitive_cell.cell)
507 phase = np.dot(q_points, cell_positions.T) # #qpoints x #unit cells
508 phase_factors = np.exp(1.0j * phase)
510 # This dict maps the offsets to an index so ndarrays can be over
511 # offset,index instead of atoms in supercell
512 offset_dict = {off: n for n, off in enumerate(set(tuple(offset) for offset in offsets))}
514 # Pick out some shapes
515 n_super, _, n_w = velocities.shape
516 n_qpts = len(q_points)
517 n_prim = len(primitive_cell)
518 n_offsets = len(offset_dict)
520 # This new array will be indexed by index and offset instead (and also transposed)
521 new_velocities = np.zeros((n_w, 3, n_prim, n_offsets), dtype=velocities.dtype)
523 for i in range(n_super):
524 j = indices[i] # atom with index i in the supercell is of basis type j ...
525 n = offset_dict[tuple(offsets[i])] # and its offset has index n
526 new_velocities[:, :, j, n] = velocities[i].T
528 velocities = new_velocities
530 # Same story with the spatial phase factors
531 new_phase_factors = np.zeros((n_qpts, n_prim, n_offsets), dtype=phase_factors.dtype)
533 for i in range(n_super):
534 j = indices[i]
535 n = offset_dict[tuple(offsets[i])]
536 new_phase_factors[:, j, n] = phase_factors[:, i]
538 phase_factors = new_phase_factors
540 # calculate the density in a numba function
541 density = _sed_inner_loop(phase_factors, velocities)
543 if not partial:
544 density = np.sum(density, axis=(2, 3))
546 # units
547 # make so that the velocities were originally in Angstrom / fs to be compatible with eV and Da
549 # the time delta in the fourier transform
550 density = density * delta_t**2
552 # Divide by the length of the time signal
553 density = density / (N_samples * delta_t)
555 # Divide by the number of primitive cells
556 density = density / (n_super / n_prim)
558 # Factor so the sed can be integrated together with the returned omega
559 # numpy fft works with ordinary/linear frequencies and not angular freqs
560 density = density / (2*np.pi)
562 # angular frequencies
563 w = 2 * np.pi * np.fft.rfftfreq(N_samples, delta_t) # rad/fs
565 return w, density
568@numba.njit(parallel=True, fastmath=True)
569def _sed_inner_loop(phase_factors, velocities):
570 """This numba function calculates the spatial
571 Fourier transform using precomputed phase factors.
573 As the use case can be one or many q-points the parallelization is over the
574 temporal frequency components instead.
575 """
577 n_qpts = phase_factors.shape[0] # q-point index
578 n_prim = phase_factors.shape[1] # basis atom index
579 n_super = phase_factors.shape[2] # unit cell index
581 n_freqs = velocities.shape[0] # frequency, direction, basis atom, unit cell
583 density = np.zeros((n_qpts, n_freqs, n_prim, 3), dtype=np.float64)
585 for w in numba.prange(n_freqs):
586 for k in range(n_qpts):
587 for a in range(3):
588 for b in range(n_prim):
589 tmp = 0.0j
590 for n in range(n_super):
591 tmp += phase_factors[k, b, n] * velocities[w, a, b, n]
592 density[k, w, b, a] += np.abs(tmp)**2
593 return density