Coverage for pySDC/tutorial/step_6/C_MPI_parallelization.py: 100%

1import sys 

2from pathlib import Path 

3 

4from mpi4py import MPI 

5 

6from pySDC.helpers.stats_helper import get_sorted 

7from pySDC.implementations.controller_classes.controller_MPI import controller_MPI 

8from pySDC.tutorial.step_6.A_run_non_MPI_controller import set_parameters_ml 

9 

10 

11def main(fname='step_6_C_out.txt'): 

12 """ 

13 Run PFASST with the MPI-parallel controller, one step per rank. 

14 

15 Run it the way you would run any MPI program, with as many ranks as you want steps:: 

16 

17 mpirun -np 4 python C_MPI_parallelization.py 

18 

19 The number of parallel steps is simply the size of ``MPI.COMM_WORLD``, so nothing in here has to 

20 be told how many there are. 

21 

22 Args: 

23 fname (str): file under ``data/`` to append the results to 

24 """ 

25 

26 # set MPI communicator 

27 comm = MPI.COMM_WORLD 

28 

29 # get parameters from Part A 

30 description, controller_params, t0, Tend = set_parameters_ml() 

31 

32 # instantiate controllers 

33 controller = controller_MPI(controller_params=controller_params, description=description, comm=comm) 

34 # get initial values on finest level 

35 P = controller.S.levels[0].prob 

36 uinit = P.u_exact(t0) 

37 

38 # call main functions to get things done... 

39 uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend) 

40 

41 # filter statistics by type (number of iterations) 

42 iter_counts = get_sorted(stats, type='niter', sortby='time') 

43 

44 # combine statistics into list of statistics 

45 iter_counts_list = comm.gather(iter_counts, root=0) 

46 

47 rank = comm.Get_rank() 

48 size = comm.Get_size() 

49 

50 if rank == 0: 

51 Path("data").mkdir(parents=True, exist_ok=True) 

52 f = open('data/' + fname, 'a') 

53 out = 'Working with %2i processes...' % size 

54 f.write(out + '\n') 

55 print(out) 

56 

57 # compute exact solutions and compare with both results 

58 uex = P.u_exact(Tend) 

59 err = abs(uex - uend) 

60 

61 out = 'Error vs. exact solution: %12.8e' % err 

62 f.write(out + '\n') 

63 print(out) 

64 

65 # build one list of statistics instead of list of lists, the sort by time 

66 iter_counts_gather = [item for sublist in iter_counts_list for item in sublist] 

67 iter_counts = sorted(iter_counts_gather, key=lambda tup: tup[0]) 

68 

69 # compute and print statistics 

70 for item in iter_counts: 

71 out = 'Number of iterations for time %4.2f: %1i ' % (item[0], item[1]) 

72 f.write(out + '\n') 

73 print(out) 

74 

75 f.write('\n') 

76 print() 

77 

78 assert all(item[1] <= 8 for item in iter_counts), "ERROR: weird iteration counts, got %s" % iter_counts 

79 

80 

81if __name__ == "__main__": 

82 main(sys.argv[1] if len(sys.argv) == 2 else 'step_6_C_out.txt')