Coverage for larch/xafs/xafsft.py: 89%

142 statements  

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

1#!/usr/bin/env python 

2""" 

3 XAFS Fourier transforms 

4""" 

5import numpy as np 

6from numpy import (pi, arange, zeros, ones, sin, cos, 

7 exp, log, sqrt, where, interp, linspace) 

8 

9# notes on FFT method choice: (April 2023) 

10# benchmarking (and tossing out the first few runs)for complex128 arrays of 

11# length 2048, as used here for XAFS, gives these kinds of run times: 

12# 

13# fft interface : mean std worst 

14# numpy : 2.2509 0.0955 2.4858 

15# scipy : 2.1156 0.0592 2.2758 

16# scipy fftpack : 1.9757 0.0233 2.0246 

17# mkl numpy : 1.1822 0.0188 1.2279 

18# mkl scipy : 1.4442 0.0688 1.6263 

19# fftw numpy : 5.4955 0.1958 6.0779 

20# fftw scipy : 5.5210 0.1620 5.9154 

21# fftw fftpack : 5.4329 0.0807 5.5917 

22# 

23# this shows a noticeable performance win for MKL numpy interface, 

24# with a modest win for scipy.fftpack over numpy or scipy if MKL 

25# is not available. These results were consistent for Linux, 

26# Windows, and MacOSX. 

27# 

28 

29try: 

30 from mkl_fft.interfaces.numpy_fft import fft, ifft 

31except (ImportError, ModuleNotFoundError): 

32 from scipy.fftpack import fft, ifft 

33 

34from scipy.special import i0 as bessel_i0 

35 

36from larch import (Group, Make_CallArgs, parse_group_args) 

37 

38from larch.math import complex_phase 

39from .xafsutils import set_xafsGroup, FT_WINDOWS_SHORT 

40 

41MODNAME = '_xafs' 

42sqrtpi = sqrt(pi) 

43 

44def ftwindow(x, xmin=None, xmax=None, dx=1, dx2=None, 

45 window='hanning', _larch=None, **kws): 

46 """ 

47 create a Fourier transform window array. 

48 

49 Parameters: 

50 ------------- 

51 x: 1-d array array to build window on. 

52 xmin: starting x for FT Window 

53 xmax: ending x for FT Window 

54 dx: tapering parameter for FT Window 

55 dx2: second tapering parameter for FT Window (=dx) 

56 window: name of window type 

57 

58 Returns: 

59 ---------- 

60 1-d window array. 

61 

62 Notes: 

63 ------- 

64 Valid Window names: 

65 hanning cosine-squared taper 

66 parzen linear taper 

67 welch quadratic taper 

68 gaussian Gaussian (normal) function window 

69 sine sine function window 

70 kaiser Kaiser-Bessel function-derived window 

71 

72 """ 

73 if window is None: 

74 window = FT_WINDOWS_SHORT[0] 

75 nam = window.strip().lower()[:3] 

76 if nam not in FT_WINDOWS_SHORT: 

77 raise RuntimeError("invalid window name %s" % window) 

78 

79 dx1 = dx 

80 if dx2 is None: dx2 = dx1 

81 if xmin is None: xmin = min(x) 

82 if xmax is None: xmax = max(x) 

83 

84 xstep = (x[-1] - x[0]) / (len(x)-1) 

85 xeps = 1.e-4 * xstep 

86 x1 = max(min(x), xmin - dx1/2.0) 

87 x2 = xmin + dx1/2.0 + xeps 

88 x3 = xmax - dx2/2.0 - xeps 

89 x4 = min(max(x), xmax + dx2/2.0) 

90 

91 if nam == 'fha': 

92 if dx1 < 0: dx1 = 0 

93 if dx2 > 1: dx2 = 1 

94 x2 = x1 + xeps + dx1*(xmax-xmin)/2.0 

95 x3 = x4 - xeps - dx2*(xmax-xmin)/2.0 

96 elif nam == 'gau': 

97 dx1 = max(dx1, xeps) 

98 

99 def asint(val): return int((val+xeps)/xstep) 

100 i1, i2, i3, i4 = asint(x1), asint(x2), asint(x3), asint(x4) 

101 i1, i2 = max(0, i1), max(0, i2) 

102 i3, i4 = min(len(x)-1, i3), min(len(x)-1, i4) 

103 if i2 == i1: i1 = max(0, i2-1) 

104 if i4 == i3: i3 = max(i2, i4-1) 

105 x1, x2, x3, x4 = x[i1], x[i2], x[i3], x[i4] 

106 if x1 == x2: x2 = x2+xeps 

107 if x3 == x4: x4 = x4+xeps 

108 # initial window 

109 fwin = zeros(len(x)) 

110 if i3 > i2: 

111 fwin[i2:i3] = ones(i3-i2) 

112 

113 # now finish making window 

114 if nam in ('han', 'fha'): 

115 fwin[i1:i2+1] = sin((pi/2)*(x[i1:i2+1]-x1) / (x2-x1))**2 

116 fwin[i3:i4+1] = cos((pi/2)*(x[i3:i4+1]-x3) / (x4-x3))**2 

117 elif nam == 'par': 

118 fwin[i1:i2+1] = (x[i1:i2+1]-x1) / (x2-x1) 

119 fwin[i3:i4+1] = 1 - (x[i3:i4+1]-x3) / (x4-x3) 

120 elif nam == 'wel': 

121 fwin[i1:i2+1] = 1 - ((x[i1:i2+1]-x2) / (x2-x1))**2 

122 fwin[i3:i4+1] = 1 - ((x[i3:i4+1]-x3) / (x4-x3))**2 

123 elif nam in ('kai', 'bes'): 

124 cen = (x4+x1)/2 

125 wid = (x4-x1)/2 

126 arg = 1 - (x-cen)**2 / (wid**2) 

127 arg[where(arg<0)] = 0 

128 if nam == 'bes': # 'bes' : ifeffit 1.0 implementation of kaiser-bessel 

129 fwin = bessel_i0(dx* sqrt(arg)) / bessel_i0(dx) 

130 fwin[where(x<=x1)] = 0 

131 fwin[where(x>=x4)] = 0 

132 else: # better version 

133 scale = max(1.e-10, bessel_i0(dx)-1) 

134 fwin = (bessel_i0(dx * sqrt(arg)) - 1) / scale 

135 elif nam == 'sin': 

136 fwin[i1:i4+1] = sin(pi*(x4-x[i1:i4+1]) / (x4-x1)) 

137 elif nam == 'gau': 

138 cen = (x4+x1)/2 

139 fwin = exp(-(((x - cen)**2)/(2*dx1*dx1))) 

140 return fwin 

141 

142 

143@Make_CallArgs(["r", "chir"]) 

144def xftr(r, chir=None, group=None, rmin=0, rmax=20, with_phase=False, 

145 dr=1, dr2=None, rw=0, window='kaiser', qmax_out=None, 

146 nfft=2048, kstep=0.05, _larch=None, **kws): 

147 """ 

148 reverse XAFS Fourier transform, from chi(R) to chi(q). 

149 

150 calculate reverse XAFS Fourier transform 

151 This assumes that chir_re and (optional chir_im are 

152 on a uniform r-grid given by r. 

153 

154 Parameters: 

155 ------------ 

156 r: 1-d array of distance, or group. 

157 chir: 1-d array of chi(R) 

158 group: output Group 

159 qmax_out: highest *k* for output data (30 Ang^-1) 

160 rweight: exponent for weighting spectra by r^rweight (0) 

161 rmin: starting *R* for FT Window 

162 rmax: ending *R* for FT Window 

163 dr: tapering parameter for FT Window 

164 dr2: second tapering parameter for FT Window 

165 window: name of window type 

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

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

168 with_phase: output the phase as well as magnitude, real, imag [False] 

169 

170 Returns: 

171 --------- 

172 None -- outputs are written to supplied group. 

173 

174 Notes: 

175 ------- 

176 Arrays written to output group: 

177 rwin window Omega(R) (length of input chi(R)). 

178 q uniform array of k, out to qmax_out. 

179 chiq complex array of chi(k). 

180 chiq_mag magnitude of chi(k). 

181 chiq_re real part of chi(k). 

182 chiq_im imaginary part of chi(k). 

183 chiq_pha phase of chi(k) if with_phase=True 

184 (a noticable performance hit) 

185 

186 Supports First Argument Group convention (with group member names 'r' and 'chir') 

187 """ 

188 if 'rweight' in kws: 

189 rw = kws['rweight'] 

190 

191 r, chir, group = parse_group_args(r, members=('r', 'chir'), 

192 defaults=(chir,), group=group, 

193 fcn_name='xftr') 

194 rstep = r[1] - r[0] 

195 kstep = pi/(rstep*nfft) 

196 scale = 1.0 

197 

198 cchir = zeros(nfft, dtype='complex128') 

199 r_ = rstep * arange(nfft, dtype='float64') 

200 

201 cchir[0:len(chir)] = chir 

202 if chir.dtype == np.dtype('complex128'): 

203 scale = 0.5 

204 

205 win = ftwindow(r_, xmin=rmin, xmax=rmax, dx=dr, dx2=dr2, window=window) 

206 out = scale * xftr_fast( cchir*win * r_**rw, kstep=kstep, nfft=nfft) 

207 if qmax_out is None: qmax_out = 30.0 

208 q = linspace(0, qmax_out, int(1.05 + qmax_out/kstep)) 

209 nkpts = len(q) 

210 

211 group = set_xafsGroup(group, _larch=_larch) 

212 group.q = q 

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

214 group.rwin = win[:len(chir)] 

215 group.chiq = out[:nkpts] 

216 group.chiq_mag = mag[:nkpts] 

217 group.chiq_re = out.real[:nkpts] 

218 group.chiq_im = out.imag[:nkpts] 

219 if with_phase: 

220 group.chiq_pha = complex_phase(out[:nkpts]) 

221 

222 

223 

224@Make_CallArgs(["k", "chi"]) 

225def xftf(k, chi=None, group=None, kmin=0, kmax=20, kweight=None, 

226 dk=1, dk2=None, with_phase=False, window='kaiser', rmax_out=10, 

227 nfft=2048, kstep=None, _larch=None, **kws): 

228 """ 

229 forward XAFS Fourier transform, from chi(k) to chi(R), using 

230 common XAFS conventions. 

231 

232 Parameters: 

233 ----------- 

234 k: 1-d array of photo-electron wavenumber in Ang^-1 or group 

235 chi: 1-d array of chi 

236 group: output Group 

237 rmax_out: highest R for output data (10 Ang) 

238 kweight: exponent for weighting spectra by k**kweight [2] 

239 kmin: starting k for FT Window 

240 kmax: ending k for FT Window 

241 dk: tapering parameter for FT Window 

242 dk2: second tapering parameter for FT Window 

243 window: name of window type 

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

245 kstep: value to use for delta_k (k[1]-k[0] Ang^-1). 

246 with_phase: output the phase as well as magnitude, real, imag [False] 

247 

248 Returns: 

249 --------- 

250 None -- outputs are written to supplied group. 

251 

252 Notes: 

253 ------- 

254 Arrays written to output group: 

255 kwin window function Omega(k) (length of input chi(k)). 

256 r uniform array of R, out to rmax_out. 

257 chir complex array of chi(R). 

258 chir_mag magnitude of chi(R). 

259 chir_re real part of chi(R). 

260 chir_im imaginary part of chi(R). 

261 chir_pha phase of chi(R) if with_phase=True 

262 (a noticable performance hit) 

263 

264 Supports First Argument Group convention (with group member names 'k' and 'chi') 

265 """ 

266 # allow kweight keyword == kw 

267 if kweight is None: 

268 if 'kw' in kws: 

269 kweight = kws['kw'] 

270 else: 

271 kweight = 2 

272 k, chi, group = parse_group_args(k, members=('k', 'chi'), 

273 defaults=(chi,), group=group, 

274 fcn_name='xftf') 

275 

276 if kstep is None: 

277 kstep = k[1] - k[0] 

278 

279 cchi, win = xftf_prep(k, chi, kmin=kmin, kmax=kmax, kweight=kweight, 

280 dk=dk, dk2=dk2, nfft=nfft, kstep=kstep, 

281 window=window, _larch=_larch) 

282 

283 out = xftf_fast(cchi*win, kstep=kstep, nfft=nfft) 

284 rstep = pi/(kstep*nfft) 

285 

286 irmax = int(min(nfft/2, 1.01 + rmax_out/rstep)) 

287 

288 group = set_xafsGroup(group, _larch=_larch) 

289 r = rstep * arange(irmax) 

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

291 group.kwin = win[:len(chi)] 

292 group.r = r[:irmax] 

293 group.chir = out[:irmax] 

294 group.chir_mag = mag[:irmax] 

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

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

297 if with_phase: 

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

299 

300 

301 

302def xftf_prep(k, chi, kmin=0, kmax=20, kweight=2, dk=1, dk2=None, 

303 window='kaiser', nfft=2048, kstep=0.05, _larch=None): 

304 """ 

305 calculate weighted chi(k) on uniform grid of len=nfft, and the 

306 ft window. 

307 

308 Returns weighted chi, window function which can easily be multiplied 

309 and used in xftf_fast. 

310 """ 

311 if dk2 is None: dk2 = dk 

312 kweight = int(kweight) 

313 npts = int(1.01 + max(k)/kstep) 

314 k_max = max(max(k), kmax+dk2) 

315 k_ = kstep * np.arange(int(1.01+k_max/kstep), dtype='float64') 

316 chi_ = interp(k_, k, chi) 

317 win = ftwindow(k_, xmin=kmin, xmax=kmax, dx=dk, dx2=dk2, window=window) 

318 return ((chi_[:npts] *k_[:npts]**kweight), win[:npts]) 

319 

320 

321def xftf_fast(chi, nfft=2048, kstep=0.05, _larch=None, **kws): 

322 """ 

323 calculate forward XAFS Fourier transform. Unlike xftf(), 

324 this assumes that: 

325 1. data is already on a uniform grid 

326 2. any windowing and/or kweighting has been applied. 

327 and simply returns the complex chi(R), not setting any larch data. 

328 

329 This is useful for repeated FTs, as inside loops. 

330 

331 Parameters: 

332 ------------ 

333 chi: 1-d array of chi to be transformed 

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

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

336 

337 Returns: 

338 -------- 

339 complex 1-d array chi(R) 

340 

341 """ 

342 cchi = zeros(nfft, dtype='complex128') 

343 cchi[0:len(chi)] = chi 

344 return (kstep / sqrtpi) * fft(cchi)[:int(nfft/2)] 

345 

346def xftr_fast(chir, nfft=2048, kstep=0.05, _larch=None, **kws): 

347 """ 

348 calculate reverse XAFS Fourier transform, from chi(R) to 

349 chi(q), using common XAFS conventions. This version demands 

350 chir be the complex chi(R) as created from xftf(). 

351 

352 It returns the complex array of chi(q) without putting any 

353 values into an output group. 

354 

355 Parameters: 

356 ------------- 

357 chir: 1-d array of chi(R) to be transformed 

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

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

360 

361 Returns: 

362 ---------- 

363 complex 1-d array for chi(q). 

364 

365 This is useful for repeated FTs, as inside loops. 

366 """ 

367 cchi = zeros(nfft, dtype='complex128') 

368 cchi[0:len(chir)] = chir 

369 return (4*sqrtpi/kstep) * ifft(cchi)[:int(nfft/2)]