--- a
+++ b/src/plotting.py
@@ -0,0 +1,902 @@
+import seaborn as sns
+import matplotlib.pyplot as plt
+import pickle
+import pandas as pd 
+import numpy as np
+from pathlib import Path
+import warnings
+warnings.simplefilter(action='ignore', category=FutureWarning)
+from PIL import Image
+import sys 
+
+# Hydra for configuration
+import hydra
+from omegaconf import DictConfig, OmegaConf
+from matplotlib.ticker import ScalarFormatter
+from matplotlib.ticker import MaxNLocator
+from matplotlib.ticker import FuncFormatter
+
+# Custom formatter function
+def custom_formatter(x, pos):
+    if x.is_integer():
+        return f'{int(x)}'
+    elif x==0.5:
+        return r'$1/2$'
+    elif x==0.25:
+        return r'$1/4$'
+    # Do a diagonal fraction instead
+
+    else:
+        return f'{x:.2f}'
+
+cblind_palete = sns.color_palette("colorblind", as_cmap=True)
+learner_colors = {
+    "Torch_SLearner": cblind_palete[0],
+    "Torch_TLearner": cblind_palete[1],
+    "Torch_XLearner": cblind_palete[2],
+    "Torch_TARNet": cblind_palete[3],
+    'Torch_CFRNet_0.01': cblind_palete[4],
+    "Torch_CFRNet_0.001": cblind_palete[6],
+    'Torch_CFRNet_0.0001': cblind_palete[9],
+    'Torch_ActionNet': cblind_palete[7],
+    "Torch_DRLearner": cblind_palete[8],
+    "Torch_RALearner": cblind_palete[9],
+    "Torch_DragonNet": cblind_palete[5],
+    "Torch_DragonNet_2": cblind_palete[5],
+    "Torch_DragonNet_4": cblind_palete[3],
+    "Torch_ULearner": cblind_palete[6],
+    "Torch_PWLearner": cblind_palete[7],
+    "Torch_RLearner": cblind_palete[8],
+    "Torch_FlexTENet": cblind_palete[9],
+    "EconML_CausalForestDML": cblind_palete[2],
+    "EconML_DML": cblind_palete[0],
+    "EconML_DMLOrthoForest": cblind_palete[1],
+    "EconML_DRLearner": cblind_palete[6],
+    "EconML_DROrthoForest": cblind_palete[9],
+    "EconML_ForestDRLearner": cblind_palete[7],
+    "EconML_LinearDML": cblind_palete[8],
+    "EconML_LinearDRLearner": cblind_palete[5],
+    "EconML_SparseLinearDML": cblind_palete[3],
+    "EconML_SparseLinearDRLearner": cblind_palete[4],
+    "EconML_XLearner_Lasso": cblind_palete[7],
+    "EconML_TLearner_Lasso": cblind_palete[8],
+    "EconML_SLearner_Lasso": cblind_palete[9],
+    "DiffPOLearner": cblind_palete[0],
+    "Truth": cblind_palete[9],
+}
+
+learner_linestyles = {
+    "Torch_SLearner": "-",
+    "Torch_TLearner": "--",
+    "Torch_XLearner": ":",
+    "Torch_TARNet": "-.",
+    "Torch_DragonNet": "--",
+    "Torch_DragonNet_2": "-",
+    "Torch_DragonNet_4": "-.",
+    "Torch_XLearner": "--",
+    "Torch_CFRNet_0.01": "-",
+    "Torch_CFRNet_0.001": ":",
+    "Torch_CFRNet_0.0001": "--",
+    "Torch_DRLearner": "-",
+    "Torch_RALearner": "--",
+    "Torch_ULearner": "-",
+    "Torch_PWLearner": "-",
+    "Torch_RLearner": "-",
+    "Torch_FlexTENet": "-",
+    'Torch_ActionNet': "-",
+    "EconML_CausalForestDML": "-",
+    "EconML_DML": "--",
+    "EconML_DMLOrthoForest": ":",
+    "EconML_DRLearner": "-.",
+    "EconML_DROrthoForest": "--",
+    "EconML_ForestDRLearner": "-.",
+    "EconML_LinearDML": ":",
+    "EconML_LinearDRLearner": "-",
+    "EconML_SparseLinearDML": "--",
+    "EconML_SparseLinearDRLearner": ":",
+    "EconML_SLearner_Lasso": "-.",
+    "EconML_TLearner_Lasso": "--",
+    "EconML_XLearner_Lasso": "-",
+    "DiffPOLearner": "-.",
+    "Truth": ":",
+}
+
+
+learner_markers = {
+    "Torch_SLearner": "d",
+    "Torch_TLearner": "o",
+    "Torch_XLearner": "^",
+    "Torch_TARNet": "*",
+    "Torch_DragonNet": "x",
+    "Torch_DragonNet_2": "o",
+    "Torch_DragonNet_4": "*",
+    "Torch_XLearner": "D",
+    "Torch_CFRNet_0.01": "8",
+    "Torch_CFRNet_0.001": "s",
+    "Torch_CFRNet_0.0001": "x",
+    "Torch_DRLearner": "x",
+    "Torch_RALearner": "H",
+    "Torch_ULearner": "x",
+    "Torch_PWLearner": "*",
+    "Torch_RLearner": "*",
+    "Torch_FlexTENet": "*",
+    'Torch_ActionNet': "*",
+    "EconML_CausalForestDML": "d",
+    "EconML_DML": "o",
+    "EconML_DMLOrthoForest": "^",
+    "EconML_DRLearner": "*",
+    "EconML_DROrthoForest": "D",
+    "EconML_ForestDRLearner": "8",
+    "EconML_LinearDML": "s",
+    "EconML_LinearDRLearner": "x",
+    "EconML_SparseLinearDML": "x",
+    "EconML_SparseLinearDRLearner": "H",
+    "EconML_TLearner_Lasso": "o",
+    "EconML_SLearner_Lasso": "^",
+    "EconML_XLearner_Lasso": "d",
+    "DiffPOLearner": "H",
+    "Truth": "<",
+}
+
+datasets_names_map = {
+    "tcga_100": "TCGA", 
+    "twins": "Twins", 
+    "news_100": "News", 
+    "all_notupro_technologies": "AllNoTuproTechnologies",
+    "all_notupro_technologies_small": "AllNoTuproTechnologiesSmall",
+    "dummy_data": "DummyData",
+    "selected_technologies_pategan_1000": "selected_technologies_pategan_1000",
+    "selected_technologies_with_fastdrug": "selected_technologies_with_fastdrug",
+    "cytof_normalized":"cytof_normalized",
+    "cytof_normalized_with_fastdrug":"cytof_normalized_with_fastdrug",
+    "cytof_pategan_1000_normalized": "cytof_pategan_1000_normalized",
+    "all_notupro_technologies_with_fastdrug": "all_notupro_technologies_with_fastdrug",
+    "acic": "ACIC2016", 
+    "depmap_drug_screen_2_drugs": "depmap_drug_screen_2_drugs",
+    "depmap_drug_screen_2_drugs_norm": "depmap_drug_screen_2_drugs_norm",
+    "depmap_drug_screen_2_drugs_all_features": "depmap_drug_screen_2_drugs_all_features",
+    "depmap_drug_screen_2_drugs_all_features_norm": "depmap_drug_screen_2_drugs_all_features_norm",
+    "depmap_crispr_screen_2_kos": "depmap_crispr_screen_2_kos",
+    "depmap_crispr_screen_2_kos_norm": "depmap_crispr_screen_2_kos_norm",
+    "depmap_crispr_screen_2_kos_all_features": "depmap_crispr_screen_2_kos_all_features",
+    "depmap_crispr_screen_2_kos_all_features_norm": "depmap_crispr_screen_2_kos_all_features_norm",
+    "depmap_drug_screen_2_drugs_100pcs_norm":"depmap_drug_screen_2_drugs_100pcs_norm",
+    "depmap_drug_screen_2_drugs_5000hv_norm":"depmap_drug_screen_2_drugs_5000hv_norm",
+    "depmap_crispr_screen_2_kos_100pcs_norm":"depmap_crispr_screen_2_kos_100pcs_norm",
+    "depmap_crispr_screen_2_kos_5000hv_norm":"depmap_crispr_screen_2_kos_5000hv_norm",
+    "independent_normally_dist": "independent_normally_dist",
+    "ovarian_semi_synthetic_l1":"ovarian_semi_synthetic_l1",
+    "ovarian_semi_synthetic_rf":"ovarian_semi_synthetic_rf",
+    "melanoma_semi_synthetic_l1": "melanoma_semi_synthetic_l1",
+    "pred": "Predictive confounding", 
+    "prog": "Prognostic confounding", 
+    "irrelevant_var": "Non-confounded propensity",
+    "selective": "General confounding"}
+
+metric_names_map = {
+    'Pred: Pred features ACC': r'Predictive $\mathrm{Attr}$', #^{\mathrm{pred}}_{\mathrm{pred}}$',
+    'Pred: Prog features ACC': r'$\mathrm{Attr}^{\mathrm{pred}}_{\mathrm{prog}}$',
+    'Pred: Select features ACC': r'$\mathrm{Attr}^{\mathrm{pred}}_{\mathrm{select}}$',
+    'Prog: Pred features ACC': r'$\mathrm{Attr}^{\mathrm{prog}}_{\mathrm{pred}}$',
+    'Prog: Prog features ACC': r'Prognostic $\mathrm{Attr}$', #^{\mathrm{prog}}$', #_{\mathrm{prog}}$',
+    'Prog: Select features ACC': r'$\mathrm{Attr}^{\mathrm{prog}}_{\mathrm{select}}$',
+    'Select: Pred features ACC': r'$\mathrm{Attr}^{\mathrm{select}}_{\mathrm{pred}}$',
+    'Select: Prog features ACC': r'$\mathrm{Attr}^{\mathrm{select}}_{\mathrm{prog}}$',
+    'Select: Select features ACC': r'$\mathrm{Attr}^{\mathrm{select}}_{\mathrm{select}}$',
+    'CI Coverage': 'CI Coverage',
+    'Normalized PEHE': 'N. PEHE',
+    'PEHE': 'PEHE',
+    'CF RMSE': 'CF-RMSE',
+    'AUROC': 'AUROC',
+    'Factual AUROC': 'Factual AUROC',
+    'CF AUROC': "CF AUROC",
+    'Factual RMSE': 'F-RMSE',
+    "Factual RMSE Y0": "F-RMSE Y0", 
+    "Factual RMSE Y1": "F-RMSE Y1",
+    "CF RMSE Y0": "CF-RMSE Y0",
+    "CF RMSE Y1": "CF-RMSE Y1",
+    'Normalized F-RMSE': 'N. F-RMSE',
+    'Normalized CF-RMSE': 'N. CF-RMSE',
+    "F-Outcome true mean":"F-Outcome true mean",
+    "CF-Outcome true mean":"CF-Outcome true mean",
+    "F-Outcome true std":"F-Outcome true std",
+    "CF-Outcome true std":"CF-Outcome true std",
+    "F-CF Outcome Diff":"F-CF Outcome Diff",
+    'Swap AUROC@1': 'AUROC@1',
+    'Swap AUPRC@1': 'AUPRC@1',
+    'Swap AUROC@5': 'AUROC@5',
+    'Swap AUPRC@5': 'AUPRC@5',
+    'Swap AUROC@tre': 'AUROC@tre',
+    'Swap AUPRC@tre': 'AUPRC@tre',
+    'Swap AUROC@all': 'AUROC',
+    'Swap AUPRC@all': 'AUPRC',
+    "GT Pred Expertise": r'$\mathrm{B}^{\pi}_{Y_1-Y_0}$',
+    "GT Prog Expertise": r'$\mathrm{B}^{\pi}_{Y_0}$',
+    "GT Tre Expertise": r'$\mathrm{B}^{\pi}_{Y_1}$',
+    "Upd. GT Pred Expertise": r'$\mathrm{B}^{\hat{\pi}}_{Y_1-Y_0}$',
+    "Upd. GT Prog Expertise": r'$\mathrm{B}^{\hat{\pi}}_{Y_0}$',
+    "Upd. GT Tre Expertise": r'$\mathrm{B}^{\hat{\pi}}_{Y_1}$',
+    "GT Expertise Ratio": r'$\mathrm{E}^{\pi}_{\mathrm{ratio}}$',
+    "GT Total Expertise": r'$\mathrm{B}^{\pi}_{Y_0,Y_1}$',
+    "ES Pred Expertise": "ES Pred Bias",
+    "ES Prog Expertise": "ES Prog Bias",
+    "ES Total Expertise": "ES Outcome Bias",
+    "Pred Precision": r'$\mathrm{Prec}^{\hat{\pi}}_{\mathrm{Ass.}}$',
+    "Policy Precision": r'$\mathrm{Prec}^{\pi}_{\mathrm{Ass.}}$',
+    'T Distribution: Train': 'T Distribution: Train',
+    'T Distribution: Test': 'T Distribution: Test',
+    'True Swap Perc': 'True Swap Perc',
+    "Normalized F-CF Diff": "Normalized F-CF Diff",
+    'Training Duration': 'Training Duration',
+    "FC PEHE":"PEHE(Model) - PEHE(TARNet)",
+    "FC F-RMSE":"Rel. N. F-RMSE",
+    "FC CF-RMSE":"Rel. N. CF-RMSE",
+    "FC Swap AUROC":"Rel. AUROC",
+    "FC Swap AUPRC":"Rel. AUPRC",
+    "GT In-context Var":r'$\mathrm{B}^{\pi}_{X}$', 
+    "ES In-context Var":"ES Total Bias",
+    "GT-ES Pred Expertise Diff":r'$\mathrm{E}^{\pi}_{\mathrm{pred}}$ Error',
+    "GT-ES Prog Expertise Diff":r'$\mathrm{E}^{\pi}_{\mathrm{prog}}$ Error',
+    "GT-ES Total Expertise Diff":r'$\mathrm{E}^{\pi}$ Error',
+    "RMSE Y0":"RMSE Y0",
+    "RMSE Y1":"RMSE Y1",
+}
+
+learners_names_map = {
+    "Torch_TLearner":"T-Learner-MLP", 
+    "Torch_SLearner": "S-Learner-MLP", 
+    "Torch_TARNet": "Baseline-TAR",  
+    "Torch_DragonNet": "DragonNet-1 (Act. Pred.)",
+    "Torch_DragonNet_2": "DragonNet-2 (Act. Pred.)",
+    "Torch_DragonNet_4": "DragonNet-4 (Act. Pred.)",
+    "Torch_DRLearner": "Direct-DR", 
+    "Torch_XLearner": "XLearner-MLP (Direct)", 
+    "Torch_CFRNet_0.001": 'CFRNet-0.001 (Balancing)', #-\gamma=0.001)$',  
+    "Torch_CFRNet_0.01": 'CFRNet-0.01 (Balancing)', 
+    "Torch_CFRNet_0.0001": 'CFRNet-0.0001 (Balancing)', 
+    'Torch_ActionNet': "ActionNet (Act. Pred.)",
+    "Torch_RALearner": "RA-Learner",
+    "Torch_ULearner": "U-Learner",
+    "Torch_PWLearner":"Torch_PWLearner",
+    "Torch_RLearner":"Torch_RLearner",
+    "Torch_FlexTENet":"Torch_FlexTENet",
+    "EconML_CausalForestDML": "CausalForestDML",
+    "EconML_DML": "APred-Prop-Lasso",
+    "EconML_DMLOrthoForest": "DMLOrthoForest",
+    "EconML_DRLearner": "DRLearner",
+    "EconML_DROrthoForest": "DROrthoForest",
+    "EconML_ForestDRLearner": "ForestDRLearner",
+    "EconML_LinearDML": "LinearDML",
+    "EconML_LinearDRLearner": "LinearDRLearner",
+    "EconML_SparseLinearDML": "SparseLinearDML",
+    "EconML_SparseLinearDRLearner": "SparseLinearDRLearner",
+    "EconML_TLearner_Lasso": "T-Learner-Lasso",
+    "EconML_SLearner_Lasso": "S-Learner-Lasso",
+    "EconML_XLearner_Lasso": "XLearnerLasso",
+    "DiffPOLearner": "DiffPOLearner",
+    
+    "Truth": "Truth"
+}
+
+compare_values_map = {
+    # Propensity
+    "none_prog": r'$\pi_{\mathrm{RCT}} \rightarrow \pi_{\mathrm{Y_0}}^\beta$',
+    "none_tre": r'$\pi_{\mathrm{RCT}} \rightarrow \pi_{\mathrm{Y_1}}^\beta$',
+    "none_pred": r'$\pi_{\mathrm{RCT}} \rightarrow \pi_{\mathrm{Y_1-Y_0}}^\beta$',
+    "rct_none": r'$\pi_{\mathrm{RCT}} \rightarrow \pi_{\mathrm{X_{irr}}}^\beta$',
+    "none_pred_overlap": r'$\pi_{\mathrm{RCT}} \rightarrow \pi_{\mathrm{X_{pred}}}^\beta$',
+
+    # Toy
+    "toy7": r'$\pi_{\mathrm{RCT}} \rightarrow \pi_{\mathrm{T_7}}^\beta$',
+    "toy8_nonlinear": r'$\pi_{\mathrm{RCT}} \rightarrow \pi_{\mathrm{T_8}}^\beta$',
+    "toy1_linear": r'$\pi_{\mathrm{RCT}} \rightarrow \pi_{\mathrm{T_1^{lin}}}^\beta$',
+    "toy1_nonlinear": r'$\mathrm{Toy 1:} \pi_{\mathrm{T_1}}^\beta$', #r'$\pi_{\mathrm{RCT}} \rightarrow \pi_{\mathrm{T_1}}^\beta$',
+    "toy2_linear": r'$\pi_{\mathrm{RCT}} \rightarrow \pi_{\mathrm{T_3^{lin}}}^\beta$',
+    "toy2_nonlinear": r'$\mathrm{Toy 3:} \pi_{\mathrm{T_3}}^\beta$',
+    "toy3_nonlinear": r'$\mathrm{Toy 2:} \pi_{\mathrm{T_2}}^\beta$',
+    "toy4_nonlinear": r'$\pi_{\mathrm{RCT}} \rightarrow \pi_{\mathrm{T_4}}^\beta$',
+    "toy5": r'$\pi_{\mathrm{RCT}} \rightarrow \pi_{\mathrm{T_5}}^\beta$',
+    "toy6_nonlinear": r'$\mathrm{Toy 4:} \pi_{\mathrm{T_4}}^\beta$', #r'$\pi_{\mathrm{RCT}} \rightarrow \pi_{\mathrm{T_6}}^\beta$',
+
+    # Expertise
+    # "prog_tre": r'$\pi_{\mathrm{Y_0}}^{\beta=4} \rightarrow \pi_{\mathrm{Y_1-Y_0}}^{\beta=4} \rightarrow \pi_{\mathrm{Y_1}}^{\beta=4}$',
+    # "none_prog": r'$\pi_{\mathrm{X_{rand}}}^{\beta=4} \rightarrow \pi_{\mathrm{Y_0}}^{\beta=4}$',
+    # "none_tre": r'$\pi_{\mathrm{X_{rand}}}^{\beta=4} \rightarrow \pi_{\mathrm{Y_1}}^{\beta=4}$',
+    # "none_pred": r'$\pi_{\mathrm{X_{rand}}}^{\beta=4} \rightarrow \pi_{\mathrm{Y_1-Y_0}}^{\beta=4}$',
+    #["prog_tre", "none_prog", "none_tre", "none_pred"]
+
+    0: r'$\pi_{\mathrm{RCT}}$',
+    2: r'$\pi_{\mathrm{Y_1-Y_0}}^{\beta=2}$',
+    100: r'$\pi_{\mathrm{Y_1-Y_0}}^{\beta=100}$',
+
+    "0": r'$\pi_{\mathrm{RCT}}$',
+    "2": r'$\pi_{\mathrm{Y_1-Y_0}}^{\beta=2}$',
+    "100": r'$\pi_{\mathrm{Y_1-Y_0}}^{\beta=100}$',
+}
+
+
+def plot_results_datasets_compare(results_df: pd.DataFrame, 
+                          model_names: list,
+                          dataset: str,
+                          compare_axis: str,
+                          compare_axis_values,
+                          x_axis, 
+                          x_label_name, 
+                          x_values_to_plot, 
+                          metrics_list, 
+                          learners_list, 
+                          figsize, 
+                          legend_position, 
+                          seeds_list, 
+                          n_splits,
+                          sharey=False, 
+                          legend_rows=1,
+                          dim_X=1,
+                          log_x_axis = False): 
+    """
+    Plot the results for a given dataset.
+    """
+    # Get the unique values of the compare axis
+    if compare_axis_values is None:
+        compare_axis_values = results_df[compare_axis].unique()
+
+    metrics_list = ["Pred Precision"]
+    # Initialize the plot
+    nrows = len(metrics_list)
+    columns = len(compare_axis_values)
+    figsize = (3*columns+2, 3*nrows)
+    #figsize = (3*columns, 3)
+
+    font_size=10
+    fig, axs = plt.subplots(len(metrics_list), len(compare_axis_values), figsize=figsize, squeeze=False, sharey=sharey, dpi=500)
+    plt.gcf().subplots_adjust(bottom=0.15)
+    
+    # Aggregate results across seeds for each metric
+    for i in range(len(compare_axis_values)):
+        cmp_value = compare_axis_values[i]
+        for metric_id, metric in enumerate(metrics_list):
+            for model_name in model_names:
+                # Extract results for individual cate models
+                sub_df = results_df.loc[(results_df["Learner"] == model_name)]
+                sub_df = sub_df.loc[(sub_df[compare_axis] == cmp_value)][[x_axis, metric]]
+                sub_df = sub_df[sub_df[x_axis].isin(x_values_to_plot)]
+                sub_df_mean = sub_df.groupby(x_axis).agg('median').reset_index()
+                sub_df_std = sub_df.groupby(x_axis).agg('std').reset_index()
+                sub_df_min = sub_df.groupby(x_axis).agg('min').reset_index()
+                sub_df_max = sub_df.groupby(x_axis).agg('max').reset_index()
+
+                # Plot the results
+                x_values = sub_df_mean.loc[:, x_axis].values
+
+                try:
+                    y_values = sub_df_mean.loc[:, metric].values
+                except:
+                    continue
+
+                y_err = sub_df_std.loc[:, metric].values / (np.sqrt(n_splits*len(seeds_list)))
+                y_min = sub_df_min.loc[:, metric].values
+                y_max = sub_df_max.loc[:, metric].values
+                
+                # axs[metric_id][i].plot(x_values, y_values, label=learners_names_map[model_name], 
+                #                                             color=learner_colors[model_name], linestyle=learner_linestyles[model_name], marker=learner_markers[model_name], markersize=5)
+                axs[metric_id][i].plot(x_values, y_values, label=learners_names_map[model_name], 
+                                                            color=learner_colors[model_name], linestyle=learner_linestyles[model_name], marker=learner_markers[model_name], markersize=3, alpha=0.5)
+                axs[metric_id][i].fill_between(x_values, y_values-y_err, y_values+y_err, alpha=0.1, color=learner_colors[model_name])
+
+            
+            
+            # if log_x_axis:
+            #     axs[metric_id][i].set_xscale('symlog', linthresh=0.5, base=2)
+            #     #axs[metric_id][i].fill_between(x_values, y_min, y_max, alpha=0.1, color=learner_colors[model_name])
+            
+            axs[metric_id][i].tick_params(axis='x', labelsize=font_size-2)
+            axs[metric_id][i].tick_params(axis='y', labelsize=font_size-1)
+
+            
+            axs[metric_id][i].set_title(compare_values_map[cmp_value], fontsize=font_size+11, y=1.04)
+
+            axs[metric_id][i].set_xlabel(x_label_name, fontsize=font_size-1)
+            if i == 0:
+                axs[metric_id][i].set_ylabel(metric_names_map[metric], fontsize=font_size-1)
+
+            if log_x_axis:
+                axs[metric_id][i].set_xscale('symlog', linthresh=0.5, base=2)
+                # Display as fractions if not integers and as integers if integers
+                # axs[0][i].xaxis.set_major_formatter(ScalarFormatter())
+                # Get the current ticks
+                current_ticks = axs[metric_id][i].get_xticks()
+                
+                # Calculate the midpoint between the first and second tick
+                if len(current_ticks) > 1:
+                    midpoint = (current_ticks[0] + current_ticks[1]) / 2
+                    # Add the midpoint to the list of ticks
+                    new_ticks = [current_ticks[0], midpoint] + list(current_ticks[1:])
+                    axs[metric_id][i].set_xticks(new_ticks)
+
+                # Add a tick at 0.25
+                axs[metric_id][i].set_xticks(sorted(set(axs[metric_id][i].get_xticks()).union({0.25})))
+                axs[metric_id][i].xaxis.set_major_formatter(FuncFormatter(custom_formatter))
+
+            if metric in ["True Swap Perc", "T Distribution: Train", "T Distribution: Test", "GT Total Expertise", "ES Total Expertise", "GT Expertise Ratio", "GT Pred Expertise", "GT Prog Expertise", "ES Pred Expertise", "ES Prog Expertise","GT In-context Var","ES In-context Var","GT-ES Pred Expertise Diff","GT-ES Prog Expertise Diff","GT-ES Total Expertise Diff", "Policy Precision", "GT In-context Var", "GT Total Expertise", "GT Prog Expertise", "GT Tre Expertise", "GT Pred Expertise", "Upd. GT Prog Expertise", "Upd. GT Tre Expertise", "Upd. GT Pred Expertise"]:
+                axs[metric_id][i].set_ylim(0, 1)
+
+            if metric == "PEHE":
+                axs[metric_id][i].set_ylim(top = 1.75)
+            #axs[metric_id][i].set_ylim(bottom=0.475)
+            #axs[metric_id][i].set_aspect(0.7/axs[metric_id][i].get_data_ratio(), adjustable='box')
+            #axs[metric_id][i].tick_params(axis='y', labelsize=font_size-1)
+
+            # axs[metric_id][i].tick_params(
+            #     axis='x',          # changes apply to the x-axis
+            #     which='both',      # both major and minor ticks are affected
+            #     bottom=False,      # ticks along the bottom edge are off
+            #     top=False,         # ticks along the top edge are off
+            #     labelbottom=False) # labels along the bottom edge are off
+
+    # Add the legend
+    lines_labels = [ax.get_legend_handles_labels() for ax in fig.axes]
+    lines, labels = [sum(lol, []) for lol in zip(*lines_labels)]
+    legend_rows = 6
+
+    # Iterate over each row of subplots
+    for row in range(len(axs)):
+        # Create a legend for each row
+        handles, labels = axs[row, -1].get_legend_handles_labels()
+        axs[row, -1].legend(
+            lines[:len(learners_list)],
+            labels[:len(learners_list)],
+            ncol=1, #len(learners_list) if legend_rows == 1 else int((len(learners_list) + 1) / legend_rows),
+            loc='center right',
+            bbox_to_anchor=(1.8, 0.5),
+            prop={'size': font_size+2}
+        )
+
+
+    #fig.tight_layout()
+    plt.subplots_adjust( wspace=0.07)
+    return fig   
+
+def plot_performance_metrics(results_df: pd.DataFrame, 
+                          model_names: list,
+                          dataset: str,
+                          compare_axis: str,
+                          compare_axis_values,
+                          x_axis, 
+                          x_label_name, 
+                          x_values_to_plot, 
+                          metrics_list, 
+                          learners_list, 
+                          figsize, 
+                          legend_position, 
+                          seeds_list, 
+                          n_splits,
+                          sharey=False, 
+                          legend_rows=1,
+                          dim_X=1,
+                          log_x_axis = False): 
+    
+    # Get the unique values of the compare axis
+    if compare_axis_values is None:
+        compare_axis_values = results_df[compare_axis].unique()
+
+    metrics_list = ['PEHE', 'FC PEHE', "Pred Precision", 'Pred: Pred features ACC', 'Prog: Prog features ACC'] #]
+    #log_x_axis=False
+    # Initialize the plot
+    nrows = len(metrics_list)
+    columns = len(compare_axis_values)
+
+    #figsize = (3*columns+2, 3*nrows) #PREV
+    figsize = (3*columns+2, 3.4*nrows)
+    #figsize = (3*columns, 3)
+
+    font_size=10
+    fig, axs = plt.subplots(len(metrics_list), len(compare_axis_values), figsize=figsize, squeeze=False, sharey=sharey, dpi=500)
+    plt.gcf().subplots_adjust(bottom=0.15)
+    
+    model_names_cpy = model_names.copy()
+    # Aggregate results across seeds for each metric
+    for i in range(len(compare_axis_values)):
+        cmp_value = compare_axis_values[i]
+        for metric_id, metric in enumerate(metrics_list):
+            # if metric in ["FC PEHE", 'Prog: Prog features ACC', 'Prog: Pred features ACC']:
+            #     model_names = ["Torch_TARNet","Torch_DragonNet","Torch_CFRNet_0.001","EconML_TLearner_Lasso"]
+            # else:
+            model_names = model_names_cpy #["Torch_TARNet","Torch_DragonNet","Torch_ActionNet", "Torch_CFRNet_0.001","EconML_TLearner_Lasso"]
+
+            for model_name in model_names:
+                # Extract results for individual cate models
+                sub_df = results_df.loc[(results_df["Learner"] == model_name)]
+                sub_df = sub_df.loc[(sub_df[compare_axis] == cmp_value)][[x_axis, metric]]
+                sub_df = sub_df[sub_df[x_axis].isin(x_values_to_plot)]
+                sub_df_mean = sub_df.groupby(x_axis).agg('median').reset_index()
+                sub_df_std = sub_df.groupby(x_axis).agg('std').reset_index()
+                sub_df_min = sub_df.groupby(x_axis).agg('min').reset_index()
+                sub_df_max = sub_df.groupby(x_axis).agg('max').reset_index()
+
+                # Plot the results
+                x_values = sub_df_mean.loc[:, x_axis].values
+
+                try:
+                    y_values = sub_df_mean.loc[:, metric].values
+                except:
+                    continue
+
+                y_err = sub_df_std.loc[:, metric].values / (np.sqrt(n_splits*len(seeds_list)))
+                y_min = sub_df_min.loc[:, metric].values
+                y_max = sub_df_max.loc[:, metric].values
+                
+                # axs[metric_id][i].plot(x_values, y_values, label=learners_names_map[model_name], 
+                #                                             color=learner_colors[model_name], linestyle=learner_linestyles[model_name], marker=learner_markers[model_name], markersize=5)
+                axs[metric_id][i].plot(x_values, y_values, label=learners_names_map[model_name], 
+                                                            color=learner_colors[model_name], linestyle=learner_linestyles[model_name], marker=learner_markers[model_name], markersize=3, alpha=0.5)
+                axs[metric_id][i].fill_between(x_values, y_values-y_err, y_values+y_err, alpha=0.1, color=learner_colors[model_name])
+
+            
+            
+            # if log_x_axis:
+            #     axs[metric_id][i].set_xscale('symlog', linthresh=0.5, base=2)
+            #     #axs[metric_id][i].fill_between(x_values, y_min, y_max, alpha=0.1, color=learner_colors[model_name])
+            
+            axs[metric_id][i].tick_params(axis='x', labelsize=font_size-2)
+            axs[metric_id][i].tick_params(axis='y', labelsize=font_size-1)
+
+            
+            if metric_id == 0:
+                axs[metric_id][i].set_title(compare_values_map[cmp_value], fontsize=font_size+1, y=1.0)
+
+
+            axs[metric_id][i].set_xlabel(x_label_name, fontsize=font_size-1)
+            if i == 0:
+                axs[metric_id][i].set_ylabel(metric_names_map[metric], fontsize=font_size-1)
+
+            if log_x_axis:
+                axs[metric_id][i].set_xscale('symlog', linthresh=0.5, base=2)
+                # Display as fractions if not integers and as integers if integers
+                # axs[0][i].xaxis.set_major_formatter(ScalarFormatter())
+                # Get the current ticks
+                current_ticks = axs[metric_id][i].get_xticks()
+                
+                # Calculate the midpoint between the first and second tick
+                if len(current_ticks) > 1:
+                    midpoint = (current_ticks[0] + current_ticks[1]) / 2
+                    # Add the midpoint to the list of ticks
+                    new_ticks = [current_ticks[0], midpoint] + list(current_ticks[1:])
+                    axs[metric_id][i].set_xticks(new_ticks)
+
+                # Add a tick at 0.25
+                axs[metric_id][i].set_xticks(sorted(set(axs[metric_id][i].get_xticks()).union({0.25})))
+                axs[metric_id][i].xaxis.set_major_formatter(FuncFormatter(custom_formatter))
+
+            if metric in ["True Swap Perc", "T Distribution: Train", "T Distribution: Test", "GT Total Expertise", "ES Total Expertise", "GT Expertise Ratio", "GT Pred Expertise", "GT Prog Expertise", "ES Pred Expertise", "ES Prog Expertise","GT In-context Var","ES In-context Var","GT-ES Pred Expertise Diff","GT-ES Prog Expertise Diff","GT-ES Total Expertise Diff", "Policy Precision", "GT In-context Var", "GT Total Expertise", "GT Prog Expertise", "GT Tre Expertise", "GT Pred Expertise", "Upd. GT Prog Expertise", "Upd. GT Tre Expertise", "Upd. GT Pred Expertise"]:
+                axs[metric_id][i].set_ylim(0, 1)
+
+            # if metric == "PEHE":
+            #     axs[metric_id][i].set_ylim(top = 1.75)
+            #axs[metric_id][i].set_ylim(bottom=0.475)
+            #axs[metric_id][i].set_aspect(0.7/axs[metric_id][i].get_data_ratio(), adjustable='box')
+            #axs[metric_id][i].tick_params(axis='y', labelsize=font_size-1)
+
+            axs[metric_id][i].tick_params(
+                axis='x',          # changes apply to the x-axis
+                which='both',      # both major and minor ticks are affected
+                bottom=False,      # ticks along the bottom edge are off
+                top=False,         # ticks along the top edge are off
+                labelbottom=False) # labels along the bottom edge are off
+
+    # Add the legend
+    lines_labels = [ax.get_legend_handles_labels() for ax in fig.axes]
+    lines, labels = [sum(lol, []) for lol in zip(*lines_labels)]
+    legend_rows = 6
+
+    # Iterate over each row of subplots
+    for row in range(len(axs)):
+        # Create a legend for each row
+        handles, labels = axs[row, -1].get_legend_handles_labels()
+        axs[row, -1].legend(
+            lines[:len(learners_list)],
+            labels[:len(learners_list)],
+            ncol=1, #len(learners_list) if legend_rows == 1 else int((len(learners_list) + 1) / legend_rows),
+            loc='center right',
+            bbox_to_anchor=(1.9, 0.5),
+            prop={'size': font_size+2}
+        )
+
+
+    plt.subplots_adjust( wspace=0.07)
+    #fig.tight_layout()
+    return fig  
+
+
+
+def plot_performance_metrics_f_cf(results_df: pd.DataFrame, 
+                          model_names: list,
+                          dataset: str,
+                          compare_axis: str,
+                          compare_axis_values,
+                          x_axis, 
+                          x_label_name, 
+                          x_values_to_plot, 
+                          metrics_list, 
+                          learners_list, 
+                          figsize, 
+                          legend_position, 
+                          seeds_list, 
+                          n_splits,
+                          sharey=False, 
+                          legend_rows=1,
+                          dim_X=1,
+                          log_x_axis = False): 
+    # Get the unique values of the compare axis
+    if compare_axis_values is None:
+        compare_axis_values = results_df[compare_axis].unique()
+
+    # Initialize the plot
+    #model_names = model_names[0] #["Torch_TARNet"] #[EconML_TLearner_Lasso"]
+    columns = len(compare_axis_values)
+    rows = len(model_names)
+    figsize = (3*columns+2, 3.3*rows)
+    #figsize = (3*columns, 3)
+    font_size=10
+    fig, axs = plt.subplots(len(model_names), len(compare_axis_values), figsize=figsize, squeeze=False, sharey=sharey, dpi=500)
+    #plt.gcf().subplots_adjust(bottom=0.15)
+    
+    # Filter results_df for first model and first seed and first split
+    
+    # results_df = results_df.loc[(results_df["Seed"] == seeds_list[0])]
+    # results_df = results_df.loc[(results_df["Split ID"] == 0)]
+
+    # Only consider expertise metrics
+    #colors = ['black', 'orange', 'darkorange', 'orchid', 'darkorchid']
+    colors = ['blue', 'lightcoral', 'lightgreen', 'red', 'green']
+
+    markers = ['o', 'D', 'D', 'x',  'x']
+    metrics_list = ["PEHE", "Factual RMSE Y0", "Factual RMSE Y1", "CF RMSE Y0", "CF RMSE Y1"]
+    filtered_df = results_df[[x_axis, compare_axis] + metrics_list]
+
+    # Aggregate results across seeds for each metric
+    for model_id, model_name in enumerate(model_names):
+        filtered_df_model = filtered_df.loc[(results_df["Learner"] == model_name)]
+        for i in range(len(compare_axis_values)):
+            cmp_value = compare_axis_values[i]
+
+            # Plot all metric outcomes as lines in a single plot for the given cmp_value and use x_axis as x-axis
+            x_values = filtered_df_model[x_axis].values
+
+            for metric_id, metric in enumerate(metrics_list):
+                # Extract results for individual cate models
+
+                sub_df = filtered_df_model.loc[(filtered_df_model[compare_axis] == cmp_value)][[x_axis, metric]]
+                sub_df = sub_df[sub_df[x_axis].isin(x_values_to_plot)]
+                sub_df_mean = sub_df.groupby(x_axis).agg('median').reset_index()
+                sub_df_std = sub_df.groupby(x_axis).agg('std').reset_index()
+                sub_df_min = sub_df.groupby(x_axis).agg('min').reset_index()
+                sub_df_max = sub_df.groupby(x_axis).agg('max').reset_index()
+
+                # Plot the results
+                x_values = sub_df_mean.loc[:, x_axis].values
+
+                try:
+                    y_values = sub_df_mean.loc[:, metric].values
+                except:
+                    continue
+
+                y_err = sub_df_std.loc[:, metric].values / (np.sqrt(n_splits*len(seeds_list)))
+                y_min = sub_df_min.loc[:, metric].values
+                y_max = sub_df_max.loc[:, metric].values
+
+                # use a different linestyle for each metric
+
+                axs[model_id][i].plot(x_values, y_values, label=metric_names_map[metric], 
+                                                                color=colors[metric_id], linestyle='-', marker=markers[metric_id], alpha=0.5, markersize=3)
+                axs[model_id][i].fill_between(x_values, y_values-y_err, y_values+y_err, alpha=0.1, color=colors[metric_id])
+
+
+                axs[model_id][i].tick_params(axis='x', labelsize=font_size-2)
+                axs[model_id][i].tick_params(axis='y', labelsize=font_size-1)
+                
+                axs[model_id][i].set_title(compare_values_map[cmp_value], fontsize=font_size+2, y=1.01)
+
+                axs[model_id][i].set_xlabel(x_label_name, fontsize=font_size-1)
+
+            # if i == 0:
+            #     axs[model_id][i].set_ylabel(metric_names_map[metric], fontsize=font_size-1)
+
+            if log_x_axis:
+                axs[model_id][i].set_xscale('symlog', linthresh=0.5, base=2)
+                # Display as fractions if not integers and as integers if integers
+                # axs[0][i].xaxis.set_major_formatter(ScalarFormatter())
+                # Get the current ticks
+                current_ticks = axs[model_id][i].get_xticks()
+                
+                # Calculate the midpoint between the first and second tick
+                if len(current_ticks) > 1:
+                    midpoint = (current_ticks[0] + current_ticks[1]) / 2
+                    # Add the midpoint to the list of ticks
+                    new_ticks = [current_ticks[0], midpoint] + list(current_ticks[1:])
+                    axs[model_id][i].set_xticks(new_ticks)
+
+                # Add a tick at 0.25
+                axs[model_id][i].set_xticks(sorted(set(axs[model_id][i].get_xticks()).union({0.25})))
+                axs[model_id][i].xaxis.set_major_formatter(FuncFormatter(custom_formatter))
+            axs[model_id][i].tick_params(
+                axis='x',          # changes apply to the x-axis
+                which='both',      # both major and minor ticks are affected
+                bottom=False,      # ticks along the bottom edge are off
+                top=False,         # ticks along the top edge are off
+                labelbottom=False) # labels along the bottom edge are off
+            
+            axs[model_id][i].tick_params(axis='y', labelsize=font_size-1)
+            #axs[model_id][i].set_aspect(0.7/axs[model_id][i].get_data_ratio(), adjustable='box')
+
+
+    # Add the legend
+    lines_labels = [ax.get_legend_handles_labels() for ax in fig.axes]
+    lines, labels = [sum(lol, []) for lol in zip(*lines_labels)]
+
+    # Add legends to the right of each row
+    for i, row in enumerate(axs):
+        lines_labels = [ax.get_legend_handles_labels() for ax in row]
+        lines, labels = [sum(lol, []) for lol in zip(*lines_labels)]
+        row[-1].legend(
+            lines[:len(metrics_list)],
+            labels[:len(metrics_list)],
+            loc='center right',
+            bbox_to_anchor=(1.8, 0.5),
+            ncol=1,
+            prop={'size': font_size+2},
+            title_fontsize=font_size+4,
+            title=learners_names_map[model_names[i]]
+        )
+
+    #fig.tight_layout()
+    plt.subplots_adjust( wspace=0.07)
+
+    return fig
+
+
+def plot_expertise_metrics(results_df: pd.DataFrame, 
+                          model_names: list,
+                          dataset: str,
+                          compare_axis: str,
+                          compare_axis_values,
+                          x_axis, 
+                          x_label_name, 
+                          x_values_to_plot, 
+                          metrics_list, 
+                          learners_list, 
+                          figsize, 
+                          legend_position, 
+                          seeds_list, 
+                          n_splits,
+                          sharey=False, 
+                          legend_rows=1,
+                          dim_X=1,
+                          log_x_axis = False): 
+    
+    if compare_axis_values is None:
+        compare_axis_values = results_df[compare_axis].unique()
+
+    # Initialize the plot
+    columns = len(compare_axis_values)
+    figsize = (3*columns+2, 3)
+    font_size=10
+    fig, axs = plt.subplots(1, len(compare_axis_values), figsize=figsize, squeeze=False, sharey=sharey, dpi=500)
+    plt.gcf().subplots_adjust(bottom=0.15)
+    
+    # Filter results_df for first model and first seed and first split
+    results_df = results_df.loc[(results_df["Learner"] == model_names[0])]
+    results_df = results_df.loc[(results_df["Seed"] == seeds_list[0])]
+    results_df = results_df.loc[(results_df["Split ID"] == 0)]
+
+
+    # Only consider expertise metrics
+    colors = ['black', 'grey', 'red', 'green', 'blue']
+    markers = ['o', 'x', 'x',  'x', 'x']
+    metrics_list = ["Policy Precision", "GT In-context Var", "GT Prog Expertise", "GT Tre Expertise", "GT Pred Expertise"]
+    sub_df = results_df[[x_axis, compare_axis, "Seed", "Split ID"] + metrics_list]
+
+    # Aggregate results across seeds for each metric
+    for i in range(len(compare_axis_values)):
+        cmp_value = compare_axis_values[i]
+
+        # Plot all metric outcomes as lines in a single plot for the given cmp_value and use x_axis as x-axis
+        filtered_df = sub_df[(sub_df[compare_axis] == cmp_value)]
+        x_values = filtered_df[x_axis].values
+
+        for metric_id, metric in enumerate(metrics_list):
+            y_values = filtered_df[metric].values
+            # use a different linestyle for each metric
+
+            axs[0][i].plot(x_values, y_values, label=metric_names_map[metric], color=colors[metric_id], linestyle='-', marker=markers[metric_id], alpha=0.5,  markersize=5)
+            
+
+            # if i == 0:
+            #     axs[0][i].set_ylabel("Selection Bias", fontsize=font_size)
+
+            axs[0][i].tick_params(axis='x', labelsize=font_size-2)
+            axs[0][i].tick_params(axis='y', labelsize=font_size-1)
+            axs[0][i].set_title(compare_values_map[cmp_value], fontsize=font_size+11, y=1.04)
+            axs[0][i].set_xlabel(x_label_name, fontsize=font_size-1)
+
+        if log_x_axis:
+            axs[0][i].set_xscale('symlog', linthresh=0.5, base=2)
+            # Display as fractions if not integers and as integers if integers
+            # axs[0][i].xaxis.set_major_formatter(ScalarFormatter())
+            # Get the current ticks
+            current_ticks = axs[0][i].get_xticks()
+            
+            # Calculate the midpoint between the first and second tick
+            if len(current_ticks) > 1:
+                midpoint = (current_ticks[0] + current_ticks[1]) / 2
+                # Add the midpoint to the list of ticks
+                new_ticks = [current_ticks[0], midpoint] + list(current_ticks[1:])
+                axs[0][i].set_xticks(new_ticks)
+
+            # Add a tick at 0.25
+            axs[0][i].set_xticks(sorted(set(axs[0][i].get_xticks()).union({0.25})))
+            axs[0][i].xaxis.set_major_formatter(FuncFormatter(custom_formatter))
+
+        # if i == 0:
+        #         axs[0][i].set_ylabel(metric_names_map[metric], fontsize=font_size-1)
+
+        axs[0][i].tick_params(
+            axis='x',          # changes apply to the x-axis
+            which='both',      # both major and minor ticks are affected
+            bottom=False,      # ticks along the bottom edge are off
+            top=False,         # ticks along the top edge are off
+            labelbottom=False) # labels along the bottom edge are off
+        #axs[0][i].set_aspect(0.7/axs[0][i].get_data_ratio(), adjustable='box')
+        axs[0][i].tick_params(axis='y', labelsize=font_size-1)
+
+    # Add the legend
+    lines_labels = [ax.get_legend_handles_labels() for ax in fig.axes]
+    lines, labels = [sum(lol, []) for lol in zip(*lines_labels)]
+
+    fig.legend(
+        lines[:len(metrics_list)],
+        labels[:len(metrics_list)],
+        loc='center right',  # Position the legend to the right
+        #bbox_to_anchor=(1, 0.5),  # Adjust the anchor point to the right center
+        ncol=1,  # Set the number of columns to 1 for a vertical legend
+        prop={'size': font_size+2}
+    )
+
+    #fig.tight_layout()
+    plt.subplots_adjust( wspace=0.07)
+
+    return fig
+    
+                           
+
+def merge_pngs(images, axis="horizontal"):
+    """
+    Merge a list of png images into a single image.
+    """
+    widths, heights = zip(*(i.size for i in images))
+
+    if axis == "vertical":
+        total_height = sum(heights)
+        max_width = max(widths)
+
+        new_im = Image.new('RGB', (max_width, total_height))
+
+        y_offset = 0
+        for im in images:
+            new_im.paste(im, (0,y_offset))
+            y_offset += im.size[1]
+        
+        return new_im
+    
+    elif axis == "horizontal":
+        total_width = sum(widths)
+        max_height = max(heights)
+
+        new_im = Image.new('RGB', (total_width, max_height))
+
+        x_offset = 0
+        for im in images:
+            new_im.paste(im, (x_offset,0))
+            x_offset += im.size[0]
+        
+        return new_im
+