Coverage for gpaw/test/poisson/test_poisson_extravacuum.py: 72%

1import pytest 

2from gpaw.mpi import world 

3import time 

4import numpy as np 

5 

6from gpaw.poisson import PoissonSolver 

7from gpaw.poisson_extravacuum import ExtraVacuumPoissonSolver 

8from gpaw.grid_descriptor import GridDescriptor 

9 

10 

11pytestmark = pytest.mark.skipif(world.size > 4, 

12 reason='world.size > 4') 

13 

14 

15def test_poisson_poisson_extravacuum(): 

16 do_output = False 

17 do_plot = False 

18 poissoneps = 1e-16 

19 

20 # Model grid 

21 N = 16 

22 N_c = np.array((1, 1, 3)) * N 

23 cell_c = N_c / float(N) 

24 gd = GridDescriptor(N_c, cell_c, False) 

25 

26 # Construct model density 

27 coord_vg = gd.get_grid_point_coordinates() 

28 z_g = coord_vg[2, :] 

29 rho_g = gd.zeros() 

30 for z0 in [1, 2]: 

31 rho_g += 10 * (z_g - z0) * \ 

32 np.exp(-20 * np.sum((coord_vg.T - np.array([.5, .5, z0])).T**2, 

33 axis=0)) 

34 

35 if do_plot: 

36 big_rho_g = gd.collect(rho_g) 

37 if gd.comm.rank == 0: 

38 import matplotlib.pyplot as plt 

39 fig, ax_ij = plt.subplots(3, 4, figsize=(20, 10)) 

40 ax_i = ax_ij.ravel() 

41 ploti = 0 

42 Ng_c = gd.get_size_of_global_array() 

43 plt.sca(ax_i[ploti]) 

44 ploti += 1 

45 plt.pcolormesh(big_rho_g[Ng_c[0] / 2]) 

46 plt.sca(ax_i[ploti]) 

47 ploti += 1 

48 plt.plot(big_rho_g[Ng_c[0] / 2, Ng_c[1] / 2]) 

49 

50 def plot_phi(phi_g): 

51 if do_plot: 

52 big_phi_g = gd.collect(phi_g) 

53 if gd.comm.rank == 0: 

54 global ploti 

55 if ploti == 4: 

56 ploti -= 2 

57 plt.sca(ax_i[ploti]) 

58 ploti += 1 

59 plt.pcolormesh(big_phi_g[Ng_c[0] / 2]) 

60 plt.sca(ax_i[ploti]) 

61 ploti += 1 

62 plt.plot(big_phi_g[Ng_c[0] / 2, Ng_c[1] / 2]) 

63 plt.ylim(np.array([-1, 1]) * 0.15) 

64 

65 def poisson_solve(gd, rho_g, poisson): 

66 phi_g = gd.zeros() 

67 rho_g = rho_g 

68 t0 = time.time() 

69 npoisson = poisson.solve(phi_g, rho_g) 

70 t1 = time.time() 

71 if do_output: 

72 if gd.comm.rank == 0: 

73 print('Iterations: %s, Time: %.3f s' % 

74 (str(npoisson), t1 - t0)) 

75 return phi_g, npoisson 

76 

77 def poisson_init_solve(gd, rho_g, poisson): 

78 poisson.set_grid_descriptor(gd) 

79 phi_g, npoisson = poisson_solve(gd, rho_g, poisson) 

80 plot_phi(phi_g) 

81 return phi_g, npoisson 

82 

83 def compare(phi1_g, phi2_g, val, eps=np.sqrt(poissoneps)): 

84 big_phi1_g = gd.collect(phi1_g) 

85 big_phi2_g = gd.collect(phi2_g) 

86 if gd.comm.rank == 0: 

87 diff = np.max(np.absolute(big_phi1_g - big_phi2_g)) 

88 else: 

89 diff = 0.0 

90 diff = np.array([diff]) 

91 gd.comm.broadcast(diff, 0) 

92 assert diff[0] == pytest.approx(val, abs=eps) 

93 

94 # Get reference from default poissonsolver 

95 poisson = PoissonSolver('fd', eps=poissoneps) 

96 phiref_g, npoisson = poisson_init_solve(gd, rho_g, poisson) 

97 

98 # Test agreement with default 

99 poisson = ExtraVacuumPoissonSolver( 

100 N_c, PoissonSolver('fd', eps=poissoneps)) 

101 phi_g, npoisson = poisson_init_solve(gd, rho_g, poisson) 

102 compare(phi_g, phiref_g, 0.0, 1e-15) 

103 

104 # New reference with extra vacuum 

105 gpts = N_c * 4 

106 poisson = ExtraVacuumPoissonSolver( 

107 gpts, PoissonSolver('fd', eps=poissoneps)) 

108 phi_g, npoisson = poisson_init_solve(gd, rho_g, poisson) 

109 # print poisson.get_description() 

110 compare(phi_g, phiref_g, 2.6485385144e-02) 

111 phiref_g = phi_g 

112 

113 # Test with single coarsening 

114 poisson = ExtraVacuumPoissonSolver( 

115 gpts, 

116 PoissonSolver('fd', eps=poissoneps), 

117 PoissonSolver('fd', eps=poissoneps), 

118 1) 

119 phi_g, npoisson = poisson_init_solve(gd, rho_g, poisson) 

120 compare(phi_g, phiref_g, 1.5043946611e-04) 

121 

122 # Test with two coarsenings 

123 poisson = ExtraVacuumPoissonSolver( 

124 gpts, 

125 PoissonSolver('fd', eps=poissoneps), 

126 PoissonSolver('fd', eps=poissoneps), 

127 2) 

128 phi_g, npoisson = poisson_init_solve(gd, rho_g, poisson) 

129 compare(phi_g, phiref_g, 1.2980906205e-03) 

130 

131 # Test with cascaded single coarsenings 

132 poisson1 = ExtraVacuumPoissonSolver(gpts / 2, 

133 PoissonSolver('fd', eps=poissoneps), 

134 PoissonSolver('fd', eps=poissoneps), 1) 

135 poisson = ExtraVacuumPoissonSolver(gpts / 2, poisson1, 

136 PoissonSolver('fd', eps=poissoneps), 1) 

137 phi_g, npoisson = poisson_init_solve(gd, rho_g, poisson) 

138 # print poisson.get_description() 

139 compare(phi_g, phiref_g, 1.7086531461e-04) 

140 

141 # Test auxgrid 

142 gpts = N_c * 4 

143 for coarses in [1, 2, 3]: 

144 for nn_refine in [1, 3]: 

145 for nn_laplace in [1, 3]: 

146 EVPS = ExtraVacuumPoissonSolver 

147 poisson = EVPS(gpts, PoissonSolver('fd', eps=poissoneps), 

148 PoissonSolver('fd', eps=poissoneps), coarses, 

149 use_aux_grid=False, 

150 nn_refine=nn_refine, nn_laplace=nn_laplace) 

151 phiref_g, npoisson = poisson_init_solve(gd, rho_g, poisson) 

152 

153 poisson = EVPS(gpts, PoissonSolver('fd', eps=poissoneps), 

154 PoissonSolver('fd', eps=poissoneps), coarses, 

155 use_aux_grid=True, 

156 nn_refine=nn_refine, nn_laplace=nn_laplace) 

157 phi_g, npoisson = poisson_init_solve(gd, rho_g, poisson) 

158 compare(phi_g, phiref_g, 0.0, 1e-14) 

159 

160 if do_plot: 

161 if gd.comm.rank == 0: 

162 plt.show()