example_gene_subset_analysis.py

#!/usr/bin/env python3
"""
Example: Using Gene Subsetting in Stability Analysis

This script demonstrates how to use the new gene subsetting functionality
in gather_stability_columns to analyze only highly variable genes.
"""

from gather_stability_columns import gather_stability_columns, get_highly_variable_genes
from plot_stability_analysis import plot_multi_tool_correlation_vs_parameter
import logging

# Set up logging to see the gene subsetting process
logging.basicConfig(level=logging.INFO)

# ==============================================================================
# Method 1: Using highly variable genes from original dataset
# ==============================================================================

# Extract highly variable genes from the original aging dataset
# Replace this path with the actual path to your original aging dataset
original_aging_path = "path/to/original_aging_data.h5ad"  # Update this path

try:
    # Get highly variable genes from the original dataset
    hvg_genes = get_highly_variable_genes(
        original_aging_path, 
        hvg_column='highly_variable'  # or 'highly_variable_genes' depending on your data
    )
    print(f"Found {len(hvg_genes)} highly variable genes")
    
    # Example of first few genes
    print(f"First 10 HVG genes: {hvg_genes[:10]}")
    
except Exception as e:
    print(f"Could not load HVG genes from original dataset: {e}")
    print("Using a mock list for demonstration purposes")
    # Mock highly variable genes for demonstration
    hvg_genes = ["GENE1", "GENE2", "GENE3"]  # Replace with actual gene names

# ==============================================================================
# Method 2: Using a predefined gene list
# ==============================================================================

# You can also manually specify genes of interest
immune_genes = [
    "CD3D", "CD3E", "CD3G", "CD4", "CD8A", "CD8B", 
    "IL2", "IL4", "IL10", "IFNG", "TNF", "CCL2"
]

# Choose which gene subset to use
selected_genes = hvg_genes  # or immune_genes for immune-specific analysis

# ==============================================================================
# Gather data with gene subsetting
# ==============================================================================

print("\n=== Gathering stability data with gene subsetting ===")

# For DE comparison with gene subsetting
kompot_regular_de_df, _ = gather_stability_columns(
    'data/processed/aging_subsampling',
    'kompot_de_mahalanobis_Young_to_Old',
    'var',
    subset_genes=selected_genes  # NEW: Apply gene subsetting
)

seurat_de_df, _ = gather_stability_columns(
    'data/processed/aging_seurat_subsampling',
    'seurat_de_neg_log10_pvalue_Young_to_Old',
    'var',
    subset_genes=selected_genes  # NEW: Apply gene subsetting
)

scanpy_de_df, _ = gather_stability_columns(
    'data/processed/aging_scanpy_subsampling',
    'scanpy_de_neg_log10_pvalue_Young_to_Old',
    'var',
    subset_genes=selected_genes  # NEW: Apply gene subsetting
)

kompot_precomputed_de_df, _ = gather_stability_columns(
    'data/processed/aging_precomputed_embeddings_subsampling',
    'kompot_de_mahalanobis_Young_to_Old',
    'var',
    subset_genes=selected_genes  # NEW: Apply gene subsetting
)

print(f"Data shapes after gene subsetting:")
print(f"  Kompot Regular: {kompot_regular_de_df.shape}")
print(f"  Seurat: {seurat_de_df.shape}")
print(f"  Scanpy: {scanpy_de_df.shape}")
print(f"  Kompot Precomputed: {kompot_precomputed_de_df.shape}")

# ==============================================================================
# Create comparison plot with subsetted genes
# ==============================================================================

print("\n=== Creating comparison plot ===")

fig = plot_multi_tool_correlation_vs_parameter(
    tool_data={
        'Kompot DE (Regular)': kompot_regular_de_df,
        'Kompot DE (Precomputed Embeddings)': kompot_precomputed_de_df,
        'Seurat': seurat_de_df,
        'Scanpy': scanpy_de_df,
    },
    correlation_method='spearman',
    title=f'Aging Dataset: DE Comparison (HVG subset, n={len(selected_genes)} genes)',
    reference_index=0,
    save_path='plots/kompot_de_embedding_comparison_hvg_subset.png'
)

print("Analysis complete!")
print(f"Plot saved to: plots/kompot_de_embedding_comparison_hvg_subset.png")

# ==============================================================================
# Additional example: Compare HVG vs All genes
# ==============================================================================

print("\n=== Bonus: Comparing HVG subset vs All genes ===")

# Load the same data without gene subsetting for comparison
kompot_all_genes_df, _ = gather_stability_columns(
    'data/processed/aging_subsampling',
    'kompot_de_mahalanobis_Young_to_Old',
    'var'
    # No subset_genes parameter = use all genes
)

print(f"Kompot data shapes:")
print(f"  All genes: {kompot_all_genes_df.shape}")
print(f"  HVG subset: {kompot_regular_de_df.shape}")
print(f"  Reduction: {kompot_all_genes_df.shape[0] - kompot_regular_de_df.shape[0]} genes filtered out")

# You could create additional plots comparing HVG vs all genes
fig2 = plot_multi_tool_correlation_vs_parameter(
    tool_data={
        'Kompot (All genes)': kompot_all_genes_df,
        'Kompot (HVG only)': kompot_regular_de_df,
    },
    correlation_method='spearman',
    title='Gene Subsetting Impact: All genes vs HVG',
    reference_index=0,
    save_path='plots/all_vs_hvg_comparison.png'
)

print("Bonus analysis complete!")