Create array/tensor of cycle shifted arrays-CodePudding

I want to create 2d tensor (or numpy array, doesn't really matter), where every row will be cycle shifted first row. I do it using for loop:

import torch
import numpy as np

a = np.random.rand(33, 11)
miss_size = 64
lp_order = a.shape[1] - 1
inv_a = -np.flip(a, axis=1)
mtx_size = miss_size lp_order   # some constant
mtx_row = torch.cat((torch.from_numpy(inv_a), torch.zeros((a.shape[0], miss_size - 1   a.shape[1]))), dim=1)
mtx_full = mtx_row.unsqueeze(1)
for i in range(mtx_size):
        mtx_row = torch.roll(mtx_row, 1, 1)
        mtx_full = torch.cat((mtx_full, mtx_row.unsqueeze(1)), dim=1)

unsqueezing is needed because I stack 2d tensors into 3d tensor

Is there more efficient way to do that? Maybe linear algebra trick or more pythonic approach.

CodePudding user response：

You can use scipy.linalg.circulant():

scipy.linalg.circulant([1, 2, 3])
# array([[1, 3, 2],
#        [2, 1, 3],
#        [3, 2, 1]])

CodePudding user response：

I believe you can achieve this using torch.gather by constructing the appropriate index tensor. This approach works with batches too.

If we take this approach, the objective is to construct an index tensor where each value refers to an index in mtx_row (along the last dimension here dim=1). In this case, it would be shaped (3, 3):

tensor([[0, 1, 2],
        [2, 0, 1],
        [1, 2, 0]])

You can achieve this by broadcasting torch.arange with its own transpose and applying modulo on the resulting matrix:

>>> idx = (n-torch.arange(n)[None].T   torch.arange(n)[None]) % n
tensor([[0, 1, 2],
        [2, 0, 1],
        [1, 2, 0]])

Let mtx_row be shaped (2, 3):

>>> mtx_row
tensor([[0.3523, 0.0170, 0.1875],
        [0.2156, 0.7773, 0.4563]])

From there you need to idx and mtx_row so they have the same shapes:

>>> idx_ = idx[None].expand(len(mtx_row), -1, -1)
>>> val_ = mtx_row[:, None].expand(-1, n, -1)

Then we can apply torch.gather on the last dimension dim=2:

>>> val_.gather(-1, idx_)
tensor([[[0.3523, 0.0170, 0.1875],
         [0.1875, 0.3523, 0.0170],
         [0.0170, 0.1875, 0.3523]],

        [[0.2156, 0.7773, 0.4563],
         [0.4563, 0.2156, 0.7773],
         [0.7773, 0.4563, 0.2156]]])