Coverage for /Users/Newville/Codes/xraylarch/larch/xafs/xafsft.py: 88%
143 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 Fourier transforms
4"""
5import numpy as np
6from numpy import (pi, arange, zeros, ones, sin, cos,
7 exp, log, sqrt, where, interp, linspace)
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#
29try:
30 from mkl_fft.interfaces.numpy_fft import fft, ifft
31except (ImportError, ModuleNotFoundError):
32 from scipy.fftpack import fft, ifft
34from scipy.special import i0 as bessel_i0
36from larch import (Group, Make_CallArgs, parse_group_args)
38from larch.math import complex_phase
39from .xafsutils import set_xafsGroup, FT_WINDOWS_SHORT
41MODNAME = '_xafs'
42sqrtpi = sqrt(pi)
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.
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
58 Returns:
59 ----------
60 1-d window array.
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
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)
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)
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)
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)
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)
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
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).
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.
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]
170 Returns:
171 ---------
172 None -- outputs are written to supplied group.
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)
186 Supports First Argument Group convention (with group member names 'r' and 'chir')
187 """
188 if 'rweight' in kws:
189 rw = kws['rweight']
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
198 cchir = zeros(nfft, dtype='complex128')
199 r_ = rstep * arange(nfft, dtype='float64')
201 cchir[0:len(chir)] = chir
202 if chir.dtype == np.dtype('complex128'):
203 scale = 0.5
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)
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])
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.
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]
248 Returns:
249 ---------
250 None -- outputs are written to supplied group.
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)
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')
276 if kstep is None:
277 kstep = k[1] - k[0]
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)
283 out = xftf_fast(cchi*win, kstep=kstep, nfft=nfft)
284 rstep = pi/(kstep*nfft)
286 irmax = int(min(nfft/2, 1.01 + rmax_out/rstep))
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])
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.
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])
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.
329 This is useful for repeated FTs, as inside loops.
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).
337 Returns:
338 --------
339 complex 1-d array chi(R)
341 """
342 cchi = zeros(nfft, dtype='complex128')
343 cchi[0:len(chi)] = chi
344 return (kstep / sqrtpi) * fft(cchi)[:int(nfft/2)]
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().
352 It returns the complex array of chi(q) without putting any
353 values into an output group.
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).
361 Returns:
362 ----------
363 complex 1-d array for chi(q).
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)]