Coverage for pySDC/helpers/transfer_helper.py: 100%

1# coding=utf-8 

2import numpy as np 

3import scipy.sparse as sprs 

4from scipy.interpolate import BarycentricInterpolator 

5 

6 

7def next_neighbors_periodic(p, ps, k): 

8 """ 

9 Function to find the next neighbors for a periodic setup 

10 

11 This function gives for a value p the k points next to it which are found in 

12 in the vector ps and the points which are found periodically. 

13 

14 Note that the period is hardcoded to one, here and in the periodic branches of the two matrix 

15 constructors below. Normalise the grid before you pass it in; the interpolation weights only 

16 depend on ratios of distances, so nothing is lost by doing so. 

17 

18 Args: 

19 p: the current point 

20 ps (np.ndarray): the grid with the potential neighbors 

21 k (int): number of neighbors to find 

22 

23 Returns: 

24 list: the k next neighbors 

25 """ 

26 p_bar = p - np.floor(p / 1.0) * 1.0 

27 ps = ps - ps[0] 

28 distance_to_p = np.asarray( 

29 list(map(lambda tk: min([np.abs(tk + 1 - p_bar), np.abs(tk - p_bar), np.abs(tk - 1 - p_bar)]), ps)) 

30 ) 

31 

32 # zip it 

33 value_index = [] 

34 for d, i in zip(distance_to_p, range(distance_to_p.size), strict=True): 

35 value_index.append((d, i)) 

36 # sort by distance 

37 value_index_sorted = sorted(value_index, key=lambda s: s[0]) 

38 # take first k indices with least distance and sort them 

39 return sorted(map(lambda s: s[1], value_index_sorted[0:k])) 

40 

41 

42def next_neighbors(p, ps, k): 

43 """ 

44 Function to find the next neighbors for a non-periodic setup 

45 

46 This function gives for a value p the k points next to it which are found in 

47 in the vector ps 

48 

49 Args: 

50 p: the current point 

51 ps (np.ndarray): the grid with the potential neighbors 

52 k (int): number of neighbors to find 

53 

54 Returns: 

55 list: the k next neighbors 

56 """ 

57 distance_to_p = np.abs(ps - p) 

58 # zip it 

59 value_index = [] 

60 for d, i in zip(distance_to_p, range(distance_to_p.size), strict=True): 

61 value_index.append((d, i)) 

62 # sort by distance 

63 value_index_sorted = sorted(value_index, key=lambda s: s[0]) 

64 # take first k indices with least distance and sort them 

65 return sorted(map(lambda s: s[1], value_index_sorted[0:k])) 

66 

67 

68def continue_periodic_array(arr, nn): 

69 """ 

70 Function to append an array for nn neighbors for periodicity 

71 

72 Args: 

73 arr (np.ndarray): the input array 

74 nn (list): the neighbors 

75 

76 Returns: 

77 np.ndarray: the continued array 

78 """ 

79 

80 nn = np.asarray(nn) 

81 d_nn = nn[1:] - nn[:-1] 

82 if np.all(d_nn == np.ones(nn.shape[0] - 1)): 

83 return arr[nn] 

84 else: 

85 cont_arr = [arr[nn[0]]] 

86 shift = 0.0 

87 for n, d in zip(nn[1:], d_nn, strict=True): 

88 if d != 1: 

89 shift = -1 

90 cont_arr.append(arr[n] + shift) 

91 

92 return np.asarray(cont_arr) 

93 

94 

95def restriction_matrix_1d(fine_grid, coarse_grid, k=2, periodic=False, pad=1): 

96 """ 

97 Function to contruct the restriction matrix in 1d using barycentric interpolation 

98 

99 Args: 

100 fine_grid (np.ndarray): a one dimensional 1d array containing the nodes of the fine grid 

101 coarse_grid (np.ndarray): a one dimensional 1d array containing the nodes of the coarse grid 

102 k (int): order of the restriction 

103 periodic (bool): flag to indicate periodicity, which assumes a grid of period one 

104 pad (int): padding parameter for boundaries 

105 

106 Returns: 

107 sprs.csc_matrix: restriction matrix 

108 """ 

109 

110 n_g = coarse_grid.size 

111 

112 if periodic: 

113 M = np.zeros((coarse_grid.size, fine_grid.size)) 

114 for i, p in zip(range(n_g), coarse_grid, strict=True): 

115 nn = next_neighbors_periodic(p, fine_grid, k) 

116 circulating_one = np.asarray([1.0] + [0.0] * (k - 1)) 

117 cont_arr = continue_periodic_array(fine_grid, nn) 

118 if p > np.mean(coarse_grid) and not (cont_arr[0] <= p <= cont_arr[-1]): 

119 cont_arr += 1 

120 bary_pol = [] 

121 for l in range(k): 

122 bary_pol.append(BarycentricInterpolator(cont_arr, np.roll(circulating_one, l))) 

123 with np.errstate(divide='ignore'): 

124 M[i, nn] = np.asarray(list(map(lambda x: x(p), bary_pol))) 

125 else: 

126 M = np.zeros((coarse_grid.size, fine_grid.size + 2 * pad)) 

127 for i, p in zip(range(n_g), coarse_grid, strict=True): 

128 padded_f_grid = border_padding(fine_grid, pad, pad) 

129 nn = next_neighbors(p, padded_f_grid, k) 

130 # construct the lagrange polynomials for the k neighbors 

131 circulating_one = np.asarray([1.0] + [0.0] * (k - 1)) 

132 bary_pol = [] 

133 for l in range(k): 

134 bary_pol.append(BarycentricInterpolator(padded_f_grid[nn], np.roll(circulating_one, l))) 

135 with np.errstate(divide='ignore'): 

136 M[i, nn] = np.asarray(list(map(lambda x: x(p), bary_pol))) 

137 if pad > 0: 

138 M = M[:, pad:-pad] 

139 

140 return sprs.csc_matrix(M) 

141 

142 

143def interpolation_matrix_1d(fine_grid, coarse_grid, k=2, periodic=False, pad=1, equidist_nested=True): 

144 """ 

145 Function to construct the restriction matrix in 1d using barycentric interpolation 

146 

147 Args: 

148 fine_grid (np.ndarray): a one dimensional 1d array containing the nodes of the fine grid 

149 coarse_grid (np.ndarray): a one dimensional 1d array containing the nodes of the coarse grid 

150 k (int): order of the restriction 

151 periodic (bool): flag to indicate periodicity, which assumes a grid of period one 

152 pad (int): padding parameter for boundaries 

153 equidist_nested (bool): shortcut possible, if nodes are equidistant and nested 

154 

155 Returns: 

156 sprs.csc_matrix: interpolation matrix 

157 """ 

158 

159 n_f = fine_grid.size 

160 

161 if periodic: 

162 M = np.zeros((fine_grid.size, coarse_grid.size)) 

163 

164 if equidist_nested: 

165 for i, p in zip(range(n_f), fine_grid, strict=True): 

166 if i % 2 == 0: 

167 M[i, int(i / 2)] = 1.0 

168 else: 

169 nn = [] 

170 cpos = int(i / 2) 

171 offset = int(k / 2) 

172 for j in range(k): 

173 nn.append(cpos - offset + 1 + j) 

174 if nn[-1] < 0: 

175 nn[-1] += coarse_grid.size 

176 elif nn[-1] > coarse_grid.size - 1: 

177 nn[-1] -= coarse_grid.size 

178 nn = sorted(nn) 

179 

180 circulating_one = np.asarray([1.0] + [0.0] * (k - 1)) 

181 if len(nn) > 0: 

182 cont_arr = continue_periodic_array(coarse_grid, nn) 

183 else: 

184 cont_arr = coarse_grid 

185 

186 if p > np.mean(fine_grid) and not (cont_arr[0] <= p <= cont_arr[-1]): 

187 cont_arr += 1 

188 

189 bary_pol = [] 

190 for l in range(k): 

191 bary_pol.append(BarycentricInterpolator(cont_arr, np.roll(circulating_one, l))) 

192 with np.errstate(divide='ignore'): 

193 M[i, nn] = np.asarray(list(map(lambda x: x(p), bary_pol))) 

194 

195 else: 

196 for i, p in zip(range(n_f), fine_grid, strict=True): 

197 nn = next_neighbors_periodic(p, coarse_grid, k) 

198 circulating_one = np.asarray([1.0] + [0.0] * (k - 1)) 

199 cont_arr = continue_periodic_array(coarse_grid, nn) 

200 

201 if p > np.mean(fine_grid) and not (cont_arr[0] <= p <= cont_arr[-1]): 

202 cont_arr += 1 

203 

204 bary_pol = [] 

205 for l in range(k): 

206 bary_pol.append(BarycentricInterpolator(cont_arr, np.roll(circulating_one, l))) 

207 with np.errstate(divide='ignore'): 

208 M[i, nn] = np.asarray(list(map(lambda x: x(p), bary_pol))) 

209 

210 else: 

211 M = np.zeros((fine_grid.size, coarse_grid.size + 2 * pad)) 

212 padded_c_grid = border_padding(coarse_grid, pad, pad) 

213 

214 if equidist_nested: 

215 for i, p in zip(range(n_f), fine_grid, strict=True): 

216 if i % 2 != 0: 

217 M[i, int((i - 1) / 2) + 1] = 1.0 

218 else: 

219 nn = [] 

220 cpos = int(i / 2) 

221 offset = int(k / 2) 

222 for j in range(k): 

223 nn.append(cpos - offset + 1 + j) 

224 if nn[-1] < 0: 

225 nn[-1] += k 

226 elif nn[-1] > coarse_grid.size + 1: 

227 nn[-1] -= k 

228 nn = sorted(nn) 

229 # construct the lagrange polynomials for the k neighbors 

230 circulating_one = np.asarray([1.0] + [0.0] * (k - 1)) 

231 bary_pol = [] 

232 for l in range(k): 

233 bary_pol.append(BarycentricInterpolator(padded_c_grid[nn], np.roll(circulating_one, l))) 

234 with np.errstate(divide='ignore'): 

235 M[i, nn] = np.asarray(list(map(lambda x: x(p), bary_pol))) 

236 

237 else: 

238 for i, p in zip(range(n_f), fine_grid, strict=True): 

239 nn = next_neighbors(p, padded_c_grid, k) 

240 # construct the lagrange polynomials for the k neighbors 

241 circulating_one = np.asarray([1.0] + [0.0] * (k - 1)) 

242 bary_pol = [] 

243 for l in range(k): 

244 bary_pol.append(BarycentricInterpolator(padded_c_grid[nn], np.roll(circulating_one, l))) 

245 with np.errstate(divide='ignore'): 

246 M[i, nn] = np.asarray(list(map(lambda x: x(p), bary_pol))) 

247 

248 if pad > 0: 

249 M = M[:, pad:-pad] 

250 

251 return sprs.csc_matrix(M) 

252 

253 

254def border_padding(grid, l, r, pad_type='mirror'): 

255 """ 

256 Function to pad/embed an array at the boundaries 

257 

258 Args: 

259 grid (np.npdarray): the input array 

260 l: left boundary 

261 r: right boundary 

262 pad_type: type of padding 

263 

264 Returns: 

265 np.npdarray: the padded array 

266 

267 """ 

268 

269 assert l < grid.size and r < grid.size 

270 padded_arr = np.zeros(grid.size + l + r) 

271 if pad_type == 'mirror': 

272 for i in range(l): 

273 padded_arr[i] = 2 * grid[0] - grid[l - i] 

274 for j in range(r): 

275 padded_arr[-j - 1] = 2 * grid[-1] - grid[-r + j - 1] 

276 padded_arr[l : l + grid.size] = grid 

277 return padded_arr