Coverage for dynasor/modes/complex_coordinate.py: 97%

1from __future__ import annotations # for 3.9 for typehint 'Band' in cc.band 

2from typing import TYPE_CHECKING 

3import math 

4 

5import numpy as np 

6 

7from ..units import radians_per_fs_to_THz 

8 

9 

10if TYPE_CHECKING: # pragma: no cover 

11 from .band import Band 

12 

13 

14class ComplexCoordinate: 

15 """Class to work with the in general complex coordinates of lattice dynamics. 

16 

17 Can be cast to complex by ``complex(cc)``. 

18 

19 Concrete subclasses :class:`Q`, :class:`P`, and :class:`F` represent the 

20 mode coordinate, momentum, and force respectively, and are accessed via 

21 :attr:`~dynasor.modes.band.Band.Q`, :attr:`~dynasor.modes.band.Band.P`, 

22 and :attr:`~dynasor.modes.band.Band.F` on a :class:`~dynasor.modes.band.Band` object. 

23 

24 Example 

25 ------- 

26 >>> # Access a coordinate via a Band object; mp is a ModeProjector instance, 

27 >>> # mp[0] selects q-point 0 and [1] selects band 1 within that q-point. 

28 >>> band = mp[0][1] # doctest: +SKIP 

29 >>> # The coordinate can be set directly as a complex number, 

30 >>> # or via its amplitude and phase independently. 

31 >>> band.Q.complex = 1.7j # doctest: +SKIP 

32 >>> band.Q.amplitude = 3.8 # doctest: +SKIP 

33 >>> band.Q.phase = 2*np.pi # doctest: +SKIP 

34 """ 

35 

36 def __init__(self, q, s, mp): 

37 """The complex coordinate belongs to a q-point and band.""" 

38 self._q = q 

39 self._s = s 

40 self._mp = mp 

41 

42 def __str__(self): 

43 return f'{self.__class__.__name__}(q={self._q}, s={self._s})' 

44 

45 def __repr__(self): 

46 return str(self) 

47 

48 def __complex__(self): 

49 return complex(getattr(self._mp, 'get_' + self.__class__.__name__)()[self._q, self._s]) 

50 

51 @property 

52 def band(self) -> 'Band': 

53 """The band to which this coordinate belongs to.""" 

54 return self._mp[self._q][self._s] 

55 

56 @property 

57 def complex(self): 

58 """The complex coordinate as a complex number.""" 

59 return complex(self) 

60 

61 @complex.setter 

62 def complex(self, C): 

63 C = complex(C) 

64 prop = self.__class__.__name__ 

65 

66 val = getattr(self._mp, 'get_' + prop)() 

67 

68 if self._mp[self._q].is_real: 68 ↛ 72line 68 didn't jump to line 72 because the condition on line 68 was always true

69 assert np.isclose(C.imag, 0) 

70 q2 = self._q 

71 else: 

72 q2 = self._mp[self._q].q_minus.index 

73 

74 val[self._q, self._s] = C 

75 val[q2, self._s] = C.conjugate() 

76 

77 getattr(self._mp, 'set_' + prop)(val) 

78 

79 @property 

80 def phase(self): 

81 """The phase of the complex coordinate.""" 

82 c = complex(self) 

83 return math.atan2(c.imag, c.real) 

84 

85 @phase.setter 

86 def phase(self, phase): 

87 C = complex(self) 

88 C = np.abs(C) * np.exp(1j * phase) 

89 self.complex = C 

90 

91 @property 

92 def amplitude(self): 

93 """The amplitude of the complex coordinate.""" 

94 c = complex(self) 

95 return np.abs(c) 

96 

97 @amplitude.setter 

98 def amplitude(self, amplitude): 

99 C = amplitude * np.exp(1j * self.phase) 

100 self.complex = C 

101 

102 

103class Q(ComplexCoordinate): 

104 def __str__(self): 

105 s = ['### Mode coordinate Q ###'] 

106 s += [f'band: [{self.band.index}] {self.band.omega * radians_per_fs_to_THz:.2f} THz'] 

107 s += [f'q-point: [{self.band.qpoint.index}] {self.band.qpoint.q_reduced}'] 

108 s += [f'Value: {self.complex:.2f} Å√dmu'] 

109 s += [f'Amplitude: {self.amplitude:.2f} Å√dmu'] 

110 s += [f'Phase: {self.phase:.2f} radians'] 

111 return '\n'.join(s) 

112 

113 

114class P(ComplexCoordinate): 

115 def __str__(self): 

116 s = ['### Mode momentum P ###'] 

117 s += [f'band: [{self.band.index}] {self.band.omega * radians_per_fs_to_THz:.2f} THz'] 

118 s += [f'q-point: [{self.band.qpoint.index}] {self.band.qpoint.q_reduced}'] 

119 s += [f'Value: {self.complex:.2f} Å√dmu/fs'] 

120 s += [f'Amplitude: {self.amplitude:.2f} Å√dmu/fs'] 

121 s += [f'Phase: {self.phase:.2f} radians'] 

122 return '\n'.join(s) 

123 

124 

125class F(ComplexCoordinate): 

126 def __str__(self): 

127 s = ['### Mode force F ###'] 

128 s += [f'band: [{self.band.index}] {self.band.omega * radians_per_fs_to_THz:.2f} THz'] 

129 s += [f'q-point: [{self.band.qpoint.index}] {self.band.qpoint.q_reduced}'] 

130 s += [f'Value: {self.complex:.2f} Å√dmu/fs²'] 

131 s += [f'Amplitude: {self.amplitude:.2f} Å√dmu/fs²'] 

132 s += [f'Phase: {self.phase:.2f} radians'] 

133 return '\n'.join(s)