Coverage for /Users/Newville/Codes/xraylarch/larch/xafs/feffit.py: 79%

563 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-11-09 10:08 -0600

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.optimize import leastsq as scipy_leastsq 

16 

17from lmfit import Parameters, Parameter, Minimizer, fit_report 

18 

19from larch import Group, isNamedClass 

20from larch.utils import fix_varname, gformat 

21from larch.utils.strutils import b32hash, random_varname 

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

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

24 group2params, params2group, isParameter) 

25 

26from .xafsutils import set_xafsGroup, gfmt 

27from .xafsft import xftf_fast, xftr_fast, ftwindow 

28from .feffdat import FeffPathGroup, ff2chi 

29 

30class TransformGroup(Group): 

31 """A Group of transform parameters. 

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

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

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

35 recalculated.... 

36 

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

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

39 recalculated. 

40 """ 

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

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

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

44 fitspace='r', wavelet_mask=None, _larch=None, **kws): 

45 Group.__init__(self, **kws) 

46 self.kmin = kmin 

47 self.kmax = kmax 

48 self.kweight = kweight 

49 if 'kw' in kws: 

50 self.kweight = kws['kw'] 

51 

52 self.dk = dk 

53 self.dk2 = dk2 

54 self.window = window 

55 self.rmin = rmin 

56 self.rmax = rmax 

57 self.dr = dr 

58 self.dr2 = dr2 

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

60 self.rwindow = rwindow 

61 self.__nfft = 0 

62 self.__kstep = None 

63 self.nfft = nfft 

64 self.kstep = kstep 

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

66 

67 self.fitspace = fitspace 

68 self.wavelet_mask = wavelet_mask 

69 self._cauchymask = None 

70 

71 self._larch = _larch 

72 

73 self.kwin = None 

74 self.rwin = None 

75 self.make_karrays() 

76 

77 def __repr__(self): 

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

79 

80 def __copy__(self): 

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

82 kweight=self.kweight, dk=self.dk, dk2=self.dk2, 

83 window=self.window, kstep=self.kstep, 

84 rmin=self.rmin, rmax=self.rmax, 

85 dr=self.dr, dr2=self.dr2, 

86 rwindow=self.rwindow, nfft=self.nfft, 

87 fitspace=self.fitspace, 

88 wavelet_mask=self.wavelet_mask, 

89 _larch=self._larch) 

90 

91 def __deepcopy__(self, memo): 

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

93 kweight=self.kweight, dk=self.dk, dk2=self.dk2, 

94 window=self.window, kstep=self.kstep, 

95 rmin=self.rmin, rmax=self.rmax, 

96 dr=self.dr, dr2=self.dr2, 

97 rwindow=self.rwindow, nfft=self.nfft, 

98 fitspace=self.fitspace, 

99 wavelet_mask=self.wavelet_mask, 

100 _larch=self._larch) 

101 

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

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

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

105 return 

106 self.__kstep = self.kstep 

107 self.__nfft = self.nfft 

108 

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

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

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

112 

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

114 "returns " 

115 for key, val in kws: 

116 if key == 'kw': 

117 key = 'kweight' 

118 setattr(self, key, val) 

119 self.make_karrays() 

120 

121 out = self.fftf(chi) 

122 

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

124 

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

126 r = self.rstep * arange(irmax) 

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

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

129 group.r = r[:irmax] 

130 group.chir = out[:irmax] 

131 group.chir_mag = mag[:irmax] 

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

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

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

135 if self.rwin is None: 

136 self.rwin = ftwindow(self.r_, xmin=self.rmin, xmax=self.rmax, 

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

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

139 

140 def get_kweight(self): 

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

142 if isinstance(self.kweight, Iterable): 

143 return self.kweight[0] 

144 return self.kweight 

145 

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

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

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

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

150 self.make_karrays() 

151 if self.kwin is None: 

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

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

154 if kweight is None: 

155 kweight = self.get_kweight() 

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

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

158 

159 def fftr(self, chir): 

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

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

162 self.make_karrays() 

163 if self.rwin is None: 

164 self.rwin = ftwindow(self.r_, xmin=self.rmin, xmax=self.rmax, 

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

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

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

168 

169 

170 def make_cwt_arrays(self, nkpts, nrpts): 

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

172 self.make_karrays() 

173 if self.kwin is None: 

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

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

176 

177 if self._cauchymask is None: 

178 if self.wavelet_mask is not None: 

179 self._cauchymask = self.wavelet_mask 

180 else: 

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

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

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

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

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

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

187 self._cauchymask = cm 

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

189 

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

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

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

193 self.make_karrays() 

194 nkpts = len(chi) 

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

196 if self.kwin is None: 

197 self.make_cwt_arrays(nkpts, nrpts) 

198 

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

200 

201 if kweight is None: 

202 kweight = self.get_kweight() 

203 if kweight != 0: 

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

205 

206 if rmax is not None: 

207 self.rmax = rmax 

208 

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

210 chix[:nkpts] = chi 

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

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

213 

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

215 r = self.rstep * arange(nrpts) 

216 r[0] = 1.e-19 

217 alpha = nrpts/(2*r) 

218 

219 self.make_cwt_arrays(nkpts, nrpts) 

220 

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

222 

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

224 

225 for i in range(nrpts): 

226 aom = alpha[i]*omega 

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

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

229 

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

231 

232class FeffitDataSet(Group): 

233 

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

235 refine_bkg=False, pathlist=None, model=None, _larch=None, 

236 **kws): 

237 

238 self._larch = _larch 

239 Group.__init__(self, **kws) 

240 self.refine_bkg = refine_bkg 

241 

242 if paths is None and pathlist is not None: # legacy 

243 paths = pathlist 

244 

245 if isinstance(paths, dict): 

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

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

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

249 else: 

250 self.paths = {} 

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

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

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

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

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

256 if hasattr(data, 'config'): 

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

258 else: 

259 self.data.config = Group() 

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

261 if epsilon_k is not None: 

262 self.data.epsilon_k = epsilon_k 

263 

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

265 if data is not None: 

266 dat_attrs.extend(dir(data)) 

267 for attr in dat_attrs: 

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

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

270 

271 if transform is None: 

272 transform = TransformGroup() 

273 else: 

274 trasform = copy(transform) 

275 self.transform = transform 

276 

277 if model is None: 

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

279 self.model.k = None 

280 else: 

281 self.model = model 

282 self.hashkey = None 

283 self.__chi = None 

284 self.__prepared = False 

285 

286 def _generate_hashkey(self, other_hashkeys=None): 

287 """generate hash for dataset""" 

288 if self.hashkey is not None: 

289 return 

290 hlen = 10 

291 dat = [self.data.ek0, self.data.e0, self.data.rbkg] 

292 for arr in (self.data.energy, self.data.norm, self.data.chi): 

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

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

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

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

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

298 # may need to look for hash collisions: if there is a collision, 

299 # it is probably the same dataset, so just go witha 

300 if other_hashkeys is not None: 

301 ntry = 0 

302 while self.hashkey in other_hashkeys: 

303 ntry += 1 

304 if ntry > 1e6: 

305 ntry = 0 

306 hlen += 1 

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

308 

309 def __repr__(self): 

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

311 

312 def __copy__(self): 

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

314 paths=self.paths, 

315 transform=self.transform, 

316 refine_bkg=self.refine_bkg, 

317 epsilon_k=self.data.epsilon_k, 

318 model=self.model, 

319 _larch=self._larch) 

320 

321 def __deepcopy__(self, memo): 

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

323 paths=self.paths, 

324 transform=self.transform, 

325 refine_bkg=self.refine_bkg, 

326 epsilon_k=self.data.epsilon_k, 

327 model=self.model, 

328 _larch=self._larch) 

329 

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

331 """prepare for fit with this dataset""" 

332 trans = self.transform 

333 trans.make_karrays() 

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

335 

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

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

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

339 self.__chi = interp(self.model.k, self.data.k, self.data.chi) 

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

341 

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

343 eps_k = self.data.epsilon_k 

344 if isinstance(eps_k, np.ndarray): 

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

346 self.set_epsilon_k(eps_k) 

347 else: 

348 self.estimate_noise(chi=self.__chi, rmin=15.0, rmax=30.0) 

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

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

351 # update epsilon_k to be this combination. 

352 if getattr(self.data, 'delta_chi', None) is not None: 

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

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

355 eps_ave = 0. 

356 for eps in cur_eps_k: 

357 eps_ave += eps 

358 cur_eps_k = eps_ave/len(cur_eps_k) 

359 

360 _dchi = self.data.delta_chi 

361 if isinstance(_dchi, np.ndarray): 

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

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

364 

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

366 # to use the current Parameters namespace 

367 if isinstance(params, Group): 

368 params = group2params(params) 

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

370 path.create_path_params(params=params) 

371 if path.spline_coefs is None: 

372 path.create_spline_coefs() 

373 

374 self._generate_hashkey(other_hashkeys=other_hashkeys) 

375 self.bkg_spline = {} 

376 if self.refine_bkg: 

377 self.n_spline = 1 + 2*(trans.rmin)*(trans.kmax-trans.kmin)/pi 

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

379 self.__prepared = True 

380 

381 

382 

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

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

385 trans = self.transform 

386 trans.make_karrays() 

387 if chi is None: 

388 chi = self.__chi 

389 

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

391 

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

393 if all_kweights: 

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

395 else: 

396 chir = [trans.fftf(chi)] 

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

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

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

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

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

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

403 eps_k = [] 

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

405 # compensating for kweight 

406 if all_kweights: 

407 kweights = trans.kweight[:] 

408 else: 

409 kweights = [trans.kweight] 

410 

411 for i, kw in enumerate(kweights): 

412 w = 2 * kw + 1 

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

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

415 

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

417 

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

419 self.epsilon_k = eps_k 

420 self.epsilon_r = eps_r 

421 if len(eps_r) == 1: 

422 self.epsilon_k = eps_k[0] 

423 self.epsilon_r = eps_r[0] 

424 if isinstance(eps_r, np.ndarray): 

425 self.epsilon_r = eps_r.mean() 

426 

427 def set_epsilon_k(self, eps_k): 

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

429 

430 eps_k = remove_nans(eps_k, 0.001) 

431 

432 trans = self.transform 

433 all_kweights = isinstance(trans.kweight, Iterable) 

434 if isinstance(trans.kweight, Iterable): 

435 self.epsilon_k = [] 

436 self.epsilon_r = [] 

437 for kw in trans.kweight: 

438 w = 2 * kw + 1 

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

440 self.epsilon_k.append(eps_k) 

441 eps_r = eps_k / scale 

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

443 self.epsilon_r.append(eps_r) 

444 

445 else: 

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

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

448 self.epsilon_k = eps_k 

449 eps_r = eps_k / scale 

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

451 self.epsilon_r = eps_r 

452 

453 # check for nans 

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

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

456 

457 

458 

459 def _residual(self, paramgroup, data_only=False, **kws): 

460 """return the residual for this data set 

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

462 where model_chi is the result of ff2chi(paths) 

463 """ 

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

465 return 

466 if not self.__prepared: 

467 self.prepare_fit(paramgroup) 

468 

469 ff2chi(self.paths, paramgroup=paramgroup, k=self.model.k, 

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

471 

472 eps_k = self.epsilon_k 

473 if isinstance(eps_k, np.ndarray): 

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

475 

476 diff = (self.__chi - self.model.chi) 

477 if data_only: # for extracting transformed data separately from residual 

478 diff = self.__chi 

479 trans = self.transform 

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

481 

482 all_kweights = isinstance(trans.kweight, Iterable) 

483 if trans.fitspace == 'k': 

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

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

486 if all_kweights: 

487 out = [] 

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

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

490 return np.concatenate(out) 

491 else: 

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

493 elif trans.fitspace == 'w': 

494 if all_kweights: 

495 out = [] 

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

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

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

499 return np.concatenate(out) 

500 else: 

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

502 return realimag(cwt).ravel() 

503 else: # 'r' space 

504 out = [] 

505 if all_kweights: 

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

507 eps_r = self.epsilon_r 

508 else: 

509 chir = [trans.fftf(diff)] 

510 eps_r = [self.epsilon_r] 

511 if trans.fitspace == 'r': 

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

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

514 for i, chir_ in enumerate(chir): 

515 chir_ = chir_ / (eps_r[i]) 

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

517 else: 

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

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

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

521 for chiq_ in chiq: 

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

523 return np.concatenate(out) 

524 

525 def save_ffts(self, rmax_out=10, path_outputs=True): 

526 "save fft outputs" 

527 xft = self.transform._xafsft 

528 xft(self.__chi, group=self.data, rmax_out=rmax_out) 

529 xft(self.model.chi, group=self.model, rmax_out=rmax_out) 

530 if path_outputs: 

531 for p in self.paths.values(): 

532 xft(p.chi, group=p, rmax_out=rmax_out) 

533 

534def feffit_dataset(data=None, paths=None, transform=None, 

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

536 """create a Feffit Dataset group. 

537 

538 Parameters: 

539 ------------ 

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

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

542 transform: Feffit Transform group. 

543 pathlist: list of FeffPathGroup [deprecated - use 'paths'] 

544 

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

546 same length as data.k) 

547 

548 Returns: 

549 ---------- 

550 a Feffit Dataset group. 

551 

552 

553 """ 

554 return FeffitDataSet(data=data, paths=paths, transform=transform, 

555 pathlist=pathlist, _larch=_larch) 

556 

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

558 """create a feffit transform group 

559 

560 Parameters: 

561 -------------- 

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

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

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

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

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

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

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

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

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

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

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

573 dr: tapering parameter for reverse FT Window 0. 

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

575 

576 Returns: 

577 ---------- 

578 a feffit transform group. 

579 

580 """ 

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

582 

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

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

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

586 

587 Parameters: 

588 ------------ 

589 paramgroup: group containing parameters for fit 

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

591 rmax_out: maximum R value to calculate output arrays. 

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

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

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

595 Returns: 

596 --------- 

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

598 

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

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

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

602 

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

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

605 k wavenumber array of k 

606 chi chi(k). 

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

608 r uniform array of R, out to rmax_out. 

609 chir complex array of chi(R). 

610 chir_mag magnitude of chi(R). 

611 chir_pha phase of chi(R). 

612 chir_re real part of chi(R). 

613 chir_im imaginary part of chi(R). 

614 """ 

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

616 if 'tol' in kws: 

617 tol = kws.pop('tol') 

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

619 

620 fit_kws.update(kws) 

621 

622 work_paramgroup = deepcopy(paramgroup) 

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

624 wpar = getattr(work_paramgroup, pname) 

625 opar = getattr(paramgroup, pname) 

626 if isinstance(wpar, Parameter): 

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

628 

629 

630 params = group2params(work_paramgroup) 

631 

632 def _resid(params, datasets=None, pargroup=None, **kwargs): 

633 """ this is the residual function""" 

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

635 

636 if isNamedClass(datasets, FeffitDataSet): 

637 datasets = [datasets] 

638 

639 # we need unique dataset hashes if refine_bkg is used 

640 dset_hashkeys = [] 

641 for ds in datasets: 

642 if not isNamedClass(ds, FeffitDataSet): 

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

644 return 

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

646 dset_hashkeys.append(ds.hashkey) 

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

648 vars, exprs = [], [] 

649 for p in params.values(): 

650 if p.vary: 

651 vars.append(p.name) 

652 elif p.expr is not None: 

653 exprs.append(p.expr) 

654 

655 for expr in exprs: 

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

657 if isinstance(node, ast.Name): 

658 if node.id in vars: 

659 vars.remove(node.id) 

660 if len(vars) > 0: 

661 if fix_unused_variables: 

662 for v in vars: 

663 params[v].vary = False 

664 else: 

665 vlist = ', '.join(vars) 

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

667 

668 # run fit 

669 fit = Minimizer(_resid, params, 

670 fcn_kws=dict(datasets=datasets, pargroup=work_paramgroup), 

671 scale_covar=False, **fit_kws) 

672 

673 result = fit.leastsq() 

674 params2group(result.params, work_paramgroup) 

675 dat = concatenate([d._residual(work_paramgroup, data_only=True) for d in datasets]) 

676 

677 n_idp = 0 

678 for ds in datasets: 

679 n_idp += ds.n_idp 

680 

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

682 npts = len(result.residual) 

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

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

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

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

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

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

689 aic = neg2_loglikel + 2 * result.nvarys 

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

691 

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

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

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

695 if covar is not None: 

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

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

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

699 par.stderr *= sqrt(err_scale) 

700 

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

702 result.covar *= err_scale 

703 vsave, vbest = {}, [] 

704 

705 # 1. save current params 

706 for vname in result.var_names: 

707 par = result.params[vname] 

708 vsave[vname] = par 

709 vbest.append(par.value) 

710 # 2. get correlated uncertainties, set params accordingly 

711 uvars = correlated_values(vbest, result.covar) 

712 # 3. evaluate constrained params, save stderr 

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

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

715 

716 # 3. evaluate path_ params, save stderr 

717 for ds in datasets: 

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

719 path.store_feffdat() 

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

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

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

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

724 # restore saved parameters again 

725 for vname in result.var_names: 

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

727 params[vname] = vsave[vname] 

728 

729 # clear any errors evaluting uncertainties 

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

731 _larch.error = [] 

732 

733 # reset the parameters group with the newly updated uncertainties 

734 params2group(result.params, work_paramgroup) 

735 

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

737 for ds in datasets: 

738 ds.save_ffts(rmax_out=rmax_out, path_outputs=path_outputs) 

739 

740 out = Group(name='feffit results', 

741 paramgroup=work_paramgroup, 

742 datasets=datasets, 

743 # fitter=fit, 

744 fit_details=result, chi_square=chi_square, 

745 n_independent=n_idp, chi2_reduced=chi2_reduced, 

746 rfactor=rfactor, aic=aic, bic=bic, covar=covar) 

747 

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

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

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

751 return out 

752 

753 

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

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

756 

757 Parameters: 

758 ------------ 

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

760 min_correl: minimum correlation to report [0.1] 

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

762 

763 Returns: 

764 --------- 

765 printable string of report. 

766 

767 """ 

768 input_ok = False 

769 try: 

770 params = result.params 

771 datasets = result.datasets 

772 input_ok = True 

773 except: 

774 pass 

775 if not input_ok: 

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

777 return 

778 

779 path_hashkeys = [] 

780 for ds in datasets: 

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

782 topline = '=================== FEFFIT RESULTS ====================' 

783 header = '[[%s]]' 

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

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

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

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

788 

789 out.append(' nvarys, npts = %i, %i' % (result.nvarys, 

790 result.ndata)) 

791 out.append(' n_independent = %.3f' % (result.n_independent)) 

792 out.append(' chi_square = %s' % gfmt(result.chi_square)) 

793 out.append(' reduced chi_square = %s' % gfmt(result.chi2_reduced)) 

794 out.append(' r-factor = %s' % gfmt(result.rfactor)) 

795 out.append(' Akaike info crit = %s' % gfmt(result.aic)) 

796 out.append(' Bayesian info crit = %s' % gfmt(result.bic)) 

797 out.append(' ') 

798 

799 if len(datasets) == 1: 

800 out.append(header % 'Data') 

801 else: 

802 out.append(header % 'Datasets (%i)' % len(datasets)) 

803 for i, ds in enumerate(datasets): 

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

805 ds.prepare_fit(params) 

806 tr = ds.transform 

807 if len(datasets) > 1: 

808 out.append(' dataset %i:' % (i+1)) 

809 if isinstance(tr.kweight, Iterable): 

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

811 msg = [] 

812 for eps in ds.epsilon_k: 

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

814 gfmt(eps.std()).strip())) 

815 eps_k = ', '.join(msg) 

816 else: 

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

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

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

820 else: 

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

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

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

824 else: 

825 eps_k = gfmt(ds.epsilon_k) 

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

827 kweigh = '%i' % tr.kweight 

828 out.append(' fit space = \'%s\'' % (tr.fitspace)) 

829 out.append(' r-range = %.3f, %.3f' % (tr.rmin, tr.rmax)) 

830 out.append(' k-range = %.3f, %.3f' % (tr.kmin, tr.kmax)) 

831 kwin = ' k window, dk = \'%s\', %.3f' % (tr.window, tr.dk) 

832 if tr.dk2 is not None: 

833 kwin = "%s, %.3f" % (kwin, tr.dk2) 

834 out.append(kwin) 

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

836 out.append(' paths used in fit = %s' % (repr(pathfiles))) 

837 out.append(' k-weight = %s' % kweigh) 

838 out.append(' epsilon_k = %s' % eps_k) 

839 out.append(' epsilon_r = %s' % eps_r) 

840 out.append(' n_independent = %.3f' % (ds.n_idp)) 

841 # 

842 out.append(' ') 

843 out.append(header % 'Variables') 

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

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

846 continue 

847 if len(name) < 14: 

848 name = (name + ' '*14)[:14] 

849 

850 if isParameter(par): 

851 if par.vary: 

852 stderr = 'unknown' 

853 if par.stderr is not None: 

854 stderr = gfmt(par.stderr) 

855 out.append(varformat % (name, gfmt(par.value), 

856 stderr, gfmt(par.init_value))) 

857 

858 elif par.expr is not None: 

859 stderr = 'unknown' 

860 if par.stderr is not None: 

861 stderr = gfmt(par.stderr) 

862 out.append(exprformat % (name, gfmt(par.value), 

863 stderr, par.expr)) 

864 else: 

865 out.append(fixformat % (name, gfmt(par.value))) 

866 

867 covar_vars = result.var_names 

868 if len(covar_vars) > 0: 

869 out.append(' ') 

870 out.append(header % 'Correlations' + 

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

872 correls = {} 

873 for i, name in enumerate(covar_vars): 

874 par = params[name] 

875 if not par.vary: 

876 continue 

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

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

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

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

881 

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

883 sort_correl.reverse() 

884 for name, val in sort_correl: 

885 if abs(val) < min_correl: 

886 break 

887 if len(name) < 20: 

888 name = (name + ' '*20)[:20] 

889 out.append(' %s = % .3f' % (name, val)) 

890 

891 if with_paths: 

892 out.append(' ') 

893 out.append(header % 'Paths') 

894 for ids, ds in enumerate(datasets): 

895 if len(datasets) > 1: 

896 out.append(' dataset %i:' % (ids+1)) 

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

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

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

900 return '\n'.join(out)