Coverage for larch/xafs/prepeaks.py: 9%

176 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 time 

6from string import printable 

7from copy import deepcopy 

8import numpy as np 

9from lmfit import Parameters, Minimizer, Model 

10from lmfit.models import (LorentzianModel, GaussianModel, VoigtModel, 

11 ConstantModel, LinearModel, QuadraticModel) 

12 

13from xraydb import guess_edge, xray_edge, core_width 

14 

15from larch import Group, Make_CallArgs, isgroup, parse_group_args 

16# now we can reliably import other std and xafs modules... 

17 

18from larch.math import (index_of, index_nearest, remove_nans2, 

19 peak_indices) 

20 

21from larch.fitting import dict2params 

22from .xafsutils import set_xafsGroup 

23 

24@Make_CallArgs(["energy", "norm"]) 

25def prepeaks_setup(energy, norm=None, arrayname=None, group=None, emin=None, emax=None, 

26 elo=None, ehi=None, _larch=None): 

27 """set up pre edge peak group. 

28 

29 This assumes that pre_edge() has been run successfully on the spectra 

30 and that the spectra has decent pre-edge subtraction and normalization. 

31 

32 Arguments: 

33 energy (ndarray or group): array of x-ray energies, in eV, or group (see note 1) 

34 norm (ndarray or None): array of normalized mu(E) to be fit (deprecated, see note 2) 

35 arrayname (string or None): name of array to use as data (see note 2) 

36 group (group or None): output group 

37 emax (float or None): max energy (eV) to use for baesline fit [e0-5] 

38 emin (float or None): min energy (eV) to use for baesline fit [e0-40] 

39 elo: (float or None) low energy of pre-edge peak region to not fit baseline [e0-20] 

40 ehi: (float or None) high energy of pre-edge peak region ot not fit baseline [e0-10] 

41 _larch (larch instance or None): current larch session. 

42 

43 A group named `prepeaks` will be created in the output group, containing: 

44 

45 ============== =========================================================== 

46 attribute meaning 

47 ============== =========================================================== 

48 energy energy array for pre-edge peaks = energy[emin:emax] 

49 norm spectrum over pre-edge peak energies 

50 ============== =========================================================== 

51 

52 Note that 'norm' will be used here even if a differnt input array was used. 

53 

54 Notes: 

55 1. Supports :ref:`First Argument Group` convention, requiring group members `energy` and `norm` 

56 2. You can pass is an array to fit with 'norm=', or you can name the array to use with 

57 `arrayname`, which can be one of ['norm', 'flat', 'deconv', 'aspassed', None] with None 

58 and 'aspassed' meaning that the argument of `norm` will be used, as passed in. 

59 3. Supports :ref:`Set XAFS Group` convention within Larch or if `_larch` is set. 

60 """ 

61 ydat = None 

62 if norm is not None and arrayname in (None, 'aspassed'): 

63 arrayname = 'aspassed' 

64 ydat = norm[:]*1.0 

65 

66 energy, norm, group = parse_group_args(energy, members=('energy', 'norm'), 

67 defaults=(norm,), group=group, 

68 fcn_name='pre_edge_baseline') 

69 if arrayname == 'flat' and hasattr(group, 'flat'): 

70 ydat = group.flat[:]*1.0 

71 elif arrayname == 'deconv' and hasattr(group, 'deconv'): 

72 ydat = group.deconv[:]*1.0 

73 if ydat is None: 

74 ydat = norm[:]*1.0 

75 

76 if len(energy.shape) > 1: 

77 energy = energy.squeeze() 

78 if len(ydat.shape) > 1: 

79 ydat = ydat.squeeze() 

80 

81 dat_emin, dat_emax = min(energy), max(energy) 

82 dat_e0 = getattr(group, 'e0', -1) 

83 

84 if dat_e0 > 0: 

85 if emin is None: 

86 emin = dat_e0 - 30.0 

87 if emax is None: 

88 emax = dat_e0 - 1.0 

89 if elo is None: 

90 elo = dat_e0 - 15.0 

91 if ehi is None: 

92 ehi = dat_e0 - 5.0 

93 if emin < 0: 

94 emin += dat_e0 

95 if elo < 0: 

96 elo += dat_e0 

97 if emax < dat_emin: 

98 emax += dat_e0 

99 if ehi < dat_emin: 

100 ehi += dat_e0 

101 

102 if emax is None or emin is None or elo is None or ehi is None: 

103 raise ValueError("must provide emin and emax to prepeaks_setup") 

104 

105 # get indices for input energies 

106 if emin > emax: 

107 emin, emax = emax, emin 

108 if emin > elo: 

109 elo, emin = emin, elo 

110 if ehi > emax: 

111 ehi, emax = emax, ehi 

112 

113 dele = 1.e-13 + min(np.diff(energy))/5.0 

114 

115 ilo = index_of(energy, elo+dele) 

116 ihi = index_of(energy, ehi+dele) 

117 imin = index_of(energy, emin+dele) 

118 imax = index_of(energy, emax+dele) 

119 

120 edat = energy[imin: imax+1] 

121 ydat = ydat[imin:imax+1] 

122 

123 if not hasattr(group, 'prepeaks'): 

124 group.prepeaks = Group(energy=edat, norm=ydat, 

125 emin=emin, emax=emax, 

126 elo=elo, ehi=ehi) 

127 else: 

128 group.prepeaks.energy = edat 

129 group.prepeaks.norm = ydat 

130 group.prepeaks.emin = emin 

131 group.prepeaks.emax = emax 

132 group.prepeaks.elo = elo 

133 group.prepeaks.ehi = ehi 

134 

135 group.prepeaks.xplot = edat 

136 group.prepeaks.yplot = norm 

137 return 

138 

139@Make_CallArgs(["energy", "norm"]) 

140def pre_edge_baseline(energy, norm=None, group=None, form='linear+lorentzian', 

141 emin=None, emax=None, elo=None, ehi=None, _larch=None): 

142 """remove baseline from main edge over pre edge peak region 

143 

144 This assumes that pre_edge() has been run successfully on the spectra 

145 and that the spectra has decent pre-edge subtraction and normalization. 

146 

147 Arguments: 

148 energy (ndarray or group): array of x-ray energies, in eV, or group (see note 1) 

149 norm (ndarray or group): array of normalized mu(E) 

150 group (group or None): output group 

151 elo (float or None): low energy of pre-edge peak region to not fit baseline [e0-20] 

152 ehi (float or None): high energy of pre-edge peak region ot not fit baseline [e0-10] 

153 emax (float or None): max energy (eV) to use for baesline fit [e0-5] 

154 emin (float or None): min energy (eV) to use for baesline fit [e0-40] 

155 form (string): form used for baseline (see description) ['linear+lorentzian'] 

156 _larch (larch instance or None): current larch session. 

157 

158 

159 A function will be fit to the input mu(E) data over the range between 

160 [emin:elo] and [ehi:emax], ignorng the pre-edge peaks in the region 

161 [elo:ehi]. The baseline function is specified with the `form` keyword 

162 argument, which can be one or a combination of 'lorentzian', 'gaussian', or 'voigt', 

163 plus one of 'constant', 'linear', 'quadratic', for example, 'linear+lorentzian', 

164 'constant+voigt', 'quadratic', 'gaussian'. 

165 

166 A group named 'prepeaks' will be used or created in the output group, containing 

167 

168 ============== =========================================================== 

169 attribute meaning 

170 ============== =========================================================== 

171 energy energy array for pre-edge peaks = energy[emin:emax] 

172 baseline fitted baseline array over pre-edge peak energies 

173 norm spectrum over pre-edge peak energies 

174 peaks baseline-subtraced spectrum over pre-edge peak energies 

175 centroid estimated centroid of pre-edge peaks (see note 3) 

176 peak_energies list of predicted peak energies (see note 4) 

177 fit_details details of fit to extract pre-edge peaks. 

178 ============== =========================================================== 

179 

180 Notes: 

181 1. Supports :ref:`First Argument Group` convention, requiring group members `energy` and `norm` 

182 2. Supports :ref:`Set XAFS Group` convention within Larch or if `_larch` is set. 

183 3. The value calculated for `prepeaks.centroid` will be found as 

184 (prepeaks.energy*prepeaks.peaks).sum() / prepeaks.peaks.sum() 

185 4. The values in the `peak_energies` list will be predicted energies 

186 of the peaks in `prepeaks.peaks` as found by peakutils. 

187 

188 """ 

189 energy, norm, group = parse_group_args(energy, members=('energy', 'norm'), 

190 defaults=(norm,), group=group, 

191 fcn_name='pre_edge_baseline') 

192 

193 prepeaks_setup(energy, norm=norm, group=group, emin=emin, emax=emax, 

194 elo=elo, ehi=ehi, _larch=_larch) 

195 

196 emin = group.prepeaks.emin 

197 emax = group.prepeaks.emax 

198 elo = group.prepeaks.elo 

199 ehi = group.prepeaks.ehi 

200 

201 dele = 1.e-11 + min(np.diff(energy))/5.0 

202 

203 imin = index_of(energy, emin+dele) 

204 ilo = index_of(energy, elo+dele) 

205 ihi = index_of(energy, ehi+dele) 

206 imax = index_of(energy, emax+dele) 

207 

208 # build xdat, ydat: dat to fit (skipping pre-edge peaks) 

209 xdat = np.concatenate((energy[imin:ilo+1], energy[ihi:imax+1])) 

210 ydat = np.concatenate((norm[imin:ilo+1], norm[ihi:imax+1])) 

211 

212 edat = energy[imin: imax+1] 

213 cen = dcen = 0. 

214 peak_energies = [] 

215 

216 # energy including pre-edge peaks, for output 

217 norm = norm[imin:imax+1] 

218 baseline = peaks = dpeaks = norm*0.0 

219 

220 # build fitting model: 

221 modelcomps = [] 

222 parvals = {} 

223 

224 MODELDAT = {'gauss': (GaussianModel, dict(amplitude=1, center=emax, sigma=2)), 

225 'loren': (LorentzianModel, dict(amplitude=1, center=emax, sigma=2)), 

226 'voigt': (VoigtModel, dict(amplitude=1, center=emax, sigma=2)), 

227 'line': (LinearModel, dict(slope=0, intercept=0)), 

228 'quad': (QuadraticModel, dict(a=0, b=0, c=0)), 

229 'const': (ConstantModel, dict(c=0))} 

230 

231 if '+' in form: 

232 forms = [f.lower() for f in form.split('+')] 

233 else: 

234 forms = [form.lower(), ''] 

235 

236 for form in forms[:2]: 

237 for key, dat in MODELDAT.items(): 

238 if form.startswith(key): 

239 modelcomps.append(dat[0]()) 

240 parvals.update(dat[1]) 

241 

242 if len(modelcomps) == 0: 

243 group.prepeaks = Group(energy=edat, norm=norm, baseline=0.0*edat, 

244 peaks=0.0*edat, delta_peaks=0.0*edat, 

245 centroid=0, delta_centroid=0, 

246 peak_energies=[], 

247 fit_details=None, 

248 emin=emin, emax=emax, elo=elo, ehi=ehi, 

249 form=form) 

250 return 

251 

252 model = modelcomps.pop() 

253 if len(modelcomps) > 0: 

254 model += modelcomps.pop() 

255 

256 params = model.make_params(**parvals) 

257 if 'amplitude' in params: 

258 params['amplitude'].min = 0.0 

259 if 'sigma' in params: 

260 params['sigma'].min = 0.05 

261 params['sigma'].max = 500.0 

262 if 'center' in params: 

263 params['center'].max = emax + 25.0 

264 params['center'].min = emin - 25.0 

265 

266 result = model.fit(ydat, params, x=xdat) 

267 

268 # get baseline and resulting norm over edat range 

269 if result is not None: 

270 baseline = result.eval(result.params, x=edat) 

271 peaks = norm-baseline 

272 

273 # estimate centroid 

274 cen = (edat*peaks).sum() / peaks.sum() 

275 

276 # uncertainty in norm includes only uncertainties in baseline fit 

277 # and uncertainty in centroid: 

278 try: 

279 dpeaks = result.eval_uncertainty(result.params, x=edat) 

280 except: 

281 dbpeaks = 0.0 

282 

283 cen_plus = (edat*(peaks+dpeaks)).sum()/ (peaks+dpeaks).sum() 

284 cen_minus = (edat*(peaks-dpeaks)).sum()/ (peaks-dpeaks).sum() 

285 dcen = abs(cen_minus - cen_plus) / 2.0 

286 

287 # locate peak positions 

288 peak_ids = peak_indices(peaks, threshold=0.05, min_dist=2) 

289 peak_energies = [edat[pid] for pid in peak_ids] 

290 

291 group = set_xafsGroup(group, _larch=_larch) 

292 group.prepeaks = Group(energy=edat, norm=norm, baseline=baseline, 

293 peaks=peaks, delta_peaks=dpeaks, 

294 centroid=cen, delta_centroid=dcen, 

295 peak_energies=peak_energies, 

296 fit_details=result, 

297 emin=emin, emax=emax, elo=elo, ehi=ehi, 

298 form=form) 

299 return 

300 

301 

302def prepeaks_fit(group, peakmodel, params, user_options=None, _larch=None): 

303 """do pre-edge peak fitting - must be done after setting up the fit 

304 returns a group with Peakfit data, including `result`, the lmfit ModelResult 

305 

306 """ 

307 prepeaks = getattr(group, 'prepeaks', None) 

308 if prepeaks is None: 

309 raise ValueError("must run prepeask_setup() for a group before doing fit") 

310 

311 if not isinstance(peakmodel, Model): 

312 raise ValueError("peakmodel must be an lmfit.Model") 

313 

314 if isinstance(params, dict): 

315 params = dict2params(params) 

316 

317 if not isinstance(params, Parameters): 

318 raise ValueError("params must be an lmfit.Parameters") 

319 

320 if not hasattr(prepeaks, 'fit_history'): 

321 prepeaks.fit_history = [] 

322 

323 pkfit = Group() 

324 

325 for k in ('energy', 'norm', 'norm_std', 'user_options'): 

326 if hasattr(prepeaks, k): 

327 setattr(pkfit, k, deepcopy(getattr(prepeaks, k))) 

328 

329 if user_options is not None: 

330 pkfit.user_options = user_options 

331 

332 pkfit.init_fit = peakmodel.eval(params, x=prepeaks.energy) 

333 pkfit.init_ycomps = peakmodel.eval_components(params=params, x=prepeaks.energy) 

334 

335 norm_std = getattr(prepeaks, 'norm_std', 1.0) 

336 if isinstance(norm_std, np.ndarray): 

337 norm_std[np.where(norm_std<1.e-13)] = 1.e-13 

338 elif norm_std < 0: 

339 norm_std = 1.0 

340 

341 pkfit.result = peakmodel.fit(prepeaks.norm, params=params, x=prepeaks.energy, 

342 weights=1.0/norm_std) 

343 

344 pkfit.ycomps = peakmodel.eval_components(params=pkfit.result.params, x=prepeaks.energy) 

345 pkfit.label = 'Fit %i' % (1+len(prepeaks.fit_history)) 

346 

347 label = now = time.strftime("%b-%d %H:%M") 

348 pkfit.timestamp = time.strftime("%Y-%b-%d %H:%M") 

349 pkfit.label = label 

350 

351 

352 fitlabels = [fhist.label for fhist in prepeaks.fit_history] 

353 if label in fitlabels: 

354 count = 1 

355 while label in fitlabels: 

356 label = f'{now:s}_{printable[count]:s}' 

357 count +=1 

358 pkfit.label = label 

359 

360 

361 prepeaks.fit_history.insert(0, pkfit) 

362 return pkfit