Numpy: ¿escribir una función en forma vectorial?

Question

Soy bastante nuevo en NumPy (o SciPy) y viniendo de Octave / Matlab, esto me parece un poco desafiante.

Estoy leyendo los documentos y escribiendo algunas funciones básicas. Me encontré con esta sección: Funciones de vectorización (vectorizar)

Define esta función:

def addsubtract(a, b):
   if a > b:
       return a - b
   else:
       return a + b

Luego lo vectoriza:

vec_addsubtract = np.vectorize(addsubtract)

Pero al final dice:

Esta función en particular podría haberse escrito en forma vectorial sin el uso de vectorizar.

No conocería ninguna otra forma de escribir tal función. Entonces, ¿cuál es la forma vectorial?

Answer 1

np.vectorizees un forbucle de python glorificado , lo que significa que elimina efectivamente cualquier optimización que ofrece numpy.

Para vectorizar realmente addsubtract, podemos usar el hecho de que numpy ofrece tres cosas: una addfunción vectorizada , una subtractfunción vectorizada y todo tipo de operaciones de máscara booleana.

La forma más simple, pero menos eficiente, de escribir esto es usando np.where:

np.where(a > b, a - b, a + b)

Esto es ineficaz porque calcula previamente a - by a + ben todos los casos, y luego selecciona uno u otro para cada elemento.

Una solución más eficiente solo calcularía los valores donde la condición lo requiriera:

result = np.empty_like(a)
mask = a > b
np.subtract(a, b, where=mask, out=result)
np.add(a, b, where=~mask, out=result)

Para arreglos muy pequeños, la sobrecarga del método complicado hace que valga menos la pena. Pero para arreglos grandes, es la solución más rápida.

Dato curioso: la página del tutorial a la que hace referencia no estará disponible en futuras versiones del tutorial SciPy exactamente porque es una introducción a NumPy, como se explica en PR # 12432 .

Answer 2

Puede hacer esto con np.where, que calcula ambos resultados ( a-by a+b) y selecciona los valores según una matriz booleana ( a>b):

def addsubtract(a, b):
    return np.where(a>b, a-b, a+b)

Puede verse como un operador ternario vectorizado: "Donde a> b, toma el valor de ab, de lo contrario toma el valor de a + b".

A pesar de calcular ambos resultados posibles, fue significativamente más rápido que la función vectorizada if / else que escribió (al menos en mi máquina).