Coverage for dynasor/modes/project_modes.py: 100%

1from typing import Optional 

2 

3import numpy as np 

4 

5from ase import Atoms 

6from numpy.typing import NDArray 

7 

8from dynasor.logging_tools import logger 

9from dynasor.trajectory import Trajectory 

10from dynasor.tools.structures import get_displacements_from_u 

11 

12 

13def project_modes( 

14 traj: Trajectory, 

15 modes: NDArray[float], 

16 ideal_supercell: Atoms, 

17 check_mic: Optional[bool] = True, 

18 logging_interval: Optional[int] = 1000, 

19) -> tuple[NDArray[float], NDArray[float]]: 

20 """Projects an atomic trajectory onto set of phonon modes. 

21 

22 Parameters 

23 ---------- 

24 traj 

25 Input trajectory. 

26 modes 

27 Modes to project on, as an array with shape ``(..., N, 3)`` where ``N`` is the 

28 number of atoms in the supercell and the leading dimensions define the output shape. 

29 ideal_supercell 

30 Ideal supercell used to find atomic displacements. It should correspond to the ideal 

31 structure. Be careful not to mess up the permutation: its atom count is checked 

32 against :attr:`traj`; a mismatched cell only triggers a warning (some thermal 

33 expansion relative to the reference structure is normal), and the atom ordering 

34 is not checked at all. 

35 check_mic 

36 Whether to wrap the displacements or not, faster if no wrap. 

37 logging_interval 

38 Log progress at ``INFO`` level every this many frames. Set to ``0`` to disable 

39 progress logging. 

40 

41 Returns 

42 ------- 

43 A tuple comprising `(Q,P)` where `Q` are the mode coordinates as a complex array 

44 with dimension (length of traj, number of modes) and `P` are the mode momenta as a 

45 complex array with dimension (length of traj, number of modes). If :attr:`traj` does 

46 not provide velocities, `P` is returned as all zeros. 

47 """ 

48 # logger 

49 logger.info('Running mode projection') 

50 if not traj.has_velocities: 

51 logger.info('traj does not provide velocities; P will be returned as zeros') 

52 

53 modes = np.asarray(modes) 

54 

55 original_mode_shape = modes.shape 

56 

57 if modes.ndim < 2 or modes.shape[-1] != 3 or modes.shape[-2] != traj.n_atoms: 

58 raise ValueError( 

59 f'modes must have shape (..., N, 3), where N = {traj.n_atoms} is the number of ' 

60 f'atoms in the trajectory, but has shape {modes.shape}.') 

61 if traj.n_atoms != len(ideal_supercell): 

62 raise ValueError('ideal_supercell must contain the same number of atoms as the trajectory.') 

63 if not np.allclose(traj.cell, ideal_supercell.cell, atol=1e-5, rtol=0.0): 

64 logger.warning('ideal_supercell cell does not match the trajectory cell.') 

65 

66 modes = modes.reshape((-1, modes.shape[-2], 3)) 

67 modes_conj = modes.conj() 

68 

69 Q_traj, P_traj = [], [] 

70 for it, frame in enumerate(traj): 

71 if logging_interval and it % logging_interval == 0: 

72 logger.info(f'Reading frame {it}') 

73 else: 

74 logger.debug(f'Reading frame {it}') 

75 

76 # Make positions into displacements 

77 x = frame.get_positions_as_array(traj._atomic_indices) 

78 u = x - ideal_supercell.positions 

79 

80 # Calculate Q 

81 u = get_displacements_from_u(u, ideal_supercell.cell, check_mic=check_mic) 

82 Q = np.einsum('mnx,nx->m', modes, u, optimize=True) 

83 

84 # Calculate P 

85 if traj.has_velocities: 

86 v = frame.get_velocities_as_array(traj._atomic_indices) 

87 P = np.einsum('mnx,nx->m', modes_conj, v, optimize=True) 

88 else: 

89 P = np.zeros_like(Q) 

90 

91 Q_traj.append(Q) 

92 P_traj.append(P) 

93 

94 Q_traj = np.array(Q_traj).reshape((len(Q_traj), *original_mode_shape[:-2])) 

95 P_traj = np.array(P_traj).reshape((len(P_traj), *original_mode_shape[:-2])) 

96 

97 return Q_traj, P_traj