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

1from typing import Tuple 

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 .tools import get_wrapped_displacements 

11 

12 

13def project_modes(traj: Trajectory, modes: NDArray[float], ideal_supercell: Atoms, check_mic=True 

14 ) -> Tuple[NDArray[float], NDArray[float]]: 

15 """Projects an atomic trajectory onto set of phonon modes 

16 

17 Parameters 

18 ---------- 

19 traj 

20 Input trajectory 

21 modes 

22 Modes to project on (1, Nlambda, N, 3) array where Nm is the number of modes 

23 and N is the number of atoms in the supercell 

24 ideal_supercell 

25 Used to find atomic displacements and should correspond to the ideal 

26 structure. Be careful not to mess up the permutation 

27 check_mic 

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

29 

30 Returns 

31 ------- 

32 Q 

33 mode coordinates, complex ndarray (length of traj, number of modes) 

34 P 

35 mode momenta, complex ndarray (length of traj, number of modes) 

36 """ 

37 # logger 

38 logger.info('Running mode projection') 

39 

40 modes = np.array(modes) 

41 

42 original_mode_shape = modes.shape 

43 

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

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

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

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

48 

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

50 

51 Q_traj, P_traj = [], [] 

52 for it, frame in enumerate(traj): 

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

54 

55 # Make positions into displacements 

56 x = frame.get_positions_as_array(traj._atomic_indices) 

57 u = x - ideal_supercell.positions 

58 

59 # Calculate Q 

60 u = get_wrapped_displacements(u, ideal_supercell.cell, check_mic=check_mic) 

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

62 

63 # Calculate P 

64 if frame.velocities_by_type is not None: 

65 v = frame.get_velocities_as_array(traj._atomic_indices) 

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

67 else: 

68 P = np.zeros_as(Q) 

69 

70 Q_traj.append(Q) 

71 P_traj.append(P) 

72 

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

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

75 

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