{"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":1171411,"sourceType":"datasetVersion","datasetId":664454},{"sourceId":2397625,"sourceType":"datasetVersion","datasetId":1449606},{"sourceId":5123458,"sourceType":"datasetVersion","datasetId":2975803},{"sourceId":5127751,"sourceType":"datasetVersion","datasetId":2978411},{"sourceId":7051993,"sourceType":"datasetVersion","datasetId":4058626},{"sourceId":7292242,"sourceType":"datasetVersion","datasetId":4229373},{"sourceId":10855324,"sourceType":"datasetVersion","datasetId":6742586},{"sourceId":10878276,"sourceType":"datasetVersion","datasetId":6758842},{"sourceId":10880419,"sourceType":"datasetVersion","datasetId":6760509},{"sourceId":10923077,"sourceType":"datasetVersion","datasetId":6785143},{"sourceId":11014650,"sourceType":"datasetVersion","datasetId":6857936},{"sourceId":11118830,"sourceType":"datasetVersion","datasetId":6933267},{"sourceId":11184784,"sourceType":"datasetVersion","datasetId":6981833},{"sourceId":11411364,"sourceType":"datasetVersion","datasetId":6896780},{"sourceId":11661237,"sourceType":"datasetVersion","datasetId":7317710},{"sourceId":11695366,"sourceType":"datasetVersion","datasetId":6791615},{"sourceId":11882705,"sourceType":"datasetVersion","datasetId":7468209},{"sourceId":11964501,"sourceType":"datasetVersion","datasetId":7234091},{"sourceId":11988491,"sourceType":"datasetVersion","datasetId":7505545},{"sourceId":157966080,"sourceType":"kernelVersion"},{"sourceId":224830487,"sourceType":"kernelVersion"}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!mkdir /kaggle/tmp","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:40:40.432836Z","iopub.execute_input":"2025-05-29T06:40:40.433199Z","iopub.status.idle":"2025-05-29T06:40:40.590158Z","shell.execute_reply.started":"2025-05-29T06:40:40.433171Z","shell.execute_reply":"2025-05-29T06:40:40.589012Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd \nimport numpy as np\n\ndf = pd.read_csv('/kaggle/input/stanford-rna-3d-folding/test_sequences.csv')\ndf.to_csv('/kaggle/working/test_sequence.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:40:40.592012Z","iopub.execute_input":"2025-05-29T06:40:40.592267Z","iopub.status.idle":"2025-05-29T06:40:40.604392Z","shell.execute_reply.started":"2025-05-29T06:40:40.592235Z","shell.execute_reply":"2025-05-29T06:40:40.603594Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!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'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:40:40.606327Z","iopub.execute_input":"2025-05-29T06:40:40.606531Z","iopub.status.idle":"2025-05-29T06:40:57.161728Z","shell.execute_reply.started":"2025-05-29T06:40:40.606513Z","shell.execute_reply":"2025-05-29T06:40:57.160577Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# boltz","metadata":{}},{"cell_type":"code","source":"%%bash\ncp -r /kaggle/input/rna-folding-codebase/boltz /kaggle/working/boltz\nset -euo pipefail\n\nWHEEL_DIR=\"/kaggle/input/boltz-dependencies\"\nFAILED=()\n\npip install --no-index /kaggle/input/fairscale-0413/*whl --no-deps\n\necho \"→ Installing wheels from $WHEEL_DIR …\"\n\nfor whl in \"$WHEEL_DIR\"/*.whl; do\n  echo -n \"  • $whl … \"\n  if pip install --no-index \"$whl\" --no-deps --upgrade; then\n    echo \"OK\"\n  else\n    echo \"FAILED, skipping\"\n    FAILED+=(\"$(basename \"$whl\")\")\n  fi\ndone\n\nif [ \"${#FAILED[@]}\" -gt 0 ]; then\n  echo\n  echo \"⚠️  The following wheels failed to install (platform mismatch?):\"\n  for f in \"${FAILED[@]}\"; do\n    echo \"    – $f\"\n  done\nelse\n  echo\n  echo \"✅ All wheels installed successfully.\"\nfi","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:40:57.163758Z","iopub.execute_input":"2025-05-29T06:40:57.164081Z","iopub.status.idle":"2025-05-29T06:41:14.505144Z","shell.execute_reply.started":"2025-05-29T06:40:57.164056Z","shell.execute_reply":"2025-05-29T06:41:14.50443Z"},"collapsed":true,"jupyter":{"outputs_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile /kaggle/working/boltz/boltz_predict.py\n\nimport os\nimport sys\nimport argparse\nimport pandas as pd\nfrom pathlib import Path\nimport shutil\nimport torch\nimport yaml\nimport shutil\nimport gc\nfrom tqdm import tqdm\nfrom Bio.PDB import MMCIFParser, PDBIO\nimport glob\nimport subprocess\n\n\nfrom inference import predict  # Boltz-1 function\n\n\nfrom Bio.PDB.PDBParser import PDBParser\nfrom Bio.PDB.PDBExceptions import PDBConstructionException\n\ndef validate_pdb_c1_auto(pdb_path: str, expected_len: int):\n    parser = PDBParser(QUIET=True, PERMISSIVE=True)\n    try:\n        structure = parser.get_structure('val', pdb_path)\n    except PDBConstructionException as e:\n        raise ValueError(f\"PDBParser could not read '{pdb_path}': {e}\")\n\n    # pick the first model and first chain\n    model = next(structure.get_models())\n    try:\n        chain = next(model.get_chains())\n    except StopIteration:\n        raise ValueError(f\"No chains found in '{pdb_path}'\")\n\n    # count C1' atoms\n    c1_count = sum(1 for res in chain if \"C1'\" in res)\n    if c1_count != expected_len:\n        raise ValueError(\n            f\"Wrong C1' count in '{pdb_path}': found {c1_count}, expected {expected_len}\"\n        )\n\n    # ensure END record\n    with open(pdb_path, 'r') as fh:\n        lines = fh.read().rstrip().splitlines()\n    if not lines or lines[-1] != 'END':\n        raise ValueError(f\"Missing 'END' record in '{pdb_path}'\")\n    \ndef fix_pdb_format(input_pdb: str, output_pdb: str, chain_id: str = 'A'):\n    \"\"\"\n    Rewrite any ATOM/HETATM lines in `input_pdb` to a fixed-column\n    PDB format in `output_pdb`, enforcing a single chain ID and\n    correct residue numbering.\n    \"\"\"\n    def write_line(record, serial, atom_name, resname, chain, resseq,\n                   x, y, z, occupancy, temp_factor, element):\n        # Columns per PDB spec: 1-6 record, 7-11 serial, 13-16 atom, 18-20 resname,\n        # 22 chain, 23-26 resseq, 31-38 x, 39-46 y, 47-54 z, 55-60 occupancy,\n        # 61-66 tempFactor, 77-78 element, newline.\n        return (\n            f\"{record:<6s}\"                           # 1-6\n            f\"{serial:5d} \"                           # 7-11 + space\n            f\"{atom_name:<4s}\"                       # 13-16\n            f\"{resname:>3s} \"                        # 18-20 + space\n            f\"{chain:1s}\"                            # 22\n            f\"{resseq:4d}\"                           # 23-26\n            f\"    \"                                  # 27-30 (insertion code + spacing)\n            f\"{x:8.3f}{y:8.3f}{z:8.3f}\"               # 31-54 coords\n            f\"{occupancy:6.2f}{temp_factor:6.2f}\"     # 55-66\n            f\"          \"                            # 67-76 padding\n            f\"{element:>2s}\"                         # 77-78\n            f\"\\n\"\n        )\n\n    with open(input_pdb, 'r') as src, open(output_pdb, 'w') as dst:\n        for line in src:\n            if not (line.startswith(\"ATOM\") or line.startswith(\"HETATM\")):\n                dst.write(line)\n                continue\n\n            parts = line.split()\n            # Expect at least: record, serial, atom_name, resname, orig_chain,\n            # orig_resseq, x, y, z, occupancy, temp_factor, [element]\n            record      = parts[0]                                           # ATOM/HETATM :contentReference[oaicite:1]{index=1}\n            serial      = int(parts[1])                                      # atom serial :contentReference[oaicite:2]{index=2}\n            atom_name   = parts[2]                                           # e.g. C1' :contentReference[oaicite:3]{index=3}\n            resname     = parts[3]                                           # residue name :contentReference[oaicite:4]{index=4}\n            orig_chain  = parts[4] if len(parts) > 5 else chain_id           # often incorrect\n            orig_resseq = int(parts[5]) if len(parts) > 5 else 1             # zero-based or wrong\n            x, y, z     = map(float, parts[6:9])                             # coords :contentReference[oaicite:5]{index=5}\n            occupancy   = float(parts[9]) if len(parts) > 9 else 1.00        # occupancy :contentReference[oaicite:6]{index=6}\n            temp_factor = float(parts[10]) if len(parts) > 10 else 0.00      # B-factor :contentReference[oaicite:7]{index=7}\n            element     = parts[11] if len(parts) > 11 else atom_name[0]     # element symbol :contentReference[oaicite:8]{index=8}\n\n            # Override chain to user-specified or default 'A', and\n            # ensure residue numbering starts at 1 if orig_resseq==0\n            chain       = chain_id\n            resseq      = orig_resseq if orig_resseq > 0 else 1\n\n            dst.write(write_line(record, serial, atom_name, resname,\n                                 chain, resseq, x, y, z,\n                                 occupancy, temp_factor, element))\n        # Ensure file ends with END if not already present\n        if not line.strip().upper().startswith(\"END\"):\n            dst.write(\"END\\n\")\n\n# def cif_to_pdb(cif_file, pdb_file):\n#     parser = MMCIFParser(QUIET=True)\n#     struct = parser.get_structure(os.path.basename(cif_file), cif_file)\n#     io = PDBIO()\n#     io.set_structure(struct)\n#     io.save(pdb_file)\nfrom pathlib import Path\nfrom Bio.PDB import MMCIFParser, PDBIO\n\ndef cif_to_pdb(cif_file, pdb_file):\n    \"\"\"\n    Convert a .cif (mmCIF) file to PDB, handling both str and Path inputs,\n    and ensuring single‐char chain IDs if necessary.\n    \"\"\"\n    # Accept either str or Path\n    cif_path = Path(cif_file)\n    pdb_path = Path(pdb_file)\n\n    # 1) Parse the mmCIF\n    parser   = MMCIFParser(QUIET=True)\n    struct_id = cif_path.stem          # filename without suffix\n    structure = parser.get_structure(str(struct_id), str(cif_path))\n\n    # 2) (Optional) Truncate any multi‐char chain IDs to their first letter:\n    for model in structure:\n        for chain in model:\n            chain.id = chain.id[0]      # enforce single‐char chain ID\n\n    # 3) Write out as PDB (pass a string so Biopython can open it)\n    io = PDBIO()\n    io.set_structure(structure)\n    io.save(str(pdb_path))\n\nfrom pathlib import Path\nfrom Bio.PDB import MMCIFParser, PDBIO\nfrom Bio.PDB.PDBParser import PDBParser, PDBConstructionException\nimport tempfile, os, re\n\ndef cif_to_clean_pdb(cif_path, pdb_path, chain_id='A'):\n    \"\"\"\n    1) Parse mmCIF → Structure\n    2) Truncate any multi-char chain IDs\n    3) Write a raw PDB with PDBIO.save(str(…))\n    4) Re-load that PDB (so malformed coords are reformatted)\n    5) Write final clean PDB at pdb_path\n    \"\"\"\n    cif_path = Path(cif_path)\n    pdb_path = Path(pdb_path)\n\n    # Step 1: parse mmCIF\n    parser = MMCIFParser(QUIET=True)\n    struct_id = cif_path.stem\n    structure = parser.get_structure(struct_id, str(cif_path))\n\n    # Step 2: truncate chains to one char\n    for model in structure:\n        for chain in model:\n            chain.id = chain.id[0]\n\n    # Step 3: dump raw PDB to temp\n    fd, raw_pdb = tempfile.mkstemp(suffix='.pdb')\n    os.close(fd)\n    io = PDBIO()\n    io.set_structure(structure)\n    io.save(str(raw_pdb))\n\n    # Step 4: re-parse raw PDB so that any jammed fields get realigned\n    #          by letting PDBParser + PDBIO do the formatting for us.\n    parser2 = PDBParser(QUIET=True, PERMISSIVE=True)\n    try:\n        struct2 = parser2.get_structure(struct_id, raw_pdb)\n    except PDBConstructionException as e:\n        os.remove(raw_pdb)\n        raise RuntimeError(f\"Error cleaning PDB coords: {e}\")\n\n    # Step 5: write final clean PDB\n    io2 = PDBIO()\n    io2.set_structure(struct2)\n    io2.save(str(pdb_path))\n\n    os.remove(raw_pdb)\n\n\ndef main():\n    parser = argparse.ArgumentParser(description=\"Boltz-1 inference\")\n    parser.add_argument('--input_csv', required=True, help='CSV with [sequence,target_id]')\n    parser.add_argument('--output_dir', required=True, help='Unified output dir for CIFs, PDBs, and summary CSV')\n    parser.add_argument('--diffusion_samples', type=int, default=2, help='Number of samples per input')\n    parser.add_argument('--seed', type=int, default=42)\n    parser.add_argument('--cache_dir', default='/kaggle/input/rna-prediction-boltz/', help='Path to Boltz model cache')\n    parser.add_argument('--ll', type=int, default=100000000,\n                        help='lengthlim')\n    args = parser.parse_args()\n\n    # set up directories\n    os.makedirs(args.output_dir, exist_ok=True)\n    input_dir = Path(args.output_dir) / 'inputs_prediction'\n    result_dir = Path(args.output_dir) / 'boltz_results'\n    os.makedirs(input_dir, exist_ok=True)\n    os.makedirs(result_dir, exist_ok=True)\n    \n    # read inputs\n    df = pd.read_csv(args.input_csv)\n    records = []\n    \n    # write YAMLs for Boltz predict\n    for _, row in df.iterrows():\n        target_id = row['target_id']\n        seq = row['sequence']\n        yaml_path = os.path.join(input_dir, f'{target_id}.yaml')\n    \n        spec = {\n            \"constraints\": [],\n            \"sequences\": [\n                {\n                    \"rna\": {\n                        \"id\": [\"A1\"],\n                        \"sequence\": seq,\n                        \"msa\": f\"/kaggle/input/stanford-rna-3d-folding/MSA_v2/{seq}.MSA.fasta\"\n                    }\n                }\n            ],\n        }\n    \n        with open(yaml_path, 'w') as f:\n            yaml.safe_dump(spec, f, sort_keys=False)\n    # # run Boltz batch inference\n    # predict(\n    #     data=input_dir,\n    #     out_dir=result_dir,\n    #     cache=args.cache_dir,\n    #     diffusion_samples=args.diffusion_samples,\n    #     output_format = 'pdb',\n    #     seed=args.seed,\n    #     override=True,\n    #     num_workers=0,\n    #     accelerator=\"gpu\",\n    #     devices=1,\n    # )\n    \n    yaml_paths = sorted(glob.glob(os.path.join(input_dir, '*.yaml')))\n    \n    # … after you build `yaml_paths` and inside your for‐batch loop …\n    \n    for batch_idx in range(0, len(yaml_paths), 20):\n        batch = yaml_paths[batch_idx:batch_idx+20]\n        batch_name = f\"batch_{batch_idx//20}\"\n        batch_dir = Path(input_dir) / batch_name\n        batch_dir.mkdir(exist_ok=True)\n    \n        # copy only this batch’s yamls\n        for y in batch:\n            shutil.copy(y, batch_dir / Path(y).name)\n    \n        # 1) run inference on this batch\n        cmd = [\n            sys.executable, 'predict_one.py',\n            '--data', str(batch_dir),\n            '--out_dir', str(result_dir),\n            '--cache', args.cache_dir,\n            '--diffusion_samples', str(args.diffusion_samples),\n            '--num_workers', '0',\n            '--seed', str(args.seed),\n        ]\n        # if args.override:\n        #     cmd.append('--override')\n    \n        print(\"Running:\", \" \".join(cmd))\n        os.system(\" \".join(cmd))\n        # res = subprocess.run(cmd)\n        # if res.returncode != 0:\n        #     print(f\"[ERROR] batch {batch_name} failed (exit {res.returncode})\")\n    \n        # 2) locate the actual Boltz output folder for this batch\n        #    inference.py writes under: result_dir/boltz_results_<data.stem>\n        out_sub = result_dir / f\"boltz_results_{batch_name}\"\n        pred_base = out_sub / \"predictions\"\n    \n        # 3) for each YAML in this batch, copy its samples\n        for y in batch:\n            tid = Path(y).stem\n            targ_pred = pred_base / tid\n            targ_pred.mkdir(exist_ok=True,parents=True)\n            if not targ_pred.exists():\n                print(f\"[WARN] no predictions dir for {tid} in {pred_base}\")\n                continue\n            \n            for sample_i in range(args.diffusion_samples):\n                # predict_one.py writes files like <tid>_model_<i>.pdb\n                src = targ_pred / f\"{tid}_model_{sample_i}.cif\"\n                if not src.is_file():\n                    print(f\"[WARN] missing file {src}\")\n                    continue\n                \n                seq = df.loc[df.target_id==tid,'sequence'].iat[0]\n                \n                dst = Path(args.output_dir) / f\"{tid}_model_{sample_i}.pdb\"\n                # print(dst)\n                try:\n                    cif_to_clean_pdb(src, dst)\n                except FileNotFoundError:\n                    print(f\"[Warn] Missing CIF: {src}\")\n                    continue\n                except ValueError as e:\n                    print(f\"[Warn] {e}\")\n                    continue\n                \n                validate_pdb_c1_auto(dst, len(seq))\n                \n                records.append({\n                    'sequence': seq,\n                    'target_id': tid,\n                    'sample': sample_i,\n                    'predicted_model_path': str(dst)\n                })\n                print(f\"Copied {src} → {dst}\")\n    \n        # 4) cleanup batch and its Boltz output\n        # break\n        shutil.rmtree(batch_dir, ignore_errors=True)\n        shutil.rmtree(out_sub,  ignore_errors=True)\n        gc.collect()\n\n\n    # save summary CSV\n    out_csv = os.path.join(args.output_dir, 'boltz_preds_summary.csv')\n    pd.DataFrame(records).to_csv(out_csv, index=False)\n    print(f\"✅ Saved summary CSV to {out_csv}\")\n\nif __name__ == '__main__':\n    main()\n","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-05-29T06:41:14.506314Z","iopub.execute_input":"2025-05-29T06:41:14.506594Z","iopub.status.idle":"2025-05-29T06:41:14.51451Z","shell.execute_reply.started":"2025-05-29T06:41:14.506571Z","shell.execute_reply":"2025-05-29T06:41:14.513699Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%bash\nset -e\nexport PYTHONPATH=/kaggle/working/boltz/boltz/src:$PYTHONPATH\n# install boltz dependencies\n# pip install --no-index /kaggle/input/boltz-dependencies/*whl --no-deps\nmkdir /kaggle/working/boltz_cache/\ncp /kaggle/input/boltz-checkpoints/ccd.pkl /kaggle/working/boltz_cache/\ncp /kaggle/input/boltz-checkpoints/boltz1_conf.ckpt /kaggle/working/boltz_cache/\n\n# move into the boltz folder\ncp -r /kaggle/input/rna-prediction-boltz/boltz /kaggle/working/boltz_cache\n\ncd /kaggle/working/boltz\n\n# run the predictor\npython boltz_predict.py \\\n    --input_csv  /kaggle/working/test_sequence.csv \\\n    --output_dir /kaggle/working/boltz_predictions \\\n    --cache_dir  /kaggle/working/boltz_cache \\\n    --diffusion_samples 2\nrm -rf /kaggle/working/boltz_cache\nrm -rf /kaggle/working/boltz","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:41:14.51561Z","iopub.execute_input":"2025-05-29T06:41:14.515979Z","iopub.status.idle":"2025-05-29T06:43:07.479236Z","shell.execute_reply.started":"2025-05-29T06:41:14.515941Z","shell.execute_reply":"2025-05-29T06:43:07.478219Z"},"collapsed":true,"jupyter":{"outputs_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Proteinx","metadata":{}},{"cell_type":"code","source":"%%time\n!cp -r /kaggle/input/rna-folding-codebase/proteinx /kaggle/working/proteinx\n!export PROTENIX_DATA_ROOT_DIR=/kaggle/input/protenix-checkpoints\n!mkdir /af3-dev \n!cp /kaggle/input/usalign/USalign /kaggle/tmp/USalign\n!ln -s /kaggle/input/protenix-checkpoints /af3-dev/release_data\n!ls /af3-dev/release_data/\n!cd /kaggle/working/proteinx\n!chmod +x /kaggle/tmp/USalign\nimport sys\n!{sys.executable} /kaggle/working/proteinx/proteinx_predict.py \\\n    --input_csv   /kaggle/working/test_sequence.csv \\\n    --output_dir  /kaggle/working/proteinx_predictions \\\n    --checkpoint  /kaggle/input/protenix-checkpoints/model_v0.2.0.pt \\\n    --tmp_dir     /kaggle/working/output","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:43:07.48038Z","iopub.execute_input":"2025-05-29T06:43:07.480698Z","iopub.status.idle":"2025-05-29T06:44:21.034249Z","shell.execute_reply.started":"2025-05-29T06:43:07.480672Z","shell.execute_reply":"2025-05-29T06:44:21.033193Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Nufold","metadata":{}},{"cell_type":"code","source":"!rm -rf /kaggle/working/nufold_predictions","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:44:21.037187Z","iopub.execute_input":"2025-05-29T06:44:21.037472Z","iopub.status.idle":"2025-05-29T06:44:21.182715Z","shell.execute_reply.started":"2025-05-29T06:44:21.037448Z","shell.execute_reply":"2025-05-29T06:44:21.181726Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Assuming you've mounted your dataset containing the wheels and Protenix code\nDATASET_PATH = '/kaggle/input/required-presets'\nimport os\n# Install dependencies from wheels\nwheel_path = os.path.join(DATASET_PATH, 'wheels')\n!pip install -q --no-index --find-links {wheel_path} torch torchvision torchaudio \\\nml-collections dm-tree deepspeed protobuf scipy biopython numpy matplotlib biotite","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:44:21.184642Z","iopub.execute_input":"2025-05-29T06:44:21.184873Z","iopub.status.idle":"2025-05-29T06:44:24.680669Z","shell.execute_reply.started":"2025-05-29T06:44:21.184852Z","shell.execute_reply":"2025-05-29T06:44:24.679336Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\n!cp -r /kaggle/input/rna-folding-codebase/nufold /kaggle/working/nufold\n%cd /kaggle/working/nufold\n!chmod +x /kaggle/working/nufold/req/ipknot-1.1.0-x86_64-linux/ipknot\n!python nufold_predict.py \\\n --input_csv /kaggle/working/test_sequence.csv \\\n --output_dir  /kaggle/working/nufold_predictions \\\n --checkpoint /kaggle/working/nufold/req/checkpoints/global_step145245.pt\\\n --run_nufold_script /kaggle/working/nufold/req/NuFold/run_nufold.py\\\n --msa_dir /kaggle/input/stanford-rna-3d-folding/MSA\\\n --ipknot_bin /kaggle/working/nufold/req/ipknot-1.1.0-x86_64-linux/ipknot\\\n --tmp_dir /kaggle/tmp/nufold\n%cd /kaggle/working","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:44:24.681923Z","iopub.execute_input":"2025-05-29T06:44:24.682237Z","iopub.status.idle":"2025-05-29T06:45:18.270413Z","shell.execute_reply.started":"2025-05-29T06:44:24.682205Z","shell.execute_reply":"2025-05-29T06:45:18.269352Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# rhofold","metadata":{}},{"cell_type":"code","source":"# !pip install /kaggle/input/openmm/OpenMM-8.2.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:45:18.271785Z","iopub.execute_input":"2025-05-29T06:45:18.27215Z","iopub.status.idle":"2025-05-29T06:45:18.276174Z","shell.execute_reply.started":"2025-05-29T06:45:18.272112Z","shell.execute_reply":"2025-05-29T06:45:18.275377Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# import shutil\n# shutil.copytree(\"/kaggle/input/rna-folding-codebase/rhofold\", \"/kaggle/working/rhofold\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:45:18.277201Z","iopub.execute_input":"2025-05-29T06:45:18.277482Z","iopub.status.idle":"2025-05-29T06:45:18.292212Z","shell.execute_reply.started":"2025-05-29T06:45:18.277449Z","shell.execute_reply":"2025-05-29T06:45:18.291594Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %%bash\n# cd /kaggle/input/rna-folding-codebase/rhofold\n# python rhofold_predict.py \\\n#   --input_csv /kaggle/working/test_sequence.csv \\\n#   --rhonet_dir /kaggle/input/rna-folding-codebase/rhofold/RhoFold-main\\\n#   --tmp_dir /kaggle/tmp/rhofold \\\n#   --output_dir /kaggle/working/rhofold_predictions \\\n#   --msa_dir /kaggle/input/stanford-rna-3d-folding/MSA/\n#   # --device cuda","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:45:18.29312Z","iopub.execute_input":"2025-05-29T06:45:18.293452Z","iopub.status.idle":"2025-05-29T06:45:18.309683Z","shell.execute_reply.started":"2025-05-29T06:45:18.293414Z","shell.execute_reply":"2025-05-29T06:45:18.308811Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# deepfold","metadata":{}},{"cell_type":"code","source":"# !pip install /kaggle/input/tensorflow-gpu-1150/tensorflow_gpu-1.15.0-cp37-cp37m-manylinux2010_x86_64.whl","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:45:18.310544Z","iopub.execute_input":"2025-05-29T06:45:18.310827Z","iopub.status.idle":"2025-05-29T06:45:18.323652Z","shell.execute_reply.started":"2025-05-29T06:45:18.310798Z","shell.execute_reply":"2025-05-29T06:45:18.322988Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!rm -rf /kaggle/working/deepfold\nimport shutil\nshutil.copytree('/kaggle/input/rna-folding-codebase/deepfold','/kaggle/working/deepfold')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:45:18.324651Z","iopub.execute_input":"2025-05-29T06:45:18.324864Z","iopub.status.idle":"2025-05-29T06:45:21.141735Z","shell.execute_reply.started":"2025-05-29T06:45:18.324842Z","shell.execute_reply":"2025-05-29T06:45:21.140672Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%bash\nchmod +x /kaggle/working/deepfold/DeepFoldRNA/fold/DeepFold\nchmod +x /kaggle/working/deepfold/DeepFoldRNA/bin/PETfold\nchmod +x /kaggle/working/deepfold/DeepFoldRNA/bin/SimRNA_64bitIntel_Linux/SimRNA\nchmod +x /kaggle/working/deepfold/DeepFoldRNA/bin/SimRNA_64bitIntel_Linux/SimRNA_trafl2pdbs\nchmod +x /kaggle/working/deepfold/DeepFoldRNA/bin/QRNAS/QRNA","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:45:21.143046Z","iopub.execute_input":"2025-05-29T06:45:21.143449Z","iopub.status.idle":"2025-05-29T06:45:21.159562Z","shell.execute_reply.started":"2025-05-29T06:45:21.143405Z","shell.execute_reply":"2025-05-29T06:45:21.15878Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %%time\n# %cd /kaggle/working/deepfold\n# !mkdir -p /kaggle/tmp/deepfold/fasta_files\n# !python -u main.py \\\n#     --input_csv /kaggle/working/test_sequence.csv \\\n#     --output_dir /kaggle/working/deepfold_predictions \\\n#     --tmp_dir /kaggle/tmp/deepfold \\\n#     --msa_dir /kaggle/input/stanford-rna-3d-folding/MSA/ \\\n#     --ll 800 \\\n#     --deepfoldrna_dir /kaggle/working/deepfold/DeepFoldRNA\n# %cd /kaggle/working","metadata":{"trusted":true,"scrolled":true,"execution":{"iopub.status.busy":"2025-05-29T06:45:21.160606Z","iopub.execute_input":"2025-05-29T06:45:21.160897Z","iopub.status.idle":"2025-05-29T06:50:59.409338Z","shell.execute_reply.started":"2025-05-29T06:45:21.16087Z","shell.execute_reply":"2025-05-29T06:50:59.408194Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# drfold","metadata":{}},{"cell_type":"code","source":"!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'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:50:59.410811Z","iopub.execute_input":"2025-05-29T06:50:59.411177Z","iopub.status.idle":"2025-05-29T06:51:17.29774Z","shell.execute_reply.started":"2025-05-29T06:50:59.411124Z","shell.execute_reply":"2025-05-29T06:51:17.296792Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import shutil\n!rm -rf /kaggle/working/drfold\nshutil.copytree('/kaggle/input/rna-folding-codebase/drfold', '/kaggle/working/drfold')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:51:17.298804Z","iopub.execute_input":"2025-05-29T06:51:17.29915Z","iopub.status.idle":"2025-05-29T06:51:19.675325Z","shell.execute_reply.started":"2025-05-29T06:51:17.299117Z","shell.execute_reply":"2025-05-29T06:51:19.674212Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile /kaggle/working/drfold/predict.py\n#!/usr/bin/env python3\n\"\"\"\nDRfold2 direct‐inference script with notebook logic, now fully parameterized.\n\nUsage:\n    python drfold2_direct.py \\\n      --mode test \\\n      --test_path train_sequences.csv \\\n      --output_path drfold2_output \\\n      --output_csv submission.csv \\\n      --model_dir /kaggle/input/drfold-model-bf16 \\\n      [--cfg_name RCLM]\n\"\"\"\n\nimport os, sys, time, argparse, glob\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom timeit import default_timer as timer\nimport torch.nn.functional as F\n\n\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\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\nBASE_COOR = None \nOTHER_COOR = None\nSIDE_COOR = None\n# -----------------------------------------------------------------------------\ndef parse_args():\n    p = argparse.ArgumentParser(description=\"DRfold2 direct inference\")\n    p.add_argument('--mode',        choices=['train','test'], default='train')\n    p.add_argument('--train_path',  type=str,   default='train_sequences.csv')\n    p.add_argument('--test_path',   type=str,   default='test_sequences.csv')\n    p.add_argument('--output_path', type=str,   required=True,\n                   help=\"Where to write pdbs/ and summary\")\n    p.add_argument('--output_csv',  type=str,   required=True,\n                   help=\"Final submission CSV\")\n    p.add_argument('--model_dir',   type=str,   required=True,\n                   help=\"Root of drfold‐model‐bf16 (contains model_hub/ etc.)\")\n    p.add_argument('--cfg_name',    type=str,   default='cfg_97',\n                   help=\"Configuration subdirectory under model_hub (e.g. cfg_97)\")\n    p.add_argument('--start_id',    type=str,   default=None,\n                   help=\"If set, only process IDs greater than this\")\n    p.add_argument('--prec',type=str,   default='bf16')\n    p.add_argument('--cfg_merge',type=bool,   default=False)\n    p.add_argument('--max_cat_length',type=int,default=2400)\n    p.add_argument('--optim_length',type=int,default=0)\n    p.add_argument('--max_length',type=int,default=480)\n    p.add_argument('--no_sort',type=bool,   default=False)\n    p.add_argument('--full_energy',type=bool,   default=False)\n    p.add_argument('--dr_score',type=bool,   default=False)\n    p.add_argument('--op_score',type=bool,   default=False)\n    p.add_argument('--get_center',type=bool,   default=False)\n    \n    return p.parse_args()\n\n# -----------------------------------------------------------------------------\ndef write_c1p_pdb(coords, seq, pdb_path):\n    \"\"\"Write minimal PDB containing only C1' atoms + TER/END.\"\"\"\n    os.makedirs(os.path.dirname(pdb_path), exist_ok=True)\n    with open(pdb_path, 'w') as fh:\n        fh.write(\"REMARK   Generated by drfold2_direct.py\\n\")\n        atom_idx = 1\n        for i,(x,y,z) in enumerate(coords, start=1):\n            res = seq[i-1]\n            fh.write(\n                f\"ATOM  {atom_idx:5d}  C1' {res:>3s} A{i:4d}\"\n                f\"{x:8.3f}{y:8.3f}{z:8.3f}\"\n                f\"{1.00:6.2f}{0.00:6.2f}           C\\n\"\n            )\n            atom_idx += 1\n        fh.write(f\"TER   {atom_idx:5d}      {'':3s}A{i:4d}\\nEND\\n\")\n\n# -----------------------------------------------------------------------------\ndef convert_c1_coords(pdb_path):\n    \"\"\"Extract C1' coords for submission.csv.\"\"\"\n    xs, ys, zs = [], [], []\n    with open(pdb_path) as fh:\n        for L in fh:\n            parts = L.split()\n            if parts[0]==\"ATOM\" and parts[2]==\"C1'\":\n                xs.append(float(parts[6])); ys.append(float(parts[7])); zs.append(float(parts[8]))\n    return xs, ys, zs\n\ndef add_model_paths(args):\n    global BASE_COOR, OTHER_COOR,SIDE_COOR\n    BASE_COOR = np.load(f'{args.model_dir}/PotentialFold/lib/base.npy')\n    OTHER_COOR = np.load(f'{args.model_dir}/PotentialFold/lib/other2.npy')\n    SIDE_COOR = np.load(f'{args.model_dir}/PotentialFold/lib/side.npy')\n    paths = [\n        args.model_dir,\n        os.path.join(args.model_dir, 'PotentialFold'),\n        os.path.join(args.model_dir, args.cfg_name),\n        os.path.join(args.model_dir, args.cfg_name, 'RNALM2'),\n    ]\n    for p in paths:\n        if not os.path.isdir(p):\n            raise FileNotFoundError(f\"Expected directory not found: {p!r}\")\n        # Prepend so your local copy wins\n        sys.path.append(p)\n    from Optimization import Structure\n    import a2b\n\ndef frame_coor_to_C1(coor, seq, BASE_COOR, OTHER_COOR):\n    \n    import a2b\n    \"\"\"\n    使用四元数批量运算将模型输出的局部 frame 坐标直接转换为每个位点 C1' 原子全局坐标。\n\n    参数\n    ----\n    coor : numpy.ndarray 或 torch.Tensor, shape=(L, 3, 3)\n        模型预测的每个位点 P, C4', N 三点的局部坐标\n    seq : str, length=L\n        核苷酸序列\n    BASE_COOR : dict\n        Get_base 所需的三个骨架原子模板坐标常量\n    OTHER_COOR : dict\n        Get_base 所需的五个糖环原子模板坐标常量\n\n    返回\n    ----\n    c1_xyz : numpy.ndarray, shape=(L, 3)\n        每个位点 C1' 原子在全局坐标系下的 XYZ\n    \"\"\"\n    # 1. 转为 torch.Tensor\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    # 2. 初始化参数向量 x：21 维，其中前 18 维用于旋转投影，后 3 维用于平移\n    x = torch.rand((L, 21), dtype=tx.dtype, device=tx.device)\n\n    # 3. 中心化处理：计算质心 biasq，并得到零均值坐标 q\n    biasq = tx.mean(dim=1, keepdim=True)            # (L, 1, 3)\n    q = tx - biasq                                  # (L, 3, 3)\n\n    # 4. 计算旋转投影 m，并填充到 x 的前两组 9 维\n    m = torch.einsum('bnz,bny->bzy', basex, q).reshape(L, -1)\n    x[:, :9] = m\n    x[:, 9:18] = m\n\n    # 5. 填入平移参数（将 biasq 作为平移量）\n    x[:, 18:] = biasq.squeeze(1)\n    rama = x.double()  # 四元数参数需要 double 精度\n\n    # 6. 调用 quat2b：对模板中所有糖环原子做旋转+平移\n    #    otherx: (L, 5, 3) -> other_xyz: (L, 5, 3)\n    other_xyz = a2b.quat2b(otherx.double(), rama[:, 9:]).float().cpu().numpy()\n\n    # 7. 提取第 5 个原子（C1'）并返回\n    c1_xyz = other_xyz[:, 4, :]\n    return c1_xyz\n\n\ndef concat_coor(out1: dict, out2: dict) -> np.ndarray:\n    \n    import a2b\n    \"\"\"\n    将 out2 通过四元数（quat2b）做旋转 + 平移对齐到 out1 末尾，并拼接（去掉重复的第一个碱基）。\n    out1/out2 都是 dict，包含\n      'coor': np.ndarray, shape=(L, 3, 3)  —— 每个碱基的 P, C4, N 坐标\n    返回拼接后的坐标数组，shape = (L1 + L2 - 1, 3, 3)\n    \"\"\"\n    # 1. 转为 torch.DoubleTensor\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    # 2. 取公共碱基的 frame：out1 最后一个、out2 第一个\n    f1 = coor1[-1]   # (3,3)\n    f2 = coor2[0]    # (3,3)\n\n    # 3. 中心化：去除各自质心\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    # 4. 用 einsum 计算旋转矩阵 R = basex · q^T\n    #    R_{ij} = sum_z basex_{iz} * q_{jz}\n    R = torch.einsum('iz,jz->ij', basex, q)   # (3,3)\n\n    # 5. 计算平移向量 t，使得 R·bias2 + t = bias1  ⇒  t = bias1 - R·bias2\n    t = bias1 - (R @ bias2)                  # (3,)\n\n    # 6. 构造 rama 参数 (L2, 12)：前 9 维放 R.flatten，后 3 维放 t\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    # 7. 调用 quat2b：对 out2 整条链做刚性变换\n    #    输入 coord=(L2,3,3)， rama=(L2,12)，输出同 shape\n    coor2_aligned = a2b.quat2b(coor2, rama)   # torch.Tensor (L2,3,3)\n\n    # 8. 拼接（跳过 coor2_aligned[0]，因为它已与 coor1[-1] 重合）\n    coor_cat = torch.cat([coor1, coor2_aligned[1:]], dim=0)  # (L1+L2-1,3,3)\n\n    return coor_cat.cpu().numpy()\n\ndef 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\ndef make_data(seq, device):\n    from data import parse_seq\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\ndef score_energy_one_simple(seq, target_id, out):\n    \"\"\"\n    计算 RNA 结构的粗粒化自由能，只基于每个碱基的 P–C4'–N (N1 or N9) 坐标帧。\n\n    参数：\n      seq:      RNA 序列字符串，长度 L\n      target_id:结构标识符（仅用于日志或输出，不影响计算）\n      out:      模型输出字典，必须包含：\n                out['coor']: numpy array, shape (L, 3, 3)\n                             coor[i,0] = P 坐标\n                             coor[i,1] = C4′ 坐标\n                             coor[i,2] = N1/N9 坐标\n\n    返回：\n      total_energy: float，自由能估计值（越低表示越稳定）\n    \"\"\"\n\n    coor = out['coor']  # (L, 3, 3)\n    L = len(coor)\n\n    # —— 一、键长能 —— #\n    # 理想键长 (Å)：P–C4′ ~ 1.60, C4′–N ~ 1.47 (典型值，可根据经验或晶体数据调整)\n    d0_P_C4  = 1.60\n    d0_C4_N  = 1.47\n    k_bond   = 100.0  # 键长力常数 (kcal/mol/Ų)\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    # —— 二、键角能 —— #\n    #  ideal angle P–C4′–N ~ 109.5°（四面体角），用简谐势\n    theta0   = np.deg2rad(109.5)\n    k_angle  = 20.0   # 键角力常数 (kcal/mol/rad²)\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    # —— 三、堆积相互作用 —— #\n    # 简单地对相邻碱基 C4′–C4′ 间距进行势能惩罚，鼓励相邻堆积：\n    #   ideal 平面内距离 ~ 3.4 Å\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    # —— 四、合计 —— #\n    total_energy = energy_bond + energy_angle + energy_stack\n\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\nimport numpy as np\n\ndef score_energy_one_full(seq, target_id, out, paired=None,\n                               # stacking parameters (optional if needed)\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                               # implicit solvation (ASA) coefficient\n                               k_solv=0.05,\n                               # Debye-Hückel electrostatics\n                               q_P=-1.0, epsilon_r=80.0, kappa=10.0,\n                               k_e=332.0637,\n                               # statistical potential map: dict of (i,j)->energy\n                               stat_potentials=None,\n                               # numerical stability\n                               r_min=1.0):\n    \"\"\"\n    以 TM-score 最大化为目的的粗粒化能量函数，包含：\n      - 主链二面角 (backbone torsion)\n      - 氢键 (hydrogen bonding)\n      - 范德华 + 隐式溶剂 (LJ + solvation)\n      - 静电相互作用 (electrostatics)\n      - 知识库统计势 (knowledge-based potentials)\n    已移除键长和键角项。\n    \"\"\"\n    coor = out['coor']\n    L = len(coor)\n    if paired is None:\n        paired = out.get('paired', [])\n\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, b, c = coor[i,0], coor[i,1], coor[i,2]\n        d = coor[i+1,0]\n        phi = torsion_angle(a, b, c, d)\n        E_dihedral += k_dihedral * (phi - phi0)**2\n\n    # —— 二、氢键势 —— #\n    E_hb = 0.0\n    for i, j in paired:\n        r = np.linalg.norm(coor[i,2] - coor[j,2])\n        E_hb += k_hb * (r - d0_hb)**2\n\n    # —— 三、范德华 + 隐式溶剂 —— #\n    # E_LJ = 0.0\n    # E_solv = 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    #         r_clamped = max(r, r_min)\n    #         sr6 = (sigma / r_clamped)**6\n    #         sr12 = sr6 * sr6\n            # E_LJ += 4 * epsilon * (sr12 - sr6)\n            # 隐式溶剂: 与 ASA 成正比, 这里简化使用 r^-2\n            # E_solv += k_solv / (r_clamped**2)\n\n    # —— 四、静电（Debye-Hückel）—— #\n    E_elec = 0.0\n    prefac = k_e * q_P * q_P / epsilon_r\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            r_clamped = max(r, r_min)\n            E_elec += prefac * np.exp(-r_clamped / kappa) / r_clamped\n\n    # —— 五、知识库统计势 —— #\n    E_stat = 0.0\n    if stat_potentials:\n        for (i, j), Eij in stat_potentials.items():\n            E_stat += Eij\n\n    # total_energy = E_dihedral + E_hb + E_LJ + E_solv + E_elec + E_stat\n    total_energy = E_dihedral + E_hb + E_elec + E_stat\n\n    # print(f\"[{target_id}] torsion={E_dihedral:.2f}, hb={E_hb:.2f}, LJ={E_LJ:.2f}, \"\n    #       f\"solv={E_solv:.2f}, elec={E_elec:.2f}, stat={E_stat:.2f} -> total={total_energy:.2f}\")\n    # print(f\"[{target_id}] torsion={E_dihedral:.2f}, hb={E_hb:.2f}, \"\n    #       f\"elec={E_elec:.2f}, stat={E_stat:.2f} -> total={total_energy:.2f}\")\n    \n    return -total_energy\nimport pyrosetta\nfrom pyrosetta import pose_from_pdb, create_score_function\nfrom pyrosetta.rosetta import core\n\n# 1. 初始化 PyRosetta（只需执行一次）\npyrosetta.init(options='-mute all')  # 静默模式，不输出多余日志\n# 3. 创建并定制 ScoreFunction\nop_score = create_score_function('ref2015')\nop_score.set_weight(core.scoring.atom_pair_constraint, 9.0)\nop_score.set_weight(core.scoring.dihedral_constraint, 4.0)\nop_score.set_weight(core.scoring.angle_constraint, 4.0)\nop_score.set_weight(core.scoring.fa_rep, 9.0)\nop_score.set_weight(core.scoring.rna_sugar_close, 9.0)\nop_score.set_weight(core.scoring.fa_intra_rep, 9.0)\nop_score.set_weight(core.scoring.rna_base_pair, 9.0)\nop_score.set_weight(core.scoring.rna_base_stack, 9.0)\nimport os\nimport json\nimport pickle\nfrom subprocess import PIPE, STDOUT, Popen\nimport numpy as np\nimport tempfile\nfrom Bio.PDB import PDBParser\n\ndef score_energy_one(seq, target_id, out):\n    \n    from Optimization import Structure\n    with tempfile.TemporaryDirectory() as tmpdirname:\n        # 输出临时fasta\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        # 输出临时ret\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    \ndef score_energy_one_op(seq, target_id, out):\n    from Optimization import Structure\n    with tempfile.TemporaryDirectory() as tmpdirname:\n        # 输出临时fasta\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        # 输出临时ret\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        \n        cgpdb = os.path.join(tmpdirname, 'tmp.pdb')\n        # print(rama)\n        newrama=torch.DoubleTensor(rama) \n        stru.outpdb(newrama, cgpdb)\n        \n        savepdb = os.path.join(tmpdirname, 'tmp_full.pdb')\n        \n        exp_dir = os.getcwd()\n        arena = os.path.join(exp_dir, 'Arena', 'Arena')\n        cmd = f'{arena} {cgpdb} {savepdb} 7'\n        p = Popen(cmd, shell=True, stdin=PIPE, stdout=PIPE, stderr=STDOUT)\n        output, error = p.communicate()\n        pose= pose_from_pdb(savepdb)\n        \n        score = op_score(pose)\n        return score\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        # 补全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    \n    from Optimization import Structure\n    with tempfile.TemporaryDirectory() as tmpdirname:\n        # 1) 写入临时 fasta\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        # 2) 写入临时 ret\n        retfile = os.path.join(tmpdirname, 'tmp.ret')\n        with open(retfile, 'wb') as f:\n            pickle.dump(out, f)\n\n        # 3) 调用 Structure 完成折叠优化\n        foldconfig = '/kaggle/input/drfold-model/DRfold2/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        # 4) 解析输出的 PDB\n        pdb_file = save_prefix + '.pdb'\n        parser = PDBParser(QUIET=True)\n        structure = parser.get_structure(target_id, pdb_file)\n\n        # 5) 收集所有标准残基（排除 HETATM）\n        residues = [\n            res for res in structure.get_residues()\n            if res.id[0] == ' '\n        ]\n        # 检查残基数和序列长度一致\n        L = len(seq)\n        if len(residues) != L:\n            raise ValueError(f\"PDB 中残基数 ({len(residues)}) 与序列长度 ({L}) 不一致\")\n\n        # 6) 提取三种原子坐标\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  # 初始化为 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                # 缺失时填 NaN 方便排查\n                coord = res[\"C4'\"].get_vector().get_array()\n                coor[i, 1, :] = coord\n            if 'N1' in res:\n            # 缺失时填 NaN 方便排查\n                coord = res['N1'].get_vector().get_array()\n                coor[i, 2, :] = coord\n            elif 'N9' in res:\n                # 缺失时填 NaN 方便排查\n                coord = res['N9'].get_vector().get_array()\n                coor[i, 2, :] = coord\n\n        return coor\n    \n\ndef write_c1p_pdb(coords: np.ndarray, sequence: str, out_path: str):\n    \"\"\"Emit a minimal PDB with only C1' atoms.\"\"\"\n    os.makedirs(os.path.dirname(out_path), exist_ok=True)\n    with open(out_path, 'w') as f:\n        for i, (x,y,z) in enumerate(coords, start=1):\n            res = sequence[i-1]\n            f.write(\n                f\"ATOM  {i:5d}  C1' {res:>3s} A{i:4d}    \"\n                f\"{x:8.3f}{y:8.3f}{z:8.3f}  1.00  0.00           C\\n\"\n            )\n        f.write(\"END\\n\")\n        \nfrom torch.cuda.amp import autocast\ndef run_submit(df, args):\n    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n    # 1) ensure all the code dirs are on PYTHONPATH\n    add_model_paths(args)\n    saved_argv = sys.argv\n    print(sys.argv[0])\n    sys.argv    = [sys.argv[0],device,'bf16']\n\n    # 2) Import model classes\n    from EvoMSA2XYZ import MSA2XYZ\n    from RNALM2.Model import RNA2nd\n    from data import parse_seq\n    \n    # sys.argv = saved_argv\n\n    r2d_ckpt = os.path.join(args.model_dir, 'model_hub', 'RCLM', 'epoch_67000')\n    rnalm = RNA2nd(dict(s_in_dim=5, z_in_dim=2, s_dim=512, z_dim=128, N_elayers=18))\n    rnalm.load_state_dict(torch.load(r2d_ckpt, map_location='cpu', weights_only=True), strict=False)\n    rnalm = rnalm.to(device).eval()\n\n    if args.prec == 'fp16': rnalm = rnalm.half()\n    if args.prec == 'bf16': rnalm = rnalm.bfloat16()\n\n    # Load MSA2XYZ ensemble\n    if args.cfg_merge:\n        cfgs = ['cfg_97', 'cfg_95', 'cfg_96', 'cfg_99']\n        msa2xyz_file = [\n            os.path.join(args.model_dir, 'model_hub', cfg, f'model_{i}')\n            for cfg in cfgs for i in range(20)\n        ]\n    else:\n        msa2xyz_file = [\n            os.path.join(args.model_dir, 'model_hub', args.cfg_name, f'model_{i}')\n            for i in range(20)\n        ]\n\n    msa2xyz = MSA2XYZ(seq_dim=6, msa_dim=7, N_ensemble=1, N_cycle=8, m_dim=64, s_dim=64, z_dim=64)\n    assert msa2xyz.msaxyzone.evmodel.evos[0].msa_row.qlinear.linear.weight.shape[0] == 64, \\\n        f\"m_dim wrong: got {msa2xyz.msaxyzone.evmodel.evos[0].msa_row.qlinear.linear.weight.shape[0]}\"\n\n    msa2xyz_state_dict = [torch.load(f, map_location='cpu', weights_only=True) for f in msa2xyz_file]\n    msa2xyz.load_state_dict(msa2xyz_state_dict[0], strict=False)\n    msa2xyz = msa2xyz.to(device).eval()\n\n    if args.prec == 'fp16': msa2xyz = msa2xyz.half()\n    if args.prec == 'bf16': msa2xyz = msa2xyz.bfloat16()\n\n    msa2xyz.msaxyzone.premsa.rnalm = rnalm\n\n    rows = []\n    total_time_taken = 0\n    max_gpu_mem_used = 0\n\n    for i, row in df.iterrows():\n        start_timer = timer()\n        target_id = row.target_id\n        sequence = row.sequence\n        seq = sequence[:args.max_cat_length] if len(sequence) > args.max_cat_length else sequence\n        L = len(seq)\n\n        print(i, target_id, L, len(seq), seq[:75] + '...')\n\n        if L > 480:\n            model_to_try = [16, 9, 1, 2, 0]\n        elif L > 200:\n            model_to_try = [13, 6, 14, 5, 3]\n        elif L > 100:\n            model_to_try = [13, 6, 14, 12, 7, 2, 5, 19, 10, 9] if not args.cfg_merge else [24, 34, 20, 13, 6, 37, 28, 25, 14, 39]\n        else:\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 args.no_sort:\n            model_to_try = model_to_try[:5]\n\n        def predict_segment(seq):\n            msa, base_x, seq_idx = make_data(seq, device)\n            if args.prec == 'fp16': msa, base_x = msa.half(), base_x.half()\n            if args.prec == 'bf16': msa, base_x = msa.bfloat16(), base_x.bfloat16()\n            with torch.no_grad():\n                with autocast(dtype=torch.bfloat16):\n                    return msa2xyz.pred(msa, seq_idx, None, base_x, np.array(list(seq)))\n\n        energy = []\n        coordinate = []\n        outputs = []\n\n        for c in model_to_try:\n            msa2xyz.load_state_dict(msa2xyz_state_dict[c], strict=False)\n            if L <= args.max_length:\n                outs = [predict_segment(seq)]\n            else:\n                step = args.max_length - 1\n                outs = [predict_segment(seq[s:s+args.max_length]) for s in range(0, L, step)]\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 args.no_sort:\n                e = 0\n            elif len(model_to_try) > 5 and args.dr_score:\n                e = score_energy_one(seq, target_id, out_cat)\n            elif len(model_to_try) > 5 and args.op_score:\n                e = score_energy_one_op(seq, target_id, out_cat)\n            elif args.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\n            energy.append(e)\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            xyz = frame_coor_to_C1(out_cat['coor'], sequence, BASE_COOR, OTHER_COOR)\n            coordinate.append(xyz)\n\n            time_taken = timer() - start_timer\n            total_time_taken += time_taken\n            max_gpu_mem_used = max(max_gpu_mem_used, gpu_memory_use())\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_memory_use()} gb')\n\n        torch.cuda.empty_cache()\n\n        if args.get_center:\n            energy = np.abs(np.array(energy) - np.mean(energy))\n\n        argsort = np.array(energy).argsort()[:8]\n\n        if L <= args.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        pdb_dir = os.path.join(args.output_path, \"pdbs\")\n        for rank, model_idx in enumerate(argsort, start=1):\n            coords = coordinate[model_idx]  # this is an (L,3) np array\n            out_pdb = os.path.join(pdb_dir, f\"{target_id}_model_{rank}.pdb\")\n            write_c1p_pdb(coords, sequence, out_pdb)\n            rows.append({\n                \"target_id\": target_id,\n                \"model_no\":   rank,\n                \"predicted_model_path\": out_pdb\n            })\n\n        # clean up GPU mem, etc.\n        torch.cuda.empty_cache()\n\n    # build and return a dataframe listing all written PDBs\n    summary_df = pd.DataFrame(rows)\n    return summary_df\n\ndef main():\n    args = parse_args()\n    df = pd.read_csv(args.test_path if args.mode == 'test' else args.train_path)\n    os.makedirs(args.output_path, exist_ok=True)\n    submission_df = run_submit(df, args)\n    summary_csv = os.path.join(args.output_path, 'drfold2_summary.csv')\n    submission_df.to_csv(summary_csv)\n    print(f\"[DONE] submission → {args.output_csv}\")\n    print(f\"[DONE] summary    → {summary_csv}\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-05-29T06:51:19.677015Z","iopub.execute_input":"2025-05-29T06:51:19.677426Z","iopub.status.idle":"2025-05-29T06:51:19.689445Z","shell.execute_reply.started":"2025-05-29T06:51:19.677398Z","shell.execute_reply":"2025-05-29T06:51:19.688469Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\n%cd /kaggle/working/drfold\n!torchrun --standalone --nproc_per_node=1 --master_port=29501 predict.py\\\n        --train_path /kaggle/working/test_sequence.csv\\\n        --test_path /kaggle/working/test_sequence.csv\\\n        --model_dir /kaggle/working/drfold/drfold-model-bf16\\\n        --cfg_name cfg_97\\\n        --output_path /kaggle/working/drfold_output\\\n        --output_csv drfold_output.csv \n%cd /kaggle/working","metadata":{"trusted":true,"scrolled":true,"execution":{"iopub.status.busy":"2025-05-29T06:51:19.690545Z","iopub.execute_input":"2025-05-29T06:51:19.690835Z","iopub.status.idle":"2025-05-29T06:52:13.212887Z","shell.execute_reply.started":"2025-05-29T06:51:19.690807Z","shell.execute_reply":"2025-05-29T06:52:13.212026Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %%time\n# %cd /kaggle/working/drfold\n# !torchrun --standalone --nproc_per_node=1 --master_port=29501 predict.py\\\n#         --train_path /kaggle/working/test_sequence.csv\\\n#         --test_path /kaggle/working/test_sequence.csv\\\n#         --model_dir /kaggle/working/drfold/drfold-model-bf16\\\n#         --cfg_name cfg_99\\\n#         --output_path /kaggle/working/drfold_output1\\\n#         --output_csv drfold_output1.csv \n# %cd /kaggle/working","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:52:13.213888Z","iopub.execute_input":"2025-05-29T06:52:13.214204Z","iopub.status.idle":"2025-05-29T06:52:13.217827Z","shell.execute_reply.started":"2025-05-29T06:52:13.214171Z","shell.execute_reply":"2025-05-29T06:52:13.216995Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# process","metadata":{}},{"cell_type":"code","source":"import os\nimport pandas as pd\n\ndef load_df(path, fallback=None):\n    if os.path.exists(path):\n        return pd.read_csv(path)\n    elif fallback is not None:\n        return fallback\n    else:\n        print(f\"[WARN] File not found: {path}\")\n        return pd.DataFrame()\n\n# Load individual model predictions\ndf1 = load_df('/kaggle/working/proteinx_predictions/proteinx_preds.csv')\ndf1.to_csv('proteinx_preds.csv', index=False)\n\ndf2 = load_df('/kaggle/working/boltz_predictions/boltz_preds_summary.csv')\ndf2.to_csv('boltz_preds.csv', index=False)\n\ndf3 = load_df('/kaggle/working/rhofold_predictions/rhofold_preds.csv')\nif not df3.empty:\n    df3['sample'] = 0\n    df3.to_csv('rhofold_preds.csv', index=False)\n\ndf4 = load_df('/kaggle/working/nufold_predictions/nufold_full/nufold_full_summary.csv')\nif not df4.empty:\n    df4['sample'] = df4['variant']\n    df4.to_csv('nufold_preds.csv', index=False)\n\ndf5 = load_df('/kaggle/working/deepfold_predictions/deepfold_preds.csv')\nif not df5.empty:\n    df5['sample'] = df5['predicted_model_path'].apply(lambda x: int(x.split('seed_')[1].split('_sample')[0]))\n    df5.to_csv('deepfold_preds.csv', index=False)\n\ndf6 = load_df('/kaggle/working/drfold_output/drfold2_summary.csv')\nif not df6.empty:\n    df6['sample'] = df6['model_no']\n    df6.to_csv('drfold_preds.csv', index=False)\n\ndf7 = load_df('/kaggle/working/drfold_output1/drfold2_summary.csv')\nif not df7.empty:\n    df7['sample'] = df7['model_no']\n    df7.to_csv('drfold_preds.csv', index=False)\n\n# Combine all available predictions\noutput_path = '/kaggle/working/all_preds.csv'\ndf_all = pd.concat([df for df in [df1, df2, df4, df5, df6, df7] if not df.empty], axis=0)\ndf_all.to_csv(output_path, index=False)\nprint(f\"[✓] Combined predictions written to {output_path}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:52:13.221022Z","iopub.execute_input":"2025-05-29T06:52:13.221223Z","iopub.status.idle":"2025-05-29T06:52:13.256803Z","shell.execute_reply.started":"2025-05-29T06:52:13.221205Z","shell.execute_reply":"2025-05-29T06:52:13.256163Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# modelling + submit","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\ndef normalize_rotation(coords):\n    \"\"\"\n    Makes coordinates invariant to translation and rotation, but not scale.\n\n    Args:\n        coords: (B, L, 3) - batch of coordinate tensors\n\n    Returns:\n        coords_aligned: (B, L, 3) - rotation-invariant coordinates\n    \"\"\"\n    B, L, _ = coords.shape\n\n    # 1. Subtract centroid (translation invariance)\n    center = coords.mean(dim=1, keepdim=True)  # (B, 1, 3)\n    coords_centered = coords - center\n\n    # 2. Align using PCA (per batch item)\n    coords_aligned = []\n    for i in range(B):\n        x = coords_centered[i]  # (L, 3)\n        cov = x.T @ x           # (3, 3)\n        U, _, _ = torch.svd(cov.float())\n        x_aligned = x @ U       # rotate to principal axes\n        coords_aligned.append(x_aligned)\n\n    return torch.stack(coords_aligned, dim=0)  # (B, L, 3)\n\n\n\nclass SequenceEncoder(nn.Module):\n    def __init__(self,d_model=128,nhead=8,layers=4,max_len=1000,rel=32):\n        super().__init__()\n        self.embed = nn.Embedding(4,d_model)\n        self.pos   = nn.Parameter(torch.randn(max_len,d_model))\n        self.rel_k = rel\n        self.rel_e = nn.Embedding(2*rel+1,d_model)\n        layer = nn.TransformerEncoderLayer(d_model,nhead,4*d_model,0.1,batch_first=True)\n        self.trans=nn.TransformerEncoder(layer,layers)\n    def forward(self,seq):\n        B,L=seq.shape\n        x=self.embed(seq)+self.pos[:L]\n        idx=torch.arange(L,device=seq.device)\n        r=(idx.unsqueeze(0)-idx.unsqueeze(1)).clamp(-self.rel_k,self.rel_k)+self.rel_k\n        rp=self.rel_e(r)\n        x=x+rp.sum(1)/L\n        return self.trans(x).mean(1)\n\ndef knn_graph(x, k):\n    \"\"\"\n    Build a k-NN graph on coordinates x (B, N, 3).\n    Returns: idx (B, N, k) of neighbor indices for each node.\n    \"\"\"\n    # Compute pairwise distances\n    dist = torch.cdist(x, x)  # (B, N, N)\n    # Exclude self-connections by setting diagonal to large\n    B, N, _ = dist.shape\n    diag = torch.arange(N, device=x.device)\n    dist[:, diag, diag] = float('inf')\n    # Select k smallest distances (largest negative)\n    _, idx = torch.topk(-dist, k, dim=-1)  # (B, N, k)\n    return idx\n\n\nclass EfficientEGNNLayer(nn.Module):\n    \"\"\"Memory-efficient EGNN layer using k-NN sparsity.\"\"\"\n    def __init__(self, dim, k=16):\n        super().__init__()\n        self.k = k\n        self.edge_mlp = nn.Sequential(\n            nn.Linear(2 * dim + 1, dim),\n            nn.SiLU(),\n            nn.Linear(dim, 1)\n        )\n        self.node_mlp = nn.Sequential(\n            nn.Linear(dim, dim),\n            nn.SiLU(),\n            nn.Linear(dim, dim)\n        )\n\n    def forward(self, h, x):\n        B, N, F = h.shape\n        # Get k-NN indices\n        idx = knn_graph(x, self.k)  # (B, N, k)\n\n        # Gather neighbor features\n        idx_feat = idx.unsqueeze(-1).expand(-1, -1, -1, F)  # (B, N, k, F)\n        h_j = torch.gather(h.unsqueeze(1).expand(B, N, N, F), 2, idx_feat)\n        h_i = h.unsqueeze(2).expand(-1, -1, self.k, -1)\n\n        # Gather neighbor coord diffs\n        idx_coord = idx.unsqueeze(-1).expand(-1, -1, -1, 3)  # (B, N, k, 3)\n        x_j = torch.gather(x.unsqueeze(1).expand(B, N, N, 3), 2, idx_coord)\n        x_i = x.unsqueeze(2).expand(-1, -1, self.k, -1)\n        dx = x_i - x_j  # (B, N, k, 3)\n        dist2 = (dx ** 2).sum(dim=-1, keepdim=True)  # (B, N, k, 1)\n\n        # Edge MLP and attention\n        edge_input = torch.cat([h_i, h_j, dist2], dim=-1)\n        e_ij = self.edge_mlp(edge_input).squeeze(-1)  # (B, N, k)\n        a_ij = torch.softmax(e_ij, dim=-1)\n\n        # Message passing\n        m = (a_ij.unsqueeze(-1) * h_j).sum(dim=2)  # (B, N, F)\n        h_out = h + self.node_mlp(m)\n\n        # Coordinate update\n        x_out = x + (a_ij.unsqueeze(-1) * dx).sum(dim=2)\n        return h_out, x_out\n\n\nclass EfficientStructureEncoder(nn.Module):\n    \"\"\"\n    Encoder with the memory-efficient EGNN layers.\n    Drops full N^2 attention in favor of k-NN sparsity.\n    \"\"\"\n    def __init__(self, d_model=128, num_layers=3, k=64):\n        super().__init__()\n        self.input_proj = nn.Linear(3, d_model)\n        self.layers = nn.ModuleList([\n            EfficientEGNNLayer(d_model, k=k) for _ in range(num_layers)\n        ])\n\n    def forward(self, coords):\n        coords = normalize_rotation(coords)\n        h = self.input_proj(coords)\n        x = coords\n        for layer in self.layers:\n            h, x = layer(h, x)\n        return h.mean(dim=1)\n\n\nclass FusionHead(nn.Module):\n    \"\"\"Combines sequence and structure representations\"\"\"\n    def __init__(self, d_model=128):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(2 * d_model, d_model),\n            nn.ReLU(),\n            nn.Linear(d_model, d_model // 2),\n            nn.ReLU(),\n            nn.Linear(d_model // 2, 1)\n        )\n\n    def forward(self, seq_emb, struct_emb):\n        x = torch.cat([seq_emb, struct_emb], dim=-1)\n        return self.net(x).squeeze(1)\n\n\nclass RNAQualityModel(nn.Module):\n    \"\"\"End-to-end model for RNA decoy quality scoring\"\"\"\n    def __init__(self, d_model=128):\n        super().__init__()\n        self.seq = SequenceEncoder(d_model)\n        self.str = EfficientStructureEncoder(d_model)\n        self.head = FusionHead(d_model)\n\n    def forward(self, seq_idx, coords):\n        seq_emb = self.seq(seq_idx)     # (B, d)\n        struct_emb = self.str(coords)   # (B, d)\n        return self.head(seq_emb, struct_emb)\n\n\nclass ActorNet(nn.Module):\n    def __init__(self, d_model):\n        super().__init__()\n        self.seq_enc = SequenceEncoder(d_model=d_model)\n        self.str_enc = EfficientStructureEncoder(d_model=d_model)\n        self.combine = nn.Linear(2 * d_model, d_model)\n\n        self.attn = nn.TransformerEncoderLayer(d_model=d_model, nhead=4, batch_first=True)\n        self.out_proj = nn.Linear(d_model, 1)\n\n    def forward(self, seq_idx, coords):\n        # coords: [B, C, L, 3]\n        B, C, L, _ = coords.shape\n        device = coords.device\n\n        # Encode sequence once per batch\n        seq_emb = self.seq_enc(seq_idx)  # [B, d_model]\n\n        # Encode each candidate structure one at a time to reduce peak memory\n        str_emb_list = []\n        for c in range(C):\n            coords_c = coords[:, c]            # [B, L, 3]\n            emb_c = self.str_enc(coords_c)     # [B, d_model]\n            str_emb_list.append(emb_c)\n        # Stack back to [B, C, d_model]\n        str_emb = torch.stack(str_emb_list, dim=1)\n\n        # Repeat sequence embeddings across candidates\n        seq_emb = seq_emb.unsqueeze(1).expand(-1, C, -1)  # [B, C, d_model]\n\n        # Combine and attend across candidates\n        x = torch.cat([seq_emb, str_emb], dim=-1)  # [B, C, 2*d_model]\n        x = self.combine(x)                        # [B, C, d_model]\n        x = self.attn(x)                           # [B, C, d_model]\n\n        logits = self.out_proj(x).squeeze(-1)      # [B, C]\n        return logits\n\n\nclass CriticNet(nn.Module):\n    def __init__(self, d_model):\n        super().__init__()\n        self.seq_enc = SequenceEncoder(d_model=d_model)\n        self.str_enc = EfficientStructureEncoder(d_model=d_model)\n        self.combine = nn.Linear(2 * d_model, d_model)\n\n        self.attn = nn.TransformerEncoderLayer(d_model=d_model, nhead=4, batch_first=True)\n        self.pool = nn.AdaptiveAvgPool1d(1)\n\n        self.value_head = nn.Sequential(\n            nn.Linear(d_model, d_model // 2),\n            nn.ReLU(),\n            nn.Linear(d_model // 2, 1)\n        )\n\n    def forward(self, seq_idx, coords):\n        # coords: [B, C, L, 3]\n        B, C, L, _ = coords.shape\n        device = coords.device\n\n        # Sequence encoding\n        seq_emb = self.seq_enc(seq_idx)  # [B, d_model]\n\n        # Structure encoding per candidate\n        str_emb_list = []\n        for c in range(C):\n            coords_c = coords[:, c]           # [B, L, 3]\n            emb_c = self.str_enc(coords_c)    # [B, d_model]\n            str_emb_list.append(emb_c)\n        str_emb = torch.stack(str_emb_list, dim=1)  # [B, C, d_model]\n\n        # Combine\n        seq_exp = seq_emb.unsqueeze(1).expand(-1, C, -1)  # [B, C, d_model]\n        x = torch.cat([seq_exp, str_emb], dim=-1)\n        x = self.combine(x)  # [B, C, d_model]\n\n        # Cross-candidate attention\n        x = self.attn(x)  # [B, C, d_model]\n\n        # Pool across candidates\n        x = x.transpose(1, 2)  # [B, d_model, C]\n        pooled = self.pool(x).squeeze(-1)  # [B, d_model]\n\n        value = self.value_head(pooled).squeeze(-1)  # [B]\n        return value","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T06:52:13.257662Z","iopub.execute_input":"2025-05-29T06:52:13.257859Z","iopub.status.idle":"2025-05-29T06:52:13.281522Z","shell.execute_reply.started":"2025-05-29T06:52:13.25783Z","shell.execute_reply":"2025-05-29T06:52:13.280814Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport torch\nimport pandas as pd\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader, Dataset\nfrom Bio.PDB import PDBParser\nimport tempfile\n\nclass RNADecoyTestDataset(Dataset):\n    def __init__(self, df, max_len=500):\n        self.df = df\n        self.max_len = max_len\n        self.b2i = {'A': 0, 'U': 1, 'G': 2, 'C': 3}\n        self.parser = PDBParser(QUIET=True)\n        self.groups = df.groupby('target_id')\n        self.ids = list(self.groups.groups.keys())\n        self.seq_map = df.groupby('target_id')['sequence'].first().to_dict()\n\n    def __len__(self):\n        return len(self.ids)\n\n    def __getitem__(self, idx):\n        tid = self.ids[idx]\n        group = self.groups.get_group(tid)\n        seq = self.seq_map[tid]\n\n        # encode sequence\n        seq_idx = torch.tensor([self.b2i.get(b, 0) for b in seq], dtype=torch.long)\n        if len(seq_idx) > self.max_len:\n            seq_idx = seq_idx[:self.max_len]\n        else:\n            seq_idx = F.pad(seq_idx, (0, self.max_len - len(seq_idx)))\n\n        coords_list, paths = [], []\n        for _, row in group.iterrows():\n            pdb_path = row['predicted_model_path']\n            if not os.path.exists(pdb_path):\n                continue\n            # ensure PDB ends correctly\n            with open(pdb_path, 'r') as f:\n                text = f.read().rstrip()\n            if not text.endswith('END'):\n                text += '\\nEND\\n'\n            with tempfile.NamedTemporaryFile('w+', delete=False, suffix='.pdb') as tmp:\n                tmp.write(text)\n                tmp_path = tmp.name\n\n            struct = self.parser.get_structure(tid, tmp_path)\n            model_item = next(struct.get_models())\n            chain = next(model_item.get_chains())\n\n            coords = []\n            for res in chain:\n                if res.get_resname().strip() in self.b2i and \"C1'\" in res:\n                    coords.append(res[\"C1'\"].get_coord())\n            coords = torch.tensor(coords, dtype=torch.float32)\n            if coords.size(0) > self.max_len:\n                coords = coords[:self.max_len]\n            else:\n                coords = F.pad(coords, (0, 0, 0, self.max_len - coords.size(0)))\n\n            coords_list.append(coords)\n            paths.append(pdb_path)\n\n        if not coords_list:\n            coords_list = [torch.zeros((self.max_len, 3), dtype=torch.float32)]\n            paths = ['']\n\n        coords = torch.stack(coords_list)  # [C, L, 3]\n        return seq_idx, coords, tid, paths\n\n\ndef rank_and_submit_actor(\n    input_csv: str,\n    actor_path: str,\n    output_csv: str,\n    d_model: int = 128,\n    max_len: int = 500,\n    top_k: int = 5\n):\n    # load all decoys\n    df_meta = pd.read_csv(input_csv)\n    ds = RNADecoyTestDataset(df_meta, max_len=max_len)\n    dl = DataLoader(ds, batch_size=1)\n\n    # load model\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    model = ActorNet(d_model).to(device)\n    model.load_state_dict(torch.load(actor_path, map_location=\"cpu\"))\n    model.eval()\n\n    parser = PDBParser(QUIET=True)\n    rows = {}\n\n    with torch.no_grad():\n        for seq_idx, coords, tid, paths in dl:\n            target_id = tid[0]\n\n            # seq_idx: [1, L], coords: [1, C, L, 3]\n            seq_idx = seq_idx.to(device)\n            coords = coords.to(device)\n\n            # forward: logits shape [1, C]\n            logits = model(seq_idx, coords)\n            # take first batch\n            scores = logits[0].cpu().numpy()\n\n            # print(scores)\n            model_paths = paths\n            ranked = sorted(zip(scores, model_paths), key=lambda x: -x[0])[:top_k]\n            # print(ranked,model_paths)\n\n            for rank, (score, pdb_path) in enumerate(ranked, start=1):\n                pdb_path = pdb_path[0]\n                struct = parser.get_structure(target_id, pdb_path)\n                model_item = next(struct.get_models())\n                chain = next(model_item.get_chains())\n\n                for residue in chain:\n                    if \"C1'\" not in residue:\n                        continue\n                    resnum = residue.get_id()[1]\n                    resname = residue.get_resname().strip()\n                    ID = f\"{target_id}_{resnum}\"\n\n                    if ID not in rows:\n                        rows[ID] = {'ID': ID, 'resname': resname, 'resid': resnum}\n                        for i in range(1, top_k+1):\n                            rows[ID][f'x_{i}'] = float('nan')\n                            rows[ID][f'y_{i}'] = float('nan')\n                            rows[ID][f'z_{i}'] = float('nan')\n\n                    x, y, z = residue[\"C1'\"].get_coord()\n                    rows[ID][f'x_{rank}'] = float(x)\n                    rows[ID][f'y_{rank}'] = float(y)\n                    rows[ID][f'z_{rank}'] = float(z)\n\n    # build and save\n    df_sub = pd.DataFrame.from_records(list(rows.values()))\n    cols = ['ID', 'resname', 'resid'] + [f\"{c}_{i}\" for i in range(1, top_k+1) for c in ('x','y','z')]\n    df_sub = df_sub[cols]\n    df_sub.to_csv(output_csv, index=False)\n    print(f\"[✓] Submission saved to {output_csv}\")\n# Usage:\n# rank_and_submit_point(\n#     input_csv=\"all_preds.csv\",\n#     model_path=\"/kaggle/input/rna-train-pred/pointwise.pth\",\n#     output_csv=\"submission.csv\"\n# )\nrank_and_submit_actor(\n    input_csv=\"all_preds.csv\",\n    actor_path=\"/kaggle/input/rna-train-pred/actor.pth\",\n    output_csv=\"submission.csv\"\n)\n# rank_and_submit_random(input_csv=\"all_preds.csv\",output_csv=\"submission.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T07:00:46.397902Z","iopub.execute_input":"2025-05-29T07:00:46.398263Z","iopub.status.idle":"2025-05-29T07:00:47.094649Z","shell.execute_reply.started":"2025-05-29T07:00:46.398223Z","shell.execute_reply":"2025-05-29T07:00:47.093883Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# remove all shit\ndirs = ['/kaggle/working/boltz_predictions',\n        '/kaggle/working/rhofold_predictions',\n        '/kaggle/working/proteinx_predictions',\n        '/kaggle/working/drfold_predictions',\n        '/kaggle/working/deepfold_predictions',\n        '/kaggle/working/nufold_predictions',\n        '/kaggle/working/drfold_output',\n        '/kaggle/working/drfold_output1',\n        '/kaggle/working/output',\n        '/kaggle/working/boltz',\n        '/kaggle/working/rhofold',\n        '/kaggle/working/proteinx',\n        '/kaggle/working/drfold',\n        '/kaggle/working/deepfold',\n        '/kaggle/working/nufold',\n        '/kaggle/working/boltz_preds.csv',\n        '/kaggle/working/rhofold_preds.csv',\n        '/kaggle/working/proteinx_preds.csv',\n        '/kaggle/working/drfold_preds.csv',\n        '/kaggle/working/deepfold_preds.csv',\n        '/kaggle/working/nufold_preds.csv',\n        '/kaggle/working/all_preds.csv',\n       ]\nfor dir in dirs:\n    try:\n        shutil.rmtree(dir)\n        print(f'rm {dir} succeded')\n    except:\n        print(f'rm {dir} failed')\n        try:\n            os.remove(dir)\n            print(f\"'{dir}' removed successfully.\")\n        except:\n            print(f\"{dir}' remove failed\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T07:00:47.121508Z","iopub.execute_input":"2025-05-29T07:00:47.121745Z","iopub.status.idle":"2025-05-29T07:00:47.125281Z","shell.execute_reply.started":"2025-05-29T07:00:47.121724Z","shell.execute_reply":"2025-05-29T07:00:47.124359Z"}},"outputs":[],"execution_count":null}]}