numpy argmax with max less than some number

Question

I have a numpy array as:

myArray
array([[ 1.    ,     nan,     nan,     nan,     nan],
       [ 1.    ,     nan,     nan,     nan,     nan],
       [ 0.63  ,  0.79  ,  1.    ,     nan,     nan],
       [ 0.25  ,  0.4   ,  0.64  ,  0.84  ,     nan]])

I need to find for each row, the column numbers for max value but the max has to be less than 1.

In the above array, row 0,1 should return Nan.

Row 2 should return 1.

Row 3 should return 3.

I am not sure how to condition this on argmax.

Answer 1

Here's one approach with np.where -

m = a < 1  # Mask of elems < 1 and non-NaNs

# Set NaNs and elems > 1 to global minimum values minus 1, 
# so that when used with argmax those would be ignored
idx0 = np.where(m, a,np.nanmin(a)-1).argmax(1)

# Look for rows with no non-NaN and < 1 elems and set those in o/p as NaNs
idx = np.where(m.any(1), idx0, np.nan)

Sample run -

In [97]: a
Out[97]: 
array([[ 1.  ,   nan,   nan,   nan,   nan],
       [ 1.  ,   nan,   nan,   nan,   nan],
       [ 0.63,  0.79,  1.  ,   nan,   nan],
       [ 0.25,  0.4 ,  0.64,  0.84,   nan]])

In [98]: m = a < 1

In [99]: idx0 = np.where(m, a,np.nanmin(a)-1).argmax(1)

In [100]: idx0
Out[100]: array([0, 0, 1, 3])

In [101]: np.where(m.any(1), idx0, np.nan)
Out[101]: array([ nan,  nan,   1.,   3.])