Сдвиг строки массива Numpy независимо

Question

Это расширение поставленного вопроса здесь (цитируется ниже)

У меня есть матрица (точнее, 2d numpy ndarray):

A = np.array([[4, 0, 0],
              [1, 2, 3],
              [0, 0, 5]])

И я хочу свернуть каждый ряд А независимо, в соответствии с броском значения в другом массиве:

r = np.array([2, 0, -1])

То есть я хочу сделать это:

print np.array([np.roll(row, x) for row,x in zip(A, r)])

[[0 0 4]
 [1 2 3]
 [0 5 0]]

Есть ли способ сделать это эффективно? Возможно, с помощью модной индексации трюки?

Принято решение:

rows, column_indices = np.ogrid[:A.shape[0], :A.shape[1]]

# Use always a negative shift, so that column_indices are valid.
# (could also use module operation)
r[r < 0] += A.shape[1]
column_indices = column_indices - r[:,np.newaxis]

result = A[rows, column_indices]

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

Может быть, с помощью np.pad как-то? Но я не могу понять, как заставить это заполнить разные строки разными значениями.

Answer 1

Вдохновленный Свернуть строки матрицы независимо решения , вот векторизованный на основе np.lib.stride_tricks.as_strided -

from skimage.util.shape import view_as_windows as viewW

def strided_indexing_roll(a, r):
    # Concatenate with sliced to cover all rolls
    p = np.full((a.shape[0],a.shape[1]-1),np.nan)
    a_ext = np.concatenate((p,a,p),axis=1)

    # Get sliding windows; use advanced-indexing to select appropriate ones
    n = a.shape[1]
    return viewW(a_ext,(1,n))[np.arange(len(r)), -r + (n-1),0]

Пробный прогон -

In [76]: a
Out[76]: 
array([[4, 0, 0],
       [1, 2, 3],
       [0, 0, 5]])

In [77]: r
Out[77]: array([ 2,  0, -1])

In [78]: strided_indexing_roll(a, r)
Out[78]: 
array([[nan, nan,  4.],
       [ 1.,  2.,  3.],
       [ 0.,  5., nan]])

Answer 2

Мне удалось взломать это вместе с линейным индексированием ... это дает правильный результат, но работает довольно медленно на больших массивах.

A = np.array([[4, 0, 0],
              [1, 2, 3],
              [0, 0, 5]]).astype(float)

r = np.array([2, 0, -1])

rows, column_indices = np.ogrid[:A.shape[0], :A.shape[1]]

# Use always a negative shift, so that column_indices are valid.
# (could also use module operation)
r_old = r.copy()
r[r < 0] += A.shape[1]
column_indices = column_indices - r[:,np.newaxis]

result = A[rows, column_indices]

# replace with NaNs
row_length = result.shape[-1]

pad_inds = []
for ind,i in np.enumerate(r_old):
    if i > 0:
        inds2pad = [np.ravel_multi_index((ind,) + (j,),result.shape) for j in range(i)]
        pad_inds.extend(inds2pad)
    if i < 0:
        inds2pad = [np.ravel_multi_index((ind,) + (j,),result.shape) for j in range(row_length+i,row_length)]
        pad_inds.extend(inds2pad)
result.ravel()[pad_inds] = nan

Дает ожидаемый результат:

print result

[[ nan  nan   4.]
 [  1.   2.   3.]
 [  0.   5.  nan]]