Coverage for gpaw/test/core/test_grid.py: 99%

1from math import pi 

2 

3import numpy as np 

4import pytest 

5from gpaw.core import UGDesc, PWDesc, UGArray 

6from gpaw.fd_operators import Laplace 

7from gpaw.mpi import world 

8 

9 

10def test_real_to_complex_fft(): 

11 L = 1.0 

12 n = 8 

13 a = UGDesc(cell=[L, L, L], size=(n, n, n), comm=world).empty() 

14 a.randomize() 

15 pw = PWDesc(ecut=50.0, cell=a.desc.cell, dtype=complex, comm=world) 

16 # real -> complex: 

17 with pytest.raises(TypeError): 

18 a.fft(pw=pw) 

19 # complex -> complex: 

20 b = a.to_complex().fft(pw=pw).gather(broadcast=True) 

21 # c(G) = c(-G)^*: 

22 coefs = {tuple(i_c): coef for i_c, coef in zip(pw.indices_cG.T, b.data)} 

23 for (i, j, k), coef in coefs.items(): 

24 assert coef == coefs[(-i, -j, -k)].conj() 

25 

26 

27@pytest.mark.ci 

28def test_redist(): 

29 a = 2.5 

30 n = 2 

31 grid1 = UGDesc(cell=[a, a, a], size=(n, n, n), comm=world) 

32 f1 = grid1.empty() 

33 f1.data[:] = world.rank + 1 

34 f2 = f1.gather() 

35 f3 = f1.gather(broadcast=True) 

36 if world.rank == 0: 

37 assert (f2.data == f3.data).all() 

38 print(f2) 

39 f4 = f1.new() 

40 f4.scatter_from(f2) 

41 assert (f4.data == f1.data).all() 

42 

43 

44def test_complex_laplace(): 

45 a = 2.5 

46 grid = UGDesc(cell=[a, a, a], 

47 size=(24, 8, 8), 

48 kpt=[1 / 3, 0, 0], 

49 comm=world) 

50 f = grid.empty() 

51 f.data[:] = 1.0 

52 f.multiply_by_eikr() 

53 lap = Laplace(grid._gd, n=2, dtype=complex) 

54 g = grid.empty() 

55 lap(f, g) 

56 k = 2 * pi / a / 3 

57 error = g.data.conj() * f.data + k**2 

58 assert abs(error).max() == pytest.approx(0.0, abs=1e-6) 

59 

60 

61def py(a: float) -> UGArray: 

62 L = 5 

63 n = 20 

64 l = 1 

65 rcut = 3.0 

66 p = (l, rcut, lambda r: np.exp(-a * r**2)) 

67 grid = UGDesc(cell=[L, L, L], size=(n, n, n), comm=world) 

68 f = grid.zeros() 

69 # Add P-y function to grid: 

70 basis = grid.atom_centered_functions( 

71 [[p]], 

72 positions=[[0.5, 0.5, 0.5]]) 

73 coefs = basis.layout.empty() 

74 if 0 in coefs: 

75 coefs[0] = [1.0, 0, 0] # y, z, x 

76 basis.add_to(f, coefs) 

77 return f 

78 

79 

80def test_moment(): 

81 # P-type Gaussian: 

82 a = 4.0 

83 f = py(a) 

84 

85 if 0: # Analytic result 

86 from sympy import integrate, exp, oo, var, Symbol, sqrt, pi 

87 x = var('x') 

88 a = Symbol('a', positive=True) 

89 m = 8 * (integrate(exp(-a * x**2) * x**2, (x, 0, oo)) * 

90 integrate(exp(-a * x**2), (x, 0, oo))**2 * 

91 sqrt(3 / (4 * pi))) 

92 print(m) # sqrt(3)*pi/(4*a**(5/2)) 

93 

94 moment = 3**0.5 * np.pi / (4 * a**(5 / 2)) 

95 

96 assert abs(f.integrate()) < 1e-14 

97 assert f.moment() == pytest.approx([0, moment, 0]) 

98 

99 pw = PWDesc(cell=f.desc.cell, ecut=f.desc.ekin_max(), comm=world) 

100 f2 = f.fft(pw=pw) 

101 

102 assert abs(f2.integrate()) < 1e-14 

103 assert f2.moment() == pytest.approx([0, moment, 0], abs=1e-5) 

104 

105 

106def test_isosurface(): 

107 pytest.importorskip('plotly') 

108 f = py(4.0) 

109 f.isosurface(show=False) 

110 

111 

112if __name__ == '__main__': 

113 py(4.0).isosurface()