{"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":"none","dataSources":[{"sourceType":"competition","sourceId":118765,"databundleVersionId":15231210,"isSourceIdPinned":false}],"dockerImageVersionId":31286,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Step 0: Imports","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom sklearn.manifold import MDS\nfrom tqdm import tqdm\nimport os","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-02-20T15:52:11.563763Z","iopub.execute_input":"2026-02-20T15:52:11.564229Z","iopub.status.idle":"2026-02-20T15:52:11.570964Z","shell.execute_reply.started":"2026-02-20T15:52:11.564194Z","shell.execute_reply":"2026-02-20T15:52:11.569436Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Step 1: Load Data","metadata":{}},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Update path based on your Kaggle mount\nbase_path = \"/kaggle/input/competitions/stanford-rna-3d-folding-2\"\n\ntrain_seq = pd.read_csv(f\"{base_path}/train_sequences.csv\", low_memory=False)\ntrain_labels = pd.read_csv(f\"{base_path}/train_labels.csv\", low_memory=False)\ntest_seq = pd.read_csv(f\"{base_path}/test_sequences.csv\", low_memory=False)\n\n# Ensure coordinates are numeric\ncoord_cols = [c for c in train_labels.columns if c.startswith((\"x_\", \"y_\", \"z_\"))]\ntrain_labels[coord_cols] = train_labels[coord_cols].apply(pd.to_numeric, errors=\"coerce\")\n\nprint(\"Train sequences:\", train_seq.shape)\nprint(\"Train labels:\", train_labels.shape)\nprint(\"Test sequences:\", test_seq.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:01:21.897974Z","iopub.execute_input":"2026-02-20T16:01:21.89829Z","iopub.status.idle":"2026-02-20T16:01:33.948019Z","shell.execute_reply.started":"2026-02-20T16:01:21.898265Z","shell.execute_reply":"2026-02-20T16:01:33.946947Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Step 2: Encode RNA Sequences","metadata":{}},{"cell_type":"code","source":"NUC_MAP = {'A':0, 'C':1, 'G':2, 'U':3}\n\ndef encode_sequence(seq):\n    \"\"\"\n    One-hot encode RNA sequence: A,C,G,U → 4-dim vector\n    \"\"\"\n    L = len(seq)\n    x = np.zeros((L, 4))\n    for i, s in enumerate(seq):\n        if s in NUC_MAP:\n            x[i, NUC_MAP[s]] = 1\n    return x\n\n# Test encoding\nprint(encode_sequence(\"ACGU\"))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:03:52.056535Z","iopub.execute_input":"2026-02-20T16:03:52.056937Z","iopub.status.idle":"2026-02-20T16:03:52.067097Z","shell.execute_reply.started":"2026-02-20T16:03:52.056907Z","shell.execute_reply":"2026-02-20T16:03:52.065037Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Step 3: Prepare Distance Matrices From Training Labels","metadata":{}},{"cell_type":"code","source":"def get_distance_matrix(df):\n    \"\"\"\n    Takes train_labels subset for a target and returns LxL distance matrix\n    \"\"\"\n    coords = df[['x_1','y_1','z_1']].values  # only first structure for now\n    L = coords.shape[0]\n    D = np.zeros((L,L))\n    for i in range(L):\n        for j in range(L):\n            D[i,j] = np.linalg.norm(coords[i] - coords[j])\n    return D\n\n# Example test\ntarget = train_seq.target_id.values[0]\nlabels_df = train_labels[train_labels.ID.str.startswith(target)]\nD_example = get_distance_matrix(labels_df)\nprint(D_example.shape)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:07:22.094938Z","iopub.execute_input":"2026-02-20T16:07:22.095278Z","iopub.status.idle":"2026-02-20T16:07:23.99183Z","shell.execute_reply.started":"2026-02-20T16:07:22.095249Z","shell.execute_reply":"2026-02-20T16:07:23.99094Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Step 4: Define the Transformer Model","metadata":{}},{"cell_type":"code","source":"class RNATransformer(nn.Module):\n    def __init__(self, d_model=128, nhead=4, num_layers=3):\n        super().__init__()\n        self.input_proj = nn.Linear(4, d_model)\n        encoder_layer = nn.TransformerEncoderLayer(\n            d_model=d_model,\n            nhead=nhead,\n            dim_feedforward=256,\n            batch_first=True\n        )\n        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)\n    \n    def forward(self, x):\n        x = self.input_proj(x)\n        x = self.encoder(x)\n        dist = torch.cdist(x, x, p=2)   # pairwise distances [batch, L, L]\n        return dist","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:22:18.772819Z","iopub.execute_input":"2026-02-20T16:22:18.77317Z","iopub.status.idle":"2026-02-20T16:22:18.780866Z","shell.execute_reply.started":"2026-02-20T16:22:18.77314Z","shell.execute_reply":"2026-02-20T16:22:18.779872Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Step 5: Device and Optimizer Setup","metadata":{}},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel = RNATransformer().to(device)\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:22:26.306773Z","iopub.execute_input":"2026-02-20T16:22:26.307165Z","iopub.status.idle":"2026-02-20T16:22:26.320916Z","shell.execute_reply.started":"2026-02-20T16:22:26.307132Z","shell.execute_reply":"2026-02-20T16:22:26.319898Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Step 6: Minimal Training Loop","metadata":{}},{"cell_type":"code","source":"model.train()\nfor epoch in range(2):\n    total_loss = 0\n    for target in tqdm(train_seq['target_id'].unique()[:5]):  # first 5 for testing\n        seq = train_seq[train_seq.target_id==target]['sequence'].values[0]\n        seq_enc = encode_sequence(seq)\n        x = torch.tensor(seq_enc, dtype=torch.float32).unsqueeze(0).to(device)\n\n        labels_df = train_labels[train_labels.ID.str.startswith(target)]\n        D_true = get_distance_matrix(labels_df)\n        D_true = torch.tensor(D_true, dtype=torch.float32).unsqueeze(0).to(device)\n\n        optimizer.zero_grad()\n        D_pred = model(x)\n        loss = F.mse_loss(D_pred, D_true)\n        loss.backward()\n        optimizer.step()\n\n        total_loss += loss.item()\n\n    print(f\"Epoch {epoch+1} total loss: {total_loss:.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:22:38.045248Z","iopub.execute_input":"2026-02-20T16:22:38.045678Z","iopub.status.idle":"2026-02-20T16:22:55.917766Z","shell.execute_reply.started":"2026-02-20T16:22:38.045641Z","shell.execute_reply":"2026-02-20T16:22:55.916958Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Step 7: Convert Distance Matrix → 3D Coordinates","metadata":{}},{"cell_type":"code","source":"def distance_to_coords(D):\n    \"\"\"\n    Convert distance matrix to 3D coordinates using classical MDS.\n    Ensure D is symmetric.\n    \"\"\"\n    # Force symmetry\n    D = (D + D.T) / 2.0\n\n    # Optional: clip negative distances to zero (numerical safety)\n    D = np.clip(D, 0, None)\n    \n    mds = MDS(n_components=3, dissimilarity='precomputed', random_state=42)\n    coords = mds.fit_transform(D)\n    return coords","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:31:22.603895Z","iopub.execute_input":"2026-02-20T16:31:22.604885Z","iopub.status.idle":"2026-02-20T16:31:22.610002Z","shell.execute_reply.started":"2026-02-20T16:31:22.604847Z","shell.execute_reply":"2026-02-20T16:31:22.608995Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Step 8: Generate 5 Predictions Per Test Sequence","metadata":{}},{"cell_type":"code","source":"# Step 8: Fast submission generation\nmodel.eval()\nsubmission_rows = []\n\nfor _, row in tqdm(test_seq.iterrows(), total=len(test_seq)):\n    target = row['target_id']\n    seq = row['sequence']\n    L = len(seq)\n    \n    # Encode sequence\n    seq_enc = encode_sequence(seq)\n    x = torch.tensor(seq_enc, dtype=torch.float32).unsqueeze(0).to(device)\n    \n    # Predict distance matrix\n    with torch.no_grad():\n        D_pred = model(x).squeeze(0).cpu().numpy()\n        # Force symmetry and clip negatives\n        D_pred = np.clip((D_pred + D_pred.T)/2, 0, None)\n    \n    # Convert distance matrix to 3D coordinates using fast classical MDS\n    def distance_to_coords(D):\n        n = D.shape[0]\n        J = np.eye(n) - np.ones((n,n))/n\n        B = -0.5 * J @ (D**2) @ J\n        eigvals, eigvecs = np.linalg.eigh(B)\n        idx = np.argsort(eigvals)[::-1][:3]\n        eigvals = np.maximum(eigvals[idx], 0)\n        eigvecs = eigvecs[:, idx]\n        coords = eigvecs * np.sqrt(eigvals)\n        return coords\n    \n    coords = distance_to_coords(D_pred)\n    \n    # Repeat coords 5 times (classical MDS is deterministic)\n    coords_list = [coords for _ in range(5)]\n    \n    # Build submission rows\n    for i in range(L):\n        entry = {\n            \"ID\": f\"{target}_{i+1}\",\n            \"resname\": seq[i],\n            \"resid\": i+1\n        }\n        for k in range(5):\n            entry[f\"x_{k+1}\"] = coords_list[k][i,0]\n            entry[f\"y_{k+1}\"] = coords_list[k][i,1]\n            entry[f\"z_{k+1}\"] = coords_list[k][i,2]\n        submission_rows.append(entry)\n\n# Create DataFrame\nsubmission = pd.DataFrame(submission_rows)\n\n# Clip coordinates to valid PDB range\nsubmission.iloc[:,3:] = submission.iloc[:,3:].clip(-999.999, 9999.999)\n\n# Save submission\nsubmission.to_csv(\"submission.csv\", index=False)\nprint(\"submission.csv saved! ✅\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T16:36:40.385525Z","iopub.execute_input":"2026-02-20T16:36:40.386227Z","iopub.status.idle":"2026-02-20T16:37:04.757641Z","shell.execute_reply.started":"2026-02-20T16:36:40.386198Z","shell.execute_reply":"2026-02-20T16:37:04.756936Z"}},"outputs":[],"execution_count":null}]}