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
« 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
18from xraydb import guess_edge, xray_edge, core_width
20from larch import Group, Make_CallArgs, isgroup, parse_group_args
21# now we can reliably import other std and xafs modules...
23from larch.math import (index_of, index_nearest,
24 remove_dups, remove_nans2)
26from larch.fitting import dict2params
27from .xafsutils import set_xafsGroup
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.
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.
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.
48 A group named `prepeaks` will be created in the output group, containing:
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 ============== ===========================================================
57 Note that 'norm' will be used here even if a differnt input array was used.
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
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
81 if len(energy.shape) > 1:
82 energy = energy.squeeze()
83 if len(ydat.shape) > 1:
84 ydat = ydat.squeeze()
86 dat_emin, dat_emax = min(energy), max(energy)
87 dat_e0 = getattr(group, 'e0', -1)
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
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")
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
118 dele = 1.e-13 + min(np.diff(energy))/5.0
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)
125 edat = energy[imin: imax+1]
126 ydat = ydat[imin:imax+1]
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
140 group.prepeaks.xdat = edat
141 group.prepeaks.ydat = norm
142 return
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
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.
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.
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'.
171 A group named 'prepeaks' will be used or created in the output group, containing
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 ============== ===========================================================
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.
193 """
194 energy, norm, group = parse_group_args(energy, members=('energy', 'norm'),
195 defaults=(norm,), group=group,
196 fcn_name='pre_edge_baseline')
198 prepeaks_setup(energy, norm=norm, group=group, emin=emin, emax=emax,
199 elo=elo, ehi=ehi, _larch=_larch)
201 emin = group.prepeaks.emin
202 emax = group.prepeaks.emax
203 elo = group.prepeaks.elo
204 ehi = group.prepeaks.ehi
206 dele = 1.e-13 + min(np.diff(energy))/5.0
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)
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]))
217 edat = energy[imin: imax+1]
218 cen = dcen = 0.
219 peak_energies = []
221 # energy including pre-edge peaks, for output
222 norm = norm[imin:imax+1]
223 baseline = peaks = dpeaks = norm*0.0
225 # build fitting model:
226 modelcomps = []
227 parvals = {}
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))}
236 if '+' in form:
237 forms = [f.lower() for f in form.split('+')]
238 else:
239 forms = [form.lower(), '']
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])
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
257 model = modelcomps.pop()
258 if len(modelcomps) > 0:
259 model += modelcomps.pop()
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
271 result = model.fit(ydat, params, x=xdat)
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
278 # estimate centroid
279 cen = (edat*peaks).sum() / peaks.sum()
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
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
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]
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
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
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")
317 if not isinstance(peakmodel, Model):
318 raise ValueError("peakmodel must be an lmfit.Model")
320 if isinstance(params, dict):
321 params = dict2params(params)
323 if not isinstance(params, Parameters):
324 raise ValueError("params must be an lmfit.Parameters")
326 if not hasattr(prepeaks, 'fit_history'):
327 prepeaks.fit_history = []
329 pkfit = Group()
331 for k in ('energy', 'norm', 'norm_std', 'user_options'):
332 if hasattr(prepeaks, k):
333 setattr(pkfit, k, deepcopy(getattr(prepeaks, k)))
335 if user_options is not None:
336 pkfit.user_options = user_options
338 pkfit.init_fit = peakmodel.eval(params, x=prepeaks.energy)
339 pkfit.init_ycomps = peakmodel.eval_components(params=params, x=prepeaks.energy)
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
347 pkfit.result = peakmodel.fit(prepeaks.norm, params=params, x=prepeaks.energy,
348 weights=1.0/norm_std)
350 pkfit.ycomps = peakmodel.eval_components(params=pkfit.result.params, x=prepeaks.energy)
351 pkfit.label = 'Fit %i' % (1+len(prepeaks.fit_history))
353 label = now = time.strftime("%b-%d %H:%M")
354 pkfit.timestamp = time.strftime("%Y-%b-%d %H:%M")
355 pkfit.label = label
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
367 prepeaks.fit_history.insert(0, pkfit)
368 return pkfit