Coverage for /Users/Newville/Codes/xraylarch/larch/math/utils.py: 69%

179 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-11-09 10:08 -0600

1#!/usr/bin/env python 

2""" 

3Some common math utilities 

4""" 

5import numpy as np 

6 

7from scipy.stats import linregress 

8from scipy.interpolate import UnivariateSpline 

9from scipy.interpolate import interp1d as scipy_interp1d 

10 

11from .lineshapes import gaussian, lorentzian, voigt 

12 

13 

14import scipy.constants as consts 

15KTOE = 1.e20*consts.hbar**2 / (2*consts.m_e * consts.e) # 3.8099819442818976 

16ETOK = 1.0/KTOE 

17def etok(energy): 

18 """convert photo-electron energy to wavenumber""" 

19 if energy < 0: return 0 

20 return np.sqrt(energy*ETOK) 

21 

22def as_ndarray(obj): 

23 """ 

24 make sure a float, int, list of floats or ints, 

25 or tuple of floats or ints, acts as a numpy array 

26 """ 

27 if isinstance(obj, (float, int)): 

28 return np.array([obj]) 

29 return np.asarray(obj) 

30 

31def index_of(array, value): 

32 """ 

33 return index of array *at or below* value 

34 returns 0 if value < min(array) 

35 

36 >> ix = index_of(array, value) 

37 

38 Arguments 

39 --------- 

40 array (ndarray-like): array to find index in 

41 value (float): value to find index of 

42 

43 Returns 

44 ------- 

45 integer for index in array at or below value 

46 """ 

47 if value < min(array): 

48 return 0 

49 return max(np.where(array<=value)[0]) 

50 

51def index_nearest(array, value): 

52 """ 

53 return index of array *nearest* to value 

54 

55 >>> ix = index_nearest(array, value) 

56 

57 Arguments 

58 --------- 

59 array (ndarray-like): array to find index in 

60 value (float): value to find index of 

61 

62 Returns 

63 ------- 

64 integer for index in array nearest value 

65 

66 """ 

67 return np.abs(array-value).argmin() 

68 

69def deriv(arr): 

70 return np.gradient(as_ndarray(arr)) 

71deriv.__doc__ = np.gradient.__doc__ 

72 

73def realimag(arr): 

74 "return real array of real/imag pairs from complex array" 

75 return np.array([(i.real, i.imag) for i in arr]).flatten() 

76 

77def complex_phase(arr): 

78 "return phase, modulo 2pi jumps" 

79 phase = np.arctan2(arr.imag, arr.real) 

80 d = np.diff(phase)/np.pi 

81 out = phase[:]*1.0 

82 out[1:] -= np.pi*(np.round(abs(d))*np.sign(d)).cumsum() 

83 return out 

84 

85def interp1d(x, y, xnew, kind='linear', fill_value=np.nan, **kws): 

86 """interpolate x, y array onto new x values, using one of 

87 linear, quadratic, or cubic interpolation 

88 

89 > ynew = interp1d(x, y, xnew, kind='linear') 

90 

91 Arguments 

92 --------- 

93 x original x values 

94 y original y values 

95 xnew new x values for values to be interpolated to 

96 kind method to use: one of 'linear', 'quadratic', 'cubic' 

97 fill_value value to use to fill values for out-of-range x values 

98 

99 Notes 

100 ----- 

101 unlike interp, this version will not extrapolate for values of `xnew` 

102 that are outside the range of `x` -- it will use NaN or `fill_value`. 

103 this is a bare-bones wrapping of scipy.interpolate.interp1d. 

104 

105 see also: interp 

106 

107 """ 

108 kwargs = {'kind': kind.lower(), 'fill_value': fill_value, 

109 'copy': False, 'bounds_error': False} 

110 kwargs.update(kws) 

111 return scipy_interp1d(x, y, **kwargs)(xnew) 

112 

113 

114def interp(x, y, xnew, kind='linear', fill_value=np.nan, **kws): 

115 """interpolate x, y array onto new x values, using one of 

116 linear, quadratic, or cubic interpolation 

117 

118 > ynew = interp(x, y, xnew, kind='linear') 

119 arguments 

120 --------- 

121 x original x values 

122 y original y values 

123 xnew new x values for values to be interpolated to 

124 kind method to use: one of 'linear', 'quadratic', 'cubic' 

125 fill_value value to use to fill values for out-of-range x values 

126 

127 note: unlike interp1d, this version will extrapolate for values of `xnew` 

128 that are outside the range of `x`, using the polynomial order `kind`. 

129 

130 see also: interp1d 

131 """ 

132 out = interp1d(x, y, xnew, kind=kind, fill_value=fill_value, **kws) 

133 

134 below = np.where(xnew<x[0])[0] 

135 above = np.where(xnew>x[-1])[0] 

136 if len(above) == 0 and len(below) == 0: 

137 return out 

138 for span, isbelow in ((below, True), (above, False)): 

139 if len(span) < 1: 

140 continue 

141 ncoef = 5 

142 if kind.startswith('lin'): 

143 ncoef = 2 

144 elif kind.startswith('quad'): 

145 ncoef = 3 

146 sel = slice(None, ncoef) if isbelow else slice(-ncoef, None) 

147 if kind.startswith('lin'): 

148 coefs = polyfit(x[sel], y[sel], 1) 

149 out[span] = coefs[0] + coefs[1]*xnew[span] 

150 elif kind.startswith('quad'): 

151 coefs = polyfit(x[sel], y[sel], 2) 

152 out[span] = coefs[0] + xnew[span]*(coefs[1] + coefs[2]*xnew[span]) 

153 elif kind.startswith('cubic'): 

154 out[span] = UnivariateSpline(x[sel], y[sel], s=0)(xnew[span]) 

155 return out 

156 

157 

158def remove_dups(arr, tiny=1.e-7, frac=1.e-6): 

159 """avoid repeated successive values of an array that is expected 

160 to be monotonically increasing. 

161 

162 For repeated values, the second encountered occurance (at index i) 

163 will be increased by an amount that is the largest of: 

164 (tiny, frac*abs(arr[i]-arr[i-1])) 

165 

166 Parameters 

167 ---------- 

168 arr : array of values expected to be monotonically increasing 

169 tiny : smallest expected absolute value of interval [1.e-7] 

170 frac : smallest expected fractional interval [1.e-6] 

171 

172 Returns 

173 ------- 

174 out : ndarray, strictly monotonically increasing array 

175 

176 Example 

177 ------- 

178 >>> x = np.array([0, 1.1, 2.2, 2.2, 3.3]) 

179 >>> print(remove_dups(x)) 

180 >>> array([ 0. , 1.1 , 2.2, 2.2000001, 3.3 ]) 

181 """ 

182 try: 

183 arr = np.asarray(arr) 

184 except Exception: 

185 print('remove_dups: argument is not an array') 

186 

187 if arr.size <= 1: 

188 return arr 

189 shape = arr.shape 

190 arr = arr.flatten() 

191 previous_value = np.nan 

192 previous_add = 0 

193 

194 add = np.zeros(arr.size) 

195 for i in range(1, len(arr)): 

196 if not np.isnan(arr[i-1]): 

197 previous_value = arr[i-1] 

198 previous_add = add[i-1] 

199 value = arr[i] 

200 if np.isnan(value) or np.isnan(previous_value): 

201 continue 

202 diff = abs(value - previous_value) 

203 if diff < tiny: 

204 add[i] = previous_add + max(tiny, frac*diff) 

205 return (arr+add).reshape(shape) 

206 

207 

208def remove_nans(val, goodval=0.0, default=0.0): 

209 """ 

210 remove nan / inf from an value (array or scalar), 

211 and replace with 'goodval'. 

212 """ 

213 if isinstance(goodval, np.ndarray): 

214 goodval = goodval.mean() 

215 if np.any(~np.isfinite(goodval)): 

216 goodval = default 

217 

218 if np.any(~np.isfinite(val)): 

219 if isinstance(val, np.ndarray): 

220 isbad = np.any(~np.isfinite(val)) 

221 val[np.where(isbad)] = goodval 

222 else: 

223 val = goodval 

224 return val 

225 

226 

227def remove_nans2(a, b): 

228 """removes NAN and INF from 2 arrays, 

229 returning 2 arrays of the same length 

230 with NANs and INFs removed 

231 

232 Parameters 

233 ---------- 

234 a : array 1 

235 b : array 2 

236 

237 Returns 

238 ------- 

239 anew, bnew 

240 

241 Example 

242 ------- 

243 >>> x = array([0, 1.1, 2.2, nan, 3.3]) 

244 >>> y = array([1, 2, 3, 4, 5) 

245 >>> emove_nans2(x, y) 

246 >>> array([ 0. , 1.1, 2.2, 3.3]), array([1, 2, 3, 5]) 

247 

248 """ 

249 if not isinstance(a, np.ndarray): 

250 try: 

251 a = np.array(a) 

252 except: 

253 print( 'remove_nans2: argument 1 is not an array') 

254 if not isinstance(b, np.ndarray): 

255 try: 

256 b = np.array(b) 

257 except: 

258 print( 'remove_nans2: argument 2 is not an array') 

259 

260 def fix_bad(isbad, x, y): 

261 if np.any(isbad): 

262 bad = np.where(isbad)[0] 

263 x, y = np.delete(x, bad), np.delete(y, bad) 

264 return x, y 

265 

266 a, b = fix_bad(~np.isfinite(a), a, b) 

267 a, b = fix_bad(~np.isfinite(b), a, b) 

268 return a, b 

269 

270 

271def safe_log(x, extreme=50): 

272 return np.log(np.clip(x, np.e**-extreme, np.e**extreme)) 

273 

274def smooth(x, y, sigma=1, gamma=None, xstep=None, npad=None, form='lorentzian'): 

275 """smooth a function y(x) by convolving wih a lorentzian, gaussian, 

276 or voigt function. 

277 

278 arguments: 

279 ------------ 

280 x input 1-d array for absicca 

281 y input 1-d array for ordinate: data to be smoothed 

282 sigma primary width parameter for convolving function 

283 gamma secondary width parameter for convolving function 

284 delx delta x to use for interpolation [mean of 

285 form name of convolving function: 

286 'lorentzian' or 'gaussian' or 'voigt' ['lorentzian'] 

287 npad number of padding pixels to use [length of x] 

288 

289 returns: 

290 -------- 

291 1-d array with same length as input array y 

292 """ 

293 # make uniform x, y data 

294 TINY = 1.e-12 

295 if xstep is None: 

296 xstep = min(np.diff(x)) 

297 if xstep < TINY: 

298 raise Warning('Cannot smooth data: must be strictly increasing ') 

299 if npad is None: 

300 npad = 5 

301 xmin = xstep * int( (min(x) - npad*xstep)/xstep) 

302 xmax = xstep * int( (max(x) + npad*xstep)/xstep) 

303 npts1 = 1 + int(abs(xmax-xmin+xstep*0.1)/xstep) 

304 npts = min(npts1, 50*len(x)) 

305 x0 = np.linspace(xmin, xmax, npts) 

306 y0 = np.interp(x0, x, y) 

307 

308 # put sigma in units of 1 for convolving window function 

309 sigma *= 1.0 / xstep 

310 if gamma is not None: 

311 gamma *= 1.0 / xstep 

312 

313 wx = np.arange(2*npts) 

314 if form.lower().startswith('gauss'): 

315 win = gaussian(wx, center=npts, sigma=sigma) 

316 elif form.lower().startswith('voig'): 

317 win = voigt(wx, center=npts, sigma=sigma, gamma=gamma) 

318 else: 

319 win = lorentzian(wx, center=npts, sigma=sigma) 

320 

321 y1 = np.concatenate((y0[npts:0:-1], y0, y0[-1:-npts-1:-1])) 

322 y2 = np.convolve(win/win.sum(), y1, mode='valid') 

323 if len(y2) > len(x0): 

324 nex = int((len(y2) - len(x0))/2) 

325 y2 = (y2[nex:])[:len(x0)] 

326 return interp(x0, y2, x) 

327 

328 

329def savitzky_golay(y, window_size, order, deriv=0): 

330 # 

331 # code from from scipy cookbook 

332 

333 """Smooth (and optionally differentiate) data with a Savitzky-Golay filter. 

334 The Savitzky-Golay filter removes high frequency noise from data. 

335 It has the advantage of preserving the original shape and 

336 features of the signal better than other types of filtering 

337 approaches, such as moving averages techhniques. 

338 Parameters 

339 ---------- 

340 y : array_like, shape (N,) 

341 the values of the time history of the signal. 

342 window_size : int 

343 the length of the window. Must be an odd integer number. 

344 order : int 

345 the order of the polynomial used in the filtering. 

346 Must be less then `window_size` - 1. 

347 deriv: int 

348 the order of the derivative to compute (default = 0 means only smoothing) 

349 Returns 

350 ------- 

351 ys : ndarray, shape (N) 

352 the smoothed signal (or it's n-th derivative). 

353 Notes 

354 ----- 

355 The Savitzky-Golay is a type of low-pass filter, particularly 

356 suited for smoothing noisy data. The main idea behind this 

357 approach is to make for each point a least-square fit with a 

358 polynomial of high order over a odd-sized window centered at 

359 the point. 

360 Examples 

361 -------- 

362 t = np.linspace(-4, 4, 500) 

363 y = np.exp( -t**2 ) + np.random.normal(0, 0.05, t.shape) 

364 ysg = savitzky_golay(y, window_size=31, order=4) 

365 import matplotlib.pyplot as plt 

366 plt.plot(t, y, label='Noisy signal') 

367 plt.plot(t, np.exp(-t**2), 'k', lw=1.5, label='Original signal') 

368 plt.plot(t, ysg, 'r', label='Filtered signal') 

369 plt.legend() 

370 plt.show() 

371 References 

372 ---------- 

373 .. [1] A. Savitzky, M. J. E. Golay, Smoothing and Differentiation of 

374 Data by Simplified Least Squares Procedures. Analytical 

375 Chemistry, 1964, 36 (8), pp 1627-1639. 

376 .. [2] Numerical Recipes 3rd Edition: The Art of Scientific Computing 

377 W.H. Press, S.A. Teukolsky, W.T. Vetterling, B.P. Flannery 

378 Cambridge University Press ISBN-13: 9780521880688 

379 """ 

380 try: 

381 window_size = abs(int(window_size)) 

382 order = abs(int(order)) 

383 except ValueError: 

384 raise ValueError("window_size and order have to be of type int") 

385 if window_size % 2 != 1 or window_size < 1: 

386 raise TypeError("window_size size must be a positive odd number") 

387 if window_size < order + 2: 

388 raise TypeError("window_size is too small for the polynomials order") 

389 order_range = range(order+1) 

390 half_window = (window_size -1) // 2 

391 # precompute coefficients 

392 b = np.mat([[k**i for i in order_range] for k in range(-half_window, half_window+1)]) 

393 m = np.linalg.pinv(b).A[deriv] 

394 # pad the signal at the extremes with 

395 # values taken from the signal itself 

396 firstvals = y[0] - abs( y[1:half_window+1][::-1] - y[0] ) 

397 lastvals = y[-1] + abs(y[-half_window-1:-1][::-1] - y[-1]) 

398 y = np.concatenate((firstvals, y, lastvals)) 

399 return np.convolve( m, y, mode='valid') 

400 

401 

402def boxcar(data, nrepeats=1): 

403 """boxcar average of an array 

404 

405 Arguments 

406 --------- 

407 data nd-array, assumed to be 1d 

408 nrepeats integer number of repeats [1] 

409 

410 Returns 

411 ------- 

412 ndarray of same size as input data 

413 

414 Notes 

415 ----- 

416 This does a 3-point smoothing, that can be repeated 

417 

418 out = data[:]*1.0 

419 for i in range(nrepeats): 

420 qdat = out/4.0 

421 left = 1.0*qdat 

422 right = 1.0*qdat 

423 right[1:] = qdat[:-1] 

424 left[:-1] = qdat[1:] 

425 out = 2*qdat + left + right 

426 return out 

427 

428 """ 

429 out = data[:]*1.0 

430 for i in range(nrepeats): 

431 qdat = out/4.0 

432 left = 1.0*qdat 

433 right = 1.0*qdat 

434 right[1:] = qdat[:-1] 

435 left[:-1] = qdat[1:] 

436 out = 2*qdat + left + right 

437 return out 

438 

439def polyfit(x, y, deg=1, reverse=False): 

440 """ 

441 simple emulation of deprecated numpy.polyfit, 

442 including its ordering of coefficients 

443 """ 

444 pfit = np.polynomial.Polynomial.fit(x, y, deg=int(deg)) 

445 coefs = pfit.convert().coef 

446 if reverse: 

447 coefs = list(reversed(coefs)) 

448 return coefs