Coverage for gpaw/test/core/test_acf.py: 96%

1import sys 

2 

3import numpy as np 

4import pytest 

5 

6from gpaw.core import PWDesc, UGDesc 

7from gpaw.gpu import cupy as cp 

8from gpaw.mpi import world 

9 

10a = 2.5 

11n = 20 

12 

13 

14@pytest.fixture 

15def grid(): 

16 # comm = world.new_communicator([world.rank]) 

17 return UGDesc(cell=[a, a, a], size=(n, n, n), comm=world, dtype=complex) 

18 

19 

20# Gussian: 

21alpha = 4.0 

22s = (0, 3.0, lambda r: np.exp(-alpha * r**2)) 

23gauss_integral = np.pi / 2 / alpha**1.5 

24 

25 

26@pytest.mark.ci 

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

28def test_acf_fd(grid, xp): 

29 

30 basis = grid.atom_centered_functions( 

31 [[s]], 

32 positions=[[0.5, 0.5, 0.5]], xp=xp) 

33 coefs = basis.layout.empty() 

34 if 0 in coefs: 

35 print(coefs[0]) 

36 coefs[0] = [1.0] 

37 f1 = grid.zeros(xp=xp) 

38 basis.add_to(f1, coefs.to_xp(xp)) 

39 

40 if 0: 

41 from sympy import exp, integrate, oo, var 

42 a0, r = var('a, r') 

43 integrate(exp(-a0 * r**2) * r**2, (r, 0, oo)) 

44 

45 assert f1.integrate() == pytest.approx(gauss_integral) 

46 

47 f2 = f1.gather(broadcast=True) 

48 x, y = f2.xy(n // 2, n // 2, ...) 

49 y0 = np.exp(-alpha * (x - a / 2)**2) / (4 * np.pi)**0.5 

50 assert abs(y - y0).max() == pytest.approx(0.0, abs=0.001) 

51 

52 

53@pytest.mark.ci 

54@pytest.mark.gpu 

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

56def test_acf_pw(grid, xp): 

57 if world.size > 1 and xp is cp: 

58 pytest.skip() 

59 if xp is cp and '_gpaw' in sys.builtin_module_names: 

60 pytest.skip() 

61 

62 pw = PWDesc(ecut=50, cell=grid.cell, dtype=complex, comm=world) 

63 

64 basis = pw.atom_centered_functions( 

65 [[s]], 

66 positions=[[0.5, 0.5, 0.5]], xp=xp) 

67 

68 coefs = basis.layout.empty() 

69 if 0 in coefs: 

70 print(coefs[0]) 

71 coefs[0] = xp.asarray([1.0]) 

72 

73 f1 = pw.zeros(xp=xp) 

74 basis.add_to(f1, coefs) 

75 assert f1.integrate() == pytest.approx(gauss_integral) 

76 f2 = f1.gather(broadcast=True) 

77 r2 = f2.ifft(grid=grid.new(comm=None)) 

78 x, y = r2.xy(10, 10, ...) 

79 y0 = np.exp(-alpha * (x - a / 2)**2) / (4 * np.pi)**0.5 

80 assert abs(y - y0).max() == pytest.approx(0.0, abs=0.002)