условный векторизованный расчет с массивами numpy без использования прямого маскирования

Question

ответ на другой вопрос

import numpy as np

repeat=int(1e5)
r_base = np.linspace(0,4,5)
a_base = 2
np.random.seed(0)
r_mat = r_base * np.random.uniform(0.9,1.1,(repeat,5))

a_array = a_base * np.random.uniform(0.9,1.1, repeat)


# original slow approach
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
def func_list_level2(r_mat, a_array):
    res_mat = np.array([func_vetorized_level1(this_r_row, this_a) 
                        for this_r_row, this_a in zip(r_mat, a_array)])
    return res_mat

# faster with direct masking, but with unnecessary more calculation
def f_faster(r_mat,a_array):
    a = a_array[:, None]  # to column vector

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

    out = np.empty_like(r_mat)

    out[row_mask & elem_mask] = (r_mat - a)[row_mask & elem_mask]
    out[~row_mask & elem_mask] = (r_mat + a)[~row_mask & elem_mask]
    out[row_mask & ~elem_mask] = np.nan
    out[~row_mask & ~elem_mask] = 0
    
    return out

# fastest with ufunc in numpy as suggested by @mad_physicist
def f_fastest(r_mat,a_array):
    a = a_array[:, None]  # to column vector

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

    out = np.empty_like(r_mat)


    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
    
    return out

Я хотел бы спросить, можно ли получить пользовательское удовольствие c, которое можно использовать или воспользоваться самого быстрого подхода? Я думал об использовании индексации, но обнаружил, что это сложно, потому что нарезанные элементы, использующие [row_ind, co_ind], представляют собой 1d-массив выбранных элементов. Я вижу, что нарезанную матрицу можно поместить в матрицу, используя reshape, но есть ли элегантный способ сделать это? В идеале эту операцию r_mat + a можно заменить пользовательской функцией.

Answer 1

У вас абсолютно может быть векторизованное решение с функцией, определяемой пользователем, если эта функция векторизована для поэлементной работы с 1D-массивом (что должно быть в случае всего, написанного с использованием numpy функций вне box).

Допустим, у вас есть r_mat как матрица (m, n) и a_array как (m,) вектор. Вы можете написать свою функцию для приема хуков. Каждый хук может быть константой или вызываемой. Если это вызываемый объект, он вызывается с двумя массивами одинаковой длины и должен возвращать третий массив такой же длины. Вы можете изменить этот контракт, включив в него индексы или что угодно, по желанию:

def f(r_mat, a_array, hook11, hook01, hook10, hook00):
    a = a_array[:, None]  # to column vector

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

    out = np.empty_like(r_mat)

    def apply_hook(mask, hook):
        r, c = np.nonzero(mask)
        out[r, c] = hook(r_mat[r, c], a_array[r]) if callable(hook) else hook

    apply_hook(row_mask & elem_mask, hook11)
    apply_hook(~row_mask & elem_mask, hook01)
    apply_hook(row_mask & ~elem_mask, hook10)
    apply_hook(~row_mask & ~elem_mask, hook00)

    return out

Текущая конфигурация в вашем коде будет называться как

f(r_mat, a_array, np.subtract, np.add, np.nan, 0)

Допустим, вы хотели что-то сделать сложнее, чем np.subtract. Например, вы можете сделать:

def my_complicated_func(r, a):
    return np.cumsum(r, a) - 3 * r // a + np.exp(a)

f(r_mat, a_array, my_complicated_func, np.add, np.nan, 0.0)

Ключ в том, что my_complicated_func работает с массивами. Ему будет передано подмножество элементов r_mat и элементов a_array, дублированных столько раз, сколько необходимо в каждой строке.

Вы также можете сделать то же самое с функцией, зная о индекс каждого места. Просто позвоните hook как hook(r_mat[r, c], a_array[r], r, c). Теперь функции ловушки должны принимать два дополнительных аргумента. Исходный код будет эквивалентен

f(r_mat, a_array, lambda r, a, *args: np.subtract(r, a), lambda r, a, *args: np.add(r, a), np.nan, 0)