{"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":"nvidiaTeslaT4","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":"import 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 Residual Connections and Enhanced Capacity\nclass RNAFoldingModel(nn.Module):\n    def __init__(self, input_dim, hidden_dim=512, num_layers=3, dropout=0.2):\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        # First fully-connected layer projects to hidden_dim\n        self.fc1 = nn.Linear(hidden_dim * 2, hidden_dim)\n        # Residual projection from lstm output (hidden_dim*2) to hidden_dim\n        self.residual_proj = 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 mechanism\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        \n        # First FC block with residual connection\n        fc1_out = self.fc1(combined)\n        # Adding residual connection (after projection)\n        fc1_residual = self.residual_proj(lstm_out)\n        x = self.relu(fc1_out + fc1_residual)\n        \n        # BatchNorm expects (B, C, L)\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    # Use CosineAnnealingLR scheduler for smoother decay\n    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)\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()\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 can be tried\n        lr=0.0005,     # You might experiment with this learning rate further\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":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}