Attempt to detect FFTW libraries in setup

Author: mikaem

README.rst

@@ -62,14 +62,14 @@ and the C-library
 
 Note that *mpi4py* requires a working MPI installation, with the compiler
 wrapper *mpicc* on your search path. All of the above dependencies are
-available and will be downlaoded through the conda-forge channel if
+available and will be downloaded through the conda-forge channel if
 conda is used for installation.
 
 For IO you need to install either `h5py <https://www.h5py.org>`_ or
 `netCDF4 <http://unidata.github.io/netcdf4-python/>`_ with support for
 MPI. These libraries are, unfortunately, not compiled with MPI on
 conda-forge. The two libraries are available, though, for both OSX and
-linux (`h5py-parallel <https://anaconda.org/spectraldns/h5py-parallel>`_
-and `netcdf4-parallel <https://anaconda.org/spectraldns/netcdf4-parallel>`_)
-from the `spectralDNS <https://anaconda.org/spectralDNS>`_ channel
-on anaconda cloud.
+linux from the `spectralDNS <https://anaconda.org/spectralDNS>`_ channel
+on anaconda cloud::
+
+    conda install -c spectralDNS h5py-parallel netcdf4-parallel

setup.py

@@ -2,45 +2,68 @@
 
 import os, sys
 import shutil
+from distutils import ccompiler
 from setuptools import setup
 from setuptools.extension import Extension
 from numpy import get_include
 
 cwd = os.path.abspath(os.path.dirname(__file__))
 fftwdir = os.path.join(cwd, 'mpi4py_fft', 'fftw')
 
-# For now assuming that all precisions are available
+include_dirs = [get_include(), os.path.join(sys.prefix, 'include')]
+library_dirs = [os.path.join(sys.prefix, 'lib')]
+for f in ('FFTW_ROOT', 'FFTW_DIR'):
+    if f in os.environ['PATH']:
+        library_dirs.append(os.path.join(os.environ[f], 'lib'))
+        include_dirs.append(os.path.join(os.environ[f], 'include'))
+
+compiler = ccompiler.new_compiler()
+assert compiler.find_library_file(library_dirs, 'fftw3'), 'Cannot find FFTW library!'
+has_threads = compiler.find_library_file(library_dirs, 'fftw3_threads')
+
+prec_map = {'float': 'fftwf_', 'double': 'fftw_', 'long double': 'fftwl_'}
 
-prec = {'fftwf_': 'float', 'fftw': 'double', 'fftwl_': 'long double'}
 libs = {
-    'fftwf_': ['m', 'fftw3f', 'fftw3f_threads'],
-    'fftw_': ['m', 'fftw3', 'fftw3_threads'],
-    'fftwl_': ['m', 'fftw3l', 'fftw3l_threads']}
+    'float': ['fftw3f'],
+    'double': ['fftw3'],
+    'long double': ['fftw3l']
+    }
+
+has_prec = ['double']
+for d in ('float', 'long double'):
+    if compiler.find_library_file(library_dirs, libs[d][0]):
+        has_prec.append(d)
+
+if has_threads:
+    for d in has_prec:
+        libs[d].append('_'.join((libs[d][0], 'threads')))
+        if sys.platform in ('unix', 'darwin'):
+            libs[d].append('m')
 
-for fl in ('fftw_planxfftn.h', 'fftw_planxfftn.c', 'fftw_xfftn.pyx', 'fftw_xfftn.pxd'):
-    for p in ('fftwf_', 'fftwl_'):
+# Generate files with float and long double if needed
+for d in has_prec[1:]:
+    p = prec_map[d]
+    for fl in ('fftw_planxfftn.h', 'fftw_planxfftn.c', 'fftw_xfftn.pyx', 'fftw_xfftn.pxd'):
         fp = fl.replace('fftw_', p)
         shutil.copy(os.path.join(fftwdir, fl), os.path.join(fftwdir, fp))
         sedcmd = "sed -i ''" if sys.platform == 'darwin' else "sed -i''"
         os.system(sedcmd + " 's/fftw_/{0}/g' {1}".format(p, os.path.join(fftwdir, fp)))
-        os.system(sedcmd + " 's/double/{0}/g' {1}".format(prec[p], os.path.join(fftwdir, fp)))
+        os.system(sedcmd + " 's/double/{0}/g' {1}".format(d, os.path.join(fftwdir, fp)))
 
 ext = [Extension("mpi4py_fft.fftw.utilities",
                  sources=[os.path.join(fftwdir, "utilities.pyx")],
-                 libraries=libs[p],
                  include_dirs=[get_include(),
-                               os.path.join(sys.prefix, 'include')],
-                 library_dirs=[os.path.join(sys.prefix, 'lib')])]
+                               os.path.join(sys.prefix, 'include')])]
 
-for p in ('fftw_', 'fftwf_', 'fftwl_'):
+for d in has_prec:
+    p = prec_map[d]
     ext.append(Extension("mpi4py_fft.fftw.{}xfftn".format(p),
                          sources=[os.path.join(fftwdir, "{}xfftn.pyx".format(p)),
                                   os.path.join(fftwdir, "{}planxfftn.c".format(p))],
                          #define_macros=[('NPY_NO_DEPRECATED_API', 'NPY_1_7_API_VERSION')],
-                         libraries=libs[p],
-                         include_dirs=[get_include(),
-                                       os.path.join(sys.prefix, 'include')],
-                         library_dirs=[os.path.join(sys.prefix, 'lib')]))
+                         libraries=libs[d],
+                         include_dirs=include_dirs,
+                         library_dirs=library_dirs))
 
 with open("README.rst", "r") as fh:
     long_description = fh.read()