Source code for pycbc.fft.fftw_pruned

"""This module provides a functions to perform a pruned FFT based on FFTW

This should be considered a test and example module, as the functionality
can and should be generalized to other FFT backends, and precisions.

These functions largely implemented the generic FFT decomposition as
described rather nicely by wikipedia.

http://en.wikipedia.org/wiki/Cooley%E2%80%93Tukey_FFT_algorithm

I use a similar naming convention here, with minor simplifications to the
twiddle factors.
"""
import numpy, ctypes, pycbc.types
from pycbc.libutils import get_ctypes_library
import logging
from .fftw_pruned_cython import second_phase_cython

logger = logging.getLogger('pycbc.events.fftw_pruned')

warn_msg = ("The FFTW_pruned module can be used to speed up computing SNR "
            "timeseries by computing first at a low sample rate and then "
            "computing at full sample rate only at certain samples. This code "
            "has not yet been used in production, and has no test case. "
            "This was also ported to Cython in this state. "
            "This code would need verification before trusting results. "
            "Please do contribute test cases.")

logger.warning(warn_msg)

# FFTW constants
FFTW_FORWARD = -1
FFTW_BACKWARD = 1
FFTW_MEASURE = 0
FFTW_PATIENT = 1 << 5
FFTW_ESTIMATE = 1 << 6
float_lib = get_ctypes_library('fftw3f', ['fftw3f'],mode=ctypes.RTLD_GLOBAL)
fexecute = float_lib.fftwf_execute_dft
fexecute.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p]

ftexecute = float_lib.fftwf_execute_dft
ftexecute.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p]

[docs] def plan_transpose(N1, N2): """ Create a plan for transposing internally to the pruned_FFT calculation. (Alex to provide a write up with more details.) Parameters ----------- N1 : int Number of rows. N2 : int Number of columns. Returns -------- plan : FFTWF plan The plan for performing the FFTW transpose. """ rows = N1 cols = N2 iodim = numpy.zeros(6, dtype=numpy.int32) iodim[0] = rows iodim[1] = 1 iodim[2] = cols iodim[3] = cols iodim[4] = rows iodim[5] = 1 N = N1*N2 vin = pycbc.types.zeros(N, dtype=numpy.complex64) vout = pycbc.types.zeros(N, dtype=numpy.complex64) f = float_lib.fftwf_plan_guru_dft f.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int] f.restype = ctypes.c_void_p return f(0, None, 2, iodim.ctypes.data, vin.ptr, vout.ptr, None, FFTW_MEASURE)
[docs] def plan_first_phase(N1, N2): """ Create a plan for the first stage of the pruned FFT operation. (Alex to provide a write up with more details.) Parameters ----------- N1 : int Number of rows. N2 : int Number of columns. Returns -------- plan : FFTWF plan The plan for performing the first phase FFT. """ N = N1*N2 vin = pycbc.types.zeros(N, dtype=numpy.complex64) vout = pycbc.types.zeros(N, dtype=numpy.complex64) f = float_lib.fftwf_plan_many_dft f.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int] f.restype = ctypes.c_void_p return f(1, ctypes.byref(ctypes.c_int(N2)), N1, vin.ptr, None, 1, N2, vout.ptr, None, 1, N2, FFTW_BACKWARD, FFTW_MEASURE)
_theplan = None
[docs] def first_phase(invec, outvec, N1, N2): """ This implements the first phase of the FFT decomposition, using the standard FFT many plans. Parameters ----------- invec : array The input array. outvec : array The output array. N1 : int Number of rows. N2 : int Number of columns. """ global _theplan if _theplan is None: _theplan = plan_first_phase(N1, N2) fexecute(_theplan, invec.ptr, outvec.ptr)
[docs] def second_phase(invec, indices, N1, N2): """ This is the second phase of the FFT decomposition that actually performs the pruning. It is an explicit calculation for the subset of points. Note that there seem to be some numerical accumulation issues at various values of N1 and N2. Parameters ---------- invec : The result of the first phase FFT indices : array of ints The index locations to calculate the FFT N1 : int The length of the second phase "FFT" N2 : int The length of the first phase FFT Returns ------- out : array of floats """ invec = numpy.array(invec.data, copy=False) NI = len(indices) # pylint:disable=unused-variable N1 = int(N1) N2 = int(N2) out = numpy.zeros(len(indices), dtype=numpy.complex64) indices = numpy.array(indices, dtype=numpy.uint32) # Note, the next step if this needs to be faster is to invert the loops second_phase_cython(N1, N2, NI, indices, out, invec) return out
_thetransposeplan = None
[docs] def fft_transpose_fftw(vec): """ Perform an FFT transpose from vec into outvec. (Alex to provide more details in a write-up.) Parameters ----------- vec : array Input array. Returns -------- outvec : array Transposed output array. """ global _thetransposeplan outvec = pycbc.types.zeros(len(vec), dtype=vec.dtype) if _theplan is None: N1, N2 = splay(vec) _thetransposeplan = plan_transpose(N1, N2) ftexecute(_thetransposeplan, vec.ptr, outvec.ptr) return outvec
fft_transpose = fft_transpose_fftw
[docs] def splay(vec): """ Determine two lengths to split stride the input vector by """ N2 = 2 ** int(numpy.log2( len(vec) ) / 2) N1 = len(vec) / N2 return N1, N2
[docs] def pruned_c2cifft(invec, outvec, indices, pretransposed=False): """ Perform a pruned iFFT, only valid for power of 2 iffts as the decomposition is easier to choose. This is not a strict requirement of the functions, but it is unlikely to the optimal to use anything but power of 2. (Alex to provide more details in write up. Parameters ----------- invec : array The input vector. This should be the correlation between the data and the template at full sample rate. Ideally this is pre-transposed, but if not this will be transposed in this function. outvec : array The output of the first phase of the pruned FFT. indices : array of ints The indexes at which to calculate the full sample-rate SNR. pretransposed : boolean, default=False Used to indicate whether or not invec is pretransposed. Returns -------- SNRs : array The complex SNRs at the indexes given by indices. """ N1, N2 = splay(invec) if not pretransposed: invec = fft_transpose(invec) first_phase(invec, outvec, N1=N1, N2=N2) out = second_phase(outvec, indices, N1=N1, N2=N2) return out