Coverage for gpaw/test/parallel/test_mpi.py: 61%

44 statements  

« prev     ^ index     » next       coverage.py v7.7.1, created at 2025-07-14 00:18 +0000

1import pytest 

2import numpy as np 

3from gpaw.mpi import world, send, receive, broadcast_array 

4from gpaw.mpi4pywrapper import MPI4PYWrapper 

5 

6worlds = [world] 

7 

8try: 

9 from mpi4py.MPI import COMM_WORLD 

10except ImportError: 

11 pass 

12else: 

13 worlds.append(MPI4PYWrapper(COMM_WORLD)) 

14 

15 

16@pytest.mark.parametrize('world', worlds) 

17@pytest.mark.ci 

18def test_send_receive_object(world): 

19 if world.size == 1: 

20 return 

21 obj = (42, 'hello') 

22 if world.rank == 0: 

23 send(obj, 1, world) 

24 elif world.rank == 1: 

25 assert obj == receive(0, world) 

26 

27 

28@pytest.mark.ci 

29def test_scalar_reduce(): 

30 assert world.sum_scalar(world.rank + 1) == world.size * \ 

31 (world.size + 1) // 2 

32 assert np.allclose(world.sum_scalar(world.rank + 1.0), 

33 world.size * (world.size + 1.0) / 2) 

34 assert world.min_scalar(world.rank + 1) == 1 

35 assert world.max_scalar(world.rank + 1) == world.size 

36 assert world.min_scalar(world.rank + 1.0) == 1.0 

37 assert world.max_scalar(world.rank + 1.0) == world.size * 1.0 

38 

39 

40def test_bcast_array(): 

41 new = world.new_communicator 

42 

43 if world.size == 2: 

44 comms = [world] 

45 elif world.size == 4: 

46 ranks = np.array([[0, 1], [2, 3]]) 

47 comms = [new(ranks[world.rank // 2]), 

48 new(ranks[:, world.rank % 2])] 

49 elif world.size == 8: 

50 ranks = [[[0, 1], [0, 2], [0, 4]], 

51 [[0, 1], [1, 3], [1, 5]], 

52 [[2, 3], [0, 2], [2, 6]], 

53 [[2, 3], [1, 3], [3, 7]], 

54 [[4, 5], [4, 6], [0, 4]], 

55 [[4, 5], [5, 7], [1, 5]], 

56 [[6, 7], [4, 6], [2, 6]], 

57 [[6, 7], [5, 7], [3, 7]]] 

58 comms = [new(r) for r in ranks[world.rank]] 

59 else: 

60 return 

61 

62 array = np.zeros(3, int) 

63 if world.rank == 0: 

64 array[:] = 42 

65 

66 out = broadcast_array(array, *comms) 

67 assert (out == 42).all()