How to combine a 2D array with a scalar being repeated as many times as necessary?-CodePudding

I have x and y numpy arrays:

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

x = np.random.rand(3, 2)
y = np.random.rand(1)

Now I want to combine x and y in a way that x's shape becomes (x.shape[0] by x.shape[1] 1). Since, y is a scalar, I want y to be repeated x.shape[0] times, so the dimensions make sense.

Is there any "broadcasted" way of doing it in numpy without hardcoding like:

np.concatenate((x, np.array([y,y,y])), axis=1)

CodePudding user response：

Manual broadcast:

>>> np.concatenate((x, np.broadcast_to(y, (x.shape[0], 1))), axis=1)
array([[4.17022005e-01, 7.20324493e-01, 1.86260211e-01],
       [1.14374817e-04, 3.02332573e-01, 1.86260211e-01],
       [1.46755891e-01, 9.23385948e-02, 1.86260211e-01]])

Or use np.repeat:

>>> np.concatenate((x, np.repeat(y, x.shape[0])[:, None]), axis=1)
array([[4.17022005e-01, 7.20324493e-01, 1.86260211e-01],
       [1.14374817e-04, 3.02332573e-01, 1.86260211e-01],
       [1.46755891e-01, 9.23385948e-02, 1.86260211e-01]])

A little test on memory usage:

import numpy as np

from memory_profiler import profile


@profile
def broadcast_to(ar):
    ar = np.broadcast_to(ar, (50000000, 1))


@profile
def repeat(ar):
    ar = np.repeat(ar, 50000000)


if __name__ == '__main__':
    a = np.array([0])
    broadcast_to(a)
    repeat(a)

Output:

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
     7     52.3 MiB     52.3 MiB           1   @profile
     8                                         def broadcast_to(ar):
     9     52.4 MiB      0.1 MiB           1       ar = np.broadcast_to(ar, (50000000, 1))



Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    12     52.4 MiB     52.4 MiB           1   @profile
    13                                         def repeat(ar):
    14    243.1 MiB    190.8 MiB           1       ar = np.repeat(ar, 50000000)