| import numpy as np |
| import torch |
| from torch.utils.data import Dataset |
|
|
| from tf_data_process import DataProcessorForPad |
|
|
|
|
| class DatasetMultiPad(Dataset): |
| def __init__(self, *args, **kwargs): |
| self.return_index = False |
| self.return_batch_label = kwargs['return_batch_label'] |
|
|
| self.cell_type_map = kwargs['cell_type_map'] |
| self.cell_type_col = kwargs['cell_type_col'] |
| if self.return_batch_label: |
| self.batch_label_map = kwargs['batch_label_map'] |
| self.batch_label_col = kwargs['batch_label_col'] |
| else: |
| self.batch_label_map = None |
| self.batch_label_col = None |
| self.data_args = kwargs |
| self.data_processor = DataProcessorForPad(**self.data_args) |
| self.adata_list = args |
| self.cumsum_lengths = np.cumsum([adata.shape[0] for adata in self.adata_list]) |
|
|
| def __len__(self): |
| return self.cumsum_lengths[-1] |
|
|
| def process_data(self, sparse_matrix, file_index): |
| value_list, chromosome_list, hg38_start, hg38_end = \ |
| self.data_processor.process( |
| value_data=sparse_matrix.toarray()[0].tolist(), |
| chromosome=self.adata_list[file_index].var["#Chromosome"].tolist(), |
| hg38_start=self.adata_list[file_index].var["hg38_Start"].tolist(), |
| hg38_end=self.adata_list[file_index].var["hg38_End"].tolist() |
| ) |
| return value_list, chromosome_list, hg38_start, hg38_end |
|
|
| def __getitem__(self, idx): |
| file_index = np.searchsorted(self.cumsum_lengths, idx, side="right") |
| if file_index == 0: |
| row_idx = idx |
| else: |
| row_idx = idx - self.cumsum_lengths[file_index - 1] |
| x_idx = self.adata_list[file_index][row_idx].X |
| value_list, chromosome_list, hg38_start, hg38_end = self.process_data( |
| x_idx, file_index) |
| cell_type = self.adata_list[file_index].obs.iloc[row_idx][self.cell_type_col] |
| if cell_type not in self.cell_type_map: |
| cell_type = 'Astrocyte 1' |
| res = [ |
| torch.tensor(value_list), |
| torch.tensor(chromosome_list), |
| torch.tensor(hg38_start), |
| torch.tensor(hg38_end), |
| torch.tensor(self.cell_type_map[cell_type]) |
| ] |
| if self.return_index: |
| res.insert(0, torch.tensor(idx)) |
| if self.return_batch_label: |
| batch_label = self.adata_list[file_index].obs.iloc[row_idx][self.batch_label_col] |
| res.append(torch.tensor(self.batch_label_map[batch_label])) |
| return res |
|
|
| def __del__(self): |
| |
| [data.file.close() for data in self.adata_list] |
|
|
