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

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_positions: 50 ↛ 51line 50 didn't jump to line 51 because the condition on line 50 was never true

51 raise ValueError('project_modes requires positions to be available in the trajectory, ' 

52 'but traj does not provide positions.') 

53 if not traj.has_velocities: 

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

55 

56 modes = np.asarray(modes) 

57 

58 original_mode_shape = modes.shape 

59 

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

61 raise ValueError( 

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

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

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

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

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

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

68 

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

70 modes_conj = modes.conj() 

71 

72 Q_traj, P_traj = [], [] 

73 for it, frame in enumerate(traj): 

74 if logging_interval and it % logging_interval == 0: 

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

76 else: 

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

78 

79 # Make positions into displacements 

80 x = frame.get_positions_as_array(traj._atomic_indices) 

81 u = x - ideal_supercell.positions 

82 

83 # Calculate Q 

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

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

86 

87 # Calculate P 

88 if traj.has_velocities: 

89 v = frame.get_velocities_as_array(traj._atomic_indices) 

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

91 else: 

92 P = np.zeros_like(Q) 

93 

94 Q_traj.append(Q) 

95 P_traj.append(P) 

96 

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

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

99 

100 return Q_traj, P_traj