Datasets
Models
Downloads: 1
Diff of
/src/iterpretability/experiments.py
[000000]
..
[6ac965]
Switch to side-by-side view
--- a +++ b/src/iterpretability/experiments.py @@ -0,0 +1,1278 @@ +from pathlib import Path +import os +import catenets.models as cate_models +import numpy as np +import pandas as pd + +import src.iterpretability.logger as log +from src.iterpretability.explain import Explainer +from src.iterpretability.datasets.data_loader import load +from src.iterpretability.synthetic_simulate import ( + SyntheticSimulatorLinear, + SyntheticSimulatorModulatedNonLinear, +) +from src.iterpretability.utils import ( + attribution_accuracy, + compute_pehe, +) + +# For contour plotting +import umap +from sklearn.decomposition import PCA +from sklearn.manifold import TSNE +from sklearn.linear_model import LogisticRegression +import matplotlib.pyplot as plt +import matplotlib.tri as tri +import imageio +import torch +import shap + +def get_learners(model_list, X_train, Y_train, n_iter, batch_size, batch_norm, discrete_treatment=True, discrete_outcome=False): + learners = { + "TLearner": cate_models.torch.TLearner( + X_train.shape[1], + binary_y=(len(np.unique(Y_train)) == 2), + n_layers_out=2, + n_units_out=100, + batch_size=batch_size, + n_iter=n_iter, + batch_norm=batch_norm, + nonlin="relu", + ), + "SLearner": cate_models.torch.SLearner( + X_train.shape[1], + binary_y=(len(np.unique(Y_train)) == 2), + n_layers_out=2, + n_units_out=100, + n_iter=n_iter, + batch_size=batch_size, + batch_norm=batch_norm, + nonlin="relu", + ), + "TARNet": cate_models.torch.TARNet( + X_train.shape[1], + binary_y=(len(np.unique(Y_train)) == 2), + n_layers_r=1, + n_layers_out=1, + n_units_out=100, + n_units_r=100, + batch_size=batch_size, + n_iter=n_iter, + batch_norm=batch_norm, + nonlin="relu", + ), + "DRLearner": cate_models.torch.DRLearner( + X_train.shape[1], + binary_y=(len(np.unique(Y_train)) == 2), + n_layers_out=2, + n_units_out=100, + n_iter=n_iter, + batch_size=batch_size, + batch_norm=batch_norm, + nonlin="relu", + ), + "XLearner": cate_models.torch.XLearner( + X_train.shape[1], + binary_y=(len(np.unique(Y_train)) == 2), + n_layers_out=2, + n_units_out=100, + n_iter=n_iter, + batch_size=batch_size, + batch_norm=batch_norm, + nonlin="relu", + ), + "CFRNet_0.01": cate_models.torch.TARNet( + X_train.shape[1], + binary_y=(len(np.unique(Y_train)) == 2), + n_layers_r=1, + n_layers_out=1, + n_units_out=100, + n_units_r=100, + batch_size=batch_size, + n_iter=n_iter, + batch_norm=batch_norm, + nonlin="relu", + penalty_disc=0.01, + ), + "CFRNet_0.001": cate_models.torch.TARNet( + X_train.shape[1], + binary_y=(len(np.unique(Y_train)) == 2), + n_layers_r=1, + n_layers_out=1, + n_units_out=100, + n_units_r=100, + batch_size=batch_size, + n_iter=n_iter, + batch_norm=batch_norm, + nonlin="relu", + penalty_disc=0.001, + ), + "CFRNet_0.0001": cate_models.torch.TARNet( + X_train.shape[1], + binary_y=(len(np.unique(Y_train)) == 2), + n_layers_r=1, + n_layers_out=1, + n_units_out=100, + n_units_r=100, + batch_size=batch_size, + n_iter=n_iter, + batch_norm=batch_norm, + nonlin="relu", + penalty_disc=0.0001, + ), + "EconML_CausalForestDML": cate_models.econml.EconMlEstimator( + model_name="EconML_CausalForestDML", + # cv_model_selection=3, + discrete_treatment=discrete_treatment, + discrete_outcome=discrete_outcome, + ), + "EconML_DMLOrthoForest": cate_models.econml.EconMlEstimator( + model_name="EconML_DMLOrthoForest", + # cv_model_selection=3, + discrete_treatment=discrete_treatment, + discrete_outcome=discrete_outcome, + ), + + "EconML_SparseLinearDML": cate_models.econml.EconMlEstimator( + model_name="EconML_SparseLinearDML", + # cv_model_selection=3, + discrete_treatment=discrete_treatment, + discrete_outcome=discrete_outcome, + ), + "EconML_SparseLinearDRLearner": cate_models.econml.EconMlEstimator( + model_name="EconML_SparseLinearDRLearner", + # cv_model_selection=3, + discrete_treatment=discrete_treatment, + discrete_outcome=discrete_outcome, + ), + "EconML_LinearDRLearner": cate_models.econml.EconMlEstimator( + model_name="EconML_LinearDRLearner", + # cv_model_selection=3, + discrete_treatment=discrete_treatment, + discrete_outcome=discrete_outcome, + ), + "EconML_DRLearner": cate_models.econml.EconMlEstimator( + model_name="EconML_DRLearner", + # cv_model_selection=3, + discrete_treatment=discrete_treatment, + discrete_outcome=discrete_outcome, + ), + "EconML_XLearner": cate_models.econml.EconMlEstimator( + model_name="EconML_XLearner", + # cv_model_selection=3, + discrete_treatment=discrete_treatment, + discrete_outcome=discrete_outcome, + ), + "EconML_SLearner": cate_models.econml.EconMlEstimator( + model_name="EconML_SLearner", + # cv_model_selection=3, + discrete_treatment=discrete_treatment, + discrete_outcome=discrete_outcome, + ), + "EconML_TLearner": cate_models.econml.EconMlEstimator( + model_name="EconML_TLearner", + # cv_model_selection=3, + discrete_treatment=discrete_treatment, + discrete_outcome=discrete_outcome, + ), + "EconML_SparseLinearDRIV": cate_models.econml.EconMlEstimator( + model_name="EconML_SparseLinearDRIV", + # cv_model_selection=3, + discrete_treatment=discrete_treatment, + discrete_outcome=discrete_outcome, + ), + + } + + for name in model_list: + if name not in learners: + raise Exception(f"Unknown model name {name}.") + + # Only return the learners that are in the model_list + learners = {name: learners[name] for name in model_list} + + return learners + +def get_learner_explanations(learners, X_test, X_train, Y_train, W_train, explainer_limit, explainer_list, return_learners=False, already_trained=False): + learner_explainers = {} + learner_explanations = {} + + for name in learners: + log.info(f"Fitting {name}.") + + if not already_trained: + learners[name].fit(X=X_train, y=Y_train, w=W_train) + + log.info(f"Explaining {name}.") + + if "EconML" in name: + shap_values = learners[name].est.shap_values(X_test[:explainer_limit], background_samples=None) + treatment_names = learners[name].est.cate_treatment_names() + output_names = learners[name].est.cate_output_names() + output_name = output_names[0] + treatment_name = treatment_names[0] + learner_explanations[name] = {"kernel_shap" : shap_values[output_name][treatment_name].values} + + else: + learner_explainers[name] = Explainer( + learners[name], + feature_names=list(range(X_train.shape[1])), + explainer_list=explainer_list, + ) + learner_explanations[name] = learner_explainers[name].explain( + X_test[: explainer_limit] + ) + + if return_learners: + return learner_explanations, learners + else: + return learner_explanations + +class PredictiveSensitivity: + """ + Sensitivity analysis for predictive scale. + """ + + def __init__( + self, + n_units_hidden: int = 50, + n_layers: int = 1, + penalty_orthogonal: float = 0.01, + batch_size: int = 1024, + batch_norm: bool = False, + n_iter: int = 1000, + seed: int = 42, + explainer_limit: int = 1000, + save_path: Path = Path.cwd(), + propensity_type: str = "pred", + predictive_scales: list = [1e-3, 1e-2, 1e-1, 0.5, 1, 2], + num_interactions: int = 1, + synthetic_simulator_type: str = "linear", + selection_type: str = "random", + non_linearity_scale: float = 0, + model_list: list = ["TLearner"] + ) -> None: + + self.n_units_hidden = n_units_hidden + self.n_layers = n_layers + self.penalty_orthogonal = penalty_orthogonal + self.batch_size = batch_size + self.batch_norm = batch_norm + self.n_iter = n_iter + self.seed = seed + self.explainer_limit = explainer_limit + self.save_path = save_path + self.predictive_scales = predictive_scales + self.propensity_type = propensity_type + self.num_interactions = num_interactions + self.synthetic_simulator_type = synthetic_simulator_type + self.selection_type = selection_type + self.non_linearity_scale = non_linearity_scale + self.model_list = model_list + + def run( + self, + dataset: str = "tcga_10", + train_ratio: float = 0.8, + num_important_features: int = 2, + binary_outcome: bool = False, + random_feature_selection: bool = True, + explainer_list: list = [ + "feature_ablation", + "feature_permutation", + "integrated_gradients", + "shapley_value_sampling", + ], + debug: bool = False, + directory_path_: str = None, + ) -> None: + log.info( + f"Using dataset {dataset} with num_important features = {num_important_features}." + ) + + X_raw_train, X_raw_test = load(dataset, train_ratio=train_ratio, debug=debug, directory_path_=directory_path_) + + if self.synthetic_simulator_type == "linear": + sim = SyntheticSimulatorLinear( + X_raw_train, + num_important_features=num_important_features, + random_feature_selection=random_feature_selection, + seed=self.seed, + ) + elif self.synthetic_simulator_type == "nonlinear": + sim = SyntheticSimulatorModulatedNonLinear( + X_raw_train, + num_important_features=num_important_features, + non_linearity_scale=self.non_linearity_scale, + seed=self.seed, + selection_type=self.selection_type, + ) + else: + raise Exception("Unknown simulator type.") + + explainability_data = [] + + for predictive_scale in self.predictive_scales: + log.info(f"Now working with predictive_scale = {predictive_scale}...") + ( + X_train, + W_train, + Y_train, + po0_train, + po1_train, + propensity_train, + ) = sim.simulate_dataset( + X_raw_train, + predictive_scale=predictive_scale, + binary_outcome=binary_outcome, + treatment_assign=self.propensity_type, + ) + + X_test, W_test, Y_test, po0_test, po1_test, _ = sim.simulate_dataset( + X_raw_test, + predictive_scale=predictive_scale, + binary_outcome=binary_outcome, + treatment_assign=self.propensity_type, + ) + + log.info("Fitting and explaining learners...") + + learners = get_learners( + model_list=self.model_list, + X_train=X_train, + Y_train=Y_train, + n_iter=self.n_iter, + batch_size=self.batch_size, + batch_norm=self.batch_norm, + discrete_outcome=binary_outcome + ) + + learner_explanations = get_learner_explanations(learners, + X_test, X_train, Y_train, W_train, + self.explainer_limit, explainer_list) + + all_important_features = sim.get_all_important_features(with_selective=True) + pred_features = sim.get_predictive_features() + prog_features = sim.get_prognostic_features() + select_features = sim.get_selective_features() + + cate_test = sim.te(X_test) + + for explainer_name in explainer_list: + for learner_name in learners: + attribution_est = np.abs( + learner_explanations[learner_name][explainer_name] + ) + acc_scores_all_features = attribution_accuracy( + all_important_features, attribution_est + ) + acc_scores_predictive_features = attribution_accuracy( + pred_features, attribution_est + ) + acc_scores_prog_features = attribution_accuracy( + prog_features, attribution_est + ) + acc_scores_selective_features = attribution_accuracy( + select_features, attribution_est + ) + cate_pred = learners[learner_name].predict(X=X_test) + + pehe_test = compute_pehe(cate_true=cate_test, cate_pred=cate_pred) + + explainability_data.append( + [ + predictive_scale, + learner_name, + explainer_name, + acc_scores_all_features, + acc_scores_predictive_features, + acc_scores_prog_features, + acc_scores_selective_features, + pehe_test, + np.mean(cate_test), + np.var(cate_test), + pehe_test / np.sqrt(np.var(cate_test)), + ] + ) + + metrics_df = pd.DataFrame( + explainability_data, + columns=[ + "Predictive Scale", + "Learner", + "Explainer", + "All features ACC", + "Pred features ACC", + "Prog features ACC", + "Select features ACC", + "PEHE", + "CATE true mean", + "CATE true var", + "Normalized PEHE", + ], + ) + + results_path = self.save_path / "results/predictive_sensitivity" + log.info(f"Saving results in {results_path}...") + if not results_path.exists(): + results_path.mkdir(parents=True, exist_ok=True) + + metrics_df.to_csv( + results_path / f"predictive_scale_{dataset}_{num_important_features}_" + f"{self.synthetic_simulator_type}_random_{random_feature_selection}_" + f"binary_{binary_outcome}-seed{self.seed}.csv" + ) + + +class NonLinearitySensitivity: + """ + Sensitivity analysis for nonlinearity in prognostic and predictive functions. + """ + + def __init__( + self, + n_units_hidden: int = 50, + n_layers: int = 1, + penalty_orthogonal: float = 0.01, + batch_size: int = 1024, + batch_norm: bool = False, + n_iter: int = 1000, + seed: int = 42, + explainer_limit: int = 1000, + save_path: Path = Path.cwd(), + propensity_type: str = "pred", + nonlinearity_scales: list = [0.0, 0.2, 0.5, 0.7, 1.0], + selection_type: str = "random", + predictive_scale: float = 1, + synthetic_simulator_type: str = "random", + model_list: list = ["TLearner"] + ) -> None: + + self.n_units_hidden = n_units_hidden + self.n_layers = n_layers + self.penalty_orthogonal = penalty_orthogonal + self.batch_size = batch_size + self.batch_norm = batch_norm + self.n_iter = n_iter + self.seed = seed + self.explainer_limit = explainer_limit + self.save_path = save_path + self.propensity_type = propensity_type + self.nonlinearity_scales = nonlinearity_scales + self.selection_type = selection_type + self.predictive_scale = predictive_scale + self.synthetic_simulator_type = synthetic_simulator_type + self.model_list = model_list + + def run( + self, + dataset: str = "tcga_100", + num_important_features: int = 15, + explainer_list: list = [ + "feature_ablation", + "feature_permutation", + "integrated_gradients", + "shapley_value_sampling", + ], + train_ratio: float = 0.8, + binary_outcome: bool = False, + debug=False, + directory_path_: str = None, + + ) -> None: + log.info( + f"Using dataset {dataset} with num_important features = {num_important_features}." + ) + X_raw_train, X_raw_test = load(dataset, train_ratio=train_ratio, debug=debug, directory_path_=directory_path_) + explainability_data = [] + + for nonlinearity_scale in self.nonlinearity_scales: + log.info(f"Now working with a nonlinearity scale {nonlinearity_scale}...") + + if self.synthetic_simulator_type == "linear": + raise Exception("Linear simulator not supported for nonlinearity sensitivity.") + + elif self.synthetic_simulator_type == "nonlinear": + sim = SyntheticSimulatorModulatedNonLinear( + X_raw_train, + num_important_features=num_important_features, + non_linearity_scale=nonlinearity_scale, + seed=self.seed, + selection_type=self.selection_type + ) + else: + raise Exception("Unknown simulator type.") + + ( + X_train, + W_train, + Y_train, + po0_train, + po1_train, + propensity_train, + ) = sim.simulate_dataset( + X_raw_train, + predictive_scale=self.predictive_scale, + binary_outcome=binary_outcome, + treatment_assign=self.propensity_type, + ) + X_test, W_test, Y_test, po0_test, po1_test, _ = sim.simulate_dataset( + X_raw_test, + predictive_scale=self.predictive_scale, + binary_outcome=binary_outcome, + treatment_assign=self.propensity_type, + ) + + log.info("Fitting and explaining learners...") + learners = get_learners( + model_list=self.model_list, + X_train=X_train, + Y_train=Y_train, + n_iter=self.n_iter, + batch_size=self.batch_size, + batch_norm=False, + discrete_outcome=binary_outcome + ) + + learner_explanations = get_learner_explanations(learners, + X_test, X_train, Y_train, W_train, + self.explainer_limit, explainer_list) + + all_important_features = sim.get_all_important_features(with_selective=True) + pred_features = sim.get_predictive_features() + prog_features = sim.get_prognostic_features() + select_features = sim.get_selective_features() + + cate_test = sim.te(X_test) + + for explainer_name in explainer_list: + for learner_name in learners: + attribution_est = np.abs( + learner_explanations[learner_name][explainer_name] + ) + acc_scores_all_features = attribution_accuracy( + all_important_features, attribution_est + ) + acc_scores_predictive_features = attribution_accuracy( + pred_features, attribution_est + ) + acc_scores_prog_features = attribution_accuracy( + prog_features, attribution_est + ) + acc_scores_selective_features = attribution_accuracy( + select_features, attribution_est + ) + + + cate_pred = learners[learner_name].predict(X=X_test) + + pehe_test = compute_pehe(cate_true=cate_test, cate_pred=cate_pred) + + explainability_data.append( + [ + nonlinearity_scale, + learner_name, + explainer_name, + acc_scores_all_features, + acc_scores_predictive_features, + acc_scores_prog_features, + acc_scores_selective_features, + pehe_test, + np.mean(cate_test), + np.var(cate_test), + pehe_test / np.sqrt(np.var(cate_test)), + ] + ) + + metrics_df = pd.DataFrame( + explainability_data, + columns=[ + "Nonlinearity Scale", + "Learner", + "Explainer", + "All features ACC", + "Pred features ACC", + "Prog features ACC", + "Select features ACC", + "PEHE", + "CATE true mean", + "CATE true var", + "Normalized PEHE", + ], + ) + + results_path = ( + self.save_path + / f"results/nonlinearity_sensitivity/{self.synthetic_simulator_type}" + ) + log.info(f"Saving results in {results_path}...") + if not results_path.exists(): + results_path.mkdir(parents=True, exist_ok=True) + + metrics_df.to_csv( + results_path + / f"{dataset}_{num_important_features}_binary_{binary_outcome}-seed{self.seed}.csv" + ) + + +class PropensitySensitivity: + """ + Sensitivity analysis for confounding. + """ + + def __init__( + self, + n_units_hidden: int = 50, + n_layers: int = 1, + penalty_orthogonal: float = 0.01, + batch_size: int = 1024, + batch_norm: bool = False, + n_iter: int = 1000, + seed: int = 42, + explainer_limit: int = 1000, + save_path: Path = Path.cwd(), + num_interactions: int = 1, + synthetic_simulator_type: str = "linear", + nonlinearity_scale: float = 0, + selection_type: str = "random", + propensity_type: str = "pred", + propensity_scales: list = [0, 0.5, 1, 2, 5, 10], + model_list: list = ["TLearner"] + ) -> None: + + self.n_units_hidden = n_units_hidden + self.n_layers = n_layers + self.penalty_orthogonal = penalty_orthogonal + self.batch_size = batch_size + self.batch_norm = batch_norm + self.n_iter = n_iter + self.seed = seed + self.explainer_limit = explainer_limit + self.save_path = save_path + self.num_interactions = num_interactions + self.synthetic_simulator_type = synthetic_simulator_type + self.nonlinearity_scale = nonlinearity_scale + self.selection_type = selection_type + self.propensity_type = propensity_type + self.propensity_scales = propensity_scales + self.model_list = model_list + + def run( + self, + dataset: str = "tcga_10", + train_ratio: float = 0.8, + num_important_features: int = 2, + binary_outcome: bool = False, + random_feature_selection: bool = True, + predictive_scale: float = 1, + nonlinearity_scale: float = 0.5, + explainer_list: list = [ + "feature_ablation", + "feature_permutation", + "integrated_gradients", + "shapley_value_sampling", + ], + debug: bool = False, + directory_path_: str = None, + ) -> None: + log.info( + f"Using dataset {dataset} with num_important features = {num_important_features} and predictive scale {predictive_scale}." + ) + + X_raw_train, X_raw_test = load(dataset, train_ratio=train_ratio, debug=debug, directory_path_=directory_path_) + + if self.synthetic_simulator_type == "linear": + sim = SyntheticSimulatorLinear( + X_raw_train, + num_important_features=num_important_features, + random_feature_selection=random_feature_selection, + seed=self.seed, + ) + elif self.synthetic_simulator_type == "nonlinear": + sim = SyntheticSimulatorModulatedNonLinear( + X_raw_train, + num_important_features=num_important_features, + non_linearity_scale=self.nonlinearity_scale, + seed=self.seed, + selection_type=self.selection_type, + ) + else: + raise Exception("Unknown simulator type.") + + explainability_data = [] + + for propensity_scale in self.propensity_scales: + log.info(f"Now working with propensity_scale = {propensity_scale}...") + ( + X_train, + W_train, + Y_train, + po0_train, + po1_train, + propensity_train, + ) = sim.simulate_dataset( + X_raw_train, + predictive_scale=predictive_scale, + binary_outcome=binary_outcome, + treatment_assign=self.propensity_type, + prop_scale=propensity_scale, + ) + + X_test, W_test, Y_test, po0_test, po1_test, _ = sim.simulate_dataset( + X_raw_test, + predictive_scale=predictive_scale, + binary_outcome=binary_outcome, + treatment_assign=self.propensity_type, + prop_scale=propensity_scale, + ) + + + log.info("Fitting and explaining learners...") + learners = get_learners( + model_list=self.model_list, + X_train=X_train, + Y_train=Y_train, + n_iter=self.n_iter, + batch_size=self.batch_size, + batch_norm=self.batch_norm, + discrete_outcome=binary_outcome + ) + + learner_explanations = get_learner_explanations(learners, + X_test, X_train, Y_train, W_train, + self.explainer_limit, explainer_list) + + all_important_features = sim.get_all_important_features(with_selective=False) + pred_features = sim.get_predictive_features() + prog_features = sim.get_prognostic_features() + + cate_test = sim.te(X_test) + + for explainer_name in explainer_list: + for learner_name in learners: + attribution_est = np.abs( + learner_explanations[learner_name][explainer_name] + ) + acc_scores_all_features = attribution_accuracy( + all_important_features, attribution_est + ) + acc_scores_predictive_features = attribution_accuracy( + pred_features, attribution_est + ) + acc_scores_prog_features = attribution_accuracy( + prog_features, attribution_est + ) + + cate_pred = learners[learner_name].predict(X=X_test) + pehe_test = compute_pehe(cate_true=cate_test, cate_pred=cate_pred) + + explainability_data.append( + [ + propensity_scale, + learner_name, + explainer_name, + acc_scores_all_features, + acc_scores_predictive_features, + acc_scores_prog_features, + pehe_test, + np.mean(cate_test), + np.var(cate_test), + pehe_test / np.sqrt(np.var(cate_test)), + ] + ) + + metrics_df = pd.DataFrame( + explainability_data, + columns=[ + "Propensity Scale", + "Learner", + "Explainer", + "All features ACC", + "Pred features ACC", + "Prog features ACC", + "PEHE", + "CATE true mean", + "CATE true var", + "Normalized PEHE", + ], + ) + + results_path = ( + self.save_path + / f"results/propensity_sensitivity/{self.synthetic_simulator_type}" + ) + log.info(f"Saving results in {results_path}...") + if not results_path.exists(): + results_path.mkdir(parents=True, exist_ok=True) + + metrics_df.to_csv( + results_path / f"propensity_scale_{dataset}_{num_important_features}_" + f"proptype_{self.propensity_type}_" + f"predscl_{predictive_scale}_" + f"nonlinscl_{nonlinearity_scale}_" + f"trainratio_{train_ratio}_" + f"binary_{binary_outcome}-seed{self.seed}.csv" + ) + + +class CohortSizeSensitivity: + """ + Sensitivity analysis for varying numbers of samples. This experiment will generate a .csv with the recorded metrics. + It will also generate a gif, showing the progression on dimensionality-reduced spaces. + """ + + def __init__( + self, + n_units_hidden: int = 50, + n_layers: int = 1, + penalty_orthogonal: float = 0.01, + batch_size: int = 1024, + batch_norm: bool = False, + n_iter: int = 1000, + seed: int = 42, + explainer_limit: int = 1000, + save_path: Path = Path.cwd(), + propensity_type: str = "selective", + cohort_sizes: list = [0.5, 0.7, 1.0], + nonlinearity_scale: float = 0.5, + predictive_scale: float = 1, + synthetic_simulator_type: str = "random", + selection_type: str = "random", + model_list: list = ["TLearner"], + num_cube_samples: int = 1000, + dim_reduction_method: str = "umap", + dim_reduction_on_important_features: bool = True, + visualize_progression: bool = True, + ) -> None: + + self.n_units_hidden = n_units_hidden + self.n_layers = n_layers + self.penalty_orthogonal = penalty_orthogonal + self.batch_size = batch_size + self.batch_norm = batch_norm + self.n_iter = n_iter + self.seed = seed + self.explainer_limit = explainer_limit + self.save_path = save_path + self.propensity_type = propensity_type + self.cohort_sizes = cohort_sizes + self.nonlinearity_scale = nonlinearity_scale + self.predictive_scale = predictive_scale + self.synthetic_simulator_type = synthetic_simulator_type + self.selection_type = selection_type + self.model_list = model_list + self.num_cube_samples = num_cube_samples + self.dim_reduction_method = dim_reduction_method + self.dim_reduction_on_important_features = dim_reduction_on_important_features + self.visualize_progression = visualize_progression + + def run( + self, + dataset: str = "tcga_100", + num_important_features: int = 15, + explainer_list: list = [ + "feature_ablation", + "feature_permutation", + "integrated_gradients", + "shapley_value_sampling", + ], + train_ratio: float = 0.8, + binary_outcome: bool = False, + debug=False, + directory_path_: str = None, + + ) -> None: + # Log setting + log.info( + f"Using dataset {dataset} with num_important features = {num_important_features}." + ) + + # Load data + X_raw_train_full, X_raw_test_full = load(dataset, train_ratio=train_ratio, debug=debug, directory_path_=directory_path_) + explainability_data = [] + + # Simulate treatment and outcome for train and test + sim = SyntheticSimulatorModulatedNonLinear( + X_raw_train_full, + num_important_features=num_important_features, + non_linearity_scale=self.nonlinearity_scale, + seed=self.seed, + selection_type=self.selection_type + ) + + ( + X_train_full, + W_train_full, + Y_train_full, + po0_train_full, + po1_train_full, + propensity_train_full + ) = sim.simulate_dataset( + X_raw_train_full, + predictive_scale=self.predictive_scale, + binary_outcome=binary_outcome, + treatment_assign=self.propensity_type, + ) + + X_test_full, W_test_full, Y_test_full, po0_test_full, po1_test_full, _ = sim.simulate_dataset( + X_raw_test_full, + predictive_scale=self.predictive_scale, + binary_outcome=binary_outcome, + treatment_assign=self.propensity_type, + ) + + # Retrieve important features + all_important_features = sim.get_all_important_features(with_selective=True) + pred_features = sim.get_predictive_features() + prog_features = sim.get_prognostic_features() + select_features = sim.get_selective_features() + + # Code for sampling from hypercube -> does not work well because data lives in very small part of that hypercube + # # Sample from a hypercube grid for the X_raw_train_full dataset - making a complete hypercube grid would be too large + # # Sample for important features only as we know these are the ones that matter and will make sampling from a hypercube more meaningful + # X_raw_train_full_important = X_raw_train_full[:, all_important_features] + # min_vals = X_raw_train_full_important.min(axis=0) + # max_vals = X_raw_train_full_important.max(axis=0) + + # # Add some relative padding to the min and max values + # min_vals = min_vals - 0.1 * np.abs(min_vals) + # max_vals = max_vals + 0.1 * np.abs(max_vals) + + # # Sample points + # grid_samples = np.zeros((self.num_cube_samples, X_raw_train_full.shape[1])) + # grid_samples[:, all_important_features] = np.random.uniform(min_vals, max_vals, (self.num_cube_samples, X_raw_train_full_important.shape[1])) + + if self.visualize_progression: + # Use full training set with added noisy samples as focused samples of space + std_dev = np.std(X_raw_train_full, axis=0) + grid_samples = X_raw_train_full + grid_samples = np.vstack([grid_samples, + X_raw_train_full + std_dev*np.random.normal(0, 1, X_raw_train_full.shape), + X_raw_train_full + 0.1*std_dev*np.random.normal(0, 1, X_raw_train_full.shape), + X_raw_train_full + 0.1*std_dev*np.random.normal(0, 1, X_raw_train_full.shape), + # X_raw_train_full + 0.3*std_dev*np.random.normal(0, 1, X_raw_train_full.shape), + X_raw_train_full + 0.1*std_dev*np.random.normal(0, 1, X_raw_train_full.shape),]) + + # Reduce samples to two dimensions for plotting using umap + if self.dim_reduction_method == "umap": + reducer = umap.UMAP(min_dist=1, n_neighbors=30, spread=1) + reducer_shap = umap.UMAP(min_dist=3, n_neighbors=40, spread=4) + reducer_shap_prop = umap.UMAP(min_dist=2, n_neighbors=30, spread=3) + + elif self.dim_reduction_method == "pca": + reducer = PCA(n_components=2) + reducer_shap = PCA(n_components=2) + reducer_shap_prop = PCA(n_components=2) + + elif self.dim_reduction_method == "tsne": + raise Exception("t-SNE not supported for this analysis. Does not offer .transform() method.") + + else: + raise Exception("Unknown dimensionality reduction method.") + + # Fit on grid samples and training data + if self.dim_reduction_on_important_features: + grid_samples_2d = reducer.fit_transform(grid_samples[:, all_important_features]) + train_samples_2d = reducer.transform(X_raw_train_full[:, all_important_features]) + else: + grid_samples_2d = reducer.fit_transform(grid_samples) + train_samples_2d = reducer.transform(X_raw_train_full) + + # Get model learners (here only one) and explanations for grid samples + learners = get_learners( + model_list=self.model_list, + X_train=X_train_full, + Y_train=Y_train_full, + n_iter=self.n_iter, + batch_size=self.batch_size, + batch_norm=False, + discrete_outcome=binary_outcome + ) + + learner_explanations, learners = get_learner_explanations(learners, + grid_samples, X_train_full, Y_train_full, W_train_full, + grid_samples.shape[0], explainer_list, + return_learners=True) + + learner_explanations_train = get_learner_explanations(learners, + X_train_full, X_train_full, Y_train_full, W_train_full, + X_train_full.shape[0], explainer_list, + already_trained=True) + + # Get shap for grid samples + shap_values_grid = learner_explanations[self.model_list[0]][explainer_list[0]] + shap_values_train = learner_explanations_train[self.model_list[0]][explainer_list[0]] + + # Perform logistic regression for propensity + est_prop = LogisticRegression().fit(X_train_full, W_train_full) + explainer_prop = shap.LinearExplainer(est_prop, X_train_full) + shap_values_prop_grid = explainer_prop.shap_values(grid_samples) + shap_values_prop_train = explainer_prop.shap_values(X_train_full) + + # Fit umap reducer for shap grid samples + if self.dim_reduction_on_important_features: + shap_values_grid_2d = reducer_shap.fit_transform(shap_values_grid[:, all_important_features]) + shap_values_train_2d = reducer_shap.transform(shap_values_train[:, all_important_features]) + shap_values_prop_grid_2d = reducer_shap_prop.fit_transform(shap_values_prop_grid[:, all_important_features]) + shap_values_prop_train_2d = reducer_shap_prop.transform(shap_values_prop_train[:, all_important_features]) + else: + shap_values_grid_2d = reducer_shap.fit_transform(shap_values_grid) + shap_values_train_2d = reducer_shap.transform(shap_values_train) + shap_values_prop_grid_2d = reducer_shap_prop.fit_transform(shap_values_prop_grid) + shap_values_prop_train_2d = reducer_shap_prop.transform(shap_values_prop_train) + + # Initialize variable for storing frames + frames = [] # To store each frame for the GIF + + cohort_size_full = X_train_full.shape[0] + for cohort_size_perc in self.cohort_sizes: + cohort_size = int(cohort_size_perc * cohort_size_full) + + # Get a subset of the training data + X_train = X_train_full[:cohort_size] + W_train = W_train_full[:cohort_size] + Y_train = Y_train_full[:cohort_size] + po0_train = po0_train_full[:cohort_size] + po1_train = po1_train_full[:cohort_size] + propensity_train = propensity_train_full[:cohort_size] + cohort_size_train = X_train.shape[0] + + # Get subsets for test data + X_test = X_test_full[:cohort_size] + W_test = W_test_full[:cohort_size] + Y_test = Y_test_full[:cohort_size] + po0_test = po0_test_full[:cohort_size] + po1_test = po1_test_full[:cohort_size] + + log.info(f"Now working with a cohort size of {cohort_size}/{X_train_full.shape[0]}...") + log.info("Fitting and explaining learners...") + learners = get_learners( + model_list=self.model_list, + X_train=X_train, + Y_train=Y_train, + n_iter=self.n_iter, + batch_size=self.batch_size, + batch_norm=False, + discrete_outcome=binary_outcome + ) + + # Get learners and explanations for training data + learner_explanations, learners = get_learner_explanations(learners, + X_train, X_train, Y_train, W_train, + X_train.shape[0], explainer_list, + return_learners=True) + + if self.visualize_progression: + # Make logistic regression for propensity + est_prop = LogisticRegression().fit(X_train, W_train) + explainer_prop = shap.LinearExplainer(est_prop, X_train) + + # Set up plot + fig, axs = plt.subplots(2, 4, figsize=(15, 10)) + eff_grid = sim.te(grid_samples) + prop_grid = sim.prop(grid_samples) + + # Get cate estimator + est_eff = learners[self.model_list[0]] + cate_pred_train = est_eff.predict(X=X_train) + + # Get predictions for the first model + p_eff_grid = est_eff.predict(X=grid_samples) + p_prop_grid = est_prop.predict_proba(grid_samples)[:, 1] + outcomes = [p_prop_grid, prop_grid, p_eff_grid, eff_grid] + titles = ['prop(x)', 'prop_true(x)', 'cate(x)', 'cate_true(x)'] + + # # Get X_Train in 2d + # shap_values_train = learner_explanations[self.model_list[0]][explainer_list[0]] + # shap_values_prop_train = explainer_prop.shap_values(X_train) + + # if self.dim_reduction_on_important_features: + # X_train_2d = reducer.transform(X_train[:, all_important_features]) + # else: + # X_train_2d = reducer.transform(X_train) + + # # Get shap values of X_Train in 2d + # if self.dim_reduction_on_important_features: + # shap_values_train_2d = reducer_shap.transform(shap_values_train[:, all_important_features]) + # shap_values_prop_train_2d = reducer_shap_prop.transform(shap_values_prop_train[:, all_important_features]) + # else: + # shap_values_train_2d = reducer_shap.transform(shap_values_train) + # shap_values_prop_train_2d = reducer_shap_prop.transform(shap_values_prop_train) + + + # If there are any non-finite elements in the outcomes, remove them from the grid_samples and the other outcomes and raise a warning + for i, outcome in enumerate(outcomes): + if type(outcome) == torch.Tensor: + outcome = outcome.cpu().detach().numpy() + + outcome = np.array(outcome) + if not np.all(np.isfinite(outcome)): + print("-----------") + print(i, np.sum(~np.isfinite(outcome))) + log.warning(f'Found non-finite elements in outcomes. Removing {np.sum(~np.isfinite(outcome))} elements.') + mask = np.isfinite(outcome) + grid_samples_2d = grid_samples_2d[mask] + for j in range(len(outcomes)): + outcomes[j] = outcomes[j][mask] + break + + for j, outcome in enumerate(outcomes): + if type(outcome) == torch.Tensor: + outcome = outcome.cpu().detach().numpy() + + # Plot settings + cmap = "viridis" + s = 15 # 4 for many samples + alpha = None + edgecolors = "w" + linewidths = 0.2 + + # Plot contours + if j == 0 or j == 1: + tcf = axs[0][j].tricontourf(grid_samples_2d[:,0], + grid_samples_2d[:,1], + outcome.ravel(), 15, cmap=cmap, levels=50) + + tcf_shap = axs[1][j].tricontourf(shap_values_prop_grid_2d[:,0], + shap_values_prop_grid_2d[:,1], + outcome.ravel(), 15, cmap=cmap, levels=50) + else: + tcf = axs[0][j].tricontourf(grid_samples_2d[:,0], + grid_samples_2d[:,1], + outcome.ravel(), 15, cmap=cmap, levels=50) + + tcf_shap = axs[1][j].tricontourf(shap_values_grid_2d[:,0], + shap_values_grid_2d[:,1], + outcome.ravel(), 15, cmap=cmap, levels=50) + + + # Version: Plot in shape space from current model + + # if j == 0: + # axs[0][j].scatter(X_train_2d[:,0], X_train_2d[:,1], c=W_train, cmap='coolwarm', edgecolors=edgecolors, s=s, label='Training data for a', alpha=alpha) + # axs[1][j].scatter(shap_values_prop_train_2d[:,0], shap_values_prop_train_2d[:,1], c=W_train, cmap='coolwarm', edgecolors=edgecolors, s=s, alpha=alpha) + # #fig.colorbar(tcf) + # if j == 2: + # axs[0][j].scatter(X_train_2d[:,0], X_train_2d[:,1], c=cate_pred_train, cmap='coolwarm', edgecolors=edgecolors, s=s, label='Training data for a', alpha=alpha) + # axs[1][j].scatter(shap_values_train_2d[:,0], shap_values_train_2d[:,1], c=Y_train, cmap='coolwarm', edgecolors=edgecolors, s=s, alpha=alpha) + # #fig.colorbar(tcf_shap) + + # Version: Always plot in same shap space from full model + + if j == 0: + axs[0][j].scatter(train_samples_2d[:cohort_size_train,0], train_samples_2d[:cohort_size_train,1], + c=W_train, cmap=cmap, s=s, edgecolors=edgecolors, linewidths=linewidths, alpha=0.5) + + axs[1][j].scatter(shap_values_prop_train_2d[:cohort_size_train,0], shap_values_prop_train_2d[:cohort_size_train,1], + c=W_train, cmap=cmap, s=s, edgecolors=edgecolors, linewidths=linewidths, alpha=0.5) + #fig.colorbar(tcf) + if j == 2: + axs[0][j].scatter(train_samples_2d[:cohort_size_train,0], train_samples_2d[:cohort_size_train,1], + c=cate_pred_train, cmap=cmap, s=s, edgecolors=edgecolors, linewidths=linewidths) + + axs[1][j].scatter(shap_values_train_2d[:cohort_size_train,0], shap_values_train_2d[:cohort_size_train,1], + c=cate_pred_train, cmap=cmap, s=s, edgecolors=edgecolors, linewidths=linewidths) + #fig.colorbar(tcf_shap) + + + # axs[0][j].set_title(f'{titles[j]} (N={cohort_size})_data_space') + # axs[0][j].set_xlim([grid_samples_2d[:,0].min(), grid_samples_2d[:,0].max()]) + # axs[0][j].set_ylim([grid_samples_2d[:,1].min(), grid_samples_2d[:,1].max()]) + # axs[0][j].legend() + + # if j == 0 or j == 1: + # axs[1][j].set_title(f'{titles[j]} (N={cohort_size})_shap_prop_space') + # axs[1][j].set_xlim([shap_values_prop_grid_2d[:,0].min(), shap_values_prop_grid_2d[:,0].max()]) + # axs[1][j].set_ylim([shap_values_prop_grid_2d[:,1].min(), shap_values_prop_grid_2d[:,1].max()]) + # axs[1][j].legend() + # else: + # axs[1][j].set_title(f'{titles[j]} (N={cohort_size})_shap_space') + # axs[1][j].set_xlim([shap_values_grid_2d[:,0].min(), shap_values_grid_2d[:,0].max()]) + # axs[1][j].set_ylim([shap_values_grid_2d[:,1].min(), shap_values_grid_2d[:,1].max()]) + # axs[1][j].legend() + + # Save the plot to a buffer + plt.tight_layout() + plt.savefig('temp_plot.png') + plt.close() + frames.append(imageio.imread('temp_plot.png')) + + for explainer_name in explainer_list: + for learner_name in learners: + attribution_est = np.abs( + learner_explanations[learner_name][explainer_name] + ) + acc_scores_all_features = attribution_accuracy( + all_important_features, attribution_est + ) + acc_scores_predictive_features = attribution_accuracy( + pred_features, attribution_est + ) + acc_scores_prog_features = attribution_accuracy( + prog_features, attribution_est + ) + acc_scores_selective_features = attribution_accuracy( + select_features, attribution_est + ) + + + cate_pred = learners[learner_name].predict(X=X_test) + cate_test = sim.te(X_test) + pehe_test = compute_pehe(cate_true=cate_test, cate_pred=cate_pred) + + explainability_data.append( + [ + cohort_size, + cohort_size_perc, + self.nonlinearity_scale, + learner_name, + explainer_name, + acc_scores_all_features, + acc_scores_predictive_features, + acc_scores_prog_features, + acc_scores_selective_features, + pehe_test, + np.mean(cate_test), + np.var(cate_test), + pehe_test / np.sqrt(np.var(cate_test)), + ] + ) + + metrics_df = pd.DataFrame( + explainability_data, + columns=[ + "Cohort Size", + "Cohort Size Perc", + "Nonlinearity Scale", + "Learner", + "Explainer", + "All features ACC", + "Pred features ACC", + "Prog features ACC", + "Select features ACC", + "PEHE", + "CATE true mean", + "CATE true var", + "Normalized PEHE", + ], + ) + + results_path = ( + self.save_path + / f"results/cohort_size_sensitivity/{self.synthetic_simulator_type}" + ) + + log.info(f"Saving results in {results_path}...") + if not results_path.exists(): + results_path.mkdir(parents=True, exist_ok=True) + + metrics_df.to_csv( + results_path + / f"{dataset}_{num_important_features}_binary_{binary_outcome}-seed{self.seed}.csv" + ) + + if self.visualize_progression: + # Create GIF + imageio.mimsave("progression.gif", frames, fps = 1) + imageio.mimsave(results_path / "progression.gif", frames, fps=1) # Set fps=1 for slower transition to observe changes clearly +