Coverage for larch/xafs/feffit.py: 75%

630 statements  

« prev     ^ index     » next       coverage.py v7.6.0, created at 2024-10-16 21:04 +0000

1#!/usr/bin/env python 

2""" 

3 feffit sums Feff paths to match xafs data 

4""" 

5try: 

6 from collections.abc import Iterable 

7except ImportError: 

8 from collections import Iterable 

9from copy import copy, deepcopy 

10from functools import partial 

11import ast 

12import numpy as np 

13from numpy import array, arange, interp, pi, zeros, sqrt, concatenate 

14 

15from scipy.interpolate import splrep, splev 

16from scipy.interpolate import InterpolatedUnivariateSpline as IUSpline 

17from lmfit import Parameters, Parameter, Minimizer, conf_interval2d 

18from lmfit.printfuncs import getfloat_attr 

19 

20from larch import Group, isNamedClass 

21from larch.utils import fix_varname, gformat 

22from larch.utils.strutils import b32hash, random_varname 

23from ..math import index_of, realimag, complex_phase, remove_nans 

24from ..fitting import (correlated_values, eval_stderr, ParameterGroup, 

25 group2params, params2group, isParameter) 

26 

27from .xafsutils import set_xafsGroup, gfmt 

28from .xafsft import xftf_fast, xftr_fast, ftwindow 

29from .autobk import autobk_delta_chi 

30from .feffdat import FeffPathGroup, ff2chi 

31 

32class TransformGroup(Group): 

33 """A Group of transform parameters. 

34 The apply() method will return the result of applying the transform, 

35 ready to use in a Fit. This caches the FT windows (k and r windows) 

36 and assumes that once created (not None), these do not need to be 

37 recalculated.... 

38 

39 That is: don't simply change the parameters and expect different results. 

40 If you do change parameters, reset kwin / rwin to None to cause them to be 

41 recalculated. 

42 """ 

43 def __init__(self, kmin=0, kmax=20, kweight=2, dk=4, dk2=None, 

44 window='kaiser', nfft=2048, kstep=0.05, 

45 rmin = 0, rmax=10, dr=0, dr2=None, rwindow='hanning', 

46 fitspace='r', wavelet_mask=None, rbkg=0, _larch=None, **kws): 

47 Group.__init__(self, **kws) 

48 self.kmin = kmin 

49 self.kmax = kmax 

50 self.kweight = kweight 

51 if 'kw' in kws: 

52 self.kweight = kws['kw'] 

53 

54 self.dk = dk 

55 self.dk2 = dk2 

56 self.window = window 

57 self.rbkg = rbkg 

58 self.rmin = rmin 

59 self.rmax = rmax 

60 self.dr = dr 

61 self.dr2 = dr2 

62 if dr2 is None: self.dr2 = self.dr 

63 self.rwindow = rwindow 

64 self.__nfft = 0 

65 self.__kstep = None 

66 self.nfft = nfft 

67 self.kstep = kstep 

68 self.rstep = pi/(self.kstep*self.nfft) 

69 

70 self.fitspace = fitspace 

71 self.wavelet_mask = wavelet_mask 

72 self._cauchymask = None 

73 

74 self._larch = _larch 

75 

76 self.kwin = None 

77 self.rwin = None 

78 self.make_karrays() 

79 

80 def __repr__(self): 

81 return '<FeffitTransform Group: %s>' % self.__name__ 

82 

83 def __copy__(self): 

84 return TransformGroup(kmin=self.kmin, kmax=self.kmax, 

85 kweight=self.kweight, dk=self.dk, dk2=self.dk2, 

86 window=self.window, kstep=self.kstep, 

87 rmin=self.rmin, rmax=self.rmax, rbkg=self.rbkg, 

88 dr=self.dr, dr2=self.dr2, 

89 rwindow=self.rwindow, nfft=self.nfft, 

90 fitspace=self.fitspace, 

91 wavelet_mask=self.wavelet_mask, 

92 _larch=self._larch) 

93 

94 def __deepcopy__(self, memo): 

95 return TransformGroup(kmin=self.kmin, kmax=self.kmax, 

96 kweight=self.kweight, dk=self.dk, dk2=self.dk2, 

97 window=self.window, kstep=self.kstep, 

98 rmin=self.rmin, rmax=self.rmax, rbkg=self.rbkg, 

99 dr=self.dr, dr2=self.dr2, 

100 rwindow=self.rwindow, nfft=self.nfft, 

101 fitspace=self.fitspace, 

102 wavelet_mask=self.wavelet_mask, 

103 _larch=self._larch) 

104 

105 def make_karrays(self, k=None, chi=None): 

106 "this should be run in kstep or nfft changes" 

107 if self.kstep == self.__kstep and self.nfft == self.__nfft: 

108 return 

109 self.__kstep = self.kstep 

110 self.__nfft = self.nfft 

111 

112 self.rstep = pi/(self.kstep*self.nfft) 

113 self.k_ = self.kstep * arange(self.nfft, dtype='float64') 

114 self.r_ = self.rstep * arange(self.nfft, dtype='float64') 

115 

116 def _xafsft(self, chi, group=None, rmax_out=10, **kws): 

117 "returns " 

118 for key, val in kws: 

119 if key == 'kw': 

120 key = 'kweight' 

121 setattr(self, key, val) 

122 self.make_karrays() 

123 

124 out = self.fftf(chi) 

125 

126 irmax = int(min(self.nfft/2, 1.01 + rmax_out/self.rstep)) 

127 

128 group = set_xafsGroup(group, _larch=self._larch) 

129 r = self.rstep * arange(irmax) 

130 mag = sqrt(out.real**2 + out.imag**2) 

131 group.kwin = self.kwin[:len(chi)] 

132 group.r = r[:irmax] 

133 group.chir = out[:irmax] 

134 group.chir_mag = mag[:irmax] 

135 group.chir_pha = complex_phase(out[:irmax]) 

136 group.chir_re = out.real[:irmax] 

137 group.chir_im = out.imag[:irmax] 

138 if self.rwin is None: 

139 xmin = max(self.rbkg, self.rmin) 

140 self.rwin = ftwindow(self.r_, xmin=xmin, xmax=self.rmax, 

141 dx=self.dr, dx2=self.dr2, window=self.rwindow) 

142 group.rwin = self.rwin[:irmax] 

143 

144 def get_kweight(self): 

145 "if kweight is a list/tuple, use only the first one here" 

146 if isinstance(self.kweight, Iterable): 

147 return self.kweight[0] 

148 return self.kweight 

149 

150 def fftf(self, chi, kweight=None): 

151 """ forward FT -- meant to be used internally. 

152 chi must be on self.k_ grid""" 

153 if self.kstep != self.__kstep or self.nfft != self.__nfft: 

154 self.make_karrays() 

155 if self.kwin is None: 

156 self.kwin = ftwindow(self.k_, xmin=self.kmin, xmax=self.kmax, 

157 dx=self.dk, dx2=self.dk2, window=self.window) 

158 if kweight is None: 

159 kweight = self.get_kweight() 

160 cx = chi * self.kwin[:len(chi)] * self.k_[:len(chi)]**kweight 

161 return xftf_fast(cx, kstep=self.kstep, nfft=self.nfft) 

162 

163 def fftr(self, chir): 

164 " reverse FT -- meant to be used internally" 

165 if self.kstep != self.__kstep or self.nfft != self.__nfft: 

166 self.make_karrays() 

167 if self.rwin is None: 

168 xmin = max(self.rbkg, self.rmin) 

169 self.rwin = ftwindow(self.r_, xmin=xmin, xmax=self.rmax, 

170 dx=self.dr, dx2=self.dr2, window=self.rwindow) 

171 cx = chir * self.rwin[:len(chir)] 

172 return xftr_fast(cx, kstep=self.kstep, nfft=self.nfft) 

173 

174 

175 def make_cwt_arrays(self, nkpts, nrpts): 

176 if self.kstep != self.__kstep or self.nfft != self.__nfft: 

177 self.make_karrays() 

178 if self.kwin is None: 

179 self.kwin = ftwindow(self.k_, xmin=self.kmin, xmax=self.kmax, 

180 dx=self.dk, dx2=self.dk2, window=self.window) 

181 

182 if self._cauchymask is None: 

183 if self.wavelet_mask is not None: 

184 self._cauchymask = self.wavelet_mask 

185 else: 

186 ikmin = int(max(0, 0.01 + self.kmin/self.kstep)) 

187 ikmax = int(min(self.nfft/2, 0.01 + self.kmax/self.kstep)) 

188 irmin = int(max(0, 0.01 + self.rmin/self.rstep)) 

189 irmax = int(min(self.nfft/2, 0.01 + self.rmax/self.rstep)) 

190 cm = np.zeros(nrpts*nkpts, dtype='int').reshape(nrpts, nkpts) 

191 cm[irmin:irmax, ikmin:ikmax] = 1 

192 self._cauchymask = cm 

193 self._cauchyslice =(slice(irmin, irmax), slice(ikmin, ikmax)) 

194 

195 def cwt(self, chi, rmax=None, kweight=None): 

196 """cauchy wavelet transform -- meant to be used internally""" 

197 if self.kstep != self.__kstep or self.nfft != self.__nfft: 

198 self.make_karrays() 

199 nkpts = len(chi) 

200 nrpts = int(np.round(self.rmax/self.rstep)) 

201 if self.kwin is None: 

202 self.make_cwt_arrays(nkpts, nrpts) 

203 

204 omega = pi*np.arange(self.nfft)/(self.kstep*self.nfft) 

205 

206 if kweight is None: 

207 kweight = self.get_kweight() 

208 if kweight != 0: 

209 chi = chi * self.kwin[:len(chi)] * self.k_[:len(chi)]**kweight 

210 

211 if rmax is not None: 

212 self.rmax = rmax 

213 

214 chix = np.zeros(int(self.nfft/2)) * self.kstep 

215 chix[:nkpts] = chi 

216 chix = chix[:int(self.nfft/2)] 

217 _ffchi = np.fft.fft(chix, n=2*self.nfft)[:self.nfft] 

218 

219 nrpts = int(np.round(self.rmax/self.rstep)) 

220 r = self.rstep * arange(nrpts) 

221 r[0] = 1.e-19 

222 alpha = nrpts/(2*r) 

223 

224 self.make_cwt_arrays(nkpts, nrpts) 

225 

226 cauchy_sum = np.log(2*pi) - np.log(1.0+np.arange(nrpts)).sum() 

227 

228 out = np.zeros(nrpts*nkpts, dtype='complex128').reshape(nrpts, nkpts) 

229 

230 for i in range(nrpts): 

231 aom = alpha[i]*omega 

232 filt = cauchy_sum + nrpts*np.log(aom) - aom 

233 out[i, :] = np.fft.ifft(np.exp(filt)*_ffchi, 2*self.nfft)[:nkpts] 

234 

235 return (out*self._cauchymask)[self._cauchyslice] 

236 

237class FeffitDataSet(Group): 

238 def __init__(self, data=None, paths=None, transform=None, epsilon_k=None, 

239 refine_bkg=False, model=None, pathlist=None, _larch=None, **kws): 

240 

241 self._larch = _larch 

242 Group.__init__(self, **kws) 

243 self.refine_bkg = refine_bkg 

244 if paths is None and pathlist is not None: 

245 paths = pathlist 

246 

247 if isinstance(paths, dict): 

248 self.paths = {key: copy(path) for key, path in paths.items()} 

249 elif isinstance(paths, (list, tuple)): 

250 self.paths = {path.label: copy(path) for path in paths} 

251 else: 

252 self.paths = {} 

253 self.pathlist = list(self.paths.values()) 

254 

255 if transform is None: 

256 transform = TransformGroup() 

257 else: 

258 trasform = copy(transform) 

259 self.transform = transform 

260 

261 if model is None: 

262 self.model = Group(__name__='Feffit Model for %s' % repr(data)) 

263 self.model.k = None 

264 else: 

265 self.model = model 

266 # make datagroup from passed in data: copy of k, chi, delta_chi, epsilon_k 

267 self.set_datagroup(data, epsilon_k=epsilon_k, refine_bkg=refine_bkg) 

268 

269 def set_datagroup(self, data, epsilon_k=None, refine_bkg=False): 

270 "set the data group for the Dataset" 

271 if data is None: 

272 self.data = Group(__name__='Feffit DatasSet (no data)', 

273 groupname=None,filename=None, 

274 k=np.arange(401)/20.0, chi=np.zeros(401)) 

275 self.has_data = False 

276 else: 

277 self.data = Group(__name__='Feffit DatasSet from %s' % repr(data), 

278 groupname=getattr(data, 'groupname', repr(data)), 

279 filename=getattr(data, 'filename', repr(data)), 

280 k=data.k[:]*1.0, chi=data.chi[:]*1.0) 

281 self.has_data = True 

282 if hasattr(data, 'config'): 

283 self.data.config = deepcopy(data.config) 

284 else: 

285 self.data.config = Group() 

286 self.data.epsilon_k = getattr(data, 'epsilon_k', epsilon_k) 

287 if epsilon_k is not None: 

288 self.data.epsilon_k = epsilon_k 

289 

290 dat_attrs = ['delta_chi', 'r', 'chir_mag', 'chir_re', 'chir_im'] 

291 if data is not None: 

292 dat_attrs.extend(dir(data)) 

293 for attr in dat_attrs: 

294 if attr not in ('feffit_history',) and not hasattr(self.data, attr): 

295 setattr(self.data, attr, getattr(data, attr, None)) 

296 self.hashkey = None 

297 self.bkg_spline = {} 

298 self._chi = None 

299 self._bkg = 0.0 

300 self._prepared = False 

301 

302 

303 def __generate_hashkey(self, other_hashkeys=None): 

304 """generate hash for dataset""" 

305 if self.hashkey is not None: 

306 return 

307 hlen = 7 

308 dat = [] 

309 for aname in ('e0', 'ek0', 'rbkg', 'edge_step'): 

310 dat.append(getattr(self.data, aname, 0.00)) 

311 

312 for aname in ('energy', 'norm', 'chi'): 

313 arr = getattr(self.data, aname, None) 

314 if isinstance(arr, np.ndarray): 

315 dat.extend([arr.min(), arr.max(), arr.mean(), 

316 (arr**2).mean(), len(arr)]) 

317 dat.extend(arr[:30:3]) 

318 s = "|".join([gformat(x) for x in dat]) 

319 self.hashkey = f"d{(b32hash(s)[:hlen].lower())}" 

320 # may need to look for hash collisions: hlen=6 gives 1e9 keys 

321 # collision are probably the same dataset, so just go with a 

322 # random string 

323 if other_hashkeys is not None: 

324 ntry = 0 

325 while self.hashkey in other_hashkeys: 

326 ntry += 1 

327 if ntry > 1e6: 

328 ntry = 0 

329 hlen += 1 

330 self.hashkey = f"d{random_varname(hlen)}" 

331 

332 def __repr__(self): 

333 return '<FeffitDataSet Group: %s>' % self.__name__ 

334 

335 def __copy__(self): 

336 return FeffitDataSet(data=copy(self.data), 

337 paths=self.paths, 

338 transform=self.transform, 

339 refine_bkg=self.refine_bkg, 

340 epsilon_k=self.data.epsilon_k, 

341 model=self.model, 

342 _larch=self._larch) 

343 

344 def __deepcopy__(self, memo): 

345 return FeffitDataSet(data=deepcopy(self.data), 

346 paths=self.paths, 

347 transform=self.transform, 

348 refine_bkg=self.refine_bkg, 

349 epsilon_k=self.data.epsilon_k, 

350 model=self.model, 

351 _larch=self._larch) 

352 

353 def prepare_fit(self, params, other_hashkeys=None): 

354 """prepare for fit with this dataset""" 

355 trans = self.transform 

356 trans.make_karrays() 

357 ikmax = int(1.01 + max(self.data.k)/trans.kstep) 

358 

359 # ikmax = index_of(trans.k_, max(self.data.k)) 

360 self.model.k = trans.k_[:ikmax] 

361 self.model.chi = np.zeros(len(self.model.k), dtype='float64') 

362 self._chi = interp(self.model.k, self.data.k, self.data.chi) 

363 self.n_idp = 1 + 2*(trans.rmax-trans.rmin)*(trans.kmax-trans.kmin)/pi 

364 

365 if getattr(self.data, 'epsilon_k', None) is not None: 

366 eps_k = self.data.epsilon_k 

367 if isinstance(eps_k, np.ndarray): 

368 eps_k = interp(self.model.k, self.data.k, self.data.epsilon_k) 

369 self.set_epsilon_k(eps_k) 

370 else: 

371 self.estimate_noise(chi=self._chi, rmin=15.0, rmax=30.0) 

372 

373 # if delta_chi (uncertainty in chi(k) from autobk or other source) 

374 # exists, add it in quadrature to high-k noise estimate, and 

375 # update epsilon_k to be this value 

376 autobk_delta_chi(self.data) 

377 if hasattr(self.data, 'delta_chi'): 

378 cur_eps_k = getattr(self, 'epsilon_k', 0.0) 

379 if isinstance(cur_eps_k, (list, tuple)): 

380 eps_ave = 0. 

381 for eps in cur_eps_k: 

382 eps_ave += eps 

383 cur_eps_k = eps_ave/len(cur_eps_k) 

384 

385 _dchi = self.data.delta_chi 

386 if _dchi is not None: 

387 if isinstance(_dchi, np.ndarray): 

388 nchi = len(self.data.k) 

389 if len(_dchi) != nchi: 

390 _dchi = np.concatenate((_dchi, np.zeros(nchi)))[:nchi] 

391 _dchi = interp(self.model.k, self.data.k, _dchi) 

392 self.set_epsilon_k(np.sqrt(_dchi**2 + cur_eps_k**2)) 

393 

394 self.__generate_hashkey(other_hashkeys=other_hashkeys) 

395 # for each path in the list of paths, setup the Path Parameters 

396 # to use the current Parameters namespace 

397 if isinstance(params, Group): 

398 params = group2params(params) 

399 for label, path in self.paths.items(): 

400 path.create_path_params(params=params, dataset=self.hashkey) 

401 if path.spline_coefs is None: 

402 path.create_spline_coefs() 

403 

404 self.bkg_spline = {} 

405 if self.refine_bkg: 

406 trans.rbkg = max(trans.rbkg, trans.rmin) 

407 trans.rmin = trans.rstep 

408 self.n_idp = 1 + 2*(trans.rmax)*(trans.kmax-trans.kmin)/pi 

409 nspline = 1 + round(2*(trans.rbkg)*(trans.kmax-trans.kmin)/pi) 

410 knots_k = np.linspace(trans.kmin, trans.kmax, nspline) 

411 # np.linspace(trans.kmax, trans.kmax+trans.kstep/10.0, 3))) 

412 knots_y = np.linspace(-1e-4, 1.e-4, nspline) 

413 knots, coefs, order = splrep(knots_k, knots_y, k=3) 

414 self.bkg_spline = {'knots': knots, 'coefs': coefs, 

415 'nspline': nspline, 'order':order} 

416 for i in range(nspline): 

417 params.add(f'bkg{i:02d}_{self.hashkey}', value=0, vary=True) 

418 self._prepared = True 

419 

420 

421 def estimate_noise(self, chi=None, rmin=15.0, rmax=30.0, all_kweights=True): 

422 """estimage noise in a chi spectrum from its high r components""" 

423 trans = self.transform 

424 trans.make_karrays() 

425 if chi is None: 

426 chi = self.data.chi 

427 

428 save = trans.rmin, trans.rmax, trans.fitspace 

429 

430 all_kweights = all_kweights and isinstance(trans.kweight, Iterable) 

431 if all_kweights: 

432 chir = [trans.fftf(chi, kweight=kw) for kw in trans.kweight] 

433 else: 

434 chir = [trans.fftf(chi)] 

435 irmin = int(0.01 + rmin/trans.rstep) 

436 irmax = int(min(trans.nfft/2, 1.01 + rmax/trans.rstep)) 

437 highr = [realimag(chir_[irmin:irmax]) for chir_ in chir] 

438 # get average of window function value, we will scale eps_r scale by this 

439 kwin_ave = trans.kwin.sum()*trans.kstep/(trans.kmax-trans.kmin) 

440 eps_r = [(sqrt((chi*chi).sum() / len(chi)) / kwin_ave) for chi in highr] 

441 eps_k = [] 

442 # use Parseval's theorem to convert epsilon_r to epsilon_k, 

443 # compensating for kweight 

444 if all_kweights: 

445 kweights = trans.kweight[:] 

446 else: 

447 kweights = [trans.kweight] 

448 

449 for i, kw in enumerate(kweights): 

450 w = 2 * kw + 1 

451 scale = sqrt((2*pi*w)/(trans.kstep*(trans.kmax**w - trans.kmin**w))) 

452 eps_k.append(scale*eps_r[i]) 

453 

454 trans.rmin, trans.rmax, trans.fitspace = save 

455 

456 ## self.n_idp = 1 + 2*(trans.rmax-trans.rmin)*(trans.kmax-trans.kmin)/pi 

457 self.epsilon_k = eps_k 

458 self.epsilon_r = eps_r 

459 if len(eps_r) == 1: 

460 self.epsilon_k = eps_k[0] 

461 self.epsilon_r = eps_r[0] 

462 if isinstance(eps_r, np.ndarray): 

463 self.epsilon_r = eps_r.mean() 

464 

465 def set_epsilon_k(self, eps_k): 

466 """set epsilon_k and epsilon_r -- ucertainties in chi(k) and chi(R)""" 

467 eps_k = remove_nans(eps_k, 0.001) 

468 trans = self.transform 

469 all_kweights = isinstance(trans.kweight, Iterable) 

470 if isinstance(trans.kweight, Iterable): 

471 self.epsilon_k = [] 

472 self.epsilon_r = [] 

473 for kw in trans.kweight: 

474 w = 2 * kw + 1 

475 scale = 2*sqrt((pi*w)/(trans.kstep*(trans.kmax**w - trans.kmin**w))) 

476 self.epsilon_k.append(eps_k) 

477 eps_r = eps_k / scale 

478 if isinstance(eps_r, np.ndarray): eps_r = eps_r.mean() 

479 self.epsilon_r.append(eps_r) 

480 

481 else: 

482 w = 2 * trans.get_kweight() + 1 

483 scale = 2*sqrt((pi*w)/(trans.kstep*(trans.kmax**w - trans.kmin**w))) 

484 self.epsilon_k = eps_k 

485 eps_r = eps_k / scale 

486 if isinstance(eps_r, np.ndarray): eps_r = eps_r.mean() 

487 self.epsilon_r = eps_r 

488 

489 # check for nans 

490 self.epsilon_k = remove_nans(self.epsilon_k, eps_k, default=0.0001) 

491 self.epsilon_r = remove_nans(self.epsilon_r, eps_r, default=0.0001) 

492 

493 

494 

495 def _residual(self, params, data_only=False, **kws): 

496 """return the residual for this data set 

497 residual = self.transform.apply(data_chi - model_chi) 

498 where model_chi is the result of ff2chi(paths) 

499 """ 

500 if not isNamedClass(self.transform, TransformGroup): 

501 return 

502 if not self._prepared: 

503 self.prepare_fit(params) 

504 

505 ff2chi(self.paths, params=params, k=self.model.k, 

506 _larch=self._larch, group=self.model) 

507 

508 self._bkg = 0.0 

509 if self.refine_bkg: 

510 knots = self.bkg_spline['knots'] 

511 order = self.bkg_spline['order'] 

512 nspline = self.bkg_spline['nspline'] 

513 coefs = [] 

514 for i in range(nspline): 

515 parname = f'bkg{i:02d}_{self.hashkey}' 

516 par = params[parname] 

517 coefs.append(par.value) 

518 self._bkg = splev(self.model.k, [knots, coefs, order]) 

519 

520 eps_k = self.epsilon_k 

521 if isinstance(eps_k, np.ndarray): 

522 eps_k[np.where(eps_k<1.e-12)[0]] = 1.e-12 

523 

524 diff = self._chi - self._bkg 

525 if not data_only: # data_only for extracting transformed data 

526 diff -= self.model.chi 

527 trans = self.transform 

528 k = trans.k_[:len(diff)] 

529 

530 all_kweights = isinstance(trans.kweight, Iterable) 

531 if trans.fitspace == 'k': 

532 iqmin = int(max(0, 0.01 + trans.kmin/trans.kstep)) 

533 iqmax = int(min(trans.nfft/2, 0.01 + trans.kmax/trans.kstep)) 

534 if all_kweights: 

535 out = [] 

536 for i, kw in enumerate(trans.kweight): 

537 out.append(((diff/eps_k[i])*k**kw)[iqmin:iqmax]) 

538 return np.concatenate(out) 

539 else: 

540 return ((diff/eps_k) * k**trans.kweight)[iqmin:iqmax] 

541 elif trans.fitspace == 'w': 

542 if all_kweights: 

543 out = [] 

544 for i, kw in enumerate(trans.kweight): 

545 cwt = trans.cwt(diff/eps_k, kweight=kw) 

546 out.append(realimag(cwt).ravel()) 

547 return np.concatenate(out) 

548 else: 

549 cwt = trans.cwt(diff/eps_k, kweight=trans.kweight) 

550 return realimag(cwt).ravel() 

551 else: # 'r' space 

552 out = [] 

553 if all_kweights: 

554 chir = [trans.fftf(diff, kweight=kw) for kw in trans.kweight] 

555 eps_r = self.epsilon_r 

556 else: 

557 chir = [trans.fftf(diff)] 

558 eps_r = [self.epsilon_r] 

559 if trans.fitspace == 'r': 

560 irmin = int(max(0, 0.01 + trans.rmin/trans.rstep)) 

561 irmax = int(min(trans.nfft/2, 0.01 + trans.rmax/trans.rstep)) 

562 for i, chir_ in enumerate(chir): 

563 chir_ = chir_ / (eps_r[i]) 

564 out.append(realimag(chir_[irmin:irmax])) 

565 else: 

566 chiq = [trans.fftr(c)/eps for c, eps in zip(chir, eps_r)] 

567 iqmin = int(max(0, 0.01 + trans.kmin/trans.kstep)) 

568 iqmax = int(min(trans.nfft/2, 0.01 + trans.kmax/trans.kstep)) 

569 for chiq_ in chiq: 

570 out.append( realimag(chiq_[iqmin:iqmax])[::2]) 

571 return np.concatenate(out) 

572 

573 def save_outputs(self, rmax_out=10, path_outputs=True): 

574 "save fft outputs, and may also map a refined _bkg to the data chi(k) arrays" 

575 def xft(dgroup): 

576 self.transform._xafsft(dgroup.chi, group=dgroup, rmax_out=rmax_out) 

577 xft(self.data) 

578 xft(self.model) 

579 if self.refine_bkg: 

580 self.data.bkgk = IUSpline(self.model.k, self._bkg)(self.data.k) 

581 gname = getattr(self.data, 'groupname', repr(self.data)) 

582 fname = getattr(self.data, 'filename', repr(self.data)) 

583 label = f"defined background data for '{gname}'" 

584 self.data_rebkg = Group(__name__=label, groupname=gname, 

585 filename=fname, k=self.data.k[:], 

586 chi=self.data.chi-self.data.bkgk) 

587 xft(self.data_rebkg) 

588 

589 if path_outputs: 

590 for path in self.paths.values(): 

591 xft(path) 

592 

593def feffit_dataset(data=None, paths=None, transform=None, refine_bkg=False, 

594 epsilon_k=None, pathlist=None, _larch=None): 

595 """create a Feffit Dataset group. 

596 

597 Parameters: 

598 ------------ 

599 data: group containing experimental EXAFS (needs arrays 'k' and 'chi'). 

600 paths: dict of {label: FeffPathGroup}, using FeffPathGroup created by feffpath() 

601 transform: Feffit Transform group. 

602 epsilon_k: Uncertainty in data (either single value or array of 

603 same length as data.k) 

604 

605 Returns: 

606 ---------- 

607 a Feffit Dataset group. 

608 

609 """ 

610 return FeffitDataSet(data=data, paths=paths, transform=transform, epsilon_k=epsilon_k, 

611 refine_bkg=refine_bkg, pathlist=pathlist, _larch=_larch) 

612 

613def feffit_transform(_larch=None, **kws): 

614 """create a feffit transform group 

615 

616 Parameters: 

617 -------------- 

618 fitspace: name of FT type for fit ('r'). 

619 kmin: starting *k* for FT Window (0). 

620 kmax: ending *k* for FT Window (20). 

621 dk: tapering parameter for FT Window (4). 

622 dk2: second tapering parameter for FT Window (None). 

623 window: name of window type ('kaiser'). 

624 nfft: value to use for N_fft (2048). 

625 kstep: value to use for delta_k (0.05). 

626 kweight: exponent for weighting spectra by k^kweight (2). 

627 rmin: starting *R* for Fit Range and/or reverse FT Window (0). 

628 rmax: ending *R* for Fit Range and/or reverse FT Window (10). 

629 dr: tapering parameter for reverse FT Window 0. 

630 rwindow: name of window type for reverse FT Window ('kaiser'). 

631 

632 Returns: 

633 ---------- 

634 a feffit transform group. 

635 

636 """ 

637 return TransformGroup(_larch=_larch, **kws) 

638 

639def _feffit_resid(params, datasets=None, **kwargs): 

640 """ this is the residual function for feffit""" 

641 return concatenate([d._residual(params) for d in datasets]) 

642 

643def feffit(paramgroup, datasets, rmax_out=10, path_outputs=True, 

644 fix_unused_variables=True, _larch=None, **kws): 

645 """execute a Feffit fit: a fit of feff paths to a list of datasets 

646 

647 Parameters: 

648 ------------ 

649 paramgroup: group containing parameters for fit 

650 datasets: Feffit Dataset group or list of Feffit Dataset group. 

651 rmax_out: maximum R value to calculate output arrays. 

652 path_output: Flag to set whether all Path outputs should be written. 

653 fix_unused_variables: Flag for whether to set `vary=False` for unused 

654 variable parameters. Otherwise, a warning will be printed. 

655 Returns: 

656 --------- 

657 a fit results group. This will contain subgroups of: 

658 

659 datasets: an array of FeffitDataSet groups used in the fit. 

660 params: This will be identical to the input parameter group. 

661 fit: an object which points to the low-level fit. 

662 

663 Statistical parameters will be put into the params group. Each 

664 dataset will have a 'data' and 'model' subgroup, each with arrays: 

665 k wavenumber array of k 

666 chi chi(k). 

667 kwin window Omega(k) (length of input chi(k)). 

668 r uniform array of R, out to rmax_out. 

669 chir complex array of chi(R). 

670 chir_mag magnitude of chi(R). 

671 chir_pha phase of chi(R). 

672 chir_re real part of chi(R). 

673 chir_im imaginary part of chi(R). 

674 """ 

675 fit_kws = dict(gtol=1.e-6, ftol=1.e-6, xtol=1.e-6, epsfcn=1.e-10) 

676 if 'tol' in kws: 

677 tol = kws.pop('tol') 

678 fit_kws['gtol'] = fit_kws['ftol'] = fit_kws['xtol'] = tol 

679 

680 fit_kws.update(kws) 

681 

682 work_paramgroup = deepcopy(paramgroup) 

683 for pname in dir(paramgroup): # explicitly copy 'skip'! 

684 wpar = getattr(work_paramgroup, pname) 

685 opar = getattr(paramgroup, pname) 

686 if isinstance(wpar, Parameter): 

687 setattr(wpar, 'skip', getattr(opar, 'skip', False)) 

688 

689 params = group2params(work_paramgroup) 

690 

691 if isNamedClass(datasets, FeffitDataSet): 

692 datasets = [datasets] 

693 

694 # we need unique dataset hashes if refine_bkg is used 

695 dset_hashkeys = [] 

696 for ds in datasets: 

697 if not isNamedClass(ds, FeffitDataSet): 

698 print( "feffit needs a list of FeffitDataSets") 

699 return 

700 ds.prepare_fit(params=params, other_hashkeys=dset_hashkeys) 

701 dset_hashkeys.append(ds.hashkey) 

702 # try to identify variable Parameters that are not actually used 

703 vars, exprs = [], [] 

704 for p in params.values(): 

705 if p.vary: 

706 nam = p.name 

707 if not any([nam.startswith('bkg') and 

708 nam.endswith(ds.hashkey) for ds in datasets]): 

709 vars.append(nam) 

710 elif p.expr is not None: 

711 exprs.append(p.expr) 

712 

713 for expr in exprs: 

714 for node in ast.walk(ast.parse(expr)): 

715 if isinstance(node, ast.Name): 

716 if node.id in vars: 

717 vars.remove(node.id) 

718 if len(vars) > 0: 

719 if fix_unused_variables: 

720 for v in vars: 

721 params[v].vary = False 

722 else: 

723 vlist = ', '.join(vars) 

724 print(f"Feffit Warning: unused variables: {vlist}") 

725 

726 # run fit 

727 fit = Minimizer(_feffit_resid, params, fcn_kws=dict(datasets=datasets), 

728 scale_covar=False, **fit_kws) 

729 

730 result = fit.leastsq() 

731 dat = concatenate([d._residual(result.params, data_only=True) 

732 for d in datasets]) 

733 

734 n_idp = 0 

735 for ds in datasets: 

736 n_idp += ds.n_idp 

737 

738 # here we rescale chi-square and reduced chi-square to n_idp 

739 npts = len(result.residual) 

740 chi_square = result.chisqr * n_idp*1.0 / npts 

741 chi2_reduced = chi_square/(n_idp*1.0 - result.nvarys) 

742 rfactor = (result.residual**2).sum() / (dat**2).sum() 

743 # calculate 'aic', 'bic' rescaled to n_idp 

744 # note that neg2_loglikel is -2*log(likelihood) 

745 neg2_loglikel = n_idp * np.log(chi_square / n_idp) 

746 aic = neg2_loglikel + 2 * result.nvarys 

747 bic = neg2_loglikel + np.log(n_idp) * result.nvarys 

748 

749 # We used scale_covar=False, so we rescale the uncertainties 

750 # by reduced chi-square * (ndata - nvarys)/(nidp - nvarys) 

751 covar = getattr(result, 'covar', None) 

752 if covar is not None: 

753 err_scale = result.redchi*(result.nfree/(n_idp - result.nvarys)) 

754 for name, par in result.params.items(): 

755 if isParameter(par) and getattr(par, 'stderr', None) is not None: 

756 par.stderr *= sqrt(err_scale) 

757 

758 # next, propagate uncertainties to constraints and path parameters. 

759 result.covar *= err_scale 

760 vsave, vbest = {}, [] 

761 

762 # 1. save current params 

763 for vname in result.var_names: 

764 par = result.params[vname] 

765 vsave[vname] = par 

766 vbest.append(par.value) 

767 # 2. get correlated uncertainties, set params accordingly 

768 uvars = correlated_values(vbest, result.covar) 

769 # 3. evaluate constrained params, save stderr 

770 for nam, obj in result.params.items(): 

771 eval_stderr(obj, uvars, result.var_names, result.params) 

772 

773 # 3. evaluate path_ params, save stderr 

774 for ds in datasets: 

775 for label, path in ds.paths.items(): 

776 path.store_feffdat() 

777 for pname in ('degen', 's02', 'e0', 'ei', 

778 'deltar', 'sigma2', 'third', 'fourth'): 

779 obj = path.params[path.pathpar_name(pname)] 

780 eval_stderr(obj, uvars, result.var_names, result.params) 

781 # restore saved parameters again 

782 for vname in result.var_names: 

783 # setattr(params, vname, vsave[vname]) 

784 params[vname] = vsave[vname] 

785 

786 # clear any errors evaluting uncertainties 

787 if _larch is not None and (len(_larch.error) > 0): 

788 _larch.error = [] 

789 

790 # reset the parameters group with the newly updated uncertainties 

791 params2group(result.params, work_paramgroup) 

792 

793 # here we create outputs arrays for chi(k), chi(r): 

794 for ds in datasets: 

795 ds.save_outputs(rmax_out=rmax_out, path_outputs=path_outputs) 

796 

797 out = Group(name='feffit results', params=result.params, 

798 paramgroup=work_paramgroup, fit_kws=fit_kws, datasets=datasets, 

799 fit_details=result, chi_square=chi_square, n_independent=n_idp, 

800 chi2_reduced=chi2_reduced, rfactor=rfactor, aic=aic, bic=bic, 

801 covar=covar) 

802 

803 for attr in ('params', 'nvarys', 'nfree', 'ndata', 'var_names', 'nfev', 

804 'success', 'errorbars', 'message', 'lmdif_message'): 

805 setattr(out, attr, getattr(result, attr, None)) 

806 return out 

807 

808def feffit_conf_map(result, xpar, ypar, nsamples=41, nsigma=3.5): 

809 """ 

810 return 2d map of confidence interval (sigma values) for a pair of variables from feffit 

811 

812 """ 

813 def show_progress(i, imax): 

814 if i > (imax-1): 

815 print('done.') 

816 elif i % round(imax//10) == 0: 

817 print(f"{i}/{imax}", flush=True, end=', ') 

818 

819 fitter = Minimizer(_feffit_resid, result.params, 

820 fcn_kws=dict(datasets=result.datasets), 

821 scale_covar=False, **result.fit_kws) 

822 

823 xvals, yvals, chi2_map = conf_interval2d(fitter, result.fit_details, xpar, ypar, 

824 nsamples, nsamples, nsigma=nsigma, chi2_out=True) 

825 

826 chi2_map = chi2_map * result.n_independent / result.fit_details.ndata 

827 chisqr0 = min(result.chi_square, chi2_map.min()) 

828 sigma_map = np.sqrt((chi2_map-chisqr0)/result.chi2_reduced) 

829 return xvals, yvals, sigma_map 

830 

831 

832 

833def feffit_report(result, min_correl=0.1, with_paths=True, _larch=None): 

834 """return a printable report of fit for feffit 

835 

836 Parameters: 

837 ------------ 

838 result: Feffit result, output group from feffit() 

839 min_correl: minimum correlation to report [0.1] 

840 wit_paths: boolean (True/False) for whether to list all paths [True] 

841 

842 Returns: 

843 --------- 

844 printable string of report. 

845 

846 """ 

847 input_ok = False 

848 try: 

849 params = result.params 

850 datasets = result.datasets 

851 input_ok = True 

852 except: 

853 pass 

854 if not input_ok: 

855 print( 'must pass output of feffit()!') 

856 return 

857 

858 path_hashkeys = [] 

859 for ds in datasets: 

860 path_hashkeys.extend([p.hashkey for p in ds.paths.values()]) 

861 topline = '=================== FEFFIT RESULTS ====================' 

862 header = '[[%s]]' 

863 varformat = ' %12s = %s +/-%s (init= %s)' 

864 fixformat = ' %12s = %s (fixed)' 

865 exprformat = ' %12s = %s +/-%s = \'%s\'' 

866 out = [topline, header % 'Statistics'] 

867 

868 def getval(attr): 

869 return getfloat_attr(result, attr) 

870 

871 def add_string(label, value, llen=20): 

872 if len(label) < llen: 

873 label = (label + ' '*llen)[:llen] 

874 out.append(f" {label} = {value}") 

875 

876 add_string('n_function_calls', getval('nfev')) 

877 add_string('n_variables', getval('nvarys')) 

878 add_string('n_data_points', getval('ndata')) 

879 add_string('n_independent', getval('n_independent')) 

880 add_string('chi_square', getval('chi_square')) 

881 add_string('reduced chi_square', getval('chi2_reduced')) 

882 add_string('r-factor', getval('rfactor')) 

883 add_string('Akaike info crit', getval('aic')) 

884 add_string('Bayesian info crit', getval('bic')) 

885 

886 out.append(' ') 

887 out.append(header % 'Variables') 

888 for name, par in params.items(): 

889 if any([name.endswith('_%s' % phash) for phash in path_hashkeys]): 

890 continue 

891 if isParameter(par): 

892 if par.vary: 

893 stderr = 'unknown' 

894 if par.stderr is not None: 

895 stderr = gfmt(par.stderr) 

896 add_string(name, f"{gfmt(par.value)} +/-{stderr} (init={gfmt(par.init_value)})") 

897 

898 elif par.expr is not None: 

899 stderr = 'unknown' 

900 if par.stderr is not None: 

901 stderr = gfmt(par.stderr) 

902 add_string(name, f"{gfmt(par.value)} +/-{stderr} = '{par.expr}'") 

903 else: 

904 add_string(name, f"{gfmt(par.value)} (fixed)") 

905 

906 covar_vars = result.var_names 

907 if len(covar_vars) > 0: 

908 out.append(' ') 

909 out.append(header % 'Correlations' + 

910 ' (unreported correlations are < % .3f)' % min_correl) 

911 correls = {} 

912 for i, name in enumerate(covar_vars): 

913 par = params[name] 

914 if not par.vary: 

915 continue 

916 if hasattr(par, 'correl') and par.correl is not None: 

917 for name2 in covar_vars[i+1:]: 

918 if name != name2 and name2 in par.correl: 

919 correls["%s, %s" % (name, name2)] = par.correl[name2] 

920 

921 sort_correl = sorted(correls.items(), key=lambda it: abs(it[1])) 

922 sort_correl.reverse() 

923 for name, val in sort_correl: 

924 if abs(val) > min_correl: 

925 vv = f"{val:+.3f}".replace('+', ' ') 

926 add_string(name, vv) 

927 

928 out.append(' ') 

929 for i, ds in enumerate(datasets): 

930 if not hasattr(ds, 'epsilon_k'): 

931 ds.prepare_fit(params) 

932 tr = ds.transform 

933 if isinstance(tr.kweight, Iterable): 

934 if isinstance(ds.epsilon_k[0], np.ndarray): 

935 msg = [] 

936 for eps in ds.epsilon_k: 

937 msg.append('Array(mean=%s, std=%s)' % (gfmt(eps.mean()).strip(), 

938 gfmt(eps.std()).strip())) 

939 eps_k = ', '.join(msg) 

940 else: 

941 eps_k = ', '.join([gfmt(eps).strip() for eps in ds.epsilon_k]) 

942 eps_r = ', '.join([gfmt(eps).strip() for eps in ds.epsilon_r]) 

943 kweigh = ', '.join(['%i' % kwe for kwe in tr.kweight]) 

944 eps_k = eps_k.strip() 

945 eps_r = eps_r.strip() 

946 kweigh = kweigh.strip() 

947 else: 

948 if isinstance(ds.epsilon_k, np.ndarray): 

949 eps_k = 'Array(mean=%s, std=%s)' % (gfmt(ds.epsilon_k.mean()).strip(), 

950 gfmt(ds.epsilon_k.std()).strip()) 

951 else: 

952 eps_k = gfmt(ds.epsilon_k).strip() 

953 eps_r = gfmt(ds.epsilon_r).strip() 

954 kweigh = '%i' % tr.kweight 

955 extra = f" {i+1} of {len(datasets)}" if len(datasets) > 1 else "" 

956 

957 out.append(f"[[Dataset{extra}]]") 

958 add_string('unique_id', f"'{ds.hashkey}'") 

959 add_string('fit space', f"'{tr.fitspace}'") 

960 if ds.refine_bkg: 

961 add_string('r_bkg (refine bkg)', f"{tr.rbkg:.3f}") 

962 add_string('r-range', f"{tr.rmin:.3f}, {tr.rmax:.3f}") 

963 add_string('k-range', f"{tr.kmin:.3f}, {tr.kmax:.3f}") 

964 kwin = f"'{tr.window}', {tr.dk:.3f}" 

965 if tr.dk2 is not None: 

966 kwin += f", {tr.dk2:.3f}" 

967 add_string('k window, dk', kwin) 

968 pathfiles = repr([p.filename for p in ds.paths.values()]) 

969 add_string('paths used in fit', pathfiles) 

970 add_string('k-weight', kweigh) 

971 add_string('epsilon_k', eps_k) 

972 add_string('epsilon_r', eps_r) 

973 add_string('n_independent', f"{ds.n_idp:.3f}") 

974 # 

975 

976 if with_paths: 

977 out.append(' ') 

978 out.append(header % 'Paths') 

979 for label, path in ds.paths.items(): 

980 out.append('%s\n' % path.report()) 

981 out.append('='*len(topline)) 

982 return '\n'.join(out)