Coverage for gpaw/test/new/test_apply_ham.py: 86%

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 

8from gpaw.new.pw.hamiltonian import PWHamiltonian 

9 

10 

11def apply(a: float, # lattice constant 

12 N: int, # number of grid-points 

13 B: int, # number of bands 

14 xp, 

15 slow=False) -> float: 

16 """Calculate density from wave functions.""" 

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

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

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

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

21 psit_nG.data[:, 0] = 1.0 

22 vt_R = grid.zeros(xp=xp) 

23 vt_R.data[:] = -2.0 

24 ham = PWHamiltonian(grid, pw, xp) 

25 t = time() 

26 ham.apply_local_potential(vt_R, psit_nG, psit_nG.new()) 

27 t = time() - t 

28 return t 

29 

30 

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

32 [np, 

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

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

35def test_apply_ham(xp, nbands): 

36 apply(a=2.5, N=6, B=nbands, xp=xp) 

37 

38 

39def main(): 

40 """Test speedup for larger system.""" 

41 for _ in range(2): 

42 t = apply(6.0, 32, 100, cp) 

43 print(t) 

44 

45 

46if __name__ == '__main__': 

47 main()