Reputation: 1507
I was trying to locate the index of a value on a numpy array (same for hundreds of them) using np.where. Its not that slow but its a bottleneck so i began experimenting with fortran and f2py and wrote this simple routine.
subroutine find_index(array1, num_elements, target_value, loc)
real, intent(in) :: array1(:)
integer, intent(in) :: num_elements, target_value
integer, intent(out) :: loc
do i = 1, num_elements
if (array1(i) .eq. target_value) then
loc = i
exit
endif
end do
end subroutine
But still no improvement (same as np.where). So i guess its about the method. any sugestions on improving the code (python or fortran)?
EDIT the values i am searching for are integers in an array of integers
Upvotes: 0
Views: 295
Reputation: 231540
It's been sometime since I worked with fortran and f2py, but I did something similar with cython
last year.
In a Hungarian algorithm search problem, I needed to find the first 0 value in a 2d array, subject to row and column masking arrays.
So using where
(argwhere
is just np.transpose(np.where(...))
, the function was:
def find_a_zero(self):
# find first uncovered 0 cost
rc = np.argwhere((self.cost + self.rc[:,None] + self.cc) == 0)
if rc.shape[0]>0:
return tuple(rc[0])
else:
return None, None
I got a good speedup with argmax
:
def find_a_zero(self):
# big help time wise, 16->10 for n=200
cond = (self.cost + self.rc[:,None] + self.cc) == 0
if np.count_nonzero(cond):
idx = np.unravel_index(np.argmax(cond), cond.shape)
return idx
return None, None
np.where
uses count_nonzero
to determine the size of its return arrays. argmax
, when operating on a boolean, short circuits on the fist True
.
I got even better speed with this cython
version:
cdef find_a_zero(int[:,:] cost, int[:] row_cover, int[:] col_cover):
n = cost.shape[0]
m = cost.shape[1]
cdef size_t r, c
for r in range(n):
for c in range(m):
if (cost[r,c]==0) and (row_cover[r]==0) and (col_cover[c]==0):
row = r
col = c
return r, c
return -1, -1
If my memory of cython
is correct, definitions like int[:,:] cost
invoke its typed memoryview. which has efficient lowlevel array operations.
http://cython.readthedocs.io/en/latest/src/userguide/memoryviews.html
http://csclab.murraystate.edu/~bob.pilgrim/445/munkres.html
Upvotes: 1