-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy path_fft.py
118 lines (99 loc) · 4.29 KB
/
_fft.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
import numpy as np
import warnings
try:
import multiprocessing
import pyfftw
from functools import partial
has_pyfftw = True
pyfftw_threads = multiprocessing.cpu_count()
pyfftw_planner_effort = 'FFTW_MEASURE'
fft2 = partial(pyfftw.interfaces.numpy_fft.fft2,
planner_effort=pyfftw_planner_effort,
threads=pyfftw_threads)
ifft2 = partial(pyfftw.interfaces.numpy_fft.ifft2,
planner_effort=pyfftw_planner_effort,
threads=pyfftw_threads)
fft = partial(pyfftw.interfaces.numpy_fft.fft,
planner_effort=pyfftw_planner_effort,
threads=pyfftw_threads)
ifft = partial(pyfftw.interfaces.numpy_fft.ifft,
planner_effort=pyfftw_planner_effort,
threads=pyfftw_threads)
fftn = partial(pyfftw.interfaces.numpy_fft.fftn,
planner_effort=pyfftw_planner_effort,
threads=pyfftw_threads)
ifftn = partial(pyfftw.interfaces.numpy_fft.ifftn,
planner_effort=pyfftw_planner_effort,
threads=pyfftw_threads)
fftshift = pyfftw.interfaces.numpy_fft.fftshift
ifftshift = pyfftw.interfaces.numpy_fft.ifftshift
fftfreq = pyfftw.interfaces.numpy_fft.fftfreq
# Turn on the cache for optimum performance
pyfftw.interfaces.cache.enable()
# increase cache preservation time from default of 0.1 seconds
pyfftw.interfaces.cache.set_keepalive_time(5)
except ImportError as e:
has_pyfftw = False
try:
warnings.warn("pyFFTW not found. will try to use mklfft instead.")
import mklfft
fft = mklfft.fftpack.fftff
ifft = mklfft.fftpack.ifft
fft2 = mklfft.fftpack.fft2
ifft2 = mklfft.fftpack.ifft2
fftn = mklfft.fftpack.fftn
ifftn = mklfft.fftpack.ifftn
fftshift = np.fft.fftshift
ifftshift = np.fft.ifftshift
fftfreq = np.fft.fftfreq
except ImportError as e:
warnings.warn("neither pyFFTW or mklfft found. will use numpy.fft.")
# Numpy's n-dimensional FFT routines may be using MKL, so prefered
# over scipy
fft = np.fft.fft
ifft = np.fft.ifft
fft2 = np.fft.fft2
ifft2 = np.fft.ifft2
fftn = np.fft.fftn
ifftn = np.fft.ifftn
fftshift = np.fft.fftshift
ifftshift = np.fft.ifftshift
fftfreq = np.fft.fftfreq
__all__ = ['fft', 'fft2', 'fftn', 'fftshift', 'fftfreq',
'ifft', 'ifft2', 'ifftn', 'ifftshift',
'fftnc', 'ifftnc', 'has_pyfftw']
if has_pyfftw:
# the following functions are PyFFTW dependent -
__all__ += ['build_fftn', 'build_ifftn', 'pyfftw_threads']
# centered versions of fftn for convenience
def fftnc(a, s=None, axes=None, pre_shift_axes=None, post_shift_axes=None):
y = ifftshift(a, axes=pre_shift_axes)
y = fftn(y, s=s, axes=axes)
return fftshift(y, axes=post_shift_axes)
def ifftnc(a, s=None, axes=None, pre_shift_axes=None, post_shift_axes=None):
y = ifftshift(a, axes=pre_shift_axes)
y = ifftn(y, s=s, axes=axes)
return fftshift(y, axes=post_shift_axes)
if has_pyfftw:
def build_fftn(a, fft_axes=None, threads=pyfftw_threads,
overwrite_input=False, planner_effort=pyfftw_planner_effort,
**kwargs):
if not has_pyfftw:
raise ValueError("pyfftw is required by plan_fftn")
return pyfftw.builders.fftn(a,
axes=fft_axes,
threads=threads,
overwrite_input=overwrite_input,
planner_effort=planner_effort,
**kwargs)
def build_ifftn(a, fft_axes=None, threads=pyfftw_threads,
overwrite_input=False,
planner_effort=pyfftw_planner_effort, **kwargs):
if not has_pyfftw:
raise ValueError("pyfftw is required by plan_fftn")
return pyfftw.builders.ifftn(a,
axes=fft_axes,
threads=threads,
overwrite_input=overwrite_input,
planner_effort=planner_effort,
**kwargs)