<s> from __future__ import absolute_import , division , print_function <EOL> from functools import partial <EOL> import numpy as np <EOL> import numpy . fft as npfft <EOL> from . core import map_blocks <EOL> chunk_error = ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> fft_preamble = """<STR_LIT>""" <EOL> def _fft_wrap ( fft_func , dtype , out_chunk_fn ) : <EOL> def func ( a , n = None , axis = - <NUM_LIT:1> ) : <EOL> if len ( a . chunks [ axis ] ) != <NUM_LIT:1> : <EOL> raise ValueError ( chunk_error % ( axis , a . chunks [ axis ] ) ) <EOL> chunks = out_chunk_fn ( a , n , axis ) <EOL> return map_blocks ( partial ( fft_func , n = n , axis = axis ) , a , dtype = dtype , <EOL> chunks = chunks ) <EOL> np_name = fft_func . __name__ <EOL> func . __doc__ = ( fft_preamble % ( np_name , np_name ) ) + fft_func . __doc__ <EOL> func . __name__ = np_name <EOL> return func <EOL> def _fft_out_chunks ( a , n , axis ) : <EOL> """<STR_LIT>""" <EOL> if n is None : <EOL> return a . chunks <EOL> chunks = list ( a . chunks ) <EOL> chunks [ axis ] = ( n , ) <EOL> return chunks <EOL> def _rfft_out_chunks ( a , n , axis ) : <EOL> if n is None : <EOL> n = a . chunks [ axis ] [ <NUM_LIT:0> ] <EOL> chunks = list ( a . chunks ) <EOL> chunks [ axis ] = ( n // <NUM_LIT:2> + <NUM_LIT:1> , ) <EOL> return chunks <EOL> def _irfft_out_chunks ( a , n , axis ) : <EOL> if n is None : <EOL> n = <NUM_LIT:2> * ( a . chunks [ axis ] [ <NUM_LIT:0> ] - <NUM_LIT:1> ) <EOL> chunks = list ( a . chunks ) <EOL> chunks [ axis ] = ( n , ) <EOL> return chunks <EOL> def _hfft_out_chunks ( a , n , axis ) : <EOL> if n is None : <EOL> n = <NUM_LIT:2> * ( a . chunks [ axis ] [ <NUM_LIT:0> ] - <NUM_LIT:1> ) <EOL> chunks = list ( a . chunks ) <EOL> chunks [ axis ] = ( n , ) <EOL> return chunks <EOL> def _ihfft_out_chunks ( a , n , axis ) : <EOL> if n is None : <EOL> n = a . chunks [ axis ] [ <NUM_LIT:0> ] <EOL> chunks = list ( a . chunks ) <EOL> if n % <NUM_LIT:2> == <NUM_LIT:0> : <EOL> m = ( n // <NUM_LIT:2> ) + <NUM_LIT:1> <EOL> else : <EOL> m = ( n + <NUM_LIT:1> ) // <NUM_LIT:2> <EOL> chunks [ axis ] = ( m , ) <EOL> return chunks <EOL> fft = _fft_wrap ( npfft . fft , np . complex_ , _fft_out_chunks ) <EOL> ifft = _fft_wrap ( npfft . ifft , np . complex_ , _fft_out_chunks ) <EOL> rfft = _fft_wrap ( npfft . rfft , np . complex_ , _rfft_out_chunks ) <EOL> irfft = _fft_wrap ( npfft . irfft , np . float_ , _irfft_out_chunks ) <EOL> hfft = _fft_wrap ( npfft . hfft , np . float_ , _hfft_out_chunks ) <EOL> ihfft = _fft_wrap ( npfft . ihfft , np . complex_ , _ihfft_out_chunks ) </s>
