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

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 

8 

9import uncertainties 

10 

11from larch import (Group, Make_CallArgs, parse_group_args, isgroup) 

12 

13from larch.math import index_of, index_nearest, realimag, remove_dups 

14 

15from .xafsutils import ETOK, set_xafsGroup 

16from .xafsft import ftwindow, xftf_fast 

17from .pre_edge import find_e0, pre_edge 

18 

19FMT_COEF = 'coef_%2.2i' 

20 

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 

26 

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): 

30 

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:])) 

44 

45 

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 

53 

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] 

81 

82 Output arrays are written to the provided group. 

83 

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() 

102 

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) 

107 

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 

116 

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 

133 

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 

138 

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 

149 

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 

170 

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)] 

174 

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)) 

180 

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] 

184 

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)) 

189 

190 initbkg, initchi = spline_eval(kraw[:iemax-iek0+1], mu[iek0:iemax+1], 

191 knots, coefs, order, kout) 

192 

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)) 

203 

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 

210 

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 

219 

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)) 

230 

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)) 

240 

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) 

251 

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) 

259 

260 cvals[i] = cval0 - cerrs[i] 

261 bkg2, chi2 = spline_eval(_k, _m, knots, cvals, order, kout) 

262 

263 cvals[i] = cval0 

264 jac_chi[i] = (chi1 - chi2) / (2*cerrs[i]) 

265 jac_bkg[i] = (bkg1 - bkg2) / (2*cerrs[i]) 

266 

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] 

273 

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) 

277 

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