Add causal_selection/meta_learner/trainer.py

causal_selection/meta_learner/trainer.py

@@ -0,0 +1,356 @@
+"""
+Meta-learner: trains models to predict algorithm performance from dataset meta-features.
+Supports multi-output regression (predict SHD per algorithm) and ranking evaluation.
+"""
+import os
+import numpy as np
+import pandas as pd
+import json
+import logging
+from collections import defaultdict
+
+from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
+from sklearn.multioutput import MultiOutputRegressor
+from sklearn.preprocessing import StandardScaler
+from sklearn.model_selection import LeaveOneGroupOut, cross_val_predict
+from sklearn.metrics import mean_squared_error, mean_absolute_error
+import joblib
+
+from causal_selection.features.extractor import FEATURE_NAMES
+from causal_selection.discovery.algorithms import ALGORITHM_POOL
+
+logger = logging.getLogger(__name__)
+
+ALGO_NAMES = list(ALGORITHM_POOL.keys())
+RESULTS_DIR = '/app/causal_selection/data/results'
+MODEL_DIR = '/app/causal_selection/models'
+
+
+def load_meta_dataset(results_dir=RESULTS_DIR):
+    """Load meta-dataset from CSV files."""
+    X = pd.read_csv(os.path.join(results_dir, 'meta_features.csv'))
+    Y_shd = pd.read_csv(os.path.join(results_dir, 'shd_matrix.csv'))
+    Y_nshd = pd.read_csv(os.path.join(results_dir, 'normalized_shd_matrix.csv'))
+    configs = pd.read_csv(os.path.join(results_dir, 'configs.csv'))
+    return X, Y_shd, Y_nshd, configs
+
+
+def train_meta_learner(X, Y, model_type='rf', **model_kwargs):
+    """Train a multi-output regression model.
+    
+    Args:
+        X: feature matrix (n_tasks, n_features)
+        Y: target matrix (n_tasks, n_algorithms) - SHD values
+        model_type: 'rf' or 'gbm'
+    
+    Returns:
+        trained model, scaler
+    """
+    scaler = StandardScaler()
+    X_scaled = scaler.fit_transform(X)
+    
+    if model_type == 'rf':
+        base = RandomForestRegressor(
+            n_estimators=model_kwargs.get('n_estimators', 200),
+            max_depth=model_kwargs.get('max_depth', None),
+            min_samples_leaf=model_kwargs.get('min_samples_leaf', 2),
+            random_state=42,
+            n_jobs=-1,
+        )
+    elif model_type == 'gbm':
+        base = GradientBoostingRegressor(
+            n_estimators=model_kwargs.get('n_estimators', 200),
+            max_depth=model_kwargs.get('max_depth', 5),
+            learning_rate=model_kwargs.get('learning_rate', 0.1),
+            min_samples_leaf=model_kwargs.get('min_samples_leaf', 3),
+            random_state=42,
+        )
+    else:
+        raise ValueError(f"Unknown model type: {model_type}")
+    
+    model = MultiOutputRegressor(base)
+    model.fit(X_scaled, Y)
+    
+    return model, scaler
+
+
+def predict_top_k(model, scaler, X_new, k=3):
+    """Predict top-k algorithms for new dataset(s).
+    
+    Args:
+        model: trained multi-output model
+        scaler: fitted StandardScaler
+        X_new: feature matrix (n_new, n_features)
+        k: number of top algorithms to return
+    
+    Returns:
+        top_k_indices: (n_new, k) array of algorithm indices (sorted by predicted SHD ascending)
+        predicted_shd: (n_new, n_algorithms) full predicted SHD matrix
+    """
+    X_scaled = scaler.transform(X_new)
+    predicted = model.predict(X_scaled)
+    
+    if predicted.ndim == 1:
+        predicted = predicted.reshape(1, -1)
+    
+    top_k_indices = np.argsort(predicted, axis=1)[:, :k]
+    return top_k_indices, predicted
+
+
+def evaluate_lono_cv(X, Y, configs, model_type='rf', k=3, **model_kwargs):
+    """Leave-One-Network-Out Cross-Validation.
+    
+    For each network, train on all other networks, test on that network.
+    This tests generalization to truly unseen graph structures.
+    
+    Returns:
+        results: dict with metrics per network and overall
+    """
+    networks = configs['network'].values
+    unique_networks = sorted(configs['network'].unique())
+    
+    results = {
+        'per_network': {},
+        'all_predictions': [],
+        'all_true': [],
+        'all_configs': [],
+    }
+    
+    scaler = StandardScaler()
+    
+    for test_net in unique_networks:
+        test_mask = networks == test_net
+        train_mask = ~test_mask
+        
+        if train_mask.sum() < 3:
+            logger.warning(f"Skipping {test_net}: only {train_mask.sum()} training samples")
+            continue
+        
+        X_train = X.values[train_mask]
+        Y_train = Y.values[train_mask]
+        X_test = X.values[test_mask]
+        Y_test = Y.values[test_mask]
+        
+        # Scale
+        scaler.fit(X_train)
+        X_train_s = scaler.transform(X_train)
+        X_test_s = scaler.transform(X_test)
+        
+        # Train
+        if model_type == 'rf':
+            base = RandomForestRegressor(
+                n_estimators=model_kwargs.get('n_estimators', 200),
+                max_depth=model_kwargs.get('max_depth', None),
+                min_samples_leaf=model_kwargs.get('min_samples_leaf', 2),
+                random_state=42, n_jobs=-1,
+            )
+        else:
+            base = GradientBoostingRegressor(
+                n_estimators=model_kwargs.get('n_estimators', 200),
+                max_depth=model_kwargs.get('max_depth', 5),
+                learning_rate=model_kwargs.get('learning_rate', 0.1),
+                min_samples_leaf=model_kwargs.get('min_samples_leaf', 3),
+                random_state=42,
+            )
+        
+        model = MultiOutputRegressor(base)
+        model.fit(X_train_s, Y_train)
+        
+        # Predict
+        Y_pred = model.predict(X_test_s)
+        
+        # Evaluate
+        net_metrics = _compute_ranking_metrics(Y_pred, Y_test, k=k)
+        net_metrics['n_test'] = int(test_mask.sum())
+        net_metrics['n_train'] = int(train_mask.sum())
+        results['per_network'][test_net] = net_metrics
+        
+        results['all_predictions'].extend(Y_pred.tolist())
+        results['all_true'].extend(Y_test.tolist())
+        results['all_configs'].extend(
+            configs[test_mask][['network', 'n_samples', 'seed']].to_dict('records')
+        )
+        
+        logger.info(f"  {test_net:15s}: top{k}_hit={net_metrics['top_k_hit_rate']:.3f} "
+                    f"regret={net_metrics['mean_regret']:.2f} "
+                    f"ndcg={net_metrics['ndcg_at_k']:.3f}")
+    
+    # Overall metrics
+    all_pred = np.array(results['all_predictions'])
+    all_true = np.array(results['all_true'])
+    overall = _compute_ranking_metrics(all_pred, all_true, k=k)
+    results['overall'] = overall
+    
+    return results
+
+
+def _compute_ranking_metrics(Y_pred, Y_true, k=3):
+    """Compute ranking metrics for algorithm selection.
+    
+    Args:
+        Y_pred: (n, n_algos) predicted SHD values
+        Y_true: (n, n_algos) true SHD values
+        k: top-k to consider
+    """
+    n = Y_pred.shape[0]
+    
+    top_k_hits = 0
+    regrets = []
+    ndcgs = []
+    
+    for i in range(n):
+        true_ranking = np.argsort(Y_true[i])  # best algo first
+        pred_ranking = np.argsort(Y_pred[i])  # predicted best first
+        
+        true_best = true_ranking[0]
+        pred_top_k = pred_ranking[:k]
+        
+        # Top-k hit rate: is the true best in predicted top-k?
+        if true_best in pred_top_k:
+            top_k_hits += 1
+        
+        # SHD regret: SHD of best in predicted top-k minus oracle best SHD
+        oracle_shd = Y_true[i, true_best]
+        selected_shds = [Y_true[i, j] for j in pred_top_k]
+        best_selected_shd = min(selected_shds)
+        regret = best_selected_shd - oracle_shd
+        regrets.append(regret)
+        
+        # NDCG@k
+        ndcg = _ndcg_at_k(Y_true[i], Y_pred[i], k)
+        ndcgs.append(ndcg)
+    
+    # Also compute: is one of the true top-3 in the predicted top-3?
+    top_k_overlap = 0
+    for i in range(n):
+        true_top_k = set(np.argsort(Y_true[i])[:k])
+        pred_top_k = set(np.argsort(Y_pred[i])[:k])
+        overlap = len(true_top_k & pred_top_k)
+        top_k_overlap += overlap / k
+    
+    return {
+        'top_k_hit_rate': top_k_hits / n,  # true best in predicted top-k
+        'top_k_overlap_rate': top_k_overlap / n,  # avg overlap between true/pred top-k
+        'mean_regret': np.mean(regrets),
+        'median_regret': np.median(regrets),
+        'max_regret': np.max(regrets),
+        'ndcg_at_k': np.mean(ndcgs),
+        'mean_pred_mse': mean_squared_error(Y_true, Y_pred),
+        'mean_pred_mae': mean_absolute_error(Y_true, Y_pred),
+    }
+
+
+def _ndcg_at_k(true_scores, pred_scores, k):
+    """Normalized Discounted Cumulative Gain at k.
+    
+    For algorithm selection: lower SHD = better, so we negate scores for ranking.
+    """
+    # Convert SHD to relevance: rel = max_shd - shd (higher = better)
+    max_shd = max(true_scores.max(), 1)
+    relevance = max_shd - true_scores
+    
+    # Predicted ranking
+    pred_order = np.argsort(pred_scores)[:k]
+    
+    # DCG
+    dcg = 0
+    for rank, idx in enumerate(pred_order):
+        dcg += relevance[idx] / np.log2(rank + 2)
+    
+    # Ideal DCG
+    ideal_order = np.argsort(-relevance)[:k]
+    idcg = 0
+    for rank, idx in enumerate(ideal_order):
+        idcg += relevance[idx] / np.log2(rank + 2)
+    
+    return dcg / idcg if idcg > 0 else 0
+
+
+def get_feature_importance(model, feature_names=FEATURE_NAMES, algo_names=ALGO_NAMES):
+    """Extract feature importance from trained model."""
+    importances = {}
+    for i, (algo, estimator) in enumerate(zip(algo_names, model.estimators_)):
+        if hasattr(estimator, 'feature_importances_'):
+            importances[algo] = dict(zip(feature_names, estimator.feature_importances_))
+    
+    # Average importance across algorithms
+    avg_importance = defaultdict(float)
+    for algo, imp in importances.items():
+        for feat, val in imp.items():
+            avg_importance[feat] += val / len(importances)
+    
+    return dict(avg_importance), importances
+
+
+def save_model(model, scaler, model_dir=MODEL_DIR):
+    """Save trained model and scaler."""
+    os.makedirs(model_dir, exist_ok=True)
+    joblib.dump(model, os.path.join(model_dir, 'meta_learner.pkl'))
+    joblib.dump(scaler, os.path.join(model_dir, 'scaler.pkl'))
+    logger.info(f"Model saved to {model_dir}")
+
+
+def load_model(model_dir=MODEL_DIR):
+    """Load trained model and scaler."""
+    model = joblib.load(os.path.join(model_dir, 'meta_learner.pkl'))
+    scaler = joblib.load(os.path.join(model_dir, 'scaler.pkl'))
+    return model, scaler
+
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+    
+    # Load meta-dataset
+    X, Y_shd, Y_nshd, configs = load_meta_dataset()
+    
+    print(f"Meta-dataset: X={X.shape}, Y_shd={Y_shd.shape}")
+    print(f"Networks: {configs['network'].unique()}")
+    print(f"Configs per network:")
+    print(configs['network'].value_counts().to_string())
+    
+    # Evaluate with LONO-CV
+    print("\n" + "=" * 80)
+    print("LEAVE-ONE-NETWORK-OUT CV (RandomForest)")
+    print("=" * 80)
+    
+    results_rf = evaluate_lono_cv(X, Y_nshd, configs, model_type='rf', k=3)
+    
+    print(f"\nOverall Results (RF):")
+    for k, v in results_rf['overall'].items():
+        print(f"  {k:25s}: {v:.4f}")
+    
+    print("\n" + "=" * 80)
+    print("LEAVE-ONE-NETWORK-OUT CV (GradientBoosting)")
+    print("=" * 80)
+    
+    results_gbm = evaluate_lono_cv(X, Y_nshd, configs, model_type='gbm', k=3)
+    
+    print(f"\nOverall Results (GBM):")
+    for k, v in results_gbm['overall'].items():
+        print(f"  {k:25s}: {v:.4f}")
+    
+    # Train final model on all data
+    print("\n" + "=" * 80)
+    print("TRAINING FINAL MODEL")
+    print("=" * 80)
+    
+    best_type = 'rf' if results_rf['overall']['top_k_hit_rate'] >= results_gbm['overall']['top_k_hit_rate'] else 'gbm'
+    print(f"Selected model type: {best_type}")
+    
+    model, scaler = train_meta_learner(X, Y_nshd, model_type=best_type)
+    save_model(model, scaler)
+    
+    # Feature importance
+    avg_imp, per_algo_imp = get_feature_importance(model)
+    print("\nTop 10 Most Important Features:")
+    for feat, imp in sorted(avg_imp.items(), key=lambda x: -x[1])[:10]:
+        print(f"  {feat:30s}: {imp:.4f}")
+    
+    # Save all evaluation results
+    with open(os.path.join(RESULTS_DIR, 'evaluation_results.json'), 'w') as f:
+        json.dump({
+            'rf': {k: v for k, v in results_rf['overall'].items()},
+            'gbm': {k: v for k, v in results_gbm['overall'].items()},
+            'feature_importance': avg_imp,
+            'selected_model': best_type,
+        }, f, indent=2)