{"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":"none","dataSources":[{"sourceId":87793,"databundleVersionId":11403143,"sourceType":"competition"}],"dockerImageVersionId":30918,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"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        pass\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,"execution":{"iopub.status.busy":"2025-03-15T03:46:35.284914Z","iopub.execute_input":"2025-03-15T03:46:35.285256Z","iopub.status.idle":"2025-03-15T03:46:38.709890Z","shell.execute_reply.started":"2025-03-15T03:46:35.285221Z","shell.execute_reply":"2025-03-15T03:46:38.708804Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport pandas as pd\n\n# Define the data folder\ndata_dir = \"/kaggle/input/stanford-rna-3d-folding\"\n\n# Define file paths for the sequence and label files\ntrain_sequences_path = os.path.join(data_dir, \"train_sequences.csv\")\ntrain_labels_path = os.path.join(data_dir, \"train_labels.csv\")\nval_sequences_path   = os.path.join(data_dir, \"validation_sequences.csv\")\nval_labels_path      = os.path.join(data_dir, \"validation_labels.csv\")\n\n# Load CSV files into DataFrames\ntrain_seq_df = pd.read_csv(train_sequences_path)\ntrain_labels_df = pd.read_csv(train_labels_path)\nval_seq_df = pd.read_csv(val_sequences_path)\nval_labels_df = pd.read_csv(val_labels_path)\n\n# Convert 'temporal_cutoff' column to datetime for proper filtering\ntrain_seq_df['temporal_cutoff'] = pd.to_datetime(train_seq_df['temporal_cutoff'], errors='coerce')\nval_seq_df['temporal_cutoff'] = pd.to_datetime(val_seq_df['temporal_cutoff'], errors='coerce')\n\n# Filter training sequences based on temporal cutoff.\n# As noted in the dataset description, only train sequences with a cutoff before 2022-05-27 should be used.\ncutoff_date = pd.to_datetime('2022-05-27')\nfiltered_train_seq_df = train_seq_df[train_seq_df['temporal_cutoff'] < cutoff_date].copy()\n\nprint(\"Number of training sequences after filtering:\", len(filtered_train_seq_df))\nprint(\"Number of validation sequences:\", len(val_seq_df))\n\n# The label files have an 'ID' column in the form 'target_id_resid'\n# We extract the target_id (everything before the last underscore) for merging.\ntrain_labels_df['target_id'] = train_labels_df['ID'].str.rsplit('_', n=1, expand=True)[0]\nval_labels_df['target_id']   = val_labels_df['ID'].str.rsplit('_', n=1, expand=True)[0]\n\n# Merge the sequences with the labels based on 'target_id'\n# Note: Each RNA sequence (in train_seq_df or val_seq_df) corresponds to multiple rows in the labels files (one per residue)\ntrain_data = pd.merge(filtered_train_seq_df, train_labels_df, on=\"target_id\", how=\"inner\")\nval_data   = pd.merge(val_seq_df, val_labels_df, on=\"target_id\", how=\"inner\")\n\nprint(\"Training data shape (after merging):\", train_data.shape)\nprint(\"Validation data shape (after merging):\", val_data.shape)\n\n# Optionally, save the processed data for later use\ntrain_data.to_csv(\"processed_train_data.csv\", index=False)\nval_data.to_csv(\"processed_val_data.csv\", index=False)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-15T04:02:43.193685Z","iopub.execute_input":"2025-03-15T04:02:43.194045Z","iopub.status.idle":"2025-03-15T04:03:20.240473Z","shell.execute_reply.started":"2025-03-15T04:02:43.194016Z","shell.execute_reply":"2025-03-15T04:03:20.239496Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\n\n# --------------------------\n# Utility Functions and Setup\n# --------------------------\ntorch.manual_seed(42)\n\n# Mapping nucleotides (with unknown token for non-ACGU)\nnucleotide_to_idx = {'A': 0, 'C': 1, 'G': 2, 'U': 3}\nunknown_token_index = 4  # for any unknown characters\nvocab_size = len(nucleotide_to_idx) + 1  # total vocab size\n\ndef sequence_to_tensor(seq):\n    \"\"\"Convert nucleotide sequence string to tensor of indices; unknown tokens get unknown_token_index.\"\"\"\n    indices = [nucleotide_to_idx.get(nt, unknown_token_index) for nt in seq]\n    return torch.tensor(indices, dtype=torch.long)\n\n# --------------------------\n# Dataset Definition with Filtering\n# --------------------------\n\nclass RNA3DDataset(Dataset):\n    \"\"\"\n    Groups processed data by target_id.\n    Returns:\n      - seq_tensor: Tensor of indices (one per residue) from the 'resname' column.\n      - coords_tensor: Tensor of ground truth coordinates (from x_1, y_1, z_1) for each residue.\n    \"\"\"\n    def __init__(self, processed_csv, coord_threshold=-1e10):\n        self.data = pd.read_csv(processed_csv)\n        \n        # Ensure coordinate columns are numeric\n        for col in ['x_1', 'y_1', 'z_1']:\n            self.data[col] = pd.to_numeric(self.data[col], errors='coerce')\n        \n        # Drop rows with NaN in coordinate columns\n        self.data.dropna(subset=['x_1', 'y_1', 'z_1'], inplace=True)\n        \n        # Filter out rows where any coordinate is extremely low (likely a placeholder)\n        mask = (self.data['x_1'] > coord_threshold) & \\\n               (self.data['y_1'] > coord_threshold) & \\\n               (self.data['z_1'] > coord_threshold)\n        self.data = self.data[mask]\n        \n        # Group by target_id so that each group represents one RNA target\n        self.groups = self.data.groupby('target_id')\n        self.target_ids = list(self.groups.groups.keys())\n    \n    def __len__(self):\n        return len(self.target_ids)\n    \n    def __getitem__(self, idx):\n        target_id = self.target_ids[idx]\n        group = self.groups.get_group(target_id)\n        \n        # Construct the input sequence using the 'resname' column\n        residue_list = group['resname'].tolist()\n        seq_tensor = torch.tensor(\n            [nucleotide_to_idx.get(nt, unknown_token_index) for nt in residue_list],\n            dtype=torch.long\n        )\n        \n        # Extract the coordinates (using only x_1, y_1, z_1)\n        coords = group[['x_1', 'y_1', 'z_1']].values  # shape: (num_residues, 3)\n        coords_tensor = torch.tensor(coords, dtype=torch.float)\n        \n        # Sanity check: the length of the sequence should equal number of coordinate rows\n        if seq_tensor.size(0) != coords_tensor.size(0):\n            raise ValueError(f\"Mismatch for target_id {target_id}: sequence length {seq_tensor.size(0)} vs. {coords_tensor.size(0)} coordinates\")\n        \n        # Optionally, you might normalize coordinates per target here (e.g., center them)\n        # coords_tensor = coords_tensor - coords_tensor.mean(dim=0, keepdim=True)\n        \n        return seq_tensor, coords_tensor\n\n# --------------------------\n# Positional Encoding Module\n# --------------------------\n\nclass PositionalEncoding(nn.Module):\n    def __init__(self, d_model, max_len=5000):\n        super(PositionalEncoding, self).__init__()\n        pe = torch.zeros(max_len, d_model)  # (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() * (-torch.log(torch.tensor(10000.0)) / d_model))\n        pe[:, 0::2] = torch.sin(position * div_term)\n        if d_model % 2 == 1:\n            pe[:, 1::2] = torch.cos(position * div_term[:pe[:, 1::2].shape[1]])\n        else:\n            pe[:, 1::2] = torch.cos(position * div_term)\n        pe = pe.unsqueeze(0)  # (1, max_len, d_model)\n        self.register_buffer('pe', pe)\n    \n    def forward(self, x):\n        # x: (batch_size, seq_len, d_model)\n        return x + self.pe[:, :x.size(1)]\n\n# --------------------------\n# RNATransformer Model Definition\n# --------------------------\n\nclass RNATransformer(nn.Module):\n    def __init__(self, d_model=128, nhead=8, num_layers=3, vocab_size=vocab_size):\n        super(RNATransformer, self).__init__()\n        self.embedding = nn.Embedding(vocab_size, d_model)\n        self.pos_encoder = PositionalEncoding(d_model)\n        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, batch_first=True)\n        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)\n        self.fc_out = nn.Linear(d_model, 3)  # Predict 3 coordinates per residue\n        \n    def forward(self, src):\n        # src: (batch_size, seq_len) with token indices\n        x = self.embedding(src)       # (batch_size, seq_len, d_model)\n        x = self.pos_encoder(x)\n        x = self.transformer_encoder(x)  # (batch_size, seq_len, d_model)\n        coords = self.fc_out(x)          # (batch_size, seq_len, 3)\n        # Center predicted coordinates for translation invariance\n        coords = coords - coords.mean(dim=1, keepdim=True)\n        return coords\n\n# --------------------------\n# Data Loading and Preparation\n# --------------------------\ntrain_csv = \"processed_train_data.csv\"\nval_csv   = \"processed_val_data.csv\"  # or you may use a filtered version if you saved one\n\ntrain_dataset = RNA3DDataset(train_csv)\nval_dataset = RNA3DDataset(val_csv)\n\n# Use batch size 1 because sequences have variable lengths\ntrain_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)\nval_loader = DataLoader(val_dataset, batch_size=1, shuffle=False)\n\n# --------------------------\n# Model Training Setup\n# --------------------------\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel = RNATransformer(d_model=128, nhead=8, num_layers=3, vocab_size=vocab_size).to(device)\ncriterion = nn.MSELoss()\noptimizer = optim.Adam(model.parameters(), lr=1e-4)\nnum_epochs = 10\nmax_grad_norm = 1.0  # For gradient clipping\n\n# --------------------------\n# Training Loop\n# --------------------------\nfor epoch in range(num_epochs):\n    model.train()\n    train_loss = 0.0\n    for seq_tensor, coords_tensor in train_loader:\n        seq_tensor = seq_tensor.to(device)         # shape: (1, seq_len)\n        coords_tensor = coords_tensor.to(device)     # shape: (1, seq_len, 3)\n        \n        optimizer.zero_grad()\n        pred_coords = model(seq_tensor)              # (1, seq_len, 3)\n        \n        # Optional: check for NaNs in predictions\n        if torch.isnan(pred_coords).any():\n            print(\"NaN detected in predictions for one batch!\")\n            continue\n        \n        loss = criterion(pred_coords, coords_tensor)\n        loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)\n        optimizer.step()\n        \n        train_loss += loss.item()\n    \n    avg_train_loss = train_loss / len(train_loader)\n    \n    # Validation loop\n    model.eval()\n    val_loss = 0.0\n    with torch.no_grad():\n        for seq_tensor, coords_tensor in val_loader:\n            seq_tensor = seq_tensor.to(device)\n            coords_tensor = coords_tensor.to(device)\n            pred_coords = model(seq_tensor)\n            loss = criterion(pred_coords, coords_tensor)\n            val_loss += loss.item()\n    avg_val_loss = val_loss / len(val_loader)\n    \n    print(f\"Epoch {epoch+1}/{num_epochs} - Train Loss: {avg_train_loss:.4f} - Val Loss: {avg_val_loss:.4f}\")\n\ntorch.save(model.state_dict(), \"rna_transformer_model.pth\")\nprint(\"Model training complete and saved.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-15T04:03:24.324227Z","iopub.execute_input":"2025-03-15T04:03:24.324585Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport pandas as pd\nimport torch\nimport numpy as np\n\n# --------------------------\n# Utility: Map nucleotides to indices\n# --------------------------\nnucleotide_to_idx = {'A': 0, 'C': 1, 'G': 2, 'U': 3}\nunknown_token_index = 4\nvocab_size = len(nucleotide_to_idx) + 1\n\ndef sequence_to_tensor(seq):\n    \"\"\"Converts a nucleotide sequence string to a tensor of indices.\n       Unknown tokens are mapped to unknown_token_index.\"\"\"\n    indices = [nucleotide_to_idx.get(nt, unknown_token_index) for nt in seq]\n    return torch.tensor(indices, dtype=torch.long)\n\n# --------------------------\n# Model Definition (should match your training code)\n# --------------------------\nimport torch.nn as nn\n\nclass PositionalEncoding(nn.Module):\n    def __init__(self, d_model, max_len=5000):\n        super(PositionalEncoding, self).__init__()\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() * (-torch.log(torch.tensor(10000.0)) / d_model))\n        pe[:, 0::2] = torch.sin(position * div_term)\n        if d_model % 2 == 1:\n            pe[:, 1::2] = torch.cos(position * div_term[:pe[:, 1::2].shape[1]])\n        else:\n            pe[:, 1::2] = torch.cos(position * div_term)\n        pe = pe.unsqueeze(0)  # shape: (1, max_len, d_model)\n        self.register_buffer('pe', pe)\n    \n    def forward(self, x):\n        return x + self.pe[:, :x.size(1)]\n\nclass RNATransformer(nn.Module):\n    def __init__(self, d_model=128, nhead=8, num_layers=3, vocab_size=vocab_size):\n        super(RNATransformer, self).__init__()\n        self.embedding = nn.Embedding(vocab_size, d_model)\n        self.pos_encoder = PositionalEncoding(d_model)\n        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, batch_first=True)\n        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)\n        self.fc_out = nn.Linear(d_model, 3)  # predict 3 coordinates per residue\n        \n    def forward(self, src):\n        # src shape: (batch_size, seq_len)\n        x = self.embedding(src)            # (batch_size, seq_len, d_model)\n        x = self.pos_encoder(x)\n        x = self.transformer_encoder(x)      # (batch_size, seq_len, d_model)\n        coords = self.fc_out(x)              # (batch_size, seq_len, 3)\n        # Center coordinates (translation invariance)\n        coords = coords - coords.mean(dim=1, keepdim=True)\n        return coords\n\n# --------------------------\n# Load the Trained Model\n# --------------------------\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel = RNATransformer(d_model=128, nhead=8, num_layers=3, vocab_size=vocab_size).to(device)\nmodel.load_state_dict(torch.load(\"rna_transformer_model.pth\", map_location=device))\n# We'll use Monte Carlo dropout to generate diverse predictions:\n# Set the model to train mode so that dropout is active at inference.\nmodel.train()\n\n# --------------------------\n# Function to Generate Five Predictions Using Monte Carlo Dropout\n# --------------------------\ndef predict_five_structures(model, seq_tensor, n_predictions=5):\n    \"\"\"\n    Generates n_predictions for the given sequence tensor.\n    seq_tensor: a tensor of shape (seq_len,) representing nucleotide indices.\n    Returns a list of numpy arrays each of shape (seq_len, 3).\n    \"\"\"\n    predictions = []\n    with torch.no_grad():\n        for i in range(n_predictions):\n            # Note: by leaving dropout on (model.train()), we sample different predictions.\n            pred = model(seq_tensor.unsqueeze(0))  # (1, seq_len, 3)\n            predictions.append(pred.squeeze(0).cpu().numpy())\n    return predictions\n\n# --------------------------\n# Generate Submission File\n# --------------------------\n# Load the test sequences.\ntest_df = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/test_sequences.csv\")\n\nsubmission_rows = []\n\n# Loop over each target in the test set.\nfor idx, row in test_df.iterrows():\n    target_id = row['target_id']\n    sequence = row['sequence'].strip()  # assume sequence is a string of A, C, G, U\n    seq_tensor = sequence_to_tensor(sequence).to(device)  # shape: (seq_len,)\n    \n    # Generate five predictions for this sequence.\n    preds = predict_five_structures(model, seq_tensor, n_predictions=5)\n    seq_len = len(sequence)\n    \n    # For each residue in the sequence, form a submission row.\n    # Residue numbering: use 1-indexing.\n    for i in range(seq_len):\n        resid = i + 1\n        # Build the unique ID for the residue: e.g. \"R1107_1\", \"R1107_2\", etc.\n        row_id = f\"{target_id}_{resid}\"\n        resname = sequence[i]  # nucleotide letter at this position\n        \n        # Gather coordinates from all five predictions.\n        # Each prediction is a (seq_len, 3) array.\n        coords = []\n        for p in range(5):\n            x, y, z = preds[p][i]  # coordinates for residue i in prediction p\n            coords.extend([x, y, z])\n        \n        # Create a dictionary for this residue.\n        submission_row = {\n            \"ID\": row_id,\n            \"resname\": resname,\n            \"resid\": resid,\n            \"x_1\": coords[0],\n            \"y_1\": coords[1],\n            \"z_1\": coords[2],\n            \"x_2\": coords[3],\n            \"y_2\": coords[4],\n            \"z_2\": coords[5],\n            \"x_3\": coords[6],\n            \"y_3\": coords[7],\n            \"z_3\": coords[8],\n            \"x_4\": coords[9],\n            \"y_4\": coords[10],\n            \"z_4\": coords[11],\n            \"x_5\": coords[12],\n            \"y_5\": coords[13],\n            \"z_5\": coords[14],\n        }\n        submission_rows.append(submission_row)\n\n# Create a DataFrame for submission.\nsubmission_df = pd.DataFrame(submission_rows)\n\n# Save to CSV.\nsubmission_df.to_csv(\"submission.csv\", index=False)\nprint(\"Submission file saved as submission.csv\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"","metadata":{}}]}