how to construct diagonal array using a 2d array in numpy?

Question

Using np.diag I'm able to construct a 2-D array where a 1-D array is input is returned on the diagonal. But how to do the same if I have n-D array as input?

This works

foo = np.random.randint(2, size=(36))
print foo
print np.diag(foo)
[1 1 1 1 1 1 0 1 0 0 1 1 0 1 1 1 1 1 0 0 0 0 0 0 0 1 0 1 0 1 0 0 0 1 1 0]
[[1 0 0 ..., 0 0 0]
 [0 1 0 ..., 0 0 0]
 [0 0 1 ..., 0 0 0]
 ..., 
 [0 0 0 ..., 1 0 0]
 [0 0 0 ..., 0 1 0]
 [0 0 0 ..., 0 0 0]]

This won't

foo = np.random.randint(2, size=(2,36))
print foo
[[1 1 0 1 1 0 1 1 1 1 1 0 0 0 0 0 1 0 0 1 0 0 0 0 1 0 0 0 1 0 0 0 0 1 0 0]
 [0 1 1 1 0 0 1 1 1 1 0 1 0 1 0 1 0 0 0 0 1 0 1 1 1 1 0 0 1 0 1 1 1 1 0 1]]
do_something(foo)

Should return

array([[[ 1.,  0.,  0., ...,  0.,  0.,  0.],
        [ 0.,  1.,  0., ...,  0.,  0.,  0.],
        [ 0.,  0.,  0., ...,  0.,  0.,  0.],
        ..., 
        [ 0.,  0.,  0., ...,  1.,  0.,  0.],
        [ 0.,  0.,  0., ...,  0.,  0.,  0.],
        [ 0.,  0.,  0., ...,  0.,  0.,  0.]],

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

EDIT

Some test based on answers of Alan and ajcr in this post and Saulo Castro and jaime where ajcr refers to. As usual, it all depends on your input. My input has usually the following shape:

M = np.random.randint(2, size=(1000, 36))

with the functions as follows:

def Alan(M):
    M = np.asarray(M)
    depth, size = M.shape
    x = np.zeros((depth,size,size))
    for i in range(depth):
        x[i].flat[slice(0,None,1+size)] = M[i]
    return x

def ajcr(M):
    return np.eye(M.shape[1]) * M[:,np.newaxis,:]

def Saulo(M):
    b = np.zeros((M.shape[0], M.shape[1], M.shape[1]))
    diag = np.arange(M.shape[1])
    b[:, diag, diag] = M
    return b

def jaime(M):
    b = np.zeros((M.shape[0], M.shape[1]*M.shape[1]))
    b[:, ::M.shape[1]+1] = M
    return b.reshape(M.shape[0], M.shape[1], M.shape[1])

result for me in the following

%timeit Alan(M)
100 loops, best of 3: 2.22 ms per loop    
%timeit ajcr(M)
100 loops, best of 3: 5.1 ms per loop    
%timeit Saulo(M)
100 loops, best of 3: 4.33 ms per loop    
%timeit jaime(M)
100 loops, best of 3: 2.07 ms per loop

Answer 1

A simple way to do this is in pure NumPy is to perform the following array multiplication:

np.eye(foo.shape[1]) * foo[:, np.newaxis]

where foo is the 2D array of diagonals.

This multiplies the NxN identity array with each row of foo to produce the required 3D matrix.

Because the syntax in this method is fairly simple, you can easily extend it to higher dimensions. For example:

>>> foo = np.array([[0, 1], [1, 1]])
>>> d = np.eye(foo.shape[1]) * foo[:, np.newaxis] # 2D to 3D
>>> d
array([[[ 0.,  0.],
        [ 0.,  1.]],

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

>>> np.eye(d.shape[1]) * d[:, :, np.newaxis] # 3D to 4D
array([[[[ 0.,  0.],
         [ 0.,  0.]],

        [[ 0.,  0.],
         [ 0.,  1.]]],

       [[[ 1.,  0.],
         [ 0.,  0.]],

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

This question may be relevant; it also shows a quicker (but slightly more verbose) way to derive the desired diagonal matrix from the 2D array.

Answer 2

def makediag3d(a):
    a = np.asarray(a)
    depth, size = a.shape
    x = np.zeros((depth,size,size))
    for i in range(depth):
        x[i].flat[slice(0,None,1+size)] = a[i]
    return x