{"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"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport collections, itertools\nimport seaborn as sns\nimport matplotlib.pyplot as plt","metadata":{"execution":{"iopub.status.busy":"2023-10-09T07:06:19.397017Z","iopub.execute_input":"2023-10-09T07:06:19.397382Z","iopub.status.idle":"2023-10-09T07:06:20.258333Z","shell.execute_reply.started":"2023-10-09T07:06:19.397356Z","shell.execute_reply":"2023-10-09T07:06:20.257434Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"To understand what we have here, we actually have **timeseries data.** \n### For each sequence of mRNA, we want to predict its reactivity at each position:\n\n```\n    C        A        U         G        A    ...\n  0.61     0.42     0.01      0.19     0.99   ...\n```\n\nOur data is formatted with many columns, but only a few of them are informative. \n### In fact, we only have one piece of data to condition our prediction on:\n1. The mRNA sequence itself.\n\nAside from that, we also have a few pieces of metadata to help us decide how to train our model.\n\n2. Reads: Number of reads in the sequencing experiment that were assigned to the RNA sequence, and whose mutations were tabulated to compile the reactivity profile. (higher = more likely that this mRNA sequence data is correct).\n3. signal_to_noise: mean(measurement value)/mean(statistical error in measurement value). Higher = smaller spread for the reactivity target.\n4. SN_filter: a boolean of whether the sequence has > 100 reads and > 1 signal to noise. Basically if the data is good quality or not.\n5. reactivity error: a measure of (I'm assuming expected) error for the reactivity at each position. High error at a position = large spread = the reactivity at that position is less likely to be good training data.\n\nOur target is a sequence of the reactivities at each position in the mRNA sequence. The columns [reactivity_0001... reactivity_n], where n is the length of your sequence, contains the training target. \n\nUnfortunately the reactivities of the positions near the starts and ends of all sequences are unable to be scanned for technical reasons, so we'll have to figure out some way to impute data for those targets.","metadata":{}},{"cell_type":"code","source":"import torch.utils.data\nfrom torch.utils.data import Dataset, DataLoader\nimport torch as t\n\nclass StreamingCSVDataset(Dataset):\n    def __init__(self, file_path, chunksize, start_idx = 0, end_idx = 806572):\n        self.file_path = file_path\n        self.chunksize = chunksize\n        colnames = pd.read_csv(self.file_path, nrows=1).columns\n        self.data_stream = pd.read_csv(self.file_path, chunksize=self.chunksize, iterator=True, skiprows=start_idx+1,  names=colnames, header = None)\n        self.mapping ={'G': [1, 0, 0, 0],\n                       'A': [0, 1, 0, 0],\n                       'U': [0, 0, 1, 0],\n                       'C': [0, 0, 0, 1]}\n        self.idx = 0; self.start_idx = start_idx; self.end_idx = end_idx\n        self.refresh_chunk()\n        \n    def refresh_chunk(self):\n        chunk = self.data_stream.get_chunk(self.chunksize)\n        chunk.fillna(0, inplace = True)\n        self.input_sequences = np.stack(chunk[\"sequence\"].apply(lambda x: t.Tensor([self.mapping[i] for i in x])))\n        self.output_sequences = t.Tensor(chunk[[\"reactivity_\"+str(i).zfill(4) for i in range(1, self.input_sequences.shape[1]+1)]].values[:, :, np.newaxis]).type(t.float32)\n        \n    def __len__(self):\n        return self.end_idx - self.start_idx + 1\n\n    def __getitem__(self, index):\n        if index >= len(self):\n            raise IndexError(f\"{index} index out of range\")\n        if index >= self.idx + self.chunksize:\n            self.refresh_chunk()\n            self.idx += self.chunksize\n        return self.input_sequences[index - self.idx], self.output_sequences[index - self.idx]","metadata":{"execution":{"iopub.status.busy":"2023-10-09T07:06:20.265518Z","iopub.execute_input":"2023-10-09T07:06:20.265836Z","iopub.status.idle":"2023-10-09T07:06:20.275869Z","shell.execute_reply.started":"2023-10-09T07:06:20.265815Z","shell.execute_reply":"2023-10-09T07:06:20.274835Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Using GPUs with Pytorch\nWe heavily recommend using GPU acceleration, which should be on by default in this notebook.","metadata":{}},{"cell_type":"code","source":"if t.cuda.is_available():    \n    device = t.device(\"cuda:0\")\n    DEVICE_IDS = [0, 1] # Parallelize across 2 gpus\n    print('There are %d GPU(s) available.' % t.cuda.device_count())\nelse:\n    print('No GPU found.')\n    device = t.device(\"cpu\")","metadata":{"execution":{"iopub.status.busy":"2023-10-09T07:06:21.118536Z","iopub.execute_input":"2023-10-09T07:06:21.118841Z","iopub.status.idle":"2023-10-09T07:06:21.125144Z","shell.execute_reply.started":"2023-10-09T07:06:21.118817Z","shell.execute_reply":"2023-10-09T07:06:21.124086Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Let's replace all the NaNs with 0s. This is bad practice - it's much better to impute missing values\n# with the sample mean or a regression-imputed value instead. Try improving this on your own time!\n\nimport torch as t\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\n\nclass NaiveRNN(nn.Module):\n    def __init__(self, input_size, hidden_size, output_size, num_layers):\n        super(NaiveRNN, self).__init__()\n        self.hidden_size = hidden_size\n        self.num_layers = num_layers\n        self.rnn = nn.RNN(input_size, hidden_size, num_layers, batch_first=True)\n        self.fc = nn.Linear(hidden_size, output_size)\n\n    def forward(self, x):\n        h0 = t.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device) \n\n        out, _ = self.rnn(x, h0)  \n        out = self.fc(out)  \n        return out\n    \nclass PositionalEncoding(nn.Module):\n    def __init__(self, dim_model, dropout_p, max_len):\n        super().__init__()\n        # Modified version from: https://pytorch.org/tutorials/beginner/transformer_tutorial.html\n        # max_len determines how far the position can have an effect on a token (window)\n        \n        self.dropout = nn.Dropout(dropout_p)\n        \n        pos_encoding = t.zeros(max_len, dim_model)\n        positions_list = t.arange(0, max_len, dtype=t.float).view(-1, 1) # 0, 1, 2, 3, 4, 5\n        division_term = t.exp(t.arange(0, dim_model, 2).float() * (-math.log(10000.0)) / dim_model) # 1000^(2i/dim_model)\n        \n        # PE(pos, 2i) = sin(pos/1000^(2i/dim_model))\n        pos_encoding[:, 0::2] = t.sin(positions_list * division_term)\n        \n        # PE(pos, 2i + 1) = cos(pos/1000^(2i/dim_model))\n        pos_encoding[:, 1::2] = t.cos(positions_list * division_term)\n        \n        # Saving buffer (same as parameter without gradients needed)\n        pos_encoding = pos_encoding.unsqueeze(0).transpose(0, 1)\n        self.register_buffer(\"pos_encoding\",pos_encoding)\n        \n    def forward(self, token_embedding: t.tensor) -> t.tensor:\n        # Residual connection + pos encoding\n        return self.dropout(token_embedding + self.pos_encoding[:token_embedding.size(0), :])\n    \nclass TransformerModel(nn.Module):\n    def __init__(self, input_dim, output_dim, hidden_dim, nhead, num_layers):\n        super(TransformerModel, self).__init__()\n        self.transformer = nn.Transformer(d_model=hidden_dim, nhead=nhead, num_encoder_layers=num_layers, num_decoder_layers=num_layers)\n        self.fc_src = nn.Linear(input_dim, hidden_dim)\n        self.fc_tgt = nn.Linear(output_dim, hidden_dim)\n        self.fc_out = nn.Linear(hidden_dim, output_dim)\n        self.positional_encoder = PositionalEncoding(hidden_dim, 0.1, 500)\n\n    def get_tgt_mask(self, tgt):\n        mask = t.tril(t.ones(tgt.size(0), tgt.size(0)) == 1).to(\"cuda\")\n        mask = mask.float()\n        mask = mask.masked_fill(mask == 0, float('-inf')) # Convert zeros to -inf\n        mask = mask.masked_fill(mask == 1, float(0.0)) # Convert ones to 0\n        return mask\n    \n    def forward(self, src, tgt):\n        src = t.transpose(src, 0, 1)\n        tgt = t.transpose(tgt, 0, 1)\n        src = self.fc_src(src); tgt = self.fc_tgt(tgt)\n        src = self.positional_encoder(src)\n        tgt = self.positional_encoder(tgt)\n        mask = self.get_tgt_mask(tgt)\n        \n        output = self.transformer(src, tgt, tgt_mask = mask)\n        output = t.transpose(output, 0, 1)\n        output = self.fc_out(output)\n        return output\n\n# Hyperparameters\ninput_dim = 4\noutput_dim = 1\nhidden_dim = 512\nnhead = 8\nnum_layers = 2\n\n# Initialize the model\nmodel = TransformerModel(input_dim, output_dim, hidden_dim, nhead, num_layers)\n\nmodel = model.to(device)\nmodel = nn.DataParallel(model, device_ids=DEVICE_IDS)","metadata":{"execution":{"iopub.status.busy":"2023-10-09T07:06:21.725934Z","iopub.execute_input":"2023-10-09T07:06:21.726993Z","iopub.status.idle":"2023-10-09T07:06:21.992374Z","shell.execute_reply.started":"2023-10-09T07:06:21.726963Z","shell.execute_reply":"2023-10-09T07:06:21.99138Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import math\n# Loss + Optimizer\ncriterion = nn.MSELoss()\noptimizer = t.optim.Adam(model.parameters(), lr=1e-4)\n    \nBATCH_SIZE = 64 * 2\nCHUNK_SIZE = 2000\n\ntrain_ds = StreamingCSVDataset(\"/kaggle/input/stanford-ribonanza-rna-folding/train_data.csv\", CHUNK_SIZE, start_idx = 30000)\ntest_ds = StreamingCSVDataset(\"/kaggle/input/stanford-ribonanza-rna-folding/train_data.csv\", CHUNK_SIZE, start_idx = 0, end_idx = 29999)\n\ntrain_dl=t.utils.data.DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=False)\ntest_dl=t.utils.data.DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=False)\n\n# Training loop\nfrom tqdm import tqdm, trange\nepochs = 1; ema_loss = None\ntrain_history, test_history = [], []\nfor epoch in range(epochs):\n    print(\"Epoch %d\"%(epoch+1))\n    with tqdm(train_dl) as pbar:\n        for input_sequence, output_sequence in pbar:\n            input_sequence = input_sequence.to(device)\n            output_sequence = output_sequence.to(device)\n            # Zero the parameter gradients\n            optimizer.zero_grad()\n\n            # Forward pass\n            outputs = model(input_sequence, output_sequence)\n            loss = criterion(outputs, output_sequence)\n            # Backward pass and optimization\n            loss.backward()\n            optimizer.step()\n            if ema_loss is None: ema_loss = loss.item()\n            else: ema_loss = ema_loss * 0.98 + loss.item() * 0.02\n            pbar.set_description(\"Training Loss: %f\"%math.sqrt(ema_loss))\n            train_history.append(loss.item())\n    with tqdm(test_dl) as pbar:\n        for input_sequence, output_sequence in pbar:\n            input_sequence = input_sequence.to(device)\n            output_sequence = output_sequence.to(device)\n            # Forward pass\n            with t.no_grad():\n                outputs = model(input_sequence, output_sequence)\n                loss = criterion(outputs, output_sequence)\n\n                pbar.set_description(\"Validation Loss: %f\"%math.sqrt(loss))\n            test_history.append(loss.item())","metadata":{"execution":{"iopub.status.busy":"2023-10-09T07:06:23.012678Z","iopub.execute_input":"2023-10-09T07:06:23.012976Z","iopub.status.idle":"2023-10-09T07:06:31.066723Z","shell.execute_reply.started":"2023-10-09T07:06:23.01295Z","shell.execute_reply":"2023-10-09T07:06:31.065497Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import seaborn as sns\ndef ema(data, smoothing = 0.98):\n    smoothed_data = []\n    val = None\n    for item in data:\n        if val is None: val = item\n        else: val = val * smoothing + item * (1 - smoothing)\n        smoothed_data.append(val)\n    return smoothed_data\n\nfig, axes = plt.subplots(2, 1, figsize = (8, 5))\nsns.lineplot(ema(train_history), ax = axes[0], label = \"Train loss\")\nsns.lineplot(ema(test_history), ax = axes[1], label = \"Val loss\")","metadata":{"execution":{"iopub.status.busy":"2023-10-09T06:08:53.484951Z","iopub.execute_input":"2023-10-09T06:08:53.485657Z","iopub.status.idle":"2023-10-09T06:08:53.976451Z","shell.execute_reply.started":"2023-10-09T06:08:53.485621Z","shell.execute_reply":"2023-10-09T06:08:53.975604Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Now that we've got a model, let's make predictions!","metadata":{}},{"cell_type":"code","source":"test_input_df = pd.read_csv(\"/kaggle/input/stanford-ribonanza-rna-folding/test_sequences.csv\")\n\ntest_input_df = test_input_df.fillna(0)\nmapping = {'G': [1, 0, 0, 0],\n           'A': [0, 1, 0, 0],\n           'U': [0, 0, 1, 0],\n           'C': [0, 0, 0, 1]}\ninput_sequences = test_input_df[\"sequence\"].apply(lambda x: t.Tensor([mapping[i] for i in x]).to(device))","metadata":{"execution":{"iopub.status.busy":"2023-10-09T07:06:33.217054Z","iopub.execute_input":"2023-10-09T07:06:33.217556Z","iopub.status.idle":"2023-10-09T07:06:33.254872Z","shell.execute_reply.started":"2023-10-09T07:06:33.217519Z","shell.execute_reply":"2023-10-09T07:06:33.25404Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = model.module.to(device).eval()","metadata":{"execution":{"iopub.status.busy":"2023-10-09T07:06:49.192615Z","iopub.execute_input":"2023-10-09T07:06:49.192926Z","iopub.status.idle":"2023-10-09T07:06:49.201178Z","shell.execute_reply.started":"2023-10-09T07:06:49.1929Z","shell.execute_reply":"2023-10-09T07:06:49.200392Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import csv\n\nINFERENCE_BATCH_SIZE = 64\nwith open(\"output.csv\", \"a\") as csv_file:\n    writer = csv.writer(csv_file, delimiter=',')\n    writer.writerow(['id','reactivity_DMS_MaP','reactivity_2A3_MaP'])\n    cum_index = 0\n    for start_idx in tqdm(range(0, len(input_sequences), INFERENCE_BATCH_SIZE)):\n        end_idx = min(len(input_sequences), start_idx + INFERENCE_BATCH_SIZE)\n        # Forward pass\n        with t.no_grad():\n            batch_input_sequences, padding_masks, seq_lengths = zip(*[(F.pad(input_sequence, (0, 0, 0, 400 - len(input_sequence)), mode='constant', value=0), t.cat([t.zeros(len(input_sequence)), t.ones(400 - len(input_sequence))]).type(t.bool), len(input_sequence)) for input_sequence in input_sequences[start_idx:end_idx]])\n            batch_input_sequences = t.stack(batch_input_sequences)\n            padding_masks = t.stack(padding_masks).to(device)\n            y_input = t.zeros((batch_input_sequences.size(0), 1, 1), dtype=t.float32, device = 'cuda')\n            while y_input.size(1) < max(seq_lengths):\n                # Predict next tokens autoregressively\n                src = batch_input_sequences; tgt = y_input\n                src = t.transpose(src, 0, 1)\n                tgt = t.transpose(tgt, 0, 1)\n                src = model.fc_src(src); tgt = model.fc_tgt(tgt)\n                src = model.positional_encoder(src)\n                tgt = model.positional_encoder(tgt)\n                mask = model.get_tgt_mask(tgt)\n\n                output = model.transformer(src, tgt, tgt_mask = mask, src_key_padding_mask = padding_masks)\n                output = t.transpose(output, 0, 1)\n                pred = model.fc_out(output)\n\n                next_item = pred[:, -1] # take latest output token\n                # Concatenate previous input with predicted best word\n                y_input = t.cat((y_input, next_item.unsqueeze(1)), dim=1)\n\n            output_list = y_input.squeeze().cpu()\n            for batch_index, row in enumerate(output_list):\n                for index, i in enumerate(row[:seq_lengths[batch_index]]):\n                    writer.writerows([[cum_index, i.item(), i.item()]])\n                    cum_index += 1","metadata":{"execution":{"iopub.status.busy":"2023-10-09T07:12:23.362856Z","iopub.execute_input":"2023-10-09T07:12:23.363238Z","iopub.status.idle":"2023-10-09T07:13:15.077093Z","shell.execute_reply.started":"2023-10-09T07:12:23.363209Z","shell.execute_reply":"2023-10-09T07:13:15.075958Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pd.read_csv('output.csv')","metadata":{"execution":{"iopub.status.busy":"2023-10-09T07:13:44.565761Z","iopub.execute_input":"2023-10-09T07:13:44.566104Z","iopub.status.idle":"2023-10-09T07:13:44.592496Z","shell.execute_reply.started":"2023-10-09T07:13:44.566078Z","shell.execute_reply":"2023-10-09T07:13:44.591463Z"},"trusted":true},"execution_count":null,"outputs":[]}]}