Numpy scalable diagonal matrices

Question

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?

Answer 1

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

---

Edit: Here's a faster solution that yields the same result

n_list = [A,B,C]
ab_list = np.cumsum([0] + n_list)
total = ab_list[-1]
ran = np.arange(total)
mats = np.zeros((len(n_list),total,total))
for k,p in enumerate(zip(ab_list[:-1],ab_list[1:])):
    idx = np.arange(p[0],p[1])
    mats[k,idx,idx] = 1
for mat in mats:
    print(mat,'\n')

This seems to yield a ~10% speedup over the currently accepted solution

---

Another with roughly equivalent performance:

n_list = [A,B,C]
m = len(n_list)
ab_list = np.cumsum([0] + n_list)
total = ab_list[-1]
ran = np.arange(total)
mats = np.zeros((m,total,total))
idx = [k for a,b in zip(ab_list[:-1],ab_list[1:]) for k in range(a,b)]
mats[[k for k,n in enumerate(n_list) for _ in range(n)],
     idx,idx] = 1
for mat in mats:
    print(mat,'\n')

Answer 2

You can achieve even better performance by just allocating the array once, then setting the values all at once by specifying the indices. The indices are fortunately easy to obtain.

import numpy as np

a = [3, 2, 1] # Put your values in a list
s = np.sum(a)
m = np.zeros((len(a), s, s), dtype=int) # Initialize array once
indices = (np.repeat(range(len(a)), a), *np.diag_indices(s, 2)) # Get indices
m[indices] = 1 # Set the diagonals at once
return m

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

Comparing to @Ben Grossmann's answer, with A=3000, B=2000, C=1000 and 100 repeats:

def A():
    '''My solution'''
    a = [3000, 2000, 1000] # Put your values in a list
    s = np.sum(a)
    m = np.zeros((len(a), s, s), dtype=int) # Initialize array once
    indices = (np.repeat(range(len(a)), a), *np.diag_indices(s, 2)) # Get indices
    m[indices] = 1 # Set the diagonals at once
    return m

def B():
    '''Bens solution'''
    A = 3000
    B = 2000
    C = 1000

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

print(f'Timings:')
timeA = timeit.timeit(A, number=100)
timeB = timeit.timeit(B, number=100)
ratio = timeA / timeB
print(f'This solution: {timeA} seconds')
print(f'Current accepted answer: {timeB} seconds')
if ratio < 1:
    print(f'This solution is {1 / ratio} times faster than Bens solution')
else:
    print(f'Bens solution is {ratio} times faster than this solution')

Output:

Timings:
This solution: 1.6834218999993027 seconds
Current accepted answer: 5.096610300000066 seconds
This solution is 3.027529997086397 times faster than Bens solution

EDIT: Changed the "indices" algorithm to use np.repeat instead of np.concatenate.

Answer 3

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]]])

Answer 4

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

Answer 5

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))