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.modes.atoms import DynasorAtoms 

10from dynasor.modes.tools import mode_coordinates, mode_forces, mode_momenta 

11from dynasor.trajectory import Trajectory 

12from dynasor.tools.structures import get_displacements_from_u 

13 

14 

15def _validated_masses(ideal_supercell: Atoms) -> NDArray[float]: 

16 """Returns the masses of the supercell in dmu, checking that every one of them can be 

17 divided by.""" 

18 masses = DynasorAtoms(ideal_supercell).masses 

19 bad = np.flatnonzero(~np.isfinite(masses) | (masses <= 0)) 

20 if len(bad) > 0: 

21 raise ValueError( 

22 'Projecting forces requires a positive and finite mass for every atom of ' 

23 f'ideal_supercell, but the masses of atoms {bad.tolist()} are not.') 

24 return masses 

25 

26 

27def project_modes( 

28 traj: Trajectory, 

29 modes: NDArray[float], 

30 ideal_supercell: Atoms, 

31 check_mic: Optional[bool] = True, 

32 logging_interval: Optional[int] = 1000, 

33) -> tuple[NDArray[complex], NDArray[complex], NDArray[complex]]: 

34 r"""Projects an atomic trajectory onto set of phonon modes. 

35 

36 The projection follows the conventions of :class:`ModeProjector 

37 <dynasor.ModeProjector>`, 

38 

39 .. math:: 

40 

41 Q = X u, \quad P = X^* v, \quad F = X^* f / m 

42 

43 where :math:`X` are the modes, :math:`u` the displacements, :math:`v` the velocities, 

44 :math:`f` the forces and :math:`m` the atomic masses. 

45 The mode forces are conjugate to the mode coordinates, meaning 

46 :math:`F = -\partial V / \partial Q`, so the virial energy of a mode is :math:`-QF/2`. 

47 See :attr:`ModeProjector.virial_energies <dynasor.ModeProjector.virial_energies>` for that 

48 quantity and for how its imaginary parts behave. 

49 

50 Parameters 

51 ---------- 

52 traj 

53 Input trajectory. 

54 modes 

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

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

57 ideal_supercell 

58 Ideal supercell used to find atomic displacements and to provide the atomic masses. 

59 The masses have to be the ones that :attr:`modes` was built with, which for modes from 

60 a :class:`ModeProjector <dynasor.ModeProjector>` are the supercell masses of that 

61 projector. 

62 Neither the atom ordering nor the masses are checked, and a mismatch changes `F` per 

63 atom instead of by an overall factor. 

64 The atom count is checked against :attr:`traj`, and a mismatched cell only triggers a 

65 warning, since some thermal expansion relative to the reference structure is normal. 

66 check_mic 

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

68 logging_interval 

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

70 progress logging. 

71 

72 Returns 

73 ------- 

74 A tuple comprising `(Q, P, F)` where `Q` are the mode coordinates in Å√dmu, `P` are 

75 the mode momenta in √eV and `F` are the mode forces in eV/Å√dmu. 

76 Each array has shape ``(n_frames, *modes.shape[:-2])``, where ``n_frames`` is the 

77 number of frames read from :attr:`traj`. 

78 `P` is all zeros if :attr:`traj` provides no velocities, and `F` is all zeros if it 

79 provides no forces. 

80 

81 Raises 

82 ------ 

83 ValueError 

84 If :attr:`traj` provides no positions, if :attr:`modes` does not have shape 

85 ``(..., N, 3)``, if :attr:`ideal_supercell` and :attr:`traj` disagree on the number of 

86 atoms, if a frame past the first provides no positions, velocities or forces, or if 

87 :attr:`traj` provides forces while :attr:`ideal_supercell` gives an atom a mass that 

88 is not positive and finite. 

89 

90 Examples 

91 -------- 

92 The mode forces give the virial energy of every mode along a trajectory:: 

93 

94 >>> Q, P, F = project_modes(traj, mp.eigenmodes, mp.supercell.to_ase()) # doctest: +SKIP 

95 >>> virial_energies = (-Q * F / 2).real # doctest: +SKIP 

96 

97 """ 

98 # logger 

99 logger.info('Running mode projection') 

100 if not traj.has_positions: 

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

102 'but traj does not provide positions.') 

103 has_velocities = traj.has_velocities 

104 has_forces = traj.has_forces 

105 if not has_velocities: 

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

107 if not has_forces: 

108 logger.info('traj does not provide forces; F will be returned as zeros') 

109 

110 modes = np.asarray(modes) 

111 

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

113 raise ValueError( 

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

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

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

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

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

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

120 

121 if has_forces: 

122 masses = _validated_masses(ideal_supercell) 

123 

124 Q_traj, P_traj, F_traj = [], [], [] 

125 for it, frame in enumerate(traj): 

126 if logging_interval and it % logging_interval == 0: 

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

128 else: 

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

130 

131 # Make positions into displacements 

132 x = frame.get_positions_as_array(traj._atomic_indices) 

133 u = x - ideal_supercell.positions 

134 

135 # Calculate Q 

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

137 Q_traj.append(mode_coordinates(modes, u)) 

138 

139 # Calculate P 

140 if has_velocities: 

141 v = frame.get_velocities_as_array(traj._atomic_indices) 

142 P_traj.append(mode_momenta(modes, v)) 

143 

144 # Calculate F 

145 if has_forces: 

146 f = frame.get_forces_as_array(traj._atomic_indices) 

147 F_traj.append(mode_forces(modes, f, masses)) 

148 

149 # The reshape only acts on a trajectory of no frames, which would otherwise lose the 

150 # leading mode dimensions. 

151 Q_traj = np.array(Q_traj).reshape((-1, *modes.shape[:-2])) 

152 P_traj = np.array(P_traj).reshape(Q_traj.shape) if has_velocities else np.zeros_like(Q_traj) 

153 F_traj = np.array(F_traj).reshape(Q_traj.shape) if has_forces else np.zeros_like(Q_traj) 

154 

155 return Q_traj, P_traj, F_traj