how to use jit on a numerical simulation

user_17750586 0
Viewed 58

I'm trying to write a physical simulation using python which requires a lot of simple math computation. I believed that numba's jit will help me run it faster, the code is mainly calculations and I'm not sure what's wrong. I tried putting @jit over each function I tried to specify @jit(nopython=True) and still its very slow.

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)
physics jit python
2 Answers

Looks like JIT is being used, the slowness is because compilation time is included into benchmarked time in your code.

If attempt to pre-compile the functions first and only then do benchmarking, it'll be much faster:

dummy_grid = np.random.choice((-1, 1), (2, 2))
energy_i_now(0, 0, dummy_grid) #poor man's pre-compilation, call each function at least once in advance
E_00(0, 0, dummy_grid)
iterate(dummy_grid)

s = time()
for i in range(10):
    iterate(GRID)
print(time() - s)

Also as per Numba doc, there is cache compilation option:

To avoid compilation times each time you invoke a Python program, you can instruct Numba to write the result of function compilation into a file-based cache. This is done by passing cache=True: @jit(cache=True)

user_10752820 
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

new = np.random.choice((-1, 1), (2, 2))
energy_i_now(0, 0, new)
E_00(0, 0, new)
iterate(new)


s = time()
for i in range(10):
    iterate(GRID)
print(time() - s)
user_171383220
Related