Cython numpy array indexer speed improvement

Question

I wrote the following code in pure python, the description of what it does is in the docstrings:

import numpy as np
from scipy.ndimage.measurements import find_objects
import itertools

def alt_indexer(arr):

    """
    Returns a dictionary with the elements of arr as key
    and the corresponding slice as value.

    Note:

        This function assumes arr is sorted.

    Example:

        >>> arr = [0,0,3,2,1,2,3]
        >>> loc = _indexer(arr)
        >>> loc
        {0: (slice(0L, 2L, None),),
        1: (slice(2L, 3L, None),),
        2: (slice(3L, 5L, None),),
        3: (slice(5L, 7L, None),)}
        >>> arr = sorted(arr)
        >>> arr[loc[3][0]]
        [3, 3]
        >>> arr[loc[2][0]]
        [2, 2]

    """

    unique, counts = np.unique(arr, return_counts=True)
    labels = np.arange(1,len(unique)+1)
    labels = np.repeat(labels,counts)

    slicearr = find_objects(labels)
    index_dict = dict(itertools.izip(unique,slicearr))

    return index_dict

Since i will be indexing very large arrays, i wanted to speed up the operations by using cython, here is the equivalent implementation:

import numpy as np
cimport numpy as np

def _indexer(arr):

    cdef tuple unique_counts = np.unique(arr, return_counts=True)
    cdef np.ndarray[np.int32_t,ndim=1] unique = unique_counts[0]
    cdef np.ndarray[np.int32_t,ndim=1] counts = unique_counts[1].astype(int)

    cdef int start=0
    cdef int end
    cdef int i
    cdef dict d ={}

    for i in xrange(len(counts)):
        if i>0:
            start = counts[i-1]+start
        end=counts[i]+start
        d[unique[i]]=slice(start,end)
    return d

Benchmarks

I compared the time it took to complete both operations:

In [26]: import numpy as np

In [27]: rr=np.random.randint(0,1000,1000000)

In [28]: %timeit _indexer(rr)
10 loops, best of 3: 40.5 ms per loop

In [29]: %timeit alt_indexer(rr) #pure python
10 loops, best of 3: 51.4 ms per loop

As you can see the speed improvements are minimal. I do realize that my code was already partly optimized since i used numpy.

Is there a bottleneck that i am not aware of? Should i not use np.unique and write my own implementation instead?

Thanks.

cython loops are faster if they can be translated to pure C. In your case the loop still uses numpy.unique and Python dictionary and slice objects. — hpaulj

Divakar Divakar · Accepted Answer · 2016-10-16T16:52:27

With arr having non-negative, not very large and many repeated int numbers, here's an alternative approach using np.bincount to simulate the same behavior as np.unique(arr, return_counts=True) -

def unique_counts(arr):
    counts = np.bincount(arr)
    mask = counts!=0
    unique = np.nonzero(mask)[0]
    return unique, counts[mask]

Runtime test

Case #1 :

In [83]: arr = np.random.randint(0,100,(1000)) # Input array

In [84]: unique, counts = np.unique(arr, return_counts=True)
    ...: unique1, counts1 = unique_counts(arr)
    ...: 

In [85]: np.allclose(unique,unique1)
Out[85]: True

In [86]: np.allclose(counts,counts1)
Out[86]: True

In [87]: %timeit np.unique(arr, return_counts=True)
10000 loops, best of 3: 53.2 µs per loop

In [88]: %timeit unique_counts(arr)
100000 loops, best of 3: 10.2 µs per loop

Case #2:

In [89]: arr = np.random.randint(0,1000,(10000)) # Input array

In [90]: %timeit np.unique(arr, return_counts=True)
1000 loops, best of 3: 713 µs per loop

In [91]: %timeit unique_counts(arr)
10000 loops, best of 3: 39.1 µs per loop

Case #3: Let's run a case with unique having some missing numbers in the min to max range and verify the results against np.unique version as a sanity check. We won't have a lot of repeated numbers in this case and as such isn't expected to be better on performance.

In [98]: arr = np.random.randint(0,10000,(1000)) # Input array

In [99]: unique, counts = np.unique(arr, return_counts=True)
    ...: unique1, counts1 = unique_counts(arr)
    ...: 

In [100]: np.allclose(unique,unique1)
Out[100]: True

In [101]: np.allclose(counts,counts1)
Out[101]: True

In [102]: %timeit np.unique(arr, return_counts=True)
10000 loops, best of 3: 61.9 µs per loop

In [103]: %timeit unique_counts(arr)
10000 loops, best of 3: 71.8 µs per loop

Cython numpy array indexer speed improvement

Benchmarks

1 Answers