{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","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":11553390,"sourceType":"competition"},{"sourceId":11841749,"sourceType":"datasetVersion","datasetId":7440088},{"sourceId":11841762,"sourceType":"datasetVersion","datasetId":7440093}],"dockerImageVersionId":31040,"isInternetEnabled":true,"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        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-05-17T01:10:00.892079Z","iopub.execute_input":"2025-05-17T01:10:00.892305Z","iopub.status.idle":"2025-05-17T01:10:17.191504Z","shell.execute_reply.started":"2025-05-17T01:10:00.892282Z","shell.execute_reply":"2025-05-17T01:10:17.190633Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport pickle\n\ndf = pd.read_csv(\"/kaggle/input/augdata-train/augmented_train_features.csv\")\nfeature_cols = [col for col in df.columns if col not in ['target_id', 'resid', 'resname', 'Unnamed: 0']]\n\nfeatures_dict_train = {\n    tid: group.sort_values(\"resid\")[feature_cols].to_numpy()\n    for tid, group in df.groupby(\"target_id\")\n}\n\nwith open(\"train_features_dict.pkl\", \"wb\") as f:\n    pickle.dump(features_dict_train, f)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T01:48:07.952142Z","iopub.execute_input":"2025-05-17T01:48:07.952438Z","iopub.status.idle":"2025-05-17T01:48:12.351927Z","shell.execute_reply.started":"2025-05-17T01:48:07.952416Z","shell.execute_reply":"2025-05-17T01:48:12.350525Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_val = pd.read_csv(\"/kaggle/input/augdata-val/augmented_validation_features.csv\")\nval_features_dict = {\n    tid: group.sort_values(\"resid\")[feature_cols].to_numpy()\n    for tid, group in df_val.groupby(\"target_id\")\n}\n\nwith open(\"val_features_dict.pkl\", \"wb\") as f:\n    pickle.dump(val_features_dict, f)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T01:48:39.565161Z","iopub.execute_input":"2025-05-17T01:48:39.565861Z","iopub.status.idle":"2025-05-17T01:48:39.667547Z","shell.execute_reply.started":"2025-05-17T01:48:39.565831Z","shell.execute_reply":"2025-05-17T01:48:39.666708Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_test = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/test_sequences.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T01:49:31.193847Z","iopub.execute_input":"2025-05-17T01:49:31.194151Z","iopub.status.idle":"2025-05-17T01:49:31.203302Z","shell.execute_reply.started":"2025-05-17T01:49:31.194127Z","shell.execute_reply":"2025-05-17T01:49:31.202473Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Explode each sequence into rows per residue\nrecords = []\nfor _, row in df_test.iterrows():\n    seq = row[\"sequence\"]\n    for i, res in enumerate(seq):\n        records.append({\n            \"target_id\": row[\"target_id\"],\n            \"resid\": i + 1,           # Residue index (1-based)\n            \"resname\": res,           # A, C, G, or U\n        })\n\ntest_residue_df = pd.DataFrame(records)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T01:51:41.945234Z","iopub.execute_input":"2025-05-17T01:51:41.945544Z","iopub.status.idle":"2025-05-17T01:51:41.962651Z","shell.execute_reply.started":"2025-05-17T01:51:41.945521Z","shell.execute_reply":"2025-05-17T01:51:41.961623Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load original train features\ntrain_df = pd.read_csv(\"/kaggle/input/augdata-train/augmented_train_features.csv\")\nfeature_cols = [col for col in train_df.columns if col not in [\"Unnamed: 0\", \"target_id\", \"resid\", \"resname\"]]\n\n# Compute column-wise means\nfeature_means = train_df[feature_cols].mean()\n\n# Broadcast to each row of the test data\nfor col in feature_cols:\n    test_residue_df[col] = feature_means[col]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T01:51:45.230552Z","iopub.execute_input":"2025-05-17T01:51:45.230846Z","iopub.status.idle":"2025-05-17T01:51:47.386365Z","shell.execute_reply.started":"2025-05-17T01:51:45.230826Z","shell.execute_reply":"2025-05-17T01:51:47.385498Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_residue_df.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T01:52:04.557833Z","iopub.execute_input":"2025-05-17T01:52:04.558160Z","iopub.status.idle":"2025-05-17T01:52:04.595747Z","shell.execute_reply.started":"2025-05-17T01:52:04.558134Z","shell.execute_reply":"2025-05-17T01:52:04.594857Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\n\nlabels_df = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/train_labels.csv\")\n\ndef extract_target_id(full_id):\n    return \"_\".join(full_id.split(\"_\")[:2])\n\nlabels_df[\"target_id\"] = labels_df[\"ID\"].apply(extract_target_id)\n\ndef extract_coords(df):\n    coords = []\n    i = 1\n    while f\"x_{i}\" in df.columns:\n        coords.append(df[[f\"x_{i}\", f\"y_{i}\", f\"z_{i}\"]].to_numpy())\n        i += 1\n    return np.stack(coords, axis=1)  # shape: [L, 5, 3] if 5 conformations\n\nlabels_dict = {\n    tid: extract_coords(group.sort_values(\"resid\"))\n    for tid, group in labels_df.groupby(\"target_id\")\n}\n\nwith open(\"train_labels_dict.pkl\", \"wb\") as f:\n    pickle.dump(labels_dict, f)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T01:52:28.714275Z","iopub.execute_input":"2025-05-17T01:52:28.714618Z","iopub.status.idle":"2025-05-17T01:52:29.703162Z","shell.execute_reply.started":"2025-05-17T01:52:28.714595Z","shell.execute_reply":"2025-05-17T01:52:29.702332Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%pip install pytorch_lightning","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T01:52:44.014584Z","iopub.execute_input":"2025-05-17T01:52:44.014925Z","iopub.status.idle":"2025-05-17T01:54:33.794051Z","shell.execute_reply.started":"2025-05-17T01:52:44.014900Z","shell.execute_reply":"2025-05-17T01:54:33.792229Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport pytorch_lightning as pl\nfrom torchmetrics import MeanMetric\nfrom torch.cuda.amp import autocast\n\nclass CleanRNADataset(torch.utils.data.Dataset):\n    def __init__(self, features_dict, labels_dict=None):\n        self.ids = list(features_dict.keys())\n        self.features = features_dict\n        self.labels = labels_dict\n\n    def __len__(self):\n        return len(self.ids)\n\n    def __getitem__(self, idx):\n        tid = self.ids[idx]\n        x = torch.tensor(self.features[tid], dtype=torch.float32)\n        sample = {\"features\": x, \"target_id\": tid}\n        if self.labels:\n            y = torch.tensor(self.labels[tid], dtype=torch.float32)\n            sample[\"labels\"] = y\n        return sample\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T01:54:52.736916Z","iopub.execute_input":"2025-05-17T01:54:52.737460Z","iopub.status.idle":"2025-05-17T01:54:52.745147Z","shell.execute_reply.started":"2025-05-17T01:54:52.737434Z","shell.execute_reply":"2025-05-17T01:54:52.744066Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torch.utils.data import Dataset\n\nclass RNATestDataset(Dataset):\n    def __init__(self, df, feature_cols):\n        self.groups = list(df.groupby(\"target_id\"))\n        self.feature_cols = feature_cols\n\n    def __len__(self):\n        return len(self.groups)\n\n    def __getitem__(self, idx):\n        tid, group = self.groups[idx]\n        group = group.sort_values(\"resid\")\n        features = torch.tensor(group[self.feature_cols].to_numpy(), dtype=torch.float32)\n        return {\n            \"features\": features,  # [L, D]\n            \"target_id\": tid,\n            \"resname\": list(group[\"resname\"]),\n            \"resid\": list(group[\"resid\"])\n        }\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T01:57:10.353995Z","iopub.execute_input":"2025-05-17T01:57:10.354328Z","iopub.status.idle":"2025-05-17T01:57:10.362419Z","shell.execute_reply.started":"2025-05-17T01:57:10.354306Z","shell.execute_reply":"2025-05-17T01:57:10.361160Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torch.utils.data import DataLoader\n\ntrain_ds = CleanRNADataset(features_dict=features_dict_train, labels_dict=labels_dict)\nval_ds = CleanRNADataset(features_dict=val_features_dict, labels_dict=None)\n\ntrain_loader = DataLoader(train_ds, batch_size=1, shuffle=True)\nval_loader = DataLoader(val_ds, batch_size=1, shuffle=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T01:57:23.771654Z","iopub.execute_input":"2025-05-17T01:57:23.771978Z","iopub.status.idle":"2025-05-17T01:57:23.778165Z","shell.execute_reply.started":"2025-05-17T01:57:23.771957Z","shell.execute_reply":"2025-05-17T01:57:23.776947Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"feature_cols = [col for col in test_residue_df.columns if col not in [\"target_id\", \"resid\", \"resname\"]]\nfrom torch.utils.data import DataLoader\n\n\ntest_ds = RNATestDataset(test_residue_df, feature_cols)\ntest_loader = DataLoader(test_ds, batch_size=1, shuffle=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T01:57:45.291851Z","iopub.execute_input":"2025-05-17T01:57:45.292161Z","iopub.status.idle":"2025-05-17T01:57:45.308921Z","shell.execute_reply.started":"2025-05-17T01:57:45.292141Z","shell.execute_reply":"2025-05-17T01:57:45.307719Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pytorch_lightning as pl\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torchmetrics import MeanMetric\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\n\nclass RNASeqRegressor(pl.LightningModule):\n    def __init__(self, input_dim, hidden_dim=128, num_structures=5, lr=1e-3):\n        super().__init__()\n        self.save_hyperparameters()\n        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers=2, batch_first=True, bidirectional=True)\n        self.proj = nn.Linear(2 * hidden_dim, hidden_dim)\n        self.heads = nn.ModuleList([\n            nn.Sequential(\n                nn.Linear(hidden_dim, hidden_dim),\n                nn.ReLU(),\n                nn.Linear(hidden_dim, 3)\n            ) for _ in range(num_structures)\n        ])\n        self.loss_tracker = MeanMetric()\n\n    def forward(self, x):\n        x, _ = self.lstm(x)              # [B, L, 2H]\n        x = self.proj(x)                 # [B, L, H]\n        return [head(x) for head in self.heads]  # List of [B, L, 3]\n\n    def calculate_rmsd(self, pred, target):\n      return torch.sqrt(torch.mean((pred.float() - target.float()) ** 2))\n\n    def _step(self, batch, mode):\n        x = batch['features']                   # [B, L, D]\n        y = batch.get('labels')                 # [B, L, S, 3]\n        if y is None:\n            return torch.tensor(0.0, device=self.device, requires_grad=True)\n\n        preds = self(x)                         # list of [B, L, 3]\n        total_loss = 0.0\n        valid = 0\n\n        for i, pred in enumerate(preds):\n          if i < y.shape[2]:  # y: [B, L_target, num_structures, 3]\n            target = y[:, :, i, :]  # [B, L_target, 3]\n\n            # ✨ Truncate to min seq length\n            min_len = min(pred.shape[1], target.shape[1])\n            pred = pred[:, :min_len, :]       # [B, min_len, 3]\n            target = target[:, :min_len, :]   # [B, min_len, 3]\n\n            loss = self.calculate_rmsd(pred, target)\n            if torch.isfinite(loss):\n                    total_loss += loss\n                    valid += 1\n\n        if valid > 0:\n            avg_loss = total_loss / valid\n        else:\n            # Return a dummy differentiable scalar to avoid error\n            avg_loss = torch.zeros(1, device=self.device, requires_grad=True).sum()\n\n            # Always log the RMSD as a separate metric\n        self.log(f\"{mode}_rmsd\", avg_loss, prog_bar=False, on_epoch=True, on_step=False, batch_size=x.size(0))\n        self.log(\"val_rmsd\", avg_loss, prog_bar=False, on_epoch=True, on_step=False, batch_size=x.size(0))\n        # Log loss under Lightning's expected name\n        if mode == \"train\":\n            self.log(\"train_loss\", avg_loss, prog_bar=True, on_epoch=True, on_step=False, batch_size=x.size(0))\n        elif mode == \"val\":\n            self.log(\"val_loss\", avg_loss, prog_bar=True, on_epoch=True, on_step=False, batch_size=x.size(0))\n\n        return avg_loss\n\n    #print(\"Loss requires grad?\", avg_loss.requires_grad)\n\n    def training_step(self, batch, batch_idx):\n        return self._step(batch, \"train\")\n\n    def validation_step(self, batch, batch_idx):\n        loss = self._step(batch, \"val\")\n        self.log(\"val_loss\", loss, prog_bar=True, on_epoch=True, on_step=False, batch_size=batch['features'].size(0))\n\n    def on_train_epoch_end(self):\n        print(\"✅ Epoch complete: train_rmsd =\", self.trainer.callback_metrics.get(\"train_rmsd\"))\n\n    def on_validation_epoch_end(self):\n        print(\"✅ Epoch complete: val_rmsd =\", self.trainer.callback_metrics.get(\"val_rmsd\"))\n\n\n\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr)\n        scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=3)\n        return {\"optimizer\": optimizer, \"lr_scheduler\": scheduler, \"monitor\": \"val_rmsd\"}\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T01:57:57.711280Z","iopub.execute_input":"2025-05-17T01:57:57.711628Z","iopub.status.idle":"2025-05-17T01:57:57.730517Z","shell.execute_reply.started":"2025-05-17T01:57:57.711602Z","shell.execute_reply":"2025-05-17T01:57:57.729370Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from pytorch_lightning import Trainer\nfrom pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint\nfrom pytorch_lightning.loggers import CSVLogger\n\ninput_dim = train_ds[0]['features'].shape[1]\nmodel = RNASeqRegressor(input_dim=input_dim)\n\ncallbacks = [\n    EarlyStopping(monitor='val_loss', patience=5, strict=False),\n    ModelCheckpoint(monitor='val_loss', save_top_k=1, mode='min')\n]\n\ntrainer = pl.Trainer(\n    max_epochs=25,\n    precision=32,  # ✅ Force full precision\n    accelerator='auto',\n    callbacks=callbacks,\n    logger=CSVLogger(\"logs\")\n)\n\ntrainer.fit(model, train_dataloaders=train_loader, val_dataloaders=val_loader)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T01:58:27.540364Z","iopub.execute_input":"2025-05-17T01:58:27.540701Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\nsample_path = \"/kaggle/input/stanford-rna-3d-folding/sample_submission.csv\"\nsample_df = pd.read_csv(sample_path)\nprint(sample_df.head())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T01:58:24.245017Z","iopub.execute_input":"2025-05-17T01:58:24.245440Z","iopub.status.idle":"2025-05-17T01:58:24.274462Z","shell.execute_reply.started":"2025-05-17T01:58:24.245410Z","shell.execute_reply":"2025-05-17T01:58:24.273214Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\n# Step 1: Run model predictions and collect coords from all heads\nmodel.eval()\nmodel.to(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nall_preds = []\n\nfor batch in test_loader:\n    features = batch[\"features\"].to(model.device)  # [B, L, D]\n    ids = batch[\"target_id\"]  # list of target_ids\n\n    with torch.no_grad():\n        outputs = model(features)  # list of 5 tensors: [B, L, 3]\n\n    for i, tid in enumerate(ids):  # loop over batch items\n        seq_len = features.shape[1]\n        entry = {\n            \"target_id\": tid,\n            \"resid\": list(range(1, seq_len + 1)),  # 1-based indexing\n            \"resname\": [\"G\"] * seq_len  # placeholder — use real nucleotide if available\n        }\n\n        # Gather coords from each head\n        for head_idx, head_output in enumerate(outputs):  # loop over 5 heads\n            coords = head_output[i].cpu().numpy()  # [L, 3]\n            entry.update({\n                f\"x_{head_idx+1}\": coords[:, 0],\n                f\"y_{head_idx+1}\": coords[:, 1],\n                f\"z_{head_idx+1}\": coords[:, 2],\n            })\n\n        df = pd.DataFrame(entry)\n        df[\"ID\"] = [f\"{tid}_{r}\" for r in df[\"resid\"]]\n        all_preds.append(df)\n\n# Step 2: Combine all sequences into one DataFrame\nsubmission_df = pd.concat(all_preds, ignore_index=True)\n\n# Step 3: Reorder columns to match Kaggle format\ncoord_cols = []\nfor i in range(1, 6):\n    coord_cols.extend([f\"x_{i}\", f\"y_{i}\", f\"z_{i}\"])\n\nsubmission_df = submission_df[[\"ID\", \"resname\", \"resid\"] + coord_cols]\n\n# Step 4: Save final CSV\nsubmission_df.to_csv(\"/kaggle/working/submission.csv\", index=False)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}