Coverage for pySDC/implementations/transfer_classes/BaseTransferMPI.py: 100%

1from mpi4py import MPI 

2 

3from pySDC.core.errors import UnlockError 

4from pySDC.core.base_transfer import BaseTransfer 

5 

6 

7class base_transfer_MPI(BaseTransfer): 

8 """ 

9 Node-parallel space-time transfer for sweepers where each MPI rank holds one collocation node. 

10 

11 Attributes: 

12 logger: custom logger for sweeper-related logging 

13 params(__Pars): parameter object containing the custom parameters passed by the user 

14 fine (pySDC.Level.level): reference to the fine level 

15 coarse (pySDC.Level.level): reference to the coarse level 

16 """ 

17 

18 def __init__(self, *args, **kwargs): 

19 """ 

20 Set up the transfer as `BaseTransfer` does and take the node communicators from the sweepers of both levels. 

21 

22 Args: 

23 *args: Passed on to `BaseTransfer` 

24 **kwargs: Passed on to `BaseTransfer` 

25 

26 Raises: 

27 NotImplementedError: If a communicator size differs from the number of collocation nodes of its level 

28 """ 

29 super().__init__(*args, **kwargs) 

30 self.comm_fine = self.fine.sweep.comm 

31 self.comm_coarse = self.coarse.sweep.comm 

32 

33 if ( 

34 self.comm_fine.size != self.fine.sweep.coll.num_nodes 

35 or self.comm_coarse.size != self.coarse.sweep.coll.num_nodes 

36 ): 

37 raise NotImplementedError( 

38 f'{type(self).__name__} only works when each rank administers one collocation node so far!' 

39 ) 

40 

41 def restrict(self): 

42 """ 

43 Space-time restriction routine 

44 

45 The routine applies the spatial restriction operator to the fine values on the fine nodes, then reevaluates f 

46 on the coarse level. This is used for the first part of the FAS correction tau via integration. The second part 

47 is the integral over the fine values, restricted to the coarse level. Finally, possible tau corrections on the 

48 fine level are restricted as well. 

49 """ 

50 

51 F, G = self.fine, self.coarse 

52 CF, CG = self.comm_fine, self.comm_coarse 

53 SG = G.sweep 

54 PG = G.prob 

55 

56 # only if the level is unlocked at least by prediction 

57 if not F.status.unlocked: 

58 raise UnlockError('fine level is still locked, cannot use data from there') 

59 

60 # restrict fine values in space 

61 tmp_u = self.space_transfer.restrict(F.u[CF.rank + 1]) 

62 

63 # restrict collocation values 

64 G.u[0] = self.space_transfer.restrict(F.u[0]) 

65 recvBuf = [None for _ in range(SG.coll.num_nodes)] 

66 recvBuf[CG.rank] = PG.u_init 

67 for n in range(SG.coll.num_nodes): 

68 CF.Reduce(self.Rcoll[n, CF.rank] * tmp_u, recvBuf[CG.rank], root=n, op=MPI.SUM) 

69 G.u[CG.rank + 1] = recvBuf[CG.rank] 

70 

71 # re-evaluate f on coarse level 

72 G.f[0] = PG.eval_f(G.u[0], G.time) 

73 G.f[CG.rank + 1] = PG.eval_f(G.u[CG.rank + 1], G.time + G.dt * SG.coll.nodes[CG.rank]) 

74 

75 # build coarse level tau correction part 

76 tauG = G.sweep.integrate() 

77 

78 # build fine level tau correction part 

79 tauF = F.sweep.integrate() 

80 

81 # restrict fine level tau correction part in space 

82 tmp_tau = self.space_transfer.restrict(tauF) 

83 

84 # restrict fine level tau correction part in collocation 

85 tauFG = tmp_tau.copy() 

86 for n in range(SG.coll.num_nodes): 

87 recvBuf = tauFG if n == CG.rank else None 

88 CF.Reduce(self.Rcoll[n, CF.rank] * tmp_tau, recvBuf, root=n, op=MPI.SUM) 

89 

90 # build tau correction 

91 G.tau[CG.rank] = tauFG - tauG 

92 

93 if F.tau[CF.rank] is not None: 

94 tmp_tau = self.space_transfer.restrict(F.tau[CF.rank]) 

95 

96 # restrict possible tau correction from fine in collocation 

97 recvBuf = [None for _ in range(SG.coll.num_nodes)] 

98 recvBuf[CG.rank] = PG.u_init 

99 for n in range(SG.coll.num_nodes): 

100 CF.Reduce(self.Rcoll[n, CF.rank] * tmp_tau, recvBuf[CG.rank], root=n, op=MPI.SUM) 

101 G.tau[CG.rank] += recvBuf[CG.rank] 

102 else: 

103 pass 

104 

105 # save u and rhs evaluations for interpolation 

106 G.uold[CG.rank + 1] = PG.dtype_u(G.u[CG.rank + 1]) 

107 G.fold[CG.rank + 1] = PG.dtype_f(G.f[CG.rank + 1]) 

108 PG.use_restricted_boundary_values([G.u[CG.rank + 1]], [G.time + G.dt * SG.coll.nodes[CG.rank]]) 

109 

110 # works as a predictor 

111 G.status.unlocked = True 

112 

113 return None 

114 

115 def prolong(self): 

116 """ 

117 Space-time prolongation routine 

118 

119 This routine applies the spatial prolongation routine to the difference between the computed and the restricted 

120 values on the coarse level and then adds this difference to the fine values as coarse correction. 

121 """ 

122 

123 # get data for easier access 

124 F, G = self.fine, self.coarse 

125 CF, CG = self.comm_fine, self.comm_coarse 

126 SF = F.sweep 

127 PF = F.prob 

128 

129 # only of the level is unlocked at least by prediction or restriction 

130 if not G.status.unlocked: 

131 raise UnlockError('coarse level is still locked, cannot use data from there') 

132 

133 # build coarse correction 

134 

135 # interpolate values in space first 

136 tmp_u = self.space_transfer.prolong(G.u[CF.rank + 1] - G.uold[CF.rank + 1]) 

137 

138 # interpolate values in collocation 

139 recvBuf = [None for _ in range(SF.coll.num_nodes)] 

140 recvBuf[CF.rank] = F.u[CF.rank + 1].copy() 

141 for n in range(SF.coll.num_nodes): 

142 

143 CG.Reduce(self.Pcoll[n, CG.rank] * tmp_u, recvBuf[n], root=n, op=MPI.SUM) 

144 F.u[CF.rank + 1] += recvBuf[CF.rank] 

145 

146 # re-evaluate f on fine level 

147 F.f[CF.rank + 1] = PF.eval_f(F.u[CF.rank + 1], F.time + F.dt * SF.coll.nodes[CF.rank]) 

148 

149 return None 

150 

151 def prolong_f(self): 

152 """ 

153 Space-time prolongation routine w.r.t. the rhs f 

154 

155 This routine applies the spatial prolongation routine to the difference between the computed and the restricted 

156 values on the coarse level and then adds this difference to the fine values as coarse correction. 

157 """ 

158 

159 # get data for easier access 

160 F, G = self.fine, self.coarse 

161 CF, CG = self.comm_fine, self.comm_coarse 

162 SF = F.sweep 

163 

164 # only of the level is unlocked at least by prediction or restriction 

165 if not G.status.unlocked: 

166 raise UnlockError('coarse level is still locked, cannot use data from there') 

167 

168 # build coarse correction 

169 

170 # interpolate values in space first 

171 tmp_u = self.space_transfer.prolong(G.u[CF.rank + 1] - G.uold[CF.rank + 1]) 

172 tmp_f = self.space_transfer.prolong(G.f[CF.rank + 1] - G.fold[CF.rank + 1]) 

173 

174 # interpolate values in collocation 

175 recvBuf_u = [None for _ in range(SF.coll.num_nodes)] 

176 recvBuf_f = [None for _ in range(SF.coll.num_nodes)] 

177 recvBuf_u[CF.rank] = F.u[CF.rank + 1].copy() 

178 recvBuf_f[CF.rank] = F.f[CF.rank + 1].copy() 

179 for n in range(SF.coll.num_nodes): 

180 

181 CG.Reduce(self.Pcoll[n, CG.rank] * tmp_u, recvBuf_u[CF.rank], root=n, op=MPI.SUM) 

182 CG.Reduce(self.Pcoll[n, CG.rank] * tmp_f, recvBuf_f[CF.rank], root=n, op=MPI.SUM) 

183 

184 F.u[CF.rank + 1] += recvBuf_u[CF.rank] 

185 F.f[CF.rank + 1] += recvBuf_f[CF.rank] 

186 

187 return None