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

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.SpaceTransfer import space_transfer 

7 

8 

9class mesh_to_mesh(space_transfer): 

10 """ 

11 Custom base_transfer class, implements Transfer.py 

12 

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

14 via matrix-vector products. 

15 

16 Attributes: 

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

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

19 """ 

20 

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

22 """ 

23 Initialization routine 

24 

25 Args: 

26 fine_prob: fine problem 

27 coarse_prob: coarse problem 

28 params: parameters for the transfer operators 

29 """ 

30 

31 # invoke super initialization 

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

33 

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

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

36 

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

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

39 

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

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

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

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

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

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

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

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

48 else: 

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

50 

51 # we have a 1d problem 

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

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

54 if self.coarse_prob.nvars == self.fine_prob.nvars: 

55 self.Rspace = sp.eye(self.coarse_prob.nvars) 

56 self.Pspace = sp.eye(self.fine_prob.nvars) 

57 # assemble restriction as transpose of interpolation 

58 else: 

59 if not self.params.periodic: 

60 fine_grid = np.array([(i + 1) * self.fine_prob.dx for i in range(self.fine_prob.nvars)]) 

61 coarse_grid = np.array([(i + 1) * self.coarse_prob.dx for i in range(self.coarse_prob.nvars)]) 

62 else: 

63 fine_grid = np.array([i * self.fine_prob.dx for i in range(self.fine_prob.nvars)]) 

64 coarse_grid = np.array([i * self.coarse_prob.dx for i in range(self.coarse_prob.nvars)]) 

65 

66 self.Pspace = th.interpolation_matrix_1d( 

67 fine_grid, 

68 coarse_grid, 

69 k=self.params.iorder, 

70 periodic=self.params.periodic, 

71 equidist_nested=self.params.equidist_nested, 

72 ) 

73 if self.params.rorder > 0: 

74 restr_factor = 0.5 

75 else: 

76 restr_factor = 1.0 

77 

78 if self.params.iorder == self.params.rorder: 

79 self.Rspace = restr_factor * self.Pspace.T 

80 

81 else: 

82 self.Rspace = ( 

83 restr_factor 

84 * th.interpolation_matrix_1d( 

85 fine_grid, 

86 coarse_grid, 

87 k=self.params.rorder, 

88 periodic=self.params.periodic, 

89 equidist_nested=self.params.equidist_nested, 

90 ).T 

91 ) 

92 

93 # we have an n-d problem 

94 else: 

95 Rspace = [] 

96 Pspace = [] 

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

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

99 if self.coarse_prob.nvars == self.fine_prob.nvars: 

100 Rspace.append(sp.eye(self.coarse_prob.nvars[i])) 

101 Pspace.append(sp.eye(self.fine_prob.nvars[i])) 

102 # assemble restriction as transpose of interpolation 

103 else: 

104 if not self.params.periodic: 

105 fine_grid = np.array([(j + 1) * self.fine_prob.dx for j in range(self.fine_prob.nvars[i])]) 

106 coarse_grid = np.array( 

107 [(j + 1) * self.coarse_prob.dx for j in range(self.coarse_prob.nvars[i])] 

108 ) 

109 else: 

110 fine_grid = np.array([j * self.fine_prob.dx for j in range(self.fine_prob.nvars[i])]) 

111 coarse_grid = np.array([j * self.coarse_prob.dx for j in range(self.coarse_prob.nvars[i])]) 

112 

113 Pspace.append( 

114 th.interpolation_matrix_1d( 

115 fine_grid, 

116 coarse_grid, 

117 k=self.params.iorder, 

118 periodic=self.params.periodic, 

119 equidist_nested=self.params.equidist_nested, 

120 ) 

121 ) 

122 if self.params.rorder > 0: 

123 restr_factor = 0.5 

124 else: 

125 restr_factor = 1.0 

126 

127 if self.params.iorder == self.params.rorder: 

128 Rspace.append(restr_factor * Pspace[-1].T) 

129 

130 else: 

131 mat = th.interpolation_matrix_1d( 

132 fine_grid, 

133 coarse_grid, 

134 k=self.params.rorder, 

135 periodic=self.params.periodic, 

136 equidist_nested=self.params.equidist_nested, 

137 ).T 

138 Rspace.append(restr_factor * mat) 

139 

140 # kronecker 1-d operators for n-d 

141 self.Pspace = Pspace[0] 

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

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

144 

145 self.Rspace = Rspace[0] 

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

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

148 

149 def restrict(self, F): 

150 """ 

151 Restriction implementation 

152 Args: 

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

154 """ 

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

156 

157 def _restrict(fine, coarse): 

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

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

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

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

162 tmpG = self.Rspace.dot(tmpF) 

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

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

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

166 tmpG = self.Rspace.dot(tmpF) 

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

168 else: 

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

170 else: 

171 tmpF = fine.flatten() 

172 tmpG = self.Rspace.dot(tmpF) 

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

174 

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

176 for comp in F.components: 

177 _restrict(F.__getattr__(comp), G.__getattr__(comp)) 

178 elif type(F).__name__ == 'mesh': 

179 _restrict(F, G) 

180 else: 

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

182 return G 

183 

184 def prolong(self, G): 

185 """ 

186 Prolongation implementation 

187 Args: 

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

189 """ 

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

191 

192 def _prolong(coarse, fine): 

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

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

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

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

197 tmpF = self.Pspace.dot(tmpG) 

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

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

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

201 tmpF = self.Pspace.dot(tmpG) 

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

203 else: 

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

205 else: 

206 tmpG = coarse.flatten() 

207 tmpF = self.Pspace.dot(tmpG) 

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

209 return fine 

210 

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

212 for comp in G.components: 

213 _prolong(G.__getattr__(comp), F.__getattr__(comp)) 

214 elif type(G).__name__ == 'mesh': 

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

216 else: 

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

218 return F