Cython code for adaptive binning

Question

Here is my Cython code used for adaptive binning. The calcAdaptiveBinnedRateMap function is called from another Python script. The script is compiled using Cython but the speed I am expecting is still not great. How can I improve the speed of execution?

import numpy as np
import scipy.ndimage.morphology as ndimmor

#define sampline rate and alpha
cdef float samplingRate = 30.0
cdef double alpha = 0.0001 #skaggs and sachin use this value while Jim uses 0.001

"""
runs the iteration for adaptive binning till the criteria mentioned in Skaggs et al 1996 is met
INPUT: spike map, occupancy map, alpha, Number of occupancy (Nocc), Euclidean distance transform (dists)
OUTPUT: Nspikes (number of spikes), Nocc (occupancy count)
"""
def mexAdaptwhile(spikeMap, occMap, alpha, Nocc, dists):
    #initialize the variable
    cdef int Nspikes2 = 1
    cdef double rsq = 0
    cdef int EnoughPoints = 0
    cdef int Nspikes = 0

    #find the row and column count for occupancy map
    cdef int rowLen = int(np.shape(occMap)[0])
    cdef int colLen = int(np.shape(occMap)[1])

    #while the radius is less than 200 and enough points are not covered
    while rsq<200.00 and EnoughPoints==0:
        r = np.sqrt(rsq)
        #need to set these Nocc and Nspikes to zero, otherwise the two for loops above add the same spikes over again.
        Nspikes = 0
        Nocc = 0
        for i in range(rowLen):
            for j in range(colLen):
                #if the distance is less than radius
                if dists[i,j]<=r:
                    #add Nspikes
                    Nspikes = Nspikes + spikeMap[i,j]
                    #add occupancy
                    Nocc = Nocc +occMap[i,j]         
        #if number spikes is greater than 0, then set nspikes2 = number of spikes
        if Nspikes > 0:
            Nspikes2 = Nspikes
        #check for the condition from skaggs et al 1996
        if (alpha*alpha*Nocc*Nocc*rsq*Nspikes2 > 1):
            EnoughPoints = 1 #set the flag to end the loop
        #keep increasing the radius
        rsq = rsq + 1
        #output occupancy = occupancy
        Noccout = Nocc
    return Nspikes, Noccout


"""calculate adaptive binned rate map    
INPUT: spikemap, occupancy map (both of them just binned for 2cm/4cm no other operation applied on them)
NOTE: occupancy is still in terms of number of frames
OUTPUT: adaptive binned rate Map which is color adjusted
"""
def calcAdaptiveBinnedRateMap(spikeMap, occMap):
    #set unoccupied occupancy and spike corresponding to unoccupied position = 0
    spikeMap[0,0] = 0
    occMap[0,0] = 0

    #find the row, col index of minimum occupied and maximum occupied pixel
    row, col = np.where(occMap)
    minrow = np.min(row)
    maxrow = np.max(row)
    mincol = np.min(col)
    maxcol = np.max(col)

    #select the best fitting rectangle according to occupied area
    occMap = occMap[minrow:maxrow+1,mincol:maxcol+1]
    spikeMap = spikeMap[minrow:maxrow+1,mincol:maxcol+1]

    #matrix of zeros same size as of occupancy map
    z = np.zeros(np.shape(occMap))
    #variable to hold adaptive binned rate map value
    abrMap = np.copy(z)
    #variale to hold adaptive binned occupancy map
    abrOcc = np.copy(z)

    #check to endure if number of spikes is greater than 1
    if np.max(np.max(spikeMap))>0:
        #iterate over the values
        for x in range(int(np.shape(occMap)[1])):
            for y in range(int(np.shape(occMap)[0])):
                if occMap[y,x] > 0:
                    #pretend there's atleas 1 spike, and 1 occ.needed to avoid 0 threshold. 
                    Nspikes2 = 1
                    Nocc = occMap[y,x]
                    d = np.copy(z)
                    d[y,x] = 1
                    #computes the Euclidean distance transform of the input matrix.
                    #For each pixel in BW, the distance transform assigns a number that is the 
                    #distance between that pixel and the nearest nonzero pixel of BW. 
                    dists = ndimmor.distance_transform_edt(d==0)
                    # function to keep on iterating while the condition mentioned in Skaggs et al 1996 is met
                    Nspikes, Nocc = mexAdaptwhile(spikeMap, occMap, alpha, Nocc, dists)
                    if Nocc < 12: #occupancy cutoff = 0.4seconds 
                        #if less than 0.4 seconds set it to 0
                        abrMap[y,x] = 0
                        abrOcc[y,x] = 0
                    else:
                        #else equal to number of spikes/occupancy map
                        abrMap[y,x] = samplingRate*float(Nspikes)/float(Nocc)
                        #adaptive binned ocuupancy map = nocc
                        abrOcc[y,x] = Nocc

    #find the maximum value of adaptive binned rate map
    cmax = np.max(np.max(abrMap))
    if cmax > 0:
        #minimum = maximum value found above/60
        cmin = -(cmax/60.0);
    else:
        cmin = -1;
    #set adaptive binned rate map = cmin wherever adaptive binned occupancy map = 0 
    abrMap[abrOcc==0] = cmin

    #return the adaptive binned rate map
    return abrMap

Answer 1

Make sure to declare all variables with Cython (otherwise the lines with such variables basically run with Python-speed). Especially important for the "for"-loop variables (i.e. i, j, x, y).

One can also use cython -a mycode.pyx to create an annotated HTML page that shows, which lines are running with C or Python speed, respectively.