<jupyter_start><jupyter_text>Setup working folders<jupyter_code># Where to save the figures and data files
PROJECT_ROOT_DIR = "Results"
FIGURE_ID = "Results/FigureFiles"
DATA_ID = "DataFiles/"
if not os.path.exists(PROJECT_ROOT_DIR):
    os.mkdir(PROJECT_ROOT_DIR)

if not os.path.exists(FIGURE_ID):
    os.makedirs(FIGURE_ID)

if not os.path.exists(DATA_ID):
    os.makedirs(DATA_ID)

def image_path(fig_id):
    return os.path.join(FIGURE_ID, fig_id)

def data_path(dat_id):
    return os.path.join(DATA_ID, dat_id)

def save_fig(fig_id):
    plt.savefig(image_path(fig_id) + ".png", format='png')
# Bounds for parametrs in the emulator are same as prior ranges so
prior_df = pd.read_csv(filepath_or_buffer="DataFiles/PbPb2760_prior", index_col=0)
# Design points
design = pd.read_csv(filepath_or_buffer="DataFiles/PbPb2760_design")
design_Au = pd.read_csv(filepath_or_buffer="DataFiles/AuAu200_design")
design_validation = pd.read_csv(filepath_or_buffer="DataFiles/AuAu200_validation_design")
#Simulation outputs at the design points
simulation_df = []
simulation_df_Au = []
for idf in range(0,4):
    simulation_df_Au.append(pd.read_csv(filepath_or_buffer=f"DataFiles/AuAu200_simulation_{idf}"))
    simulation_df.append(pd.read_csv(filepath_or_buffer=f"DataFiles/PbPb2760_simulation_{idf}"))
    simulation_df[idf]=simulation_df[idf][simulation_df_Au[0].keys()]

    #simulation_sd_df.append(pd.read_csv(filepath_or_buffer=f"DataFiles/PbPb2760_simulation_error_{idf}"))
simulation_df_Au_validation=pd.read_csv(filepath_or_buffer=f"DataFiles/AuAU200_validation_simulation")

simulation_df_Au_validation.head()
X = design.values
XAu = design_Au.values
Xv = design_validation.values
X
XAu
df_clms=simulation_df[1].keys()<jupyter_output><empty_output><jupyter_text>Normalize all the other models using mean and variance of observable in the PbPb 2.76 TeV Grad model<jupyter_code>#normalize data with respect to lower fidelity
s_l = StandardScaler()
x = simulation_df[0].values
s_l.fit(x)
for idf in range(0,4):
    x_tmp = simulation_df[idf].values
    simulation_df[idf]= pd.DataFrame(s_l.transform(x_tmp),columns=df_clms)
    
    x_tmp = simulation_df_Au[idf].values
    simulation_df_Au[idf]= pd.DataFrame(s_l.transform(x_tmp),columns=df_clms)
    

#diff = np.array(prior_df.loc['max'].values - prior_df.loc['min'].values ).reshape(1,-1)
#diff_mat = np.repeat(diff,X.shape[0],axis=0)
#print(f'Shape of diff matt {diff_mat.shape}')
#X= np.divide(X,diff_mat)
print('normalizing validation data')
x_tmp = simulation_df_Au_validation.values
print(x_tmp)
simulation_df_Au_validation=pd.DataFrame(s_l.transform(x_tmp),columns=df_clms)
simulation_df_Au_validation.head()
simulation_df[0].head()
simulation_df_Au[0].head()
for idf in range(0,4):
    Y = simulation_df[idf].values
    print('###########################\n for PbPb 2760')
    print( "X.shape : "+ str(X.shape) )
    print( "Y.shape : "+ str(Y.shape) )

    print('###########################\n for AuAU 200')
    Y_Au = simulation_df_Au[idf].values
    print( "X.shape : "+ str(XAu.shape) )
    print( "Y.shape : "+ str(Y.shape) )
#Model parameter names in Latex compatble form
model_param_dsgn = ['$N$[$200$GeV]',
 '$p$',
 '$\\sigma_k$',
 '$w$ [fm]',
 '$d_{\\mathrm{min}}$ [fm]',
 '$\\tau_R$ [fm/$c$]',
 '$\\alpha$',
 '$T_{\\eta,\\mathrm{kink}}$ [GeV]',
 '$a_{\\eta,\\mathrm{low}}$ [GeV${}^{-1}$]',
 '$a_{\\eta,\\mathrm{high}}$ [GeV${}^{-1}$]',
 '$(\\eta/s)_{\\mathrm{kink}}$',
 '$(\\zeta/s)_{\\max}$',
 '$T_{\\zeta,c}$ [GeV]',
 '$w_{\\zeta}$ [GeV]',
 '$\\lambda_{\\zeta}$',
 '$b_{\\pi}$',
 '$T_{\\mathrm{sw}}$ [GeV]']
observables = ['dNch_deta',
 'dET_deta',
 'dN_dy_pion',
 'dN_dy_kaon',
 'dN_dy_proton',
 'mean_pT_pion',
 'mean_pT_kaon',
 'mean_pT_proton',
 'pT_fluct',
 'v22',
 'v32',
 'v42']
observables_latex_2 = ['$\\frac{dN_{ch}}{d\\eta}$',
 '$\\frac{dE_T}{d\\eta}$',
 '$\\frac{dN_{\\pi}}{dy}$',
 '$\\frac{dN_{K}}{dy}$',
 '$\\frac{dN_{P}}{dy}$',
 '$\\langle pT_{\\pi} \\rangle$',
 '$\\langle pT_{K} \\rangle$',
 '$\\langle pT_{P} \\rangle$',
 '$\\frac{\\delta p_T}{\\langle p_T \\rangle}$',
 '$v_2${2}',
 '$v_3${2}',
 '$v_4${2}']
observables_latex = ['$dN_{ch} / d\\eta$',
 '$dE_T / d\\eta$',
 '${dN_{\\pi}} / {dy}$',
 '${dN_{K}} / {dy}$',
 '${dN_{P}} / {dy}$',
 '$\\langle p_{T, \\pi} \\rangle$',
 '$\\langle p_{T, K} \\rangle$',
 '$\\langle p_{T, P} \\rangle$',
 '${\\delta p_T} / {\\langle p_T \\rangle}$',
 '$v_2${2}',
 '$v_3${2}',
 '$v_4${2}']<jupyter_output><empty_output><jupyter_text>Observables considered in this analysis<jupyter_code>simulation_df[0].keys()
observables_choosen = ['dN_dy_pion[0 5]',
 'dN_dy_pion[40 50]',
# 'dN_dy_kaon[0 5]',
# 'dN_dy_kaon[60 70]',
# 'dN_dy_proton[0 5]',
# 'dN_dy_proton[60 70]',
 'mean_pT_pion[0 5]',
 'mean_pT_pion[40 50]',
# 'mean_pT_kaon[0 5]',
# 'mean_pT_kaon[60 70]',               
# 'mean_pT_proton[0 5]',
# 'mean_pT_proton[60 70]',
 'v22[0 5]',
  'v22[40 50]']<jupyter_output><empty_output><jupyter_text>Build the source emulators and find length scales and white noise that we can use in multifidelity code to fix hyper parameters.<jupyter_code>import GPy
import emukit.multi_fidelity
from emukit.model_wrappers.gpy_model_wrappers import GPyMultiOutputWrapper
from emukit.multi_fidelity.models import GPyLinearMultiFidelityModel
## Convert lists of arrays to ndarrays augmented with fidelity indicators
from sklearn.model_selection import train_test_split
from emukit.multi_fidelity.convert_lists_to_array import convert_x_list_to_array, convert_xy_lists_to_arrays
from sklearn.model_selection import KFold
#from emukit.multi_fidelity.models.non_linear_multi_fidelity_model import make_non_linear_kernels, NonLinearMultiFidelityModel
from sklearn.metrics import r2_score, mean_squared_error
import matplotlib
#low_gp_model = GPy.models.GPRegression.optimize_restarts()

    
# GP_dic={}
# n_opt = 10
# n_proc=10
# for selected_observable in observables_choosen:
#     Y_l=simulation_df[0][selected_observable].values.reshape(-1,1)
#     Y_h=simulation_df_Au[0][selected_observable].values.reshape(-1,1)
#     Y_v=simulation_df_Au_validation[selected_observable].values.reshape(-1,1)
    
#     x_train_h, x_test_h, y_train_h, y_test_h = XAu, Xv, Y_h, Y_v
#     x_train_l, y_train_l = X, Y_l
#     print(selected_observable)
#     kernel = GPy.kern.RBF(input_dim=17, ARD=True)
#     low_gp_model = GPy.models.GPRegression(x_train_l,y_train_l, kernel)
#         #high_gp_model.Gaussian_noise.fix(0)

#         ## Fit the GP model
    
#     low_gp_model.optimize_restarts(n_opt, verbose=True, parallel=True, num_processes=n_proc)
#     print(low_gp_model)
#     GP_dic[selected_observable]=low_gp_model
#import pickle
#with open(data_path('PbPbGradforAusobol'),"wb") as f:
#    pickle.dump(GP_dic,f)
import pickle
GP_dic = pickle.load(open(data_path('PbPbGradforAusobol'),"rb"))<jupyter_output><empty_output><jupyter_text>Transfer Learning using Emukit<jupyter_code>def run_train_and_validation(x_train_s, x_train_t, x_test_t, y_train_s, y_train_t, y_test_t, obs_name):
    """ Train two types of emulators. And perform validation on a given set
            1. Linear multifidelity GPs
            2. Standard GP
        =============================
        Return:
            [3 x 2 ndarray,rho value], [residue] of Transfer Learning model
            First row has MSE
            Second row has R2 scores
            Third row has the White noise variance
        """
    print('########################')
    print(obs_name)
    X_train, Y_train = convert_xy_lists_to_arrays([x_train_s, x_train_t], [y_train_s, y_train_t])
    n_opt = 10
    n_proc = 10
    
    ## Construct a linear multi-fidelity model

    #kernels = [GPy.kern.RBF(17, ARD=True), GPy.kern.RBF(17, ARD=True)]
    kernels = [GPy.kern.RBF(17, ARD=True), GPy.kern.RBF(17, ARD=True)]
    lin_tl_kernel = emukit.multi_fidelity.kernels.LinearMultiFidelityKernel(kernels)
    gpy_lin_tl_model = GPyLinearMultiFidelityModel(X_train, Y_train, lin_tl_kernel, n_fidelities=2)
    
    #Set the parameters in the Source Gaussian Process
    
    rbf_l_scales = GP_dic[obs_name].rbf.lengthscale.values
    rbf_var = GP_dic[obs_name].rbf.variance.values[0]
    wn = GP_dic[obs_name].Gaussian_noise.variance.values[0]
    
    gpy_lin_tl_model.multifidelity.rbf.lengthscale.fix(value=rbf_l_scales)
    gpy_lin_tl_model.multifidelity.rbf.variance.fix(value=rbf_var)
    gpy_lin_tl_model.mixed_noise.Gaussian_noise.fix(wn)
    #gpy_lin_mf_model.mixed_noise.Gaussian_noise_1.fix(0)


    ## Wrap the model using the given 'GPyMultiOutputWrapper'
    lin_tl_model = GPyMultiOutputWrapper(gpy_lin_tl_model, 2, n_optimization_restarts=n_opt,
                                         parallel=True, num_processes=n_proc)

    ## Fit the model
    st_l = time.time()
    lin_tl_model.optimize()
    et_l = time.time()
    
    #print("parameter array provided")
    #print(f'{rbf_l_scales} noise {wn}')
    
    if np.allclose(lin_tl_model.gpy_model.multifidelity.rbf.lengthscale.values, rbf_l_scales):
          print('RBF lengthscales are the same for low fidelity after optimization')
    if wn == lin_tl_model.gpy_model.mixed_noise.Gaussian_noise.variance.values[0]:
          print('White noise are the same after optimization')
    print('parameter array of multifidelity linear')
    print(lin_tl_model.gpy_model.param_array)
    print(f'time for linear mf optimization {et_l-st_l}')
    ## Create standard GP model using only high-fidelity data

    #kernel = GPy.kern.RBF(input_dim=17, ARD=True)
    #sd_gp_model = GPy.models.GPRegression(x_train_t,y_train_t, kernel)
    #high_gp_model.Gaussian_noise.fix(0)

    ## Fit the GP model
    
    #sd_gp_model.optimize_restarts(n_opt, verbose=True, parallel=True, num_processes=n_proc)
    
    x_temp = convert_x_list_to_array([x_test_t,x_test_t])
    x_test_t_idf_index = x_temp[x_test_t.shape[0]:,:]
    x_test_s_idf_index = x_temp[0:x_test_t.shape[0],:]
    
    
    #Transfer learning predictions and validation score calculations
    t_mean_lin_tl_model, t_var_lin_tl_model = lin_tl_model.predict(x_test_t_idf_index)
    s_mean_lin_tl_model, s_var_lin_tl_model = lin_tl_model.predict(x_test_s_idf_index)


    print(t_mean_lin_tl_model.shape)
    r2_tl = r2_score(y_test_t,t_mean_lin_tl_model )
    mse_tl = mean_squared_error(y_test_t,t_mean_lin_tl_model )
    print(f'r2 score for multifidelity linear {r2_score(y_test_t,t_mean_lin_tl_model )}')
    print(f'mse for multifidelity linear {mean_squared_error(y_test_t,t_mean_lin_tl_model )}')

    

    return lin_tl_model
# residue_df = pd.DataFrame()
# n_batch = 10
# df_results = pd.DataFrame(columns=["observable", "fold","n_training","r2_lin_tl","r2_stand_GP",
#                                   "mse_lin_tl","mse_stand_GP","wn_lin_tl","wn_stand_GP","rho"])
# split_i = 0
# lin_tl_emulators=[]
# for selected_observable in observables_choosen:
#     start_time = time.time()
#     Y_s=simulation_df[0][selected_observable].values.reshape(-1,1)
#     Y_t=simulation_df_Au[0][selected_observable].values.reshape(-1,1)
#     Y_v=simulation_df_Au_validation[selected_observable].values.reshape(-1,1)
    
#     x_train_t, x_test_t, y_train_t, y_test_t = XAu, Xv, Y_t, Y_v
#     x_train_s, y_train_s = X, Y_s
    
#     for i in range(0,n_batch):
#         if i in [5]:
#             l=0
#             h=(x_train_t.shape[0]//n_batch)*(i+1)
#         else:
#             continue
#         if i == n_batch-1:
#             h=x_train_t.shape[0]
#         lin_tl_emulators.append(run_train_and_validation(x_train_s, x_train_t[l:h,:],x_test_t, y_train_s, y_train_t[l:h,:], y_test_t, selected_observable))
#     end_time =time.time()
#     print(f'Time it take for {selected_observable} is {end_time-start_time}')
#    time_ar.append(end_time-start_time)
#time_ar= np.array(time_ar)

#import pickle
#with open(data_path('AuAusobol'),"wb") as f:
#    pickle.dump(lin_tl_emulators,f)
import pickle
lin_tl_emulators = pickle.load(open(data_path('AuAusobol'),"rb"))
lin_tl_emulators[0].Y<jupyter_output><empty_output><jupyter_text>We have the emulators. Now let's do a sensitivity study. (Global Sobel index)<jupyter_code>prior_df_Au = prior_df
prior_df_Au['norm']['max'] = 10
prior_df_Au['norm']['min'] = 3


prior_df_Au
prior_df_Au.values.T.shape
from SALib.sample import saltelli
from SALib.analyze import sobol
problem = {
    'groups' : ['$N$[$200$GeV]', 'IC', 'IC', 'IC','IC', 'FS', 'FS', '$(\\eta/s)$','$(\\eta/s)$','$(\\eta/s)$','$(\\eta/s)$',
               '$(\\zeta/s)$','$(\\zeta/s)$','$(\\zeta/s)$','$(\\zeta/s)$','$(\\eta/s)$', '$T_{\\mathrm{sw}}$ [GeV]'],
    'num_vars': 17,
    'names': model_param_dsgn,
    'bounds': prior_df_Au.values.T
}
param_values = saltelli.sample(problem, 16384)

param_values.shape
#Y = np.zeros([param_values.shape[0]])
sns.set_style("whitegrid")
sns.set_context('poster')
Si_array = []
X_train_sb, Y_train_sb = convert_xy_lists_to_arrays([param_values, param_values], [param_values, param_values])
for i, obs in enumerate(observables_choosen):
    print(obs)
    Yt,_ = lin_tl_emulators[i].predict(X_train_sb[param_values.shape[0]:,:])
    Ys,_ = lin_tl_emulators[i].predict(X_train_sb[0:param_values.shape[0],:])
    Ys = Ys *lin_tl_emulators[i].gpy_model.multifidelity.param_array[-1]
    discrp = Yt-Ys
    Y = discrp.flatten()
    Si = sobol.analyze(problem, Y)
    print(Si['S1'])
    print(Si['ST'])
    Si_array.append(Si)
    #sns.set_context('talk')
    fig ,ax = plt.subplots(3,1, figsize=(30,30))
    fig.suptitle(obs)
    axes = Si.plot(ax=ax)
    plt.tight_layout()
    fig.savefig(f'AuAUGrad_{i}_sobolgroup.png', dpi = 50)
with open(data_path('AuAusobolgroup_index'),"wb") as f:
    pickle.dump(Si_array,f)
#with open(data_path('AuAusobolgroup_index'),"rb") as f:
#    Si_array_loaded=pickle.load(f)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1603979088004647, 'fraction_numerical': 0.016192274179494627, 'mean_word_length': 4.808941374947279, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '348135', 'license': 'MIT', 'n_tokens_mistral': 5778, 'n_tokens_neox': 5499, 'n_words': 852, 'path': 'Transfer_learning_AuAu+sobol_group_sensitivity.ipynb', 'repo_name': 'danOSU/TransferLearningEmulation', 'size': '13772'}