Mejora el rendimiento de un bucle for en Python (posiblemente con numpy o numba)

Quiero mejorar el rendimiento del bucle for en esta función.

 import numpy as np import random def play_game(row, n=1000000): """Play the game! This game is a kind of random walk. Arguments: row (int[]): row index to use in the p matrix for each step in the walk. Then length of this array is the same as n. n (int): number of steps in the random walk """ p = np.array([[ 0.499, 0.499, 0.499], [ 0.099, 0.749, 0.749]]) X0 = 100 Y0 = X0 % 3 X = np.zeros(n) tempX = X0 Y = Y0 for j in range(n): tempX = X[j] = tempX + 2 * (random.random() < p.item(row.item(j), Y)) - 1 Y = tempX % 3 return np.r_[X0, X] 

La dificultad radica en el hecho de que el valor de Y se calcula en cada paso en función del valor de X y que Y se usa en el siguiente paso para actualizar el valor de X

Me pregunto si hay algún truco numpy que podría hacer una gran diferencia. Usar Numba es un juego justo (lo probé pero sin mucho éxito). Sin embargo, no quiero usar Cython.

Una observación rápida nos dice que hay dependencia de datos entre iteraciones en el código de función. Ahora, hay diferentes tipos de dependencias de datos. El tipo de dependencia de datos que está viendo es la dependencia de indexación que es la selección de datos en cualquier iteración depende de los cálculos de iteración anteriores. Esta dependencia parecía difícil de rastrear entre iteraciones, por lo que esta publicación no es realmente una solución vectorizada. Más bien, trataríamos de pre-calcular los valores que se usarían dentro del bucle, tanto como sea posible. La idea básica es hacer el trabajo mínimo dentro del bucle.

Aquí hay una breve explicación de cómo podemos proceder con los cálculos previos y así tener una solución más eficiente:

  • Dada la forma relativamente pequeña de p de la cual se extraen los elementos de la fila en función de la row entrada, puede preseleccionar todas esas filas de p con p[row] .

  • Para cada iteración, estás calculando un número aleatorio. Puede reemplazarlo con una matriz aleatoria que puede configurar antes del bucle y, por lo tanto, también habría calculado previamente esos valores aleatorios.

  • Según los valores precalculados hasta el momento, tendría los índices de columna para todas las filas en p . Tenga en cuenta que estos índices de columna serían un ndarray grande que contiene todos los índices de columna posibles y dentro de nuestro código, solo se seleccionaría uno basado en los cálculos de iteración. Usando los índices de la columna por iteración, usted incrementaría o disminuiría X0 para obtener la salida por iteración.

La implementación se vería así:

 randarr = np.random.rand(n) p = np.array([[ 0.499, 0.419, 0.639], [ 0.099, 0.749, 0.319]]) def play_game_partvect(row,n,randarr,p): X0 = 100 Y0 = X0 % 3 signvals = 2*(randarr[:,None] < p[row]) - 1 col_idx = (signvals + np.arange(3)) % 3 Y = Y0 currval = X0 out = np.empty(n+1) out[0] = X0 for j in range(n): currval = currval + signvals[j,Y] out[j+1] = currval Y = col_idx[j,Y] return out 

Para la verificación con respecto al código original, tendría el código original modificado así:

 def play_game(row,n,randarr,p): X0 = 100 Y0 = X0 % 3 X = np.zeros(n) tempX = X0 Y = Y0 for j in range(n): tempX = X[j] = tempX + 2 * (randarr[j] < p.item(row.item(j), Y)) - 1 Y = tempX % 3 return np.r_[X0, X] 

Tenga en cuenta que dado que este código precomputa esos valores aleatorios, esto ya le daría una buena aceleración sobre el código en la pregunta.

Pruebas de tiempo de ejecución y verificación de salida -

 In [2]: # Inputs ...: n = 1000 ...: row = np.random.randint(0,2,(n)) ...: randarr = np.random.rand(n) ...: p = np.array([[ 0.499, 0.419, 0.639], ...: [ 0.099, 0.749, 0.319]]) ...: In [3]: np.allclose(play_game_partvect(row,n,randarr,p),play_game(row,n,randarr,p)) Out[3]: True In [4]: %timeit play_game(row,n,randarr,p) 100 loops, best of 3: 11.6 ms per loop In [5]: %timeit play_game_partvect(row,n,randarr,p) 1000 loops, best of 3: 1.51 ms per loop In [6]: # Inputs ...: n = 10000 ...: row = np.random.randint(0,2,(n)) ...: randarr = np.random.rand(n) ...: p = np.array([[ 0.499, 0.419, 0.639], ...: [ 0.099, 0.749, 0.319]]) ...: In [7]: np.allclose(play_game_partvect(row,n,randarr,p),play_game(row,n,randarr,p)) Out[7]: True In [8]: %timeit play_game(row,n,randarr,p) 10 loops, best of 3: 116 ms per loop In [9]: %timeit play_game_partvect(row,n,randarr,p) 100 loops, best of 3: 14.8 ms per loop 

Por lo tanto, estamos viendo una aceleración de aproximadamente 7.5x+ , ¡no está mal!