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

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) -> tuple[NDArray[float], NDArray[float]]: 

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

20 

21 Parameters 

22 ---------- 

23 traj 

24 Input trajectory. 

25 modes 

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

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

28 ideal_supercell 

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

30 structure. Be careful not to mess up the permutation. 

31 check_mic 

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

33 

34 Returns 

35 ------- 

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

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

38 complex array with dimension (length of traj, number of modes). 

39 """ 

40 # logger 

41 logger.info('Running mode projection') 

42 

43 modes = np.array(modes) 

44 

45 original_mode_shape = modes.shape 

46 

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

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

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

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

51 

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

53 

54 Q_traj, P_traj = [], [] 

55 for it, frame in enumerate(traj): 

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

57 

58 # Make positions into displacements 

59 x = frame.get_positions_as_array(traj._atomic_indices) 

60 u = x - ideal_supercell.positions 

61 

62 # Calculate Q 

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

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

65 

66 # Calculate P 

67 if frame.velocities_by_type is not None: 

68 v = frame.get_velocities_as_array(traj._atomic_indices) 

69 P = np.einsum('mna,na->m', modes.conj(), v) 

70 else: 

71 P = np.zeros_as(Q) 

72 

73 Q_traj.append(Q) 

74 P_traj.append(P) 

75 

76 Q_traj = Q_traj.reshape((len(Q_traj), *original_mode_shape[:-2])) 

77 P_traj = P_traj.reshape((len(P_traj), *original_mode_shape[:-2])) 

78 

79 return np.array(Q_traj), np.array(P_traj)