#!/usr/bin/env python """mpi4py-fft -- Parallel Fast Fourier Transforms (FFTs) using MPI for Python""" import os import sys import re import platform import sysconfig from distutils import ccompiler from distutils.errors import DistutilsPlatformError from setuptools import setup from setuptools.dist import Distribution from setuptools.extension import Extension import numpy cwd = os.path.abspath(os.path.dirname(__file__)) fftwdir = os.path.join(cwd, 'mpi4py_fft', 'fftw') prec_map = {'float': 'f', 'double': '', 'long double': 'l'} triplet = sysconfig.get_config_var('MULTIARCH') or '' bits = platform.architecture()[0][:-3] def append(dirlist, *args): entry = os.path.join(*args) entry = os.path.normpath(entry) if os.path.isdir(entry): if entry not in dirlist: dirlist.append(entry) def get_prefix_dirs(): dirs = [] for envvar in ('FFTW_ROOT', 'FFTW_DIR'): if envvar in os.environ: prefix = os.environ[envvar] append(dirs, prefix) append(dirs, sys.prefix) if 'CONDA_BUILD' not in os.environ: append(dirs, '/usr') append(dirs, '/usr/local') append(dirs, '/opt/homebrew') return dirs def get_include_dirs(): dirs = [] if 'FFTW_INCLUDE_DIR' in os.environ: entry = os.environ['FFTW_INCLUDE_DIR'] append(dirs, entry) for prefix in get_prefix_dirs(): append(dirs, prefix, 'include', triplet) append(dirs, prefix, 'include') dirs.append(numpy.get_include()) return dirs def get_library_dirs(): dirs = [] if 'FFTW_LIBRARY_DIR' in os.environ: entry = os.environ['FFTW_LIBRARY_DIR'] append(dirs, entry) for prefix in get_prefix_dirs(): append(dirs, prefix, 'lib' + bits) append(dirs, prefix, 'lib', triplet) append(dirs, prefix, 'lib') return dirs def get_fftw_libs(): """Return FFTW libraries""" compiler = ccompiler.new_compiler() library_dirs = get_library_dirs() libs = {} for d in ('float', 'double', 'long double'): lib = 'fftw3'+prec_map[d] tlib = lib+'_threads' if compiler.find_library_file(library_dirs, lib): libs[d] = [lib] if compiler.find_library_file(library_dirs, tlib): libs[d].append(tlib) if os.name == 'posix': libs[d].append('m') if not libs: message = "No FFTW libraries found in {}".format(library_dirs) raise DistutilsPlatformError(message) return libs def generate_extensions(fftwlibs, force=True): """Generate files with float and long double""" try: from setuptools.modified import newer_group except ImportError: try: from setuptools.dep_util import newer_group except ImportError: from distutils.dep_util import newer_group for d in fftwlibs: if d == 'double': continue p = 'fftw'+prec_map[d]+'_' for fname in ( 'fftw_planxfftn.h', 'fftw_planxfftn.c', 'fftw_xfftn.pyx', 'fftw_xfftn.pxd', ): src = os.path.join(fftwdir, fname) dst = os.path.join(fftwdir, fname.replace('fftw_', p)) if force or newer_group([src], dst, 'newer'): with open(src, 'r') as fin: code = fin.read() code = re.sub('fftw_', p, code) code = re.sub('double', d, code) with open(dst, 'w') as fout: fout.write(code) def remove_extensions(fftwlibs): """Remove generated files""" for fname in ( 'utilities.c', 'fftw_xfftn.c', 'fftwf_xfftn.c', 'fftwl_xfftn.c', ): dst = os.path.join(fftwdir, fname) try: os.remove(dst) except OSError: pass for d in fftwlibs: if d == 'double': continue p = 'fftw'+prec_map[d]+'_' for fname in ( 'fftw_planxfftn.h', 'fftw_planxfftn.c', 'fftw_xfftn.pyx', 'fftw_xfftn.pxd', ): dst = os.path.join(fftwdir, fname.replace('fftw_', p)) try: os.remove(dst) except OSError: pass def get_extensions(): """Return list of extension modules""" include_dirs = get_include_dirs() library_dirs = get_library_dirs() ext = [ Extension( "mpi4py_fft.fftw.utilities", sources=[os.path.join(fftwdir, "utilities.pyx")], define_macros=[('NPY_NO_DEPRECATED_API', 'NPY_1_7_API_VERSION')], include_dirs=include_dirs, ), ] sdist = 'sdist' in sys.argv egg_info = 'egg_info' in sys.argv fftwlibs = ( get_fftw_libs() if not (sdist or egg_info) else {d: [] for d in ('float', 'double', 'long double')} ) for d, libs in fftwlibs.items(): p = 'fftw' + 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, include_dirs=include_dirs, library_dirs=library_dirs, ) ) return ext class Dist(Distribution): def get_command_class(self, command): get_command_class = Distribution.get_command_class if 'build_ext' not in self.cmdclass: _build_ext = get_command_class(self, 'build_ext') class build_ext(_build_ext): def run(self): fftw_libs = get_fftw_libs() generate_extensions(fftw_libs, self.force) _build_ext.run(self) self.cmdclass['build_ext'] = build_ext if 'clean' not in self.cmdclass: _clean = get_command_class(self, 'clean') class clean(_clean): def run(self): fftw_libs = get_fftw_libs() remove_extensions(fftw_libs) _clean.run(self) self.cmdclass['clean'] = clean return get_command_class(self, command) def version(): srcdir = os.path.join(cwd, 'mpi4py_fft') with open(os.path.join(srcdir, '__init__.py')) as f: m = re.search(r"__version__\s*=\s*'(.*)'", f.read()) return m.groups()[0] with open("README.rst", "r") as fh: long_description = fh.read() if __name__ == '__main__': setup(name="mpi4py-fft", version=version(), description=__doc__.strip(), long_description=long_description, long_description_content_type='text/x-rst', author="Lisandro Dalcin and Mikael Mortensen", url="https://github.com/mpi4py/mpi4py-fft", packages=[ "mpi4py_fft", "mpi4py_fft.fftw", "mpi4py_fft.io", ], package_dir={ "mpi4py_fft": "mpi4py_fft", }, classifiers=[ 'Development Status :: 4 - Beta', 'Environment :: Console', 'Intended Audience :: Developers', 'Intended Audience :: Science/Research', 'Programming Language :: Python', 'Programming Language :: Python :: 2', 'Programming Language :: Python :: 3', 'License :: OSI Approved :: BSD License', 'Topic :: Scientific/Engineering :: Mathematics', 'Topic :: Software Development :: Libraries :: Python Modules', ], keywords=['Python', 'FFTW', 'FFT', 'DCT', 'DST', 'MPI'], distclass=Dist, ext_modules=get_extensions(), install_requires=["mpi4py", "numpy"], )