Как выполнить вычисление векторизации между 1D и 2D numpy массивом с условиями if

Question

У меня есть расчет с использованием 1D и 2D массивов numpy. Он имеет два уровня if -условий. Я смог использовать np.where, чтобы избежать одного if -выражения и в дальнейшем использовать медленное понимание списка для итерации по каждой строке.

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

Вот мой код:

import numpy as np

r_base = np.linspace(0, 4, 5)
np.random.seed(0)
r_mat = np.array([r_base * np.random.uniform(0.9, 1.1, 5),
                  r_base * np.random.uniform(0.9, 1.1, 5),
                  r_base * np.random.uniform(0.9, 1.1, 5)])

a_array = np.linspace(1, 3, 3)

def func_vetorized_level1(r_row, a):
    if r_row.mean() > 2:
        result = np.where((r_row >= a), r_row - a, np.nan)
    else:
        result = np.where((r_row >= a), r_row + a, 0)
    return result
# try to broadcast this func to every row of r_mat using list comprehension
res_mat = np.array([func_vetorized_level1(this_r_row, this_a) 
                    for this_r_row, this_a in zip(r_mat, a_array)])

результат

res_mat =
array([[       nan, 0.04303787, 1.04110535, 2.02692991, 2.93892384],
       [       nan,        nan, 0.1567092 , 1.27819766, 1.90675322],
       [0.        , 0.        , 0.        , 6.25535798, 6.65682885]])

Answer 1

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

Для генерации целочисленного диапазона np.arange работает лучше, чем np.linspace:

r_base = np.arange(5.)
a_array = np.arange(1., 4.)

Случайные числа могут можно сделать за один вызов с одним умножением:

np.random.seed(0)
r_mat = r_base * np.random.uniform(0.9, 1.1, (3, 5))

Я думаю, что проще всего было бы создать выходной массив и заполнить его на основе различных условий:

out = np.empty_like(r_mat)

Было бы полезно превратить a_array в столбец, который соответствует количеству строк в r_mat:

a = a_array[:, None]

Следующее, что вам нужно сделать, это сделать маски для условий. Первый - это построчная маска для r_row.mean() > 2. Второе - это поэлементное r_row >= a условие:

row_mask = (r_mat.mean(axis=1) > 2)[:, None]
elem_mask = r_mat >= a

Индекс [:, None] на row_mask превращает его в вектор-столбец для целей широковещания. Теперь вы можете оценить выбор, используя прямое маскирование и ключевое слово where для соответствующего ufunc s:

np.subtract(r_mat, a, out=out, where=row_mask & elem_mask)
np.add(r_mat, a, out=out, where=~row_mask & elem_mask)
out[row_mask & ~elem_mask] = np.nan
out[~row_mask & ~elem_mask] = 0