{"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":118765,"databundleVersionId":15231210,"sourceType":"competition"},{"sourceId":10880374,"sourceType":"datasetVersion","datasetId":6760482},{"sourceId":10880419,"sourceType":"datasetVersion","datasetId":6760509},{"sourceId":11230242,"sourceType":"datasetVersion","datasetId":7014687},{"sourceId":11451236,"sourceType":"datasetVersion","datasetId":7174725},{"sourceId":11899194,"sourceType":"datasetVersion","datasetId":7479946},{"sourceId":13282339,"sourceType":"datasetVersion","datasetId":7162026}],"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":"2026-01-09T17:14:05.748164Z","iopub.execute_input":"2026-01-09T17:14:05.748527Z","iopub.status.idle":"2026-01-09T17:14:38.770407Z","shell.execute_reply.started":"2026-01-09T17:14:05.748497Z","shell.execute_reply":"2026-01-09T17:14:38.769007Z"},"scrolled":true,"_kg_hide-output":true,"_kg_hide-input":true},"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":"2026-01-09T17:14:38.772281Z","iopub.execute_input":"2026-01-09T17:14:38.772647Z","iopub.status.idle":"2026-01-09T17:14:57.052686Z","shell.execute_reply.started":"2026-01-09T17:14:38.772596Z","shell.execute_reply":"2026-01-09T17:14:57.051739Z"},"_kg_hide-output":true,"scrolled":true,"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import modelgenerator\nimport os","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T17:14:57.054373Z","iopub.execute_input":"2026-01-09T17:14:57.054614Z","iopub.status.idle":"2026-01-09T17:14:57.062196Z","shell.execute_reply.started":"2026-01-09T17:14:57.054593Z","shell.execute_reply":"2026-01-09T17:14:57.061376Z"},"_kg_hide-input":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,"execution":{"iopub.status.busy":"2026-01-09T17:14:57.063565Z","iopub.execute_input":"2026-01-09T17:14:57.064025Z","iopub.status.idle":"2026-01-09T17:15:00.939971Z","shell.execute_reply.started":"2026-01-09T17:14:57.064001Z","shell.execute_reply":"2026-01-09T17:15:00.938878Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import json\nimport numpy as np\nimport os\nimport pandas as pd\nimport sys\nimport time\nimport warnings\n\nfrom biotite.structure.io.pdbx import CIFFile, get_structure\nfrom pathlib import Path\n\nwarnings.simplefilter(\"ignore\", category=FutureWarning)\nwarnings.filterwarnings(\"ignore\", category=UserWarning, module=\"torch.utils.data.dataloader\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T17:15:00.941055Z","iopub.execute_input":"2026-01-09T17:15:00.941466Z","iopub.status.idle":"2026-01-09T17:15:01.642307Z","shell.execute_reply.started":"2026-01-09T17:15:00.941428Z","shell.execute_reply":"2026-01-09T17:15:01.641681Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Utils","metadata":{}},{"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\n\ndef create_base_df(sequence: str):\n    df = pd.DataFrame()\n    df[\"res_name\"] = list(sequence)\n    df[\"res_id\"] = [i+1 for i in range(len(sequence))]\n    df[[\"x\", \"y\", \"z\"]] = np.zeros((len(df),3))\n    return df\n\n\ndef format_submission(output_path, df, seed):\n    # create submission\n    subm = []\n    for _, row in df.iterrows():\n        try:\n            target_id = row[\"target_id\"]\n            sequence = row[\"sequence\"]\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                    # L = len(sequence)\n                    # N = 5\n                    # coordinate = [np.random.rand(L, 3) for _ in range(N)]\n                    # pred_df = pd.DataFrame()\n                    # pred_df[\"x\"] = c1_atoms.coord[:, 0]\n                    # pred_df[\"y\"] = c1_atoms.coord[:, 1]\n                    # pred_df[\"z\"] = c1_atoms.coord[:, 2]\n                    # df = df[[\"res_name\", \"res_id\", \"x\", \"y\", \"z\"]]\n                    continue  \n                pred_df = extract_c1_atoms(cif_path)\n                aux = create_base_df(sequence)\n                aux.loc[:len(pred_df)-1, [\"x\", \"y\", \"z\"]] = pred_df[[\"x\", \"y\", \"z\"]].values\n                all_preds.append(aux)\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        except:\n            target_id = row[\"target_id\"]\n            sequence = row[\"sequence\"]\n            res = pd.DataFrame({\n                \"ID\" : [target_id + \"_\" + str(x) for x in range(1, len(sequence) + 1)],\n                \"resname\": list(sequence),\n                \"resid\": [x for x in range(1, len(sequence) + 1)],\n            })\n            \n\n        subm.append(res)\n    submission = pd.concat(subm, axis=0)\n    submission = submission.fillna(0)\n    return submission\n\n\ndef sep():\n    print(\"—\"*100)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T17:15:01.643049Z","iopub.execute_input":"2026-01-09T17:15:01.643488Z","iopub.status.idle":"2026-01-09T17:15:01.654593Z","shell.execute_reply.started":"2026-01-09T17:15:01.643464Z","shell.execute_reply":"2026-01-09T17:15:01.653715Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Prepare data","metadata":{}},{"cell_type":"code","source":"input_path = Path(\"/kaggle/input/stanford-rna-3d-folding-2\")\nprotenix_path = Path(\"/kaggle/working/Protenix\")\noutput_path = Path(\"/kaggle/working/output\")\noutput_path.mkdir(parents=True, exist_ok=True)\n\ndf = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding-2/test_sequences.csv\")\ndf.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T17:15:01.655546Z","iopub.execute_input":"2026-01-09T17:15:01.655929Z","iopub.status.idle":"2026-01-09T17:15:03.703091Z","shell.execute_reply.started":"2026-01-09T17:15:01.655899Z","shell.execute_reply":"2026-01-09T17:15:03.702182Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df.sequence.apply(len)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T17:15:03.704534Z","iopub.execute_input":"2026-01-09T17:15:03.704827Z","iopub.status.idle":"2026-01-09T17:15:03.710379Z","shell.execute_reply.started":"2026-01-09T17:15:03.704791Z","shell.execute_reply":"2026-01-09T17:15:03.709778Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def prepare_data(df, output_path, max_length: int = 820):\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        if len(sequence) <= max_length:\n            input_json= [(\n                {\n                    \"sequences\": [\n                        {\n                            \"rnaSequence\": {\n                                \"sequence\": sequence,\n                                \"count\": 1, \n                                \"msa\": {\n                                    \"precomputed_msa_dir\": f\"{input_path}/MSA/{target_id}.MSA.fasta\",\n                                    \"pairing_db\": \"rnacentral\"\n                                }\n                            }\n                        }\n                    ],\n                    \"name\": target_id,\n                }\n            )]\n        else:\n            print(\"===== Sequence too long! ======\")\n            input_json= [(\n                {\n                    \"sequences\": [\n                        {\n                            \"rnaSequence\": {\n                                \"sequence\": sequence[:max_length],\n                                \"count\": 1, \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}\")\n\n\nprepare_data(df, output_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T17:15:03.712317Z","iopub.execute_input":"2026-01-09T17:15:03.712525Z","iopub.status.idle":"2026-01-09T17:15:03.742665Z","shell.execute_reply.started":"2026-01-09T17:15:03.712507Z","shell.execute_reply":"2026-01-09T17:15:03.742098Z"},"_kg_hide-output":true,"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"cat /kaggle/working/output/input_json/8ZNQ.json","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T17:15:03.743572Z","iopub.execute_input":"2026-01-09T17:15:03.743884Z","iopub.status.idle":"2026-01-09T17:15:03.862398Z","shell.execute_reply.started":"2026-01-09T17:15:03.743846Z","shell.execute_reply":"2026-01-09T17:15:03.861584Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Inference function","metadata":{}},{"cell_type":"code","source":"%cd $protenix_path\n\ndef inference(target_id, output_path, sequence, seed=101, n_cycle=10, n_sample=5, n_step=200, max_length=820):\n    checkpoint_path = \"/kaggle/input/protenix-finetuned-rna3db-all-1599/1599_ema_0.999.pt\"\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    if len(sequence) <= max_length:\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    else:\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 False\",\n                f\"--load_checkpoint_path={checkpoint_path}\",\n                f\"\"\n            ]\n\n    from runner.inference import run\n\n    run()\n  ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T17:15:03.863781Z","iopub.execute_input":"2026-01-09T17:15:03.864121Z","iopub.status.idle":"2026-01-09T17:15:03.872337Z","shell.execute_reply.started":"2026-01-09T17:15:03.864094Z","shell.execute_reply":"2026-01-09T17:15:03.871571Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Inference ","metadata":{}},{"cell_type":"code","source":"print(f\"Current working directory: {os.getcwd()}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T17:15:03.87328Z","iopub.execute_input":"2026-01-09T17:15:03.873657Z","iopub.status.idle":"2026-01-09T17:15:03.891667Z","shell.execute_reply.started":"2026-01-09T17:15:03.873613Z","shell.execute_reply":"2026-01-09T17:15:03.890991Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# IS_SUBMISSION = os.environ.get(\"KAGGLE_IS_COMPETITION_RERUN\")\n# print(IS_SUBMISSION)\nIS_SUBMISSION = True\n\nif IS_SUBMISSION:\n    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, seq, seed=101, n_cycle=10, n_sample=5, n_step=200, max_length=820)\n        end = time.time()\n        print(f\"Time used: {((end-start)/60):2f} minutes.\")\n        sep()\n\n    submission = format_submission(output_path, df, seed=101)\n    submission = submission.fillna(0.0)\n    submission.to_csv(f\"/kaggle/working/submission.csv\", index = False)\n    display(submission.head())\n\nelse:\n    submission = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding-2/sample_submission.csv\")\n    submission.to_csv(f'/kaggle/working/submission.csv', index=False)\n    display(submission.head())","metadata":{"execution":{"iopub.status.busy":"2026-01-09T17:15:03.892429Z","iopub.execute_input":"2026-01-09T17:15:03.892692Z"},"trusted":true,"scrolled":true},"outputs":[],"execution_count":null}]}