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

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

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 __repr__(self): 

43 return str(self) 

44 

45 def __complex__(self): 

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

47 

48 @property 

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

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

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

52 

53 @property 

54 def complex(self): 

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

56 return complex(self) 

57 

58 @complex.setter 

59 def complex(self, C): 

60 C = complex(C) 

61 prop = self.__class__.__name__ 

62 

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

64 

65 if self._mp[self._q].is_real: 

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

67 q2 = self._q 

68 else: 

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

70 

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

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

73 

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

75 

76 @property 

77 def phase(self): 

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

79 c = complex(self) 

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

81 

82 @phase.setter 

83 def phase(self, phase): 

84 C = complex(self) 

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

86 self.complex = C 

87 

88 @property 

89 def amplitude(self): 

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

91 c = complex(self) 

92 return np.abs(c) 

93 

94 @amplitude.setter 

95 def amplitude(self, amplitude): 

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

97 self.complex = C 

98 

99 

100class Q(ComplexCoordinate): 

101 def __str__(self): 

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

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

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

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

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

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

108 return '\n'.join(s) 

109 

110 

111class P(ComplexCoordinate): 

112 def __str__(self): 

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

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

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

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

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

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

119 return '\n'.join(s) 

120 

121 

122class F(ComplexCoordinate): 

123 def __str__(self): 

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

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

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

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

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

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

130 return '\n'.join(s)