Una forma que encontré para mejorar su código (es decir, reducir la cantidad de bucles) es usar np.meshgrid .

Aquí está la mejora que encontré. Tomó un poco de manipulación, pero da el mismo resultado que su código de bucle triple. Mantuve la misma estructura de código para que puedas ver qué partes corresponden a qué parte. ¡Espero que esto te sea de utilidad!

 for sigma in range(M): xx, yy = np.meshgrid(samples[sigma], samples[sigma]) results += (1/N) * (xx * yy - yy * DT - xx * D) print(results) # or return results

.

Editar: aquí hay un pequeño script para verificar que los resultados sean los esperados:

 import numpy as np M, N = 3, 4 rng = np.random.default_rng(seed=42) samples = rng.random((M, N)) D = rng.random((N, N)) results = rng.random((N, N)) results_old = results.copy() results_new = results.copy() for sigma in range(M): for i in range(N): for j in range(N): results_old[i, j] += (1/N) * (samples[sigma, i]*samples[sigma, j] - samples[sigma, i]*D[j, i] - samples[sigma, j]*D[i, j]) print('\n\nresults_old', results_old, sep='\n') for sigma in range(M): xx, yy = np.meshgrid(samples[sigma], samples[sigma]) results_new += (1/N) * (xx * yy - yy * DT - xx * D) print('\n\nresults_new', results_new, sep='\n')

Edición 2 : deshacerse por completo de los bucles: es un poco complicado pero esencialmente hace lo mismo.

 M, N = samples.shape xxx, yyy = np.meshgrid(samples, samples) split_x = np.array(np.hsplit(np.vsplit(xxx, M)[0], M)) split_y = np.array(np.vsplit(np.hsplit(yyy, M)[0], M)) results += np.sum( (1/N) * (split_x*split_y - split_y*DT - split_x*D), axis=0) print(results) # or return results

Para vectorizar bucles for , podemos hacer uso de la transmisión y luego reducir a lo largo de los ejes que no se reflejan en la matriz de salida. Para hacerlo, podemos "asignar" un eje a cada uno de los índices del bucle for (como convención). Para su ejemplo, esto significa que todas las matrices de entrada se pueden remodelar para tener una dimensión 3 (es decir len(a.shape) == 3 ); los ejes corresponden entonces a sigma, i, j respectivamente. Luego podemos realizar todas las operaciones con las matrices transmitidas y finalmente reducir (sumar) el resultado a lo largo del eje sigma (ya que solo i, j se reflejan en el resultado):

 # Ordering of axes: (sigma, i, j) samples_i = samples[:, :, np.newaxis] samples_j = samples[:, np.newaxis, :] D_ij = D[np.newaxis, :, :] D_ji = DT[np.newaxis, :, :] return (samples_i*samples_j - samples_i*D_ji - samples_j*D_ij).sum(axis=0) / N

El siguiente es un ejemplo completo que compara el código de referencia (utilizando bucles for ) con la versión anterior; tenga en cuenta que eliminé la parte 1/N para mantener los cálculos en el dominio de los números enteros y, por lo tanto, hacer que la prueba de igualdad de matrices sea exacta.

 import time import numpy as np def timeit(func): def wrapper(*args): t_start = time.process_time() res = func(*args) t_total = time.process_time() - t_start print(f'{func.__name__}: {t_total:.3f} seconds') return res return wrapper rng = np.random.default_rng() M, N = 100, 200 samples = rng.integers(0, 100, size=(M, N)) D = rng.integers(0, 100, size=(N, N)) @timeit def reference(samples, D): results = np.zeros(shape=(N, N)) for sigma in range(M): for i in range(N): for j in range(N): results[i, j] += (samples[sigma, i]*samples[sigma, j] - samples[sigma, i]*D[j, i] - samples[sigma, j]*D[i, j]) return results @timeit def new(samples, D): # Ordering of axes: (sigma, i, j) samples_i = samples[:, :, np.newaxis] samples_j = samples[:, np.newaxis, :] D_ij = D[np.newaxis, :, :] D_ji = DT[np.newaxis, :, :] return (samples_i*samples_j - samples_i*D_ji - samples_j*D_ij).sum(axis=0) assert np.array_equal(reference(samples, D), new(samples, D))

Esto me da los siguientes resultados de referencia:

 reference: 6.465 seconds new: 0.133 seconds

Encontré más fácil dividir el problema en pasos más pequeños y trabajar en él, hasta que tengamos una sola ecuación.

Pasando de su formulación original:

 for sigma in range(M): for i in range(N): for j in range(N): results[i, j] += (1/N) * (samples[sigma, i]*samples[sigma, j] - samples[sigma, i]*D[j, i] - samples[sigma, j]*D[i, j])

Lo primero es eliminar el índice j en el bucle más interno. Para esto comenzamos a trabajar con vectores en lugar de elementos individuales:

 for sigma in range(M): for i in range(N): results[i, :] += (1/N) * (samples[sigma, i]*samples[sigma, :] - samples[sigma, i]*D[:, i] - samples[sigma, :]*D[i, :])

Luego, eliminamos el segundo bucle, el que tiene índice i . En este paso empezamos a pensar en matrices. Por lo tanto, cada ciclo es la suma directa de "matrices sigma".

 for sigma in range(M): results += (1/N) * (samples[sigma, :, np.newaxis] * samples[sigma] - samples[sigma, :, np.newaxis] * DT - samples[sigma, :] * D)

Recomiendo encarecidamente usar este paso como la solución, ya que vectorizar aún más requeriría demasiada memoria para un gran valor de M . Pero, solo para saber...

piense en las matrices como objetos tridimensionales. Hacemos los cálculos y sumamos al final en el índice cero como:

 results = (1/N) * (samples[:, :, np.newaxis] * samples[:,np.newaxis] - samples[:, :, np.newaxis] * DT - samples[:, np.newaxis, :] * D).sum(axis=0)