numpy column_stack in loop

Question

I don't know if the title is apprpriate or not, but let me show you what I want to do,

In [56]: import numpy as np

In [57]: a= np.random.rand(2,2,2); a
Out[57]: 
array([[[0.4300565 , 0.82251319],
        [0.56113378, 0.83284255]],

       [[0.00822414, 0.28256243],
        [0.16648411, 0.33381438]]])

In [58]: b=np.random.rand(2); b
Out[58]: array([0.8035224 , 0.09884653])

In [59]: np.stack(( np.column_stack((b,a[:,i,:])) for i in range(a.shape[1])))
Out[59]: 
array([[[0.8035224 , 0.4300565 , 0.82251319],
        [0.09884653, 0.00822414, 0.28256243]],

       [[0.8035224 , 0.56113378, 0.83284255],
        [0.09884653, 0.16648411, 0.33381438]]])

So, I want to stack an array as column to an inner axis. Is it possible to do the looping structure more efficiently and concisely in numpy? I tried with numpy insert but could not do it.

EDIT:

another example

In [110]: a= np.random.rand(5,3,3); a
Out[110]: 
array([[[0.27506756, 0.82334411, 0.7004287 ],
        [0.6834928 , 0.28457133, 0.6275462 ],
        [0.49744358, 0.25131814, 0.56422852]],

       [[0.82591597, 0.92367306, 0.04652992],
        [0.98545051, 0.92813944, 0.14360307],
        [0.85454081, 0.8254149 , 0.5637401 ]],

       [[0.59545519, 0.41563571, 0.41937218],
        [0.90980491, 0.30169504, 0.96630809],
        [0.06713389, 0.64357544, 0.12901734]],

       [[0.47566444, 0.33476802, 0.26635363],
        [0.4678913 , 0.53028241, 0.03112231],
        [0.68445959, 0.07113376, 0.86651669]],

       [[0.66951982, 0.01827502, 0.43831829],
        [0.02798567, 0.36880876, 0.55029074],
        [0.40127051, 0.6311474 , 0.51015882]]])

In [111]:  b= np.random.rand(5,2); b
Out[111]: 
array([[0.01659589, 0.15320541],
       [0.79025065, 0.28041334],
       [0.56024173, 0.49317082],
       [0.28229119, 0.46010724],
       [0.72239851, 0.62075004]])

In [112]: np.stack(( np.column_stack((b,a[:,i,:])) for i in range(a.shape[1])))
Out[112]: 
array([[[0.01659589, 0.15320541, 0.27506756, 0.82334411, 0.7004287 ],
        [0.79025065, 0.28041334, 0.82591597, 0.92367306, 0.04652992],
        [0.56024173, 0.49317082, 0.59545519, 0.41563571, 0.41937218],
        [0.28229119, 0.46010724, 0.47566444, 0.33476802, 0.26635363],
        [0.72239851, 0.62075004, 0.66951982, 0.01827502, 0.43831829]],

       [[0.01659589, 0.15320541, 0.6834928 , 0.28457133, 0.6275462 ],
        [0.79025065, 0.28041334, 0.98545051, 0.92813944, 0.14360307],
        [0.56024173, 0.49317082, 0.90980491, 0.30169504, 0.96630809],
        [0.28229119, 0.46010724, 0.4678913 , 0.53028241, 0.03112231],
        [0.72239851, 0.62075004, 0.02798567, 0.36880876, 0.55029074]],

       [[0.01659589, 0.15320541, 0.49744358, 0.25131814, 0.56422852],
        [0.79025065, 0.28041334, 0.85454081, 0.8254149 , 0.5637401 ],
        [0.56024173, 0.49317082, 0.06713389, 0.64357544, 0.12901734],
        [0.28229119, 0.46010724, 0.68445959, 0.07113376, 0.86651669],
        [0.72239851, 0.62075004, 0.40127051, 0.6311474 , 0.51015882]]])

In your output the original arrays is changed, is that what is expected or is a bug in your code? — Dani Mesejo
– Dani Mesejo, Commented Jan 5, 2019 at 13:43
I'm confused, I don't see any change in the original array, can you show me with an example — Eular
– Eular, Commented Jan 5, 2019 at 13:47
The order of the elements is different, for example 0.56113378 is below 0.4300565 in array a but in the output is not. — Dani Mesejo
– Dani Mesejo, Commented Jan 5, 2019 at 13:50

hpaulj · Accepted Answer · 2019-01-05 17:53:00Z

A variation on concatenating is indexed assignment:

For the first example:

In [245]: a=np.arange(8).reshape(2,2,2); b=np.array([100,200])
In [246]: c = np.zeros((2,2,3), a.dtype)
In [247]: c[:,:,0]=b
In [248]: c[:,:,1:]=a.transpose(1,0,2)
In [249]: c
Out[249]: 
array([[[100,   0,   1],
        [200,   4,   5]],

       [[100,   2,   3],
        [200,   6,   7]]])

And for the second:

In [250]: a1 = np.arange(5*3*3).reshape(5,3,3)
In [251]: b1 = np.arange(10).reshape(5,2)
In [252]: c1 = np.zeros((3,5,5),a.dtype)
In [253]: c1[:,:,:2]=b1
In [254]: c1[:,:,2:]=a1.transpose(1,0,2)
In [255]: c1
Out[255]: 
array([[[ 0,  1,  0,  1,  2],
        [ 2,  3,  9, 10, 11],
        [ 4,  5, 18, 19, 20],
        [ 6,  7, 27, 28, 29],
        [ 8,  9, 36, 37, 38]],

       [[ 0,  1,  3,  4,  5],
        [ 2,  3, 12, 13, 14],
        [ 4,  5, 21, 22, 23],
        [ 6,  7, 30, 31, 32],
        [ 8,  9, 39, 40, 41]],

       [[ 0,  1,  6,  7,  8],
        [ 2,  3, 15, 16, 17],
        [ 4,  5, 24, 25, 26],
        [ 6,  7, 33, 34, 35],
        [ 8,  9, 42, 43, 44]]])

Deriving the shape of c from a and b is left as an exercise for the reader. :)

np.stack (or np.array) over the iteration on 2nd axis is effectively a partial transpose (or interchange of the first 2 axes):

In [261]: np.stack([a[:,i,:] for i in range(a.shape[1])])
Out[261]: 
array([[[0, 1],
        [4, 5]],

       [[2, 3],
        [6, 7]]])
In [262]: a.transpose(1,0,2)
Out[262]: 
array([[[0, 1],
        [4, 5]],

       [[2, 3],
        [6, 7]]])

We could also iterate on the first axis, and join on the second with:

In [263]: np.stack(a, axis=1)
Out[263]: 
array([[[0, 1],
        [4, 5]],

       [[2, 3],
        [6, 7]]])

A refinement on Ankit's answer using concatenate is:

np.concatenate([np.repeat(b[None,:,None], 2, axis=0), a.transpose(1,0,2)], axis=2)
np.concatenate([np.repeat(b1[None,:,:], 3, axis=0), a1.transpose(1,0,2)], axis=2)

Ankit Goyal · Accepted Answer · 2019-01-06 06:03:04Z

0

The below code worked for me.

>>> a= np.random.rand(2,2,2); a
array([[[0.52706506, 0.48344319],
        [0.79027196, 0.90581149]],

       [[0.25930158, 0.59498346],
        [0.02164495, 0.63081622]]])
>>> b=np.random.rand(2); b
array([0.96890722, 0.93670425])
>>> a1 = a.transpose(1, 0, 2); a1
array([[[0.52706506, 0.48344319],
        [0.25930158, 0.59498346]],

       [[0.79027196, 0.90581149],
        [0.02164495, 0.63081622]]])
>>> c = np.tile(b, (2, 1)); c
array([[0.43134454, 0.4042494 ],
       [0.43134454, 0.4042494 ]])
>>> c = np.expand_dims(c,2); c
array([[[0.43134454],
        [0.4042494 ]],

       [[0.43134454],
        [0.4042494 ]]])
>>> np.concatenate((c, a1), axis=2)
array([[[0.43134454, 0.52706506, 0.48344319],
        [0.4042494 , 0.25930158, 0.59498346]],

       [[0.43134454, 0.79027196, 0.90581149],
        [0.4042494 , 0.02164495, 0.63081622]]])

Here I first repeated b using tile by same number as 2nd dimention of a in a new dimention. Then I used concatication to concat b and a array.

For the 2nd example

>>> a= np.random.rand(5,3,3)
>>> a1 = a.transpose(1, 0, 2)
>>> b=np.random.rand(5, 2)
>>> c = np.tile(b, (3, 1, 1))
>>> np.concatenate((c, a1), axis=2)

edited Jan 6, 2019 at 6:03

answered Jan 5, 2019 at 14:53

Ankit Goyal

1196 bronze badges

2 Comments

hpaulj Over a year ago

The OP has an axis swap on a that you didn't replicate. With random numbers like this it's easy to miss details like that.

Ankit Goyal Over a year ago

You are right. I had missed that completly. I have updated my annswer.

Collectives™ on Stack Overflow

numpy column_stack in loop

2 Answers 2

Comments

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related