updates in response to reviewers

Author: wongdanr

analysis.py

@@ -5,9 +5,9 @@
 import pickle
 from scgpt.utils import compute_perturbation_metrics
 
-def run_gears(runs=1, mode="train"):
+def run_gears(runs=1, mode="train", cross_validation=False):
     for run_number in range(0, runs):
-        for data_name in ["adam_corrected", "adam_corrected_upr", "adamson", "norman", "replogle_k562_essential"]:
+        for data_name in ["adam_corrected_upr", "norman", "replogle_k562_essential", "adam_corrected", "adamson"]:
             ##setup PertData object
             if data_name == "adam_corrected":
                 pert_data = get_adam_corrected_dataset(split="simulation", batch_size=64, test_batch_size=64, generate_new=False, just_upr=False)
@@ -26,28 +26,42 @@ def run_gears(runs=1, mode="train"):
                 pert_data.get_dataloader(batch_size = 64, test_batch_size = 64)
             if "replogle" in data_name:
                 modify_pertdata_dataloaders(pert_data, logger=None)
+            if cross_validation:
+                cross_validate_split(pert_data, run_number + 1)
+                prefix = "pickles/gears_results_cv/"
+                suffix = f"cross_val_{run_number + 1}"
+            else:
+                prefix = "pickles/gears_results/"
+                suffix = f"{run_number}"
             gears_model = GEARS(pert_data, device = 'cuda:0')
             if mode == "train":
                 # set up and train a model
                 gears_model.model_initialize(hidden_size = 64)
                 gears_model.train(epochs = 20) #20 originally
-                gears_model.save_model(f'gears_models/gears_trained_{data_name}_{run_number}')
+                print("finished training, save model")
+                gears_model.save_model(f'gears_models/gears_trained_{data_name}_{suffix}')
             #load model
-            gears_model.load_pretrained(f'gears_models/gears_trained_{data_name}_{run_number}')
+            gears_model.load_pretrained(f'gears_models/gears_trained_{data_name}_{suffix}')
             ##evaluate
             eval_results = evaluate(loader=pert_data.dataloader['test_loader'], model=gears_model.model, uncertainty=gears_model.config['uncertainty'], device=torch.device("cuda:0"))
             ##get rank score
             ranks = get_gears_rank(eval_results)
             print("avg rank: ", np.mean(list(ranks.values())), np.std(list(ranks.values())))
-            pickle.dump(ranks, open(f"pickles/gears_results/gears_rank_metrics_{data_name}_{run_number}.pkl", "wb"))
+            pickle.dump(ranks, open(f"{prefix}gears_rank_metrics_{data_name}_{suffix}.pkl", "wb"))
             ##get pearson scores
             metrics, metrics_pert = compute_metrics(eval_results)
             test_metrics = compute_perturbation_metrics(eval_results, pert_data.adata[pert_data.adata.obs["condition"] == "ctrl"])
             print(f"metrics: {metrics}")
             print(f"metrics_pert: {metrics_pert}")
             print(f"test metrics: {test_metrics}")
-            pickle.dump((metrics, metrics_pert), open(f"pickles/gears_results/gears_results_{data_name}_{run_number}.pkl", "wb"))
-            pickle.dump(test_metrics, open(f"pickles/gears_results/gears_pert_delta_results_{data_name}_{run_number}.pkl", "wb"))
+            pickle.dump((metrics, metrics_pert), open(f"{prefix}gears_results_{data_name}_{suffix}.pkl", "wb"))
+            pickle.dump(test_metrics, open(f"{prefix}gears_pert_delta_results_{data_name}_{suffix}.pkl", "wb"))
+            ##condition specific performance 
+            condition_map = get_condition_performance_breakdown(eval_results, pert_data.adata[pert_data.adata.obs["condition"] == "ctrl"]) 
+            pickle.dump(condition_map, open(f"{prefix}gears_condition_specific_results_{data_name}_{suffix}.pkl", "wb"))
+            ##gene specific performance 
+            gene_to_pearson_map = get_gene_performance_breakdown(eval_results, pert_data.adata[pert_data.adata.obs["condition"] == "ctrl"]) 
+            pickle.dump(gene_to_pearson_map, open(f"{prefix}gears_gene_specific_results_{data_name}_{suffix}.pkl", "wb"))
 
 def get_gears_rank(eval_results):
     pert_map = {} ##key: condition, value: (actual avg truth vector, predicted avg vector)
@@ -62,6 +76,13 @@ def get_gears_rank(eval_results):
     return ranks
 
 print("running gears")
+##no cross validation 
 ##mode = "train" for training or "eval" for just evaluating models
-run_gears(runs=10, mode="train")
+run_gears(runs=10, mode="train", cross_validation=False)
+run_gears(runs=10, mode="eval", cross_validation=False)
+
+##using cross validation 
+##if cross_validate, then the fold will be the run_number + 1
+run_gears(runs=4, mode="train", cross_validation=True)
+
 

gears_runner.py

@@ -332,6 +332,59 @@ def eval_perturb(
     results["avg_loss"] = total_loss / float(len(loader.dataset))
     return results
 
+def get_condition_performance_breakdown(results, ctrl_adata):
+    """
+    Given output from eval_perturb, compute pertubation specific performance 
+    """
+    from scgpt.utils import find_DE_genes
+    mean_ctrl = np.array(ctrl_adata.X.mean(0)).flatten() 
+    conditions = np.unique(results["pert_cat"])
+    geneid2idx = dict(zip(ctrl_adata.var.index.values, range(len(ctrl_adata.var))))
+    de_idx = {c: find_DE_genes(ctrl_adata, c, geneid2idx, non_zero_genes=False)[0] for c in conditions}
+    condition2idx = {c: np.where(results["pert_cat"] == c)[0] for c in conditions} #condition to indices where condition occurs in conditions array 
+    condition_map = {}
+    for pert in conditions: 
+        pert_indices = condition2idx[pert]
+        pert_de_idx = de_idx[pert]
+        pert_pred = np.mean(results["pred"][pert_indices, :], axis=0)
+        pert_truth = np.mean(results["truth"][pert_indices, :], axis=0)
+        if np.sum(pert_pred[pert_de_idx]) == 0 or np.sum(pert_truth[pert_de_idx]) == 0: ##vector of all zeros will result in NaN pearson_de, skip 
+            print(f"WARNING: {pert} has 0 vector, will result in NaN pearson_de")
+        pearson = scipy.stats.pearsonr(pert_pred, pert_truth)[0]
+        pearson_delta =  scipy.stats.pearsonr(pert_pred - mean_ctrl , pert_truth - mean_ctrl)[0]
+        pearson_de = scipy.stats.pearsonr(pert_pred[pert_de_idx], pert_truth[pert_de_idx])[0]
+        pearson_de_delta = scipy.stats.pearsonr(pert_pred[pert_de_idx] - mean_ctrl[pert_de_idx], pert_truth[pert_de_idx] - mean_ctrl[pert_de_idx])[0]
+        condition_map[pert] = {"pearson": pearson, "pearson_de": pearson_de, "pearson_delta": pearson_delta, "pearson_de_delta": pearson_de_delta}
+    print("average results in get_condition_performance_breakdown: ")
+    for metric in ["pearson", "pearson_de", "pearson_delta", "pearson_de_delta"]:
+        print("    ", metric, np.nanmean([condition_map[pert][metric] for pert in condition_map]))
+    return condition_map
+
+def get_gene_performance_breakdown(results, ctrl_adata):
+    """
+    Given output from eval_perturb, compute gene specific performance
+    pearson between actual and predicted for each gene
+    note: no such concept as delta scores because pearson (x,y) == pearson (x-k, y-k) for constant k and vectors x,y
+    """
+    pred = results["pred"]
+    truth = results["truth"]
+    gene_list = ctrl_adata.var["gene_name"].tolist()
+    assert(len(gene_list) == len(pred[0]) == len(truth[0]))
+    gene_to_pearson_map = {gene: "" for gene in gene_list}
+    for i in range(0, len(gene_list)):
+        gene = gene_list[i]
+        if np.std(pred[:,i]) < 0.000001: ##mode collapse for models like mean, if std == 0 we cannot compute pearson (will be NaN), add very small random noise to prediction
+            pred_vector = pred[:,i] + np.random.rand(len(pred[:, i])) * 0.0000001
+        else: 
+            pred_vector = pred[:,i]
+        if np.std(pred[:,i]) < 0.000001:
+            truth_vector = truth[:,i] + np.random.rand(len(truth[:, i])) * 0.0000001
+        else:
+            truth_vector = truth[:,i]
+        corr_i = scipy.stats.pearsonr(pred_vector, truth_vector)[0]
+        gene_to_pearson_map[gene] = corr_i
+    return gene_to_pearson_map    
+
 def get_variables(load_model=None, config_path=None):
     """
     Reads config file and returns dictionary of variables
@@ -555,9 +608,88 @@ def modify_pertdata_dataloaders(pert_data, logger=None):
         for load_type in ["train", "val", "test"]:
             logger.info(f"    new {load_type} loader length: {len(pert_data.dataloader[f'{load_type}_loader'])}")
 
+def get_split(pert, pert_map):
+    """
+    Given a perturbation pert 
+    and a dictionary with key: split, value: list of perts
+    will return which split pert is found in
+    """
+    for split in pert_map:
+        if pert in pert_map[split]:
+            return split
+    return -1
+
+def cross_validate_split(pert_data, cross_validation_fold):
+    """
+    will modify PertData loaders to conform to the cross_validation fold 
+    4-fold, will have two be training, 1 be val, and 1 be test
+    each split has unique perturbations (minus control, which will be in just train)
+    """
+    print(f"WARNING: splitting data into cross validation fold {cross_validation_fold}")
+    ##get all perturbations in train/val/test, sort and then shuffle them by fixed seed so deterministic
+    old_dataloaders = pert_data.dataloader
+    all_perts = []
+    splits = ["train", "val", "test"]
+    for load_type in splits:
+        old_loader = old_dataloaders[f"{load_type}_loader"]
+        for batch, batch_data in enumerate(old_loader):
+            all_perts = all_perts + list(batch_data.pert)
+    all_perts = sorted(list(set(all_perts)))
+    all_perts.remove("ctrl") ##let's add ctrl back later to just the train perturbations
+    ##deterministically shuffle all_perts
+    g = torch.Generator()
+    g.manual_seed(0)
+    rand_indices = torch.randperm(len(all_perts), generator=g).tolist()
+    shuffled_perts = [all_perts[rand_index] for rand_index in rand_indices]
+    print(shuffled_perts)
+    ##chunk the list into folds
+    divisor = int(len(shuffled_perts) / 4)
+    chunk_1 = shuffled_perts[0: divisor]
+    chunk_2 = shuffled_perts[divisor: divisor * 2]
+    chunk_3 = shuffled_perts[divisor * 2: divisor * 3]
+    chunk_4 = shuffled_perts[divisor * 3: ]
+    ##assign train/val/test depending on fold
+    pert_map = {split: set() for split in splits}
+    if cross_validation_fold == 1:
+        pert_map["train"] = chunk_1 + chunk_2
+        pert_map["val"] = chunk_3
+        pert_map["test"] = chunk_4
+    if cross_validation_fold == 2:
+        pert_map["train"] = chunk_4 + chunk_1
+        pert_map["val"] = chunk_2
+        pert_map["test"] = chunk_3
+    if cross_validation_fold == 3:
+        pert_map["train"] = chunk_3 + chunk_4
+        pert_map["val"] = chunk_1
+        pert_map["test"] = chunk_2
+    if cross_validation_fold == 4:
+        pert_map["train"] = chunk_2 + chunk_3
+        pert_map["val"] = chunk_4
+        pert_map["test"] = chunk_1
+    pert_map["train"].append("ctrl")
+    ##now create new loaders and assign
+    new_data_map = {split: [] for split in splits}
+    for load_type in splits:
+        new_data = []
+        skipped = set()
+        old_loader = old_dataloaders[f"{load_type}_loader"]
+        for batch, batch_data in enumerate(old_loader): ##batch_data is of type torch_geometric.data.batch.DataBatch, batch_data[i] is of type torch_geometric.data.data.Data
+            for i in range(0, len(batch_data)):  
+                pert = batch_data.pert[i]  
+                my_split = get_split(pert, pert_map)  
+                new_data_map[my_split].append(batch_data[i])  
+    shuffle = {"train": True, "val":True, "test":False}
+    new_dataloaders = {}
+    for split in splits: 
+        new_loader = DataLoader(new_data_map[split], batch_size=old_loader.batch_size, shuffle=shuffle[load_type])
+        new_dataloaders[f"{split}_loader"] = new_loader 
+    pert_data.dataloader = new_dataloaders
+
 def check_args(opt):
     if opt.pretrain_control == True and opt.mode == "test":
         raise Exception("opt.pretrain_control == True and opt.mode == test")
+    if opt.cross_validation == True:
+        assert(opt.cross_validation_fold != None)
 
 def convert_ENSG_to_gene(ensg_list):
     """

library.py

@@ -22,10 +22,13 @@ def main():
     parser.add_argument("--use_lora", type=bool_flag, default=False, help="True if we want to use LoRa for finetuning")
     parser.add_argument("--lora_rank", type=int, default=8, help="if use_lora, specifies the inner dimension of the low-rank matrices to train")
     parser.add_argument("--config_path", type=str, default="config/default_config.json", help="path to JSON configuration file to use for setting up model")
-    parser.add_argument("--model_type", type=str, default="scGPT", help="scGPT, simple_affine, mean_control, mean_perturbed, mean_control+perturbed, smart_mean_control, smart_mean_perturbed, smart_mean_control+perturbed")
+    parser.add_argument("--model_type", type=str, default="scGPT", help="scGPT, simple_affine, simple_affine_large, mean_control, mean_perturbed, mean_control+perturbed, smart_mean_control, smart_mean_perturbed, smart_mean_control+perturbed")
     parser.add_argument("--validation_selection", type=str, default="pearson", help="how to select the best model during training, if 'pearson' will be by pearson correlation between predicted and actual expression over validation set, if 'loss' will be by minimal loss")
     parser.add_argument("--loss_type", type=str, default="mse", help="mse, mse+triplet, mse+pearson")
     parser.add_argument("--fixed_seed", type=bool_flag, default=True, help="True if we want to use a constant fixed seed")
+    parser.add_argument("--cross_validation", type=bool_flag, default=False, help="if True will cross validate instead of single random split, --cross_validation_fold arg must also be set")
+    parser.add_argument("--cross_validation_fold", type=int, default=None, help="which fold to train")
+
     opt = parser.parse_args()
     check_args(opt)
     matplotlib.rcParams["savefig.transparent"] = False
@@ -70,7 +73,10 @@ def main():
         logger.info("WARNING: filtering dataloaders! but keeping pert_data.adata the same")
         modify_pertdata_dataloaders(pert_data, logger)
 
-    check_pert_split(opt.data_name, pert_data)
+    if opt.cross_validation:
+        cross_validate_split(pert_data, opt.cross_validation_fold)
+    else:
+        check_pert_split(opt.data_name, pert_data)
 
     logger.info(f"adata.obs: {pert_data.adata.obs}")
     logger.info(f"|conditions|: {len(set(pert_data.adata.obs['condition']))}")
@@ -103,7 +109,7 @@ def main():
     ##mean predictor models - compute after data loaders are set (after a possible application of filter_perturbations)
     if opt.mode in ["train", "test", "analysis"]:
         for baseline in ["smart", "baseline"]:
-            for mean_type in ["perturbed", "control", "control+perturbed"]:
+            for mean_type in ["perturbed"]:#, "control", "control+perturbed"]:
                 ##baseline mean 
                 if baseline == "baseline":
                     mean_pred_model = MeanPredictor(pert_data, opt.data_name, mean_type=mean_type)
@@ -115,7 +121,8 @@ def main():
                 ##get rank metrics 
                 test_perts = pickle.load(open(f"pickles/{opt.data_name}_perturbation_splits.pkl", "rb"))["test"]   
                 ranks = get_rank(mean_pred_model, test_perts, pert_data=pert_data, var=var, gene_ids=gene_ids, gene_idx_map=gene_idx_map)
-                pickle.dump(ranks, open(rank_save, "wb"))
+                if opt.cross_validation == False: ##save ranks to global pickles/ only on main split 
+                    pickle.dump(ranks, open(rank_save, "wb"))
                 ##GEARS-type metrics
                 mean_res = eval_perturb(pert_data.dataloader["test_loader"], mean_pred_model, gene_ids=[], gene_idx_map={}, var={"device":"cpu"}, loss_type=opt.loss_type) ##keep on cpu, no need to shuttle to gpu for mean pred model 
                 mean_metrics, mean_metrics_pert  = compute_metrics(mean_res) ##from GEARS library
@@ -125,6 +132,12 @@ def main():
                 mean_metrics = compute_perturbation_metrics(mean_res, pert_data.adata[pert_data.adata.obs["condition"] == "ctrl"]) ##from scGPT library
                 logger.info(f"{opt.data_name} {baseline} mean {mean_type} delta test metrics: {mean_metrics}")
                 pickle.dump(mean_metrics, open(save_dir / f"{baseline}_mean_{mean_type}_pert_delta_results_{opt.data_name}.pkl", "wb"))
+                ##condition specific performance 
+                condition_map = get_condition_performance_breakdown(mean_res, pert_data.adata[pert_data.adata.obs["condition"] == "ctrl"]) 
+                pickle.dump(condition_map, open(save_dir / f"{baseline}_mean_{mean_type}_condition_specific_results_{opt.data_name}.pkl", "wb"))
+                ##gene specific performance 
+                gene_to_pearson_map = get_gene_performance_breakdown(mean_res, pert_data.adata[pert_data.adata.obs["condition"] == "ctrl"]) 
+                pickle.dump(gene_to_pearson_map, open(save_dir / f"{baseline}_mean_{mean_type}_gene_specific_results_{opt.data_name}.pkl", "wb"))
 
     if opt.model_type == "scGPT": 
         model = TransformerGenerator(
@@ -142,9 +155,10 @@ def main():
             pert_pad_id=var["pert_pad_id"],
             use_fast_transformer=var["use_fast_transformer"],
         )
-    
-    elif opt.model_type == "simple_affine":
+    elif "simple_affine" in opt.model_type:
         from simple_affine import SimpleAffine 
+        is_large = True if "large" in opt.model_type else False
+        print("LARGE: ", is_large)
         model = SimpleAffine(
             ntoken=ntokens,
             d_model=var["embsize"],
@@ -154,6 +168,7 @@ def main():
             dropout=var["dropout"],
             pad_token=var["pad_token"],
             pert_pad_id=var["pert_pad_id"],
+            is_large=is_large
         )
     elif "mean" in opt.model_type:
         if "smart" in opt.model_type:
@@ -165,7 +180,7 @@ def main():
     else:
         raise Exception("model_type must be one of scGPT, simple_affine, mean_control, mean_perturbed, mean_control+perturbed, smart_mean_control, smart_mean_perturbed, smart_mean_control+perturbed")
 
-    if opt.model_type in ["scGPT", "simple_affine"]:
+    if opt.model_type in ["scGPT", "simple_affine", "simple_affine_large"]:
         model = load_model(var, model, model_file, logger, attention_control=opt.attention_control, freeze_input_encoder=opt.freeze_input_encoder, freeze_transformer_encoder=opt.freeze_transformer_encoder, mode=opt.mode, use_lora=opt.use_lora, lora_rank=opt.lora_rank, pretrain_control=opt.pretrain_control, transformer_encoder_control=opt.transformer_encoder_control, input_encoder_control=opt.input_encoder_control)
         model.to(var["device"])
 
@@ -244,6 +259,11 @@ def main():
         test_metrics = compute_perturbation_metrics(test_res, pert_data.adata[pert_data.adata.obs["condition"] == "ctrl"]) ##from scGPT utils library
         logger.info(f"{opt.data_name} delta test metrics: {test_metrics}")
         pickle.dump(test_metrics, open(save_dir / f"{opt.model_type}_pert_delta_results_{opt.data_name}.pkl", "wb"))
+        ##condition specific performance 
+        condition_map = get_condition_performance_breakdown(test_res, pert_data.adata[pert_data.adata.obs["condition"] == "ctrl"]) 
+        pickle.dump(condition_map, open(save_dir / f"{opt.model_type}_condition_specific_results_{opt.data_name}.pkl", "wb"))
+        gene_to_pearson_map = get_gene_performance_breakdown(test_res, pert_data.adata[pert_data.adata.obs["condition"] == "ctrl"]) 
+        pickle.dump(gene_to_pearson_map, open(save_dir / f"{opt.model_type}_gene_specific_results_{opt.data_name}.pkl", "wb"))
 
     if opt.mode == "analysis":    
         for plot_type in ["boxplots", "scatterplots"]: