CODEWITHSUNDEEP
pythonarraysnumpy
ascripter
ascripter

Reputation: 6213

numpy: broadcast array by rolling along new axis with variable shift given in 2nd array

I know that numpy.roll can shift an array along one or more existing axes. How would I create a new axis on array x along which I want to have views or copies of itself rolled by array shift?

Example:

x = np.arange(10)
shift = np.array([2, 4])

#input
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

#output
array(
  [[8, 6],
   [9, 7],
   [0, 8],
   [1, 9],
   [2, 0],
   [3, 1],
   [4, 2],
   [5, 3],
   [6, 4],
   [7, 5]])

Edit: I'm looking for a general solution (ideally without looping) that can also be applied on higher dimensional arrays. Another example:

x = np.arange(20).reshape(2, 10)
shift = np.array([2, 4])

#input
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14, 15, 16, 17, 18, 19]])

#output
array(
  [[[ 8,  6],
    [ 9,  7],
    [ 0,  8],
    [ 1,  9],
    [ 2,  0],
    [ 3,  1],
    [ 4,  2],
    [ 5,  3],
    [ 6,  4],
    [ 7,  5]],

   [[18, 16],
    [19, 17],
    [10, 18],
    [11, 19],
    [12, 10],
    [13, 11],
    [14, 12],
    [15, 13],
    [16, 14],
    [17, 15]]])

Upvotes: 1

Views: 392

Answers (2)

Divakar
Divakar

Reputation: 221514

Here's a vectorized solution leveraging broadcasting that covers generic n-dim array cases -

np.take(x,(-shift + np.arange(x.shape[-1])[:,None]),axis=-1)

Sample runs

1) x as 1D -

In [114]: x = np.arange(10)
     ...: shift = np.array([2, 4])

In [115]: np.take(x,(-shift + np.arange(x.shape[-1])[:,None]),axis=-1)
Out[115]: 
array([[8, 6],
       [9, 7],
       [0, 8],
       [1, 9],
       [2, 0],
       [3, 1],
       [4, 2],
       [5, 3],
       [6, 4],
       [7, 5]])

2) x as 2D -

In [116]: x = np.arange(20).reshape(2, 10)
     ...: shift = np.array([2, 4])

In [117]: np.take(x,(-shift + np.arange(x.shape[-1])[:,None]),axis=-1)
Out[117]: 
array([[[ 8,  6],
        [ 9,  7],
        [ 0,  8],
        [ 1,  9],
        [ 2,  0],
        [ 3,  1],
        [ 4,  2],
        [ 5,  3],
        [ 6,  4],
        [ 7,  5]],

       [[18, 16],
        [19, 17],
        [10, 18],
        [11, 19],
        [12, 10],
        [13, 11],
        [14, 12],
        [15, 13],
        [16, 14],
        [17, 15]]])

Upvotes: 2

hpaulj
hpaulj

Reputation: 231355

I almost hate to provide this alternative because I think @BenT's answer is simple and logical

np.array([np.roll(x,sh) for sh in shift]).T
np.stack([np.roll(x,sh) for sh in shift], axis=1)  # may be easier to generalize

but I can do the original x=np.arange(10) case with as_strided:

Perform all shifts:

In [352]: arr = np.lib.stride_tricks.as_strided(np.hstack((x,x)),shape=(10,10), strides=(8,8))
In [353]: arr
Out[353]: 
array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
       [1, 2, 3, 4, 5, 6, 7, 8, 9, 0],
       [2, 3, 4, 5, 6, 7, 8, 9, 0, 1],
       [3, 4, 5, 6, 7, 8, 9, 0, 1, 2],
       [4, 5, 6, 7, 8, 9, 0, 1, 2, 3],
       [5, 6, 7, 8, 9, 0, 1, 2, 3, 4],
       [6, 7, 8, 9, 0, 1, 2, 3, 4, 5],
       [7, 8, 9, 0, 1, 2, 3, 4, 5, 6],
       [8, 9, 0, 1, 2, 3, 4, 5, 6, 7],
       [9, 0, 1, 2, 3, 4, 5, 6, 7, 8]])

Then select the ones you want:

In [358]: arr[::-1][shift-1]
Out[358]: 
array([[8, 9, 0, 1, 2, 3, 4, 5, 6, 7],
       [6, 7, 8, 9, 0, 1, 2, 3, 4, 5]])

I wrote and tested the stack version with one try, but had to try several things to get the as_strided right.

I'd also prefer generalizing the list comprehension to higher dimensions.

For your 2d x:

np.stack([np.roll(x,sh, axis=1) for sh in shift],2)

Upvotes: 0

Related Questions