CODEWITHSUNDEEP
pythonnumpy
Ziad Fakhoury
Ziad Fakhoury

Reputation: 197

Want to define an ndarray in numpy elementwise

I have 2 2d numpy arrays, A with shape (i,j) and B (i, k) where j >> k. I want to define a new 3d array C such that each element in C is the broadcasted element wise product of each column in A with the whole matrix B. In other words as a normal python loop I would do it like this

for x in range(j):
    C[x] = A[:,x]*B

However j is very large in this case and it would benefit me a lot if I am able to use Numpy's functionality to maybe define an ndarray C elementwise like in my loop above.

Thank you for your help

Upvotes: 1

Views: 55

Answers (1)

Brad Solomon
Brad Solomon

Reputation: 40878

You can use broadcasting like this:

a.T[:, :, None] * b

Example:

import numpy as np
np.random.seed(444)

i, j, k = 2, 10, 3
a = np.random.randn(i, j)
b = np.random.randn(i, k)
c = a.T[:, :, None] * b
print(c.shape)
# (10, 2, 3)

Transposing stems from the fact that you want to internally operate for each column in a, and [:, :, None] expands the dimensionality to enable broadcasting, as explained in NumPy's broadcasting rules.

Upvotes: 1

Related Questions