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

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, interp, smooth, 

12 polyfit, remove_dups, remove_nans, remove_nans2) 

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 E0, the energy threshold of absorption, or 'edge energy', given mu(E). 

21 

22 E0 is found as the point with maximum derivative with some checks to avoid spurious glitches. 

23 

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. 

29 

30 Returns: 

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

32 

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 

49 

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 

62 

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

72 

73 

74def _finde0(energy, mu_input, estep=None, use_smooth=True): 

75 "internally used by find e0 " 

76 

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) 

87 

88 

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 

99 

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] 

108 

109 if len(high_deriv_pts) < 3: 

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

111 

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 

121 

122def flat_resid(pars, en, mu): 

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

124 

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) 

128 

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 

135 

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 

157 

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 

164 

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

172 

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] 

181 

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

197 

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

205 

206 if norm1 > norm2: 

207 norm1, norm2 = norm2, norm1 

208 

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) 

220 

221 omu = mu*energy**nvict 

222 ex = remove_nans(energy[p1:p2], interp=True) 

223 mx = remove_nans(omu[p1:p2], interp=True) 

224 

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 

234 

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} 

252 

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 

258 

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 

265 

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. 

281 

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) 

294 

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

296 

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

322 

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) 

329 

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) 

336 

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. 

343 

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) 

349 

350 group.e0 = e0 

351 group.norm = norm 

352 group.flat = 1.0*norm 

353 group.norm_poly = 1.0*norm 

354 

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 

363 

364 enx = remove_nans(energy[p1:p2], interp=True) 

365 mux = remove_nans(norm[p1:p2], interp=True) 

366 

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) 

377 

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 

384 

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 

392 

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

399 

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

403 

404 group.pre_edge_details.pre_slope = pre_dat['precoefs'][1] 

405 group.pre_edge_details.pre_offset = pre_dat['precoefs'][0] 

406 

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) 

412 

413 # guess element and edge 

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

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

416 

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 

422 

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

424 """ 

425 align XAFS data group to a reference group 

426 

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] 

434 

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 

439 

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. 

444 

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. 

447 

448 """ 

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

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

451 

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) 

456 

457 

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

461 

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) 

466 

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 

477 

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) 

482 

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

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

485 

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] 

492 

493 params = Parameters() 

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

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

496 

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 

504 

505 group.eshift = eshift 

506 return eshift