{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.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":14786962,"sourceType":"datasetVersion","datasetId":9447634},{"sourceId":14787388,"sourceType":"datasetVersion","datasetId":9453383}],"dockerImageVersionId":31260,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Boltz2 baseline \nThis is a simple notebook using boltz in a basic config, with \"sequence_only\" modality","metadata":{}},{"cell_type":"code","source":"!ls /kaggle/input/datasets/lbugnon/","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-12T06:54:04.186172Z","iopub.execute_input":"2026-02-12T06:54:04.186569Z","iopub.status.idle":"2026-02-12T06:54:04.30595Z","shell.execute_reply.started":"2026-02-12T06:54:04.186531Z","shell.execute_reply":"2026-02-12T06:54:04.304356Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings('ignore')\n\n!cp -r /kaggle/input/datasets/lbugnon/boltz-src-minimal ./\n# Install boltz from src (some weird bug in kaggle notebooks with fairscale so it was removed)\n!pip install --no-index --no-build-isolation -e ./boltz-src-minimal","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# need to create tar so botlz do not try to download, also clean the trash dataset creation does\n!mkdir boltz_cache\n!cp -r /kaggle/input/datasets/lbugnon/boltz2 boltz_cache\n!mv boltz_cache/boltz2/mols/mols/* boltz_cache/boltz2/mols/\n!rm -r boltz_cache/boltz2/mols/mols/\n!tar -cf boltz_cache/boltz2/mols.tar boltz_cache/boltz2/mols","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Run predictions using all chains in the input (not sure how this is handled in the test input)","metadata":{}},{"cell_type":"code","source":"import time\nimport pandas as pd\nfrom random import shuffle\nimport gc\nimport torch\n        \nsequences = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding-2/test_sequences.csv\")\n!mkdir input_fasta\n!rm -r boltz_results*\n\nuse_all_chains = False\nmax_tokens = 900 # aprox for 16Gb gram\nmax_repeats = 999\npred_repeats = 5\n\nnucleotides = {\"A\", \"G\", \"C\", \"U\"}\naminoacids = {\"A\", \"R\", \"N\", \"D\", \"C\", \"E\", \"Q\", \"G\", \"H\", \"I\", \"L\", \"K\", \"M\", \"F\", \"P\", \"S\", \"T\", \"W\", \"Y\", \"V\"}\n\nfor k in range(len(sequences)):\n\n    # ignore super large seqs\n    if len(sequences.iloc[k].sequence)>max_tokens:\n        continue\n    \n    t0 = time.time()\n    name = sequences.iloc[k].target_id\n    print(f\"preparing {name} ({k+1} of {len(sequences)})\")\n    \n    pred_seq = sequences.iloc[k].sequence\n\n    # A for the chain to predict. Firs load all chains. This may fail if all_sequences is empty or different format\n    try:\n        chains, chain_ind = {}, 1\n        for entry in sequences.iloc[k].all_sequences.split(\"\\n\"):\n            if entry.startswith(\">\"):\n                # get number of repeats\n                repeats = min(len(entry.split(\"|\")[1].split(\",\")), max_repeats)\n            else:\n                if entry == pred_seq:\n                    chains[0] = entry, \"rna\"\n                elif use_all_chains:\n                    entry_type = \"rna\" if set(entry)<=nucleotides else \"protein\"  if set(entry) <= aminoacids else None\n                    if entry_type is None:\n                        continue\n                    chains[chain_ind] = entry, entry_type\n                    chain_ind += 1\n                if use_all_chains:\n                    for i in range(1, repeats):\n                        chains[chain_ind] = entry, entry_type\n                        chain_ind += 1\n    \n        # Now keep the chain 0 plus all chains randomly picked to fit max_tokens (asuming chain 0 is ok under tokens, otherwise will fail)\n        filtered_chains, ntokens = {}, 0\n        other_chains = [c for c in chains if c!=0]\n        shuffle(other_chains)\n        for c in [0] + other_chains:\n            ntokens += len(chains[c][0])\n            if ntokens >= max_tokens:\n                break\n            filtered_chains[c] = chains[c]\n        chains = filtered_chains\n    except:\n        # use only target sequence \n        chains = {0: (pred_seq, \"rna\")}\n        \n    # save fasta, chain 0 first\n    with open(f\"input_fasta/{name}.fasta\", \"w\") as fout:\n        for chain in sorted(chains.keys()):\n            seq, seq_type = chains[chain]\n            msa = \"\" if seq_type == \"rna\" else \"empty\"\n            # if single seq, save only chain A otherwise save all\n            fout.write(f\">{chain}|{seq_type}|{msa}\\n{seq}\\n\")\n    print(name, \":\")\n    !cat input_fasta/{name}.fasta\n    print()\n    for repeat in range(pred_repeats):\n        !boltz predict input_fasta/{name}.fasta --num_workers 1\t--max_parallel_samples 1 --output_format pdb\t--cache boltz_cache/boltz2/ --out_dir boltz_repeat_{repeat} \n        gc.collect()\n        torch.cuda.empty_cache()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"raw","source":"Read predictions and generate submission file","metadata":{}},{"cell_type":"code","source":"from biopandas.pdb import PandasPdb\nimport numpy as np \n\npreds = {}\nfor k in range(len(sequences)):\n    name = sequences.iloc[k].target_id\n\n    for model in range(5):\n        model_pdb = model\n        if pred_repeats == 1:\n            model_pdb = 0\n        # TODO: run 5 models, here we just use the same 5 times\n        fname = f\"boltz_repeat_{model_pdb}/boltz_results_{name}/predictions/{name}/{name}_model_0.pdb\"\n        try:\n            pl1 = PandasPdb().read_pdb(fname)\n            # Keep only C1' from chain A\n            res_coords = pl1.df[\"ATOM\"][(pl1.df[\"ATOM\"].chain_id==\"0\") & (pl1.df[\"ATOM\"].atom_name==\"C1'\")]\n        except FileNotFoundError:\n            # Randomly fill failed sequence\n            res_coords = []\n            for i, nt in enumerate(sequences.iloc[k].sequence):\n                res_coords.append([i+1, nt, np.random.rand(1).item(), np.random.rand(1).item(), np.random.rand(1).item()])\n            res_coords = pd.DataFrame(res_coords, columns=[\"residue_number\", \"residue_name\", \"x_coord\", \"y_coord\", \"z_coord\"])\n        for id, res in res_coords.iterrows():\n            if model == 0:\n                preds[f\"{name}_{res.residue_number}\"] = [f\"{name}_{res.residue_number}\", res.residue_name, res.residue_number]\n            preds[f\"{name}_{res.residue_number}\"] += [res.x_coord, res.y_coord, res.z_coord]\n        \n\n\npreds = pd.DataFrame(preds).T\npreds.columns = [\"ID\", \"resname\",\"resid\",\"x_1\",\"y_1\",\"z_1\",\"x_2\",\"y_2\",\"z_2\",\"x_3\",\"y_3\",\n                 \"z_3\",\"x_4\",\"y_4\",\"z_4\",\"x_5\",\"y_5\",\"z_5\"]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#rm all files but submisison.csv\n!rm -r *\npreds.to_csv(\"submission.csv\", index=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}