3

I want to multiply an array along it's first axis by some vector.

For instance, if a is 2D, b is 1D, and a.shape[0] == b.shape[0], we can do:

a *= b[:, np.newaxis]

What if a has an arbitrary shape? In numpy, the ellipsis "..." can be interpreted as "fill the remaining indices with ':'". Is there an equivalent for filling the remaining axes with None/np.newaxis?

The code below generates the desired result, but I would prefer a general vectorized way to accomplish this without falling back to a for loop.

from __future__ import print_function

import numpy as np


def foo(a, b):
    """
    Multiply a along its first axis by b
    """
    if len(a.shape) == 1:
        a *= b
    elif len(a.shape) == 2:
        a *= b[:, np.newaxis]
    elif len(a.shape) == 3:
        a *= b[:, np.newaxis, np.newaxis]
    else:
        n = a.shape[0]
        for i in range(n):
           a[i, ...] *= b[i]          

n = 10
b = np.arange(n)

a = np.ones((n, 3))
foo(a, b)
print(a)

a = np.ones((n, 3, 3))
foo(a, b)
print(a)
Improve this question
4
  • 2
    There's a lot of code here, would you be able to reduce this to a minimal reproducible example? Commented Mar 22, 2018 at 17:54
  • OK I've tried to remove as much boilerplate as possible while still conveying my desired intent. Commented Mar 22, 2018 at 19:29
  • I don't see any loops in your code. There's some if/else testing, but that's not a big deal. Many of the numpy functions are just as complicated, if not more so (e.g. np.atleast_2d). Getting rid of shape tests (or burying them in a function) is not vectorization. Commented Mar 22, 2018 at 21:09
  • @hpaulj - The else condition uses a loop (in my actual code this entire block is buried in another loop. Indeed, the vectorization provided by the first three cases in the if/elif block are just as good from a performance perspective as the accepted solution (in my tests thus far). The accepted solution is a lot more brief...I prefer those two lines of code to the 10 lines I had been using. Commented Mar 23, 2018 at 1:13

2 Answers 2

Reset to default
5

Just reverse the order of the axes:

transpose = a.T
transpose *= b

a.T is a transposed view of a, where "transposed" means reversing the order of the dimensions for arbitrary-dimensional a. We assign a.T to a separate variable so the *= doesn't try to set the a.T attribute; the results still apply to a, since the transpose is a view.

Demo:

In [55]: a = numpy.ones((2, 2, 3))

In [56]: a
Out[56]: 
array([[[1., 1., 1.],
        [1., 1., 1.]],

       [[1., 1., 1.],
        [1., 1., 1.]]])

In [57]: transpose = a.T

In [58]: transpose *= [2, 3]

In [59]: a
Out[59]: 
array([[[2., 2., 2.],
        [2., 2., 2.]],

       [[3., 3., 3.],
        [3., 3., 3.]]])
Improve this answer
Sign up to request clarification or add additional context in comments.

Comments

2

Following the idea of the accepted answer, you could skip the variable assignment to the transpose as follows:

arr = np.tile(np.arange(10, dtype=float), 3).reshape(3, 10)
print(arr)
factors = np.array([0.1, 1, 10])
arr.T[:, :] *= factors
print(arr)

Which would print

[[0. 1. 2. 3. 4. 5. 6. 7. 8. 9.]
 [0. 1. 2. 3. 4. 5. 6. 7. 8. 9.]
 [0. 1. 2. 3. 4. 5. 6. 7. 8. 9.]]
[[ 0.   0.1  0.2  0.3  0.4  0.5  0.6  0.7  0.8  0.9]
 [ 0.   1.   2.   3.   4.   5.   6.   7.   8.   9. ]
 [ 0.  10.  20.  30.  40.  50.  60.  70.  80.  90. ]]
Improve this answer

1 Comment

Good suggestion. However, your answer only works in 2D, while the question asks for arbitrary dimensions. To make it general you need to use the ellipsis arr.T[...] *= factors

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.