Python vectorization, how to get all index for every row with numpy

Question

I'm having a hard time trying to solve this problem, the main issue is I'm running a simulation, so for lops are mainly forbidden, I have a numpy array NxN, in this case mine is about (10000x20).

stoploss = 19.9 # condition to apply
monte_carlo_simulation(20,1.08,10000,20) #which gives me that 10000x20 np array
mask_trues = np.where(np.any((simulation <= stoploss) == True, axis=1)) # boolean mask

I need some code to make a new vector of len(10000) which returns an array with all the positions for every row, lets suppose:

function([[False,True,True],[False,False,True]])
output = [[1,2],[2]]

Again, the main problem resides in not using loops.

Answer 1

Simply this:

list(map(np.where, my_array))

performance comparison against Kasrâmvd's solution:

def f(a):
    return list(map(np.where, a))

def g(a):
    x, y = np.where(a)
    return np.split(y, np.where(np.diff(x) != 0)[0] + 1)

a = np.random.randint(2, size=(10000,20))

%timeit f(a)
%timeit g(a)


7.66 ms ± 38.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
13.3 ms ± 188 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Answer 2

For completeness I'll demonstrate a sparse matrix approach:

In [57]: A = np.array([[False,True,True],[False,False,True]])
In [58]: A
Out[58]: 
array([[False,  True,  True],
       [False, False,  True]])
In [59]: M = sparse.lil_matrix(A)
In [60]: M
Out[60]: 
<2x3 sparse matrix of type '<class 'numpy.bool_'>'
    with 3 stored elements in LInked List format>
In [61]: M.data
Out[61]: array([list([True, True]), list([True])], dtype=object)
In [62]: M.rows
Out[62]: array([list([1, 2]), list([2])], dtype=object)

And to make a large sparse one:

In [63]: BM = sparse.random(10000,20,.05, 'lil')
In [64]: BM
Out[64]: 
<10000x20 sparse matrix of type '<class 'numpy.float64'>'
    with 10000 stored elements in LInked List format>
In [65]: BM.rows
Out[65]: 
array([list([3]), list([]), list([6, 15]), ..., list([]), list([11]),
       list([])], dtype=object)

Rough time tests:

In [66]: arr = BM.A
In [67]: timeit sparse.lil_matrix(arr)
19.5 ms ± 421 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [68]: timeit list(map(np.where,arr))
11 ms ± 55.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [69]: %%timeit
    ...: x,y = np.where(arr)
    ...: np.split(y, np.where(np.diff(x) != 0)[0] + 1)
    ...: 
13.8 ms ± 24.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Generating a csr sparse format matrix is faster:

In [70]: timeit sparse.csr_matrix(arr)
2.68 ms ± 120 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [71]: Mr = sparse.csr_matrix(arr)
In [72]: Mr.indices
Out[72]: array([ 3,  6, 15, ...,  8, 16, 11], dtype=int32)
In [73]: Mr.indptr
Out[73]: array([    0,     1,     1, ...,  9999, 10000, 10000], dtype=int32)
In [74]: np.where(arr)[1]
Out[74]: array([ 3,  6, 15, ...,  8, 16, 11])

It's indices is just like the column where, while the indptr is like the split indices.