diff options
Diffstat (limited to 'numpy/fft/helper.py')
| -rw-r--r-- | numpy/fft/helper.py | 116 |
1 files changed, 8 insertions, 108 deletions
diff --git a/numpy/fft/helper.py b/numpy/fft/helper.py index 4b698bb4d..a920a4ac0 100644 --- a/numpy/fft/helper.py +++ b/numpy/fft/helper.py @@ -4,14 +4,9 @@ Discrete Fourier Transforms - helper.py """ from __future__ import division, absolute_import, print_function -import collections -try: - import threading -except ImportError: - import dummy_threading as threading from numpy.compat import integer_types from numpy.core import integer, empty, arange, asarray, roll -from numpy.core.overrides import array_function_dispatch +from numpy.core.overrides import array_function_dispatch, set_module # Created by Pearu Peterson, September 2002 @@ -24,7 +19,7 @@ def _fftshift_dispatcher(x, axes=None): return (x,) -@array_function_dispatch(_fftshift_dispatcher) +@array_function_dispatch(_fftshift_dispatcher, module='numpy.fft') def fftshift(x, axes=None): """ Shift the zero-frequency component to the center of the spectrum. @@ -52,7 +47,7 @@ def fftshift(x, axes=None): -------- >>> freqs = np.fft.fftfreq(10, 0.1) >>> freqs - array([ 0., 1., 2., 3., 4., -5., -4., -3., -2., -1.]) + array([ 0., 1., 2., ..., -3., -2., -1.]) >>> np.fft.fftshift(freqs) array([-5., -4., -3., -2., -1., 0., 1., 2., 3., 4.]) @@ -81,7 +76,7 @@ def fftshift(x, axes=None): return roll(x, shift, axes) -@array_function_dispatch(_fftshift_dispatcher) +@array_function_dispatch(_fftshift_dispatcher, module='numpy.fft') def ifftshift(x, axes=None): """ The inverse of `fftshift`. Although identical for even-length `x`, the @@ -128,6 +123,7 @@ def ifftshift(x, axes=None): return roll(x, shift, axes) +@set_module('numpy.fft') def fftfreq(n, d=1.0): """ Return the Discrete Fourier Transform sample frequencies. @@ -161,7 +157,7 @@ def fftfreq(n, d=1.0): >>> timestep = 0.1 >>> freq = np.fft.fftfreq(n, d=timestep) >>> freq - array([ 0. , 1.25, 2.5 , 3.75, -5. , -3.75, -2.5 , -1.25]) + array([ 0. , 1.25, 2.5 , ..., -3.75, -2.5 , -1.25]) """ if not isinstance(n, integer_types): @@ -174,9 +170,9 @@ def fftfreq(n, d=1.0): p2 = arange(-(n//2), 0, dtype=int) results[N:] = p2 return results * val - #return hstack((arange(0,(n-1)/2 + 1), arange(-(n/2),0))) / (n*d) +@set_module('numpy.fft') def rfftfreq(n, d=1.0): """ Return the Discrete Fourier Transform sample frequencies @@ -214,7 +210,7 @@ def rfftfreq(n, d=1.0): >>> sample_rate = 100 >>> freq = np.fft.fftfreq(n, d=1./sample_rate) >>> freq - array([ 0., 10., 20., 30., 40., -50., -40., -30., -20., -10.]) + array([ 0., 10., 20., ..., -30., -20., -10.]) >>> freq = np.fft.rfftfreq(n, d=1./sample_rate) >>> freq array([ 0., 10., 20., 30., 40., 50.]) @@ -226,99 +222,3 @@ def rfftfreq(n, d=1.0): N = n//2 + 1 results = arange(0, N, dtype=int) return results * val - - -class _FFTCache(object): - """ - Cache for the FFT twiddle factors as an LRU (least recently used) cache. - - Parameters - ---------- - max_size_in_mb : int - Maximum memory usage of the cache before items are being evicted. - max_item_count : int - Maximum item count of the cache before items are being evicted. - - Notes - ----- - Items will be evicted if either limit has been reached upon getting and - setting. The maximum memory usages is not strictly the given - ``max_size_in_mb`` but rather - ``max(max_size_in_mb, 1.5 * size_of_largest_item)``. Thus the cache will - never be completely cleared - at least one item will remain and a single - large item can cause the cache to retain several smaller items even if the - given maximum cache size has been exceeded. - """ - def __init__(self, max_size_in_mb, max_item_count): - self._max_size_in_bytes = max_size_in_mb * 1024 ** 2 - self._max_item_count = max_item_count - self._dict = collections.OrderedDict() - self._lock = threading.Lock() - - def put_twiddle_factors(self, n, factors): - """ - Store twiddle factors for an FFT of length n in the cache. - - Putting multiple twiddle factors for a certain n will store it multiple - times. - - Parameters - ---------- - n : int - Data length for the FFT. - factors : ndarray - The actual twiddle values. - """ - with self._lock: - # Pop + later add to move it to the end for LRU behavior. - # Internally everything is stored in a dictionary whose values are - # lists. - try: - value = self._dict.pop(n) - except KeyError: - value = [] - value.append(factors) - self._dict[n] = value - self._prune_cache() - - def pop_twiddle_factors(self, n): - """ - Pop twiddle factors for an FFT of length n from the cache. - - Will return None if the requested twiddle factors are not available in - the cache. - - Parameters - ---------- - n : int - Data length for the FFT. - - Returns - ------- - out : ndarray or None - The retrieved twiddle factors if available, else None. - """ - with self._lock: - if n not in self._dict or not self._dict[n]: - return None - # Pop + later add to move it to the end for LRU behavior. - all_values = self._dict.pop(n) - value = all_values.pop() - # Only put pack if there are still some arrays left in the list. - if all_values: - self._dict[n] = all_values - return value - - def _prune_cache(self): - # Always keep at least one item. - while len(self._dict) > 1 and ( - len(self._dict) > self._max_item_count or self._check_size()): - self._dict.popitem(last=False) - - def _check_size(self): - item_sizes = [sum(_j.nbytes for _j in _i) - for _i in self._dict.values() if _i] - if not item_sizes: - return False - max_size = max(self._max_size_in_bytes, 1.5 * max(item_sizes)) - return sum(item_sizes) > max_size |