{"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":11228175,"sourceType":"competition"}],"dockerImageVersionId":30918,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### This notebook is rewritten for PyTorch with DeepSeek","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T15:18:21.426705Z","iopub.execute_input":"2025-03-06T15:18:21.426952Z","iopub.status.idle":"2025-03-06T15:18:22.540250Z","shell.execute_reply.started":"2025-03-06T15:18:21.426932Z","shell.execute_reply":"2025-03-06T15:18:22.539394Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Dataset Overview","metadata":{}},{"cell_type":"code","source":"train_labels = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/train_labels.csv\")\ntrain_sequence = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/train_sequences.csv\")\nval_labels = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/validation_labels.csv\")\nval_sequence = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/validation_sequences.csv\")\ntest_sequence = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/test_sequences.csv\")\n\nprint(\"Train Seq: \" + str(train_sequence.shape))\nprint(\"Train Label: \" + str(train_labels.shape))\nprint(\"Validation Seq: \" + str(val_sequence.shape))\nprint(\"Validation Label: \" + str(val_labels.shape))\nprint(\"Test: \"+str(test_sequence.shape))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T15:18:22.541009Z","iopub.execute_input":"2025-03-06T15:18:22.541356Z","iopub.status.idle":"2025-03-06T15:18:23.071298Z","shell.execute_reply.started":"2025-03-06T15:18:22.541335Z","shell.execute_reply":"2025-03-06T15:18:23.070410Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_labels.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T15:18:23.072171Z","iopub.execute_input":"2025-03-06T15:18:23.072443Z","iopub.status.idle":"2025-03-06T15:18:23.097419Z","shell.execute_reply.started":"2025-03-06T15:18:23.072418Z","shell.execute_reply":"2025-03-06T15:18:23.096367Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Visualization\n- use plotly for interactive plots\n- target is the first label in labels","metadata":{}},{"cell_type":"code","source":"import plotly.express as px\nimport plotly.io as pio\npio.renderers.default = 'iframe'\n\nplot_labels = train_labels[['ID', 'x_1', 'y_1', 'z_1']].copy()\nplot_labels['label'] = plot_labels.ID.str.rsplit('_', n=1, expand=True).iloc[:,0]\n\ntarget = '1SCL_A'\n\nfig = px.line_3d(\n    plot_labels.loc[plot_labels.label==target],\n    x='x_1', y='y_1', z='z_1',\n    title=f'{target}'\n)\n\nfig.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T15:18:23.099020Z","iopub.execute_input":"2025-03-06T15:18:23.099283Z","iopub.status.idle":"2025-03-06T15:18:26.177990Z","shell.execute_reply.started":"2025-03-06T15:18:23.099260Z","shell.execute_reply":"2025-03-06T15:18:26.176992Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_sequence.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T15:18:26.178682Z","iopub.execute_input":"2025-03-06T15:18:26.179018Z","iopub.status.idle":"2025-03-06T15:18:26.190592Z","shell.execute_reply.started":"2025-03-06T15:18:26.178989Z","shell.execute_reply":"2025-03-06T15:18:26.188916Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"val_labels.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T15:18:26.193682Z","iopub.execute_input":"2025-03-06T15:18:26.194052Z","iopub.status.idle":"2025-03-06T15:18:26.241351Z","shell.execute_reply.started":"2025-03-06T15:18:26.194022Z","shell.execute_reply":"2025-03-06T15:18:26.240379Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Preprocessing","metadata":{}},{"cell_type":"code","source":"seq_dict = {'A': 1, 'C': 2, 'G': 3, 'U': 4}\ndef seq_map(seq):\n    return [seq_dict.get(char, 0) for char in seq]\n\ntrain_sequence['encoded_seq'] = train_sequence['sequence'].apply(seq_map)\ntest_sequence['encoded_seq'] = test_sequence['sequence'].apply(seq_map)\nval_sequence['encoded_seq'] = val_sequence['sequence'].apply(seq_map)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T15:18:26.242745Z","iopub.execute_input":"2025-03-06T15:18:26.243129Z","iopub.status.idle":"2025-03-06T15:18:26.260022Z","shell.execute_reply.started":"2025-03-06T15:18:26.243096Z","shell.execute_reply":"2025-03-06T15:18:26.258928Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from collections import defaultdict\n\ndef generate_label_coord(frame):\n    output = defaultdict(list)\n\n    frame['label'] = frame.ID.str.rsplit('_', n=1, expand=True).iloc[:, 0]\n\n    for _, row in frame.iterrows():\n        label = row['label']\n        resid = row['resid']\n\n        coord = np.array((row['x_1'], row['y_1'], row['z_1']), dtype=np.float32)\n        output[label].append((resid, coord))\n\n    for key, value in output.items():\n        coords = np.stack([c for r, c in value])\n        masks = np.isnan(coords) | np.isclose(coords, -1.0000e+18)\n\n        coords[masks] = 0.\n\n        output[key] = {\n            'coords': coords,\n            'masks': ~masks,\n        }\n\n    return output\n\ntrain_stacked_coords = generate_label_coord(train_labels)\nval_stacked_coords = generate_label_coord(val_labels)\n\ntrain_stacked_coords[list(train_stacked_coords.keys())[0]]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T15:18:26.260791Z","iopub.execute_input":"2025-03-06T15:18:26.261062Z","iopub.status.idle":"2025-03-06T15:18:32.643203Z","shell.execute_reply.started":"2025-03-06T15:18:26.261043Z","shell.execute_reply":"2025-03-06T15:18:32.642288Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Create Dataset for CNN","metadata":{}},{"cell_type":"code","source":"def generate_dataset(seq, stacked_coords):\n    X, y, my, tids = [], [], [], []\n\n    for idx, row in seq.iterrows():\n        tid = row['target_id']\n        if tid in stacked_coords:\n            X.append(row['encoded_seq'])\n            y.append(stacked_coords[tid]['coords'])\n            my.append(stacked_coords[tid]['masks'])\n            tids.append(tid)\n\n    return X, y, my, tids\n\ntrain_X, train_y, train_my, train_tids = generate_dataset(train_sequence, train_stacked_coords)\nval_X, val_y, val_my, val_tids = generate_dataset(val_sequence, val_stacked_coords)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T15:18:32.644024Z","iopub.execute_input":"2025-03-06T15:18:32.644247Z","iopub.status.idle":"2025-03-06T15:18:32.685520Z","shell.execute_reply.started":"2025-03-06T15:18:32.644228Z","shell.execute_reply":"2025-03-06T15:18:32.684269Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"- coordinates and seqences are padded to avoid different lengths, all with constant 0","metadata":{}},{"cell_type":"code","source":"import torch\nfrom torch.nn.utils.rnn import pad_sequence\n\ndef pad_sequences_torch(sequences, max_len, padding='post', value=0):\n    padded = pad_sequence([\n        torch.cat([seq, torch.full((max_len - len(seq),), value, dtype=seq.dtype)])\n        if len(seq) < max_len else seq[:max_len]\n        for seq in sequences\n    ], batch_first=True, padding_value=value)\n    return padded\n\nmax_len = max(len(seq) for seq in train_X)\n\ntrain_X_pad = pad_sequences_torch([torch.tensor(x) for x in train_X], max_len)\nval_X_pad = pad_sequences_torch([torch.tensor(x) for x in val_X], max_len)\ntest_X = test_sequence['encoded_seq'].tolist()\ntest_X_pad = pad_sequences_torch([torch.tensor(x) for x in test_sequence['encoded_seq'].tolist()], max_len)\n\ntrain_X_pad.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T15:18:32.686461Z","iopub.execute_input":"2025-03-06T15:18:32.686725Z","iopub.status.idle":"2025-03-06T15:18:36.836060Z","shell.execute_reply.started":"2025-03-06T15:18:32.686703Z","shell.execute_reply":"2025-03-06T15:18:36.835201Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch.nn.functional as F\n\ndef pad_coords_torch(coords, max_len):\n    L = coords.size(0)\n    if L < max_len:\n        pad = (0, 0, 0, max_len - L)  # (left, right, top, bottom)\n        return F.pad(coords, pad, \"constant\", 0)\n    else:\n        return coords[:max_len]\n\ntrain_y_pad = torch.stack([pad_coords_torch(torch.tensor(y), max_len) for y in train_y])\nval_y_pad = torch.stack([pad_coords_torch(torch.tensor(y), max_len) for y in val_y])\n\ntrain_my_pad = torch.stack([pad_coords_torch(torch.tensor(my), max_len) for my in train_my])\nval_my_pad = torch.stack([pad_coords_torch(torch.tensor(my), max_len) for my in val_my])\n\ntrain_y_pad.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T15:18:36.836765Z","iopub.execute_input":"2025-03-06T15:18:36.837278Z","iopub.status.idle":"2025-03-06T15:18:36.942790Z","shell.execute_reply.started":"2025-03-06T15:18:36.837251Z","shell.execute_reply":"2025-03-06T15:18:36.941770Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## CNN model\n- first conv and second conv layers are seperated for clarity\n- main hyperparams are defined in self.config\n- a model must first be built then trained","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader, TensorDataset\nimport pytorch_lightning as pl\nfrom pytorch_lightning.callbacks import EarlyStopping\n\nclass CNN(pl.LightningModule):\n    def __init__(self, config):\n        super().__init__()\n        self.config = config\n        \n        # Embedding layer\n        self.embedding = nn.Embedding(\n            num_embeddings=config['seq_mapping_size'],\n            embedding_dim=config['embedding_dim'],\n            padding_idx=0\n        )\n        \n        # First Conv Block\n        self.conv1 = nn.Conv1d(\n            in_channels=config['embedding_dim'],\n            out_channels=config['num_filters'],\n            kernel_size=config['kernel_size'],\n            padding='same'\n        )\n        self.bn1 = nn.BatchNorm1d(config['num_filters'])\n        self.drop1 = nn.Dropout(config['drop_rate'])\n        \n        # Second Conv Block\n        self.conv2 = nn.Conv1d(\n            in_channels=config['num_filters'],\n            out_channels=config['num_filters'],\n            kernel_size=config['kernel_size'],\n            padding='same'\n        )\n        self.bn2 = nn.BatchNorm1d(config['num_filters'])\n        self.drop2 = nn.Dropout(config['drop_rate'])\n        \n        # Final Prediction Layer\n        self.final_conv = nn.Conv1d(\n            in_channels=config['num_filters'],\n            out_channels=3,\n            kernel_size=1,\n            padding='same'\n        )\n        \n        self.loss_fn = nn.MSELoss(reduction='none')\n\n    def forward(self, x):\n        # Embedding\n        x = self.embedding(x)  # (batch, seq_len, embedding_dim)\n        x = x.permute(0, 2, 1)  # (batch, embedding_dim, seq_len)\n        \n        # First Conv Block\n        x = F.relu(self.conv1(x))\n        x = self.bn1(x)\n        x = self.drop1(x)\n        \n        # Second Conv Block\n        x = F.relu(self.conv2(x))\n        x = self.bn2(x)\n        x = self.drop2(x)\n        \n        # Final Prediction\n        x = self.final_conv(x)\n        x = x.permute(0, 2, 1)  # (batch, seq_len, 3)\n        return x\n    \n    def training_step(self, batch, batch_idx):\n        x, y, my = batch\n        y_hat = self(x)\n        loss = self.loss_fn(y_hat, y)\n        loss = (my * loss).mean()\n\n        \n        self.log('train_loss', loss, prog_bar=True)\n        return loss\n    \n    def validation_step(self, batch, batch_idx):\n        x, y, my = batch\n        y_hat = self(x)\n        loss = self.loss_fn(y_hat, y)\n        loss = (my * loss).mean()\n\n        self.log('val_loss', loss, prog_bar=True)\n        return loss\n    \n    def configure_optimizers(self):\n        return torch.optim.Adam(self.parameters())\n\n# Конфигурация модели\nconfig = {\n    'seq_mapping_size': max(seq_dict.values()) + 1,\n    'embedding_dim': 16,\n    'num_filters': 64,\n    'kernel_size': 3,\n    'drop_rate': 0.2,\n    'train_epochs': 50,\n    'batch_size': 16\n}\n\n# Подготовка данных\ntrain_dataset = TensorDataset(train_X_pad, train_y_pad, train_my_pad.long())\nval_dataset = TensorDataset(val_X_pad, val_y_pad, val_my_pad.long())\n\ntrain_loader = DataLoader(train_dataset, batch_size=config['batch_size'], shuffle=True)\nval_loader = DataLoader(val_dataset, batch_size=config['batch_size'])\n\n# Инициализация и обучение модели\nmodel = CNN(config)\nearly_stop = EarlyStopping(monitor='val_loss', patience=3, mode='min')\n\ntrainer = pl.Trainer(\n    max_epochs=config['train_epochs'],\n    callbacks=[early_stop],\n    accelerator='auto',\n    enable_progress_bar=True,\n    log_every_n_steps=1,\n)\ntrainer.fit(model, train_loader, val_loader)\n\n# Предсказание на тестовых данных\ntest_tensor = torch.LongTensor(test_X_pad)\nmodel.eval()\nwith torch.no_grad():\n    pred = model(test_tensor).cpu().numpy()\n\nprint(pred.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T15:19:17.322494Z","iopub.execute_input":"2025-03-06T15:19:17.322776Z","iopub.status.idle":"2025-03-06T15:20:40.826073Z","shell.execute_reply.started":"2025-03-06T15:19:17.322754Z","shell.execute_reply":"2025-03-06T15:20:40.825147Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Submission\n- the target 3 dimensional coordinates are duplicated 5 times","metadata":{}},{"cell_type":"code","source":"rows = []\n\nfor idx, row in test_sequence.iterrows():\n    target_id = row['target_id']\n    coords = pred[idx]\n    seq_length = len(row['encoded_seq'])\n    coords = coords[:seq_length, :]\n\n    for i in range(seq_length):\n        x, y, z = coords[i, :]\n        rows.append(\n            {\n                'ID': f\"{target_id}_{i+1}\",\n                'resname': row['sequence'][i],\n                'resid': i+1,\n                 **{f\"x_{j+1}\": x for j in range(5)},\n                 **{f\"y_{j+1}\": y for j in range(5)},\n                 **{f\"z_{j+1}\": z for j in range(5)}\n            }\n        )\n\nsubmission = pd.DataFrame(rows)\nsubmission.to_csv(\"submission.csv\", index=False)\n\nsubmission.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T15:20:51.637317Z","iopub.execute_input":"2025-03-06T15:20:51.637588Z","iopub.status.idle":"2025-03-06T15:20:51.711641Z","shell.execute_reply.started":"2025-03-06T15:20:51.637567Z","shell.execute_reply":"2025-03-06T15:20:51.710760Z"}},"outputs":[],"execution_count":null}]}