Coverage for larch/math/utils.py: 66%

201 statements  

« prev     ^ index     » next       coverage.py v7.6.0, created at 2024-10-16 21:04 +0000

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 InterpolatedUnivariateSpline as IUSpline 

10from scipy.interpolate import interp1d as scipy_interp1d 

11 

12from .lineshapes import gaussian, lorentzian, voigt 

13 

14 

15import scipy.constants as consts 

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

17ETOK = 1.0/KTOE 

18def etok(energy): 

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

20 if energy < 0: return 0 

21 return np.sqrt(energy*ETOK) 

22 

23def as_ndarray(obj): 

24 """ 

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

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

27 """ 

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

29 return np.array([obj]) 

30 return np.asarray(obj) 

31 

32def index_of(array, value): 

33 """ 

34 return index of array *at or below* value 

35 returns 0 if value < min(array) 

36 

37 >> ix = index_of(array, value) 

38 

39 Arguments 

40 --------- 

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

42 value (float): value to find index of 

43 

44 Returns 

45 ------- 

46 integer for index in array at or below value 

47 """ 

48 if value < min(array): 

49 return 0 

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

51 

52def index_nearest(array, value): 

53 """ 

54 return index of array *nearest* to value 

55 

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

57 

58 Arguments 

59 --------- 

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

61 value (float): value to find index of 

62 

63 Returns 

64 ------- 

65 integer for index in array nearest value 

66 

67 """ 

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

69 

70def deriv(arr): 

71 return np.gradient(as_ndarray(arr)) 

72deriv.__doc__ = np.gradient.__doc__ 

73 

74def realimag(arr): 

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

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

77 

78def complex_phase(arr): 

79 "return phase, modulo 2pi jumps" 

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

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

82 out = phase[:]*1.0 

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

84 return out 

85 

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

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

88 linear, quadratic, or cubic interpolation 

89 

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

91 

92 Arguments 

93 --------- 

94 x original x values 

95 y original y values 

96 xnew new x values for values to be interpolated to 

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

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

99 

100 Notes 

101 ----- 

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

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

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

105 

106 see also: interp 

107 

108 """ 

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

110 'copy': False, 'bounds_error': False} 

111 kwargs.update(kws) 

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

113 

114 

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

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

117 linear, quadratic, or cubic interpolation 

118 

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

120 arguments 

121 --------- 

122 x original x values 

123 y original y values 

124 xnew new x values for values to be interpolated to 

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

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

127 

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

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

130 

131 see also: interp1d 

132 """ 

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

134 

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

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

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

138 return out 

139 

140 if (len(np.where(np.diff(np.argsort(x))!=1)[0]) > 0 or 

141 len(np.where(np.diff(np.argsort(xnew))!=1)[0]) > 0): 

142 return out 

143 

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

145 if len(span) < 1: 

146 continue 

147 ncoef = 5 

148 if kind.startswith('lin'): 

149 ncoef = 2 

150 elif kind.startswith('quad'): 

151 ncoef = 3 

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

153 if kind.startswith('lin'): 

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

155 out[span] = coefs[0] 

156 if len(coefs) > 1: 

157 out[span] += coefs[1]*xnew[span] 

158 elif kind.startswith('quad'): 

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

160 out[span] = coefs[0] 

161 if len(coefs) > 1: 

162 out[span] += coefs[1]*xnew[span] 

163 if len(coefs) > 2: 

164 out[span] += coefs[2]*xnew[span]**2 

165 elif kind.startswith('cubic'): 

166 out[span] = IUSpline(x[sel], y[sel])(xnew[span]) 

167 return out 

168 

169 

170def remove_dups(arr, tiny=1.e-6): 

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

172 to be monotonically increasing. 

173 

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

175 will be increased by an amount given by tiny 

176 

177 Parameters 

178 ---------- 

179 arr : array of values expected to be monotonically increasing 

180 tiny : smallest expected absolute value of interval [1.e-6] 

181 

182 Returns 

183 ------- 

184 out : ndarray, strictly monotonically increasing array 

185 

186 Example 

187 ------- 

188 >>> x = np.array([1, 2, 3, 3, 3, 4, 5, 6, 7, 7, 8]) 

189 >>> remove_dups(x) 

190 array([1. , 2. , 3. , 3.000001, 3.000002, 4. , 

191 5. , 6. , 7. , 7.000001, 8. ]) 

192 """ 

193 try: 

194 work = np.asarray(arr) 

195 except Exception: 

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

197 

198 if work.size <= 1: 

199 return arr 

200 shape = work.shape 

201 work = work.flatten() 

202 

203 min_step = min(np.diff(work)) 

204 tval = (abs(min(work)) + abs(max(work))) /2.0 

205 if min_step > 10*tiny: 

206 return work 

207 previous_val = np.nan 

208 previous_add = 0 

209 

210 npts = len(work) 

211 add = np.zeros(npts) 

212 for i in range(1, npts): 

213 if not np.isnan(work[i-1]): 

214 previous_val = work[i-1] 

215 previous_add = add[i-1] 

216 val = work[i] 

217 if np.isnan(val) or np.isnan(previous_val): 

218 continue 

219 diff = abs(val - previous_val) 

220 if diff < tiny: 

221 add[i] = previous_add + tiny 

222 return work+add 

223 

224 

225def remove_nans(val, goodval=0.0, default=0.0, interp=False): 

226 """ 

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

228 and replace with 'goodval'. 

229 

230 """ 

231 isbad = ~np.isfinite(val) 

232 if not np.any(isbad): 

233 return val 

234 

235 if isinstance(goodval, np.ndarray): 

236 goodval = goodval.mean() 

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

238 goodval = default 

239 

240 if not isinstance(val, np.ndarray): 

241 return goodval 

242 if interp: 

243 for i in np.where(isbad)[0]: 

244 if i == 0: 

245 val[i] = 2.0*val[1] - val[2] 

246 elif i == len(val)-1: 

247 val[i] = 2.0*val[i-1] - val[i-2] 

248 else: 

249 val[i] = 0.5*(val[i+1] + val[i-1]) 

250 isbad = ~np.isfinite(val) 

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

252 return val 

253 

254 

255def remove_nans2(a, b): 

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

257 returning 2 arrays of the same length 

258 with NANs and INFs removed 

259 

260 Parameters 

261 ---------- 

262 a : array 1 

263 b : array 2 

264 

265 Returns 

266 ------- 

267 anew, bnew 

268 

269 Example 

270 ------- 

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

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

273 >>> emove_nans2(x, y) 

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

275 

276 """ 

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

278 try: 

279 a = np.array(a) 

280 except: 

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

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

283 try: 

284 b = np.array(b) 

285 except: 

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

287 

288 def fix_bad(isbad, x, y): 

289 if np.any(isbad): 

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

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

292 return x, y 

293 

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

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

296 return a, b 

297 

298 

299def safe_log(x, extreme=50): 

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

301 

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

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

304 or voigt function. 

305 

306 arguments: 

307 ------------ 

308 x input 1-d array for absicca 

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

310 sigma primary width parameter for convolving function 

311 gamma secondary width parameter for convolving function 

312 delx delta x to use for interpolation [mean of 

313 form name of convolving function: 

314 'lorentzian' or 'gaussian' or 'voigt' ['lorentzian'] 

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

316 

317 returns: 

318 -------- 

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

320 """ 

321 # make uniform x, y data 

322 TINY = 1.e-12 

323 if xstep is None: 

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

325 if xstep < TINY: 

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

327 if npad is None: 

328 npad = 5 

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

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

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

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

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

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

335 

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

337 sigma *= 1.0 / xstep 

338 if gamma is not None: 

339 gamma *= 1.0 / xstep 

340 

341 wx = np.arange(2*npts) 

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

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

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

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

346 else: 

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

348 

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

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

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

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

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

354 return interp(x0, y2, x) 

355 

356 

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

358 # 

359 # code from from scipy cookbook 

360 

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

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

363 It has the advantage of preserving the original shape and 

364 features of the signal better than other types of filtering 

365 approaches, such as moving averages techhniques. 

366 Parameters 

367 ---------- 

368 y : array_like, shape (N,) 

369 the values of the time history of the signal. 

370 window_size : int 

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

372 order : int 

373 the order of the polynomial used in the filtering. 

374 Must be less then `window_size` - 1. 

375 deriv: int 

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

377 Returns 

378 ------- 

379 ys : ndarray, shape (N) 

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

381 Notes 

382 ----- 

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

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

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

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

387 the point. 

388 Examples 

389 -------- 

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

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

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

393 import matplotlib.pyplot as plt 

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

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

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

397 plt.legend() 

398 plt.show() 

399 References 

400 ---------- 

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

402 Data by Simplified Least Squares Procedures. Analytical 

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

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

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

406 Cambridge University Press ISBN-13: 9780521880688 

407 """ 

408 try: 

409 window_size = abs(int(window_size)) 

410 order = abs(int(order)) 

411 except ValueError: 

412 raise ValueError("window_size and order must be integers") 

413 if window_size < order + 2: 

414 window_size = order + 3 

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

416 window_size = window_size + 1 

417 order_range = range(order+1) 

418 half_window = (window_size -1) // 2 

419 # precompute coefficients 

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

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

422 # pad the signal at the extremes with 

423 # values taken from the signal itself 

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

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

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

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

428 

429 

430def boxcar(data, nrepeats=1): 

431 """boxcar average of an array 

432 

433 Arguments 

434 --------- 

435 data nd-array, assumed to be 1d 

436 nrepeats integer number of repeats [1] 

437 

438 Returns 

439 ------- 

440 ndarray of same size as input data 

441 

442 Notes 

443 ----- 

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

445 

446 out = data[:]*1.0 

447 for i in range(nrepeats): 

448 qdat = out/4.0 

449 left = 1.0*qdat 

450 right = 1.0*qdat 

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

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

453 out = 2*qdat + left + right 

454 return out 

455 

456 """ 

457 out = data[:]*1.0 

458 for i in range(nrepeats): 

459 qdat = out/4.0 

460 left = 1.0*qdat 

461 right = 1.0*qdat 

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

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

464 out = 2*qdat + left + right 

465 return out 

466 

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

468 """ 

469 simple emulation of deprecated numpy.polyfit, 

470 including its ordering of coefficients 

471 """ 

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

473 coefs = pfit.convert().coef 

474 if reverse: 

475 coefs = list(reversed(coefs)) 

476 return coefs