通过在 ndimensions 中改变数量来滚动阵列轴（概括跨步索引滚动）

Question

我有一个任意形状的数组，但是假设 (A, B, C)，我想为每个元素（即每个 (A, B)）将最后一个轴滚动不同的量.

我正在尝试将@Divakar 的漂亮 solution here 归纳为二维数组，但我真的不明白 skimage.util.shape.view_as_windows 在做什么，所以我最终遇到了索引问题.

我的尝试：

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

但这失败了：

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)

Answer 1

对于3D，它会修改成这样-

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,:]

对于 n-dim 数组，要沿最后一个轴滚动，它将是 -

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,:]

通过在 ndimensions 中改变数量来滚动阵列轴（概括跨步索引滚动）

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

python

arrays

numpy

numpy-slicing