numpy np.where scan array once to get both sets of indices corresponding to T/F condition

Question

I have an array and I need to get the indices satisfying both where a condition is true, and where the same condition is false, e.g.:

x = np.random.rand(100000000)
true_inds = np.where(x < 0.5)
false_inds = np.where(x >= 0.5)

In my use case x is quite large, this code is called inside a loop, and profiling has revealed that np.where is actually the bottleneck. I'm currently doing something analogous to the above code, which unnecessarily scans x twice to get the two sets of indices. Is it possible to get both true_inds and false_inds with just one scan of x without implementing a specialized replacement for np.where from scratch?

Answer 1

For operand sizes from about 1500 upwards splitting the result of a stable argsort appears a good solution (maybe two times faster, though towards very large sizes it appears to become less).

If you have pythran installed a more consistent speedup can be obtained (numba should be similar).

Notes:

1. It is important to use stable sort, i.e. kind="stable", when omitted perfomance gets much worse on top of returned indices being unordered

2. I suspect this requires a fairly recent numpy version, since they just added new and type specific sort algorithms.

3. Some of the solutions plotted return indices in arbitrary order, but speedups gained if any are fairly small

Code to produce the plot, comment out the pythran related stuff if necessary:

from simple_benchmark import BenchmarkBuilder, MultiArgument
import numpy as np
from scipy.sparse import csr_matrix
from biwhere_pthr import biwhere as use_pythran, \
    biwhere_halfordered as use_pythran_half, \
    biwhere_unordered as use_pythran_un

B = BenchmarkBuilder()

@B.add_function(alias="nonzero")
def use_nonzero(mask):
    return mask.nonzero()[0], (~mask).nonzero()[0]

@B.add_function(alias="argpartition")
def use_partition(mask):
    nf = mask.size - np.count_nonzero(mask)
    ft = mask.argpartition(nf)
    return ft[nf:],ft[:nf]

@B.add_function(alias="argsort")
def use_sort(mask):
    nf = mask.size - np.count_nonzero(mask)
    ft = mask.argsort()
    return ft[nf:],ft[:nf]

@B.add_function(alias="argsort stable")
def use_stable_sort(mask):
    nf = mask.size - np.count_nonzero(mask)
    ft = mask.argsort(kind="stable")
    return ft[nf:],ft[:nf]

@B.add_function(alias="sparse")
def use_sparse(mask):
    aux = csr_matrix((mask,mask,np.arange(mask.size+1)),(mask.size,2)).tocsc()
    return aux.indices[aux.indptr[1]:],aux.indices[:aux.indptr[1]]

B.add_function(alias="pythran")(use_pythran)
B.add_function(alias="pythran up down")(use_pythran_half)
B.add_function(alias="pythran unordered")(use_pythran_un)

@B.add_arguments('array size')
def argument_provider():
    for exp in range(8, 27):
        sz = int(2**exp)
        yield sz,np.random.randint(0,2,sz,dtype=bool)

# checks
for sz,mask in argument_provider():
    ref = use_nonzero(mask)
    for f in use_stable_sort,use_sparse,use_pythran:
        if not all(map(np.array_equal,f(mask),ref)):
            print(sz,f.__name__)
    for f in (use_partition,use_sort,use_pythran_un):
        if not all(map(np.array_equal,map(np.sort,f(mask)),ref)):
            print(sz,f.__name__)
    ht,hf = use_pythran_half(mask)
    if not all(map(np.array_equal,(ht[::-1],hf),ref)):
        print(sz,"use_pythran_half")

r = B.run()
r.plot(relative_to=use_nonzero)

import pylab
pylab.savefig('biwhere.png')

pythran module compile using `pythran -O3 :

import numpy as np

#pythran export biwhere(bool[:])
#pythran export biwhere_halfordered(bool[:])
#pythran export biwhere_unordered(bool[:])

def biwhere(mask):
    nt = np.count_nonzero(mask)
    f,t = np.empty(mask.size-nt,int),np.empty(nt,int)
    i = 0
    j = 0
    for k,m in enumerate(mask):
        if m:
            t[j] = k
            j += 1
        else:
            f[i] = k
            i += 1
    return t,f

def biwhere_halfordered(mask):
    ft = np.empty(mask.size,int)
    i = 0
    j = mask.size-1
    for k,m in enumerate(mask):
        if m:
            ft[j] = k
            j -= 1
        else:
            ft[i] = k
            i += 1
    return ft[i:],ft[:i]

def biwhere_unordered(mask):
    ft = np.empty(mask.size,int)
    i = 0
    j = mask.size-1
    while i<=j:
        if not mask[i]:
            ft[i] = i
            i += 1
        elif mask[j]:
            ft[j] = j
            j -= 1
        else:
            ft[i] = j
            ft[j] = i
            i += 1
            j -= 1
    return ft[i:],ft[:i]