Coverage for pythia/basis.py: 93%

85 statements  

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

1""" 

2File: pythia/basis.py 

3Author: Nando Hegemann 

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

5Description: Assemble sparse univariate and multivariate basis polynomials. 

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

7""" 

8import math 

9from typing import Callable 

10import numpy as np 

11import scipy.integrate 

12import scipy.special 

13import pythia as pt 

14 

15 

16def univariate_basis( 

17 params: list[pt.parameter.Parameter], 

18 degs: list[int] | tuple[int] | np.ndarray, 

19) -> list[list[Callable]]: 

20 """Assemble a univariate polynomial basis. 

21 

22 Set polynomial basis up to deg for each parameter in `params` according to 

23 the parameter distribution and area of definition. 

24 

25 Parameters 

26 ---------- 

27 params : list of `pythia.parameter.Parameter` 

28 Parameters to compute univariate basis function for. 

29 degs : array_like 

30 Max. degrees of univariate polynomials for each parameter. 

31 

32 Returns 

33 ------- 

34 : 

35 List of normalized univariate polynomials w.r.t. parameter domain and 

36 distribution up to specified degree for each parameter in `params`. 

37 """ 

38 basis = [] 

39 param_pdfs = [pt.sampler.assign_sampler(param).pdf for param in params] 

40 for param, pdf, deg in zip(params, param_pdfs, degs): 

41 # Set the polynomial basis with corresponding area of support and 

42 # proper normalization. 

43 if param.distribution == "uniform": 

44 polynomials = normalize_polynomial( 

45 pdf, set_legendre_basis(param, deg), param 

46 ) 

47 elif param.distribution == "normal": 

48 polynomials = normalize_polynomial( 

49 pdf, set_hermite_basis(param, deg), param 

50 ) 

51 elif param.distribution == "gamma": 

52 polynomials = normalize_polynomial( 

53 pdf, set_laguerre_basis(param, deg), param 

54 ) 

55 elif param.distribution == "beta": 

56 polynomials = normalize_polynomial(pdf, set_jacobi_basis(param, deg), param) 

57 else: 

58 raise ValueError( 

59 f'Unsupported distribution "{param.distribution}"' f" for {param.name}" 

60 ) 

61 basis += [polynomials] 

62 return basis 

63 

64 

65def multivariate_basis( 

66 univariate_bases: list[list[Callable]], 

67 indices: np.ndarray, 

68 partial: list[int] | None = None, 

69) -> list[Callable]: 

70 """Assemble multivariate polynomial basis. 

71 

72 Set the (partial derivative of the) multivariate (product) polynomial basis 

73 functions. 

74 

75 Parameters 

76 ---------- 

77 univariate_bases : list of list of callable 

78 Univariate basis functions for parameters. Is called by 

79 `univariate_bases[paramIdx][deg]()`. 

80 indices : array_like 

81 Array of multiindices for multivariate basis functions. 

82 partial : list of int 

83 Number of partial derivatives for each dimension. Length is same as 

84 `univariate_bases`. 

85 

86 Returns 

87 ------- 

88 : 

89 List of multivariate product polynomials with univariate degrees as 

90 specified in `indices`. 

91 """ 

92 assert len(univariate_bases) == indices.shape[1] 

93 if partial is not None: 

94 assert len(partial) == indices.shape[1] 

95 basis_list = [] 

96 for index in indices: 

97 

98 def fun(x: np.ndarray, index: np.ndarray | None = index) -> np.ndarray: 

99 if not 1 <= x.ndim <= 2: 

100 raise ValueError(f"Wrong ndim '{x.ndim}'") 

101 if x.ndim == 1: 

102 x.shape = 1, -1 

103 if partial is None: 

104 basis = [univariate_bases[k][mu_k] for k, mu_k in enumerate(index)] 

105 else: 

106 basis = [ 

107 univariate_bases[k][mu_k].deriv(partial[k]) 

108 for k, mu_k in enumerate(index) 

109 ] 

110 return np.prod([basis[k](x[:, k]) for k, _ in enumerate(index)], axis=0) 

111 

112 basis_list += [fun] 

113 return basis_list 

114 

115 

116def normalize_polynomial( 

117 weight: Callable, 

118 basis: list[Callable], 

119 param: pt.parameter.Parameter, 

120) -> list[Callable]: 

121 """Normalize orthogonal polynomials. 

122 

123 Normalize a polynomial of an orthogonal system with respect to the scalar 

124 product 

125 

126 .. math:: 

127 a(u,v)_\\mathrm{pdf} = \\int u(p) v(p) \\mathrm{pdf}(p) \\mathrm{d}p. 

128 

129 The normalized polynomial :math:`\\phi_j` for any given polynomial 

130 :math:`P_j` is given by :math:`\\phi_j = P_j / \\sqrt{c_j}` 

131 for the constant 

132 :math:`c_j = \\int \\mathrm{pdf}(p) * P_j(p)^2 \\mathrm{d}p`. 

133 

134 Parameters 

135 ---------- 

136 weight : callable 

137 Probability density function. 

138 basis : list of `numpy.polynomial.Polynomial` 

139 Polynomials to normalize w.r.t. weight. 

140 param : `pythia.parameter.Parameter` 

141 Parameter used for distribution and domain information. 

142 

143 Returns 

144 ------- 

145 : 

146 List of normalized univariate polynomials. 

147 """ 

148 cs = np.zeros(len(basis)) 

149 for j, p in enumerate(basis): 

150 if param.distribution == "normal": 

151 cs[j] = float(math.factorial(j)) 

152 else: 

153 

154 def integrand(x): 

155 return weight(x) * p(x) ** 2 

156 

157 cs[j], _ = scipy.integrate.quad(integrand, param.domain[0], param.domain[1]) 

158 return [p / np.sqrt(c) for c, p in zip(cs, basis)] 

159 

160 

161def set_legendre_basis(param: pt.parameter.Parameter, deg: int) -> list[Callable]: 

162 """Generate list of the Legendre Polynomials. 

163 

164 Generate the Legendre Polynomials up to certain degree on the interval 

165 specified by the parameter. 

166 

167 Parameters 

168 ---------- 

169 param : `pythia.parameters.Parameter` 

170 Parameter for basis function. Needs to be uniformly distributed. 

171 deg : int 

172 Maximum degree for polynomials. 

173 

174 Returns 

175 ------- 

176 : 

177 List of Legendre polynomials up to (including) degree specified in 

178 `deg`. 

179 """ 

180 return [ 

181 np.polynomial.legendre.Legendre([0] * j + [1], param.domain) 

182 for j in range(deg + 1) 

183 ] 

184 

185 

186def set_hermite_basis(param: pt.parameter.Parameter, deg: int) -> list[Callable]: 

187 """Generate list of probabilists Hermite polynomials. 

188 

189 Generate the Hermite Polynomials up to certain degree according to the 

190 mean and variance of the specified parameter. 

191 

192 Parameters 

193 ---------- 

194 param : `pythia.parameters.Parameter` 

195 Parameter for basis function. Needs to be normal distributed. 

196 deg : int 

197 Maximum degree for polynomials. 

198 

199 Returns 

200 ------- 

201 : 

202 List of probabilists Hermite polynomials up to (including) degree 

203 specified in `deg`. 

204 """ 

205 assert isinstance(param.mean, (int, float)) 

206 assert isinstance(param.var, (int, float)) 

207 p_list = [] 

208 std = np.sqrt(param.var) 

209 a = -param.mean / (std * np.sqrt(2)) 

210 b = 1 / (np.sqrt(2) * std) 

211 shift = np.polynomial.polynomial.Polynomial([a, b]) 

212 for j in range(deg + 1): 

213 p = np.polynomial.hermite.Hermite([0] * j + [1]) 

214 p_list.append(2 ** (-j / 2) * p(shift)) 

215 return p_list 

216 

217 

218def set_jacobi_basis(param: pt.parameter.Parameter, deg: int) -> list[Callable]: 

219 """Generate list of Jacobi polynomials. 

220 

221 Generate the Jacobi Polynomials up to certain degree on the interval 

222 and DoFs specified by the parameter. 

223 

224 .. note:: 

225 The Jacobi polynomials have leading coefficient 1. 

226 

227 Parameters 

228 ---------- 

229 param : `pythia.parameters.Parameter` 

230 Parameter for basis function. Needs to be Beta-distributed. 

231 deg : int 

232 Maximum degree for polynomials. 

233 

234 Returns 

235 ------- 

236 : 

237 List of Jacobi polynomials up to (including) degree specified in `deg`. 

238 """ 

239 assert isinstance(param.alpha, (int, float)) 

240 assert isinstance(param.beta, (int, float)) 

241 p_list = [np.polynomial.polynomial.Polynomial(1)] 

242 

243 a = pt.misc.shift_coord(0.0, [-1, 1], param.domain) 

244 b = pt.misc.shift_coord(1.0, [-1, 1], param.domain) - a 

245 shift = np.polynomial.polynomial.Polynomial([a, b]) 

246 

247 for j in range(1, deg + 1): 

248 roots, _ = scipy.special.roots_jacobi(j, param.beta - 1, param.alpha - 1) 

249 coeff = np.polynomial.polynomial.polyfromroots(shift(roots)) 

250 p = np.polynomial.polynomial.Polynomial(coeff) 

251 p_list.append(p) 

252 return p_list 

253 

254 

255def set_laguerre_basis(param: pt.parameter.Parameter, deg: int) -> list[Callable]: 

256 """Generate list of Leguerre polynomials. 

257 

258 Generate the generalized Laguerre polynomials up to certain degree on 

259 the interval and DoFs specified by the parameter. 

260 

261 Parameters 

262 ---------- 

263 param : `pythia.parameters.Parameter` 

264 Parameter for basis function. Needs to be Gamma-distributed. 

265 deg : int 

266 Maximum degree for polynomials. 

267 

268 Returns 

269 ------- 

270 : 

271 List of Laguerre polynomials up to (including) degree specified in 

272 `deg`. 

273 """ 

274 assert isinstance(param.alpha, (int, float)) 

275 assert isinstance(param.beta, (int, float)) 

276 p_list = [np.polynomial.polynomial.Polynomial(1)] 

277 shift = np.polynomial.polynomial.Polynomial([param.domain[0], 1 / param.beta]) 

278 for j in range(1, deg + 1): 

279 roots, _ = scipy.special.roots_genlaguerre(j, param.alpha - 1) 

280 coeff = np.polynomial.polynomial.polyfromroots(shift(roots)) 

281 p = np.polynomial.polynomial.Polynomial(coeff) 

282 p_list.append(p) 

283 return p_list 

284 

285 

286if __name__ == "__main__": 

287 pass