Сдвиньте массив numpy

Я бы хотел «сдвинуть» массив с множеством цифр. Я не уверен, что правильно использую термин «сдвиг»; с помощью сдвига, я имею в виду что-то вроде:

Сдвиньте первый столбец на 0 мест
Сдвиньте второй столбец на 1 место
Сдвиг третьего кола на 2 места
и т.д…

Итак, этот массив:

array([[11, 12, 13], [17, 18, 19], [35, 36, 37]]) 

превратится в этот массив:

 array([[11, 36, 19], [17, 12, 37], [35, 18, 13]]) 

или что-то вроде этого массива:

 array([[11, 0, 0], [17, 12, 0], [35, 18, 13]]) 

в зависимости от того, как мы обрабатываем края. Я не слишком разбираюсь в поведении.

Вот моя попытка выполнить функцию:

 import numpy def shear(a, strength=1, shift_axis=0, increase_axis=1, edges='clip'): strength = int(strength) shift_axis = int(shift_axis) increase_axis = int(increase_axis) if shift_axis == increase_axis: raise UserWarning("Shear can't shift in the direction it increases") temp = numpy.zeros(a.shape, dtype=int) indices = [] for d, num in enumerate(a.shape): coords = numpy.arange(num) shape = [1] * len(a.shape) shape[d] = num coords = coords.reshape(shape) + temp indices.append(coords) indices[shift_axis] -= strength * indices[increase_axis] if edges == 'clip': indices[shift_axis][indices[shift_axis] < 0] = -1 indices[shift_axis][indices[shift_axis] >= a.shape[shift_axis]] = -1 res = a[indices] res[indices[shift_axis] == -1] = 0 elif edges == 'roll': indices[shift_axis] %= a.shape[shift_axis] res = a[indices] return res if __name__ == '__main__': a = numpy.random.random((3,4)) print a print shear(a) 

Кажется, это работает. Пожалуйста, скажите мне, если это не так!

Это также кажется неуклюжим и нелогичным. Могу ли я игнорировать встроенную функцию numpy / scipy, которая делает это? Есть ли более чистый / лучший / более эффективный способ сделать это в numpy? Я изобретаю колесо?

РЕДАКТИРОВАТЬ:
Бонусные очки, если это работает в N-мерном массиве, а не только в 2D-случае.

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

ВТОРАЯ РЕДАКТИРОВКА: Я, наконец, сделал бенчмаркинг. Похоже, numpy.roll – это путь, несмотря на цикл. Спасибо, tom10 и Свен Марнах!

Код бенчмаркинга: (запускать в Windows, не использовать time.clock на Linux, я думаю)

 import time, numpy def shear_1(a, strength=1, shift_axis=0, increase_axis=1, edges='roll'): strength = int(strength) shift_axis = int(shift_axis) increase_axis = int(increase_axis) if shift_axis == increase_axis: raise UserWarning("Shear can't shift in the direction it increases") temp = numpy.zeros(a.shape, dtype=int) indices = [] for d, num in enumerate(a.shape): coords = numpy.arange(num) shape = [1] * len(a.shape) shape[d] = num coords = coords.reshape(shape) + temp indices.append(coords) indices[shift_axis] -= strength * indices[increase_axis] if edges == 'clip': indices[shift_axis][indices[shift_axis] < 0] = -1 indices[shift_axis][indices[shift_axis] >= a.shape[shift_axis]] = -1 res = a[indices] res[indices[shift_axis] == -1] = 0 elif edges == 'roll': indices[shift_axis] %= a.shape[shift_axis] res = a[indices] return res def shear_2(a, strength=1, shift_axis=0, increase_axis=1, edges='roll'): indices = numpy.indices(a.shape) indices[shift_axis] -= strength * indices[increase_axis] indices[shift_axis] %= a.shape[shift_axis] res = a[tuple(indices)] if edges == 'clip': res[indices[shift_axis] < 0] = 0 res[indices[shift_axis] >= a.shape[shift_axis]] = 0 return res def shear_3(a, strength=1, shift_axis=0, increase_axis=1): if shift_axis > increase_axis: shift_axis -= 1 res = numpy.empty_like(a) index = numpy.index_exp[:] * increase_axis roll = numpy.roll for i in range(0, a.shape[increase_axis]): index_i = index + (i,) res[index_i] = roll(a[index_i], i * strength, shift_axis) return res numpy.random.seed(0) for a in ( numpy.random.random((3, 3, 3, 3)), numpy.random.random((50, 50, 50, 50)), numpy.random.random((300, 300, 10, 10)), ): print 'Array dimensions:', a.shape for sa, ia in ((0, 1), (1, 0), (2, 3), (0, 3)): print 'Shift axis:', sa print 'Increase axis:', ia ref = shear_1(a, shift_axis=sa, increase_axis=ia) for shear, label in ((shear_1, '1'), (shear_2, '2'), (shear_3, '3')): start = time.clock() b = shear(a, shift_axis=sa, increase_axis=ia) end = time.clock() print label + ': %0.6f seconds'%(end-start) if (b - ref).max() > 1e-9: print "Something's wrong." print 

 def shear3(a, strength=1, shift_axis=0, increase_axis=1): if shift_axis > increase_axis: shift_axis -= 1 res = numpy.empty_like(a) index = numpy.index_exp[:] * increase_axis roll = numpy.roll for i in range(0, a.shape[increase_axis]): index_i = index + (i,) res[index_i] = roll(a[index_i], -i * strength, shift_axis) return res 

 for i in range(x.shape[1]): x[:,i] = np.roll(x[:,i], i) 

 [[11 36 19] [17 12 37] [35 18 13]] 

 from numpy.lib.stride_tricks import as_strided a = numpy.array([[11, 12, 13], [17, 18, 19], [35, 36, 37]]) shift_axis = 0 increase_axis = 1 b = numpy.vstack((a, a)) strides = list(b.strides) strides[increase_axis] -= strides[shift_axis] strides = (b.strides[0], b.strides[1] - b.strides[0]) as_strided(b, shape=b.shape, strides=strides)[a.shape[0]:] # array([[11, 36, 19], # [17, 12, 37], # [35, 18, 13]]) 

 b = numpy.vstack((numpy.zeros(a.shape, int), a)) 

 def shear2(a, strength=1, shift_axis=0, increase_axis=1, edges='clip'): indices = numpy.indices(a.shape) indices[shift_axis] -= strength * indices[increase_axis] indices[shift_axis] %= a.shape[shift_axis] res = a[tuple(indices)] if edges == 'clip': res[indices[shift_axis] < 0] = 0 res[indices[shift_axis] >= a.shape[shift_axis]] = 0 return res 

 r = lambda l, n: l[n:]+l[:n] transpose(map(r, transpose(a), range(0, len(a)))