""" 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)