{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":14516276,"sourceType":"datasetVersion","datasetId":9271406},{"sourceId":14516480,"sourceType":"datasetVersion","datasetId":9271521},{"sourceId":14519000,"sourceType":"datasetVersion","datasetId":9273016},{"sourceId":14595063,"sourceType":"datasetVersion","datasetId":9322806}],"dockerImageVersionId":31259,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import polars as pl\nfrom pathlib import Path\n\nFILE1 = Path(r\"/kaggle/input/submission-v/submission_V.tsv\")\nFILE2 = Path(r\"/kaggle/input/sub0-313\")\n\nOUT = Path(\"/kaggle/working/submission.tsv\")\nOUT.parent.mkdir(parents=True, exist_ok=True)  # กันพลาด\n\n\nPRINT_EVERY_PROTS = 250\nN_PARTS = 4\n\nSCORE_MIN = 1e-4\nSCORE_MAX = 1.0\nS2_MINA = 0.0\nS2_MINB = 0.4\nTOPK_PER_PROTEIN = 100\ndef scan_sub(path: Path, score_col: str) -> pl.LazyFrame:\n    return (\n        pl.scan_csv(\n            str(path),\n            separator=\"\\t\",\n            has_header=False,\n            new_columns=[\"protein_id\", \"go_term\", score_col],\n            schema_overrides={\"protein_id\": pl.Utf8, \"go_term\": pl.Utf8, score_col: pl.Float64},\n            ignore_errors=True,\n        )\n        .with_columns(\n            pl.col(\"protein_id\").str.strip_chars(),\n            pl.col(\"go_term\").str.strip_chars(),\n        )\n    )\n\ndef dedup_keep_max(lf: pl.LazyFrame, score_col: str) -> pl.LazyFrame:\n    return (\n        lf.group_by([\"protein_id\", \"go_term\"])\n          .agg(pl.col(score_col).max().alias(score_col))\n    )\n\ndef main():\n    print(\"[LOAD] FILE1:\", FILE1)\n    print(\"[LOAD] FILE2:\", FILE2)\n\n    # 1) ดึงรายชื่อ protein จาก FILE1 (เบามากเมื่อเทียบกับทั้งไฟล์)\n    # 1) โปรตีนจาก FILE1\n    prot1 = scan_sub(FILE1, \"s1\").select(\"protein_id\").unique()\n    \n    # 2) โปรตีนจาก FILE2\n    prot2 = scan_sub(FILE2, \"s2\").select(\"protein_id\").unique()\n    \n    # 3) เอา intersection: อยู่ทั้ง FILE1 และ FILE2 เท่านั้น\n    prot_list = (\n        prot1.join(prot2, on=\"protein_id\", how=\"inner\")\n             .collect(engine=\"streaming\")\n             .get_column(\"protein_id\")\n             .to_list()\n    )\n    \n    prot_list.sort()\n    n_prots = len(prot_list)\n    print(f\"[DBG] proteins kept (FILE1 ∩ FILE2): {n_prots:,} | parts={N_PARTS}\")\n\n    # กระจายแบบ round-robin ให้บาลานซ์แต่ละพาร์ท\n    parts = [prot_list[i::N_PARTS] for i in range(N_PARTS)]\n    for i, p in enumerate(parts):\n        print(f\"[DBG] part {i+1}/{N_PARTS}: {len(p):,} proteins\")\n\n    # 2) เคลียร์ไฟล์ output ก่อน\n    if OUT.exists():\n        OUT.unlink()\n    with open(OUT, \"wb\") as f:\n        pass\n\n    lf1_all = scan_sub(FILE1, \"s1\")\n    lf2_all = scan_sub(FILE2, \"s2\")\n\n    total_out_rows = 0\n\n    # 3) ทำทีละพาร์ท\n    for part_idx in range(N_PARTS):\n        batch_prots = parts[part_idx]\n        if not batch_prots:\n            continue\n\n        print(f\"\\n[PART] {part_idx+1}/{N_PARTS} | proteins={len(batch_prots):,}\")\n\n        # ตารางรายชื่อโปรตีนของพาร์ทนี้ (ใช้ join แทน is_in(list) จะนิ่งกว่า)\n        prot_df = pl.DataFrame({\"protein_id\": batch_prots})\n        prot_lf = prot_df.lazy()\n\n        # ---- FILE1 subset + dedup\n        lf1 = dedup_keep_max(\n            lf1_all.join(prot_lf, on=\"protein_id\", how=\"inner\"),\n            \"s1\"\n        )\n\n        lf2a = dedup_keep_max(\n            lf2_all.join(prot_lf, on=\"protein_id\", how=\"inner\")\n                  .filter(pl.col(\"s2\") > S2_MINA),\n            \"s2\"\n        )\n        \n        lf2b = dedup_keep_max(\n            lf2_all.join(prot_lf, on=\"protein_id\", how=\"inner\")\n                  .filter(pl.col(\"s2\") > S2_MINB),\n            \"s2\"\n        )\n    # ---- merge + score (ตาม logic เดิมของนิว: ใช้เฉพาะกรณีมี s1)\n        merged = (\n            lf1.join(lf2a, on=[\"protein_id\", \"go_term\"], how=\"left\")\n              .with_columns(\n                  pl.when(pl.col(\"s2\").is_not_null())\n                    .then(0.9 * pl.col(\"s1\") + 0.2 * pl.col(\"s2\"))\n                    .otherwise(0.9 * pl.col(\"s1\"))\n                    .alias(\"score\")\n              )\n              .with_columns(pl.col(\"score\").clip(lower_bound=SCORE_MIN, upper_bound=SCORE_MAX))\n              .select([\"protein_id\", \"go_term\", \"score\"])\n              .group_by([\"protein_id\", \"go_term\"])\n              .agg(pl.col(\"score\").max().alias(\"score\"))\n        )\n\n        # ==========================================================\n        # ✅ F2-only fill-up-to TOPK (ไม่ตัดของ FILE1)\n        #    - need คำนวณจาก FILE1 หลัง dedup (lf1) เท่านั้น\n        #    - fill ใช้ FILE2 หลัง dedup (แต่ไม่ filter S2_MIN เพื่อให้มีของเติม)\n        # ==========================================================\n        \n        # base count จาก FILE1 (หลัง dedup)\n        base_cnt = (\n            lf1.group_by(\"protein_id\")\n               .len()\n               .rename({\"len\": \"base_n\"})\n        )\n        \n\n        need_tbl = (\n            base_cnt.with_columns(\n                (\n                    pl.lit(TOPK_PER_PROTEIN).cast(pl.Int64)\n                    - pl.col(\"base_n\").cast(pl.Int64)\n                )\n                .clip(lower_bound=0)\n                .alias(\"need\")\n            )\n            .select([\"protein_id\", \"base_n\", \"need\"])\n        )\n                \n\n        # --- เอาเฉพาะ F2-only ที่ยังไม่อยู่ในฐาน FILE1 (lf1) และเฉพาะโปรตีนที่ need>0\n        # NOTE: anti join เทียบกับ lf1 (ฐานจริง) ไม่ใช่ merged\n        f2_cand = (\n            lf2b.join(need_tbl.filter(pl.col(\"need\") > 0), on=\"protein_id\", how=\"inner\")\n                      .join(lf1.select([\"protein_id\", \"go_term\"]), on=[\"protein_id\", \"go_term\"], how=\"anti\")\n                      .with_columns(\n                    (pl.col(\"s2\") * 0.675)\n                      .clip(lower_bound=SCORE_MIN, upper_bound=SCORE_MAX)\n                      .alias(\"score\")\n                )\n                      .select([\"protein_id\", \"go_term\", \"score\", \"need\"])\n        )\n        \n        # --- จัดอันดับในแต่ละโปรตีน แล้วเลือกมาเท่าที่ need\n        f2_fill = (\n            f2_cand.sort([\"protein_id\", \"score\"], descending=[False, True])\n                   .with_columns(\n                       (pl.int_range(1, pl.len() + 1).over(\"protein_id\")).alias(\"rn\")\n                   )\n                   .filter(pl.col(\"rn\") <= pl.col(\"need\"))\n                   .select([\"protein_id\", \"go_term\", \"score\"])\n        )\n        \n        # --- merged2 = ฐานเดิม (merged) + เติม (f2_fill)\n        merged2 = (\n            pl.concat([merged, f2_fill], how=\"vertical\")\n              .group_by([\"protein_id\", \"go_term\"])\n              .agg(pl.col(\"score\").max().alias(\"score\"))\n        )\n        \n        final_df = (\n            pl.concat([merged, f2_fill], how=\"vertical\")\n              .group_by([\"protein_id\", \"go_term\"])\n              .agg(pl.col(\"score\").max().alias(\"score\"))\n              .sort([\"protein_id\", \"score\"], descending=[False, True])\n              .collect(engine=\"streaming\")\n        )\n        out_rows = final_df.height\n        total_out_rows += out_rows\n\n        # ---- stats (ของพาร์ทนี้)\n        df1_part = lf1.collect(streaming=True)\n        base_n = df1_part.group_by(\"protein_id\").len().rename({\"len\": \"base_n\"})\n        out_n  = final_df.group_by(\"protein_id\").len().rename({\"len\": \"out_n\"})\n        new_n  = (\n            final_df.join(df1_part.select([\"protein_id\", \"go_term\"]), on=[\"protein_id\", \"go_term\"], how=\"anti\")\n                    .group_by(\"protein_id\").len()\n                    .rename({\"len\": \"new_n\"})\n        )\n        stats = (\n            out_n.join(base_n, on=\"protein_id\", how=\"left\")\n                .join(new_n, on=\"protein_id\", how=\"left\")\n                .with_columns(\n                    pl.col(\"base_n\").fill_null(0),\n                    pl.col(\"new_n\").fill_null(0),\n                    (pl.col(\"out_n\") - pl.col(\"new_n\")).alias(\"kept_from_file1_or_overlap\"),\n                )\n                .sort(\"protein_id\")\n        )\n\n        print(f\"[DBG] part_out_rows={out_rows:,} | part_proteins={stats.height:,}\")\n        lens = stats.select(\n            pl.col(\"base_n\").min().alias(\"base_min\"),\n            pl.col(\"base_n\").max().alias(\"base_max\"),\n            pl.col(\"out_n\").min().alias(\"out_min\"),\n            pl.col(\"out_n\").max().alias(\"out_max\"),\n            pl.col(\"new_n\").sum().alias(\"new_total\"),\n        ).row(0)\n        print(f\"[DBG] base(min/max)={lens[0]}/{lens[1]} | out(min/max)={lens[2]}/{lens[3]} | new_total={lens[4]:,}\")\n\n        # ---- write append ทีละพาร์ท\n        with open(OUT, \"ab\") as f:\n            # ถ้านิวอยาก log แบบเดิมทุก 100 โปรตีนตอนเขียน\n            prots = stats.get_column(\"protein_id\").to_list()\n            for start in range(0, len(prots), PRINT_EVERY_PROTS):\n                chunk_prots = prots[start:start+PRINT_EVERY_PROTS]\n                part = final_df.filter(pl.col(\"protein_id\").is_in(chunk_prots))\n                part.write_csv(f, separator=\"\\t\", include_header=False, float_precision=4)\n\n                st_chunk = stats.filter(pl.col(\"protein_id\").is_in(chunk_prots))\n                tail = st_chunk.tail(1).row(0, named=True)\n\n                base_sum = int(st_chunk[\"base_n\"].sum())\n                new_sum  = int(st_chunk[\"new_n\"].sum())\n                out_sum  = int(st_chunk[\"out_n\"].sum())\n\n                print(\n                    f\"[WRITE][P{part_idx+1}] prots {start+1:,}-{min(start+PRINT_EVERY_PROTS, len(prots)):,} | \"\n                    f\"chunk_out_rows={out_sum:,} | chunk_base_rows={base_sum:,} | chunk_new_rows={new_sum:,} | \"\n                    f\"last_prot={tail['protein_id']} base={tail['base_n']} new={tail['new_n']} out={tail['out_n']}\"\n                )\n\n    print(f\"\\n[OK] wrote: {OUT}\")\n    print(f\"[OK] total_out_rows: {total_out_rows:,}\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-19T08:59:47.193802Z","iopub.execute_input":"2026-01-19T08:59:47.194070Z","iopub.status.idle":"2026-01-19T09:10:31.860770Z","shell.execute_reply.started":"2026-01-19T08:59:47.194038Z","shell.execute_reply":"2026-01-19T09:10:31.859871Z"}},"outputs":[],"execution_count":null}]}