Cython: How to declare numpy.argwhere()

Question

I tried to Cythonize part of my code as following to hopefully gain some speed:

# cython: boundscheck=False
import numpy as np
cimport numpy as np
import time

cpdef object my_function(np.ndarray[np.double_t, ndim = 1] array_a,
                     np.ndarray[np.double_t, ndim = 1] array_b,
                     int n_rows,
                     int n_columns):
    cdef double minimum_of_neighbours, difference, change
    cdef int i
    cdef np.ndarray[np.int_t, ndim =1] locations
    locations = np.argwhere(array_a > 0)

    for i in locations:
        minimum_of_neighbours = min(array_a[i - n_columns], array_a[i+1], array_a[i + n_columns], array_a[i-1])
        if array_a[i] - minimum_of_neighbours < 0:
            difference = minimum_of_neighbours - array_a[i]
            change = min(difference, array_a[i] / 5.)
            array_a[i] += change
            array_b[i] -= change * 5.
        print time.time()

return array_a, array_b

I can compile it without an error but when I use the function I got this error:

from cythonized_code import my_function
import numpy as np

array_a = np.random.uniform(low=-100, high=100, size = 100).astype(np.double)
array_b = np.random.uniform(low=0, high=20, size = 100).astype(np.double)

a, b = my_function(array_a,array_b,5,20)

# which gives me this error:    
# locations = np.argwhere(array_a > 0)
# ValueError: Buffer has wrong number of dimensions (expected 1, got 2)

Do I need to declare locations type here? The reason I wanted to declare it is that it has yellow colour in the annotated HTML file generated by compiling the code.

Answer 1

It's a good idea not to use the python-functionality locations[i], because it is too expensive: Python would create a full-fledged Python-integer* from the lowly c-integer (which is what is stored in the locations-numpy array), register it in the garbage collector, then cast it back to int, destroy the Python-object - quite an overhead.

To get a direct access to the lowly c-integers one needs to bind locations to a type. The normal course of action would be too look up, which properties locations has:

>>> locations.ndim
2
>>> locations.dtype
dtype('int64')

which translates to cdef np.ndarray[np.int64_t, ndim =2] locations.

However, this will (probably, can not check it right now) not be enough to get rid of Python-overhead because of a Cython-quirk:

for i in locations:
    ...

will not be interpreted as a raw-array access but will invoke the Python-machinery. See for example here.

So you will have to change it to:

for index in range(len(locations)):
      i=locations[index][0]

then Cython "understands", that you want the access to the raw c-int64 array.

---

• Actually, it is not completely true: In this case first an nd.array is created (e.g. locations[0] or locations[1]) and then __Pyx_PyInt_As_int (which is more or less an alias for [PyLong_AsLongAndOverflow][2]) is called, which creates a PyLongObject, from which C-int value is obtained before the temporary PyLongObject and nd.array are destructed.

Here we get lucky, because length-1 numpy-arrays can be converted to Python scalars. The code would not work if the second dimension of locations would be >1.