有效地旋转3D阵列中的块/窗口（矢量化扩散？）

Question

我有一个大型3D np.array，可以说尺寸为（200,200,7）。

我想对第一个轴上的每个2 * 2子数组应用np.rot90。另一个问题是以随机方式旋转每个子阵列。像这样：

颜色仅用于显示不同的2 * 2数组，箭头表示每个数组都是根据numpy.rot90(m, k=RND(1,2,3), axes=(0, 1))中作为参数生成的随机数旋转的。

是否可以在一个快速步骤中实现，而无需在每个子阵列上循环？

根据Divakar的回答，我还尝试制作一个扩展子，其中只有x％的子阵列一步移动，其余部分保持不变，这有望像一个2D扩散系统。

def vectorized_diffusion(a,H,W,pD):
   #pD - chance that a sub-array is rotated in a random direction
   rand_shift = np.random.randint(-1,2)
   rand_axis = np.random.randint(0,2)
   a = np.roll(a, shift = randshift, axis = rand_axis)
   # Since the 2*2 subgrid system is fixed, I decided to ocassionally
   #disturb the grid by rolling the whole array by one in a given 
   #direction, as in my work the array is a toroid grid i considered every direction
   m,n,r = a.shape
   a5D = a.reshape(m//H,H,n//W,W,-1)
   cw0 = a5D[:,::-1,:,:,:].transpose(0,2,3,1,4)
   ccw0 = a5D[:,:,:,::-1,:].transpose(0,2,3,1,4)
   original = a5D[:,:,:,:,:].transpose(0,2,1,3,4)
   mask_clockdirection = np.random.choice([False,True],size=(m//H,n//W))
   mask_stationary = np.random.choice([True,False],size=(m//H,n//W), p=[1-pD,pD])
   w0 = np.where(mask_clockdirection[:,:,None,None,None],cw0,ccw0)
   out = np.where(mask_stationary[:,:,None,None,None],original,w0)
   out = out.swapaxes(1,2).reshape(a.shape)
   out_rerolled = np.roll(out, shift = -1*randshift, axis = rand_axis)
   #this way the disturbed grid is rerolled into its original position
   return out_rerolled

我知道这可能不是解决这个问题的最优雅的解决方案，但它似乎可行，我对此表示满意。

Answer 1

使用翻转和排列轴进行旋转（顺时针和逆时针）的通用方式-

# Input array
In [176]: k
Out[176]: 
array([[26, 48, 71],
       [54, 96, 82],
       [87, 21,  2]])

# Clockwise
In [178]: k[::-1,:].T
Out[178]: 
array([[87, 54, 26],
       [21, 96, 48],
       [ 2, 82, 71]])

# Anti-clockwise
In [177]: k[:,::-1].T
Out[177]: 
array([[71, 82,  2],
       [48, 96, 21],
       [26, 54, 87]])

通过窗口旋转扩展到2D数组

In [204]: np.random.seed(0)

In [205]: a = np.random.randint(0,100,(6,6))

In [206]: a
Out[206]: 
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],
       [72,  9, 20, 80, 69, 79],
       [47, 64, 82, 99, 88, 49]])

# Clockwise
In [207]: a.reshape(3,2,3,2)[:,::-1,:,:].swapaxes(1,3).reshape(a.shape)
Out[207]: 
array([[83, 44, 36, 64, 70, 67],
       [21, 47, 87, 67, 88,  9],
       [46, 88, 81, 58, 25, 39],
       [88, 12, 37, 65, 77, 87],
       [47, 72, 82, 20, 88, 69],
       [64,  9, 99, 80, 49, 79]])

# Anti-clockwise
In [209]: a.reshape(3,2,3,2)[:,:,:,::-1].swapaxes(1,3).reshape(a.shape)
Out[209]: 
array([[47, 21, 67, 87,  9, 88],
       [44, 83, 64, 36, 67, 70],
       [12, 88, 65, 37, 87, 77],
       [88, 46, 58, 81, 39, 25],
       [ 9, 64, 80, 99, 79, 49],
       [72, 47, 20, 82, 69, 88]])

扩展为3D数组，并在每个2D切片上进行窗口旋转-

In [223]: a = np.random.randint(0,100,(6,6,2))

# Clockwise
In [224]: cw = a.reshape(3,2,3,2,2)[:,::-1,:,:,:].swapaxes(1,3).reshape(a.shape)

# Anti-clockwise
In [233]: ccw = a.reshape(3,2,3,2,2)[:,:,:,::-1,:].swapaxes(1,3).reshape(a.shape)

In [225]: a[...,0]
Out[225]: 
array([[44, 64, 67, 83, 36, 70],
       [88, 58, 39, 46, 81, 25],
       [72, 20, 69, 47, 82, 88],
       [29, 19, 39, 65, 57, 31],
       [23, 75, 28,  0, 36,  5],
       [17,  4, 58,  1, 41, 35]])

In [226]: cw[...,0]
Out[226]: 
array([[88, 44, 39, 67, 81, 36],
       [58, 64, 46, 83, 25, 70],
       [29, 72, 39, 69, 57, 82],
       [19, 20, 65, 47, 31, 88],
       [17, 23, 58, 28, 41, 36],
       [ 4, 75,  1,  0, 35,  5]])


In [236]: ccw[...,0]
Out[236]: 
array([[64, 58, 83, 46, 70, 25],
       [44, 88, 67, 39, 36, 81],
       [20, 19, 47, 65, 88, 31],
       [72, 29, 69, 39, 82, 57],
       [75,  4,  0,  1,  5, 35],
       [23, 17, 28, 58, 36, 41]])

解决我们的情况，并选择戴着面具的这两种情况

我们需要使其适合我们的情况。我们将使用遮罩在顺时针和逆时针版本之间进行选择-

cw0 = a.reshape(3,2,3,2,2)[:,::-1,:,:,:].swapaxes(1,3)
ccw0 = a.reshape(3,2,3,2,2)[:,:,:,::-1,:].swapaxes(1,3)
mask = np.random.choice([False,True],size=(3,3))
out = np.where(mask[:,:,None,None,None],cw0.swapaxes(1,2),ccw0.swapaxes(1,2))

我们可以优化/使其更紧凑-

cw0 = a.reshape(3,2,3,2,2)[:,::-1,:,:,:].transpose(0,2,3,1,4)
ccw0 = a.reshape(3,2,3,2,2)[:,:,:,::-1,:].transpose(0,2,3,1,4)
out = np.where(mask[:,:,None,None,None],cw0,ccw0)

最后，让我们处理一般情况-

def random_rotate_windows(a,H,W):
    m,n,r = a.shape
    a5D = a.reshape(m//H,H,n//W,W,-1)
    cw0 = a5D[:,::-1,:,:,:].transpose(0,2,3,1,4)
    ccw0 = a5D[:,:,:,::-1,:].transpose(0,2,3,1,4)
    mask = np.random.choice([False,True],size=(m//H,n//W))
    out = np.where(mask[:,:,None,None,None],cw0,ccw0)
    return out.swapaxes(1,2).reshape(a.shape)

以示例运行结束-

In [332]: np.random.seed(0)
     ...: a = np.random.randint(0,100,(6,6,2))

In [333]: a[...,0]
Out[333]: 
array([[44, 64, 67, 83, 36, 70],
       [88, 58, 39, 46, 81, 25],
       [72, 20, 69, 47, 82, 88],
       [29, 19, 39, 65, 57, 31],
       [23, 75, 28,  0, 36,  5],
       [17,  4, 58,  1, 41, 35]])

In [334]: out = random_rotate_windows(a,2,2)

In [335]: out[...,0]
Out[335]: 
array([[64, 58, 83, 46, 81, 36],
       [44, 88, 67, 39, 25, 70],
       [20, 19, 47, 65, 57, 82],
       [72, 29, 69, 39, 31, 88],
       [17, 23,  0,  1, 41, 36],
       [ 4, 75, 28, 58, 35,  5]])