Coverage for gpaw/test/core/test_abs_square.py: 80%

1from time import time 

2 

3import numpy as np 

4import pytest 

5 

6from gpaw.core import PWDesc, UGDesc 

7from gpaw.gpu import cupy as cp 

8 

9 

10def abs_square(a: float, # lattice constant 

11 N: int, # number of grid-points 

12 B: int, # number of bands 

13 xp, 

14 slow=False) -> float: 

15 """Calculate density from wave functions.""" 

16 grid = UGDesc(cell=[a, a, a], size=[N, N, N]) 

17 ecut = 0.5 * (np.pi * N / a) 

18 pw = PWDesc(ecut=ecut, cell=grid.cell, dtype=complex) 

19 psit_nG = pw.zeros(B, xp=xp) 

20 psit_nG.data[:, 0] = 1.0 

21 weight_n = np.linspace(1, 0, B) 

22 nt_R = grid.zeros(xp=xp) 

23 

24 t = time() 

25 psit_nG.abs_square(weight_n, nt_R, _slow=slow) 

26 t = time() - t 

27 

28 assert nt_R.integrate() == pytest.approx(a**3 * weight_n.sum()) 

29 

30 return t 

31 

32 

33@pytest.mark.parametrize('xp', 

34 [np, 

35 pytest.param(cp, marks=pytest.mark.gpu)]) 

36@pytest.mark.parametrize('nbands', [2, 17]) 

37def test_pw_abs_square(xp, nbands): 

38 abs_square(a=2.5, N=6, B=nbands, xp=xp) 

39 

40 

41def main(): 

42 """Test speedup for larger system.""" 

43 abs_square(6.0, 32, 100, cp) # GPU-warmup 

44 t = abs_square(6.0, 32, 100, cp) 

45 print('Fast:', t) 

46 t = abs_square(6.0, 32, 100, cp, slow=True) 

47 print('Slow:', t) 

48 

49 

50if __name__ == '__main__': 

51 main()