Coverage for pythia/index.py: 98%

85 statements  

« prev     ^ index     » next       coverage.py v7.3.1, created at 2023-09-08 17:13 +0000

1""" 

2File: pythia/index.py 

3Author: Nando Hegemann 

4Gitlab: https://gitlab.com/Nando-Hegemann 

5Description: Create, manipulate and store information about multiindices. 

6SPDX-License-Identifier: LGPL-3.0-or-later OR Hippocratic-3.0-ECO-MEDIA-MIL 

7""" 

8import itertools 

9import numpy as np 

10import pythia as pt 

11 

12 

13class IndexSet: 

14 """Generate index set object for sparse PC expansion. 

15 

16 A general polynomial chaos expansion of a function 

17 :math:`f\\colon\\Gamma\\subset\\mathbb{R}^M\\to\\mathbb{R}^J` 

18 with :math:`y\\sim\\pi` is given by 

19 

20 .. math:: 

21 f(y) = \\sum_{\\mu\\in\\mathbb{N}_0^M} \\mathbf{f}[\\mu]P_{\\mu}(y) 

22 \\quad\\mbox{for}\\quad 

23 \\mathbf{f}[\\mu] = \\int_\\Gamma f(y)P_\\mu(y)\\ \\mathrm{d}y, 

24 

25 where :math:`\\mu` is a multiindex, 

26 :math:`\\mathbf{f}[\\mu]\\in\\mathbb{R}^J` is a coefficient vector and 

27 :math:`\\{P_\\mu\\}_{\\mu\\in\\mathbb{N}_0^M}` is an orthonormal basis in 

28 :math:`L^2(\\Gamma,\\pi)`. 

29 To approximate the infinite expansion choose an index set 

30 :math:`\\Lambda\\subset\\mathbb{N}_0^M` of multiindices and consider 

31 

32 .. math:: 

33 f(y) \\approx \\sum_{\\mu\\in\\Lambda} \\mathbf{f}[\\mu]P_{\\mu}(y), 

34 

35 Parameters 

36 ---------- 

37 indices : np.ndarray 

38 Array of multiindices with shape (#indices, param dim). 

39 

40 Examples 

41 -------- 

42 Create the sparse index set 

43 

44 .. math:: 

45 \\Lambda = \\{ (0,0), (1,0), (2,0), (0,1) \\} \\subset \\mathbb{N}_0^2 

46 

47 >>> import pythia as pt 

48 >>> indices = np.array([[0, 0], [1, 0], [2, 0], [0, 1]], dtype=int) 

49 >>> index_set = pt.index.IndexSet(indices) 

50 """ 

51 

52 def __init__(self, indices: np.ndarray) -> None: 

53 """Initialize sparse multiindex object.""" 

54 assert indices.ndim == 2 and indices.shape[0] > 0 

55 assert indices.dtype == int 

56 assert np.all(indices >= 0) 

57 self.indices = sort_index_array(indices) 

58 self.shape = self.indices.shape 

59 self.max = np.array(np.max(self.indices, axis=0), dtype=int) 

60 self.min = np.array(np.min(self.indices, axis=0), dtype=int) 

61 self.sobol_tuples = self._get_sobol_tuple_list() 

62 

63 def _get_sobol_tuple_list(self) -> list: 

64 """Generate list of all possible Sobol index id tuples (subscripts). 

65 

66 Returns 

67 ------- 

68 : 

69 List of Sobol tuples. 

70 """ 

71 sobol_tuples = [] 

72 for r in range(1, self.shape[1] + 1): 

73 sobol_tuples += list(itertools.combinations(range(1, self.shape[1] + 1), r)) 

74 return sobol_tuples 

75 

76 def get_index_number(self, indices: np.ndarray) -> np.ndarray: 

77 """Get enumeration number of indices. 

78 

79 Get the row indices of the given multiindices such that 

80 `self.indices[rows] = indices`. 

81 

82 Parameters 

83 ---------- 

84 indices : np.ndarray 

85 Indices to get the number of. 

86 

87 Returns 

88 ------- 

89 : 

90 Array containing the enumeration numbers of the indices. 

91 """ 

92 return np.array( 

93 [np.where((self.indices == index).all(axis=1))[0] for index in indices], 

94 dtype=int, 

95 ).flatten() 

96 

97 def get_sobol_tuple_number(self, sobol_tuples: list[tuple]) -> np.ndarray: 

98 """Get enumeration indices of Sobol tuples. 

99 

100 Parameters 

101 ---------- 

102 sobol_tuples : list of tuple 

103 List of Sobol tuples. 

104 

105 Returns 

106 ------- 

107 : 

108 Array containing the enumeration number of the Sobol tuples. 

109 """ 

110 return np.array([self.sobol_tuples.index(s) for s in sobol_tuples], dtype=int) 

111 

112 def index_to_sobol_tuple(self, indices: np.ndarray) -> list[tuple]: 

113 """Map array of indices to their respective Sobol tuples. 

114 

115 Parameters 

116 ---------- 

117 indices : np.ndarray 

118 Array of multiindices. 

119 

120 Returns 

121 ------- 

122 : 

123 List of Sobol tuples. 

124 """ 

125 sobol_tuples = [tuple(np.flatnonzero(index) + 1) for index in indices] 

126 return sobol_tuples 

127 

128 def sobol_tuple_to_indices( 

129 self, sobol_tuples: tuple | list[tuple] 

130 ) -> list[np.ndarray]: 

131 """Map Sobol tuples to their respective indices. 

132 

133 Parameters 

134 ---------- 

135 sobol_tuples : tuple or list of tuple 

136 List of Sobol tuples. 

137 

138 Returns 

139 ------- 

140 : 

141 List of index arrays for each given Sobol tuple. 

142 """ 

143 if isinstance(sobol_tuples, tuple): 

144 sobol_tuples = [sobol_tuples] 

145 assert isinstance(sobol_tuples, list) 

146 ret = [] 

147 lookup_dict = {sobol_tuple: [] for sobol_tuple in self.sobol_tuples} 

148 index_sobol_tuple_list = self.index_to_sobol_tuple(self.indices) 

149 for sobol_tuple, index in zip(index_sobol_tuple_list, self.indices): 

150 if len(sobol_tuple) > 0: 

151 lookup_dict[sobol_tuple] += [index] 

152 for sobol_tuple in sobol_tuples: 

153 ret += [np.array(lookup_dict[sobol_tuple], dtype=int)] 

154 return ret 

155 

156 

157def sort_index_array(indices: np.ndarray) -> np.ndarray: 

158 """Sort multiindices and remove duplicates. 

159 

160 Sort rows of `indices` by sum of multiindex and remove duplicate 

161 multiindices. 

162 

163 Parameters 

164 ---------- 

165 indices : np.ndarray 

166 Index list before sorting. 

167 

168 Returns 

169 ------- 

170 : 

171 Sorted index array. 

172 """ 

173 sorted_indices = np.unique(indices, axis=0) 

174 if sorted_indices.size == 0: 

175 return sorted_indices 

176 idx = np.argsort(np.sum(sorted_indices, axis=1)) 

177 sorted_indices = sorted_indices[idx] 

178 return np.array(sorted_indices, dtype=int) 

179 

180 

181def union(index_list: list[np.ndarray]) -> np.ndarray: 

182 """Build union of multiindex sets. 

183 

184 Given sparse index sets :math:`\\Lambda_1, \\dots, \\Lambda_N`, 

185 compute :math:`\\Lambda=\\Lambda_1\\cup\\dots\\cup\\Lambda_N`. 

186 

187 Parameters 

188 ---------- 

189 index_list : list of np.ndarray 

190 List of multiindex arrays. 

191 

192 Returns 

193 ------- 

194 : 

195 Array with all multiindices. 

196 """ 

197 all_indices = np.concatenate(index_list, axis=0) 

198 return sort_index_array(all_indices) 

199 

200 

201def intersection(index_list: list[np.ndarray]) -> np.ndarray: 

202 """Intersect list of multiindex sets. 

203 

204 Given sparse index sets :math:`\\Lambda_1, \\dots, \\Lambda_N`, 

205 compute :math:`\\Lambda=\\Lambda_1\\cap\\dots\\cap\\Lambda_N`. 

206 

207 Parameters 

208 ---------- 

209 index_list : list[np.ndarray] 

210 List of index sets. 

211 

212 Returns 

213 ------- 

214 : 

215 Intersection of index sets. 

216 """ 

217 assert np.all(s.shape[1] == index_list[0].shape[1] for s in index_list) 

218 ret = index_list[0] 

219 for indices in index_list[1:]: 

220 ret = np.array( 

221 [x for x in set(tuple(x) for x in ret) & set(tuple(x) for x in indices)] 

222 ) 

223 return sort_index_array(np.array(ret)) 

224 

225 

226def set_difference(indices: np.ndarray, subtract: np.ndarray) -> np.ndarray: 

227 """Set difference of two index arrays. 

228 

229 Given two sparse index sets :math:`\\Lambda_1` and :math:`\\Lambda_2`, 

230 compute :math:`\\Lambda=\\Lambda_1\\setminus\\Lambda_2`. 

231 

232 Parameters 

233 ---------- 

234 indices : np.ndarray 

235 Index array multiindices are taken out of. 

236 subtract : np.ndarray 

237 Indices that are taken out of the original set. 

238 

239 Returns 

240 ------- 

241 : 

242 Set difference of both index arrays. 

243 """ 

244 indices = sort_index_array(indices) 

245 subtract = sort_index_array(subtract) 

246 assert indices.shape[1] == subtract.shape[1] 

247 idxs = [] 

248 for mdx in subtract: 

249 idx = np.where((indices == mdx).all(axis=1))[0] 

250 assert idx.size < 2 

251 if idx.size == 1: 

252 idxs += [idx] 

253 return np.delete(indices, np.array(idxs, dtype=int).flatten(), axis=0) 

254 

255 

256def tensor_set( 

257 shape: list[int] | tuple[int] | np.ndarray, 

258 lower: list[int] | tuple[int] | np.ndarray | None = None, 

259) -> np.ndarray: 

260 """Create a tensor index set. 

261 

262 For given upper and lower bounds 

263 :math:`0 \\leq \\ell_m < u_m \\in \\mathbb{N}_0` with 

264 :math:`m=1,\\dots,M\\in\\mathbb{N}`, the tensor index set (n-D cube) is 

265 given by 

266 

267 .. math:: 

268 \\Lambda = \\{ \\mu\\in\\mathbb{N}_0^M \\ \\vert\\ \\ell_m \\leq \\mu_m \\leq u_m \\mbox{ for } m=1,\\dots,M\\}. 

269 

270 Parameters 

271 ---------- 

272 shape : array_like 

273 Shape of the tensor, enumeration starting from 0. 

274 lower : array_like, default = None 

275 Starting values for each dimension of the tensor set. If None, all 

276 dimensions start with 0. 

277 

278 Returns 

279 ------- 

280 : 

281 Array with all possible multiindices in tensor set. 

282 

283 See Also 

284 -------- 

285 pythia.misc.cart_prod, pythia.index.lq_bound_set, pythia.index.simplex_set 

286 

287 Examples 

288 -------- 

289 Create the tensor product multiindices :math:`\\{0, 1\\}\\times\\{0, 1\\}` 

290 

291 >>> pt.index.tensor_set([2, 2]) 

292 array([[0, 0], 

293 [0, 1], 

294 [1, 0], 

295 [1, 1]]) 

296 

297 Create 3D univariate multiindices :math:`\\{0\\}\\times\\{1,\\dots, 4\\}\\times\\{0\\}` 

298 

299 >>> pt.index.tensor_set([1, 5, 1], [0, 1, 0]) 

300 array([[0, 1, 0], 

301 [0, 2, 0], 

302 [0, 3, 0], 

303 [0, 4, 0]]) 

304 

305 Create 1D indices similar to ``np.arange(1, 5, dtype=int).reshape(-1, 1)`` 

306 

307 >>> pt.index.tensor_set([5], [1]) 

308 array([[1], 

309 [2], 

310 [3], 

311 [4]]) 

312 """ 

313 shape = np.array(shape, dtype=int) 

314 if lower is None: 

315 lower = np.zeros(shape.size, dtype=int) 

316 elif isinstance(lower, (list, tuple)): 

317 lower = np.array(lower, dtype=int) 

318 assert shape.size == lower.size 

319 univariate_dims = [np.arange(low, up, dtype=int) for low, up in zip(lower, shape)] 

320 if shape.size > 1: 

321 ret = pt.misc.cart_prod(univariate_dims) 

322 else: 

323 ret = np.array(univariate_dims).T 

324 return sort_index_array(ret) 

325 

326 

327def lq_bound_set( 

328 dimensions: list[int] | tuple[int] | np.ndarray, bound: float, q: float = 1 

329) -> np.ndarray: 

330 """Create set of multiindices with bounded :math:`\\ell^q`-norm. 

331 

332 For given dimensions :math:`d \\in \\mathbb{N}^M`, bound 

333 :math:`b \\in \\mathbb{R}_{>0}` and norm factor 

334 :math:`q \\in \\mathbb{R}_{>0}`, the :math:`\\ell^q`-norm index set is 

335 given by 

336 

337 .. math:: 

338 \\Lambda = \\{ \\mu\\in [d_1]\\times\\dots\\times [d_M] \\ \\vert\\ \\Vert \\mu \\Vert_{\\ell^q} \\leq b\\}, 

339 

340 where :math:`[d_m]=\\{0, \\dots, d_m-1\\}` and 

341 

342 .. math:: 

343 \\Vert \\mu \\Vert_{\\ell^q} = \\Bigl(\\sum_{m=1}^M \\mu_m^q\\Bigr)^{\\frac{1}{q}}. 

344 

345 Parameters 

346 ---------- 

347 dimensions : list[int] | tuple[int] | np.ndarray 

348 Dimensions for each component, i.e., indices from ``0`` to 

349 ``dimension-1``. 

350 bound : float 

351 Bound for the :math:`\\ell^q`-norm. 

352 q : float, optional 

353 Norm factor. 

354 

355 Returns 

356 ------- 

357 : 

358 Array with all possible multiindices with bounded :math:`\\ell^q`-norm. 

359 

360 See Also 

361 -------- 

362 pythia.index.tensor_set, pythia.index.simplex_set 

363 

364 Examples 

365 -------- 

366 >>> pt.index.lq_bound_set([5, 5], 4, 0.5) 

367 array([[0, 0], 

368 [0, 1], 

369 [1, 0], 

370 [0, 2], 

371 [1, 1], 

372 [2, 0], 

373 [0, 3], 

374 [3, 0], 

375 [0, 4], 

376 [4, 0]]) 

377 """ 

378 assert np.all([d > 0 for d in dimensions]) and bound > 0 and q > 0 

379 all_indices = tensor_set(dimensions) 

380 rows = np.where(np.power(np.sum(all_indices**q, axis=1), 1 / q) <= bound)[0] 

381 return sort_index_array(all_indices[rows]) 

382 

383 

384def simplex_set(dimension: int, maximum: int) -> np.ndarray: 

385 """Create a simplex index set. 

386 

387 For given dimension :math:`M\\in\\mathbb{N}` and maximum 

388 :math:`d\\in\\mathbb{N}` the simplex index set is given by 

389 

390 .. math:: 

391 \\Lambda = \\{ \\mu\\in\\mathbb{N}_0^M \\ \\vert\\ \\sum_{m=1}^M \\mu_m \\leq d\\}. 

392 

393 Notes 

394 ----- 

395 Limiting the absolute value of the multiindices creates a simplex in 

396 :math:`\\mathbb{N}_0^M`, which motivates the name of the function. 

397 As an example, in two dimensions this gives us points inside a triangle 

398 limited by the axes and the line :math:`x_1 + x_2 = d`. 

399 

400 Parameters 

401 ---------- 

402 dimension : int 

403 Dimension of the multiindices. 

404 maximum : int 

405 Maximal sum value for the multiindices. 

406 

407 Returns 

408 ------- 

409 : 

410 Array with all possible multiindices in simplex set. 

411 

412 See Also 

413 -------- 

414 pythia.index.lq_bound_set, pythia.index.tensor_set 

415 

416 Examples 

417 -------- 

418 >>> pt.index.simplex(2, 2) 

419 array([[0, 0], 

420 [0, 1], 

421 [1, 0], 

422 [0, 2], 

423 [1, 1], 

424 [2, 0]]) 

425 """ 

426 assert dimension > 0 and maximum > 0 

427 return lq_bound_set([maximum + 1] * dimension, maximum, 1)