CODEWITHSUNDEEP
arraysnumpymultidimensional-arrayarray-broadcasting
cheyp
cheyp

Reputation: 391

Broadcast function that changes dimension of the input array

Given some function f that accepts 1D array and gives 2D array, is it possible to apply it efficiently for each row of the NxM array A?

More specifically, I want to apply np.triu for each of the row of the NxM array A and then concatenate all the results. I can achieve this by

B = np.dstack(map(np.triu, A))

which gives MxMxN matrix. However, this is not very efficiently for large N. Unfortunately, the function np.apply_along_axis cannot be employed here because f changes dimension.

Knowing the power of NumPy for efficient broadcasting, I am almost sure that there exists a better solution for my problem.

Upvotes: 1

Views: 114

Answers (1)

Divakar
Divakar

Reputation: 221504

Here's a vectorized approach using broadcasting -

Bout = A.T*(np.tri(A.shape[1],dtype=bool).T[...,None])

Runtime test and output verification -

In [319]: A = np.random.randint(0,20,(400,100))

In [320]: %timeit np.dstack(map(np.triu, A))
10 loops, best of 3: 69.9 ms per loop

In [321]: %timeit A.T*(np.tri(A.shape[1],dtype=bool).T[...,None])
10 loops, best of 3: 24.8 ms per loop

In [322]: B = np.dstack(map(np.triu, A))

In [323]: Bout = A.T*(np.tri(A.shape[1],dtype=bool).T[...,None])

In [324]: np.allclose(B,Bout)
Out[324]: True

Upvotes: 2

Related Questions