GEARS base model class
| 226 | return torch.stack(out) |
| 227 | |
| 228 | class GEARS: |
| 229 | """ |
| 230 | GEARS base model class |
| 231 | """ |
| 232 | |
| 233 | def __init__(self, pert_data, |
| 234 | device = 'cuda', |
| 235 | weight_bias_track = True, |
| 236 | proj_name = 'GEARS', |
| 237 | exp_name = 'GEARS'): |
| 238 | |
| 239 | self.weight_bias_track = weight_bias_track |
| 240 | |
| 241 | if self.weight_bias_track: |
| 242 | import wandb |
| 243 | wandb.init(project=proj_name, name=exp_name) |
| 244 | self.wandb = wandb |
| 245 | else: |
| 246 | self.wandb = None |
| 247 | |
| 248 | self.device = device |
| 249 | self.config = None |
| 250 | |
| 251 | self.dataloader = pert_data.dataloader |
| 252 | self.adata = pert_data.adata |
| 253 | self.node_map = pert_data.node_map |
| 254 | self.node_map_pert = pert_data.node_map_pert |
| 255 | self.data_path = pert_data.data_path |
| 256 | self.dataset_name = pert_data.dataset_name |
| 257 | self.split = pert_data.split |
| 258 | self.seed = pert_data.seed |
| 259 | self.train_gene_set_size = pert_data.train_gene_set_size |
| 260 | self.set2conditions = pert_data.set2conditions |
| 261 | self.subgroup = pert_data.subgroup |
| 262 | self.gene_list = pert_data.gene_names.values.tolist() |
| 263 | self.pert_list = pert_data.pert_names.tolist() |
| 264 | self.num_genes = len(self.gene_list) |
| 265 | self.num_perts = len(self.pert_list) |
| 266 | self.default_pert_graph = pert_data.default_pert_graph |
| 267 | self.saved_pred = {} |
| 268 | self.saved_logvar_sum = {} |
| 269 | |
| 270 | self.ctrl_expression = torch.tensor( |
| 271 | np.mean(self.adata.X[self.adata.obs['condition'].values == 'ctrl'], |
| 272 | axis=0)).reshape(-1, ).to(self.device) |
| 273 | pert_full_id2pert = dict(self.adata.obs[['condition_name', 'condition']].values) |
| 274 | self.dict_filter = {pert_full_id2pert[i]: j for i, j in |
| 275 | self.adata.uns['non_zeros_gene_idx'].items() if |
| 276 | i in pert_full_id2pert} |
| 277 | self.ctrl_adata = self.adata[self.adata.obs['condition'] == 'ctrl'] |
| 278 | |
| 279 | gene_dict = {g:i for i,g in enumerate(self.gene_list)} |
| 280 | self.pert2gene = {p: gene_dict[pert] for p, pert in |
| 281 | enumerate(self.pert_list) if pert in self.gene_list} |
| 282 | |
| 283 | def model_initialize(self, hidden_size = 64, |
| 284 | num_go_gnn_layers = 1, |
| 285 | num_gene_gnn_layers = 1, |