{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":87793,"databundleVersionId":12276181,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":5123458,"sourceType":"datasetVersion","datasetId":2975803},{"sourceId":11084789,"sourceType":"datasetVersion","datasetId":6908845},{"sourceId":11118830,"sourceType":"datasetVersion","datasetId":6933267},{"sourceId":11413730,"sourceType":"datasetVersion","datasetId":7148547},{"sourceId":11421128,"sourceType":"datasetVersion","datasetId":7152794},{"sourceId":11428421,"sourceType":"datasetVersion","datasetId":7157768},{"sourceId":11593585,"sourceType":"datasetVersion","datasetId":7270084},{"sourceId":11695366,"sourceType":"datasetVersion","datasetId":6791615},{"sourceId":11729214,"sourceType":"datasetVersion","datasetId":7362800},{"sourceId":11916147,"sourceType":"datasetVersion","datasetId":7491229},{"sourceId":11960429,"sourceType":"datasetVersion","datasetId":7520552},{"sourceId":11989607,"sourceType":"datasetVersion","datasetId":7541175},{"sourceId":11990395,"sourceType":"datasetVersion","datasetId":7541717},{"sourceId":11999509,"sourceType":"datasetVersion","datasetId":7548267}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import sys\nimport os\nimport shutil\nfrom pathlib import Path\nimport subprocess\nimport glob\nimport gc","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:00:36.063162Z","iopub.execute_input":"2025-05-26T09:00:36.063447Z","iopub.status.idle":"2025-05-26T09:00:36.068465Z","shell.execute_reply.started":"2025-05-26T09:00:36.063413Z","shell.execute_reply":"2025-05-26T09:00:36.067457Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"val = False\nif val:\n    !pip install --no-deps protenix\n    !pip install biopython\n    !pip install ml-collections\n    !pip install biotite==1.0.1\n    !pip install rdkit\n\nsrc = \"/kaggle/input/arena-with-permission/Arena\"\ndst = \"/kaggle/working/Arena\"\n\nshutil.copytree(src, dst)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:00:36.069768Z","iopub.execute_input":"2025-05-26T09:00:36.070344Z","iopub.status.idle":"2025-05-26T09:00:57.558121Z","shell.execute_reply.started":"2025-05-26T09:00:36.070315Z","shell.execute_reply":"2025-05-26T09:00:57.557164Z"}},"outputs":[],"execution_count":null},{"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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:01:27.718402Z","iopub.execute_input":"2025-05-26T09:01:27.718686Z","iopub.status.idle":"2025-05-26T09:01:37.689642Z","shell.execute_reply.started":"2025-05-26T09:01:27.718663Z","shell.execute_reply":"2025-05-26T09:01:37.688834Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"os.chmod('/kaggle/working/Arena/Arena', 0o755)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:01:37.690943Z","iopub.execute_input":"2025-05-26T09:01:37.691253Z","iopub.status.idle":"2025-05-26T09:01:37.695508Z","shell.execute_reply.started":"2025-05-26T09:01:37.691219Z","shell.execute_reply":"2025-05-26T09:01:37.694699Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport os, sys\nimport pandas as pd\nimport numpy as np\nfrom Bio import SeqIO, AlignIO\nfrom Bio.Seq import Seq\nfrom Bio.SeqRecord import SeqRecord\nfrom Bio.PDB import Atom, Model, Chain, Residue, Structure, PDBParser\nimport json\nimport time\n\nprint('torch',torch.__version__)\nprint('torch.cuda',torch.version.cuda)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:01:37.697177Z","iopub.execute_input":"2025-05-26T09:01:37.697374Z","iopub.status.idle":"2025-05-26T09:01:41.131665Z","shell.execute_reply.started":"2025-05-26T09:01:37.697356Z","shell.execute_reply":"2025-05-26T09:01:41.1308Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"local = False","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:01:41.133215Z","iopub.execute_input":"2025-05-26T09:01:41.133667Z","iopub.status.idle":"2025-05-26T09:01:41.137113Z","shell.execute_reply.started":"2025-05-26T09:01:41.133642Z","shell.execute_reply":"2025-05-26T09:01:41.136387Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"PROTENIX_DATA_ROOT_DIR = '/kaggle/input/protenix-checkpoints'\nLONG_MODEL_PATH = f'{PROTENIX_DATA_ROOT_DIR}/model_v0.2.0.pt'\nSHORT_MODEL_PATH = f'{PROTENIX_DATA_ROOT_DIR}/model_v0.2.0.pt'\n# LONG_MODEL_PATH = '/kaggle/input/rna16999/16999_ema_0.995.pt'\n# SHORT_MODEL_PATH = '/kaggle/input/rna16999/16999_ema_0.995.pt'\nLONG_MODEL_PATH = '/kaggle/input/allrnaonly/4019_ema_0.995.pt'\nSHORT_MODEL_PATH = '/kaggle/input/rnaonly19799/19799_ema_0.995.pt'\nSEQUENCE_LENGTH_THRESHOLD = 1","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:01:41.138018Z","iopub.execute_input":"2025-05-26T09:01:41.138205Z","iopub.status.idle":"2025-05-26T09:01:41.152845Z","shell.execute_reply.started":"2025-05-26T09:01:41.138188Z","shell.execute_reply":"2025-05-26T09:01:41.152112Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"! mkdir /af3-dev \n! ln -s /kaggle/input/protenix-checkpoints /af3-dev/release_data\n! ls /af3-dev/release_data/","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:01:41.153472Z","iopub.execute_input":"2025-05-26T09:01:41.153676Z","iopub.status.idle":"2025-05-26T09:01:41.535628Z","shell.execute_reply.started":"2025-05-26T09:01:41.153657Z","shell.execute_reply":"2025-05-26T09:01:41.534823Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if not local:\n    KAGGLE_INPUT = '/kaggle/input/drfold2-12/DRfold2_optimized'\n    WORKING_DIR = '/kaggle/working'\n    DRFOLD_DIR = WORKING_DIR\n    \n\n    # if not os.path.exists(DRFOLD_DIR):\n    #     print(f\"Copying DRfold2 from {KAGGLE_INPUT} to {DRFOLD_DIR}\")\n    #     shutil.copytree(KAGGLE_INPUT, DRFOLD_DIR)\n\n\n    os.chdir(DRFOLD_DIR)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:01:41.536623Z","iopub.execute_input":"2025-05-26T09:01:41.536898Z","iopub.status.idle":"2025-05-26T09:01:41.540935Z","shell.execute_reply.started":"2025-05-26T09:01:41.536863Z","shell.execute_reply":"2025-05-26T09:01:41.540165Z"}},"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\nfrom boltz.main import BoltzSteeringParams\n\nsteering_args = BoltzSteeringParams(\n    fk_steering=True,\n    num_particles=3,\n    fk_lambda=4.0,\n    fk_resampling_interval=3,\n    guidance_update=True,\n    num_gd_steps=16,\n)\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        steering_args=asdict(steering_args),\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=\"./boltz\",\n            diffusion_samples=1,\n            seed=42,\n            override=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:01:47.370479Z","iopub.execute_input":"2025-05-26T09:01:47.370785Z","iopub.status.idle":"2025-05-26T09:01:47.379369Z","shell.execute_reply.started":"2025-05-26T09:01:47.37076Z","shell.execute_reply":"2025-05-26T09:01:47.378514Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%cp -rf /kaggle/input/boltz-codes/boltz .\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:01:52.272603Z","iopub.execute_input":"2025-05-26T09:01:52.272943Z","iopub.status.idle":"2025-05-26T09:02:22.091732Z","shell.execute_reply.started":"2025-05-26T09:01:52.272914Z","shell.execute_reply":"2025-05-26T09:02:22.090401Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BOLTZ_INPUT_DIR = '/kaggle/working/inputs_prediction'\nBOLTZ_OUTPUT_DIR = '/kaggle/working/outputs_prediction'\nBOLTZ_SCRIPT = '/kaggle/working/inference.py'\nBOLTZ_CONFIDENCE_KEY = 'confidence_score'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:02:22.093061Z","iopub.execute_input":"2025-05-26T09:02:22.093382Z","iopub.status.idle":"2025-05-26T09:02:22.09718Z","shell.execute_reply.started":"2025-05-26T09:02:22.093345Z","shell.execute_reply":"2025-05-26T09:02:22.096533Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from os import listdir\nfrom os.path import isfile, join\nfrom copy import deepcopy\n\n# sto generation\nmsa_dir = \"/kaggle/input/stanford-rna-3d-folding/MSA/\"\nsto_dir = \"/kaggle/working/MSA-STO/\"\n# os.makedirs(sto_dir, exist_ok=True)\n\nsrc = \"/kaggle/input/sto-files/MSA-STO\"\ndst = \"/kaggle/working/MSA-STO\"\n\nshutil.copytree(src, dst)\n\nmsa_files = [f for f in listdir(msa_dir) if isfile(join(msa_dir, f))]\nfor file_name in msa_files:\n    target_name = file_name[:-6]\n\n    if isfile(f\"{sto_dir}/{target_name}.sto\"):\n        continue\n\n    # Read the FASTA MSA\n    alignment = AlignIO.read(f\"{msa_dir}{file_name}\", \"fasta\")\n    # Write to Stockholm format\n    AlignIO.write(alignment, f\"{sto_dir}/{target_name}.sto\", \"stockholm\")\n\nsto_cut_bound = 0\nfor file_name in msa_files:\n    target_name = file_name[:-6]\n    sto_file_path = f\"{sto_dir}/{target_name}.sto\"\n\n    query_count = 0\n    with open(sto_file_path, \"r\") as f:\n        lines = f.readlines()\n        query_count = min((len(lines) - 3) // 3, 0)\n\n    if query_count < sto_cut_bound:\n        # delete file\n        print(f\"delete {target_name}.sto under sto_cut_bound\")\n        os.remove(sto_file_path)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:02:22.098842Z","iopub.execute_input":"2025-05-26T09:02:22.099099Z","iopub.status.idle":"2025-05-26T09:02:45.126189Z","shell.execute_reply.started":"2025-05-26T09:02:22.09908Z","shell.execute_reply":"2025-05-26T09:02:45.125447Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_boltz_c1_coords(tmp_id, idx):\n    \"\"\"Extract C1' coordinates from Boltz CIF output file\"\"\"\n    from Bio.PDB.MMCIF2Dict import MMCIF2Dict\n    cif_file = f\"{BOLTZ_OUTPUT_DIR}/boltz_results_inputs_prediction/predictions/{tmp_id}/{tmp_id}_model_{idx}.cif\"\n    \n    # Check if file exists\n    if not os.path.exists(cif_file):\n        print(f\"Warning: Could not find CIF file {cif_file}\")\n        return []\n        \n    mmcif_dict = MMCIF2Dict(cif_file)\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    \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    \n    return c1_coords\n\ndef get_boltz_top5_indices(tmp_id, top_num):\n    topwhat = top_num\n    \"\"\"Get top 5 model indices from Boltz by confidence score\"\"\"\n    conf_dir = f\"{BOLTZ_OUTPUT_DIR}/boltz_results_inputs_prediction/predictions/{tmp_id}\"\n    conf_scores = []\n    \n    for idx in range(topwhat): #######################################################################################change here\n        conf_path = os.path.join(conf_dir, f\"confidence_{tmp_id}_model_{idx}.json\")\n        if os.path.exists(conf_path):\n            with open(conf_path) as f:\n                conf = json.load(f)\n                conf_scores.append((idx, conf[BOLTZ_CONFIDENCE_KEY]))\n    \n    # If no confidence files found\n    if not conf_scores:\n        print(f\"Warning: No confidence scores found for {tmp_id}, using default order\")\n        return list(range(topwhat))  # Use first 5 models\n        \n    # Sort by confidence (highest first) and take top 5\n    top5 = sorted(conf_scores, key=lambda x: x[1], reverse=True)[:topwhat]\n    return [idx for idx, _ in top5]\n\n\ndef prepare_boltz_inputs(names, sequences):\n    \"\"\"Prepare YAML input files for Boltz\"\"\"\n    os.makedirs(BOLTZ_INPUT_DIR, exist_ok=True)\n    for tmp_id, tmp_sequence in zip(names, sequences):\n        with open(f'{BOLTZ_INPUT_DIR}/{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}\")\n\n\ndef extract_c1_coords(pdb_file, target_id, model_idx=0):\n    \"\"\"Extract C1' atom coordinates from PDB file for a specific model index.\n    \n    Args:\n        pdb_file: Path to the PDB file\n        target_id: Target ID for the RNA sequence\n        model_idx: Model index (0-4) corresponding to which model this is (will be mapped to model_idx+1 in output)\n    \"\"\"\n    coords = []\n    residue_ids = []\n    residue_names = []\n    \n    with open(pdb_file, 'r') as f:\n        for line in f:\n            if line.startswith('ATOM'):\n                atom_name = line[12:16].strip()\n                residue_name = line[17:20].strip()\n                residue_id = int(line[22:26].strip())\n                \n                if atom_name == \"C1'\":\n                    x = float(line[30:38].strip())\n                    y = float(line[38:46].strip())\n                    z = float(line[46:54].strip())\n                    \n                    residue_ids.append(residue_id)\n                    residue_names.append(residue_name)\n                    coords.append([x, y, z])\n    \n    # Process results\n    results = {}\n    for i, (res_id, res_name, coord) in enumerate(zip(residue_ids, residue_names, coords)):\n        # Use the first letter of the residue name\n        res_code = res_name[0]\n        \n        # Create key for each residue\n        key = f\"{target_id}_{res_id}\"\n        \n        # Initialize entry or update existing entry\n        if key not in results:\n            results[key] = {\n                'ID': key,\n                'resname': res_code,\n                'resid': res_id\n            }\n            \n        # Add coordinates for this specific model (model_idx + 1 because models are 1-indexed in output)\n        results[key][f'x_{model_idx+1}'] = coord[0]\n        results[key][f'y_{model_idx+1}'] = coord[1]\n        results[key][f'z_{model_idx+1}'] = coord[2]\n    \n    return results\n\ndef create_fasta_from_sequence(target_id, sequence, fasta_path, max_length=400):\n    \"\"\"Create a FASTA file from sequence, with sequence on a single line\"\"\"\n    # Truncate sequence if it exceeds max_length\n    truncated = False\n    if len(sequence) > max_length:\n        truncated_sequence = sequence[:max_length]\n        truncated = True\n    else:\n        truncated_sequence = sequence\n        \n    record = SeqRecord(Seq(truncated_sequence), id=target_id, description=\"\")\n    with open(fasta_path, \"w\") as f:\n        # Write header and sequence manually\n        f.write(f\">{record.id} {record.description}\\n\")\n        f.write(f\"{str(record.seq)}\\n\")\n    \n    return fasta_path, truncated, len(sequence)\n\ndef clean_output_dirs(outdir, ret_dir, folddir, refdir):\n    \"\"\"Clean output directories between runs\"\"\"\n    for dir_path in [ret_dir, folddir, refdir]:\n        if os.path.exists(dir_path):\n            for file in os.listdir(dir_path):\n                file_path = os.path.join(dir_path, file)\n                if os.path.isfile(file_path):\n                    os.remove(file_path)\n                    \ndef parse_output_to_df(output, seq, target_id):\n    \"\"\"Convert Protenix output to DataFrame format\"\"\"\n    df = []\n    chain_data = []\n    for i, res in enumerate(seq):\n        d = dict(ID = f\"{target_id}_{i+1}\",\n                 resname=res,\n                 resid=i+1)\n        for n in range(len(output)):\n            d = {**d, \n                 f'x_{n+1}': round(output[n,i,0].item(),3),\n                 f'y_{n+1}': round(output[n,i,1].item(),3),\n                 f'z_{n+1}': round(output[n,i,2].item(),3)}\n        chain_data.append(d)\n\n    if len(chain_data) != 0:\n        chain_df = pd.DataFrame(chain_data)\n        df.append(chain_df)\n    return df\n\ndef run_with_time(cmd, description):\n    \"\"\"Run a command and measure its execution time\"\"\"\n    print(f\"Running {description}...\")\n    print(f\"Command: {cmd}\")\n    start_time = time.time()\n    ret_code = os.system(cmd)\n    end_time = time.time()\n    elapsed = end_time - start_time\n    print(f\"✓ {description} completed in {elapsed:.2f} seconds (return code: {ret_code})\")\n    return ret_code, elapsed","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:02:45.127558Z","iopub.execute_input":"2025-05-26T09:02:45.127884Z","iopub.status.idle":"2025-05-26T09:02:45.144498Z","shell.execute_reply.started":"2025-05-26T09:02:45.127852Z","shell.execute_reply":"2025-05-26T09:02:45.143782Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"exp_dir = KAGGLE_INPUT\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(device)\n#dlexps = ['cfg_97']\ndlexps = ['cfg_95', 'cfg_96', 'cfg_97', 'cfg_99']\ndirs = [os.path.join(exp_dir, 'model_hub', one_exp) for one_exp in dlexps]\n\n# Setting up the fasta path\nfasta_path = './test/seq.fasta'\nos.makedirs(os.path.dirname(fasta_path), exist_ok=True)\n\n# Setup directories (use the original structure)\noutdir = './outputs'\nif not os.path.isdir(outdir):\n    os.makedirs(outdir)\n    \nret_dir = os.path.join(outdir, 'rets_dir')\nif not os.path.isdir(ret_dir):\n    os.makedirs(ret_dir)\n\nfolddir = os.path.join(outdir, 'folds')\nif not os.path.isdir(folddir):\n    os.makedirs(folddir)\n\nrefdir = os.path.join(outdir, 'relax')\nif not os.path.isdir(refdir):\n    os.makedirs(refdir)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:02:45.145399Z","iopub.execute_input":"2025-05-26T09:02:45.14575Z","iopub.status.idle":"2025-05-26T09:02:45.210498Z","shell.execute_reply.started":"2025-05-26T09:02:45.145699Z","shell.execute_reply":"2025-05-26T09:02:45.209792Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if local:\n    csv_path = '../data/rna_fold_kaggle_data/test_sequences.csv'\nelse:\n    csv_path = '/kaggle/input/stanford-rna-3d-folding/test_sequences.csv'\n\nsequences_df = pd.read_csv(csv_path)\nprint(f\"Loaded {len(sequences_df)} sequences from {csv_path}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:02:45.21128Z","iopub.execute_input":"2025-05-26T09:02:45.211485Z","iopub.status.idle":"2025-05-26T09:02:45.234654Z","shell.execute_reply.started":"2025-05-26T09:02:45.211467Z","shell.execute_reply":"2025-05-26T09:02:45.234027Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"exp_dir","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:02:45.235467Z","iopub.execute_input":"2025-05-26T09:02:45.235747Z","iopub.status.idle":"2025-05-26T09:02:45.240397Z","shell.execute_reply.started":"2025-05-26T09:02:45.235701Z","shell.execute_reply":"2025-05-26T09:02:45.239776Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"all_results = {}\ndlmains = [os.path.join(exp_dir, one_exp, 'test_modeldir.py') for one_exp in dlexps]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:02:45.242236Z","iopub.execute_input":"2025-05-26T09:02:45.242478Z","iopub.status.idle":"2025-05-26T09:02:45.252667Z","shell.execute_reply.started":"2025-05-26T09:02:45.242459Z","shell.execute_reply":"2025-05-26T09:02:45.251979Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import json\nimport os\nimport shutil\nfrom abc import ABC, abstractmethod\nfrom collections import defaultdict\nfrom copy import deepcopy\nfrom os.path import exists as opexists\nfrom os.path import join as opjoin\nfrom typing import Any, Mapping, Optional, Sequence, Union\n\nimport numpy as np\nimport torch\nfrom biotite.structure import AtomArray\n\nfrom protenix.data.constants import STD_RESIDUES, rna_order_with_x\nfrom protenix.data.msa_utils import (\n    PROT_TYPE_NAME,\n    FeatureDict,\n    add_assembly_features,\n    clip_msa,\n    convert_monomer_features,\n    get_identifier_func,\n    load_and_process_msa,\n    make_sequence_features,\n    merge_features_from_prot_rna,\n    msa_parallel,\n    pair_and_merge,\n    rna_merge,\n)\nfrom typing import Any, Mapping\nfrom biotite.structure import AtomArray\nfrom protenix.data.json_to_feature import SampleDictToFeatures\nfrom protenix.data.data_pipeline import DataPipeline\nfrom protenix.data.msa_featurizer import InferenceMSAFeaturizer, process_single_sequence, SEQ_LIMITS, tokenize_msa, convert_monomer_features\nfrom protenix.data.utils import data_type_transform, make_dummy_feature\nfrom protenix.utils.distributed import DIST_WRAPPER\nfrom protenix.utils.torch_utils import dict_to_tensor\nfrom protenix.data.tokenizer import TokenArray\nfrom protenix.utils.logger import get_logger\n\n# Common function for train and inference\ndef merge_all_chain_features(\n    pdb_id: str,\n    all_chain_features: dict[str, FeatureDict],\n    asym_to_entity_id: dict,\n    is_homomer_or_monomer: bool = False,\n    merge_method: str = \"dense_max\",\n    max_size: int = 16384,\n    msa_entity_type: str = \"prot\",\n) -> dict[str, np.ndarray]:\n    \"\"\"\n    Merges features from all chains in the bioassembly.\n\n    Args:\n        pdb_id (str): The PDB ID of the bioassembly.\n        all_chain_features (dict[str, FeatureDict]): Features for each chain in the bioassembly.\n        asym_to_entity_id (dict): Mapping from asym ID to entity ID.\n        is_homomer_or_monomer (bool): Indicates if the bioassembly is a homomer or monomer. Defaults to False.\n        merge_method (str): Method used for merging features. Defaults to \"dense_max\".\n        max_size (int): Maximum size of the MSA. Defaults to 16384.\n        msa_entity_type (str): Type of MSA entity, either \"prot\" or \"rna\". Defaults to \"prot\".\n\n    Returns:\n        dict[str, np.ndarray]: Merged features for the bioassembly.\n    \"\"\"\n    all_chain_features = add_assembly_features(\n        pdb_id,\n        all_chain_features,\n        asym_to_entity_id=asym_to_entity_id,\n    )\n    if msa_entity_type == \"rna\":\n        np_example = rna_merge(\n            all_chain_features=all_chain_features,\n            merge_method=merge_method,\n            msa_crop_size=max_size,\n        )\n    elif msa_entity_type == \"prot\":\n        np_example = pair_and_merge(\n            is_homomer_or_monomer=is_homomer_or_monomer,\n            all_chain_features=all_chain_features,\n            merge_method=merge_method,\n            msa_crop_size=max_size,\n        )\n    np_example = clip_msa(np_example, max_num_msa=max_size)\n    return np_example","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:02:45.253792Z","iopub.execute_input":"2025-05-26T09:02:45.254087Z","iopub.status.idle":"2025-05-26T09:02:46.609276Z","shell.execute_reply.started":"2025-05-26T09:02:45.254057Z","shell.execute_reply":"2025-05-26T09:02:46.608595Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sequence_info = {}\n    \nprotenix_initialized = False\n#has_long_sequences = any(len(seq) >= 400 for seq in sequences_df['sequence'])\nhas_long_sequences = True\nif has_long_sequences:\n    try:\n        # Import Protenix modules\n        from runner.batch_inference import get_default_runner\n        from runner.inference import update_inference_configs, InferenceRunner\n        from protenix.data.infer_data_pipeline import InferenceDataset\n        from configs.configs_base import configs as configs_base\n        from configs.configs_data import data_configs\n        from configs.configs_inference import inference_configs\n        from protenix.config.config import parse_configs\n        \n        # Set random seeds for reproducibility\n        np.random.seed(0)\n        torch.random.manual_seed(0)\n        torch.cuda.manual_seed_all(0)\n        \n        # Define DictDataset for Protenix\n        class DictDataset(InferenceDataset):\n            def __init__(\n                self,\n                seq_list: list,\n                dump_dir: str,\n                id_list: list = None,\n                use_msa: bool = True,\n            ) -> None:\n                self.dump_dir = dump_dir\n                self.use_msa = use_msa\n                if isinstance(id_list, type(None)):\n                    self.inputs = [{\"sequences\": \n                                    [{\"rnaSequence\": \n                                      {\"sequence\": seq, \n                                       \"count\": 1}}],\n                                    \"name\": \"query\"} for seq in seq_list]\n                else:\n                    self.inputs = [\n                        {\n                            \"sequences\": [\n                                {\n                                    \"rnaSequence\": {\n                                        \"sequence\": seq,\n                                        \"count\": 1,\n                                        \"msa\": {\n                                            \"precomputed_msa_dir\": f\"{sto_dir}/\",\n                                            \"pairing_db\": \"no_pairing_db\"\n                                        }\n                                    }\n                                }\n                            ],\n                            \"name\": i\n                        } for i, seq in zip(id_list,seq_list)\n                    ]\n\n            # override for RNA MSA functionality\n            def process_one(\n                self,\n                single_sample_dict: Mapping[str, Any],\n            ) -> tuple[dict[str, torch.Tensor], AtomArray, dict[str, float]]:\n                \"\"\"\n                Processes a single sample from the input JSON to generate features and statistics.\n    \n                Args:\n                    single_sample_dict: A dictionary containing the sample data.\n    \n                Returns:\n                    A tuple containing:\n                        - A dictionary of features.\n                        - An AtomArray object.\n                        - A dictionary of time tracking statistics.\n                \"\"\"\n                # general features\n                t0 = time.time()\n                sample2feat = SampleDictToFeatures(\n                    single_sample_dict,\n                )\n                features_dict, atom_array, token_array = sample2feat.get_feature_dict()\n                features_dict[\"distogram_rep_atom_mask\"] = torch.Tensor(\n                    atom_array.distogram_rep_atom_mask\n                ).long()\n                entity_poly_type = sample2feat.entity_poly_type\n                t1 = time.time()\n    \n                # Msa features\n                entity_to_asym_id = DataPipeline.get_label_entity_id_to_asym_id_int(atom_array)\n    \n                try:\n                    if not self.use_msa:\n                        msa_features = {}\n                    else:\n                        sequence_to_features: dict[str, dict[str, Any]] = {}\n                        name = single_sample_dict[\"name\"]\n                        sequence = single_sample_dict[\"sequences\"][0][\"rnaSequence\"][\"sequence\"]\n                        msa_dir = single_sample_dict[\"sequences\"][0][\"rnaSequence\"][\"msa\"][\"precomputed_msa_dir\"]\n                        sequence_feat = process_single_sequence(\n                            pdb_name=name,\n                            sequence=sequence,\n                            raw_msa_paths=[f\"{msa_dir}{name}.MSA.sto\"], \n                            seq_limits=SEQ_LIMITS, # Optional[list[str]],\n                            msa_entity_type =\"rna\",\n                            msa_type = \"non_pairing\",\n                        )\n                        sequence_feat = convert_monomer_features(sequence_feat)\n                        sequence_to_features[sequence] = sequence_feat\n                        all_chain_features = {\n                            0: deepcopy(sequence_to_features[sequence])\n                        }\n                        del sequence_to_features\n    \n                        asym_to_entity_id = {\n                            0: '1'\n                        }\n    \n                        msa_feats = merge_all_chain_features(\n                            pdb_id=\"test_assembly\",\n                            all_chain_features=all_chain_features,\n                            asym_to_entity_id=asym_to_entity_id,\n                            is_homomer_or_monomer=False,\n                            msa_entity_type=\"rna\",\n                        )\n                        if msa_feats is None:\n                            return {}\n    \n                        msa_feats = tokenize_msa(\n                            msa_feats=msa_feats,\n                            token_array=token_array,\n                            atom_array=atom_array,\n                        )\n    \n                        msa_features = {\n                            k: v\n                            for (k, v) in msa_feats.items()\n                            if k\n                            in [\"msa\", \"has_deletion\", \"deletion_value\", \"deletion_mean\", \"profile\"]\n                        }\n                except Exception as e:\n                    print(\"msa feature skipped as error occur, run without msa\")\n                    msa_features = {}\n    \n                # Make dummy features for not implemented features\n                dummy_feats = [\"template\"]\n                if len(msa_features) == 0:\n                    dummy_feats.append(\"msa\")\n                else:\n                    msa_features = dict_to_tensor(msa_features)\n                    features_dict.update(msa_features)\n                features_dict = make_dummy_feature(\n                    features_dict=features_dict,\n                    dummy_feats=dummy_feats,\n                )\n    \n                # Transform to right data type\n                feat = data_type_transform(feat_or_label_dict=features_dict)\n    \n                t2 = time.time()\n    \n                data = {}\n                data[\"input_feature_dict\"] = feat\n    \n                # Add dimension related items\n                N_token = feat[\"token_index\"].shape[0]\n                N_atom = feat[\"atom_to_token_idx\"].shape[0]\n                N_msa = feat[\"msa\"].shape[0]\n    \n                stats = {}\n                for mol_type in [\"ligand\", \"protein\", \"dna\", \"rna\"]:\n                    mol_type_mask = feat[f\"is_{mol_type}\"].bool()\n                    stats[f\"{mol_type}/atom\"] = int(mol_type_mask.sum(dim=-1).item())\n                    stats[f\"{mol_type}/token\"] = len(\n                        torch.unique(feat[\"atom_to_token_idx\"][mol_type_mask])\n                    )\n    \n                N_asym = len(torch.unique(data[\"input_feature_dict\"][\"asym_id\"]))\n                data.update(\n                    {\n                        \"N_asym\": torch.tensor([N_asym]),\n                        \"N_token\": torch.tensor([N_token]),\n                        \"N_atom\": torch.tensor([N_atom]),\n                        \"N_msa\": torch.tensor([N_msa]),\n                    }\n                )\n    \n                def formatted_key(key):\n                    type_, unit = key.split(\"/\")\n                    if type_ == \"protein\":\n                        type_ = \"prot\"\n                    elif type_ == \"ligand\":\n                        type_ = \"lig\"\n                    else:\n                        pass\n                    return f\"N_{type_}_{unit}\"\n    \n                data.update(\n                    {\n                        formatted_key(k): torch.tensor([stats[k]])\n                        for k in [\n                            \"protein/atom\",\n                            \"ligand/atom\",\n                            \"dna/atom\",\n                            \"rna/atom\",\n                            \"protein/token\",\n                            \"ligand/token\",\n                            \"dna/token\",\n                            \"rna/token\",\n                        ]\n                    }\n                )\n                data.update({\"entity_poly_type\": entity_poly_type})\n                t3 = time.time()\n                time_tracker = {\n                    \"crop\": t1 - t0,\n                    \"featurizer\": t2 - t1,\n                    \"added_feature\": t3 - t2,\n                }\n    \n                return data, atom_array, time_tracker\n        \n        # Setup Protenix model configuration\n        configs_base[\"use_deepspeed_evo_attention\"] = (\n            os.environ.get(\"USE_DEEPSPEED_EVO_ATTENTION\", False) == \"true\")\n        configs_base[\"model\"][\"N_cycle\"] = 10\n        configs_base[\"sample_diffusion\"][\"N_sample\"] = 5  # Generate 5 models\n        configs_base[\"sample_diffusion\"][\"N_step\"] = 200\n        inference_configs['load_checkpoint_path'] = SHORT_MODEL_PATH\n        configs = {**configs_base, **{\"data\": data_configs}, **inference_configs}\n        \n        configs = parse_configs(\n            configs=configs,\n            fill_required_with_null=True,\n        )\n        \n        protenix_runner = InferenceRunner(configs)\n        protenix_initialized = True\n        print(\"Protenix model initialized successfully!\")\n        \n    except Exception as e:\n        print(f\"Error initializing Protenix model: {e}\")\n        print(\"Will attempt to use DRfold for all sequences.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:02:46.610136Z","iopub.execute_input":"2025-05-26T09:02:46.610675Z","iopub.status.idle":"2025-05-26T09:03:35.720373Z","shell.execute_reply.started":"2025-05-26T09:02:46.610643Z","shell.execute_reply":"2025-05-26T09:03:35.719568Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"csv_df = pd.read_csv('/kaggle/input/testingcsv/run_test.csv')\n\n# Model paths for three different Protenix models\nPROTENIX1_MODEL_PATH = \"/kaggle/input/rnaonly19799/19799_ema_0.995.pt\"  \nPROTENIX2_MODEL_PATH = \"/kaggle/input/37989all/37989_ema_0.995.pt\"\nPROTENIX3_MODEL_PATH = \"/kaggle/input/protenix-checkpoints/model_v0.2.0.pt\"\n\n# Constants\nMAX_SEQUENCE_LENGTH = 1000  # New constant for truncation\n\nfor idx, row in sequences_df.iterrows():\n    \n    sequence_results = {}\n    \n    clean_output_dirs(outdir, ret_dir, folddir, refdir)\n    start_time = time.time()\n    target_id = row['target_id']\n    sequence = row['sequence']\n    optimize = False\n\n    target_rows = csv_df[csv_df['ID'].str.startswith(target_id + '_')]\n    \n    if not target_rows.empty:\n        for _, csv_row in target_rows.iterrows():\n            key = csv_row['ID']  # e.g., R1107_1\n            resid = int(csv_row['ID'].split('_')[-1])  # Extract resid from ID\n            resname = csv_row['resname']\n            \n            # Create dictionary for this residue\n            result_entry = {\n                'ID': key,\n                'resname': resname,\n                'resid': resid,\n                'x_1': csv_row['x_1'],\n                'y_1': csv_row['y_1'],\n                'z_1': csv_row['z_1'],\n                'x_2': csv_row['x_2'],\n                'y_2': csv_row['y_2'],\n                'z_2': csv_row['z_2'],\n                'x_3': csv_row['x_3'],\n                'y_3': csv_row['y_3'],\n                'z_3': csv_row['z_3'],\n                'x_4': csv_row['x_4'],\n                'y_4': csv_row['y_4'],\n                'z_4': csv_row['z_4'],\n                'x_5': csv_row['x_5'],\n                'y_5': csv_row['y_5'],\n                'z_5': csv_row['z_5']\n            }\n            \n            sequence_results[key] = result_entry\n        \n        # Update all_results with CSV data\n        all_results.update(sequence_results)\n        print(f\"Added {len(sequence_results)} residue positions from CSV for {target_id}\")\n        print(f\"Time taken for {target_id}: {time.time() - start_time:.2f} seconds\")\n        continue\n        \n    \n    seq_len = len(sequence)\n    use_drfold = False  # No longer using DRfold\n    use_protenix = True\n    use_boltz = True\n\n    # Always use 3 Protenix models + 1 Boltz model regardless of length\n    protenix_models_needed = [2, 1, 1]  # 2 from first model, 1 each from second and third\n    boltz_models_needed = 1\n\n    print(f\"Strategy for {target_id} (length: {seq_len}):\")\n    print(f\"- Using 3 Protenix models (2+1+1 outputs) + 1 Boltz model\")\n\n    # Handle sequence truncation for sequences > 1000\n    original_sequence = sequence\n    was_truncated = len(sequence) > MAX_SEQUENCE_LENGTH\n    if was_truncated:\n        sequence = sequence[:MAX_SEQUENCE_LENGTH]\n        print(f\"Truncated sequence from {len(original_sequence)} to {len(sequence)} residues\")\n\n    sequence_info[target_id] = {\n        'was_truncated': was_truncated,\n        'original_length': len(original_sequence)\n    }\n\n    protenix_results = {}\n    boltz_results = {}\n    \n    if use_protenix:\n        print(f\"Using Protenix pipeline for {target_id} (length: {len(sequence)})\")\n        protenix_start = time.time()\n        \n        try:\n            # Model paths for three different Protenix models\n            PROTENIX_MODEL_PATHS = [\n                PROTENIX1_MODEL_PATH,  # protenix1\n                PROTENIX2_MODEL_PATH,  # protenix2\n                PROTENIX3_MODEL_PATH   # protenix3\n            ]\n            \n            # Run 3 separate Protenix predictions with different sample counts\n            all_protenix_predictions = []\n            \n            for model_run, (model_path, n_samples) in enumerate(zip(PROTENIX_MODEL_PATHS, protenix_models_needed)):\n                print(f\"Running Protenix model {model_run + 1} ({model_path}) with {n_samples} samples...\")\n                \n                # Update model checkpoint path\n                inference_configs['load_checkpoint_path'] = model_path\n                \n                configs_base[\"sample_diffusion\"][\"N_sample\"] = n_samples\n                \n                configs = {**configs_base, **{\"data\": data_configs}, **inference_configs}\n                configs = parse_configs(\n                    configs=configs,\n                    fill_required_with_null=True,\n                )\n                \n                # Reinitialize the model with the appropriate checkpoint\n                protenix_runner = InferenceRunner(configs)\n                \n                # Create dataset with single sequence\n                dataset = DictDataset([sequence], dump_dir='output', id_list=[target_id], use_msa=True)\n                \n                # Process with Protenix\n                data, atom_array, data_error_message = dataset[0]\n                \n                # Check for errors\n                if data_error_message:\n                    raise ValueError(f\"Protenix data error: {data_error_message}\")\n                \n                # Update model configs for this sequence\n                new_configs = update_inference_configs(configs, data[\"N_token\"].item())\n                protenix_runner.update_model_configs(new_configs)\n                \n                # Run prediction\n                prediction = protenix_runner.predict(data)\n                print(prediction['summary_confidence'])\n                ptm_scores = [summary['gpde'].item() for summary in prediction['summary_confidence']]\n                \n                protenix_indices_by_score = torch.tensor(ptm_scores).argsort(descending=True)\n                best_protenix_indices = protenix_indices_by_score.tolist()\n                print(f\"Protenix model {model_run + 1} ranked by GPDE: {best_protenix_indices}\")\n                \n                prediction['coordinate'] = prediction['coordinate'][protenix_indices_by_score]\n                prediction = prediction['coordinate'][:, data['input_feature_dict']['atom_to_tokatom_idx']==12]\n                \n                # Store all predictions from this run\n                all_protenix_predictions.extend(prediction)\n            \n            # Now we have 4 total predictions (2+1+1), process them\n            sequence_results_list = parse_output_to_df(torch.stack(all_protenix_predictions), sequence, target_id)\n\n            if sequence_results_list:\n                protenix_results_df = sequence_results_list[0]\n                \n                # Convert to dictionary format\n                protenix_results_dict = {}\n                for _, row in protenix_results_df.iterrows():\n                    key = row['ID']\n                    protenix_results_dict[key] = row.to_dict()\n                \n                protenix_results = protenix_results_dict\n                \n                print(f\"Successfully processed {target_id} with 3 Protenix models in {time.time() - protenix_start:.2f} seconds\")\n            \n        except Exception as e:\n            print(f\"Error processing {target_id} with Protenix: {e}\")\n            use_protenix = False\n\n    if use_boltz:\n        print(f\"Using Boltz pipeline for {target_id}\")\n        boltz_start = time.time()\n        \n        # Prepare directories\n        os.makedirs(BOLTZ_INPUT_DIR, exist_ok=True)\n        os.makedirs(BOLTZ_OUTPUT_DIR, exist_ok=True)\n        \n        # Check if inference.py exists\n        if not os.path.exists(BOLTZ_SCRIPT):\n            print(f\"ERROR: Cannot find Boltz script at {BOLTZ_SCRIPT}\")\n        else:\n            print(f\"Found Boltz script at {BOLTZ_SCRIPT}\")\n\n        # Clear previous input files\n        for yaml_file in glob.glob(f\"{BOLTZ_INPUT_DIR}/*.yaml\"):\n            os.remove(yaml_file)\n            print(f\"Removed previous input file: {yaml_file}\")\n\n        if os.path.exists(BOLTZ_OUTPUT_DIR):\n            import shutil\n            try:\n                shutil.rmtree(BOLTZ_OUTPUT_DIR)\n                print(f\"Cleared previous output directory: {BOLTZ_OUTPUT_DIR}\")\n                # Recreate the directory\n                os.makedirs(BOLTZ_OUTPUT_DIR, exist_ok=True)\n            except Exception as e:\n                print(f\"Warning: Could not clear output directory: {str(e)}\")\n        \n        # Prepare YAML input for this sequence\n        with open(f'{BOLTZ_INPUT_DIR}/{target_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: {sequence}\")\n        \n        print(f\"Created input file at {BOLTZ_INPUT_DIR}/{target_id}.yaml\")\n        \n        # Clear CUDA cache before running inference\n        torch.cuda.empty_cache()\n        if torch.cuda.is_available():\n            torch.cuda.synchronize()\n        gc.collect()\n        \n        print(\"Starting Boltz inference (generating 1 model)...\")\n        try:\n            # Run the modified script\n            result = subprocess.run(\n                ['python', BOLTZ_SCRIPT], \n                capture_output=True,\n                text=True, \n            )\n            \n            # Print the command output for debugging\n            print(\"Boltz inference stdout:\")\n            print(result.stdout)\n            \n            if result.stderr:\n                print(\"Boltz inference stderr:\")\n                print(result.stderr)\n                \n            print(f\"Boltz inference completed with return code: {result.returncode}\")\n            \n            # Check if output directories were created\n            pred_dir = f\"{BOLTZ_OUTPUT_DIR}/boltz_results_inputs_prediction/predictions/{target_id}\"\n            if os.path.exists(pred_dir):\n                print(f\"Output directory created: {pred_dir}\")\n                print(f\"Contents: {os.listdir(pred_dir)}\")\n            else:\n                print(f\"ERROR: Output directory not created: {pred_dir}\")\n                \n        except Exception as e:\n            print(f\"Error running Boltz inference: {str(e)}\")\n        \n        # Get model indices sorted by confidence scores\n        print(\"Getting Boltz model rankings by confidence score...\")\n        boltz_top_indices = get_boltz_top5_indices(target_id, 1)\n        print(f\"Boltz models ranked by score: {boltz_top_indices}\")\n        \n        # Extract coordinates for the single model\n        boltz_results = {}\n        model_idx = 0  # Only process the first (and only) model\n        print(f\"Processing Boltz model {model_idx}...\")\n        c1_coords = get_boltz_c1_coords(target_id, model_idx)\n        \n        # Skip if no coordinates found\n        if not c1_coords:\n            print(f\"Warning: No C1' coordinates found for {target_id} model {model_idx}\")\n        else:\n            print(f\"Found {len(c1_coords)} C1' coordinates for Boltz model\")\n                \n            # Process coordinates and assign to residues\n            for i, coords in enumerate(c1_coords):\n                res_id = i + 1\n                key = f\"{target_id}_{res_id}\"\n                \n                # Determine residue code\n                res_code = sequence[i] if i < len(sequence) else 'X'\n                \n                boltz_results[key] = {\n                    'ID': key,\n                    'resname': res_code,\n                    'resid': res_id,\n                    'x_1': coords[0],\n                    'y_1': coords[1],\n                    'z_1': coords[2]\n                }\n        \n        print(f\"Boltz processing completed in {time.time() - boltz_start:.2f} seconds\")\n\n    # Combine results from Protenix + Boltz\n    all_keys = set()\n    if protenix_results:\n        all_keys.update(protenix_results.keys())\n    if boltz_results:\n        all_keys.update(boltz_results.keys())\n        \n    for key in all_keys:\n        # Initialize with base info from either source\n        if key in protenix_results:\n            result_entry = {\n                'ID': protenix_results[key]['ID'],\n                'resname': protenix_results[key]['resname'],\n                'resid': protenix_results[key]['resid']\n            }\n        elif key in boltz_results:\n            result_entry = {\n                'ID': boltz_results[key]['ID'],\n                'resname': boltz_results[key]['resname'],\n                'resid': boltz_results[key]['resid']\n            }\n        else:\n            continue\n        \n        # Fill positions 1-4 from Protenix models (4 total outputs: 2+1+1)\n        for out_pos in range(1, 5):\n            model_idx = out_pos\n            if key in protenix_results and f'x_{model_idx}' in protenix_results[key]:\n                result_entry[f'x_{out_pos}'] = protenix_results[key][f'x_{model_idx}']\n                result_entry[f'y_{out_pos}'] = protenix_results[key][f'y_{model_idx}']\n                result_entry[f'z_{out_pos}'] = protenix_results[key][f'z_{model_idx}']\n            else:\n                # Use zeros if missing\n                result_entry[f'x_{out_pos}'] = 0.0\n                result_entry[f'y_{out_pos}'] = 0.0\n                result_entry[f'z_{out_pos}'] = 0.0\n        \n        # Fill position 5 from Boltz model\n        if key in boltz_results:\n            result_entry['x_5'] = boltz_results[key]['x_1']\n            result_entry['y_5'] = boltz_results[key]['y_1']\n            result_entry['z_5'] = boltz_results[key]['z_1']\n        else:\n            result_entry['x_5'] = 0.0\n            result_entry['y_5'] = 0.0\n            result_entry['z_5'] = 0.0\n        \n        sequence_results[key] = result_entry\n\n    # Handle truncated sequences by adding dummy entries with zero coordinates\n    if sequence_info[target_id]['was_truncated']:\n        orig_length = sequence_info[target_id]['original_length']\n        for res_id in range(MAX_SEQUENCE_LENGTH + 1, orig_length + 1):\n            key = f\"{target_id}_{res_id}\"\n            # Determine residue name from original sequence\n            res_code = original_sequence[res_id - 1]\n            \n            # Create dummy entry with zero coordinates\n            sequence_results[key] = {\n                'ID': key,\n                'resname': res_code,\n                'resid': res_id\n            }\n            \n            # Add zero coordinates for all 5 models\n            for model_idx in range(1, 6):\n                sequence_results[key][f'x_{model_idx}'] = 0.0\n                sequence_results[key][f'y_{model_idx}'] = 0.0\n                sequence_results[key][f'z_{model_idx}'] = 0.0\n        \n        print(f\"Added {orig_length - MAX_SEQUENCE_LENGTH} dummy entries with zero coordinates for {target_id}\")\n    \n    # ADD RESULTS TO MAIN DICTIONARY\n    if sequence_results:\n        all_results.update(sequence_results)\n        print(f\"Successfully processed {target_id} with multi-model approach\")\n        print(f\"Added {len(sequence_results)} residue entries to all_results\")\n    else:\n        print(f\"ERROR: No results available for {target_id}\")        \n    \n    # Clean up\n    torch.cuda.empty_cache()\n    if torch.cuda.is_available():\n        torch.cuda.synchronize()\n    \n    print(f\"Total time taken for {target_id}: {time.time() - start_time:.2f} seconds\")\n    print(f\"Current all_results size: {len(all_results)}\")\n    print(\"-\" * 50)\n  ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:03:47.614132Z","iopub.execute_input":"2025-05-26T09:03:47.614453Z","iopub.status.idle":"2025-05-26T09:07:06.928235Z","shell.execute_reply.started":"2025-05-26T09:03:47.614425Z","shell.execute_reply":"2025-05-26T09:07:06.927291Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# results_list = []\n# for key, result in all_results.items():\n#     results_list.append(result)\n\n# submit_df = pd.DataFrame(results_list)\nresults_df = pd.DataFrame(all_results.values())\n# Sort by target ID first, then by residue ID\nresults_df['target_id'] = results_df['ID'].str.split('_').str[0]\nresults_df['residue_num'] = results_df['ID'].str.split('_').str[1].astype(int)\nresults_df = results_df.sort_values(['target_id', 'residue_num'])\n# Remove helper columns if needed\nresults_df = results_df.drop(['target_id', 'residue_num'], axis=1)\nsubmit_df = results_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:07:12.936195Z","iopub.execute_input":"2025-05-26T09:07:12.936484Z","iopub.status.idle":"2025-05-26T09:07:12.970606Z","shell.execute_reply.started":"2025-05-26T09:07:12.936462Z","shell.execute_reply":"2025-05-26T09:07:12.969918Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submit_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:07:13.838329Z","iopub.execute_input":"2025-05-26T09:07:13.838598Z","iopub.status.idle":"2025-05-26T09:07:13.871207Z","shell.execute_reply.started":"2025-05-26T09:07:13.838576Z","shell.execute_reply":"2025-05-26T09:07:13.870524Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submit_df = submit_df.fillna(0)\n# Save submission file\nif local:\n    submit_df.to_csv('submission.csv', index=False)\nelse:\n    submit_df.to_csv('/kaggle/working/submission.csv', index=False)\n\nprint(f\"Saved submission file with {len(submit_df)} predictions\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T09:07:24.25077Z","iopub.execute_input":"2025-05-26T09:07:24.25107Z","iopub.status.idle":"2025-05-26T09:07:24.298625Z","shell.execute_reply.started":"2025-05-26T09:07:24.251046Z","shell.execute_reply":"2025-05-26T09:07:24.297792Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"directory = '/kaggle/working'\n\n# Loop through all files and directories\nfor filename in os.listdir(directory):\n    file_path = os.path.join(directory, filename)\n    if filename != 'submission.csv':\n        try:\n            if os.path.isfile(file_path) or os.path.islink(file_path):\n                os.remove(file_path)  # remove file or symbolic link\n            elif os.path.isdir(file_path):\n                import shutil\n                shutil.rmtree(file_path)  # remove directory and its contents\n        except Exception as e:\n            print(f'Failed to delete {file_path}. Reason: {e}')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T13:58:29.530598Z","iopub.execute_input":"2025-05-22T13:58:29.530933Z","iopub.status.idle":"2025-05-22T13:58:30.392424Z","shell.execute_reply.started":"2025-05-22T13:58:29.530902Z","shell.execute_reply":"2025-05-22T13:58:30.391628Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}