python: shift each matrix inside a tensor independantly

Question

This question is quite similar for matrices Roll rows of a matrix independently

But i'm failing to adapt it for 3D tensors

I'm given a tensor

and a vector that specifies by how much I want to shift my matrices column wise

1 2

I want a new tensor where each matrix has been shiften column wise like so

So far I have been able to get a potential mapping indices

import numpy as np

# Input
A = np.array([[0, 0, 0],
              [1, 1, 1],
              [0, 0, 0],
              [0, 0, 0]])
B = np.array([[0, 0, 0],
              [1, 1, 1],
              [0, 0, 0],
              [0, 0, 0]])
AB = np.array([A, B])

# Shifting amount
r = np.array([-1, 1])

d1, d2, d3 = np.ogrid[:AB.shape[0], :AB.shape[1], :AB.shape[2]]

r[r < 0] += AB.shape[1]
r = np.array([r, ]*AB.shape[1]).transpose()
r = r[:, np.newaxis]

# New column indices?
d2 = d2 - r
d2[d2 < 0] += AB.shape[1]

result = AB[d2]
print(result)

But I get this error :

~/Work/ethz/iml/task2 $ python test.py
Traceback (most recent call last):
  File "test.py", line 27, in <module>
    result = AB[d2]
IndexError: index 2 is out of bounds for axis 0 with size 2

This is what d2 looks like :

[[[1 1 1 1]
  [2 2 2 2]
  [3 3 3 3]
  [0 0 0 0]]

 [[3 3 3 3]
  [0 0 0 0]
  [1 1 1 1]
  [2 2 2 2]]]

yes, the approach #2 from @Divakar 's answer. but i'll admit I don't quite understand it :( — truvaking, May 01 '20 at 12:50

Divakar · Accepted Answer · 2020-04-27T20:04:49.583

Approach #1

Adapting strided-based solution from the same linked Q&A for performance -

from skimage.util.shape import view_as_windows

def roll_along_second_axis_3dar(a, r):
    r = np.asarray(r)
    a_ext = np.concatenate((a,a[:,:-1,:]),axis=1)
    n = a.shape[1]
    idx = (n-r)%n
    w = view_as_windows(a_ext,(1,n,1))[...,0,:,0]
    return w[np.arange(len(idx)),idx].swapaxes(1,2)

Sample run -

In [11]: a
Out[11]: 
array([[[44, 47, 64],
        [67, 67,  9],
        [83, 21, 36],
        [87, 70, 88]],

       [[88, 12, 58],
        [65, 39, 87],
        [46, 88, 81],
        [37, 25, 77]]])

In [12]: roll_along_second_axis_3dar(a, r=[-1,1])
Out[12]: 
array([[[67, 67,  9],
        [83, 21, 36],
        [87, 70, 88],
        [44, 47, 64]],

       [[37, 25, 77],
        [88, 12, 58],
        [65, 39, 87],
        [46, 88, 81]]])

Approach #2

Going with your try, seems you were close enough. We could get to the final output with few modifications/corrections -

d1, d2, d3 = np.ogrid[:AB.shape[0], :AB.shape[1], :AB.shape[2]]
r[r < 0] += AB.shape[1]
D2 = ((d2 - r).transpose(2,1,0))%AB.shape[1]
out = AB[d1,D2,d3]

For the second approach, `np.indices(AB.shape, sparse=True)` might be cleaner — Mad Physicist, Apr 27 '20 at 20:37

hilberts_drinking_problem · Answer 2 · 2020-04-27T20:22:09.177

2

Here is an indexing approach with take_along_axis:

idx = (-r[:,None] + np.arange(AB.shape[1])[None]) % AB.shape[1]
res = np.take_along_axis(AB, idx[...,None], 1)

edited Apr 27 '20 at 20:22

answered Apr 27 '20 at 20:12

hilberts_drinking_problem

11,322
3
22
51

Mad Physicist · Answer 3 · 2020-04-27T20:48:10.923

Since rolling will require moving the data into a new array anyway, you can create a list of the indices that each row will map to. We can permute the elements of A along an axis you select as the "row" axis, keeping all the other axes in-place.

def roll_axis_elements(A, r, raxis=-2, caxis=-1):
    A = np.array(A, copy=False, subok=True)
    *index, = np.indices(A.shape, sparse=True)
    shape = list(A.shape)
    shape[raxis] = shape[caxis] = 1
    r = np.reshape(r, shape)
    index[raxis] = (index[raxis] - r) % A.shape[raxis]
    return A[tuple(index)]

This will work on any input shape of A. raxis is the "row" axis that you want to permute, and caxis are the matrix columns that move as a unit. You can trivially expand this to multiple axes by setting all of them to 1 before reshaping r, which will have to exclude them from its shape as well. All other axes remain as-is. This is basically a generalized version of Divakar's answer.

Sample run:

>>> A = [[[0, 0, 0], [1, 1, 1], [0, 0, 0], [2, 2, 2]],
         [[0, 0, 0], [1, 1, 1], [0, 0, 0], [2, 2, 2]]]
>>> r = [1, -2]
>>> roll_axis_elements(A, r)
array([[[2, 2, 2],
        [0, 0, 0],
        [1, 1, 1],
        [0, 0, 0]],

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

@Divakar. Definitely. `np.ogrid` gets messy if you need to expand it to an indeterminate number of dimensions. — Mad Physicist, Apr 27 '20 at 20:43

python: shift each matrix inside a tensor independantly

3 Answers3