Coverage for gpaw/test/test_gauss_wave.py: 100%

1# flake8: noqa 

2import time 

3import numpy as np 

4from gpaw.utilities.gauss import gaussian_wave 

5 

6 

7def test_gauss_wave(rng): 

8 sigma = 2.4 

9 C = 3 

10 G = 50 

11 

12 t = time.time() 

13 r_cG = rng.normal(size=C * G**3).reshape((C, G, G, G)) 

14 r0_c = rng.normal(size=C) 

15 k_c = rng.normal(size=C) 

16 A = rng.uniform() * np.exp(1j * rng.uniform(0, 2 * np.pi)) 

17 print('Allocation: %8.5f s' % (time.time() - t)) 

18 

19 # ------------------------------------------------------------------- 

20 

21 # Test case for real-part of gamma-point wave with normalized amplitude 

22 _gaussRGN = lambda r_cG, r0_c, sigma: 1/(sigma*np.pi**0.5)**1.5 \ 

23 * np.exp(-np.sum((r_cG-r0_c[:,np.newaxis,np.newaxis,np.newaxis])**2, axis=0)/(2*sigma**2)) 

24 

25 t = time.time() 

26 gs0_G = _gaussRGN(r_cG, r0_c, sigma) 

27 print('_gaussRGN: %8.5f s' % (time.time() - t)) 

28 

29 t = time.time() 

30 gs1_G = gaussian_wave(r_cG, r0_c, sigma) 

31 print('+gaussRGN: %8.5f s' % (time.time() - t)) 

32 

33 assert np.abs(gs0_G - 

34 gs1_G).max() < 1e-12, 'Max error %g' % np.abs(gs0_G - 

35 gs1_G).max() 

36 del gs0_G, gs1_G 

37 

38 # Test case for real-part of gamma-point wave with complex amplitude 

39 _gaussRGA = lambda r_cG, r0_c, sigma, A: np.real(A) \ 

40 * np.exp(-np.sum((r_cG-r0_c[:,np.newaxis,np.newaxis,np.newaxis])**2, axis=0)/(2*sigma**2)) 

41 

42 t = time.time() 

43 gs0_G = _gaussRGA(r_cG, r0_c, sigma, A) 

44 print('_gaussRGA: %8.5f s' % (time.time() - t)) 

45 

46 t = time.time() 

47 gs1_G = gaussian_wave(r_cG, r0_c, sigma, None, A) 

48 print('+gaussRGA: %8.5f s' % (time.time() - t)) 

49 

50 assert np.abs(gs0_G - 

51 gs1_G).max() < 1e-12, 'Max error %g' % np.abs(gs0_G - 

52 gs1_G).max() 

53 del gs0_G, gs1_G 

54 

55 # Test case for real-part of kpoint-point wave with normalized amplitude 

56 _gaussRKN = lambda r_cG, r0_c, sigma, k_c: 1/(sigma*np.pi**0.5)**1.5 \ 

57 * np.exp(-np.sum((r_cG-r0_c[:,np.newaxis,np.newaxis,np.newaxis])**2, axis=0)/(2*sigma**2)) \ 

58 * np.cos(np.sum(r_cG*k_c[:,np.newaxis,np.newaxis,np.newaxis], axis=0)) 

59 

60 t = time.time() 

61 gs0_G = _gaussRKN(r_cG, r0_c, sigma, k_c) 

62 print('_gaussRKN: %8.5f s' % (time.time() - t)) 

63 

64 t = time.time() 

65 gs1_G = gaussian_wave(r_cG, r0_c, sigma, k_c) 

66 print('+gaussRKN: %8.5f s' % (time.time() - t)) 

67 

68 assert np.abs(gs0_G - 

69 gs1_G).max() < 1e-12, 'Max error %g' % np.abs(gs0_G - 

70 gs1_G).max() 

71 del gs0_G, gs1_G 

72 

73 # Test case for real-part of kpoint-point wave with complex amplitude 

74 _gaussRKA = lambda r_cG, r0_c, sigma, k_c, A: \ 

75 np.exp(-np.sum((r_cG-r0_c[:,np.newaxis,np.newaxis,np.newaxis])**2, axis=0)/(2*sigma**2)) \ 

76 * np.real(A*np.exp(1j*np.sum(r_cG*k_c[:,np.newaxis,np.newaxis,np.newaxis], axis=0))) 

77 

78 t = time.time() 

79 gs0_G = _gaussRKA(r_cG, r0_c, sigma, k_c, A) 

80 print('_gaussRKA: %8.5f s' % (time.time() - t)) 

81 

82 t = time.time() 

83 gs1_G = gaussian_wave(r_cG, r0_c, sigma, k_c, A) 

84 print('+gaussRKA: %8.5f s' % (time.time() - t)) 

85 

86 assert np.abs(gs0_G - 

87 gs1_G).max() < 1e-12, 'Max error %g' % np.abs(gs0_G - 

88 gs1_G).max() 

89 del gs0_G, gs1_G 

90 

91 # ------------------------------------------------------------------- 

92 

93 # Test case for complex case of gamma-point wave with normalized amplitude 

94 _gaussCGN = lambda r_cG, r0_c, sigma: (1+0j)/(sigma*np.pi**0.5)**1.5 \ 

95 * np.exp(-np.sum((r_cG-r0_c[:,np.newaxis,np.newaxis,np.newaxis])**2, axis=0)/(2*sigma**2)) 

96 

97 t = time.time() 

98 gs0_G = _gaussCGN(r_cG, r0_c, sigma) 

99 print('_gaussCGN: %8.5f s' % (time.time() - t)) 

100 

101 t = time.time() 

102 gs1_G = gaussian_wave(r_cG, r0_c, sigma, dtype=complex) 

103 print('+gaussCGN: %8.5f s' % (time.time() - t)) 

104 

105 assert np.abs(gs0_G - 

106 gs1_G).max() < 1e-12, 'Max error %g' % np.abs(gs0_G - 

107 gs1_G).max() 

108 del gs0_G, gs1_G 

109 

110 # Test case for complex case of gamma-point wave with complex amplitude 

111 _gaussCGA = lambda r_cG, r0_c, sigma, A: A \ 

112 * np.exp(-np.sum((r_cG-r0_c[:,np.newaxis,np.newaxis,np.newaxis])**2, axis=0)/(2*sigma**2)) 

113 

114 t = time.time() 

115 gs0_G = _gaussCGA(r_cG, r0_c, sigma, A) 

116 print('_gaussCGA: %8.5f s' % (time.time() - t)) 

117 

118 t = time.time() 

119 gs1_G = gaussian_wave(r_cG, r0_c, sigma, None, A, dtype=complex) 

120 print('+gaussCGA: %8.5f s' % (time.time() - t)) 

121 

122 assert np.abs(gs0_G - 

123 gs1_G).max() < 1e-12, 'Max error %g' % np.abs(gs0_G - 

124 gs1_G).max() 

125 del gs0_G, gs1_G 

126 

127 # Test case for complex case of kpoint-point wave with normalized amplitude 

128 _gaussCKN = lambda r_cG, r0_c, sigma, k_c: (1+0j)/(sigma*np.pi**0.5)**1.5 \ 

129 * np.exp(-np.sum((r_cG-r0_c[:,np.newaxis,np.newaxis,np.newaxis])**2, axis=0)/(2*sigma**2)) \ 

130 * np.exp(1j*np.sum(r_cG*k_c[:,np.newaxis,np.newaxis,np.newaxis], axis=0)) 

131 

132 t = time.time() 

133 gs0_G = _gaussCKN(r_cG, r0_c, sigma, k_c) 

134 print('_gaussCKN: %8.5f s' % (time.time() - t)) 

135 

136 t = time.time() 

137 gs1_G = gaussian_wave(r_cG, r0_c, sigma, k_c, dtype=complex) 

138 print('+gaussCKN: %8.5f s' % (time.time() - t)) 

139 

140 assert np.abs(gs0_G - 

141 gs1_G).max() < 1e-12, 'Max error %g' % np.abs(gs0_G - 

142 gs1_G).max() 

143 del gs0_G, gs1_G 

144 

145 # Test case for complex case of kpoint-point wave with complex amplitude 

146 _gaussCKA = lambda r_cG, r0_c, sigma, k_c, A: A \ 

147 * np.exp(-np.sum((r_cG-r0_c[:,np.newaxis,np.newaxis,np.newaxis])**2, axis=0)/(2*sigma**2)) \ 

148 * np.exp(1j*np.sum(r_cG*k_c[:,np.newaxis,np.newaxis,np.newaxis], axis=0)) 

149 

150 t = time.time() 

151 gs0_G = _gaussCKA(r_cG, r0_c, sigma, k_c, A) 

152 print('_gaussCKA: %8.5f s' % (time.time() - t)) 

153 

154 t = time.time() 

155 gs1_G = gaussian_wave(r_cG, r0_c, sigma, k_c, A, dtype=complex) 

156 print('+gaussCKA: %8.5f s' % (time.time() - t)) 

157 

158 assert np.abs(gs0_G - 

159 gs1_G).max() < 1e-12, 'Max error %g' % np.abs(gs0_G - 

160 gs1_G).max() 

161 del gs0_G, gs1_G