# Example of hard disk model script : transfer-matrix and tensor renormalization group methods¶

import numpy as np
from math import sqrt, log
from susmost import make_single_point_zmatrix, make_triangular_Oxy_plane_cell_neg, \
import os
from scipy.misc import derivative
from scipy.optimize import minimize_scalar
from time import time

#constants
INF_E = 1E6                       # prohibitive energy level
k_B = 1.0                                 # 1/(kT)=1
T = 1.0                                   # 1/(kT)=1
a = 1.0                                   # lattice parameter
r_cut_off = 8.0 + .0001

#parameters
cell,atoms = make_triangular_Oxy_plane_cell_neg(a)

#warious types of NN models

#h  = a * 0.5   * 1.001/2.0 # 0NN radius
h  = a * 1.     * 1.001/2.0 # 1NN radius
#h = a * 3.**.5 * 1.001/2.0 # 2NN radius
#h = a * 2.     * 1.001/2.0 # 3NN radius
#h = a * 7.**.5 * 1.001/2.0 # 4NN radius
#h = a * 3.     * 1.001/2.0 # 5NN radius
join_x = 1
join_y = 1

#interaction function
interaction = lambda cc1,cc2: nearest_int(cc1, cc2, h=h)

mono_state = SiteStateType('mono', make_single_point_zmatrix(), -1, ['N'], coverage=1.)
empty_state = SiteStateType('Empty',make_single_point_zmatrix(), 0.0, ['H'],coverage=0.)
lt = normal_lattice_task([mono_state, empty_state], cell, atoms, interaction, r_cut_off, INF_E)

#Build W tensor
def build_W(mu, T, force=True):
beta = 1. / (k_B * T)
W = make_tensor(lt, beta, force)
return W

#Grand potential calculation by transfer-matrix method
def GrandPotential_TM(mu,T,M):
W = build_W(mu, T, True)
beta = 1. / (k_B * T)
tm_sol = transferm.solve_TM(W, M)
GP = -log(tm_sol.lr)/beta/M
print ('TM', "M", M, 'mu', mu, 'GP', GP)
return GP/(join_x*join_y)

#Grand potential calculation by TRG
def GrandPotential_TRG(mu, T, D):
W = build_W(mu, T, True)
lnZ = solve_TRG(W,contractions_count = 512, max_count = D,reduce_factor = 0)
GP = -k_B*T*lnZ
print ('TRG', "D", D, 'mu', mu, 'GP', GP)
return GP/(join_x*join_y)

#Density calculation by TM method
def density_TM(mu, T, M):
density = derivative(lambda x: GrandPotential_TM(x, T, M), mu, n=1, dx=0.0001)
print ("M", M, "mu", mu, "density", density)
return density

#Heat capacity calculation by TM method
def heat_cap_TM(mu, T, M):
HC = derivative(lambda x: GrandPotential_TM(mu, x, M), T, n=2, dx=0.0001)
print ("TM", "M", M, "mu", mu, "HC", HC)
return HC

#Density calculation by TRG method
def density_TRG(mu, T, D):
density = derivative(lambda x: GrandPotential_TRG(x, T, D), mu, n=1, dx=0.0001)
print ("TRG", "D", D, "mu", mu, "density", density)
return density

#Heat capacity calculation by TRG method
def heat_cap_TRG(mu, T, D):
HC = derivative(lambda x: GrandPotential_TRG(mu, x, D), T, n=2, dx=0.0001)
print ("TRG", "D", D, "mu", mu, "HC", HC)
return HC

#find critical chemical potential by TRG method
if False:
prev = None
for D in range(16, 17,5):
opt_res = minimize_scalar(lambda mu: heat_cap_TRG(mu, T, D=D), bounds=(-2.5, -2.4), method='bounded', tol=0.00001)
omu = opt_res.x
print ("ISO TRG", "D",D, "omu", omu, (log(abs(prev -omu)) if prev is not None else ""))
prev = omu

#find critical chemical potential by TM method
if False:
prev = None
for M in range(9, 16):
t0 = time()
opt_res = minimize_scalar(lambda mu: heat_cap_TM(mu, T,M=M), bounds=(-2.5, -2.4), method='bounded', tol=0.00001)
omu = opt_res.x
print ("ISO TM", "M", M*join_y, "omu", omu, (log(abs(prev -omu)) if prev is not None else ""))
prev = omu

#calculate density or/and entropy and heat capacity curves
if True:
GrandPotential_function = GrandPotential_TRG                                                            #calculation method
Entropy_and_heat_capacity = False                                                                                       #calculation entropy and heat capacity.
dmu = 0.0001                                                                                                                            #differentiation step for chemical potential
dT = 0.0001                                                                                                                                     #differentiation step for temperature
calc_param = 16                                                                                                                         #calculation parametr of method. D for TRG and M for TM.
for mu in np.arange(-8.0,8.0001,0.1):                                                                           #chemical potential for investigation
GP_dmu = []
for dop_mu in np.arange(mu-dmu,mu+dmu+1e-5,2.0*dmu):
GP_dmu.append(GrandPotential_function(dop_mu,T,calc_param))
density = (GP_dmu-GP_dmu)/(dmu*2.0)                                                               #first derivative of grand partition on chemical potential
GP_dT = []
if Entropy_and_heat_capacity :
for dop_T in np.arange(T-dT,T+dT+1e-5,dT):
GP_dT.append(GrandPotential_function(mu,dop_T,calc_param))
entropy = -(-0.5*GP_dT+0.5*GP_dT)/dT                                                      #first derivative of grand partition on temperature
heat_capacity = -(GP_dT-2.0*GP_dT+GP_dT)/(dT**2)                       #second derivative of grand partition on temperature
print ("ISO mu", -mu," GP", GP_dT,"density", density,"entropy", entropy,"heat_capacity",heat_capacity)
else:
print ("ISO mu", -mu, "density=",density)