Source

pypy / lib_pypy / numpypy / fft / fftpack_lite.py

#!/usr/bin/env python

import numpy as np

# This is a replacement for numpy's fftpack_litemodule.c
# written by mikefc for numpypy. targetting fftpack with cffi.

#-------------------------------------------------------------------
# CFFI interface to fftpack.c
#-------------------------------------------------------------------
import cffi
ffi = cffi.FFI()

fftpack_h = """
void cfftf(int N, double data[], const double wrk[]);
void cfftb(int N, double data[], const double wrk[]);
void cffti(int N, double wrk[]);
void rfftf(int N, double data[], const double wrk[]);
void rfftb(int N, double data[], const double wrk[]);
void rffti(int N, double wrk[]);
"""

ffi.cdef(fftpack_h)
fftpack_c = ffi.dlopen("./libfftpack.so.1.0")

#-------------------------------------------------------------------
# 1D FFT of real values
#-------------------------------------------------------------------
def rffti(n):
    """Init function for real-valued fft
    @param n length of fft to be performed.
    @return appropriately sized working array 
    """
    wrk = np.zeros(2 * n + 15, dtype='f8')
    wrk_ptr = ffi.cast("double *", wrk.__array_interface__['data'][0])
    fftpack_c.rffti(n, wrk_ptr)
    return wrk

def rfftf(a, wrk):
    """Perform an fft on 'a' assuming that it only contains real values
    @param a real-valued 1D array
    @wrk a previously initialised working array
    @return complex ndarray 
    """
    a = a.copy() # C code works in-place. Avoid overwriting
    n = len(a)
    a_ptr   = ffi.cast("double *",   a.__array_interface__['data'][0])
    wrk_ptr = ffi.cast("double *", wrk.__array_interface__['data'][0])
    fftpack_c.rfftf(n, a_ptr, wrk_ptr)
    # The result from fftpack_c.rfftf needs to be unscrambled to create
    # the complex array. Each pair of doubles represents a complex
    # number, except for the first value.
    # i.e. a is [r, r1, i1, r2, i2, ....]
    #    res = [r, r1+i1, r2+i2, ... r]
    if n % 2 == 0:
        # Even length ffts have a real term at 0 and n//2
        # the terms 1..n//2 are the conj of the terms n//2+1..n
        res = np.zeros(n//2+1, dtype='complex128')
        res[0]    = a[0]
        res[1:n//2] = a[1:n-1:2] + 1j * a[2:n:2]
        res[n//2] = a[-1]
    else:
        # Odd length ffts have a real term at 0
        # r, c1, c2 ... cn, conj(cn), ... conj(c2), conj(c1)
        res = np.zeros((n+1)//2, dtype='complex128')
        res[0]    = a[0]
        res[1:(n+1)//2] = a[1:n-1:2] + 1j * a[2:n:2]
    return res

def rfftb(a, wrk):
    # reconfigure input for fftpack
    # a = [r, r1+i1, r2+i2, ...] ==> [r, r1, i1, r2, i2, ...]

    n = (len(wrk) - 15) // 2
    res = np.zeros(n, dtype='f8')
    res[0] = np.real(a[0])
    if n % 2 != 0:
        res[1:n  :2] = np.real(a[1:len(a)])
        res[2:n  :2] = np.imag(a[1:len(a)])
    else:
        res[1:n-1:2] = np.real(a[1:len(a)-1])
        res[2:n  :2] = np.imag(a[1:len(a)-1])
        res[-1] = np.real(a[-1])
    res_ptr = ffi.cast("double *", res.__array_interface__['data'][0])
    wrk_ptr = ffi.cast("double *", wrk.__array_interface__['data'][0])
    fftpack_c.rfftb(n, res_ptr, wrk_ptr)
    return res 

#-------------------------------------------------------------------
# 1D FFT of general (complex) values
#-------------------------------------------------------------------
def cffti(n):
    """Init function for real-valued fft
    @param n length of fft to be performed.
    @return appropriately sized working array 
    """
    wrk = np.zeros(4 * n + 15, dtype='f8')
    wrk_ptr = ffi.cast("double *", wrk.__array_interface__['data'][0])
    fftpack_c.cffti(n, wrk_ptr)
    return wrk

def cfftf(a, wrk):
    """ Compute the general forwards complex fft """
    n = len(a)  # n complex numbers
    n2 = 2 * n  # represented by 2*n doubles
    
    # unpack complex values for c
    # [r1+i1, r2+i2, .... rn+in]
    #   ==>  [r1, i1, r2, i2, .... rn, in]
    res = np.zeros(n2, dtype='f8')
    res[0:n2-1:2] = np.real(a)
    res[1:n2  :2] = np.imag(a)

    res_ptr = ffi.cast("double *", res.__array_interface__['data'][0])
    wrk_ptr = ffi.cast("double *", wrk.__array_interface__['data'][0])
    fftpack_c.cfftf(n, res_ptr, wrk_ptr)

    # pack complex values for python
    # [r1, i1, r2, i2, .... rn, in]
    #   ==>  [r1+i1, r2+i2, .... rn+in]
    res2 = np.zeros(n, dtype='complex128')
    res2 = res[0:n2-1:2] + 1j * res[1:n2  :2]
    return res2

def cfftb(a, wrk):
    """Compute the general backwards complex fft
    Identical to cfftf except for fftpack_c function call
    """
    n = len(a)  # n complex numbers
    n2 = 2 * n  # represented by 2*n doubles

    # unpack complex values for c
    # [r1+i1, r2+i2, .... rn+in]
    #   ==>  [r1, i1, r2, i2, .... rn, in]
    res = np.zeros(n2, dtype='f8')
    res[0:n2-1:2] = np.real(a)
    res[1:n2  :2] = np.imag(a)

    res_ptr = ffi.cast("double *", res.__array_interface__['data'][0])
    wrk_ptr = ffi.cast("double *", wrk.__array_interface__['data'][0])
    fftpack_c.cfftb(n, res_ptr, wrk_ptr)

    # pack complex values for python
    # [r1, i1, r2, i2, .... rn, in]
    #   ==>  [r1+i1, r2+i2, .... rn+in]
    res2 = np.zeros(n, dtype='complex128')
    res2 = res[0:n2-1:2] + 1j * res[1:n2  :2]
    return res2
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.