1
\$\begingroup\$

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
  1. 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?
\$\endgroup\$
2
  • \$\begingroup\$ 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. \$\endgroup\$
    – Reinderien
    Dec 4, 2022 at 1:04
  • \$\begingroup\$ using np.random in jitted numba may cause it inefficient. \$\endgroup\$
    – Ali_Sh
    Jan 1 at 8:41

1 Answer 1

1
\$\begingroup\$

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()
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.