Эффективно вращайте блоки / окна в трехмерном массиве (векторизованная диффузия?)

Question

У меня большой 3D np.array, скажем, с размером (200,200,7).

Я бы хотел применить np.rot90 для каждого 2 * 2 подмассива в первомоси.Другой проблемой является случайное вращение каждого подмассива.Например:

Цвета предназначены только для отображения отдельных массивов 2 * 2, стрелки показывают, что каждый массив вращается на основе случайного числа, сгенерированного как аргумент в numpy.rot90(m, k=RND(1,2,3), axes=(0, 1)).

Возможно ли это за один быстрый шаг, без зацикливания на каждом отдельном подмассиве?

Основываясь на ответе Дивакара, я также попытался создать экстенсон, в котором только х процентов отПодмассивы перемещаются за один шаг, остальные остаются без изменений, и мы надеемся, что это ведет себя как двумерная диффузионная система.

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]])