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

geneval/deg/evaluator.py

@@ -0,0 +1,538 @@
+"""
+DEG-focused evaluator for GenGeneEval.
+
+Computes metrics only on differentially expressed genes, with support for:
+- Per-context evaluation (covariates × perturbations)
+- Fast DEG detection with GPU acceleration
+- Aggregated and expanded result reporting
+"""
+from __future__ import annotations
+
+from typing import Optional, List, Dict, Union, Any, Literal
+from dataclasses import dataclass, field
+from pathlib import Path
+import numpy as np
+import pandas as pd
+import warnings
+
+from .detection import (
+    compute_degs_fast,
+    compute_degs_gpu,
+    compute_degs_auto,
+    DEGResult,
+    DEGMethod,
+)
+from .context import (
+    ContextEvaluator,
+    ContextResult,
+    get_context_id,
+    get_contexts,
+)
+
+# Import metrics
+from ..metrics.base_metric import BaseMetric
+from ..metrics.correlation import PearsonCorrelation, SpearmanCorrelation
+from ..metrics.distances import (
+    Wasserstein1Distance,
+    Wasserstein2Distance,
+    MMDDistance,
+    EnergyDistance,
+)
+from ..metrics.accelerated import (
+    get_available_backends,
+    vectorized_wasserstein1,
+    vectorized_mmd,
+)
+
+
+@dataclass
+class DEGEvaluationResult:
+    """Complete DEG evaluation results.
+    
+    Attributes
+    ----------
+    context_results : List[ContextResult]
+        Results for each context
+    aggregated_metrics : pd.DataFrame
+        Aggregated metrics across contexts
+    expanded_metrics : pd.DataFrame
+        Per-context expanded metrics
+    deg_summary : pd.DataFrame
+        Summary of DEG detection per context
+    gene_names : np.ndarray
+        All gene names
+    settings : Dict
+        Evaluation settings
+    """
+    context_results: List[ContextResult]
+    aggregated_metrics: pd.DataFrame
+    expanded_metrics: pd.DataFrame
+    deg_summary: pd.DataFrame
+    gene_names: np.ndarray
+    settings: Dict[str, Any]
+    
+    def save(self, output_dir: Union[str, Path]) -> None:
+        """Save results to directory."""
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+        
+        self.aggregated_metrics.to_csv(output_dir / "deg_aggregated_metrics.csv")
+        self.expanded_metrics.to_csv(output_dir / "deg_expanded_metrics.csv")
+        self.deg_summary.to_csv(output_dir / "deg_summary.csv")
+        
+        # Save per-context DEG results
+        deg_dir = output_dir / "deg_per_context"
+        deg_dir.mkdir(exist_ok=True)
+        for ctx_result in self.context_results:
+            if ctx_result.deg_result is not None:
+                ctx_result.deg_result.to_dataframe().to_csv(
+                    deg_dir / f"{ctx_result.context_id}_degs.csv"
+                )
+    
+    def __repr__(self) -> str:
+        return (
+            f"DEGEvaluationResult(n_contexts={len(self.context_results)}, "
+            f"metrics={list(self.aggregated_metrics.columns)})"
+        )
+
+
+class DEGEvaluator:
+    """
+    Evaluator that computes metrics on DEGs only.
+    
+    This evaluator:
+    1. Detects DEGs for each perturbation context
+    2. Computes distributional metrics only on DEG genes
+    3. Reports per-context and aggregated results
+    
+    Parameters
+    ----------
+    real_data : np.ndarray
+        Real expression matrix (n_samples, n_genes)
+    generated_data : np.ndarray
+        Generated expression matrix (n_samples, n_genes)
+    real_obs : pd.DataFrame
+        Real data observation metadata
+    generated_obs : pd.DataFrame
+        Generated data observation metadata
+    condition_columns : List[str]
+        Columns defining contexts (e.g., ["cell_type", "perturbation"])
+    gene_names : np.ndarray, optional
+        Gene names
+    control_key : str
+        Value indicating control samples (default: "control")
+    perturbation_column : str, optional
+        Column containing perturbation info. If None, uses first condition column.
+    deg_method : str
+        DEG detection method: "welch", "student", "wilcoxon", "logfc"
+    pval_threshold : float
+        P-value threshold for DEG significance
+    lfc_threshold : float
+        Log2 fold change threshold
+    min_degs : int
+        Minimum DEGs required to compute metrics
+    metrics : List[str], optional
+        Metrics to compute. Default: all supported metrics.
+    n_jobs : int
+        Number of parallel CPU jobs
+    device : str
+        Compute device: "cpu", "cuda", "mps", "auto"
+    verbose : bool
+        Print progress
+        
+    Examples
+    --------
+    >>> evaluator = DEGEvaluator(
+    ...     real_data, generated_data,
+    ...     real_obs, generated_obs,
+    ...     condition_columns=["perturbation"],
+    ...     deg_method="welch",
+    ...     device="cuda",
+    ... )
+    >>> results = evaluator.evaluate()
+    >>> results.save("output/")
+    """
+    
+    # Supported metrics
+    SUPPORTED_METRICS = [
+        "wasserstein_1",
+        "wasserstein_2", 
+        "mmd",
+        "energy",
+        "pearson",
+        "spearman",
+    ]
+    
+    def __init__(
+        self,
+        real_data: np.ndarray,
+        generated_data: np.ndarray,
+        real_obs: pd.DataFrame,
+        generated_obs: pd.DataFrame,
+        condition_columns: List[str],
+        gene_names: Optional[np.ndarray] = None,
+        control_key: str = "control",
+        perturbation_column: Optional[str] = None,
+        deg_method: DEGMethod = "welch",
+        pval_threshold: float = 0.05,
+        lfc_threshold: float = 0.5,
+        min_degs: int = 5,
+        metrics: Optional[List[str]] = None,
+        n_jobs: int = 1,
+        device: str = "cpu",
+        verbose: bool = True,
+    ):
+        self.real_data = np.asarray(real_data, dtype=np.float32)
+        self.generated_data = np.asarray(generated_data, dtype=np.float32)
+        self.real_obs = real_obs.reset_index(drop=True)
+        self.generated_obs = generated_obs.reset_index(drop=True)
+        self.condition_columns = condition_columns
+        self.gene_names = gene_names if gene_names is not None else np.array(
+            [f"Gene_{i}" for i in range(real_data.shape[1])]
+        )
+        self.control_key = control_key
+        self.perturbation_column = perturbation_column or condition_columns[0]
+        self.deg_method = deg_method
+        self.pval_threshold = pval_threshold
+        self.lfc_threshold = lfc_threshold
+        self.min_degs = min_degs
+        self.metrics = metrics or self.SUPPORTED_METRICS
+        self.n_jobs = n_jobs
+        self.device = device
+        self.verbose = verbose
+        
+        # Create context evaluator
+        self.context_evaluator = ContextEvaluator(
+            real_data=self.real_data,
+            generated_data=self.generated_data,
+            real_obs=self.real_obs,
+            generated_obs=self.generated_obs,
+            condition_columns=condition_columns,
+            gene_names=self.gene_names,
+            control_key=control_key,
+            perturbation_column=self.perturbation_column,
+        )
+        
+        # Initialize metric objects
+        self._metric_objects = {
+            "wasserstein_1": Wasserstein1Distance(),
+            "wasserstein_2": Wasserstein2Distance(),
+            "mmd": MMDDistance(),
+            "energy": EnergyDistance(),
+            "pearson": PearsonCorrelation(),
+            "spearman": SpearmanCorrelation(),
+        }
+        
+        self._log(f"DEGEvaluator initialized with {len(self.context_evaluator)} contexts")
+        self._log(f"Perturbation contexts: {len(self.context_evaluator.get_perturbation_contexts())}")
+    
+    def _log(self, msg: str) -> None:
+        """Print if verbose."""
+        if self.verbose:
+            print(msg)
+    
+    def _compute_degs(
+        self,
+        control: np.ndarray,
+        perturbed: np.ndarray,
+    ) -> DEGResult:
+        """Compute DEGs using configured method and device."""
+        return compute_degs_auto(
+            control=control,
+            perturbed=perturbed,
+            gene_names=self.gene_names,
+            method=self.deg_method,
+            pval_threshold=self.pval_threshold,
+            lfc_threshold=self.lfc_threshold,
+            n_jobs=self.n_jobs,
+            device=self.device,
+        )
+    
+    def _compute_metrics_on_degs(
+        self,
+        real: np.ndarray,
+        generated: np.ndarray,
+        deg_indices: np.ndarray,
+    ) -> Dict[str, float]:
+        """Compute metrics on DEG genes only."""
+        if len(deg_indices) < self.min_degs:
+            return {m: np.nan for m in self.metrics}
+        
+        # Slice to DEGs only
+        real_degs = real[:, deg_indices]
+        gen_degs = generated[:, deg_indices]
+        
+        results = {}
+        
+        for metric_name in self.metrics:
+            if metric_name not in self._metric_objects:
+                continue
+            
+            metric = self._metric_objects[metric_name]
+            
+            try:
+                # Compute per-gene and aggregate
+                per_gene = metric.compute_per_gene(real_degs, gen_degs)
+                results[metric_name] = float(np.nanmean(per_gene))
+            except Exception as e:
+                if self.verbose:
+                    self._log(f"Warning: {metric_name} failed: {e}")
+                results[metric_name] = np.nan
+        
+        return results
+    
+    def _compute_metrics_accelerated(
+        self,
+        real: np.ndarray,
+        generated: np.ndarray,
+        deg_indices: np.ndarray,
+    ) -> Dict[str, float]:
+        """Compute metrics using accelerated implementations."""
+        if len(deg_indices) < self.min_degs:
+            return {m: np.nan for m in self.metrics}
+        
+        # Slice to DEGs only
+        real_degs = real[:, deg_indices]
+        gen_degs = generated[:, deg_indices]
+        
+        results = {}
+        backends = get_available_backends()
+        
+        # Use vectorized implementations where available
+        if "wasserstein_1" in self.metrics:
+            try:
+                w1_per_gene = vectorized_wasserstein1(real_degs, gen_degs)
+                results["wasserstein_1"] = float(np.nanmean(w1_per_gene))
+            except Exception:
+                results["wasserstein_1"] = np.nan
+        
+        if "mmd" in self.metrics:
+            try:
+                mmd_per_gene = vectorized_mmd(real_degs, gen_degs)
+                results["mmd"] = float(np.nanmean(mmd_per_gene))
+            except Exception:
+                results["mmd"] = np.nan
+        
+        # Fall back to standard computation for other metrics
+        for metric_name in self.metrics:
+            if metric_name in results:
+                continue
+            if metric_name not in self._metric_objects:
+                continue
+            
+            metric = self._metric_objects[metric_name]
+            try:
+                per_gene = metric.compute_per_gene(real_degs, gen_degs)
+                results[metric_name] = float(np.nanmean(per_gene))
+            except Exception:
+                results[metric_name] = np.nan
+        
+        return results
+    
+    def evaluate(self) -> DEGEvaluationResult:
+        """
+        Run DEG-focused evaluation on all contexts.
+        
+        Returns
+        -------
+        DEGEvaluationResult
+            Complete evaluation results with per-context and aggregated metrics.
+        """
+        context_results = []
+        
+        perturbation_contexts = self.context_evaluator.get_perturbation_contexts()
+        n_contexts = len(perturbation_contexts)
+        
+        self._log(f"Evaluating {n_contexts} perturbation contexts...")
+        
+        for i, context in enumerate(perturbation_contexts):
+            context_id = get_context_id(context)
+            
+            if self.verbose:
+                print(f"  [{i+1}/{n_contexts}] {context_id}", end="... ")
+            
+            try:
+                # Get perturbed data
+                real_pert, gen_pert = self.context_evaluator.get_context_data(context)
+                
+                # Get control data
+                real_ctrl, gen_ctrl = self.context_evaluator.get_control_data(context)
+                
+                if len(real_ctrl) < 2 or len(real_pert) < 2:
+                    if self.verbose:
+                        print("skipped (insufficient samples)")
+                    continue
+                
+                # Compute DEGs using real data (control vs perturbed)
+                deg_result = self._compute_degs(real_ctrl, real_pert)
+                
+                if self.verbose:
+                    print(f"{deg_result.n_degs} DEGs", end="... ")
+                
+                # Compute metrics on DEGs
+                metrics = self._compute_metrics_accelerated(
+                    real_pert, gen_pert, deg_result.deg_indices
+                )
+                
+                ctx_result = ContextResult(
+                    context_id=context_id,
+                    context_values=context,
+                    n_samples_real=len(real_pert),
+                    n_samples_gen=len(gen_pert),
+                    deg_result=deg_result,
+                    metrics=metrics,
+                )
+                context_results.append(ctx_result)
+                
+                if self.verbose:
+                    print("done")
+                    
+            except Exception as e:
+                if self.verbose:
+                    print(f"error: {e}")
+                continue
+        
+        # Build result DataFrames
+        expanded_data = []
+        for ctx_result in context_results:
+            row = {
+                "context_id": ctx_result.context_id,
+                **ctx_result.context_values,
+                "n_samples_real": ctx_result.n_samples_real,
+                "n_samples_gen": ctx_result.n_samples_gen,
+                "n_degs": ctx_result.deg_result.n_degs if ctx_result.deg_result else 0,
+                **ctx_result.metrics,
+            }
+            expanded_data.append(row)
+        
+        expanded_metrics = pd.DataFrame(expanded_data)
+        
+        # Aggregated metrics
+        if len(expanded_metrics) > 0:
+            agg_data = {
+                "n_contexts": len(context_results),
+                "total_samples_real": expanded_metrics["n_samples_real"].sum(),
+                "total_samples_gen": expanded_metrics["n_samples_gen"].sum(),
+                "mean_n_degs": expanded_metrics["n_degs"].mean(),
+                "median_n_degs": expanded_metrics["n_degs"].median(),
+            }
+            for metric in self.metrics:
+                if metric in expanded_metrics.columns:
+                    agg_data[f"{metric}_mean"] = expanded_metrics[metric].mean()
+                    agg_data[f"{metric}_std"] = expanded_metrics[metric].std()
+                    agg_data[f"{metric}_median"] = expanded_metrics[metric].median()
+            
+            aggregated_metrics = pd.DataFrame([agg_data])
+        else:
+            aggregated_metrics = pd.DataFrame()
+        
+        # DEG summary
+        deg_summary_data = []
+        for ctx_result in context_results:
+            if ctx_result.deg_result is not None:
+                deg_summary_data.append({
+                    "context_id": ctx_result.context_id,
+                    **ctx_result.context_values,
+                    "n_degs": ctx_result.deg_result.n_degs,
+                    "n_upregulated": (ctx_result.deg_result.log_fold_changes[ctx_result.deg_result.is_deg] > 0).sum(),
+                    "n_downregulated": (ctx_result.deg_result.log_fold_changes[ctx_result.deg_result.is_deg] < 0).sum(),
+                    "mean_abs_lfc": np.abs(ctx_result.deg_result.log_fold_changes[ctx_result.deg_result.is_deg]).mean() if ctx_result.deg_result.n_degs > 0 else np.nan,
+                })
+        deg_summary = pd.DataFrame(deg_summary_data)
+        
+        self._log(f"\nEvaluation complete: {len(context_results)} contexts evaluated")
+        
+        return DEGEvaluationResult(
+            context_results=context_results,
+            aggregated_metrics=aggregated_metrics,
+            expanded_metrics=expanded_metrics,
+            deg_summary=deg_summary,
+            gene_names=self.gene_names,
+            settings={
+                "deg_method": self.deg_method,
+                "pval_threshold": self.pval_threshold,
+                "lfc_threshold": self.lfc_threshold,
+                "min_degs": self.min_degs,
+                "metrics": self.metrics,
+                "device": self.device,
+                "n_jobs": self.n_jobs,
+            },
+        )
+
+
+def evaluate_degs(
+    real_data: np.ndarray,
+    generated_data: np.ndarray,
+    real_obs: pd.DataFrame,
+    generated_obs: pd.DataFrame,
+    condition_columns: List[str],
+    gene_names: Optional[np.ndarray] = None,
+    control_key: str = "control",
+    perturbation_column: Optional[str] = None,
+    deg_method: DEGMethod = "welch",
+    pval_threshold: float = 0.05,
+    lfc_threshold: float = 0.5,
+    metrics: Optional[List[str]] = None,
+    n_jobs: int = 1,
+    device: str = "auto",
+    verbose: bool = True,
+) -> DEGEvaluationResult:
+    """
+    Convenience function for DEG-focused evaluation.
+    
+    Parameters
+    ----------
+    real_data : np.ndarray
+        Real expression matrix
+    generated_data : np.ndarray
+        Generated expression matrix
+    real_obs : pd.DataFrame
+        Real data metadata
+    generated_obs : pd.DataFrame
+        Generated data metadata
+    condition_columns : List[str]
+        Columns defining contexts
+    gene_names : np.ndarray, optional
+        Gene names
+    control_key : str
+        Control condition identifier
+    perturbation_column : str, optional
+        Column containing perturbation info. If None, uses first condition column.
+    deg_method : str
+        DEG detection method
+    pval_threshold : float
+        P-value threshold
+    lfc_threshold : float
+        Log fold change threshold
+    metrics : List[str], optional
+        Metrics to compute
+    n_jobs : int
+        Parallel CPU jobs
+    device : str
+        Compute device
+    verbose : bool
+        Print progress
+        
+    Returns
+    -------
+    DEGEvaluationResult
+        Evaluation results
+    """
+    evaluator = DEGEvaluator(
+        real_data=real_data,
+        generated_data=generated_data,
+        real_obs=real_obs,
+        generated_obs=generated_obs,
+        condition_columns=condition_columns,
+        gene_names=gene_names,
+        control_key=control_key,
+        perturbation_column=perturbation_column,
+        deg_method=deg_method,
+        pval_threshold=pval_threshold,
+        lfc_threshold=lfc_threshold,
+        metrics=metrics,
+        n_jobs=n_jobs,
+        device=device,
+        verbose=verbose,
+    )
+    return evaluator.evaluate()