Coverage for gpaw/test/gpu/test_fd_transformers.py: 27%

1import pytest 

2import numpy as np 

3from gpaw.grid_descriptor import GridDescriptor 

4from gpaw.transformers import Transformer 

5from gpaw.mpi import world 

6from gpaw.gpu import cupy as cp, cupy_is_fake 

7 

8 

9@pytest.mark.gpu 

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

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

12@pytest.mark.parametrize('nn', [1, 2, 3, 4]) 

13def test_fd_transformers(pbc, nn): 

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 

42 a_gpu = cp.asarray(a) 

43 

44 # Transformers 

45 coarsegd = gd.coarsen() 

46 a_coarse = coarsegd.zeros(dtype=dtype) 

47 a_coarse_gpu = cp.zeros_like(a_coarse) 

48 

49 # Restrict 

50 Transformer(gd, coarsegd, nn, dtype=dtype).apply(a, a_coarse, phases=phase) 

51 Transformer(gd, coarsegd, nn, dtype=dtype, xp=cp).apply( 

52 a_gpu, a_coarse_gpu, phases=phase) 

53 a_coarse_ref = a_coarse_gpu.get() 

54 assert a_coarse == pytest.approx(a_coarse_ref, abs=1e-14) 

55 

56 # Interpolate 

57 Transformer(coarsegd, gd, nn, dtype=dtype).apply(a_coarse, a, phases=phase) 

58 Transformer(coarsegd, gd, nn, dtype=dtype, xp=cp).apply( 

59 a_coarse_gpu, a_gpu, phases=phase) 

60 a_ref = a_gpu.get() 

61 assert a == pytest.approx(a_ref, abs=1e-14)