CODEWITHSUNDEEP
pythonnumpyslice
Tohiko
Tohiko

Reputation: 1982

Append np.newaxis as needed

I want to slice an array so that I can use it to perform an operation with another array of arbitrary dimension. In other words, I am doing the following:

A = np.random.rand(5)
B = np.random.rand(5,2,3,4)
slicer = [slice(None)] + [None]*(len(B.shape)-1)
result = B*A[slicer]

Is there some syntax that I can use so that I do not have to construct slicer?

Upvotes: 1

Views: 102

Answers (1)

Daniel F
Daniel F

Reputation: 14399

In this specific case you can use np.einsum with an ellipsis.

result2 = np.einsum('i,i...->i...', A, B)

np.allclose(result, result2)

Out[232]: True

Although, as @hpaulj points out this only works for multiplication (or division if you use 1/B).

Since broadcasting works from the other end normally, you can use np.transpose twice get the axes in the right order.

result3 = np.transpose(np.transpose(B) * A)

But that's also not a general case

Upvotes: 3

Related Questions