Coverage for gpaw/test/gpu/test_fd_laplace.py: 26%

1import numpy as np 

2import pytest 

3 

4from gpaw.fd_operators import Laplace 

5from gpaw.gpu import cupy as cp, cupy_is_fake 

6from gpaw.grid_descriptor import GridDescriptor 

7from gpaw.mpi import world 

8 

9 

10@pytest.mark.gpu 

11@pytest.mark.skipif(cupy_is_fake, reason='No cupy') 

12@pytest.mark.parametrize('pbc', [True, False]) 

13def test_fd_laplace(pbc): 

14 if world.size > 4: 

15 # Grid is so small that domain decomposition cannot exceed 4 domains 

16 assert world.size % 4 == 0 

17 group, other = divmod(world.rank, 4) 

18 ranks = np.arange(4 * group, 4 * (group + 1)) 

19 domain_comm = world.new_communicator(ranks) 

20 else: 

21 domain_comm = world 

22 

23 lat = 8.0 

24 gd = GridDescriptor((32, 32, 32), (lat, lat, lat), 

25 pbc_c=pbc, comm=domain_comm) 

26 

27 if pbc: 

28 dtype = complex 

29 phase = np.ones((3, 2), complex) 

30 else: 

31 dtype = float 

32 phase = None 

33 

34 # Use Gaussian as input 

35 x, y, z = gd.get_grid_point_coordinates() 

36 sigma = 1.5 

37 mu = lat / 2.0 

38 

39 a = gd.zeros(dtype=dtype) 

40 a[:] = np.exp(-((x - mu)**2 + (y - mu)**2 + (z - mu)**2) / (2.0 * sigma)) 

41 # analytic solution 

42 b_analytic = np.zeros_like(a) 

43 b_analytic[:] = (((x - mu)**2 + (y - mu)**2 + (z - mu)**2) / sigma**2 

44 - 3.0 / sigma) * a 

45 

46 b = np.zeros_like(a) 

47 a_gpu = cp.asarray(a) 

48 b_gpu = cp.zeros_like(a_gpu) 

49 

50 # Laplace 

51 Laplace(gd, 1.0, 3, dtype=dtype).apply(a, b, phase_cd=phase) 

52 Laplace(gd, 1.0, 3, dtype=dtype, xp=cp).apply(a_gpu, b_gpu, phase_cd=phase) 

53 b_ref = b_gpu.get() 

54 

55 assert b == pytest.approx(b_ref, abs=1e-12) 

56 # Neglect boundaries in check to analytic solution 

57 assert (b_analytic[2:-2, 2:-2, 2:-2] 

58 == pytest.approx(b_ref[2:-2, 2:-2, 2:-2], abs=1e-2))