{"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":5123458,"sourceType":"datasetVersion","datasetId":2975803},{"sourceId":10880419,"sourceType":"datasetVersion","datasetId":6760509},{"sourceId":10923077,"sourceType":"datasetVersion","datasetId":6785143},{"sourceId":11056335,"sourceType":"datasetVersion","datasetId":6888333},{"sourceId":11056379,"sourceType":"datasetVersion","datasetId":6888367},{"sourceId":11695366,"sourceType":"datasetVersion","datasetId":6791615},{"sourceId":11990166,"sourceType":"datasetVersion","datasetId":7541592}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Credits:\n\n# https://www.kaggle.com/code/youhanlee/boltz-1-inference-submission","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T08:08:09.095393Z","iopub.execute_input":"2026-01-10T08:08:09.095818Z","iopub.status.idle":"2026-01-10T08:08:09.099255Z","shell.execute_reply.started":"2026-01-10T08:08:09.095788Z","shell.execute_reply":"2026-01-10T08:08:09.098323Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Boltz-1","metadata":{}},{"cell_type":"code","source":"!pip install --no-index /kaggle/input/boltz-dependencies/*whl --no-deps\n!pip install --no-index /kaggle/input/fairscale-0413/*whl --no-deps\n!pip install --no-index /kaggle/input/biopython/*whl --no-deps","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T08:08:09.100465Z","iopub.execute_input":"2026-01-10T08:08:09.100798Z","iopub.status.idle":"2026-01-10T08:08:20.582571Z","shell.execute_reply.started":"2026-01-10T08:08:09.100768Z","shell.execute_reply":"2026-01-10T08:08:20.581458Z"},"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pwd","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T08:08:20.584297Z","iopub.execute_input":"2026-01-10T08:08:20.584527Z","iopub.status.idle":"2026-01-10T08:08:20.590969Z","shell.execute_reply.started":"2026-01-10T08:08:20.584507Z","shell.execute_reply":"2026-01-10T08:08:20.590052Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%mkdir inputs_prediction\n%mkdir outputs_prediction\n%cp -rf /kaggle/input/rna-prediction-boltz/boltz/src/boltz .\n%ls boltz","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T08:08:20.592346Z","iopub.execute_input":"2026-01-10T08:08:20.592593Z","iopub.status.idle":"2026-01-10T08:08:22.042829Z","shell.execute_reply.started":"2026-01-10T08:08:20.59257Z","shell.execute_reply":"2026-01-10T08:08:22.042015Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile inference.py\n\nimport pickle\nimport urllib.request\nfrom dataclasses import asdict, dataclass\nfrom pathlib import Path\nfrom typing import Literal, Optional\n\nimport click\nimport torch\nfrom pytorch_lightning import Trainer, seed_everything\nfrom pytorch_lightning.strategies import DDPStrategy\nfrom pytorch_lightning.utilities import rank_zero_only\nfrom tqdm import tqdm\n\nfrom boltz.data import const\nfrom boltz.data.module.inference import BoltzInferenceDataModule\nfrom boltz.data.msa.mmseqs2 import run_mmseqs2\nfrom boltz.data.parse.a3m import parse_a3m\nfrom boltz.data.parse.csv import parse_csv\nfrom boltz.data.parse.fasta import parse_fasta\nfrom boltz.data.parse.yaml import parse_yaml\nfrom boltz.data.types import MSA, Manifest, Record\nfrom boltz.data.write.writer import BoltzWriter\nfrom boltz.model.model import Boltz1\n\nCCD_URL = \"https://huggingface.co/boltz-community/boltz-1/resolve/main/ccd.pkl\"\nMODEL_URL = (\n    \"https://huggingface.co/boltz-community/boltz-1/resolve/main/boltz1_conf.ckpt\"\n)\n\n\n@dataclass\nclass BoltzProcessedInput:\n    \"\"\"Processed input data.\"\"\"\n\n    manifest: Manifest\n    targets_dir: Path\n    msa_dir: Path\n\n\n@dataclass\nclass BoltzDiffusionParams:\n    \"\"\"Diffusion process parameters.\"\"\"\n\n    gamma_0: float = 0.605\n    gamma_min: float = 1.107\n    noise_scale: float = 0.901\n    rho: float = 8\n    step_scale: float = 1.638\n    sigma_min: float = 0.0004\n    sigma_max: float = 160.0\n    sigma_data: float = 16.0\n    P_mean: float = -1.2\n    P_std: float = 1.5\n    coordinate_augmentation: bool = True\n    alignment_reverse_diff: bool = True\n    synchronize_sigmas: bool = True\n    use_inference_model_cache: bool = True\n\n\n@rank_zero_only\ndef download(cache: Path) -> None:\n    \"\"\"Download all the required data.\n\n    Parameters\n    ----------\n    cache : Path\n        The cache directory.\n\n    \"\"\"\n    # Download CCD\n    ccd = cache / \"ccd.pkl\"\n    if not ccd.exists():\n        click.echo(\n            f\"Downloading the CCD dictionary to {ccd}. You may \"\n            \"change the cache directory with the --cache flag.\"\n        )\n        urllib.request.urlretrieve(CCD_URL, str(ccd))  # noqa: S310\n\n    # Download model\n    model = cache / \"boltz1_conf.ckpt\"\n    if not model.exists():\n        click.echo(\n            f\"Downloading the model weights to {model}. You may \"\n            \"change the cache directory with the --cache flag.\"\n        )\n        urllib.request.urlretrieve(MODEL_URL, str(model))  # noqa: S310\n\n\ndef check_inputs(\n    data: Path,\n    outdir: Path,\n    override: bool = False,\n) -> list[Path]:\n    \"\"\"Check the input data and output directory.\n\n    If the input data is a directory, it will be expanded\n    to all files in this directory. Then, we check if there\n    are any existing predictions and remove them from the\n    list of input data, unless the override flag is set.\n\n    Parameters\n    ----------\n    data : Path\n        The input data.\n    outdir : Path\n        The output directory.\n    override: bool\n        Whether to override existing predictions.\n\n    Returns\n    -------\n    list[Path]\n        The list of input data.\n\n    \"\"\"\n    click.echo(\"Checking input data.\")\n\n    # Check if data is a directory\n    if data.is_dir():\n        data: list[Path] = list(data.glob(\"*\"))\n\n        # Filter out non .fasta or .yaml files, raise\n        # an error on directory and other file types\n        filtered_data = []\n        for d in data:\n            if d.suffix in (\".fa\", \".fas\", \".fasta\", \".yml\", \".yaml\"):\n                filtered_data.append(d)\n            elif d.is_dir():\n                msg = f\"Found directory {d} instead of .fasta or .yaml.\"\n                raise RuntimeError(msg)\n            else:\n                msg = (\n                    f\"Unable to parse filetype {d.suffix}, \"\n                    \"please provide a .fasta or .yaml file.\"\n                )\n                raise RuntimeError(msg)\n\n        data = filtered_data\n    else:\n        data = [data]\n\n    # Check if existing predictions are found\n    existing = (outdir / \"predictions\").rglob(\"*\")\n    existing = {e.name for e in existing if e.is_dir()}\n\n    # Remove them from the input data\n    if existing and not override:\n        data = [d for d in data if d.stem not in existing]\n        num_skipped = len(existing) - len(data)\n        msg = (\n            f\"Found some existing predictions ({num_skipped}), \"\n            f\"skipping and running only the missing ones, \"\n            \"if any. If you wish to override these existing \"\n            \"predictions, please set the --override flag.\"\n        )\n        click.echo(msg)\n    elif existing and override:\n        msg = \"Found existing predictions, will override.\"\n        click.echo(msg)\n\n    return data\n\n\ndef compute_msa(\n    data: dict[str, str],\n    target_id: str,\n    msa_dir: Path,\n    msa_server_url: str,\n    msa_pairing_strategy: str,\n) -> None:\n    \"\"\"Compute the MSA for the input data.\n\n    Parameters\n    ----------\n    data : dict[str, str]\n        The input protein sequences.\n    target_id : str\n        The target id.\n    msa_dir : Path\n        The msa directory.\n    msa_server_url : str\n        The MSA server URL.\n    msa_pairing_strategy : str\n        The MSA pairing strategy.\n\n    \"\"\"\n    if len(data) > 1:\n        paired_msas = run_mmseqs2(\n            list(data.values()),\n            msa_dir / f\"{target_id}_paired_tmp\",\n            use_env=True,\n            use_pairing=True,\n            host_url=msa_server_url,\n            pairing_strategy=msa_pairing_strategy,\n        )\n    else:\n        paired_msas = [\"\"] * len(data)\n\n    unpaired_msa = run_mmseqs2(\n        list(data.values()),\n        msa_dir / f\"{target_id}_unpaired_tmp\",\n        use_env=True,\n        use_pairing=False,\n        host_url=msa_server_url,\n        pairing_strategy=msa_pairing_strategy,\n    )\n\n    for idx, name in enumerate(data):\n        # Get paired sequences\n        paired = paired_msas[idx].strip().splitlines()\n        paired = paired[1::2]  # ignore headers\n        paired = paired[: const.max_paired_seqs]\n\n        # Set key per row and remove empty sequences\n        keys = [idx for idx, s in enumerate(paired) if s != \"-\" * len(s)]\n        paired = [s for s in paired if s != \"-\" * len(s)]\n\n        # Combine paired-unpaired sequences\n        unpaired = unpaired_msa[idx].strip().splitlines()\n        unpaired = unpaired[1::2]\n        unpaired = unpaired[: (const.max_msa_seqs - len(paired))]\n        if paired:\n            unpaired = unpaired[1:]  # ignore query is already present\n\n        # Combine\n        seqs = paired + unpaired\n        keys = keys + [-1] * len(unpaired)\n\n        # Dump MSA\n        csv_str = [\"key,sequence\"] + [f\"{key},{seq}\" for key, seq in zip(keys, seqs)]\n\n        msa_path = msa_dir / f\"{name}.csv\"\n        with msa_path.open(\"w\") as f:\n            f.write(\"\\n\".join(csv_str))\n\n\n@rank_zero_only\ndef process_inputs(  # noqa: C901, PLR0912, PLR0915\n    data: list[Path],\n    out_dir: Path,\n    ccd_path: Path,\n    msa_server_url: str,\n    msa_pairing_strategy: str,\n    max_msa_seqs: int = 4096,\n    use_msa_server: bool = False,\n) -> None:\n    \"\"\"Process the input data and output directory.\n\n    Parameters\n    ----------\n    data : list[Path]\n        The input data.\n    out_dir : Path\n        The output directory.\n    ccd_path : Path\n        The path to the CCD dictionary.\n    max_msa_seqs : int, optional\n        Max number of MSA sequences, by default 4096.\n    use_msa_server : bool, optional\n        Whether to use the MMSeqs2 server for MSA generation, by default False.\n\n    Returns\n    -------\n    BoltzProcessedInput\n        The processed input data.\n\n    \"\"\"\n    click.echo(\"Processing input data.\")\n    existing_records = None\n\n    # Check if manifest exists at output path\n    manifest_path = out_dir / \"processed\" / \"manifest.json\"\n    if manifest_path.exists():\n        click.echo(f\"Found a manifest file at output directory: {out_dir}\")\n\n        manifest: Manifest = Manifest.load(manifest_path)\n        input_ids = [d.stem for d in data]\n        existing_records, processed_ids = zip(\n            *[\n                (record, record.id)\n                for record in manifest.records\n                if record.id in input_ids\n            ]\n        )\n\n        if isinstance(existing_records, tuple):\n            existing_records = list(existing_records)\n\n        # Check how many examples need to be processed\n        missing = len(input_ids) - len(processed_ids)\n        if not missing:\n            click.echo(\"All examples in data are processed. Updating the manifest\")\n            # Dump updated manifest\n            updated_manifest = Manifest(existing_records)\n            updated_manifest.dump(out_dir / \"processed\" / \"manifest.json\")\n            return\n\n        click.echo(f\"{missing} missing ids. Preprocessing these ids\")\n        missing_ids = list(set(input_ids).difference(set(processed_ids)))\n        data = [d for d in data if d.stem in missing_ids]\n        assert len(data) == len(missing_ids)\n\n    # Create output directories\n    msa_dir = out_dir / \"msa\"\n    structure_dir = out_dir / \"processed\" / \"structures\"\n    processed_msa_dir = out_dir / \"processed\" / \"msa\"\n    predictions_dir = out_dir / \"predictions\"\n\n    out_dir.mkdir(parents=True, exist_ok=True)\n    msa_dir.mkdir(parents=True, exist_ok=True)\n    structure_dir.mkdir(parents=True, exist_ok=True)\n    processed_msa_dir.mkdir(parents=True, exist_ok=True)\n    predictions_dir.mkdir(parents=True, exist_ok=True)\n\n    # Load CCD\n    with ccd_path.open(\"rb\") as file:\n        ccd = pickle.load(file)  # noqa: S301\n\n    if existing_records is not None:\n        click.echo(f\"Found {len(existing_records)} records. Adding them to records\")\n\n    # Parse input data\n    records: list[Record] = existing_records if existing_records is not None else []\n    for path in tqdm(data):\n        try:\n            # Parse data\n            if path.suffix in (\".fa\", \".fas\", \".fasta\"):\n                target = parse_fasta(path, ccd)\n            elif path.suffix in (\".yml\", \".yaml\"):\n                target = parse_yaml(path, ccd)\n            elif path.is_dir():\n                msg = f\"Found directory {path} instead of .fasta or .yaml, skipping.\"\n                raise RuntimeError(msg)\n            else:\n                msg = (\n                    f\"Unable to parse filetype {path.suffix}, \"\n                    \"please provide a .fasta or .yaml file.\"\n                )\n                raise RuntimeError(msg)\n\n            # Get target id\n            target_id = target.record.id\n\n            # Get all MSA ids and decide whether to generate MSA\n            to_generate = {}\n            prot_id = const.chain_type_ids[\"PROTEIN\"]\n            for chain in target.record.chains:\n                # Add to generate list, assigning entity id\n                if (chain.mol_type == prot_id) and (chain.msa_id == 0):\n                    entity_id = chain.entity_id\n                    msa_id = f\"{target_id}_{entity_id}\"\n                    to_generate[msa_id] = target.sequences[entity_id]\n                    chain.msa_id = msa_dir / f\"{msa_id}.csv\"\n\n                # We do not support msa generation for non-protein chains\n                elif chain.msa_id == 0:\n                    chain.msa_id = -1\n\n            # Generate MSA\n            if to_generate and not use_msa_server:\n                msg = \"Missing MSA's in input and --use_msa_server flag not set.\"\n                raise RuntimeError(msg)\n\n            if to_generate:\n                msg = f\"Generating MSA for {path} with {len(to_generate)} protein entities.\"\n                click.echo(msg)\n                compute_msa(\n                    data=to_generate,\n                    target_id=target_id,\n                    msa_dir=msa_dir,\n                    msa_server_url=msa_server_url,\n                    msa_pairing_strategy=msa_pairing_strategy,\n                )\n\n            # Parse MSA data\n            msas = sorted({c.msa_id for c in target.record.chains if c.msa_id != -1})\n            msa_id_map = {}\n            for msa_idx, msa_id in enumerate(msas):\n                # Check that raw MSA exists\n                msa_path = Path(msa_id)\n                if not msa_path.exists():\n                    msg = f\"MSA file {msa_path} not found.\"\n                    raise FileNotFoundError(msg)\n\n                # Dump processed MSA\n                processed = processed_msa_dir / f\"{target_id}_{msa_idx}.npz\"\n                msa_id_map[msa_id] = f\"{target_id}_{msa_idx}\"\n                if not processed.exists():\n                    # Parse A3M\n                    if msa_path.suffix == \".a3m\":\n                        msa: MSA = parse_a3m(\n                            msa_path,\n                            taxonomy=None,\n                            max_seqs=max_msa_seqs,\n                        )\n                    elif msa_path.suffix == \".csv\":\n                        msa: MSA = parse_csv(msa_path, max_seqs=max_msa_seqs)\n                    else:\n                        msg = f\"MSA file {msa_path} not supported, only a3m or csv.\"\n                        raise RuntimeError(msg)\n\n                    msa.dump(processed)\n\n            # Modify records to point to processed MSA\n            for c in target.record.chains:\n                if (c.msa_id != -1) and (c.msa_id in msa_id_map):\n                    c.msa_id = msa_id_map[c.msa_id]\n\n            # Keep record\n            records.append(target.record)\n\n            # Dump structure\n            struct_path = structure_dir / f\"{target.record.id}.npz\"\n            target.structure.dump(struct_path)\n\n        except Exception as e:\n            if len(data) > 1:\n                print(f\"Failed to process {path}. Skipping. Error: {e}.\")\n            else:\n                raise e\n\n    # Dump manifest\n    manifest = Manifest(records)\n    manifest.dump(out_dir / \"processed\" / \"manifest.json\")\n\ndef predict(\n    data: str,\n    out_dir: str,\n    cache: str = \"~/.boltz\",\n    checkpoint: Optional[str] = None,\n    devices: int = 1,\n    accelerator: str = \"gpu\",\n    recycling_steps: int = 3,\n    sampling_steps: int = 200,\n    diffusion_samples: int = 1,\n    step_scale: float = 1.638,\n    write_full_pae: bool = False,\n    write_full_pde: bool = False,\n    output_format: Literal[\"pdb\", \"mmcif\"] = \"mmcif\",\n    num_workers: int = 2,\n    override: bool = False,\n    seed: Optional[int] = None,\n    use_msa_server: bool = False,\n    msa_server_url: str = \"https://api.colabfold.com\",\n    msa_pairing_strategy: str = \"greedy\",\n) -> None:\n    \"\"\"Run predictions with Boltz-1.\"\"\"\n    # If cpu, write a friendly warning\n    if accelerator == \"cpu\":\n        msg = \"Running on CPU, this will be slow. Consider using a GPU.\"\n        click.echo(msg)\n\n    # Set no grad\n    torch.set_grad_enabled(False)\n\n    # Ignore matmul precision warning\n    torch.set_float32_matmul_precision(\"highest\")\n\n    # Set seed if desired\n    if seed is not None:\n        seed_everything(int(seed))\n\n    # Set cache path\n    cache = Path(cache).expanduser()\n    cache.mkdir(parents=True, exist_ok=True)\n\n    # Create output directories\n    data = Path(data).expanduser()\n    out_dir = Path(out_dir).expanduser()\n    out_dir = out_dir / f\"boltz_results_{data.stem}\"\n    out_dir.mkdir(parents=True, exist_ok=True)\n\n    # Download necessary data and model\n    download(cache)\n\n    # Validate inputs\n    data = check_inputs(data, out_dir, override)\n    if not data:\n        click.echo(\"No predictions to run, exiting.\")\n        return\n\n    # Set up trainer\n    strategy = \"auto\"\n    if (isinstance(devices, int) and devices > 1) or (\n        isinstance(devices, list) and len(devices) > 1\n    ):\n        strategy = DDPStrategy()\n        if len(data) < devices:\n            msg = (\n                \"Number of requested devices is greater \"\n                \"than the number of predictions.\"\n            )\n            raise ValueError(msg)\n\n    msg = f\"Running predictions for {len(data)} structure\"\n    msg += \"s\" if len(data) > 1 else \"\"\n    click.echo(msg)\n\n    # Process inputs\n    ccd_path = cache / \"ccd.pkl\"\n    process_inputs(\n        data=data,\n        out_dir=out_dir,\n        ccd_path=ccd_path,\n        use_msa_server=use_msa_server,\n        msa_server_url=msa_server_url,\n        msa_pairing_strategy=msa_pairing_strategy,\n    )\n\n    # Load processed data\n    processed_dir = out_dir / \"processed\"\n    processed = BoltzProcessedInput(\n        manifest=Manifest.load(processed_dir / \"manifest.json\"),\n        targets_dir=processed_dir / \"structures\",\n        msa_dir=processed_dir / \"msa\",\n    )\n\n    # Create data module\n    data_module = BoltzInferenceDataModule(\n        manifest=processed.manifest,\n        target_dir=processed.targets_dir,\n        msa_dir=processed.msa_dir,\n        num_workers=num_workers,\n    )\n\n    # Load model\n    if checkpoint is None:\n        checkpoint = cache / \"boltz1_conf.ckpt\"\n\n    predict_args = {\n        \"recycling_steps\": recycling_steps,\n        \"sampling_steps\": sampling_steps,\n        \"diffusion_samples\": diffusion_samples,\n        \"write_confidence_summary\": True,\n        \"write_full_pae\": write_full_pae,\n        \"write_full_pde\": write_full_pde,\n    }\n    diffusion_params = BoltzDiffusionParams()\n    diffusion_params.step_scale = step_scale\n    model_module: Boltz1 = Boltz1.load_from_checkpoint(\n        checkpoint,\n        strict=True,\n        predict_args=predict_args,\n        map_location=\"cpu\",\n        diffusion_process_args=asdict(diffusion_params),\n        ema=False,\n    )\n    model_module.eval()\n\n    # Create prediction writer\n    pred_writer = BoltzWriter(\n        data_dir=processed.targets_dir,\n        output_dir=out_dir / \"predictions\",\n        output_format=output_format,\n    )\n\n    trainer = Trainer(\n        default_root_dir=out_dir,\n        strategy=strategy,\n        callbacks=[pred_writer],\n        accelerator=accelerator,\n        devices=devices,\n        precision=32,\n    )\n\n    # Compute predictions\n    trainer.predict(\n        model_module,\n        datamodule=data_module,\n        return_predictions=False,\n    )\n\n\nif __name__ == \"__main__\":\n    predict(data=\"./inputs_prediction\",\n            out_dir=\"./outputs_prediction\",\n            cache=\"/kaggle/input/rna-prediction-boltz/\",\n            diffusion_samples=1,\n            recycling_steps=10,\n            accelerator=\"gpu\",\n            sampling_steps=500,\n            seed=42,\n            override=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T08:08:22.043894Z","iopub.execute_input":"2026-01-10T08:08:22.044148Z","iopub.status.idle":"2026-01-10T08:08:22.05272Z","shell.execute_reply.started":"2026-01-10T08:08:22.044118Z","shell.execute_reply":"2026-01-10T08:08:22.051752Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport pandas as pd\nimport numpy as np\n\nsub_file = pd.read_csv('/kaggle/input/stanford-rna-3d-folding-2/test_sequences.csv')\n\nnames = sub_file['target_id'].tolist()\nsequences = sub_file['sequence'].tolist()\n\n# Inference\nidx = 0 \nfor tmp_id, tmp_sequence in zip(names, sequences):\n    with open(f'/kaggle/working/inputs_prediction/{tmp_id}.yaml', 'w') as f:\n        f.write(\"constraints: []\\n\")\n        f.write(\"sequences:\\n\")\n        f.write(\"- rna:\\n\")\n        f.write(\"    id:\\n\")\n        f.write(\"    - A1\\n\")\n        f.write(f\"    sequence: {tmp_sequence}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T08:08:22.053567Z","iopub.execute_input":"2026-01-10T08:08:22.053822Z","iopub.status.idle":"2026-01-10T08:08:22.415296Z","shell.execute_reply.started":"2026-01-10T08:08:22.053801Z","shell.execute_reply":"2026-01-10T08:08:22.414396Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\ntorch.cuda.empty_cache()\nimport gc\ngc.collect()\n\nimport subprocess\nimport logging\nimport pandas as pd\nfrom Bio.PDB.MMCIF2Dict import MMCIF2Dict\nimport os\n\nlogging.basicConfig(level=logging.INFO)\nlogger = logging.getLogger(__name__)\n\nresult = subprocess.run(['python', 'inference.py'], capture_output=True, text=True)\nlogger.info(f\"Command output: {result.stdout}\")\nlogger.error(f\"Command error: {result.stderr}\")\n\ndef get_coords(tmp_id, idx):\n    cif_file = f\"outputs_prediction/boltz_results_inputs_prediction/predictions/{tmp_id}/{tmp_id}_model_{idx}.cif\"\n    mmcif_dict = MMCIF2Dict(cif_file)\n    entity_poly_seq = mmcif_dict.get(\"_entity_poly_seq.mon_id\", [])\n    sequence = \"\".join(entity_poly_seq)\n    print(\"RNA sequence:\", sequence)\n    x_coords = mmcif_dict[\"_atom_site.Cartn_x\"]\n    y_coords = mmcif_dict[\"_atom_site.Cartn_y\"]\n    z_coords = mmcif_dict[\"_atom_site.Cartn_z\"]\n    atom_names = mmcif_dict[\"_atom_site.label_atom_id\"]\n    c1_coords = []\n    for i, atom in enumerate(atom_names):\n        if atom == \"C1'\":\n            c1_coords.append((float(x_coords[i]), float(y_coords[i]), float(z_coords[i])))\n    return c1_coords\n\nall_preds = os.listdir('outputs_prediction/boltz_results_inputs_prediction/predictions')\nsubmission = pd.read_csv('/kaggle/input/stanford-rna-3d-folding-2/sample_submission.csv')\n\nidx = 0\nfor tmp_id in all_preds:\n    print('#'*20, f'inferences for {tmp_id}')\n    for idx in range(1):\n        c1_coords = get_coords(tmp_id, idx)\n        submission.loc[submission['ID'].apply(lambda x: tmp_id in x), [f'x_{idx+1}', f'y_{idx+1}', f'z_{idx+1}']] = c1_coords\n    print()\n\n# Save the submission\nsubmission.to_csv('submission_boltz.csv', index=False)\nlogger.info(\"Created submission from predictions\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T08:08:22.416192Z","iopub.execute_input":"2026-01-10T08:08:22.416452Z","execution_failed":"2026-01-10T08:14:04.031Z"},"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pwd","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.032Z"},"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!ls","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.033Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.033Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %rm -rf boltz\n# %rm -rf inputs_prediction\n# %rm -rf outputs_prediction\n# %rm -rf inference.py","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.033Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Free up memory\nif torch.cuda.is_available():\n    torch.cuda.empty_cache()","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.034Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## DRfold2 + Template  \n","metadata":{}},{"cell_type":"code","source":"import os\nimport time\nimport torch\nimport random\nimport shutil\nimport numpy as np\nimport pandas as pd\n\nfrom Bio.Seq import Seq\nfrom Bio import pairwise2\n\nfrom tqdm import tqdm\nfrom scipy.spatial import distance_matrix\nfrom scipy.spatial.transform import Rotation as R\n\nimport warnings\nwarnings.filterwarnings('ignore')\n\n\ntest_sequences = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding-2/test_sequences.csv\")\n\nis_submission_mode = len(test_sequences) != 12","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.034Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_seqs = pd.read_csv('/kaggle/input/rna-all-data/merged_sequences_final.csv')\ntrain_labels = pd.read_csv('/kaggle/input/rna-all-data/merged_labels_final.csv')","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.034Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_seqs.head()","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.034Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_labels.head()","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.034Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !ls /kaggle/input/stanford-rna-3d-folding/PDB_RNA","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.035Z"},"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Check for new CIF files that need processing\nimport os\nimport pandas as pd\nfrom pathlib import Path\n\n# Get existing target_ids (without chain suffix)\nexisting_pdb_ids = set()\nfor target_id in train_seqs['target_id']:\n    pdb_id = target_id.rsplit('_', 1)[0]  # Remove chain suffix\n    existing_pdb_ids.add(pdb_id.lower())\n\nprint(f\"Existing PDB IDs in train_seqs: {len(existing_pdb_ids)}\")\n\n# Get all CIF files in directory\ncif_dir = '/kaggle/input/stanford-rna-3d-folding-2/PDB_RNA'\nall_cif_files = [f for f in os.listdir(cif_dir) if f.endswith('.cif')]\nall_pdb_ids = set(Path(f).stem.lower() for f in all_cif_files)\n\nprint(f\"Total CIF files found: {len(all_cif_files)}\")\n\n# Find new files to process\nnew_pdb_ids = all_pdb_ids - existing_pdb_ids\nnew_cif_files = [f\"{pdb_id}.cif\" for pdb_id in new_pdb_ids]\n\nprint(f\"New files to process: {len(new_cif_files)}\")\nprint(f\"New PDB IDs: {sorted(list(new_pdb_ids))}\")\n\nif new_cif_files:\n    print(f\"\\nFirst 10 new CIF files: {new_cif_files[:10]}\")\nelse:\n    print(\"No new files to process!\")","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.035Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Comprehensive modified nucleotide mapping\nnucleotide_mapping = {\n    # Standard nucleotides\n    'A': 'A', 'U': 'U', 'G': 'G', 'C': 'C',\n    \n    # === ADENOSINE MODIFICATIONS ===\n    'I': 'A',      # Inosine (hypoxanthine)\n    '1MA': 'A',    # 1-methyladenosine\n    '2MA': 'A',    # 2-methyladenosine\n    '6MA': 'A',    # N6-methyladenosine (m6A)\n    'M2A': 'A',    # N2-methyladenosine\n    'MS2': 'A',    # 2-methylthio-N6-isopentenyladenosine\n    'AET': 'A',    # 2-aminoethylthio-adenosine\n    'A2L': 'A',    # 2'-O-methyladenosine\n    'A44': 'A',    # Modified adenosine\n    '6OP': 'A',    # Modified adenosine\n    '8XA': 'A',    # Modified adenosine\n    'ZAD': 'A',    # Modified adenosine\n    \n    # === URIDINE MODIFICATIONS ===\n    'PSU': 'U',    # Pseudouridine (most common)\n    'H2U': 'U',    # Dihydrouridine\n    '5MU': 'U',    # 5-methyluridine (ribothymidine)\n    '4SU': 'U',    # 4-thiouridine\n    '2MU': 'U',    # 2'-O-methyluridine\n    'OMU': 'U',    # O-methyluridine\n    'T': 'U',      # Thymine (in RNA)\n    'RT': 'U',     # Ribothymidine\n    'DHU': 'U',    # Dihydrouridine\n    'UMS': 'U',    # 5-methoxycarbonylmethyluridine\n    'U2L': 'U',    # Modified uridine\n    'U36': 'U',    # Modified uridine\n    'Y5P': 'U',    # Modified uridine\n    'P5P': 'U',    # Modified uridine\n    'UFT': 'U',    # Modified uridine\n    'F2T': 'U',    # Modified uridine\n    '0U': 'U',     # Modified uridine\n    '8XU': 'U',    # Modified uridine\n    'ZBU': 'U',    # Modified uridine\n    'ZTH': 'U',    # Modified uridine\n    'ZHP': 'U',    # Modified uridine\n    'SSU': 'U',    # Modified uridine\n    \n    # === GUANOSINE MODIFICATIONS ===\n    'M2G': 'G',    # N2-methylguanosine\n    'M7G': 'G',    # 7-methylguanosine (cap structure)\n    'OMG': 'G',    # O-methylguanosine\n    '1MG': 'G',    # 1-methylguanosine\n    '2MG': 'G',    # 2'-O-methylguanosine\n    'YYG': 'G',    # Modified guanosine\n    'QUO': 'G',    # Queuosine\n    'G7M': 'G',    # 7-methylguanosine\n    'GTP': 'G',    # Guanosine triphosphate\n    'GDP': 'G',    # Guanosine diphosphate\n    'GMP': 'G',    # Guanosine monophosphate\n    'G2L': 'G',    # Modified guanosine\n    'G48': 'G',    # Modified guanosine\n    '6OO': 'G',    # Modified guanosine\n    '0G': 'G',     # Modified guanosine\n    '8XG': 'G',    # Modified guanosine\n    'ZGU': 'G',    # Modified guanosine\n    'LCG': 'G',    # Modified guanosine\n    \n    # === CYTIDINE MODIFICATIONS ===\n    '5MC': 'C',    # 5-methylcytidine\n    'OMC': 'C',    # O-methylcytidine\n    '2MC': 'C',    # 2'-O-methylcytidine\n    'M5C': 'C',    # 5-methylcytidine\n    'CBV': 'C',    # Carbovir cytidine\n    'C2L': 'C',    # Modified cytidine\n    'C43': 'C',    # Modified cytidine\n    '6NW': 'C',    # Modified cytidine\n    '0C': 'C',     # Modified cytidine\n    '8XC': 'C',    # Modified cytidine\n    'ZCY': 'C',    # Modified cytidine\n    'ZBC': 'C',    # Modified cytidine\n    \n    # === RARE/SYNTHETIC MODIFICATIONS ===\n    'ADP': 'A',    # Adenosine diphosphate\n    'ATP': 'A',    # Adenosine triphosphate\n    'AMP': 'A',    # Adenosine monophosphate\n    'UDP': 'U',    # Uridine diphosphate\n    'UTP': 'U',    # Uridine triphosphate\n    'UMP': 'U',    # Uridine monophosphate\n    'CDP': 'C',    # Cytidine diphosphate\n    'CTP': 'C',    # Cytidine triphosphate\n    'CMP': 'C',    # Cytidine monophosphate\n    \n    # === WYOSINE DERIVATIVES ===\n    'YW1': 'G',    # Wybutosine\n    'YW2': 'G',    # Wybutosine derivative\n    'YW3': 'G',    # Wybutosine derivative\n    \n    # === HYPERMODIFIED BASES ===\n    'Q': 'G',      # Queuosine\n    'X': 'G',      # Xanthosine\n    'D': 'U',      # Dihydrouridine\n    'P': 'U',      # Pseudouridine\n    \n    # === METHYLATION VARIANTS ===\n    'M1A': 'A',    # 1-methyladenosine\n    'M1G': 'G',    # 1-methylguanosine\n    'M3C': 'C',    # 3-methylcytidine\n    'M5U': 'U',    # 5-methyluridine\n    'M6A': 'A',    # N6-methyladenosine\n    \n    # === THIO MODIFICATIONS ===\n    'S2C': 'C',    # 2-thiocytidine\n    'S2U': 'U',    # 2-thiouridine\n    'S4U': 'U',    # 4-thiouridine\n    \n    # === CAP STRUCTURES ===\n    '7MG': 'G',    # 7-methylguanosine (5' cap)\n    'M7G': 'G',    # 7-methylguanosine\n    'G7M': 'G',    # 7-methylguanosine\n}","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.038Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Complete final code for processing new CIF files with comprehensive support\nfrom Bio.PDB import MMCIFParser\nimport pandas as pd\nfrom pathlib import Path\nimport os\nfrom tqdm import tqdm\n\n# Original function (for standard nucleotides)\ndef extract_rna_data_from_cif(cif_file_path):\n    \"\"\"Extract unique RNA sequences and C1' coordinates from a CIF file\"\"\"\n    parser = MMCIFParser(QUIET=True)\n    \n    try:\n        structure = parser.get_structure('structure', cif_file_path)\n        pdb_id = Path(cif_file_path).stem.upper()\n        \n        sequences_data = []\n        coordinates_data = []\n        seen_sequences = set()  # Track unique sequences\n        \n        for model in structure:\n            for chain in model:\n                chain_id = chain.id\n                target_id = f\"{pdb_id}_{chain_id}\"\n                \n                # Check if chain contains RNA residues\n                rna_residues = []\n                for residue in chain:\n                    if residue.get_resname() in ['A', 'U', 'G', 'C']:  # RNA nucleotides\n                        rna_residues.append(residue)\n                \n                if rna_residues:  # Only process if RNA residues found\n                    # Build sequence\n                    sequence = ''.join([res.get_resname() for res in rna_residues])\n                    \n                    # Only add if sequence is unique\n                    if sequence not in seen_sequences:\n                        seen_sequences.add(sequence)\n                        sequences_data.append({\n                            'target_id': target_id,\n                            'sequence': sequence\n                        })\n                        \n                        # Extract C1' coordinates for this unique sequence\n                        for i, residue in enumerate(rna_residues, 1):\n                            if \"C1'\" in residue:\n                                atom = residue[\"C1'\"]\n                                coordinates_data.append({\n                                    'ID': f\"{target_id}_{i}\",\n                                    'resname': residue.get_resname(),\n                                    'resid': i,\n                                    'x_1': atom.coord[0],\n                                    'y_1': atom.coord[1], \n                                    'z_1': atom.coord[2]\n                                })\n        \n        return sequences_data, coordinates_data\n        \n    except Exception as e:\n        print(f\"Error processing {cif_file_path}: {e}\")\n        return [], []\n\n# Disorder-aware glycosidic carbon detection\ndef get_glycosidic_carbon_disorder_aware(residue):\n    \"\"\"\n    Find C1' or C1{suffix} atoms, handling DisorderedAtom objects\n    \"\"\"\n    # Get all available atom names\n    available_atoms = [atom.get_name() for atom in residue]\n    \n    # Look for C1' first (most common)\n    if \"C1'\" in available_atoms:\n        atom = residue[\"C1'\"]\n        # Handle DisorderedAtom by getting the first conformation\n        if hasattr(atom, 'selected_child'):\n            return atom.selected_child\n        return atom\n    \n    # Look for any C1{suffix} pattern\n    c1_variants = [atom_name for atom_name in available_atoms if atom_name.startswith('C1') and len(atom_name) > 2]\n    \n    if c1_variants:\n        # If multiple C1 variants, prefer the shortest one\n        best_variant = min(c1_variants, key=len)\n        atom = residue[best_variant]\n        # Handle DisorderedAtom by getting the first conformation\n        if hasattr(atom, 'selected_child'):\n            return atom.selected_child\n        return atom\n    \n    return None\n\n# Comprehensive extraction function with disorder handling\ndef extract_rna_data_from_cif_comprehensive_final(cif_file_path):\n    \"\"\"Extract RNA with comprehensive modified nucleotide recognition and disorder handling\"\"\"\n    parser = MMCIFParser(QUIET=True)\n    \n    try:\n        structure = parser.get_structure('structure', cif_file_path)\n        pdb_id = Path(cif_file_path).stem.upper()\n        \n        sequences_data = []\n        coordinates_data = []\n        seen_sequences = set()\n        \n        for model in structure:\n            for chain in model:\n                chain_id = chain.id\n                target_id = f\"{pdb_id}_{chain_id}\"\n                \n                # Check if chain contains RNA residues (including all modified ones)\n                rna_residues = []\n                for residue in chain:\n                    res_name = residue.get_resname()\n                    if res_name in nucleotide_mapping:\n                        rna_residues.append(residue)\n                \n                if rna_residues:  # Only process if RNA residues found\n                    # Build sequence using standard nucleotides\n                    sequence = ''.join([nucleotide_mapping[res.get_resname()] for res in rna_residues])\n                    \n                    # Only add if sequence is unique\n                    if sequence not in seen_sequences:\n                        seen_sequences.add(sequence)\n                        sequences_data.append({\n                            'target_id': target_id,\n                            'sequence': sequence\n                        })\n                        \n                        # Extract coordinates using disorder-aware detection\n                        for i, residue in enumerate(rna_residues, 1):\n                            carbon_atom = get_glycosidic_carbon_disorder_aware(residue)\n                            \n                            if carbon_atom is not None:  # Use 'is not None' to avoid DisorderedAtom issues\n                                coordinates_data.append({\n                                    'ID': f\"{target_id}_{i}\",\n                                    'resname': nucleotide_mapping[residue.get_resname()],  # Use standard name\n                                    'resid': i,\n                                    'x_1': carbon_atom.coord[0],\n                                    'y_1': carbon_atom.coord[1], \n                                    'z_1': carbon_atom.coord[2]\n                                })\n        \n        return sequences_data, coordinates_data\n        \n    except Exception as e:\n        print(f\"Error processing {cif_file_path}: {e}\")\n        return [], []\n\n# Smart extraction function\ndef extract_rna_data_smart_final(cif_file_path):\n    \"\"\"\n    Final smart extraction: standard nucleotides first, then comprehensive with disorder handling\n    \"\"\"\n    # First try the original function (standard nucleotides only)\n    sequences_std, coordinates_std = extract_rna_data_from_cif(cif_file_path)\n    \n    # If we found RNA data with standard function, use it\n    if sequences_std:\n        return sequences_std, coordinates_std, \"standard\"\n    \n    # If no standard RNA found, try the comprehensive function with disorder handling\n    sequences_mod, coordinates_mod = extract_rna_data_from_cif_comprehensive_final(cif_file_path)\n    \n    if sequences_mod:\n        return sequences_mod, coordinates_mod, \"modified\"\n    else:\n        return [], [], \"none\"\n\n# Main processing function\ndef process_new_cif_files_final(train_seqs, train_labels, cif_dir, new_cif_files):\n    \"\"\"Final processing function with comprehensive nucleotide support and disorder handling\"\"\"\n    \n    if not new_cif_files:\n        print(\"No new files to process - all CIF files have already been processed!\")\n        return train_seqs, train_labels\n    \n    print(f\"Processing {len(new_cif_files)} new CIF files with final comprehensive extraction...\")\n    print(f\"Nucleotide mapping includes {len(nucleotide_mapping)} variants\")\n    print(f\"Includes disorder handling for DisorderedAtom objects\")\n    \n    new_sequences = []\n    new_coordinates = []\n    processing_stats = {\"standard\": 0, \"modified\": 0, \"none\": 0}\n    \n    for cif_file in tqdm(new_cif_files):\n        cif_path = os.path.join(cif_dir, cif_file)\n        sequences, coordinates, extraction_type = extract_rna_data_smart_final(cif_path)\n        \n        processing_stats[extraction_type] += 1\n        \n        if sequences:  # Only add if we found RNA data\n            new_sequences.extend(sequences)\n            new_coordinates.extend(coordinates)\n            print(f\"✅ {cif_file} ({extraction_type}): {len(sequences)} sequences, {len(coordinates)} coordinates\")\n        else:\n            print(f\"❌ {cif_file}: No RNA data found\")\n    \n    print(f\"\\nFINAL PROCESSING SUMMARY:\")\n    print(f\"Files with standard nucleotides: {processing_stats['standard']}\")\n    print(f\"Files with modified nucleotides: {processing_stats['modified']}\")\n    print(f\"Files with no RNA data: {processing_stats['none']}\")\n    print(f\"Success rate: {processing_stats['standard'] + processing_stats['modified']} / {len(new_cif_files)} = {((processing_stats['standard'] + processing_stats['modified']) / len(new_cif_files) * 100):.1f}%\")\n    \n    if new_sequences:\n        # Create DataFrames for new data\n        new_sequences_df = pd.DataFrame(new_sequences)\n        new_coordinates_df = pd.DataFrame(new_coordinates)\n        \n        print(f\"\\nFINAL NEW DATA SUMMARY:\")\n        print(f\"Total new sequences: {len(new_sequences)}\")\n        print(f\"Total new coordinates: {len(new_coordinates)}\")\n        \n        # Add to existing dataframes\n        train_seqs_updated = pd.concat([train_seqs, new_sequences_df], ignore_index=True)\n        train_labels_updated = pd.concat([train_labels, new_coordinates_df], ignore_index=True)\n        \n        print(f\"\\nFINAL UPDATED DATAFRAMES:\")\n        print(f\"train_seqs: {train_seqs.shape} -> {train_seqs_updated.shape}\")\n        print(f\"train_labels: {train_labels.shape} -> {train_labels_updated.shape}\")\n        \n        \n        print(f\"\\nFinal improvement summary:\")\n        print(f\"Sequences added: {len(train_seqs_updated) - len(train_seqs)}\")\n        print(f\"Coordinates added: {len(train_labels_updated) - len(train_labels)}\")\n        \n        return train_seqs_updated, train_labels_updated\n        \n    else:\n        print(\"No new RNA sequences found in any of the new files\")\n        return train_seqs, train_labels\n\n# EXECUTE THE FINAL COMPREHENSIVE PROCESSING\nprint(\"=\"*90)\nprint(\"FINAL COMPREHENSIVE PROCESSING WITH DISORDER HANDLING AND 93 NUCLEOTIDE VARIANTS\")\nprint(\"=\"*90)\n\n# Process the new files\ntrain_seqs_final, train_labels_final = process_new_cif_files_final(\n    train_seqs=train_seqs, \n    train_labels=train_labels, \n    cif_dir=cif_dir, \n    new_cif_files=new_cif_files\n)\n\nprint(\"\\n\" + \"=\"*90)\nprint(\"FINAL COMPREHENSIVE PROCESSING COMPLETE\")\nprint(\"=\"*90)\nprint(f\"Final train_seqs shape: {train_seqs_final.shape}\")\nprint(f\"Final train_labels shape: {train_labels_final.shape}\")\n\n# Show which files were successfully processed\nsuccessful_files = []\nfailed_files = []\n\nfor cif_file in new_cif_files:\n    cif_path = os.path.join(cif_dir, cif_file)\n    sequences, coordinates, extraction_type = extract_rna_data_smart_final(cif_path)\n    if sequences:\n        successful_files.append(cif_file)\n    else:\n        failed_files.append(cif_file)\n\nprint(f\"\\nSuccessfully processed files ({len(successful_files)}):\")\nfor f in successful_files:\n    print(f\"  ✅ {f}\")\n\nif failed_files:\n    print(f\"\\nFiles with no RNA data ({len(failed_files)}):\")\n    for f in failed_files:\n        print(f\"  ❌ {f}\")\n\nprint(f\"\\nFinal statistics:\")\nprint(f\"Original dataset: {len(train_seqs)} sequences, {len(train_labels)} coordinates\")\nprint(f\"Final dataset: {len(train_seqs_final)} sequences, {len(train_labels_final)} coordinates\")","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.039Z"},"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Success rate: 29 / 38 files contained RNA","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.041Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set up directories\npredictions_dir = \"/kaggle/working/predictions\"\nos.makedirs(predictions_dir, exist_ok=True)\nfasta_dir = \"/kaggle/working/fasta_files\"\nos.makedirs(fasta_dir, exist_ok=True)\n\n# Set time limit for DRfold2 (in seconds)\nDRFOLD_TIME_LIMIT = 7 * 60 * 60  # 7 hours\nstart_time_global = time.time()","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.041Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!cp -r /kaggle/input/drfold2-repo/DRfold2 /kaggle/working/\n%cd DRfold2\n%cd Arena\n!make Arena\n%cd ..","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.043Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!cp -r /kaggle/input/drfold2/model_hub /kaggle/working/DRfold2/","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.043Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile /kaggle/working/DRfold2/DRfold_infer.py\nimport os, sys\nimport torch\nimport numpy as np\nfrom subprocess import Popen, PIPE, STDOUT\n\n# Get the directory where the script is located\nexp_dir = os.path.dirname(os.path.abspath(__file__))\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n# dlexps = ['cfg_95','cfg_96','cfg_97','cfg_99']\ndlexps = ['cfg_97']\n\nprint(f\"[DRfold2] Starting prediction pipeline on {device} device\")\n\n# Get input FASTA file and output directory from command line arguments\nfastafile = os.path.realpath(sys.argv[1])\noutdir = os.path.realpath(sys.argv[2])\n\nprint(f\"[DRfold2] Input: {fastafile}\")\nprint(f\"[DRfold2] Output: {outdir}\")\n\n# Initialize clustering flag and AF3 file path\npclu = False\naf3file = None\n\n# Parse command line arguments\n# Acceptable formats:\n# python DRfold_infer.py input.fasta output_dir\n# python DRfold_infer.py input.fasta output_dir 1\n# python DRfold_infer.py input.fasta output_dir --af3 af3_model.pdb\n# python DRfold_infer.py input.fasta output_dir 1 --af3 af3_model.pdb\n\nfor i in range(3, len(sys.argv)):\n    if sys.argv[i] == \"1\" and i == 3:\n        pclu = True\n        print('[DRfold2] Clustering enabled - will generate multiple models')\n    elif sys.argv[i] == \"--af3\" and i+1 < len(sys.argv):\n        af3file = os.path.realpath(sys.argv[i+1])\n        print(f'[DRfold2] Using AlphaFold3 structure: {af3file}')\n\nif not pclu:\n    print('[DRfold2] Clustering disabled - will generate single model')\n\n# Create output directory if it doesn't exist\nif not os.path.isdir(outdir):\n    os.makedirs(outdir)\n    print(f\"[DRfold2] Created output directory: {outdir}\")  \n\n# Create subdirectories for different outputs\nret_dir = os.path.join(outdir,'rets_dir')  # For return files\nif not os.path.isdir(ret_dir):\n    os.makedirs(ret_dir)\n    print(f\"[DRfold2] Created returns directory: {ret_dir}\")\n\nfolddir = os.path.join(outdir,'folds')     # For folded structures\nif not os.path.isdir(folddir):\n    os.makedirs(folddir)\n    print(f\"[DRfold2] Created folds directory: {folddir}\")\n\nrefdir = os.path.join(outdir,'relax')      # For relaxed structures\nif not os.path.isdir(refdir):\n    os.makedirs(refdir)\n    print(f\"[DRfold2] Created relaxation directory: {refdir}\")\n\n\n# Helper function to run commands and capture output\ndef run_cmd(cmd, description):\n    print(f\"[DRfold2] {description}\")\n    print(f\"[DRfold2] Command: {cmd}\")\n    \n    # Execute the command and capture output in real-time\n    process = Popen(cmd, shell=True, stdout=PIPE, stderr=STDOUT, universal_newlines=True, bufsize=1)\n    \n    # Print output line by line as it becomes available\n    for line in iter(process.stdout.readline, ''):\n        line = line.strip()\n        if line:\n            print(f\"[DRfold2 subprocess] {line}\")\n    \n    # Get return code\n    return_code = process.wait()\n    if return_code == 0:\n        print(f\"[DRfold2] {description} completed successfully\")\n    else:\n        print(f\"[DRfold2] {description} failed with return code {return_code}\")\n    return return_code\n\n\n# Create paths for model directories and test scripts\ndlmains = [os.path.join(exp_dir, one_exp, 'test_modeldir.py') for one_exp in dlexps]\ndirs = [os.path.join(exp_dir, 'model_hub', one_exp) for one_exp in dlexps]\n\n\n# Check if processing has been done before\nif not os.path.isfile(ret_dir + '/done'): \n    print(\"[DRfold2] Step 1/4: GENERATING INITIAL PREDICTIONS\")\n    print(f\"[DRfold2] No previous predictions found, will generate e2e and geo files\")\n    \n    # Run each model configuration\n    for idx, (dlmain, one_exp, mdir) in enumerate(zip(dlmains, dlexps, dirs)):\n        # Construct command to run the model\n        cmd = f'python {dlmain} {device} {fastafile} {ret_dir}/{one_exp}_ {mdir}'\n        description = f\"Running model {idx+1}/{len(dlexps)}: {one_exp}\"\n        run_cmd(cmd, description)\n\n    # Mark processing as complete\n    wfile = open(ret_dir+'/done','w')\n    wfile.write('1')\n    wfile.close()\n    print(\"[DRfold2] Initial predictions generation completed\")\nelse:\n    print(\"[DRfold2] Step 1/4: USING EXISTING PREDICTIONS\")\n    print(f\"[DRfold2] Found previous predictions in {ret_dir}, using existing e2e and geo files\")\n\n# Helper function to get model PDB file\ndef get_model_pdb(tdir,opt):\n    files = os.listdir(tdir)\n    files = [afile for afile in files if afile.startswith(opt)][0]\n    return files\n\n\n# Set up directory paths and configuration files\ncso_dir = folddir                                                    # Directory for coarse-grained structures\nclufile = os.path.join(folddir,'clu.txt')                            # Clustering results file\nconfig_sel = os.path.join(exp_dir,'cfg_for_selection.json')          # Selection configuration\nfoldconfig = os.path.join(exp_dir,'cfg_for_folding.json')            # Folding configuration\nselpython = os.path.join(exp_dir,'PotentialFold','Selection.py')     # Selection script\noptpython = os.path.join(exp_dir,'PotentialFold','Optimization.py')  # Optimization script\nclupy = os.path.join(exp_dir,'PotentialFold','Clust.py')             # Clustering script\narena = os.path.join(exp_dir,'Arena','Arena')                        # Arena executable for structure refinement\n\n\n# Set up initial save prefixes for optimization and selection\noptsaveprefix = os.path.join(cso_dir, f'opt_0')\nsave_prefix = os.path.join(cso_dir, f'sel_0')\n\n# Get all .ret files from the return directory\nrets = os.listdir(ret_dir)\nrets = [afile for afile in rets if afile.endswith('.ret')]\nrets = [os.path.join(ret_dir,aret) for aret in rets ]\nret_str = ' '.join(rets)\n\nprint(\"[DRfold2] Step 2/4: SELECTION PROCESS\")\nprint(f\"[DRfold2] Found {len(rets)} return files for selection\")\nprint(f\"[DRfold2] Using selection config: {config_sel}\")\nprint(f\"[DRfold2] Output prefix: {save_prefix}\")\n\n\n# Run selection process\ncmd = f'python {selpython} {fastafile} {config_sel} {save_prefix} {ret_str}'\nrun_cmd(cmd, \"Running selection process\")\n\nprint(\"[DRfold2] Step 3/4: OPTIMIZATION PROCESS\")\nprint(f\"[DRfold2] Using fold config: {foldconfig}\")\nprint(f\"[DRfold2] Optimization output prefix: {optsaveprefix}\")\n\n# Run optimization process with optional AF3 file\ncmd = f'python {optpython} {fastafile} {optsaveprefix} {ret_dir} {save_prefix} {foldconfig}'\nif af3file and os.path.exists(af3file):\n    cmd += f' {af3file}'\nrun_cmd(cmd, \"Running optimization process\")\n\n# Get the coarse-grained PDB and save refined structure\ncgpdb = os.path.join(folddir,get_model_pdb(folddir,'opt_0'))\nsavepdb = os.path.join(refdir,'model_1.pdb')\n\nprint(\"[DRfold2] Step 4/4: STRUCTURE REFINEMENT\")\nprint(f\"[DRfold2] Found optimized structure: {cgpdb}\")\nprint(f\"[DRfold2] Final output will be saved to: {savepdb}\")\n\ncmd = f'{arena} {cgpdb} {savepdb} 7'\nrun_cmd(cmd, \"Running structure refinement\")\n\n# If clustering is enabled (pclu=True)\nif pclu:\n    print(\"[DRfold2] ADDITIONAL STEP: CLUSTERING\")\n    print(f\"[DRfold2] Running clustering process, output: {clufile}\")\n    \n    # Run clustering process\n    cmd = f'python {clupy} {ret_dir} {clufile}'\n    run_cmd(cmd, \"Running clustering\")\n\n    # Read clustering results\n    lines = open(clufile).readlines()\n    lines = [aline.strip() for aline in lines]\n    lines = [aline for aline in lines if aline]\n    \n    cluster_count = len(lines) - 1\n    print(f\"[DRfold2] Found {cluster_count} additional clusters to process\")\n\n    # Process each cluster\n    for i in range(1,len(lines)):\n        print(f\"[DRfold2] PROCESSING CLUSTER {i}/{cluster_count}\")\n        \n        # Get return files for this cluster\n        rets = lines[i].split()\n        rets = [os.path.join(ret_dir,aret.replace('.pdb','.ret')) for aret in rets ]\n        ret_str = ' '.join(rets)\n\n        # Set up save prefixes for this cluster\n        optsaveprefix =  os.path.join(cso_dir,f'opt_{str(i+1)}')\n        save_prefix = os.path.join(cso_dir,f'sel_{str(i+1)}')\n        \n        print(f\"[DRfold2] Cluster {i} Selection Process\")\n        print(f\"[DRfold2] Found {len(rets)} return files for selection\")\n        print(f\"[DRfold2] Selection output prefix: {save_prefix}\")\n\n        # Run selection process for this cluster\n        cmd = f'python {selpython} {fastafile} {config_sel} {save_prefix} {ret_str}'\n        run_cmd(cmd, f\"Running selection for cluster {i}\")\n        \n        print(f\"[DRfold2] Cluster {i} Optimization Process\")\n        print(f\"[DRfold2] Optimization output prefix: {optsaveprefix}\")\n\n        # Run optimization process for this cluster with optional AF3 file\n        cmd = f'python {optpython} {fastafile} {optsaveprefix} {ret_dir} {save_prefix} {foldconfig}'\n        if af3file and os.path.exists(af3file):\n            cmd += f' {af3file}'\n        run_cmd(cmd, f\"Running optimization for cluster {i}\")\n\n        # Get the coarse-grained PDB and save refined structure for this cluster\n        cgpdb = os.path.join(folddir,get_model_pdb(folddir,f'opt_{str(i+1)}'))\n        savepdb = os.path.join(refdir,f'model_{str(i+1)}.pdb')\n        \n        print(f\"[DRfold2] Cluster {i} Refinement Process\")\n        print(f\"[DRfold2] Found optimized structure: {cgpdb}\")\n        print(f\"[DRfold2] Final output will be saved to: {savepdb}\")\n\n        cmd = f'{arena} {cgpdb} {savepdb} 7'\n        run_cmd(cmd, f\"Running refinement for cluster {i}\")\n\nprint(\"[DRfold2] PREDICTION PIPELINE COMPLETED SUCCESSFULLY\")","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.044Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile /kaggle/working/DRfold2/PotentialFold/operations.py\n\"\"\"\noperations.py: Core Mathematical Operations for RNA Structure Analysis\n\nThis module provides essential mathematical operations for manipulating and analyzing\nRNA 3D structures, organized into four main categories:\n\n1. Basic Vector Operations:\n   Functions for selecting coordinates and calculating distances between points,\n   which form the foundation for all structural calculations.\n\n2. Angle Calculations:\n   Functions for computing bond angles and dihedral (torsion) angles between atoms,\n   with differentiable implementations suitable for gradient-based optimization.\n\n3. Rigid Body Transformations:\n   Functions for determining optimal rotations and translations between sets of\n   coordinates, enabling structure alignment and manipulation.\n\n4. Sequence Utilities:\n   Functions for converting RNA sequence data into standard 3D coordinate templates,\n   allowing sequence-structure mapping.\n\nThese operations support the core functionality of RNA structure prediction, analysis,\nand optimization throughout the codebase.\n\"\"\"\n\nimport os\nimport torch\nimport torch.nn as nn\nimport numpy as np \nimport math, sys, math\nfrom io import BytesIO\nimport torch.nn.functional as F\nfrom torch.autograd import Function\nfrom torch.nn.parameter import Parameter\nfrom subprocess import Popen, PIPE, STDOUT\n\n# Use consistent epsilon value across all functions\nEPS = 1e-8\n\n\n# === Basic Vector Operations ===\ndef coor_selection(coor,mask):\n    #[L,n,3],[L,n],byte\n    return torch.masked_select(coor,mask.bool()).view(-1,3)\n\ndef pair_distance(x1, x2, eps=1e-6, p=2):\n    # Use torch.cdist for p=2 (Euclidean) which is highly optimized\n    if p == 2:\n        return torch.cdist(x1, x2, p=2)\n    \n    # For other p-norms, avoid memory expansion with broadcasting\n    x1_ = x1.unsqueeze(1)  # [n1, 1, dim]\n    x2_ = x2.unsqueeze(0)  # [1, n2, dim]\n    diff = torch.abs(x1_ - x2_)\n    out = torch.pow(diff + eps, p).sum(dim=2)\n    return torch.pow(out, 1. / p)\n\n\n# === Angle Calculations ===\ndef angle(p0, p1, p2):\n    # [b 3] \n    b0 = p0-p1\n    b1 = p2-p1\n\n    b0 = b0 / (torch.norm(b0, dim =-1, keepdim=True) + EPS)\n    b1 = b1 / (torch.norm(b1, dim =-1, keepdim=True) + EPS)\n    \n    recos = torch.sum(b0*b1, -1)\n    recos = torch.clamp(recos, -0.9999, 0.9999)\n    return torch.acos(recos)\n\nclass torsion(Function):\n    #PyTorch class to calculate differentiable torsion angle\n    #https://stackoverflow.com/questions/20305272/dihedral-torsion-angle-from-four-points-in-cartesian-coordinates-in-python\n    #https://salilab.org/modeller/manual/node492.html\n    @staticmethod\n    def forward(ctx, p0, p1, p2, p3):\n        # Save input points for backward pass\n        ctx.save_for_backward(p0, p1, p2, p3)\n\n        # Calculate bond vectors\n        b0 = p0 - p1\n        b1 = p2 - p1\n        b2 = p3 - p2\n\n        # Normalize the middle bond vector\n        b1_norm = torch.norm(b1, dim=-1, keepdim=True) + 1e-8\n        b1_unit = b1 / b1_norm\n\n        # Project the other bonds onto the plane perpendicular to middle bond\n        v = b0 - torch.sum(b0 * b1_unit, dim=-1, keepdim=True) * b1_unit\n        w = b2 - torch.sum(b2 * b1_unit, dim=-1, keepdim=True) * b1_unit\n\n        # Calculate torsion using the arctan2 formula (more stable than arccos)\n        x = torch.sum(v * w, dim=-1)                                # cosine component\n        y = torch.sum(torch.cross(b1_unit, v, dim=-1) * w, dim=-1)  # sine component\n\n        return torch.atan2(y, x)\n\n    \n    @staticmethod\n    def backward(ctx, grad_output):\n        # Retrieve saved tensors from forward pass\n        p0, p1, p2, p3 = ctx.saved_tensors\n\n        # Calculate bond vectors\n        r01 = p0 - p1\n        r12 = p2 - p1\n        r23 = p3 - p2\n\n        # Calculate bond lengths with numerical stability\n        d01 = torch.norm(r01, dim=-1, keepdim=True) + 1e-8\n        d12 = torch.norm(r12, dim=-1, keepdim=True) + 1e-8\n        d23 = torch.norm(r23, dim=-1, keepdim=True) + 1e-8\n\n        # Normalize bond vectors\n        e01 = r01 / d01\n        e12 = r12 / d12\n        e23 = r23 / d23\n\n        # Calculate normal vectors to the two planes\n        n1 = torch.cross(e01, e12, dim=-1)\n        n2 = torch.cross(e12, e23, dim=-1)\n\n        # Normalize normal vectors\n        n1_norm = torch.norm(n1, dim=-1, keepdim=True) + 1e-8\n        n2_norm = torch.norm(n2, dim=-1, keepdim=True) + 1e-8\n        n1 = n1 / n1_norm\n        n2 = n2 / n2_norm\n\n        # Calculate gradients for each atom\n        # These are based on the analytical derivatives of dihedral angles\n        g0 = torch.cross(e01, n1, dim=-1) / d01\n        g1 = -g0 - torch.cross(e12, n1, dim=-1) / d12\n        g2 = torch.cross(e12, n2, dim=-1) / d12 - torch.cross(e23, n2, dim=-1) / d23\n        g3 = torch.cross(e23, n2, dim=-1) / d23\n\n        # Apply chain rule with incoming gradient\n        g0 = g0 * grad_output.unsqueeze(-1)\n        g1 = g1 * grad_output.unsqueeze(-1)\n        g2 = g2 * grad_output.unsqueeze(-1)\n        g3 = g3 * grad_output.unsqueeze(-1)\n\n        return g0, g1, g2, g3\n\n\ndef dihedral(input1, input2, input3, input4):\n    return torsion.apply(input1, input2, input3, input4)\n\n\n\n# === Rigid Body Transformations ===\ndef rigidFrom3Points(x):    \n    x1, x2, x3 = x[:, 0], x[:, 1], x[:, 2]\n    v1 = x3 - x2\n    v2 = x1 - x2\n    \n    # Normalize v1 to get e1\n    e1 = F.normalize(v1, p=2, dim=-1)\n    \n    # Project v2 onto e1 and subtract to get the component orthogonal to e1\n    u2 = v2 - e1 * (torch.einsum('bn,bn->b', e1, v2)[:, None])\n    \n    # Normalize u2 to get e2\n    e2 = F.normalize(u2, p=2, dim=-1)\n    \n    # Cross product to get e3\n    e3 = torch.cross(e1, e2, dim=-1)\n    \n    return torch.stack([e1, e2, e3], dim=1)\n\n\n# return the direction from to_q to from_p\ndef Kabsch_rigid(bases,x1,x2,x3):\n    # Early return for empty input\n    if x1.shape[0] == 0:\n        return torch.empty(0, 3, 3), torch.empty(0, 3)\n    \n    the_dim=1\n    to_q = torch.stack([x1,x2,x3],dim=the_dim)\n    biasq=torch.mean(to_q,dim=the_dim,keepdim=True)\n    q=to_q-biasq\n    m = torch.einsum('bnz,bny->bzy',bases,q)\n    u, s, v = torch.svd(m)\n    vt = torch.transpose(v, 1, 2)\n    det = torch.det(torch.matmul(u, vt))\n    det = det.view(-1, 1, 1)\n    vt = torch.cat((vt[:, :2, :], vt[:, -1:, :] * det), 1)\n    r = torch.matmul(u, vt)\n    return r,biasq.squeeze()\n\n\n\n# === Sequence Utilities ===\ndef Get_base(seq,basenpy_standard):\n    base_num = basenpy_standard.shape[1]\n    basenpy = np.zeros([len(seq),base_num,3])\n    seqnpy = np.array(list(seq))\n    basenpy[seqnpy=='A']=basenpy_standard[0]\n    basenpy[seqnpy=='a']=basenpy_standard[0]\n\n    basenpy[seqnpy=='G']=basenpy_standard[1]\n    basenpy[seqnpy=='g']=basenpy_standard[1]\n\n    basenpy[seqnpy=='C']=basenpy_standard[2]\n    basenpy[seqnpy=='c']=basenpy_standard[2]\n\n    basenpy[seqnpy=='U']=basenpy_standard[3]\n    basenpy[seqnpy=='u']=basenpy_standard[3]\n\n    basenpy[seqnpy=='T']=basenpy_standard[3]\n    basenpy[seqnpy=='t']=basenpy_standard[3]\n    \n    return torch.from_numpy(basenpy).double()","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.044Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile /kaggle/working/DRfold2/PotentialFold/Optimization.py\n#! /nfs/amino-home/liyangum/miniconda3/bin/python\nimport torch\nimport random\nimport numpy as np \nimport os, json, sys\n\nimport Cubic, Potential\nimport operations\nimport a2b, rigid\nimport torch.optim as opt\nfrom scipy.optimize import minimize\nimport pickle\n\ntorch.manual_seed(6)\nnp.random.seed(9)\nrandom.seed(9)\n\n\nScale_factor = 1.0\nUSEGEO = False\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n\ndef readconfig(configfile=''):\n    config=[]\n    expdir=os.path.dirname(os.path.abspath(__file__))\n    if configfile=='':\n        configfile=os.path.join(expdir,'lib','ddf.json')\n    config=json.load(open(configfile,'r'))\n    return config \n\n    \nclass Structure:\n    def __init__(self,fastafile,geofiles,saveprefix,initial_ret,foldconfig,af3file=None):\n        self.config=readconfig(foldconfig)\n        self.seqfile=fastafile\n        self.init_ret = initial_ret\n        self.foldconfig = foldconfig\n        self.geofiles = geofiles\n        self.rets = [pickle.load(open(refile,'rb')) for refile  in geofiles]\n        self.txs=[]\n        for ret in self.rets:\n            self.txs.append( torch.from_numpy(ret['coor']).double().to(device))\n        self.handle_geo()\n        self.pair = []\n        for ret in self.rets:\n            self.pair.append(torch.from_numpy(ret['plddt']).double().to(device))\n        self.saveprefix=saveprefix\n        self.seq=open(fastafile).readlines()[1].strip()\n        self.L=len(self.seq)\n        basenpy = np.load(os.path.join(os.path.dirname(os.path.abspath(__file__)),'lib','base.npy'))\n        self.basex = operations.Get_base(self.seq,basenpy).to(device)\n        othernpy = np.load(os.path.join(os.path.dirname(os.path.abspath(__file__)),'lib','other2.npy'))\n        self.otherx = operations.Get_base(self.seq,othernpy).to(device)\n        sidenpy = np.load(os.path.join(os.path.dirname(os.path.abspath(__file__)),'lib','side.npy'))\n        self.sidex = operations.Get_base(self.seq,sidenpy).to(device)\n        \n        self.init_mask()\n        self.init_paras()\n        self._init_fape()\n        self.tx2ds = [td.to(device) for td in self.tx2ds]\n        self.local_weight = torch.ones(self.L,self.L).to(device)\n        \n        for i in range(self.L):\n            for j in range(i+1,min(self.L,i+2)):\n                self.local_weight[i,j] = self.local_weight[j,i] = 4\n            for j in range(i+2,min(self.L,i+3)):\n                self.local_weight[i,j] = self.local_weight[j,i] = 3\n            for j in range(i+3,min(self.L,i+4)):\n                self.local_weight[i,j] = self.local_weight[j,i] = 2\n\n        # Load AlphaFold3 prediction if provided\n        if af3file and os.path.exists(af3file):\n            self.af3_coords = self.load_af3_structure(af3file)\n            print(f\"[DRfold2] Loaded AlphaFold3 structure from {af3file}\")\n            self._init_af3()\n        elif self.config.get('weight_af3', 0) > 0:\n            print(f\"[DRfold2] Warning: AlphaFold3 weight is set but no AF3 file provided\")\n\n    def load_af3_structure(self, pdbfile):\n        \"\"\"Load AlphaFold3 predicted structure from PDB file\"\"\"\n        print(\"[DRfold2] Loading AlphaFold3 prediction from {}\".format(pdbfile))\n        af3_coords = torch.zeros((self.L, 3, 3), device=device, dtype=torch.double)  # L residues, 3 atoms (P, C4', N1/N9), 3 coordinates\n        \n        # Read PDB file and extract P, C4', N1/N9 coordinates\n        atom_types = [' P  ', \" C4'\", ' N1 ']\n        purine_bases = ['A', 'G', 'a', 'g']\n        \n        with open(pdbfile, 'r') as f:\n            for line in f:\n                if not line.startswith(\"ATOM\"):\n                    continue\n                    \n                atom_type = line[12:16]\n                res_id = int(line[22:26]) - 1  # PDB is 1-indexed\n                if res_id < 0 or res_id >= self.L:\n                    continue\n                \n                # Handle purines (A, G) which have N9 instead of N1\n                if atom_type == ' N1 ' and self.seq[res_id].upper() in ['A', 'G']:\n                    continue\n                if atom_type == ' N9 ' and self.seq[res_id].upper() in ['A', 'G']:\n                    atom_idx = 2  # Index for the third atom (N1/N9)\n                elif atom_type in atom_types:\n                    atom_idx = atom_types.index(atom_type)\n                else:\n                    continue\n                \n                x = float(line[30:38])\n                y = float(line[38:46])\n                z = float(line[46:54])\n                af3_coords[res_id, atom_idx] = torch.tensor([x, y, z], device=device, dtype=torch.double)\n        \n        # Check if structure is complete\n        if torch.any(torch.sum(af3_coords, dim=-1) == 0):\n            print(\"[DRfold2] Warning: AlphaFold3 structure is incomplete\")\n        \n        return af3_coords\n\n    def _init_af3(self):\n        \"\"\"Initialize AlphaFold3 aligned coordinates for FAPE calculation\"\"\"\n        if not hasattr(self, 'af3_coords'):\n            return\n            \n        self.af3_rot, self.af3_trans = operations.Kabsch_rigid(\n            self.basex, \n            self.af3_coords[:, 0],  # P atoms\n            self.af3_coords[:, 1],  # C4' atoms\n            self.af3_coords[:, 2]   # N1/N9 atoms\n        )\n        \n        # Create aligned coordinates for fast energy calculation\n        # Shape: [L, L, 3, 3]\n        self.af3_aligned = self.af3_coords.unsqueeze(0).repeat(self.L, 1, 1, 1)\n        \n        # Translate by af3_trans\n        self.af3_aligned = self.af3_aligned - self.af3_trans.unsqueeze(1).unsqueeze(1)\n        \n        # Rotate by af3_rot\n        self.af3_aligned = torch.einsum('ijkl,ild->ijkd', self.af3_aligned, self.af3_rot.transpose(-1, -2))\n        \n        print(f\"[DRfold2] AlphaFold3 alignment initialized (shape: {self.af3_aligned.shape})\")\n\n    def _init_fape(self):\n        self.tx2ds = []\n        for tx in self.txs:\n            true_rot,true_trans = operations.Kabsch_rigid(self.basex,tx[:,0],tx[:,1],tx[:,2])\n            true_x2 = tx[:,None,:,:] - true_trans[None,:,None,:]\n            true_x2 = torch.einsum('ijnd,jde->ijne',true_x2,true_rot.transpose(-1,-2))\n            self.tx2ds.append(true_x2)\n    \n    def handle_geo(self):\n        oldkeys=['dist_p','dist_c','dist_n']\n        newkeys=['pp','cc','nn']\n        self.geos=[]\n        for ret in self.rets:\n            geo = {}\n            for nk,ok in zip(newkeys,oldkeys):\n                geo[nk] = torch.from_numpy(ret[ok].astype(np.float64)).to(device) + 0\n            self.geos.append(geo)\n\n\n    def init_mask(self):\n        halfmask=np.zeros([self.L,self.L])\n        fullmask=np.zeros([self.L,self.L])\n        for i in range(self.L):\n            for j in range(i+1,self.L):\n                halfmask[i,j]=1\n                fullmask[i,j]=1\n                fullmask[j,i]=1\n        self.halfmask=(torch.DoubleTensor(halfmask) > 0.5).to(device)\n        self.fullmask=(torch.DoubleTensor(fullmask) > 0.5).to(device)\n        self.clash_mask = torch.zeros([self.L,self.L,22,22], device=device)\n        for i in range(self.L):\n            for j in range(i+1,self.L):\n                self.clash_mask[i,j]=1\n\n        for i in range(self.L):\n             self.clash_mask[i,i,:6,7:]=1\n\n        for i in range(self.L-1):\n            self.clash_mask[i,i+1,:,0]=0\n            self.clash_mask[i,i+1,0,:]=0\n            self.clash_mask[i,i+1,:,5]=0\n            self.clash_mask[i,i+1,5,:]=0\n\n        self.side_mask = rigid.side_mask(self.seq).to(device)\n        self.side_mask = (self.side_mask[:,None,:,None] * self.side_mask[None,:,None,:]).to(device)\n        self.clash_mask = ((self.clash_mask > 0.5) * (self.side_mask > 0.5)).to(device)\n\n        self.geo_confimask_cc = []\n        self.geo_confimask_pp = []\n        self.geo_confimask_nn = []\n        for geo in self.geos:\n            confimask_cc = geo['cc'][:,:,-1] < 0.5\n            confimask_pp = geo['pp'][:,:,-1] < 0.5\n            confimask_nn = geo['nn'][:,:,-1] < 0.5\n            self.geo_confimask_cc.append(confimask_cc)\n            self.geo_confimask_pp.append(confimask_pp)\n            self.geo_confimask_nn.append(confimask_nn)\n\n\n    def init_paras(self):\n        self.geo_cc = []\n        self.geo_pp = []\n        self.geo_nn = []\n        self.cs_coefs = {'cc': [], 'pp': [], 'nn': []}\n        self.cs_knots = {'cc': [], 'pp': [], 'nn': []}\n        for geo in self.geos:\n            cc_cs,cc_decs=Cubic.dis_cubic(geo['cc'],2,40,36)\n            pp_cs,pp_decs=Cubic.dis_cubic(geo['pp'],2,40,36)\n            nn_cs,nn_decs=Cubic.dis_cubic(geo['nn'],2,40,36)\n            self.geo_cc.append([cc_cs,cc_decs])\n            self.geo_pp.append([pp_cs,pp_decs])\n            self.geo_nn.append([nn_cs,nn_decs])\n            \n            L = self.L\n            cc_coefs_np  = np.stack([[cc_cs[i,j].c for j in range(L)] for i in range(L)], axis=0)\n            cc_knots_np  = np.stack([[cc_cs[i,j].x for j in range(L)] for i in range(L)], axis=0)\n            self.cs_coefs['cc'].append(torch.from_numpy(cc_coefs_np).to(device))\n            self.cs_knots['cc'].append(torch.from_numpy(cc_knots_np).to(device))\n            \n            pp_coefs_np  = np.stack([[pp_cs[i,j].c for j in range(L)] for i in range(L)], axis=0)\n            pp_knots_np  = np.stack([[pp_cs[i,j].x for j in range(L)] for i in range(L)], axis=0)\n            self.cs_coefs['pp'].append(torch.from_numpy(pp_coefs_np).to(device))\n            self.cs_knots['pp'].append(torch.from_numpy(pp_knots_np).to(device))\n            \n            nn_coefs_np  = np.stack([[nn_cs[i,j].c for j in range(L)] for i in range(L)], axis=0)\n            nn_knots_np  = np.stack([[nn_cs[i,j].x for j in range(L)] for i in range(L)], axis=0)\n            self.cs_coefs['nn'].append(torch.from_numpy(nn_coefs_np).to(device))\n            self.cs_knots['nn'].append(torch.from_numpy(nn_knots_np).to(device))\n\n\n    def compute_bb_clash(self,coor,other_coor):\n        com_coor = torch.cat([coor,other_coor],dim=1)\n        com_dis  = (com_coor[:,None,:,None,:] - com_coor[None,:,None,:,:]).norm(dim=-1)\n        dynamicmask2_vdw= (com_dis <= 3.15) * (self.clash_mask)\n        vdw_dynamic = Potential.LJpotential(com_dis[dynamicmask2_vdw],3.15)\n        return vdw_dynamic.sum()*self.config['weight_vdw']\n\n    def compute_full_clash(self,coor,other_coor,side_coor):\n        com_coor = torch.cat([coor[:,:2],other_coor,side_coor],dim=1)\n        com_dis  = (com_coor[:,None,:,None,:] - com_coor[None,:,None,:,:]).norm(dim=-1)\n        dynamicmask2_vdw= (com_dis <= 2.5) * (self.clash_mask)\n        vdw_dynamic = Potential.LJpotential(com_dis[dynamicmask2_vdw],2.5)\n        return vdw_dynamic.sum()*self.config['weight_vdw']\n\n\n    def _cubic_pair_energy(self, atom_map, geo_cs, geo_confimask, weight_key):\n        \"\"\"General cubic-spline energy for CC/PP/NN pairs.\"\"\"\n        min_dis, max_dis, bin_num = 2, 40, 36\n        dev = atom_map.device\n        upper_th = max_dis - ((max_dis - min_dis) / bin_num) * 0.5\n        lower_th = 2.5\n        total = torch.zeros((), device=dev, dtype=torch.double)\n        spline_key   = weight_key.split('_')[1]  # 'cc', 'pp', or 'nn'\n        coeffs_list  = self.cs_coefs[spline_key]\n        knots_list   = self.cs_knots[spline_key]\n        for block_idx, mask_block in enumerate(geo_confimask):\n            mask = (atom_map <= upper_th) & mask_block & self.fullmask & (atom_map >= lower_th)\n            idx = mask.nonzero(as_tuple=True)\n            if idx[0].numel() > 1:\n                coef  = coeffs_list[block_idx][idx]\n                knots = knots_list[block_idx][idx]\n                part1 = Potential.cubic_distance(atom_map[mask], coef, knots, min_dis, max_dis, bin_num).sum() * self.config[weight_key] * 0.5\n            else:\n                part1 = torch.zeros((), device=dev)\n            part2 = ((atom_map <= lower_th) & mask_block & self.fullmask).sum() * self.config[weight_key]\n            total = total + part1 + part2\n        return total\n\n    def compute_cc_energy(self, coor):\n        atom_map = operations.pair_distance(coor[:,1], coor[:,1])\n        return self._cubic_pair_energy(atom_map, self.geo_cc, self.geo_confimask_cc, 'weight_cc')\n    \n    def compute_pp_energy(self, coor):\n        atom_map = operations.pair_distance(coor[:,0], coor[:,0])\n        return self._cubic_pair_energy(atom_map, self.geo_pp, self.geo_confimask_pp, 'weight_pp')\n    \n    def compute_nn_energy(self, coor):\n        atom_map = operations.pair_distance(coor[:,-1], coor[:,-1])\n        return self._cubic_pair_energy(atom_map, self.geo_nn, self.geo_confimask_nn, 'weight_nn')\n\n    def compute_pccp_energy(self,coor):\n        p_atoms=coor[:,0]\n        c_atoms=coor[:,1]\n        pccpmap=operations.dihedral( p_atoms[self.pccpi], c_atoms[self.pccpi], c_atoms[self.pccpj] ,p_atoms[self.pccpj]                  )\n        neg_log = Potential.cubic_torsion(pccpmap,self.pccp_coe,self.pccp_x,36)\n        return neg_log.sum()*self.config['weight_pccp']\n\n    def compute_cnnc_energy(self,coor):\n        n_atoms=coor[:,-1]\n        c_atoms=coor[:,1]\n        pccpmap=operations.dihedral( c_atoms[self.cnnci], n_atoms[self.cnnci], n_atoms[self.cnncj] ,c_atoms[self.cnncj]                  )\n        neg_log = Potential.cubic_torsion(pccpmap,self.cnnc_coe,self.cnnc_x,36)\n        return neg_log.sum()*self.config['weight_cnnc']\n\n    def compute_pnnp_energy(self,coor):\n        n_atoms=coor[:,-1]\n        p_atoms=coor[:,0]\n        pccpmap=operations.dihedral( p_atoms[self.pnnpi], n_atoms[self.pnnpi], n_atoms[self.pnnpj] ,p_atoms[self.pnnpj]                  )\n        neg_log = Potential.cubic_torsion(pccpmap,self.pnnp_coe,self.pnnp_x,36)\n        return neg_log.sum()*self.config['weight_pnnp']\n\n    def compute_pcc_energy(self,coor):\n        p_atoms=coor[:,1]\n        c_atoms=coor[:,2]\n        pccmap=operations.angle( p_atoms[self.pcci], c_atoms[self.pcci], c_atoms[self.pccj]                   )\n        neg_log = Potential.cubic_angle(pccmap,self.pcc_coe,self.pcc_x,12)\n        return neg_log.sum()*self.config['weight_pcc']\n\n    def compute_fape_energy(self,coor,ep=1e-3,epmax=20):\n        energy= 0\n        for tx in self.tx2ds:\n            px_mean = coor[:,[1]]\n            p_rot   = operations.rigidFrom3Points(coor)\n            p_tran  = px_mean[:,0]\n            pred_x2 = coor[:,None,:,:] - p_tran[None,:,None,:] # Lx Lrot N , 3\n            pred_x2 = torch.einsum('ijnd,jde->ijne',pred_x2,p_rot.transpose(-1,-2)) # transpose should be equal to inverse\n            errmap=torch.sqrt( ((pred_x2 - tx)**2).sum(dim=-1) + ep )\n            energy = energy + torch.sum(  torch.clamp(errmap,max=epmax)        )\n        return energy * self.config['weight_fape']\n\n    def compute_bond_energy(self,coor,other_coor):\n        # 3.87\n        o3 = other_coor[:-1,-2]\n        p  = coor[1:,0]\n        dis = (o3-p).norm(dim=-1)\n        energy = ((dis-1.607)**2).sum()\n        return energy * self.config['weight_bond']\n\n    def tooth_func(self,errmap, ep = 0.05):\n        return -1/(errmap/10+ep) + (1/ep)\n\n    def reweight_func(self,ww):\n        reweighting = torch.pow(ww,self.config['pair_weight_power'])\n        reweighting[ww < self.config['pair_weight_min']] = 0\n        return reweighting\n\n    def compute_fape_energy_fromquat(self,x,coor,ep=1e-6,epmax=100):\n        energy= 0\n        p_rot,px_mean = a2b.Non2rot(x[:,:9],x.shape[0]),x[:,9:]\n        pred_x2 = coor[:,None,:,:] - px_mean[None,:,None,:] # Lx Lrot N , 3\n        pred_x2 = torch.einsum('ijnd,jde->ijne',pred_x2,p_rot.transpose(-1,-2)) # transpose should be equal to inverse\n        for tx,weightplddt in zip(self.tx2ds,self.pair):\n\n            tamplate_dist_map = torch.min( tx.norm(dim=-1), dim=2   )[0]\n            errmap=torch.sqrt( ((pred_x2 - tx)**2).sum(dim=-1) + ep ) \n            energy = energy + torch.sum( ( (torch.clamp(errmap,max=self.config['FAPE_max'])**self.config['pair_error_power'])  * self.reweight_func(weightplddt[...,None]) * self.local_weight[...,None] )[tamplate_dist_map>self.config['pair_rest_min_dist']]    )\n\n        return energy * self.config['weight_fape']\n\n    def compute_af3_energy(self, coor, ep=1e-6):\n        \"\"\"Compute energy based on deviation from AlphaFold3 structure\"\"\"\n        # Skip if no AF3 structure is available\n        if not hasattr(self, 'af3_coords'):\n            return 0\n        \n        # Print status message to indicate AF3 integration is active\n        print(f\"[DRfold2] Computing AlphaFold3 energy contribution (weight: {self.config['weight_af3']})\")\n        \n        # Calculate rigid transformation\n        p_rot = operations.rigidFrom3Points(coor)\n        p_trans = coor[:, 1].clone()  # Use C4' as center of transform\n        \n        # Apply transformation to coordinates\n        pred_coords = coor.clone().unsqueeze(1)  # Shape: [L, 1, 3, 3]\n        pred_coords = pred_coords - p_trans.unsqueeze(1).unsqueeze(1)  # Translate\n        pred_coords = torch.einsum('ijkl,ild->ijkd', pred_coords, p_rot.transpose(-1, -2))  # Rotate\n        \n        # Calculate error map between aligned structures\n        af3_aligned = self.af3_aligned.clone()  # Shape: [L, L, 3, 3]\n        errmap = torch.sqrt(((pred_coords - af3_aligned)**2).sum(dim=-1) + ep)  # Distance error\n        \n        # Apply error power and clamping\n        max_dist = self.config.get('AF3_max', 20.0)\n        error_power = self.config.get('af3_error_power', 2.0)\n        \n        # Calculate per-atom distances to filter out distant pairs\n        atom_dists = torch.min(af3_aligned.norm(dim=-1), dim=2)[0]  # Min distance between atoms\n        \n        # Get pair weighting based on distance threshold\n        pair_min_dist = self.config.get('pair_rest_min_dist', 2.0)\n        mask = atom_dists > pair_min_dist\n        \n        # Apply error function\n        energy = torch.sum((torch.clamp(errmap, max=max_dist)**error_power)[mask])\n        \n        return energy * self.config['weight_af3']\n\n    def energy(self,rama):\n        coor=a2b.quat2b(self.basex,rama[:,9:])\n        other_coor = a2b.quat2b(self.otherx,rama[:,9:])\n        side_coor = a2b.quat2b(self.sidex,torch.cat([rama[:,:9],coor[:,-1]],dim=-1))\n        \n        if self.config['weight_cc']>0:\n            E_cc= self.compute_cc_energy(coor) / len(self.rets)\n        else:\n            E_cc=0\n        if self.config['weight_pp']>0:\n            E_pp= self.compute_pp_energy(coor) / len(self.rets)\n        else:\n            E_pp=0\n        if self.config['weight_nn']>0:\n            E_nn= self.compute_nn_energy(coor) / len(self.rets)\n        else:\n            E_nn=0\n\n        if self.config['weight_pccp']>0:\n            E_pccp= self.compute_pccp_energy(coor) / len(self.rets)\n        else:\n            E_pccp=0\n\n        if self.config['weight_cnnc']>0:\n            E_cnnc= self.compute_cnnc_energy(coor)  / len(self.rets)\n        else:\n            E_cnnc=0\n\n        if self.config['weight_pnnp']>0:\n            E_pnnp= self.compute_pnnp_energy(coor) / len(self.rets)\n        else:\n            E_pnnp=0\n\n        if self.config['weight_vdw']>0:\n            E_vdw= self.compute_full_clash(coor,other_coor,side_coor)\n        else:\n            E_vdw=0\n\n        if self.config['weight_fape']>0:\n            E_fape= self.compute_fape_energy_fromquat(rama[:,9:],coor) / len(self.rets)\n        else:\n            E_fape=0\n            \n        if self.config.get('weight_af3', 0) > 0 and hasattr(self, 'af3_coords'):\n            E_af3 = self.compute_af3_energy(coor)\n        else:\n            E_af3 = 0\n            \n        if self.config['weight_bond']>0:\n            E_bond= self.compute_bond_energy(coor,other_coor)\n        else:\n            E_bond=0\n            \n        return E_vdw + E_fape + E_bond + E_pp + E_cc + E_nn + E_pccp + E_cnnc + E_pnnp + E_af3\n\n\n    def obj_func_grad_np(self,rama_):\n        rama=torch.DoubleTensor(rama_)\n        rama.requires_grad=True\n        if rama.grad:\n            rama.grad.zero_()\n        f=self.energy(rama.view(self.L,21))*Scale_factor\n        grad_value=autograd.grad(f,rama)[0]\n        return grad_value.data.numpy().astype(np.float64)\n    \n    def obj_func_np(self,rama_):\n        rama=torch.DoubleTensor(rama_)\n        rama=rama.view(self.L,21)\n        with torch.no_grad():\n            f=self.energy(rama)*Scale_factor\n            return f.item()\n\n\n    def foldning(self):\n        ilter = self.init_ret\n        # 1) get initial quaternions (double precision)\n        try:\n            init_q = self.init_quat(ilter).double()\n        except:\n            init_q = self.init_quat_safe(ilter).double()\n\n        # 2) move to target device (GPU if available), enable grad\n        param = init_q.to(device).clone().detach().requires_grad_(True)\n\n        # 3) set up PyTorch LBFGS optimizer over `param`\n        optimizer = opt.LBFGS(\n            [param],\n            max_iter=self.config.get('max_iter', 300),\n            tolerance_grad=1e-6,\n            tolerance_change=1e-9,\n            history_size=10,\n            line_search_fn='strong_wolfe'\n        )\n\n        # 4) define the \"closure\" that LBFGS will call to reevaluate loss + gradients\n        def closure():\n            optimizer.zero_grad()                                 # clear old grads\n            E = self.energy(param.view(self.L,21)) * Scale_factor # compute ∂E/∂param\n            E.backward()\n            return E\n\n        # 5) run LBFGS until convergence (it calls closure repeatedly)\n        optimizer.step(closure)\n\n        # 6) write out final PDB\n        final_energy = self.energy(param.view(self.L,21)).item()\n        self.outpdb(param, self.saveprefix + '.pdb', energystr=str(final_energy))\n\n\n    def outpdb(self,rama,savefile,start=0,end=10000,energystr=''):\n        # bring baseframes and quaternion data onto CPU to prevent device mismatch\n        basex_cpu = self.basex.detach().cpu()\n        otherx_cpu = self.otherx.detach().cpu()\n        sidex_cpu = self.sidex.detach().cpu()\n        shaped_rama = rama.view(self.L,21).detach().cpu()\n        # compute backbone and other coords\n        coor_np = a2b.quat2b(basex_cpu, shaped_rama[:,9:]).detach().cpu().numpy()\n        other_np = a2b.quat2b(otherx_cpu, shaped_rama[:,9:]).detach().cpu().numpy()\n        coor = torch.FloatTensor(coor_np)\n        # compute side atom coords\n        side_coor_NP = a2b.quat2b(sidex_cpu, torch.cat([shaped_rama[:,:9], coor[:,-1]], dim=-1)).detach().cpu().numpy()\n        \n        Atom_name=[' P  ',\" C4'\",' N1 ']\n        Other_Atom_name = [\" O5'\",\" C5'\",\" C3'\",\" O3'\",\" C1'\"]\n        other_last_name = ['O',\"C\",\"C\",\"O\",\"C\"]\n\n        side_atoms=         [' N1 ',' C2 ',' O2 ',' N2 ',' N3 ',' N4 ',' C4 ',' O4 ',' C5 ',' C6 ',' O6 ',' N6 ',' N7 ',' N8 ',' N9 ']\n        side_last_name =    ['N',      \"C\",   \"O\",   \"N\",   \"N\",   'N',   'C',   'O',   'C',   'C',   'O',   'N',    'N', 'N','N']\n\n        base_dict = rigid.base_table()\n        last_name=['P','C','N']\n        wstr=[f'REMARK {str(energystr)}']\n        templet='%6s%5d %4s %3s %1s%4d    %8.3f%8.3f%8.3f%6.2f%6.2f          %2s%2s'\n        count=1\n        for i in range(self.L):\n            if self.seq[i] in ['a','g','A','G']:\n                Atom_name = [' P  ',\" C4'\",' N9 ']\n                #atoms = ['P','C4']\n\n            elif self.seq[i] in ['c','u','C','U']:\n                Atom_name = [' P  ',\" C4'\",' N1 ']\n            for j in range(coor_np.shape[1]):\n                outs=('ATOM  ',count,Atom_name[j],self.seq[i],'A',i+1,coor_np[i][j][0],coor_np[i][j][1],coor_np[i][j][2],0,0,last_name[j],'')\n                if i>=start-1 and i < end:\n                    wstr.append(templet % outs)\n                    count+=1\n\n            for j in range(other_np.shape[1]):\n                outs=('ATOM  ',count,Other_Atom_name[j],self.seq[i],'A',i+1,other_np[i][j][0],other_np[i][j][1],other_np[i][j][2],0,0,other_last_name[j],'')\n                if i>=start-1 and i < end:\n                    wstr.append(templet % outs)\n                    count+=1\n            \n        wstr='\\n'.join(wstr)\n        wfile=open(savefile,'w')\n        wfile.write(wstr)\n        wfile.close()\n    \n    def outpdb_coor(self,coor_np,savefile,start=0,end=1000,energystr=''):\n        Atom_name=[' P  ',\" C4'\",' N1 ']\n        last_name=['P','C','N']\n        wstr=[f'REMARK {str(energystr)}']\n        templet='%6s%5d %4s %3s %1s%4d    %8.3f%8.3f%8.3f%6.2f%6.2f          %2s%2s'\n        count=1\n        for i in range(self.L):\n            if self.seq[i] in ['a','g','A','G']:\n                Atom_name = [' P  ',\" C4'\",' N9 ']\n\n            elif self.seq[i] in ['c','u','C','U']:\n                Atom_name = [' P  ',\" C4'\",' N1 ']\n            for j in range(coor_np.shape[1]):\n                outs=('ATOM  ',count,Atom_name[j],self.seq[i],'A',i+1,coor_np[i][j][0],coor_np[i][j][1],coor_np[i][j][2],0,0,last_name[j],'')\n                if i>=start-1 and i < end:\n                    wstr.append(templet % outs)\n                count+=1\n            \n        wstr='\\n'.join(wstr)\n        wfile=open(savefile,'w')\n        wfile.write(wstr)\n        wfile.close()\n\n\n    def init_quat(self,ii):\n        x = torch.rand([self.L,21])\n        x[:,18:] = self.txs[ii].mean(dim=1)\n        init_coor = self.txs[ii]\n        biasq = torch.mean(init_coor,dim=1,keepdim=True)\n        q = init_coor - biasq\n        m = torch.einsum('bnz,bny->bzy',self.basex,q).reshape([self.L,-1])\n        x[:,:9] = x[:,9:18] = m\n        x.requires_grad_()\n        return x\n\n    def init_quat_safe(self,ii):\n        x = torch.rand([self.L,21])\n        x[:,18:] = self.txs[ii].mean(dim=1)\n        init_coor = self.txs[ii]\n        biasq = torch.mean(init_coor,dim=1,keepdim=True)\n        q = init_coor - biasq + torch.rand([self.L,3,3])\n        m = (torch.einsum('bnz,bny->bzy',self.basex,q) + torch.eye(3)[None,:,:]).reshape([self.L,-1])\n        x[:,:9] = x[:,9:18] = m\n        x.requires_grad_()\n        return x\n\n\nif __name__ == '__main__': \n\n    fastafile=sys.argv[1]\n    saveprefix=sys.argv[2]\n    retdirs  =sys.argv[3]\n    ret_score = sys.argv[4]\n    foldconfig = sys.argv[5]\n    \n    # Check for optional AF3 file\n    af3file = None\n    if len(sys.argv) > 6:\n        af3file = sys.argv[6]\n        if os.path.exists(af3file):\n            print(f\"[DRfold2] Using AlphaFold3 structure: {af3file}\")\n        else:\n            print(f\"[DRfold2] Warning: AlphaFold3 file not found: {af3file}\")\n            af3file = None\n\n    savepare = os.path.dirname(saveprefix)\n    if not os.path.isdir(savepare):\n        os.makedirs(savepare)\n\n    num_of_models = readconfig(foldconfig)['num_of_models']\n\n    score_dict = readconfig(ret_score)\n    sorted_items = sorted(score_dict.items(), key=lambda x: x[1])\n    lowest_n_keys = [item[0] for item in sorted_items][:num_of_models]\n    bestkey = lowest_n_keys[0] + ''\n    print(\"Before sort:\", lowest_n_keys)\n    lowest_n_keys.sort()\n    print(\"After sort:\", lowest_n_keys)\n    bestindex = lowest_n_keys.index(bestkey)\n\n    current_ret = bestkey\n    retfiles = [os.path.join(retdirs, afile) for afile in lowest_n_keys]\n    stru = Structure(fastafile, retfiles, saveprefix + '_from_' + current_ret, bestindex, foldconfig, af3file)\n    stru.foldning()","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.044Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile /kaggle/working/DRfold2/PotentialFold/Selection.py\n#! /nfs/amino-home/liyangum/miniconda3/bin/python\nimport numpy\nimport torch\nimport torch.autograd as autograd\nimport numpy as np \n\nimport random\nimport Cubic, Potential\nimport operations\nimport os, json, sys\n\nimport a2b, rigid\nimport torch.optim as opt\nfrom torch.nn.parameter import Parameter\nimport torch.nn as nn\nimport math\nfrom scipy.optimize import fmin_l_bfgs_b,fmin_cg,fmin_bfgs\nfrom scipy.optimize import minimize\nimport lbfgs_rosetta\nimport pickle\nimport shutil\n\ntorch.manual_seed(6)\ntorch.set_num_threads(4)\nnp.random.seed(9)\nrandom.seed(9)\n\nScale_factor = 1.0\nUSEGEO = False\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\ndef readconfig(configfile=''):\n    config=[]\n    expdir=os.path.dirname(os.path.abspath(__file__))\n    if configfile=='':\n        configfile=os.path.join(expdir,'lib','ddf.json')\n    config=json.load(open(configfile,'r'))\n    return config \n\n    \nclass Structure:\n    def __init__(self, fastafile, geofiles, foldconfig, saveprefix):\n        # Load Configuration and Inputs\n        self.config = readconfig(foldconfig)\n        self.seqfile = fastafile\n        self.foldconfig = foldconfig\n        self.geofiles = geofiles\n\n        # Load Model Results\n        self.rets = [pickle.load(open(refile, 'rb')) for refile  in geofiles]\n        \n        # Extract Coordinates\n        self.txs = []\n        for ret in self.rets:\n            self.txs.append(torch.from_numpy(ret['coor']).double().to(device))\n        \n        # Handle Geometrical Data\n        self.handle_geo()\n\n        # Extract pLDDT Scores\n        self.pair = []\n        for ret in self.rets:\n            self.pair.append( torch.from_numpy(ret['plddt']).double().to(device))\n        \n        # Store Output and Sequence Info\n        self.saveprefix = saveprefix\n        self.seq = open(fastafile).readlines()[1].strip()\n        self.L = len(self.seq)\n        \n        # Load Reference Arrays for Structure Construction\n        basenpy = np.load(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'lib', 'base.npy'))\n        self.basex = operations.Get_base(self.seq, basenpy).double().to(device)\n        \n        othernpy = np.load(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'lib', 'other2.npy'))\n        self.otherx = operations.Get_base(self.seq, othernpy).double().to(device)\n        \n        sidenpy = np.load(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'lib', 'side.npy'))\n        self.sidex = operations.Get_base(self.seq, sidenpy).double().to(device)        \n        \n        # Initialize Masks, Parameters, and FAPE\n        self.init_mask()\n        self.init_paras()\n        self._init_fape()\n    \n\n    def _init_fape(self):\n        self.tx2ds = []\n        for tx in self.txs:\n            true_rot, true_trans = operations.Kabsch_rigid(self.basex, tx[:, 0], tx[:, 1], tx[:, 2])\n            true_x2 = tx[:, None, :, :] - true_trans[None, :, None, :]\n            true_x2 = torch.einsum('ijnd,jde->ijne', true_x2, true_rot.transpose(-1,-2))\n            self.tx2ds.append(true_x2)\n    \n\n    def handle_geo(self):\n        oldkeys = ['dist_p', 'dist_c', 'dist_n']\n        newkeys = ['pp', 'cc', 'nn']\n        self.geos = []\n        geo = {'pp':0, 'cc':0, 'nn':0}\n        \n        for ret in self.rets:    \n            for nk, ok in zip(newkeys, oldkeys):\n                geo[nk] = geo[nk] + (ret[ok].astype(np.float64) /(len(self.rets)))\n        self.geos.append(geo)\n\n\n    def init_mask(self):\n        halfmask=np.zeros([self.L,self.L])\n        fullmask=np.zeros([self.L,self.L])\n        for i in range(self.L):\n            for j in range(i+1,self.L):\n                halfmask[i,j]=1\n                fullmask[i,j]=1\n                fullmask[j,i]=1\n        self.halfmask=torch.DoubleTensor(halfmask) > 0.5\n        self.fullmask=torch.DoubleTensor(fullmask) > 0.5\n        self.clash_mask = torch.zeros([self.L,self.L,22,22])\n        for i in range(self.L):\n            for j in range(i+1,self.L):\n                self.clash_mask[i,j]=1\n\n        for i in range(self.L):\n             self.clash_mask[i,i,:6,7:]=1\n\n        for i in range(self.L-1):\n            self.clash_mask[i,i+1,:,0]=0\n            self.clash_mask[i,i+1,0,:]=0\n            self.clash_mask[i,i+1,:,5]=0\n            self.clash_mask[i,i+1,5,:]=0\n\n        self.side_mask = rigid.side_mask(self.seq)\n        self.side_mask = self.side_mask[:,None,:,None] * self.side_mask[None,:,None,:]\n        self.clash_mask = (self.clash_mask > 0.5) * (self.side_mask > 0.5)\n\n        self.geo_confimask_cc = []\n        self.geo_confimask_pp = []\n        self.geo_confimask_nn = []\n        for geo in self.geos:\n            confimask_cc = torch.DoubleTensor(geo['cc'][:,:,-1]) < 0.5\n            confimask_pp = torch.DoubleTensor(geo['pp'][:,:,-1]) < 0.5\n            confimask_nn = torch.DoubleTensor(geo['nn'][:,:,-1]) < 0.5\n            self.geo_confimask_cc.append(confimask_cc)\n            self.geo_confimask_pp.append(confimask_pp)\n            self.geo_confimask_nn.append(confimask_nn)\n\n        # Move masks and confimasks to the GPU/CPU device\n        self.halfmask = self.halfmask.to(device)\n        self.fullmask = self.fullmask.to(device)\n        self.clash_mask = self.clash_mask.to(device)\n        self.side_mask = self.side_mask.to(device)\n        # geo_confimasks are lists\n        self.geo_confimask_cc = [m.to(device) for m in self.geo_confimask_cc]\n        self.geo_confimask_pp = [m.to(device) for m in self.geo_confimask_pp]\n        self.geo_confimask_nn = [m.to(device) for m in self.geo_confimask_nn]\n\n\n    def init_paras(self):\n        self.geo_cc = []\n        self.geo_pp = []\n        self.geo_nn = []\n        self.cs_coefs = {'cc': [], 'pp': [], 'nn': []}\n        self.cs_knots = {'cc': [], 'pp': [], 'nn': []}\n        for geo in self.geos:\n            cc_cs, cc_decs = Cubic.dis_cubic(geo['cc'], 2, 40, 36)\n            pp_cs, pp_decs = Cubic.dis_cubic(geo['pp'], 2, 40, 36)\n            nn_cs, nn_decs = Cubic.dis_cubic(geo['nn'], 2, 40, 36)\n            self.geo_cc.append([cc_cs, cc_decs])\n            self.geo_pp.append([pp_cs, pp_decs])\n            self.geo_nn.append([nn_cs, nn_decs])\n            L = self.L\n            cc_coefs_np = np.stack([[cc_cs[i,j].c for j in range(L)] for i in range(L)], axis=0)\n            cc_knots_np = np.stack([[cc_cs[i,j].x for j in range(L)] for i in range(L)], axis=0)\n            self.cs_coefs['cc'].append(torch.from_numpy(cc_coefs_np).to(device))\n            self.cs_knots['cc'].append(torch.from_numpy(cc_knots_np).to(device))\n            pp_coefs_np = np.stack([[pp_cs[i,j].c for j in range(L)] for i in range(L)], axis=0)\n            pp_knots_np = np.stack([[pp_cs[i,j].x for j in range(L)] for i in range(L)], axis=0)\n            self.cs_coefs['pp'].append(torch.from_numpy(pp_coefs_np).to(device))\n            self.cs_knots['pp'].append(torch.from_numpy(pp_knots_np).to(device))\n            nn_coefs_np = np.stack([[nn_cs[i,j].c for j in range(L)] for i in range(L)], axis=0)\n            nn_knots_np = np.stack([[nn_cs[i,j].x for j in range(L)] for i in range(L)], axis=0)\n            self.cs_coefs['nn'].append(torch.from_numpy(nn_coefs_np).to(device))\n            self.cs_knots['nn'].append(torch.from_numpy(nn_knots_np).to(device))\n     \n\n    def _cubic_pair_energy(self, atom_map, geo_cs, geo_confimask, weight_key):\n        \"\"\"General cubic-spline energy for CC/PP/NN pairs.\"\"\"\n        min_dis, max_dis, bin_num = 2, 40, 36\n        dev = atom_map.device\n        upper_th = max_dis - ((max_dis - min_dis) / bin_num) * 0.5\n        lower_th = 2.5\n        total = torch.zeros((), device=dev, dtype=torch.double)\n        spline_key = weight_key.split('_')[1]\n        coeffs_list = self.cs_coefs[spline_key]\n        knots_list = self.cs_knots[spline_key]\n        for block_idx, mask_block in enumerate(geo_confimask):\n            mask = (atom_map <= upper_th) & mask_block & self.fullmask & (atom_map >= lower_th)\n            idx = mask.nonzero(as_tuple=True)\n            if idx[0].numel() > 1:\n                coef = coeffs_list[block_idx][idx]\n                knots = knots_list[block_idx][idx]\n                part1 = Potential.cubic_distance(atom_map[mask], coef, knots, min_dis, max_dis, bin_num).sum() * self.config[weight_key] * 0.5\n            else:\n                part1 = torch.zeros((), device=dev, dtype=torch.double)\n            part2 = ((atom_map <= lower_th) & mask_block & self.fullmask).sum() * self.config[weight_key]\n            total = total + part1 + part2\n        return total\n\n    # GPU-friendly torsion and angle energy helpers\n    def _cubic_torsion_energy(self, atom_map, coef, x_vals, weight_key, num_bin):\n        energy = Potential.cubic_torsion(atom_map, coef, x_vals, num_bin)\n        return energy.sum() * self.config[weight_key]\n\n    def _cubic_angle_energy(self, atom_map, coef, x_vals, weight_key, num_bin):\n        energy = Potential.cubic_angle(atom_map, coef, x_vals, num_bin)\n        return energy.sum() * self.config[weight_key]\n\n    def compute_cc_energy(self, coor):\n        atom_map = operations.pair_distance(coor[:,1], coor[:,1])\n        return self._cubic_pair_energy(atom_map, self.geo_cc, self.geo_confimask_cc, 'weight_cc')\n\n    def compute_pp_energy(self, coor):\n        atom_map = operations.pair_distance(coor[:,0], coor[:,0])\n        return self._cubic_pair_energy(atom_map, self.geo_pp, self.geo_confimask_pp, 'weight_pp')\n\n    def compute_nn_energy(self, coor):\n        atom_map = operations.pair_distance(coor[:,-1], coor[:,-1])\n        return self._cubic_pair_energy(atom_map, self.geo_nn, self.geo_confimask_nn, 'weight_nn')\n\n    def compute_pccp_energy(self, coor):\n        # P-C-C-P dihedral energy on GPU\n        p = coor[:, 0]\n        c = coor[:, 1]\n        dia = operations.dihedral(\n            p[self.pccpi], c[self.pccpi], c[self.pccpj], p[self.pccpj]\n        )\n        return self._cubic_torsion_energy(dia, self.pccp_coe, self.pccp_x, 'weight_pccp', 36)\n\n    def compute_cnnc_energy(self, coor):\n        # C-N-N-C dihedral energy on GPU\n        n = coor[:, -1]\n        c = coor[:, 1]\n        dia = operations.dihedral(\n            c[self.cnnci], n[self.cnnci], n[self.cnncj], c[self.cnncj]\n        )\n        return self._cubic_torsion_energy(dia, self.cnnc_coe, self.cnnc_x, 'weight_cnnc', 36)\n\n    def compute_pnnp_energy(self, coor):\n        # P-N-N-P dihedral energy on GPU\n        n = coor[:, -1]\n        p = coor[:, 0]\n        dia = operations.dihedral(\n            p[self.pnnpi], n[self.pnnpi], n[self.pnnpj], p[self.pnnpj]\n        )\n        return self._cubic_torsion_energy(dia, self.pnnp_coe, self.pnnp_x, 'weight_pnnp', 36)\n\n    def compute_pcc_energy(self, coor):\n        # P-C-C angle energy on GPU\n        p = coor[:, 1]\n        c = coor[:, 2]\n        ang = operations.angle(\n            p[self.pcci], c[self.pcci], c[self.pccj]\n        )\n        return self._cubic_angle_energy(ang, self.pcc_coe, self.pcc_x, 'weight_pcc', 12)\n\n    def compute_fape_energy(self,coor,ep=1e-3,epmax=20):\n        energy= 0\n        for tx in self.tx2ds:\n            px_mean = coor[:,[1]]\n            p_rot   = operations.rigidFrom3Points(coor)\n            p_tran  = px_mean[:,0]\n            pred_x2 = coor[:,None,:,:] - p_tran[None,:,None,:] # Lx Lrot N , 3\n            pred_x2 = torch.einsum('ijnd,jde->ijne',pred_x2,p_rot.transpose(-1,-2)) # transpose should be equal to inverse\n            errmap=torch.sqrt( ((pred_x2 - tx)**2).sum(dim=-1) + ep )\n            energy = energy + torch.sum(  torch.clamp(errmap,max=epmax)        )\n        return energy * self.config['weight_fape']\n\n    def compute_bond_energy(self,coor,other_coor):\n        # 3.87\n        o3 = other_coor[:-1,-2]\n        p  = coor[1:,0]\n        dis = (o3-p).norm(dim=-1)\n        energy = ((dis-1.607)**2).sum()\n        return energy * self.config['weight_bond']\n\n    def tooth_func(self,errmap, ep = 0.05):\n        return -1/(errmap/10+ep) + (1/ep)\n    \n    def reweight_func(self,ww):\n        reweighting = torch.pow(ww,self.config['pair_weight_power'])\n        reweighting[ww < self.config['pair_weight_min']] = 0\n        return reweighting\n    \n    def compute_fape_energy_fromquat(self,x,coor,ep=1e-6,epmax=100):\n        energy= 0\n        p_rot,px_mean = a2b.Non2rot(x[:,:9],x.shape[0]),x[:,9:]\n        pred_x2 = coor[:,None,:,:] - px_mean[None,:,None,:] # Lx Lrot N , 3\n        pred_x2 = torch.einsum('ijnd,jde->ijne',pred_x2,p_rot.transpose(-1,-2)) # transpose should be equal to inverse\n\n        for tx,weightplddt in zip(self.tx2ds,self.pair):\n            tamplate_dist_map = torch.min( tx.norm(dim=-1), dim=2   )[0]\n            errmap=torch.sqrt( ((pred_x2 - tx)**2).sum(dim=-1) + ep ) \n            energy = energy + torch.sum( ( (torch.clamp(errmap,max=self.config['FAPE_max'])**self.config['pair_error_power'])  * self.reweight_func(weightplddt[...,None]) )[tamplate_dist_map>self.config['pair_rest_min_dist']]    )\n\n        return energy * self.config['weight_fape']\n    \n    def compute_fape_energy_fromcoor(self,coor,ep=1e-6,epmax=100):\n        energy= 0\n        \n        p_rot,px_mean = operations.Kabsch_rigid(self.basex,coor[:,0],coor[:,1],coor[:,2])\n        pred_x2 = coor[:,None,:,:] - px_mean[None,:,None,:] # Lx Lrot N , 3\n        pred_x2 = torch.einsum('ijnd,jde->ijne',pred_x2,p_rot.transpose(-1,-2)) # transpose should be equal to inverse\n        \n        for tx,weightplddt in zip(self.tx2ds,self.pair):\n            tamplate_dist_map = torch.min( tx.norm(dim=-1), dim=2   )[0]\n            errmap=torch.sqrt( ((pred_x2 - tx)**2).sum(dim=-1) + ep ) \n            energy = energy + torch.sum( ( (torch.clamp(errmap,max=self.config['FAPE_max'])**self.config['pair_error_power'])  * self.reweight_func(weightplddt[...,None]) )[tamplate_dist_map>self.config['pair_rest_min_dist']]    )\n\n        return energy * self.config['weight_fape']\n    \n    \n    def energy(self, rama):\n        coor = a2b.quat2b(self.basex, rama[:, 9:])\n        other_coor = a2b.quat2b(self.otherx, rama[:, 9:])\n        side_coor = a2b.quat2b(self.sidex, torch.cat([rama[:, :9], coor[:, -1]], dim=-1))\n\n        E_cc = self.compute_cc_energy(coor) / len(self.geofiles) if self.config['weight_cc'] > 0 else 0\n        E_pp = self.compute_pp_energy(coor) / len(self.geofiles) if self.config['weight_pp'] > 0 else 0\n        E_nn = self.compute_nn_energy(coor) / len(self.geofiles) if self.config['weight_nn'] > 0 else 0\n        E_pccp = self.compute_pccp_energy(coor) / len(self.geofiles) if self.config['weight_pccp'] > 0 else 0\n        E_cnnc = self.compute_cnnc_energy(coor) / len(self.geofiles) if self.config['weight_cnnc'] > 0 else 0\n        E_pnnp = self.compute_pnnp_energy(coor) / len(self.geofiles) if self.config['weight_pnnp'] > 0 else 0\n        E_vdw = self.compute_full_clash(coor, other_coor, side_coor) if self.config['weight_vdw'] > 0 else 0\n        E_fape = self.compute_fape_energy_fromquat(rama[:, 9:], coor) / len(self.geofiles) if self.config['weight_fape'] > 0 else 0\n        E_bond = self.compute_bond_energy(coor, other_coor) if self.config['weight_bond'] > 0 else 0\n\n        return E_vdw + E_fape + E_bond + E_pp + E_cc + E_nn + E_pccp + E_cnnc + E_pnnp\n\n\n    def energy_from_coor(self, coor):\n        E_cc = self.compute_cc_energy(coor) if self.config['weight_cc'] > 0 else 0\n        E_pp = self.compute_pp_energy(coor) if self.config['weight_pp'] > 0 else 0\n        E_nn = self.compute_nn_energy(coor) if self.config['weight_nn'] > 0 else 0\n        E_fape = (self.compute_fape_energy_fromcoor(coor) / len(self.geofiles)) if self.config['weight_fape'] > 0 else 0\n        print(E_fape, E_pp, E_cc, E_nn)\n        return E_fape + E_pp + E_cc + E_nn \n\n    def obj_func_grad_np(self,rama_):\n        rama=torch.DoubleTensor(rama_)\n        rama.requires_grad=True\n        if rama.grad:\n            rama.grad.zero_()\n        f=self.energy(rama.view(self.L,21))*Scale_factor\n        grad_value=autograd.grad(f,rama)[0]\n        return grad_value.data.numpy().astype(np.float64)\n    \n    def obj_func_np(self,rama_):\n        rama=torch.DoubleTensor(rama_)\n        rama=rama.view(self.L,21)\n        with torch.no_grad():\n            f = self.energy(rama)*Scale_factor\n            return f.item()\n\n    def saveconfig(self,dict,confile):\n        json_object = json.dumps(dict, indent = 4)\n        wfile = open(confile,'w')\n        wfile.write(json_object)\n        wfile.close()\n    \n    def scoring(self):\n        geoscale = self.config['geo_scale']\n        self.config['weight_pp'] = geoscale * self.config['weight_pp']\n        self.config['weight_cc'] = geoscale * self.config['weight_cc']\n        self.config['weight_nn'] = geoscale * self.config['weight_nn']\n        self.config['weight_pccp'] = geoscale * self.config['weight_pccp']\n        self.config['weight_cnnc'] = geoscale * self.config['weight_cnnc']\n        self.config['weight_pnnp'] = geoscale * self.config['weight_pnnp']  \n        \n        energy_dict = {}\n        saveenergy_dict  = {}\n        \n        with torch.no_grad():\n            for retfile, tx in zip(self.geofiles, self.txs):\n                one = self.energy_from_coor(tx)\n                aaretfile = os.path.basename(retfile) \n                energy_dict[aaretfile] = one.item()\n                saveenergy_dict[retfile] = one.item()\n            self.saveconfig(energy_dict, self.saveprefix)\n\n\n    def foldning(self):\n        minenergy=1e16\n        count=0\n        for tx in self.txs:\n            count+=1\n        \n        minirama=None\n\n        ilter = self.init_ret\n        selected_ret = self.geofiles[ilter]\n        try:\n            rama=self.init_quat(ilter).data.numpy()\n            self.config=readconfig(os.path.join(os.path.dirname(os.path.abspath(__file__)),'lib','vdw.json'))\n            rama = fmin_l_bfgs_b(func=self.obj_func_np, x0=rama,  fprime=self.obj_func_grad_np,iprint=10,maxfun=100)[0]\n            rama = rama.flatten()\n        except:\n            rama=self.init_quat_safe(ilter).data.numpy()\n            self.config=readconfig(os.path.join(os.path.dirname(os.path.abspath(__file__)),'lib','vdw.json'))\n            rama = fmin_l_bfgs_b(func=self.obj_func_np, x0=rama,  fprime=self.obj_func_grad_np,iprint=10,maxfun=100)[0]\n            rama = rama.flatten()\n            \n        self.config=readconfig(self.foldconfig)\n        geoscale = self.config['geo_scale']\n        self.config['weight_pp'] =geoscale * self.config['weight_pp']\n        self.config['weight_cc'] =geoscale * self.config['weight_cc']\n        self.config['weight_nn'] =geoscale * self.config['weight_nn']\n        self.config['weight_pccp'] =geoscale * self.config['weight_pccp']\n        self.config['weight_cnnc'] =geoscale * self.config['weight_cnnc']\n        self.config['weight_pnnp'] =geoscale * self.config['weight_pnnp']\n        for i in range(3):\n            line_min = lbfgs_rosetta.ArmijoLineMinimization(self.obj_func_np,self.obj_func_grad_np,True,len(rama),120)\n            lbfgs_opt = lbfgs_rosetta.lbfgs(self.obj_func_np,self.obj_func_grad_np)\n            rama=lbfgs_opt.run(rama,256,lbfgs_rosetta.absolute_converge_test,line_min,8000,self.obj_func_np,self.obj_func_grad_np,1e-9)\n        newrama=rama+0.0\n        newrama=torch.DoubleTensor(newrama) \n        current_energy =self.obj_func_np(rama)\n\n        if current_energy < minenergy:\n            print(current_energy,minenergy)\n            minenergy=current_energy\n            self.outpdb(newrama,self.saveprefix+'.pdb',energystr=str(current_energy))\n\n\n    def outpdb(self,rama,savefile,start=0,end=10000,energystr=''):\n        coor_np=a2b.quat2b(self.basex,rama.view(self.L,21)[:,9:]).data.numpy()\n        other_np=a2b.quat2b(self.otherx,rama.view(self.L,21)[:,9:]).data.numpy()\n        shaped_rama=rama.view(self.L,21)\n        coor = torch.FloatTensor(coor_np)\n        side_coor_NP = a2b.quat2b(self.sidex,torch.cat([shaped_rama[:,:9],coor[:,-1]],dim=-1)).data.numpy()\n        \n        Atom_name=[' P  ',\" C4'\",' N1 ']\n        Other_Atom_name = [\" O5'\",\" C5'\",\" C3'\",\" O3'\",\" C1'\"]\n        other_last_name = ['O',\"C\",\"C\",\"O\",\"C\"]\n\n        side_atoms=         [' N1 ',' C2 ',' O2 ',' N2 ',' N3 ',' N4 ',' C4 ',' O4 ',' C5 ',' C6 ',' O6 ',' N6 ',' N7 ',' N8 ',' N9 ']\n        side_last_name =    ['N',      \"C\",   \"O\",   \"N\",   \"N\",   'N',   'C',   'O',   'C',   'C',   'O',   'N',    'N', 'N','N']\n\n        base_dict = rigid.base_table()\n        \n        last_name=['P','C','N']\n        wstr=[f'REMARK {str(energystr)}']\n        templet='%6s%5d %4s %3s %1s%4d    %8.3f%8.3f%8.3f%6.2f%6.2f          %2s%2s'\n        count=1\n        for i in range(self.L):\n            if self.seq[i] in ['a','g','A','G']:\n                Atom_name = [' P  ',\" C4'\",' N9 ']\n\n            elif self.seq[i] in ['c','u','C','U']:\n                Atom_name = [' P  ',\" C4'\",' N1 ']\n            for j in range(coor_np.shape[1]):\n                outs=('ATOM  ',count,Atom_name[j],self.seq[i],'A',i+1,coor_np[i][j][0],coor_np[i][j][1],coor_np[i][j][2],0,0,last_name[j],'')\n                if i>=start-1 and i < end:\n                    wstr.append(templet % outs)\n                    count+=1\n\n            for j in range(other_np.shape[1]):\n                outs=('ATOM  ',count,Other_Atom_name[j],self.seq[i],'A',i+1,other_np[i][j][0],other_np[i][j][1],other_np[i][j][2],0,0,other_last_name[j],'')\n                if i>=start-1 and i < end:\n                    wstr.append(templet % outs)\n                    count+=1\n            \n        wstr='\\n'.join(wstr)\n        wfile=open(savefile,'w')\n        wfile.write(wstr)\n        wfile.close()\n    \n    \n    def outpdb_coor(self,coor_np,savefile,start=0,end=1000,energystr=''):\n        Atom_name=[' P  ',\" C4'\",' N1 ']\n        last_name=['P','C','N']\n        wstr=[f'REMARK {str(energystr)}']\n        templet='%6s%5d %4s %3s %1s%4d    %8.3f%8.3f%8.3f%6.2f%6.2f          %2s%2s'\n        count=1\n        for i in range(self.L):\n            if self.seq[i] in ['a','g','A','G']:\n                Atom_name = [' P  ',\" C4'\",' N9 ']\n\n            elif self.seq[i] in ['c','u','C','U']:\n                Atom_name = [' P  ',\" C4'\",' N1 ']\n            \n            for j in range(coor_np.shape[1]):\n                outs=('ATOM  ',count,Atom_name[j],self.seq[i],'A',i+1,coor_np[i][j][0],coor_np[i][j][1],coor_np[i][j][2],0,0,last_name[j],'')\n                if i>=start-1 and i < end:\n                    wstr.append(templet % outs)\n                count+=1\n            \n        wstr='\\n'.join(wstr)\n        wfile=open(savefile,'w')\n        wfile.write(wstr)\n        wfile.close()\n\n\nif __name__ == '__main__': \n\n    fastafile = sys.argv[1]\n    foldconfig = sys.argv[2]\n    save_prefix = sys.argv[3]\n    retfiles = sys.argv[4:]\n\n    save_parent_dir = os.path.dirname(save_prefix)\n    if not os.path.isdir(save_parent_dir):\n        os.makedirs(save_parent_dir)\n\n    retfiles.sort()\n    print(retfiles)\n\n    stru = Structure(fastafile, retfiles, foldconfig, save_prefix)    \n    stru.scoring()","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.044Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile /kaggle/working/DRfold2/PotentialFold/Cubic.py\nimport numpy as np \nfrom scipy.interpolate import CubicSpline,UnivariateSpline\nimport os\nfrom torch.autograd import Function\nimport torch\nimport math\n\ndef fit_dis_cubic(dis_matrix,min_dis,max_dis,num_bin):\n    # convert torch Tensor on GPU to numpy array for SciPy\n    if isinstance(dis_matrix, torch.Tensor):\n        dis_matrix = dis_matrix.detach().cpu().numpy()\n    dis_region=np.zeros(num_bin)\n    for i in range(num_bin):\n        dis_region[i]=min_dis+(i+0.5)*(max_dis-min_dis)*1.0/num_bin\n    L=dis_matrix.shape[0]\n    csnp=[]\n    decsnp=[]\n    for i in range(L):\n        css=[]\n        decss=[]\n        for j in range(L):\n            y=-np.log(      (dis_matrix[i,j,1:-1]+1e-8) / (dis_matrix[i,j,[-2]]+1e-8)              )\n            x=dis_region\n            x[0]=-0.0001\n            y[0]= max(10,y[1]+4)\n            cs= CubicSpline(x,y)\n            decs=cs.derivative()\n            css.append(cs)\n            decss.append(decs)\n        csnp.append(css)\n        decsnp.append(decss)\n    return np.array(csnp),np.array(decsnp)\n\ndef dis_cubic(out,min_dis,max_dis,num_bin):\n    print('fitting cubic distance')\n    cs,decs=fit_dis_cubic(out,min_dis,max_dis,num_bin)\n    return cs,decs\n\n\n\ndef cubic_matrix_torsion(dis_matrix,min_dis,max_dis,num_bin):\n    dis_region=np.zeros(num_bin)\n    bin_size=(max_dis-min_dis)/num_bin\n    for i in range(num_bin):\n        dis_region[i]=min_dis+(i+0.5)*(max_dis-min_dis)*1.0/num_bin\n    L=dis_matrix.shape[0]\n    csnp=[]\n    decsnp=[]\n    for i in range(L):\n        css=[]\n        decss=[]\n        for j in range(L):\n            y=-np.log(      dis_matrix[i,j,:-1]+1e-8             )\n            x=dis_region\n            x=np.append(x,x[-1]+bin_size)\n            y=np.append(y,y[0])\n            cs= CubicSpline(x,y,bc_type='periodic')\n            decs=cs.derivative()\n            css.append(cs)\n            decss.append(decs)\n        csnp.append(css)\n        decsnp.append(decss)\n    return np.array(csnp),np.array(decsnp)\ndef torsion_cubic(out,min_dis,max_dis,num_bin):\n    print('fitting cubic')\n    cs,decs=cubic_matrix_torsion(out,min_dis,max_dis,num_bin)\n    return cs,decs\n\ndef cubic_matrix_angle(dis_matrix,min_dis,max_dis,num_bin): # 0 - np.pi 12\n    dis_region=np.zeros(num_bin)\n    bin_size=(max_dis-min_dis)/num_bin\n    for i in range(num_bin):\n        dis_region[i]=min_dis+(i+0.5)*(max_dis-min_dis)*1.0/num_bin\n    L=dis_matrix.shape[0]\n    csnp=[]\n    decsnp=[]\n    for i in range(L):\n        css=[]\n        decss=[]\n        for j in range(L):\n            y=-np.log(      dis_matrix[i,j,:-1]+1e-8             )\n            x=dis_region\n\n            x=np.concatenate([[x[0]-bin_size*3,x[0]-bin_size*2,x[0]-bin_size], x,[x[-1]+bin_size,x[-1]+bin_size*2,x[-1]+bin_size*3]               ])\n            y=np.concatenate([ [y[2],y[1],y[0]],y,[y[-1],y[-2],y[-3]]                                                                                                                    ])\n\n            cs= CubicSpline(x,y)\n            decs=cs.derivative()\n\n            css.append(cs)\n            decss.append(decs)\n        csnp.append(css)\n        decsnp.append(decss)\n\n    return np.array(csnp),np.array(decsnp)\ndef angle_cubic(out,min_dis,max_dis,num_bin):\n\n    print('fitting angle cubic')\n    cs,decs=cubic_matrix_angle(out,min_dis,max_dis,num_bin)\n\n    return cs,decs","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.045Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile /kaggle/working/DRfold2/cfg_for_folding.json\n{\n    \"weight_pp\": 1,\n    \"weight_cc\": 1,\n    \"weight_nn\": 1,\n    \"weight_pccp\": 0,\n    \"weight_cnnc\": 0,\n    \"weight_pnnp\": 0,\n    \"weight_pcc\": 0,\n    \"weight_cnn\": 0,\n    \"weight_pnn\": 0,\n    \"weight_vdw\": 1,\n    \"weight_nn_contact\": 0,\n    \"weight_cc_contact\": 0,\n    \"weight_beta\": 0,\n    \"weight_fape\": 2,\n    \"weight_af3\": 2.5,\n    \"weight_bond\": 5000,\n    \"pair_weight_power\": 0.25,\n    \"pair_weight_min\": 0.2,\n    \"pair_error_power\": 3,\n    \"af3_error_power\": 2.0,\n    \"pair_rest_min_dist\": 2,\n    \"FAPE_max\": 30,\n    \"AF3_max\": 20,\n    \"geo_scale\": 450,\n    \"num_of_models\": 5\n}","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.045Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile /kaggle/working/DRfold2/cfg_for_selection.json\n{\n    \"weight_pp\": 1,\n    \"weight_cc\": 1,\n    \"weight_nn\": 1,\n    \"weight_pccp\": 0,\n    \"weight_cnnc\": 0,\n    \"weight_pnnp\": 0,\n    \"weight_pcc\": 0,\n    \"weight_cnn\": 0,\n    \"weight_pnn\": 0,\n    \"weight_vdw\": 1,\n    \"weight_nn_contact\": 0,\n    \"weight_cc_contact\": 0,\n    \"weight_beta\": 0,\n    \"weight_fape\": 1,\n    \"weight_af3\": 2.0,\n    \"weight_bond\": 1000,\n    \"pair_weight_power\": 0.5,\n    \"pair_weight_min\": 0.3,\n    \"pair_error_power\": 3.5,\n    \"af3_error_power\": 2.0,\n    \"pair_rest_min_dist\": 2,\n    \"FAPE_max\": 30,\n    \"AF3_max\": 20,\n    \"geo_scale\": 450,\n    \"num_of_models\": 5\n}","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.045Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Hybrid Method Follows  \n---","metadata":{}},{"cell_type":"code","source":"from Bio.PDB import MMCIFParser, PDBIO, Select\n\nclass RNAAtomSelect(Select):\n    \"\"\"Select only RNA backbone and base atoms needed by DRfold2\"\"\"\n    def accept_atom(self, atom):\n        # Check if this is an atom DRfold2 needs\n        atom_name = atom.name\n        residue = atom.get_parent()\n        resname = residue.get_resname()\n        \n        # Main backbone atoms needed by DRfold2\n        if atom_name in [\"P\", \"C4'\"]:\n            return True\n        \n        # For purines (A, G) we need N9, for pyrimidines (C, U) we need N1\n        if atom_name == \"N9\" and resname in [\"A\", \"G\"]:\n            return True\n        if atom_name == \"N1\" and resname in [\"C\", \"U\"]:\n            return True\n        \n        return False\n\ndef convert_cif_to_pdb(cif_file, pdb_file):\n    \"\"\"Convert mmCIF file to PDB format, fixing chain IDs and keeping only needed atoms.\n    \n    Args:\n        cif_file: Path to input mmCIF file\n        pdb_file: Path to output PDB file\n        \n    Returns:\n        bool: True if conversion was successful, False otherwise\n    \"\"\"\n    try:\n        # Parse the mmCIF file\n        parser = MMCIFParser(QUIET=True)\n        structure = parser.get_structure('', cif_file)\n        \n        # Fix chain IDs (map multi-character IDs like 'A1' to single characters)\n        for model in structure:\n            for chain in model:\n                if len(chain.id) > 1:\n                    # Just take the first character of the chain ID\n                    print(f\"Fixing chain ID: {chain.id} -> {chain.id[0]}\")\n                    chain.id = chain.id[0]\n        \n        # Write to PDB format, selecting only atoms needed by DRfold2\n        io = PDBIO()\n        io.set_structure(structure)\n        io.save(pdb_file, RNAAtomSelect())\n        print(f\"Successfully converted {cif_file} to {pdb_file}\")\n        return True\n    except Exception as e:\n        print(f\"Error converting {cif_file} to PDB: {str(e)}\")\n        \n        # Attempt alternative approach if primary method fails\n        try:\n            print(\"Trying alternative method...\")\n            parser = MMCIFParser(QUIET=True)\n            structure = parser.get_structure('', cif_file)\n            \n            # Create a new PDB file manually\n            with open(pdb_file, 'w') as f:\n                atom_num = 1\n                \n                for model in structure:\n                    for chain in model:\n                        chain_id = 'A'  # Use 'A' regardless of original ID\n                        \n                        for residue in chain:\n                            resname = residue.get_resname()\n                            resnum = residue.id[1]\n                            \n                            # Select atoms based on residue type\n                            needed_atoms = [\"P\", \"C4'\"]\n                            if resname in [\"A\", \"G\"]:\n                                needed_atoms.append(\"N9\")\n                            else:  # C, U\n                                needed_atoms.append(\"N1\")\n                                \n                            for atom_name in needed_atoms:\n                                if atom_name in residue:\n                                    atom = residue[atom_name]\n                                    x, y, z = atom.coord\n                                    \n                                    # Format as PDB ATOM line\n                                    line = f\"ATOM  {atom_num:5d} {atom_name:<4s} {resname:3s} {chain_id:1s}{resnum:4d}    {x:8.3f}{y:8.3f}{z:8.3f}  1.00  0.00           {atom.element:>2s}  \\n\"\n                                    f.write(line)\n                                    atom_num += 1\n                \n                f.write(\"END\\n\")\n            \n            print(f\"Successfully created PDB file using alternative method\")\n            return True\n            \n        except Exception as e2:\n            print(f\"Alternative method also failed: {str(e2)}\")\n            return False","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.045Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Define a function to run DRfold2 that captures output\ndef predict_rna_structures_drfold2(sequence, target_id, af3_pdb=None):\n    \"\"\"\n    Use DRfold2 to predict RNA structures with proper output capture\n    \n    Parameters:\n    -----------\n    sequence : str\n        RNA sequence to predict\n    target_id : str\n        Identifier for the target\n    af3_pdb : str, optional\n        Path to AlphaFold3 PDB file for hybrid prediction mode\n    \n    Returns:\n    --------\n    list of list of tuples\n        Five predictions, each containing C1' coordinates for each residue\n    \"\"\"\n    import subprocess\n    from subprocess import PIPE, STDOUT\n    \n    # Create FASTA file for this sequence\n    fasta_path = os.path.join(fasta_dir, f\"{target_id}.fasta\")\n    with open(fasta_path, \"w\") as f:\n        f.write(f\">{target_id}\\n{sequence}\\n\")\n    \n    # Run DRfold2 with proper output capture\n    output_dir = os.path.join(predictions_dir, target_id)\n    \n    # Build command with optional AF3 integration\n    cmd = f\"python /kaggle/working/DRfold2/DRfold_infer.py {fasta_path} {output_dir} 1\"\n    if af3_pdb and os.path.exists(af3_pdb):\n        cmd += f\" --af3 {af3_pdb}\"\n        print(f\"Using AlphaFold3 structure from: {af3_pdb}\")\n    elif af3_pdb:\n        print(f\"Warning: AlphaFold3 file not found: {af3_pdb}\")\n    \n    print(f\"Running command: {cmd}\")\n    process = subprocess.Popen(\n        cmd, \n        shell=True, \n        stdout=PIPE, \n        stderr=STDOUT,\n        universal_newlines=True,\n        bufsize=1\n    )\n    \n    # Print output in real-time\n    for line in iter(process.stdout.readline, ''):\n        line = line.strip()\n        if line:\n            print(line)\n    \n    # Get return code and check success\n    return_code = process.wait()\n    if return_code != 0:\n        print(f\"DRfold2 failed with return code {return_code}\")\n        return None\n    \n    # Clean up FASTA file to save space\n    os.remove(fasta_path)\n    \n    # Extract coordinates\n    relax_dir = os.path.join(output_dir, \"relax\")\n    if not os.path.isdir(relax_dir):\n        print(f\"Warning: No relax directory found for {target_id}\")\n        relax_dir = output_dir\n    \n    # Get up to 5 PDB files\n    pdb_files = sorted([f for f in os.listdir(relax_dir) if f.endswith(\".pdb\")])[:5]\n    \n    if not pdb_files:\n        print(f\"Warning: No PDB files found for {target_id}\")\n        # Return None to indicate failure\n        return None\n    \n    # Parse PDB files to extract C1' coordinates\n    predictions = []\n    for pdb_file in pdb_files:\n        file_path = os.path.join(relax_dir, pdb_file)\n        \n        # Read PDB file\n        coords = []\n        with open(file_path, \"r\") as f:\n            residue_map = {}\n            for line in f:\n                if line.startswith(\"ATOM\") and \" C1' \" in line:\n                    parts = line.split()\n                    resid = int(parts[5])  # Residue ID as integer\n                    x, y, z = float(parts[6]), float(parts[7]), float(parts[8])\n                    residue_map[resid] = (x, y, z)\n            \n            # Ensure we have coordinates for all residues\n            for j in range(1, len(sequence) + 1):\n                if j in residue_map:\n                    coords.append(residue_map[j])\n                else:\n                    # If residue not found, use zeros\n                    print(f\"Warning: Residue {j} not found in {pdb_file} for {target_id}\")\n                    coords.append((0.0, 0.0, 0.0))\n        \n        predictions.append(coords)\n    \n    # Clean up PDB files to save space\n    if is_submission_mode:\n        shutil.rmtree(output_dir)\n    \n    # If we have fewer than 5 predictions, duplicate the last one\n    while len(predictions) < 5:\n        predictions.append(predictions[-1] if predictions else [(0.0, 0.0, 0.0) for _ in range(len(sequence))])\n    \n    return predictions[:5]  # Return exactly 5 predictions","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.046Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Vectorized version of process_labels function\ndef process_labels_vectorized(labels_df):\n    # Extract target_id from ID column (remove last part after underscore)\n    labels_df = labels_df.copy()\n    labels_df['target_id'] = labels_df['ID'].str.rsplit('_', n=1).str[0]\n    \n    # Sort by target_id and resid for proper ordering\n    labels_df = labels_df.sort_values(['target_id', 'resid'])\n    \n    # Group by target_id and convert coordinates to arrays\n    coords_dict = {}\n    for target_id, group in labels_df.groupby('target_id'):\n        # Extract coordinates as numpy array in one operation\n        coords_dict[target_id] = group[['x_1', 'y_1', 'z_1']].values\n    \n    return coords_dict\n\ndef find_similar_sequences(query_seq, train_seqs_df, train_coords_dict, top_n=5):\n    similar_seqs = []\n    query_seq_obj = Seq(query_seq)\n\n    for _, row in train_seqs_df.iterrows():\n        target_id = row['target_id']\n        train_seq = row['sequence']\n\n        # Skip if coordinates not available\n        if target_id not in train_coords_dict:\n            continue\n\n        # Skip if sequence is too different in length (more than 40% difference)\n        if abs(len(train_seq) - len(query_seq)) / max(len(train_seq), len(query_seq)) > 0.4:\n            continue\n\n        # Perform sequence alignment\n        alignments = pairwise2.align.globalms(query_seq_obj, train_seq, 2.9, -1, -10, -0.5, one_alignment_only=True)\n\n        if alignments:\n            alignment = alignments[0]\n            similarity_score = alignment.score / (2 * min(len(query_seq), len(train_seq)))\n            similar_seqs.append((target_id, train_seq, similarity_score, train_coords_dict[target_id]))\n\n    # Sort by similarity score (higher is better) and return top N\n    similar_seqs.sort(key=lambda x: x[2], reverse=True)\n    return similar_seqs[:top_n]\n\n\n\n# ======= adaptive_rna_constraints =================\ndef adaptive_rna_constraints(coordinates, sequence, confidence=1.0):\n    \"\"\"Apply realistic RNA structural constraints\"\"\"\n    # Make a copy of coordinates to refine\n    refined_coords = coordinates.copy()\n    n_residues = len(sequence)\n    \n    # Calculate constraint strength (inverse of confidence)\n    constraint_strength = 0.8 * (1.0 - min(confidence, 0.8))\n    \n    # 1. Sequential distance constraints (consecutive nucleotides)\n    seq_min_dist = 5.5  # Minimum sequential distance\n    seq_max_dist = 6.5  # Maximum sequential distance\n    \n    for i in range(n_residues - 1):\n        current_pos = refined_coords[i]\n        next_pos = refined_coords[i+1]\n        \n        # Calculate current distance\n        current_dist = np.linalg.norm(next_pos - current_pos)\n        \n        # Only adjust if significantly outside expected range\n        if current_dist < seq_min_dist or current_dist > seq_max_dist:\n            # Calculate target distance (midpoint of range)\n            target_dist = (seq_min_dist + seq_max_dist) / 2\n            \n            # Get direction vector\n            direction = next_pos - current_pos\n            direction = direction / (np.linalg.norm(direction) + 1e-10)\n            \n            # Apply partial adjustment based on constraint strength\n            adjustment = (target_dist - current_dist) * constraint_strength\n            \n            # Only adjust the next position to preserve the overall fold\n            refined_coords[i+1] = current_pos + direction * (current_dist + adjustment)\n    \n    # 2. Steric clash prevention\n    min_allowed_distance = 3.8  # Minimum distance between non-consecutive C1' atoms\n    \n    # Calculate all pairwise distances\n    dist_matrix = distance_matrix(refined_coords, refined_coords)\n    \n    # Find severe clashes (atoms too close)\n    severe_clashes = np.where((dist_matrix < min_allowed_distance) & (dist_matrix > 0))\n    \n    # Fix severe clashes\n    for idx in range(len(severe_clashes[0])):\n        i, j = severe_clashes[0][idx], severe_clashes[1][idx]\n        \n        # Skip consecutive nucleotides and previously processed pairs\n        if abs(i - j) <= 1 or i >= j:\n            continue\n            \n        # Get current positions and distance\n        pos_i = refined_coords[i]\n        pos_j = refined_coords[j]\n        current_dist = dist_matrix[i, j]\n        \n        # Calculate necessary adjustment but scale by constraint strength\n        direction = pos_j - pos_i\n        direction = direction / (np.linalg.norm(direction) + 1e-10)\n        \n        # Calculate partial adjustment\n        adjustment = (min_allowed_distance - current_dist) * constraint_strength\n        \n        # Move points apart\n        refined_coords[i] = pos_i - direction * (adjustment / 2)\n        refined_coords[j] = pos_j + direction * (adjustment / 2)\n    \n    return refined_coords\n\ndef adapt_template_to_query(query_seq, template_seq, template_coords, alignment=None):\n    if alignment is None:\n        from Bio.Seq import Seq\n        from Bio import pairwise2\n        \n        query_seq_obj = Seq(query_seq)\n        template_seq_obj = Seq(template_seq)\n        alignments = pairwise2.align.globalms(query_seq_obj, template_seq_obj, 2.9, -1, -10, -0.5, one_alignment_only=True)\n        \n        if not alignments:\n            return generate_improved_rna_structure(query_seq)\n            \n        alignment = alignments[0]\n    \n    aligned_query = alignment.seqA\n    aligned_template = alignment.seqB\n    \n    query_coords = np.zeros((len(query_seq), 3))\n    query_coords.fill(np.nan)\n    \n    # Map template coordinates to query\n    query_idx = 0\n    template_idx = 0\n    \n    for i in range(len(aligned_query)):\n        query_char = aligned_query[i]\n        template_char = aligned_template[i]\n        \n        if query_char != '-' and template_char != '-':\n            if template_idx < len(template_coords):\n                query_coords[query_idx] = template_coords[template_idx]\n            template_idx += 1\n            query_idx += 1\n        elif query_char != '-' and template_char == '-':\n            query_idx += 1\n        elif query_char == '-' and template_char != '-':\n            template_idx += 1\n    \n    # IMPROVED GAP FILLING - maintains RNA backbone geometry\n    backbone_distance = 5.9  # Typical C1'-C1' distance\n    \n    # Fill gaps by maintaining realistic backbone connectivity\n    for i in range(len(query_coords)):\n        if np.isnan(query_coords[i, 0]):\n            # Find nearest valid neighbors\n            prev_valid = next_valid = None\n            \n            for j in range(i-1, -1, -1):\n                if not np.isnan(query_coords[j, 0]):\n                    prev_valid = j\n                    break\n                    \n            for j in range(i+1, len(query_coords)):\n                if not np.isnan(query_coords[j, 0]):\n                    next_valid = j\n                    break\n            \n            if prev_valid is not None and next_valid is not None:\n                # Interpolate along realistic RNA backbone path\n                gap_size = next_valid - prev_valid\n                total_distance = np.linalg.norm(query_coords[next_valid] - query_coords[prev_valid])\n                expected_distance = gap_size * backbone_distance\n                \n                # If gap is compressed, extend it realistically\n                if total_distance < expected_distance * 0.7:\n                    direction = query_coords[next_valid] - query_coords[prev_valid]\n                    direction = direction / (np.linalg.norm(direction) + 1e-10)\n                    \n                    # Place intermediate points along extended path\n                    for k, idx in enumerate(range(prev_valid + 1, next_valid)):\n                        progress = (k + 1) / gap_size\n                        base_pos = query_coords[prev_valid] + direction * expected_distance * progress\n                        \n                        # Add slight curvature for realism\n                        perpendicular = np.cross(direction, [0, 0, 1])\n                        if np.linalg.norm(perpendicular) < 1e-6:\n                            perpendicular = np.cross(direction, [1, 0, 0])\n                        perpendicular = perpendicular / (np.linalg.norm(perpendicular) + 1e-10)\n                        \n                        curve_amplitude = 2.0 * np.sin(progress * np.pi)\n                        query_coords[idx] = base_pos + perpendicular * curve_amplitude\n                else:\n                    # Linear interpolation for normal gaps\n                    for k, idx in enumerate(range(prev_valid + 1, next_valid)):\n                        weight = (k + 1) / gap_size\n                        query_coords[idx] = (1 - weight) * query_coords[prev_valid] + weight * query_coords[next_valid]\n            \n            elif prev_valid is not None:\n                # Extend from previous position\n                if prev_valid > 0 and not np.isnan(query_coords[prev_valid-1, 0]):\n                    direction = query_coords[prev_valid] - query_coords[prev_valid-1]\n                    direction = direction / (np.linalg.norm(direction) + 1e-10)\n                else:\n                    direction = np.array([1.0, 0.0, 0.0])\n                \n                steps_needed = i - prev_valid\n                for step in range(1, steps_needed + 1):\n                    pos_idx = prev_valid + step\n                    if pos_idx < len(query_coords):\n                        query_coords[pos_idx] = query_coords[prev_valid] + direction * backbone_distance * step\n            \n            elif next_valid is not None:\n                # Work backwards from next position\n                direction = np.array([-1.0, 0.0, 0.0])  # Default backward direction\n                steps_needed = next_valid - i\n                for step in range(steps_needed, 0, -1):\n                    pos_idx = next_valid - step\n                    if pos_idx >= 0:\n                        query_coords[pos_idx] = query_coords[next_valid] - direction * backbone_distance * step\n    \n    # Final cleanup\n    query_coords = np.nan_to_num(query_coords)\n    return query_coords\n\n\n# ========== generate_improved_rna_structure ========================\ndef generate_improved_rna_structure(sequence):\n    \"\"\"\n    Generate a more realistic RNA structure fallback based on sequence patterns\n    and basic RNA structure principles.\n    \n    Args:\n        sequence: RNA sequence string\n        \n    Returns:\n        Array of 3D coordinates\n    \"\"\"\n    n_residues = len(sequence)\n    coordinates = np.zeros((n_residues, 3))\n    \n    # Analyze sequence to predict structural elements\n    # Look for complementary regions that could form base pairs\n    potential_stems = identify_potential_stems(sequence)\n    \n    # Default parameters\n    radius_helix = 10.0\n    radius_loop = 15.0\n    rise_per_residue_helix = 2.5\n    rise_per_residue_loop = 1.5\n    angle_per_residue_helix = 0.6\n    angle_per_residue_loop = 0.3\n    \n    # Assign structural classifications\n    structure_types = assign_structure_types(sequence, potential_stems)\n    \n    # Generate coordinates based on predicted structure\n    current_pos = np.array([0.0, 0.0, 0.0])\n    current_direction = np.array([0.0, 0.0, 1.0])\n    current_angle = 0.0\n    \n    for i in range(n_residues):\n        if structure_types[i] == 'stem':\n            # Part of a helical stem\n            current_angle += angle_per_residue_helix\n            coordinates[i] = [\n                radius_helix * np.cos(current_angle), \n                radius_helix * np.sin(current_angle), \n                current_pos[2] + rise_per_residue_helix\n            ]\n            current_pos = coordinates[i]\n        elif structure_types[i] == 'loop':\n            # Part of a loop\n            current_angle += angle_per_residue_loop\n            z_shift = rise_per_residue_loop * np.sin(current_angle * 0.5)\n            coordinates[i] = [\n                radius_loop * np.cos(current_angle), \n                radius_loop * np.sin(current_angle), \n                current_pos[2] + z_shift\n            ]\n            current_pos = coordinates[i]\n        else:\n            # Single-stranded region\n            # Add some randomness to make it look more realistic\n            jitter = np.random.normal(0, 1, 3) * 2.0\n            coordinates[i] = current_pos + jitter\n            current_pos = coordinates[i]\n            \n    return coordinates\n\ndef identify_potential_stems(sequence):\n    \"\"\"\n    Identify potential stem regions by looking for self-complementary segments.\n    \n    Args:\n        sequence: RNA sequence string\n        \n    Returns:\n        List of tuples (start1, end1, start2, end2) representing potentially paired regions\n    \"\"\"\n    complementary_bases = {'A': 'U', 'U': 'A', 'G': 'C', 'C': 'G'}\n    min_stem_length = 3\n    potential_stems = []\n    \n    # Simple stem identification\n    for i in range(len(sequence) - min_stem_length):\n        for j in range(i + min_stem_length + 3, len(sequence) - min_stem_length + 1):\n            # Check if regions could form a stem\n            potential_stem_len = min(min_stem_length, len(sequence) - j)\n            is_stem = True\n            \n            for k in range(potential_stem_len):\n                if sequence[i+k] not in complementary_bases or \\\n                   complementary_bases[sequence[i+k]] != sequence[j+potential_stem_len-k-1]:\n                    is_stem = False\n                    break\n            \n            if is_stem:\n                potential_stems.append((i, i+potential_stem_len-1, j, j+potential_stem_len-1))\n    \n    return potential_stems\n\ndef assign_structure_types(sequence, potential_stems):\n    \"\"\"\n    Assign each nucleotide to a structural element type.\n    \n    Args:\n        sequence: RNA sequence string\n        potential_stems: List of tuples representing stem regions\n        \n    Returns:\n        List of structure types ('stem', 'loop', 'single')\n    \"\"\"\n    structure_types = ['single'] * len(sequence)\n    \n    # Mark stem regions\n    for stem in potential_stems:\n        start1, end1, start2, end2 = stem\n        for i in range(end1 - start1 + 1):\n            structure_types[start1 + i] = 'stem'\n            structure_types[end2 - i] = 'stem'\n    \n    # Mark loop regions (regions between paired regions)\n    for i in range(len(potential_stems) - 1):\n        _, end1, start2, _ = potential_stems[i]\n        next_start1, _, _, _ = potential_stems[i+1]\n        \n        if next_start1 > end1 + 1 and start2 > next_start1:\n            for j in range(end1 + 1, next_start1):\n                structure_types[j] = 'loop'\n    \n    return structure_types\n\n\n# =========== generate_rna_structure ======================\ndef generate_rna_structure(sequence, seed=None):\n    \"\"\"Generate a more realistic RNA structure when no good templates are found\"\"\"\n    if seed is not None:\n        np.random.seed(seed)\n        random.seed(seed)\n    \n    n_residues = len(sequence)\n    coordinates = np.zeros((n_residues, 3))\n    \n    # Initialize the first few residues in a helix\n    for i in range(min(3, n_residues)):\n        angle = i * 0.6\n        coordinates[i] = [10.0 * np.cos(angle), 10.0 * np.sin(angle), i * 2.5]\n    \n    # Add more complex folding patterns\n    current_direction = np.array([0.0, 0.0, 1.0])  # Start moving along z-axis\n    \n    # Define base-pairing tendencies (G-C and A-U pairs)\n    for i in range(3, n_residues):\n        # Check for potential base-pairing in the sequence\n        has_pair = False\n        pair_idx = -1\n        \n        # Simple detection of complementary bases (G-C, A-U)\n        complementary = {'G': 'C', 'C': 'G', 'A': 'U', 'U': 'A'}\n        current_base = sequence[i]\n        \n        # Look for potential base-pairing within a window before the current position\n        window_size = min(i, 15)  # Look back up to 15 bases\n        for j in range(i-window_size, i):\n            if j >= 0 and sequence[j] == complementary.get(current_base, 'X'):\n                # Found a potential pair\n                has_pair = True\n                pair_idx = j\n                break\n        \n        if has_pair and i - pair_idx <= 10 and random.random() < 0.7:\n            # Try to create a base-pair by positioning this nucleotide near its pair\n            pair_pos = coordinates[pair_idx]\n            \n            # Create a position that's roughly opposite to the pair\n            random_offset = np.random.normal(0, 1, 3) * 2.0\n            base_pair_distance = 10.0 + random.uniform(-1.0, 1.0)\n            \n            # Calculate a vector from base-pair toward center of structure\n            center = np.mean(coordinates[:i], axis=0)\n            direction = center - pair_pos\n            direction = direction / (np.linalg.norm(direction) + 1e-10)\n            \n            # Position new nucleotide in the general direction of the \"center\"\n            coordinates[i] = pair_pos + direction * base_pair_distance + random_offset\n            \n            # Update direction for next nucleotide\n            current_direction = np.random.normal(0, 0.3, 3)\n            current_direction = current_direction / (np.linalg.norm(current_direction) + 1e-10)\n            \n        else:\n            # No base-pairing detected, continue with the current fold direction\n            # Randomly rotate current direction to simulate RNA flexibility\n            if random.random() < 0.3:\n                # More significant direction change\n                angle = random.uniform(0.2, 0.6)\n                axis = np.random.normal(0, 1, 3)\n                axis = axis / (np.linalg.norm(axis) + 1e-10)\n                rotation = R.from_rotvec(angle * axis)\n                current_direction = rotation.apply(current_direction)\n            else:\n                # Small random changes in direction\n                current_direction += np.random.normal(0, 0.15, 3)\n                current_direction = current_direction / (np.linalg.norm(current_direction) + 1e-10)\n            \n            # Distance between consecutive nucleotides (3.5-4.5Å is typical)\n            step_size = random.uniform(3.5, 4.5)\n            \n            # Update position\n            coordinates[i] = coordinates[i-1] + step_size * current_direction\n    \n    return coordinates\n\n\n# ========== predict_rna_structures ==================\ndef predict_rna_structures(sequence, target_id, train_seqs_df, train_coords_dict, n_predictions=5):\n    predictions = []\n    \n    # Find similar sequences in the training data\n    similar_seqs = find_similar_sequences(sequence, train_seqs_df, train_coords_dict, top_n=n_predictions)\n    \n    # If we found any similar sequences, use them as templates\n    if similar_seqs:\n        for i, (template_id, template_seq, similarity_score, template_coords) in enumerate(similar_seqs):\n            # Adapt template coordinates to the query sequence\n            adapted_coords = adapt_template_to_query(sequence, template_seq, template_coords)\n            \n            if adapted_coords is not None:\n                # Apply adaptive constraints based on template similarity\n                # For high similarity templates, apply very gentle constraints\n                refined_coords = adaptive_rna_constraints(adapted_coords, sequence, confidence=similarity_score)\n                \n                # Add some randomness (less for better templates)\n                random_scale = max(0.05, 0.8 - similarity_score)  # Reduced randomness\n                randomized_coords = refined_coords.copy()\n                randomized_coords += np.random.normal(0, random_scale, randomized_coords.shape)\n                \n                predictions.append(randomized_coords)\n                \n                if len(predictions) >= n_predictions:\n                    break\n    \n    # If we don't have enough predictions from templates, generate de novo structures\n    while len(predictions) < n_predictions:\n        seed_value = hash(target_id) % 10000 + len(predictions) * 1000\n        de_novo_coords = generate_rna_structure(sequence, seed=seed_value)\n        \n        # Apply stronger constraints to de novo structures (lower confidence)\n        refined_de_novo = adaptive_rna_constraints(de_novo_coords, sequence, confidence=0.2)\n        \n        predictions.append(refined_de_novo)\n    \n    return predictions[:n_predictions]","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.046Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Initialize counters and range settings\nif is_submission_mode:\n    DRFOLD_START_IDX = 14\n    DRFOLD_END_IDX = len(test_sequences) - 1\nelse:\n    DRFOLD_START_IDX = 0\n    DRFOLD_END_IDX = 0\n\ndrfold_processed = 0\ntemplate_processed = 0\n\ntrain_coords_dict = process_labels_vectorized(train_labels_final)","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.046Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# from IPython.display import clear_output","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.046Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Sort test sequences by length to process shorter ones with DRfold2\ntest_sequences = test_sequences.sort_values(by=['sequence'], key=lambda x: x.str.len())\n\n# List to store all prediction records\nall_predictions = []\n\n# Set up time tracking\nstart_time = time.time()\ntotal_targets = len(test_sequences)\n\n# For each sequence in the test set\nfor idx, row in test_sequences.iterrows():\n    target_id = row['target_id']\n    sequence = row['sequence']\n    \n    # Progress tracking\n    elapsed = time.time() - start_time\n    targets_processed = idx\n    if targets_processed > 0:\n        avg_time_per_target = elapsed / targets_processed\n        est_time_remaining = avg_time_per_target * (total_targets - targets_processed)\n        time_left = DRFOLD_TIME_LIMIT - (time.time() - start_time_global)\n        print(f\"Processing target {targets_processed+1}/{total_targets}: {target_id} ({len(sequence)} nt), \"\n              f\"elapsed: {elapsed:.1f}s, est. remaining: {est_time_remaining:.1f}s, time left: {time_left:.1f}s\")\n    \n    # Check if we should use DRfold2 or template-based approach\n    use_drfold = (DRFOLD_START_IDX <= idx <= DRFOLD_END_IDX and \n                 (time.time() - start_time_global) < DRFOLD_TIME_LIMIT)\n    \n    # Generate 5 different structure predictions\n    if use_drfold:\n        print(f\"Using DRfold2 for target {target_id} (index {idx})\")\n\n        cif_file = f\"/kaggle/working/outputs_prediction/boltz_results_inputs_prediction/predictions/{target_id}/{target_id}_model_0.cif\"\n        pdb_file = cif_file.replace('.cif', '.pdb')\n        \n        if convert_cif_to_pdb(cif_file, pdb_file):\n            # Use the converted PDB file with DRfold2\n            print(\"af3file generated\")\n            af3file = pdb_file\n        else:\n            print(\"Failed to convert mmCIF file, skipping AF3 integration\")\n            af3file = None\n\n\n        print(f\"Using af3file: {af3file}\")\n\n        # Without AlphaFold3 (original mode)\n        # predictions = predict_rna_structures_drfold2(sequence, target_id)\n\n        # With AlphaFold3 integration\n        predictions = predict_rna_structures_drfold2(sequence, target_id, af3_pdb=af3file)\n        \n        # If DRfold2 fails, fall back to template approach\n        if predictions is None:\n            print(f\"DRfold2 failed for {target_id}, falling back to template approach\")\n            predictions = predict_rna_structures(sequence, target_id, train_seqs_final, train_coords_dict)\n            template_processed += 1\n        else:\n            drfold_processed += 1\n    else:\n        if idx > DRFOLD_END_IDX:\n            reason = \"index out of DRfold range\"\n        elif idx < DRFOLD_START_IDX:\n            reason = \"index before DRfold start range\"\n        else:\n            reason = \"time limit reached\"\n        print(f\"Using template approach for {target_id} ({reason})\")\n        predictions = predict_rna_structures(sequence, target_id, train_seqs_final, train_coords_dict)\n        template_processed += 1\n    \n    # For each residue in the sequence\n    for j in range(len(sequence)):\n        pred_row = {\n            'ID': f\"{target_id}_{j+1}\",\n            'resname': sequence[j],\n            'resid': j + 1\n        }\n        \n        # Add coordinates from all 5 predictions\n        for i in range(5):\n            pred_row[f'x_{i+1}'] = predictions[i][j][0]\n            pred_row[f'y_{i+1}'] = predictions[i][j][1]\n            pred_row[f'z_{i+1}'] = predictions[i][j][2]\n        \n        all_predictions.append(pred_row)\n\n    # clear_output(wait=False)\n    \n    # Free up memory\n    if torch.cuda.is_available():\n        torch.cuda.empty_cache()\n\n# Create DataFrame with predictions\nsubmission_df = pd.DataFrame(all_predictions)\n\n# Ensure the submission file has the correct format\ncolumn_order = ['ID', 'resname', 'resid']\nfor i in range(1, 6):\n    for coord in ['x', 'y', 'z']:\n        column_order.append(f'{coord}_{i}')\n        \nsubmission_df = submission_df[column_order]\n\n# Save the submission\nsubmission_df.to_csv('/kaggle/working/submission_dr.csv', index=False)\nprint(f\"Generated predictions for {len(test_sequences)} RNA sequences\")\nprint(f\"Used DRfold2 for {drfold_processed} targets and template approach for {template_processed} targets\")\nprint(f\"Total runtime: {time.time() - start_time_global:.1f} seconds\")","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.046Z"},"scrolled":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.046Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_df","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.047Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Post Processing:  \n---","metadata":{}},{"cell_type":"code","source":"%cd /kaggle/working/","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.047Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!ls","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.047Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\n# Load the data\nsubmission_boltz = pd.read_csv('submission_boltz.csv')\nsubmission_dr = pd.read_csv('submission_dr.csv')\n\n# Function to extract target_id from ID column\ndef extract_target_id(id_string):\n    \"\"\"Extract target_id by removing the suffix (everything after the last underscore)\"\"\"\n    return id_string.rsplit('_', 1)[0]\n\n# Add target_id column to both dataframes\nsubmission_boltz['target_id'] = submission_boltz['ID'].apply(extract_target_id)\nsubmission_dr['target_id'] = submission_dr['ID'].apply(extract_target_id)\n\n# Calculate sequence lengths\ntest_sequences['sequence_length'] = test_sequences['sequence'].str.len()\n\n# Get target_ids with sequence length > 600\nlong_sequences = test_sequences[test_sequences['sequence_length'] > 600]['target_id'].tolist()\n\nprint(f\"Found {len(long_sequences)} sequences with length > 600:\")\nfor seq_id in long_sequences:\n    seq_len = test_sequences[test_sequences['target_id'] == seq_id]['sequence_length'].iloc[0]\n    print(f\"  {seq_id}: {seq_len} nucleotides\")\n\n# Create a copy of submission_dr for modification\nsubmission_processed = submission_dr.copy()\n\n# Define a threshold for detecting placeholder values (e.g., values < -1e10)\nPLACEHOLDER_THRESHOLD = -1e17\n\n# For each row in submission_dr, check if it belongs to a long sequence\nfor idx, row in submission_dr.iterrows():\n    target_id = row['target_id']\n    \n    if target_id in long_sequences:\n        # Find corresponding row in submission_boltz\n        boltz_row = submission_boltz[submission_boltz['ID'] == row['ID']]\n        \n        if not boltz_row.empty:\n            boltz_x1 = boltz_row.iloc[0]['x_1']\n            boltz_y1 = boltz_row.iloc[0]['y_1']\n            boltz_z1 = boltz_row.iloc[0]['z_1']\n            \n            # Always replace the first conformation\n            submission_processed.loc[idx, 'x_1'] = boltz_x1\n            submission_processed.loc[idx, 'y_1'] = boltz_y1\n            submission_processed.loc[idx, 'z_1'] = boltz_z1\n            \n            # Check and replace placeholder values in other conformations\n            for conf_num in [2, 3, 4, 5]:\n                x_col = f'x_{conf_num}'\n                y_col = f'y_{conf_num}'\n                z_col = f'z_{conf_num}'\n                \n                # If any coordinate in this conformation is a placeholder, replace all three\n                if (row[x_col] < PLACEHOLDER_THRESHOLD or \n                    row[y_col] < PLACEHOLDER_THRESHOLD or \n                    row[z_col] < PLACEHOLDER_THRESHOLD):\n                    \n                    submission_processed.loc[idx, x_col] = boltz_x1\n                    submission_processed.loc[idx, y_col] = boltz_y1\n                    submission_processed.loc[idx, z_col] = boltz_z1\n\n# Remove the temporary target_id column\nsubmission_processed = submission_processed.drop('target_id', axis=1)","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.047Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_processed","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.047Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!ls","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.047Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%cd /kaggle/working/","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.047Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!rm -rf ./*","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.048Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!ls","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.048Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Save the processed sub\nsubmission_processed.to_csv('submission.csv', index=False)\n\n# Print summary of changes\ntotal_rows_modified = 0\nfor target_id in long_sequences:\n    rows_for_target = submission_dr[submission_dr['target_id'] == target_id].shape[0]\n    total_rows_modified += rows_for_target\n    print(f\"Modified {rows_for_target} rows for {target_id}\")\n\nprint(f\"\\nTotal rows modified: {total_rows_modified}\")\nprint(f\"Total rows in submission_dr: {len(submission_dr)}\")\nprint(f\"Processed submission_dr saved as 'submission.csv'\")","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.048Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!ls","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.048Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_processed","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-10T08:14:04.048Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}