{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"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":10880374,"sourceType":"datasetVersion","datasetId":6760482},{"sourceId":10880419,"sourceType":"datasetVersion","datasetId":6760509},{"sourceId":11151583,"sourceType":"datasetVersion","datasetId":6957452},{"sourceId":11230242,"sourceType":"datasetVersion","datasetId":7014687},{"sourceId":11451236,"sourceType":"datasetVersion","datasetId":7174725},{"sourceId":11452308,"sourceType":"datasetVersion","datasetId":7162026},{"sourceId":11899194,"sourceType":"datasetVersion","datasetId":7479946}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Prepare environment","metadata":{}},{"cell_type":"code","source":"!cp -r /kaggle/input/protenix-packages/packages /kaggle/working\n%cd /kaggle/working/packages\n!pip install --no-deps --exists-action=i *.whl\n%cd /kaggle/working\n\n!mv /kaggle/working/packages/ihm-2.3/ihm-2.3 /kaggle/working\n!mv /kaggle/working/packages/modelcif-0.7/modelcif-0.7 /kaggle/working\n\n!pip install /kaggle/working/ihm-2.3\n!pip install /kaggle/working/modelcif-0.7\n\n!rm -rf /kaggle/working/ihm-2.3\n!rm -rf /kaggle/working/modelcif-0.7\n\n!pip install /kaggle/input/biopython/biopython-1.85-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n!pip install /kaggle/input/ml-collections/ml_collections-1.0.0-py3-none-any.whl\n\n!rm -rf /kaggle/working/packages","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-06T11:24:40.575729Z","iopub.execute_input":"2025-10-06T11:24:40.576138Z","iopub.status.idle":"2025-10-06T11:25:13.489787Z","shell.execute_reply.started":"2025-10-06T11:24:40.576104Z","shell.execute_reply":"2025-10-06T11:25:13.488606Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#!pip install --dry-run modelgenerator\n!cp -r /kaggle/input/protenix-mg-packages/protenix_mg_packages /kaggle/working\n%cd /kaggle/working/protenix_mg_packages\n\n!pip install --no-deps --exists-action=i *.whl\n\n%cd /kaggle/working\n\n!rm -rf /kaggle/working/protenix_mg_packages","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-06T11:25:13.491066Z","iopub.execute_input":"2025-10-06T11:25:13.491315Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import modelgenerator","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!cp -R /kaggle/input/protenix-rmsa-repo/protenix_kaggle /kaggle/working/\n!mv protenix_kaggle Protenix\n%cd /kaggle/working/Protenix","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Inference","metadata":{}},{"cell_type":"code","source":"import os\nimport sys\nimport json\nimport time\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nimport warnings\nfrom biotite.structure.io.pdbx import CIFFile, get_structure\n\nwarnings.simplefilter(\"ignore\", category=FutureWarning)\nwarnings.filterwarnings(\"ignore\", category=UserWarning, module=\"torch.utils.data.dataloader\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"input_path = Path(\"/kaggle/input/stanford-rna-3d-folding\")\nprotenix_path = Path(\"/kaggle/working/Protenix\")\noutput_path = Path(\"/kaggle/working/output\")\noutput_path.mkdir(parents=True, exist_ok=True)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/test_sequences.csv\")\ndf.head()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def prepare_data(df, output_path):\n    \"\"\"\n    convert df to json files\n    \"\"\"\n    for _, row in df.iterrows():\n        target_id = row[\"target_id\"]\n        sequence = row[\"sequence\"]\n        input_json= [(\n            {\n                \"sequences\": [\n                    {\"rnaSequence\": {\"sequence\": sequence, \"count\": 1, \n                                    \"msa\": {\"precomputed_msa_dir\": f\"{input_path}/MSA/{target_id}.MSA.fasta\",\n                                    \"pairing_db\": \"rnacentral\"}\n                                    }\n                    }\n                ],\n                \"name\": target_id,\n            }\n        )]\n        # save json file\n        json_path = output_path / \"input_json\" / f\"{target_id}.json\"\n        json_path.parent.mkdir(parents=True, exist_ok=True)\n        with open(json_path, \"w\") as f:\n            json.dump(input_json, f, indent=4)\n        print(f\"Input JSON file saved to {json_path}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"prepare_data(df, output_path)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"cat /kaggle/working/output/input_json/R1190.json","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%cd $protenix_path\n\ndef inference(target_id, output_path, seed=101, n_cycle=10, n_sample=5, n_step=200):\n\n    checkpoint_path = \"/kaggle/input/protenix-finetuned-rna3db-all-1599/1599_ema_0.999.pt\"\n    \n    input_json_path = output_path / \"input_json\" / f\"{target_id}.json\"\n    dump_dir = output_path / target_id\n    dump_dir.mkdir(parents=True, exist_ok=True)\n    \n    sys.argv = [\n        f\"runner/inference.py\",\n        f\"--seeds={seed}\",\n        f\"--dump_dir={dump_dir}\",\n        f\"--input_json_path={input_json_path}\",\n        f\"--model.N_cycle={n_cycle}\",\n        f\"--sample_diffusion.N_sample={n_sample}\",\n        f\"--sample_diffusion.N_step={n_step}\",\n        f\"--augment.use_rnalm True\",\n        f\"--use_msa True\",\n        f\"--load_checkpoint_path={checkpoint_path}\",\n        f\"\"\n    ]\n\n    from runner.inference import run\n\n    run()\n  ","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for i, row in df.iterrows():\n    print(f\"Running Protenix for {i + 1}/{len(df)}: {row['target_id']}\")\n    target_id, seq = row[\"target_id\"], row[\"sequence\"]\n    print(target_id, seq)\n    start = time.time()\n    inference(target_id, output_path, seed=101, n_cycle=10, n_sample=5, n_step=200)\n    end = time.time()\n    print(f\"Time used: {((end-start)/60):2f} minutes.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def extract_c1_atoms(cif_path) -> pd.DataFrame:\n    cif_file = CIFFile.read(cif_path)\n    model = get_structure(cif_file, model=1)\n    if len(np.unique(model.chain_id)) > 1:\n        print(f\"Found multiple chains in {cif_path}: {len(np.unique(model.chain_id))}\")\n    # select the first chain # TODO: multi-chain cases\n    chain = model[model.chain_id == \"A\"]\n    mask = chain.atom_name == \"C1'\"\n    c1_atoms = chain[mask]\n    df = pd.DataFrame.from_dict(c1_atoms._annot)\n    df[\"x\"] = c1_atoms.coord[:, 0]\n    df[\"y\"] = c1_atoms.coord[:, 1]\n    df[\"z\"] = c1_atoms.coord[:, 2]\n    df = df[[\"res_name\", \"res_id\", \"x\", \"y\", \"z\"]]\n    return df\n\ndef format_submission(output_path, df, seed):\n    # create submission\n    subm = []\n    for _, row in df.iterrows():\n        target_id = row[\"target_id\"]\n\n        all_preds = []\n        for i in range(5):\n            cif_path = output_path/f\"{target_id}/{target_id}/seed_{seed}/predictions/{target_id}_seed_{seed}_sample_{i}.cif\"\n            if not Path(cif_path).exists():\n                continue  \n            pred_df = extract_c1_atoms(cif_path)\n            all_preds.append(pred_df)\n        \n        res = pd.DataFrame({\n            \"ID\" : [target_id + \"_\" + x for x in all_preds[0].res_id.values.astype(\"str\")],\n            \"resname\": all_preds[0].res_name,\n            \"resid\": all_preds[0].res_id,\n            \"x_1\": all_preds[0].x,\n            \"y_1\": all_preds[0].y,\n            \"z_1\": all_preds[0].z,\n            \"x_2\": all_preds[1].x,\n            \"y_2\": all_preds[1].y,\n            \"z_2\": all_preds[1].z,\n            \"x_3\": all_preds[2].x,\n            \"y_3\": all_preds[2].y,\n            \"z_3\": all_preds[2].z,\n            \"x_4\": all_preds[3].x,\n            \"y_4\": all_preds[3].y,\n            \"z_4\": all_preds[3].z,\n            \"x_5\": all_preds[4].x,\n            \"y_5\": all_preds[4].y,\n            \"z_5\": all_preds[4].z,\n        })\n\n        subm.append(res)\n    submission = pd.concat(subm, axis=0)\n    return submission","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission = format_submission(output_path, df, seed=101)\nsubmission","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission.to_csv(\"/kaggle/working/submission.csv\", index = False)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%cd /kaggle/working/\n!rm -rf /kaggle/working/output\n!rm -rf /kaggle/working/Protenix","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}