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
« 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
7from scipy.stats import linregress
8from scipy.interpolate import UnivariateSpline
9from scipy.interpolate import interp1d as scipy_interp1d
11from .lineshapes import gaussian, lorentzian, voigt
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)
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)
31def index_of(array, value):
32 """
33 return index of array *at or below* value
34 returns 0 if value < min(array)
36 >> ix = index_of(array, value)
38 Arguments
39 ---------
40 array (ndarray-like): array to find index in
41 value (float): value to find index of
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])
51def index_nearest(array, value):
52 """
53 return index of array *nearest* to value
55 >>> ix = index_nearest(array, value)
57 Arguments
58 ---------
59 array (ndarray-like): array to find index in
60 value (float): value to find index of
62 Returns
63 -------
64 integer for index in array nearest value
66 """
67 return np.abs(array-value).argmin()
69def deriv(arr):
70 return np.gradient(as_ndarray(arr))
71deriv.__doc__ = np.gradient.__doc__
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()
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
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
89 > ynew = interp1d(x, y, xnew, kind='linear')
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
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.
105 see also: interp
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)
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
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
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`.
130 see also: interp1d
131 """
132 out = interp1d(x, y, xnew, kind=kind, fill_value=fill_value, **kws)
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
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.
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]))
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]
172 Returns
173 -------
174 out : ndarray, strictly monotonically increasing array
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')
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
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)
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
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
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
232 Parameters
233 ----------
234 a : array 1
235 b : array 2
237 Returns
238 -------
239 anew, bnew
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])
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')
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
266 a, b = fix_bad(~np.isfinite(a), a, b)
267 a, b = fix_bad(~np.isfinite(b), a, b)
268 return a, b
271def safe_log(x, extreme=50):
272 return np.log(np.clip(x, np.e**-extreme, np.e**extreme))
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.
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]
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)
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
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)
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)
329def savitzky_golay(y, window_size, order, deriv=0):
330 #
331 # code from from scipy cookbook
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')
402def boxcar(data, nrepeats=1):
403 """boxcar average of an array
405 Arguments
406 ---------
407 data nd-array, assumed to be 1d
408 nrepeats integer number of repeats [1]
410 Returns
411 -------
412 ndarray of same size as input data
414 Notes
415 -----
416 This does a 3-point smoothing, that can be repeated
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
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
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