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

1#!/usr/bin/env python 

2""" 

3 XAFS pre-edge subtraction, normalization algorithms 

4""" 

5import numpy as np 

6 

7from lmfit import Parameters, Minimizer, report_fit 

8from xraydb import guess_edge 

9from larch import Group, Make_CallArgs, parse_group_args 

10 

11from larch.math import (index_of, index_nearest, remove_dups, remove_nans2, 

12 interp, smooth, polyfit) 

13from .xafsutils import set_xafsGroup, TINY_ENERGY 

14 

15MODNAME = '_xafs' 

16MAX_NNORM = 5 

17 

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

22 

23 :math:`E_0` is found as the point with maximum derivative with 

24 some checks to avoid spurious glitches. 

25 

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. 

31 

32 Returns: 

33 float: Value of e0. If a group is provided, group.e0 will also be set. 

34 

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 

53 

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

62 

63 

64def _finde0(energy, mu, estep=None, use_smooth=True): 

65 "internally used by find e0 " 

66 

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 

84 

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] 

93 

94 if len(high_deriv_pts) < 3: 

95 high_deriv_pts = np.where(np.isfinite(dmu))[0] 

96 

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 

106 

107def flat_resid(pars, en, mu): 

108 return pars['c0'] + en * (pars['c1'] + en * pars['c2']) - mu 

109 

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) 

113 

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 

120 

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 

142 

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 

149 

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

157 

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] 

165 

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) 

172 

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 

178 

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

186 

187 if norm1 > norm2+5: 

188 norm1, norm2 = norm2, norm1 

189 

190 norm1 = min(norm1, norm2 - 10) 

191 

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) 

202 

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 

214 

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} 

232 

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 

238 

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 

245 

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. 

261 

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) 

274 

275 (if the output group is None, _sys.xafsGroup will be written to) 

276 

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

301 

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) 

309 

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. 

316 

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) 

322 

323 group.e0 = e0 

324 group.norm = norm 

325 group.flat = 1.0*norm 

326 group.norm_poly = 1.0*norm 

327 

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 

336 

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) 

348 

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 

355 

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 

363 

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

370 

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

374 

375 group.pre_edge_details.pre_slope = pre_dat['precoefs'][0] 

376 group.pre_edge_details.pre_offset = pre_dat['precoefs'][1] 

377 

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) 

383 

384 # guess element and edge 

385 group.atsym = getattr(group, 'atsym', None) 

386 group.edge = getattr(group, 'edge', None) 

387 

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 

393 

394def energy_align(group, reference, array='dmude', emin=-15, emax=35): 

395 """ 

396 align XAFS data group to a reference group 

397 

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] 

405 

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 

410 

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. 

415 

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. 

418 

419 """ 

420 if not (hasattr(group, 'energy') and hasattr(group, 'mu')): 

421 raise ValueError("group must have attributes 'energy' and 'mu'") 

422 

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) 

427 

428 

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

432 

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) 

437 

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) 

450 

451 i1 = index_of(xref, reference.e0-emin) 

452 i2 = index_of(xref, reference.e0+emax) 

453 

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] 

460 

461 params = Parameters() 

462 params.add('eshift', value=0, min=-50, max=50) 

463 params.add('scale', value=1, min=0, max=50) 

464 

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 

472 

473 group.eshift = eshift 

474 return eshift