Apply a function to only array slices which satisfy a condition (NumPy)

Question

I have an array A[i,j]. The last index contains various input values for a function myfunc that will be applied for each i and produce an output B[i]. However, many values indexed by j will not contribute to B, so I would like to avoid unnecessary calls of myfunc. This can be accomplished by slicing out the relevant values with conditional indexing such as C = C[C>mythreshold] with for loops relatively easily as in the MWE below:

def myfunc(X):
    return np.square(X).sum()
A = np.floor(np.random.rand(3,4)*100)
mythreshold = 10
(N1, N2) = A.shape
B = np.zeros(N1)
for i in range(N1):
    C = A[i,:]
    C = C[C>mythreshold]
    B[i] = myfunc(C)

I had to break this up into for loops so I could remove slices of A without removing slices of the full array. This was since I cannot drop elements of A[i,:] for one i without dropping the corresponding elements for another i. For speed, however, I would like to vectorize wherever possible - to avoid for loops and do this for all i in one go. How can I do this?

Note: That was an MWE; the actual case has larger array dimensions, so that my arrays would be A[i,j,k,l] and B[i,j], so the for loop example would be something like the code below. I think the extra dimensions wouldn't complicate things, but it's worth mentioning just in case.

(N1, N2, N3, N4) = A.shape
for i in range(N1):
    for j in range(N2):
        C = A[i,j,:,:].flatten()
        C = C[C>mythreshold]
        B[i,j] = myfunc(C)

Answer 1

In [10]: A = np.floor(np.random.rand(3,4)*2*mythreshold)
In [11]: A
Out[11]: 
array([[14.,  4.,  1.,  8.],
       [11., 11.,  4.,  2.],
       [ 8.,  6., 18., 12.]])
In [12]: (N1, N2) = A.shape
    ...: B = np.zeros(N1)
    ...: for i in range(N1):
    ...:     C = A[i,:]
    ...:     C = C[C>mythreshold]
    ...:     B[i] = myfunc(C)
    ...: 
In [13]: B
Out[13]: array([196., 242., 468.])

A threshold test for the whole array:

In [14]: A>mythreshold
Out[14]: 
array([[ True, False, False, False],
       [ True,  True, False, False],
       [False, False,  True,  True]])

Make a copy, and set the other values to 0 (or something innocuous):

In [15]: A1 = A.copy(); A1[A<=mythreshold]=0
In [16]: np.square(A1).sum(axis=1)
Out[16]: array([196., 242., 468.])

That's doesn't avoid applying the function to all elements, but it avoids iterating on rows. Usually avoiding python level loops speeds up numpy code. But if your function can't be "vectorized" as I do with the axis parameter, or where it is so complex that including those "0" values is expensive, then this is not the way to go.

If your function includes ufunc, you could use its where parameter

In [17]: mask = A>mythreshold
In [18]: out = np.zeros_like(A)
In [19]: np.square(A, out=out, where=mask)
Out[19]: 
array([[196.,   0.,   0.,   0.],
       [121., 121.,   0.,   0.],
       [  0.,   0., 324., 144.]])
In [20]: _.sum(axis=1)
Out[20]: array([196., 242., 468.])

Usually this where is used when certain values give bad results, like divide by 0 or log of a negative. I don't think it's a time saver, but haven't done timings to confirm this.