interweave 3 numpy matrices?

Question

How do i interweave numpy matrices columwise.

given this example:

>>> import numpy as np
>>> a = np.zeros((3,3))
>>> b = np.ones((3,3))
>>> c = b*2

should the interweaved output be

[[ a[0,0].  b[0,0].  c[0,0].  a[0,1]  b[0,1]  c[0,1] .. ]
 [ a[1,0].  b[1,0].  c[1,0].  a[1,1]  b[1,1]  c[1,1] .. ]
 [ a[2,0].  b[2,0].  c[2,0].  a[2,1]  b[2,1]  c[2,1] .. ]]

The end shape should be (3,9)

Answer 1

Another option, you can use np.dstack + reshape, according to the docs, dstack stacks arrays in sequence depth wise (along third axis) so it can be convenient for this case:

np.dstack([a, b, c]).reshape(3,-1)
#array([[ 0.,  1.,  2.,  0.,  1.,  2.,  0.,  1.,  2.],
#       [ 0.,  1.,  2.,  0.,  1.,  2.,  0.,  1.,  2.],
#       [ 0.,  1.,  2.,  0.,  1.,  2.,  0.,  1.,  2.]])

With some sample data less ambiguous:

import numpy as np
a = np.arange(9).reshape(3,3)
b = (np.arange(9) + 9).reshape(3,3)
c = b*2

#array([[ 0,  9, 18,  1, 10, 20,  2, 11, 22],
#       [ 3, 12, 24,  4, 13, 26,  5, 14, 28],
#       [ 6, 15, 30,  7, 16, 32,  8, 17, 34]])

Answer 2

If all matrices have the same dimension, you can use:

np.array([a,b,c]).transpose(1,2,0).reshape(3,-1)

Or a generic function to merge n matrices:

def merge(*args):
    (m,_) = args[0].shape
    return np.array(args).transpose(1,2,0).reshape(m,-1)

(and call it with merge(a,b,c))

You can even let it work for arbitrary dimensions, with:

def merge_arbitrary(*args):
    (*m,_) = args[0].shape
    return np.array(args).transpose(tuple(range(1,len(m)+2))+(0,)). \
                          reshape(m+[-1])

The code works as follows. We first construct a 3×3×3 matrix that has the shape:

array([[[ a00,  a01,  a02],
        [ a10,  a11,  a12],
        [ a20,  a21,  a22]],

       [[ b00,  b01,  b02],
        [ b10,  b11,  b12],
        [ b20,  b21,  b22]],

       [[ c00,  c01,  c02],
        [ c10,  c11,  c12],
        [ c20,  c21,  c22]]])

Next we make a transpose(1,2,0) such that now [aij,bij,cij] are the lowest dimension. So from now on the matrix has shape:

array([[[a00, b00, c00],
        [a01, b01, c01],
        [a02, b02, c02]],

       [[a10, b10, c10],
        [a11, b11, c11],
        [a12, b12, c12]],

       [[a20, b20, c20],
        [a21, b21, c21],
        [a22, b22, c22]]])

Finally by calling reshape(3,-1) we "drop" the lowest dimension and concatenate:

array([[a00, b00, c00, a01, b01, c01, a02, b02, c02],
       [a10, b10, c10, a11, b11, c11, a12, b12, c12],
       [a20, b20, c20, a21, b21, c21, a22, b22, c22]])