scripts/analyze_results.py

#!/usr/bin/env python3
"""Analyze raw results and produce processed CSVs and summary statistics."""
import os
import sys
import json
import glob
import argparse
import yaml
import numpy as np
import pandas as pd
from scipy import stats
from datetime import datetime

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils import ensure_dir


def load_all_results(raw_dir='results/raw'):
    """Load all JSONL result files."""
    records = []
    for fpath in sorted(glob.glob(os.path.join(raw_dir, '*.jsonl'))):
        with open(fpath) as f:
            for line in f:
                line = line.strip()
                if line:
                    try:
                        records.append(json.loads(line))
                    except json.JSONDecodeError:
                        pass
    return records


def process_synthetic(df):
    """Process synthetic results into summary tables."""
    if len(df) == 0:
        return pd.DataFrame()
    
    # Regime summary
    group_cols = ['graph_type', 'prior_strength', 'K']
    available = [c for c in group_cols if c in df.columns]
    if not available:
        return pd.DataFrame()
    
    agg_dict = {}
    for col in ['chi_seed_max', 'chi_seed_sum', 'empirical_decay_mu', 'empirical_decay_r2',
                 'rel_error_R1', 'rel_error_R2', 'rel_error_R3', 'rel_error_R4',
                 'runtime_exact', 'runtime_local_R2', 'runtime_local_R4',
                 'interference_cosine_R2', 'error_warm_start', 'rel_error_warm_start',
                 'error_one_step', 'rel_error_one_step']:
        if col in df.columns:
            agg_dict[col] = ['mean', 'median', 'std', 'count']
    
    if not agg_dict:
        return pd.DataFrame()
    
    summary = df.groupby(available).agg(agg_dict)
    summary.columns = ['_'.join(c) for c in summary.columns]
    summary = summary.reset_index()
    
    return summary


def compute_correlations(df):
    """Compute correlation table between proxies and error metrics.
    
    Computes within-regime correlations (controlling for graph structure)
    and also log-transformed chi correlations.
    """
    rows = []
    
    proxy_cols = ['chi_seed_max', 'chi_seed_sum', 'seed_degree']
    log_proxy_cols = ['log_chi_max', 'log_chi_sum']
    target_cols = ['rel_error_R2', 'rel_error_R4', 'interference_cosine_R2']
    
    # Add log-chi columns
    df_copy = df.copy()
    if 'chi_seed_max' in df_copy.columns:
        df_copy['log_chi_max'] = np.log1p(df_copy['chi_seed_max'].clip(lower=0))
    if 'chi_seed_sum' in df_copy.columns:
        df_copy['log_chi_sum'] = np.log1p(df_copy['chi_seed_sum'].clip(lower=0))
    
    all_proxies = proxy_cols + log_proxy_cols
    
    # Within-regime correlations (most informative)
    if 'regime' in df_copy.columns:
        regime_groups = df_copy.groupby('regime')
    elif 'dataset_name' in df_copy.columns:
        regime_groups = df_copy.groupby('dataset_name')
    else:
        regime_groups = [('all', df_copy)]
    
    for grp_name, grp_df in regime_groups:
        for proxy in all_proxies:
            for target in target_cols:
                if proxy in grp_df.columns and target in grp_df.columns:
                    x = grp_df[proxy].dropna()
                    y = grp_df[target].dropna()
                    common = x.index.intersection(y.index)
                    x, y = x.loc[common], y.loc[common]
                    
                    mask = np.isfinite(x) & np.isfinite(y)
                    x, y = x[mask], y[mask]
                    
                    if len(x) >= 5:
                        try:
                            pr, pp = stats.pearsonr(x, y)
                            sr, sp = stats.spearmanr(x, y)
                        except:
                            continue
                        if np.isnan(pr) or np.isnan(sr):
                            continue
                        rows.append({
                            'dataset_regime': grp_name,
                            'model_family': grp_df['model_family'].iloc[0] if 'model_family' in grp_df.columns else 'unknown',
                            'proxy': proxy,
                            'target': target,
                            'pearson_r': round(pr, 4),
                            'spearman_r': round(sr, 4),
                            'pearson_p': round(pp, 6),
                            'spearman_p': round(sp, 6),
                            'n_deletions': len(x),
                        })
    
    return pd.DataFrame(rows)


def build_method_comparison(df):
    """Build method comparison table."""
    rows = []
    
    if 'dataset_name' in df.columns:
        groups = df.groupby('dataset_name')
    else:
        groups = [('all', df)]
    
    for grp_name, grp_df in groups:
        # Exact
        if 'runtime_exact' in grp_df.columns:
            rows.append({
                'dataset_regime': grp_name,
                'method': 'exact',
                'radius': None,
                'mean_error': 0.0,
                'median_error': 0.0,
                'mean_runtime': grp_df['runtime_exact'].mean(),
                'speedup_vs_exact': 1.0,
            })
            t_exact = grp_df['runtime_exact'].mean()
        else:
            t_exact = 1.0
        
        # Local radii
        for R in [1, 2, 3, 4]:
            err_col = f'rel_error_R{R}'
            rt_col = f'runtime_local_R{R}'
            if err_col in grp_df.columns and rt_col in grp_df.columns:
                mean_rt = grp_df[rt_col].mean()
                rows.append({
                    'dataset_regime': grp_name,
                    'method': 'local',
                    'radius': R,
                    'mean_error': grp_df[err_col].mean(),
                    'median_error': grp_df[err_col].median(),
                    'mean_runtime': mean_rt,
                    'speedup_vs_exact': t_exact / max(mean_rt, 1e-6),
                })
        
        # Warm start
        if 'rel_error_warm_start' in grp_df.columns:
            mean_rt_ws = grp_df['runtime_warm_start'].mean() if 'runtime_warm_start' in grp_df.columns else 0
            rows.append({
                'dataset_regime': grp_name,
                'method': 'warm_start',
                'radius': None,
                'mean_error': grp_df['rel_error_warm_start'].mean(),
                'median_error': grp_df['rel_error_warm_start'].median(),
                'mean_runtime': mean_rt_ws,
                'speedup_vs_exact': t_exact / max(mean_rt_ws, 1e-6),
            })
        
        # One-step
        if 'rel_error_one_step' in grp_df.columns:
            mean_rt_os = grp_df['runtime_one_step'].mean() if 'runtime_one_step' in grp_df.columns else 0
            rows.append({
                'dataset_regime': grp_name,
                'method': 'one_step',
                'radius': None,
                'mean_error': grp_df['rel_error_one_step'].dropna().mean(),
                'median_error': grp_df['rel_error_one_step'].dropna().median(),
                'mean_runtime': mean_rt_os,
                'speedup_vs_exact': t_exact / max(mean_rt_os, 1e-6) if mean_rt_os > 0 else float('inf'),
            })
    
    return pd.DataFrame(rows)


def save_table(df, base_path, table_name):
    """Save table as CSV and markdown."""
    ensure_dir(os.path.dirname(base_path))
    
    csv_path = base_path + '.csv'
    md_path = base_path + '.md'
    
    df.to_csv(csv_path, index=False)
    
    with open(md_path, 'w') as f:
        f.write(f'# {table_name}\n\n')
        f.write(df.to_markdown(index=False))
        f.write('\n\n')
        
        # LaTeX version
        f.write('## LaTeX\n\n```latex\n')
        f.write(df.to_latex(index=False, float_format='%.4f'))
        f.write('```\n')
    
    print(f"  Saved: {csv_path}, {md_path}")
    return csv_path, md_path


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, default='config/default.yaml')
    args = parser.parse_args()
    
    print("Loading results...")
    records = load_all_results()
    if not records:
        print("No results found in results/raw/")
        return
    
    df = pd.DataFrame(records)
    print(f"Loaded {len(df)} records")
    
    # Remove complex columns for CSV
    drop_cols = ['influence_by_distance', 'influence_by_distance_full', 'deletion_edge']
    for col in drop_cols:
        if col in df.columns:
            df_clean = df.drop(columns=[col])
        else:
            df_clean = df
    
    # Save processed full dataframe
    proc_dir = ensure_dir('results/processed')
    df_clean.to_csv(os.path.join(proc_dir, 'all_results.csv'), index=False)
    print(f"Saved processed CSV: results/processed/all_results.csv")
    
    # Split by dataset type
    syn_df = df[df['dataset_type'] == 'synthetic'] if 'dataset_type' in df.columns else df
    real_df = df[df['dataset_type'] == 'real'] if 'dataset_type' in df.columns else pd.DataFrame()
    
    tables_dir = ensure_dir('results/tables')
    
    # Table 1: Synthetic regime summary
    if len(syn_df) > 0:
        syn_summary = process_synthetic(syn_df)
        if len(syn_summary) > 0:
            save_table(syn_summary, os.path.join(tables_dir, 'table_synthetic_regimes'),
                      'Synthetic Regime Summary')
    
    # Table 2: Real dataset summary
    if len(real_df) > 0:
        real_summary = process_synthetic(real_df)
        if len(real_summary) > 0:
            save_table(real_summary, os.path.join(tables_dir, 'table_real_datasets'),
                      'Real Dataset Summary')
    
    # Bootstrap CIs for key metrics
    from src.metrics import compute_bootstrap_summary
    
    metric_cols = ['empirical_decay_mu', 'rel_error_R1', 'rel_error_R2', 'rel_error_R3', 
                   'rel_error_R4', 'chi_seed_max', 'interference_cosine_R2',
                   'rel_error_warm_start', 'rel_error_one_step']
    
    if len(syn_df) > 0:
        boot_syn = compute_bootstrap_summary(
            syn_df, ['graph_type', 'prior_strength', 'K'], metric_cols)
        if len(boot_syn) > 0:
            save_table(boot_syn, os.path.join(tables_dir, 'table_synthetic_bootstrap'),
                      'Synthetic Bootstrap CIs')
    
    if len(real_df) > 0:
        boot_real = compute_bootstrap_summary(
            real_df, ['dataset_name', 'K'], metric_cols)
        if len(boot_real) > 0:
            save_table(boot_real, os.path.join(tables_dir, 'table_real_bootstrap'),
                      'Real Data Bootstrap CIs')
    
    if 'model_family' in df.columns and df['model_family'].nunique() > 1:
        boot_mf = compute_bootstrap_summary(
            df[df['dataset_type'] == 'synthetic'], 
            ['model_family', 'graph_type'], metric_cols)
        if len(boot_mf) > 0:
            save_table(boot_mf, os.path.join(tables_dir, 'table_model_family_bootstrap'),
                      'Model Family Bootstrap CIs')
    
    # Table 3: Correlations
    corr_df = compute_correlations(df)
    if len(corr_df) > 0:
        save_table(corr_df, os.path.join(tables_dir, 'table_correlations'),
                  'Correlation Summary')
    
    # Table 4: Method comparison
    method_df = build_method_comparison(df)
    if len(method_df) > 0:
        save_table(method_df, os.path.join(tables_dir, 'table_method_comparison'),
                  'Method Comparison')
    
    # Model family tables
    if 'model_family' in df.columns:
        mf_df = df[df['model_family'].notna()]
        if len(mf_df) > 0:
            # Model family summary
            mf_group_cols = ['model_family', 'graph_type', 'prior_strength']
            avail = [c for c in mf_group_cols if c in mf_df.columns]
            
            agg_cols = {}
            for col in ['chi_seed_max', 'empirical_decay_mu', 'rel_error_R2', 'rel_error_R4',
                         'runtime_exact', 'runtime_local_R2']:
                if col in mf_df.columns:
                    agg_cols[col] = ['mean', 'median']
            
            if agg_cols and avail:
                mf_summary = mf_df.groupby(avail).agg(agg_cols)
                mf_summary.columns = ['_'.join(c) for c in mf_summary.columns]
                mf_summary = mf_summary.reset_index()
                save_table(mf_summary, os.path.join(tables_dir, 'table_model_family_summary'),
                          'Model Family Summary')
            
            # Model family correlations
            mf_corr = compute_correlations(mf_df)
            if len(mf_corr) > 0:
                save_table(mf_corr, os.path.join(tables_dir, 'table_model_family_correlations'),
                          'Model Family Correlations')
    
    print("\nAnalysis complete.")


if __name__ == '__main__':
    main()