Cython numpy array indexer speed improvement
<p>I wrote the following code in pure python, the description of what it does is in the docstrings:</p>

<pre><code>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:

        &gt;&gt;&gt; arr = [0,0,3,2,1,2,3]
        &gt;&gt;&gt; loc = _indexer(arr)
        &gt;&gt;&gt; loc
        {0: (slice(0L, 2L, None),),
        1: (slice(2L, 3L, None),),
        2: (slice(3L, 5L, None),),
        3: (slice(5L, 7L, None),)}
        &gt;&gt;&gt; arr = sorted(arr)
        &gt;&gt;&gt; arr[loc[3][0]]
        [3, 3]
        &gt;&gt;&gt; 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
</code></pre>

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

<pre><code>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&gt;0:
            start = counts[i-1]+start
        end=counts[i]+start
        d[unique[i]]=slice(start,end)
    return d
</code></pre>

<h1>Benchmarks</h1>

<p>I compared the time it took to complete both operations:</p>

<pre><code>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
</code></pre>

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

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

<p>Thanks.</p>

<p>With <code>arr</code> having non-negative, not very large and many repeated <code>int</code> numbers, here's an alternative approach using <a href="https://docs.scipy.org/doc/numpy/reference/generated/numpy.bincount.html" rel="nofollow"><code>np.bincount</code></a> to simulate the same behavior as <code>np.unique(arr, return_counts=True)</code> -</p>

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

<p><strong>Runtime test</strong></p>

<p>Case #1 :</p>

<pre><code>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
</code></pre>

<p>Case #2:</p>

<pre><code>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
</code></pre>

<p>Case #3: Let's run a case with <code>unique</code> having some missing numbers in the min to max range and verify the results against <code>np.unique</code> 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.</p>

<pre><code>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
</code></pre>

python performance numpy cython