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