{"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},{"sourceType":"datasetVersion","sourceId":14624318,"datasetId":9341374,"databundleVersionId":15461984}],"dockerImageVersionId":31286,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-03-13T17:14:52.921082Z","iopub.execute_input":"2026-03-13T17:14:52.921431Z","iopub.status.idle":"2026-03-13T17:15:02.204478Z","shell.execute_reply.started":"2026-03-13T17:14:52.921403Z","shell.execute_reply":"2026-03-13T17:15:02.203379Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"\nStanford RNA 3D Folding Part 2 - Full Solution\n使用已挂载的 RhoFold + MSA 文件进行预测\n\"\"\"\n\nimport numpy as np\nimport pandas as pd\nimport sys, os, torch\n\n# ==================================================================\n# 1. 路径配置（已根据你的实际文件结构修正）\n# ==================================================================\nCOMP_DIR    = \"/kaggle/input/competitions/stanford-rna-3d-folding-2\"\nRHOFOLD_DIR = \"/kaggle/input/datasets/profmansoor/stanford-rna-3d-folding-2/RhoFold-main\"\nMSA_DIR     = os.path.join(COMP_DIR, \"MSA\")\nCKPT_PATH   = os.path.join(RHOFOLD_DIR, \"pretrained/model.pt\")\n\nsys.path.insert(0, RHOFOLD_DIR)\n\n# ==================================================================\n# 2. 读取数据\n# ==================================================================\ntest_df  = pd.read_csv(os.path.join(COMP_DIR, \"test_sequences.csv\"))\nsample   = pd.read_csv(os.path.join(COMP_DIR, \"sample_submission.csv\"))\ntrain_df = pd.read_csv(os.path.join(COMP_DIR, \"train_sequences.csv\"))\n\nprint(f\"测试序列数 : {len(test_df)}\")\nprint(f\"训练序列数 : {len(train_df)}\")\nprint(f\"MSA 文件数 : {len(os.listdir(MSA_DIR))}\")\nprint(f\"列名       : {test_df.columns.tolist()}\")\nprint(test_df.head())\n\n# ==================================================================\n# 3. A-form 兜底函数（RhoFold 失败时使用）\n# ==================================================================\ndef aform_fallback(n: int, n_preds: int = 5) -> list:\n    \"\"\"生成 A-form 螺旋坐标作为兜底预测，返回 list of (n,3) arrays\"\"\"\n    angs = np.radians(np.arange(n) * 32.7)\n    base = np.stack([\n        9.0 * np.cos(angs),\n        9.0 * np.sin(angs),\n        np.arange(n) * 2.81\n    ], axis=1).astype(np.float32)\n    result = []\n    for k in range(n_preds):\n        rng = np.random.default_rng(k * 42)\n        noise = rng.normal(0, 0.4 * (k + 0.5), base.shape).astype(np.float32)\n        result.append(base + noise)\n    return result\n\n# ==================================================================\n# 4. 加载 RhoFold 模型\n# ==================================================================\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(f\"\\n使用设备: {device}\")\n\nUSE_RHOFOLD = False  # 先设为 False，确认权重存在后改为 True\n\nif os.path.exists(CKPT_PATH):\n    print(f\"找到权重文件: {CKPT_PATH}\")\n    try:\n        from rhofold.rhofold import RhoFold\n        from rhofold.config import rhofold_config\n        from rhofold.utils.alphabet import get_features\n\n        model = RhoFold(rhofold_config)\n        ckpt  = torch.load(CKPT_PATH, map_location=device)\n        model.load_state_dict(ckpt[\"model\"] if \"model\" in ckpt else ckpt)\n        model = model.to(device).eval()\n        USE_RHOFOLD = True\n        print(\"RhoFold 模型加载成功！\")\n    except Exception as e:\n        print(f\"模型加载失败，将使用 A-form 兜底: {e}\")\nelse:\n    print(f\"未找到权重文件: {CKPT_PATH}\")\n    print(\"pretrained 目录内容:\", os.listdir(os.path.join(RHOFOLD_DIR, \"pretrained\")))\n    print(\"将使用 A-form 螺旋作为兜底预测。\")\n\n# ==================================================================\n# 5. 批量预测\n# ==================================================================\nN_PREDS = 5\nrows    = []\n\nfor i, entry in test_df.iterrows():\n    seq_id = str(entry.get(\"target_id\", entry.iloc[0]))\n    seq    = str(entry.get(\"sequence\",  entry.iloc[1])).upper().replace(\"T\", \"U\")\n    n      = len(seq)\n\n    # 查找 MSA 文件：文件名格式为 XXXX.MSA.fasta（取序列ID前4位大写）\n    pdb_id   = seq_id[:4].upper()\n    msa_path = os.path.join(MSA_DIR, f\"{pdb_id}.MSA.fasta\")\n\n    all_coords = []\n\n    if USE_RHOFOLD and os.path.exists(msa_path):\n        try:\n            data = get_features(msa_path, seq)\n            data = {k: v.to(device) if hasattr(v, \"to\") else v\n                    for k, v in data.items()}\n            with torch.no_grad():\n                output = model(**data)\n\n            # cord_tns_pred: (recycling, batch, n_res, n_atoms, 3)\n            # C1' 原子 index = 4\n            base_coords = output[\"cord_tns_pred\"][-1][0, :n, 4, :3].cpu().numpy()\n\n            for k in range(N_PREDS):\n                rng   = np.random.default_rng(k * 7)\n                noise = rng.normal(0, 0.25 * (k + 0.5), base_coords.shape).astype(np.float32)\n                all_coords.append(base_coords + noise)\n\n            if i % 5 == 0:\n                print(f\"[{i+1}/{len(test_df)}] RhoFold OK: {seq_id}  len={n}\")\n\n        except Exception as e:\n            print(f\"[{i+1}/{len(test_df)}] RhoFold 失败 ({seq_id}): {e}\")\n            all_coords = aform_fallback(n, N_PREDS)\n    else:\n        if i % 10 == 0:\n            reason = \"no MSA\" if USE_RHOFOLD else \"no model\"\n            print(f\"[{i+1}/{len(test_df)}] A-form fallback: {seq_id} ({reason})\")\n        all_coords = aform_fallback(n, N_PREDS)\n\n    # 写入每个残基的行\n    for res_idx in range(n):\n        row = {\n            \"ID\":      f\"{seq_id}_{res_idx + 1}\",\n            \"resname\": seq[res_idx],\n            \"resid\":   res_idx + 1,\n        }\n        for k in range(N_PREDS):\n            row[f\"x_{k+1}\"] = round(float(all_coords[k][res_idx, 0]), 3)\n            row[f\"y_{k+1}\"] = round(float(all_coords[k][res_idx, 1]), 3)\n            row[f\"z_{k+1}\"] = round(float(all_coords[k][res_idx, 2]), 3)\n        rows.append(row)\n\n# ==================================================================\n# 6. 保存 submission.csv\n# ==================================================================\ncol_order = [\"ID\", \"resname\", \"resid\"]\nfor k in range(1, N_PREDS + 1):\n    col_order += [f\"x_{k}\", f\"y_{k}\", f\"z_{k}\"]\n\nsubmission = pd.DataFrame(rows, columns=col_order)\nsubmission.to_csv(\"submission.csv\", index=False)\n\nprint(f\"\\n✅ 完成！\")\nprint(f\"   总行数  : {len(submission):,}\")\nprint(f\"   总列数  : {len(submission.columns)}\")\nprint(f\"   预测方式: {'RhoFold' if USE_RHOFOLD else 'A-form 螺旋 (兜底)'}\")\nprint(submission.head(3).to_string())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-13T17:17:05.781278Z","iopub.execute_input":"2026-03-13T17:17:05.781641Z","iopub.status.idle":"2026-03-13T17:17:10.25785Z","shell.execute_reply.started":"2026-03-13T17:17:05.781614Z","shell.execute_reply":"2026-03-13T17:17:10.256657Z"}},"outputs":[],"execution_count":null}]}