{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":118765,"databundleVersionId":15231210},{"sourceType":"kernelVersion","sourceId":306402703}],"dockerImageVersionId":31329,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# =================================================\n# RNA 3D Folding Challenge - Complete Pipeline\n# =================================================\n\nimport os\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm import tqdm\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# Install Biopython if needed\ntry:\n    from Bio import AlignIO\nexcept ImportError:\n    !pip install biopython -q\n    from Bio import AlignIO\n\n# Set device\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\n\n# ------------------------------\n# 1. Data loading\n# ------------------------------\ntrain_seq = pd.read_csv('/kaggle/input/competitions/stanford-rna-3d-folding-2/train_sequences.csv')\ntrain_labels = pd.read_csv('/kaggle/input/competitions/stanford-rna-3d-folding-2/train_labels.csv')\nval_seq = pd.read_csv('/kaggle/input/competitions/stanford-rna-3d-folding-2/validation_sequences.csv')\nval_labels = pd.read_csv('/kaggle/input/competitions/stanford-rna-3d-folding-2/validation_labels.csv')\ntest_seq = pd.read_csv('/kaggle/input/competitions/stanford-rna-3d-folding-2/test_sequences.csv')\n\nprint(\"Train columns:\", train_seq.columns.tolist())\nprint(\"Label columns:\", train_labels.columns.tolist())\nprint(f\"Number of train sequences: {len(train_seq)}\")\nprint(f\"Number of train residues: {len(train_labels)}\")\n\n# ------------------------------\n# 2. Preprocessing helpers\n# ------------------------------\nnuc_to_int = {'A':0, 'U':1, 'G':2, 'C':3, 'N':4, 'X':4}\ndef seq_to_tensor(seq):\n    return torch.tensor([nuc_to_int.get(c, 4) for c in seq], dtype=torch.long)\n\n# ------------------------------\n# 3. MSA feature extractor\n# ------------------------------\ndef compute_msa_features(msa_file, seq):\n    \"\"\"\n    Compute position-specific frequencies of A,U,G,C and gap.\n    Returns tensor of shape (L, 5) where last dimension is [freq_A, freq_U, freq_G, freq_C, freq_gap]\n    \"\"\"\n    if not os.path.exists(msa_file):\n        return torch.zeros((len(seq), 5))\n    try:\n        alignment = AlignIO.read(msa_file, \"fasta\")\n    except:\n        return torch.zeros((len(seq), 5))\n    L = len(seq)\n    counts = np.zeros((L, 5))  # A,U,G,C,gap\n    for record in alignment:\n        seq_msa = str(record.seq).upper()\n        for i, c in enumerate(seq_msa):\n            if i >= L: break\n            if c == 'A': counts[i,0] += 1\n            elif c == 'U': counts[i,1] += 1\n            elif c == 'G': counts[i,2] += 1\n            elif c == 'C': counts[i,3] += 1\n            else: counts[i,4] += 1  # gap or other\n    # Normalize to frequencies\n    row_sums = counts.sum(axis=1, keepdims=True)\n    row_sums[row_sums == 0] = 1\n    counts = counts / row_sums\n    return torch.tensor(counts, dtype=torch.float32)\n\n# ------------------------------\n# 4. Dataset\n# ------------------------------\nclass RNADataset(Dataset):\n    def __init__(self, seq_df, label_df=None, msa_dir=None):\n        self.seq_df = seq_df\n        self.label_df = label_df\n        self.msa_dir = msa_dir\n        # Get sequences\n        self.seqs = seq_df['sequence'].tolist()\n        # Get IDs: train/val use 'target_id', test uses 'id'\n        if 'target_id' in seq_df.columns:\n            self.ids = seq_df['target_id'].tolist()\n        elif 'id' in seq_df.columns:\n            self.ids = seq_df['id'].tolist()\n        else:\n            # fallback to index\n            self.ids = [f\"seq_{i}\" for i in range(len(seq_df))]\n        # Build label dictionary if labels provided\n        self.labels_dict = None\n        if label_df is not None:\n            self.labels_dict = {}\n            grouped = label_df.groupby('ID')\n            print(\"Building label dictionary...\")\n            for seq_id, group in tqdm(grouped, total=len(grouped), desc=\"Labels\"):\n                # Sort by chain, copy, residue number to ensure correct order\n                group = group.sort_values(['chain', 'copy', 'resid'])\n                coords = group[['x_1', 'y_1', 'z_1']].values.astype(np.float32)\n                self.labels_dict[seq_id] = torch.tensor(coords, dtype=torch.float32)\n\n    def __len__(self):\n        return len(self.seqs)\n\n    def __getitem__(self, idx):\n        seq = self.seqs[idx]\n        seq_id = self.ids[idx]\n        seq_tensor = seq_to_tensor(seq)\n        # Get coordinates if available\n        if self.labels_dict is not None and seq_id in self.labels_dict:\n            coords = self.labels_dict[seq_id]\n            # Ensure length matches sequence length (truncate/pad if needed)\n            if coords.shape[0] != len(seq):\n                if coords.shape[0] > len(seq):\n                    coords = coords[:len(seq)]\n                else:\n                    pad = torch.zeros((len(seq) - coords.shape[0], 3), dtype=torch.float32)\n                    coords = torch.cat([coords, pad], dim=0)\n        else:\n            coords = torch.zeros((len(seq), 3), dtype=torch.float32)\n        # MSA features\n        msa_features = None\n        if self.msa_dir is not None:\n            msa_file = os.path.join(self.msa_dir, f\"{seq_id}.MSA.fasta\")\n            msa_features = compute_msa_features(msa_file, seq)\n        return seq_tensor, coords, msa_features, len(seq)\n\ndef collate_fn(batch):\n    seq_tensors, coords, msa_features, lengths = zip(*batch)\n    max_len = max(lengths)\n    B = len(batch)\n    padded_seqs = torch.zeros(B, max_len, dtype=torch.long)\n    padded_coords = torch.zeros(B, max_len, 3)\n    for i, (s, c, l) in enumerate(zip(seq_tensors, coords, lengths)):\n        padded_seqs[i, :l] = s\n        padded_coords[i, :l] = c\n    # Pad MSA features\n    if msa_features[0] is not None:\n        feat_dim = msa_features[0].size(-1)\n        padded_msa = torch.zeros(B, max_len, feat_dim)\n        for i, m in enumerate(msa_features):\n            if m is not None:\n                padded_msa[i, :lengths[i]] = m\n    else:\n        padded_msa = None\n    return padded_seqs, padded_coords, padded_msa, torch.tensor(lengths)\n\n# ------------------------------\n# 5. Model definition\n# ------------------------------\nclass RNATransformer(nn.Module):\n    def __init__(self, vocab_size=5, embed_dim=128, num_layers=6, num_heads=8, dropout=0.1, use_msa=True):\n        super().__init__()\n        self.use_msa = use_msa\n        self.embed = nn.Embedding(vocab_size, embed_dim)\n        self.pos_embed = nn.Parameter(torch.randn(1, 1024, embed_dim))\n        if use_msa:\n            self.msa_linear = nn.Linear(5, embed_dim)\n        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dropout=dropout, batch_first=True)\n        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)\n        self.coord_head = nn.Linear(embed_dim, 3)\n\n    def forward(self, seq_tokens, msa_features=None):\n        L = seq_tokens.size(1)\n        x = self.embed(seq_tokens)                       # (B, L, D)\n        x = x + self.pos_embed[:, :L, :]\n        if self.use_msa and msa_features is not None:\n            msa_emb = self.msa_linear(msa_features)\n            x = x + msa_emb\n        x = self.transformer(x)\n        coords = self.coord_head(x)\n        return coords\n\n# ------------------------------\n# 6. RMSD loss with Kabsch algorithm\n# ------------------------------\ndef kabsch_rmsd(pred, true):\n    # pred, true: (L, 3)\n    pred_center = pred - pred.mean(dim=0)\n    true_center = true - true.mean(dim=0)\n    H = pred_center.T @ true_center\n    U, S, Vt = torch.linalg.svd(H)\n    R = Vt.T @ U.T\n    pred_aligned = pred_center @ R.T\n    diff = pred_aligned - true_center\n    rmsd = torch.sqrt((diff ** 2).sum() / pred.size(0))\n    return rmsd\n\ndef rmsd_loss(pred, true):\n    batch_loss = 0.0\n    for i in range(pred.size(0)):\n        batch_loss += kabsch_rmsd(pred[i], true[i])\n    return batch_loss / pred.size(0)\n\n# ------------------------------\n# 7. Training setup\n# ------------------------------\nprint(\"\\nCreating datasets...\")\ntrain_dataset = RNADataset(train_seq, train_labels, msa_dir='/kaggle/input/competitions/stanford-rna-3d-folding-2/MSA')\nval_dataset = RNADataset(val_seq, val_labels)  # no MSA for validation\ntest_dataset = RNADataset(test_seq)\n\nprint(\"Creating dataloaders...\")\ntrain_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, collate_fn=collate_fn)\nval_loader = DataLoader(val_dataset, batch_size=4, shuffle=False, collate_fn=collate_fn)\ntest_loader = DataLoader(test_dataset, batch_size=4, shuffle=False, collate_fn=collate_fn)\n\nmodel = RNATransformer(use_msa=True).to(device)\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\nscheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, factor=0.5)\n\n# ------------------------------\n# 8. Training loop\n# ------------------------------\nnum_epochs = 30\nprint(\"\\nStarting training...\")\nfor epoch in range(num_epochs):\n    model.train()\n    train_loss = 0.0\n    for seqs, coords, msa, lengths in tqdm(train_loader, desc=f\"Epoch {epoch+1} train\"):\n        seqs = seqs.to(device)\n        coords = coords.to(device)\n        msa = msa.to(device) if msa is not None else None\n        optimizer.zero_grad()\n        pred = model(seqs, msa)\n        loss = rmsd_loss(pred, coords)\n        loss.backward()\n        optimizer.step()\n        train_loss += loss.item()\n    train_loss /= len(train_loader)\n\n    # Validation\n    model.eval()\n    val_loss = 0.0\n    with torch.no_grad():\n        for seqs, coords, msa, lengths in tqdm(val_loader, desc=f\"Epoch {epoch+1} val\"):\n            seqs = seqs.to(device)\n            coords = coords.to(device)\n            pred = model(seqs, None)   # no MSA for validation\n            loss = rmsd_loss(pred, coords)\n            val_loss += loss.item()\n    val_loss /= len(val_loader)\n    print(f\"Epoch {epoch+1:2d}: train_loss = {train_loss:.4f}, val_loss = {val_loss:.4f}\")\n    scheduler.step(val_loss)\n\n# ------------------------------\n# 9. Inference and submission\n# ------------------------------\nprint(\"\\nGenerating predictions...\")\nmodel.eval()\npredictions = []\nwith torch.no_grad():\n    for seqs, _, _, lengths in tqdm(test_loader, desc=\"Inference\"):\n        seqs = seqs.to(device)\n        pred = model(seqs, None)   # no MSA for test\n        for i, l in enumerate(lengths):\n            coords = pred[i, :l].cpu().numpy()\n            coord_str = ' '.join([f\"{c:.6f}\" for coord in coords for c in coord])\n            predictions.append(coord_str)\n\n# Determine test ID column\nif 'target_id' in test_seq.columns:\n    test_ids = test_seq['target_id'].tolist()\nelif 'id' in test_seq.columns:\n    test_ids = test_seq['id'].tolist()\nelse:\n    test_ids = [f\"seq_{i}\" for i in range(len(test_seq))]\n\nsubmission = pd.DataFrame({'id': test_ids, 'predicted_coordinates': predictions})\nsubmission.to_csv('submission.csv', index=False)\nprint(\"Submission saved as submission.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T12:14:48.958134Z","iopub.execute_input":"2026-03-25T12:14:48.958407Z","iopub.status.idle":"2026-03-25T12:18:15.773193Z","shell.execute_reply.started":"2026-03-25T12:14:48.958381Z","shell.execute_reply":"2026-03-25T12:18:15.771962Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 2. Preprocess sequences and labels\n# Convert sequences to integer tokens: A=0, U=1, G=2, C=3, other=4\nnuc_to_int = {'A':0, 'U':1, 'G':2, 'C':3, 'N':4, 'X':4}\ndef seq_to_tensor(seq):\n    return torch.tensor([nuc_to_int.get(c, 4) for c in seq], dtype=torch.long)\n\n# Labels: each row contains coordinates as a space-separated string of length 3*L\ndef parse_coords(s):\n    coords = np.array(s.split(), dtype=np.float32).reshape(-1, 3)\n    return torch.tensor(coords, dtype=torch.float32)\n\n# Create dataset class\nclass RNADataset(Dataset):\n    def __init__(self, seq_df, label_df=None, msa_dir=None):\n        self.seq_df = seq_df\n        self.label_df = label_df\n        self.msa_dir = msa_dir\n        self.seqs = seq_df['sequence'].tolist()\n        self.ids = seq_df['sequence_id'].tolist()\n        if label_df is not None:\n            self.labels = [parse_coords(row) for row in label_df['coordinates']]\n        else:\n            self.labels = None\n\n    def __len__(self):\n        return len(self.seqs)\n\n    def __getitem__(self, idx):\n        seq = self.seqs[idx]\n        seq_id = self.ids[idx]\n        seq_tensor = seq_to_tensor(seq)\n        if self.labels is not None:\n            coords = self.labels[idx]\n        else:\n            coords = torch.zeros((len(seq), 3), dtype=torch.float32)  # placeholder\n        # Load MSA features if available\n        msa_features = None\n        if self.msa_dir is not None:\n            msa_file = os.path.join(self.msa_dir, f\"{seq_id}.MSA.fasta\")\n            if os.path.exists(msa_file):\n                # parse MSA and compute features\n                msa_features = self.compute_msa_features(msa_file, seq)\n            else:\n                # no MSA for this sequence, use zeros\n                msa_features = torch.zeros((len(seq), 5))  # e.g., frequency of 4 bases + gap\n        return seq_tensor, coords, msa_features, len(seq)\n\n    def compute_msa_features(self, msa_file, seq):\n        # Parse MSA, compute position-specific frequencies and entropy\n        # This is a placeholder; implement actual parsing\n        # We'll return a tensor of shape (L, 5) for A,U,G,C,gap frequencies\n        return torch.randn(len(seq), 5)  # dummy","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 1. Load data\ntrain_seq = pd.read_csv('/kaggle/input/competitions/stanford-rna-3d-folding-2/train_sequences.csv')\ntrain_labels = pd.read_csv('/kaggle/input/competitions/stanford-rna-3d-folding-2/train_labels.csv')\nval_seq = pd.read_csv('/kaggle/input/competitions/stanford-rna-3d-folding-2/validation_sequences.csv')\nval_labels = pd.read_csv('/kaggle/input/competitions/stanford-rna-3d-folding-2/validation_labels.csv')\ntest_seq = pd.read_csv('/kaggle/input/competitions/stanford-rna-3d-folding-2/test_sequences.csv')\nsample_sub = pd.read_csv('/kaggle/input/competitions/stanford-rna-3d-folding-2/sample_submission.csv')\n\n# Inspect data\nprint(train_seq.head())\nprint(train_labels.head())\nprint(f\"Train sequences: {len(train_seq)}, Test sequences: {len(test_seq)}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 3. Build model\nclass RNATransformer(nn.Module):\n    def __init__(self, vocab_size=5, embed_dim=128, num_layers=6, num_heads=8, dropout=0.1, use_msa=False):\n        super().__init__()\n        self.use_msa = use_msa\n        self.embed = nn.Embedding(vocab_size, embed_dim)\n        self.pos_embed = nn.Parameter(torch.randn(1, 1024, embed_dim))  # max length 1024\n        if use_msa:\n            self.msa_linear = nn.Linear(5, embed_dim)  # map MSA features to embedding\n        self.transformer = nn.TransformerEncoder(\n            nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dropout=dropout, batch_first=True),\n            num_layers=num_layers\n        )\n        self.coord_head = nn.Linear(embed_dim, 3)\n\n    def forward(self, seq_tokens, msa_features=None):\n        L = seq_tokens.size(1)\n        x = self.embed(seq_tokens)  # (B, L, embed_dim)\n        x = x + self.pos_embed[:, :L, :]\n        if self.use_msa and msa_features is not None:\n            msa_emb = self.msa_linear(msa_features)  # (B, L, embed_dim)\n            x = x + msa_emb\n        x = self.transformer(x)\n        coords = self.coord_head(x)  # (B, L, 3)\n        return coords\n\n# Loss function: RMSD after optimal superposition (Kabsch algorithm)\ndef kabsch_rmsd(pred, true):\n    # pred, true: (L, 3)\n    # Center both\n    pred_center = pred - pred.mean(dim=0)\n    true_center = true - true.mean(dim=0)\n    # Compute covariance matrix\n    H = pred_center.T @ true_center\n    # SVD\n    U, S, Vt = torch.linalg.svd(H)\n    # Compute rotation\n    R = Vt.T @ U.T\n    # Apply rotation\n    pred_aligned = pred_center @ R.T\n    # Compute RMSD\n    diff = pred_aligned - true_center\n    rmsd = torch.sqrt((diff ** 2).sum() / pred.size(0))\n    return rmsd\n\ndef rmsd_loss(pred, true):\n    batch_loss = 0.0\n    for i in range(pred.size(0)):\n        batch_loss += kabsch_rmsd(pred[i], true[i])\n    return batch_loss / pred.size(0)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class RNADataset(Dataset):\n    def __init__(self, seq_df, label_df=None, msa_dir=None):\n        self.seq_df = seq_df\n        self.label_df = label_df\n        self.msa_dir = msa_dir\n        # Determine sequence column\n        if 'sequence' in seq_df.columns:\n            self.seqs = seq_df['sequence'].tolist()\n        else:\n            raise KeyError(\"Sequence column 'sequence' not found\")\n        # Determine ID column in seq_df (could be 'target_id' or 'id' or 'sequence_id')\n        if 'target_id' in seq_df.columns:\n            self.ids = seq_df['target_id'].tolist()\n        elif 'id' in seq_df.columns:\n            self.ids = seq_df['id'].tolist()\n        elif 'sequence_id' in seq_df.columns:\n            self.ids = seq_df['sequence_id'].tolist()\n        else:\n            self.ids = [f\"seq_{i}\" for i in range(len(seq_df))]\n        # If labels are provided, build mapping from ID to coordinates\n        self.labels_dict = None\n        if label_df is not None:\n            # group by ID\n            grouped = label_df.groupby('ID')\n            self.labels_dict = {}\n            for seq_id, group in grouped:\n                # sort by resid (and possibly chain, copy)\n                group = group.sort_values(['resid', 'chain', 'copy'])\n                coords = group[['x_1', 'y_1', 'z_1']].values.astype(np.float32)\n                # Convert to torch tensor\n                self.labels_dict[seq_id] = torch.tensor(coords, dtype=torch.float32)\n\n    def __len__(self):\n        return len(self.seqs)\n\n    def __getitem__(self, idx):\n        seq = self.seqs[idx]\n        seq_id = self.ids[idx]\n        seq_tensor = seq_to_tensor(seq)\n        # Look up coordinates for this sequence\n        if self.labels_dict is not None and seq_id in self.labels_dict:\n            coords = self.labels_dict[seq_id]\n            # Ensure length matches sequence length (should, but we can truncate/pad if needed)\n            if coords.shape[0] != len(seq):\n                # This might happen if the label file has extra residues? Usually should match.\n                # We'll pad or truncate to the sequence length (though it's better to debug)\n                if coords.shape[0] > len(seq):\n                    coords = coords[:len(seq)]\n                else:\n                    pad = torch.zeros((len(seq) - coords.shape[0], 3), dtype=torch.float32)\n                    coords = torch.cat([coords, pad], dim=0)\n        else:\n            coords = torch.zeros((len(seq), 3), dtype=torch.float32)  # placeholder\n        # MSA features\n        msa_features = None\n        if self.msa_dir is not None:\n            msa_file = os.path.join(self.msa_dir, f\"{seq_id}.MSA.fasta\")\n            msa_features = compute_msa_features(msa_file, seq)\n        return seq_tensor, coords, msa_features, len(seq)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Import necessary libraries\n!pip install biopython\n!pip install transformers scipy\n\nimport os\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom transformers import AutoModel, AutoTokenizer  # for RNA language model (optional)\nfrom scipy.spatial.transform import Rotation\nfrom Bio import AlignIO\nfrom Bio.PDB import PDBParser, Selection\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# Set device\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}