Coverage for gpaw/test/parallel/test_submatrix_redist.py: 96%

1from gpaw.mpi import world 

2from gpaw.blacs import BlacsGrid, Redistributor 

3 

4 

5def test_parallel_submatrix_redist(scalapack): 

6 if world.size > 1: 

7 shape = (2, world.size // 2) 

8 else: 

9 shape = (1, 1) 

10 

11 grid = BlacsGrid(world, *shape) 

12 

13 desc = grid.new_descriptor(12, 8, 2, 3) 

14 

15 a = desc.zeros() 

16 a[:] = world.rank 

17 

18 subdesc = grid.new_descriptor(7, 7, 2, 2) 

19 b = subdesc.zeros() 

20 

21 r = Redistributor(grid.comm, desc, subdesc) 

22 

23 ia = 3 

24 ja = 2 

25 ib = 1 

26 jb = 1 

27 M = 4 

28 N = 5 

29 

30 r.redistribute(a, b, M, N, ia, ja, ib, jb) 

31 

32 a0 = desc.collect_on_master(a) 

33 b0 = subdesc.collect_on_master(b) 

34 if world.rank == 0: 

35 print(a0) 

36 print(b0) 

37 xa = a0[ia:ia + M, ja:ja + N] 

38 xb = b0[ib:ib + M, jb:jb + N] 

39 assert (xa == xb).all()