Cómo utilizar el JIT en una simulación numérica.
Estoy tratando de escribir una simulación física utilizando python que requiere muchos cálculos matemáticos simples. Creí que el jit de numba me ayudaría a ejecutarlo más rápido, el código es principalmente cálculos y no estoy seguro de qué está mal. Intenté poner @jit sobre cada función, también intenté especificar @jit(nopython=True) y aún así es muy lento.
import numpy as np
from time import time
from numba import jit
GRID = np.random.choice((-1, 1), (32, 32))
J = 1
UB = 1
BETA = 1
@jit
def energy_i_now(i, j, grid):
E = -(UB * grid[i, j] +
J * grid[i, j] * (grid[(i + 1) % 32, j] + grid[(i - 1) % 32, j]
+ grid[i, (j + 1) % 32] + grid[i, (j - 1) % 32]))
return E
@jit
def E_00(i, j, grid):
for i in range(32):
for j in range(32):
E = -0.5 * (UB * grid[i, j] +
J * grid[i, j] * (grid[(i + 1) % 32, j] + grid[(i - 1) % 32, j]
+ grid[i, (j + 1) % 32] + grid[i, (j - 1) % 32]))
return E - energy_i_now(i, j, grid)
@jit
def iterate(grid):
E_tot = E_00(0, 0, grid) + energy_i_now(0, 0, grid)
for i in range(32):
for j in range(32):
e_i_now = energy_i_now(i, j, grid)
p_flip = 1 / (1 + np.exp(-2 * BETA * e_i_now))
rand = np.random.random()
if rand <= p_flip:
grid[i, j] = grid[i, j] * -1
E_tot -= 2 * e_i_now
return grid
s = time()
for i in range(10):
iterate(GRID)
print(time() - s)
davy.ai
Hay algunas cosas que podrían estar causando que el código siga ejecutándose lento a pesar de usar el jit de numba.
En primer lugar, el decorador @jit por sí solo puede no ser suficiente para acelerar el código. También es importante especificar el modo de jit agregando el argumento ‘nopython=True’. En el código proporcionado, esto ya se ha hecho, pero vale la pena verificar que se esté utilizando correctamente.
Otra sugerencia sería considerar el uso de operaciones vectorizadas de numpy cuando sea posible. Por ejemplo, en lugar de usar bucles anidados en la función E_00, podemos escribirla de la siguiente manera:
@jit(nopython=True)
def E_00(i, j, grid):
E = -0.5 * (UB * grid +
J * grid * (np.roll(grid, 1, axis=0) + np.roll(grid, -1, axis=0)
+ np.roll(grid, 1, axis=1) + np.roll(grid, -1, axis=1)))
return np.sum(E) – energy_i_now(i, j, grid)
Esto aprovecha la capacidad de numpy para aplicar operaciones a matrices completas en lugar de recorrerlas elemento por elemento.
Por último, es posible que el cuello de botella no esté en el código en sí, sino en el tamaño de la cuadrícula o en la cantidad de iteraciones que se están ejecutando. En ese caso, puede ser necesario realizar una optimización adicional, como dividir las simulaciones en fragmentos o ejecutarlas en paralelo.