Coverage for pySDC/implementations/transfer_classes/TransferMesh.py: 82%

1import numpy as np 

2import scipy.sparse as sp 

3 

4import pySDC.helpers.transfer_helper as th 

5from pySDC.core.errors import TransferError 

6from pySDC.core.space_transfer import SpaceTransfer 

7 

8 

9def _unit_grid(nvars, periodic): 

10 """ 

11 Grid of one spatial dimension, scaled to a domain of length one. 

12 

13 The interpolation weights depend only on ratios of distances, so building the grids on the unit 

14 domain instead of on the problem's own ``dx`` gives the same operator for any domain length. 

15 For periodic grids it is also what makes the operator correct at all: the helpers in 

16 ``transfer_helper`` hardcode a period of one. 

17 

18 Args: 

19 nvars (int): number of degrees of freedom in this dimension 

20 periodic (bool): whether this dimension is periodic 

21 

22 Returns: 

23 np.ndarray: the grid, within [0, 1) 

24 """ 

25 return np.arange(nvars) / nvars if periodic else (np.arange(nvars) + 1) / (nvars + 1) 

26 

27 

28def _transfer_matrices_1d(nvars_fine, nvars_coarse, params): 

29 """ 

30 Interpolation and restriction matrix of one spatial dimension. 

31 

32 The ``equidist_nested`` shortcut of ``interpolation_matrix_1d`` and the restriction scale of 0.5 

33 are only right for grids that are nested with coarsening factor two, that is ``nf = 2 nc`` for 

34 periodic and ``nf = 2 nc + 1`` for Dirichlet grids. Any other pair of sizes takes the general 

35 interpolation, and its restriction is the transposed interpolation normalised to reproduce 

36 constants, which is what the factor 0.5 does on the nested grids. 

37 

38 Args: 

39 nvars_fine (int): number of degrees of freedom on the fine level 

40 nvars_coarse (int): number of degrees of freedom on the coarse level 

41 params: the parameters of the transfer class 

42 

43 Returns: 

44 tuple: the interpolation matrix (nf x nc) and the restriction matrix (nc x nf) 

45 """ 

46 # if number of variables is the same on both levels, Rspace and Pspace are identity 

47 if nvars_coarse == nvars_fine: 

48 return sp.eye(nvars_fine), sp.eye(nvars_coarse) 

49 

50 nested = nvars_fine == 2 * nvars_coarse + (not params.periodic) 

51 if not nested and min(params.iorder, params.rorder) == 0: 

52 raise TransferError( 

53 f'Order 0 (injection) needs grids nested with factor 2, got {nvars_fine} and {nvars_coarse} points' 

54 ) 

55 

56 fine_grid = _unit_grid(nvars_fine, params.periodic) 

57 coarse_grid = _unit_grid(nvars_coarse, params.periodic) 

58 

59 def interpolation_matrix(k): 

60 return th.interpolation_matrix_1d( 

61 fine_grid, 

62 coarse_grid, 

63 k=k, 

64 periodic=params.periodic, 

65 equidist_nested=params.equidist_nested and nested, 

66 ) 

67 

68 # assemble restriction as transpose of interpolation 

69 Pspace = interpolation_matrix(params.iorder) 

70 RspaceT = Pspace.T if params.iorder == params.rorder else interpolation_matrix(params.rorder).T 

71 

72 if not nested: 

73 return Pspace, sp.diags(1.0 / np.asarray(RspaceT.sum(axis=1)).ravel()) @ RspaceT 

74 

75 if params.rorder > 0: 

76 restr_factor = 0.5 

77 else: 

78 restr_factor = 1.0 

79 return Pspace, restr_factor * RspaceT 

80 

81 

82class mesh_to_mesh(SpaceTransfer): 

83 """ 

84 Space transfer between nd meshes by sparse interpolation matrices of even order, restricting by their transpose. 

85 

86 This implementation can restrict and prolong between nd meshes with dirichlet-0 or periodic boundaries 

87 via matrix-vector products. 

88 

89 Attributes: 

90 Rspace: spatial restriction matrix, dim. Nc x Nf 

91 Pspace: spatial prolongation matrix, dim. Nf x Nc 

92 """ 

93 

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

95 """ 

96 Initialization routine 

97 

98 Args: 

99 fine_prob: fine problem 

100 coarse_prob: coarse problem 

101 params: parameters for the transfer operators 

102 """ 

103 

104 # invoke super initialization 

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

106 

107 if self.params.rorder % 2 != 0: 

108 raise TransferError('Need even order for restriction') 

109 

110 if self.params.iorder % 2 != 0: 

111 raise TransferError('Need even order for interpolation') 

112 

113 if type(self.fine_prob.nvars) is tuple: 

114 if type(self.coarse_prob.nvars) is not tuple: 

115 raise TransferError('nvars parameter of coarse problem needs to be a tuple') 

116 if not len(self.fine_prob.nvars) == len(self.coarse_prob.nvars): 

117 raise TransferError('nvars parameter of fine and coarse level needs to have the same length') 

118 elif type(self.fine_prob.nvars) is int: 

119 if type(self.coarse_prob.nvars) is not int: 

120 raise TransferError('nvars parameter of coarse problem needs to be an int') 

121 else: 

122 raise TransferError("unknow type of nvars for transfer, got %s" % self.fine_prob.nvars) 

123 

124 # we have a 1d problem 

125 if type(self.fine_prob.nvars) is int: 

126 self.Pspace, self.Rspace = _transfer_matrices_1d(self.fine_prob.nvars, self.coarse_prob.nvars, self.params) 

127 

128 # we have an n-d problem 

129 else: 

130 Rspace = [] 

131 Pspace = [] 

132 for i in range(len(self.fine_prob.nvars)): 

133 P, R = _transfer_matrices_1d(self.fine_prob.nvars[i], self.coarse_prob.nvars[i], self.params) 

134 Pspace.append(P) 

135 Rspace.append(R) 

136 

137 # kronecker 1-d operators for n-d 

138 self.Pspace = Pspace[0] 

139 for i in range(1, len(Pspace)): 

140 self.Pspace = sp.kron(self.Pspace, Pspace[i], format='csc') 

141 

142 self.Rspace = Rspace[0] 

143 for i in range(1, len(Rspace)): 

144 self.Rspace = sp.kron(self.Rspace, Rspace[i], format='csc') 

145 

146 # Carry the operators at the precision of what they produce, rather than always at float64. 

147 # A float64 operator applied to a reduced-precision vector upcasts, so the transfer would do 

148 # its work in double and only round on the way into the destination, which is the one place 

149 # a reduced-precision level would silently keep paying full freight. `promote_types` with 

150 # float32 keeps it legal for SciPy, which has no half-precision sparse matrix. 

151 self.Rspace = self.Rspace.astype(np.promote_types(self.coarse_prob.init[-1], np.float32)) 

152 self.Pspace = self.Pspace.astype(np.promote_types(self.fine_prob.init[-1], np.float32)) 

153 

154 # Which side of the PCI bus this runs on is a property of the problem, not something the 

155 # transfer is told -- the same way `TransferMesh_MPIFFT` decides it. The operators are 

156 # assembled with SciPy either way, since that work is small, one-off and full of host-side 

157 # index arithmetic; only the finished matrices move. 

158 if 'cupy' in self.fine_prob.dtype_u.__name__.lower(): 

159 import cupyx.scipy.sparse as csp 

160 

161 self.Rspace = csp.csr_matrix(self.Rspace) 

162 self.Pspace = csp.csr_matrix(self.Pspace) 

163 

164 def restrict(self, F): 

165 """ 

166 Restriction implementation 

167 Args: 

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

169 """ 

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

171 

172 def _restrict(fine, coarse): 

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

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

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

176 tmpF = fine[..., i].flatten() 

177 tmpG = self.Rspace.dot(tmpF) 

178 coarse[..., i] = tmpG.reshape(self.coarse_prob.nvars) 

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

180 tmpF = fine[i, ...].flatten() 

181 tmpG = self.Rspace.dot(tmpF) 

182 coarse[i, ...] = tmpG.reshape(self.coarse_prob.nvars) 

183 else: 

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

185 else: 

186 tmpF = fine.flatten() 

187 tmpG = self.Rspace.dot(tmpF) 

188 coarse[:] = tmpG.reshape(self.coarse_prob.nvars) 

189 

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

191 for comp in F.components: 

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

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

194 _restrict(F, G) 

195 else: 

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

197 return G 

198 

199 def prolong(self, G): 

200 """ 

201 Prolongation implementation 

202 Args: 

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

204 """ 

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

206 

207 def _prolong(coarse, fine): 

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

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

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

211 tmpG = coarse[..., i].flatten() 

212 tmpF = self.Pspace.dot(tmpG) 

213 fine[..., i] = tmpF.reshape(self.fine_prob.nvars) 

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

215 tmpG = coarse[i, ...].flatten() 

216 tmpF = self.Pspace.dot(tmpG) 

217 fine[i, ...] = tmpF.reshape(self.fine_prob.nvars) 

218 else: 

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

220 else: 

221 tmpG = coarse.flatten() 

222 tmpF = self.Pspace.dot(tmpG) 

223 fine[:] = tmpF.reshape(self.fine_prob.nvars) 

224 return fine 

225 

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

227 for comp in G.components: 

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

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

230 F[:] = _prolong(G, F) 

231 else: 

232 raise TransferError('Wrong data type for prolongation, got %s' % type(G)) 

233 return F