signal.py 文件源码

def fft_1d_loop(arr, axis=-1):
    """Like scipy.fft.pack.fft and numpy.fft.fft, perform fft along an axis.
    Here do this by looping over remaining axes and perform 1D FFTs.

    This was implemented as a low-memory version like
    :func:`~pwtools.crys.smooth` to be used in :func:`~pwtools.pydos.pdos`,
    which fills up the memory for big MD data. But actually it has the same
    memory footprint as the plain scipy fft routine. Keep it here anyway as a
    nice reference for how to loop over remaining axes in the ndarray case.
    """
    if axis < 0:
        axis = arr.ndim - 1
    axes = [ax for ax in range(arr.ndim) if ax != axis]
    # tuple here is 3x faster than generator expression
    #   idxs = (range(arr.shape[ax]) for ax in axes)  
    idxs = tuple(range(arr.shape[ax]) for ax in axes)
    out = np.empty(arr.shape, dtype=complex)
    for idx_tup in product(*idxs):
        sl = [slice(None)] * arr.ndim
        for idx,ax in izip(idx_tup, axes):
            sl[ax] = idx
        out[sl] = fft(arr[sl])
    return out