notebook_de_tool_comparison.py

# Cell 1: Imports and setup
from gather_stability_columns import gather_stability_columns, plot_multi_tool_correlation_vs_parameter
import os

# Create plots directory if it doesn't exist
os.makedirs('plots', exist_ok=True)

# Cell 2: Define the flexible loading function
def load_de_tool_data(tools_config, reference_column='mahalanobis', reference_index=0):
    """
    Load DE data for multiple tools with flexible embedding strategies.
    
    Parameters:
    -----------
    tools_config : dict
        Configuration for each tool with embedding strategy
    reference_column : str
        Type of column to use ('mahalanobis', 'log_fold_change', etc.)
    reference_index : int
        Which dataset size to use as reference for correlations
    
    Returns:
    --------
    dict : Tool data ready for plotting
    """
    
    tool_data = {}
    
    for tool_name, config in tools_config.items():
        use_precomputed = config.get('use_precomputed_embeddings', False)
        column_name = config['column_name']
        table_type = config.get('table_type', 'var')
        
        # Determine directory based on embedding strategy
        if use_precomputed:
            directory = f"data/processed/aging_precomputed_embeddings_{tool_name.lower()}_subsampling"
            display_name = f"{tool_name} (Precomputed)"
        else:
            directory = f"data/processed/aging_{tool_name.lower()}_subsampling"
            display_name = f"{tool_name} (Regular)"
        
        try:
            print(f"Loading {display_name} from {directory}...")
            df, _ = gather_stability_columns(directory, column_name, table_type)
            tool_data[display_name] = df
            print(f"✓ Loaded {df.shape[0]} genes, {df.shape[1]} cell counts")
        except Exception as e:
            print(f"✗ Failed to load {display_name}: {e}")
    
    return tool_data

# Cell 3: Quick comparison function
def quick_comparison(tools_with_embeddings):
    """
    Quick comparison function.
    
    Parameters:
    -----------
    tools_with_embeddings : dict
        {tool_name: use_precomputed_embeddings}
        
    Example:
    --------
    quick_comparison({
        'kompot': True,      # Use precomputed
        'scanpy': False,     # Use regular
        'mrvi': True,        # Use precomputed
        'lemur': False       # Use regular
    })
    """
    config = {}
    column_mapping = {
        'kompot': 'kompot_de_mahalanobis_Young_to_Old',
        'scanpy': 'scanpy_de_log_fold_change_Young_to_Old',
        'mrvi': 'mrvi_de_log_fold_change_Young_to_Old',
        'lemur': 'lemur_de_log_fold_change_Young_to_Old',
        'seurat': 'seurat_de_log_fold_change_Young_to_Old',
        'deseq2': 'deseq2_de_log_fold_change_Young_to_Old'
    }
    
    for tool, use_precomputed in tools_with_embeddings.items():
        config[tool] = {
            'use_precomputed_embeddings': use_precomputed,
            'column_name': column_mapping[tool],
            'table_type': 'var'
        }
    
    return load_de_tool_data(config)

# Cell 4: Configuration examples
# Just Kompot with precomputed embeddings vs others with regular
kompot_precomputed_config = {
    'kompot': {
        'use_precomputed_embeddings': True,
        'column_name': 'kompot_de_mahalanobis_Young_to_Old',
        'table_type': 'var'
    },
    'scanpy': {
        'use_precomputed_embeddings': False,
        'column_name': 'scanpy_de_log_fold_change_Young_to_Old',
        'table_type': 'var'
    },
    'mrvi': {
        'use_precomputed_embeddings': False,
        'column_name': 'mrvi_de_log_fold_change_Young_to_Old',
        'table_type': 'var'
    },
    'lemur': {
        'use_precomputed_embeddings': False,
        'column_name': 'lemur_de_log_fold_change_Young_to_Old',
        'table_type': 'var'
    },
    'seurat': {
        'use_precomputed_embeddings': False,
        'column_name': 'seurat_de_log_fold_change_Young_to_Old',
        'table_type': 'var'
    }
}

# Cell 5: Run Kompot precomputed vs others regular comparison
print("=== Kompot Precomputed vs Others Regular ===")
tool_data_1 = load_de_tool_data(kompot_precomputed_config)

if tool_data_1:
    fig1 = plot_multi_tool_correlation_vs_parameter(
        tool_data=tool_data_1,
        correlation_method='spearman',
        title='DE Comparison: Kompot (Precomputed) vs Others (Regular)',
        reference_index=0,
        save_path='plots/kompot_precomputed_vs_others_regular.png'
    )
    print("✓ Saved: plots/kompot_precomputed_vs_others_regular.png")

# Cell 6: Quick comparison example
print("\n=== Quick Comparison: Mixed Strategies ===")
quick_data = quick_comparison({
    'kompot': True,      # Kompot with precomputed embeddings
    'scanpy': False,     # Scanpy with regular embeddings  
    'mrvi': False,       # mrVI with regular embeddings
    'lemur': True        # LEMUR with precomputed embeddings
})

if quick_data:
    fig2 = plot_multi_tool_correlation_vs_parameter(
        tool_data=quick_data,
        correlation_method='spearman',
        title='Quick DE Comparison: Mixed Strategies',
        reference_index=0,
        save_path='plots/quick_mixed_comparison.png'
    )
    print("✓ Saved: plots/quick_mixed_comparison.png")

# Cell 7: Custom configuration (modify as needed)
# Example: Compare just Kompot regular vs precomputed
custom_config = {
    'kompot_regular': {
        'use_precomputed_embeddings': False,
        'column_name': 'kompot_de_mahalanobis_Young_to_Old',
        'table_type': 'var'
    },
    'kompot_precomputed': {
        'use_precomputed_embeddings': True,
        'column_name': 'kompot_de_mahalanobis_Young_to_Old',
        'table_type': 'var'
    }
}

print("\n=== Kompot: Regular vs Precomputed Embeddings ===")
custom_data = load_de_tool_data(custom_config)

if custom_data:
    fig3 = plot_multi_tool_correlation_vs_parameter(
        tool_data=custom_data,
        correlation_method='spearman',
        title='Kompot: Regular vs Precomputed Embeddings (DE)',
        reference_index=0,
        save_path='plots/kompot_regular_vs_precomputed.png'
    )
    print("✓ Saved: plots/kompot_regular_vs_precomputed.png")

print("\nAll plots completed!")