Coverage for gpaw/test/gpu/test_apply_local_pot.py: 97%

1import pytest 

2from gpaw.core import UGDesc, PWDesc 

3from gpaw.new.pw.hamiltonian import apply_local_potential_gpu 

4from gpaw.gpu import cupy as cp 

5 

6 

7@pytest.mark.gpu 

8@pytest.mark.serial 

9@pytest.mark.parametrize('dtype', [float, complex]) 

10@pytest.mark.parametrize('nbands', [1, 2, 3, 5]) 

11def test_apply_loc_pot(dtype, nbands): 

12 a = 1.5 

13 n = 8 

14 vt_R = UGDesc(cell=[a, a, a], size=(n, n, n)).empty(xp=cp) 

15 v0 = 1.3 

16 vt_R.data[:] = v0 

17 if dtype == complex: 

18 kpt = [0, 0, 0.5] 

19 else: 

20 kpt = None 

21 pw = PWDesc(cell=vt_R.desc.cell, ecut=15.0, kpt=kpt) 

22 psit_nG = pw.empty(nbands, xp=cp) 

23 p0 = 1.2 - 2.0j 

24 psit_nG.data[:] = p0 

25 if dtype == float: 

26 psit_nG.data[:, 0] = -1.1 

27 out_nG = pw.empty(nbands, xp=cp) 

28 apply_local_potential_gpu(vt_R, 

29 psit_nG, 

30 out_nG, 

31 blocksize=3) 

32 error_nG = (cp.asnumpy(out_nG.data) - 

33 (v0 + pw.ekin_G) * cp.asnumpy(psit_nG.data)) 

34 assert abs(error_nG).max() == pytest.approx(0.0) 

35 

36 

37if __name__ == '__main__': 

38 test_apply_loc_pot(float, 1)