{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","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":14084779,"sourceType":"competition"}],"dockerImageVersionId":31153,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip -q install optuna\n!pip install optuna-integration","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:13:43.825237Z","iopub.execute_input":"2025-10-25T17:13:43.825556Z","iopub.status.idle":"2025-10-25T17:13:52.038479Z","shell.execute_reply.started":"2025-10-25T17:13:43.825531Z","shell.execute_reply":"2025-10-25T17:13:52.037682Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from pathlib import Path\nimport re\nimport pandas as pd\nimport numpy as np\nimport optuna\nimport gc\n\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers, regularizers\nfrom optuna.integration import TFKerasPruningCallback\nfrom sklearn.metrics import f1_score\nfrom tensorflow.keras import backend as K\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:13:52.040441Z","iopub.execute_input":"2025-10-25T17:13:52.040657Z","iopub.status.idle":"2025-10-25T17:14:06.605531Z","shell.execute_reply.started":"2025-10-25T17:13:52.040638Z","shell.execute_reply":"2025-10-25T17:14:06.604680Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from pathlib import Path\n\nFASTA_PATH = Path(\"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\")\nTAX_PATH   = Path(\"/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv\")   # not used here, kept for reference\nTRAIN_PATH = Path(\"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\")\n\n# Output directory & files\nOUT_DIR = Path(\"outputs\")\nOUT_DIR.mkdir(parents=True, exist_ok=True)\nOUT_TSV = OUT_DIR / \"features_propy.tsv\"\nOUT_PQ  = OUT_DIR / \"features_propy.parquet\"\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:14:06.606397Z","iopub.execute_input":"2025-10-25T17:14:06.607035Z","iopub.status.idle":"2025-10-25T17:14:06.611887Z","shell.execute_reply.started":"2025-10-25T17:14:06.607007Z","shell.execute_reply":"2025-10-25T17:14:06.611325Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"AA_STD = set(\"ACDEFGHIKLMNPQRSTVWY\")\n\nGROUPS = {\n    # === Basic (classic physicochemical) ===\n    \"AROMATIC\":       set(\"FWY\"),\n    \"POSITIVE\":       set(\"KRH\"),\n    \"NEGATIVE\":       set(\"DE\"),\n    \"POLAR\":          set(\"STNQYC\"),\n    \"NON_POLAR\":      set(\"AVLIMFWP\"),\n    \"HYDROPHOBIC\":    set(\"AILMFWV\"),\n    \"ALIPHATIC\":      set(\"AILV\"),\n    \"SMALL\":          set(\"AGSV\"),\n    \"PROLINE\":        set(\"P\"),\n\n    # === Physicochemical / functional ===\n    \"SULFUR\":         set(\"CM\"),\n    \"BASIC\":          set(\"KRH\"),          # alias of POSITIVE\n    \"ACIDIC\":         set(\"DE\"),           # alias of NEGATIVE\n    \"SMALL_FLEX\":     set(\"AGST\"),\n    \"BULKY\":          set(\"WFRYK\"),\n    \"TURN_PRONE\":     set(\"GPND\"),\n    \"FLEXIBLE\":       set(\"GSDN\"),\n    \"RIGID\":          set(\"CWYF\"),\n\n    # === Structural propensities ===\n    \"HELIX_FAVORING\": set(\"ALMQEKR\"),      # α-helix promoters\n    \"SHEET_FAVORING\": set(\"VIFYTW\"),       # β-sheet promoters\n    \"HELIX_BREAKERS\": set(\"PG\"),           # α-helix breakers\n    \"SHEET_BREAKERS\": set(\"DEKR\"),         # β-sheet breakers\n    \"TURN_FAVORING\":  set(\"GPNDS\"),        # β-turn prone\n    \"HELIX_CAPPERS\":  set(\"NDST\"),         # often found at helix termini\n    \"BETA_BRIDGING\":  set(\"CFYW\"),         # can bridge β-strands\n\n    # === Size / volume ===\n    \"TINY\":           set(\"AGSC\"),\n    \"SMALL_SIZE\":     set(\"AGSVTP\"),\n    \"MEDIUM_SIZE\":    set(\"NDQEC\"),\n    \"LARGE_SIZE\":     set(\"WFYRKH\"),\n    \"BULKY_SIZE\":     set(\"WFYRK\"),\n\n    # === Hydrophobicity ===\n    \"HYDROPHOBIC_STRONG\":     set(\"ILVFWCM\"),\n    \"HYDROPHILIC_STRONG\":     set(\"DEKRNHQ\"),\n    \"MODERATELY_HYDROPHOBIC\": set(\"ATY\"),\n    \"MODERATELY_HYDROPHILIC\": set(\"SGP\"),\n\n    # === Accessibility and structural position ===\n    \"SURFACE_PRONE\":  set(\"DEKNRQHSTY\"),   # typically surface-exposed residues\n    \"CORE_PRONE\":     set(\"AILMVFWY\"),     # typically buried residues\n\n    # === Specific chemical groups ===\n    \"HYDROXYL\":       set(\"STY\"),\n    \"AMIDE\":          set(\"NQ\"),\n    \"CATIONIC\":       set(\"KRH\"),\n    \"ANIONIC\":        set(\"DE\"),\n    \"NEUTRAL\":        set(\"ACFGILMNPQSTVWY\"),\n\n    # === Functional / biological groups ===\n    \"METAL_BINDING\":         set(\"CHDE\"),\n    \"PHOSPHORYLATION_SITES\": set(\"STY\"),\n    \"GLYCOSYLATION_SITES\":   set(\"NST\"),\n    \"DISULFIDE_FORMING\":     set(\"C\"),\n    \"ZINC_FINGER_CORE\":      set(\"CH\"),\n\n    # === Evolutionary / conservative properties ===\n    \"HYDROPHOBIC_SET\":      set(\"AVLIMFWY\"),\n    \"POLAR_UNCHARGED_SET\":  set(\"STNQC\"),\n    \"CHARGED_SET\":          set(\"DEKR\"),\n    \"SMALL_FLEXIBLE_SET\":   set(\"AGP\"),\n}\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:14:06.612782Z","iopub.execute_input":"2025-10-25T17:14:06.613107Z","iopub.status.idle":"2025-10-25T17:14:06.672832Z","shell.execute_reply.started":"2025-10-25T17:14:06.613082Z","shell.execute_reply":"2025-10-25T17:14:06.672085Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def parse_fasta_idseq(path: Path) -> dict:\n    \"\"\"\n    Parse a FASTA file and return {EntryID: cleaned_sequence}.\n    - Handles UniProt headers like 'sp|/tr|ACC|...'\n    - Keeps only the 20 standard amino acids\n    \"\"\"\n    if not path.exists():\n        raise FileNotFoundError(f\"FASTA not found: {path}\")\n    seqs = {}\n    pid = None\n    with open(path, encoding=\"utf-8\") as f:\n        for raw in f:\n            line = raw.strip()\n            if not line:\n                continue\n            if line.startswith(\">\"):\n                header = line[1:].strip()\n                m = re.match(r\"(?:sp|tr)\\|([^|]+)\\|\", header)\n                pid = m.group(1) if m else header.split()[0].split(\"|\")[0]\n                seqs[pid] = \"\"\n            else:\n                seqs[pid] += \"\".join(c for c in line if c in AA_STD)\n    return seqs\n\ndef frac_and_count(seq: str, aa_set: set) -> tuple[float, int]:\n    n = len(seq)\n    if n == 0:\n        return 0.0, 0\n    cnt = sum(1 for c in seq if c in aa_set)\n    return cnt / n, cnt\n\ndef build_group_features(\n    fasta_path: Path,\n    groups: dict[str, set],\n    include_fractions: bool = True,\n    include_counts: bool = True,\n) -> pd.DataFrame:\n    \"\"\"\n    Compute per-sequence length + group fractions and/or counts.\n    Output columns: len, <group>_frac, <group>_count\n    \"\"\"\n    seqs = parse_fasta_idseq(fasta_path)\n    rows = []\n    for pid, seq in seqs.items():\n        row = {\"EntryID\": pid, \"len\": len(seq)}\n        for gname, gset in groups.items():\n            f, c = frac_and_count(seq, gset)\n            if include_fractions:\n                row[f\"{gname.lower()}_frac\"] = f\n            if include_counts:\n                row[f\"{gname.lower()}_count\"] = c\n        rows.append(row)\n    df = pd.DataFrame(rows).set_index(\"EntryID\").sort_index()\n    return df\n\ndef safe_read_terms(train_path: Path) -> pd.DataFrame:\n    \"\"\"\n    Read term annotations and normalize column names to: EntryID, term.\n    Adjust here if your file uses different names.\n    \"\"\"\n    df = pd.read_csv(train_path, sep=\"\\t\")\n    low = {c.lower(): c for c in df.columns}\n\n    # normalize EntryID\n    if \"entryid\" in low and low[\"entryid\"] != \"EntryID\":\n        df.rename(columns={low[\"entryid\"]: \"EntryID\"}, inplace=True)\n    elif \"uniprot\" in low:\n        df.rename(columns={low[\"uniprot\"]: \"EntryID\"}, inplace=True)\n\n    # normalize term\n    if \"term\" in low and low[\"term\"] != \"term\":\n        df.rename(columns={low[\"term\"]: \"term\"}, inplace=True)\n    elif \"goterm\" in low:\n        df.rename(columns={low[\"goterm\"]: \"term\"}, inplace=True)\n\n    return df\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:14:06.673673Z","iopub.execute_input":"2025-10-25T17:14:06.673940Z","iopub.status.idle":"2025-10-25T17:14:06.688017Z","shell.execute_reply.started":"2025-10-25T17:14:06.673899Z","shell.execute_reply":"2025-10-25T17:14:06.687107Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Build features\ndf_groups = build_group_features(\n    FASTA_PATH,\n    GROUPS,\n    include_fractions=True,\n    include_counts=True,\n)\n\n# Save outputs\ndf_groups.to_csv(OUT_TSV, sep=\"\\t\")\nprint(f\"Saved TSV: {OUT_TSV.resolve()} | shape={df_groups.shape}\")\n\ntry:\n    df_groups.to_parquet(OUT_PQ, index=True)\n    print(f\"Saved Parquet: {OUT_PQ.resolve()}\")\nexcept Exception as e:\n    print(f\"Parquet not written: {e}\")\n\ndf_groups.head()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:14:06.688716Z","iopub.execute_input":"2025-10-25T17:14:06.689033Z","iopub.status.idle":"2025-10-25T17:15:41.551640Z","shell.execute_reply.started":"2025-10-25T17:14:06.689016Z","shell.execute_reply":"2025-10-25T17:15:41.550880Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Merge with GO terms (optional but useful for analysis)\ndf_merged = None\nif TRAIN_PATH.exists():\n    df_terms = safe_read_terms(TRAIN_PATH)\n    if \"EntryID\" in df_terms.columns:\n        df_merged = (\n            df_groups.reset_index()\n                     .merge(df_terms, on=\"EntryID\", how=\"left\")\n                     .set_index(\"EntryID\")\n        )\n        print(\"Merge completed:\", df_merged.shape)\n    else:\n        print(\"Column 'EntryID' not found in train_terms.tsv — skipping merge.\")\nelse:\n    print(\"TRAIN_PATH not found — skipping merge.\")\n\ndf_merged.head() if df_merged is not None else None\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:15:41.553681Z","iopub.execute_input":"2025-10-25T17:15:41.553898Z","iopub.status.idle":"2025-10-25T17:15:42.528561Z","shell.execute_reply.started":"2025-10-25T17:15:41.553881Z","shell.execute_reply":"2025-10-25T17:15:42.527787Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Compute per-term statistics (mean, std, count) for all feature columns\nif df_merged is not None and \"term\" in df_merged.columns:\n    feature_cols = [c for c in df_merged.columns if c.endswith(\"_frac\") or c.endswith(\"_count\") or c == \"len\"]\n    print(f\"Total feature columns: {len(feature_cols)}\")\n\n    term_stats = (\n        df_merged.groupby(\"term\")[feature_cols]\n                 .agg([\"mean\", \"std\", \"count\"])\n                 .reset_index()\n    )\n\n    # Flatten MultiIndex columns\n    term_stats.columns = [f\"{a}_{b}\" if b else a for a, b in term_stats.columns]\n    print(\"Computed term-level stats:\", term_stats.shape)\n    term_stats.head()\nelse:\n    print(\"Skipping: df_merged or 'term' not available.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:15:42.529399Z","iopub.execute_input":"2025-10-25T17:15:42.529684Z","iopub.status.idle":"2025-10-25T17:15:43.906270Z","shell.execute_reply.started":"2025-10-25T17:15:42.529659Z","shell.execute_reply":"2025-10-25T17:15:43.905414Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Heatmap of Z-scores (means vs global means) for top N frequent GO terms\nif df_merged is not None and \"term\" in df_merged.columns:\n    import matplotlib.pyplot as plt\n    import seaborn as sns\n\n    feature_cols = [c for c in df_merged.columns if c.endswith(\"_frac\") or c.endswith(\"_count\") or c == \"len\"]\n\n    # Global stats\n    global_means = df_merged[feature_cols].mean()\n    global_stds  = df_merged[feature_cols].std().replace(0, np.nan)\n\n    # Build term-level means matrix\n    means = (\n        df_merged.groupby(\"term\")[feature_cols]\n                 .mean()\n    )\n\n    # Z = (mean_term - mean_global) / std_global\n    Z = (means - global_means) / global_stds\n\n    # Pick top N terms by frequency\n    top_terms = df_merged[\"term\"].value_counts().head(15).index\n    Z_sub = Z.loc[Z.index.intersection(top_terms)]\n\n    plt.figure(figsize=(14, 8))\n    sns.heatmap(Z_sub, cmap=\"coolwarm\", center=0)\n    plt.title(\"Feature Z-scores by GO term (Top 15 by frequency)\")\n    plt.xlabel(\"Feature\")\n    plt.ylabel(\"GO Term\")\n    plt.tight_layout()\n    plt.show()\nelse:\n    print(\"Skipping heatmap: df_merged or 'term' not available.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:15:43.907131Z","iopub.execute_input":"2025-10-25T17:15:43.907499Z","iopub.status.idle":"2025-10-25T17:15:46.236463Z","shell.execute_reply.started":"2025-10-25T17:15:43.907474Z","shell.execute_reply":"2025-10-25T17:15:46.235633Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TEST_FASTA = Path(\"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\")\n\nOUT_DIR = Path(\"outputs\")\nOUT_DIR.mkdir(exist_ok=True, parents=True)\nSUB_PATH = OUT_DIR / \"submission_cafa6.tsv\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:15:46.237190Z","iopub.execute_input":"2025-10-25T17:15:46.237668Z","iopub.status.idle":"2025-10-25T17:15:46.242412Z","shell.execute_reply.started":"2025-10-25T17:15:46.237648Z","shell.execute_reply":"2025-10-25T17:15:46.241657Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Read training annotations; expect columns: EntryID, term\ndf_terms = pd.read_csv(TRAIN_PATH, sep=\"\\t\")\nassert {\"EntryID\", \"term\"}.issubset(df_terms.columns)\n\n# Ensure every training EntryID in df_groups has a row in the label table (multi-label)\n# 1) Collect terms per EntryID\nterms_by_id = df_terms.groupby(\"EntryID\")[\"term\"].apply(lambda s: sorted(set(s))).reindex(df_groups.index, fill_value=[])\n\n# 2) Build vocabulary of GO terms from training\ngo_terms = sorted(df_terms[\"term\"].unique())\nterm_to_idx = {t:i for i,t in enumerate(go_terms)}\n\n# 3) Build Y matrix (multi-hot)\nY = np.zeros((len(df_groups), len(go_terms)), dtype=np.float32)\nfor r, entry in enumerate(df_groups.index):\n    for t in terms_by_id.loc[entry]:\n        Y[r, term_to_idx[t]] = 1.0\n\n# 4) Feature matrix X\nfeature_cols = [c for c in df_groups.columns if c.endswith(\"_frac\") or c.endswith(\"_count\") or c == \"len\"]\nX = df_groups[feature_cols].astype(np.float32).values\n\nX.shape, Y.shape, len(go_terms)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:15:46.243259Z","iopub.execute_input":"2025-10-25T17:15:46.243668Z","iopub.status.idle":"2025-10-25T17:15:49.813285Z","shell.execute_reply.started":"2025-10-25T17:15:46.243645Z","shell.execute_reply":"2025-10-25T17:15:49.812658Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Read training annotations; expect columns: EntryID, term\ndf_terms = pd.read_csv(TRAIN_PATH, sep=\"\\t\")\nassert {\"EntryID\", \"term\"}.issubset(df_terms.columns)\n\n# Ensure every training EntryID in df_groups has a row in the label table (multi-label)\n# 1) Collect terms per EntryID\nterms_by_id = df_terms.groupby(\"EntryID\")[\"term\"].apply(lambda s: sorted(set(s))).reindex(df_groups.index, fill_value=[])\n\n# 2) Build vocabulary of GO terms from training\ngo_terms = sorted(df_terms[\"term\"].unique())\nterm_to_idx = {t:i for i,t in enumerate(go_terms)}\n\n# 3) Build Y matrix (multi-hot)\nY = np.zeros((len(df_groups), len(go_terms)), dtype=np.float32)\nfor r, entry in enumerate(df_groups.index):\n    for t in terms_by_id.loc[entry]:\n        Y[r, term_to_idx[t]] = 1.0\n\n# 4) Feature matrix X\nfeature_cols = [c for c in df_groups.columns if c.endswith(\"_frac\") or c.endswith(\"_count\") or c == \"len\"]\nX = df_groups[feature_cols].astype(np.float32).values\n\nX.shape, Y.shape, len(go_terms)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:15:49.814012Z","iopub.execute_input":"2025-10-25T17:15:49.814192Z","iopub.status.idle":"2025-10-25T17:15:53.303445Z","shell.execute_reply.started":"2025-10-25T17:15:49.814178Z","shell.execute_reply":"2025-10-25T17:15:53.302780Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.model_selection import train_test_split\nfrom sklearn.preprocessing import StandardScaler\n\n# Standardize features (helps NN training)\nscaler = StandardScaler(with_mean=True, with_std=True)\nX_scaled = scaler.fit_transform(X)\n\nX_tr, X_va, Y_tr, Y_va = train_test_split(\n    X_scaled, Y, test_size=0.15, random_state=42, stratify=(Y.sum(axis=1) > 0)\n)\n\nX_tr.shape, X_va.shape, Y_tr.shape, Y_va.shape\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:15:53.304198Z","iopub.execute_input":"2025-10-25T17:15:53.304433Z","iopub.status.idle":"2025-10-25T17:16:00.929676Z","shell.execute_reply.started":"2025-10-25T17:15:53.304407Z","shell.execute_reply":"2025-10-25T17:16:00.929097Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"'''\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\n\nn_in  = X_tr.shape[1]\nn_out = Y_tr.shape[1]\n\ndef make_model(n_in, n_out, hidden=256, dropout=0.25):\n    inp = keras.Input(shape=(n_in,))\n    x = layers.BatchNormalization()(inp)\n    x = layers.Dense(hidden, activation=\"relu\")(x)\n    x = layers.Dropout(dropout)(x)\n    x = layers.Dense(hidden//2, activation=\"relu\")(x)\n    x = layers.Dropout(dropout)(x)\n    out = layers.Dense(n_out, activation=\"sigmoid\")(x)  # one prob per GO term\n    model = keras.Model(inp, out)\n    model.compile(\n        optimizer=keras.optimizers.Adam(learning_rate=1e-3),\n        loss=\"binary_crossentropy\",\n        metrics=[keras.metrics.AUC(curve=\"PR\", multi_label=True, num_labels=n_out, name=\"PR-AUC\")]\n    )\n    return model\n\nmodel = make_model(n_in, n_out)\ncb = [\n    keras.callbacks.EarlyStopping(monitor=\"val_PR-AUC\", mode=\"max\", patience=5, restore_best_weights=True),\n    keras.callbacks.ReduceLROnPlateau(monitor=\"val_PR-AUC\", mode=\"max\", factor=0.5, patience=2, min_lr=1e-5),\n]\nhist = model.fit(\n    X_tr, Y_tr,\n    validation_data=(X_va, Y_va),\n    epochs=50,\n    batch_size=512,\n    callbacks=cb,\n    verbose=2\n)\n\nfrom sklearn.metrics import f1_score\n\nY_va_pred = (model.predict(X_va, batch_size=1024) > 0.2).astype(int)  # threshold can be tuned\nmicro_f1 = f1_score(Y_va.ravel(), Y_va_pred.ravel(), average=\"binary\")\nprint(\"Validation micro-F1 (th=0.2):\", round(micro_f1, 4))\n'''","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:16:00.930332Z","iopub.execute_input":"2025-10-25T17:16:00.930592Z","iopub.status.idle":"2025-10-25T17:16:00.935961Z","shell.execute_reply.started":"2025-10-25T17:16:00.930574Z","shell.execute_reply":"2025-10-25T17:16:00.935182Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# safe evaluation helpers","metadata":{}},{"cell_type":"code","source":"def predict_in_batches(model, X, batch_size=1024):\n    \"\"\"Predict in batches to reduce memory footprint.\"\"\"\n    n = X.shape[0]\n    outs = []\n    for start in range(0, n, batch_size):\n        end = min(start + batch_size, n)\n        outs.append(model.predict(X[start:end], batch_size=batch_size, verbose=0))\n    return np.vstack(outs)\n\ndef best_threshold_by_f1(y_true, y_prob, thresholds=None, batch_size=1024):\n    \"\"\"\n    Find the probability threshold in [0,1] that maximizes micro-F1.\n    y_true: (N, C) binary\n    y_prob: (N, C) float in [0,1]\n    thresholds: list or np.array of thresholds to scan\n    \"\"\"\n    if thresholds is None:\n        thresholds = np.concatenate([\n            np.linspace(0.02, 0.5, 25),\n            np.linspace(0.5, 0.9, 9),\n            np.linspace(0.9, 0.99, 10)\n        ])\n    best_f1, best_t = -1.0, 0.5\n    y_true_flat = y_true.ravel()\n    for t in thresholds:\n        y_pred = (y_prob >= t).astype(np.uint8)\n        f1 = f1_score(y_true_flat, y_pred.ravel(), average=\"binary\", zero_division=0)\n        if f1 > best_f1:\n            best_f1, best_t = f1, float(t)\n    return best_t, best_f1\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:16:00.936554Z","iopub.execute_input":"2025-10-25T17:16:00.936718Z","iopub.status.idle":"2025-10-25T17:16:00.951213Z","shell.execute_reply.started":"2025-10-25T17:16:00.936706Z","shell.execute_reply":"2025-10-25T17:16:00.950381Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Objective for Optuna","metadata":{}},{"cell_type":"code","source":"def build_model(n_in, n_out, params):\n    \"\"\"\n    Build a Keras multilabel model using sampled hyperparameters.\n    params keys:\n      - n_layers: 1..3\n      - hidden: 128..1024\n      - dropout: 0.0..0.6\n      - l2: 0..1e-3\n      - lr: 1e-4..5e-3\n      - bn_first: True/False\n      - activation: 'relu' or 'gelu' (via tf.nn.gelu)\n    \"\"\"\n    inp = keras.Input(shape=(n_in,))\n    x = inp\n    if params[\"bn_first\"]:\n        x = layers.BatchNormalization()(x)\n\n    for li in range(params[\"n_layers\"]):\n        x = layers.Dense(\n            params[\"hidden\"] if li == 0 else max(params[\"hidden\"] // 2, 64),\n            activation=None,\n            kernel_regularizer=regularizers.l2(params[\"l2\"]) if params[\"l2\"] > 0 else None,\n        )(x)\n        x = layers.Activation(tf.nn.gelu if params[\"activation\"] == \"gelu\" else \"relu\")(x)\n        if params[\"dropout\"] > 0:\n            x = layers.Dropout(params[\"dropout\"])(x)\n\n    out = layers.Dense(n_out, activation=\"sigmoid\")(x)\n    model = keras.Model(inp, out)\n\n    opt = keras.optimizers.Adam(learning_rate=params[\"lr\"])\n    # PR-AUC multi_label=True works in TF 2.5+; specify num_labels for clarity\n    pr_auc = keras.metrics.AUC(curve=\"PR\", multi_label=True, num_labels=n_out, name=\"PR-AUC\")\n    model.compile(optimizer=opt, loss=\"binary_crossentropy\", metrics=[pr_auc])\n    return model\n\n\ndef objective(trial: optuna.Trial):\n    # Use globals prepared earlier\n    global X_tr, Y_tr, X_va, Y_va\n\n    n_in = X_tr.shape[1]\n    n_out = Y_tr.shape[1]\n\n    params = {\n        \"n_layers\":  trial.suggest_int(\"n_layers\", 1, 3),\n        \"hidden\": trial.suggest_categorical(\"hidden\", [128, 256, 384]),\n        \"dropout\":   trial.suggest_float(\"dropout\", 0.0, 0.6),\n        \"l2\":        trial.suggest_float(\"l2\", 1e-8, 1e-3, log=True),  # << fix\n        \"lr\":        trial.suggest_float(\"lr\", 1e-4, 5e-3, log=True),\n        \"bn_first\":  trial.suggest_categorical(\"bn_first\", [True, False]),\n        \"activation\":trial.suggest_categorical(\"activation\", [\"relu\", \"gelu\"]),\n        \"batch_size\": trial.suggest_categorical(\"batch_size\", [128, 256]),\n        \"epochs\":    trial.suggest_int(\"epochs\", 10, 60),\n        \"patience\":  trial.suggest_int(\"patience\", 3, 8),\n    }\n\n\n    model = build_model(n_in, n_out, params)\n\n    callbacks = [\n        keras.callbacks.EarlyStopping(\n            monitor=\"val_PR-AUC\",\n            mode=\"max\",\n            patience=params[\"patience\"],\n            restore_best_weights=True\n        ),\n        keras.callbacks.ReduceLROnPlateau(\n            monitor=\"val_PR-AUC\",\n            mode=\"max\",\n            factor=0.5,\n            patience=max(1, params[\"patience\"] // 2),\n            min_lr=1e-5\n        ),\n        TFKerasPruningCallback(trial, monitor=\"val_PR-AUC\"),\n    ]\n\n    history = model.fit(\n        X_tr, Y_tr,\n        validation_data=(X_va, Y_va),\n        epochs=params[\"epochs\"],\n        batch_size=params[\"batch_size\"],\n        verbose=0,\n        callbacks=callbacks\n    )\n\n    # Evaluate PR-AUC on validation set (higher is better)\n    # model.evaluate returns [loss, PR-AUC]\n    _, val_pr_auc = model.evaluate(X_va, Y_va, batch_size=1024, verbose=0)\n\n    # cleanup to release GPU memory\n    del model\n    K.clear_session()\n    gc.collect()\n\n    # Report to Optuna\n    return float(val_pr_auc)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:16:00.952212Z","iopub.execute_input":"2025-10-25T17:16:00.952447Z","iopub.status.idle":"2025-10-25T17:16:00.972805Z","shell.execute_reply.started":"2025-10-25T17:16:00.952431Z","shell.execute_reply":"2025-10-25T17:16:00.972078Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Run the study","metadata":{}},{"cell_type":"code","source":"# Optional: make TF use memory growth to reduce OOM risk\ntry:\n    gpus = tf.config.list_physical_devices('GPU')\n    for g in gpus:\n        tf.config.experimental.set_memory_growth(g, True)\nexcept Exception as e:\n    print(f\"GPU memory growth not set: {e}\")\n\nstudy_name = \"keras_cafa6_pr_auc\"\nstorage = None  # e.g., \"sqlite:///optuna.db\" if you want persistence\n\npruner = optuna.pruners.MedianPruner(n_warmup_steps=3)\nsampler = optuna.samplers.TPESampler(seed=42, multivariate=True)\n\nstudy = optuna.create_study(\n    study_name=study_name,\n    direction=\"maximize\",\n    sampler=sampler,\n    pruner=pruner,\n    storage=storage,\n    load_if_exists=False\n)\n\nN_TRIALS = 8  \nstudy.optimize(objective, n_trials=N_TRIALS, show_progress_bar=True)\n\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:16:00.973523Z","iopub.execute_input":"2025-10-25T17:16:00.973762Z","iopub.status.idle":"2025-10-25T17:30:46.779751Z","shell.execute_reply.started":"2025-10-25T17:16:00.973747Z","shell.execute_reply":"2025-10-25T17:30:46.778862Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import json\n\nprint(\"Best value (val PR-AUC):\", study.best_value)\nprint(\"Best trial params:\")\nprint(json.dumps(study.best_trial.params, indent=2))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:30:46.780754Z","iopub.execute_input":"2025-10-25T17:30:46.781052Z","iopub.status.idle":"2025-10-25T17:30:46.788626Z","shell.execute_reply.started":"2025-10-25T17:30:46.781027Z","shell.execute_reply":"2025-10-25T17:30:46.787888Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Retrain best model, find best threshold on validation","metadata":{}},{"cell_type":"code","source":"best_params = study.best_trial.params\nn_in, n_out = X_tr.shape[1], Y_tr.shape[1]\n\nbest_model = build_model(n_in, n_out, best_params)\n\ncallbacks = [\n    keras.callbacks.EarlyStopping(\n        monitor=\"val_PR-AUC\", mode=\"max\",\n        patience=best_params.get(\"patience\", 5),\n        restore_best_weights=True\n    ),\n    keras.callbacks.ReduceLROnPlateau(\n        monitor=\"val_PR-AUC\", mode=\"max\",\n        factor=0.5,\n        patience=max(1, best_params.get(\"patience\", 5) // 2),\n        min_lr=1e-5\n    ),\n]\n\nhistory = best_model.fit(\n    X_tr, Y_tr,\n    validation_data=(X_va, Y_va),\n    epochs=best_params.get(\"epochs\", 30),\n    batch_size=best_params.get(\"batch_size\", 512),\n    verbose=2,\n    callbacks=callbacks\n)\n\n\n\nmodel = best_model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:30:46.789339Z","iopub.execute_input":"2025-10-25T17:30:46.789562Z","iopub.status.idle":"2025-10-25T17:35:50.402005Z","shell.execute_reply.started":"2025-10-25T17:30:46.789536Z","shell.execute_reply":"2025-10-25T17:35:50.401335Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Build group features for the test set; reuse your helper\ndf_test = build_group_features(TEST_FASTA, GROUPS, include_fractions=True, include_counts=True)\n\n# Transform with the fitted scaler and select the same feature columns\nX_test = df_test[feature_cols].astype(np.float32).values\nX_test_scaled = scaler.transform(X_test)\n\ndf_test.shape\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:35:50.402879Z","iopub.execute_input":"2025-10-25T17:35:50.403211Z","iopub.status.idle":"2025-10-25T17:39:10.893283Z","shell.execute_reply.started":"2025-10-25T17:35:50.403192Z","shell.execute_reply":"2025-10-25T17:39:10.892557Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = best_model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:39:10.894146Z","iopub.execute_input":"2025-10-25T17:39:10.894442Z","iopub.status.idle":"2025-10-25T17:39:10.897901Z","shell.execute_reply.started":"2025-10-25T17:39:10.894424Z","shell.execute_reply":"2025-10-25T17:39:10.897329Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Save model, scaler, and metadata","metadata":{}},{"cell_type":"code","source":"# --- SAVE MODEL, SCALER, AND TEST DATA --------------------------------------\n\nfrom pathlib import Path\nimport joblib\nimport json\nimport numpy as np\nimport pandas as pd\n\nSAVE_DIR = Path(\"outputs/checkpoints\")\nSAVE_DIR.mkdir(parents=True, exist_ok=True)\n\n# 1) Save Keras model\nmodel_path = SAVE_DIR / \"cafa6_model.h5\"\nmodel.save(model_path)\nprint(f\"Saved model: {model_path.resolve()}\")\n\n# 2) Save fitted StandardScaler\nscaler_path = SAVE_DIR / \"scaler.joblib\"\njoblib.dump(scaler, scaler_path)\nprint(f\"Saved scaler: {scaler_path.resolve()}\")\n\n# 3) Save metadata (feature names, GO terms, and amino acid groups)\nmeta = {\n    \"feature_cols\": feature_cols,\n    \"go_terms\": go_terms,\n    \"term_to_idx\": term_to_idx,\n    \"idx_to_term\": go_terms,\n    \"groups\": {k: sorted(list(v)) for k, v in GROUPS.items()},\n}\nmeta_path = SAVE_DIR / \"metadata.json\"\nwith open(meta_path, \"w\") as f:\n    json.dump(meta, f, indent=2)\nprint(f\"Saved metadata: {meta_path.resolve()}\")\n\n# 4) Save test feature matrices and dataframe\nnp.save(SAVE_DIR / \"X_test.npy\", X_test)\nnp.save(SAVE_DIR / \"X_test_scaled.npy\", X_test_scaled)\ndf_test.to_csv(SAVE_DIR / \"df_test.tsv\", sep=\"\\t\")\nprint(\"Saved X_test, X_test_scaled, and df_test.\")\n\nprint(\"\\nAll training artifacts and test features have been saved successfully.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:39:10.902991Z","iopub.execute_input":"2025-10-25T17:39:10.903222Z","iopub.status.idle":"2025-10-25T17:39:29.985673Z","shell.execute_reply.started":"2025-10-25T17:39:10.903206Z","shell.execute_reply":"2025-10-25T17:39:29.984732Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Load model, scaler, metadata, and test data","metadata":{}},{"cell_type":"code","source":"# --- LOAD MODEL, SCALER, METADATA, AND TEST DATA ----------------------------\n\nfrom pathlib import Path\nimport joblib\nimport json\nimport numpy as np\nimport pandas as pd\nimport tensorflow as tf\nfrom tensorflow import keras\n\nLOAD_DIR = Path(\"outputs/checkpoints\")\n\n# 1) Load Keras model\nmodel = keras.models.load_model(LOAD_DIR / \"cafa6_model.h5\")\nprint(\"Model loaded.\")\n\n# 2) Load fitted StandardScaler\nscaler = joblib.load(LOAD_DIR / \"scaler.joblib\")\nprint(\"Scaler loaded.\")\n\n# 3) Load metadata\nwith open(LOAD_DIR / \"metadata.json\") as f:\n    meta = json.load(f)\n\nfeature_cols = meta[\"feature_cols\"]\ngo_terms = meta[\"go_terms\"]\nterm_to_idx = meta[\"term_to_idx\"]\nGROUPS = {k: set(v) for k, v in meta[\"groups\"].items()}\n\nprint(\"Metadata loaded.\")\n\n# 4) Load test matrices and dataframe\nX_test = np.load(LOAD_DIR / \"X_test.npy\")\nX_test_scaled = np.load(LOAD_DIR / \"X_test_scaled.npy\")\ndf_test = pd.read_csv(LOAD_DIR / \"df_test.tsv\", sep=\"\\t\", index_col=0)\n\nprint(\"Test data loaded.\")\nprint(f\"X_test: {X_test.shape}, X_test_scaled: {X_test_scaled.shape}, df_test: {df_test.shape}\")\nprint(\"All components successfully reloaded. Ready for inference.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:39:29.986752Z","iopub.execute_input":"2025-10-25T17:39:29.987150Z","iopub.status.idle":"2025-10-25T17:39:33.100946Z","shell.execute_reply.started":"2025-10-25T17:39:29.987124Z","shell.execute_reply":"2025-10-25T17:39:33.099648Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- PREDICT ON TEST IN STREAMING, POST-PROCESS, WRITE TSV ON THE FLY --------\nfrom pathlib import Path\nimport numpy as np\nimport pandas as pd\nimport tensorflow as tf\nfrom tensorflow import keras\n\n# (Optional) enable GPU memory growth to avoid pre-allocating all VRAM\ntry:\n    gpus = tf.config.list_physical_devices('GPU')\n    for g in gpus:\n        tf.config.experimental.set_memory_growth(g, True)\nexcept Exception as e:\n    print(f\"GPU memory growth not set: {e}\")\n\n# 0) Required globals and paths\nOUT_DIR = Path(\"outputs\")\nOUT_DIR.mkdir(parents=True, exist_ok=True)\nSUB_PATH =  Path(\"submission.tsv\")        # no header, tab-separated\n\nrequired = [\"model\", \"df_test\", \"X_test\", \"X_test_scaled\", \"feature_cols\", \"go_terms\"]\nmissing = [k for k in required if k not in globals()]\nif missing:\n    raise RuntimeError(f\"Missing required variables: {missing}. Ensure training/feature cells ran.\")\n\n# 1) Sanity checks and shape guards\nassert all(c in df_test.columns for c in feature_cols), \"Feature columns mismatch between train and test.\"\nX_test = np.asarray(X_test, dtype=np.float32, order=\"C\")\nX_test_scaled = np.asarray(X_test_scaled, dtype=np.float32, order=\"C\")\nassert X_test.shape == X_test_scaled.shape, \"X_test and X_test_scaled must have identical shapes.\"\nn_test, n_feat = X_test_scaled.shape\ngo_terms = list(go_terms)              # ensure list\nidx_to_term = np.asarray(go_terms, dtype=str)\nn_terms = len(idx_to_term)\ntest_ids = df_test.index.astype(str).to_numpy()\nassert len(test_ids) == n_test, \"df_test index length must match X_test rows.\"\n\n# 2) CAFA post-processing parameters\nTHRESH = 0.1                  # drop probabilities <= threshold to avoid zeros\nTOP_K = None                  # set e.g. 500 to pre-cap per target; None keeps all above THRESH\nMAX_PAIRS_PER_TARGET = 1500   # global cap per target across MF/BP/CC\nBATCH = 128                   # reduce if you still hit OOM\n\n# 3) Streamed inference and direct writing to disk\nwith SUB_PATH.open(\"w\", encoding=\"utf-8\") as fout:\n    for start in range(0, n_test, BATCH):\n        end = min(start + BATCH, n_test)\n        xb = X_test_scaled[start:end]                            # view without extra copy\n        # Predict a small batch to limit RAM/VRAM usage\n        preds = model.predict(xb, batch_size=BATCH, verbose=0)  # shape (batch, n_terms)\n        if preds.dtype != np.float32:\n            preds = preds.astype(np.float32, copy=False)\n        # Process each row and emit lines directly\n        for i in range(end - start):\n            pid = test_ids[start + i]\n            probs = np.nan_to_num(preds[i], nan=0.0)          # guard against NaNs\n            mask = probs > THRESH\n            if not np.any(mask):\n                # If a target is not listed, evaluators assume zero for all terms\n                continue\n            sel_terms = idx_to_term[mask]\n            sel_probs = probs[mask]\n            order = np.argsort(-sel_probs)                    # descending by prob\n            if TOP_K is not None and len(order) > TOP_K:\n                order = order[:TOP_K]\n            # Enforce CAFA global cap per target\n            order = order[:min(len(order), MAX_PAIRS_PER_TARGET)]\n            # Write lines: target \\t GO \\t score (3 significant figures, (0,1])\n            for j in order:\n                p = float(sel_probs[j])\n                if p <= 0.0:\n                    continue\n                if p > 1.0:\n                    p = 1.0\n                fout.write(f\"{pid}\\t{sel_terms[j]}\\t{p:.3g}\\n\")\n\nprint(f\"Submission written to: {SUB_PATH.resolve()}\")\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T17:49:11.401573Z","iopub.execute_input":"2025-10-25T17:49:11.402385Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.summary()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T18:11:36.486089Z","iopub.execute_input":"2025-10-25T18:11:36.486660Z","iopub.status.idle":"2025-10-25T18:11:36.503999Z","shell.execute_reply.started":"2025-10-25T18:11:36.486636Z","shell.execute_reply":"2025-10-25T18:11:36.502963Z"}},"outputs":[],"execution_count":null}]}