Заполнить маскированный массив по столбцам

Question

Кажется, я скучаю по быстрому способу заполнения замаскированной области в порядке столбцов в Python.

Вот пример, который работает, но очень неэффективно:

      def fill_masked_array_colmajor(arr_to_fill, mask, fill_values):
    c = 0
    for h in range(arr_to_fill.shape[1]):
        for v in range(arr_to_fill.shape[0]):
            if mask[v,h]:
                arr_to_fill[v,h] = fill_values[c,0]
                c+=1
    return arr_to_fill

test_arr = np.arange(1,65).reshape(8,8)
test_mask = np.zeros((8,8))
test_mask[2:6,3:5] = 1
test_mask = test_mask.astype(bool)
test_vals = np.arange(1,9)[..., np.newaxis]  # added axis because my data looks like that

test_res = fill_masked_array_colmajor(test_arr, test_mask, test_vals)

Правильный вывод становится следующим:

       [ 1  2  3  4  5  6  7  8]
 [ 9 10 11 12 13 14 15 16]
 [17 18 19  1  5 22 23 24]
 [25 26 27  2  6 30 31 32]
 [33 34 35  3  7 38 39 40]
 [41 42 43  4  8 46 47 48]
 [49 50 51 52 53 54 55 56]
 [57 58 59 60 61 62 63 64]

Если я попробую что-то вроде:

      test_arr2 = np.arange(1,65).reshape(8,8)
test_mask2 = np.zeros((8,8))
test_mask2[2:6,3:5] = 1
test_mask2 = test_mask2.astype(bool)
test_vals = np.arange(1,9)[..., np.newaxis] # added axis because my data looks like that

test_arr2[test_mask2] = test_vals[:,0]

тогда неправильный основной вывод строки становится:

      [ 1  2  3  4  5  6  7  8]
[ 9 10 11 12 13 14 15 16]
[17 18 19  1  2 22 23 24]
[25 26 27  3  4 30 31 32]
[33 34 35  5  6 38 39 40]
[41 42 43  7  8 46 47 48]
[49 50 51 52 53 54 55 56]
[57 58 59 60 61 62 63 64]

Важно отметить, что замаскированная область не обязательно имеет прямоугольную форму, как показано здесь, но может иметь причудливую форму и множество пятен. Длинаtest_valsвсегда равно количеству маскируемых элементов. Мне просто нужно заполнить их в порядке столбцов.

Я понимаю такое поведение, поскольку Python является основным по строкам. Возможно, я вижу решение, используя.flatten(order="f")и сделать несколько копий, а затем изменить форму - но должна ли быть быстрая и простая реализация этого? Я чувствую, что упускаю что-то очевидное. Спасибо.

Новая неудачная попытка

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

      idx = np.argwhere(test_mask2.ravel("F"))
tmp = test_arr2.ravel("F")
tmp[idx] = test_vals
out = tmp.reshape(test_arr2.shape).T