Coverage for /Users/Newville/Codes/xraylarch/larch/xafs/autobk.py: 91%
168 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
2import sys
3import numpy as np
4from scipy.interpolate import splrep, splev, UnivariateSpline
5from scipy.stats import t
6from scipy.special import erf
7from lmfit import Parameter, Parameters, minimize, fit_report
9import uncertainties
11from larch import (Group, Make_CallArgs, parse_group_args, isgroup)
13from larch.math import index_of, index_nearest, realimag, remove_dups
15from .xafsutils import ETOK, set_xafsGroup
16from .xafsft import ftwindow, xftf_fast
17from .pre_edge import find_e0, pre_edge
19FMT_COEF = 'coef_%2.2i'
21def spline_eval(kraw, mu, knots, coefs, order, kout):
22 """eval bkg(kraw) and chi(k) for knots, coefs, order"""
23 bkg = splev(kraw, [knots, coefs, order])
24 chi = UnivariateSpline(kraw, (mu-bkg), s=0)(kout)
25 return bkg, chi
27def __resid(pars, ncoefs=1, knots=None, order=3, irbkg=1, nfft=2048,
28 kraw=None, mu=None, kout=None, ftwin=1, kweight=1, chi_std=None,
29 nclamp=0, clamp_lo=1, clamp_hi=1, **kws):
31 # coefs = [getattr(pars, FMT_COEF % i) for i in range(ncoefs)]
32 coefs = [pars[FMT_COEF % i].value for i in range(ncoefs)]
33 bkg, chi = spline_eval(kraw, mu, knots, coefs, order, kout)
34 if chi_std is not None:
35 chi = chi - chi_std
36 out = realimag(xftf_fast(chi*ftwin, nfft=nfft)[:irbkg])
37 if nclamp == 0:
38 return out
39 # spline clamps:
40 scale = 1.0 + 100*(out*out).mean()
41 return np.concatenate((out,
42 abs(clamp_lo)*scale*chi[:nclamp],
43 abs(clamp_hi)*scale*chi[-nclamp:]))
46@Make_CallArgs(["energy" ,"mu"])
47def autobk(energy, mu=None, group=None, rbkg=1, nknots=None, e0=None, ek0=None,
48 edge_step=None, kmin=0, kmax=None, kweight=1, dk=0.1,
49 win='hanning', k_std=None, chi_std=None, nfft=2048, kstep=0.05,
50 pre_edge_kws=None, nclamp=3, clamp_lo=0, clamp_hi=1,
51 calc_uncertainties=True, err_sigma=1, _larch=None, **kws):
52 """Use Autobk algorithm to remove XAFS background
54 Parameters:
55 -----------
56 energy: 1-d array of x-ray energies, in eV, or group
57 mu: 1-d array of mu(E)
58 group: output group (and input group for e0 and edge_step).
59 rbkg: distance (in Ang) for chi(R) above
60 which the signal is ignored. Default = 1.
61 e0: edge energy, in eV. (deprecated: use ek0)
62 ek0: edge energy, in eV. If None, it will be determined.
63 edge_step: edge step. If None, it will be determined.
64 pre_edge_kws: keyword arguments to pass to pre_edge()
65 nknots: number of knots in spline. If None, it will be determined.
66 kmin: minimum k value [0]
67 kmax: maximum k value [full data range].
68 kweight: k weight for FFT. [1]
69 dk: FFT window window parameter. [0.1]
70 win: FFT window function name. ['hanning']
71 nfft: array size to use for FFT [2048]
72 kstep: k step size to use for FFT [0.05]
73 k_std: optional k array for standard chi(k).
74 chi_std: optional chi array for standard chi(k).
75 nclamp: number of energy end-points for clamp [3]
76 clamp_lo: weight of low-energy clamp [0]
77 clamp_hi: weight of high-energy clamp [1]
78 calc_uncertaintites: Flag to calculate uncertainties in
79 mu_0(E) and chi(k) [True]
80 err_sigma: sigma level for uncertainties in mu_0(E) and chi(k) [1]
82 Output arrays are written to the provided group.
84 Follows the 'First Argument Group' convention.
85 """
86 msg = sys.stdout
87 if _larch is not None:
88 msg = _larch.writer.write
89 if 'kw' in kws:
90 kweight = kws.pop('kw')
91 if len(kws) > 0:
92 msg('Unrecognized a:rguments for autobk():\n')
93 msg(' %s\n' % (', '.join(kws.keys())))
94 return
95 energy, mu, group = parse_group_args(energy, members=('energy', 'mu'),
96 defaults=(mu,), group=group,
97 fcn_name='autobk')
98 if len(energy.shape) > 1:
99 energy = energy.squeeze()
100 if len(mu.shape) > 1:
101 mu = mu.squeeze()
103 energy = remove_dups(energy)
104 # if e0 or edge_step are not specified, get them, either from the
105 # passed-in group or from running pre_edge()
106 group = set_xafsGroup(group, _larch=_larch)
108 if edge_step is None and isgroup(group, 'edge_step'):
109 edge_step = group.edge_step
110 if e0 is not None and ek0 is None: # command-line e0 still valid
111 ek0 = e0
112 if ek0 is None and isgroup(group, 'ek0'):
113 ek0 = group.ek0
114 if ek0 is None and isgroup(group, 'e0'):
115 ek0 = group.e0
117 if ek0 is not None and (ek0 < energy.min() or ek0 > energy.max()):
118 ek0 = None
119 if ek0 is None or edge_step is None:
120 # need to run pre_edge:
121 pre_kws = dict(nnorm=None, nvict=0, pre1=None,
122 pre2=None, norm1=None, norm2=None)
123 if pre_edge_kws is not None:
124 pre_kws.update(pre_edge_kws)
125 pre_edge(energy, mu, group=group, _larch=_larch, **pre_kws)
126 if ek0 is None:
127 ek0 = group.e0
128 if edge_step is None:
129 edge_step = group.edge_step
130 if ek0 is None or edge_step is None:
131 msg('autobk() could not determine ek0 or edge_step!: trying running pre_edge first\n')
132 return
134 # get array indices for rkbg and ek0: irbkg, iek0
135 iek0 = index_of(energy, ek0)
136 rgrid = np.pi/(kstep*nfft)
137 if rbkg < 2*rgrid: rbkg = 2*rgrid
139 # save ungridded k (kraw) and grided k (kout)
140 # and ftwin (*k-weighting) for FT in residual
141 enpe = energy[iek0:] - ek0
142 kraw = np.sign(enpe)*np.sqrt(ETOK*abs(enpe))
143 if kmax is None:
144 kmax = max(kraw)
145 else:
146 kmax = max(0, min(max(kraw), kmax))
147 kout = kstep * np.arange(int(1.01+kmax/kstep), dtype='float64')
148 iemax = min(len(energy), 2+index_of(energy, ek0+kmax*kmax/ETOK)) - 1
150 # interpolate provided chi(k) onto the kout grid
151 if chi_std is not None and k_std is not None:
152 chi_std = np.interp(kout, k_std, chi_std)
153 # pre-load FT window
154 ftwin = kout**kweight * ftwindow(kout, xmin=kmin, xmax=kmax,
155 window=win, dx=dk, dx2=dk)
156 # calc k-value and initial guess for y-values of spline params
157 nspl = 1 + int(2*rbkg*(kmax-kmin)/np.pi)
158 irbkg = int(1 + (nspl-1)*np.pi/(2*rgrid*(kmax-kmin)))
159 if nknots is not None:
160 nspl = nknots
161 nspl = max(5, min(128, nspl))
162 spl_y, spl_k = np.ones(nspl), np.zeros(nspl)
163 for i in range(nspl):
164 q = kmin + i*(kmax-kmin)/(nspl - 1)
165 ik = index_nearest(kraw, q)
166 i1 = min(len(kraw)-1, ik + 5)
167 i2 = max(0, ik - 5)
168 spl_k[i] = kraw[ik]
169 spl_y[i] = (2*mu[ik+iek0] + mu[i1+iek0] + mu[i2+iek0] ) / 4.0
171 order = 3
172 qmin, qmax = spl_k[0], spl_k[nspl-1]
173 knots = [spl_k[0] - 1.e-4*(order-i) for i in range(order)]
175 for i in range(order, nspl):
176 knots.append((i-order)*(qmax - qmin)/(nspl-order+1))
177 qlast = knots[-1]
178 for i in range(order+1):
179 knots.append(qlast + 1.e-4*(i+1))
181 # coefs = [mu[index_nearest(energy, ek0 + q**2/ETOK)] for q in knots]
182 knots, coefs, order = splrep(spl_k, spl_y, k=order)
183 coefs[nspl:] = coefs[nspl-1]
185 # set fit parameters from initial coefficients
186 params = Parameters()
187 for i in range(len(coefs)):
188 params.add(name = FMT_COEF % i, value=coefs[i], vary=i<len(spl_y))
190 initbkg, initchi = spline_eval(kraw[:iemax-iek0+1], mu[iek0:iemax+1],
191 knots, coefs, order, kout)
193 # do fit
194 result = minimize(__resid, params, method='leastsq',
195 gtol=1.e-6, ftol=1.e-6, xtol=1.e-6, epsfcn=1.e-6,
196 kws = dict(ncoefs=len(coefs), chi_std=chi_std,
197 knots=knots, order=order,
198 kraw=kraw[:iemax-iek0+1],
199 mu=mu[iek0:iemax+1], irbkg=irbkg, kout=kout,
200 ftwin=ftwin, kweight=kweight,
201 nfft=nfft, nclamp=nclamp,
202 clamp_lo=clamp_lo, clamp_hi=clamp_hi))
204 # write final results
205 coefs = [result.params[FMT_COEF % i].value for i in range(len(coefs))]
206 bkg, chi = spline_eval(kraw[:iemax-iek0+1], mu[iek0:iemax+1],
207 knots, coefs, order, kout)
208 obkg = mu[:]*1.0
209 obkg[iek0:iek0+len(bkg)] = bkg
211 # outputs to group
212 group = set_xafsGroup(group, _larch=_larch)
213 group.bkg = obkg
214 group.chie = (mu-obkg)/edge_step
215 group.k = kout
216 group.chi = chi/edge_step
217 group.ek0 = ek0
218 group.rbkg = rbkg
220 # now fill in 'autobk_details' group
221 details = Group(kmin=kmin, kmax=kmax, irbkg=irbkg, nknots=len(spl_k),
222 knots_k=knots, init_knots_y=spl_y, nspl=nspl,
223 init_chi=initchi/edge_step, report=fit_report(result))
224 details.init_bkg = mu[:]*1.0
225 details.init_bkg[iek0:iek0+len(bkg)] = initbkg
226 details.knots_y = np.array([coefs[i] for i in range(nspl)])
227 group.autobk_details = details
228 for attr in ('nfev', 'redchi', 'chisqr', 'aic', 'bic', 'params'):
229 setattr(details, attr, getattr(result, attr, None))
231 # uncertainties in mu0 and chi: can be fairly slow.
232 covar = getattr(result, 'covar', None)
233 if calc_uncertainties and covar is not None:
234 nchi = len(chi)
235 nmue = iemax-iek0 + 1
236 redchi = result.redchi
237 covar = result.covar / redchi
238 jac_chi = np.zeros(nchi*nspl).reshape((nspl, nchi))
239 jac_bkg = np.zeros(nmue*nspl).reshape((nspl, nmue))
241 cvals, cerrs = [], []
242 for i in range(len(coefs)):
243 par = result.params[FMT_COEF % i]
244 cvals.append(getattr(par, 'value', 0.0))
245 cdel = getattr(par, 'stderr', 0.0)
246 if cdel is None:
247 cdel = 0.0
248 cerrs.append(cdel/2.0)
249 cvals = np.array(cvals)
250 cerrs = np.array(cerrs)
252 # find derivatives by hand!
253 _k = kraw[:nmue]
254 _m = mu[iek0:iemax+1]
255 for i in range(nspl):
256 cval0 = cvals[i]
257 cvals[i] = cval0 + cerrs[i]
258 bkg1, chi1 = spline_eval(_k, _m, knots, cvals, order, kout)
260 cvals[i] = cval0 - cerrs[i]
261 bkg2, chi2 = spline_eval(_k, _m, knots, cvals, order, kout)
263 cvals[i] = cval0
264 jac_chi[i] = (chi1 - chi2) / (2*cerrs[i])
265 jac_bkg[i] = (bkg1 - bkg2) / (2*cerrs[i])
267 dfchi = np.zeros(nchi)
268 dfbkg = np.zeros(nmue)
269 for i in range(nspl):
270 for j in range(nspl):
271 dfchi += jac_chi[i]*jac_chi[j]*covar[i,j]
272 dfbkg += jac_bkg[i]*jac_bkg[j]*covar[i,j]
274 prob = 0.5*(1.0 + erf(err_sigma/np.sqrt(2.0)))
275 dchi = t.ppf(prob, nchi-nspl) * np.sqrt(dfchi*redchi)
276 dbkg = t.ppf(prob, nmue-nspl) * np.sqrt(dfbkg*redchi)
278 if not any(np.isnan(dchi)):
279 group.delta_chi = dchi
280 group.delta_bkg = 0.0*mu
281 group.delta_bkg[iek0:iek0+len(dbkg)] = dbkg