Coverage for pySDC/implementations/sweeper_classes/imex_1st_order_MPI.py: 100%

1from mpi4py import MPI 

2from pySDC.implementations.sweeper_classes.generic_implicit_MPI import SweeperMPI 

3from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order 

4 

5 

6class imex_1st_order_MPI(SweeperMPI, imex_1st_order): 

7 """ 

8 Node-parallel IMEX-SDC sweeper, one collocation node per MPI rank, so far only with Picard for the explicit part. 

9 """ 

10 

11 def integrate(self, last_only=False): 

12 """ 

13 Integrates the right-hand side (here impl + expl) 

14 

15 Args: 

16 last_only (bool): Integrate only the last node for the residual or all of them 

17 

18 Returns: 

19 list of dtype_u: containing the integral as values 

20 """ 

21 

22 L = self.level 

23 P = L.prob 

24 

25 me = P.dtype_u(P.init, val=0.0) 

26 for m in [self.coll.num_nodes - 1] if last_only else range(self.coll.num_nodes): 

27 recvBuf = me if m == self.rank else None 

28 self.comm.Reduce( 

29 L.dt * self.coll.Qmat[m + 1, self.rank + 1] * (L.f[self.rank + 1].impl + L.f[self.rank + 1].expl), 

30 recvBuf, 

31 root=m, 

32 op=MPI.SUM, 

33 ) 

34 

35 return me 

36 

37 def update_nodes(self): 

38 """ 

39 Update the u- and f-values at the collocation nodes -> corresponds to a single sweep over all nodes 

40 

41 Returns: 

42 None 

43 """ 

44 

45 L = self.level 

46 P = L.prob 

47 

48 # only if the level has been touched before 

49 assert L.status.unlocked 

50 

51 # get number of collocation nodes for easier access 

52 

53 # gather all terms which are known already (e.g. from the previous iteration) 

54 # this corresponds to u0 + QF(u^k) - QdF(u^k) + tau 

55 

56 # get QF(u^k) 

57 rhs = self.integrate() 

58 

59 # subtract QdF(u^k) 

60 rhs -= L.dt * (self.QI[self.rank + 1, self.rank + 1] * L.f[self.rank + 1].impl) 

61 

62 # add initial conditions 

63 rhs += L.u[0] 

64 # add tau if associated 

65 if L.tau[self.rank] is not None: 

66 rhs += L.tau[self.rank] 

67 

68 # implicit solve with prefactor stemming from the diagonal of Qd 

69 L.u[self.rank + 1] = P.solve_system( 

70 rhs, 

71 L.dt * self.QI[self.rank + 1, self.rank + 1], 

72 L.u[self.rank + 1], 

73 L.time + L.dt * self.coll.nodes[self.rank], 

74 ) 

75 # update function values 

76 L.f[self.rank + 1] = P.eval_f(L.u[self.rank + 1], L.time + L.dt * self.coll.nodes[self.rank]) 

77 

78 # indicate presence of new values at this level 

79 L.status.updated = True 

80 

81 return None 

82 

83 def compute_end_point(self): 

84 """ 

85 Compute u at the right point of the interval 

86 

87 Returns: 

88 None 

89 """ 

90 

91 L = self.level 

92 P = L.prob 

93 

94 # check if Mth node is equal to right point and do_coll_update is false, perform a simple copy 

95 if self.coll.right_is_node and not self.params.do_coll_update: 

96 super().compute_end_point() 

97 else: 

98 L.uend = P.dtype_u(L.u[0]) 

99 self.comm.Allreduce( 

100 L.dt * self.coll.weights[self.rank] * (L.f[self.rank + 1].impl + L.f[self.rank + 1].expl), 

101 L.uend, 

102 op=MPI.SUM, 

103 ) 

104 L.uend += L.u[0] 

105 

106 # add up tau correction of the full interval (last entry) 

107 if L.tau[self.rank] is not None: 

108 self.communicate_tau_correction_for_full_interval() 

109 L.uend += L.tau[-1] 

110 return None