Further optimizing the ISING model

Question

I've implemented the 2D ISING model in Python, using NumPy and Numba's JIT:

from timeit import default_timer as timer
import matplotlib.pyplot as plt
import numba as nb
import numpy as np
# TODO for Dict optimization.
# from numba import types
# from numba.typed import Dict
@nb.njit(nogil=True)
def initialstate(N): 
 ''' 
 Generates a random spin configuration for initial condition
 '''
 state = np.empty((N,N),dtype=np.int8)
 for i in range(N):
 for j in range(N):
 state[i,j] = 2*np.random.randint(2)-1
 return state
@nb.njit(nogil=True)
def mcmove(lattice, beta, N):
 '''
 Monte Carlo move using Metropolis algorithm 
 '''
 
 # # TODO* Dict optimization 
 # dict_param = Dict.empty(
 # key_type=types.int64,
 # value_type=types.float64,
 # )
 # dict_param = {cost : np.exp(-cost*beta) for cost in [-8, -4, 0, 4, 8] }
 for _ in range(N):
 for __ in range(N):
 a = np.random.randint(0, N)
 b = np.random.randint(0, N)
 s = lattice[a, b]
 dE = lattice[(a+1)%N,b] + lattice[a,(b+1)%N] + lattice[(a-1)%N,b] + lattice[a,(b-1)%N]
 cost = 2*s*dE
 if cost < 0:
 s *= -1
 
 #TODO* elif np.random.rand() < dict_param[cost]:
 elif np.random.rand() < np.exp(-cost*beta):
 s *= -1
 lattice[a, b] = s
 return lattice
@nb.njit(nogil=True)
def calcEnergy(lattice, N):
 '''
 Energy of a given configuration
 '''
 energy = 0 
 
 for i in range(len(lattice)):
 for j in range(len(lattice)):
 S = lattice[i,j]
 nb = lattice[(i+1)%N, j] + lattice[i,(j+1)%N] + lattice[(i-1)%N, j] + lattice[i,(j-1)%N]
 energy += -nb*S
 return energy/2
@nb.njit(nogil=True)
def calcMag(lattice):
 '''
 Magnetization of a given configuration
 '''
 mag = np.sum(lattice, dtype=np.int32)
 return mag
@nb.njit(nogil=True)
def ISING_model(nT, N, burnin, mcSteps):
 """ 
 nT : Number of temperature points.
 N : Size of the lattice, N x N.
 burnin : Number of MC sweeps for equilibration (Burn-in).
 mcSteps : Number of MC sweeps for calculation.
 """
 T = np.linspace(1.2, 3.8, nT); 
 E,M,C,X = np.zeros(nT), np.zeros(nT), np.zeros(nT), np.zeros(nT)
 n1, n2 = 1.0/(mcSteps*N*N), 1.0/(mcSteps*mcSteps*N*N) 
 for temperature in range(nT):
 lattice = initialstate(N) # initialise
 E1 = M1 = E2 = M2 = 0
 iT = 1/T[temperature]
 iT2= iT*iT
 
 for _ in range(burnin): # equilibrate
 mcmove(lattice, iT, N) # Monte Carlo moves
 for _ in range(mcSteps):
 mcmove(lattice, iT, N) 
 Ene = calcEnergy(lattice, N) # calculate the Energy
 Mag = calcMag(lattice,) # calculate the Magnetisation
 E1 += Ene
 M1 += Mag
 M2 += Mag*Mag 
 E2 += Ene*Ene
 E[temperature] = n1*E1
 M[temperature] = n1*M1
 C[temperature] = (n1*E2 - n2*E1*E1)*iT2
 X[temperature] = (n1*M2 - n2*M1*M1)*iT
 return T,E,M,C,X
def main():
 
 N = 32
 start_time = timer()
 T,E,M,C,X = ISING_model(nT = 64, N = N, burnin = 8 * 10**4, mcSteps = 8 * 10**4)
 end_time = timer()
 print("Elapsed time: %g seconds" % (end_time - start_time))
 f = plt.figure(figsize=(18, 10)); # 
 # figure title
 f.suptitle(f"Ising Model: 2D Lattice\nSize: {N}x{N}", fontsize=20)
 _ = f.add_subplot(2, 2, 1 )
 plt.plot(T, E, '-o', color='Blue') 
 plt.xlabel("Temperature (T)", fontsize=20)
 plt.ylabel("Energy ", fontsize=20)
 plt.axis('tight')
 _ = f.add_subplot(2, 2, 2 )
 plt.plot(T, abs(M), '-o', color='Red')
 plt.xlabel("Temperature (T)", fontsize=20)
 plt.ylabel("Magnetization ", fontsize=20)
 plt.axis('tight')
 _ = f.add_subplot(2, 2, 3 )
 plt.plot(T, C, '-o', color='Green')
 plt.xlabel("Temperature (T)", fontsize=20)
 plt.ylabel("Specific Heat ", fontsize=20)
 plt.axis('tight')
 _ = f.add_subplot(2, 2, 4 )
 plt.plot(T, X, '-o', color='Black')
 plt.xlabel("Temperature (T)", fontsize=20)
 plt.ylabel("Susceptibility", fontsize=20)
 plt.axis('tight')
 plt.show()
if __name__ == '__main__':
 main()

Which of course, works:

enter image description here

I have two main questions:

1. Is there anything left to optimize? I knew ISING model is hard to simulate, but looking at the following table, it seems like I'm missing something...

 lattice size : 32x32 
 burnin = 8 * 10**4
 mcSteps = 8 * 10**4
 Simulation time = 365.98 seconds
 lattice size : 64x64 
 burnin = 10**5
 mcSteps = 10**5
 Simulation time = 1869.58 seconds

2. I tried implementing another optimization based on not calculating the exponential over and over again using a dictionary, yet on my tests, it seems like its slower. What am I doing wrong?

• Reinderien: There's no saving this. You need to throw away the works and replace it with something better like C/C++. Numpy will not ever perform well on iterative algorithms no matter how much JIT you attempt.
• Ali_Sh: using np.random in jitted numba may cause it inefficient.

Answer 1

Rename functions like initialstate to initial_state by PEP8.

Add PEP484 type hints to your function signatures.

Consider adding a stateful RNG (default_rng) instance so that you can choose initial conditions and have them repeatable. This is important for scientific work.

Replace this loop:

state = np.empty((N,N),dtype=np.int8)
for i in range(N):
 for j in range(N):
 state[i,j] = 2*np.random.randint(2)-1
return state

with a vectorised version:

return np.random.randint(2, size=(N, N), dtype=np.int8)*2 - 1

Replace the nested loop having a __ with a single loop over N*N elements.

Reformat your dE assignment to be over multiple lines for legibility.

Replace your s *= -1 multiplication with the assignment lattice[a, b] = -s.

Combine the cost comparison ifs into one using an or, since their effect is the same.

Vectorise calcEnergy to a sum of a product of lattice with four rolled versions of itself.

n2 is just n1/mc_steps.

Your _ = is non-effectual so delete it.

mc_move mutates lattice, so there's no point in returning it.

Suggested

All of this is window dressing and won't in any significant way improve your performance; for that you need to move to a better language.

from timeit import default_timer as timer
import matplotlib.pyplot as plt
import numpy as np
def initial_state(N: int) -> np.ndarray:
 """
 Generates a random spin configuration for initial condition
 """
 return np.random.randint(2, size=(N, N), dtype=np.int8)*2 - 1
def mc_move(lattice: np.ndarray, beta: float, N: int) -> None:
 """
 Monte Carlo move using Metropolis algorithm
 """
 for a, b in np.random.randint(0, N, size=(N*N, 2)):
 s = lattice[a, b]
 dE = (
 lattice[(a + 1) % N, b] +
 lattice[(a - 1) % N, b] +
 lattice[a, (b + 1) % N] +
 lattice[a, (b - 1) % N]
 )
 cost = 2*s*dE
 if cost < 0 or np.random.rand() < np.exp(-cost * beta):
 lattice[a, b] = -s
def calc_energy(lattice: np.ndarray) -> int:
 """
 Energy of a given configuration
 """
 return (
 -lattice * (
 np.roll(lattice, axis=0, shift=+1) +
 np.roll(lattice, axis=0, shift=-1) +
 np.roll(lattice, axis=1, shift=+1) +
 np.roll(lattice, axis=1, shift=-1)
 )
 ).sum() / 2
def calc_mag(lattice: np.ndarray) -> int:
 """
 Magnetization of a given configuration
 """
 return np.sum(lattice, dtype=np.int32)
def ising_model(nT: int, N: int, burnin: int, mc_steps: int) -> tuple[
 np.ndarray, # T
 np.ndarray, # E
 np.ndarray, # M
 np.ndarray, # C
 np.ndarray, # X
]:
 """
 nT : Number of temperature points.
 N : Size of the lattice, N x N.
 burnin : Number of MC sweeps for equilibration (Burn-in).
 mcSteps : Number of MC sweeps for calculation.
 """
 T = np.linspace(1.2, 3.8, nT)
 E, M, C, X = np.zeros((4, nT))
 n1 = 1/mc_steps/N/N
 n2 = n1/mc_steps
 for temperature in range(nT):
 lattice = initial_state(N) # initialise
 E1 = M1 = E2 = M2 = 0
 iT = 1 / T[temperature]
 iT2 = iT * iT
 for _ in range(burnin): # equilibrate
 mc_move(lattice, iT, N) # Monte Carlo moves
 for _ in range(mc_steps):
 mc_move(lattice, iT, N)
 Ene = calc_energy(lattice) # calculate the Energy
 Mag = calc_mag(lattice) # calculate the Magnetisation
 E1 += Ene
 M1 += Mag
 E2 += Ene * Ene
 M2 += Mag * Mag
 E[temperature] = n1*E1
 M[temperature] = n1*M1
 C[temperature] = (n1*E2 - n2*E1*E1)*iT2
 X[temperature] = (n1*M2 - n2*M1*M1)*iT
 return T, E, M, C, X
def main() -> None:
 N = 32
 start_time = timer()
 T, E, M, C, X = ising_model(nT=8, N=N, burnin=8, mc_steps=32)
 end_time = timer()
 print(f'Elapsed time: {end_time - start_time:g} seconds')
 f = plt.figure()
 # figure title
 f.suptitle(f'Ising Model: 2D Lattice\nSize: {N}x{N}')
 f.add_subplot(2, 2, 1)
 plt.plot(T, E, '-o', color='Blue')
 plt.xlabel('Temperature (T)')
 plt.ylabel('Energy ')
 plt.axis('tight')
 f.add_subplot(2, 2, 2)
 plt.plot(T, abs(M), '-o', color='Red')
 plt.xlabel('Temperature (T)')
 plt.ylabel('Magnetization ')
 plt.axis('tight')
 f.add_subplot(2, 2, 3)
 plt.plot(T, C, '-o', color='Green')
 plt.xlabel('Temperature (T)')
 plt.ylabel('Specific Heat ')
 plt.axis('tight')
 f.add_subplot(2, 2, 4)
 plt.plot(T, X, '-o', color='Black')
 plt.xlabel('Temperature (T)')
 plt.ylabel('Susceptibility')
 plt.axis('tight')
 plt.show()
if __name__ == '__main__':
 main()