gengeneeval 0.2.1__py3-none-any.whl → 0.4.0__py3-none-any.whl

geneval/deg/visualization.py

@@ -0,0 +1,376 @@
+"""
+Visualization module for DEG evaluation results.
+
+Provides:
+- Distribution plots for metrics across contexts
+- Heatmaps for context × metric results
+- Comprehensive DEG reports
+"""
+from __future__ import annotations
+
+from typing import Optional, List, Dict, Union, Any, TYPE_CHECKING
+from pathlib import Path
+import numpy as np
+
+if TYPE_CHECKING:
+    import matplotlib.pyplot as plt
+    import pandas as pd
+    from .evaluator import DEGEvaluationResult
+
+
+def _check_matplotlib():
+    """Check if matplotlib is available."""
+    try:
+        import matplotlib.pyplot as plt
+        return plt
+    except ImportError:
+        raise ImportError("matplotlib is required for visualization. Install with: pip install matplotlib")
+
+
+def _check_seaborn():
+    """Check if seaborn is available."""
+    try:
+        import seaborn as sns
+        return sns
+    except ImportError:
+        return None  # Seaborn is optional
+
+
+def plot_deg_distributions(
+    results: "DEGEvaluationResult",
+    metrics: Optional[List[str]] = None,
+    figsize: tuple = (12, 4),
+    save_path: Optional[Union[str, Path]] = None,
+    show: bool = True,
+) -> "plt.Figure":
+    """
+    Plot distribution of metrics across contexts.
+    
+    Creates violin/box plots showing metric distributions across all contexts.
+    
+    Parameters
+    ----------
+    results : DEGEvaluationResult
+        DEG evaluation results
+    metrics : List[str], optional
+        Metrics to plot. If None, plots all available.
+    figsize : tuple
+        Figure size per metric
+    save_path : str or Path, optional
+        Path to save figure
+    show : bool
+        Whether to display the figure
+        
+    Returns
+    -------
+    plt.Figure
+        Matplotlib figure
+    """
+    plt = _check_matplotlib()
+    sns = _check_seaborn()
+    
+    df = results.expanded_metrics
+    
+    if metrics is None:
+        # Get all numeric columns that are metrics
+        metrics = [c for c in df.columns if c in results.settings.get("metrics", [])]
+    
+    if len(metrics) == 0:
+        raise ValueError("No metrics to plot")
+    
+    n_metrics = len(metrics)
+    fig, axes = plt.subplots(1, n_metrics, figsize=(figsize[0], figsize[1]))
+    
+    if n_metrics == 1:
+        axes = [axes]
+    
+    for ax, metric in zip(axes, metrics):
+        values = df[metric].dropna()
+        
+        if sns is not None:
+            sns.violinplot(y=values, ax=ax, inner="box")
+        else:
+            ax.boxplot(values, vert=True)
+        
+        ax.set_title(metric)
+        ax.set_ylabel("Value")
+        
+        # Add mean annotation
+        mean_val = values.mean()
+        ax.axhline(mean_val, color='red', linestyle='--', alpha=0.5, label=f'Mean: {mean_val:.3f}')
+        ax.legend(loc='upper right', fontsize=8)
+    
+    plt.suptitle("Metric Distributions Across Contexts (DEGs only)", y=1.02)
+    plt.tight_layout()
+    
+    if save_path:
+        fig.savefig(save_path, dpi=150, bbox_inches='tight')
+    
+    if show:
+        plt.show()
+    
+    return fig
+
+
+def plot_context_heatmap(
+    results: "DEGEvaluationResult",
+    metrics: Optional[List[str]] = None,
+    figsize: tuple = (12, 8),
+    cmap: str = "RdYlBu_r",
+    save_path: Optional[Union[str, Path]] = None,
+    show: bool = True,
+) -> "plt.Figure":
+    """
+    Plot heatmap of metrics across contexts.
+    
+    Parameters
+    ----------
+    results : DEGEvaluationResult
+        DEG evaluation results
+    metrics : List[str], optional
+        Metrics to include
+    figsize : tuple
+        Figure size
+    cmap : str
+        Colormap name
+    save_path : str or Path, optional
+        Path to save figure
+    show : bool
+        Whether to display
+        
+    Returns
+    -------
+    plt.Figure
+        Matplotlib figure
+    """
+    plt = _check_matplotlib()
+    sns = _check_seaborn()
+    
+    df = results.expanded_metrics
+    
+    if metrics is None:
+        metrics = [c for c in df.columns if c in results.settings.get("metrics", [])]
+    
+    # Create matrix for heatmap
+    heatmap_data = df.set_index("context_id")[metrics]
+    
+    fig, ax = plt.subplots(figsize=figsize)
+    
+    if sns is not None:
+        sns.heatmap(
+            heatmap_data,
+            ax=ax,
+            cmap=cmap,
+            annot=True,
+            fmt=".3f",
+            cbar_kws={"label": "Metric Value"},
+        )
+    else:
+        im = ax.imshow(heatmap_data.values, cmap=cmap, aspect='auto')
+        plt.colorbar(im, ax=ax, label="Metric Value")
+        ax.set_xticks(range(len(metrics)))
+        ax.set_xticklabels(metrics, rotation=45, ha='right')
+        ax.set_yticks(range(len(heatmap_data)))
+        ax.set_yticklabels(heatmap_data.index)
+    
+    ax.set_title("Metrics per Context (DEGs only)")
+    ax.set_xlabel("Metric")
+    ax.set_ylabel("Context")
+    
+    plt.tight_layout()
+    
+    if save_path:
+        fig.savefig(save_path, dpi=150, bbox_inches='tight')
+    
+    if show:
+        plt.show()
+    
+    return fig
+
+
+def plot_deg_counts(
+    results: "DEGEvaluationResult",
+    figsize: tuple = (10, 6),
+    save_path: Optional[Union[str, Path]] = None,
+    show: bool = True,
+) -> "plt.Figure":
+    """
+    Plot DEG counts per context.
+    
+    Parameters
+    ----------
+    results : DEGEvaluationResult
+        DEG evaluation results
+    figsize : tuple
+        Figure size
+    save_path : str or Path, optional
+        Path to save figure
+    show : bool
+        Whether to display
+        
+    Returns
+    -------
+    plt.Figure
+        Matplotlib figure
+    """
+    plt = _check_matplotlib()
+    
+    df = results.deg_summary
+    
+    if len(df) == 0:
+        raise ValueError("No DEG data to plot")
+    
+    fig, ax = plt.subplots(figsize=figsize)
+    
+    x = range(len(df))
+    width = 0.35
+    
+    if "n_upregulated" in df.columns and "n_downregulated" in df.columns:
+        ax.bar(x, df["n_upregulated"], width, label='Upregulated', color='red', alpha=0.7)
+        ax.bar(x, -df["n_downregulated"], width, label='Downregulated', color='blue', alpha=0.7)
+        ax.axhline(0, color='black', linewidth=0.5)
+        ax.set_ylabel("Number of DEGs")
+    else:
+        ax.bar(x, df["n_degs"], color='purple', alpha=0.7)
+        ax.set_ylabel("Number of DEGs")
+    
+    ax.set_xlabel("Context")
+    ax.set_xticks(x)
+    ax.set_xticklabels(df["context_id"], rotation=45, ha='right')
+    ax.set_title("DEG Counts per Context")
+    ax.legend()
+    
+    plt.tight_layout()
+    
+    if save_path:
+        fig.savefig(save_path, dpi=150, bbox_inches='tight')
+    
+    if show:
+        plt.show()
+    
+    return fig
+
+
+def create_deg_report(
+    results: "DEGEvaluationResult",
+    output_dir: Union[str, Path],
+    include_plots: bool = True,
+) -> None:
+    """
+    Create a comprehensive DEG evaluation report.
+    
+    Creates:
+    - Summary statistics (CSV)
+    - Aggregated metrics (CSV)
+    - Expanded per-context metrics (CSV)
+    - DEG summary per context (CSV)
+    - Visualization plots (PNG, if include_plots=True)
+    - Markdown report
+    
+    Parameters
+    ----------
+    results : DEGEvaluationResult
+        DEG evaluation results
+    output_dir : str or Path
+        Output directory
+    include_plots : bool
+        Whether to generate plots
+    """
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Save CSVs
+    results.save(output_dir)
+    
+    # Generate plots
+    if include_plots:
+        try:
+            plt = _check_matplotlib()
+            
+            # Distribution plot
+            if len(results.expanded_metrics) > 0:
+                plot_deg_distributions(
+                    results,
+                    save_path=output_dir / "deg_metric_distributions.png",
+                    show=False,
+                )
+                plt.close()
+                
+                # Heatmap
+                plot_context_heatmap(
+                    results,
+                    save_path=output_dir / "deg_context_heatmap.png",
+                    show=False,
+                )
+                plt.close()
+            
+            # DEG counts
+            if len(results.deg_summary) > 0:
+                plot_deg_counts(
+                    results,
+                    save_path=output_dir / "deg_counts.png",
+                    show=False,
+                )
+                plt.close()
+                
+        except ImportError:
+            pass  # Skip plots if matplotlib not available
+    
+    # Create markdown report
+    report_lines = [
+        "# DEG Evaluation Report",
+        "",
+        "## Summary",
+        "",
+        f"- **Number of contexts evaluated**: {len(results.context_results)}",
+        f"- **DEG detection method**: {results.settings.get('deg_method', 'N/A')}",
+        f"- **P-value threshold**: {results.settings.get('pval_threshold', 'N/A')}",
+        f"- **Log fold change threshold**: {results.settings.get('lfc_threshold', 'N/A')}",
+        f"- **Compute device**: {results.settings.get('device', 'N/A')}",
+        "",
+        "## Aggregated Metrics",
+        "",
+    ]
+    
+    if len(results.aggregated_metrics) > 0:
+        for col in results.aggregated_metrics.columns:
+            val = results.aggregated_metrics[col].iloc[0]
+            if isinstance(val, float):
+                report_lines.append(f"- **{col}**: {val:.4f}")
+            else:
+                report_lines.append(f"- **{col}**: {val}")
+    else:
+        report_lines.append("No aggregated metrics available.")
+    
+    report_lines.extend([
+        "",
+        "## DEG Summary",
+        "",
+    ])
+    
+    if len(results.deg_summary) > 0:
+        report_lines.append(results.deg_summary.to_markdown(index=False))
+    else:
+        report_lines.append("No DEG data available.")
+    
+    report_lines.extend([
+        "",
+        "## Files Generated",
+        "",
+        "- `deg_aggregated_metrics.csv`: Summary metrics across all contexts",
+        "- `deg_expanded_metrics.csv`: Per-context metrics",
+        "- `deg_summary.csv`: DEG detection summary",
+        "- `deg_per_context/`: Individual DEG results per context",
+    ])
+    
+    if include_plots:
+        report_lines.extend([
+            "- `deg_metric_distributions.png`: Distribution plots",
+            "- `deg_context_heatmap.png`: Context × metric heatmap",
+            "- `deg_counts.png`: DEG counts per context",
+        ])
+    
+    with open(output_dir / "DEG_REPORT.md", "w") as f:
+        f.write("\n".join(report_lines))
+    
+    print(f"DEG report saved to {output_dir}")

geneval/evaluator.py

@@ -66,6 +66,10 @@ class GeneEvalEvaluator:
         Whether to include multivariate (whole-space) metrics
     verbose : bool
         Whether to print progress
+    n_jobs : int
+        Number of parallel CPU jobs. -1 uses all cores. Default is 1.
+    device : str
+        Compute device: "cpu", "cuda", "cuda:0", "auto". Default is "cpu".
         
     Examples
     --------
@@ -73,6 +77,10 @@ class GeneEvalEvaluator:
     >>> evaluator = GeneEvalEvaluator(loader)
     >>> results = evaluator.evaluate()
     >>> results.save("output/")
+    
+    >>> # With acceleration
+    >>> evaluator = GeneEvalEvaluator(loader, n_jobs=8, device="cuda")
+    >>> results = evaluator.evaluate()
     """
     
     def __init__(
@@ -82,11 +90,15 @@ class GeneEvalEvaluator:
         aggregate_method: str = "mean",
         include_multivariate: bool = True,
         verbose: bool = True,
+        n_jobs: int = 1,
+        device: str = "cpu",
     ):
         self.data_loader = data_loader
         self.aggregate_method = aggregate_method
         self.include_multivariate = include_multivariate
         self.verbose = verbose
+        self.n_jobs = n_jobs
+        self.device = device
         
         # Initialize metrics
         self.metrics: List[BaseMetric] = []
@@ -106,6 +118,25 @@ class GeneEvalEvaluator:
                 MultivariateWasserstein(),
                 MultivariateMMD(),
             ])
+        
+        # Initialize accelerated computer if using parallelization or GPU
+        self._parallel_computer = None
+        if n_jobs != 1 or device != "cpu":
+            try:
+                from .metrics.accelerated import ParallelMetricComputer
+                self._parallel_computer = ParallelMetricComputer(
+                    n_jobs=n_jobs,
+                    device=device,
+                    verbose=verbose,
+                )
+                if verbose:
+                    from .metrics.accelerated import get_available_backends
+                    backends = get_available_backends()
+                    self._log(f"Acceleration enabled: n_jobs={n_jobs}, device={device}")
+                    self._log(f"Available backends: {backends}")
+            except ImportError as e:
+                if verbose:
+                    self._log(f"Warning: Could not enable acceleration: {e}")
     
     def _log(self, msg: str):
         """Print message if verbose."""
@@ -262,6 +293,8 @@ def evaluate(
     metrics: Optional[List[Union[BaseMetric, Type[BaseMetric]]]] = None,
     include_multivariate: bool = True,
     verbose: bool = True,
+    n_jobs: int = 1,
+    device: str = "cpu",
     **loader_kwargs
 ) -> EvaluationResult:
     """
@@ -285,6 +318,10 @@ def evaluate(
         Whether to include multivariate metrics
     verbose : bool
         Print progress
+    n_jobs : int
+        Number of parallel CPU jobs. -1 uses all cores. Default is 1.
+    device : str
+        Compute device: "cpu", "cuda", "cuda:0", "auto". Default is "cpu".
     **loader_kwargs
         Additional arguments for data loader
         
@@ -295,6 +332,7 @@ def evaluate(
         
     Examples
     --------
+    >>> # Standard CPU evaluation
     >>> results = evaluate(
     ...     "real.h5ad",
     ...     "generated.h5ad",
@@ -302,6 +340,12 @@ def evaluate(
     ...     split_column="split",
     ...     output_dir="evaluation_output/"
     ... )
+    
+    >>> # Parallel CPU evaluation (8 cores)
+    >>> results = evaluate(..., n_jobs=8)
+    
+    >>> # GPU-accelerated evaluation
+    >>> results = evaluate(..., device="cuda")
     """
     # Load data
     loader = load_data(
@@ -318,6 +362,8 @@ def evaluate(
         metrics=metrics,
         include_multivariate=include_multivariate,
         verbose=verbose,
+        n_jobs=n_jobs,
+        device=device,
     )
     
     # Run evaluation

geneval/metrics/__init__.py

@@ -35,6 +35,20 @@ from .reconstruction import (
     R2Score,
 )
 
+# Accelerated computation
+from .accelerated import (
+    AccelerationConfig,
+    ParallelMetricComputer,
+    get_available_backends,
+    compute_metrics_accelerated,
+    GPUWasserstein1,
+    GPUWasserstein2,
+    GPUMMD,
+    GPUEnergyDistance,
+    vectorized_wasserstein1,
+    vectorized_mmd,
+)
+
 # All available metrics
 ALL_METRICS = [
     # Reconstruction
@@ -81,4 +95,15 @@ __all__ = [
     "MultivariateMMD",
     # Collections
     "ALL_METRICS",
+    # Acceleration
+    "AccelerationConfig",
+    "ParallelMetricComputer",
+    "get_available_backends",
+    "compute_metrics_accelerated",
+    "GPUWasserstein1",
+    "GPUWasserstein2",
+    "GPUMMD",
+    "GPUEnergyDistance",
+    "vectorized_wasserstein1",
+    "vectorized_mmd",
 ]