Маскирование массива Numpy и применение расчета по маске без использования цикла for

Question

Допустим, у нас есть следующий массив данных:

data_array = np.array([[1, 1, 1], [1, 1, 2], [2, 2, 2], [3, 3, 3], [4, 4, 4]], np.int16)

data_array
array([[1, 1, 1],
       [1, 1, 2],
       [2, 2, 2],
       [3, 3, 3],
       [4, 4, 4]])

И мы хотим замаскировать массив в соответствии со следующими диапазонами, чтобы иметь возможность применить вычисления к замаскированным частям:

intervals = [[1, 2], [2, 3], [3, 4]]

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

init = np.zeros((data_array.shape[0], data_array.shape[1]))
result_array = np.ma.masked_where((init == 0), init)

result_array
masked_array(
data=[[--, --, --],
      [--, --, --],
      [--, --, --],
      [--, --, --],
      [--, --, --]],
mask=[[ True,  True,  True],
      [ True,  True,  True],
      [ True,  True,  True],
      [ True,  True,  True],
      [ True,  True,  True]]

С этим мы можем запустить цикл for, который маскирует массив в соответствии синтервал интервалов, выполняет вычисление для маскированного массива и объединяет результаты в один массив результатов:

for inter in intervals:

    # Extact the start and en values for interval range
    start_inter = inter[0]
    end_inter = inter[1]

    # Mask the array based on interval range
    mask_init = np.ma.masked_where((data_array > end_inter), data_array)
    masked_array = np.ma.masked_where((mask_init < start_inter), mask_init)

    # Perform a dummy calculation on masked array
    outcome = (masked_array + end_inter) * 100

    # Combine the outcome arrays
    result_array[result_array.mask] = outcome[result_array.mask]

со следующим результатом:

array([[300.0, 300.0, 300.0],
      [300.0, 300.0, 400.0],
      [400.0, 400.0, 400.0],
      [600.0, 600.0, 600.0],
      [800.0, 800.0, 800.0]])

У меня есть вопросКак можно достичь того же результата без использования цикла for?Таким образом, применяя маскирование и расчет для всего data_array в одной операции.Обратите внимание, что переменные расчета меняются с каждой маской.Можно ли применить векторизованный подход к этой проблеме?Я думаю, что numpy_indexed может помочь.Спасибо.

Answer 1

Если интервалы могут быть сделаны не перекрывающимися, то вы можете использовать функцию, подобную этой:

import numpy as np

def func(data_array, intervals):
    data_array = np.asarray(data_array)
    start, end = np.asarray(intervals).T
    data_array_exp = data_array[..., np.newaxis]
    mask = (data_array_exp >= start) & (data_array_exp <= end)
    return np.sum((data_array_exp + end) * mask * 100, axis=-1)

Результат должен быть таким же, как и в исходном коде в этом случае:

import numpy as np

def func_orig(data_array, intervals):
    init = np.zeros((data_array.shape[0], data_array.shape[1]))
    result_array = np.ma.masked_where((init == 0), init)
    for inter in intervals:
        start_inter = inter[0]
        end_inter = inter[1]
        mask_init = np.ma.masked_where((data_array > end_inter), data_array)
        masked_array = np.ma.masked_where((mask_init < start_inter), mask_init)
        outcome = (masked_array + end_inter) * 100
        result_array[result_array.mask] = outcome[result_array.mask]
    return result_array.data

data_array = np.array([[1, 1, 1], [1, 1, 2], [2, 2, 2], [3, 3, 3], [4, 4, 4]], np.int16)
intervals = [[1, 1.9], [2, 2.9], [3, 4]]
print(np.allclose(func(data_array, intervals), func_orig(data_array, intervals)))
# True