{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.12"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":118765,"databundleVersionId":15231210},{"sourceType":"datasetVersion","sourceId":14604295,"datasetId":9328538,"databundleVersionId":15440074},{"sourceType":"modelInstanceVersion","sourceId":781970,"databundleVersionId":16019818,"modelInstanceId":596642,"modelId":608905}],"dockerImageVersionId":31287,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":3722.480652,"end_time":"2026-03-10T22:50:32.865712","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2026-03-10T21:48:30.38506","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"ba4ca4e8","cell_type":"code","source":"import torch\nimport sys\nimport os\nos.environ['CUDA_VISIBLE_DEVICES'] = '0'\n\n# sys.argv: [脚本名, 设备, 精度]\nsys.argv = ['notebook', 'cuda' if torch.cuda.is_available() else 'cpu', 'fp32'] #'fp16' #'fp32'# 'bf16' #\ndevice = sys.argv[1]\nsys_dtype = sys.argv[2] if len(sys.argv) > 2 else 'fp32'\n\nprint('Using device:', device)\nprint('Using dtype:', sys_dtype)\n\n# dr settings\nNUM_CONF=5\nMAX_LENGTH=480\nMAX_CAT_LENGTH=2400\n\nCFG_DIR='cfg_97'\nCFG_MERGE=False\nDR_SCORE=False\nNO_SORT=False\nGET_CENTER=True\n\nFULL_ENERGY=False\n\nOPTIM_LENGTH=0\n\nDEVICE=device #'cuda' #'cpu'#\nPREC=sys_dtype ","metadata":{"execution":{"iopub.status.busy":"2026-03-11T08:15:26.072954Z","iopub.execute_input":"2026-03-11T08:15:26.073246Z","iopub.status.idle":"2026-03-11T08:15:29.855309Z","shell.execute_reply.started":"2026-03-11T08:15:26.073214Z","shell.execute_reply":"2026-03-11T08:15:29.854482Z"},"papermill":{"duration":3.840767,"end_time":"2026-03-10T21:48:36.783433","exception":false,"start_time":"2026-03-10T21:48:32.942666","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"d8ae669f","cell_type":"code","source":"if PREC=='fp16':\n    torch.set_default_dtype(torch.float16)\nif PREC=='bf16':\n    torch.set_default_dtype(torch.bfloat16)","metadata":{"execution":{"iopub.status.busy":"2026-03-11T08:15:29.857195Z","iopub.execute_input":"2026-03-11T08:15:29.85775Z","iopub.status.idle":"2026-03-11T08:15:29.86376Z","shell.execute_reply.started":"2026-03-11T08:15:29.857703Z","shell.execute_reply":"2026-03-11T08:15:29.862973Z"},"papermill":{"duration":0.010625,"end_time":"2026-03-10T21:48:36.798726","exception":false,"start_time":"2026-03-10T21:48:36.788101","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"8af62ced","cell_type":"markdown","source":"## Install requirements ","metadata":{"papermill":{"duration":0.002661,"end_time":"2026-03-10T21:48:36.804174","exception":false,"start_time":"2026-03-10T21:48:36.801513","status":"completed"},"tags":[]}},{"id":"1a7a85a4","cell_type":"code","source":"#rhofold+\n!pip install --no-deps '/kaggle/input/datasets/kami1976/biopython-cp312/biopython-1.86-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl'\n","metadata":{"execution":{"iopub.status.busy":"2026-03-11T08:15:29.864481Z","iopub.execute_input":"2026-03-11T08:15:29.864753Z","iopub.status.idle":"2026-03-11T08:15:33.111872Z","shell.execute_reply.started":"2026-03-11T08:15:29.864707Z","shell.execute_reply":"2026-03-11T08:15:33.111162Z"},"papermill":{"duration":3.302349,"end_time":"2026-03-10T21:48:40.109093","exception":false,"start_time":"2026-03-10T21:48:36.806744","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"d9564564","cell_type":"markdown","source":"# Drfold2","metadata":{"papermill":{"duration":0.002952,"end_time":"2026-03-10T21:48:40.115158","exception":false,"start_time":"2026-03-10T21:48:40.112206","status":"completed"},"tags":[]}},{"id":"73702b13","cell_type":"code","source":"from datetime import datetime\nimport pytz\nprint('LOGGING TIME OF START:',  datetime.strftime(datetime.now(pytz.timezone('Asia/Singapore')), \"%Y-%m-%d %H:%M:%S\"))\n\n\nprint('PIP INSTALL OK !!!!')\nimport os,sys\n\nimport pandas as pd\npd.set_option('display.max_columns', 20)\npd.set_option('display.expand_frame_repr', False)\n\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nfrom timeit import default_timer as timer\n\n\n\n# helper--\nclass dotdict(dict):\n\t__setattr__ = dict.__setitem__\n\t__delattr__ = dict.__delitem__\n\n\tdef __getattr__(self, name):\n\t\ttry:\n\t\t\treturn self[name]\n\t\texcept KeyError:\n\t\t\traise AttributeError(name)\n\ndef time_to_str(t, mode='min'):\n\tif mode=='min':\n\t\tt  = int(t)/60\n\t\thr = t//60\n\t\tmin = t%60\n\t\treturn '%2d hr %02d min'%(hr,min) \n\telif mode=='sec':\n\t\tt   = int(t)\n\t\tmin = t//60\n\t\tsec = t%60\n\t\treturn '%2d min %02d sec'%(min,sec)\n\n\telse:\n\t\traise NotImplementedError\n\ndef gpu_memory_use():\n    if torch.cuda.is_available():\n        device = torch.device(0)\n        free, total = torch.cuda.mem_get_info(device)\n        used= (total - free) / 1024 ** 3\n        return round(used,2)\n    else:\n        return 0\n\ndef set_aspect_equal(ax):\n\tx_limits = ax.get_xlim()\n\ty_limits = ax.get_ylim()\n\tz_limits = ax.get_zlim()\n\n\t# Compute the mean of each axis\n\tx_middle = np.mean(x_limits)\n\ty_middle = np.mean(y_limits)\n\tz_middle = np.mean(z_limits)\n\n\t# Compute the max range across all axes\n\tmax_range = max(x_limits[1] - x_limits[0],\n\t\t\t\t\ty_limits[1] - y_limits[0],\n\t\t\t\t\tz_limits[1] - z_limits[0]) / 2.0\n\n\t# Set the new limits to ensure equal scaling\n\tax.set_xlim(x_middle - max_range, x_middle + max_range)\n\tax.set_ylim(y_middle - max_range, y_middle + max_range)\n\tax.set_zlim(z_middle - max_range, z_middle + max_range)\n\n\nprint('torch',torch.__version__)\nprint('torch.cuda',torch.version.cuda)\n\nprint('IMPORT OK!!!')\nMODE = 'submit' #'local' # submit\n\nDATA_KAGGLE_DIR = '/kaggle/input/competitions/stanford-rna-3d-folding-2'\n\nif MODE == 'local':\n    valid_df = pd.read_csv(f'{DATA_KAGGLE_DIR}/validation_sequences.csv')\n    label_df = pd.read_csv(f'{DATA_KAGGLE_DIR}/validation_labels.csv')\n    label_df['target_id'] = label_df['ID'].apply(lambda x: '_'.join(x.split('_')[:-1]))\n\nif MODE == 'submit':\n\tvalid_df = pd.read_csv(f'{DATA_KAGGLE_DIR}/test_sequences.csv')\n\nprint('len(valid_df)',len(valid_df))\nprint(valid_df.iloc[0])\nprint('')\n\n\nprint('MODE:', MODE)\nprint('SETTING OK!!!')\n","metadata":{"execution":{"iopub.status.busy":"2026-03-11T08:15:33.113125Z","iopub.execute_input":"2026-03-11T08:15:33.11356Z","iopub.status.idle":"2026-03-11T08:15:33.438276Z","shell.execute_reply.started":"2026-03-11T08:15:33.113534Z","shell.execute_reply":"2026-03-11T08:15:33.437482Z"},"papermill":{"duration":0.845618,"end_time":"2026-03-10T21:48:40.963454","exception":false,"start_time":"2026-03-10T21:48:40.117836","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"e9e8bcd1","cell_type":"code","source":"sys.path.append('/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/PotentialFold')\nimport a2b\ndef frame_coor_to_C1(coor, seq, BASE_COOR, OTHER_COOR):\n    \"\"\"\n    使用四元数批量运算将模型输出的局部 frame 坐标直接转换为每个位点 C1' 原子全局坐标。\n\n    参数\n    ----\n    coor : numpy.ndarray 或 torch.Tensor, shape=(L, 3, 3)\n        模型预测的每个位点 P, C4', N 三点的局部坐标\n    seq : str, length=L\n        核苷酸序列\n    BASE_COOR : dict\n        Get_base 所需的三个骨架原子模板坐标常量\n    OTHER_COOR : dict\n        Get_base 所需的五个糖环原子模板坐标常量\n\n    返回\n    ----\n    c1_xyz : numpy.ndarray, shape=(L, 3)\n        每个位点 C1' 原子在全局坐标系下的 XYZ\n    \"\"\"\n    # 1. 转为 torch.Tensor\n    tx = torch.as_tensor(coor, dtype=torch.float32)\n    basex = torch.from_numpy(Get_base(seq, BASE_COOR)).to(tx.dtype)\n    otherx = torch.from_numpy(Get_base(seq, OTHER_COOR)).to(tx.dtype)\n    L = len(seq)\n\n    # 2. 初始化参数向量 x：21 维，其中前 18 维用于旋转投影，后 3 维用于平移\n    x = torch.rand((L, 21), dtype=tx.dtype, device=tx.device)\n\n    # 3. 中心化处理：计算质心 biasq，并得到零均值坐标 q\n    biasq = tx.mean(dim=1, keepdim=True)            # (L, 1, 3)\n    q = tx - biasq                                  # (L, 3, 3)\n\n    # 4. 计算旋转投影 m，并填充到 x 的前两组 9 维\n    m = torch.einsum('bnz,bny->bzy', basex, q).reshape(L, -1)\n    x[:, :9] = m\n    x[:, 9:18] = m\n\n    # 5. 填入平移参数（将 biasq 作为平移量）\n    x[:, 18:] = biasq.squeeze(1)\n    rama = x.double()  # 四元数参数需要 double 精度\n\n    # 6. 调用 quat2b：对模板中所有糖环原子做旋转+平移\n    #    otherx: (L, 5, 3) -> other_xyz: (L, 5, 3)\n    other_xyz = a2b.quat2b(otherx.double(), rama[:, 9:]).float().cpu().numpy()\n\n    # 7. 提取第 5 个原子（C1'）并返回\n    c1_xyz = other_xyz[:, 4, :]\n    return c1_xyz\n\ndef kabsch_rotation(P: torch.Tensor, Q: torch.Tensor) -> torch.Tensor:\n    \"\"\"\n    Kabsch算法：计算将点集P对齐到点集Q的最优旋转矩阵R\n\n    参数:\n        P: (N, 3) 源点集\n        Q: (N, 3) 目标点集\n\n    返回:\n        R: (3, 3) 旋转矩阵\n    \"\"\"\n    # 1. 计算质心\n    centroid_P = P.mean(dim=0, keepdim=True)  # (1, 3)\n    centroid_Q = Q.mean(dim=0, keepdim=True)  # (1, 3)\n\n    # 2. 中心化\n    P_centered = P - centroid_P  # (N, 3)\n    Q_centered = Q - centroid_Q  # (N, 3)\n\n    # 3. 计算协方差矩阵 H = P^T * Q\n    H = P_centered.T @ Q_centered  # (3, 3)\n\n    # 4. SVD分解\n    U, S, Vt = torch.linalg.svd(H)\n\n    # 5. 计算旋转矩阵 R = V * U^T\n    R = Vt.T @ U.T\n\n    # 6. 处理反射（确保行列式为1，表示纯旋转）\n    det_R = torch.det(R)\n    if det_R < 0:\n        print(f\"Warning: det(R) = {det_R}, 正在修正反射...\")\n        Vt[-1, :] *= -1\n        R = Vt.T @ U.T\n\n    return R\n\n\ndef kabsch_align(coor1: torch.Tensor, coor2: torch.Tensor,\n                n_overlap: int = 5) -> torch.Tensor:\n    \"\"\"\n    使用Kabsch算法对齐两个片段\n\n    参数:\n        coor1: (L1, 3, 3) 前一段坐标\n        coor2: (L2, 3, 3) 后一段坐标\n        n_overlap: 用于对齐的碱基数目（建议3-10）\n\n    返回:\n        coor2_aligned: (L2, 3, 3) 对齐后的坐标\n    \"\"\"\n    L1, L2 = coor1.shape[0], coor2.shape[0]\n\n    # 限制n_overlap不能超过两段的长度\n    n_overlap = min(n_overlap, L1, L2)\n\n    # 取前一段最后n_overlap个碱基的C4'坐标\n    # coor[:, 1] 是 C4' 坐标\n    P = coor1[-n_overlap:, 1, :]  # (n_overlap, 3)\n    Q = coor2[:n_overlap, 1, :]   # (n_overlap, 3)\n\n    # 计算旋转矩阵\n    R = kabsch_rotation(P, Q)  # (3, 3)\n\n    # 计算平移向量（使用质心）\n    centroid_P = P.mean(dim=0)  # (3,)\n    centroid_Q = Q.mean(dim=0)  # (3,)\n    t = centroid_Q - R @ centroid_P  # (3,)\n\n    # 对coor2整个链做刚性变换\n    # 先中心化\n    coor2_centered = coor2 - centroid_Q  # (L2, 3, 3)\n\n    # 旋转\n    coor2_rotated = coor2_centered @ R.T  # (L2, 3, 3)\n\n    # 平移（对齐到coor1末尾）\n    coor2_aligned = coor2_rotated + centroid_Q  # 先回到原始位置，再应用变换\n\n    # 实际上需要将coor2对齐到coor1末尾：\n    # R * (x - centroid_Q) + centroid_Q -> R*x - R*centroid_Q + centroid_Q\n    # 然后我们希望 R*x' + t = coor1[-1]\n    # 正确做法：\n    coor2_aligned = coor2 @ R.T + (centroid_Q - R @ centroid_Q)  # (L2, 3, 3)\n\n    return coor2_aligned\n\n\ndef concat_coor_kabsch(out1: dict, out2: dict, n_overlap: int = 5) -> np.ndarray:\n    \"\"\"\n    使用Kabsch算法将out2对齐到out1末尾，并拼接（去掉重复的第一个碱基）\n\n    参数:\n        out1: dict, 包含 'coor': (L1, 3, 3)\n        out2: dict, 包含 'coor': (L2, 3, 3)\n        n_overlap: 用于对齐的碱基数目（建议3-10）\n\n    返回:\n        coor_cat: (L1 + L2 - 1, 3, 3) 拼接后的坐标\n    \"\"\"\n    # 转为 torch.Tensor\n    coor1 = torch.as_tensor(out1['coor'], dtype=torch.float64)  # (L1, 3, 3)\n    coor2 = torch.as_tensor(out2['coor'], dtype=torch.float64)  # (L2, 3, 3)\n\n    # 使用Kabsch算法对齐\n    coor2_aligned = kabsch_align(coor1, coor2, n_overlap=n_overlap)\n\n    # 拼接（跳过coor2_aligned[0]，因为它已与coor1[-1]重合）\n    coor_cat = torch.cat([coor1, coor2_aligned[1:]], dim=0)  # (L1+L2-1, 3, 3)\n\n    return coor_cat.cpu().numpy()\n\ndef concat_coor_original(out1: dict, out2: dict) -> np.ndarray:\n    \"\"\"\n    将 out2 通过四元数（quat2b）做旋转 + 平移对齐到 out1 末尾，并拼接（去掉重复的第一个碱基）。\n    out1/out2 都是 dict，包含\n      'coor': np.ndarray, shape=(L, 3, 3)  —— 每个碱基的 P, C4, N 坐标\n    返回拼接后的坐标数组，shape = (L1 + L2 - 1, 3, 3)\n    \"\"\"\n    # 1. 转为 torch.DoubleTensor\n    coor1 = torch.as_tensor(out1['coor'], dtype=torch.float64)   # (L1,3,3)\n    coor2 = torch.as_tensor(out2['coor'], dtype=torch.float64)   # (L2,3,3)\n\n    # 2. 取公共碱基的 frame：out1 最后一个、out2 第一个\n    f1 = coor1[-1]   # (3,3)\n    f2 = coor2[0]    # (3,3)\n\n    # 3. 中心化：去除各自质心\n    bias1 = f1.mean(dim=0)   # (3,)\n    bias2 = f2.mean(dim=0)   # (3,)\n    basex = f1 - bias1       # (3,3)\n    q     = f2 - bias2       # (3,3)\n\n    # 4. 用 einsum 计算旋转矩阵 R = basex · q^T\n    #    R_{ij} = sum_z basex_{iz} * q_{jz}\n    R = torch.einsum('iz,jz->ij', basex, q)   # (3,3)\n\n    # 5. 计算平移向量 t，使得 R·bias2 + t = bias1  ⇒  t = bias1 - R·bias2\n    t = bias1 - (R @ bias2)                  # (3,)\n\n    # 6. 构造 rama 参数 (L2, 12)：前 9 维放 R.flatten，后 3 维放 t\n    L2 = coor2.shape[0]\n    rama = torch.empty((L2, 12), dtype=torch.float64, device=coor2.device)\n    R_flat = R.reshape(1, 9).repeat(L2, 1)    # (L2,9)\n    t_rep  = t.reshape(1, 3).repeat(L2, 1)    # (L2,3)\n    rama[:, :9] = R_flat\n    rama[:, 9:] = t_rep\n\n    # 7. 调用 quat2b：对 out2 整条链做刚性变换\n    #    输入 coord=(L2,3,3)， rama=(L2,12)，输出同 shape\n    coor2_aligned = a2b.quat2b(coor2, rama)   # torch.Tensor (L2,3,3)\n\n    # 8. 拼接（跳过 coor2_aligned[0]，因为它已与 coor1[-1] 重合）\n    coor_cat = torch.cat([coor1, coor2_aligned[1:]], dim=0)  # (L1+L2-1,3,3)\n\n    return coor_cat.cpu().numpy()\n\n# ============ Switch配置 ============\nUSE_KABSCH = True  # True: 使用Kabsch多碱基对齐, False: 使用原始单碱基frame方法\nN_OVERLAP = 5  # Kabsch模式下用于对齐的碱基数目 (建议3-10)\n\ndef concat_coor(out1: dict, out2: dict) -> np.ndarray:\n    \"\"\"\n    对齐并拼接坐标的接口函数，支持切换Kabsch/原始方法\n    \"\"\"\n    if USE_KABSCH:\n        return concat_coor_kabsch(out1, out2, n_overlap=N_OVERLAP)\n    else:\n        return concat_coor_original(out1, out2)","metadata":{"execution":{"iopub.status.busy":"2026-03-11T08:15:33.439236Z","iopub.execute_input":"2026-03-11T08:15:33.439864Z","iopub.status.idle":"2026-03-11T08:15:33.508798Z","shell.execute_reply.started":"2026-03-11T08:15:33.439838Z","shell.execute_reply":"2026-03-11T08:15:33.508087Z"},"papermill":{"duration":0.07923,"end_time":"2026-03-10T21:48:41.045744","exception":false,"start_time":"2026-03-10T21:48:40.966514","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"ff170e8d","cell_type":"code","source":"def Get_base(seq, basenpy_standard):\n    n_atoms = basenpy_standard.shape[1]\n    basenpy = np.zeros([len(seq), n_atoms, 3])\n    seqnpy = np.array(list(seq))\n    basenpy[seqnpy=='A'] = basenpy_standard[0]\n    basenpy[seqnpy=='a'] = basenpy_standard[0]\n    basenpy[seqnpy=='G'] = basenpy_standard[1]\n    basenpy[seqnpy=='g'] = basenpy_standard[1]\n    basenpy[seqnpy=='C'] = basenpy_standard[2]\n    basenpy[seqnpy=='c'] = basenpy_standard[2]\n    basenpy[seqnpy=='U'] = basenpy_standard[3]\n    basenpy[seqnpy=='u'] = basenpy_standard[3]\n    basenpy[seqnpy=='T'] = basenpy_standard[3]\n    basenpy[seqnpy=='t'] = basenpy_standard[3]\n    return basenpy","metadata":{"execution":{"iopub.status.busy":"2026-03-11T08:15:33.509853Z","iopub.execute_input":"2026-03-11T08:15:33.510145Z","iopub.status.idle":"2026-03-11T08:15:33.515286Z","shell.execute_reply.started":"2026-03-11T08:15:33.510113Z","shell.execute_reply":"2026-03-11T08:15:33.514638Z"},"papermill":{"duration":0.010419,"end_time":"2026-03-10T21:48:41.059817","exception":false,"start_time":"2026-03-10T21:48:41.049398","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"8d0278e0","cell_type":"code","source":"import numpy as np\n\ndef score_energy_one_simple(seq, target_id, out):\n    \"\"\"\n    计算 RNA 结构的粗粒化自由能，只基于每个碱基的 P–C4'–N (N1 or N9) 坐标帧。\n\n    参数：\n      seq:      RNA 序列字符串，长度 L\n      target_id:结构标识符（仅用于日志或输出，不影响计算）\n      out:      模型输出字典，必须包含：\n                out['coor']: numpy array, shape (L, 3, 3)\n                             coor[i,0] = P 坐标\n                             coor[i,1] = C4′ 坐标\n                             coor[i,2] = N1/N9 坐标\n\n    返回：\n      total_energy: float，自由能估计值（越低表示越稳定）\n    \"\"\"\n\n    coor = out['coor']  # (L, 3, 3)\n    L = len(seq)\n\n    # —— 一、键长能 —— #\n    # 理想键长 (Å)：P–C4′ ~ 1.60, C4′–N ~ 1.47 (典型值，可根据经验或晶体数据调整)\n    d0_P_C4  = 1.60\n    d0_C4_N  = 1.47\n    k_bond   = 100.0  # 键长力常数 (kcal/mol/Ų)\n\n    energy_bond = 0.0\n    for i in range(L):\n        P  = coor[i,0]\n        C4 = coor[i,1]\n        N  = coor[i,2]\n        d_PC4 = np.linalg.norm(P - C4)\n        d_C4N = np.linalg.norm(C4 - N)\n        energy_bond += k_bond * (d_PC4 - d0_P_C4)**2\n        energy_bond += k_bond * (d_C4N - d0_C4_N)**2\n\n    # —— 二、键角能 —— #\n    #  ideal angle P–C4′–N ~ 109.5°（四面体角），用简谐势\n    theta0   = np.deg2rad(109.5)\n    k_angle  = 20.0   # 键角力常数 (kcal/mol/rad²)\n\n    energy_angle = 0.0\n    for i in range(L):\n        P  = coor[i,0]\n        C4 = coor[i,1]\n        N  = coor[i,2]\n        v1 = P  - C4\n        v2 = N  - C4\n        cos_theta = np.dot(v1, v2) / (np.linalg.norm(v1)*np.linalg.norm(v2) + 1e-8)\n        theta = np.arccos(np.clip(cos_theta, -1.0, 1.0))\n        energy_angle += k_angle * (theta - theta0)**2\n\n    # —— 三、堆积相互作用 —— #\n    # 简单地对相邻碱基 C4′–C4′ 间距进行势能惩罚，鼓励相邻堆积：\n    #   ideal 平面内距离 ~ 3.4 Å\n    d0_stack = 3.4\n    k_stack  = 5.0   # (kcal/mol/Ų)\n\n    energy_stack = 0.0\n    for i in range(L-1):\n        C4_i   = coor[i  ,1]\n        C4_ip1 = coor[i+1,1]\n        d = np.linalg.norm(C4_i - C4_ip1)\n        energy_stack += k_stack * (d - d0_stack)**2\n\n    # —— 四、合计 —— #\n    total_energy = energy_bond + energy_angle + energy_stack\n\n    # （可选）打印调试\n    # print(f\"[{target_id}] bond={energy_bond:.2f}, angle={energy_angle:.2f}, stack={energy_stack:.2f} → total={total_energy:.2f}\")\n\n    return total_energy","metadata":{"execution":{"iopub.status.busy":"2026-03-11T08:15:33.517426Z","iopub.execute_input":"2026-03-11T08:15:33.517895Z","iopub.status.idle":"2026-03-11T08:15:33.530695Z","shell.execute_reply.started":"2026-03-11T08:15:33.517873Z","shell.execute_reply":"2026-03-11T08:15:33.530012Z"},"papermill":{"duration":0.012564,"end_time":"2026-03-10T21:48:41.075312","exception":false,"start_time":"2026-03-10T21:48:41.062748","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"8e3da50d","cell_type":"code","source":"import numpy as np\n\ndef score_energy_one_full(seq, target_id, out, paired=None,\n                          # bond parameters\n                          d0_P_C4=1.60, d0_C4_N=1.47, k_bond=100.0,\n                          # angle parameters\n                          theta0=np.deg2rad(109.5), k_angle=20.0,\n                          # stacking parameters\n                          d0_stack=3.4, k_stack=5.0,\n                          # dihedral parameters\n                          phi0=np.deg2rad(180.0), k_dihedral=5.0,\n                          # hydrogen-bond parameters\n                          d0_hb=2.9, k_hb=10.0,\n                          # Lennard-Jones parameters\n                          sigma=4.0, epsilon=0.1,\n                          # Debye-Hückel electrostatics\n                          q_P=-1.0, epsilon_r=80.0, kappa=10.0,\n                          k_e=332.0637):\n    \"\"\"\n    计算 RNA 结构的粗粒化自由能，包含键长、键角、堆积、二面角、氢键、范德华和静电相互作用。\n\n    参数：\n      seq:      RNA 序列字符串，长度 L\n      target_id:结构标识符（仅用于日志或输出，不影响计算）\n      out:      模型输出字典，包含：\n                out['coor']: numpy array, shape (L, 3, 3)\n                可选 out['paired']: 碱基配对列表 [(i,j), ...]\n    可选参数：\n      paired:   碱基配对索引列表，默认从 out['paired'] 读取\n      其余为各项势的参数。\n\n    返回：\n      total_energy: float，自由能估计值（越低表示越稳定）\n    \"\"\"\n    coor = out['coor']\n    L = len(seq)\n    if paired is None:\n        paired = out.get('paired', [])\n\n    # —— 一、键长能 —— #\n    E_bond = 0.0\n    for i in range(L):\n        P  = coor[i,0]; C4 = coor[i,1]; N = coor[i,2]\n        d_PC4 = np.linalg.norm(P - C4)\n        d_C4N = np.linalg.norm(C4 - N)\n        E_bond += k_bond * (d_PC4 - d0_P_C4)**2\n        E_bond += k_bond * (d_C4N - d0_C4_N)**2\n\n    # —— 二、键角能 —— #\n    E_angle = 0.0\n    for i in range(L):\n        P, C4, N = coor[i]\n        v1 = P - C4; v2 = N - C4\n        cost = np.dot(v1, v2) / (np.linalg.norm(v1)*np.linalg.norm(v2) + 1e-8)\n        theta = np.arccos(np.clip(cost, -1.0, 1.0))\n        E_angle += k_angle * (theta - theta0)**2\n\n    # —— 三、堆积相互作用 —— #\n    E_stack = 0.0\n    for i in range(L-1):\n        d = np.linalg.norm(coor[i,1] - coor[i+1,1])\n        E_stack += k_stack * (d - d0_stack)**2\n\n    # —— 四、二面角（扭转）—— #\n    def torsion_angle(a, b, c, d):\n        b1, b2, b3 = b-a, c-b, d-c\n        n1 = np.cross(b1, b2); n2 = np.cross(b2, b3)\n        n1 /= np.linalg.norm(n1) + 1e-8; n2 /= np.linalg.norm(n2) + 1e-8\n        cos_phi = np.dot(n1, n2)\n        return np.arccos(np.clip(cos_phi, -1, 1))\n\n    E_dihedral = 0.0\n    for i in range(L-1):\n        a = coor[i,0]; b = coor[i,1]; c = coor[i,2]; d = coor[i+1,0]\n        phi = torsion_angle(a, b, c, d)\n        E_dihedral += k_dihedral * (phi - phi0)**2\n\n    # —— 五、氢键势 —— #\n    E_hb = 0.0\n    for i, j in paired:\n        d = np.linalg.norm(coor[i,2] - coor[j,2])\n        E_hb += k_hb * (d - d0_hb)**2\n\n    # —— 六、范德华（Lennard-Jones）—— #\n    E_LJ = 0.0\n    for i in range(L):\n        for j in range(i+2, L):  # 排除1-2相邻\n            r = np.linalg.norm(coor[i,1] - coor[j,1])\n            sr6 = (sigma / (r + 1e-8))**6\n            sr12 = sr6 * sr6\n            E_LJ += 4 * epsilon * (sr12 - sr6)\n\n    # —— 七、静电（Debye-Hückel）—— #\n    E_elec = 0.0\n    for i in range(L):\n        for j in range(i+1, L):\n            r = np.linalg.norm(coor[i,0] - coor[j,0])\n            prefac = k_e * q_P * q_P / epsilon_r\n            E_elec += prefac * np.exp(-r / kappa) / (r + 1e-8)\n\n\n    weights = {\n        'bond': 1.0,\n        'angle': 1.0,\n        'stack': 1.2,\n        'dihedral': 1.0,\n        'hb': 1.5,\n        'LJ': 0.8,\n        'elec': 0.6\n    }\n    \n    total_energy = (\n        weights['bond'] * E_bond +\n        weights['angle'] * E_angle +\n        weights['stack'] * E_stack +\n        weights['dihedral'] * E_dihedral +\n        weights['hb'] * E_hb +\n        weights['LJ'] * E_LJ +\n        weights['elec'] * E_elec\n    )\n\n    # total_energy = (E_bond + E_angle + E_stack +\n    #                 E_dihedral + E_hb + E_LJ + E_elec)\n\n    # print(f\"[{target_id}] bond={E_bond:.2f}, angle={E_angle:.2f}, stack={E_stack:.2f}, \\\n#          dihedral={E_dihedral:.2f}, hb={E_hb:.2f}, LJ={E_LJ:.2f}, elec={E_elec:.2f} -> total={total_energy:.2f}\")\n    return total_energy","metadata":{"execution":{"iopub.status.busy":"2026-03-11T08:15:33.531568Z","iopub.execute_input":"2026-03-11T08:15:33.531803Z","iopub.status.idle":"2026-03-11T08:15:33.548064Z","shell.execute_reply.started":"2026-03-11T08:15:33.531783Z","shell.execute_reply":"2026-03-11T08:15:33.547445Z"},"papermill":{"duration":0.018366,"end_time":"2026-03-10T21:48:41.096699","exception":false,"start_time":"2026-03-10T21:48:41.078333","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"63cecc07","cell_type":"code","source":"import os\nimport json\nimport pickle\nimport numpy as np\nimport tempfile\nfrom Bio.PDB import PDBParser\nsys.path.append('/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/PotentialFold')\nfrom Optimization import Structure\n\ndef score_energy_one(seq, target_id, out):\n    with tempfile.TemporaryDirectory() as tmpdirname:\n        # 输出临时fasta\n        fastafile = os.path.join(tmpdirname, 'tmp.fasta')\n        with open(fastafile, 'w') as f:\n            f.write(f'>{target_id}\\n{seq}\\n')\n        # 输出临时ret\n        retfile = os.path.join(tmpdirname, 'tmp.ret')\n        with open(retfile, 'wb') as f:\n            f.write(pickle.dumps(out))\n        # foldconfig = '/kaggle/input/drfold-model/DRfold2/cfg_for_selection.json'\n        # foldconfig = '/kaggle/input/drfold-model/DRfold2/cfg_for_folding.json'\n        foldconfig = '/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/cfg_for_folding.json'\n        # foldconfig = 'cfg_for_folding.json'\n        save_prefix = os.path.join(tmpdirname, 'tmp.json')\n        stru=Structure(fastafile,[retfile],save_prefix,0,foldconfig)\n        rama=stru.init_quat(0).data.numpy()\n        energy=stru.obj_func_np(rama)\n        return energy\n\n\n\ndef optimize_coor(seq, target_id, out):\n    print('Optimizing structure for ', target_id)\n    if len(out['coor'])>len(out['plddt']):\n        # 补全plddt\n        if len(out['plddt'])>0:\n            mean_plddt = np.mean(out['plddt'])\n            out['plddt'] = np.concatenate([out['plddt'], np.full((len(out['coor'])-len(out['plddt'])), mean_plddt)])\n        else:\n            out['plddt'] = np.full((len(out['coor'])), 0.0)\n            \n    if len(out['coor'])<len(out['plddt']):\n        mean_plddt = np.mean(out['plddt'])\n        out['plddt'] = np.full((len(out['coor'])), mean_plddt)\n    \n    with tempfile.TemporaryDirectory() as tmpdirname:\n        # 1) 写入临时 fasta\n        fastafile = os.path.join(tmpdirname, 'tmp.fasta')\n        with open(fastafile, 'w') as f:\n            f.write(f'>{target_id}\\n{seq}\\n')\n\n        # 2) 写入临时 ret\n        retfile = os.path.join(tmpdirname, 'tmp.ret')\n        with open(retfile, 'wb') as f:\n            pickle.dump(out, f)\n\n        # 3) 调用 Structure 完成折叠优化\n        # foldconfig = '/kaggle/input/drfold-model/DRfold2/cfg_for_folding.json'\n        foldconfig = '/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/cfg_for_folding.json'\n        # foldconfig = 'cfg_for_folding.json'\n        save_prefix = os.path.join(tmpdirname, 'tmp')\n        stru = Structure(fastafile, [retfile], save_prefix, 0, foldconfig)\n        stru.foldning()\n\n        # 4) 解析输出的 PDB\n        pdb_file = save_prefix + '.pdb'\n        parser = PDBParser(QUIET=True)\n        structure = parser.get_structure(target_id, pdb_file)\n\n        # 5) 收集所有标准残基（排除 HETATM）\n        residues = [\n            res for res in structure.get_residues()\n            if res.id[0] == ' '\n        ]\n        # 检查残基数和序列长度一致\n        L = len(seq)\n        if len(residues) != L:\n            raise ValueError(f\"PDB 中残基数 ({len(residues)}) 与序列长度 ({L}) 不一致\")\n\n        # 6) 提取三种原子坐标\n        atom_order = ['P', \"C4'\", 'N1/N9']\n        coor = np.zeros((L, 3, 3), dtype=float)\n\n        for i, res in enumerate(residues):\n            coor[i, :, :] = np.nan  # 初始化为 NaN\n            \n            if 'P' in res:\n                coord = res['P'].get_vector().get_array()\n                coor[i, 0, :] = coord\n            if \"C4'\" in res:\n                coord = res[\"C4'\"].get_vector().get_array()\n                coor[i, 1, :] = coord\n            if 'N1' in res:\n                coord = res['N1'].get_vector().get_array()\n                coor[i, 2, :] = coord\n            elif 'N9' in res:\n                coord = res['N9'].get_vector().get_array()\n                coor[i, 2, :] = coord\n\n        return coor","metadata":{"execution":{"iopub.status.busy":"2026-03-11T08:15:33.549042Z","iopub.execute_input":"2026-03-11T08:15:33.549307Z","iopub.status.idle":"2026-03-11T08:15:34.106133Z","shell.execute_reply.started":"2026-03-11T08:15:33.549286Z","shell.execute_reply":"2026-03-11T08:15:34.105581Z"},"papermill":{"duration":0.957717,"end_time":"2026-03-10T21:48:42.057441","exception":false,"start_time":"2026-03-10T21:48:41.099724","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"5f12a45e","cell_type":"code","source":"import sys\nimport os\nimport numpy as np\nimport sys\nsys.path.append('/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/')\nsys.path.append('/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/PotentialFold')\n# 添加自定义代码路径\nsys.path.append(f'/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/{CFG_DIR}')\nsys.path.append(f'/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/{CFG_DIR}/RNALM2')\n\nBASE_COOR = np.load('/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/PotentialFold/lib/base.npy')\nOTHER_COOR = np.load('/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/PotentialFold/lib/other2.npy')\nSIDE_COOR = np.load('/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/PotentialFold/lib/side.npy')\n\n\nfrom EvoMSA2XYZ import MSA2XYZ\nfrom RNALM2.Model import RNA2nd\nfrom data import parse_seq\n\n\n\n# data helper\ndef make_data(seq, device):\n    aa_type = parse_seq(seq)\n    base = Get_base(seq, BASE_COOR)\n    seq_idx = np.arange(len(seq)) + 1\n\n    msa = aa_type[None, :]\n    msa = torch.from_numpy(msa)\n    msa = torch.cat([msa, msa], 0)  # ???\n    msa = F.one_hot(msa.long(), 6).float()\n\n    base_x = torch.from_numpy(base).float()\n    seq_idx = torch.from_numpy(seq_idx).long()\n\n    msa, base_x, seq_idx = msa.to(device), base_x.to(device), seq_idx.to(device)\n    return msa, base_x, seq_idx\n\n\ndef solution_to_submit_df(solution):\n    submit_df = []\n    for k,s in solution.items():\n        df = coord_to_df(s.sequence, s.coord, s.target_id)\n        submit_df.append(df)\n    \n    submit_df = pd.concat(submit_df)\n    return submit_df\n \n\ndef coord_to_df(sequence, coord, target_id):\n    L = len(sequence)\n    df = pd.DataFrame()\n    df['ID'] = [f'{target_id}_{i + 1}' for i in range(L)]\n    df['resname'] = [s for s in sequence]\n    df['resid'] = [i + 1 for i in range(L)]\n\n    num_coord = len(coord)\n    for j in range(num_coord):\n        df[f'x_{j+1}'] = coord[j][:, 0]\n        df[f'y_{j+1}'] = coord[j][:, 1]\n        df[f'z_{j+1}'] = coord[j][:, 2]\n    return df\n\n\nout_dir = '/kaggle/working/model-output'\nos.makedirs(out_dir, exist_ok=True)\n","metadata":{"execution":{"iopub.status.busy":"2026-03-11T08:15:34.107098Z","iopub.execute_input":"2026-03-11T08:15:34.107526Z","iopub.status.idle":"2026-03-11T08:15:37.978589Z","shell.execute_reply.started":"2026-03-11T08:15:34.107501Z","shell.execute_reply":"2026-03-11T08:15:37.977762Z"},"papermill":{"duration":4.138093,"end_time":"2026-03-10T21:48:46.198551","exception":false,"start_time":"2026-03-10T21:48:42.060458","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"8dbcdedf","cell_type":"code","source":"import warnings\nwarnings.filterwarnings(\"ignore\")\n\ndef run_submit(valid_df):\n    \n    #load model (these are moified versions, not the same from their github repo)\n    rnalm = RNA2nd(dict(\n        s_in_dim=5,\n        z_in_dim=2,\n        s_dim= 512,\n        z_dim= 128,\n        N_elayers=18,\n    ))\n    rnalm_file = '/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/model_hub/RCLM/epoch_67000'\n    print(rnalm_file)\n    print(\n        rnalm.load_state_dict(torch.load(rnalm_file, map_location='cpu', weights_only=True), strict=False)\n        #Unexpected key(s) in state_dict: \"ss_head.linear.weight\", \"ss_head.linear.bias\".\n    )\n    rnalm = rnalm.to(DEVICE)\n    if(PREC=='fp16'):\n        rnalm=rnalm.half()\n        \n    if PREC=='bf16':\n        rnalm = rnalm.bfloat16()\n        \n    rnalm = rnalm.eval()\n\n    #---\n    msa2xyz = MSA2XYZ(\n        seq_dim=6,\n        msa_dim=7,\n        N_ensemble=1,\n        N_cycle=8,  # 8\n        m_dim=64,\n        s_dim=64,\n        z_dim=64,\n    )\n    msa2xyz_file = [f'/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/model_hub/{CFG_DIR}/model_{i}' for i in range(20)]\n    if CFG_MERGE:\n        msa2xyz_file = [\n            f'/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/model_hub/cfg_97/model_{i}'\n            for i in range(20)\n        ] + [\n            f'/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/model_hub/cfg_95/model_{i}'\n            for i in range(20)\n        ] + [\n            f'/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/model_hub/cfg_96/model_{i}'\n            for i in range(20)\n        ] + [\n            f'/kaggle/input/models/dfgtde/drfold/pytorch/default/1/DRfold2/model_hub/cfg_99/model_{i}'\n            for i in range(20)\n        ]\n    num_msa2xyz = len(msa2xyz_file) \n    msa2xyz_state_dict = []\n    for c in range(num_msa2xyz):\n        if c==0: print(msa2xyz_file[c])\n        m = torch.load(msa2xyz_file[c], map_location='cpu', weights_only=True)\n        msa2xyz_state_dict.append(m)\n        \n    #print(msa2xyz.load_state_dict(msa2xyz_state_dict[0], strict=True))\n    print(msa2xyz.load_state_dict(msa2xyz_state_dict[0], strict=False))\n    msa2xyz = msa2xyz.to(DEVICE)\n    if(PREC=='fp16'):\n        msa2xyz=msa2xyz.half()\n\n    if PREC=='bf16':\n        msa2xyz = msa2xyz.bfloat16()\n    msa2xyz = msa2xyz.eval()\n    \n    msa2xyz.msaxyzone.premsa.rnalm = rnalm\n\n    #---\n    # start here !!!!!!!!!!!!!!!!!!!!!!\n    #valid_df = valid_df.iloc[[0,1]].reset_index(drop=True)\n\n\n    submit_df = [] \n    total_time_taken = 0\n    max_gpu_mem_used = 0\n\n    for i, row in valid_df.iterrows():\n        start_timer = timer()\n        target_id = row.target_id  # 'R1116' #casp15 R1116: len(157)\n        sequence = row.sequence\n        seq = row.sequence  \n        L = len(seq)\n        if L > MAX_CAT_LENGTH:\n            seq = seq[:MAX_CAT_LENGTH]\n        # else:\n        #     continue\n        print(i, target_id, L, len(seq), seq[:75] + '...')\n\n        \n        if len(seq)>480:\n            model_to_try=[16, 9, 1, 2, 0]\n        elif len(seq)>200:\n            # model_to_try = [0,1,2,8,9]\n            model_to_try = [13, 6, 14, 5, 3]\n        elif  len(seq)>100:\n            # model_to_try = [0,2,4,6,8,10,12,14,16,18]#list(range(min(num_msa2xyz,10)))\n            model_to_try = [13, 6, 14, 12, 7, 2, 5, 19, 10, 9]\n            if CFG_MERGE:\n                model_to_try = [24, 34, 20, 13, 6, 37, 28, 25, 14, 39]\n        else:\n            # model_to_try = list(range(num_msa2xyz))\n            \n            # model_to_try = list(range(20))\n            model_to_try = [1, 2, 0, 8, 7, 5, 6, 14, 10, 18, 4, 13, 3, 17, 19, 11, 12, 15, 16, 9]\n            \n            # if CFG_MERGE:\n            #     model_to_try = list(range(20)) + list(range(40,60)) + list(range(60, 80))\n            \n        if NO_SORT:\n            model_to_try=model_to_try[:5]\n\n\n        # 分段预测\n        def predict_segment(seq):\n            msa, base_x, seq_idx = make_data(seq, DEVICE)\n            with torch.no_grad():\n                if PREC=='fp16':\n                    msa, base_x = msa.half(), base_x.half()\n                if PREC=='bf16':\n                    msa, base_x = msa.bfloat16(), base_x.bfloat16()\n                return msa2xyz.pred(msa, seq_idx, None, base_x, np.array(list(seq)))\n\n                \n        energy = []\n        coordinate=[]\n        outputs=[]\n        for c in model_to_try:\n            msa2xyz.load_state_dict(msa2xyz_state_dict[c], strict=False)\n\n            if len(seq) <= MAX_LENGTH:\n                outs = [ predict_segment(seq) ]\n            else:\n                step = MAX_LENGTH - 1\n                outs = []\n                for s in range(0, len(seq), step):\n                    seg = seq[s : min(s+MAX_LENGTH, len(seq))]\n                    outs.append(predict_segment(seg))\n                    \n            out_cat = outs[0]\n            for out_seg in outs[1:]:\n                out_cat = {'coor': concat_coor(out_cat, out_seg)}\n                    \n            if NO_SORT:\n                e=0\n            elif len(model_to_try)>5 and DR_SCORE:\n                e = score_energy_one(seq, target_id, out_cat)\n            elif FULL_ENERGY:\n                e = score_energy_one_full(seq, target_id, out_cat)\n            else:\n                e = score_energy_one_simple(seq, target_id, out_cat)\n            energy.append(e) #tranucated sequence\n            \n            if L != len(seq):\n                out_cat['coor'] = np.pad(out_cat['coor'], ((0, L - len(seq)), (0, 0), (0, 0)), 'constant', constant_values=0)\n                \n            outputs.append(out_cat)\n            \n            \n            xyz = frame_coor_to_C1(out_cat['coor'], sequence, BASE_COOR, OTHER_COOR)\n            \n            coordinate.append(xyz)\n            \n\n            time_taken = timer() - start_timer\n            total_time_taken += time_taken\n            #print('time_taken:', time_to_str(time_taken, mode='sec'))\n\n            gpu_mem_used = gpu_memory_use()\n            max_gpu_mem_used = max(max_gpu_mem_used,gpu_mem_used)\n            #print('gpu_mem_used:', gpu_mem_used, 'GB')\n\n            print(f'{c:02d}   energy:{e:10.0f}   out_cat{str(out_cat[\"coor\"].shape)}  time:{time_to_str(time_taken, mode=\"sec\")}   gpu={gpu_mem_used} gb')\n\n            \n        #------- \n        torch.cuda.empty_cache()\n        \n        if GET_CENTER:\n            energy = np.array(energy)\n            energy_mean = np.mean(energy)\n            energy= np.abs(energy - energy_mean)\n        #select top5\n        argsort = np.array(energy).argsort()\n        argsort = argsort[:5]\n        \n        if L <= OPTIM_LENGTH:\n            out_opt= outputs[argsort[0]]\n            out_opt['coor'] = optimize_coor(seq, target_id, out_opt)\n            coordinate[argsort[0]] = frame_coor_to_C1(out_opt['coor'], sequence, BASE_COOR, OTHER_COOR)\n            torch.cuda.empty_cache()\n            \n        df = coord_to_df(row.sequence, [coordinate[k] for k in argsort], row.target_id)\n        submit_df.append(df)\n    \n    print('----------------------------------------')\n    print('MAX_LENGTH', MAX_LENGTH)\n    print('### total_time_taken:', time_to_str(total_time_taken, mode='min'))\n    print('### max_gpu_mem_used:', max_gpu_mem_used, 'GB')\n    print('')\n\n    submit_df = pd.concat(submit_df)\n    submit_df.to_csv(f'submission.csv', index=False)\n    print(submit_df)\n    return submit_df\n\nrun_submit(valid_df)\n\nprint('SUBMIT OK!!!')","metadata":{"execution":{"iopub.status.busy":"2026-03-11T08:15:37.979687Z","iopub.execute_input":"2026-03-11T08:15:37.980701Z"},"papermill":{"duration":3704.82581,"end_time":"2026-03-10T22:50:31.027674","exception":false,"start_time":"2026-03-10T21:48:46.201864","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null}]}