Coverage for gpaw/test/core/test_pw.py: 100%

1from math import pi 

2 

3import numpy as np 

4import pytest 

5from gpaw.core import PWDesc, UGDesc 

6from gpaw.core.plane_waves import find_reciprocal_vectors 

7from gpaw.gpu import cupy as cp 

8from gpaw.gpu.mpi import CuPyMPI 

9from gpaw.mpi import world 

10from gpaw.new.c import GPU_AWARE_MPI 

11 

12 

13@pytest.mark.ci 

14def test_pw_redist(): 

15 a = 2.5 

16 pw = PWDesc(ecut=10, cell=[a, a, a], comm=world) 

17 f1 = pw.empty() 

18 f1.data[:] = 1.0 

19 f2 = f1.gather() 

20 if f2 is not None: 

21 assert (f2.data == 1.0).all() 

22 assert f2.desc.comm.size == 1 

23 

24 

25@pytest.mark.ci 

26def test_pw_map(): 

27 """Test mapping from 5 to 9 G-vectors.""" 

28 pw5 = PWDesc(ecut=0.51 * (2 * pi)**2, 

29 cell=[1, 1, 0.1], 

30 dtype=complex) 

31 pw9 = PWDesc(ecut=1.01 * (2 * pi)**2, 

32 cell=[1, 1, 0.1], 

33 dtype=complex, 

34 comm=world) 

35 my_G_g, g_r = pw9.map_indices(pw5) 

36 print(world.rank, my_G_g, g_r) 

37 expected = { 

38 1: ([[0, 1, 2, 3, 6]], 

39 [[0, 1, 2, 3, 4]]), 

40 2: ([[0, 1, 2, 3], [1]], 

41 [[0, 1, 2, 3], [4]]), 

42 4: ([[0, 1, 2], [0], [0], []], 

43 [[0, 1, 2], [3], [4], []]), 

44 8: ([[0, 1], [0, 1], [], [0], [], [], [], []], 

45 [[0, 1], [2, 3], [], [4], [], [], [], []])} 

46 assert not (my_G_g - expected[world.size][0][world.rank]).any() 

47 for g, g0 in zip(g_r, expected[world.size][1]): 

48 assert not (np.array(g) - g0).any() 

49 

50 

51def grids(): 

52 comm = world 

53 if not GPU_AWARE_MPI: 

54 comm = CuPyMPI(comm) 

55 a = 1.0 

56 decomp = {1: [[0, 4], [0, 4], [0, 4]], 

57 2: [[0, 2, 4], [0, 4], [0, 4]], 

58 4: [[0, 2, 4], [0, 2, 4], [0, 4]], 

59 8: [[0, 1, 2, 3, 4], [0, 2, 4], [0, 4]]}[world.size] 

60 grid = UGDesc(cell=[a, a, a], size=(4, 4, 4), comm=comm, decomp=decomp) 

61 gridc = grid.new(dtype=complex) 

62 

63 g1 = grid.empty() 

64 g1.data[:] = 1.0 

65 

66 g2 = grid.empty() 

67 g2.data[:] = 1.0 

68 g2.data += [0, 1, 0, -1] 

69 

70 g3 = gridc.empty() 

71 g3.data[:] = 1.0 

72 

73 g4 = gridc.empty() 

74 g4.data[:] = 1.0 

75 g4.data += [0, 1, 0, -1] 

76 

77 g5 = gridc.empty() 

78 g5.data[:] = 1.0 + 1.0j 

79 g5.data += [0, 1, 0, -1] 

80 

81 return [g1, g2, g3, g4, g5] 

82 

83 

84@pytest.mark.gpu 

85@pytest.mark.parametrize('xp', [np, cp]) 

86@pytest.mark.parametrize('grid', grids()) 

87def test_pw_integrate(xp, grid): 

88 g = grid 

89 a = g.desc.cell[0, 0] 

90 ecut = 0.5 * (2 * np.pi / a)**2 * 1.01 

91 if xp is cp: 

92 g = g.to_xp(cp) 

93 pw = PWDesc(cell=g.desc.cell, dtype=g.desc.dtype, 

94 ecut=ecut, comm=g.desc.comm) 

95 f = g.fft(pw=pw) 

96 

97 gg = g.new() 

98 gg.scatter_from(f.gather(broadcast=True) 

99 .ifft(grid=g.desc.new(comm=None))) 

100 assert xp.allclose(g.data, gg.data, rtol=1e-14) 

101 

102 i1 = g.integrate() 

103 i2 = f.integrate() 

104 assert i1 == i2 

105 assert np.dtype(type(i1)) == g.desc.dtype 

106 

107 i1 = g.integrate(g) 

108 i2 = f.integrate(f) 

109 assert i1 == i2 

110 assert np.dtype(type(i1)) == g.desc.dtype 

111 

112 g1 = g.desc.empty(1, xp=xp) 

113 g1.data[:] = g.data 

114 m1 = g1.matrix_elements(g1) 

115 assert i1 == m1.data.item() 

116 

117 f1 = f.desc.empty(1, xp=xp) 

118 f1.data[:] = f.data 

119 m2 = f1.matrix_elements(f1) 

120 assert i2 == pytest.approx(m2.data[0, 0].item(), abs=1e-13) 

121 

122 

123def test_grr(): 

124 from ase.units import Bohr, Ha 

125 grid = UGDesc(cell=[2 / Bohr, 2 / Bohr, 2.737166 / Bohr], 

126 size=(9, 9, 12), 

127 comm=world) 

128 pw = PWDesc(ecut=340 / Ha, cell=grid.cell, comm=world) 

129 print(pw.G_plus_k_Gv.shape) 

130 from gpaw.grid_descriptor import GridDescriptor 

131 from gpaw.pw.descriptor import PWDescriptor 

132 g = GridDescriptor((9, 9, 12), [2 / Bohr, 2 / Bohr, 2.737166 / Bohr]) 

133 p = PWDescriptor(340 / Ha, g) 

134 print(p.get_reciprocal_vectors().shape) 

135 assert (p.get_reciprocal_vectors() == pw.G_plus_k_Gv).all() 

136 

137 

138def test_find_g(): 

139 G, e, i = find_reciprocal_vectors(0.501 * (2 * pi)**2, 

140 np.eye(3)) 

141 assert i.T.tolist() == [[0, 0, 0], 

142 [0, 0, 1], 

143 [0, 0, -1], 

144 [0, 1, 0], 

145 [0, -1, 0], 

146 [1, 0, 0], 

147 [-1, 0, 0]] 

148 G, e, i = find_reciprocal_vectors(0.501 * (2 * pi)**2, 

149 np.eye(3), 

150 dtype=float) 

151 assert i.T.tolist() == [[0, 0, 0], 

152 [0, 0, 1], 

153 [0, 1, 0], 

154 [1, 0, 0]] 

155 G, e, i = find_reciprocal_vectors(0.5 * (2 * pi)**2, 

156 np.eye(3), 

157 kpt=np.array([0.1, 0, 0])) 

158 assert i.T.tolist() == [[0, 0, 0], 

159 [-1, 0, 0]] 

160 

161 

162def test_random(): 

163 pw = PWDesc(ecut=20, cell=[1, 1, 1], comm=world) 

164 a = pw.empty(2) 

165 a.randomize() 

166 assert world.rank != 0 or (a.data[:, 0].imag == 0.0).all() 

167 

168 

169def test_morph(): 

170 pw1 = PWDesc(ecut=20, cell=[1, 1, 1], comm=world) 

171 a = pw1.empty() 

172 a.randomize() 

173 pw2 = PWDesc(ecut=20, cell=[1, 1, 1.1], comm=world) 

174 b = a.morph(pw2) 

175 c = b.morph(pw1) 

176 assert (a.data == c.data).all()