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. 

32 check_mic 

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

34 logging_interval 

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

36 progress logging. 

37 

38 Returns 

39 ------- 

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

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

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

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

44 """ 

45 # logger 

46 logger.info('Running mode projection') 

47 if not traj.has_velocities: 

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

49 

50 modes = np.asarray(modes) 

51 

52 original_mode_shape = modes.shape 

53 

54 if modes.shape[-2] != traj.n_atoms: 

55 raise ValueError('Second dim in modes must be same len as number of atoms in trajectory') 

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

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

58 

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

60 modes_conj = modes.conj() 

61 

62 Q_traj, P_traj = [], [] 

63 for it, frame in enumerate(traj): 

64 if logging_interval and it % logging_interval == 0: 

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

66 else: 

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

68 

69 # Make positions into displacements 

70 x = frame.get_positions_as_array(traj._atomic_indices) 

71 u = x - ideal_supercell.positions 

72 

73 # Calculate Q 

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

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

76 

77 # Calculate P 

78 if traj.has_velocities: 

79 v = frame.get_velocities_as_array(traj._atomic_indices) 

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

81 else: 

82 P = np.zeros_like(Q) 

83 

84 Q_traj.append(Q) 

85 P_traj.append(P) 

86 

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

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

89 

90 return Q_traj, P_traj