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

1from pySDC.core.errors import TransferError 

2from pySDC.core.space_transfer import SpaceTransfer 

3from mpi4py_fft import newDistArray 

4 

5from pySDC.helpers.fft_helper import PFFT 

6 

7 

8class fft_to_fft(SpaceTransfer): 

9 """ 

10 Space transfer between distributed periodic mpi4py-fft meshes: injection to restrict, spectral padding to prolong. 

11 

12 This implementation can restrict and prolong between ``mesh`` or ``cupy_mesh`` datatypes, or datatypes made of 

13 several such components, with FFT for periodic boundaries 

14 

15 """ 

16 

17 def __init__(self, fine_prob, coarse_prob, params): 

18 """ 

19 Initialization routine 

20 

21 Args: 

22 fine_prob: fine problem 

23 coarse_prob: coarse problem 

24 params: parameters for the transfer operators 

25 """ 

26 # invoke super initialization 

27 super().__init__(fine_prob, coarse_prob, params) 

28 

29 assert self.fine_prob.spectral == self.coarse_prob.spectral 

30 

31 self.spectral = self.fine_prob.spectral 

32 

33 Nf = list(self.fine_prob.fft.global_shape()) 

34 Nc = list(self.coarse_prob.fft.global_shape()) 

35 self.ratio = [int(nf / nc) for nf, nc in zip(Nf, Nc, strict=True)] 

36 # slice for injection, one entry per dimension rather than hardcoded for 2d 

37 self.injection = tuple(slice(None, None, r) for r in self.ratio) 

38 axes = tuple(range(len(Nf))) 

39 

40 fft_args = {} 

41 useGPU = 'cupy' in self.fine_prob.dtype_u.__name__.lower() 

42 if useGPU: 

43 fft_args['backend'] = 'cupy' 

44 fft_args['comm_backend'] = 'NCCL' 

45 

46 self.fft_pad = PFFT( 

47 self.coarse_prob.comm, 

48 Nc, 

49 padding=self.ratio, 

50 axes=axes, 

51 dtype=self.coarse_prob.fft.dtype(False), 

52 slab=True, 

53 **fft_args, 

54 ) 

55 

56 def restrict(self, F): 

57 """ 

58 Restriction implementation 

59 

60 Args: 

61 F: the fine level data (easier to access than via the fine attribute) 

62 """ 

63 G = type(F)(self.coarse_prob.init) 

64 

65 def _restrict(fine, coarse): 

66 if self.spectral: 

67 if hasattr(self.fine_prob, 'ncomp'): 

68 for i in range(self.fine_prob.ncomp): 

69 if fine.shape[-1] == self.fine_prob.ncomp: 

70 tmpF = newDistArray(self.fine_prob.fft, False) 

71 tmpF = self.fine_prob.fft.backward(fine[..., i], tmpF) 

72 tmpG = tmpF[self.injection] 

73 coarse[..., i] = self.coarse_prob.fft.forward(tmpG, coarse[..., i]) 

74 elif fine.shape[0] == self.fine_prob.ncomp: 

75 tmpF = newDistArray(self.fine_prob.fft, False) 

76 tmpF = self.fine_prob.fft.backward(fine[i, ...], tmpF) 

77 tmpG = tmpF[self.injection] 

78 coarse[i, ...] = self.coarse_prob.fft.forward(tmpG, coarse[i, ...]) 

79 else: 

80 raise TransferError('Don\'t know how to restrict for this problem with multiple components') 

81 else: 

82 tmpF = self.fine_prob.fft.backward(fine) 

83 tmpG = tmpF[self.injection] 

84 coarse[:] = self.coarse_prob.fft.forward(tmpG, coarse) 

85 elif hasattr(self.fine_prob, 'ncomp') and fine.shape[0] == self.fine_prob.ncomp: 

86 # the component axis comes first here, so inject on the spatial axes behind it. A 

87 # trailing component axis needs no special case: the slice simply does not reach it. 

88 coarse[:] = fine[(slice(None),) + self.injection] 

89 else: 

90 coarse[:] = fine[self.injection] 

91 

92 if hasattr(type(F), 'components'): 

93 for comp in F.components: 

94 _restrict(getattr(F, comp), getattr(G, comp)) 

95 elif type(F).__name__ in ['mesh', 'cupy_mesh']: 

96 _restrict(F, G) 

97 else: 

98 raise TransferError('Wrong data type for restriction, got %s' % type(F)) 

99 

100 return G 

101 

102 def prolong(self, G): 

103 """ 

104 Prolongation implementation 

105 

106 Args: 

107 G: the coarse level data (easier to access than via the coarse attribute) 

108 """ 

109 F = type(G)(self.fine_prob.init) 

110 

111 def _prolong(coarse, fine): 

112 if self.spectral: 

113 if hasattr(self.fine_prob, 'ncomp'): 

114 for i in range(self.fine_prob.ncomp): 

115 if coarse.shape[-1] == self.fine_prob.ncomp: 

116 tmpF = self.fft_pad.backward(coarse[..., i]) 

117 fine[..., i] = self.fine_prob.fft.forward(tmpF, fine[..., i]) 

118 elif coarse.shape[0] == self.fine_prob.ncomp: 

119 tmpF = self.fft_pad.backward(coarse[i, ...]) 

120 fine[i, ...] = self.fine_prob.fft.forward(tmpF, fine[i, ...]) 

121 else: 

122 raise TransferError('Don\'t know how to prolong for this problem with multiple components') 

123 

124 else: 

125 tmpF = self.fft_pad.backward(coarse) 

126 fine[:] = self.fine_prob.fft.forward(tmpF, fine) 

127 else: 

128 if hasattr(self.fine_prob, 'ncomp'): 

129 for i in range(self.fine_prob.ncomp): 

130 if coarse.shape[-1] == self.fine_prob.ncomp: 

131 G_hat = self.coarse_prob.fft.forward(coarse[..., i]) 

132 fine[..., i] = self.fft_pad.backward(G_hat, fine[..., i]) 

133 elif coarse.shape[0] == self.fine_prob.ncomp: 

134 G_hat = self.coarse_prob.fft.forward(coarse[i, ...]) 

135 fine[i, ...] = self.fft_pad.backward(G_hat, fine[i, ...]) 

136 else: 

137 raise TransferError('Don\'t know how to prolong for this problem with multiple components') 

138 else: 

139 G_hat = self.coarse_prob.fft.forward(coarse) 

140 fine[:] = self.fft_pad.backward(G_hat, fine) 

141 

142 if hasattr(type(F), 'components'): 

143 for comp in F.components: 

144 _prolong(getattr(G, comp), getattr(F, comp)) 

145 elif type(G).__name__ in ['mesh', 'cupy_mesh']: 

146 _prolong(G, F) 

147 

148 else: 

149 raise TransferError('Unknown data type, got %s' % type(G)) 

150 

151 return F