From 517e3e6ce1142132ff6b83978d5ce5736da29b7a Mon Sep 17 00:00:00 2001 From: smallnamespace Date: Mon, 16 Nov 2015 17:56:43 -0800 Subject: [PATCH 1/2] Add support for pyFFTW; minor restructuring to fft.py. --- nsgt/fft.py | 187 +++++++++++++++++++++++---------------------- tests/cq_test.py | 2 +- tests/nsgt_test.py | 2 +- 3 files changed, 99 insertions(+), 92 deletions(-) diff --git a/nsgt/fft.py b/nsgt/fft.py index 41c89cd..62d320c 100644 --- a/nsgt/fft.py +++ b/nsgt/fft.py @@ -14,102 +14,109 @@ import numpy as np from warnings import warn -realized = False - -if not realized: - # try to use FFT3 if available, else use numpy.fftpack +# Try engines in order of: +# PyFFTW3 +# pyFFTW +# numpy.fftpack +try: + import fftw3, fftw3f +except ImportError: + fftw3 = None + fftw3f = None try: - import fftw3 + import pyfftw except ImportError: - fftw3 = None - - try: - import fftw3f - except ImportError: - fftw3f = None + pyfftw = None + + +if fftw3 is not None and fftw3f is not None: + ENGINE = "PYFFTW3" + # Use fftw3 methods + class fftpool: + def __init__(self, measure, dtype=float): + self.measure = measure + self.dtype = np.dtype(dtype) + dtsz = self.dtype.itemsize + if dtsz == 4: + self.tpfloat = np.float32 + self.tpcplx = np.complex64 + self.fftw = fftw3f + elif dtsz == 8: + self.tpfloat = np.float64 + self.tpcplx = np.complex128 + self.fftw = fftw3 + else: + raise TypeError("nsgt.fftpool: dtype '%s' not supported"%repr(self.dtype)) + self.pool = {} + + def __call__(self, x, outn=None, ref=False): + lx = len(x) + try: + transform = self.pool[lx] + except KeyError: + transform = self.init(lx, measure=self.measure, outn=outn) + self.pool[lx] = transform + plan,pre,post = transform + if pre is not None: + x = pre(x) + plan.inarray[:] = x + plan() + if not ref: + tx = plan.outarray.copy() + else: + tx = plan.outarray + if post is not None: + tx = post(tx) + return tx + + class fftp(fftpool): + def __init__(self, measure=False, dtype=float): + fftpool.__init__(self, measure, dtype=dtype) + def init(self, n, measure, outn): + inp = self.fftw.create_aligned_array(n, dtype=self.tpcplx) + outp = self.fftw.create_aligned_array(n, dtype=self.tpcplx) + plan = self.fftw.Plan(inp, outp, direction='forward', flags=('measure' if measure else 'estimate',)) + return (plan,None,None) + + class rfftp(fftpool): + def __init__(self, measure=False, dtype=float): + fftpool.__init__(self, measure, dtype=dtype) + def init(self, n, measure, outn): + inp = self.fftw.create_aligned_array(n, dtype=self.tpfloat) + outp = self.fftw.create_aligned_array(n//2+1, dtype=self.tpcplx) + plan = self.fftw.Plan(inp, outp, direction='forward', realtypes='halfcomplex r2c',flags=('measure' if measure else 'estimate',)) + return (plan,None,None) - if fftw3 is not None: - # fftw3 methods - class fftpool: - def __init__(self, measure, dtype=float): - self.measure = measure - self.dtype = np.dtype(dtype) - dtsz = self.dtype.itemsize - if dtsz == 4: - self.tpfloat = np.float32 - self.tpcplx = np.complex64 - self.fftw = fftw3f - elif dtsz == 8: - self.tpfloat = np.float64 - self.tpcplx = np.complex128 - self.fftw = fftw3 - else: - raise TypeError("nsgt.fftpool: dtype '%s' not supported"%repr(self.dtype)) - self.pool = {} - - def __call__(self, x, outn=None, ref=False): - lx = len(x) - try: - transform = self.pool[lx] - except KeyError: - transform = self.init(lx, measure=self.measure, outn=outn) - self.pool[lx] = transform - plan,pre,post = transform - if pre is not None: - x = pre(x) - plan.inarray[:] = x - plan() - if not ref: - tx = plan.outarray.copy() - else: - tx = plan.outarray - if post is not None: - tx = post(tx) - return tx - - class fftp(fftpool): - def __init__(self, measure=False, dtype=float): - fftpool.__init__(self, measure, dtype=dtype) - def init(self, n, measure, outn): - inp = self.fftw.create_aligned_array(n, dtype=self.tpcplx) - outp = self.fftw.create_aligned_array(n, dtype=self.tpcplx) - plan = self.fftw.Plan(inp, outp, direction='forward', flags=('measure' if measure else 'estimate',)) - return (plan,None,None) - - class rfftp(fftpool): - def __init__(self, measure=False, dtype=float): - fftpool.__init__(self, measure, dtype=dtype) - def init(self, n, measure, outn): - inp = self.fftw.create_aligned_array(n, dtype=self.tpfloat) - outp = self.fftw.create_aligned_array(n//2+1, dtype=self.tpcplx) - plan = self.fftw.Plan(inp, outp, direction='forward', realtypes='halfcomplex r2c',flags=('measure' if measure else 'estimate',)) - return (plan,None,None) - - class ifftp(fftpool): - def __init__(self, measure=False, dtype=float): - fftpool.__init__(self, measure, dtype=dtype) - def init(self, n, measure, outn): - inp = self.fftw.create_aligned_array(n, dtype=self.tpcplx) - outp = self.fftw.create_aligned_array(n, dtype=self.tpcplx) - plan = self.fftw.Plan(inp, outp, direction='backward', flags=('measure' if measure else 'estimate',)) - return (plan,None,lambda x: x/len(x)) - - class irfftp(fftpool): - def __init__(self, measure=False, dtype=float): - fftpool.__init__(self, measure, dtype=dtype) - def init(self, n, measure, outn): - inp = self.fftw.create_aligned_array(n, dtype=self.tpcplx) - outp = self.fftw.create_aligned_array(outn if outn is not None else (n-1)*2, dtype=self.tpfloat) - plan = self.fftw.Plan(inp, outp, direction='backward', realtypes='halfcomplex c2r', flags=('measure' if measure else 'estimate',)) - return (plan,lambda x: x[:n],lambda x: x/len(x)) - - realized = True - + class ifftp(fftpool): + def __init__(self, measure=False, dtype=float): + fftpool.__init__(self, measure, dtype=dtype) + def init(self, n, measure, outn): + inp = self.fftw.create_aligned_array(n, dtype=self.tpcplx) + outp = self.fftw.create_aligned_array(n, dtype=self.tpcplx) + plan = self.fftw.Plan(inp, outp, direction='backward', flags=('measure' if measure else 'estimate',)) + return (plan,None,lambda x: x/len(x)) -if not realized: + class irfftp(fftpool): + def __init__(self, measure=False, dtype=float): + fftpool.__init__(self, measure, dtype=dtype) + def init(self, n, measure, outn): + inp = self.fftw.create_aligned_array(n, dtype=self.tpcplx) + outp = self.fftw.create_aligned_array(outn if outn is not None else (n-1)*2, dtype=self.tpfloat) + plan = self.fftw.Plan(inp, outp, direction='backward', realtypes='halfcomplex c2r', flags=('measure' if measure else 'estimate',)) + return (plan,lambda x: x[:n],lambda x: x/len(x)) +elif pyfftw is not None: + ENGINE = "PYFFTW" + # Monkey patch in pyFFTW Numpy interface + np.fft = pyfftw.interfaces.numpy_fft + original_fft = np.fft + # Enable cache to keep wisdom, etc. + pyfftw.interfaces.cache.enable() +else: # fall back to numpy methods warn("nsgt.fft falling back to numpy.fft") - + ENGINE = "NUMPY" + +if ENGINE in ["PYFFTW", "NUMPY"]: class fftp: def __init__(self, measure=False, dtype=float): pass diff --git a/tests/cq_test.py b/tests/cq_test.py index 3d95167..4b837a9 100644 --- a/tests/cq_test.py +++ b/tests/cq_test.py @@ -14,7 +14,7 @@ def test_oct(self): nsgt = NSGT(scale, fs=44100, Ls=len(sig)) c = nsgt.forward(sig) s_r = nsgt.backward(c) - self.assertTrue(np.allclose(sig, s_r)) + self.assertTrue(np.allclose(sig, s_r, atol=1e-07)) def load_tests(*_): test_cases = unittest.TestSuite() diff --git a/tests/nsgt_test.py b/tests/nsgt_test.py index ea4f778..4013a07 100644 --- a/tests/nsgt_test.py +++ b/tests/nsgt_test.py @@ -13,7 +13,7 @@ def test_transform(self, length=100000, fmin=50, fmax=22050, bins=12, fs=44100): # inverse transform s_r = nsgt.backward(c) - self.assertTrue(np.allclose(s, s_r)) + self.assertTrue(np.allclose(s, s_r, atol=1e-07)) def load_tests(*_): test_cases = unittest.TestSuite() From b9ebadc7b5aeadb260d22858b82f011d87f6314d Mon Sep 17 00:00:00 2001 From: smallnamespace Date: Tue, 17 Nov 2015 15:03:15 -0800 Subject: [PATCH 2/2] Implement saving and loading of pyFFTW wisdom for substantial speedup --- nsgt/fft.py | 38 ++++++++++++++++++++++++++++++-------- tests/cq_test.py | 4 ++++ tests/fft_test.py | 4 ++++ tests/nsgt_test.py | 4 ++++ 4 files changed, 42 insertions(+), 8 deletions(-) diff --git a/nsgt/fft.py b/nsgt/fft.py index 62d320c..9df8760 100644 --- a/nsgt/fft.py +++ b/nsgt/fft.py @@ -11,7 +11,11 @@ AudioMiner project, supported by Vienna Science and Technology Fund (WWTF) """ +import atexit import numpy as np +import os.path +import pickle +from threading import Timer from warnings import warn # Try engines in order of: @@ -111,29 +115,47 @@ def init(self, n, measure, outn): original_fft = np.fft # Enable cache to keep wisdom, etc. pyfftw.interfaces.cache.enable() + + # Load stored wisdom + PYFFTW_WISDOM_FILENAME = os.path.join(os.path.expanduser("~"), ".nsgt_pyfftw_wisdom.p") + if os.path.isfile(PYFFTW_WISDOM_FILENAME): + with open(PYFFTW_WISDOM_FILENAME, 'rb') as f: + pyfftw.import_wisdom(pickle.load(f)) + print("Loaded pyFFTW wisdom from %s" % PYFFTW_WISDOM_FILENAME) + + def save_wisdom(): + print("Saving pyFFTW wisdom to %s" % PYFFTW_WISDOM_FILENAME) + with open(PYFFTW_WISDOM_FILENAME, 'wb') as f: + pickle.dump(pyfftw.export_wisdom(), f) + + # Save wisdom on exit + atexit.register(save_wisdom) else: # fall back to numpy methods warn("nsgt.fft falling back to numpy.fft") ENGINE = "NUMPY" if ENGINE in ["PYFFTW", "NUMPY"]: + def get_kwargs(measure): + return ({'planner_effort': 'FFTW_MEASURE' if measure else 'FFTW_ESTIMATE'} + if ENGINE=="PYFFTW" else {}) class fftp: def __init__(self, measure=False, dtype=float): - pass + self.kwargs = get_kwargs(measure) def __call__(self,x, outn=None, ref=False): - return np.fft.fft(x) + return np.fft.fft(x, **self.kwargs) class ifftp: def __init__(self, measure=False, dtype=float): - pass + self.kwargs = get_kwargs(measure) def __call__(self,x, outn=None, n=None, ref=False): - return np.fft.ifft(x,n=n) + return np.fft.ifft(x,n=n,**self.kwargs) class rfftp: def __init__(self, measure=False, dtype=float): - pass + self.kwargs = get_kwargs(measure) def __call__(self,x, outn=None, ref=False): - return np.fft.rfft(x) + return np.fft.rfft(x,**self.kwargs) class irfftp: def __init__(self, measure=False, dtype=float): - pass + self.kwargs = get_kwargs(measure) def __call__(self,x,outn=None,ref=False): - return np.fft.irfft(x,n=outn) + return np.fft.irfft(x,n=outn,**self.kwargs) diff --git a/tests/cq_test.py b/tests/cq_test.py index 4b837a9..2647936 100644 --- a/tests/cq_test.py +++ b/tests/cq_test.py @@ -2,6 +2,10 @@ from nsgt import NSGT, OctScale import unittest +# Make test deterministic +np.random.seed(666) + + class TestNSGT(unittest.TestCase): def test_oct(self): diff --git a/tests/fft_test.py b/tests/fft_test.py index 69be02b..9007370 100644 --- a/tests/fft_test.py +++ b/tests/fft_test.py @@ -2,6 +2,10 @@ from nsgt.fft import rfftp, irfftp, fftp, ifftp import unittest +# Make test deterministic +np.random.seed(666) + + class TestFFT(unittest.TestCase): def __init__(self, methodName, n=10000): super(TestFFT, self).__init__(methodName) diff --git a/tests/nsgt_test.py b/tests/nsgt_test.py index 4013a07..7d0598c 100644 --- a/tests/nsgt_test.py +++ b/tests/nsgt_test.py @@ -2,6 +2,10 @@ from nsgt import CQ_NSGT import unittest +# Make test deterministic +np.random.seed(666) + + class Test_CQ_NSGT(unittest.TestCase): def test_transform(self, length=100000, fmin=50, fmax=22050, bins=12, fs=44100):