Numpy scalable diagonal matrices-CodePudding

Assuming I have the variables:

A = 3
B = 2
C = 1

How can i transform them into diagonal matrices in the following form:

np.diag([1, 1, 1, 0, 0, 0])
Out[0]: 
array([[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, 0, 0],
       [0, 0, 0, 0, 0, 0]])

np.diag([0,0,0,1,1,0])
Out[1]: 
array([[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, 0]])

np.diag([0,0,0,0,0,1])
Out[2]: 
array([[0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 1]])

I would like this to be scalable, so for instance with 4 variables a = 500, b = 20, c = 300, d = 200 the size of the matrix will be 500 20 300 200 = 1020. What is the easiest way to do this?

CodePudding user response：

Here's one approach. The resulting array mats contains the matrices you're looking for.

A = 3
B = 2
C = 1

n_list = [A,B,C]
ab_list = np.cumsum([0]   n_list)
ran = np.arange(ab_list[-1])
mats = [np.diag(((a <= ran) & (ran < b)).astype('int')) 
        for a,b in zip(ab_list[:-1],ab_list[1:])]
for mat in mats:
    print(mat,'\n')

Result:

[[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 0 0]
 [0 0 0 0 0 0]] 

[[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 0]] 

[[0 0 0 0 0 0]
 [0 0 0 0 0 0]
 [0 0 0 0 0 0]
 [0 0 0 0 0 0]
 [0 0 0 0 0 0]
 [0 0 0 0 0 1]]

CodePudding user response：

One posible method ( don't think it's optimal but it works):



import numpy as np

a = 3
b = 2
c = 1
values = [a,b,c] #create a list with values

n = sum(values) #calc total length of diagnal

#create an array with cumulative sums but starting from 0 to use as index
idx_vals = np.zeros(len(values) 1,dtype=int)
np.cumsum(values,out=idx_vals[1:]); 

#create every diagonal using values, then create diagonal matrices and 
#save them in `matrices` list
matrices = []
for idx,v in enumerate(values):

    diag = np.zeros(n)
    diag[idx_vals[idx]:idx_vals[idx] v] = np.ones(v)
    print(diag)
    matrices.append(np.diag(diag))

CodePudding user response：

Yet another possibility:

import numpy as np

# your constants here
constants = [3, 2, 1] # [A, B, C]

size = sum(constants)
cumsum = np.cumsum([0]   constants)

for i in range(len(cumsum) - 1):
    inputVector = np.zeros(size, dtype=int)
    inputVector[cumsum[i]:cumsum[i 1]] = 1
    matrix = np.diag(inputVector)
    
    print(matrix, '\n')

Output:

[[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 0 0]
 [0 0 0 0 0 0]] 

[[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 0]] 

[[0 0 0 0 0 0]
 [0 0 0 0 0 0]
 [0 0 0 0 0 0]
 [0 0 0 0 0 0]
 [0 0 0 0 0 0]
 [0 0 0 0 0 1]]

CodePudding user response：

The obligatory solution with np.einsum, about ~2.25x slower than the accepted answer for the [500,20,200,300] arrays on a 2-core colab instance.

import numpy as np

A = 3
B = 2
C = 1

r = [A,B,C]
m = np.arange(len(r))
np.einsum('ij,kj->ijk', m.repeat(r) == m[:,None], np.eye(np.sum(r), dtype='int'))

Output

array([[[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, 0, 0],
        [0, 0, 0, 0, 0, 0]],

       [[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, 0]],

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

CodePudding user response：

it is not clear how you intend to use A,B,C, could you please clarify ?