Numpy: get only elements on odd or even diagonal offsets in a matrix, change the rest to zeros

user_11235205 0
Viewed 416

I have a 2D (square) matrix, for example it can be like this:

1 2 3
4 5 6
7 8 9

I want to get only the elements on the odd or even diagonal offsets of it, and let the rest be zeros. For example, with even diagonal offsets (%2 = 0), the resulting matrix is:

1 0 3
0 5 0
7 0 9

Explanation: main diagonal has offset 0, which is 1 5 9. The next diagonal offsets are 2, 6 and 4, 8 thus they are changed to zeros. Repeat the process until we reach the last diagonal.

And with odd diagonal index, the resulting matrix is:

0 2 0
4 0 6
0 8 0

I look at the np.diag(np.diag(x)) but it only returns the main diagonal and the rest are zeros. How can I extend it to odd/even offsets?

I can also use PyTorch.

pytorch numpy python
4 Answers

I would do it following way using numpy

import numpy as np
arr = np.array([[1,2,3],[4,5,6],[7,8,9]])
masktile = np.array([[True,False],[False,True]])
mask = np.tile(masktile, (2,2))  # this must be at least as big as arr
arr0 = np.where(mask[:arr.shape[0],:arr.shape[1]], arr, 0)
arr1 = np.where(mask[:arr.shape[0],:arr.shape[1]], 0, arr)
print(arr0)
print(arr1)

output:

[[1 0 3]
 [0 5 0]
 [7 0 9]]
[[0 2 0]
 [4 0 6]
 [0 8 0]]

Explanation: I am creating mask which is array of Trues and Falses to use to decide if given element is to remain or should be replaced by 0. I create single tile which then I feed into np.tile to get "chessboard" of sufficient size, then I use part of apprioate size of it together with np.where to replace selected elements with 0.

user_107859750

You can use the offset argument of diagonal in PyTorch like so:

x = torch.arange(1,10).view(3,3)

for i in range(offset, x.shape[0], 2):
    x.diagonal(i).fill_(0)
    x.diagonal(-i).fill_(0)

# offset = 0 (Even)
tensor([[0, 2, 0],
        [4, 0, 6],
        [0, 8, 0]])

# offset = 1 (Odd)
tensor([[1, 0, 3],
        [0, 5, 0],
        [7, 0, 9]])
user_90676150

You can also use np.indices. If the row and column indexes % 2 are equal, you are on an "even" diagonal. If they are not equal you are on an "odd" diagonal.

arr = np.array([[1,2,3],[4,5,6],[7,8,9]])
rows, cols = np.indices(arr.shape)

Set the even indexes to 0

>>> arr[(rows%2)==(cols%2)] = 0
>>> arr
... array([[0, 2, 0],
           [4, 0, 6],
           [0, 8, 0]])

Set the odd indexes to 0

>>> arr[(rows%2)!=(cols%2)] = 0
>>> arr
... array([[1, 0, 3],
           [0, 5, 0],
           [7, 0, 9]])
user_57631650

np.diag has an offset parameter, for example:

import numpy as np
a = np.ones((3,3))
print(a - np.diag(np.diag(a, 1), 1))

[[1. 0. 1.]
 [1. 1. 0.]
 [1. 1. 1.]]
user_52119220
Related