def update2d(N, spin, kT, E, M):    # 2D Ising model, N x N lattice
    i, j, flip = rnd.randint(0, N-1), rnd.randint(0, N-1), 0
    dE = 2*spin[i][j]*( spin[i-1][j] + spin[(i+1)%N][j]
                      + spin[i][j-1] + spin[i][(j+1)%N] )
    if (dE < 0.0): flip=1           # flip if dE<0, else flip 
    else:                           # according to exp(-dE/kT)
        p = np.exp(-dE/kT)
        if (rnd.random() < p): flip=1
    if (flip == 1):
        E = E + dE
        M = M - 2*spin[i][j]
        spin[i][j] = -spin[i][j]
    return E, M

#
# Program 11.6: Ising model in 1D (ising1d.py)
# J Wang, Computational modeling and visualization with Python
#

import random as rnd, numpy as np
import matplotlib.pyplot as plt

def initialize2d(N):  # set initial spins
    p, spin, E, M = 0.5, [[1]*N]*N, 0., 0.      # p = order para.
    for i in range(1, N):
        for j in range(1,N):
            if (rnd.random() < p): spin[i][j]=-1
            E = E - 2*( spin[i-1][j] + spin[(i+1)%N][j]
                      + spin[i][j-1] + spin[i][(j+1)%N] )
            M = M + spin[i][j]
    return spin, E - spin[N-1][0]*spin[0][N-1]*spin[1][0]*spin[0][1]*spin[0][0], M+spin[0][0]

###2d update function renamed to update
def update2d(N, spin, kT, E, M):    # 2D Ising model, N x N lattice
    i, j, flip = rnd.randint(0, N-1), rnd.randint(0, N-1), 0
    dE = 2*spin[i][j]*( spin[i-1][j] + spin[(i+1)%N][j]
                      + spin[i][j-1] + spin[i][(j+1)%N] )
    if (dE < 0.0): flip=1           # flip if dE<0, else flip 
    else:                           # according to exp(-dE/kT)
        p = np.exp(-dE/kT)
        if (rnd.random() < p): flip=1
    if (flip == 1):
        E = E + dE
        M = M - 2*spin[i][j]
        spin[i][j] = -spin[i][j]
    return E, M

N, passes = 32, 100
loop, Nmc = passes*N, passes*N
T, Eavg, Mavg = [], [], []
for i in range(1,41):               # temperature loop
    kT = 0.1*i                      # kT = reservoir temperature
    spin, E, M  = initialize2d(N)
    for k in range(loop):           # let it equilibrate
        E, M = update2d(N, spin, kT, E, M)
        
    E1, M1 = 0., 0.
    for k in range(Nmc):            # take averages
        E, M = update2d(N, spin, kT, E, M)
        E1, M1 = E1 + E, M1 + M
    E1, M1 = E1/Nmc, M1/Nmc
    T.append(kT), Eavg.append(E1/N**2), Mavg.append(M1/N**2)


plt.plot(T, Eavg, 'o', T, -np.tanh(1./np.array(T)))
plt.xlabel('$kT/\epsilon$'), plt.ylabel(r'$\langle E \rangle/N\epsilon$')
plt.figure()
plt.plot(T, Mavg,'o')
plt.xlabel('$kT/\epsilon$'), plt.ylabel(r'$\langle M \rangle/N$')
plt.show()