{"metadata":{"kernelspec":{"display_name":".venv (3.13.5)","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.13.5"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":118765,"databundleVersionId":15231210,"sourceType":"competition"}],"dockerImageVersionId":31259,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Stanford RNA 3D Folding Part 2 - Improved Solution","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nimport warnings\nwarnings.filterwarnings('ignore')\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.model_selection import train_test_split\nimport math","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-08T07:54:33.986776Z","iopub.execute_input":"2026-02-08T07:54:33.987491Z","iopub.status.idle":"2026-02-08T07:54:33.991581Z","shell.execute_reply.started":"2026-02-08T07:54:33.987462Z","shell.execute_reply":"2026-02-08T07:54:33.990827Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"IS_KAGGLE = os.path.exists('/kaggle/input')\nDATA_PATH = '/kaggle/input/stanford-rna-3d-folding-2' if IS_KAGGLE else './data'\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n\nNUCLEOTIDE_ENCODING = {'A': 0, 'U': 1, 'G': 2, 'C': 3, 'N': 4}\n\nprint(f\"Device: {DEVICE}\")\nif DEVICE == 'cuda':\n    print(f\"GPU Name: {torch.cuda.get_device_name(0)}\")\n    print(f\"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB\")\n    print(f\"CUDA Version: {torch.version.cuda}\")\nelse:\n    print(\"WARNING: CUDA not available, using CPU\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-08T07:54:33.99281Z","iopub.execute_input":"2026-02-08T07:54:33.993041Z","iopub.status.idle":"2026-02-08T07:54:34.011243Z","shell.execute_reply.started":"2026-02-08T07:54:33.993022Z","shell.execute_reply":"2026-02-08T07:54:34.01069Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_seq_df = pd.read_csv(f'{DATA_PATH}/train_sequences.csv')\ntrain_labels_df = pd.read_csv(f'{DATA_PATH}/train_labels.csv')\ntest_seq_df = pd.read_csv(f'{DATA_PATH}/test_sequences.csv')\nsample_submission = pd.read_csv(f'{DATA_PATH}/sample_submission.csv')\n\nprint(f\"Train sequences: {len(train_seq_df)}\")\nprint(f\"Train labels: {len(train_labels_df)}\")\nprint(f\"Test sequences: {len(test_seq_df)}\")\nprint(f\"\\nTrain labels columns: {list(train_labels_df.columns)}\")\nprint(f\"Train labels sample:\")\nprint(train_labels_df.head())\n\n# Check sequence lengths\nif 'sequence' in train_seq_df.columns:\n    train_lengths = train_seq_df['sequence'].str.len()\n    print(f\"\\nTrain sequence lengths - Min: {train_lengths.min()}, Max: {train_lengths.max()}, Mean: {train_lengths.mean():.1f}\")\nif 'sequence' in test_seq_df.columns:\n    test_lengths = test_seq_df['sequence'].str.len()\n    print(f\"Test sequence lengths - Min: {test_lengths.min()}, Max: {test_lengths.max()}, Mean: {test_lengths.mean():.1f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-08T07:54:34.012404Z","iopub.execute_input":"2026-02-08T07:54:34.01263Z","iopub.status.idle":"2026-02-08T07:54:41.285366Z","shell.execute_reply.started":"2026-02-08T07:54:34.012606Z","shell.execute_reply":"2026-02-08T07:54:41.284592Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def encode_sequence(sequence):\n    return np.array([NUCLEOTIDE_ENCODING.get(nuc, 4) for nuc in sequence])\n\n\nclass ImprovedRNAFolding:\n    def __init__(self):\n        # More realistic RNA helix parameters based on A-form RNA\n        self.rise_per_residue = 2.56  # A-form RNA rise\n        self.rotation_per_residue = 32.7  # ~11 bp per turn\n        self.helix_radius = 10.0  # Slightly larger radius\n        self.base_pair_distance = 20.0  # Watson-Crick base pair distance in Angstroms\n        \n    def predict_secondary_structure(self, sequence):\n        \"\"\"Improved base pairing with GU wobble pairs\"\"\"\n        pairs = {\n            'A': ['U'], \n            'U': ['A', 'G'],  # U can pair with A or G (wobble)\n            'G': ['C', 'U'],  # G can pair with C or U (wobble)\n            'C': ['G']\n        }\n        n = len(sequence)\n        paired = [False] * n\n        pair_map = {}\n        \n        # Use dynamic programming approach for better pairing\n        # Simplified Nussinov algorithm\n        for length in range(4, n + 1):\n            for i in range(n - length + 1):\n                j = i + length - 1\n                if not paired[i] and not paired[j]:\n                    if sequence[i] in pairs and sequence[j] in pairs.get(sequence[i], []):\n                        paired[i] = True\n                        paired[j] = True\n                        pair_map[i] = j\n                        pair_map[j] = i\n        \n        return pair_map\n    \n    def build_structure(self, sequence, variant=0):\n        \"\"\"Build 3D structure with improved physics\"\"\"\n        n = len(sequence)\n        coords = np.zeros((n, 3))\n        \n        pair_map = self.predict_secondary_structure(sequence)\n        \n        # Vary parameters per variant for diversity\n        variant_rotation = variant * 15  # More rotation variation\n        variant_radius = self.helix_radius * (1 + 0.15 * (variant - 2))\n        variant_rise = self.rise_per_residue * (1 + 0.1 * (variant - 2))\n        \n        # Initial helix placement\n        for i in range(n):\n            angle = np.radians(i * self.rotation_per_residue + variant_rotation)\n            coords[i, 0] = variant_radius * np.cos(angle)\n            coords[i, 1] = variant_radius * np.sin(angle)\n            coords[i, 2] = i * variant_rise\n        \n        # Apply secondary structure constraints with improved refinement\n        learning_rate = 0.05\n        for iteration in range(100):  # More iterations\n            forces = np.zeros_like(coords)\n            \n            for i, j in pair_map.items():\n                if i < j:\n                    current_vec = coords[j] - coords[i]\n                    current_dist = np.linalg.norm(current_vec)\n                    \n                    if current_dist > 0.1:  # Avoid division by zero\n                        target_dist = self.base_pair_distance\n                        direction = current_vec / current_dist\n                        error = current_dist - target_dist\n                        \n                        # Apply force proportional to error\n                        force = direction * error * learning_rate\n                        forces[i] += force\n                        forces[j] -= force\n            \n            coords += forces\n            \n            # Add local smoothing to maintain backbone continuity\n            for i in range(1, n - 1):\n                smoothed = (coords[i-1] + coords[i] + coords[i+1]) / 3\n                coords[i] = 0.7 * coords[i] + 0.3 * smoothed\n        \n        # Add small random perturbation for diversity\n        noise = np.random.normal(0, 0.5, coords.shape)\n        coords += noise\n        \n        return coords\n    \n    def predict_with_template(self, sequence, template_coords, variant=0):\n        \"\"\"Use template with better transformations\"\"\"\n        n = len(sequence)\n        \n        # Handle length mismatch\n        if len(template_coords) != n:\n            if len(template_coords) > n:\n                coords = template_coords[:n].copy()\n            else:\n                # Fill remaining with physics-based prediction\n                remaining_seq = sequence[len(template_coords):]\n                remaining_coords = self.build_structure(remaining_seq, variant)\n                # Translate to connect with template\n                if len(template_coords) > 0:\n                    offset = template_coords[-1] - remaining_coords[0]\n                    remaining_coords += offset\n                coords = np.vstack([template_coords, remaining_coords])\n        else:\n            coords = template_coords.copy()\n        \n        # Apply diverse transformations\n        # Scale\n        scale = 1.0 + 0.15 * (variant - 2)\n        coords = coords * scale\n        \n        # Rotation around multiple axes\n        angle_x = np.radians(15 * variant)\n        angle_y = np.radians(20 * variant)\n        angle_z = np.radians(25 * variant)\n        \n        # Rotation matrices\n        Rx = np.array([\n            [1, 0, 0],\n            [0, np.cos(angle_x), -np.sin(angle_x)],\n            [0, np.sin(angle_x), np.cos(angle_x)]\n        ])\n        \n        Ry = np.array([\n            [np.cos(angle_y), 0, np.sin(angle_y)],\n            [0, 1, 0],\n            [-np.sin(angle_y), 0, np.cos(angle_y)]\n        ])\n        \n        Rz = np.array([\n            [np.cos(angle_z), -np.sin(angle_z), 0],\n            [np.sin(angle_z), np.cos(angle_z), 0],\n            [0, 0, 1]\n        ])\n        \n        # Apply rotations\n        coords = coords @ Rx.T @ Ry.T @ Rz.T\n        \n        # Add controlled noise\n        noise = np.random.normal(0, 0.8 * (1 + variant * 0.2), coords.shape)\n        coords += noise\n        \n        return coords","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-08T07:54:41.286805Z","iopub.execute_input":"2026-02-08T07:54:41.287146Z","iopub.status.idle":"2026-02-08T07:54:41.302494Z","shell.execute_reply.started":"2026-02-08T07:54:41.287125Z","shell.execute_reply":"2026-02-08T07:54:41.301729Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class RNADataset(Dataset):\n    def __init__(self, sequences_df, labels_df, max_length=500):\n        self.sequences_df = sequences_df\n        self.max_length = max_length\n        self.id_col = 'target_id' if 'target_id' in sequences_df.columns else 'ID'\n        \n        # Pre-process labels EFFICIENTLY\n        print(f\"Pre-processing labels for {len(sequences_df)} sequences...\")\n        self.labels_dict = {}\n        \n        # Get the list of sequence IDs we actually need\n        needed_ids = set(sequences_df[self.id_col].unique())\n        \n        # Filter labels_df to only needed sequences FIRST (major speedup!)\n        print(\"  Filtering labels...\")\n        labels_df['target_id_extracted'] = labels_df['ID'].str.rsplit('_', n=1).str[0]\n        labels_subset = labels_df[labels_df['target_id_extracted'].isin(needed_ids)]\n        print(f\"  Filtered from {len(labels_df)} to {len(labels_subset)} labels\")\n        \n        # Group by target_id and process\n        print(\"  Grouping and extracting coordinates...\")\n        grouped = labels_subset.groupby('target_id_extracted')\n        \n        for seq_id in needed_ids:\n            if seq_id in grouped.groups:\n                group = grouped.get_group(seq_id).sort_values('resid')\n                coords = np.zeros((len(group), 3))\n                \n                if 'x_1' in group.columns:\n                    coords[:, 0] = group['x_1'].values\n                    coords[:, 1] = group['y_1'].values\n                    coords[:, 2] = group['z_1'].values\n                elif 'x' in group.columns:\n                    coords[:, 0] = group['x'].values\n                    coords[:, 1] = group['y'].values\n                    coords[:, 2] = group['z'].values\n                \n                self.labels_dict[seq_id] = coords\n        \n        print(f\"  Done! Pre-processed {len(self.labels_dict)} sequences\")\n        \n    def __len__(self):\n        return len(self.sequences_df)\n    \n    def __getitem__(self, idx):\n        row = self.sequences_df.iloc[idx]\n        seq_id = row[self.id_col]\n        sequence = row['sequence']\n        \n        # Encode sequence\n        encoded = encode_sequence(sequence)\n        if len(encoded) < self.max_length:\n            encoded = np.pad(encoded, (0, self.max_length - len(encoded)), constant_values=4)\n        else:\n            encoded = encoded[:self.max_length]\n        \n        # Get pre-processed coordinates (MUCH FASTER!)\n        coordinates = np.zeros((self.max_length, 3))\n        if seq_id in self.labels_dict:\n            coords = self.labels_dict[seq_id]\n            n_coords = min(len(coords), self.max_length)\n            coordinates[:n_coords] = coords[:n_coords]\n        \n        return torch.LongTensor(encoded), torch.FloatTensor(coordinates), len(sequence)\n\n\n# IMPROVED: Transformer-based architecture (better than LSTM for long-range dependencies)\nclass PositionalEncoding(nn.Module):\n    def __init__(self, d_model, dropout=0.1, max_len=5000):\n        super().__init__()\n        self.dropout = nn.Dropout(p=dropout)\n        \n        pe = torch.zeros(max_len, d_model)\n        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n        pe[:, 0::2] = torch.sin(position * div_term)\n        pe[:, 1::2] = torch.cos(position * div_term)\n        pe = pe.unsqueeze(0)\n        self.register_buffer('pe', pe)\n        \n    def forward(self, x):\n        x = x + self.pe[:, :x.size(1), :]\n        return self.dropout(x)\n\n\nclass RNA3DTransformer(nn.Module):\n    \"\"\"Transformer-based RNA 3D predictor - significantly better than LSTM\"\"\"\n    def __init__(self, vocab_size=5, embedding_dim=384, num_heads=8, num_layers=6, \n                 dim_feedforward=1536, dropout=0.1, max_length=500):\n        super().__init__()\n        \n        # Store embedding_dim for forward pass\n        self.embedding_dim = embedding_dim\n        \n        # Embedding layer\n        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=4)\n        self.pos_encoder = PositionalEncoding(embedding_dim, dropout, max_length)\n        \n        # Transformer encoder layers\n        encoder_layer = nn.TransformerEncoderLayer(\n            d_model=embedding_dim,\n            nhead=num_heads,\n            dim_feedforward=dim_feedforward,\n            dropout=dropout,\n            activation='gelu',  # GELU activation (better than ReLU)\n            batch_first=True,\n            norm_first=True  # Pre-LayerNorm architecture (more stable)\n        )\n        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)\n        \n        # Prediction head with residual connections\n        self.fc1 = nn.Linear(embedding_dim, embedding_dim)\n        self.layer_norm1 = nn.LayerNorm(embedding_dim)\n        self.dropout1 = nn.Dropout(dropout)\n        \n        self.fc2 = nn.Linear(embedding_dim, embedding_dim // 2)\n        self.layer_norm2 = nn.LayerNorm(embedding_dim // 2)\n        self.dropout2 = nn.Dropout(dropout)\n        \n        self.fc3 = nn.Linear(embedding_dim // 2, 3)  # Output: x, y, z coordinates\n        \n        self.gelu = nn.GELU()\n        \n        # Initialize weights properly\n        self._init_weights()\n        \n    def _init_weights(self):\n        for p in self.parameters():\n            if p.dim() > 1:\n                nn.init.xavier_uniform_(p)\n                \n    def forward(self, x, src_key_padding_mask=None):\n        # x: (batch_size, seq_length)\n        # Embedding with scaling (standard for transformers)\n        x = self.embedding(x) * math.sqrt(self.embedding_dim)\n        x = self.pos_encoder(x)\n        \n        # Transformer encoding\n        x = self.transformer(x, src_key_padding_mask=src_key_padding_mask)\n        \n        # Prediction head with residual connections\n        out = self.fc1(x)\n        out = self.layer_norm1(out)\n        out = self.gelu(out)\n        out = self.dropout1(out)\n        \n        out = self.fc2(out)\n        out = self.layer_norm2(out)\n        out = self.gelu(out)\n        out = self.dropout2(out)\n        \n        coords = self.fc3(out)\n        \n        return coords\n\n\n# Keep old LSTM model for comparison\nclass RNA3DPredictor(nn.Module):\n    def __init__(self, vocab_size=5, embedding_dim=256, hidden_dim=512, num_layers=3, dropout=0.3):\n        super().__init__()\n        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=4)\n        \n        # Bidirectional LSTM\n        self.lstm = nn.LSTM(\n            embedding_dim, \n            hidden_dim, \n            num_layers, \n            batch_first=True, \n            bidirectional=True, \n            dropout=dropout if num_layers > 1 else 0\n        )\n        \n        # Simplified attention mechanism\n        self.attention = nn.MultiheadAttention(\n            embed_dim=hidden_dim * 2,\n            num_heads=8,\n            dropout=dropout,\n            batch_first=True\n        )\n        \n        # Streamlined prediction head\n        self.fc1 = nn.Linear(hidden_dim * 2, hidden_dim)\n        self.dropout1 = nn.Dropout(dropout)\n        \n        self.fc2 = nn.Linear(hidden_dim, 3)\n        \n        self.relu = nn.ReLU()\n        self.layer_norm = nn.LayerNorm(hidden_dim * 2)\n        \n    def forward(self, x):\n        # Embedding\n        embedded = self.embedding(x)\n        \n        # LSTM\n        lstm_out, _ = self.lstm(embedded)\n        \n        # Layer normalization\n        lstm_out = self.layer_norm(lstm_out)\n        \n        # Self-attention\n        attended, _ = self.attention(lstm_out, lstm_out, lstm_out)\n        \n        # Residual connection\n        combined = lstm_out + attended\n        \n        # Prediction head\n        out = self.fc1(combined)\n        out = self.relu(out)\n        out = self.dropout1(out)\n        \n        # Final layer\n        coords = self.fc2(out)\n        \n        return coords","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-08T07:54:41.303928Z","iopub.execute_input":"2026-02-08T07:54:41.304562Z","iopub.status.idle":"2026-02-08T07:54:41.325892Z","shell.execute_reply.started":"2026-02-08T07:54:41.304531Z","shell.execute_reply":"2026-02-08T07:54:41.325287Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Training neural network with IMPROVED architecture...\")\n\n# CRITICAL: Split data into train and validation (90/10 split)\nimport os\nimport time\n\nprint(f\"Total available training sequences: {len(train_seq_df)}\")\n\n# Split by sequence ID to avoid data leakage\ntrain_df, val_df = train_test_split(train_seq_df, test_size=0.1, random_state=42)\nprint(f\"Training on {len(train_df)} sequences, validating on {len(val_df)} sequences\")\n\n# Create datasets\ntrain_dataset = RNADataset(train_df, train_labels_df, max_length=500)\nval_dataset = RNADataset(val_df, train_labels_df, max_length=500)\n\n# DataLoaders with optimized settings\ntrain_loader = DataLoader(\n    train_dataset, \n    batch_size=16,  # Reduced for Transformer (uses more memory)\n    shuffle=True, \n    num_workers=0,\n    pin_memory=True if DEVICE == 'cuda' else False\n)\n\nval_loader = DataLoader(\n    val_dataset,\n    batch_size=16,\n    shuffle=False,\n    num_workers=0,\n    pin_memory=True if DEVICE == 'cuda' else False\n)\n\n# IMPROVED MODEL: Use Transformer instead of LSTM\nmodel = RNA3DTransformer(\n    embedding_dim=384,  # Larger embedding\n    num_heads=8,\n    num_layers=6,  # 6 transformer layers\n    dim_feedforward=1536,\n    dropout=0.1,  # Lower dropout for transformers\n    max_length=500\n).to(DEVICE)\n\nprint(f\"Model type: Transformer (6 layers, 8 heads)\")\nprint(f\"Model parameters: {sum(p.numel() for p in model.parameters()):,}\")\n\nif DEVICE == 'cuda':\n    print(f\"Model is on device: {next(model.parameters()).device}\")\n\n# Better optimizer and scheduler\noptimizer = torch.optim.AdamW(\n    model.parameters(), \n    lr=5e-4,  # Lower learning rate for transformer\n    betas=(0.9, 0.98),  # Better betas for transformers\n    eps=1e-9,\n    weight_decay=1e-4\n)\n\n# Warmup + Cosine annealing scheduler\nnum_epochs = 50  # More epochs with early stopping\nwarmup_epochs = 5\n\ndef get_lr_multiplier(epoch):\n    if epoch < warmup_epochs:\n        return (epoch + 1) / warmup_epochs\n    else:\n        progress = (epoch - warmup_epochs) / (num_epochs - warmup_epochs)\n        return 0.5 * (1 + math.cos(math.pi * progress))\n\nscheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, get_lr_multiplier)\n\n# Loss function\ncriterion = nn.MSELoss(reduction='none')\n\n# Early stopping\nbest_val_loss = float('inf')\npatience = 10\npatience_counter = 0\nbest_epoch = -1\n\nprint(f\"Training for up to {num_epochs} epochs with early stopping (patience={patience})...\")\nprint(\"=\" * 80)\n\n# Mixed precision training (faster and uses less memory)\nscaler = torch.cuda.amp.GradScaler() if DEVICE == 'cuda' else None\n\n# Enable optimization\nif DEVICE == 'cuda':\n    torch.backends.cudnn.benchmark = True\n\nfor epoch in range(num_epochs):\n    epoch_start = time.time()\n    \n    # ========== TRAINING ==========\n    model.train()\n    train_loss = 0\n    train_batches = 0\n    \n    for batch_idx, (sequences, coordinates, lengths) in enumerate(train_loader):\n        sequences = sequences.to(DEVICE, non_blocking=True)\n        coordinates = coordinates.to(DEVICE, non_blocking=True)\n        \n        # Create padding mask for transformer\n        batch_size = sequences.size(0)\n        seq_length = sequences.size(1)\n        lengths_tensor = torch.tensor(lengths, dtype=torch.long, device=DEVICE)\n        positions = torch.arange(seq_length, device=DEVICE).unsqueeze(0).expand(batch_size, -1)\n        \n        # Transformer mask (True = ignore)\n        padding_mask = positions >= lengths_tensor.unsqueeze(1)\n        # Loss mask (1.0 = use)\n        loss_mask = (positions < lengths_tensor.unsqueeze(1)).float().unsqueeze(-1)\n        \n        optimizer.zero_grad(set_to_none=True)\n        \n        # Mixed precision training\n        if scaler is not None:\n            with torch.cuda.amp.autocast():\n                predictions = model(sequences, src_key_padding_mask=padding_mask)\n                loss_per_coord = criterion(predictions, coordinates) * loss_mask\n                loss = loss_per_coord.sum() / (loss_mask.sum() * 3 + 1e-8)\n            \n            scaler.scale(loss).backward()\n            scaler.unscale_(optimizer)\n            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n            scaler.step(optimizer)\n            scaler.update()\n        else:\n            predictions = model(sequences, src_key_padding_mask=padding_mask)\n            loss_per_coord = criterion(predictions, coordinates) * loss_mask\n            loss = loss_per_coord.sum() / (loss_mask.sum() * 3 + 1e-8)\n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n            optimizer.step()\n        \n        train_loss += loss.item()\n        train_batches += 1\n        \n        if batch_idx % 100 == 0 and batch_idx > 0:\n            print(f\"  Epoch {epoch+1} - Batch {batch_idx}/{len(train_loader)} - Loss: {loss.item():.4f}\")\n    \n    avg_train_loss = train_loss / train_batches\n    \n    # ========== VALIDATION ==========\n    model.eval()\n    val_loss = 0\n    val_batches = 0\n    \n    with torch.no_grad():\n        for sequences, coordinates, lengths in val_loader:\n            sequences = sequences.to(DEVICE, non_blocking=True)\n            coordinates = coordinates.to(DEVICE, non_blocking=True)\n            \n            batch_size = sequences.size(0)\n            seq_length = sequences.size(1)\n            lengths_tensor = torch.tensor(lengths, dtype=torch.long, device=DEVICE)\n            positions = torch.arange(seq_length, device=DEVICE).unsqueeze(0).expand(batch_size, -1)\n            \n            padding_mask = positions >= lengths_tensor.unsqueeze(1)\n            loss_mask = (positions < lengths_tensor.unsqueeze(1)).float().unsqueeze(-1)\n            \n            if scaler is not None:\n                with torch.cuda.amp.autocast():\n                    predictions = model(sequences, src_key_padding_mask=padding_mask)\n                    loss_per_coord = criterion(predictions, coordinates) * loss_mask\n                    loss = loss_per_coord.sum() / (loss_mask.sum() * 3 + 1e-8)\n            else:\n                predictions = model(sequences, src_key_padding_mask=padding_mask)\n                loss_per_coord = criterion(predictions, coordinates) * loss_mask\n                loss = loss_per_coord.sum() / (loss_mask.sum() * 3 + 1e-8)\n            \n            val_loss += loss.item()\n            val_batches += 1\n    \n    avg_val_loss = val_loss / val_batches\n    \n    # Update learning rate\n    scheduler.step()\n    \n    epoch_time = time.time() - epoch_start\n    current_lr = optimizer.param_groups[0]['lr']\n    \n    # Print progress\n    if DEVICE == 'cuda':\n        gpu_mem = torch.cuda.max_memory_allocated() / 1e9\n        print(f\"Epoch {epoch+1}/{num_epochs} - Train Loss: {avg_train_loss:.4f} - Val Loss: {avg_val_loss:.4f} - LR: {current_lr:.6f} - Time: {epoch_time:.1f}s - GPU: {gpu_mem:.2f}GB\")\n        torch.cuda.reset_peak_memory_stats()\n    else:\n        print(f\"Epoch {epoch+1}/{num_epochs} - Train Loss: {avg_train_loss:.4f} - Val Loss: {avg_val_loss:.4f} - LR: {current_lr:.6f} - Time: {epoch_time:.1f}s\")\n    \n    # Early stopping and model saving\n    if avg_val_loss < best_val_loss:\n        best_val_loss = avg_val_loss\n        best_epoch = epoch\n        patience_counter = 0\n        print(f\"  ✓ New best validation loss: {best_val_loss:.4f} - Model saved!\")\n        # Save best model checkpoint\n        torch.save({\n            'epoch': epoch,\n            'model_state_dict': model.state_dict(),\n            'optimizer_state_dict': optimizer.state_dict(),\n            'val_loss': best_val_loss,\n        }, 'best_model.pth')\n    else:\n        patience_counter += 1\n        print(f\"  No improvement ({patience_counter}/{patience})\")\n        \n        if patience_counter >= patience:\n            print(f\"\\nEarly stopping triggered at epoch {epoch+1}\")\n            break\n\nprint(\"=\" * 80)\nprint(f\"Training complete! Best validation loss: {best_val_loss:.4f}\")\n\n# Load best model for predictions (if it was saved)\nif os.path.exists('best_model.pth'):\n    print(\"Loading best model checkpoint for inference...\")\n    checkpoint = torch.load('best_model.pth')\n    model.load_state_dict(checkpoint['model_state_dict'])\n    print(f\"Loaded model from epoch {checkpoint['epoch']+1} with val loss {checkpoint['val_loss']:.4f}\")\nelse:\n    print(\"No checkpoint saved (using current model state)\")\n    if best_epoch >= 0:\n        print(f\"Note: Best validation loss was {best_val_loss:.4f} at epoch {best_epoch+1}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-08T07:54:41.327236Z","iopub.execute_input":"2026-02-08T07:54:41.327508Z","iopub.status.idle":"2026-02-08T08:00:23.985458Z","shell.execute_reply.started":"2026-02-08T07:54:41.327489Z","shell.execute_reply":"2026-02-08T08:00:23.984785Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def generate_predictions(test_df, sample_submission, model, device):\n    \"\"\"Generate 5 diverse predictions per sequence using improved strategies\"\"\"\n    test_id_col = 'target_id' if 'target_id' in test_df.columns else 'ID'\n    model.eval()\n    \n    results = []\n    max_model_length = 500\n    \n    # Helper function for ML prediction with options\n    def get_ml_prediction(seq, enable_dropout=False, temperature=1.0):\n        \"\"\"Get ML prediction with optional dropout and temperature scaling\"\"\"\n        if len(seq) <= max_model_length:\n            encoded = encode_sequence(seq)\n            encoded = np.pad(encoded, (0, max_model_length - len(encoded)), constant_values=4)\n            seq_tensor = torch.LongTensor(encoded).unsqueeze(0).to(device)\n            \n            # Create padding mask\n            seq_len = len(seq)\n            padding_mask = torch.zeros(1, max_model_length, dtype=torch.bool, device=device)\n            padding_mask[0, seq_len:] = True\n            \n            if enable_dropout:\n                model.train()  # Enable dropout\n                with torch.no_grad():\n                    coords = model(seq_tensor, src_key_padding_mask=padding_mask).cpu().numpy()[0, :seq_len]\n                model.eval()\n            else:\n                with torch.no_grad():\n                    coords = model(seq_tensor, src_key_padding_mask=padding_mask).cpu().numpy()[0, :seq_len]\n            \n            # Apply temperature scaling (diversity control)\n            if temperature != 1.0:\n                center = coords.mean(axis=0)\n                coords = center + (coords - center) * temperature\n            \n            return coords\n        else:\n            # Handle long sequences with overlapping chunks\n            coords_list = []\n            overlap = 50\n            \n            for start_idx in range(0, len(seq), max_model_length - overlap):\n                end_idx = min(start_idx + max_model_length, len(seq))\n                chunk = seq[start_idx:end_idx]\n                chunk_len = len(chunk)\n                \n                encoded = encode_sequence(chunk)\n                encoded = np.pad(encoded, (0, max_model_length - len(encoded)), constant_values=4)\n                seq_tensor = torch.LongTensor(encoded).unsqueeze(0).to(device)\n                \n                padding_mask = torch.zeros(1, max_model_length, dtype=torch.bool, device=device)\n                padding_mask[0, chunk_len:] = True\n                \n                if enable_dropout:\n                    model.train()\n                    with torch.no_grad():\n                        chunk_coords = model(seq_tensor, src_key_padding_mask=padding_mask).cpu().numpy()[0, :chunk_len]\n                    model.eval()\n                else:\n                    with torch.no_grad():\n                        chunk_coords = model(seq_tensor, src_key_padding_mask=padding_mask).cpu().numpy()[0, :chunk_len]\n                \n                # Average overlapping regions\n                if coords_list and start_idx > 0:\n                    overlap_size = min(overlap, chunk_len)\n                    if overlap_size > 0:\n                        prev_end = coords_list[-1][-overlap_size:]\n                        curr_start = chunk_coords[:overlap_size]\n                        averaged = (prev_end + curr_start) / 2\n                        coords_list[-1] = coords_list[-1][:-overlap_size]\n                        chunk_coords[:overlap_size] = averaged\n                \n                coords_list.append(chunk_coords)\n            \n            coords = np.vstack(coords_list)[:len(seq)]\n            \n            # Apply temperature scaling\n            if temperature != 1.0:\n                center = coords.mean(axis=0)\n                coords = center + (coords - center) * temperature\n            \n            return coords\n    \n    for idx, row in test_df.iterrows():\n        seq_id = row[test_id_col]\n        sequence = row['sequence']\n        n_residues = len(sequence)\n        \n        if (idx + 1) % 5 == 0 or idx == 0:\n            print(f\"Processing {idx+1}/{len(test_df)}: {seq_id} (length: {n_residues})\")\n        \n        predictions = []\n        \n        # IMPROVED DIVERSITY STRATEGY\n        # Instead of random noise, use principled approaches that create meaningful structural alternatives\n        \n        # Prediction 1: Base ML prediction (best single model output)\n        coords_base = get_ml_prediction(sequence, enable_dropout=False, temperature=1.0)\n        coords_base = np.nan_to_num(coords_base, nan=0.0, posinf=1000.0, neginf=-1000.0)\n        predictions.append(coords_base)\n        \n        # Prediction 2: ML with dropout (Monte Carlo dropout for uncertainty)\n        # This captures model uncertainty and produces a different but valid structure\n        coords_dropout = get_ml_prediction(sequence, enable_dropout=True, temperature=1.0)\n        coords_dropout = np.nan_to_num(coords_dropout, nan=0.0, posinf=1000.0, neginf=-1000.0)\n        predictions.append(coords_dropout)\n        \n        # Prediction 3: Temperature-scaled prediction (more compact structure)\n        # Temperature < 1.0 = more conservative/compact\n        coords_compact = get_ml_prediction(sequence, enable_dropout=False, temperature=0.85)\n        coords_compact = np.nan_to_num(coords_compact, nan=0.0, posinf=1000.0, neginf=-1000.0)\n        predictions.append(coords_compact)\n        \n        # Prediction 4: Temperature-scaled prediction (more expanded structure)\n        # Temperature > 1.0 = more expanded/diverse\n        coords_expanded = get_ml_prediction(sequence, enable_dropout=False, temperature=1.15)\n        coords_expanded = np.nan_to_num(coords_expanded, nan=0.0, posinf=1000.0, neginf=-1000.0)\n        predictions.append(coords_expanded)\n        \n        # Prediction 5: Dropout + expanded (combining uncertainty with expansion)\n        coords_diverse = get_ml_prediction(sequence, enable_dropout=True, temperature=1.1)\n        coords_diverse = np.nan_to_num(coords_diverse, nan=0.0, posinf=1000.0, neginf=-1000.0)\n        predictions.append(coords_diverse)\n        \n        # Create submission rows\n        for res_idx in range(n_residues):\n            row_data = {\n                'ID': f\"{seq_id}_{res_idx + 1}\",\n                'resname': sequence[res_idx],\n                'resid': res_idx + 1\n            }\n            \n            for pred_idx in range(5):\n                x, y, z = predictions[pred_idx][res_idx]\n                row_data[f'x_{pred_idx+1}'] = x\n                row_data[f'y_{pred_idx+1}'] = y\n                row_data[f'z_{pred_idx+1}'] = z\n            \n            results.append(row_data)\n    \n    submission_df = pd.DataFrame(results)\n    submission_df = submission_df[sample_submission.columns]\n    \n    submission_df['ID'] = submission_df['ID'].astype(str)\n    submission_df['resname'] = submission_df['resname'].astype(str)\n    submission_df['resid'] = submission_df['resid'].astype('int64')\n    \n    # Convert coordinates to integers\n    coord_cols = [f'{c}_{i}' for c in ['x', 'y', 'z'] for i in range(1, 6)]\n    for col in coord_cols:\n        submission_df[col] = submission_df[col].fillna(0)\n        submission_df[col] = submission_df[col].replace([np.inf, -np.inf], [1000, -1000])\n        submission_df[col] = submission_df[col].astype('int64')\n    \n    return submission_df\n\n\nprint(\"\\nGenerating predictions with improved diversity strategy...\")\nsubmission = generate_predictions(test_seq_df, sample_submission, model, DEVICE)\n\nprint(f\"\\nValidation:\")\nprint(f\"Shape: {submission.shape}\")\nprint(f\"Expected shape: {sample_submission.shape}\")\nprint(f\"Missing values: {submission.isna().sum().sum()}\")\nprint(f\"\\nFirst few rows:\")\nprint(submission.head())\n\nsubmission.to_csv('submission.csv', index=False)\nprint(\"\\nSubmission saved to submission.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-08T08:00:23.986288Z","iopub.execute_input":"2026-02-08T08:00:23.986503Z","iopub.status.idle":"2026-02-08T08:00:25.61938Z","shell.execute_reply.started":"2026-02-08T08:00:23.986483Z","shell.execute_reply":"2026-02-08T08:00:25.618812Z"}},"outputs":[],"execution_count":null}]}