CODEWITHSUNDEEP
pythonnumpy
TeenyTinySparkles
TeenyTinySparkles

Reputation: 419

Elementwise operations in numpy

I have a 4000 x 16 matrix A. I have 256 x 1 vector B. I need to get the elementwise addition for each element in A with every element of B and get a 3D array of size 4000 x 16 x 256. What is the most efficient way to achieve this without loops with numpy?

Upvotes: 2

Views: 252

Answers (2)

DarkDrassher34
DarkDrassher34

Reputation: 59

This is where numpy's broadcasting and np.transpose() features really shine! 

A[:,:, np.newaxis] + B[np.newaxis,:,:].transpose(2,0,1)

Upvotes: 4

willeM_ Van Onsem
willeM_ Van Onsem

Reputation: 476594

You can first reshape A to a 4000×16×1 matrix, and B to 1×1×256 matrix. Then you can perform addition:

A[:,:,None] + B.reshape(1, 1, -1)

For A a 4×5 matrix, and B a 3×1 matrix, for example, we get:

>>> A
array([[1, 2, 2, 2, 2],
       [2, 1, 0, 3, 2],
       [4, 1, 2, 1, 3],
       [3, 2, 3, 2, 2]])
>>> B
array([[0],
       [3],
       [3],
       [6],
       [4],
       [3]])
>>> A[:,:,None] + B.reshape(1, 1, -1)
array([[[ 1,  4,  4,  7,  5,  4],
        [ 2,  5,  5,  8,  6,  5],
        [ 2,  5,  5,  8,  6,  5],
        [ 2,  5,  5,  8,  6,  5],
        [ 2,  5,  5,  8,  6,  5]],

       [[ 2,  5,  5,  8,  6,  5],
        [ 1,  4,  4,  7,  5,  4],
        [ 0,  3,  3,  6,  4,  3],
        [ 3,  6,  6,  9,  7,  6],
        [ 2,  5,  5,  8,  6,  5]],

       [[ 4,  7,  7, 10,  8,  7],
        [ 1,  4,  4,  7,  5,  4],
        [ 2,  5,  5,  8,  6,  5],
        [ 1,  4,  4,  7,  5,  4],
        [ 3,  6,  6,  9,  7,  6]],

       [[ 3,  6,  6,  9,  7,  6],
        [ 2,  5,  5,  8,  6,  5],
        [ 3,  6,  6,  9,  7,  6],
        [ 2,  5,  5,  8,  6,  5],
        [ 2,  5,  5,  8,  6,  5]]])

Performance: If I run the above 100 times with A a 4000×16 matrix, and B a 256×1 matrix of floating points, I get the following results:

>>> timeit(lambda: A[:,:,None] + B.reshape(1, 1, -1), number=100)
5.949456596048549

It thus takes roughly 59.49457ms for a single run. This looks reasonable to calculate 16'384'000 elements.

Upvotes: 4

Related Questions