Coverage for /Users/Newville/Codes/xraylarch/larch/xafs/pre_edge.py: 72%
242 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"""
3 XAFS pre-edge subtraction, normalization algorithms
4"""
5import numpy as np
7from lmfit import Parameters, Minimizer, report_fit
8from xraydb import guess_edge
9from larch import Group, Make_CallArgs, parse_group_args
11from larch.math import (index_of, index_nearest, remove_dups, remove_nans2,
12 interp, smooth, polyfit)
13from .xafsutils import set_xafsGroup, TINY_ENERGY
15MODNAME = '_xafs'
16MAX_NNORM = 5
18@Make_CallArgs(["energy","mu"])
19def find_e0(energy, mu=None, group=None, _larch=None):
20 """calculate :math:`E_0`, the energy threshold of absorption, or
21 'edge energy', given :math:`\mu(E)`.
23 :math:`E_0` is found as the point with maximum derivative with
24 some checks to avoid spurious glitches.
26 Arguments:
27 energy (ndarray or group): array of x-ray energies, in eV, or group
28 mu (ndaarray or None): array of mu(E) values
29 group (group or None): output group
30 _larch (larch instance or None): current larch session.
32 Returns:
33 float: Value of e0. If a group is provided, group.e0 will also be set.
35 Notes:
36 1. Supports :ref:`First Argument Group` convention, requiring group members `energy` and `mu`
37 2. Supports :ref:`Set XAFS Group` convention within Larch or if `_larch` is set.
38 """
39 energy, mu, group = parse_group_args(energy, members=('energy', 'mu'),
40 defaults=(mu,), group=group,
41 fcn_name='find_e0')
42 # first find e0 without smoothing, then refine with smoothing
43 e1, ie0, estep = _finde0(energy, mu, estep=None, use_smooth=False)
44 istart = max(2, ie0-75)
45 istop = min(ie0+75, len(energy)-2)
46 e0, ix, ex = _finde0(energy[istart:istop], mu[istart:istop], estep=estep, use_smooth=True)
47 if ix < 1 :
48 e0 = energy[istart+2]
49 if group is not None:
50 group = set_xafsGroup(group, _larch=_larch)
51 group.e0 = e0
52 return e0
54def find_energy_step(energy, frac_ignore=0.01, nave=10):
55 """robustly find energy step in XAS energy array,
56 ignoring the smallest fraction of energy steps (frac_ignore),
57 and averaging over the next `nave` values
58 """
59 ediff = np.diff(energy)
60 nskip = int(frac_ignore*len(energy))
61 return ediff[np.argsort(ediff)][nskip:nskip+nave].mean()
64def _finde0(energy, mu, estep=None, use_smooth=True):
65 "internally used by find e0 "
67 en = remove_dups(energy, tiny=TINY_ENERGY)
68 if len(en.shape) > 1:
69 en = en.squeeze()
70 if len(mu.shape) > 1:
71 mu = mu.squeeze()
72 if estep is None:
73 estep = find_energy_step(en)/2.0
74 nmin = max(2, int(len(en)*0.01))
75 if use_smooth:
76 dmu = smooth(en, np.gradient(mu)/np.gradient(en), xstep=estep, sigma=3*estep)
77 else:
78 dmu = np.gradient(mu)/np.gradient(en)
79 # find points of high derivative
80 dmu[np.where(~np.isfinite(dmu))] = -1.0
81 dm_min = dmu[nmin:-nmin].min()
82 dm_ptp = max(1.e-10, dmu[nmin:-nmin].ptp())
83 dmu = (dmu - dm_min)/dm_ptp
85 dhigh = 0.60 if len(en) > 20 else 0.30
86 high_deriv_pts = np.where(dmu > dhigh)[0]
87 if len(high_deriv_pts) < 3:
88 for _ in range(2):
89 if len(high_deriv_pts) > 3:
90 break
91 dhigh *= 0.5
92 high_deriv_pts = np.where(dmu > dhigh)[0]
94 if len(high_deriv_pts) < 3:
95 high_deriv_pts = np.where(np.isfinite(dmu))[0]
97 imax, dmax = 0, 0
98 for i in high_deriv_pts:
99 if i < nmin or i > len(en) - nmin:
100 continue
101 if (dmu[i] > dmax and
102 (i+1 in high_deriv_pts) and
103 (i-1 in high_deriv_pts)):
104 imax, dmax = i, dmu[i]
105 return en[imax], imax, estep
107def flat_resid(pars, en, mu):
108 return pars['c0'] + en * (pars['c1'] + en * pars['c2']) - mu
110def preedge(energy, mu, e0=None, step=None, nnorm=None, nvict=0, pre1=None,
111 pre2=None, norm1=None, norm2=None):
112 """pre edge subtraction, normalization for XAFS (straight python)
114 This performs a number of steps:
115 1. determine E0 (if not supplied) from max of deriv(mu)
116 2. fit a line to the region below the edge
117 3. fit a polymonial to the region above the edge
118 4. extrapolate the two curves to E0 and take their difference
119 to determine the edge jump
121 Arguments
122 ----------
123 energy: array of x-ray energies, in eV
124 mu: array of mu(E)
125 e0: edge energy, in eV. If None, it will be determined here.
126 step: edge jump. If None, it will be determined here.
127 pre1: low E range (relative to E0) for pre-edge fit
128 pre2: high E range (relative to E0) for pre-edge fit
129 nvict: energy exponent to use for pre-edg fit. See Note
130 norm1: low E range (relative to E0) for post-edge fit
131 norm2: high E range (relative to E0) for post-edge fit
132 nnorm: degree of polynomial (ie, nnorm+1 coefficients will be found) for
133 post-edge normalization curve. Default=None -- see note.
134 Returns
135 -------
136 dictionary with elements (among others)
137 e0 energy origin in eV
138 edge_step edge step
139 norm normalized mu(E)
140 pre_edge determined pre-edge curve
141 post_edge determined post-edge, normalization curve
143 Notes
144 -----
145 1 pre_edge: a line is fit to mu(energy)*energy**nvict over the region,
146 energy=[e0+pre1, e0+pre2]. pre1 and pre2 default to None, which will set
147 pre1 = e0 - 2nd energy point, rounded to 5 eV
148 pre2 = roughly pre1/3.0, rounded to 5 eV
150 2 post-edge: a polynomial of order nnorm is fit to mu(energy)*energy**nvict
151 between energy=[e0+norm1, e0+norm2]. nnorm, norm1, norm2 default to None,
152 which will set:
153 nnorm = 2 in norm2-norm1>350, 1 if norm2-norm1>50, or 0 if less.
154 norm2 = max energy - e0, rounded to 5 eV
155 norm1 = roughly min(150, norm2/3.0), rounded to 5 eV
156 """
158 energy = remove_dups(energy, tiny=TINY_ENERGY)
159 if energy.size <= 1:
160 raise ValueError("energy array must have at least 2 points")
161 if e0 is None or e0 < energy[1] or e0 > energy[-2]:
162 e0 = find_e0(energy, mu)
163 ie0 = index_nearest(energy, e0)
164 e0 = energy[ie0]
166 if pre1 is None:
167 # skip first energy point, often bad
168 if ie0 > 20:
169 pre1 = 5.0*round((energy[1] - e0)/5.0)
170 else:
171 pre1 = 2.0*round((energy[1] - e0)/2.0)
173 pre1 = max(pre1, (min(energy) - e0))
174 if pre2 is None:
175 pre2 = 5.0*round(pre1/15.0)
176 if pre1 > pre2:
177 pre1, pre2 = pre2, pre1
179 if norm2 is None:
180 norm2 = 5.0*round((max(energy) - e0)/5.0)
181 if norm2 < 0:
182 norm2 = max(energy) - e0 - norm2
183 norm2 = min(norm2, (max(energy) - e0))
184 if norm1 is None:
185 norm1 = min(25, 5.0*round(norm2/15.0))
187 if norm1 > norm2+5:
188 norm1, norm2 = norm2, norm1
190 norm1 = min(norm1, norm2 - 10)
192 if nnorm is None:
193 nnorm = 2
194 if norm2-norm1 < 350: nnorm = 1
195 if norm2-norm1 < 50: nnorm = 0
196 nnorm = max(min(nnorm, MAX_NNORM), 0)
197 # preedge
198 p1 = index_of(energy, pre1+e0)
199 p2 = index_nearest(energy, pre2+e0)
200 if p2-p1 < 2:
201 p2 = min(len(energy), p1 + 2)
203 omu = mu*energy**nvict
204 ex, mx = remove_nans2(energy[p1:p2], omu[p1:p2])
205 precoefs = polyfit(ex, mx, 1)
206 pre_edge = (precoefs[0] + energy*precoefs[1]) * energy**(-nvict)
207 # normalization
208 p1 = index_of(energy, norm1+e0)
209 p2 = index_nearest(energy, norm2+e0)
210 if p2-p1 < 2:
211 p2 = min(len(energy), p1 + 2)
212 if p2-p1 < 2:
213 p1 = p1-2
215 presub = (mu-pre_edge)[p1:p2]
216 coefs = polyfit(energy[p1:p2], presub, nnorm)
217 post_edge = 1.0*pre_edge
218 norm_coefs = []
219 for n, c in enumerate(coefs):
220 post_edge += c * energy**(n)
221 norm_coefs.append(c)
222 edge_step = step
223 if edge_step is None:
224 edge_step = post_edge[ie0] - pre_edge[ie0]
225 edge_step = max(1.e-12, abs(float(edge_step)))
226 norm = (mu - pre_edge)/edge_step
227 return {'e0': e0, 'edge_step': edge_step, 'norm': norm,
228 'pre_edge': pre_edge, 'post_edge': post_edge,
229 'norm_coefs': norm_coefs, 'nvict': nvict,
230 'nnorm': nnorm, 'norm1': norm1, 'norm2': norm2,
231 'pre1': pre1, 'pre2': pre2, 'precoefs': precoefs}
233@Make_CallArgs(["energy","mu"])
234def pre_edge(energy, mu=None, group=None, e0=None, step=None, nnorm=None,
235 nvict=0, pre1=None, pre2=None, norm1=None, norm2=None,
236 make_flat=True, _larch=None):
237 """pre edge subtraction, normalization for XAFS
239 This performs a number of steps:
240 1. determine E0 (if not supplied) from max of deriv(mu)
241 2. fit a line of polymonial to the region below the edge
242 3. fit a polymonial to the region above the edge
243 4. extrapolate the two curves to E0 and take their difference
244 to determine the edge jump
246 Arguments
247 ----------
248 energy: array of x-ray energies, in eV, or group (see note 1)
249 mu: array of mu(E)
250 group: output group
251 e0: edge energy, in eV. If None, it will be determined here.
252 step: edge jump. If None, it will be determined here.
253 pre1: low E range (relative to E0) for pre-edge fit
254 pre2: high E range (relative to E0) for pre-edge fit
255 nvict: energy exponent to use for pre-edg fit. See Notes.
256 norm1: low E range (relative to E0) for post-edge fit
257 norm2: high E range (relative to E0) for post-edge fit
258 nnorm: degree of polynomial (ie, nnorm+1 coefficients will be found) for
259 post-edge normalization curve. See Notes.
260 make_flat: boolean (Default True) to calculate flattened output.
262 Returns
263 -------
264 None: The following attributes will be written to the output group:
265 e0 energy origin
266 edge_step edge step
267 norm normalized mu(E), using polynomial
268 norm_area normalized mu(E), using integrated area
269 flat flattened, normalized mu(E)
270 pre_edge determined pre-edge curve
271 post_edge determined post-edge, normalization curve
272 dmude derivative of normalized mu(E)
273 d2mude second derivative of normalized mu(E)
275 (if the output group is None, _sys.xafsGroup will be written to)
277 Notes
278 -----
279 1. Supports `First Argument Group` convention, requiring group members `energy` and `mu`.
280 2. Support `Set XAFS Group` convention within Larch or if `_larch` is set.
281 3. pre_edge: a line is fit to mu(energy)*energy**nvict over the region,
282 energy=[e0+pre1, e0+pre2]. pre1 and pre2 default to None, which will set
283 pre1 = e0 - 2nd energy point, rounded to 5 eV
284 pre2 = roughly pre1/3.0, rounded to 5 eV
285 4. post-edge: a polynomial of order nnorm is fit to mu(energy)*energy**nvict
286 between energy=[e0+norm1, e0+norm2]. nnorm, norm1, norm2 default to None,
287 which will set:
288 norm2 = max energy - e0, rounded to 5 eV
289 norm1 = roughly min(150, norm2/3.0), rounded to 5 eV
290 nnorm = 2 in norm2-norm1>350, 1 if norm2-norm1>50, or 0 if less.
291 5. flattening fits a quadratic curve (no matter nnorm) to the post-edge
292 normalized mu(E) and subtracts that curve from it.
293 """
294 energy, mu, group = parse_group_args(energy, members=('energy', 'mu'),
295 defaults=(mu,), group=group,
296 fcn_name='pre_edge')
297 if len(energy.shape) > 1:
298 energy = energy.squeeze()
299 if len(mu.shape) > 1:
300 mu = mu.squeeze()
302 energy, mu = remove_nans2(energy, mu)
303 if group is not None and e0 is None:
304 e0 = getattr(group, 'e0', None)
305 pre_dat = preedge(energy, mu, e0=e0, step=step, nnorm=nnorm,
306 nvict=nvict, pre1=pre1, pre2=pre2, norm1=norm1,
307 norm2=norm2)
308 group = set_xafsGroup(group, _larch=_larch)
310 e0 = pre_dat['e0']
311 norm = pre_dat['norm']
312 norm1 = pre_dat['norm1']
313 norm2 = pre_dat['norm2']
314 # generate flattened spectra, by fitting a quadratic to .norm
315 # and removing that.
317 ie0 = index_nearest(energy, e0)
318 p1 = index_of(energy, norm1+e0)
319 p2 = index_nearest(energy, norm2+e0)
320 if p2-p1 < 2:
321 p2 = min(len(energy), p1 + 2)
323 group.e0 = e0
324 group.norm = norm
325 group.flat = 1.0*norm
326 group.norm_poly = 1.0*norm
328 if make_flat:
329 pre_edge = pre_dat['pre_edge']
330 post_edge = pre_dat['post_edge']
331 edge_step = pre_dat['edge_step']
332 flat_residue = (post_edge - pre_edge)/edge_step
333 flat = norm - flat_residue + flat_residue[ie0]
334 flat[:ie0] = norm[:ie0]
335 group.flat = flat
337 enx, mux = remove_nans2(energy[p1:p2], norm[p1:p2])
338 # enx, mux = (energy[p1:p2], norm[p1:p2])
339 fpars = Parameters()
340 ncoefs = len(pre_dat['norm_coefs'])
341 fpars.add('c0', value=1.0, vary=True)
342 fpars.add('c1', value=0.0, vary=False)
343 fpars.add('c2', value=0.0, vary=False)
344 if ncoefs > 1:
345 fpars['c1'].set(value=1.e-5, vary=True)
346 if ncoefs > 2:
347 fpars['c2'].set(value=1.e-5, vary=True)
349 try:
350 fit = Minimizer(flat_resid, fpars, fcn_args=(enx, mux))
351 result = fit.leastsq()
352 fc0 = result.params['c0'].value
353 fc1 = result.params['c1'].value
354 fc2 = result.params['c2'].value
356 flat_diff = fc0 + energy * (fc1 + energy * fc2)
357 flat_alt = norm - flat_diff + flat_diff[ie0]
358 flat_alt[:ie0] = norm[:ie0]
359 group.flat_coefs = (fc0, fc1, fc2)
360 group.flat_alt = flat_alt
361 except:
362 pass
364 group.dmude = np.gradient(norm)/np.gradient(energy)
365 group.d2mude = np.gradient(group.dmude)/np.gradient(energy)
366 group.edge_step = pre_dat['edge_step']
367 group.edge_step_poly = pre_dat['edge_step']
368 group.pre_edge = pre_dat['pre_edge']
369 group.post_edge = pre_dat['post_edge']
371 group.pre_edge_details = Group()
372 for attr in ('pre1', 'pre2', 'norm1', 'norm2', 'nnorm', 'nvict'):
373 setattr(group.pre_edge_details, attr, pre_dat.get(attr, None))
375 group.pre_edge_details.pre_slope = pre_dat['precoefs'][0]
376 group.pre_edge_details.pre_offset = pre_dat['precoefs'][1]
378 for i in range(MAX_NNORM):
379 if hasattr(group, 'norm_c%i' % i):
380 delattr(group, 'norm_c%i' % i)
381 for i, c in enumerate(pre_dat['norm_coefs']):
382 setattr(group.pre_edge_details, 'norm_c%i' % i, c)
384 # guess element and edge
385 group.atsym = getattr(group, 'atsym', None)
386 group.edge = getattr(group, 'edge', None)
388 if group.atsym is None or group.edge is None:
389 _atsym, _edge = guess_edge(group.e0)
390 if group.atsym is None: group.atsym = _atsym
391 if group.edge is None: group.edge = _edge
392 return
394def energy_align(group, reference, array='dmude', emin=-15, emax=35):
395 """
396 align XAFS data group to a reference group
398 Arguments
399 ---------
400 group Larch group for spectrum to be aligned (see Note 1)
401 reference Larch group for reference spectrum (see Note 1)
402 array string of 'dmude', 'norm', or 'mu' (see Note 2) ['dmude']
403 emin float, min energy relative to e0 of reference for alignment [-15]
404 emax float, max energy relative to e0 of reference for alignment [+35]
406 Returns
407 -------
408 eshift energy shift to add to group.energy to match reference.
409 This value will also be written to group.eshift
411 Notes
412 -----
413 1. Both group and reference must be XAFS data, with arrays of 'energy' and 'mu'.
414 The reference group must already have an e0 value set.
416 2. The alignment can be done with 'mu' or 'dmude'. If it does not exist, the
417 dmude array will be built for group and reference.
419 """
420 if not (hasattr(group, 'energy') and hasattr(group, 'mu')):
421 raise ValueError("group must have attributes 'energy' and 'mu'")
423 if not hasattr(group, 'dmude'):
424 mu = getattr(group, 'norm', getattr(group, 'mu'))
425 en = getattr(group, 'energy')
426 group.dmude = gradient(mu)/gradient(en)
429 if not (hasattr(reference, 'energy') and hasattr(reference, 'mu')
430 and hasattr(reference, 'e0') ):
431 raise ValueError("reference must have attributes 'energy', 'mu', and 'e0'")
433 if not hasattr(reference, 'dmude'):
434 mu = getattr(reference, 'norm', getattr(reference, 'mu'))
435 en = getattr(reference, 'energy')
436 reference.dmude = gradient(mu)/gradient(en)
438 xdat = group.energy[:]*1.0
439 xref = reference.energy[:]*1.0
440 ydat = group.dmude[:]*1.0
441 yref = reference.dmude[:]*1.0
442 if array == 'mu':
443 ydat = group.mu[:]*1.0
444 yref = reference.mu[:]*1.0
445 elif array == 'norm':
446 ydat = group.norm[:]*1.0
447 yref = reference.norm[:]*1.0
448 xdat, ydat = remove_nans2(xdat, ydat)
449 xref, yref = remove_nans2(xref, yref)
451 i1 = index_of(xref, reference.e0-emin)
452 i2 = index_of(xref, reference.e0+emax)
454 def align_resid(params, xdat, ydat, xref, yref, i1, i2):
455 "fit residual"
456 newx = xdat + params['eshift'].value
457 scale = params['scale'].value
458 ytmp = interp(newx, ydat, xref, kind='cubic')
459 return (ytmp*scale - yref)[i1:i2]
461 params = Parameters()
462 params.add('eshift', value=0, min=-50, max=50)
463 params.add('scale', value=1, min=0, max=50)
465 try:
466 fit = Minimizer(align_resid, params,
467 fcn_args=(xdat, ydat, xref, yref, i1, i2))
468 result = fit.leastsq()
469 eshift = result.params['eshift'].value
470 except:
471 eshift = 0
473 group.eshift = eshift
474 return eshift