Coverage for /Users/Newville/Codes/xraylarch/larch/xafs/prepeaks.py: 11%

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

12try: 

13 import peakutils 

14 HAS_PEAKUTILS = True 

15except ImportError: 

16 HAS_PEAKUTILS = False 

17 

18from xraydb import guess_edge, xray_edge, core_width 

19 

20from larch import Group, Make_CallArgs, isgroup, parse_group_args 

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

22 

23from larch.math import (index_of, index_nearest, 

24 remove_dups, remove_nans2) 

25 

26from larch.fitting import dict2params 

27from .xafsutils import set_xafsGroup 

28 

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

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

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

32 """set up pre edge peak group. 

33 

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

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

36 

37 Arguments: 

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

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

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

41 group (group or None): output group 

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

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

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

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

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

47 

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

49 

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

51 attribute meaning 

52 ============== =========================================================== 

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

54 norm spectrum over pre-edge peak energies 

55 ============== =========================================================== 

56 

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

58 

59 Notes: 

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

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

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

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

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

65 """ 

66 ydat = None 

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

68 arrayname = 'aspassed' 

69 ydat = norm[:]*1.0 

70 

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

72 defaults=(norm,), group=group, 

73 fcn_name='pre_edge_baseline') 

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

75 ydat = group.flat[:]*1.0 

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

77 ydat = group.deconv[:]*1.0 

78 if ydat is None: 

79 ydat = norm[:]*1.0 

80 

81 if len(energy.shape) > 1: 

82 energy = energy.squeeze() 

83 if len(ydat.shape) > 1: 

84 ydat = ydat.squeeze() 

85 

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

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

88 

89 if dat_e0 > 0: 

90 if emin is None: 

91 emin = dat_e0 - 30.0 

92 if emax is None: 

93 emax = dat_e0 - 1.0 

94 if elo is None: 

95 elo = dat_e0 - 15.0 

96 if ehi is None: 

97 ehi = dat_e0 - 5.0 

98 if emin < 0: 

99 emin += dat_e0 

100 if elo < 0: 

101 elo += dat_e0 

102 if emax < dat_emin: 

103 emax += dat_e0 

104 if ehi < dat_emin: 

105 ehi += dat_e0 

106 

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

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

109 

110 # get indices for input energies 

111 if emin > emax: 

112 emin, emax = emax, emin 

113 if emin > elo: 

114 elo, emin = emin, elo 

115 if ehi > emax: 

116 ehi, emax = emax, ehi 

117 

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

119 

120 ilo = index_of(energy, elo+dele) 

121 ihi = index_of(energy, ehi+dele) 

122 imin = index_of(energy, emin+dele) 

123 imax = index_of(energy, emax+dele) 

124 

125 edat = energy[imin: imax+1] 

126 ydat = ydat[imin:imax+1] 

127 

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

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

130 emin=emin, emax=emax, 

131 elo=elo, ehi=ehi) 

132 else: 

133 group.prepeaks.energy = edat 

134 group.prepeaks.norm = ydat 

135 group.prepeaks.emin = emin 

136 group.prepeaks.emax = emax 

137 group.prepeaks.elo = elo 

138 group.prepeaks.ehi = ehi 

139 

140 group.prepeaks.xdat = edat 

141 group.prepeaks.ydat = norm 

142 return 

143 

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

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

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

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

148 

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

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

151 

152 Arguments: 

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

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

155 group (group or None): output group 

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

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

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

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

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

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

162 

163 

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

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

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

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

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

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

170 

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

172 

173 ============== =========================================================== 

174 attribute meaning 

175 ============== =========================================================== 

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

177 baseline fitted baseline array over pre-edge peak energies 

178 norm spectrum over pre-edge peak energies 

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

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

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

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

183 ============== =========================================================== 

184 

185 Notes: 

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

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

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

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

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

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

192 

193 """ 

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

195 defaults=(norm,), group=group, 

196 fcn_name='pre_edge_baseline') 

197 

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

199 elo=elo, ehi=ehi, _larch=_larch) 

200 

201 emin = group.prepeaks.emin 

202 emax = group.prepeaks.emax 

203 elo = group.prepeaks.elo 

204 ehi = group.prepeaks.ehi 

205 

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

207 

208 imin = index_of(energy, emin+dele) 

209 ilo = index_of(energy, elo+dele) 

210 ihi = index_of(energy, ehi+dele) 

211 imax = index_of(energy, emax+dele) 

212 

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

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

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

216 

217 edat = energy[imin: imax+1] 

218 cen = dcen = 0. 

219 peak_energies = [] 

220 

221 # energy including pre-edge peaks, for output 

222 norm = norm[imin:imax+1] 

223 baseline = peaks = dpeaks = norm*0.0 

224 

225 # build fitting model: 

226 modelcomps = [] 

227 parvals = {} 

228 

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

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

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

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

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

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

235 

236 if '+' in form: 

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

238 else: 

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

240 

241 for form in forms[:2]: 

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

243 if form.startswith(key): 

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

245 parvals.update(dat[1]) 

246 

247 if len(modelcomps) == 0: 

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

249 peaks=0.0*edat, delta_peaks=0.0*edat, 

250 centroid=0, delta_centroid=0, 

251 peak_energies=[], 

252 fit_details=None, 

253 emin=emin, emax=emax, elo=elo, ehi=ehi, 

254 form=form) 

255 return 

256 

257 model = modelcomps.pop() 

258 if len(modelcomps) > 0: 

259 model += modelcomps.pop() 

260 

261 params = model.make_params(**parvals) 

262 if 'amplitude' in params: 

263 params['amplitude'].min = 0.0 

264 if 'sigma' in params: 

265 params['sigma'].min = 0.05 

266 params['sigma'].max = 500.0 

267 if 'center' in params: 

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

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

270 

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

272 

273 # get baseline and resulting norm over edat range 

274 if result is not None: 

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

276 peaks = norm-baseline 

277 

278 # estimate centroid 

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

280 

281 # uncertainty in norm includes only uncertainties in baseline fit 

282 # and uncertainty in centroid: 

283 try: 

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

285 except: 

286 dbpeaks = 0.0 

287 

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

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

290 dcen = abs(cen_minus - cen_plus) / 2.0 

291 

292 # locate peak positions 

293 if HAS_PEAKUTILS: 

294 peak_ids = peakutils.peak.indexes(peaks, thres=0.05, min_dist=2) 

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

296 

297 group = set_xafsGroup(group, _larch=_larch) 

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

299 peaks=peaks, delta_peaks=dpeaks, 

300 centroid=cen, delta_centroid=dcen, 

301 peak_energies=peak_energies, 

302 fit_details=result, 

303 emin=emin, emax=emax, elo=elo, ehi=ehi, 

304 form=form) 

305 return 

306 

307 

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

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

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

311 

312 """ 

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

314 if prepeaks is None: 

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

316 

317 if not isinstance(peakmodel, Model): 

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

319 

320 if isinstance(params, dict): 

321 params = dict2params(params) 

322 

323 if not isinstance(params, Parameters): 

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

325 

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

327 prepeaks.fit_history = [] 

328 

329 pkfit = Group() 

330 

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

332 if hasattr(prepeaks, k): 

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

334 

335 if user_options is not None: 

336 pkfit.user_options = user_options 

337 

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

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

340 

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

342 if isinstance(norm_std, np.ndarray): 

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

344 elif norm_std < 0: 

345 norm_std = 1.0 

346 

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

348 weights=1.0/norm_std) 

349 

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

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

352 

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

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

355 pkfit.label = label 

356 

357 

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

359 if label in fitlabels: 

360 count = 1 

361 while label in fitlabels: 

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

363 count +=1 

364 pkfit.label = label 

365 

366 

367 prepeaks.fit_history.insert(0, pkfit) 

368 return pkfit