Make special diagonal matrix in Numpy

Question

I am trying to make a numpy array that looks like this:

[a b c       ]
[  a b c     ]
[    a b c   ]
[      a b c ] 

So this involves updating the main diagonal and the two diagonals above it.

What would be an efficient way of doing this?

Answer 1

You can use np.indices to get the indices of your array and then assign the values where you want.

a = np.zeros((5,10))
i,j = np.indices(a.shape)

i,j are the line and column indices, respectively.

a[i==j] = 1.
a[i==j-1] = 2.
a[i==j-2] = 3.

will result in:

array([[ 1.,  2.,  3.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  1.,  2.,  3.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  1.,  2.,  3.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  1.,  2.,  3.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  1.,  2.,  3.,  0.,  0.,  0.]])

Answer 2

This is an example of a Toeplitz matrix - you can construct it using scipy.linalg.toeplitz:

import numpy as np
from scipy.linalg import toeplitz

first_row = np.array([1, 2, 3, 0, 0, 0])
first_col = np.array([1, 0, 0, 0])

print(toeplitz(first_col, first_row))
# [[1 2 3 0 0 0]
#  [0 1 2 3 0 0]
#  [0 0 1 2 3 0]
#  [0 0 0 1 2 3]]

Answer 3

import numpy as np

def using_tile_and_stride():
    arr = np.tile(np.array([10,20,30,0,0,0], dtype='float'), (4,1))
    row_stride, col_stride = arr.strides
    arr.strides = row_stride-col_stride, col_stride
    return arr

In [108]: using_tile_and_stride()
Out[108]: 
array([[ 10.,  20.,  30.,   0.,   0.,   0.],
       [  0.,  10.,  20.,  30.,   0.,   0.],
       [  0.,   0.,  10.,  20.,  30.,   0.],
       [  0.,   0.,   0.,  10.,  20.,  30.]])

Other, slower alternatives include:

import numpy as np

import numpy.lib.stride_tricks as stride

def using_put():
    arr = np.zeros((4,6), dtype='float')
    a, b, c = 10, 20, 30
    nrows, ncols = arr.shape
    ind = (np.arange(3) + np.arange(0,(ncols+1)*nrows,ncols+1)[:,np.newaxis]).ravel()
    arr.put(ind, [a, b, c])
    return arr

def using_strides():
    return np.flipud(stride.as_strided(
        np.array([0, 0, 0, 10, 20, 30, 0, 0, 0], dtype='float'), 
        shape=(4, 6), strides = (8, 8)))

---

If you use using_tile_and_stride, note that the array is only appropriate for read-only purposes. Otherwise, if you were to try to modify the array, you might be surprised when multiple array locations change simultaneously:

In [32]: arr = using_tile_and_stride()

In [33]: arr[0, -1] = 100

In [34]: arr
Out[34]: 
array([[  10.,   20.,   30.,    0.,  100.],
       [ 100.,   10.,   20.,   30.,    0.],
       [   0.,    0.,   10.,   20.,   30.],
       [  30.,    0.,    0.,   10.,   20.]])

You could work around this by returning np.ascontiguousarray(arr) instead of just arr, but then using_tile_and_stride would be slower than using_put. So if you intend to modify the array, using_put would be a better choice.

Answer 4

I can't comment yet, but I want to bump that ali_m's answer is by far the most efficient as scipy takes care of things for you.

For example, with a matrix of size n,m = 1200, repeatedly adding np.diag() calls takes ~6.14s, Saullo G. P. Castro's answer takes ~7.7s, and scipy.linalg.toeplitz(np.arange(N), np.arange(N)) takes 1.57ms.

Answer 5

Using my answer to this question: changing the values of the diagonal of a matrix in numpy , you can do some tricky slicing to get a view of each diagonal, then do the assignment. In this case it would just be:

import numpy as np
A = np.zeros((4,6))
# main diagonal
A.flat[:A.shape[1]**2:A.shape[1]+1] = a
# first superdiagonal
A.flat[1:max(0,A.shape[1]-1)*A.shape[1]:A.shape[1]+1] = b
# second superdiagonal
A.flat[2:max(0,A.shape[1]-2)*A.shape[1]:A.shape[1]+1] = c