CODEWITHSUNDEEP
pythonarraysnumpy
ollydbg23
ollydbg23

Reputation: 1210

python numpy: how to construct a big diagonal array(matrix) from two small array

import numpy as np

A = np.array([[1, 2], 
              [3, 4]])    
B = np.array([[5, 6], 
              [7, 8]])

C = np.array([[1, 2, 0, 0],
              [3, 4, 0, 0],
              [0, 0, 5, 6],
              [0, 0, 7, 8]])

I would like to make C directly from A and B, are there any simply ways to construct a diagonal array C? Thanks.

Upvotes: 12

Views: 4151

Answers (2)

user7328723
user7328723

Reputation:

Here's a recursive solution that does does not require that the output array is square. The input is a list of 2-D arrays.

import numpy as np


def diag_mat(rem=[], result=np.empty((0, 0))):
    if not rem:
        return result
    m = rem.pop(0)
    result = np.block(
        [
            [result, np.zeros((result.shape[0], m.shape[1]))],
            [np.zeros((m.shape[0], result.shape[1])), m],
        ]
    )
    return diag_mat(rem, result)

Testing the output:

>>> a = np.array([[2, 1, 5], [7, 3, 1]])
>>> b = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
>>> diag_mat([a, b])
array([[2., 1., 5., 0., 0., 0.],
       [7., 3., 1., 0., 0., 0.],
       [0., 0., 0., 1., 2., 3.],
       [0., 0., 0., 4., 5., 6.],
       [0., 0., 0., 7., 8., 9.]])

Upvotes: 0

Divakar
Divakar

Reputation: 221614

Approach #1 : One easy way would be with np.bmat -

Z = np.zeros((2,2),dtype=int) # Create off-diagonal zeros array
out = np.asarray(np.bmat([[A, Z], [Z, B]]))

Sample run -

In [24]: Z = np.zeros((2,2),dtype=int)

In [25]: np.asarray(np.bmat([[A, Z], [Z, B]]))
Out[25]: 
array([[1, 2, 0, 0],
       [3, 4, 0, 0],
       [0, 0, 5, 6],
       [0, 0, 7, 8]])

Approach #2 : For generic number of arrays, we can use masking -

def diag_block_mat_boolindex(L):
    shp = L[0].shape
    mask = np.kron(np.eye(len(L)), np.ones(shp))==1
    out = np.zeros(np.asarray(shp)*len(L),dtype=int)
    out[mask] = np.concatenate(L).ravel()
    return out

Approach #3 : For generic number of arrays, another way with multi-dimensional slicing -

def diag_block_mat_slicing(L):
    shp = L[0].shape
    N = len(L)
    r = range(N)
    out = np.zeros((N,shp[0],N,shp[1]),dtype=int)
    out[r,:,r,:] = L
    return out.reshape(np.asarray(shp)*N)

Sample runs -

In [137]: A = np.array([[1, 2], 
     ...:               [3, 4]])    
     ...: B = np.array([[5, 6], 
     ...:               [7, 8]])
     ...: C = np.array([[11, 12], 
     ...:               [13, 14]])  
     ...: D = np.array([[15, 16], 
     ...:               [17, 18]])
     ...: 

In [138]: diag_block_mat_boolindex((A,B,C,D))
Out[138]: 
array([[ 1,  2,  0,  0,  0,  0,  0,  0],
       [ 3,  4,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  5,  6,  0,  0,  0,  0],
       [ 0,  0,  7,  8,  0,  0,  0,  0],
       [ 0,  0,  0,  0, 11, 12,  0,  0],
       [ 0,  0,  0,  0, 13, 14,  0,  0],
       [ 0,  0,  0,  0,  0,  0, 15, 16],
       [ 0,  0,  0,  0,  0,  0, 17, 18]])

In [139]: diag_block_mat_slicing((A,B,C,D))
Out[139]: 
array([[ 1,  2,  0,  0,  0,  0,  0,  0],
       [ 3,  4,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  5,  6,  0,  0,  0,  0],
       [ 0,  0,  7,  8,  0,  0,  0,  0],
       [ 0,  0,  0,  0, 11, 12,  0,  0],
       [ 0,  0,  0,  0, 13, 14,  0,  0],
       [ 0,  0,  0,  0,  0,  0, 15, 16],
       [ 0,  0,  0,  0,  0,  0, 17, 18]])

Upvotes: 14

Related Questions