Power spectrum with Cython

Question

I am trying to optimize my code with Cython. It is doing a a power spectrum, not using FFT, because this is what we were told to do in class. I've tried to write to code in Cython, but do not see any difference. Here is my code

#! /usr/bin/env python
# -*- coding: utf8 -*-

from __future__ import division
cimport numpy as np
import numpy as np
cimport cython

@cython.boundscheck(False)
def power_spectrum(time, data, double f_min, double f_max, double df,w=1 ):

    cdef double com,f
    cdef double s,c,sc,cc,ss
    cdef np.ndarray[double, ndim=1] power
    cdef np.ndarray[double, ndim=1] freq

    alfa, beta = [],[] 
    m = np.mean(data)
    data -= m       

    freq = np.arange( f_min,f_max,df )
    for f in freq:
            sft = np.sin(2*np.pi*f*time)
            cft = np.cos(2*np.pi*f*time)
            s   = np.sum( w*data*sft )
            c   = np.sum( w*data*cft )
            ss  = np.sum( w*sft**2  )
            cc  = np.sum( w*cft**2  )
            sc  = np.sum( w*sft*cft )

            alfa.append( ( s*cc-c*sc )/( ss*cc-sc**2 ))
            beta.append( ( c*ss-s*sc )/( ss*cc-sc**2 ))
            com = -(f-f_min)/(f_min-f_max)*100
            print "%0.3f%% complete" %com

    power = np.array(alfa)**2 + np.array(beta)**2
    return freq,power,alfa,beta

The time and data is loaded via numpy.loadtxt and sent to this function. When I do

cython -a power_spectrum.pyx

the .html file is very yellow, so not very efficient. Especially the entire for-loop and the calulation of power and returning everything.

I have tried to read the official guide to Cython, but as I have never coded in C it is somewhat hard to understand.

All help is very preciated :)

Answer 1

Cython can read numpy arrays according to this but it won't magically compile stuff like np.sum - you are still just calling the numpy methods.

What you need to do is re-write your inner loop in pure cython which can then compile it for you. So you will need to re-implement np.sum, np.sin etc. Pre-allocating aplfa and beta is a good idea so you don't use append and try to cdef as many variables as possible.

EDIT

Here is an complete example showing the inner loop completely C compiled (no yellow). I've no idea whether the code is correct but it should be a good starting point! In particular note use of cdef everywhere, turning on cdivision and import of sin and cos from the standard library.

from __future__ import division
cimport numpy as np
import numpy as np
cimport cython
from math import pi

cdef extern from "math.h":
    double cos(double theta)
    double sin(double theta)

@cython.boundscheck(False)
@cython.cdivision(True)
def power_spectrum(np.ndarray[double, ndim=1] time, np.ndarray[double, ndim=1] data, double f_min, double f_max, double df, double w=1 ):

    cdef double com,f
    cdef double s,c,sc,cc,ss,t,d
    cdef double twopi = 6.283185307179586
    cdef np.ndarray[double, ndim=1] power
    cdef np.ndarray[double, ndim=1] freq = np.arange( f_min,f_max,df )
    cdef int n = len(freq)
    cdef np.ndarray[double, ndim=1] alfa = np.zeros(n)
    cdef np.ndarray[double, ndim=1] beta = np.zeros(n)
    cdef int ndata = len(data)
    cdef int i, j

    m = np.mean(data)
    data -= m       

    for i in range(ndata):
        f = freq[i]

        s = 0.0
        c = 0.0
        ss = 0.0
        cc = 0.0
        sc = 0.0
        for j in range(n):
            t = time[j]
            d = data[j]
            sf = sin(twopi*f*t)
            cf = cos(twopi*f*t)
            s += w*d*sf
            c += w*d*cf
            ss += w*sf**2
            cc += w*cf**2
            sc += w*sf*cf

        alfa[i] = ( s*cc-c*sc )/( ss*cc-sc**2 )
        beta[i] = ( c*ss-s*sc )/( ss*cc-sc**2 )

    power = np.array(alfa)**2 + np.array(beta)**2
    return freq,power,alfa,beta