Source Code for Module fftw3.lib

  1  #   This file is part of PyFFTW. 
  2  # 
  3  #    Copyright (C) 2009 Jochen Schroeder 
  4  #    Email: jschrod@berlios.de 
  5  # 
  6  #    PyFFTW is free software: you can redistribute it and/or modify 
  7  #    it under the terms of the GNU General Public License as published by 
  8  #    the Free Software Foundation, either version 3 of the License, or 
  9  #    (at your option) any later version. 
 10  # 
 11  #    PyFFTW is distributed in the hope that it will be useful, 
 12  #    but WITHOUT ANY WARRANTY; without even the implied warranty of 
 13  #    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the 
 14  #    GNU General Public License for more details. 
 15  # 
 16  #    You should have received a copy of the GNU General Public License 
 17  #    along with PyFFTW.  If not, see <http://www.gnu.org/licenses/>. 
 18   
 19  import ctypes 
 20  from ctypes import pythonapi, util, py_object 
 21  from numpy import ctypeslib, typeDict 
 22  from platform import system as psystem 
 23  from os.path import splitext, join, isfile, dirname, abspath, basename 
 24  from os.path import join as joinpath 
 25  from os import name as osname 
 26  from os import environ 
 27  from warnings import warn 
 28   
 29  try: 
 30      fftw_path = environ['FFTW_PATH'] 
 31      libfullpath = joinpath(abspath(fftw_path),r'libfftw3.so') 
 32      if not isfile(libfullpath): 
 33          raise IOError 
 34  except KeyError: 
 35      libfullpath = r'libfftw3.so.3' 
 36  except IOError: 
 37      warn('could not find %s in FFTW_PATH using installtime path' 
 38               %'libfftw3.so') 
 39      libfullpath = r'libfftw3.so.3' 
 40   
 41  if not isfile(libfullpath) and (osname=='nt' or psystem=='Windows'): 
 42      if isfile(joinpath(dirname(__file__), libfullpath)): 
 43          libfullpath = joinpath(dirname(__file__), libfullpath) 
 44   
 45  # must use ctypes.RTLD_GLOBAL for threading support 
 46  ctypes._dlopen(libfullpath, ctypes.RTLD_GLOBAL) 
 47  lib = ctypes.cdll.LoadLibrary(libfullpath) 
 48  #check if library is actually loaded there doesn't seem to be a better way to 
 49  #do this in ctypes 
 50  if not hasattr(lib, 'fftw_plan_dft_1d'): 
 51      raise OSError('Could not load libfftw3.so') 
 52   
 53  if osname == 'nt' or psystem() == 'Windows': 
 54      lib_threads = lib 
 55  else: 
 56      libbase, dot, ext = basename(libfullpath).partition('.') 
 57      libdir = dirname(libfullpath) 
 58      lib_threads = joinpath(libdir, libbase + '_threads.'+ ext) 
 59      try: 
 60          lib_threads = ctypes.cdll.LoadLibrary(lib_threads) 
 61      except OSError, e: 
 62          warn("Could not load threading library %s, threading support is disabled" 
 63              %lib_threads) 
 64          lib_threads = None 
 65   
 66   
 67  _typelist =    [('fftw_plan_dft_1d', (typeDict['complex'], typeDict['complex'], 1)), 
 68                         ('fftw_plan_dft_2d', (typeDict['complex'], typeDict['complex'], 2)), 
 69                         ('fftw_plan_dft_3d', (typeDict['complex'], typeDict['complex'], 3)), 
 70                         ('fftw_plan_dft', (typeDict['complex'], typeDict['complex'])), 
 71                         ('fftw_plan_dft_c2r_1d', (typeDict['complex'], typeDict['double'], 1)), 
 72                         ('fftw_plan_dft_c2r_2d', (typeDict['complex'], typeDict['double'], 2)), 
 73                         ('fftw_plan_dft_c2r_3d', (typeDict['complex'], typeDict['double'], 3)), 
 74                         ('fftw_plan_dft_c2r', (typeDict['complex'], typeDict['double'])), 
 75                         ('fftw_plan_dft_r2c_1d', (typeDict['double'], typeDict['complex'], 1)), 
 76                         ('fftw_plan_dft_r2c_2d', (typeDict['double'], typeDict['complex'], 2)), 
 77                         ('fftw_plan_dft_r2c_3d', (typeDict['double'], typeDict['complex'], 3)), 
 78                         ('fftw_plan_dft_r2c', (typeDict['double'], typeDict['complex'])), 
 79                         ('fftw_plan_r2r_1d', (typeDict['double'], typeDict['double'], 1)), 
 80                         ('fftw_plan_r2r_2d', (typeDict['double'], typeDict['double'], 2)), 
 81                         ('fftw_plan_r2r_3d', (typeDict['double'], typeDict['double'], 3)), 
 82                         ('fftw_plan_r2r', (typeDict['double'], typeDict['double']))] 
 83   
 84  _adv_typelist = [('fftw_plan_many_dft', (typeDict['complex'], 
 85                                                typeDict['complex'])), 
 86                    ('fftw_plan_many_dft_c2r', (typeDict['complex'], 
 87                                                     typeDict['double'])), 
 88                    ('fftw_plan_many_dft_r2c', (typeDict['double'], 
 89                                                     typeDict['complex'])), 
 90                    ('fftw_plan_many_r2r', (typeDict['double'], 
 91                                                     typeDict['double']))] 
 92   
 93   
 94 -def set_argtypes(val, types): 
 95      if types[0] == typeDict['complex'] and types[1] == typeDict['complex']: 
 96          set_argtypes_c2c(val,types) 
 97      elif types[0] == typeDict['complex'] or types[1] == typeDict['complex']: 
 98          set_argtypes_c2r(val,types) 
 99      else: 
100          set_argtypes_r2r(val,types) 
101   
102 -def set_argtypes_c2c(val,types): 
103      if len(types) >2: 
104          val.argtypes = [ctypes.c_int for i in range(types[2])] +\ 
105                         [ctypeslib.ndpointer(dtype=types[0],ndim=types[2], \ 
106                                              flags='contiguous, writeable, '\ 
107                                                    'aligned'), 
108                          ctypeslib.ndpointer(dtype=types[1], ndim=types[2],\ 
109                                              flags='contiguous, writeable, '\ 
110                                                    'aligned'), 
111                          ctypes.c_int, ctypes.c_uint] 
112      else: 
113          val.argtypes = [ctypes.c_int, ctypeslib.ndpointer(dtype=int, ndim=1,\ 
114                                                            flags='contiguous, '\ 
115                                                                  'aligned'), 
116                          ctypeslib.ndpointer(dtype=types[0], flags='contiguous,'\ 
117                                                                   ' writeable, '\ 
118                                                                    'aligned'), 
119                          ctypeslib.ndpointer(dtype=types[1],flags='contiguous, '\ 
120                                                                   'writeable,'\ 
121                                                                   'aligned'), 
122                          ctypes.c_int, ctypes.c_uint] 
123   
124 -def set_argtypes_c2r(val,types): 
125      if len(types) >2: 
126          val.argtypes = [ctypes.c_int for i in range(types[2])] +\ 
127                         [ctypeslib.ndpointer(dtype=types[0],ndim=types[2], \ 
128                                              flags='contiguous, writeable, '\ 
129                                                    'aligned'), 
130                          ctypeslib.ndpointer(dtype=types[1], ndim=types[2],\ 
131                                              flags='contiguous, writeable, '\ 
132                                                    'aligned'), 
133                          ctypes.c_uint] 
134      else: 
135          val.argtypes = [ctypes.c_int, ctypeslib.ndpointer(dtype=int, ndim=1,\ 
136                                                            flags='contiguous, '\ 
137                                                                  'aligned'), 
138                          ctypeslib.ndpointer(dtype=types[0], flags='contiguous,'\ 
139                                                                   ' writeable, '\ 
140                                                                    'aligned'), 
141                          ctypeslib.ndpointer(dtype=types[1],flags='contiguous, '\ 
142                                                                   'writeable,'\ 
143                                                                   'aligned'), 
144                          ctypes.c_uint] 
145   
146 -def set_argtypes_r2r(val, types): 
147      if len(types) > 2: 
148          val.argtypes = [ctypes.c_int for i in range(types[2])] +\ 
149                         [ctypeslib.ndpointer(dtype=types[0], ndim=types[2], 
150                                              flags='contiguous, writeable, '\ 
151                                                    'aligned'), 
152                          ctypeslib.ndpointer(dtype=types[1], ndim=types[2], 
153                                              flags='contiguous, writeable, '\ 
154                                                    'aligned')] +\ 
155                          [ctypes.c_int for i in range(types[2])] +\ 
156                          [ctypes.c_uint] 
157      else: 
158          val.argtypes = [ctypes.c_int, ctypeslib.ndpointer(dtype=int, ndim=1, 
159                                                            flags='contiguous, '\ 
160                                                                  'aligned'), 
161                          ctypeslib.ndpointer(dtype=types[0], flags='contiguous,'\ 
162                                                                    'writeable, '\ 
163                                                                    'aligned'), 
164                          ctypeslib.ndpointer(dtype=types[1], flags='contiguous,'\ 
165                                                                    'writeable, '\ 
166                                                                    'aligned'), 
167                          ctypeslib.ndpointer(dtype=int, ndim=1, 
168                                              flags='contiguous, aligned'), 
169                          ctypes.c_uint] 
170   
171 -def set_argtypes_adv(val, types): 
172      if types[0] == typeDict['complex'] and types[1] == typeDict['complex']: 
173          val.argtypes = [ctypes.c_int, ctypeslib.ndpointer(dtype=int, ndim=1, 
174                                                            flags='contiguous, '\ 
175                                                                  'aligned'), 
176                          ctypes.c_int, 
177                          ctypeslib.ndpointer(dtype=types[0], flags='contiguous,'\ 
178                                                                    'aligned,'\ 
179                                                                    'writeable'), 
180                          ctypeslib.ndpointer(dtype=int, ndim=1, 
181                                              flags='contiguous,aligned'), 
182                          ctypes.c_int, ctypes.c_int, 
183                          ctypeslib.ndpointer(dtype=types[1], flags='contiguous,'\ 
184                                                                    'aligned,'\ 
185                                                                    'writeable'), 
186                          ctypeslib.ndpointer(dtype=int, ndim=1, 
187                                              flags='contiguous,aligned'), 
188                          ctypes.c_int, ctypes.c_int, 
189                          ctypes.c_int, ctypes.c_uint] 
190      elif types[0] == typeDict['complex'] or types[1]==typeDict['complex']: 
191          val.argtypes = [ctypes.c_int, ctypeslib.ndpointer(dtype=int, ndim=1, 
192                                                            flags='contiguous, '\ 
193                                                                  'aligned'), 
194                          ctypes.c_int, 
195                          ctypeslib.ndpointer(dtype=types[0], flags='contiguous,'\ 
196                                                                    'aligned,'\ 
197                                                                    'writeable'), 
198                          ctypeslib.ndpointer(dtype=int, ndim=1, 
199                                              flags='contiguous,aligned'), 
200                          ctypes.c_int, ctypes.c_int, 
201                          ctypeslib.ndpointer(dtype=types[1], flags='contiguous,'\ 
202                                                                    'aligned,'\ 
203                                                                    'writeable'), 
204                          ctypeslib.ndpointer(dtype=int, ndim=1, 
205                                              flags='contiguous,aligned'), 
206                          ctypes.c_int, ctypes.c_int, 
207                          ctypes.c_uint] 
208      elif types[0] == typeDict['double'] and types[1]==typeDict['double']: 
209          val.argtypes = [ctypes.c_int, ctypeslib.ndpointer(dtype=int, ndim=1, 
210                                                            flags='contiguous, '\ 
211                                                                  'aligned'), 
212                          ctypes.c_int, 
213                          ctypeslib.ndpointer(dtype=types[0], flags='contiguous,'\ 
214                                                                    'aligned,'\ 
215                                                                    'writeable'), 
216                          ctypeslib.ndpointer(dtype=int, ndim=1, 
217                                              flags='contiguous,aligned'), 
218                          ctypes.c_int, ctypes.c_int, 
219                          ctypeslib.ndpointer(dtype=types[1], flags='contiguous,'\ 
220                                                                    'aligned,'\ 
221                                                                    'writeable'), 
222                          ctypeslib.ndpointer(dtype=int, ndim=1, 
223                                              flags='contiguous, aligned'), 
224                          ctypes.c_int, ctypes.c_int, 
225                          ctypeslib.ndpointer(dtype=int, ndim=1, 
226                                              flags='contiguous, aligned'), 
227                          ctypes.c_uint] 
228   
229   
230   
231  # set the return and argument types on the plan functions 
232  for name, types in _typelist: 
233      val = getattr(lib, name) 
234      val.restype = ctypes.c_void_p 
235      set_argtypes(val,types) 
236   
237  ##do the same for advanced plans 
238  for name, types in _adv_typelist: 
239      val = getattr(lib, name) 
240      val.restype = ctypes.c_void_p 
241      set_argtypes_adv(val,types) 
242   
243  #malloc and free 
244  lib.fftw_malloc.restype = ctypes.c_void_p 
245  lib.fftw_malloc.argtypes = [ctypes.c_int] 
246  lib.fftw_free.restype = None 
247  lib.fftw_free.argtypes = [ctypes.c_void_p] 
248   
249  #create a buffer from memory (necessary for array allocation) 
250  PyBuffer_FromReadWriteMemory = pythonapi.PyBuffer_FromReadWriteMemory 
251  PyBuffer_FromReadWriteMemory.restype = py_object 
252  PyBuffer_FromReadWriteMemory.argtypes = [ctypes.c_void_p, ctypes.c_int] 
253   
254  #executing arrays 
255  lib.fftw_execute.restype = None 
256  lib.fftw_execute.argtypes = [ctypes.c_void_p] 
257   
258  #guru execution 
259  lib.fftw_execute_dft.restype = None 
260  lib.fftw_execute_dft.argtypes = [ctypes.c_void_p, 
261                          ctypeslib.ndpointer(flags='aligned, contiguous, '\ 
262                                                          'writeable'),\ 
263                          ctypeslib.ndpointer(flags='aligned, contiguous, '\ 
264                                                          'writeable')] 
265   
266  #destroy plans 
267  lib.fftw_destroy_plan.restype = None 
268  lib.fftw_destroy_plan.argtypes = [ctypes.c_void_p] 
269   
270  #enable threading for plans 
271  if lib_threads is not None: 
272      lib_threads.fftw_init_threads.restype = ctypes.c_int 
273      lib_threads.fftw_init_threads.argtypes = [] 
274      lib_threads.fftw_plan_with_nthreads.restype = None 
275      lib_threads.fftw_plan_with_nthreads.argtypes = [ctypes.c_int] 
276      lib_threads.fftw_cleanup_threads.restype = None 
277      lib_threads.fftw_cleanup_threads.argtypes = [] 
278   
279      s = lib_threads.fftw_init_threads() 
280      if not s: 
281          sys.stderr.write('fftw_init_threads call failed, disabling threads support\n') 
282          lib_threads = None 
283   
284  #wisdom 
285   
286  # create c-file object from python 
287  PyFile_AsFile = pythonapi.PyFile_AsFile 
288  PyFile_AsFile.argtypes = [ctypes.py_object] 
289  PyFile_AsFile.restype = ctypes.c_void_p 
290   
291  #export to file 
292  lib.fftw_export_wisdom_to_file.argtypes = [ctypes.c_void_p] 
293  lib.fftw_export_wisdom_to_file.restype = None 
294   
295  #export to string 
296  lib.fftw_export_wisdom_to_string.argtypes = None 
297  lib.fftw_export_wisdom_to_string.restype = ctypes.c_char_p 
298   
299  #import from file 
300  lib.fftw_import_wisdom_from_file.argtypes = [ctypes.c_void_p] 
301  lib.fftw_import_wisdom_from_file.restype = ctypes.c_int 
302   
303  #import from string 
304  lib.fftw_import_wisdom_from_string.argtypes = [ctypes.c_char_p] 
305  lib.fftw_import_wisdom_from_string.restype = ctypes.c_int 
306   
307  #import system wisdom 
308  lib.fftw_import_system_wisdom.restype = ctypes.c_int 
309  lib.fftw_import_system_wisdom.argtypes = None 
310   
311  #forget wisdom 
312  lib.fftw_forget_wisdom.restype = None 
313  lib.fftw_forget_wisdom.argtype = None 
314