mRNAdataset.py

import os
import os.path as osp
import pickle
import shutil
import pandas as pd
from sklearn.preprocessing import MultiLabelBinarizer
import torch
from torch.utils.data import Dataset
from torch_geometric.data import (
    Data,
    InMemoryDataset,
    extract_gz,
)
from stru_decompose3 import *
import gdown
from utils.data_utils import cnn_tokenizer, kmer_encode
from tqdm import tqdm

def edge_order(e_idx):
    e_order = []
    for i in range(len(set(e_idx))):
        e_order.append(e_idx.count(i))
    return e_order
class mRNAdataset(InMemoryDataset):
    r"""
    Args:
        root (string): Root directory where the dataset should be saved.
        polymer (bool): whether it's polymeric tasks
        partition (string): which dataset split
        transform (callable, optional): A function/transform that takes in an
            :obj:`torch_geometric.data.Data` object and returns a transformed
            version. The data object will be transformed before every access.
            (default: :obj:`None`)
        pre_transform (callable, optional): A function/transform that takes in
            an :obj:`torch_geometric.data.Data` object and returns a
            transformed version. The data object will be transformed before
            being saved to disk. (default: :obj:`None`)
        pre_filter (callable, optional): A function that takes in an
            :obj:`torch_geometric.data.Data` object and returns a boolean
            value, indicating whether the data object should be included in the
            final dataset. (default: :obj:`None`)
    """

    raw_url0 = ('https://drive.google.com/uc?export=download&id=1k7S3pJU0GUIwqTAvUFKUfwbb68folDqW')
    raw_url1 = ('https://drive.google.com/uc?export=download&id=125B-q13oiURCkt9s1nCmYyqwJQqcgLI8')
    raw_url2 = ('https://drive.google.com/uc?export=download&id=125B-q13oiURCkt9s1nCmYyqwJQqcgLI8')

    def __init__(self, root, partition='train',
                 transform=None, pre_transform=None, pre_filter=None):
        assert partition in ['train','val','test']
        self.partition = partition
        self.label_binarizer = MultiLabelBinarizer()

        unique_labels = ['Exosome', 'Nucleus', 'Nucleoplasm', 'Chromatin',
                         'Nucleolus', 'Cytosol', 'Membrane', 'Ribosome', 'Cytoplasm']

        self.label_binarizer.fit([unique_labels])

        super().__init__(root, transform, pre_transform, pre_filter)

        if self.partition == 'train':
            self.data, self.slices = torch.load(self.processed_paths[0])
        elif self.partition == 'val':
            self.data, self.slices = torch.load(self.processed_paths[1])
        else:
            self.data, self.slices = torch.load(self.processed_paths[2])
        self.ids = self.data.id

    @property
    def raw_file_names(self):
        return ['train.csv', 'val.csv', 'test.csv']

    @property
    def processed_file_names(self):
        return ['train.pt','val.pt', 'test.pt']

    def download(self):
        print('Downloading mRNA train dataset...')
        try:
            gdown.download(self.raw_url0, self.raw_dir)
            file_path = osp.join(self.raw_dir, 'train.csv.gz')
            extract_gz(file_path, self.raw_dir)
            os.unlink(file_path)
        except:
            shutil.copy('./data/train.csv', self.raw_dir)

        print('Downloading mRNA val dataset...')
        try:
            gdown.download(self.raw_url0, self.raw_dir)
            file_path = osp.join(self.raw_dir, 'val.csv.gz')
            extract_gz(file_path, self.raw_dir)
            os.unlink(file_path)
        except:
            shutil.copy('./data/val.csv', self.raw_dir)

        print('Downloading mRNA test dataset...')
        try:
            gdown.download(self.raw_url2, self.raw_dir)
            file_path = osp.join(self.raw_dir, 'test.csv.gz')
            extract_gz(file_path, self.raw_dir)
            os.unlink(file_path)
        except:
            shutil.copy('./data/test.csv', self.raw_dir)

    def compute_hgraph_data(self, df):
        ids = df['id'].values.tolist()
        seqs = df['sequence'].values.tolist()
        dotbrackets = df['dot'].values.tolist()
        labels_l = [label.split('|') for label in df['label'].values.tolist()]
        labels = self.label_binarizer.transform(labels_l)
        hyper_res = process_rna_parallel(dotbrackets, seqs)
        data_list = []
        for data, hyper in tqdm(zip(zip(ids, seqs, labels), hyper_res), total=len(ids), desc="set up graph data"):
            id, seq, label = data
            node_fvs, n_idx, e_idx, edge_fvs = hyper
            y = torch.from_numpy(label).unsqueeze(0).float()
            x = torch.tensor(node_fvs, dtype=torch.long)
            edge_index0 = torch.tensor(n_idx, dtype=torch.long)
            edge_index1 = torch.tensor(e_idx, dtype=torch.long)
            edge_attr = torch.tensor(edge_fvs, dtype=torch.long)
            n_e = len(edge_index1.unique())
            e_order = torch.tensor(edge_order(e_idx), dtype=torch.long)
            data = HData(x=x, y=y, n_e=n_e, seq=seq, id=id,
                         edge_index0=edge_index0,
                         edge_index1=edge_index1,
                         edge_attr=edge_attr,
                         e_order=e_order)

            if self.pre_filter is not None and not self.pre_filter(data):
                continue
            if self.pre_transform is not None:
                data = self.pre_transform(data)

            data_list.append(data)
        return data_list

    def process(self):
        # for train set
        if not os.path.exists(self.processed_paths[0]):
            df = pd.read_csv(self.raw_paths[0])
            data_list = self.compute_hgraph_data(df)
            torch.save(self.collate(data_list), self.processed_paths[0])
        else:
            print(f"Skipping train set processing, {self.processed_paths[0]} already exists.")

        # for val set
        if not os.path.exists(self.processed_paths[1]):
            df = pd.read_csv(self.raw_paths[1])
            data_list = self.compute_hgraph_data(df)
            torch.save(self.collate(data_list), self.processed_paths[1])
        else:
            print(f"Skipping val set processing, {self.processed_paths[1]} already exists.")

        # for test set
        if not os.path.exists(self.processed_paths[2]):
            df = pd.read_csv(self.raw_paths[2])
            data_list = self.compute_hgraph_data(df)
            torch.save(self.collate(data_list), self.processed_paths[2])
        else:
            print(f"Skipping test set processing, {self.processed_paths[2]} already exists.")


class HData(Data):
    """ PyG data class for molecular hypergraphs
    """
    def __init__(self, x=None, edge_index=None, edge_attr=None, y=None, pos=None,
                 edge_index0=None, edge_index1=None, n_e=None, seq=None, id=None, **kwargs):
        super().__init__(x, edge_index, edge_attr, y, pos, **kwargs)
        self.edge_index0 = edge_index0
        self.edge_index1 = edge_index1
        self.n_e = n_e
        self.seq = seq
        self.id = id

    def __inc__(self, key, value, *args, **kwargs):
        if key == 'edge_index0':
            return self.x.size(0)
        if key == 'edge_index1':
            return self.n_e
        else:
            return super().__inc__(key, value, *args, **kwargs)


class mRNAdataset2(Dataset):
    def __init__(self, root, partition='train', model_name=None):
        assert partition in ['train', 'val', 'test']
        self.root = root
        self.partition = partition
        self.model_name = model_name
        self.max_seq_len = 6000
        self.save_path = osp.join(root, 'processed', f'{partition}_tokenized.pkl')
        self.raw_path = osp.join(root, 'raw', f'{partition}.csv')
        self.data = pd.read_csv(self.raw_path)
        self.ids = self.data['id'].values.tolist()
        self.seqs = self.data['sequence'].values.tolist()
        self.seq_data = self.load_or_extract_data()


    # def extract_(self):
    #     tokenized_seqs = []
    #     for seq in tqdm(self.seqs, desc="Tokenizing sequences"):
    #         tokenized_seqs.append(cnn_tokenizer(seq, max_len=self.max_seq_len))
    #
    #     with open(self.save_path, 'wb') as f:
    #         pickle.dump(tokenized_seqs, f)
    #     print(f"{self.save_path} saved. Tokenization complete.")
    #     return tokenized_seqs

    def extract_(self, k_values=[3, 4, 5, 6]):
        tokenized_seqs = {
            'one_hot': [],
            'k_mer_3': [],
            'k_mer_4': [],
            'k_mer_5': [],
            'k_mer_6': []
        }

        for seq in tqdm(self.seqs, desc="Tokenizing sequences"):
            seq = seq.replace("U", "T")
            if len(seq) > self.max_seq_len:
                seq = seq[:self.max_seq_len // 2] + seq[-(self.max_seq_len // 2):]
            tokenized_seqs['one_hot'].append(cnn_tokenizer(seq, self.max_seq_len))
            for k in k_values:
                tokenized_seqs[f'k_mer_{k}'].append(kmer_encode([seq], k)[0])

        with open(self.save_path, 'wb') as f:
            pickle.dump(tokenized_seqs, f)
        print(f"{self.save_path} saved. Tokenization complete.")
        return tokenized_seqs

    def load_or_extract_data(self):
        if os.path.exists(self.save_path):
            with open(self.save_path, 'rb') as f:
                tokenized_seqs = pickle.load(f)
            print(f"{self.save_path} loaded")
            return tokenized_seqs
        else:
            return self.extract_()

    def __len__(self):
        return len(self.ids)

    def __getitem__(self, idx):
        seq_data = {
            'one_hot': torch.tensor(self.seq_data['one_hot'][idx], dtype=torch.float32),
            'k_mer_3': torch.tensor(self.seq_data['k_mer_3'][idx], dtype=torch.float32),
            'k_mer_4': torch.tensor(self.seq_data['k_mer_4'][idx], dtype=torch.float32),
            'k_mer_5': torch.tensor(self.seq_data['k_mer_5'][idx], dtype=torch.float32),
            'k_mer_6': torch.tensor(self.seq_data['k_mer_6'][idx], dtype=torch.float32)
        }
        return self.ids[idx], seq_data

    # def load_or_extract_data(self):
    #     if os.path.exists(self.save_path):
    #         with open(self.save_path, 'rb') as f:
    #             tokenized_seqs = pickle.load(f)
    #         print(f"{self.save_path} loaded")
    #         return tokenized_seqs
    #     else:
    #         return self.extract_()
    #
    # def __len__(self):
    #     return len(self.ids)
    #
    # def __getitem__(self, idx):
    #     return self.ids[idx], self.seq_data[idx]

class CombinedDataset(Dataset):
    def __init__(self, graph_dataset, seq_dataset):
        self.graph_dataset = graph_dataset
        self.seq_dataset = seq_dataset
        assert len(self.graph_dataset) == len(self.seq_dataset), "data sizes do not match"
        for id1, (id2, _) in zip(self.graph_dataset.ids, self.seq_dataset):
            assert id1 == id2, f"ID no match: {id1} vs {id2}"

    def __len__(self):
        return len(self.graph_dataset)

    def __getitem__(self, idx):
        graph_data = self.graph_dataset[idx]
        id_, seq_data = self.seq_dataset[idx]
        return graph_data, seq_data