Using Cython to speed up connected components algorithm

Question

First off, I am using python[2.7.2], numpy[1.6.2rc1], cython[0.16], gcc[MinGW] compiler, on a windows xp machine.

I needed a 3D connected components algorithm to process some 3D binary data (i.e. 1s and 0s) stored in numpy arrays. Unfortunately, I could not find any existing code so I adapted the code found here to work with 3D arrays. Everything works great, however speed is desirable for processing huge data sets. As a result I stumbled upon cython and decided to give it a try.

So far cython has improved the speed: Cython: 0.339 s Python: 0.635 s

Using cProfile, my time consuming line in the pure python version is:

new_region = min(filter(lambda i: i > 0, array_region[xMin:xMax,yMin:yMax,zMin:zMax].ravel()))

---

The Question: What is the correct way to "cythonize" the lines:

new_region = min(filter(lambda i: i > 0, array_region[xMin:xMax,yMin:yMax,zMin:zMax].ravel()))
for x,y,z in zip(ind[0],ind[1],ind[2]):

Any help would be appreciated and hopefully this work will help others.

---

Pure python version [*.py]:

import numpy as np

def find_regions_3D(Array):
    x_dim=np.size(Array,0)
    y_dim=np.size(Array,1)
    z_dim=np.size(Array,2)
    regions = {}
    array_region = np.zeros((x_dim,y_dim,z_dim),)
    equivalences = {}
    n_regions = 0
    #first pass. find regions.
    ind=np.where(Array==1)
    for x,y,z in zip(ind[0],ind[1],ind[2]):

        # get the region number from all surrounding cells including diagnols (27) or create new region                        
        xMin=max(x-1,0)
        xMax=min(x+1,x_dim-1)
        yMin=max(y-1,0)
        yMax=min(y+1,y_dim-1)
        zMin=max(z-1,0)
        zMax=min(z+1,z_dim-1)

        max_region=array_region[xMin:xMax+1,yMin:yMax+1,zMin:zMax+1].max()

        if max_region > 0:
            #a neighbour already has a region, new region is the smallest > 0
            new_region = min(filter(lambda i: i > 0, array_region[xMin:xMax+1,yMin:yMax+1,zMin:zMax+1].ravel()))
            #update equivalences
            if max_region > new_region:
                if max_region in equivalences:
                    equivalences[max_region].add(new_region)
                else:
                    equivalences[max_region] = set((new_region, ))
        else:
            n_regions += 1
            new_region = n_regions

        array_region[x,y,z] = new_region


    #Scan Array again, assigning all equivalent regions the same region value.
    for x,y,z in zip(ind[0],ind[1],ind[2]):
        r = array_region[x,y,z]
        while r in equivalences:
            r= min(equivalences[r])
        array_region[x,y,z]=r

    #return list(regions.itervalues())
    return array_region

Pure python speedups:

#Original line:
new_region = min(filter(lambda i: i > 0, array_region[xMin:xMax+1,yMin:yMax+1,zMin:zMax+1].ravel()))

#ver A:
new_region = array_region[xMin:xMax+1,yMin:yMax+1,zMin:zMax+1]
min(new_region[new_region>0])

#ver B:
new_region = min( i for i in array_region[xMin:xMax,yMin:yMax,zMin:zMax].ravel() if i>0)

#ver C:
sub=array_region[xMin:xMax,yMin:yMax,zMin:zMax]
nlist=np.where(sub>0)
minList=[]
for x,y,z in zip(nlist[0],nlist[1],nlist[2]):
    minList.append(sub[x,y,z])
new_region=min(minList)

Time results:
O: 0.0220445
A: 0.0002161
B: 0.0173195
C: 0.0002560

---

Cython version [*.pyx]:

import numpy as np
cimport numpy as np

DTYPE = np.int
ctypedef np.int_t DTYPE_t

cdef inline int int_max(int a, int b): return a if a >= b else b
cdef inline int int_min(int a, int b): return a if a <= b else b

def find_regions_3D(np.ndarray Array not None):
    cdef int x_dim=np.size(Array,0)
    cdef int y_dim=np.size(Array,1)
    cdef int z_dim=np.size(Array,2)
    regions = {}
    cdef np.ndarray array_region = np.zeros((x_dim,y_dim,z_dim),dtype=DTYPE)
    equivalences = {}
    cdef int n_regions = 0
    #first pass. find regions.
    ind=np.where(Array==1)
    cdef int xMin, xMax, yMin, yMax, zMin, zMax, max_region, new_region, x, y, z
    for x,y,z in zip(ind[0],ind[1],ind[2]):

        # get the region number from all surrounding cells including diagnols (27) or create new region                        
        xMin=int_max(x-1,0)
        xMax=int_min(x+1,x_dim-1)+1
        yMin=int_max(y-1,0)
        yMax=int_min(y+1,y_dim-1)+1
        zMin=int_max(z-1,0)
        zMax=int_min(z+1,z_dim-1)+1

        max_region=array_region[xMin:xMax,yMin:yMax,zMin:zMax].max()

        if max_region > 0:
            #a neighbour already has a region, new region is the smallest > 0
            new_region = min(filter(lambda i: i > 0, array_region[xMin:xMax,yMin:yMax,zMin:zMax].ravel()))
            #update equivalences
            if max_region > new_region:
                if max_region in equivalences:
                    equivalences[max_region].add(new_region)
                else:
                    equivalences[max_region] = set((new_region, ))
        else:
            n_regions += 1
            new_region = n_regions

        array_region[x,y,z] = new_region


    #Scan Array again, assigning all equivalent regions the same region value.
    cdef int r
    for x,y,z in zip(ind[0],ind[1],ind[2]):
        r = array_region[x,y,z]
        while r in equivalences:
            r= min(equivalences[r])
        array_region[x,y,z]=r

    #return list(regions.itervalues())
    return array_region

Cython speedups:

Using:

cdef np.ndarray region = np.zeros((3,3,3),dtype=DTYPE)
...
        region=array_region[xMin:xMax,yMin:yMax,zMin:zMax]
        new_region=np.min(region[region>0])

Time: 0.170, original: 0.339 s

---

Results

After considering the many useful comments and answers provided, my current algorithms are running at:
Cython: 0.0219
Python: 0.4309

Cython is providing a 20x increase in speed over the pure python.

Current Cython Code:

import numpy as np
import cython
cimport numpy as np
cimport cython

from libcpp.map cimport map

DTYPE = np.int
ctypedef np.int_t DTYPE_t

cdef inline int int_max(int a, int b): return a if a >= b else b
cdef inline int int_min(int a, int b): return a if a <= b else b

@cython.boundscheck(False)
def find_regions_3D(np.ndarray[DTYPE_t,ndim=3] Array not None):
    cdef unsigned int x_dim=np.size(Array,0),y_dim=np.size(Array,1),z_dim=np.size(Array,2)
    regions = {}
    cdef np.ndarray[DTYPE_t,ndim=3] array_region = np.zeros((x_dim,y_dim,z_dim),dtype=DTYPE)
    cdef np.ndarray region = np.zeros((3,3,3),dtype=DTYPE)
    cdef map[int,int] equivalences
    cdef unsigned int n_regions = 0

    #first pass. find regions.
    ind=np.where(Array==1)
    cdef np.ndarray[DTYPE_t,ndim=1] ind_x = ind[0], ind_y = ind[1], ind_z = ind[2]
    cells=range(len(ind_x))
    cdef unsigned int xMin, xMax, yMin, yMax, zMin, zMax, max_region, new_region, x, y, z, i, xi, yi, zi, val
    for i in cells:

        x=ind_x[i]
        y=ind_y[i]
        z=ind_z[i]

        # get the region number from all surrounding cells including diagnols (27) or create new region                        
        xMin=int_max(x-1,0)
        xMax=int_min(x+1,x_dim-1)+1
        yMin=int_max(y-1,0)
        yMax=int_min(y+1,y_dim-1)+1
        zMin=int_max(z-1,0)
        zMax=int_min(z+1,z_dim-1)+1

        max_region = 0
        new_region = 2000000000 # huge number
        for xi in range(xMin, xMax):
            for yi in range(yMin, yMax):
                for zi in range(zMin, zMax):
                    val = array_region[xi,yi,zi]
                    if val > max_region: # val is the new maximum
                        max_region = val

                    if 0 < val < new_region: # val is the new minimum
                        new_region = val

        if max_region > 0:
           if max_region > new_region:
                if equivalences.count(max_region) == 0 or new_region < equivalences[max_region]:
                    equivalences[max_region] = new_region
        else:
           n_regions += 1
           new_region = n_regions

        array_region[x,y,z] = new_region


    #Scan Array again, assigning all equivalent regions the same region value.
    cdef int r
    for i in cells:
        x=ind_x[i]
        y=ind_y[i]
        z=ind_z[i]

        r = array_region[x,y,z]
        while equivalences.count(r) > 0:
            r= equivalences[r]
        array_region[x,y,z]=r

    return array_region

Setup file [setup.py]

from distutils.core import setup
from distutils.extension import Extension
from Cython.Distutils import build_ext
import numpy

setup(
    cmdclass = {'build_ext': build_ext},
    ext_modules = [Extension("ConnectComp", ["ConnectedComponents.pyx"],
                             include_dirs =[numpy.get_include()],
                             language="c++",
                             )]
)

Build command:

python setup.py build_ext --inplace

Answer 1

As @gotgenes points out, you should definitely be using cython -a <file>, and trying to reduce the amount of yellow you see. Yellow corresponds to worse and worse generated C.

Things I found that reduced the amount of yellow:

1. This looks like a situation where there will never be any out of bounds array access, as long as the input Array has 3 dimensions, so one can turn off bounds checking:

   cimport cython
   
   @cython.boundscheck(False)
   def find_regions_3d(...):
   

2. Give the compiler more information for efficient indexing, i.e. whenever you cdef an ndarray give as much information as you can:

    def find_regions_3D(np.ndarray[DTYPE_t,ndim=3] Array not None):
        [...]
        cdef np.ndarray[DTYPE_t,ndim=3] array_region = ...
        [etc.]
   

3. Give the compiler more information about positive/negative-ness. I.e. if you know a certain variable is always going to be positive, cdef it as unsigned int rather than int, as this means that Cython can eliminate any negative-indexing checks.

4. Unpack the ind tuple immediately, i.e.

   ind = np.where(Array==1)
   cdef np.ndarray[DTYPE_t,ndim=1] ind_x = ind[0], ind_y = ind[1], ind_z = ind[2]
   

5. Avoid using the for x,y,z in zip(..[0],..[1],..[2]) construct. In both cases, replace it with

   cdef int i
   for i in range(len(ind_x)):
       x = ind_x[i]
       y = ind_y[i]
       z = ind_z[i]
   

6. Avoid doing the fancy indexing/slicing. And especially avoid doing it twice! And avoid using filter! I.e. replace

   max_region=array_region[xMin:xMax,yMin:yMax,zMin:zMax].max()
   if max_region > 0:
       new_region = min(filter(lambda i: i > 0, array_region[xMin:xMax,yMin:yMax,zMin:zMax].ravel()))
       if max_region > new_region:
           if max_region in equivalences:
               equivalences[max_region].add(new_region)
           else:
               equivalences[max_region] = set((new_region, ))
   

   with the more verbose

   max_region = 0
   new_region = 2000000000 # "infinity"
   for xi in range(xMin, xMax):
       for yi in range(yMin, yMax):
           for zi in range(zMin, zMax):
               val = array_region[xi,yi,zi]
               if val > max_region: # val is the new maximum
                   max_region = val
   
               if 0 < val < new_region: # val is the new minimum
                   new_region = val
   
   if max_region > 0:
      if max_region > new_region:
          if max_region in equivalences:
              equivalences[max_region].add(new_region)
          else:
              equivalences[max_region] = set((new_region, ))
   else:
      n_regions += 1
      new_region = n_regions
   

   This doesn't look so nice, but the triple loop compiles down to about 10 or so lines of C, while the compiled version of the original is hundreds of lines long and has a lot of Python object manipulation.

   (Obviously you must cdef all the variables you use, especially xi, yi, zi and val in this code.)

7. You don't need to store all the equivalences, since the only thing you do with the set is find the minimum element. So if you instead have equivalences mapping int to int, you can replace

   if max_region in equivalences:
       equivalences[max_region].add(new_region)
   else:
       equivalences[max_region] = set((new_region, ))
   
   [...]
   
   while r in equivalences:
       r = min(equivalences[r])
   

   with

   if max_region not in equivalences or new_region < equivalences[max_region]:
       equivalences[max_region] = new_region
   
   [...]
   
   while r in equivalences:
       r = equivalences[r]
   

8. The last thing to do after all that would be to not use any Python objects at all, specifically, don't use a dictionary for equivalences. This is now easy, since it is mapping int to int, so one could use from libcpp.map cimport map and then cdef map[int,int] equivalences, and replace .. not in equivalences with equivalences.count(..) == 0 and .. in equivalences with equivalences.count(..) > 0. (Note that it will then require a C++ compiler.)

Answer 2

(copied from the above comment for others ease of reading)

I believe scipy's ndimage.label does what you want (I did not test it against your code but it should be quite efficient). Note that you have to import it explicitely:

from scipy import ndimage 
ndimage.label(your_data, connectivity_struct)

then later you can apply other built-in functions (like finding the bounding rectangle, centre-of-mass, etc)