{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"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":"code","source":"# Stanford RNA 3D Folding Competition\n# Revised Notebook for RNA 3D Structure Prediction (Improved V3)\n\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nimport os\nimport warnings\nwarnings.filterwarnings('ignore')\n\nprint(\"Starting Stanford RNA 3D Folding notebook...\")\n\n# 1. Data Loading and Exploration\nprint(\"Loading datasets...\")\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# Fill missing coordinate values to avoid NaNs.\ntrain_labels.fillna(0, inplace=True)\nvalidation_labels.fillna(0, inplace=True)\n\nprint(\"\\nBasic dataset information:\")\nprint(f\"Training sequences: {train_sequences.shape}\")\nprint(f\"Training labels: {train_labels.shape}\")\nprint(f\"Validation sequences: {validation_sequences.shape}\")\nprint(f\"Validation labels: {validation_labels.shape}\")\nprint(f\"Test sequences: {test_sequences.shape}\")\nprint(f\"Sample submission: {sample_submission.shape}\")\n\n# 2. Data Analysis (optional visualizations)\nprint(\"\\nAnalyzing RNA sequence lengths...\")\ntrain_sequences['length'] = train_sequences['sequence'].str.len()\nplt.figure(figsize=(12, 6))\nsns.histplot(train_sequences['length'], bins=50)\nplt.title('Distribution of RNA Sequence Lengths')\nplt.xlabel('Sequence Length')\nplt.ylabel('Count')\nplt.savefig('sequence_length_distribution.png')\nplt.close()\n\n# 3. Data Preprocessing\ndef preprocess_sequence_data(sequences_df, labels_df=None, is_train=True):\n    \"\"\"\n    Preprocess RNA sequence data.\n    Convert sequences to numerical form and normalize coordinate targets per sequence.\n    \"\"\"\n    nucleotide_map = {'A': 0, 'C': 1, 'G': 2, 'U': 3, 'T': 3}\n    processed_data = []\n    \n    for idx, row in sequences_df.iterrows():\n        seq_id = row['target_id']\n        sequence = row['sequence']\n        numerical_seq = [nucleotide_map.get(nuc, 4) for nuc in sequence]\n        \n        structures = None\n        if is_train and labels_df is not None:\n            sequence_labels = labels_df[labels_df['ID'].str.startswith(seq_id + '_')]\n            if not sequence_labels.empty:\n                num_structures = (len(sequence_labels.columns) - 3) // 3\n                structures = []\n                for i in range(1, num_structures + 1):\n                    coords = []\n                    for _, label_row in sequence_labels.iterrows():\n                        x = label_row[f'x_{i}']\n                        y = label_row[f'y_{i}']\n                        z = label_row[f'z_{i}']\n                        coords.append([x, y, z])\n                    coords = np.array(coords)\n                    # Normalize coordinates per sequence (center and scale)\n                    mean = np.mean(coords, axis=0)\n                    std = np.std(coords, axis=0) + 1e-8\n                    coords_norm = (coords - mean) / std\n                    structures.append(coords_norm)\n        processed_data.append({\n            'id': seq_id,\n            'sequence': numerical_seq,\n            'structures': structures\n        })\n    return processed_data\n\nprint(\"Preprocessing training data...\")\ntrain_data = preprocess_sequence_data(train_sequences, train_labels)\nprint(\"Preprocessing validation data...\")\nvalidation_data = preprocess_sequence_data(validation_sequences, validation_labels)\nprint(\"Preprocessing test data...\")\ntest_data = preprocess_sequence_data(test_sequences, is_train=False)\n\n# 4. Feature Engineering\ndef extract_sequence_features(sequence):\n    \"\"\"\n    Extract one-hot encoding, positional encoding, and GC-content as features.\n    \"\"\"\n    one_hot = np.zeros((len(sequence), 5))\n    for i, nucleotide in enumerate(sequence):\n        one_hot[i, nucleotide] = 1\n    gc_content = []\n    window_size = 5\n    for i in range(len(sequence)):\n        start = max(0, i - window_size // 2)\n        end = min(len(sequence), i + window_size // 2 + 1)\n        window = sequence[start:end]\n        gc_count = sum(1 for n in window if n in [1, 2])\n        gc_content.append(gc_count / len(window))\n    positions = np.array([[i / len(sequence)] for i in range(len(sequence))])\n    features = np.hstack((one_hot, positions, np.array(gc_content).reshape(-1, 1)))\n    return features\n\nprint(\"Extracting sequence features...\")\nfor i, data in enumerate(train_data):\n    train_data[i]['features'] = extract_sequence_features(data['sequence'])\nfor i, data in enumerate(validation_data):\n    validation_data[i]['features'] = extract_sequence_features(data['sequence'])\nfor i, data in enumerate(test_data):\n    test_data[i]['features'] = extract_sequence_features(data['sequence'])\n\n# 5. RNA Secondary Structure Prediction (simple rule-based)\ndef predict_rna_secondary_structure(sequence):\n    nucleotide_map_inv = {0: 'A', 1: 'C', 2: 'G', 3: 'U', 4: 'X'}\n    seq_chars = [nucleotide_map_inv[n] for n in sequence]\n    structure = ['.' for _ in range(len(seq_chars))]\n    complementary = {'A': 'U', 'U': 'A', 'G': 'C', 'C': 'G', 'X': None}\n    for i in range(len(seq_chars)):\n        if structure[i] != '.':\n            continue\n        for j in range(len(seq_chars) - 1, i + 3, -1):\n            if structure[j] != '.':\n                continue\n            if complementary[seq_chars[i]] == seq_chars[j]:\n                structure[i] = '('\n                structure[j] = ')'\n                break\n    return ''.join(structure)\n\ndef enhance_features_with_ss(data):\n    for i, item in enumerate(data):\n        seq = item['sequence']\n        ss = predict_rna_secondary_structure(seq)\n        ss_features = np.zeros((len(ss), 3))\n        for j, char in enumerate(ss):\n            if char == '.':\n                ss_features[j, 0] = 1\n            elif char == '(':\n                ss_features[j, 1] = 1\n            elif char == ')':\n                ss_features[j, 2] = 1\n        data[i]['features'] = np.hstack((item['features'], ss_features))\n    return data\n\nprint(\"Enhancing features with secondary structure information...\")\ntrain_data = enhance_features_with_ss(train_data)\nvalidation_data = enhance_features_with_ss(validation_data)\ntest_data = enhance_features_with_ss(test_data)\n\n# 6. Custom Dataset and DataLoader\nclass RNADataset(Dataset):\n    def __init__(self, data):\n        self.data = data\n    def __len__(self):\n        return len(self.data)\n    def __getitem__(self, idx):\n        item = self.data[idx]\n        features = torch.tensor(item['features'], dtype=torch.float32)\n        if item['structures'] is not None:\n            target = torch.tensor(item['structures'][0], dtype=torch.float32)\n            return features, target, item['id']\n        else:\n            return features, None, item['id']\n\ndef collate_fn(batch):\n    batch = sorted(batch, key=lambda x: x[0].shape[0], reverse=True)\n    features = [item[0] for item in batch]\n    targets = [item[1] for item in batch]\n    ids = [item[2] for item in batch]\n    max_length = features[0].shape[0]\n    feature_dim = features[0].shape[1]\n    padded_features = []\n    padded_targets = []\n    for i, feature in enumerate(features):\n        length = feature.shape[0]\n        padding = torch.zeros((max_length - length, feature_dim), dtype=torch.float32)\n        padded_feature = torch.cat([feature, padding], dim=0)\n        padded_features.append(padded_feature)\n        if targets[i] is not None:\n            target_padding = torch.zeros((max_length - length, 3), dtype=torch.float32)\n            padded_target = torch.cat([targets[i], target_padding], dim=0)\n            padded_targets.append(padded_target)\n    features_tensor = torch.stack(padded_features)\n    if all(target is not None for target in targets):\n        targets_tensor = torch.stack(padded_targets)\n        return features_tensor, targets_tensor, ids, [f.shape[0] for f in features]\n    else:\n        return features_tensor, None, ids, [f.shape[0] for f in features]\n\ntrain_dataset = RNADataset(train_data)\nvalidation_dataset = RNADataset(validation_data)\ntest_dataset = RNADataset(test_data)\n\nbatch_size = 4\ntrain_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)\nvalidation_loader = DataLoader(validation_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)\ntest_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)\n\n# 7. Model Architecture with BatchNorm and increased capacity\nclass RNAFoldingModel(nn.Module):\n    def __init__(self, input_dim, hidden_dim=256, num_layers=3, dropout=0.3):\n        super(RNAFoldingModel, self).__init__()\n        self.hidden_dim = hidden_dim\n        self.lstm = nn.LSTM(\n            input_dim, \n            hidden_dim, \n            num_layers=num_layers, \n            bidirectional=True, \n            batch_first=True,\n            dropout=dropout\n        )\n        self.attention = nn.Linear(hidden_dim * 2, 1)\n        self.fc1 = nn.Linear(hidden_dim * 2, hidden_dim)\n        self.bn1 = nn.BatchNorm1d(hidden_dim)\n        self.fc2 = nn.Linear(hidden_dim, hidden_dim // 2)\n        self.bn2 = nn.BatchNorm1d(hidden_dim // 2)\n        self.fc3 = nn.Linear(hidden_dim // 2, 3)  # 3D coordinates\n        self.dropout = nn.Dropout(dropout)\n        self.relu = nn.ReLU()\n    \n    def forward(self, x, seq_lengths=None):\n        batch_size, seq_len, _ = x.size()\n        if seq_lengths is not None:\n            packed_input = nn.utils.rnn.pack_padded_sequence(x, seq_lengths, batch_first=True, enforce_sorted=True)\n            packed_output, _ = self.lstm(packed_input)\n            lstm_out, _ = nn.utils.rnn.pad_packed_sequence(packed_output, batch_first=True)\n        else:\n            lstm_out, _ = self.lstm(x)\n        attention_scores = self.attention(lstm_out)\n        attention_weights = torch.softmax(attention_scores, dim=1)\n        context_vector = torch.sum(lstm_out * attention_weights, dim=1)\n        context_vector = context_vector.unsqueeze(1).expand(-1, seq_len, -1)\n        combined = lstm_out + context_vector\n        x = self.relu(self.fc1(combined))\n        # BatchNorm expects input as (B, C, L), so transpose, apply, then transpose back\n        x = self.bn1(x.transpose(1, 2)).transpose(1, 2)\n        x = self.dropout(x)\n        x = self.relu(self.fc2(x))\n        x = self.bn2(x.transpose(1, 2)).transpose(1, 2)\n        x = self.dropout(x)\n        x = self.fc3(x)\n        return x\n\n# 8. Training Functions using Smooth L1 Loss\ndef smooth_l1_loss(output, target, seq_lengths):\n    mask = torch.zeros_like(target, dtype=torch.bool)\n    for i, length in enumerate(seq_lengths):\n        mask[i, :length, :] = True\n    loss = nn.SmoothL1Loss(reduction='none')(output, target)\n    masked_loss = loss * mask.float()\n    return masked_loss.sum() / mask.sum() if mask.sum() > 0 else 0\n\ndef train_epoch(model, dataloader, optimizer, device):\n    model.train()\n    epoch_loss = 0\n    batches = 0\n    for features, targets, _, seq_lengths in dataloader:\n        if targets is None:\n            continue\n        optimizer.zero_grad()\n        features = features.to(device)\n        targets = targets.to(device)\n        outputs = model(features, seq_lengths)\n        loss = smooth_l1_loss(outputs, targets, seq_lengths)\n        loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)\n        optimizer.step()\n        epoch_loss += loss.item()\n        batches += 1\n    return epoch_loss / batches if batches > 0 else float('inf')\n\ndef validate(model, dataloader, device):\n    model.eval()\n    val_loss = 0\n    batches = 0\n    with torch.no_grad():\n        for features, targets, _, seq_lengths in dataloader:\n            if targets is None:\n                continue\n            features = features.to(device)\n            targets = targets.to(device)\n            outputs = model(features, seq_lengths)\n            loss = smooth_l1_loss(outputs, targets, seq_lengths)\n            val_loss += loss.item()\n            batches += 1\n    return val_loss / batches if batches > 0 else float('inf')\n\ndef calculate_tm_score(predicted, reference):\n    l_ref = len(reference)\n    if l_ref >= 30:\n        d0 = 0.6 * (l_ref - 0.5) ** 0.5 - 2.5\n    elif l_ref >= 24:\n        d0 = 0.7\n    elif l_ref >= 20:\n        d0 = 0.6\n    elif l_ref >= 16:\n        d0 = 0.5\n    elif l_ref >= 12:\n        d0 = 0.4\n    else:\n        d0 = 0.3\n    tm_score = 0\n    for i in range(min(len(predicted), l_ref)):\n        di = np.linalg.norm(predicted[i] - reference[i])\n        tm_score += 1 / (1 + (di/d0)**2)\n    return tm_score / l_ref\n\ndef evaluate_model(model, dataloader, device):\n    model.eval()\n    tm_scores = []\n    with torch.no_grad():\n        for features, targets, _, seq_lengths in dataloader:\n            if targets is None:\n                continue\n            features = features.to(device)\n            outputs = model(features, seq_lengths)\n            outputs = outputs.cpu().numpy()\n            targets = targets.cpu().numpy()\n            for i, length in enumerate(seq_lengths):\n                pred_coords = outputs[i, :length, :]\n                target_coords = targets[i, :length, :]\n                tm_score = calculate_tm_score(pred_coords, target_coords)\n                tm_scores.append(tm_score)\n    return np.mean(tm_scores) if tm_scores else 0\n\ndef train_model(model, train_loader, val_loader, num_epochs=20, lr=0.0005, device='cpu'):\n    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)\n    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=3, factor=0.5)\n    \n    train_losses = []\n    val_losses = []\n    tm_scores = []\n    \n    best_model_state = model.state_dict().copy()\n    best_val_loss = float('inf')\n    best_tm_score = 0\n    \n    print(\"Starting training...\")\n    for epoch in range(num_epochs):\n        train_loss = train_epoch(model, train_loader, optimizer, device)\n        val_loss = validate(model, val_loader, device)\n        tm_score = evaluate_model(model, val_loader, device)\n        \n        train_losses.append(train_loss)\n        val_losses.append(val_loss)\n        tm_scores.append(tm_score)\n        \n        scheduler.step(val_loss)\n        \n        print(f'Epoch {epoch+1}/{num_epochs}:')\n        print(f'  Train Loss: {train_loss:.4f}')\n        print(f'  Val Loss: {val_loss:.4f}')\n        print(f'  TM-Score: {tm_score:.4f}')\n        \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 (Val Loss: {val_loss:.4f})')\n        if tm_score > best_tm_score:\n            best_tm_score = tm_score\n            print(f'  New best TM-score: {tm_score:.4f}')\n    \n    if best_model_state is not None:\n        model.load_state_dict(best_model_state)\n    else:\n        print(\"Warning: Best model state not found; using current parameters.\")\n    \n    plt.figure(figsize=(15, 5))\n    plt.subplot(1, 2, 1)\n    plt.plot(train_losses, label='Training Loss')\n    plt.plot(val_losses, label='Validation Loss')\n    plt.title('Training and Validation Loss')\n    plt.xlabel('Epoch')\n    plt.ylabel('Loss')\n    plt.legend()\n    plt.subplot(1, 2, 2)\n    plt.plot(tm_scores, label='TM-Score')\n    plt.title('TM-Score Evolution')\n    plt.xlabel('Epoch')\n    plt.ylabel('TM-Score')\n    plt.legend()\n    plt.tight_layout()\n    plt.savefig('training_history.png')\n    plt.close()\n    \n    print(f\"Training complete. Best Val Loss: {best_val_loss:.4f}, Best TM-Score: {best_tm_score:.4f}\")\n    return model\n\n# 9. Model Inference and Multiple Structure Generation\ndef generate_diverse_structures(model, features, seq_length, num_structures=5, noise_scale=0.05, device='cpu'):\n    model.eval()\n    structures = []\n    for i in range(num_structures):\n        with torch.no_grad():\n            if i > 0:\n                noise = torch.randn_like(features) * noise_scale\n                features_with_noise = features + noise\n            else:\n                features_with_noise = features\n            output = model(features_with_noise.unsqueeze(0))\n            coords = output[0, :seq_length, :].cpu().numpy()\n            structures.append(coords)\n    return structures\n\ndef generate_predictions(model, dataloader, device, num_predictions=5):\n    model.eval()\n    all_predictions = {}\n    for features, _, ids, seq_lengths in dataloader:\n        features = features.to(device)\n        for i, (seq_id, length) in enumerate(zip(ids, seq_lengths)):\n            seq_features = features[i, :length, :]\n            predictions = generate_diverse_structures(\n                model, \n                seq_features, \n                length, \n                num_structures=num_predictions,\n                device=device\n            )\n            all_predictions[seq_id] = predictions\n    return all_predictions\n\n# 10. Submission File Generation\ndef create_submission_file(predictions, test_sequences_df, output_file='submission.csv'):\n    submission_rows = []\n    for _, row in test_sequences_df.iterrows():\n        seq_id = row['target_id']\n        sequence = row['sequence']\n        if seq_id in predictions:\n            pred_structures = predictions[seq_id]\n            num_structures = len(pred_structures)\n            for i in range(len(sequence)):\n                submission_row = {\n                    'ID': f\"{seq_id}_{i+1}\",\n                    'resname': sequence[i],\n                    'resid': i+1\n                }\n                for j in range(5):\n                    if j < num_structures:\n                        coords = pred_structures[j][i]\n                        submission_row[f'x_{j+1}'] = coords[0]\n                        submission_row[f'y_{j+1}'] = coords[1]\n                        submission_row[f'z_{j+1}'] = coords[2]\n                    else:\n                        submission_row[f'x_{j+1}'] = submission_row[f'x_{j}']\n                        submission_row[f'y_{j+1}'] = submission_row[f'y_{j}']\n                        submission_row[f'z_{j+1}'] = submission_row[f'z_{j}']\n                submission_rows.append(submission_row)\n    submission_df = pd.DataFrame(submission_rows)\n    submission_df.to_csv(output_file, index=False)\n    return submission_df\n\n# 11. Visualization Functions\ndef visualize_3d_structure(coords, title=\"RNA 3D Structure\"):\n    import matplotlib.pyplot as plt\n    fig = plt.figure(figsize=(10, 8))\n    ax = fig.add_subplot(111, projection='3d')\n    ax.scatter(coords[:, 0], coords[:, 1], coords[:, 2], c='blue', marker='o', s=30, label=\"C1' atoms\")\n    for i in range(len(coords) - 1):\n        ax.plot([coords[i, 0], coords[i+1, 0]], \n                [coords[i, 1], coords[i+1, 1]], \n                [coords[i, 2], coords[i+1, 2]], 'k-', lw=1)\n    ax.set_title(title)\n    ax.set_xlabel('X (Å)')\n    ax.set_ylabel('Y (Å)')\n    ax.set_zlabel('Z (Å)')\n    ax.legend()\n    plt.savefig(f\"{title.replace(' ', '_')}.png\")\n    plt.close()\n\n# 12. (Optional) Ensemble Modeling\nclass ModelEnsemble:\n    def __init__(self, models, weights=None):\n        self.models = models\n        self.weights = weights if weights is not None else [1/len(models)] * len(models)\n    def predict(self, features, seq_lengths=None):\n        all_predictions = []\n        for i, model in enumerate(self.models):\n            model.eval()\n            with torch.no_grad():\n                output = model(features, seq_lengths)\n                all_predictions.append(output * self.weights[i])\n        return sum(all_predictions)\n\n# 13. Main Execution\ndef main():\n    print(\"\\n--- Main execution ---\")\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    print(f\"Using device: {device}\")\n    \n    input_dim = train_data[0]['features'].shape[1]\n    model = RNAFoldingModel(input_dim=input_dim).to(device)\n    print(\"\\nModel instantiated.\")\n    \n    print(\"\\nStarting model training...\")\n    trained_model = train_model(\n        model,\n        train_loader,\n        validation_loader,\n        num_epochs=20,  # Increased epochs\n        lr=0.0005,     # Lower learning rate for stability\n        device=device\n    )\n    print(\"\\nModel training finished.\")\n    \n    print(\"\\nGenerating predictions on test data...\")\n    test_predictions = generate_predictions(trained_model, test_loader, device, num_predictions=5)\n    print(\"\\nPredictions generated.\")\n    \n    print(\"\\nCreating submission file...\")\n    submission_file = create_submission_file(test_predictions, test_sequences)\n    print(f\"\\nSubmission file created: submission.csv\")\n    print(submission_file.head())\n    \n    print(\"\\nVisualizing a sample prediction (first test sequence)...\")\n    sample_seq_id = test_sequences['target_id'].iloc[0]\n    if sample_seq_id in test_predictions:\n        sample_prediction = test_predictions[sample_seq_id][0]\n        visualize_3d_structure(sample_prediction, title=f\"Predicted 3D Structure - {sample_seq_id}\")\n        print(f\"Visualization saved for {sample_seq_id}.\")\n    else:\n        print(\"No prediction found for the first test sequence for visualization.\")\n    \n    print(\"\\n--- Main execution completed ---\")\n\nif __name__ == '__main__':\n    main()\n\nprint(\"\\nNotebook execution finished.\")","metadata":{"_uuid":"4a46c934-fe3f-460c-88cf-b1f1a9d250bb","_cell_guid":"2027a68e-272a-4b5d-9711-089e5e0a70da","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-02-28T09:33:24.077386Z","iopub.execute_input":"2025-02-28T09:33:24.077774Z","iopub.status.idle":"2025-02-28T09:40:40.322560Z","shell.execute_reply.started":"2025-02-28T09:33:24.077748Z","shell.execute_reply":"2025-02-28T09:40:40.321579Z"}},"outputs":[],"execution_count":null}]}