{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":51294,"databundleVersionId":6923401,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import pandas as pd\n\ndf = pd.read_csv('/kaggle/input/stanford-ribonanza-rna-folding/train_data.csv')\n\ndf_2A3 = df[df['experiment_type'] == '2A3_MaP']\ndf_DMS = df[df['experiment_type'] == 'DMS_MaP']\ndf_2A3 = df_2A3.sample(n=10000, replace=False)\ndf_DMS = df_DMS.sample(n=10000, replace=False)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"### train 2A3 reactivity\nimport pandas as pd\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch import nn, optim\nfrom torch.nn import TransformerEncoder, TransformerEncoderLayer\nimport time\nimport math\nfrom sklearn.model_selection import KFold\nfrom torch.nn.utils.rnn import pad_sequence\n\nbase_to_int = {'A': 0, 'C': 1, 'G': 2, 'U': 3, 'PAD': 4}\n\nclass RNADataset(Dataset):\n    def __init__(self, dataframe, max_length, context_size=3):\n        self.dataframe = dataframe\n        self.max_length = max_length\n        self.context_size = context_size\n        self.pad_id = base_to_int['PAD']\n\n    def __len__(self):\n        return len(self.dataframe)\n\n    def __getitem__(self, idx):\n        sequence = self.dataframe.iloc[idx, 1]\n        context_padded_sequence = self.get_context_padded_sequence(sequence)\n        padded_sequence = context_padded_sequence + [self.pad_id] * (self.max_length - len(context_padded_sequence))\n        padded_sequence = padded_sequence[:self.max_length]\n\n        reactivity = self.dataframe.iloc[idx, 6:212].fillna(0).values.astype(float)\n        reactivity_padded = list(reactivity) + [0] * (self.max_length - len(reactivity))\n        reactivity_padded = reactivity_padded[:self.max_length]\n\n        mask = [1 if i < len(context_padded_sequence) else 0 for i in range(self.max_length)]\n\n        return torch.tensor(padded_sequence, dtype=torch.long), torch.tensor(reactivity_padded, dtype=torch.float), torch.tensor(mask, dtype=torch.float)\n\n    def get_context_padded_sequence(self, sequence):\n        # Add context to each nucleotide\n        padded_sequence = []\n        for i in range(len(sequence)):\n            context = sequence[max(0, i - self.context_size):min(len(sequence), i + self.context_size + 1)]\n            context = [self.pad_id] * (self.context_size - i) + [base_to_int[base] for base in context] + [self.pad_id] * (i + self.context_size + 1 - len(sequence))\n            padded_sequence.extend(context)\n\n        return padded_sequence\n\ndef collate_fn(batch):\n    sequences, reactivities, masks = zip(*batch)\n\n    lengths = [len(seq) for seq in sequences]\n\n    sequences_padded = pad_sequence(sequences, batch_first=True, padding_value=base_to_int['PAD'])\n    reactivities_padded = pad_sequence(reactivities, batch_first=True, padding_value=0)\n    masks_padded = pad_sequence(masks, batch_first=True, padding_value=0)\n\n    return sequences_padded, reactivities_padded, masks_padded, torch.tensor(lengths)\n\n\nclass PositionalEncoding(nn.Module):\n    def __init__(self, d_model, dropout=0.6, max_len=5000):\n        super(PositionalEncoding, self).__init__()\n        self.dropout = nn.Dropout(p=dropout)\n        position = torch.arange(max_len).unsqueeze(1)\n        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))\n        pe = torch.zeros(max_len, 1, d_model)\n        pe[:, 0, 0::2] = torch.sin(position * div_term)\n        pe[:, 0, 1::2] = torch.cos(position * div_term)\n        self.register_buffer('pe', pe)\n\n    def forward(self, x):\n        x = x + self.pe[:x.size(0)]\n        return self.dropout(x)\n\n\nclass TransformerModel(nn.Module):\n    def __init__(self, ntoken, ninp, nhead, nhid, nlayers, dropout=0.6):\n        super(TransformerModel, self).__init__()\n        self.model_type = 'Transformer'\n        self.pos_encoder = PositionalEncoding(ninp, dropout)\n\n        encoder_layers = TransformerEncoderLayer(ninp, nhead, nhid, dropout)\n        self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers)\n\n        self.encoder = nn.Embedding(ntoken, ninp)\n        self.ninp = ninp\n        self.decoder = nn.Linear(ninp, 1)\n        self.init_weights()\n\n    def init_weights(self):\n        initrange = 0.1\n        self.encoder.weight.data.uniform_(-initrange, initrange)\n        self.decoder.bias.data.zero_()\n        self.decoder.weight.data.uniform_(-initrange, initrange)\n\n    def forward(self, src):\n        src = self.encoder(src) * math.sqrt(self.ninp)\n        src = self.pos_encoder(src)\n        output = self.transformer_encoder(src)\n        output = self.decoder(output)\n        return output.squeeze(-1)\n\nclass MaskedHuberLoss(nn.Module):\n    def __init__(self, delta=1.0):\n        super(MaskedHuberLoss, self).__init__()\n        self.delta = delta\n\n    def forward(self, input, target, mask):\n        loss = torch.where(torch.abs(input - target) < self.delta,\n                           0.5 * (input - target) ** 2,\n                           self.delta * (torch.abs(input - target) - 0.5 * self.delta))\n        loss = loss * mask\n        return loss.sum() / mask.sum()\n\n\nntokens = len(base_to_int) \nemsize = 300  \nnhid = 300   \nnlayers = 2 \nnhead = 100    \ndropout = 0.6 \nweight_decay = 0.0001\n\nmax_length = max(df_2A3['sequence'].apply(len))\ndataset = RNADataset(df_2A3, max_length)\n\nkf = KFold(n_splits=3)\n\n\nfor fold, (train_index, test_index) in enumerate(kf.split(dataset)):\n    print(f\"Fold {fold + 1}\")\n    train_dataset = torch.utils.data.Subset(dataset, train_index)\n    test_dataset = torch.utils.data.Subset(dataset, test_index)\n    \n    train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)\n    test_dataloader = DataLoader(test_dataset, batch_size=32, shuffle=False, collate_fn=collate_fn)\n\n\n    model = TransformerModel(ntokens, emsize, nhead, nhid, nlayers, dropout)\n    loss_fn = MaskedHuberLoss()\n    optimizer = optim.Adam(model.parameters(), weight_decay=0.0001)\n\n   \n    num_epochs = 3\n\n    for epoch in range(num_epochs):\n        model.train()\n        total_loss = 0.0\n        for batch, (seq, target, mask, lengths) in enumerate(train_dataloader):\n            optimizer.zero_grad()\n            output = model(seq)\n            loss = loss_fn(output, target, mask)\n            loss.backward()\n            optimizer.step()\n            total_loss += loss.item()\n            start_time = time.time()\n\n            if batch % 10 == 0 and batch > 0:\n                current = batch * len(seq)\n                total = len(train_dataloader.dataset)\n                elapsed = time.time() \n                print(f'| Epoch {epoch+1:3d}/{num_epochs:3d} | {current:5d}/{total:5d} sequences | Elapsed time {elapsed:.2f}s')\n\n            elapsed_time = time.time() - start_time\n            print(f'Epoch {epoch+1:3d} completed in {elapsed_time:.2f}s, Total loss {total_loss:.4f}')\n            \n    torch.save(model.state_dict(), '/kaggle/working/model_state_dict_2A3.pth')\n\n\n    \n    model.eval()\n    total_test_loss = 0\n    with torch.no_grad():\n        for batch, (data, target, mask, lengths) in enumerate(test_dataloader): \n            output = model(data)\n            loss = loss_fn(output, target, mask)\n            total_test_loss += loss.item()\n\n        print(f\"Fold {fold + 1} Test Loss: {total_test_loss / len(test_dataloader)}\")\n\n\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"### predict 2A3 reactivity\nimport pandas as pd\nimport numpy as np\nimport torch\nfrom tqdm import tqdm\nimport pandas as pd\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch import nn, optim\nimport time\nimport math\nfrom sklearn.model_selection import KFold\nfrom torch.nn.utils.rnn import pad_sequence\n\ndf_test = pd.read_csv('/kaggle/input/stanford-ribonanza-rna-folding/test_sequences.csv')\n\nnum_splits = 20\nsplit_size = len(df_test) // num_splits\n\nfor i in range(num_splits):\n    start_index = i * split_size\n    end_index = start_index + split_size\n    if i == num_splits - 1:\n        end_index = len(df_test) \n\n    df_split = df_test.iloc[start_index:end_index]\n    df_split.to_csv(f'/kaggle/working/split_{i}.csv', index=False)\n\n\n    \nimport pandas as pd\nimport time\nimport torch\nfrom tqdm import tqdm\n\nbase_to_int = {'A': 0, 'C': 1, 'G': 2, 'U': 3, 'PAD': 4}\n\ndef get_context_padded_sequence(sequence, context_size=2, pad_id=base_to_int['PAD']):\n    padded_sequence = []\n    for i in range(len(sequence)):\n        context = sequence[max(0, i - context_size):min(len(sequence), i + context_size + 1)]\n        context = [pad_id] * (context_size - i) + [base_to_int[base] for base in context] + [pad_id] * (i + context_size + 1 - len(sequence))\n        padded_sequence.extend(context)\n    return padded_sequence\n\n\nmax_length = max(df_2A3['sequence'].apply(len))\n\nntokens = len(base_to_int) \nemsize = 300 \nnhid = 300   \nnlayers = 2   \nnhead = 100   \ndropout = 0.6 \n\nmodel = TransformerModel(ntokens, emsize, nhead, nhid, nlayers, dropout)\nmodel.load_state_dict(torch.load('/kaggle/input/model-state-dict-2a3-pth/model_state_dict_2A3.pth'))\nmodel.eval()\n\nstart_time = time.time()\nnum_splits = 20\nfor i in range(num_splits):\n    # Load split file\n    df_split = pd.read_csv(f'/kaggle/working/split_{i}.csv')\n    predictions = []\n\n    for sequence in tqdm(df_split['sequence'], desc=f'Processing split {i}'):\n        # Preprocess sequence\n        context_padded_sequence = get_context_padded_sequence(sequence)\n        padded_sequence = context_padded_sequence + [base_to_int['PAD']] * (max_length - len(context_padded_sequence))\n        sequence_tensor = torch.tensor(padded_sequence, dtype=torch.long).unsqueeze(0)\n\n        # Prediction\n        with torch.no_grad():\n            prediction = model(sequence_tensor)\n        \n        # Get and process prediction results\n        predicted_reactivity = prediction.squeeze(0).tolist()\n        predicted_reactivity_trimmed = predicted_reactivity[:len(sequence)]\n        predictions.append(predicted_reactivity_trimmed)\n\n    # Update DataFrame and save file after processing each split\n    df_split['reactivity_2A3_MaP'] = predictions\n    df_split.to_csv(f'/kaggle/working/predicted_sequences_2A3_{i}.csv', index=False)\n\n    # Calculate elapsed and estimated remaining time\n    elapsed_time = time.time() - start_time\n    estimated_remaining_time = (elapsed_time / (i + 1)) * (num_splits - i - 1)\n    print(f\"Completed split {i}. Estimated remaining time: {estimated_remaining_time:.2f} seconds\")","metadata":{},"execution_count":null,"outputs":[]}]}