Enrichment Analysis¶

def run_deseq2_analysis(
    sample_data: pd.DataFrame,
    control_data: pd.DataFrame,
    control_metadata: pd.DataFrame,
    gene_list: np.ndarray,
    sample_name: str,
    stage_transition: str
) -> pd.Series:
    """
    Perform DESeq2 differential expression analysis between sample and controls.

    This function properly:
    1. Converts log2(RSEM+1) data back to RSEM integer counts
    2. Runs PyDESeq2 analysis to get statistically valid log2FoldChange values
    3. Returns ranked gene list for GSEA

    Args:
        sample_data: Sample expression data (genes x 1) in log2(RSEM+1) format
        control_data: Control expression data (genes x samples) in log2(RSEM+1) format
        control_metadata: Control clinical metadata with stage information
        gene_list: List of genes
        sample_name: Name/ID of the sample
        stage_transition: Stage transition label (e.g., 'early_to_late', 'I_to_II')

    Returns:
        Series of log2FoldChange values sorted for GSEA input
    """
    # Debug logging
    logger.debug(f"DESeq2 analysis for {sample_name}")
    logger.debug(f"  Sample data shape: {sample_data.shape}")
    logger.debug(f"  Control data shape: {control_data.shape}")
    logger.debug(f"  Gene list length: {len(gene_list)}")

    # 1. Convert from log2(RSEM+1) back to RSEM counts
    # The preprocessed data is in log2(RSEM+1) format, so we reverse this transformation
    # and round to integers as required by DESeq2
    sample_counts = 2 ** sample_data.values.flatten() - 1
    control_counts = 2 ** control_data.values - 1

    # Clip negative values to 0 (can occur due to numerical precision or interpolation)
    sample_counts = np.clip(sample_counts, 0, None)
    control_counts = np.clip(control_counts, 0, None)

    # Round to integers as required by DESeq2
    sample_counts = np.round(sample_counts).astype(int)
    control_counts = np.round(control_counts).astype(int)

    # Ensure control_counts is 2D (genes x samples)
    if control_counts.ndim == 1:
        control_counts = control_counts[:, np.newaxis]

    # Validate counts
    if np.any(sample_counts < 0):
        logger.error(f"Sample has negative counts after conversion: min={sample_counts.min()}")
        raise ValueError("Sample counts contain negative values after conversion")
    if np.any(control_counts < 0):
        logger.error(f"Controls have negative counts after conversion: min={control_counts.min()}")
        raise ValueError("Control counts contain negative values after conversion")

    logger.debug(f"  Sample counts: min={sample_counts.min()}, max={sample_counts.max()}, mean={sample_counts.mean():.1f}")
    logger.debug(f"  Control counts: min={control_counts.min()}, max={control_counts.max()}, mean={control_counts.mean():.1f}")

    # Ensure sample_counts has the same number of genes as control_counts
    if len(sample_counts) != control_counts.shape[0]:
        raise ValueError(
            f"Sample has {len(sample_counts)} genes but controls have {control_counts.shape[0]} genes. "
            f"Sample shape: {sample_data.shape}, Control shape: {control_data.shape}"
        )

    # 2. Combine sample with controls
    # Create count matrix: genes (rows) x samples (columns)
    # sample_counts is 1D (n_genes,), control_counts is 2D (n_genes, n_controls)
    # We need to reshape sample_counts to (n_genes, 1) to stack horizontally
    counts_matrix = np.column_stack([sample_counts.reshape(-1, 1), control_counts])

    counts_df = pd.DataFrame(
        counts_matrix,
        index=gene_list,
        columns=[sample_name] + list(control_data.columns)
    )

    # 3. Create metadata DataFrame with condition labels
    metadata_df = pd.DataFrame({
        'condition': [stage_transition] + list(control_metadata['ajcc_pathologic_tumor_stage'].values)
    }, index=counts_df.columns)

    # 4. Run PyDESeq2 analysis
    # Initialize DESeqDataSet
    dds = DeseqDataSet(
        counts=counts_df.T,  # DESeq2 expects samples as rows, genes as columns
        metadata=metadata_df,
        design_factors='condition',
        refit_cooks=True,
        quiet=True  # Suppress verbose output in parallel processing
    )

    # Fit dispersions and log-fold changes
    dds.deseq2()

    # Statistical analysis
    stat_res = DeseqStats(
        dds,
        alpha=0.05,
        cooks_filter=True,
        independent_filter=True,
        quiet=True
    )
    stat_res.summary()

    # 5. Extract log2FoldChange values
    results_df = stat_res.results_df

    # Ensure we have log2FoldChange column
    if 'log2FoldChange' not in results_df.columns:
        raise ValueError(f"DESeq2 results missing 'log2FoldChange' column. Available columns: {results_df.columns.tolist()}")

    # Extract log2FoldChange and sort by absolute value (for GSEA ranking)
    log2fc = results_df['log2FoldChange'].copy()

    # Replace NaN values with 0 (genes with no change or insufficient data)
    log2fc = log2fc.fillna(0)

    # Sort by absolute value descending (GSEA expects ranked list)
    log2fc = log2fc.sort_values(ascending=False)

    logger.debug(f"  DESeq2 complete: {len(log2fc)} genes, range [{log2fc.min():.2f}, {log2fc.max():.2f}]")

    return log2fc

def run_gsea_command(cmd: str) -> bool:
    """
    Run a single GSEA command.

    Args:
        cmd: GSEA command string

    Returns:
        True if successful, False otherwise
    """
    try:
        result = subprocess.run(
            cmd,
            shell=True,
            capture_output=True,
            text=True,
            timeout=600  # 10 minute timeout per command
        )

        if result.returncode != 0:
            logger.error(f"GSEA command failed with return code {result.returncode}")
            logger.error(f"Command: {cmd[:200]}...")  # Truncate long commands
            logger.error(f"STDERR: {result.stderr[:500]}")  # First 500 chars
            if result.stdout:
                logger.error(f"STDOUT: {result.stdout[:500]}")
            return False

        # Check for Java errors in stdout (GSEA sometimes returns 0 even on error)
        if result.stdout and ("error" in result.stdout.lower() or "exception" in result.stdout.lower()):
            logger.warning(f"GSEA command may have failed (found error in output)")
            logger.warning(f"Command: {cmd[:200]}...")
            logger.warning(f"Output: {result.stdout[:500]}")  # First 500 chars
            return False

        # Check for GSEA-specific success indicators
        if result.stdout and "Enrichment score" in result.stdout:
            logger.debug(f"GSEA command completed successfully")
            return True
        elif result.stdout:
            # Log a sample of output for debugging
            logger.debug(f"GSEA output sample: {result.stdout[:200]}")

        return True

    except subprocess.TimeoutExpired:
        logger.error(f"GSEA command timed out after 600s")
        logger.error(f"Command: {cmd[:200]}...")
        return False
    except Exception as e:
        logger.error(f"Error running GSEA command: {e}")
        logger.error(f"Command: {cmd[:200]}...")
        return False

def generate_pathway_heatmap(
    enrichment_df: pd.DataFrame,
    output_dir: str,
    fdr_threshold: float = 0.05,
    colorbar: bool = True,
    legend: bool = False,
    yticks_fontsize: int = 12,
    show: bool = False
) -> Tuple[pd.DataFrame, Dict[str, 'matplotlib.figure.Figure']]:
    """
    Generate multiple pathway enrichment heatmaps from GSEA results.

    This function creates several heatmaps showing the sum of NES (Normalized Enrichment Score)
    across all trajectories for each pathway at each timepoint:

    1. Top 50 most changing pathways (first vs last timepoint)
    2. Top 50 most upregulated pathways (average NES > 0)
    3. Top 50 most downregulated pathways (average NES < 0)
    4. Selected pathways (high-level Reactome + literature pathways)

    The heatmaps have:
    - Rows: Pathway names
    - Columns: Timepoints (pseudo-time from early to late)
    - Values: Sum of NES across all trajectories at each timepoint

    Args:
        enrichment_df: DataFrame with columns [Patient, Idx, Transition, NAME, ES, NES, FDR q-val]
        output_dir: Output directory for heatmap files
        fdr_threshold: FDR q-value threshold for significance (default: 0.05)
        colorbar: Whether to show colorbar (default: True)
        legend: Whether to show legend (default: False)
        yticks_fontsize: Font size for y-axis tick labels (default: 12)
        show: Whether to display the plot (default: False)

    Returns:
        Tuple of (heatmap_data, figures_dict):
            - heatmap_data: DataFrame with summed NES values (pathways × timepoints)
            - figures_dict: Dictionary mapping figure names to Matplotlib Figure objects

    Example:
        >>> enrichment_df = pd.read_csv('trajectory_enrichment.csv')
        >>> heatmap_data, figs = generate_pathway_heatmap(
        ...     enrichment_df=enrichment_df,
        ...     output_dir='results/',
        ...     fdr_threshold=0.05
        ... )
        >>> print(f"Generated {len(figs)} heatmaps")
    """
    import matplotlib.pyplot as plt
    import matplotlib.patches as mpatches
    import matplotlib.colors as mcolors

    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    logger.info(f"Generating pathway enrichment heatmaps (FDR < {fdr_threshold})...")

    # Define pathway lists
    highest_pathways = [
        "Autophagy",
        "Cell Cycle",
        "Cell-Cell communication",
        "Cellular responses to stimuli",
        "Chromatin organization",
        "Circadian Clock",
        "DNA Repair",
        "DNA Replication",
        "Developmental Biology",
        "Digestion and absorption",
        "Disease",
        "Drug ADME",
        "Extracellular matrix organization",
        "Gene expression (Transcription)",
        "Hemostasis",
        "Immune System",
        "Metabolism",
        "Metabolism of RNA",
        "Metabolism of proteins",
        "Muscle contraction",
        "Neuronal System",
        "Organelle biogenesis and maintenance",
        "Programmed Cell Death",
        "Protein localization",
        "Reproduction",
        "Sensory Perception",
        "Signal Transduction",
        "Transport of small molecules",
        "Vesicle-mediated transport",
    ]

    pathways_literature = [
        # VHL/HIF pathway
        "CELLULAR RESPONSE TO HYPOXIA",
        "OXYGEN-DEPENDENT PROLINE HYDROXYLATION OF HYPOXIA-INDUCIBLE FACTOR ALPHA",
        "REGULATION OF GENE EXPRESSION BY HYPOXIA-INDUCIBLE FACTOR",
        # PI3K/AKT/MTOR Pathway
        "PI3K/AKT ACTIVATION",
        "PI3K/AKT SIGNALING IN CANCER",
        "MTOR SIGNALLING",
        # Warburg effect
        "TP53 REGULATES METABOLIC GENES",
        "GLYCOLYSIS",
        "GLUCOSE METABOLISM",
        # TCA/Krebs cycle
        "CITRIC ACID CYCLE (TCA CYCLE)",
        "THE CITRIC ACID (TCA) CYCLE AND RESPIRATORY ELECTRON TRANSPORT",
        # Pentose phosphate pathway
        "NFE2L2 REGULATES PENTOSE PHOSPHATE PATHWAY GENES",
        "PENTOSE PHOSPHATE PATHWAY",
        "PENTOSE PHOSPHATE PATHWAY DISEASE",
        # Fatty Acid Metabolism
        "FATTY ACID METABOLISM",
        # Glutamine metabolism
        "GLUTAMATE AND GLUTAMINE METABOLISM",
        # EGFR
        "SIGNALING BY EGFR",
        "SIGNALING BY EGFR IN CANCER",
        "EGFR DOWNREGULATION",
        # TGF-β signaling
        "SIGNALING BY TGF-BETA RECEPTOR COMPLEX",
        "TGF-BETA RECEPTOR SIGNALING IN EMT (EPITHELIAL TO MESENCHYMAL TRANSITION)",
        "SIGNALING BY TGF-BETA RECEPTOR COMPLEX IN CANCER",
        "SIGNALING BY TGFB FAMILY MEMBERS",
        "TGF-BETA RECEPTOR SIGNALING ACTIVATES SMADS",
        # Wnt/β-catenin pathway
        "BETA-CATENIN INDEPENDENT WNT SIGNALING",
        "SIGNALING BY WNT",
        # SLIT-2-ROBO1 pathways
        "REGULATION OF EXPRESSION OF SLITS AND ROBOS",
        # DNA repair
        "DNA REPAIR",
        # Energy homeostasis
        "ION HOMEOSTASIS",
        # Apoptosis
        "APOPTOSIS",
        # Angiogenesis
        "SIGNALING BY VEGF"
    ]

    # Step 1: Ensure numeric types for NES and FDR q-val
    enrichment_df = enrichment_df.copy()
    enrichment_df['NES'] = pd.to_numeric(enrichment_df['NES'], errors='coerce')
    enrichment_df['FDR q-val'] = pd.to_numeric(enrichment_df['FDR q-val'], errors='coerce')
    enrichment_df['Idx'] = pd.to_numeric(enrichment_df['Idx'], errors='coerce')

    # Log any rows that had non-numeric values
    invalid_nes = enrichment_df['NES'].isna().sum()
    invalid_fdr = enrichment_df['FDR q-val'].isna().sum()
    if invalid_nes > 0:
        logger.warning(f"Found {invalid_nes} rows with non-numeric NES values (converted to NaN)")
    if invalid_fdr > 0:
        logger.warning(f"Found {invalid_fdr} rows with non-numeric FDR q-val values (converted to NaN)")

    # Step 2: Filter by FDR threshold
    significant = enrichment_df[enrichment_df['FDR q-val'] < fdr_threshold].copy()

    logger.info(f"Found {significant.shape[0]} significant pathway enrichments (FDR < {fdr_threshold})")

    if significant.empty:
        logger.warning("No significant pathways found. Cannot generate heatmap.")
        # Return empty results
        empty_df = pd.DataFrame()
        empty_dict = {}
        return empty_df, empty_dict

    # Step 3: Group by Timepoint (Idx) and Pathway (NAME), sum NES across all trajectories
    pathway_summary = significant.groupby(['Idx', 'NAME'])['NES'].sum().reset_index()

    logger.info(f"Aggregated results for {pathway_summary['Idx'].nunique()} timepoints "
                f"and {pathway_summary['NAME'].nunique()} pathways")

    # Step 4: Pivot to create matrix (pathways × timepoints)
    heatmap_data = pathway_summary.pivot(
        index='NAME',
        columns='Idx',
        values='NES'
    ).fillna(0)  # Fill missing with 0

    logger.info(f"Full heatmap dimensions: {heatmap_data.shape[0]} pathways × {heatmap_data.shape[1]} timepoints")

    # Save full summary data
    summary_file = output_dir / 'pathway_nes_summary.csv'
    heatmap_data.to_csv(summary_file)
    logger.info(f"Saved pathway NES summary to: {summary_file}")

    # Dictionary to store all figures
    figures = {}

    # Helper function to create and save a heatmap
    def plot_heatmap_regulation(df_plot, unique_pathways, cmap_here='viridis',
                               save_name=None, colorbar_title='Sum of NES'):
        """Plot heatmap following paper_figures.ipynb style"""
        # Generate a range of locations for the ticks
        tick_locations = range(len(unique_pathways))
        z_min, z_max = df_plot.min().min(), df_plot.max().max()

        # Make the range symmetric around 0
        if z_min < 0 and z_max > 0:
            abs_max = max(abs(z_min), abs(z_max))
            z_min, z_max = -abs_max, abs_max
            norm = mcolors.TwoSlopeNorm(vmin=z_min, vcenter=0, vmax=z_max)
        else:
            # If all values are positive or all negative, use regular normalization
            norm = mcolors.Normalize(vmin=z_min, vmax=z_max)
            logger.warning(f"Cannot center colormap at zero: range [{z_min:.3f}, {z_max:.3f}] does not cross zero")

        fig, ax = plt.subplots(figsize=(30, 10))

        # Make heatmap
        cax = ax.imshow(df_plot.values, cmap=cmap_here, norm=norm, aspect='auto')

        # Set the y-ticks
        plt.yticks(tick_locations, unique_pathways, fontsize=yticks_fontsize)

        # Set the x-ticks at specific positions
        num_timepoints = df_plot.shape[1]
        ax.set_xticks([0, num_timepoints - 1])
        ax.set_xticklabels(['early', 'late'], fontsize=yticks_fontsize*1.33, rotation=45)
        ax.set_xlabel('Pseudo-Time', fontsize=yticks_fontsize*1.33)

        # Get x and y axis range
        ymin, ymax = ax.get_ylim()
        xmin, xmax = ax.get_xlim()

        # Custom x and y ticks
        x_custom = np.arange(xmin, xmax, step=1)
        y_custom = np.arange(ymax, ymin, step=1)

        # set minor ticks at custom locations:
        ax.set_xticks(x_custom, minor=True)
        ax.set_yticks(y_custom, minor=True)

        # Add grid lines at both major and minor ticks
        plt.grid(False, which='major')
        plt.grid(True, which='minor', color='black', linestyle='-', linewidth=1)

        # Remove ticks
        ax.tick_params(axis='both', which='minor', length=0)

        # Add colorbar
        if colorbar:
            cbar = plt.colorbar(cax, shrink=0.7)
            cbar.set_label(colorbar_title, rotation=270, labelpad=20, fontsize=16)

        # Add legend
        if legend:
            colors = np.append(plt.get_cmap(cmap_here)([0, 0.5, 1]), np.array([[1, 1, 1, 1]]), axis=0)
            labels = ['Downregulated', 'No change', 'Upregulated', 'No data']
            patches = [
                mpatches.Patch(facecolor=colors[i], label=labels[i], edgecolor='black')
                for i in range(len(labels))
            ]
            ax.legend(handles=patches, bbox_to_anchor=(1.05, 1),
                      loc=2, borderaxespad=0., title='Regulation',
                      fontsize=yticks_fontsize*2, title_fontsize=24)

        # Save figures
        if save_name:
            plt.savefig(output_dir / f'{save_name}.pdf', bbox_inches='tight')
            plt.savefig(output_dir / f'{save_name}.png', bbox_inches='tight', dpi=600)
            plt.savefig(output_dir / f'{save_name}.svg', bbox_inches='tight',
                       format='svg', transparent=True)
            logger.info(f"Saved heatmap to: {output_dir / save_name}.{{pdf,png,svg}}")

        if show:
            plt.show()
        else:
            plt.close()

        return fig

    # 1. Top 50 most changing pathways (first vs last timepoint)
    logger.info("Creating heatmap 1/5: Top 50 most changing pathways...")
    first_col = heatmap_data.columns[0]
    last_col = heatmap_data.columns[-1]
    change = (heatmap_data[last_col] - heatmap_data[first_col]).abs()
    top_changing = change.nlargest(50).index.tolist()

    df_top_changing = heatmap_data.loc[top_changing]
    fig1 = plot_heatmap_regulation(
        df_top_changing,
        top_changing,
        cmap_here='RdBu_r',
        save_name='top50_most_changing_pathways',
        colorbar_title='Sum of NES'
    )
    figures['top50_changing'] = fig1

    # 2. Top 50 most upregulated pathways (average NES > 0)
    logger.info("Creating heatmap 2/5: Top 50 most upregulated pathways...")
    avg_nes = heatmap_data.mean(axis=1)
    upregulated = avg_nes[avg_nes > 0].nlargest(50).index.tolist()

    df_upregulated = heatmap_data.loc[upregulated]
    fig2 = plot_heatmap_regulation(
        df_upregulated,
        upregulated,
        cmap_here='YlGn',
        save_name='top50_most_upregulated_pathways',
        colorbar_title='Sum of NES'
    )
    figures['top50_upregulated'] = fig2

    # 3. Top 50 most downregulated pathways (average NES < 0)
    logger.info("Creating heatmap 3/5: Top 50 most downregulated pathways...")
    downregulated = avg_nes[avg_nes < 0].nsmallest(50).index.tolist()

    df_downregulated = heatmap_data.loc[downregulated]
    fig3 = plot_heatmap_regulation(
        df_downregulated,
        downregulated,
        cmap_here='YlOrBr',
        save_name='top50_most_downregulated_pathways',
        colorbar_title='Sum of NES'
    )
    figures['top50_downregulated'] = fig3

    # 4. High-level pathways (29 pathways from Reactome highest level)
    logger.info("Creating heatmap 4/5: High-level pathways...")
    available_highest = [p for p in highest_pathways if p in heatmap_data.index]

    if available_highest:
        df_highest = heatmap_data.loc[available_highest]
        fig4 = plot_heatmap_regulation(
            df_highest,
            available_highest,
            cmap_here='RdBu_r',
            save_name='selected_pathways_highest_level',
            colorbar_title='Sum of NES'
        )
        figures['selected_highest_level'] = fig4
        logger.info(f"Found {len(available_highest)}/{len(highest_pathways)} high-level pathways in data")
    else:
        logger.warning("No high-level pathways found in the data")

    # 5. Literature pathways (33 pathways from literature review)
    logger.info("Creating heatmap 5/5: Literature pathways...")
    available_literature = [p for p in pathways_literature if p in heatmap_data.index]

    if available_literature:
        df_literature = heatmap_data.loc[available_literature]
        fig5 = plot_heatmap_regulation(
            df_literature,
            available_literature,
            cmap_here='RdBu_r',
            save_name='selected_pathways_literature',
            colorbar_title='Sum of NES'
        )
        figures['selected_literature'] = fig5
        logger.info(f"Found {len(available_literature)}/{len(pathways_literature)} literature pathways in data")
    else:
        logger.warning("No literature pathways found in the data")

    logger.info(f"Pathway heatmap generation complete. Created {len(figures)} heatmaps.")

    return heatmap_data, figures