{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"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":"gpu","dataSources":[{"sourceId":87793,"databundleVersionId":12276181,"sourceType":"competition"},{"sourceId":11661237,"sourceType":"datasetVersion","datasetId":7317710},{"sourceId":11759715,"sourceType":"datasetVersion","datasetId":7382011},{"sourceId":11775973,"sourceType":"datasetVersion","datasetId":7393299},{"sourceId":11837219,"sourceType":"datasetVersion","datasetId":7436926},{"sourceId":224830487,"sourceType":"kernelVersion"}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# RNA 3D Structure Prediction Project\n# 📋 Project Overview\nThis project predicts the 3D structure of RNA molecules using deep learning models. Given an RNA sequence, the system generates atomic coordinates that represent how the RNA folds in 3D space.\n\n# 🧬 What is RNA Folding?\nRNA (Ribonucleic Acid) molecules don't just exist as straight chains - they fold into complex 3D shapes that determine their function. Predicting this 3D structure from the sequence alone is a challenging computational problem.\n# 🛠️ Technical Approach\n# 1. Model Architecture\nWe use two main neural network models working together:\n\n# RNA Language Model (RNA2nd)\nPurpose: Understands RNA sequence patterns and relationships\n\nInput: RNA sequence (A, G, C, U nucleotides)\n\nOutput: Learned representations of the sequence\n\n# Key Features:\n\n512-dimensional hidden states\n\n18 encoder layers\n\nProcesses sequences up to 2400 nucleotides long\n# Structure Prediction Model (MSA2XYZ)\nPurpose: Converts sequence information into 3D coordinates\n\nInput: Processed sequence from RNA language model\n\nOutput: 3D atomic coordinates for each nucleotide\n\nKey Features:\n\n8 prediction cycles for refinement\n\nEnsemble of multiple trained models (up to 20 different versions)\n\nOutputs coordinates for key atoms (P, C4', N1/N9)\n\n\n\n\n","metadata":{}},{"cell_type":"markdown","source":"# 2. Coordinate System\nThe model predicts local coordinate frames for each nucleotide, then converts them to global 3D positions:\n\n# Local Frames:\nEach nucleotide has its own coordinate system\n\n# Key Atoms Predicted:\n\nP: Phosphate atom (backbone)\n\nC4': Sugar ring atom\n\nN1/N9: Nitrogen atoms in the base (distinguishes purines/pyrimidines)\n\n","metadata":{}},{"cell_type":"markdown","source":"# 3. Structure Optimization\nFor longer sequences, we use multiple strategies:\n\n# Segmented Prediction\nSequences longer than 480 nucleotides are split into overlapping segments\n\nEach segment is predicted separately\n\nSegments are stitched together using mathematical transformations\n\n# Energy-Based Scoring\n\nWe evaluate predicted structures using physical energy functions:\n\n# Bond Energy\nChecks if atom distances are realistic\n# Angle Energy:\nEnsures proper bond angles\n\n# Stacking Energy:\nEncourages proper base stacking\n\n# Hydrogen Bonds:\nModels base pairing interactions\n\n","metadata":{}},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"markdown","source":"# 🧩 1. Importing Libraries\n\nBefore running the model, some Python tools are imported.\n\ntorch is the main deep learning library used to run the model.\n\nsys helps handle system arguments, such as command-line inputs.\n\nos allows control over the system environment, including GPU selection and file management.\n\n# ⚙️ 2. Choosing Which GPU to Use\n\nThe setup tells the system to use GPU number “0”.\nIf there are several GPUs, this number can be changed to choose another one.\nOn Kaggle, there is usually only one GPU, so “0” is correct.\nThis step ensures the model runs on the GPU instead of the CPU for faster speed.\n\n💻 3. Setting Device and Precision Type\n\nThe script checks if a GPU is available.\nIf it is, it sets the device to CUDA, which means it will use the GPU.\nIf not, it will use the CPU instead.\n\nIt also chooses the precision type, which decides how detailed the numbers are when the model runs.\n\nfp32 means 32-bit floating-point (standard accuracy).\n\nfp16 means 16-bit (faster but slightly less accurate).\n\nbf16 means balanced precision for special GPUs.\n\nThis allows a balance between speed and accuracy depending on the hardware.\n\n","metadata":{}},{"cell_type":"markdown","source":"# 🔧 3. Defining Configuration Settings\n\nSeveral parameters are defined to control how the model runs.\nEach one has a specific purpose:\n\nSetting\tMeaning\tWhy It's Important\n# NUM_CONF\nNumber of 3D structures to generate\tLets the model try multiple folding options (crucial for ensemble selection)\n# MAX_LENGTH\nMaximum allowed single sequence length\tPrevents memory overload and defines chunking size\n# MAX_CAT_LENGTH\nMaximum length after combining sequences\tHandles very long or merged inputs\n# CFG_DIR\nFolder for configuration and checkpoints\tTells the model where to read configuration and weights from\n# CFG_MERGE\nWhether to merge configuration files\tUsually left as False\n# DR_SCORE\t\nEnables an extra deep learning scoring check\tOptional accuracy improvement\n# NO_SORT\nDisables sorting of inputs or predictions\tUsually left as False (meaning sorting, often by energy, is enabled)\n# ET_CENTER\nCenters the 3D structure\tKeeps geometry properly aligned for consistent output\n# FULL_ENERGY\nCalculates full structure energy\tSlower, but gives more detailed results for stability\n# OPTIM_LENGTH\nNumber of optimization steps\tZero means skip the final coordinate refinement optimization","metadata":{}},{"cell_type":"markdown","source":"# 🧰 4. Installing Required Libraries\n# Library,Purpose,Why It's Used\n# rhofold\n,Main RNA folding model,Predicts the 3D RNA structure (likely the core engine).\n# gprofiler_official,\nBiological profiling tool,Helps identify gene-related data and functional context.\n# mygene and biothings_client,\nConnect to biological databases,Retrieve sequence and gene information from public repositories.\n# einops,\nHelps reshape and combine tensors,Makes model computations easier and more readable by simplifying tensor transformations.\n# \"bio, biopython, biopandas\",\nHandle biological data,\"Read, write, and analyze sequence files and biological formats.\"\n# OpenMM,Molecular simulation engine,\nCalculates physical movements and forces for energy minimization and dynamics.\n# \"protenix, rdkit\",\nChemistry and molecule tools,Manage and manipulate atomic structures and chemical properties.\nml_collections,Manage model settings,Simplifies configuration handling and parameter management.\n# \"biotraj, biotite\",\nAnalyze molecular movements,Study RNA dynamics and trajectory data from simulations.\n# blosc,\nCompress data efficiently,Speeds up file reading and writing by applying fast compression.\n# looseversion,\nCompare version numbers,Keeps library versions compatible and avoids unexpected errors.","metadata":{}},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport sys\nimport os\nos.environ['CUDA_VISIBLE_DEVICES'] = '0'\n\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":{"trusted":true},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-18T03:31:51.710568Z","iopub.execute_input":"2025-05-18T03:31:51.711025Z","iopub.status.idle":"2025-05-18T03:31:51.714833Z","shell.execute_reply.started":"2025-05-18T03:31:51.710993Z","shell.execute_reply":"2025-05-18T03:31:51.714189Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Install requirements ","metadata":{}},{"cell_type":"code","source":"#rhofold+\n#!pip install --no-deps '/kaggle/input/rhofold/wheels_py311/python_box-7.3.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl'\n!pip install --no-deps '/kaggle/input/rhofold/wheels_py311/gprofiler_official-1.0.0-py3-none-any.whl'\n!pip install --no-deps '/kaggle/input/rhofold/wheels_py311/mygene-3.2.2-py2.py3-none-any.whl'\n!pip install --no-deps '/kaggle/input/rhofold/wheels_py311/biothings_client-0.4.1-py3-none-any.whl'\n#!pip install --no-deps '/kaggle/input/rhofold/wheels_py311/dm_tree-0.1.9-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl'\n!pip install --no-deps '/kaggle/input/rhofold/wheels_py311/einops-0.8.1-py3-none-any.whl'\n#!pip install --no-deps '/kaggle/input/rhofold/wheels_py311/openmm-7.7.0-py3-none-any.whl'\n!pip install --no-deps '/kaggle/input/rhofold/wheels_py311/bio-1.8.0-py3-none-any.whl'\n!pip install --no-deps '/kaggle/input/rhofold/wheels_py311/absl_py-2.2.2-py3-none-any.whl'\n!pip install --no-deps '/kaggle/input/rhofold/wheels_py311/OpenMM-8.2.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl'\nprint(\"-------------------------------------------------------------------------------\")\nprint(\"-------------------------------------------------------------------------------\")\nprint(\"-------------------------------------------------------------------------------\")\n\n!pip install --no-deps '/kaggle/input/dependencies-tr-pr/protenix-0.4.6-py3-none-any.whl'\n!pip install --no-deps '/kaggle/input/dependencies-tr-pr/rdkit-2024.9.6-cp310-cp310-manylinux_2_28_x86_64.whl'\n!pip install --no-deps '/kaggle/input/dependencies-tr-pr/ml_collections-1.1.0-py3-none-any.whl'\n!pip install --no-deps '/kaggle/input/dependencies-tr-pr/biopython-1.85-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl'\n# !pip install --no-deps '/kaggle/input/dependencies-tr-pr/pyrosetta-2025.13-cp310-cp310-linux_x86_64.whl'\n!pip install --no-deps '/kaggle/input/dependencies-tr-pr/blosc-1.11.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl'\n!pip install --no-deps '/kaggle/input/dependencies-tr-pr/ml_collections-1.1.0-py3-none-any.whl'\n!pip install --no-deps '/kaggle/input/dependencies-tr-pr/biotraj-1.2.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl'\n!pip install --no-deps '/kaggle/input/dependencies-tr-pr/biotite-1.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl'\n!pip install --no-deps '/kaggle/input/dependencies-tr-pr/biopandas-0.5.1-py3-none-any.whl'\n!pip install --no-deps '/kaggle/input/dependencies-tr-pr/looseversion-1.1.2-py3-none-any.whl'\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-18T03:31:51.716513Z","iopub.execute_input":"2025-05-18T03:31:51.716757Z","iopub.status.idle":"2025-05-18T03:32:16.586184Z","shell.execute_reply.started":"2025-05-18T03:31:51.716736Z","shell.execute_reply":"2025-05-18T03:32:16.585187Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python --version","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-18T03:32:16.587985Z","iopub.execute_input":"2025-05-18T03:32:16.588335Z","iopub.status.idle":"2025-05-18T03:32:16.71558Z","shell.execute_reply.started":"2025-05-18T03:32:16.588301Z","shell.execute_reply":"2025-05-18T03:32:16.71459Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"! mkdir /af3-dev \n! ln -s /kaggle/input/protenix-checkpoints /af3-dev/release_data\n! ls /af3-dev/release_data/","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-18T03:32:16.716739Z","iopub.execute_input":"2025-05-18T03:32:16.717076Z","iopub.status.idle":"2025-05-18T03:32:17.071108Z","shell.execute_reply.started":"2025-05-18T03:32:16.717043Z","shell.execute_reply":"2025-05-18T03:32:17.070355Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Drfold2","metadata":{}},{"cell_type":"markdown","source":"🕒 Step 1: Start Log Time\n\nThe code prints the current time (in Singapore time zone).\nThis helps to know when the program started running.\nIt also prints a message saying “PIP INSTALL OK” — meaning all required packages were installed successfully.\n\n📦 Step 2: Import Libraries\n\nSeveral important tools are imported:\n\nos, sys → control the system and files\n\npandas → read and work with CSV data\n\nnumpy → handle numbers and arrays\n\ntorch → use GPU and deep learning models\n\ntimeit → measure how long things take to run\n\nThen, pandas is told to show more columns and keep tables in one line for easier reading.\n\n🧩 Step 3: Helper Class — dotdict\n\nThis small helper lets you use dot notation for dictionaries.\nExample:\nInstead of data['id'], you can write data.id.\nIt just makes the code easier to read.\n\n⏱ Step 4: Helper Function — time_to_str\n\nThis turns seconds into human-readable time:\n\nShows hours and minutes, or\n\nShows minutes and seconds.\n\nUseful for timing how long code takes to run.\n\n💾 Step 5: Helper Function — gpu_memory_use\n\nThis checks if a GPU is available.\nIf yes, it returns how much GPU memory is being used (in GB).\nIf not, it returns 0.\nThis helps you monitor GPU usage.\n\n📊 Step 6: Helper Function — set_aspect_equal\n\nThis helps when drawing 3D plots.\nIt keeps all axes (X, Y, Z) at the same scale so the shape looks correct (not stretched).\n\n🔥 Step 7: Check PyTorch Setup\n\nThe script prints:\n\nThe PyTorch version\n\nThe CUDA version (used for GPU support)\n\nIf both print correctly, it means PyTorch is ready to use the GPU.\n\nThen it prints “IMPORT OK!!!”, confirming everything works.\n\n# ⚙️ Step 8: Set Program Mode\n\nThere are two ways to run:\n\nlocal → for testing on your own computer (uses validation data)\n\nsubmit → for running on Kaggle (uses test data)\n\nHere, the mode is set to “submit”.# ","metadata":{}},{"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/stanford-rna-3d-folding'\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":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-18T03:32:17.072173Z","iopub.execute_input":"2025-05-18T03:32:17.072487Z","iopub.status.idle":"2025-05-18T03:32:17.482353Z","shell.execute_reply.started":"2025-05-18T03:32:17.072454Z","shell.execute_reply":"2025-05-18T03:32:17.481528Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# ⚙️ Function 1: frame_coor_to_C1(coor, seq, BASE_COOR, OTHER_COOR)\n\nThis function finds the position of a specific atom (C1′) for each base in an RNA sequence.\n\n# 🧠 Step-by-step meaning:\n\n1-Convert input to tensors\nTurns coordinate data into PyTorch tensors so GPU can process them fast.\n\n2-Get base and other coordinates\nUses helper data (BASE_COOR and OTHER_COOR) to get basic 3D reference shapes for each nucleotide.\n\n3-Center and normalize\nRemoves bias (average position) to center each base.\n\n3-Create a combined vector (x)\nStores rotation and position info for every base.\n\n3-Use quaternion rotation (quat2b)\nConverts rotation info into actual 3D coordinates.\n\n4-Extract the fifth atom (C1′)\nEach base has multiple atoms — it takes only the C1′ atom, which is used to represent the backbone structure.\n\n5-Return final result\nGives back the C1′ positions as a NumPy array.","metadata":{}},{"cell_type":"markdown","source":"# 🔄 Function 2: concat_coor(out1, out2)\n\nThis function joins two pieces of RNA together smoothly in 3D space.\nEach piece has coordinates (coor) for its atoms.\n\n# 🧠 Step-by-step meaning:\n\n1-Convert to tensors\nTurns the coordinate arrays into double precision tensors for accuracy.\n\n2-Find connection points\nTakes:\n\nThe last base of out1\n\nThe first base of out2\n\n3-Center both frames\nRemoves the mean (average) position from each frame — makes rotation easier.\n\n4-Find rotation matrix (R)\nCalculates how much to rotate out2 so its first base matches the last base of out1.\n\n5-Find translation vector (t)\nCalculates how much to move (shift) out2 so it connects perfectly to out1.\n\n6-Combine rotation + translation\nBuilds a set of transformation parameters (called rama) for the whole chain.\n\n7-Apply transformation\nUses a2b.quat2b() to rotate and move every atom in out2 according to the calculated values.\n\n8-Join both chains\nConnects the two coordinate sets together — skipping the repeated first base of out2.\n\n9-Return the result\nReturns the full, continuous coordinate array as a NumPy matrix.","metadata":{}},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"code","source":"sys.path.append('/kaggle/input/drfold-model-bf16/PotentialFold')\nimport a2b\ndef frame_coor_to_C1(coor, seq, BASE_COOR, OTHER_COOR):\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    x = torch.rand((L, 21), dtype=tx.dtype, device=tx.device)\n\n    biasq = tx.mean(dim=1, keepdim=True)            # (L, 1, 3)\n    q = tx - biasq                                  # (L, 3, 3)\n\n    m = torch.einsum('bnz,bny->bzy', basex, q).reshape(L, -1)\n    x[:, :9] = m\n    x[:, 9:18] = m\n\n    x[:, 18:] = biasq.squeeze(1)\n    rama = x.double()  \n\n    #    otherx: (L, 5, 3) -> other_xyz: (L, 5, 3)\n    other_xyz = a2b.quat2b(otherx.double(), rama[:, 9:]).float().cpu().numpy()\n\n    c1_xyz = other_xyz[:, 4, :]\n    return c1_xyz\n\n\ndef concat_coor(out1: dict, out2: dict) -> np.ndarray:\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    f1 = coor1[-1]   # (3,3)\n    f2 = coor2[0]    # (3,3)\n\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    #    R_{ij} = sum_z basex_{iz} * q_{jz}\n    R = torch.einsum('iz,jz->ij', basex, q)   # (3,3)\n\n    t = bias1 - (R @ bias2)                  # (3,)\n\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    coor2_aligned = a2b.quat2b(coor2, rama)   # torch.Tensor (L2,3,3)\n\n    coor_cat = torch.cat([coor1, coor2_aligned[1:]], dim=0)  # (L1+L2-1,3,3)\n\n    return coor_cat.cpu().numpy()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-18T03:32:17.483104Z","iopub.execute_input":"2025-05-18T03:32:17.48338Z","iopub.status.idle":"2025-05-18T03:32:17.544333Z","shell.execute_reply.started":"2025-05-18T03:32:17.483358Z","shell.execute_reply":"2025-05-18T03:32:17.543733Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 🧬 Function: Get_base\n\nThis function is used to build a 3D template (a shape) for every base in an RNA or DNA sequence.\nIt takes your sequence — like “AGCU” — and gives back a 3D version of it using standard base structures.\n\n# 💡 What This Function Does\n\nImagine you have standard 3D shapes for the four bases:\n\nA = Adenine\n\nG = Guanine\n\nC = Cytosine\n\nU = Uracil (for RNA)\n\nT = Thymine (for DNA, same shape as Uracil)\n\nThe function looks at each letter in your sequence and chooses the matching shape.\nIt then combines all these shapes into one large 3D list, so you can use it for modeling or visualization.\n\n# 🧠 Step-by-Step Explanation\n\nThe function first checks how many atoms are in one base (for example, 10 atoms in Adenine).\nThis helps to know how big each 3D piece is.\n\nIt then creates an empty space to hold all base coordinates for the full sequence.\nThink of it like an empty 3D box waiting to be filled with the atoms of each base.\n\nIt reads your sequence one letter at a time — for example, “A”, “G”, “C”, “U”.\n\nFor each base:\n\nIf the base is “A” (or lowercase “a”), it copies the standard Adenine shape.\n\nIf it’s “G” (or “g”), it copies the Guanine shape.\n\nIf it’s “C” (or “c”), it copies the Cytosine shape.\n\nIf it’s “U” or “u”, it copies the Uracil shape.\n\nIf it’s “T” or “t”, it also uses the Uracil shape (because Thymine and Uracil are similar).\n\nAfter going through every base, it returns the complete 3D structure —\none for each letter of your sequence.# ","metadata":{}},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-18T03:32:17.546425Z","iopub.execute_input":"2025-05-18T03:32:17.546638Z","iopub.status.idle":"2025-05-18T03:32:17.551239Z","shell.execute_reply.started":"2025-05-18T03:32:17.54662Z","shell.execute_reply":"2025-05-18T03:32:17.550518Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# ⚗️ Function: score_energy_one_simple\n\nThis function estimates the stability energy (also called “free energy”) of an RNA 3D structure.\nIt checks how good or bad the shape of the RNA is based on distances and angles between atoms.\n\n# 💡 What This Function Does\n\nIt looks at the 3D coordinates of each RNA base (the atoms P, C4′, and N)\nand calculates an energy score.\n\nA lower energy value means the structure is more stable,\nand a higher value means it’s less stable or physically unrealistic.\n# 📦 What the Inputs Mean\n\nseq → the RNA sequence (like “AGCU”).\n\ntarget_id → the structure’s name or label (used only for logging).\n\nout → a dictionary that contains coordinates of each base:\n\ncoor[i,0] = the P atom position\n\ncoor[i,1] = the C4′ atom position\n\ncoor[i,2] = the N atom position (N1 or N9 depending on base type)\n# 🧠 Step-by-Step Explanation\n\nThe function calculates three types of energy:\n\n1. 🔗 Bond Length Energy\n\nChecks how far apart connected atoms are.\n\nFor each base, it measures:\n\nDistance between P and C4′\n\nDistance between C4′ and N\n\nEach has an ideal distance (like 1.6 and 1.47 angstroms).\n\nIf the actual distance is too long or too short, the energy increases.\n\n👉 This part ensures each base has correct internal bond lengths.\n\n2. 📐 Bond Angle Energy\n\nLooks at the angle between P, C4′, and N atoms.\n\nThe ideal angle is about 109.5°, which is typical for tetrahedral carbon bonds.\n\nThe energy increases when this angle bends too much or too little.\n\n👉 This keeps the base geometry realistic and prevents distorted shapes.\n\n3. 🧱 Stacking Interaction Energy\n\nChecks the distance between neighboring bases.\n\nMeasures how close the C4′ atoms of two consecutive bases are.\n\nIdeal distance is around 3.4 angstroms, similar to spacing in stacked RNA or DNA bases.\n\nIf the distance is too large or too small, it adds a penalty.\n\n👉 This encourages bases to stay neatly stacked like in real RNA helices.\n\n4. ⚖️ Total Energy\n\nAdds up all three parts:\n\nBond energy\n\nAngle energy\n\nStacking energy\n\nThe result is a single number called total_energy.\n\n👉 The lower this number, the more stable and realistic the RNA shape is.","metadata":{}},{"cell_type":"code","source":"import numpy as np\n\ndef score_energy_one_simple(seq, target_id, out):\n\n    coor = out['coor']  # (L, 3, 3)\n    L = len(seq)\n\n    d0_P_C4  = 1.60\n    d0_C4_N  = 1.47\n    k_bond   = 100.0 \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    theta0   = np.deg2rad(109.5)\n    k_angle  = 20.0   \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    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    total_energy = energy_bond + energy_angle + energy_stack\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":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-18T03:32:17.55231Z","iopub.execute_input":"2025-05-18T03:32:17.552498Z","iopub.status.idle":"2025-05-18T03:32:17.566658Z","shell.execute_reply.started":"2025-05-18T03:32:17.552481Z","shell.execute_reply":"2025-05-18T03:32:17.565864Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# ⚗️ Function: score_energy_one_full\n\nThis function calculates a detailed stability score (free energy) for an RNA 3D structure.\nIt’s like checking how “realistic” or “stable” the shape of the RNA is,\nbased on the positions of its main atoms (P, C4′, N).\n\nCompared to the simpler version, this one includes more physics-based effects,\nso it gives a more accurate idea of how the RNA behaves in real life.\n\n# 💡 What the Function Does\n\nIt takes an RNA’s 3D coordinates and calculates its total energy,\nwhich is made up of several smaller energies (bond, angle, stacking, etc.).\nThe smaller the total energy, the more stable the RNA shape.\n📦 What You Give to It\n\nseq → The RNA sequence, like “AGCUUGA”.\n\ntarget_id → The structure’s name or ID (just for labeling).\n\nout → A dictionary that must have:\n\nout['coor']: an array of shape (L, 3, 3), where each base has:\n\nP atom position\n\nC4′ atom position\n\nN atom position (N1 or N9)\n\noptionally: out['paired']: a list of base pairs like [(0,10), (1,9)].\n\nYou can also give paired separately (a list of base pairs).\n\nOther parameters (like d0_P_C4, k_angle, etc.) are constants that control how strong each energy term is.\n\n# ⚙️ What It Calculates (Step by Step)\n\nThe total energy comes from seven parts:\nEach part checks a different aspect of RNA structure.\n\n1. 🔗 Bond Length Energy\n\nEach RNA base has connections between its atoms:\n\nP → C4′\n\nC4′ → N\n\nThe function measures these distances and compares them to their ideal lengths\n(1.60 Å and 1.47 Å).\n\nIf the distances are too long or too short, energy increases.\n\n👉 Keeps each base atom connected properly and avoids “stretched” atoms.\n\n2. 📐 Bond Angle Energy\n\nChecks the angle between the three atoms:\nP – C4′ – N.\n\nThe best angle is around 109.5°, like in real molecular geometry.\nIf this angle bends too far, it costs energy.\n\n👉 Keeps the shape of each base realistic.\n\n3. 🧱 Base Stacking Energy\n\nRNA bases stack one on top of another.\nThis part looks at how far apart neighboring bases are —\nby checking the distance between the C4′ atoms of two consecutive bases.\n\nIdeal stacking distance ≈ 3.4 Å.\n\nIf they’re too far or too close, energy rises.\n\n👉 Encourages bases to stay neatly stacked like in helices.\n\n4. 🔄 Dihedral (Torsion) Energy\n\nThis measures how much the backbone twists between bases.\n\nIt uses four points (atoms from two bases) to find the torsion angle,\nwhich describes how the RNA chain rotates.\n\nThe best twist angle is 180° (flat).\nIf it twists too much, energy increases.\n\n👉 Controls the smoothness and rotation of the RNA chain.\n\n5. 🧬 Hydrogen Bond Energy\n\nHydrogen bonds hold base pairs together (like A–U or G–C).\n\nIf you give the function a list of paired bases (paired),\nit checks how far apart their N atoms are.\n\nThe ideal hydrogen bond distance is about 2.9 Å.\n\nToo close or too far = higher energy.\n\n👉 Encourages correct base pairing (A–U, G–C).\n\n6. ⚛️ Van der Waals (Lennard–Jones) Energy\n\nThis models how atoms attract or repel each other at short distances.\n\nFor every pair of bases that are not neighbors, it measures how close their C4′ atoms are.\nIf they’re too close, they repel; if they’re at a good distance, they attract slightly.\n\n👉 Prevents atoms from overlapping and adds realistic molecular spacing.\n\n7. ⚡ Electrostatic (Debye–Hückel) Energy\n\nRNA’s phosphate (P) atoms have negative charges.\n\nThis part estimates how these charges interact —\nthey repel each other, but the effect is weaker in water (because of screening).\n\nIt uses the Debye–Hückel equation,\nwhich decreases the effect of charge repulsion as distance increases.\n\n👉 Simulates electrical forces between phosphate groups.# ","metadata":{}},{"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    coor = out['coor']\n    L = len(seq)\n    if paired is None:\n        paired = out.get('paired', [])\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    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    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    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    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    E_LJ = 0.0\n    for i in range(L):\n        for j in range(i+2, L): \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    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    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":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# ⚙️ Functions: score_energy_one & optimize_coor\n\n(Used for energy scoring and structure optimization in RNA folding)\n\nThese two functions work together to handle realistic RNA structure evaluation and improvement.\nThey help the model check how stable a predicted RNA shape is — and even make it better.","metadata":{}},{"cell_type":"markdown","source":"# 🧮 1. score_energy_one — “Check how stable the RNA is”\n# 🔍 What this function does\n\nThis function measures how much energy (or instability) is inside a predicted RNA structure.\nIt uses a ready-made structure analysis tool called DRfold’s Structure class, which can compute accurate energy scores.\n\nYou can think of it like this:\n\n“Take the RNA, check how well it’s folded, and give me a number —\nlower number = more stable shape.”\n\n🧩 What it uses\n\nIt needs:\n\nseq → The RNA sequence, for example: AGCUUGA.\n\ntarget_id → The RNA name or ID (like a label).\n\nout → The predicted structure data (includes coordinates).\n\n⚙️ What it does step by step\n\nCreates a temporary folder —\nIt works inside a short-lived directory so nothing stays on disk afterward.\n\nWrites the sequence to a small FASTA file —\nThis file holds the RNA letters (A, G, C, U) for the model to read.\n\nWrites the structure data (out) into a temporary file —\nThe coordinates and predictions are saved as a small .ret (pickle) file.\n\nLoads a “folding configuration” file —\nThis file (cfg_for_folding.json) tells DRfold how to measure the structure’s energy\n(which formulas, constants, and physics rules to use).\n\nCreates a Structure object —\nThe Structure class reads the RNA and uses the configuration file to prepare energy calculations.\n\nInitializes the internal “rotation data” (rama) —\nThis describes how each base in the RNA chain is oriented in space.\n\nCalculates the total energy —\nThe obj_func_np() method returns the energy value — a single number that shows how stable the RNA structure is.\n\nReturns the energy value —\nThe smaller this number, the more natural and stable the RNA shape.","metadata":{}},{"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/drfold-model/DRfold2/PotentialFold')\nfrom Optimization import Structure\n\ndef score_energy_one(seq, target_id, out):\n    with tempfile.TemporaryDirectory() as tmpdirname:\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        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/drfold-model-bf16/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        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        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        retfile = os.path.join(tmpdirname, 'tmp.ret')\n        with open(retfile, 'wb') as f:\n            pickle.dump(out, f)\n\n        # foldconfig = '/kaggle/input/drfold-model/DRfold2/cfg_for_folding.json'\n        foldconfig = '/kaggle/input/drfold-model-bf16/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        pdb_file = save_prefix + '.pdb'\n        parser = PDBParser(QUIET=True)\n        structure = parser.get_structure(target_id, pdb_file)\n\n        residues = [\n            res for res in structure.get_residues()\n            if res.id[0] == ' '\n        ]\n        L = len(seq)\n        if len(residues) != L:\n            raise ValueError(f\"PDB 中残基数 ({len(residues)}) 与序列长度 ({L}) 不一致\")\n\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  \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":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-18T03:32:17.567718Z","iopub.execute_input":"2025-05-18T03:32:17.56801Z","iopub.status.idle":"2025-05-18T03:32:17.830473Z","shell.execute_reply.started":"2025-05-18T03:32:17.567982Z","shell.execute_reply":"2025-05-18T03:32:17.829696Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 🧬 RNA Structure Prediction Setup — (Final Part)\n\nThis section prepares everything needed before the RNA folding model runs.\nIt loads required data files, defines helper functions, and organizes how the model will output predictions.","metadata":{}},{"cell_type":"markdown","source":"# 🧩 1. make_data(seq, device)\n\nThis function prepares the input for the RNA model.\nIt changes the RNA letters (A, U, G, C) into numbers, makes coordinates for each base,\nand turns them into tensors so the model can read them.\nIt returns:\n\nmsa: sequence in one-hot form\n\nbase_x: base coordinates\n\nseq_idx: index numbers for each base\n\n# 📄 2. solution_to_submit_df(solution)\n\nThis makes one big table from all predicted RNA structures.\nIt loops through each structure and calls coord_to_df() to format it.\nFinally, it joins all of them into one combined table for output or saving.\n\n# 📐 3. coord_to_df(sequence, coord, target_id)\n\nThis turns one RNA prediction into a clean table.\nEach row is one base and has:\n\nbase name (A, U, G, C)\n\nresidue number\n\n3D coordinates (x, y, z for 3 atoms)\n\nIt’s used to make data easy to view or export.\n\n# 💾 4. out_dir\n\nCreates a folder named model-output inside /kaggle/working/.\nThis is where all results and prediction files will be saved.","metadata":{}},{"cell_type":"code","source":"import sys\nimport os\nimport numpy as np\nimport sys\nsys.path.append('/kaggle/input/drfold-model-bf16')\nsys.path.append('/kaggle/input/drfold-model-bf16/PotentialFold')\nsys.path.append(f'/kaggle/input/drfold-model-bf16/{CFG_DIR}')\nsys.path.append(f'/kaggle/input/drfold-model-bf16/{CFG_DIR}/RNALM2')\n\nBASE_COOR = np.load('/kaggle/input/drfold-model-bf16/PotentialFold/lib/base.npy')\nOTHER_COOR = np.load('/kaggle/input/drfold-model-bf16/PotentialFold/lib/other2.npy')\nSIDE_COOR = np.load('/kaggle/input/drfold-model-bf16/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":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-18T03:32:17.831311Z","iopub.execute_input":"2025-05-18T03:32:17.831647Z","iopub.status.idle":"2025-05-18T03:32:20.817004Z","shell.execute_reply.started":"2025-05-18T03:32:17.831624Z","shell.execute_reply":"2025-05-18T03:32:20.816067Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 🧩 1. Load the RNA Language Model (RNA2nd)\n\nThe function first loads a pre-trained RNA model called RNA2nd.\n\nIt uses a saved file from /kaggle/input/drfold-model-bf16/model_hub/...\n\nThe model helps understand the RNA sequence — kind of like how language models understand text.\n\nThe model is moved to the right computing device (CPU or GPU).\n\nDepending on precision settings (fp16 or bf16), it uses lower memory and faster math.\n# ⚙️ 2. Load the 3D Prediction Model (MSA2XYZ)\n\nThis model takes the RNA sequence and predicts its 3D shape (coordinates).\n\nSeveral model files are loaded (up to 20 or more) to create an ensemble — meaning many models predict, and the best ones are chosen.\n\nEach model’s settings are loaded into memory.\n\nThe model is also moved to GPU and set to evaluation mode (not training).\n# 🔗 3. Connect Both Models\n\nThe 3D model (msa2xyz) uses the RNA language model (rnalm) inside it.\n\nThis connection helps it understand RNA patterns better before predicting coordinates.\n# 🧬 4. Loop Through All RNA Sequences\n\nThe function loops through each row in valid_df.\nEach row has a target ID and a sequence (the RNA letters).\n\nFor every RNA sequence:\n\nIt gets the sequence length (L).\n\nIf it’s too long, it cuts it to a maximum allowed size.\n\nBased on the length, it picks a list of model numbers to try (model_to_try).\n# 🔮 5. Predict 3D Structures\n\nInside the loop:\n\nFor each selected model, it runs a helper function called predict_segment():\n\nConverts sequence into tensors (make_data)\n\nFeeds it into the model\n\nGets predicted coordinates\n\nIf the RNA is too long, it is split into small pieces and predicted part by part.\n\nThe pieces are joined back together (concat_coor).\n# ⚖️ 6. Compute Energy (Stability Check)\n\nFor every predicted structure, the function calculates its energy score.\n\nLower energy = more stable structure.\n\nIt uses one of three energy functions:\n\nscore_energy_one_simple → quick version\n\nscore_energy_one_full → detailed version\n\nscore_energy_one → official version with folding model\n\nThese help decide which predicted model is best.\n\n# 🧱 7. Fix Sequence Length (Padding)\n\nIf the predicted structure is shorter than the original RNA sequence,\nit adds zeros to make lengths match.\n# 📐 8. Convert Frames to C1′ Atom Coordinates\n\nConverts the predicted frame coordinates into real atom positions using frame_coor_to_C1.\nThis step turns model data into real 3D atomic points.\n\n# ⚙️ 9. Optional Optimization\n\nIf the RNA is short enough, the best structure is refined again using optimize_coor().\nThis makes the 3D structure smoother and more natural.\n\n# 📊#  10. Save All Predictions\n\nThe best 5 predictions (lowest energy) are kept.\n\nTheir coordinates are turned into a DataFrame using coord_to_df().\n\nAll DataFrames are combined into one final table.\n# 💾 11. Export Final Submission\n\nAll results are merged and saved into a CSV file called submission.csv.\n\nThe CSV contains:\n\nTarget ID\n\nResidue name (A, U, G, C)\n\nAtom coordinates (x, y, z)\n\nThis is the file you submit to Kaggle.","metadata":{}},{"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/drfold-model-bf16/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/drfold-model-bf16/model_hub/{CFG_DIR}/model_{i}' for i in range(20)]\n    if CFG_MERGE:\n        msa2xyz_file = [\n            f'/kaggle/input/drfold-model-bf16/model_hub/cfg_97/model_{i}'\n            for i in range(20)\n        ] + [\n            f'/kaggle/input/drfold-model-bf16/model_hub/cfg_95/model_{i}'\n            for i in range(20)\n        ] + [\n            f'/kaggle/input/drfold-model-bf16/model_hub/cfg_96/model_{i}'\n            for i in range(20)\n        ] + [\n            f'/kaggle/input/drfold-model-bf16/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":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-18T03:32:20.818027Z","iopub.execute_input":"2025-05-18T03:32:20.818539Z","execution_failed":"2025-05-18T04:22:52.299Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"display(pd.read_csv('submission.csv'))","metadata":{"trusted":true,"execution":{"execution_failed":"2025-05-18T04:22:52.3Z"}},"outputs":[],"execution_count":null}]}