I wanna find the indices of the rows that all have entries smaller than 1e-6 or where the number of nonzero values is less than 3. Something like this would be nice:
import numpy as np
prob = np.random.rand(15, 500)
all_zero = np.where(prob.max(1) < 1e-6 | np.nonzero(prob, axis=1) < 3)
I tried to measure the execution times of the solutions proposed so far:
Benchmark data:
prob = np.random.rand(10000, 500)
%%timeit
[i for i, val in enumerate(prob>1e-6)if val.sum() < 3]
# 39.5 ms ± 1.4 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%%timeit
np.where(np.sum(prob>1e-6, axis=1) < 3)
# 9.92 ms ± 199 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%%timeit
all_zero = np.logical_or(prob.max(axis=1) < 1e-6, np.sum(prob != 0, axis=1) < 3)
np.where(all_zero)
# 13.9 ms ± 150 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
The most efficient solution seems to be the second one.
You can use np.logical_or and np.sum the non-zero values to check which row has fewer than 3 non-zero elements:
all_zero = np.logical_or(prob.max(axis=1) < 1e-6, np.sum(prob != 0, axis=1) < 3)
This code returns the list of index of rows with less than 3 values other than 0 (less than 1e-6):
[i for i, val in enumerate(prob>1e-6) if val.sum()<3]
or using only numpy functions:
np.where(np.sum(prob>1e-6, axis=1)<3)
nonzero, i.e.np.nonzero(prob).sum(axis=1).