通过在n维上变化量来滚动数组轴(推广步进分度辊)

Question

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

我试图将@Divakar的漂亮的solution here推广到2D数组，但我并不真正理解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维阵列，沿着最后一个轴滚动，它将是-

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