np.concatenate a ND tensor/array with a 1D array

Question

I have two arrays a & b

a.shape
(5, 4, 3)
array([[[ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.10772717,  0.604584  ,  0.41664413]],

       [[ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.10772717,  0.604584  ,  0.41664413],
        [ 0.95879616,  0.85575133,  0.46135877]],

       [[ 0.        ,  0.        ,  0.        ],
        [ 0.10772717,  0.604584  ,  0.41664413],
        [ 0.95879616,  0.85575133,  0.46135877],
        [ 0.70442301,  0.74126523,  0.88965603]],

       [[ 0.10772717,  0.604584  ,  0.41664413],
        [ 0.95879616,  0.85575133,  0.46135877],
        [ 0.70442301,  0.74126523,  0.88965603],
        [ 0.8039435 ,  0.62802183,  0.58885027]],

       [[ 0.95879616,  0.85575133,  0.46135877],
        [ 0.70442301,  0.74126523,  0.88965603],
        [ 0.8039435 ,  0.62802183,  0.58885027],
        [ 0.95848603,  0.72429311,  0.71461332]]])

and b

array([ 0.79212707,  0.66629398,  0.58676553], dtype=float32)
b.shape
(3,)

I want to get array

ab.shape
(5,5,3)

I do as below first

b = b.reshape(1,1,3)

then

b=np.concatenate((b, b,b, b, b), axis = 0)

And

ab=np.concatenate((a, b), axis = 1)
ab.shape
(5, 5, 3)

I get the right result, but it's not very convenient especially at the step

b=np.concatenate((b, b,b, b, b), axis = 0)

when I have to type many times (the real dataset has much dimensions). Are there any faster ways to come to this result?

Answer 1

Simply broadcast b to 3D and then concatenate along second axis -

b3D = np.broadcast_to(b,(a.shape[0],1,len(b)))
out = np.concatenate((a,b3D),axis=1)

The broadcasting part with np.broadcast_to doesn't actual replicate or make copies and is simply a replicated view and then in the next step, we do the concatenation that does the replication on-the-fly.

Benchmarking

We are comparing np.repeat version from @cᴏʟᴅsᴘᴇᴇᴅ's solution against np.broadcast_to one in this section with focus on performance. The broadcasting based one does the replication and concatenation in the second step, as a merged command so to speak, while np.repeat version makes copy and then concatenates in two separate steps.

Timing the approaches as whole :

Case #1 : a = (500,400,300) and b = (300,)

In [321]: a = np.random.rand(500,400,300)

In [322]: b = np.random.rand(300)

In [323]: %%timeit
     ...: b3D = b.reshape(1, 1, -1).repeat(a.shape[0], axis=0)
     ...: r = np.concatenate((a, b3D), axis=1)
10 loops, best of 3: 72.1 ms per loop

In [325]: %%timeit
     ...: b3D = np.broadcast_to(b,(a.shape[0],1,len(b)))
     ...: out = np.concatenate((a,b3D),axis=1)
10 loops, best of 3: 72.5 ms per loop

For smaller input shapes, call to np.broadcast_to would take a bit longer than np.repeat given the work needed for setting up the broadcasting is apparently more complicated, as the timings suggest below :

In [360]: a = np.random.rand(5,4,3)

In [361]: b = np.random.rand(3)

In [366]: %timeit np.broadcast_to(b,(a.shape[0],1,len(b)))
100000 loops, best of 3: 3.12 µs per loop

In [367]: %timeit b.reshape(1, 1, -1).repeat(a.shape[0], axis=0)
1000000 loops, best of 3: 957 ns per loop

But, the broadcasting part would have a constant time irrepective of the shapes of the inputs, i.e. the 3 u-sec part would stay around that mark. The timing for the counterpart : b.reshape(1, 1, -1).repeat(a.shape[0], axis=0) would depend on the input shapes. So, let's dig deeper and see how the concatenation steps for the two approaches fair/behave.

Diging deeper

Trying to dig deeper to see how much the concatenation part is consuming :

In [353]: a = np.random.rand(500,400,300)

In [354]: b = np.random.rand(300)

In [355]: b3D = np.broadcast_to(b,(a.shape[0],1,len(b)))

In [356]: %timeit np.concatenate((a,b3D),axis=1)
10 loops, best of 3: 72 ms per loop

In [357]: b3D = b.reshape(1, 1, -1).repeat(a.shape[0], axis=0)

In [358]: %timeit np.concatenate((a,b3D),axis=1)
10 loops, best of 3: 72 ms per loop

Conclusion : Doesn't seem too different.

Now, let's try a case where the replication needed for b is a bigger number and b has noticeably high number of elements as well.

In [344]: a = np.random.rand(10000, 10, 1000)

In [345]: b = np.random.rand(1000)

In [346]: b3D = np.broadcast_to(b,(a.shape[0],1,len(b)))

In [347]: %timeit np.concatenate((a,b3D),axis=1)
10 loops, best of 3: 130 ms per loop

In [348]: b3D = b.reshape(1, 1, -1).repeat(a.shape[0], axis=0)

In [349]: %timeit np.concatenate((a,b3D),axis=1)
10 loops, best of 3: 141 ms per loop

Conclusion : Seems like the merged concatenate+replication with np.broadcast_to is doing a bit better here.

Let's try the original case of (5,4,3) shape :

In [360]: a = np.random.rand(5,4,3)

In [361]: b = np.random.rand(3)

In [362]: b3D = np.broadcast_to(b,(a.shape[0],1,len(b)))

In [363]: %timeit np.concatenate((a,b3D),axis=1)
1000000 loops, best of 3: 948 ns per loop

In [364]: b3D = b.reshape(1, 1, -1).repeat(a.shape[0], axis=0)

In [365]: %timeit np.concatenate((a,b3D),axis=1)
1000000 loops, best of 3: 950 ns per loop

Conclusion : Again, not too different.

So, the final conclusion is that if there are a lot of elements in b and if the first axis of a is also a big number (as the replication number is that one), np.broadcast_to would be a good option, otherwise np.repeat based version takes care of the other cases pretty well.