from pylab import *
import scipy.interpolate

"""
Script to plot equilibrium curves of AGN1/2, which
allows assessing thermal instability properties.

Also plotted are contours of constant cooling; these
rates determine the sizes of the clouds that can form.

Contours are evaluated more accurately using 
interpolation of the data tables.

Either unpack the tarball DPKW19_tables.tar in
the same dir as this script, or update datapath
to the location of the tables

To run from a terminal, do
>> python equilb_contours.py

To run using ipython, do
>> run equilb_contour.py


- A script originally due to Randall Dannen that has been 
modified and documented by Tim Waters on 12/5/2018
"""

#----------------------------------------------------
# INPUTS
#----------------------------------------------------
N = 400 # number of interpolation points

datapath = 'DPKW19_tables/'
datafile1 = datapath + 'AGN1_net.dat'
datafile2 = datapath + 'AGN2_net.dat'
datafile1c = datapath + 'AGN1_cooling.dat'
datafile2c = datapath + 'AGN2_cooling.dat'

title = 'AGN1/2 Equilibrium and Cooling Contours'

# flag to make x-axis log(Xi), with Xi the
# 'pressure photoionization parameter'
USE_BIG_XI = 0


#----------------------------------------------------
# load tables and unpack T and xi
#----------------------------------------------------
# load net cooling rates
data1  = np.genfromtxt(datafile1)
data2  = np.genfromtxt(datafile2)

# extract axes - same for all tables
Nx = data1.shape[1] - 1
Ny = data1.shape[0]  - 1
x_data = np.log10(data1[0,1:])   # first row of values
y_data = np.log10(data1[1:,0]) # first column of values

# load cooling rates also
data1c  = np.genfromtxt(datafile1c)
data2c  = np.genfromtxt(datafile2c)


#----------------------------------------------------
# duplicate T and xi many times over (neeeded
# for interpolation routine) and upack the rates
#----------------------------------------------------
Nv = Nx*Ny
x_data_train = np.zeros(Nv)
y_data_train = np.zeros(Nv)
z1 = np.zeros(Nv)
z2 = np.zeros(Nv)
z1c = np.zeros(Nv)
z2c = np.zeros(Nv)
for i in range(Ny):
	x_data_train[i*Nx:(i+1)*Nx] = x_data[:]
	ys = np.ones(Nx)*y_data[i]
	y_data_train[i*Nx:(i+1)*Nx] = ys

	# net cooling rates 
	z1[i*Nx:(i+1)*Nx] = data1[1+i,1:]
	z2[i*Nx:(i+1)*Nx] = data2[1+i,1:]

	# cooling rates
	z1c[i*Nx:(i+1)*Nx] = data1c[1+i,1:]
	z2c[i*Nx:(i+1)*Nx] = data2c[1+i,1:]

# define pressure photoionization parameter, big Xi
if USE_BIG_XI:
	c = 3e10
	kb = 1.380658e-16
	bigXi_values = 10**x_data_train/(4.*np.pi*c*kb*10**y_data_train)
	x_data_train = np.log10(bigXi_values)


#----------------------------------------------------
# generate the interpolation points
#----------------------------------------------------
# new axes on which to interpolate	
x_vals = np.linspace(x_data_train.min(),x_data_train.max(),N)
y_vals = np.linspace(y_data_train.min(),y_data_train.max(),N)

# Generate grid data for our interpolation points.
z1_grid = scipy.interpolate.griddata((x_data_train,y_data_train), z1, (x_vals[None,:], y_vals[:,None]), method='linear')
z2_grid = scipy.interpolate.griddata((x_data_train,y_data_train), z2, (x_vals[None,:], y_vals[:,None]), method='linear')

z1c_grid = scipy.interpolate.griddata((x_data_train,y_data_train), z1c, (x_vals[None,:], y_vals[:,None]), method='linear')
z2c_grid = scipy.interpolate.griddata((x_data_train,y_data_train), z2c, (x_vals[None,:], y_vals[:,None]), method='linear')


#----------------------------------------------------
# make the contour plot
#----------------------------------------------------
# net cooling contours
eq1_contour = contour(10**x_vals,10**y_vals,z1_grid,levels=[0],colors='k')
eq2_contour = contour(10**x_vals,10**y_vals,z2_grid,levels=[0],colors='k',linestyles='--')

# cooling contours
c_rates = [1e-26, 1e-25]
eq1c_contour = contour(10**x_vals,10**y_vals,z1c_grid,levels=c_rates,colors=['c','b'])
eq2c_contour = contour(10**x_vals,10**y_vals,z2c_grid,levels=c_rates,colors=['c','b'],linestyles='--')

# step that shouldn't be necessary 
eq1_contour = eq1_contour.collections[0]
eq2_contour = eq2_contour.collections[0]
eq1c_contour = eq1c_contour.collections[0]
eq2c_contour = eq2c_contour.collections[0]

plt.setp(eq1_contour,linewidth=2) 
plt.setp(eq2_contour,linewidth=2) 
plt.setp(eq1c_contour,linewidth=1,label='AGN1, cooling =' + str(c_rates[-1]))
plt.setp(eq2c_contour,linewidth=1,label='AGN2, cooling = ' + str(c_rates[-1]))


plt.legend(loc='best')

# limits
xmin,xmax = 10**x_vals[0],10**x_vals[-1]
ymin,ymax = 10**y_vals[0],10**y_vals[-1]
xlim(xmin,xmax)
ylim(ymin,ymax)

# labels 
if USE_BIG_XI:
	tag = '_Xi'
	x_string = r'$\Xi$'
else:
	tag = '_xi'
	x_string = r'$\xi$'
xlabel(x_string,fontsize=16)
ylabel('Temperature [K]',fontsize=16)
xscale("log")
yscale("log")
plt.title('AGN1 (solid lines) and AGN2 (dashed lines)')
savefig('AGN1and2_contours' + tag + '.png')
show()