# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Cooling Fraction Function

# - Using **scipy.interpolate.RectBivariateSpline** to interpolate function:
# - Linear interpolation (kx=1, ky=1)

# +
import numpy as np
import pandas as pd
import sympy as s
import math
import glob

import matplotlib
import matplotlib.pyplot as plt
from cycler import cycler

from scipy import interpolate

from astropy import units as u
from astropy import constants as const
from astropy.units import imperial
imperial.enable()

# %matplotlib inline
import os
#home_dir = os.environ['/Users/eriksolhaug'] + '/'
import pyCloudy as pc
# -

pd.set_option("display.max_rows", None, "display.max_columns", None)
pd.set_option('precision', 16)
pd.set_option('display.float_format', lambda x: '%.3e' % x)

# +
# The directory in which we will have the model
# You may want to change this to a different place so that the current directory
# will not receive all the Cloudy files.
dir_ = '/Users/eriksolhaug/cloudy/c17.02/cloudoutputs/'

user_dir = '/Users/eriksolhaug/cloudy/c17.02/'

# Define verbosity to high level (will print errors, warnings and messages)
pc.log_.level = 3


# -

def calccoolingfrac(model_name, model_dir):
    """
    Returns FUNCTION object f(X, Y) where X is a hydrogen density (LOG) and Y is a temperature (LOG). The output fraction is in LOG.
    
    The temperature range given to the Cloudy input file is 10^4 - 10^7 K
    and the hydrogen density range given to the Cloudy input file is 10^-5.0 - 10^-3.0 cm^-3

    Inputs:
    model_name - a string, name of model run in CLOUDY (f.ex. 'model_41')
    model_dir - a string, directory containing model output files
    
    Output:
    f - an object, function yielding cooling fraction for the requested for the given model in LOG as a function of hydrogen density and temperature
    To get the cooling fraction for a certain hden X and temperature Y, simply call f(X, Y) and take the exponent 10**f(X, Y) of this to find 
    the cooling fraction for the model at the specified hydrogen density and temperature.
    
    An example: 
                In[]:
                f = calccolfrac('model_47', '/Users/eriksolhaug/cloudy/c17.02/cloudoutputs/')
                10**f(-5.0, 5.5)
                Out[]:
                array([[4.50693322e-33]]) -- where 4.50693322e-33 is the cooling fraction of OVI for hydrogen density of 10**-5.0 and a temperature of 10**5.5 Kelvin
    """
    
    # Interpolating function for the data computed in Cloudy
    grid_df = pd.read_csv(f'{model_dir}/{model_name}.grid', sep='\t')
    hdengrid = grid_df.loc[:, 'HDEN=%f L']
    tempgrid = grid_df.loc[:, 'CONSTANT ']
    x = hdengrid
    y = tempgrid
    
    # Loading .cool file
    cols = ['#depth cm','Temp K','Htot erg/cm3/s','Ctot erg/cm3/s','cool fracs']
    cool_df = pd.read_csv(f'{model_dir}{model_name}.cool', sep = '\t', usecols=cols)
    
    cool_df = cool_df.iloc[:, 0:4] # Excluding cool fracs column which contains unnecessary info, we want the H... and C... columns
    
    # Getting fractional columns
    z_array = []
    for index in range(0, x.size):
        frac_cool = float(cool_df.iloc[index*2, -1]) # Need to convert this to float as values are reported as strings
        z_array.append(np.log10(frac_cool+1e-50))  #Adding a small value to avoid case of log(0)
    z = pd.DataFrame(z_array, columns=['z'])
    
    if model_name == 'model_43':
        step = 5
    elif model_name == 'model_45':
        step = 41
    elif model_name == 'model_46':
        step = 401
    elif model_name == 'model_47':
        step = 101
    else:
        step = 11
    
    # Putting vectors in dataframe representation
    xyz = pd.DataFrame({'x' : x, 'y' : y, 'z' : z['z']})
    
    # Simplifying x and y inputs
    xi = xyz['x'][:step]
    yi = xyz['y'][::step]
    
    # Preparing spline arrays
    twoDarray = []
    for i in range(len(xi)):
        array = []
        for j in range(len(yi)):
            idx = i + j*step
            array.append(xyz['z'][idx])
        twoDarray.append(array)
    
    # Simplifying z inputs
    zi = twoDarray
    
    print(xi, yi, zi)
    print(len(xi), len(yi), len(zi))
    
    # INTERPOLATION
    f = interpolate.RectBivariateSpline(xi, yi, zi, kx=1, ky=1) # Linear determined by kx, ky
    
    # Displaying match between old fractions and interpolated function
    interpolated_z = []
    for temp in yi:
        for hden in xi:
            interpolated_z.append(f(hden, temp))
    interpolated_z = np.concatenate(interpolated_z)
    
    print(interpolated_z)
    print(min(interpolated_z))
    
    return f


def coolingplot(keyword, second_val_arr, model_name, model_dir, plot_dir):
    '''
    Function used to make and save plots of the fractional columns for different elements
    
    Input:
    keyword - a string, either 'temp' or 'hden' for what needs to be plotted against
    second_val - a float in LOG, value for either temp or hden (whatever not requested by keyword) - the plot needs a set temp or hden and this is set by this parameter
    model_name - a string, name of model run in CLOUDY (f.ex. 'model_42')
    model_dir - a string, directory containing model output files
    plot_dir - a string, directory for where to save plot
    '''    
    
    # Plotting
    tick_fontsize = 16
    axis_fontsize = 22
    lwidth = 3
    lstyle = '-'
    
    fig, ax = plt.subplots(1, 1)
    fig.set_size_inches(8, 5)
    fig.tight_layout(w_pad = 10.0)

    if keyword == 'temp':
        constant_arr = second_val_arr
        vary = np.arange(4, 7, 0.01)
        xaxis = 10**vary
    elif keyword == 'hden':
        vary = np.arange(-5.0, -3.0, 0.01)
        constant_arr = second_val_arr
        xaxis = 10**vary
    else:
        print('Not a valid keyword. Needs either "temp" or "hden".')
    
    f = calccoolingfrac(model_name, model_dir)
    if keyword == 'temp':
        other_keyword = 'hden'
        for constant in constant_arr:
            ax.plot(xaxis, 10**f(constant, vary)[0], linewidth=lwidth, label=f'10^{constant}cm-3', linestyle=lstyle)
    elif keyword == 'hden':
        other_keyword = 'temp'
        for constant in constant_arr:
            ax.plot(xaxis, 10**f(vary, constant), linewidth=lwidth, label=f'10^{constant}K', linestyle=lstyle)

    ax.set_xscale('log')
    ax.set_yscale('log')
    if keyword == 'temp':
        ax.set_xlabel(r'Temperature ($K$)', fontsize = 12)
    else:
        ax.set_xlabel(r'Hydrogen density ($cm^{-3}$)', fontsize = 12)
    ax.set_ylabel('Cooling Fraction', fontsize = 12)
#     ax.set_ylim(1e-36, 1e-25)
    ax.set_title(f'Cooling fractions | Constant {other_keyword}', fontsize=18, fontweight='bold', pad=15)
    for tick in ax.xaxis.get_major_ticks():
        tick.label.set_fontsize(fontsize=tick_fontsize)
    for tick in ax.yaxis.get_major_ticks():
        tick.label.set_fontsize(fontsize=tick_fontsize)
    ax.xaxis.label.set_fontsize(axis_fontsize)
    ax.yaxis.label.set_fontsize(axis_fontsize)
    ax.tick_params(which='major', width=2, length=8)
    ax.tick_params(which='minor', width=1, length=5)
    ax.grid(linestyle='--')
    ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left', prop={'size': 18})
    ax.set_prop_cycle(cycler('color', ['c', 'm', 'y', 'k']))
    
    fig.savefig(f'{plot_dir}{keyword}vary_coolingfractions.pdf', bbox_inches="tight")


# ### Example Executions:

model_name = 'model_47'
model_dir = '/Users/eriksolhaug/cloudy/c17.02/cloudoutputs/'
f = calccoolingfrac(model_name, model_dir)

10**f(-5.0,5.5)

# $\: \uparrow$ This is the cooling fraction for the model at a given hydrogen density and temperature.

# One can use the same procedure for other ions:

# ### Plotting fractional columns for all ions (C, N, O, Si)

model_name = 'model_47'
plot_dir = '/Users/eriksolhaug/cloudy/c17.02/es/es_data/cooling_fractions/' + model_name + '/'
keyword_array = ['temp', 'hden']
temp_val_array = np.arange(-5.0, -2.5, 0.5)
hden_val_array = np.arange(4.0, 6.5, 0.5)
for keyword in keyword_array:
    if keyword == 'temp':
        second_val_array = temp_val_array
    elif keyword == 'hden':
        second_val_array = hden_val_array
    plot = coolingplot(keyword, second_val_array, model_name, model_dir, plot_dir)



# +
# Individual plots:

# def coolingplot(keyword, second_val, model_name, model_dir, plot_dir):
#     '''
#     Function used to make and save plots of the fractional columns for different elements
    
#     Input:
#     keyword - a string, either 'temp' or 'hden' for what needs to be plotted against
#     second_val - a float in LOG, value for either temp or hden (whatever not requested by keyword) - the plot needs a set temp or hden and this is set by this parameter
#     model_name - a string, name of model run in CLOUDY (f.ex. 'model_42')
#     model_dir - a string, directory containing model output files
#     plot_dir - a string, directory for where to save plot
#     '''    
    
#     # Plotting
#     tick_fontsize = 16
#     axis_fontsize = 22
#     lwidth = 3
#     lstyle = '-'
    
#     fig, ax = plt.subplots(1, 1)
#     fig.set_size_inches(8, 5)
#     fig.tight_layout(w_pad = 10.0)

#     if keyword == 'temp':
#         constant = second_val
#         vary = np.arange(4, 7, 0.01)
#         xaxis = 10**vary
#     elif keyword == 'hden':
#         vary = np.arange(-5.0, -3.0, 0.01)
#         constant = second_val
#         xaxis = 10**vary
#     else:
#         print('Not a valid keyword. Needs either "temp" or "hden".')
    
#     f = calccoolingfrac(model_name, model_dir)
#     if keyword == 'temp':
#         other_keyword = 'hden'
#         ax.plot(xaxis, 10**f(constant, vary)[0], linewidth=lwidth, label=f'{model_name}', linestyle=lstyle)
#     elif keyword == 'hden':
#         other_keyword = 'temp'
#         ax.plot(xaxis, 10**f(vary, constant), linewidth=lwidth, label=f'{model_name}', linestyle=lstyle)

#     ax.set_xscale('log')
#     ax.set_yscale('log')
#     if keyword == 'temp':
#         ax.set_xlabel(r'Temperature ($K$)', fontsize = 12)
#     else:
#         ax.set_xlabel(r'Hydrogen density ($cm^{-3}$)', fontsize = 12)
#     ax.set_ylabel('Cooling Fraction', fontsize = 12)
# #     ax.set_ylim(1e-36, 1e-25)
#     ax.set_title(f'Cooling fractions | Constant {other_keyword}: 10^{second_val}', fontsize=18, fontweight='bold', pad=15)
#     for tick in ax.xaxis.get_major_ticks():
#         tick.label.set_fontsize(fontsize=tick_fontsize)
#     for tick in ax.yaxis.get_major_ticks():
#         tick.label.set_fontsize(fontsize=tick_fontsize)
#     ax.xaxis.label.set_fontsize(axis_fontsize)
#     ax.yaxis.label.set_fontsize(axis_fontsize)
#     ax.tick_params(which='major', width=2, length=8)
#     ax.tick_params(which='minor', width=1, length=5)
#     ax.grid(linestyle='--')
#     ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left', prop={'size': 18})
#     ax.set_prop_cycle(cycler('color', ['c', 'm', 'y', 'k']))
    
#     fig.savefig(f'{plot_dir}{keyword}vary_{second_val}_coolingfractions.pdf', bbox_inches="tight")

StarcoderdataJupyterScriptsDedupFiltered