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
Datasets
Models
