{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.13.2"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":87793,"databundleVersionId":11228175,"sourceType":"competition"}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Transformer Model for Stanford RNA 3D Folding Kaggle Competition\nThis notebook implements a deep learning model for predicting the 3D structure of RNA molecules based on their sequences. The competition aims to advance RNA-based medicine and biotechnology by improving our understanding of RNA folding. The model, a simplified version of RhoFold+, is a transformer-based architecture designed to predict RNA 3D structures. It includes data preprocessing, a custom dataset class, a transformer encoder, and a training loop. The predictions are evaluated using the TM-score metric, which measures the alignment of predicted and experimental structures. The final submission file contains the x, y, z coordinates of the C1' atom for each residue across five predicted structures.","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\n\n# Define constants\nMAX_SEQ_LENGTH = 512  # Maximum sequence length\nEMBEDDING_DIM = 256   # Dimension of embeddings\nNUM_HEADS = 4         # Number of attention heads\nNUM_LAYERS = 4        # Number of transformer layers\nHIDDEN_DIM = 512      # Hidden dimension for feed-forward layers\nBATCH_SIZE = 4       # Batch size\nLEARNING_RATE = 3e-4  # Learning rate\nNUM_EPOCHS = 20       # Number of training epochs\nNUM_PREDICTIONS = 5   # Number of structure predictions required\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","execution":{"iopub.execute_input":"2025-03-09T06:03:02.512379Z","iopub.status.busy":"2025-03-09T06:03:02.512185Z","iopub.status.idle":"2025-03-09T06:03:02.517257Z","shell.execute_reply":"2025-03-09T06:03:02.516436Z","shell.execute_reply.started":"2025-03-09T06:03:02.512361Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Load and Preprocess Data\nIn this section, the RNA sequence data and labels are loaded from CSV files. The maximum sequence length is calculated to ensure consistent input dimensions for the model.","metadata":{}},{"cell_type":"code","source":"    # Load data\ntrain_sequences = pd.read_csv('/kaggle/input/stanford-rna-3d-folding/train_sequences.csv')\ntrain_labels = pd.read_csv('/kaggle/input/stanford-rna-3d-folding/train_labels.csv')\nvalidation_sequences = pd.read_csv('/kaggle/input/stanford-rna-3d-folding/validation_sequences.csv')\nvalidation_labels = pd.read_csv('/kaggle/input/stanford-rna-3d-folding/validation_labels.csv')\ntest_sequences = pd.read_csv('/kaggle/input/stanford-rna-3d-folding/test_sequences.csv')\nsample_submission = pd.read_csv('/kaggle/input/stanford-rna-3d-folding/sample_submission.csv')\n\n\n#set Max_SEQ_length\nmax_len = max(train_sequences['sequence'].str.len().max(), \n              validation_sequences['sequence'].str.len().max(),\n              test_sequences['sequence'].str.len().max())\nprint(f\"Maximum sequence length in dataset: {max_len}\")\nMAX_SEQ_LENGTH = max_len","metadata":{"execution":{"iopub.execute_input":"2025-03-09T06:03:02.522114Z","iopub.status.busy":"2025-03-09T06:03:02.521855Z","iopub.status.idle":"2025-03-09T06:03:03.047163Z","shell.execute_reply":"2025-03-09T06:03:03.046264Z","shell.execute_reply.started":"2025-03-09T06:03:02.522094Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Define RNA Dataset Class\nThe `RNADataset` class is implemented to handle RNA sequence data. It encodes sequences, creates attention masks, and prepares labels for training and validation.","metadata":{}},{"cell_type":"code","source":"# RNA Dataset class\nclass RNADataset(Dataset):\n    def __init__(self, sequences, labels=None, is_test=False):\n        self.sequences = sequences\n        self.labels = labels\n        self.is_test = is_test\n        self.nucleotide_map = {'A': 1, 'C': 2, 'G': 3, 'U': 4, 'T': 4}\n        \n        if not is_test and labels is not None:\n            # Extract target_id from ID in labels\n            self.labels['target_id'] = self.labels['ID'].apply(\n                lambda x: '_'.join(x.split('_')[:-1]) if '_' in x else x\n            )\n    \n    def __len__(self):\n        return len(self.sequences)\n    \n    def __getitem__(self, idx):\n        sequence_row = self.sequences.iloc[idx]\n        target_id = sequence_row['target_id']\n        sequence = sequence_row['sequence']\n        \n        # Encode sequence\n        encoded_seq = np.zeros(MAX_SEQ_LENGTH, dtype=np.int64)\n        for i, nuc in enumerate(sequence[:MAX_SEQ_LENGTH]):\n            encoded_seq[i] = self.nucleotide_map.get(nuc, 0)\n        \n        # Create attention mask\n        seq_length = min(len(sequence), MAX_SEQ_LENGTH)\n        attention_mask = np.zeros(MAX_SEQ_LENGTH, dtype=np.int64)\n        attention_mask[:seq_length] = 1\n        \n        result = {\n            'target_id': target_id,\n            'sequence': sequence,\n            'encoded_seq': torch.tensor(encoded_seq),\n            'attention_mask': torch.tensor(attention_mask),\n            'seq_length': seq_length\n        }\n        \n        if not self.is_test and self.labels is not None:\n            # Find all labels for this target_id\n            target_labels = self.labels[self.labels['target_id'] == target_id]\n            \n            if len(target_labels) > 0:\n                # Extract coordinates\n                coords = np.zeros((MAX_SEQ_LENGTH, 3), dtype=np.float32)\n                mask = np.zeros(MAX_SEQ_LENGTH, dtype=np.float32)\n                \n                for i in range(1, seq_length + 1):\n                    # Find row with matching resid\n                    label_row = target_labels[target_labels['resid'] == i]\n                    \n                    if len(label_row) > 0:\n                        x = label_row['x_1'].values[0]\n                        y = label_row['y_1'].values[0]\n                        z = label_row['z_1'].values[0]\n                        \n                        if not (pd.isna(x) or pd.isna(y) or pd.isna(z)):\n                            coords[i-1] = [x, y, z]\n                            mask[i-1] = 1.0\n                \n                result['coords'] = torch.tensor(coords)\n                result['mask'] = torch.tensor(mask)\n        \n        return result\n","metadata":{"execution":{"iopub.execute_input":"2025-03-09T06:03:03.048540Z","iopub.status.busy":"2025-03-09T06:03:03.048268Z","iopub.status.idle":"2025-03-09T06:03:03.057618Z","shell.execute_reply":"2025-03-09T06:03:03.056664Z","shell.execute_reply.started":"2025-03-09T06:03:03.048501Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Transformer Encoder Layer\nThis section defines the Transformer encoder layer, which includes self-attention and feed-forward layers with residual connections and layer normalization.","metadata":{}},{"cell_type":"code","source":"# Transformer encoder layer\nclass TransformerEncoderLayer(nn.Module):\n    def __init__(self, d_model, nhead, dim_feedforward=HIDDEN_DIM, dropout=0.1):\n        super(TransformerEncoderLayer, self).__init__()\n        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)\n        self.linear1 = nn.Linear(d_model, dim_feedforward)\n        self.dropout = nn.Dropout(dropout)\n        self.linear2 = nn.Linear(dim_feedforward, d_model)\n        self.norm1 = nn.LayerNorm(d_model)\n        self.norm2 = nn.LayerNorm(d_model)\n        self.activation = nn.ReLU()\n    \n    def forward(self, src, src_mask=None):\n        # Self attention with residual connection and layer norm\n        src2, _ = self.self_attn(src, src, src, key_padding_mask=src_mask)\n        src = src + self.dropout(src2)\n        src = self.norm1(src)\n        \n        # Feed forward with residual connection and layer norm\n        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))\n        src = src + self.dropout(src2)\n        src = self.norm2(src)\n        \n        return src\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# RhoFold+ Model\nThe RhoFold+ model is a simplified transformer-based architecture designed for RNA 3D structure prediction. It includes embedding layers, positional encodings, and multiple transformer layers.","metadata":{}},{"cell_type":"code","source":"# RhoFold+ model (simplified)\nclass RhoFoldPlus(nn.Module):\n    def __init__(self, vocab_size=5):\n        super(RhoFoldPlus, self).__init__()\n        \n        # Embedding layer\n        self.embedding = nn.Embedding(vocab_size, EMBEDDING_DIM, padding_idx=0)\n        \n        # Positional encoding\n        self.pos_encoder = nn.Embedding(MAX_SEQ_LENGTH, EMBEDDING_DIM)\n        \n        # Transformer encoder layers\n        self.transformer_layers = nn.ModuleList([\n            TransformerEncoderLayer(EMBEDDING_DIM, NUM_HEADS)\n            for _ in range(NUM_LAYERS)\n        ])\n        \n        # Output heads for coordinate prediction (multiple heads for diverse predictions)\n        self.coordinate_heads = nn.ModuleList([\n            nn.Sequential(\n                nn.Linear(EMBEDDING_DIM, HIDDEN_DIM),\n                nn.ReLU(),\n                nn.Dropout(0.1),\n                nn.Linear(HIDDEN_DIM, HIDDEN_DIM // 2),\n                nn.ReLU(),\n                nn.Linear(HIDDEN_DIM // 2, 3)  # x, y, z coordinates\n            )\n            for _ in range(NUM_PREDICTIONS)\n        ])\n    \n    def forward(self, input_ids, attention_mask):\n        # Create padding mask for attention (True for padding positions)\n        padding_mask = (attention_mask == 0)\n        \n        # Create position indices\n        batch_size, seq_len = input_ids.size()\n        positions = torch.arange(0, seq_len, device=input_ids.device).unsqueeze(0).expand(batch_size, -1)\n        \n        # Embedding lookup\n        x = self.embedding(input_ids)\n        pos_emb = self.pos_encoder(positions)\n        \n        # Add positional embeddings\n        x = x + pos_emb\n        \n        # Apply transformer layers\n        for layer in self.transformer_layers:\n            x = layer(x, padding_mask)\n        \n        # Predict coordinates using different heads for diverse predictions\n        coordinates = [head(x) for head in self.coordinate_heads]\n        \n        return coordinates\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Loss Function\nThe `MaskedMSELoss` class implements a custom loss function that calculates the mean squared error for valid positions in the RNA sequence.","metadata":{}},{"cell_type":"code","source":"# Loss function with masking for valid positions\nclass MaskedMSELoss(nn.Module):\n    def __init__(self):\n        super(MaskedMSELoss, self).__init__()\n    \n    def forward(self, pred, target, mask):\n        # Expand mask to match dimensions\n        mask = mask.unsqueeze(-1).expand_as(pred)\n        \n        # Calculate squared error\n        squared_error = (pred - target) ** 2\n        \n        # Apply mask and calculate mean\n        masked_error = squared_error * mask\n        loss = masked_error.sum() / (mask.sum() + 1e-8)\n        \n        return loss\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Training the Model\nThis section defines the training loop for the RhoFold+ model. It includes forward and backward passes, loss calculation, and model evaluation on the validation set.","metadata":{}},{"cell_type":"code","source":"# Training function\ndef train_model(model, train_loader, val_loader):\n    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)\n    criterion = MaskedMSELoss()\n    \n    best_val_loss = float('inf')\n    best_model_state = None\n    \n    for epoch in range(NUM_EPOCHS):\n        # Training\n        model.train()\n        train_loss = 0.0\n        \n        for batch in train_loader:\n            input_ids = batch['encoded_seq'].to(DEVICE)\n            attention_mask = batch['attention_mask'].to(DEVICE)\n            target_coords = batch['coords'].to(DEVICE)\n            mask = batch['mask'].to(DEVICE)\n            \n            # Forward pass\n            pred_coords_list = model(input_ids, attention_mask)\n            \n            # Calculate loss for all prediction heads\n            loss = 0.0\n            for pred_coords in pred_coords_list:\n                loss += criterion(pred_coords, target_coords, mask)\n            loss /= len(pred_coords_list)\n            \n            # Backward pass\n            optimizer.zero_grad()\n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n            optimizer.step()\n            \n            train_loss += loss.item()\n        \n        train_loss /= len(train_loader)\n        \n        # Validation\n        model.eval()\n        val_loss = 0.0\n        \n        with torch.no_grad():\n            for batch in val_loader:\n                input_ids = batch['encoded_seq'].to(DEVICE)\n                attention_mask = batch['attention_mask'].to(DEVICE)\n                target_coords = batch['coords'].to(DEVICE)\n                mask = batch['mask'].to(DEVICE)\n                \n                # Forward pass\n                pred_coords_list = model(input_ids, attention_mask)\n                \n                # Calculate loss\n                batch_loss = 0.0\n                for pred_coords in pred_coords_list:\n                    batch_loss += criterion(pred_coords, target_coords, mask)\n                batch_loss /= len(pred_coords_list)\n                \n                val_loss += batch_loss.item()\n        \n        val_loss /= len(val_loader)\n        \n        print(f'Epoch {epoch+1}/{NUM_EPOCHS}, Train Loss: {train_loss:.6f}, Val Loss: {val_loss:.6f}')\n        \n        # Save best model\n        if val_loss < best_val_loss:\n            best_val_loss = val_loss\n            best_model_state = model.state_dict().copy()\n            print(f'New best model saved with val loss: {val_loss:.6f}')\n    \n    # Load best model\n    model.load_state_dict(best_model_state)\n    return model\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Generate Test Predictions\nThe `generate_predictions` function generates 3D coordinate predictions for RNA sequences in the test set. It formats the predictions to match the competition's submission requirements.","metadata":{}},{"cell_type":"code","source":"# Generate test predictions\ndef generate_predictions(model, test_loader, test_sequences, sample_submission):\n    model.eval()\n    \n    # Create test_clean dataframe (similar to original code)\n    def parse_target(tmp_ID, tmp_sequence):\n        seq_length = len(tmp_sequence)\n        tmp_df = pd.DataFrame(columns=['ID', 'resname', 'resid'], index=range(seq_length))\n        tmp_df['resname'] = list(tmp_sequence)\n        tmp_df['ID'] = tmp_ID\n        tmp_df['resid'] = range(1, seq_length + 1)\n        return tmp_df\n    \n    test_id_seq = test_sequences[['target_id', 'sequence']]\n    test_clean = pd.DataFrame(columns=['ID', 'resname', 'resid'])\n    \n    for index, row in test_id_seq.iterrows():\n        tmp_df = parse_target(row['target_id'], row['sequence'])\n        test_clean = pd.concat([test_clean, tmp_df], ignore_index=True)\n    \n    # Generate predictions\n    predictions = {}\n    \n    with torch.no_grad():\n        for batch in test_loader:\n            target_ids = batch['target_id']\n            sequences = batch['sequence']\n            input_ids = batch['encoded_seq'].to(DEVICE)\n            attention_mask = batch['attention_mask'].to(DEVICE)\n            seq_lengths = batch['seq_length']\n            \n            # Get predictions\n            pred_coords_list = model(input_ids, attention_mask)\n            \n            # Process each sequence\n            for i, (target_id, seq_len) in enumerate(zip(target_ids, seq_lengths)):\n                for j in range(seq_len):\n                    key = f\"{target_id}_{j+1}\"\n                    predictions[key] = {}\n                    \n                    # Save all 5 predictions\n                    for k, pred_coords in enumerate(pred_coords_list):\n                        predictions[key][f'x_{k+1}'] = pred_coords[i, j, 0].item()\n                        predictions[key][f'y_{k+1}'] = pred_coords[i, j, 1].item()\n                        predictions[key][f'z_{k+1}'] = pred_coords[i, j, 2].item()\n    \n    # Create submission dataframe\n    for idx, row in test_clean.iterrows():\n        key = f\"{row['ID']}_{row['resid']}\"\n        if key in predictions:\n            for col, value in predictions[key].items():\n                test_clean.loc[idx, col] = value\n    \n    # Format submission to match sample submission\n    submission = test_clean.copy()\n    submission['ID'] = submission['ID'] + '_' + submission['resid'].astype(str)\n    \n    # Store original columns before adding sort_order\n    original_columns = sample_submission.columns.tolist()\n    \n    # Add sort_order for sorting\n    sample_submission['sort_order'] = range(len(sample_submission))\n    \n    # Use only the original columns for the merge\n    submission = pd.merge(\n        submission[original_columns], \n        sample_submission[['ID', 'sort_order']], \n        on='ID', \n        how='left'\n    ).sort_values('sort_order').drop(columns=['sort_order'])\n    \n    return submission\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Main Function\nThe main function initializes the dataset, data loaders, and model. It trains the model and generates predictions for submission.","metadata":{}},{"cell_type":"code","source":"# Main function\nprint(f\"Using device: {DEVICE}\")\n    \n\n    # Create datasets\ntrain_dataset = RNADataset(train_sequences, train_labels)\nval_dataset = RNADataset(validation_sequences, validation_labels)\ntest_dataset = RNADataset(test_sequences, is_test=True)\n    \n    # Create data loaders\ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\nval_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE)\ntest_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE)\n    \n    # Initialize model\nmodel = RhoFoldPlus().to(DEVICE)\nprint(f\"Model initialized with {sum(p.numel() for p in model.parameters())} parameters\")\n    \n    # Train model\nmodel = train_model(model, train_loader, val_loader)\n    \n    # Generate predictions\nsubmission = generate_predictions(model, test_loader, test_sequences, sample_submission)\n    \n    # Save submission\nsubmission.to_csv('submission.csv', index=False)\nprint(\"Submission saved to submission.csv\")\n\n\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nfor dirname, _, filenames in os.walk('/kaggle/working'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"subs = pd.read_csv(\"/kaggle/working/submission.csv\")\nsubs","metadata":{"execution":{"iopub.execute_input":"2025-03-09T06:52:46.520817Z","iopub.status.busy":"2025-03-09T06:52:46.520601Z","iopub.status.idle":"2025-03-09T06:52:46.570364Z","shell.execute_reply":"2025-03-09T06:52:46.569462Z","shell.execute_reply.started":"2025-03-09T06:52:46.520798Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# After generating submission but before saving\nprint(submission[submission['ID'].str.startswith('R1138_')].head())\n","metadata":{"execution":{"iopub.execute_input":"2025-03-09T07:07:12.054165Z","iopub.status.busy":"2025-03-09T07:07:12.053823Z","iopub.status.idle":"2025-03-09T07:07:12.067997Z","shell.execute_reply":"2025-03-09T07:07:12.067199Z","shell.execute_reply.started":"2025-03-09T07:07:12.054140Z"},"trusted":true},"outputs":[],"execution_count":null}]}