{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceType":"competition","sourceId":118765,"databundleVersionId":15231210},{"sourceType":"datasetVersion","sourceId":14604295,"datasetId":9328538,"databundleVersionId":15440074},{"sourceType":"datasetVersion","sourceId":14962460,"datasetId":9577079,"databundleVersionId":15833819},{"sourceType":"datasetVersion","sourceId":14874339,"datasetId":9502242,"databundleVersionId":15736806},{"sourceType":"datasetVersion","sourceId":14962495,"datasetId":9577097,"databundleVersionId":15833858},{"sourceType":"datasetVersion","sourceId":10855324,"datasetId":6742586,"databundleVersionId":11219268},{"sourceType":"kernelVersion","sourceId":148788846},{"sourceType":"kernelVersion","sourceId":227713340},{"sourceType":"kernelVersion","sourceId":239534907},{"sourceType":"kernelVersion","sourceId":290929004},{"sourceType":"kernelVersion","sourceId":291182720},{"sourceType":"kernelVersion","sourceId":291833135},{"sourceType":"kernelVersion","sourceId":292115284}],"dockerImageVersionId":31287,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Stanford RNA 3D Folding Part 2 - v10\n\n**Changes in v10 (timeout fix from v9):**\n- **FIX 1:**  reduced from 200 to 50 — 4x faster per sequence with modest quality tradeoff. Short sequences (~30-500 res) still get good predictions.\n- **FIX 2:**  — sequences longer than 800 residues skip Protenix and use TBM instead. This excludes 9MME (4640 res, ~24 chunks × 9.5min = 4hr alone) and 9ZCC (1460 res, ~7 chunks). TBM has comp=0.55-0.56 for these, which is reasonable.\n- All v9 fixes retained (cache prep, rnaSequence JSON, name key, CIF parser, self-match exclusion).\n- **Expected runtime:** ~25-40 minutes total for 26 short sequences at 50 steps.\n","metadata":{}},{"cell_type":"code","source":"\n# ================================================================\n# Stanford RNA 3D Folding Part 2 - v8\n# Key fix over v7:\n#   Protenix cache pre-population so no network downloads needed.\n#   PROTENIX_ROOT_DIR -> writable /kaggle/working/protenix_cache/\n#   with all 4 required cache files already present.\n# ================================================================\n\nimport os, sys, gc, json, glob, math, random, subprocess, warnings, shutil\nwarnings.filterwarnings(\"ignore\", category=FutureWarning)\nwarnings.filterwarnings(\"ignore\", category=UserWarning)\nfrom pathlib import Path\nfrom collections import defaultdict\n\n# -------------------- package bootstrap --------------------\ndef _install_wheel_only(patterns, pkg_name=None):\n    # Wheel-only install - no pip network fallback (no internet on Kaggle inference)\n    for pattern in patterns:\n        matches = glob.glob(pattern)\n        if matches:\n            wheel = sorted(matches)[0]\n            print(f\"Installing {wheel}\")\n            subprocess.run([\"pip\",\"install\",\"--no-index\",\"--no-deps\",wheel], check=False)\n            return True\n    if pkg_name: print(f\"Wheel not found for {pkg_name} - skipping\")\n    return False\n\n_install_wheel_only([\"/kaggle/input/biopython-cp312/*.whl\",\"/kaggle/input/datasets/*/biopython-cp312/*.whl\"],\"biopython\")\n_install_wheel_only([\"/kaggle/input/biotite/*.whl\",\"/kaggle/input/datasets/*/biotite/*.whl\"],\"biotite\")\n_install_wheel_only([\"/kaggle/input/rdkit-2025-9-5/*.whl\",\"/kaggle/input/datasets/*/rdkit-2025-9-5/*.whl\"],\"rdkit\")\n_install_wheel_only([\"/kaggle/input/parasail*.whl\",\"/kaggle/input/datasets/*/parasail*.whl\"],\"parasail\")\n\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom Bio.Align import PairwiseAligner\n\ntry:\n    import parasail; HAS_PARASAIL = True; print(\"parasail: available\")\nexcept Exception:\n    parasail = None; HAS_PARASAIL = False; print(\"parasail: not available, using BioPython aligner\")\n\n# -------------------- config --------------------\nos.environ[\"LAYERNORM_TYPE\"] = \"torch\"\nos.environ.setdefault(\"RNA_MSA_DEPTH_LIMIT\", \"512\")\n\nSEED     = 42\nN_SAMPLE = 5\n\nMAX_SEQ_LEN            = int(os.environ.get(\"MAX_SEQ_LEN\",            \"512\"))\nCHUNK_OVERLAP          = int(os.environ.get(\"CHUNK_OVERLAP\",          \"320\"))\nTBM_TOP_K_SHORTLIST    = int(os.environ.get(\"TBM_TOP_K_SHORTLIST\",    \"180\"))\nTBM_TOP_N              = int(os.environ.get(\"TBM_TOP_N\",              \"60\"))\nTBM_MAX_CANDIDATES     = int(os.environ.get(\"TBM_MAX_CANDIDATES\",     \"10\"))\nMIN_TEMPLATE_COMPOSITE = float(os.environ.get(\"MIN_TEMPLATE_COMPOSITE\",\"0.35\"))\nMIN_TEMPLATE_IDENTITY  = float(os.environ.get(\"MIN_TEMPLATE_IDENTITY\", \"20.0\"))\nMAX_LEN_RATIO          = float(os.environ.get(\"MAX_LEN_RATIO\",         \"1.65\"))\n\nUSE_PROTENIX        = os.environ.get(\"USE_PROTENIX\",        \"true\").lower() == \"true\"\nPROTENIX_MODEL_NAME = os.environ.get(\"PROTENIX_MODEL_NAME\", \"protenix_base_20250630_v1.0.0\")\nPROTENIX_N_CYCLE    = int(os.environ.get(\"PROTENIX_N_CYCLE\",  \"10\"))\n# FIX: Reduced from 200->50 (4x faster, modest quality loss on short seqs)\nPROTENIX_N_STEP     = int(os.environ.get(\"PROTENIX_N_STEP\",   \"50\"))\nPROTENIX_N_SAMPLE   = int(os.environ.get(\"PROTENIX_N_SAMPLE\", str(N_SAMPLE)))\nPROTENIX_TRI_ATT    = os.environ.get(\"PROTENIX_TRI_ATT\",    \"torch\")\nPROTENIX_TRI_MUL    = os.environ.get(\"PROTENIX_TRI_MUL\",    \"torch\")\n# FIX: Skip Protenix for sequences longer than this (use TBM instead)\n# 9MME=4640res and 9ZCC=1460res each take hours when chunked - not worth it\nPROTENIX_MAX_SEQ_LEN = int(os.environ.get(\"PROTENIX_MAX_SEQ_LEN\", \"800\"))\nUSE_MSA             = os.environ.get(\"USE_MSA\",             \"false\").lower() == \"true\"\nUSE_TEMPLATE        = os.environ.get(\"USE_TEMPLATE\",        \"false\").lower() == \"true\"\nUSE_RNA_MSA         = os.environ.get(\"USE_RNA_MSA\",         \"false\").lower() == \"true\"\n\nUSE_EXTERNAL_NOTEBOOK_PRIORS      = os.environ.get(\"USE_EXTERNAL_NOTEBOOK_PRIORS\",      \"true\").lower() == \"true\"\nUSE_EXTERNAL_NOTEBOOK_PREDICTIONS = os.environ.get(\"USE_EXTERNAL_NOTEBOOK_PREDICTIONS\", \"true\").lower() == \"true\"\n\nDEFAULT_OUTPUT = \"/kaggle/working/submission.csv\"\n\ndef seed_everything(seed=42):\n    random.seed(seed); np.random.seed(seed)\n    torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)\nseed_everything(SEED)\n\n# -------------------- paths --------------------\ndef find_competition_base():\n    for c in [\n        \"/kaggle/input/competitions/stanford-rna-3d-folding-2\",\n        \"/kaggle/input/stanford-rna-3d-folding-2\",\n        \"/kaggle/input/stanford-rna-3d-folding-part-2\",\n        \"/kaggle/input/stanford-rna-3d-folding\",\n        \"/kaggle/input/competitions/stanford-rna-3d-folding-part-2\",\n    ]:\n        if os.path.isdir(c) and os.path.exists(os.path.join(c,\"test_sequences.csv\")):\n            return c\n    hits = glob.glob(\"/kaggle/input/**/test_sequences.csv\", recursive=True)\n    if not hits: raise FileNotFoundError(\"test_sequences.csv not found\")\n    return os.path.dirname(hits[0])\n\nDATA_BASE          = find_competition_base()\nDEFAULT_TEST_CSV   = f\"{DATA_BASE}/test_sequences.csv\"\nDEFAULT_TRAIN_CSV  = f\"{DATA_BASE}/train_sequences.csv\"\nDEFAULT_TRAIN_LBLS = f\"{DATA_BASE}/train_labels.csv\"\nDEFAULT_VAL_CSV    = f\"{DATA_BASE}/validation_sequences.csv\"\nDEFAULT_VAL_LBLS   = f\"{DATA_BASE}/validation_labels.csv\"\nprint(\"DATA_BASE =\", DATA_BASE)\nprint(\"Data files:\", sorted(os.listdir(DATA_BASE))[:20])\n\ndef resolve_protenix_root():\n    for c in [\n        \"/kaggle/input/protenix-v1-adjusted/Protenix-v1-adjust-v2/Protenix-v1-adjust-v2/Protenix-v1\",\n        \"/kaggle/input/protenix-v1-adjusted/Protenix-v1-adjust/Protenix-v1\",\n        \"/kaggle/input/protenix-v1-adjusted/Protenix-v1\",\n        \"/kaggle/input/protenix-v1-adjusted\",\n        \"/kaggle/input/datasets/qiweiyin/protenix-v1-adjusted/Protenix-v1-adjust-v2/Protenix-v1-adjust-v2/Protenix-v1\",\n        \"/kaggle/input/datasets/qiweiyin/protenix-v1-adjusted/Protenix-v1-adjust/Protenix-v1\",\n        \"/kaggle/input/datasets/qiweiyin/protenix-v1-adjusted/Protenix-v1\",\n        \"/kaggle/input/datasets/qiweiyin/protenix-v1-adjusted\",\n    ]:\n        if os.path.isdir(c): return c\n    for c in glob.glob(\"/kaggle/input/**/*protenix*\", recursive=True):\n        if os.path.isdir(c) and (os.path.isdir(os.path.join(c,\"runner\")) or\n                                  os.path.isdir(os.path.join(c,\"checkpoint\"))):\n            return c\n    raise FileNotFoundError(\"Protenix root not found\")\n\nPROTENIX_ROOT = resolve_protenix_root()\nprint(\"PROTENIX_ROOT =\", PROTENIX_ROOT)\nprint(\"Checkpoints  :\", os.listdir(os.path.join(PROTENIX_ROOT,\"checkpoint\")))\n\nALL_INPUT_DIRS = [d for d in glob.glob(\"/kaggle/input/*\") if os.path.isdir(d)]\n\ndef find_attached_notebook_dirs():\n    keys = {\"rhofold\":[],\"mmseqs\":[],\"tbm\":[],\"parasail\":[],\"templates\":[]}\n    for d in ALL_INPUT_DIRS:\n        base = os.path.basename(d).lower()\n        if \"rhofold\" in base or \"ribonanzanet\" in base: keys[\"rhofold\"].append(d)\n        if \"mmseqs\" in base:   keys[\"mmseqs\"].append(d)\n        if \"tbm\" in base:      keys[\"tbm\"].append(d)\n        if \"parasail\" in base: keys[\"parasail\"].append(d)\n        if \"template\" in base: keys[\"templates\"].append(d)\n    return keys\n\nATTACHED_NOTEBOOK_DIRS = find_attached_notebook_dirs()\nprint(\"Attached notebook dirs:\", {k:len(v) for k,v in ATTACHED_NOTEBOOK_DIRS.items()})\n\n\n# ==================== PROTENIX CACHE PREPARATION ====================\n\ndef prepare_protenix_cache():\n    # Pre-populate a writable cache directory so Protenix finds all required\n    # files without attempting any network downloads.\n\n    # Protenix's download_inference_cache() checks:\n    # configs['data']['ccd_components_file']           -> components.cif\n    # configs['data']['ccd_components_rdkit_mol_file'] -> components.cif.rdkit_mol.pkl\n    # configs['data']['pdb_cluster_file']              -> clusters-by-entity-40.txt\n    # configs['data']['obsolete_release_data_csv']     -> obsolete_release_date.csv\n\n    # All paths resolve to PROTENIX_ROOT_DIR/common/<filename>.\n    # The first two already exist in the dataset's common/ dir.\n    # The last two are missing - we create valid empty stubs.\n    # The checkpoint dir is symlinked so load_checkpoint_dir works.\n    cache_dir    = \"/kaggle/working/protenix_cache\"\n    common_dir   = os.path.join(cache_dir, \"common\")\n    ckpt_dir     = os.path.join(cache_dir, \"checkpoint\")\n    src_common   = os.path.join(PROTENIX_ROOT, \"common\")\n    src_ckpt     = os.path.join(PROTENIX_ROOT, \"checkpoint\")\n\n    os.makedirs(common_dir, exist_ok=True)\n\n    # --- Copy existing cache files ---\n    for fname in [\"components.cif\", \"components.cif.rdkit_mol.pkl\", \"components.txt\"]:\n        src = os.path.join(src_common, fname)\n        dst = os.path.join(common_dir, fname)\n        if os.path.exists(src) and not os.path.exists(dst):\n            shutil.copy2(src, dst)\n            print(f\"  Copied {fname}\")\n        elif os.path.exists(dst):\n            print(f\"  Already exists: {fname}\")\n        else:\n            print(f\"  WARNING: source not found: {src}\")\n\n    # --- Create stub for clusters-by-entity-40.txt ---\n    # This file maps PDB chains to sequence clusters. Protenix reads it during\n    # data loading to avoid training/test leakage - for inference it's just\n    # checked for existence. An empty file is sufficient.\n    clusters_path = os.path.join(common_dir, \"clusters-by-entity-40.txt\")\n    if not os.path.exists(clusters_path):\n        with open(clusters_path, \"w\") as f:\n            f.write(\"\")  # empty stub - valid, no clusters\n        print(\"  Created stub: clusters-by-entity-40.txt\")\n\n    # --- Create stub for obsolete_release_date.csv ---\n    # Contains PDB release dates for obsolete entries - not needed for inference\n    # on new sequences. Empty CSV with header is sufficient.\n    obsolete_path = os.path.join(common_dir, \"obsolete_release_date.csv\")\n    if not os.path.exists(obsolete_path):\n        with open(obsolete_path, \"w\") as f:\n            f.write(\"pdb_id,release_date\\n\")  # valid empty CSV\n        print(\"  Created stub: obsolete_release_date.csv\")\n\n    # --- Symlink checkpoint directory ---\n    # Checkpoints are large .pt files - symlink instead of copying\n    if not os.path.exists(ckpt_dir):\n        os.symlink(src_ckpt, ckpt_dir)\n        print(f\"  Symlinked checkpoint/ -> {src_ckpt}\")\n    else:\n        print(\"  checkpoint/ already linked\")\n\n    # Verify all required files are present\n    required = [\n        os.path.join(common_dir, \"components.cif\"),\n        os.path.join(common_dir, \"components.cif.rdkit_mol.pkl\"),\n        os.path.join(common_dir, \"clusters-by-entity-40.txt\"),\n        os.path.join(common_dir, \"obsolete_release_date.csv\"),\n        os.path.join(ckpt_dir,   f\"{PROTENIX_MODEL_NAME}.pt\"),\n    ]\n    all_ok = True\n    for p in required:\n        exists = os.path.exists(p)\n        print(f\"  {'OK' if exists else 'MISSING'}: {os.path.basename(p)}\")\n        if not exists: all_ok = False\n\n    if not all_ok:\n        raise RuntimeError(\"Protenix cache preparation failed - missing required files\")\n\n    print(f\"Protenix cache ready at {cache_dir}\")\n    return cache_dir\n\n\n# ==================== GEOMETRY ====================\n\ndef kabsch_align(P, Q):\n    P=np.asarray(P,dtype=np.float64); Q=np.asarray(Q,dtype=np.float64)\n    Pc=P.mean(0,keepdims=True); Qc=Q.mean(0,keepdims=True)\n    H=(P-Pc).T@(Q-Qc); U,S,Vt=np.linalg.svd(H)\n    d=np.sign(np.linalg.det(Vt.T@U.T))\n    R=Vt.T@np.diag([1.,1.,d])@U.T\n    return R, Qc[0]-Pc[0]@R.T\n\ndef apply_rt(X,R,t): return X@R.T+t\ndef rmsd(P,Q): return float(np.sqrt(np.mean(np.sum((np.asarray(P)-np.asarray(Q))**2,axis=1))))\n\n\n# ==================== SUBMISSION HELPERS ====================\n\ndef coords_to_rows(target_id, sequence, samples):\n    rows=[]\n    for i,base in enumerate(sequence,start=1):\n        row={\"ID\":f\"{target_id}_{i}\",\"resname\":base,\"resid\":i}\n        for s in range(samples.shape[0]):\n            row[f\"x_{s+1}\"]=float(samples[s,i-1,0])\n            row[f\"y_{s+1}\"]=float(samples[s,i-1,1])\n            row[f\"z_{s+1}\"]=float(samples[s,i-1,2])\n        rows.append(row)\n    return rows\n\ndef pad_or_trim_samples(sample_list, n_sample):\n    sample_list=list(sample_list)\n    if not sample_list: raise ValueError(\"empty list\")\n    while len(sample_list)<n_sample: sample_list.append(sample_list[-1].copy())\n    return np.stack(sample_list[:n_sample],axis=0)\n\n\n# ==================== CHUNKING ====================\n\ndef split_into_chunks(seq_len, max_len=512, overlap=320):\n    if seq_len<=max_len: return [(0,seq_len)]\n    chunks,start=[],0\n    while start<seq_len:\n        end=min(seq_len,start+max_len); chunks.append((start,end))\n        if end==seq_len: break\n        start=end-overlap\n    return chunks\n\ndef blend_weights(length):\n    return np.clip(1.-np.abs(np.linspace(-1,1,length)),1e-6,None)\n\ndef stitch_chunk_coords(chunk_coords, chunk_ranges, full_len):\n    acc=np.zeros((full_len,3),dtype=np.float64); wgt=np.zeros(full_len,dtype=np.float64)\n    prev_range=prev_chunk=None\n    for X,(s,e) in zip(chunk_coords,chunk_ranges):\n        X=np.asarray(X,dtype=np.float64).copy()\n        if prev_chunk is not None:\n            ps,pe=prev_range; ov_s,ov_e=max(ps,s),min(pe,e)\n            if ov_e-ov_s>=8:\n                R,t=kabsch_align(X[ov_s-s:ov_e-s],prev_chunk[ov_s-ps:ov_e-ps])\n                X=apply_rt(X,R,t)\n        ww=blend_weights(e-s); acc[s:e]+=X*ww[:,None]; wgt[s:e]+=ww\n        prev_range=(s,e); prev_chunk=X\n    return (acc/np.maximum(wgt[:,None],1e-8)).astype(np.float32)\n\n\n# ==================== TRAINING DATA ====================\n\ndef process_labels(labels_df):\n    coords={}; prefixes=labels_df[\"ID\"].str.rsplit(\"_\",n=1).str[0]\n    for prefix,grp in labels_df.groupby(prefixes):\n        coords[prefix]=grp.sort_values(\"resid\")[[\"x_1\",\"y_1\",\"z_1\"]].values.astype(np.float32)\n    return coords\n\n\n# ==================== TEMPLATE RETRIEVAL ====================\n\ndef build_kmer_set(seq, k=4):\n    return {seq} if len(seq)<k else {seq[i:i+k] for i in range(len(seq)-k+1)}\n\ndef kmer_jaccard(a, b):\n    return len(a&b)/max(len(a|b),1)\n\ndef _make_bio_aligner():\n    a=PairwiseAligner(); a.mode=\"global\"\n    a.match_score=2.0; a.mismatch_score=-1.0\n    a.open_gap_score=-2.5; a.extend_gap_score=-0.5\n    return a\n\nBIO_ALIGNER=_make_bio_aligner()\nRNA_MATRIX=None\nif HAS_PARASAIL:\n    try: RNA_MATRIX=parasail.matrix_create(\"ACGU\",2,-1)\n    except: pass\n\ndef _biopython_align(query, template):\n    aln=next(iter(BIO_ALIGNER.align(query,template)))\n    q_segs,t_segs=aln.aligned; aq,at,qi,ti=[],[],0,0\n    for (qs,qe),(ts,te) in zip(q_segs,t_segs):\n        while qi<qs: aq.append(query[qi]);    at.append(\"-\");           qi+=1\n        while ti<ts: aq.append(\"-\");          at.append(template[ti]);  ti+=1\n        for qp,tp in zip(range(qs,qe),range(ts,te)):\n            aq.append(query[qp]); at.append(template[tp])\n        qi,ti=qe,te\n    while qi<len(query):    aq.append(query[qi]);    at.append(\"-\");           qi+=1\n    while ti<len(template): aq.append(\"-\");           at.append(template[ti]); ti+=1\n    q_aln,t_aln=\"\".join(aq),\"\".join(at)\n    matches=sum(a==b for a,b in zip(q_aln,t_aln) if a!=\"-\" and b!=\"-\")\n    aligned=sum(a!=\"-\" and b!=\"-\" for a,b in zip(q_aln,t_aln))\n    ident=100.*matches/max(aligned,1)\n    score=aln.score/max(2*min(len(query),len(template)),1)\n    cover=aligned/max(len(query),1)\n    return score,ident,cover,q_aln,t_aln\n\ndef parasail_global_stats(query, template):\n    if not HAS_PARASAIL or RNA_MATRIX is None: return _biopython_align(query,template)\n    result=parasail.nw_trace_striped_16(query,template,8,1,RNA_MATRIX)\n    cigar=result.cigar.decode.decode() if result.cigar is not None else None\n    if cigar is None: return _biopython_align(query,template)\n    q_aln,t_aln,qi,ti,num=[],[],0,0,\"\"\n    for ch in cigar:\n        if ch.isdigit(): num+=ch; continue\n        n=int(num) if num else 1; num=\"\"\n        if ch in (\"=\",\"X\",\"M\"): q_aln+=list(query[qi:qi+n]); t_aln+=list(template[ti:ti+n]); qi+=n; ti+=n\n        elif ch==\"I\": q_aln+=list(query[qi:qi+n]); t_aln+=[\"-\"]*n; qi+=n\n        elif ch==\"D\": q_aln+=[\"-\"]*n; t_aln+=list(template[ti:ti+n]); ti+=n\n    q_aln,t_aln=\"\".join(q_aln),\"\".join(t_aln)\n    matches=sum(a==b for a,b in zip(q_aln,t_aln) if a!=\"-\" and b!=\"-\")\n    aligned=sum(a!=\"-\" and b!=\"-\" for a,b in zip(q_aln,t_aln))\n    ident=100.*matches/max(aligned,1)\n    score=result.score/max(2*min(len(query),len(template)),1)\n    cover=aligned/max(len(query),1)\n    return score,ident,cover,q_aln,t_aln\n\ndef external_template_priors(test_target_id):\n    if not USE_EXTERNAL_NOTEBOOK_PRIORS: return []\n    priors,files=[],[]\n    for dirs in ATTACHED_NOTEBOOK_DIRS.values():\n        for d in dirs:\n            files+=glob.glob(os.path.join(d,\"**\",\"*.csv\"),recursive=True)\n            files+=glob.glob(os.path.join(d,\"**\",\"*.json\"),recursive=True)\n    for fp in files[:600]:\n        if not any(k in fp.lower() for k in [\"template\",\"mmseqs\",\"candidate\",\"match\",\"tbm\"]): continue\n        try:\n            if fp.endswith(\".csv\"):\n                df=pd.read_csv(fp); cols={c.lower():c for c in df.columns}\n                tc=next((cols[c] for c in cols if c in [\"target_id\",\"query_id\",\"id\",\"target\"]),None)\n                mc=next((cols[c] for c in cols if c in [\"template_id\",\"hit_id\",\"match_id\",\"template\",\"candidate_id\"]),None)\n                if tc and mc: priors+=df[df[tc].astype(str)==str(test_target_id)][mc].astype(str).tolist()\n            elif fp.endswith(\".json\"):\n                obj=json.load(open(fp))\n                if isinstance(obj,dict) and test_target_id in obj and isinstance(obj[test_target_id],list):\n                    priors+=[str(x) for x in obj[test_target_id]]\n        except: pass\n    out,seen=[],set()\n    for x in priors:\n        if x not in seen: out.append(x); seen.add(x)\n    return out\n\ndef build_train_metadata(train_df, coords_dict, k=4):\n    return [{\"target_id\":r[\"target_id\"],\"sequence\":r[\"sequence\"],\n             \"coords\":coords_dict[r[\"target_id\"]],\"length\":len(r[\"sequence\"]),\n             \"kmers\":build_kmer_set(r[\"sequence\"],k=k)}\n            for _,r in train_df.iterrows() if r[\"target_id\"] in coords_dict]\n\ndef shortlist_templates(query_seq, train_meta, target_id=None):\n    qk=build_kmer_set(query_seq,k=4); qlen=len(query_seq)\n    prior_ids=set(external_template_priors(target_id)) if target_id else set()\n    scored=[]\n    for item in train_meta:\n        if target_id and item[\"target_id\"]==target_id: continue  # skip self-match\n        tlen=item[\"length\"]\n        if max(qlen,tlen)/max(1,min(qlen,tlen))>MAX_LEN_RATIO: continue\n        jac=kmer_jaccard(qk,item[\"kmers\"])\n        rough=jac+(0.12 if item[\"target_id\"] in prior_ids else 0)-0.015*abs(math.log(max(tlen,1)/max(qlen,1)))\n        scored.append((rough,item))\n    scored.sort(key=lambda x:x[0],reverse=True)\n    return [x[1] for x in scored[:TBM_TOP_K_SHORTLIST]]\n\ndef composite_template_search(query_seq, train_meta, target_id=None, top_n=TBM_TOP_N):\n    shortlist=shortlist_templates(query_seq,train_meta,target_id=target_id)\n    qk=build_kmer_set(query_seq,k=4); out=[]\n    for item in shortlist:\n        score,ident,cover,aq,at=parasail_global_stats(query_seq,item[\"sequence\"])\n        jac=kmer_jaccard(qk,item[\"kmers\"])\n        lr=min(len(query_seq),len(item[\"sequence\"]))/max(len(query_seq),len(item[\"sequence\"]))\n        comp=0.50*score+0.20*(ident/100)+0.15*cover+0.10*jac+0.05*lr\n        out.append({\"template_id\":item[\"target_id\"],\"template_seq\":item[\"sequence\"],\n                    \"template_coords\":item[\"coords\"],\"score_norm\":score,\n                    \"identity\":ident,\"coverage\":cover,\"kmer_jaccard\":jac,\n                    \"composite\":comp,\"aq\":aq,\"at\":at})\n    out.sort(key=lambda x:x[\"composite\"],reverse=True)\n    return [x for x in out if x[\"composite\"]>=MIN_TEMPLATE_COMPOSITE and x[\"identity\"]>=MIN_TEMPLATE_IDENTITY][:top_n]\n\n\n# ==================== TEMPLATE ADAPTATION ====================\n\ndef adapt_template_to_query(query_seq, template_seq, template_coords, q_aln=None, t_aln=None):\n    if q_aln is None: _,_,_,q_aln,t_aln=parasail_global_stats(query_seq,template_seq)\n    new_coords=np.full((len(query_seq),3),np.nan,dtype=np.float32)\n    qi=ti=-1\n    for qa,ta in zip(q_aln,t_aln):\n        if qa!=\"-\": qi+=1\n        if ta!=\"-\": ti+=1\n        if qa!=\"-\" and ta!=\"-\" and 0<=qi<len(query_seq) and 0<=ti<len(template_coords):\n            new_coords[qi]=template_coords[ti]\n    known=np.where(~np.isnan(new_coords[:,0]))[0]\n    if len(known)==0:\n        return np.array([[i*5.9,0.,0.] for i in range(len(query_seq))],dtype=np.float32)\n    for i in range(len(query_seq)):\n        if not np.isnan(new_coords[i,0]): continue\n        left=known[known<i]; right=known[known>i]\n        if len(left) and len(right):\n            l,r=left[-1],right[0]; w=(i-l)/max(r-l,1)\n            new_coords[i]=(1-w)*new_coords[l]+w*new_coords[r]\n        elif len(left):\n            l=left[-1]; step=np.array([5.9,0,0],dtype=np.float32)\n            if l-1>=0 and not np.isnan(new_coords[l-1,0]): step=new_coords[l]-new_coords[l-1]\n            new_coords[i]=new_coords[l]+step\n        else:\n            r=right[0]; step=np.array([-5.9,0,0],dtype=np.float32)\n            if r+1<len(query_seq) and not np.isnan(new_coords[r+1,0]): step=new_coords[r]-new_coords[r+1]\n            new_coords[i]=new_coords[r]+step\n    return new_coords.astype(np.float32)\n\ndef geometry_score(coords):\n    X=np.asarray(coords,dtype=np.float64)\n    if len(X)<3: return -1e9\n    d1=np.linalg.norm(X[1:]-X[:-1],axis=1); d2=np.linalg.norm(X[2:]-X[:-2],axis=1)\n    bond_pen=np.mean(np.abs(d1-5.95)); next_pen=np.mean(np.abs(d2-10.2))\n    smooth_pen=np.mean(np.linalg.norm(X[2:]-2*X[1:-1]+X[:-2],axis=1))\n    idx=np.arange(len(X))\n    if len(X)>220: idx=np.linspace(0,len(X)-1,200).astype(int)\n    P=X[idx]; diff=P[:,None,:]-P[None,:,:]\n    dm=np.linalg.norm(diff,axis=2)+np.eye(len(P))*9e9\n    sep=np.abs(idx[:,None]-idx[None,:]); close=(sep>2)&(dm<3.2)\n    clash_pen=float((3.2-dm[close]).sum()) if np.any(close) else 0.\n    return float(-(1.9*bond_pen+1.2*next_pen+0.08*smooth_pen+0.02*clash_pen))\n\ndef adaptive_rna_constraints(coords, passes=3, strength=0.25):\n    X=np.asarray(coords,dtype=np.float64).copy()\n    for _ in range(passes):\n        if len(X)>=2:\n            d=X[1:]-X[:-1]; dist=np.linalg.norm(d,axis=1)+1e-8\n            adj=d*((5.95-dist)/dist)[:,None]*strength\n            X[:-1]-=0.5*adj; X[1:]+=0.5*adj\n        if len(X)>=3:\n            d2=X[2:]-X[:-2]; dist2=np.linalg.norm(d2,axis=1)+1e-8\n            adj2=d2*((10.2-dist2)/dist2)[:,None]*(0.4*strength)\n            X[:-2]-=0.5*adj2; X[2:]+=0.5*adj2\n            X[1:-1]+=0.08*strength*(0.5*(X[:-2]+X[2:])-X[1:-1])\n        idx=np.arange(len(X))\n        if len(X)>220: idx=np.linspace(0,len(X)-1,200).astype(int)\n        P=X[idx]; diff=P[:,None,:]-P[None,:,:]\n        dm=np.linalg.norm(diff,axis=2)+1e-8\n        sep=np.abs(idx[:,None]-idx[None,:]); mask=(sep>2)&(dm<3.2)\n        if np.any(mask): X[idx]+=0.01*strength*(diff*((3.2-dm)/dm)[:,:,None]*mask[:,:,None]).sum(axis=1)\n    return X.astype(np.float32)\n\ndef diversify_template_candidate(base_coords, variant_idx, rng):\n    X=np.asarray(base_coords,dtype=np.float32).copy(); L=len(X)\n    if variant_idx==0: return adaptive_rna_constraints(X,passes=2,strength=0.18)\n    if variant_idx==1: X+=rng.normal(0,0.12,X.shape).astype(np.float32)\n    elif variant_idx==2 and L>24:\n        pivot=int(rng.integers(8,max(9,L-8))); angle=np.deg2rad(float(rng.uniform(-18,18)))\n        axis=rng.normal(size=3); axis/=np.linalg.norm(axis)+1e-8; ux,uy,uz=axis; c,s=np.cos(angle),np.sin(angle)\n        R=np.array([[c+ux*ux*(1-c),ux*uy*(1-c)-uz*s,ux*uz*(1-c)+uy*s],\n                    [uy*ux*(1-c)+uz*s,c+uy*uy*(1-c),uy*uz*(1-c)-ux*s],\n                    [uz*ux*(1-c)-uy*s,uz*uy*(1-c)+ux*s,c+uz*uz*(1-c)]],dtype=np.float32)\n        anchor=X[pivot].copy(); X[pivot+1:]=(X[pivot+1:]-anchor)@R.T+anchor\n    elif variant_idx==3 and L>16:\n        t=np.linspace(0,1,L).astype(np.float32)\n        X+=np.stack([np.sin(2*np.pi*t),np.cos(2*np.pi*t),np.sin(4*np.pi*t)],axis=1)*float(rng.uniform(0.2,0.8))\n    else: X+=rng.normal(0,0.18,X.shape).astype(np.float32)\n    return adaptive_rna_constraints(X,passes=3,strength=0.22)\n\ndef candidate_distance(a, b):\n    a=np.asarray(a,dtype=np.float64); b=np.asarray(b,dtype=np.float64)\n    return rmsd(apply_rt(a,*kabsch_align(a,b)),b)\n\ndef select_diverse_top_candidates(candidates, n_sample=5, diversity_threshold=2.25):\n    candidates=sorted(candidates,key=lambda x:x[\"quality\"],reverse=True); chosen=[]\n    for cand in candidates:\n        if not chosen: chosen.append(cand); continue\n        if min(candidate_distance(cand[\"coords\"],c[\"coords\"]) for c in chosen)>=diversity_threshold:\n            chosen.append(cand)\n        if len(chosen)>=n_sample: break\n    used=set(id(c) for c in chosen)\n    for cand in candidates:\n        if len(chosen)>=n_sample: break\n        if id(cand) not in used: chosen.append(cand)\n    return chosen[:n_sample]\n\n\n# ==================== EXTERNAL PREDICTIONS ====================\n\ndef load_external_predictions(target_id, seq):\n    if not USE_EXTERNAL_NOTEBOOK_PREDICTIONS: return []\n    preds,files=[],[]\n    for dirs in ATTACHED_NOTEBOOK_DIRS.values():\n        for d in dirs:\n            files+=glob.glob(os.path.join(d,\"**\",\"*.csv\"),recursive=True)\n            files+=glob.glob(os.path.join(d,\"**\",\"*.parquet\"),recursive=True)\n    for fp in files[:700]:\n        if not any(k in fp.lower() for k in [\"submission\",\"pred\",\"coord\"]): continue\n        try:\n            df=pd.read_csv(fp) if fp.endswith(\".csv\") else pd.read_parquet(fp)\n            if \"ID\" not in df.columns: continue\n            sub=df[df[\"ID\"].astype(str).str.startswith(f\"{target_id}_\")].copy()\n            if len(sub)!=len(seq): continue\n            sub[\"resid\"]=sub[\"ID\"].str.rsplit(\"_\",n=1).str[-1].astype(int)\n            sub=sub.sort_values(\"resid\")\n            for s in range(1,N_SAMPLE+1):\n                cols=[f\"x_{s}\",f\"y_{s}\",f\"z_{s}\"]\n                if all(c in sub.columns for c in cols):\n                    preds.append(sub[cols].values.astype(np.float32))\n        except: pass\n    uniq=[]\n    for P in preds:\n        if all(candidate_distance(P,Q)>=0.6 for Q in uniq): uniq.append(P)\n    return uniq[:max(0,N_SAMPLE+2)]\n\n\n# ==================== PROTENIX SUBPROCESS ====================\n\ndef build_protenix_input_json(tasks, output_path):\n    # Correct Protenix JSON format (confirmed from examples/input.json):\n    #   - top-level key is \"name\" (NOT \"sample_name\")\n    #   - rnaSequence needs \"modifications\": [] field\n    records = []\n    for t in tasks:\n        records.append({\n            \"name\":          t[\"target_id\"],   # <-- KEY FIX: must be \"name\" not \"sample_name\"\n            \"covalent_bonds\": [],\n            \"sequences\": [{\n                \"rnaSequence\": {\n                    \"sequence\":      t[\"sequence\"],\n                    \"count\":         1,\n                    \"modifications\": []           # required field\n                }\n            }]\n        })\n    with open(output_path, \"w\") as f:\n        json.dump(records, f, indent=2)\n    return output_path\n\n\ndef run_protenix_subprocess(input_json_path, dump_dir, cache_dir,\n                            model_name, n_sample, n_cycle, n_step, seed):\n    # FIX: set PROTENIX_ROOT_DIR to the writable cache dir so all\n    # data paths resolve there (where files are pre-populated)\n    env=os.environ.copy()\n    env[\"PROTENIX_ROOT_DIR\"] = cache_dir          # <-- KEY FIX\n    env[\"PYTHONPATH\"]        = PROTENIX_ROOT+\":\"+env.get(\"PYTHONPATH\",\"\")\n\n    inference_script=os.path.join(PROTENIX_ROOT,\"runner\",\"inference.py\")\n    os.makedirs(dump_dir,exist_ok=True)\n\n    cmd=[\n        sys.executable, inference_script,\n        \"--model_name\",               model_name,\n        \"--seeds\",                     str(seed),\n        \"--dump_dir\",                  dump_dir,\n        \"--input_json_path\",           input_json_path,\n        \"--model.N_cycle\",             str(n_cycle),\n        \"--sample_diffusion.N_sample\", str(n_sample),\n        \"--sample_diffusion.N_step\",   str(n_step),\n        \"--triangle_attention\",        PROTENIX_TRI_ATT,\n        \"--triangle_multiplicative\",   PROTENIX_TRI_MUL,\n        \"--use_msa\",                   str(USE_MSA).lower(),\n        \"--use_template\",              str(USE_TEMPLATE).lower(),\n        \"--use_rna_msa\",               str(USE_RNA_MSA).lower(),\n    ]\n    print(\"Running Protenix subprocess...\")\n    print(\"CMD:\", \" \".join(cmd[:8]),\"...\")\n    result=subprocess.run(cmd,env=env,capture_output=False,text=True)\n    if result.returncode!=0:\n        print(f\"Protenix subprocess exited with code {result.returncode}\")\n        return False\n    return True\n\n\ndef parse_cif_coords(cif_path, seq_len):\n    # Parse Protenix mmCIF output. Priority: C1' > C4' > P per residue\n    try:\n        with open(cif_path) as f: lines=f.readlines()\n        in_atom_loop=False; col_map={}; col_idx=0; atom_records=[]\n        i=0\n        while i<len(lines):\n            line=lines[i].strip()\n            if line==\"loop_\":\n                j=i+1\n                while j<len(lines) and lines[j].strip().startswith(\"_\"): j+=1\n                if any(\"_atom_site.\" in lines[k] for k in range(i+1,j)):\n                    in_atom_loop=True; col_map={}; col_idx=0; i+=1; continue\n                else: in_atom_loop=False\n            if in_atom_loop:\n                if line.startswith(\"_atom_site.\"):\n                    col_map[line.split(\".\")[1].strip()]=col_idx; col_idx+=1\n                elif line.startswith(\"_\") or line==\"#\": in_atom_loop=False\n                elif line and not line.startswith(\"#\"):\n                    parts=line.split()\n                    if len(parts)>=col_idx and col_map: atom_records.append(parts)\n            i+=1\n        if not atom_records or not col_map: return None\n        ci_name =col_map.get(\"label_atom_id\",col_map.get(\"auth_atom_id\",None))\n        ci_resid=col_map.get(\"label_seq_id\",col_map.get(\"auth_seq_id\",None))\n        ci_x=col_map.get(\"Cartn_x\",None); ci_y=col_map.get(\"Cartn_y\",None); ci_z=col_map.get(\"Cartn_z\",None)\n        ci_group=col_map.get(\"group_PDB\",None)\n        if any(c is None for c in [ci_name,ci_resid,ci_x,ci_y,ci_z]): return None\n        PRIORITY={\"C1'\":0,\"C1*\":0,\"C4'\":1,\"C4*\":1,\"P\":2}\n        best={}\n        for parts in atom_records:\n            try:\n                if ci_group is not None and parts[ci_group] not in (\"ATOM\",\"HETATM\"): continue\n                atom_name=parts[ci_name]; prio=PRIORITY.get(atom_name,99)\n                if prio==99: continue\n                resid=int(parts[ci_resid])\n                xyz=np.array([float(parts[ci_x]),float(parts[ci_y]),float(parts[ci_z])],dtype=np.float32)\n                if resid not in best or prio<best[resid][0]: best[resid]=(prio,xyz)\n            except (ValueError,IndexError): continue\n        if not best: return None\n        coords=np.zeros((seq_len,3),dtype=np.float32)\n        residues_found=sorted(best.keys()); offset=residues_found[0]-1\n        for resid,(_,xyz) in best.items():\n            arr_idx=resid-1-offset\n            if 0<=arr_idx<seq_len: coords[arr_idx]=xyz\n        known=np.where(np.any(coords!=0,axis=1))[0]\n        if len(known)<2: return None\n        for idx in range(seq_len):\n            if not np.any(coords[idx]!=0):\n                left=known[known<idx]; right=known[known>idx]\n                if len(left) and len(right):\n                    l,r=left[-1],right[0]; w=(idx-l)/max(r-l,1)\n                    coords[idx]=(1-w)*coords[l]+w*coords[r]\n                elif len(left): coords[idx]=coords[left[-1]]\n                elif len(right): coords[idx]=coords[right[0]]\n        return coords\n    except Exception as e:\n        print(f\"  CIF parse error ({os.path.basename(cif_path)}): {e}\"); return None\n\n\ndef collect_protenix_predictions(dump_dir, tasks, n_sample):\n    # Read Protenix CIF outputs from dump_dir/sample_name/**/*.cif\n    results={}\n    for task in tasks:\n        tid=task[\"target_id\"]; seq_len=len(task[\"sequence\"])\n        cif_files=sorted(glob.glob(os.path.join(dump_dir,tid,\"**\",\"*.cif\"),recursive=True))\n        if not cif_files:\n            cif_files=sorted(glob.glob(os.path.join(dump_dir,\"**\",f\"{tid}*.cif\"),recursive=True))\n        if not cif_files:\n            print(f\"  No CIF output for {tid}\"); results[tid]=None; continue\n        sample_coords=[]\n        for cif_path in cif_files:\n            coords=parse_cif_coords(cif_path,seq_len)\n            if coords is not None: sample_coords.append(coords)\n            if len(sample_coords)>=n_sample: break\n        if not sample_coords:\n            print(f\"  CIF parse failed for {tid} ({len(cif_files)} files)\"); results[tid]=None\n        else:\n            print(f\"  {tid}: {len(sample_coords)}/{len(cif_files)} CIFs parsed OK\")\n            while len(sample_coords)<n_sample: sample_coords.append(sample_coords[-1].copy())\n            results[tid]=np.stack(sample_coords[:n_sample],axis=0)\n    return results\n\n\n# ==================== TBM CANDIDATES ====================\n\ndef build_template_candidates(target_id, seq, train_meta):\n    hits=composite_template_search(seq,train_meta,target_id=target_id,top_n=TBM_TOP_N)\n    candidates=[]; n_unique_hits=min(len(hits),TBM_MAX_CANDIDATES)\n    for rank,hit in enumerate(hits[:TBM_MAX_CANDIDATES]):\n        adapted=adapt_template_to_query(seq,hit[\"template_seq\"],hit[\"template_coords\"],\n                                        q_aln=hit[\"aq\"],t_aln=hit[\"at\"])\n        for j in range(2 if rank<3 else 1):\n            rng=np.random.default_rng(SEED+1009*rank+97*j+len(seq))\n            X=diversify_template_candidate(adapted,j,rng)\n            candidates.append({\"source\":f\"tbm:{hit['template_id']}\",\"coords\":X,\n                                \"quality\":float(geometry_score(X)+2.4*hit[\"composite\"]+\n                                               0.25*hit[\"coverage\"]+0.15*(hit[\"identity\"]/100)),\n                                \"template_hit\":hit})\n    return candidates,hits,n_unique_hits\n\ndef generate_denovo(seq, seed):\n    rng=np.random.default_rng(seed); n=len(seq); coords=np.zeros((n,3),dtype=np.float32)\n    radius=float(rng.uniform(7.5,11.)); pitch=float(rng.uniform(2.1,2.8))\n    for i in range(n):\n        ang=0.58*i+rng.normal(0,0.03)\n        coords[i]=[radius*np.cos(ang),radius*np.sin(ang),pitch*i]\n    return adaptive_rna_constraints(coords,passes=3,strength=0.25)\n\ndef quality_boost_by_consensus(candidate, all_candidates):\n    if len(all_candidates)<=1: return candidate[\"quality\"]\n    dists=[candidate_distance(candidate[\"coords\"],o[\"coords\"]) for o in all_candidates if o is not candidate]\n    return candidate[\"quality\"]-0.10*max(float(np.median(dists))-4.,0.) if dists else candidate[\"quality\"]\n\n\n# ==================== MAIN ====================\n\ndef main():\n    test_df=pd.read_csv(DEFAULT_TEST_CSV)\n    train_df=pd.concat([pd.read_csv(DEFAULT_TRAIN_CSV),pd.read_csv(DEFAULT_VAL_CSV)],ignore_index=True)\n    train_labels=pd.concat([pd.read_csv(DEFAULT_TRAIN_LBLS,low_memory=False),\n                             pd.read_csv(DEFAULT_VAL_LBLS,low_memory=False)],ignore_index=True)\n    coords_dict=process_labels(train_labels)\n    train_meta=build_train_metadata(train_df,coords_dict,k=4)\n    print(f\"Loaded template pool: {len(train_meta)} structures\")\n\n    per_target_candidates={}\n    protenix_tasks=[]\n\n    # ---- Phase 1: TBM ----\n    for _,row in test_df.iterrows():\n        target_id=row[\"target_id\"]; seq=row[\"sequence\"]\n        tbm_cands,hits,n_hits=build_template_candidates(target_id,seq,train_meta)\n        ext_cands=[]\n        for i,X in enumerate(load_external_predictions(target_id,seq)):\n            X=adaptive_rna_constraints(X,passes=2,strength=0.18)\n            ext_cands.append({\"source\":f\"external:{i}\",\"coords\":X,\"quality\":geometry_score(X)+0.35})\n        per_target_candidates[target_id]=tbm_cands+ext_cands\n        # FIX: skip Protenix for very long sequences (too slow when chunked)\n        if len(seq) > PROTENIX_MAX_SEQ_LEN:\n            print(f\"  {target_id}: len={len(seq)} > PROTENIX_MAX_SEQ_LEN={PROTENIX_MAX_SEQ_LEN}, using TBM only\")\n        elif len(seq)<=MAX_SEQ_LEN:\n            protenix_tasks.append({\"target_id\":target_id,\"sequence\":seq})\n        else:\n            for ci,(s,e) in enumerate(split_into_chunks(len(seq),MAX_SEQ_LEN,CHUNK_OVERLAP)):\n                protenix_tasks.append({\"target_id\":f\"{target_id}_chunk{ci}\",\"sequence\":seq[s:e],\n                                       \"_parent\":target_id,\"_range\":(s,e)})\n        top_hit=hits[0][\"template_id\"] if hits else \"none\"\n        top_comp=hits[0][\"composite\"] if hits else -1.\n        print(f\"{target_id:>20s} | len={len(seq):4d} | tbm_hits={n_hits:2d} | tbm_vars={len(tbm_cands):2d}\"\n              f\" | ext={len(ext_cands):2d} | top={top_hit} | comp={top_comp:.3f} | protenix=Y\")\n\n    # ---- Phase 2: Protenix ----\n    ptx_preds={}\n    if USE_PROTENIX and protenix_tasks:\n        work_dir=Path(\"/kaggle/working\")\n        ptx_input_json=str(work_dir/\"protenix_input_v10.json\")\n        ptx_dump_dir  =str(work_dir/\"protenix_out_v10\")\n\n        # FIX: pre-populate writable cache dir before subprocess\n        print(\"\\nPreparing Protenix cache...\")\n        try:\n            cache_dir=prepare_protenix_cache()\n        except Exception as e:\n            print(f\"Cache prep failed: {e} - falling back to TBM-only\")\n            cache_dir=None\n\n        if cache_dir:\n            build_protenix_input_json(protenix_tasks,ptx_input_json)\n            print(f\"Running Protenix on {len(protenix_tasks)} tasks...\")\n            success=run_protenix_subprocess(\n                ptx_input_json, ptx_dump_dir, cache_dir,\n                PROTENIX_MODEL_NAME, PROTENIX_N_SAMPLE,\n                PROTENIX_N_CYCLE, PROTENIX_N_STEP, SEED\n            )\n            if success:\n                print(\"\\nCollecting Protenix CIF outputs...\")\n                raw_ptx=collect_protenix_predictions(ptx_dump_dir,protenix_tasks,PROTENIX_N_SAMPLE)\n                for _,row in test_df.iterrows():\n                    tid=row[\"target_id\"]; seq=row[\"sequence\"]\n                    chunks=[t for t in protenix_tasks if t.get(\"_parent\",t[\"target_id\"])==tid]\n                    if not chunks: continue\n                    if \"_parent\" not in chunks[0]:\n                        ptx_preds[tid]=raw_ptx.get(tid)\n                    else:\n                        sample_lists=defaultdict(list); ok=True\n                        for c in chunks:\n                            pred=raw_ptx.get(c[\"target_id\"])\n                            if pred is None: ok=False; break\n                            for s in range(pred.shape[0]): sample_lists[s].append((pred[s],c[\"_range\"]))\n                        if ok and sample_lists:\n                            stitched=[stitch_chunk_coords([x for x,_ in sample_lists[s]],\n                                                          [r for _,r in sample_lists[s]],len(seq))\n                                      for s in range(len(sample_lists))]\n                            ptx_preds[tid]=np.stack(stitched,axis=0)\n                        else: ptx_preds[tid]=None\n                n_ok=sum(1 for v in ptx_preds.values() if v is not None)\n                print(f\"Protenix: {n_ok}/{len(test_df)} sequences predicted successfully\")\n            else:\n                print(\"Protenix subprocess failed - TBM-only fallback\")\n    else:\n        print(f\"Protenix skipped: USE_PROTENIX={USE_PROTENIX}\")\n\n    # ---- Phase 3: Reranking + submission ----\n    all_rows=[]\n    for row_idx,row in test_df.reset_index(drop=True).iterrows():\n        target_id=row[\"target_id\"]; seq=row[\"sequence\"]\n        candidates=list(per_target_candidates.get(target_id,[]))\n        if target_id in ptx_preds and ptx_preds[target_id] is not None:\n            for s_idx in range(ptx_preds[target_id].shape[0]):\n                X=adaptive_rna_constraints(ptx_preds[target_id][s_idx],passes=2,strength=0.18)\n                candidates.append({\"source\":f\"protenix:{s_idx}\",\"coords\":X,\n                                    \"quality\":geometry_score(X)+0.20})\n        while len(candidates)<N_SAMPLE+2:\n            X=generate_denovo(seq,seed=SEED+7919*row_idx+37*len(candidates))\n            candidates.append({\"source\":\"denovo\",\"coords\":X,\"quality\":geometry_score(X)-0.30})\n        for cand in candidates: cand[\"quality\"]=quality_boost_by_consensus(cand,candidates)\n        selected=select_diverse_top_candidates(candidates,n_sample=N_SAMPLE,diversity_threshold=2.10)\n        samples=pad_or_trim_samples([c[\"coords\"] for c in selected],N_SAMPLE)\n        sources=[c[\"source\"] for c in selected]\n        has_ptx=any(\"protenix\" in s for s in sources)\n        print(f\"Selected for {target_id}: {sources}{'  <- Protenix' if has_ptx else ''}\")\n        all_rows.extend(coords_to_rows(target_id,seq,samples))\n\n    sub=pd.DataFrame(all_rows)\n    cols=[\"ID\",\"resname\",\"resid\"]+[f\"{c}_{i}\" for i in range(1,N_SAMPLE+1) for c in [\"x\",\"y\",\"z\"]]\n    coord_cols=[c for c in cols if c not in (\"ID\",\"resname\",\"resid\")]\n    sub[coord_cols]=sub[coord_cols].clip(-999.999,9999.999)\n    sub=sub[cols]; sub.to_csv(DEFAULT_OUTPUT,index=False)\n    print(\"\\nSaved submission to\",DEFAULT_OUTPUT)\n    print(\"Submission shape:\",sub.shape)\n    print(sub.head(10).to_string(index=False))\n\n\nif __name__==\"__main__\":\n    main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T21:26:04.138849Z","iopub.execute_input":"2026-03-25T21:26:04.139064Z","iopub.status.idle":"2026-03-25T21:50:33.518306Z","shell.execute_reply.started":"2026-03-25T21:26:04.139043Z","shell.execute_reply":"2026-03-25T21:50:33.517527Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---- Verification ----\nimport pandas as pd, os, glob\n\ndf = pd.read_csv(\"/kaggle/working/submission.csv\")\nexpected_cols = [\"ID\",\"resname\",\"resid\"] + [f\"{c}_{i}\" for i in range(1,6) for c in [\"x\",\"y\",\"z\"]]\ncoord_cols    = [c for c in df.columns if c not in (\"ID\",\"resname\",\"resid\")]\nmissing       = [c for c in expected_cols if c not in df.columns]\nnan_cnt       = df.isnull().sum().sum()\n\nptx_cifs  = glob.glob(\"/kaggle/working/protenix_out_v10/**/*.cif\", recursive=True)\ncache_ok  = os.path.exists(\"/kaggle/working/protenix_cache/common/clusters-by-entity-40.txt\")\nptx_ran   = os.path.exists(\"/kaggle/working/protenix_input_v10.json\")\n\nprint(f\"Shape          : {df.shape}\")\nprint(f\"Missing cols   : {missing}\")\nprint(f\"NaN count      : {nan_cnt}\")\nprint(f\"Coord range    : [{df[coord_cols].min().min():.2f}, {df[coord_cols].max().max():.2f}]\")\nprint(f\"Cache prepared : {'Yes' if cache_ok else 'No'}\")\nprint(f\"Protenix input : {'Created' if ptx_ran else 'Not found'}\")\nprint(f\"Protenix CIFs  : {len(ptx_cifs)} files {'(Protenix ran!)' if ptx_cifs else '(TBM-only fallback)'}\")\nif ptx_cifs:\n    print(f\"  Sample       : {ptx_cifs[0]}\")\nprint()\nprint(df.head(10).to_string(index=False))\nprint()\nprint(\"All checks passed!\" if not missing and nan_cnt==0 else \"WARNING: issues found!\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T21:50:33.519714Z","iopub.execute_input":"2026-03-25T21:50:33.52007Z","iopub.status.idle":"2026-03-25T21:50:33.579729Z","shell.execute_reply.started":"2026-03-25T21:50:33.520046Z","shell.execute_reply":"2026-03-25T21:50:33.579039Z"}},"outputs":[],"execution_count":null}]}