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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
|
#pylint: disable=no-name-in-module
import numpy as np
from mpi4py import MPI
from .utilities import FFTW_FORWARD, FFTW_MEASURE
def get_fftw_lib(dtype):
"""Return compiled fftw module interfacing the FFTW library
Parameters
----------
dtype : dtype
Data precision
Returns
-------
Module or ``None``
Module can be either :mod:`.fftwf_xfftn`, :mod:`.fftw_xfftn` or
:mod:`.fftwl_xfftn`, depending on precision.
"""
dtype = np.dtype(dtype).char.upper()
if dtype == 'G':
try:
from . import fftwl_xfftn
return fftwl_xfftn
except ImportError: #pragma: no cover
return None
elif dtype == 'D':
try:
from . import fftw_xfftn
return fftw_xfftn
except ImportError: #pragma: no cover
return None
elif dtype == 'F':
try:
from . import fftwf_xfftn
return fftwf_xfftn
except ImportError: #pragma: no cover
return None
else: #pragma: no cover
return None
fftlib = {}
for t in 'fdg':
fftw_lib = get_fftw_lib(t)
if fftw_lib is not None:
fftlib[t.upper()] = fftw_lib
comm = MPI.COMM_WORLD
def get_planned_FFT(input_array, output_array, axes=(-1,), kind=FFTW_FORWARD,
threads=1, flags=(FFTW_MEASURE,), normalization=1.0):
"""Return instance of transform class
Parameters
----------
input_array : array
real or complex input array
output_array : array
real or complex output array
axes : sequence of ints, optional
The axes to transform over, starting from the last
kind : int or sequence of ints, optional
Any one of (or possibly several for real-to-real)
- FFTW_FORWARD (-1)
- FFTW_R2HC (0)
- FFTW_BACKWARD (1)
- FFTW_HC2R (1)
- FFTW_DHT (2)
- FFTW_REDFT00 (3)
- FFTW_REDFT01 (4)
- FFTW_REDFT10 (5)
- FFTW_REDFT11 (6)
- FFTW_RODFT00 (7)
- FFTW_RODFT01 (8)
- FFTW_RODFT10 (9)
- FFTW_RODFT11 (10)
threads : int, optional
Number of threads to use in transforms
flags : int or sequence of ints, optional
Any one of, but not necessarily for all transforms or all combinations
- FFTW_MEASURE (0)
- FFTW_DESTROY_INPUT (1)
- FFTW_UNALIGNED (2)
- FFTW_CONSERVE_MEMORY (4)
- FFTW_EXHAUSTIVE (8)
- FFTW_PRESERVE_INPUT (16)
- FFTW_PATIENT (32)
- FFTW_ESTIMATE (64)
- FFTW_WISDOM_ONLY (2097152)
normalization : float, optional
Normalization factor
Returns
-------
:class:`.fftwf_xfftn.FFT`, :class:`.fftw_xfftn.FFT` or :class:`.fftwl_xfftn.FFT`
An instance of the return type configured for the desired transforms
"""
dtype = input_array.dtype.char
assert dtype.upper() in fftlib
_fft = fftlib[dtype.upper()]
return _fft.FFT(input_array, output_array, axes, kind, threads, flags,
normalization)
def export_wisdom(filename):
"""Export FFTW wisdom
Parameters
----------
filename : str
Name of file used to export wisdom to
Note
----
Wisdom is stored for all precisions available: float, double and long
double, using, respectively, prefix ``Fn_``, ``Dn_`` and ``Gn_``, where
n is the rank of the processor.
Wisdom is imported using :func:`.import_wisdom`, which must be called
with the same MPI configuration as used with :func:`.export_wisdom`.
See also
--------
:func:`.import_wisdom`
"""
rank = str(comm.Get_rank())
e = []
for key, lib in fftlib.items():
e.append(lib.export_wisdom(bytearray(key+rank+'_'+filename, 'utf-8')))
assert np.all(np.array(e) == 1), "Not able to export wisdom {}".format(filename)
def import_wisdom(filename):
"""Import FFTW wisdom
Parameters
----------
filename : str
Name of file used to import wisdom from
Note
----
Wisdom is imported for all available precisions: float, double and long
double, using, respectively, prefix ``Fn_``, ``Dn_`` and ``Gn_``, where
n is the rank of the processor.
Wisdom is exported using :func:`.export_wisdom`.
Note that importing wisdom only works when using the same MPI configuration
as used with :func:`.export_wisdom`.
See also
--------
:func:`.export_wisdom`
"""
rank = str(comm.Get_rank())
e = []
for key, lib in fftlib.items():
e.append(lib.import_wisdom(bytearray(key+rank+'_'+filename, 'utf-8')))
assert np.all(np.array(e) == 1), "Not able to import wisdom {}".format(filename)
def forget_wisdom():
for lib in fftlib.values():
lib.forget_wisdom()
def set_timelimit(limit):
"""Set time limit for planning
Parameters
----------
limit : number
The new time limit set for planning of serial transforms
"""
for lib in fftlib.values():
lib.set_timelimit(limit) # limit's precision handled by cython
def cleanup():
for lib in fftlib.values():
lib.cleanup()
|
