Numpy, Numba y compatible con nopythonmode, manera de evitar el ciclo for.
Tengo dos matrices, una en 2-D (N x N) y la otra en 3-D (N x N x 3), respectivamente xi y R.
Tengo que realizar el siguiente cálculo:
h = 0
for i in range(N):
h += xi[i].dot(R[i].dot(R[i].T).dot(xi[i]))
Esto está incluido en una función decorada con numba (@njit). Estoy tratando de optimizar la función y me gustaría encontrar una manera de evitar el bucle for. ¿Hay alguna manera de hacerlo con solo funciones compatibles con numpy de Numba?
PD: sé que podría usar prange y ejecutar el bucle for en paralelo para obtener algo de velocidad adicional, pero como N es bastante pequeño (~100), la función en realidad funciona peor decorada con @njit(parallel=True).
davy.ai
Sí, es posible evitar el ciclo for utilizando simplemente las funciones de numpy y las funciones soportadas por Numba. Puedes usar el siguiente código:
Aquí,
np.einsumrealiza una multiplicación de matrices entreRyxi, y luego toma el producto interno consigo misma a lo largo del último eje. El arreglo resultante se suma a lo largo de las filas para obtener el resultado final. Esta función utiliza solo funciones de numpy soportadas por Numba y debería ser más rápida que el loop original, especialmente si se decora con@njit.