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:

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?

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()