""" ================================================================================ CAFA-6 V7 - TUNED LOGISTIC REGRESSION + GO PROPAGATION Improvements over V5: - More PCA components (200 vs 150) - Stronger regularization (C=0.3) - More iterations (1000) - Lower confidence threshold (0.005) for better recall Target: 0.28+ (current best 0.245) ================================================================================ """ # Install biopython import subprocess subprocess.run(['pip', 'install', 'biopython', '-q'], check=True) import os import gc import numpy as np import pandas as pd from collections import defaultdict from tqdm.auto import tqdm import warnings warnings.filterwarnings('ignore') from sklearn.decomposition import PCA from sklearn.preprocessing import MultiLabelBinarizer from sklearn.linear_model import LogisticRegression from sklearn.multiclass import OneVsRestClassifier # ============================================================================ # CONFIG - TUNED PARAMETERS # ============================================================================ CONFIG = { 'PCA_COMPONENTS': 200, # More components = more information (was 150) 'REGULARIZATION': 0.3, # Stronger regularization (was 0.5) 'MAX_ITER': 1000, # More iterations (was 800) 'MIN_CONFIDENCE': 0.005, # Lower threshold for more recall (was 0.01) 'MAX_PREDS_PER_PROTEIN': 1500, 'BASE_PATH': '/kaggle/input/cafa-6-protein-function-prediction', 'EMB_PATH': '/kaggle/input/esm2-t6-8m-ur50d-cafa6', } print(f"CAFA-6 V7 - Tuned LogReg + GO Propagation") print(f"Config: PCA={CONFIG['PCA_COMPONENTS']}, C={CONFIG['REGULARIZATION']}, iter={CONFIG['MAX_ITER']}") # ============================================================================ # GO ONTOLOGY # ============================================================================ def parse_go_obo(obo_path): """Parse GO ontology for parent relationships.""" go_parents = defaultdict(set) current_id = None with open(obo_path, 'r') as f: for line in f: line = line.strip() if line == '[Term]': current_id = None elif line.startswith('id: GO:'): current_id = line[4:] elif line.startswith('is_a:') and current_id: parent = line.split()[1] go_parents[current_id].add(parent) elif line.startswith('relationship: part_of') and current_id: parent = line.split()[2] go_parents[current_id].add(parent) return go_parents def get_all_ancestors(go_term, go_parents, cache=None): """Get all ancestors of a GO term.""" if cache is None: cache = {} if go_term in cache: return cache[go_term] ancestors = set() for parent in go_parents.get(go_term, set()): ancestors.add(parent) ancestors.update(get_all_ancestors(parent, go_parents, cache)) cache[go_term] = ancestors return ancestors def propagate_predictions(predictions_df, go_parents): """Apply GO propagation - if child has score, parent gets max(parent, child).""" print("Applying GO propagation...") # Build ancestors cache cache = {} all_terms = predictions_df['term'].unique() for term in tqdm(all_terms, desc="Building ancestors"): get_all_ancestors(term, go_parents, cache) # Propagate using groupby for speed protein_preds = defaultdict(dict) # Group by protein for efficiency grouped = predictions_df.groupby('protein') for protein, group in tqdm(grouped, desc="Propagating"): term_scores = {} for _, row in group.iterrows(): term = row['term'] score = row['score'] # Set score for term if term not in term_scores or score > term_scores[term]: term_scores[term] = score # Propagate to ancestors for ancestor in cache.get(term, set()): if ancestor not in term_scores or score > term_scores[ancestor]: term_scores[ancestor] = score protein_preds[protein] = term_scores # Convert to DataFrame rows = [] for protein, term_scores in protein_preds.items(): for term, score in term_scores.items(): if score >= CONFIG['MIN_CONFIDENCE']: rows.append({'protein': protein, 'term': term, 'score': score}) result = pd.DataFrame(rows) print(f"After propagation: {len(result):,} predictions") return result # ============================================================================ # MAIN # ============================================================================ def main(): # 1. Load GO ontology print("\n[1/6] Loading GO ontology...") go_parents = parse_go_obo(f"{CONFIG['BASE_PATH']}/Train/go-basic.obo") print(f"GO terms: {len(go_parents):,}") # 2. Load sequences and IDs print("\n[2/6] Loading sequence IDs...") from Bio import SeqIO train_seqs = {rec.id: str(rec.seq) for rec in SeqIO.parse( f"{CONFIG['BASE_PATH']}/Train/train_sequences.fasta", 'fasta')} test_seqs = {rec.id: str(rec.seq) for rec in SeqIO.parse( f"{CONFIG['BASE_PATH']}/Test/testsuperset.fasta", 'fasta')} train_ids = [i.split('|')[1] for i in train_seqs.keys()] test_ids = list(test_seqs.keys()) print(f"Train: {len(train_ids)}, Test: {len(test_ids)}") # 3. Load embeddings and apply PCA print("\n[3/6] Loading embeddings and applying PCA...") X_train = np.load(f"{CONFIG['EMB_PATH']}/protein_embeddings_train.npy") X_test = np.load(f"{CONFIG['EMB_PATH']}/protein_embeddings_test.npy") print(f"Train shape: {X_train.shape}, Test shape: {X_test.shape}") pca = PCA(n_components=CONFIG['PCA_COMPONENTS'], random_state=42) X_train_pca = pca.fit_transform(X_train) X_test_pca = pca.transform(X_test) print(f"After PCA: {X_train_pca.shape}") print(f"Explained variance: {pca.explained_variance_ratio_.sum():.3f}") del X_train, X_test gc.collect() # 4. Train models print("\n[4/6] Training models...") train_terms_df = pd.read_csv( f"{CONFIG['BASE_PATH']}/Train/train_terms.tsv", sep="\t") mlb_dict = {} models = {} for aspect in tqdm(['P', 'C', 'F'], desc="Training"): ont_terms = train_terms_df[train_terms_df['aspect'] == aspect] protein_terms = ont_terms.groupby('EntryID')['term'].apply(list).to_dict() labels = [protein_terms.get(eid, []) for eid in train_ids] mlb = MultiLabelBinarizer(sparse_output=True) y_train = mlb.fit_transform(labels) print(f" {aspect}: {y_train.shape[1]} classes") model = OneVsRestClassifier( LogisticRegression( max_iter=CONFIG['MAX_ITER'], solver='lbfgs', C=CONFIG['REGULARIZATION'], random_state=42 ), n_jobs=-1 ) model.fit(X_train_pca, y_train) mlb_dict[aspect] = mlb models[aspect] = model del X_train_pca gc.collect() # 5. Generate predictions print("\n[5/6] Generating predictions...") BATCH_SIZE = 5000 submission_list = [] for i in tqdm(range(0, len(test_ids), BATCH_SIZE), desc="Predicting"): batch_ids = test_ids[i:i+BATCH_SIZE] X_batch = X_test_pca[i:i+BATCH_SIZE] for aspect, model in models.items(): mlb = mlb_dict[aspect] y_pred = model.predict_proba(X_batch) for j, eid in enumerate(batch_ids): probs = y_pred[j] indices = np.where(probs > CONFIG['MIN_CONFIDENCE'])[0] for idx in indices: submission_list.append({ 'protein': eid, 'term': mlb.classes_[idx], 'score': probs[idx] }) del X_test_pca gc.collect() predictions_df = pd.DataFrame(submission_list) print(f"Raw predictions: {len(predictions_df):,}") # 6. Apply GO propagation and save print("\n[6/6] Applying GO propagation and saving...") predictions_df = propagate_predictions(predictions_df, go_parents) # Aggregate and clean predictions_df = predictions_df.groupby(['protein', 'term'], as_index=False)['score'].max() predictions_df['score'] = predictions_df['score'].clip(0.001, 0.999).round(3) # Keep top predictions per protein predictions_df = predictions_df.sort_values(['protein', 'score'], ascending=[True, False]) predictions_df = predictions_df.groupby('protein').head(CONFIG['MAX_PREDS_PER_PROTEIN']) predictions_df.to_csv('submission.tsv', sep='\t', header=False, index=False) print(f"\n{'='*60}") print("CAFA-6 V7 COMPLETE!") print(f"{'='*60}") print(f"Total: {len(predictions_df):,} predictions") print(f"Proteins: {predictions_df['protein'].nunique():,}") print(f"Terms: {predictions_df['term'].nunique():,}") print(predictions_df.head(10)) if __name__ == "__main__": main()