| 974 | self.node_map_pert = {x: it for it, x in enumerate(self.pert_names)} |
| 975 | |
| 976 | def load(self, data_name = None, data_path = None): |
| 977 | if data_name in ['norman', 'adamson', 'dixit', |
| 978 | 'replogle_k562_essential', |
| 979 | 'replogle_rpe1_essential']: |
| 980 | data_path = os.path.join(self.data_path, data_name) |
| 981 | #zip_data_download_wrapper(url, data_path, self.data_path) |
| 982 | self.dataset_name = data_path.split('/')[-1] |
| 983 | self.dataset_path = data_path |
| 984 | adata_path = os.path.join(data_path, 'perturb_processed.h5ad') |
| 985 | self.adata = sc.read_h5ad(adata_path) |
| 986 | |
| 987 | elif os.path.exists(data_path): |
| 988 | adata_path = os.path.join(data_path, 'perturb_processed.h5ad') |
| 989 | self.adata = sc.read_h5ad(adata_path) |
| 990 | self.dataset_name = data_path.split('/')[-1] |
| 991 | self.dataset_path = data_path |
| 992 | else: |
| 993 | raise ValueError("data attribute is either norman, adamson, dixit " |
| 994 | "replogle_k562 or replogle_rpe1 " |
| 995 | "or a path to an h5ad file") |
| 996 | |
| 997 | self.set_pert_genes() |
| 998 | print_sys('These perturbations are not in the GO graph and their ' |
| 999 | 'perturbation can thus not be predicted') |
| 1000 | not_in_go_pert = np.array(self.adata.obs[ |
| 1001 | self.adata.obs.condition.apply( |
| 1002 | lambda x:not filter_pert_in_go(x, |
| 1003 | self.pert_names))].condition.unique()) |
| 1004 | print_sys(not_in_go_pert) |
| 1005 | |
| 1006 | filter_go = self.adata.obs[self.adata.obs.condition.apply( |
| 1007 | lambda x: filter_pert_in_go(x, self.pert_names))] |
| 1008 | self.adata = self.adata[filter_go.index.values, :] |
| 1009 | pyg_path = os.path.join(data_path, 'data_pyg') |
| 1010 | if not os.path.exists(pyg_path): |
| 1011 | os.mkdir(pyg_path) |
| 1012 | dataset_fname = os.path.join(pyg_path, 'cell_graphs.pkl') |
| 1013 | |
| 1014 | if os.path.isfile(dataset_fname): |
| 1015 | print_sys("Local copy of pyg dataset is detected. Loading...") |
| 1016 | self.dataset_processed = pickle.load(open(dataset_fname, "rb")) |
| 1017 | print_sys("Done!") |
| 1018 | else: |
| 1019 | self.ctrl_adata = self.adata[self.adata.obs['condition'] == 'ctrl'] |
| 1020 | self.gene_names = self.adata.var.gene_name |
| 1021 | |
| 1022 | |
| 1023 | print_sys("Creating pyg object for each cell in the data...") |
| 1024 | self.create_dataset_file() |
| 1025 | print_sys("Saving new dataset pyg object at " + dataset_fname) |
| 1026 | pickle.dump(self.dataset_processed, open(dataset_fname, "wb")) |
| 1027 | print_sys("Done!") |
| 1028 | |
| 1029 | |
| 1030 | def prepare_split(self, split = 'simulation', |