Coverage for pySDC/implementations/transfer_classes/BaseTransferDeltaMPI.py: 90%

1r""" 

2Node-parallel counterpart of :mod:`pySDC.implementations.transfer_classes.BaseTransferDelta`. 

3 

4Same two identities -- the coarse residual is the restricted fine residual, and the coarse-grid 

5correction is the sum of the sweep's own increments -- with every loop over collocation nodes 

6written as the reduction the node-parallel layout needs, since one rank holds one node. 

7 

8One difference from the serial transfer: this one still builds the FAS :math:`\tau`, which the 

9delta-form hierarchy then never reads. Skipping it is the saving 

10:meth:`BaseTransferDelta.delta_transfer.restrict_state` takes, and it has not been written for the 

11node-parallel layout yet. 

12""" 

13 

14from mpi4py import MPI 

15 

16from pySDC.implementations.transfer_classes.BaseTransferMPI import base_transfer_MPI 

17 

18 

19class delta_transfer_MPI(base_transfer_MPI): 

20 """ 

21 Node-parallel counterpart of 

22 :class:`~pySDC.implementations.transfer_classes.BaseTransferDelta.delta_transfer`. 

23 

24 Same two identities -- the coarse residual is the restricted fine residual, and the coarse-grid 

25 correction is the sum of the sweep's own increments -- with the collocation transfer written as 

26 a reduction, since each rank holds one node. 

27 """ 

28 

29 def restrict(self): 

30 """ 

31 Restrict as usual, then hand the coarse level the restricted fine residual. 

32 

33 Returns 

34 ------- 

35 None 

36 """ 

37 SF, SG = self.fine.sweep, self.coarse.sweep 

38 SF._delta_setup() 

39 SG._delta_setup() 

40 eps_F = SF._residual_nodes()[0] 

41 super().restrict() 

42 

43 CF, CG, PG = self.comm_fine, self.comm_coarse, self.coarse.prob 

44 tmp = self.space_transfer.restrict(eps_F) 

45 received = PG.u_init 

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

47 CF.Reduce(self.Rcoll[n, CF.rank] * tmp, received if n == CG.rank else None, root=n, op=MPI.SUM) 

48 

49 SG.eps_in = [SG._to_work(PG, received)] 

50 SG.delta_acc = None 

51 self.coarse.u0_reference = PG.dtype_u(self.coarse.u[0]) 

52 return None 

53 

54 def prolong(self): 

55 """ 

56 Add the coarse level's accumulated correction to the fine level. 

57 

58 Returns 

59 ------- 

60 None 

61 """ 

62 if self.coarse.sweep.delta_acc is None: 

63 return super().prolong() 

64 

65 F, PF = self.fine, self.fine.prob 

66 SF, CF, CG = self.fine.sweep, self.comm_fine, self.comm_coarse 

67 SF._delta_setup() 

68 tmp = self.space_transfer.prolong(self.coarse.sweep.delta_acc[0]) 

69 

70 correction = PF.u_init 

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

72 CG.Reduce(self.Pcoll[n, CG.rank] * tmp, correction if n == CF.rank else None, root=n, op=MPI.SUM) 

73 # quantised, then back to backend units before it is applied -- see the serial transfer 

74 applied = SF._to_backend(PF, SF._to_work(PF, correction)) 

75 correction = SF._to_work(PF, applied) 

76 

77 rank = CF.rank 

78 t_node = F.time + F.dt * SF.coll.nodes[rank] 

79 u_old, f_old = PF.dtype_u(F.u[rank + 1]), PF.dtype_f(F.f[rank + 1]) 

80 F.u[rank + 1] += applied 

81 F.f[rank + 1] = PF.eval_f(F.u[rank + 1], t_node) 

82 

83 if SF.eps_in is not None: 

84 # records into SF._dfs, which is what advance_residual reads 

85 SF._f_increment(PF, F.f[rank + 1], f_old, u_old, applied, t_node) 

86 SF.eps_in = SF.advance_residual(SF.eps_in, [correction], SF._dfs) 

87 SF.accumulate([correction]) 

88 return None