gengeneeval 0.1.0__py3-none-any.whl

geneval/evaluators/gene_expression_evaluator.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional, Dict, Any
+
+import numpy as np
+import pandas as pd
+from anndata import AnnData
+import scipy.stats as sstats
+from sklearn.metrics import mean_squared_error
+
+from ..metrics.metrics import compute_metrics
+from ..utils.preprocessing import to_dense as _to_dense
+from .base_evaluator import BaseEvaluator
+from ..visualization import EvaluationPlotter
+
+if TYPE_CHECKING:
+    from ..data.gene_expression_datamodule import GeneExpressionDataModule
+
+
+class GeneExpressionEvaluator(BaseEvaluator):
+    """
+    Evaluator for gene expression data.
+    """
+
+    def __init__(self, data: "GeneExpressionDataModule", output: AnnData):
+        super().__init__(data, output)
+
+    def evaluate(
+        self,
+        delta: bool = False,
+        plot: bool = False,
+        DEG: Optional[Dict[str, Any]] = None,
+        save_dir: Optional[str] = None,
+        max_panels: int = 12,
+        dpi: int = 150,
+    ):
+        """
+        Run evaluation. If plot=True, returns and optionally saves figures.
+        """
+        data = self.data.gene_expression_dataset.adata.copy()
+        generated = self.output.copy()
+        data, generated = self._align_varnames_like(data, generated)
+
+        pert_col = self.data.perturbation_key
+        split_key = self.data.split_key
+        control = self.data.control
+
+        order_cols = []
+        if "cell_type" in data.obs.columns and "cell_type" in generated.obs.columns:
+            order_cols.append("cell_type")
+        for c in (getattr(self.data, "condition_keys", None) or []):
+            if c in data.obs.columns and c in generated.obs.columns:
+                order_cols.append(c)
+
+        # Baseline handling
+        if delta:
+            b = self._compute_control_means(data, pert_col, control, strata_cols=order_cols)
+            data.X = self._apply_baseline_per_strata(data.X, data.obs, b, strata_cols=order_cols, mode="subtract")
+        else:
+            b = self._compute_control_means(data, pert_col, control, strata_cols=order_cols)
+            generated.X = self._apply_baseline_per_strata(
+                generated.X, generated.obs, b, strata_cols=order_cols, mode="add"
+            )
+
+        is_test = (data.obs[split_key].astype(str) == "test").to_numpy()
+        test_data = data[is_test].copy()
+
+        if "perturbation" not in generated.obs.columns and pert_col not in generated.obs.columns:
+            raise KeyError("'perturbation' column not found in generated data.")
+        if pert_col not in generated.obs.columns and "perturbation" in generated.obs.columns:
+            generated.obs[pert_col] = generated.obs["perturbation"].astype(test_data.obs[pert_col].dtype)
+
+        def _means_masks(adata, cols):
+            means, masks = {}, {}
+            df = adata.obs[[pert_col] + cols].astype(str)
+            for _, row in df.drop_duplicates().iterrows():
+                pert = row[pert_col]
+                key = "####".join([pert] + [row[c] for c in cols])
+                mask = (adata.obs[pert_col].astype(str) == pert).to_numpy()
+                for c in cols:
+                    mask &= (adata.obs[c].astype(str) == str(row[c])).to_numpy()
+                if mask.any():
+                    masks[key] = mask
+                    means[key] = _to_dense(adata[mask].X).mean(axis=0)
+            return means, masks
+
+        real_means, real_masks = _means_masks(test_data, order_cols)
+        gen_means, gen_masks = _means_masks(generated, order_cols)
+        common = sorted(set(real_means).intersection(gen_means))
+        if not common:
+            raise ValueError("No common (pert + covariates) between real TEST and generated.")
+
+        # Metric accumulators
+        w1 = []; w2 = []; mmd = []; energy = []
+        pearson_corr = []; pearson_p = []
+        spearman_corr = []; spearman_p = []
+        mse_val = []
+
+        vnames = pd.Index(test_data.var_names.astype(str))
+
+        # For plotting
+        plot_means = {}
+        residuals_per_key = {}
+        stats_per_key = {}
+        deg_map = {}
+
+        def maybe_filter(om, gm, td, gd, key):
+            if DEG is None:
+                return om, gm, td, gd
+            deg = DEG.get(key) or DEG.get(key.split("####", 1)[0])
+            if deg is None:
+                return om, gm, td, gd
+            names = None
+            if isinstance(deg, dict):
+                names = deg.get("names", None)
+            elif hasattr(deg, "columns") and "names" in deg.columns:
+                names = deg["names"]
+            else:
+                names = deg
+            if hasattr(names, "tolist"):
+                names = names.tolist()
+            if not names:
+                return om, gm, td, gd
+            mask = np.asarray(vnames.isin([str(x) for x in names]), dtype=bool)
+            if not mask.any():
+                return om, gm, td, gd
+            return om[mask], gm[mask], td[:, mask], gd[:, mask]
+
+        for key in common:
+            td = _to_dense(test_data.X[real_masks[key], :])
+            gd = _to_dense(generated.X[gen_masks[key], :])
+            om = real_means[key]; gm = gen_means[key]
+            om_f, gm_f, td_f, gd_f = maybe_filter(om, gm, td, gd, key)
+
+            # distributional metrics
+            w1.append({key: compute_metrics(td_f, gd_f, 'w1')})
+            w2.append({key: compute_metrics(td_f, gd_f, 'w2')})
+            mmd.append({key: compute_metrics(td_f, gd_f, 'mmd')})
+            energy.append({key: compute_metrics(td_f, gd_f, 'energy')})
+
+            # mean-wise metrics
+            pc, pcp = sstats.pearsonr(om_f, gm_f)
+            sc, scp = sstats.spearmanr(om_f, gm_f)
+            pearson_corr.append({key: pc}); pearson_p.append({key: pcp})
+            spearman_corr.append({key: sc}); spearman_p.append({key: scp})
+            mse = mean_squared_error(om_f, gm_f)
+            mse_val.append({key: mse})
+
+            # for plots
+            plot_means[key] = (om, gm, vnames.tolist())
+            residuals_per_key[key] = (gm - om)
+            stats_per_key[key] = {"pearson": float(pc), "spearman": float(sc), "mse": float(mse)}
+            if DEG is not None:
+                deg_map[key] = DEG.get(key) or DEG.get(key.split("####", 1)[0])
+
+        def _m(lst):
+            return float("nan") if not lst else float(np.mean([list(d.values())[0] for d in lst]))
+
+        print(f"Mean Pearson: {_m(pearson_corr):.4f} (p={_m(pearson_p):.4g})")
+        print(f"Mean Spearman: {_m(spearman_corr):.4f} (p={_m(spearman_p):.4g})")
+        print(f"Mean MSE: {_m(mse_val):.4f}")
+        print(f"Wasserstein-1: {_m(w1):.4f}")
+        print(f"Wasserstein-2: {_m(w2):.4f}")
+        print(f"MMD: {_m(mmd):.4f}")
+        print(f"Energy: {_m(energy):.4f}")
+
+        results = dict(
+            pearson_corr=pearson_corr,
+            spearman_corr=spearman_corr,
+            mse_val=mse_val,
+            w1=w1,
+            w2=w2,
+            mmd=mmd,
+            energy=energy,
+        )
+
+        # Plotting
+        figures = {}
+        if plot:
+            plotter = EvaluationPlotter()
+            # scatter grid
+            fig_scatter = plotter.scatter_means_grid(
+                data=plot_means,
+                stats=stats_per_key,
+                deg_map=deg_map if deg_map else None,
+                max_panels=max_panels,
+            )
+            figures["scatter_means"] = fig_scatter
+
+            # residual distributions
+            fig_residuals = plotter.residuals_violin(residuals=residuals_per_key)
+            figures["residuals"] = fig_residuals
+
+            # metrics bar: combine main metrics
+            metrics_pk = {}
+            for k in common:
+                metrics_pk[k] = {
+                    "pearson": stats_per_key[k]["pearson"],
+                    "spearman": stats_per_key[k]["spearman"],
+                    "MSE": stats_per_key[k]["mse"],
+                    "W1": float([d[k] for d in w1 if k in d][0]),
+                    "W2": float([d[k] for d in w2 if k in d][0]),
+                    "MMD": float([d[k] for d in mmd if k in d][0]),
+                    "Energy": float([d[k] for d in energy if k in d][0]),
+                }
+            fig_metrics = plotter.metrics_bar(metrics_per_key=metrics_pk)
+            figures["metrics_bar"] = fig_metrics
+
+            if save_dir:
+                import os
+                os.makedirs(save_dir, exist_ok=True)
+                fig_scatter.savefig(os.path.join(save_dir, "scatter_means.png"), dpi=dpi, bbox_inches="tight")
+                fig_residuals.savefig(os.path.join(save_dir, "residuals.png"), dpi=dpi, bbox_inches="tight")
+                fig_metrics.savefig(os.path.join(save_dir, "metrics_bar.png"), dpi=dpi, bbox_inches="tight")
+
+            results["figures"] = figures
+
+        return results

geneval/metrics/__init__.py

@@ -0,0 +1,65 @@
+"""
+Metrics module for gene expression evaluation.
+
+Provides per-gene and aggregate metrics for comparing distributions:
+- Correlation metrics (Pearson, Spearman)
+- Distribution distances (Wasserstein, MMD, Energy)
+- Multivariate distances
+"""
+
+from .base_metric import (
+    BaseMetric,
+    MetricResult,
+    DistributionMetric,
+    CorrelationMetric,
+)
+from .correlation import (
+    PearsonCorrelation,
+    SpearmanCorrelation,
+    MeanPearsonCorrelation,
+    MeanSpearmanCorrelation,
+)
+from .distances import (
+    Wasserstein1Distance,
+    Wasserstein2Distance,
+    MMDDistance,
+    EnergyDistance,
+    MultivariateWasserstein,
+    MultivariateMMD,
+)
+
+# All available metrics
+ALL_METRICS = [
+    PearsonCorrelation,
+    SpearmanCorrelation,
+    MeanPearsonCorrelation,
+    MeanSpearmanCorrelation,
+    Wasserstein1Distance,
+    Wasserstein2Distance,
+    MMDDistance,
+    EnergyDistance,
+    MultivariateWasserstein,
+    MultivariateMMD,
+]
+
+__all__ = [
+    # Base classes
+    "BaseMetric",
+    "MetricResult",
+    "DistributionMetric",
+    "CorrelationMetric",
+    # Correlation metrics
+    "PearsonCorrelation",
+    "SpearmanCorrelation",
+    "MeanPearsonCorrelation",
+    "MeanSpearmanCorrelation",
+    # Distance metrics
+    "Wasserstein1Distance",
+    "Wasserstein2Distance",
+    "MMDDistance",
+    "EnergyDistance",
+    "MultivariateWasserstein",
+    "MultivariateMMD",
+    # Collections
+    "ALL_METRICS",
+]

geneval/metrics/base_metric.py

@@ -0,0 +1,229 @@
+"""
+Base metric classes for gene expression evaluation.
+
+Provides abstract interface for all metrics with per-gene and aggregate computation.
+"""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Union, Any, Callable
+import numpy as np
+
+
+@dataclass
+class MetricResult:
+    """
+    Container for metric computation results.
+    
+    Stores both per-gene and aggregate values.
+    """
+    name: str
+    per_gene_values: np.ndarray  # Shape: (n_genes,)
+    gene_names: List[str]
+    aggregate_value: float
+    aggregate_method: str = "mean"  # mean, median, etc.
+    condition: Optional[str] = None
+    split: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    
+    @property
+    def as_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {
+            "name": self.name,
+            "aggregate_value": float(self.aggregate_value),
+            "aggregate_method": self.aggregate_method,
+            "per_gene_mean": float(np.nanmean(self.per_gene_values)),
+            "per_gene_std": float(np.nanstd(self.per_gene_values)),
+            "per_gene_median": float(np.nanmedian(self.per_gene_values)),
+            "n_genes": len(self.gene_names),
+            "condition": self.condition,
+            "split": self.split,
+            **self.metadata
+        }
+    
+    def top_genes(self, n: int = 10, ascending: bool = True) -> Dict[str, float]:
+        """Get top n genes by metric value."""
+        order = np.argsort(self.per_gene_values)
+        if not ascending:
+            order = order[::-1]
+        indices = order[:n]
+        return {self.gene_names[i]: float(self.per_gene_values[i]) for i in indices}
+
+
+class BaseMetric(ABC):
+    """
+    Abstract base class for all evaluation metrics.
+    
+    Metrics can be computed per-gene (returning a vector) or as aggregates.
+    All metrics should inherit from this class.
+    """
+    
+    def __init__(
+        self,
+        name: str,
+        description: str = "",
+        higher_is_better: bool = True,
+        requires_distribution: bool = False,
+    ):
+        """
+        Initialize metric.
+        
+        Parameters
+        ----------
+        name : str
+            Unique identifier for the metric
+        description : str
+            Human-readable description
+        higher_is_better : bool
+            Whether higher values indicate better performance
+        requires_distribution : bool
+            Whether metric needs full distribution (not just means)
+        """
+        self.name = name
+        self.description = description
+        self.higher_is_better = higher_is_better
+        self.requires_distribution = requires_distribution
+    
+    @abstractmethod
+    def compute_per_gene(
+        self,
+        real: np.ndarray,
+        generated: np.ndarray,
+    ) -> np.ndarray:
+        """
+        Compute metric for each gene.
+        
+        Parameters
+        ----------
+        real : np.ndarray
+            Real data matrix, shape (n_samples_real, n_genes)
+        generated : np.ndarray
+            Generated data matrix, shape (n_samples_gen, n_genes)
+            
+        Returns
+        -------
+        np.ndarray
+            Metric value per gene, shape (n_genes,)
+        """
+        pass
+    
+    def compute_aggregate(
+        self,
+        per_gene_values: np.ndarray,
+        method: str = "mean",
+    ) -> float:
+        """
+        Aggregate per-gene values to single metric.
+        
+        Parameters
+        ----------
+        per_gene_values : np.ndarray
+            Per-gene metric values
+        method : str
+            Aggregation method: "mean", "median", "std", "min", "max"
+            
+        Returns
+        -------
+        float
+            Aggregated metric value
+        """
+        methods = {
+            "mean": np.nanmean,
+            "median": np.nanmedian,
+            "std": np.nanstd,
+            "min": np.nanmin,
+            "max": np.nanmax,
+        }
+        if method not in methods:
+            raise ValueError(f"Unknown aggregation method: {method}")
+        return float(methods[method](per_gene_values))
+    
+    def compute(
+        self,
+        real: np.ndarray,
+        generated: np.ndarray,
+        gene_names: Optional[List[str]] = None,
+        aggregate_method: str = "mean",
+        condition: Optional[str] = None,
+        split: Optional[str] = None,
+    ) -> MetricResult:
+        """
+        Compute full metric result with per-gene and aggregate values.
+        
+        Parameters
+        ----------
+        real : np.ndarray
+            Real data matrix, shape (n_samples_real, n_genes)
+        generated : np.ndarray
+            Generated data matrix, shape (n_samples_gen, n_genes)
+        gene_names : List[str], optional
+            Names of genes (columns)
+        aggregate_method : str
+            How to aggregate per-gene values
+        condition : str, optional
+            Condition identifier
+        split : str, optional
+            Split identifier (train/test)
+            
+        Returns
+        -------
+        MetricResult
+            Complete metric result
+        """
+        n_genes = real.shape[1] if real.ndim > 1 else 1
+        if gene_names is None:
+            gene_names = [f"gene_{i}" for i in range(n_genes)]
+        
+        per_gene = self.compute_per_gene(real, generated)
+        aggregate = self.compute_aggregate(per_gene, method=aggregate_method)
+        
+        return MetricResult(
+            name=self.name,
+            per_gene_values=per_gene,
+            gene_names=gene_names,
+            aggregate_value=aggregate,
+            aggregate_method=aggregate_method,
+            condition=condition,
+            split=split,
+            metadata={
+                "higher_is_better": self.higher_is_better,
+                "description": self.description,
+            }
+        )
+    
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(name='{self.name}')"
+
+
+class DistributionMetric(BaseMetric):
+    """
+    Base class for distribution-based metrics (Wasserstein, MMD, Energy).
+    
+    These metrics require the full sample distributions, not just means.
+    """
+    
+    def __init__(self, name: str, description: str = "", higher_is_better: bool = False):
+        super().__init__(
+            name=name,
+            description=description,
+            higher_is_better=higher_is_better,
+            requires_distribution=True,
+        )
+
+
+class CorrelationMetric(BaseMetric):
+    """
+    Base class for correlation-based metrics (Pearson, Spearman).
+    
+    These compare mean profiles between real and generated data.
+    """
+    
+    def __init__(self, name: str, description: str = ""):
+        super().__init__(
+            name=name,
+            description=description,
+            higher_is_better=True,
+            requires_distribution=False,
+        )