#!/usr/bin/python3
#
# Markov Chain Monte Carlo for 2d Ising model

import numpy as np

############################
# Set up various constants
############################

burnin_steps = 10000   # steps for burn-in procedure
mcmc_steps = 100000    # length of Markov chain 
L = 16                 # for an L x L lattice of spins
coupling = 1.0         # coupling constant J
filename = "ising.dat" # save data here

norm = 1 / L**2        # normalization constant
lat = np.ones((L, L), dtype=int)  # the spin lattice

min_T = 0.1            # minimal temperature for scan
max_T = 3              # maximal temperature
T_steps = 30           # number of intermediate points

###########################
# Some useful functions
###########################

# Flip spin at site (i, j)    
def flip(i, j):
        lat[i, j] *= -1

# Flip spin at site (i, j) if energy gained or due to thermal fluctuation.
# Return energy difference
def flip_maybe(T, i, j):
    DeltaE = deltaE(i, j)
    if DeltaE < 0:         # energy gained
        flip(i, j)
        return DeltaE
    elif np.random.random() < np.exp(- DeltaE / T): # en. lost, flip anyway
        flip(i, j)
        return DeltaE
    else:                  # don't flip, keep old configuration
        return 0.0
    
# T -> infinity: all sites random    
def heat():
    for i in range(L):
        for j in range(L):
            if np.random.choice([True, False]):
                flip(i,j)
    
# T -> 0: all sites equal
def freeze():
    spin = np.random.choice([-1, 1])
    for i in range(L):
        for j in range(L):
            lat[i, j] = spin
    
# total energy of entire lattice               
def energy():
    E = 0.0
    for i in range(L):
        next_i = (i + 1) % L     # 0 if i == L-1, periodic BC
        for j in range(L):    
            next_j = (j + 1) % L # periodic BC
            E -= coupling * lat[i, j] * lat[next_i, j]
            E -= coupling * lat[i, j] * lat[i, next_j]
    return E
    
# energy change that would result if spin (i, j) was flipped
def deltaE(i, j):
    prev_i, next_i = (i - 1) % L, (i + 1) % L
    prev_j, next_j = (j - 1) % L, (j + 1) % L
    DeltaE = 2 * coupling * lat[i, j] *      \
                (lat[prev_i, j] + lat[next_i, j]
                    + lat[i, prev_j] + lat[i, next_j])
    return DeltaE
    
# generate random index    
def random_index():
    return np.random.randint(L)
        
# initial phase: let the chain find thermal equilibrium  
def burnin(T, burnin_steps):
    for n in range(burnin_steps):
        i, j = random_index(), random_index()
        flip_maybe(T, i, j)
    
# total magnetisation        
def magnetisation():
    return lat.sum()

################################
# Main program starts here:
################################

# initialize arrays to store <E>, <M>, <E²>, <M²>
Epoints = np.zeros(T_steps)
Mpoints = np.zeros(T_steps)
E2points = np.zeros(T_steps)
M2points = np.zeros(T_steps)

# initialize array of temperatures to scan over
T = min_T
dT = (max_T - min_T) / T_steps
Tpoints = np.arange(min_T, max_T, dT)

# temperature loop
for step in range(T_steps):
    freeze()                 # start in the ground state
    burnin(T, burnin_steps)  # burn-in without recording the state
    
    en = energy()            # energy after equilibrium has been reached
    magn = magnetisation()   # magnetisation after equilibrium has been reached
    E = M = E2 = M2 = 0.0    # current value of <E(T)>, <M(T)>, <E²(T)>, <M²(T)>
    
    for n in range(mcmc_steps):  # MCMC loop
        i, j = random_index(), random_index() # choose random spin
        oldstate = lat[i, j]                  # save its value
        en += flip_maybe(T, i, j)    # flip it or not according to Metropolis
        E += en                      # update mean energy
        E2 += en**2                  # update mean E²
        magn += lat[i, j] - oldstate # recalculate magnetization
        M += magn                    # update mean M
        M2 += magn**2                # update mean M²
        
    E *= norm      # once the loop has finished: normalize <E>, <M>, <M²>, <E²>
    M *= norm      
    M2 *= norm**2  
    E2 *= norm**2  
    E /= mcmc_steps
    M /= mcmc_steps
    E2 /= mcmc_steps
    M2 /= mcmc_steps
    
    Epoints[step] = E    # add <E>, <M>, <M²>, <E²> to respective list
    E2points[step] = E2
    Mpoints[step] = M
    M2points[step] = M2
    
    T += dT  # go to next temperature

np.savetxt(filename, 
            np.transpose([Epoints, E2points, Mpoints, M2points, Tpoints]), 
            header = "# Data file generated from ising.py")