Combine Numpy "where" Statements

I am trying to speed up a code that is using Numpy's where() function. There are two calls to where(), which return an array of indices for where the statement is evaluated as True

Solution 1:

Does this help?

deffind_match2(x, y, z):
    returnlen(np.nonzero(np.logical_and(x == z, y == z))[0])

Sample run:

In [227]: print(find_match(X,Y,Z))
1000896

In [228]: print(find_match2(X,Y,Z))
1000896

In [229]: %timeit find_match(X,Y,Z)
2.37 s ± 70.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [230]: %timeit find_match2(X,Y,Z)
272 ms ± 9.64 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

I've added np.random.seed(210) before creating the arrays for the sake of reproducibility.

Solution 2:

Two versions that scale differently depending on size:

deffind_match1(x,y,z):
    return (x==y).astype(int) @ (y==z).astype(int) #equality and summation in one stepdeffind_match2(x,y,z):
    out = np.zeros_like(x)
    np.equal(x, y, out = out, where = np.equal(y, z)) #only calculates x==y if y==zreturn out.sum()

Testing different data sizes:

N = np.power(10, 7)
...

%timeit find_match(X,Y,Z)
206 ms ± 12.8 ms per loop (mean ± std. dev. of7 runs, 1 loop each)

%timeit find_match1(X,Y,Z)
70.7 ms ± 1.67 ms per loop (mean ± std. dev. of7 runs, 10 loops each)

%timeit find_match2(X,Y,Z)
74.7 ms ± 3.66 ms per loop (mean ± std. dev. of7 runs, 10 loops each)

N = np.power(10, 8)
...

%timeit find_match(X,Y,Z)
2.51 s ± 168 ms per loop (mean ± std. dev. of7 runs, 1 loop each)

%timeit find_match1(X,Y,Z)
886 ms ± 154 ms per loop (mean ± std. dev. of7 runs, 1 loop each)

%timeit find_match2(X,Y,Z)
776 ms ± 26.1 ms per loop (mean ± std. dev. of7 runs, 1 loop each)

EDIT: since @Tonechas's is faster than both, here's a numba method:

from numba import njit

@njitdeffind_match_jit(x, y, z):
    out = 0for i, j, k inzip(x, y, z):
        if i == j and j == k:
            out += 1return out

find_match_jit(X,Y,Z)  #run it once to compile
Out[]: 1001426

%timeit find_match_jit(X,Y,Z) # N = 10**8204 ms ± 13.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

If threading is allowed:

@njit(parallel = True)
def find_match_jit_p(x, y, z):
    xy = x == yyz= y == z    
    return np.logical_and(xy, yz).sum()
        

find_match_jit_p(X,Y,Z)
Out[]: 1001426

%timeit find_match_jit_p(X,Y,Z)84.6 ms ± 2.31 ms per loop(mean ± std. dev. of 7 runs, 10 loops each)