Implement saving and loading of pyFFTW wisdom for substantial speedup

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)

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):

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)

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):