CODEWITHSUNDEEP
pythonarraysnumpynumpy-slicing
DilithiumMatrix
DilithiumMatrix

Reputation: 18627

Rolling an array axis by varying amounts in ndimensions (generalizing strided indexing roll)

I have an array of arbitrary shape, but let's say, (A, B, C), and I'd like to roll the last axis by a different amount for each element (i.e. for each (A, B)).

I'm trying to generalize @Divakar's beautiful solution here for 2D arrays, but I don't really understand what the skimage.util.shape.view_as_windows is doing, so I'm ending up with an indexing problem.

My attempt:

import numpy as np
from skimage.util.shape import view_as_windows as viewW

def strided_indexing_roll(a, r, axis=-1):
    a = np.asarray(a)
    r = np.asarray(r)
    a = np.moveaxis(a, axis, -1)

    ndim = np.ndim(a)

    # Repeat along the given axis to cover all rolls
    cut = [slice(None) for ii in range(ndim)]
    cut[-1] = slice(None, -1)
    cut = tuple(cut)
    a_ext = np.concatenate((a, a[cut]), axis=-1)

    # Get sliding windows; use advanced-indexing to select appropriate ones
    n = a.shape[-1]
    shape = np.ones(ndim, dtype=int)
    shape[-1] = n
    shape = tuple(shape)

    cut = [np.arange(jj) for jj in np.shape(r)]
    cut = cut + [(n - r) % n,]
    cut = cut + [0 for ii in range(ndim-1)]
    cut = tuple(cut)

    res = viewW(a_ext, shape)
    res = res[cut]
    res = np.moveaxis(res, -1, axis)
    return res

But this fails with:

aa = np.random.uniform(0.0, 1.0, 10)
bb = np.random.randint(0, aa.size, (2, 4))
shape = np.shape(bb) + (np.size(aa),)
aa = aa[np.newaxis, np.newaxis, :] * np.ones(shape)

strided_indexing_roll(aa, bb)

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-322-5f9c871acf06> in <module>
     99 aa = aa[np.newaxis, np.newaxis, :] * np.ones(shape)
    100 
--> 101 strided_indexing_roll(aa, bb)

<ipython-input-322-5f9c871acf06> in strided_indexing_roll(a, r, axis)
     75 
     76     res = viewW(a_ext, shape)
---> 77     res = res[cut]
     78     res = np.moveaxis(res, -1, axis)
     79     return res

IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (2,) (4,) (2,4)

Upvotes: 1

Views: 114

Answers (1)

Divakar
Divakar

Reputation: 221524

For 3D, it would modify to something like this -

def strided_indexing_roll_3d_lastaxis(a, r):
    # Concatenate with sliced to cover all rolls
    a_ext = np.concatenate((a,a[...,:-1]),axis=-1)

    # Get sliding windows; use advanced-indexing to select appropriate ones
    n = a.shape[-1]
    w = viewW(a_ext,(1,1,n))
    idx = (n-r)%n
    return np.take_along_axis(w,idx[:,:,None,None,None,None],axis=2)[:,:,0,0,0,:]

For n-dim arrays, to roll along last axis, it would be -

def strided_indexing_roll_nd_lastaxis(a, r):
    # Concatenate with sliced to cover all rolls
    a_ext = np.concatenate((a,a[...,:-1]),axis=-1)

    # Get sliding windows; use advanced-indexing to select appropriate ones
    n = a.shape[-1]
    w = viewW(a_ext,(1,)*r.ndim + (n,)).reshape(a.shape+(n,))
    idx = (n-r)%n    
    idxr = idx.reshape(idx.shape+(1,)*(w.ndim-r.ndim))
    return np.take_along_axis(w,idxr,axis=r.ndim)[...,0,:]

Upvotes: 1

Related Questions