runtime_utils.py

#!/usr/bin/env python3
"""
Runtime Scaling Analysis for Kompot Parameter Sweeps

This script analyzes the runtime behavior of kompot across different parameter sweeps:
- ls_factor: Length scale factor for Gaussian process
- sigma: Sigma parameter
- n_components: Number of diffusion map components
- n_landmarks: Number of inducing points for Gaussian process (critical for performance!)
- topn_genes: Number of genes used for differential expression (linear scaling!)
- subsampling: Number of cells (subsampling stability)

We examine runtime scaling for:
- DA (no sample variance): Differential abundance analysis
- DE (no sample variance): Differential expression analysis on ALL genes

IMPORTANT NOTES:
- DE without sample variance uses ALL genes (16,285 for aging, 5,000 for COVID/BCR-XL)
- DE with sample variance uses top 200 highest-variance genes (TOPN_GENES=200)
- Gene selection is based on variance across cells
- The genes_processed field in runinfo shows TOPN_GENES (200), NOT the number used for no-sample-variance DE
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
from pathlib import Path
import yaml
from typing import Dict, List, Tuple, Optional
import warnings
warnings.filterwarnings('ignore')

# Set transparent background for all plots
matplotlib.rcParams['figure.facecolor'] = (1, 1, 1, 0)
matplotlib.rcParams['axes.facecolor'] = (1, 1, 1, 0)
matplotlib.rcParams['savefig.facecolor'] = (1, 1, 1, 0)

# High quality plots
matplotlib.rcParams['figure.dpi'] = 150
matplotlib.rcParams['savefig.dpi'] = 300
matplotlib.rcParams['font.size'] = 10
matplotlib.rcParams['axes.labelsize'] = 11
matplotlib.rcParams['axes.titlesize'] = 12
matplotlib.rcParams['legend.fontsize'] = 9


def extract_runtime_data(stability_dir: Path, force_param_type: str = None) -> pd.DataFrame:
    """
    Extract runtime data from all .runinfo.yml files in a stability directory.

    Parameters:
    -----------
    stability_dir : Path
        Directory containing runinfo.yml files
    force_param_type : str, optional
        Force a parameter type (used for old data without parameter_type field)

    Returns DataFrame with columns:
    - parameter_value: The parameter value for this run
    - da_no_sample_seconds: DA runtime without sample variance
    - de_no_sample_seconds: DE runtime without sample variance (uses ALL genes!)
    - n_cells: Number of cells in dataset
    - n_vars: TOTAL number of genes in dataset (used for DE no sample variance)
    - topn_genes: Top N genes for DE WITH sample variance (typically 200)
    """
    data = []

    runinfo_files = list(stability_dir.glob('*.runinfo.yml'))

    for runinfo_file in runinfo_files:
        try:
            with open(runinfo_file, 'r') as f:
                runinfo = yaml.safe_load(f)

            # Extract parameter info
            param_info = runinfo.get('parameter_variation', {})
            param_type = param_info.get('parameter_type', force_param_type)
            param_value = param_info.get('current_value')

            # Extract dataset metadata
            adata_info = runinfo.get('adata_info', {})
            n_cells = adata_info.get('n_obs')
            n_vars = adata_info.get('n_vars')  # TOTAL genes (used for DE no sample)

            # Special handling for subsampling: use n_cells as parameter value
            # Old data: parameter_type is missing but we infer from directory name
            # New data: current_value is a string like 'subsampled_round_195'
            if param_type == 'subsampling' or force_param_type == 'subsampling':
                param_value = n_cells

            if param_value is None:
                continue

            # Extract timing data
            timing = runinfo.get('timing', {})
            da_time = timing.get('da_no_sample_seconds')
            de_time = timing.get('de_no_sample_seconds')

            # genes_processed is TOPN_GENES (for DE with sample variance only!)
            topn_genes = timing.get('genes_processed')

            # Extract configuration info
            slurm_info = runinfo.get('slurm', {})
            config_info = runinfo.get('config', {})
            cpus = slurm_info.get('cpus_per_task')
            batch_size = config_info.get('BATCH_SIZE')

            # Extract node information to track hardware heterogeneity
            node = slurm_info.get('nodelist', 'unknown')
            node_type = 'gizmoj' if 'gizmoj' in node else 'gizmok' if 'gizmok' in node else 'other'

            # Extract null_genes from config_overrides or config
            # Default to 2000 if not specified (kompot default for older runs)
            config_overrides = runinfo.get('config_overrides', {})
            null_genes = config_overrides.get('NULL_GENES', config_info.get('NULL_GENES', 2000))

            # Only include if we have valid data
            if param_value is not None and (da_time is not None or de_time is not None):
                data.append({
                    'parameter_value': float(param_value),
                    'da_no_sample_seconds': da_time,
                    'de_no_sample_seconds': de_time,
                    'n_cells': n_cells,
                    'n_vars': n_vars,  # Total genes
                    'topn_genes': topn_genes,  # For sample variance DE only
                    'cpus': int(cpus) if cpus is not None else None,
                    'batch_size': int(batch_size) if batch_size is not None else None,
                    'null_genes': int(null_genes) if null_genes is not None else None,
                    'node': node,
                    'node_type': node_type
                })
        except Exception as e:
            print(f"Error processing {runinfo_file.name}: {e}")
            continue

    df = pd.DataFrame(data)
    if len(df) > 0:
        df = df.sort_values('parameter_value')

    return df


def get_dataset_metadata(df: pd.DataFrame, is_subsampling: bool = False) -> Dict:
    """
    Extract dataset metadata from runtime DataFrame.

    Parameters:
    -----------
    df : pd.DataFrame
        Runtime data
    is_subsampling : bool
        If True, n_cells varies and we return min/max range

    Returns dict with:
    - n_cells: Number of cells (single value or dict with min/max for subsampling)
    - n_vars: Total number of genes (used for DE no sample variance)
    - topn_genes: Top N genes (for DE with sample variance)
    - cpus: Number of CPUs used
    - batch_size: Batch size used (0 = no batching)
    - null_genes: Number of null genes used for FDR correction (default 2000)
    """
    if len(df) == 0:
        return {
            'n_cells': None,
            'n_vars': None,
            'topn_genes': None,
            'cpus': None,
            'batch_size': None,
            'null_genes': None
        }

    # For subsampling, n_cells varies - return range
    if is_subsampling and 'n_cells' in df.columns:
        n_cells = {
            'min': int(df['n_cells'].min()),
            'max': int(df['n_cells'].max())
        }
    else:
        # For other parameters, n_cells is constant
        n_cells = int(df['n_cells'].iloc[0]) if 'n_cells' in df.columns else None

    n_vars = df['n_vars'].iloc[0] if 'n_vars' in df.columns else None
    topn_genes = df['topn_genes'].iloc[0] if 'topn_genes' in df.columns else None

    # Extract configuration info
    cpus = df['cpus'].iloc[0] if 'cpus' in df.columns else None
    batch_size = df['batch_size'].iloc[0] if 'batch_size' in df.columns else None
    null_genes = df['null_genes'].iloc[0] if 'null_genes' in df.columns else None

    return {
        'n_cells': n_cells,
        'n_vars': n_vars,
        'topn_genes': topn_genes,
        'cpus': cpus,
        'batch_size': batch_size,
        'null_genes': null_genes
    }


def plot_runtime_scaling(
    runtime_data: Dict[str, pd.DataFrame],
    parameter_type: str,
    analysis_type: str,
    output_dir: Optional[Path] = None,
    show_legend_separately: bool = True,
    min_param_value: Optional[float] = None
) -> Tuple[plt.Figure, Optional[plt.Figure]]:
    """
    Plot runtime scaling across parameter values for multiple datasets.

    Parameters:
    -----------
    runtime_data : Dict[str, pd.DataFrame]
        Dictionary mapping dataset names to runtime DataFrames
    parameter_type : str
        Type of parameter ('ls_factor', 'sigma', 'n_components', 'n_landmarks', 'subsampling')
    analysis_type : str
        Type of analysis ('da' or 'de')
    output_dir : Optional[Path]
        Directory to save plots (if None, don't save)
    show_legend_separately : bool
        If True, create a separate legend plot
    min_param_value : Optional[float]
        If specified, exclude parameter values below this threshold
        (e.g., min_param_value=2 to exclude n_landmarks=1 or very small cell counts)

    Returns:
    --------
    (main_fig, legend_fig) : Tuple of matplotlib figures
    """
    # Parameter display configuration
    param_config = {
        'ls_factor': {
            'display_name': 'LS factor',
            'use_log_scale': True,
        },
        'sigma': {
            'display_name': 'Sigma',
            'use_log_scale': True,
        },
        'n_components': {
            'display_name': 'Number of diffusion components',
            'use_log_scale': False,
        },
        'n_landmarks': {
            'display_name': 'Number of landmarks (inducing points)',
            'use_log_scale': True,
        },
        'topn_genes': {
            'display_name': 'Number of genes for DE',
            'use_log_scale': True,  # Log-log scale
        },
        'subsampling': {
            'display_name': 'Number of cells',
            'use_log_scale': True,
        },
    }

    config = param_config.get(parameter_type, {
        'display_name': parameter_type,
        'use_log_scale': False
    })

    # Analysis type configuration
    analysis_config = {
        'da': {
            'column': 'da_no_sample_seconds',
            'ylabel': 'Log runtime (seconds)',
            'title': f'DA Runtime Scaling vs {config["display_name"]}',
            'include_genes': False,
        },
        'de': {
            'column': 'de_no_sample_seconds',
            'ylabel': 'Log runtime (seconds)',
            'title': f'DE Runtime Scaling vs {config["display_name"]}',
            'include_genes': True,
        },
    }

    aconfig = analysis_config[analysis_type]

    # Create main plot
    fig, ax = plt.subplots(figsize=(8, 5))

    # Color palette for datasets
    colors = plt.cm.tab10(np.linspace(0, 1, len(runtime_data)))

    legend_entries = []

    for (dataset_name, df), color in zip(runtime_data.items(), colors):
        if len(df) == 0:
            continue

        # Filter out missing values
        valid_data = df.dropna(subset=[aconfig['column']])

        # Apply minimum parameter value filter if specified
        if min_param_value is not None:
            valid_data = valid_data[valid_data['parameter_value'] >= min_param_value]

        if len(valid_data) == 0:
            continue

        x = valid_data['parameter_value'].values
        y = valid_data[aconfig['column']].values

        # Plot line
        line, = ax.plot(x, y, marker='o', markersize=3, linewidth=1.5,
                       color=color, alpha=0.8)

        # Create legend label
        is_subsampling = (parameter_type == 'subsampling')
        metadata = get_dataset_metadata(df, is_subsampling=is_subsampling)
        n_cells = metadata['n_cells']
        n_vars = metadata['n_vars']  # Total genes (used for DE no sample)
        cpus = metadata['cpus']
        batch_size = metadata['batch_size']
        null_genes = metadata['null_genes']

        # Format n_cells for display
        if isinstance(n_cells, dict):
            # Subsampling: show range
            n_cells_str = f"{n_cells['min']:,}-{n_cells['max']:,}"
        else:
            # Fixed n_cells
            n_cells_str = f"{n_cells:,}"

        # Build label with configuration info
        label_parts = [dataset_name]

        # Add cell count and genes
        if aconfig['include_genes'] and n_vars is not None:
            # For topn_genes plots, show max available genes (n_vars), not current topn value
            if parameter_type == 'topn_genes':
                label_parts.append(f"n={n_cells_str}, max_genes={n_vars:,}")
            else:
                label_parts.append(f"n={n_cells_str}, genes={n_vars:,}")
        else:
            label_parts.append(f"n={n_cells_str}")

        # Add configuration details
        config_parts = []
        if cpus is not None:
            config_parts.append(f"{cpus}CPU")

        # Only show batch_size for DE plots (batching only affects gene processing)
        if analysis_type == 'de' and batch_size is not None:
            if batch_size == 0:
                config_parts.append("no-batch")
            else:
                config_parts.append(f"batch={batch_size}")

        # Show null_genes for DE plots (null genes add to computational cost)
        if analysis_type == 'de' and null_genes is not None:
            if null_genes == 0:
                config_parts.append("null=0")
            else:
                config_parts.append(f"null={null_genes}")

        if config_parts:
            label_parts.append(", ".join(config_parts))

        label = f"{label_parts[0]} ({', '.join(label_parts[1:])})"

        legend_entries.append((line, label))

    # Configure axes
    if config['use_log_scale']:
        ax.set_xscale('log')

    # Always use log scale for y-axis (runtime)
    ax.set_yscale('log')

    ax.set_xlabel(config['display_name'], fontsize=11)
    ylabel = aconfig['ylabel']
    ax.set_ylabel(ylabel, fontsize=11)

    ax.set_title(aconfig['title'], fontsize=12, fontweight='bold')
    ax.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)

    # Add legend to main plot if not showing separately
    if not show_legend_separately:
        lines, labels = zip(*legend_entries)
        ax.legend(lines, labels, fontsize=8, framealpha=0.9, loc='best')

    plt.tight_layout()

    # Save main plot
    if output_dir is not None:
        output_dir.mkdir(parents=True, exist_ok=True)
        filename = f"runtime_{analysis_type}_{parameter_type}.png"
        fig.savefig(output_dir / filename, dpi=300, bbox_inches='tight',
                   facecolor=(1, 1, 1, 0))
        print(f"Saved: {output_dir / filename}")

    # Create separate legend plot if requested
    legend_fig = None
    if show_legend_separately and len(legend_entries) > 0:
        legend_fig = plt.figure(figsize=(6, len(legend_entries) * 0.3 + 0.5))
        legend_ax = legend_fig.add_subplot(111)
        legend_ax.axis('off')

        lines, labels = zip(*legend_entries)
        legend_ax.legend(lines, labels, loc='center', fontsize=9,
                        frameon=False, ncol=1)

        plt.tight_layout()

        # Save legend
        if output_dir is not None:
            filename = f"runtime_{analysis_type}_{parameter_type}_legend.png"
            legend_fig.savefig(output_dir / filename, dpi=300, bbox_inches='tight',
                             facecolor=(1, 1, 1, 0))
            print(f"Saved: {output_dir / filename}")

    return fig, legend_fig


def collect_all_runtime_data(base_dir: Path = Path('data/processed')) -> Dict[str, Dict[str, pd.DataFrame]]:
    """
    Collect runtime data from all parameter stability directories.

    Parameters:
    -----------
    base_dir : Path
        Base directory containing processed data

    Returns:
    --------
    Dict mapping parameter_type -> dataset_name -> DataFrame
    """
    # Find all stability directories
    stability_dirs = {
        'ls_factor': [],
        'sigma': [],
        'n_components': [],
        'n_landmarks': [],
        'topn_genes': [],
        'subsampling': []
    }

    for param_type in stability_dirs.keys():
        if param_type == 'subsampling':
            # Special pattern for subsampling (only kompot, not tool-specific)
            pattern = "*_kompot_subsampling"
        else:
            pattern = f"*_{param_type}"

        dirs = sorted(base_dir.glob(pattern))
        stability_dirs[param_type] = dirs

    # Collect runtime data for all parameter types
    all_runtime_data = {}

    for param_type, dirs in stability_dirs.items():
        all_runtime_data[param_type] = {}

        for stability_dir in dirs:
            # Extract dataset name from directory
            if param_type == 'subsampling':
                dataset_name = stability_dir.name.replace('_kompot_subsampling', '')
            else:
                dataset_name = stability_dir.name.replace(f'_{param_type}', '')

            print(f"Processing {dataset_name} - {param_type}...", end=' ')

            # Extract runtime data
            # For subsampling, force the param_type (old data doesn't have it)
            if param_type == 'subsampling':
                df = extract_runtime_data(stability_dir, force_param_type='subsampling')
            else:
                df = extract_runtime_data(stability_dir)

            if len(df) > 0:
                all_runtime_data[param_type][dataset_name] = df
                is_subsampling_type = (param_type == 'subsampling')
                metadata = get_dataset_metadata(df, is_subsampling=is_subsampling_type)

                # Format n_cells for display
                n_cells = metadata['n_cells']
                if isinstance(n_cells, dict):
                    n_cells_str = f"{n_cells['min']:,}-{n_cells['max']:,}"
                else:
                    n_cells_str = f"{n_cells:,}"

                # Format config for display
                config_str = f"{metadata['cpus']}CPU"
                if metadata['batch_size'] == 0:
                    config_str += ", no-batch"
                elif metadata['batch_size'] is not None:
                    config_str += f", batch={metadata['batch_size']}"

                print(f"✓ ({len(df)} points, n_cells={n_cells_str}, n_vars={metadata['n_vars']:,}, {config_str})")
            else:
                print("✗ (no data)")

    return all_runtime_data


def print_summary_statistics(all_runtime_data: Dict[str, Dict[str, pd.DataFrame]]):
    """Print summary statistics for all collected runtime data."""

    print("\n" + "="*80)
    print("RUNTIME SUMMARY STATISTICS")
    print("="*80)

    for param_type in ['ls_factor', 'sigma', 'n_components', 'n_landmarks', 'topn_genes', 'subsampling']:
        if param_type not in all_runtime_data or len(all_runtime_data[param_type]) == 0:
            continue

        print(f"\n{param_type.upper()}:")
        print("-" * 80)

        for dataset_name, df in all_runtime_data[param_type].items():
            is_subsampling_type = (param_type == 'subsampling')
            metadata = get_dataset_metadata(df, is_subsampling=is_subsampling_type)

            print(f"\n  {dataset_name}:")

            # Format n_cells for display
            n_cells = metadata['n_cells']
            if isinstance(n_cells, dict):
                print(f"    n_cells (range)={n_cells['min']:,} to {n_cells['max']:,}")
            else:
                print(f"    n_cells={n_cells:,}")

            print(f"    n_vars (total genes)={metadata['n_vars']:,}")
            print(f"    topn_genes (for DE with sample var)={metadata['topn_genes']}")

            # Show configuration
            print(f"    config: {metadata['cpus']}CPU", end="")
            if metadata['batch_size'] == 0:
                print(", no batching")
            elif metadata['batch_size'] is not None:
                print(f", batch_size={metadata['batch_size']}")
            else:
                print()

            # DA statistics
            da_valid = df.dropna(subset=['da_no_sample_seconds'])
            if len(da_valid) > 0:
                da_min = da_valid['da_no_sample_seconds'].min()
                da_max = da_valid['da_no_sample_seconds'].max()
                da_mean = da_valid['da_no_sample_seconds'].mean()
                print(f"    DA: min={da_min:.1f}s, max={da_max:.1f}s, mean={da_mean:.1f}s")

            # DE statistics
            de_valid = df.dropna(subset=['de_no_sample_seconds'])
            if len(de_valid) > 0:
                de_min = de_valid['de_no_sample_seconds'].min()
                de_max = de_valid['de_no_sample_seconds'].max()
                de_mean = de_valid['de_no_sample_seconds'].mean()
                print(f"    DE (ALL {metadata['n_vars']:,} genes): min={de_min:.1f}s, max={de_max:.1f}s, mean={de_mean:.1f}s")

    print("\n" + "="*80)