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

1import pytest 

2import numpy as np 

3 

4from gpaw.core import UGDesc 

5from gpaw.gpu import cupy as cp 

6from gpaw.mpi import world 

7from gpaw.spline import Spline 

8from gpaw.new.c import GPU_AWARE_MPI 

9 

10 

11@pytest.mark.gpu 

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

13def test_acf(dtype): 

14 if world.size > 1 and not GPU_AWARE_MPI: 

15 pytest.skip('MPI not GPU-aware') 

16 s = Spline.from_data(0, 1.0, [1.0, 0.5, 0.0]) 

17 n = 40 

18 a = 8.0 

19 relpos_ac = [(0.5, 0.5, 0.25 + 0.25 * i) for i in [0, 1, 2]] 

20 

21 grid = UGDesc(cell=[a, a, a], size=(n, n, n), comm=world, dtype=dtype) 

22 basis_cpu = grid.atom_centered_functions([[s], [s], [s]], 

23 positions=relpos_ac, xp=np) 

24 basis_gpu = grid.atom_centered_functions([[s], [s], [s]], 

25 positions=relpos_ac, xp=cp) 

26 

27 P_cpu_ani = basis_cpu.layout.empty() 

28 P_gpu_ani = basis_gpu.layout.empty() 

29 P_cpu_ani.data[:] = 1.0 

30 P_gpu_ani.data[:] = 1.0 

31 

32 b_cpu = grid.zeros(xp=np) 

33 b_gpu = grid.zeros(xp=cp) 

34 basis_cpu.add_to(b_cpu, P_cpu_ani) 

35 basis_gpu.add_to(b_gpu, P_gpu_ani) 

36 

37 assert b_cpu.data == pytest.approx(b_gpu.data.get(), abs=1e-12) 

38 

39 out_cpu_ani = basis_cpu.integrate(b_cpu) 

40 out_gpu_ani = basis_gpu.integrate(b_gpu) 

41 

42 for a, out_cpu_ni in out_cpu_ani.items(): 

43 assert out_cpu_ni == pytest.approx(out_gpu_ani[a].get(), abs=1e-12)