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
« 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
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
12from .lineshapes import gaussian, lorentzian, voigt
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)
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)
32def index_of(array, value):
33 """
34 return index of array *at or below* value
35 returns 0 if value < min(array)
37 >> ix = index_of(array, value)
39 Arguments
40 ---------
41 array (ndarray-like): array to find index in
42 value (float): value to find index of
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])
52def index_nearest(array, value):
53 """
54 return index of array *nearest* to value
56 >>> ix = index_nearest(array, value)
58 Arguments
59 ---------
60 array (ndarray-like): array to find index in
61 value (float): value to find index of
63 Returns
64 -------
65 integer for index in array nearest value
67 """
68 return np.abs(array-value).argmin()
70def deriv(arr):
71 return np.gradient(as_ndarray(arr))
72deriv.__doc__ = np.gradient.__doc__
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()
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
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
90 > ynew = interp1d(x, y, xnew, kind='linear')
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
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.
106 see also: interp
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)
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
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
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`.
131 see also: interp1d
132 """
133 out = interp1d(x, y, xnew, kind=kind, fill_value=fill_value, **kws)
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
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
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
170def remove_dups(arr, tiny=1.e-6):
171 """avoid repeated successive values of an array that is expected
172 to be monotonically increasing.
174 For repeated values, the second encountered occurance (at index i)
175 will be increased by an amount given by tiny
177 Parameters
178 ----------
179 arr : array of values expected to be monotonically increasing
180 tiny : smallest expected absolute value of interval [1.e-6]
182 Returns
183 -------
184 out : ndarray, strictly monotonically increasing array
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')
198 if work.size <= 1:
199 return arr
200 shape = work.shape
201 work = work.flatten()
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
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
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'.
230 """
231 isbad = ~np.isfinite(val)
232 if not np.any(isbad):
233 return val
235 if isinstance(goodval, np.ndarray):
236 goodval = goodval.mean()
237 if np.any(~np.isfinite(goodval)):
238 goodval = default
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
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
260 Parameters
261 ----------
262 a : array 1
263 b : array 2
265 Returns
266 -------
267 anew, bnew
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])
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')
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
294 a, b = fix_bad(~np.isfinite(a), a, b)
295 a, b = fix_bad(~np.isfinite(b), a, b)
296 return a, b
299def safe_log(x, extreme=50):
300 return np.log(np.clip(x, np.e**-extreme, np.e**extreme))
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.
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]
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)
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
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)
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)
357def savitzky_golay(y, window_size, order, deriv=0):
358 #
359 # code from from scipy cookbook
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')
430def boxcar(data, nrepeats=1):
431 """boxcar average of an array
433 Arguments
434 ---------
435 data nd-array, assumed to be 1d
436 nrepeats integer number of repeats [1]
438 Returns
439 -------
440 ndarray of same size as input data
442 Notes
443 -----
444 This does a 3-point smoothing, that can be repeated
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
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
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