# Initializing some things
import numpy as np
from scipy.optimize import curve_fit

import matplotlib.pyplot as plt
import matplotlib as mpl
import matplotlib.animation as animation
%matplotlib inline

mpl.rcParams['lines.linewidth'] = 1.5
mpl.rcParams['lines.marker'] = "o"
mpl.rcParams['axes.grid'] = True
mpl.rcParams['legend.labelspacing'] = 0.0
cmap = mpl.cm.get_cmap(name='YlOrRd')
colors3 = [cmap(x) for x in [0.4, 0.6, 0.9]]

ax = plt.subplot(xlabel="Energy", ylabel="Probability", title="Boltzmann distribution")
E = np.linspace(0, 4.5, num=50)

for i, T in enumerate([0.8, 1.5, 3.5]):
    ax.plot(E, 1.0/T*np.exp(-E/T), "-", label="T={}".format(T), color=colors3[i])
ax.legend();

# !magneto -N1=20000 -N2=0 -TMin=1.3 -TMax=3.3 -L=62 -TSteps=3 -states=states

def draw_state(ax, grid, title):
    ax.grid(False)
    ax.imshow(grid, cmap=plt.cm.Greys, interpolation="none")
    ax.set_title(title)

def load_states(filename):
    with open(filename) as f:
        T = float(f.readline())
    grid = np.loadtxt(filename, delimiter=",", skiprows=1)
    return T, grid

def draw_states(filename_base, n=3):
    fig = plt.figure(figsize=(3*n, n))
    for i in range(n):
        T, grid = load_states(filename="{}{}.txt".format(filename_base, i))
        ax = fig.add_subplot(1, n, i+1)
        draw_state(ax, grid, "T={:.2f}".format(T))
    fig.tight_layout()

draw_states("states", 3)

# !magneto -L=30 -TSteps=15 -en=energy -dist=normal -alg=sw -N1=10 -N2=100
Tc = 2.0/np.log(1+np.sqrt(2))
ax = plt.subplot(xlabel="T/T_c", ylabel="Energy", ylim=(-2,0), xlim=(0,2))
T, E = np.loadtxt("energy.txt", delimiter=", ", unpack=True)
ax.plot(T/Tc, E);

# !magneto -L=30 -TSteps=12 -cv=cv -dist=normal
ax = plt.subplot(xlabel="T/T_c", ylabel="Heat capacity")

# variance
T, cv = np.loadtxt("cv.txt", delimiter=", ", unpack=True)
ax.plot(T/Tc, cv, label="by variance")

# derivation
T, E = np.loadtxt("energy.txt", delimiter=", ", unpack=True)
ax.plot(T[:-1]/Tc, np.diff(E)/np.diff(T), label="by derivation")

ax.legend();

# !magneto -L=30 -TSteps=12 -N2=200 -mag=mag -dist=normal
ax = plt.subplot(xlabel="T/T_C", ylabel="Magnetization |m|", xlim=(0,2), ylim=(0,1))
T, m = np.loadtxt("mag.txt",  delimiter=", ", unpack=True)
ax.plot(T/Tc, m);

# !magneto -L=60 -TSteps=9 -TMin=2.0 -TMax=2.6 -N2=100 -N3=1000 -corr=chi_corrfun -alg=metro -N1=5000

def corr_fun(x, corr_len, a):
    return np.exp(-x/corr_len)*a

def ln_corr_fun(x, corr_len, a):
    return -x/corr_len + np.log(a)

def lower_range(data, threshold):
    end_index = np.argmax(data<threshold)
    if end_index == 0:
        end_index = len(data)
    return data[:end_index]

ax = plt.subplot(xlabel="Distance d", ylabel="Correlation function G(d)")
data = np.loadtxt("chi_corrfun.txt",  delimiter=", ")

for i, T_index in enumerate([0,4,-1]):
    [T], corr_fun_data = np.split(data[T_index], [1])    
    corr_fun_data = lower_range(corr_fun_data, threshold = 0.01)
    
    distance = np.arange(corr_fun_data.size)
    corr_fun_data = np.log(corr_fun_data)
    [xi, a], _ = curve_fit(ln_corr_fun, distance, corr_fun_data)
    ax.semilogy(distance, np.exp(corr_fun_data), "o", color=colors3[i], label="T={:.2f}, $\\xi$={:.1f}".format(T, xi))
    ax.semilogy(distance, np.exp(ln_corr_fun(distance, xi, a)), "-", color=colors3[i])

plt.legend(loc="lower right");

# !magneto -N1=100 -N2=0 -TMin=2.0 -TMax=2.6 -L=60 -TSteps=3 -states=states_corr -alg=sw
draw_states("states_corr", n=3)

# !magneto -L=60 -TMin=2.0 -TMax=2.6 -TSteps=9 -N2=100 -N3=10 -chi=chi -dist=normal -N1=10 -alg=sw

ax = plt.subplot(xlabel="T/T_c", ylabel="Magnetic susceptibility $\chi$")
T, chi = np.loadtxt("chi.txt",  delimiter=", ", unpack=True)
ax.plot(T/Tc, chi, label="by variation")

data = np.loadtxt("chi_corrfun.txt",  delimiter=", ")
T = data[:,0]
chi_corrfun = []
for i in range(9):
    _, corr_fun_data = np.split(data[i], [1])
    corr_fun_data = lower_range(corr_fun_data, threshold = 0.01)
    distance = np.arange(corr_fun_data.size)
    [xi, a], _ = curve_fit(ln_corr_fun, distance, corr_fun_data)
    chi_corrfun.append(xi/T[i])

ax.plot(np.asarray(T/Tc), np.asarray(chi_corrfun), "red", label="by correlation")
plt.legend();

# !magneto -L=60 -TSteps=15 -N2=200 -N3=10 -chi=chi_exp -TMin=2.7 -TMax=3.5 -dist=normal -alg=sw -N1=10
ax = plt.subplot(111, xlabel="Reduced temperature $\\tau$", ylabel="Magnetic susceptibility $\chi$")

# Fit
def chi_fit(tau, gamma, a):
    return a*np.power(tau,-gamma)
T, chi = np.loadtxt("chi_exp.txt", delimiter=", ", unpack=True)
tau = (T-Tc)/Tc
popt, pcov = curve_fit(chi_fit, tau, chi)

# Plot
ax.plot(tau, chi, "o", label="data")
ax.plot(tau, chi_fit(tau, *popt), "--", label="fit")
ax.legend(loc="upper right")

# Result
gamma = popt[0]; gamma0 = 1.75
print("gamma: {:.2f}, deviation: {:.1f}%".format(gamma, 100.0*(gamma-gamma0)/gamma0))

# !magneto -L=8  -TSteps=12 -TMin=2.0 -TMax=3.5 -alg=sw -N2=500 -N3=30 -chi=chi_finite_8
# !magneto -L=10 -TSteps=12 -TMin=1.8 -TMax=2.8 -alg=sw -N2=400 -N3=30 -chi=chi_finite_10
# !magneto -L=12 -TSteps=15 -TMin=2.0 -TMax=3.5 -alg=sw -N2=300 -N3=25 -chi=chi_finite_12
# !magneto -L=16 -TSteps=15 -TMin=2.0 -TMax=3.5 -alg=sw -N2=200 -N3=20 -chi=chi_finite_16
# !magneto -L=32 -TSteps=15 -TMin=1.8 -TMax=3.0 -alg=sw -N2=200 -N3=20 -chi=chi_finite_32

from scipy.stats import norm

def chi_fun(x, x0, sigma, a, b):
    return norm.pdf(x, x0, sigma)*b + x*a

L = 8
T, cv = np.loadtxt("chi_finite_{}.txt".format(L),  delimiter=", ", unpack=True)
popt, _ = curve_fit(chi_fun, T, cv, p0=[ 2.45, 0.1,  0.8,  0.6])
plt.title("Example of fit for L={}. Resulting Tc={:.2f}".format(L, popt[0]))
plt.plot(T, cv, "o")
plt.plot(T, chi_fun(T, *popt), "-");

# Plot critical temps vs inverse grid sizes
grid_sizes = np.array([8,10,12,16,32])
L_inv = 1.0/grid_sizes

T_critical_array = []
for L in grid_sizes:
    T, cv = np.loadtxt("chi_finite_{}.txt".format(L),  delimiter=", ", unpack=True)
    popt, _ = curve_fit(chi_fun, T, cv, p0=[ 2.45, 0.1,  0.8,  0.6])
    Tc0 = popt[0]
    T_critical_array.append(Tc0)

ax = plt.subplot(xlabel="1/L", ylabel="Critical temperature", title="$T_c$ for different L")
ax.plot(L_inv, T_critical_array, "o")
ax.set_xticks(np.arange(0, 0.16, 0.05))

# Fit with a linear function
def T_fit(x,a,b):
    return a+b*x
popt, pcov = curve_fit(T_fit, L_inv, T_critical_array)
x = np.linspace(0, 0.14)
ax.plot(x, T_fit(x, *popt), "--")

# Result
gamma = popt[0]; gamma0 = 1.75
T_c_fit = T_fit(0, *popt)
print("Calculated Tc: {:.3f}, deviation: {:.2f}%".format(T_c_fit, 100.0*(T_c_fit-Tc)/Tc))
ax.axhline(y=Tc, marker="");

ax = plt.subplot(xlabel="$\Delta E$", ylabel="Probability", ylim=(0,1.1), title="Metropolis acceptance probability")
dE = np.linspace(-2, 2)

for i, T in enumerate([0.8, 1.5, 3.5]):
    ax.plot(dE, np.minimum(1, np.exp(-dE/T)), "-", label="T={}".format(T), color=colors3[i])
ax.legend();

L=80; a=L//2; r=30
y,x = np.ogrid[-a:L-a, -a:L-a]
mask = x*x + y*y <= r*r
array = np.zeros((L,L), dtype=int)
array[mask] = 1

plt.grid(False)
plt.imshow(array, cmap=plt.cm.Greys, interpolation="none")
np.savetxt("circle.txt", array, fmt='%1u', delimiter=",")

# !magneto -N1=0 -N3=1 -N2=100 -TMin=2.0 -L=80 -TSteps=1 -states=slowdown_metro -alg=metro -initial=circle.txt -record=main
def writeIsingMovie(file_in, file_out):
    fig = plt.figure(figsize=(1,1))
    data = np.loadtxt(file_in,  delimiter=",", skiprows=1)
    L = data.shape[1]
    grids = np.split(data, data.shape[0]/L)
    images = []
    for grid in grids:
        images.append([plt.imshow(grid, cmap=plt.cm.Greys, interpolation ="none")])
        plt.axis('off')
        plt.subplots_adjust(left=0.0, right=1.0, top=1.0, bottom=0.0)
    ani = animation.ArtistAnimation(fig, images, blit=True)
    ani.save(file_out, dpi=L*2, fps=25, extra_args=['-vcodec', 'libvpx', "-quality", "good", "-b:v", "350k"])

writeIsingMovie("slowdown_metro0.txt", "slowdown.webm")

%%html
<video controls loop autoplay>
  <source src="slowdown.webm">
</video>

# !magneto -N1=0 -N3=1 -N2=5 -TMin=2.3 -L=80 -TSteps=1 -states=slowdown_sw -initial=circle.txt -record=main -alg=sw
fig = plt.figure(figsize=(9,3))
data = np.loadtxt("slowdown_sw0.txt", delimiter=",", skiprows=1)
L = data.shape[1]
grids = np.split(data, data.shape[0]/L)
for i in range(3):
    ax = fig.add_subplot(1, 3, i+1)
    ax.grid(False)
    ax.imshow(grids[i], cmap=plt.cm.Greys, interpolation="none")
fig.tight_layout()