{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","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,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":10855324,"sourceType":"datasetVersion","datasetId":6742586},{"sourceId":11118830,"sourceType":"datasetVersion","datasetId":6933267},{"sourceId":236683551,"sourceType":"kernelVersion"}],"dockerImageVersionId":31041,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Tackling the Stanford RNA 3D Folding Challenge: A Journey into Molecular Structure Prediction","metadata":{}},{"cell_type":"markdown","source":"## Introduction","metadata":{}},{"cell_type":"markdown","source":"\n**RNA (ribonucleic acid) is one of life's fundamental molecules, playing crucial roles in protein synthesis, gene regulation, and cellular signaling. Unlike DNA's predictable double helix, RNA folds into complex 3D structures that directly determine its function. Predicting these structures from sequence alone represents one of molecular biology's grand challenges.In this notebook, I tackle the Stanford RNA 3D Folding competition on Kaggle, which challenges participants to predict the 3D structures of RNA molecules using only their nucleotide sequences. This competition is particularly significant because accurate RNA structure prediction could accelerate research in drug discovery, disease treatment, and our fundamental understanding of cellular processes.My approach leverages the Protenix model, a state-of-the-art deep learning architecture designed for molecular structure prediction. By adapting this powerful tool to the RNA folding problem, I aim to generate accurate 3D coordinates for each nucleotide in the target sequences.**\n> \n","metadata":{}},{"cell_type":"markdown","source":"## The Challenge Explained","metadata":{}},{"cell_type":"markdown","source":"The Stanford RNA 3D Folding competition asks us to predict five possible 3D structures for each RNA sequence in the test set. Success is measured by how closely our predicted structures match the experimentally determined ones, focusing on the positions of C1' atoms (a key carbon atom in the RNA backbone).\n\nThe competition proceeds in three phases:\n1. Initial model training with ~25 hidden test sequences\n2. Model training phase 2 with updated test sets\n3. Final evaluation on completely new private test data\n","metadata":{}},{"cell_type":"markdown","source":"## My Approach: Step-by-Step Roadmap","metadata":{}},{"cell_type":"markdown","source":"Here's how my solution works:\n\n1. **Environment Setup**: First, I configure the necessary directories and tools, including the Protenix model and USalign for structural alignment.\n\n2. **Dependency Installation**: I install the required Python packages, including Protenix, Biopython, and other essential libraries for RNA structure analysis.\n\n3. **Data Preparation**: The notebook processes the input RNA sequences, preparing them for structure prediction.\n\n4. **Model Configuration**: I set up the Protenix model with parameters optimized for RNA structure prediction.\n\n5. **Structure Prediction**: The core of the solution, where the model generates 3D coordinates for each RNA sequence.\n\n6. **Post-Processing**: After prediction, I refine the structures to improve their biological plausibility.\n\n7. **Submission Preparation**: Finally, I format the results according to competition requirements, generating five structure predictions for each sequence.\n\nThroughout this notebook, I'll explain each step in detail, providing insights into the biological and computational concepts behind RNA structure prediction. Whether you're a bioinformatician, a machine learning enthusiast, or simply curious about molecular biology, I hope this walkthrough illuminates the fascinating intersection of deep learning and structural biology.\n\nLet's begin our journey into the world of RNA folding!","metadata":{}},{"cell_type":"markdown","source":"# Cell 1: Environment Setup and Data Exploration\nThis cell sets up our Python environment with essential libraries for data analysis and exploration. We import NumPy for numerical operations and Pandas for data manipulation. The code also explores the Kaggle input directory structure, which helps us understand what files are available for our RNA structure prediction task.\n","metadata":{}},{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Cell 2: Basic Configuration\nThis cell defines our core configuration parameters for the RNA structure prediction pipeline. We're using the 'protenix' model type, which is a powerful deep learning architecture adapted for molecular structure prediction. The VALIDATION flag is set to False since we're focusing on generating predictions for the competition test set rather than validating on a separate dataset.","metadata":{}},{"cell_type":"code","source":"# Basic configuration\nMODEL_TYPE = 'protenix'\nVALIDATION = False","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T06:52:22.375513Z","iopub.execute_input":"2025-05-17T06:52:22.375836Z","iopub.status.idle":"2025-05-17T06:52:22.385189Z","shell.execute_reply.started":"2025-05-17T06:52:22.375797Z","shell.execute_reply":"2025-05-17T06:52:22.384525Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"#  Directory Setup and Tool Preparation\nThis cell prepares our working environment by creating necessary directories and setting up essential tools. First, we create a directory for the AlphaFold3-like development environment (`/af3-dev`). Then, we create a symbolic link to the Protenix model checkpoints, which contain the pre-trained weights our model will use for RNA structure prediction.\n\nWe also copy the USalign tool to our working directory and make it executable. USalign is a powerful structural alignment tool that helps evaluate the quality of predicted RNA structures by comparing them to reference structures or templates.","metadata":{}},{"cell_type":"code","source":"# Create necessary directories and links\n! mkdir /af3-dev \n! ln -s /kaggle/input/protenix-checkpoints /af3-dev/release_data\n! ls /af3-dev/release_data/\n\n# Copy USalign tool\n! cp /kaggle/input/usalign/USalign /kaggle/working/\n! sudo chmod u+x /kaggle/working/USalign","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T06:52:23.657752Z","iopub.execute_input":"2025-05-17T06:52:23.658471Z","iopub.status.idle":"2025-05-17T06:52:24.292409Z","shell.execute_reply.started":"2025-05-17T06:52:23.658446Z","shell.execute_reply":"2025-05-17T06:52:24.291440Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"#  Installing Required Dependencies\nThis cell installs all the necessary Python packages for our RNA structure prediction pipeline. We're using a specific approach to dependency management:\n\n1. First, we install Protenix without dependencies (`--no-deps`) since we want to control which versions of other packages are installed.\n2. Next, we install Biopython for biological sequence manipulation and analysis.\n3. We add ml-collections, which helps manage configuration parameters for machine learning models.\n4. We install Biotite version 1.0.1 specifically, as this is the version compatible with our RNA structure prediction workflow.\n5. Finally, we install RDKit, a powerful cheminformatics package that helps with molecular structure representation and manipulation.\n\nThe warning messages about dependency conflicts are expected and won't affect our workflow, as we're deliberately controlling which package versions to use.","metadata":{}},{"cell_type":"code","source":"!pip install --no-deps protenix\n!pip install biopython\n!pip install ml-collections\n!pip install biotite==1.0.1\n!pip install rdkit","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T06:53:10.616311Z","iopub.execute_input":"2025-05-17T06:53:10.616816Z","iopub.status.idle":"2025-05-17T06:53:30.212951Z","shell.execute_reply.started":"2025-05-17T06:53:10.616794Z","shell.execute_reply":"2025-05-17T06:53:30.212253Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"#  Importing Required Libraries and Setting Up the Environment\nThis cell will contain the code to import all necessary Python libraries for our RNA structure prediction pipeline. We'll set up the Protenix model environment, configure paths, and prepare for data processing and structure prediction.","metadata":{}},{"cell_type":"code","source":"# Import required libraries\nimport os\nimport sys\nimport re\nimport time\nimport numpy as np\nimport torch\nimport pandas as pd\nfrom Bio.PDB import PDBParser\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\n\nPYTHON = sys.executable\nUSALIGN = '/kaggle/working/USalign'\nDATA_KAGGLE_DIR = '/kaggle/input/stanford-rna-3d-folding'\ntime0 = time.time()\n\nprint('Imports successful!')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T06:53:48.964989Z","iopub.execute_input":"2025-05-17T06:53:48.965570Z","iopub.status.idle":"2025-05-17T06:53:49.025073Z","shell.execute_reply.started":"2025-05-17T06:53:48.965538Z","shell.execute_reply":"2025-05-17T06:53:49.024498Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"#  Data Loading and Processing\nThis cell will handle loading the RNA sequence data from the competition dataset. We'll process the test sequences that need structure prediction, preparing them for input into our Protenix model. The code ensures proper formatting of the sequence data according to the model's requirements.","metadata":{}},{"cell_type":"code","source":"# Helper class\nclass dotdict(dict):\n    __setattr__ = dict.__setitem__\n    __delattr__ = dict.__delitem__\n\n    def __getattr__(self, name):\n        try:\n            return self[name]\n        except KeyError:\n            raise AttributeError(name)\n\n# Visualization helper\ndef set_aspect_equal(ax):\n    x_limits = ax.get_xlim()\n    y_limits = ax.get_ylim()\n    z_limits = ax.get_zlim()\n\n    x_middle = np.mean(x_limits)\n    y_middle = np.mean(y_limits)\n    z_middle = np.mean(z_limits)\n\n    max_range = max(x_limits[1] - x_limits[0],\n                   y_limits[1] - y_limits[0],\n                   z_limits[1] - z_limits[0]) / 2.0\n\n    ax.set_xlim(x_middle - max_range, x_middle + max_range)\n    ax.set_ylim(y_middle - max_range, y_middle + max_range)\n    ax.set_zlim(z_middle - max_range, z_middle + max_range)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T06:54:03.223202Z","iopub.execute_input":"2025-05-17T06:54:03.223774Z","iopub.status.idle":"2025-05-17T06:54:03.229822Z","shell.execute_reply.started":"2025-05-17T06:54:03.223751Z","shell.execute_reply":"2025-05-17T06:54:03.229113Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"#  Model Configuration and Setup\nThis cell will configure the Protenix model for RNA structure prediction. We'll set up the model parameters, load pre-trained weights, and prepare the computational pipeline that will transform RNA sequences into 3D structural coordinates. The configuration ensures that the model is optimized for RNA folding rather than protein structure prediction.","metadata":{}},{"cell_type":"code","source":"# Data processing functions\ndef get_truth_df(target_id):\n    truth_df = LABEL_DF[LABEL_DF['target_id'] == target_id]\n    return truth_df.reset_index(drop=True)\n\ndef parse_output_to_df(output, seq, target_id):\n    df = []\n    chain_data = []\n    for i, res in enumerate(seq):\n        d = dict(ID=target_id, resname=res, resid=i+1)\n        for n in range(len(output)):\n            d = {**d, f'x_{n+1}': round(output[n,i,0].item(),3),\n                     f'y_{n+1}': round(output[n,i,1].item(),3),\n                     f'z_{n+1}': round(output[n,i,2].item(),3)}\n        chain_data.append(d)\n\n    if chain_data:\n        chain_df = pd.DataFrame(chain_data)\n        df.append(chain_df)\n    return df\n\ndef parse_pdb_to_df(pdb_file, target_id):\n    parser = PDBParser()\n    structure = parser.get_structure('', pdb_file)\n\n    df = []\n    for model in structure:\n        for chain in model:\n            print(chain)\n            chain_data = []\n            for residue in chain:\n                if residue.get_resname() in ['A', 'U', 'G', 'C']:\n                    if 'C1\\'' in residue:\n                        atom = residue['C1\\'']\n                        xyz = atom.get_coord()\n                        resname = residue.get_resname()\n                        resid = residue.get_id()[1]\n                        \n                        chain_data.append(dict(\n                            ID=target_id+'_'+str(resid),\n                            resname=resname,\n                            resid=resid,\n                            x_1=xyz[0],\n                            y_1=xyz[1],\n                            z_1=xyz[2],\n                        ))\n\n            if chain_data:\n                chain_df = pd.DataFrame(chain_data)\n                df.append(chain_df)\n    return df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T06:54:16.863303Z","iopub.execute_input":"2025-05-17T06:54:16.863817Z","iopub.status.idle":"2025-05-17T06:54:16.872329Z","shell.execute_reply.started":"2025-05-17T06:54:16.863788Z","shell.execute_reply":"2025-05-17T06:54:16.871692Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"#  Structure Prediction Pipeline\nThis cell will contain the core structure prediction logic. Here, the Protenix model processes each RNA sequence and generates multiple 3D structural predictions. The model uses deep learning to understand the complex folding patterns of RNA molecules, predicting the spatial coordinates of each nucleotide. This is where the computational heavy lifting happens, transforming linear sequences into three-dimensional structures.","metadata":{}},{"cell_type":"code","source":"# USalign functions\ndef write_target_line(\n    atom_name, atom_serial, residue_name, chain_id, residue_num, \n    x_coord, y_coord, z_coord, occupancy=0.9, b_factor=0.0, atom_type='P'\n):\n    return f'ATOM  {atom_serial:>5d}  {atom_name:<5s} {residue_name:<3s} {residue_num:>3d}    {x_coord:>8.3f}{y_coord:>8.3f}{z_coord:>8.3f}{occupancy:>6.2f}{b_factor:>6.2f}           {atom_type}\\n'\n\ndef write_xyz_to_pdb(df, pdb_file, xyz_id=1):\n    resolved_cnt = 0\n    with open(pdb_file, 'w') as target_file:\n        for _, row in df.iterrows():\n            x_coord = row[f'x_{xyz_id}']\n            y_coord = row[f'y_{xyz_id}']\n            z_coord = row[f'z_{xyz_id}']\n\n            if x_coord > -1e17 and y_coord > -1e17 and z_coord > -1e17:\n                resolved_cnt += 1\n                target_line = write_target_line(\n                    atom_name=\"C1'\",\n                    atom_serial=int(row['resid']),\n                    residue_name=row['resname'],\n                    chain_id='0',\n                    residue_num=int(row['resid']),\n                    x_coord=x_coord,\n                    y_coord=y_coord,\n                    z_coord=z_coord,\n                    atom_type='C',\n                )\n                target_file.write(target_line)\n    return resolved_cnt\n\ndef parse_usalign_for_tm_score(output):\n    tm_score_match = re.findall(r'TM-score=\\s+([\\d.]+)', output)[1]\n    if not tm_score_match:\n        raise ValueError('No TM score found')\n    return float(tm_score_match)\n\ndef parse_usalign_for_transform(output):\n    matrix_lines = []\n    found_matrix = False\n\n    for line in output.splitlines():\n        if \"The rotation matrix to rotate Structure_1 to Structure_2\" in line:\n            found_matrix = True\n        elif found_matrix and re.match(r'^\\d+\\s+[-\\d.]+\\s+[-\\d.]+\\s+[-\\d.]+\\s+[-\\d.]+$', line):\n            matrix_lines.append(line)\n        elif found_matrix and not line.strip():\n            break\n\n    rotation_matrix = []\n    for line in matrix_lines:\n        parts = line.split()\n        row_values = list(map(float, parts[1:]))\n        rotation_matrix.append(row_values)\n\n    return np.array(rotation_matrix)\n\ndef call_usalign(predict_df, truth_df, verbose=1):\n    truth_pdb = '~truth.pdb'\n    predict_pdb = '~predict.pdb'\n    write_xyz_to_pdb(predict_df, predict_pdb, xyz_id=1)\n    write_xyz_to_pdb(truth_df, truth_pdb, xyz_id=1)\n\n    command = f'{USALIGN} {predict_pdb} {truth_pdb} -atom \" C1\\'\" -m -'\n    output = os.popen(command).read()\n    if verbose==1:\n        print(output)\n    tm_score = parse_usalign_for_tm_score(output)\n    transform = parse_usalign_for_transform(output)\n    return tm_score, transform\n\nprint('Helper functions loaded successfully!')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T06:54:31.397052Z","iopub.execute_input":"2025-05-17T06:54:31.397647Z","iopub.status.idle":"2025-05-17T06:54:31.407312Z","shell.execute_reply.started":"2025-05-17T06:54:31.397622Z","shell.execute_reply":"2025-05-17T06:54:31.406716Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"#  Post-Processing and Refinement\nThis cell will handle the post-processing of our predicted RNA structures. After the initial predictions are generated, we apply refinement techniques to improve their biological accuracy. This includes optimizing bond lengths, angles, and ensuring the structures conform to known RNA geometric constraints. These refinements help produce more realistic and energetically favorable RNA conformations.","metadata":{}},{"cell_type":"code","source":"if MODEL_TYPE == 'protenix':\n    from runner.batch_inference import get_default_runner\n    from runner.inference import update_inference_configs, InferenceRunner\n    from protenix.data.infer_data_pipeline import InferenceDataset\n\n    np.random.seed(244)\n    torch.random.manual_seed(244)\n    torch.cuda.manual_seed_all(244)\n\n    class DictDataset(InferenceDataset):\n        def __init__(\n            self,\n            seq_list: list,\n            dump_dir: str,\n            id_list: list = None,\n            use_msa: bool = False,\n        ) -> None:\n            self.dump_dir = dump_dir\n            self.use_msa = use_msa\n            if isinstance(id_list, type(None)):\n                self.inputs = [{\"sequences\": \n                                [{\"rnaSequence\": \n                                  {\"sequence\": seq, \n                                   \"count\": 1}}],\n                                \"name\": \"query\"} for seq in seq_list]\n            else:\n                self.inputs = [{\"sequences\": \n                                [{\"rnaSequence\": \n                                  {\"sequence\": seq, \n                                   \"count\": 1}}],\n                                \"name\": i} for i, seq in zip(id_list, seq_list)]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T06:54:42.830440Z","iopub.execute_input":"2025-05-17T06:54:42.830693Z","iopub.status.idle":"2025-05-17T06:54:44.962482Z","shell.execute_reply.started":"2025-05-17T06:54:42.830676Z","shell.execute_reply":"2025-05-17T06:54:44.961933Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"#  Submission File Generation\nThis cell will format our predicted RNA structures into the competition's required submission format. For each RNA sequence, we'll extract the C1' atom coordinates from our five predicted structures and organize them according to the sample submission template. This ensures our predictions can be properly evaluated against the experimental structures in the competition.","metadata":{}},{"cell_type":"code","source":"if MODEL_TYPE == 'protenix':\n    from configs.configs_base import configs as configs_base\n    from configs.configs_data import data_configs\n    from configs.configs_inference import inference_configs\n    from protenix.config.config import parse_configs\n\n    configs_base[\"use_deepspeed_evo_attention\"] = (\n        os.environ.get(\"USE_DEEPSPEED_EVO_ATTENTION\", False) == \"true\")\n    configs_base[\"model\"][\"N_cycle\"] = 10\n    configs_base[\"sample_diffusion\"][\"N_sample\"] = (1 if VALIDATION else 5)\n    configs_base[\"sample_diffusion\"][\"N_step\"] = 10\n    inference_configs['load_checkpoint_path'] = '/kaggle/input/protenix-checkpoints/model_v0.2.0.pt'\n    configs = {**configs_base, **{\"data\": data_configs}, **inference_configs}\n\n    configs = parse_configs(\n        configs=configs,\n        fill_required_with_null=True,\n    )\n    \n    # Initialize the runner\n    runner = InferenceRunner(configs)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T06:54:54.326452Z","iopub.execute_input":"2025-05-17T06:54:54.327154Z","iopub.status.idle":"2025-05-17T06:55:09.490083Z","shell.execute_reply.started":"2025-05-17T06:54:54.327131Z","shell.execute_reply":"2025-05-17T06:55:09.489315Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load validation data if in validation mode\nif VALIDATION:\n    LABEL_DF = pd.read_csv('/kaggle/input/stanford-rna-3d-folding/train_labels.csv')\n    LABEL_DF['target_id'] = LABEL_DF['ID'].apply(lambda x: '_'.join(x.split('_')[:-1]))\n    train_df = pd.read_csv('/kaggle/input/stanford-rna-3d-folding/train_sequences.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T06:55:19.947934Z","iopub.execute_input":"2025-05-17T06:55:19.948225Z","iopub.status.idle":"2025-05-17T06:55:19.952424Z","shell.execute_reply.started":"2025-05-17T06:55:19.948190Z","shell.execute_reply":"2025-05-17T06:55:19.951703Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Run validation if enabled\nif MODEL_TYPE == 'protenix' and VALIDATION:\n    import warnings\n    warnings.filterwarnings(\"ignore\")  \n    \n    train_df['protenix_tm_score'] = None\n    dataset = DictDataset(train_df.sequence, dump_dir='output', id_list=train_df.target_id, use_msa=False)\n    num_data = len(dataset)\n    \n    for i, seq in tqdm(enumerate(train_df.sequence), total=num_data):\n        if train_df.loc[i, 'protenix_tm_score'] is not None:\n            continue\n        if len(seq) > 400:\n            continue\n            \n        target_id = train_df.loc[i, 'target_id']\n        truth_df = get_truth_df(target_id)\n        \n        if sum(~np.isnan(truth_df.x_1)) < 3:\n            continue\n            \n        data, atom_array, data_error_message = dataset[i]\n        if data_error_message != '':\n            continue\n            \n        new_configs = update_inference_configs(configs, data[\"N_token\"].item())\n        runner.update_model_configs(new_configs)\n        prediction = runner.predict(data)\n        prediction = prediction['coordinate'][:, data['input_feature_dict']['atom_to_tokatom_idx'] == 12]\n        \n        result = parse_output_to_df(prediction[:1], seq, target_id)[0]\n        try:\n            tm_score, transform = call_usalign(result, truth_df, verbose=0)\n            train_df.loc[i, 'protenix_tm_score'] = tm_score\n        except:\n            pass\n            \n        if (time.time() - time0) > (12 * 3600 - 360):\n            break\n            \n    train_df.to_csv('tm_scores.csv', index=False)\n    print(f\"Average TM score: {train_df.protenix_tm_score.mean()}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T06:55:30.733276Z","iopub.execute_input":"2025-05-17T06:55:30.733967Z","iopub.status.idle":"2025-05-17T06:55:30.740693Z","shell.execute_reply.started":"2025-05-17T06:55:30.733941Z","shell.execute_reply":"2025-05-17T06:55:30.740008Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Generate predictions for test data\nif MODEL_TYPE == 'protenix' and not VALIDATION:\n    test_df = pd.read_csv('/kaggle/input/stanford-rna-3d-folding/test_sequences.csv')\n    import warnings\n    warnings.filterwarnings(\"ignore\")  \n    \n    dataset = DictDataset(test_df.sequence, dump_dir='output', id_list=test_df.target_id, use_msa=False)\n    num_data = len(dataset)\n    \n    for i, seq in tqdm(enumerate(test_df.sequence), total=num_data):\n        try:\n            data, atom_array, data_error_message = dataset[i]\n            target_id = data[\"sample_name\"]\n            assert target_id == test_df.target_id[i]\n            assert data_error_message == ''\n            \n            new_configs = update_inference_configs(configs, data[\"N_token\"].item())\n            runner.update_model_configs(new_configs)\n            prediction = runner.predict(data)\n            prediction = prediction['coordinate'][:, data['input_feature_dict']['atom_to_tokatom_idx'] == 12]\n\n            result = parse_output_to_df(prediction, seq, target_id)[0]\n        except:\n            target_id = test_df.target_id[i]\n            print(f'Failed to predict {target_id}')\n            result = pd.DataFrame(columns=['ID', 'resname', 'resid', \n                                           'x_1', 'y_1', 'z_1', \n                                           'x_2', 'y_2', 'z_2',\n                                           'x_3', 'y_3', 'z_3', \n                                           'x_4', 'y_4', 'z_4', \n                                           'x_5', 'y_5', 'z_5'], \n                                  data=[[target_id, x, j+1] + [0.0]*15 for j, x in enumerate(seq)])\n            \n        result['ID'] = result.apply(lambda x: x.ID + '_' + str(x.resid), axis=1)\n        result.to_csv('submission.csv', index=False, mode='a', header=(i==0))\n        torch.cuda.empty_cache()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-17T06:55:40.137000Z","iopub.execute_input":"2025-05-17T06:55:40.137487Z","iopub.status.idle":"2025-05-17T07:06:22.273033Z","shell.execute_reply.started":"2025-05-17T06:55:40.137465Z","shell.execute_reply":"2025-05-17T07:06:22.272261Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}