gengeneeval 0.1.0__py3-none-any.whl

geneval/evaluator.py

@@ -0,0 +1,359 @@
+"""
+Comprehensive evaluator for gene expression data.
+
+Computes all metrics between real and generated data, organized by conditions and splits.
+"""
+from __future__ import annotations
+
+from typing import Dict, List, Optional, Union, Type, Any
+from pathlib import Path
+import numpy as np
+import warnings
+from dataclasses import dataclass
+
+from .data.loader import GeneExpressionDataLoader, load_data
+from .metrics.base_metric import BaseMetric, MetricResult
+from .metrics.correlation import (
+    PearsonCorrelation,
+    SpearmanCorrelation,
+    MeanPearsonCorrelation,
+    MeanSpearmanCorrelation,
+)
+from .metrics.distances import (
+    Wasserstein1Distance,
+    Wasserstein2Distance,
+    MMDDistance,
+    EnergyDistance,
+    MultivariateWasserstein,
+    MultivariateMMD,
+)
+from .results import EvaluationResult, SplitResult, ConditionResult
+
+
+# Default metrics to compute
+DEFAULT_METRICS = [
+    PearsonCorrelation,
+    SpearmanCorrelation,
+    MeanPearsonCorrelation,
+    MeanSpearmanCorrelation,
+    Wasserstein1Distance,
+    Wasserstein2Distance,
+    MMDDistance,
+    EnergyDistance,
+]
+
+
+class GeneEvalEvaluator:
+    """
+    Main evaluator class for gene expression data.
+    
+    Computes comprehensive metrics between real and generated datasets,
+    supporting multiple conditions, splits, and metric types.
+    
+    Parameters
+    ----------
+    data_loader : GeneExpressionDataLoader
+        Loaded and aligned data loader
+    metrics : List[BaseMetric or Type[BaseMetric]], optional
+        Metrics to compute. If None, uses default set.
+    aggregate_method : str
+        How to aggregate per-gene values (mean, median, etc.)
+    include_multivariate : bool
+        Whether to include multivariate (whole-space) metrics
+    verbose : bool
+        Whether to print progress
+        
+    Examples
+    --------
+    >>> loader = load_data("real.h5ad", "generated.h5ad", ["perturbation"])
+    >>> evaluator = GeneEvalEvaluator(loader)
+    >>> results = evaluator.evaluate()
+    >>> results.save("output/")
+    """
+    
+    def __init__(
+        self,
+        data_loader: GeneExpressionDataLoader,
+        metrics: Optional[List[Union[BaseMetric, Type[BaseMetric]]]] = None,
+        aggregate_method: str = "mean",
+        include_multivariate: bool = True,
+        verbose: bool = True,
+    ):
+        self.data_loader = data_loader
+        self.aggregate_method = aggregate_method
+        self.include_multivariate = include_multivariate
+        self.verbose = verbose
+        
+        # Initialize metrics
+        self.metrics: List[BaseMetric] = []
+        metric_classes = metrics or DEFAULT_METRICS
+        
+        for m in metric_classes:
+            if isinstance(m, type):
+                # It's a class, instantiate it
+                self.metrics.append(m())
+            else:
+                # It's already an instance
+                self.metrics.append(m)
+        
+        # Add multivariate metrics if requested
+        if include_multivariate:
+            self.metrics.extend([
+                MultivariateWasserstein(),
+                MultivariateMMD(),
+            ])
+    
+    def _log(self, msg: str):
+        """Print message if verbose."""
+        if self.verbose:
+            print(msg)
+    
+    def evaluate(
+        self,
+        splits: Optional[List[str]] = None,
+        save_dir: Optional[Union[str, Path]] = None,
+    ) -> EvaluationResult:
+        """
+        Run full evaluation on all conditions and splits.
+        
+        Parameters
+        ----------
+        splits : List[str], optional
+            Splits to evaluate. If None, evaluates all available splits.
+        save_dir : str or Path, optional
+            If provided, save results to this directory
+            
+        Returns
+        -------
+        EvaluationResult
+            Complete evaluation results
+        """
+        # Get available splits
+        available_splits = self.data_loader.get_splits()
+        
+        if splits is None:
+            splits = available_splits
+        else:
+            # Validate requested splits
+            invalid = set(splits) - set(available_splits)
+            if invalid:
+                warnings.warn(f"Requested splits not found: {invalid}")
+                splits = [s for s in splits if s in available_splits]
+        
+        self._log(f"Evaluating {len(splits)} splits: {splits}")
+        self._log(f"Using {len(self.metrics)} metrics: {[m.name for m in self.metrics]}")
+        
+        # Create result container
+        result = EvaluationResult(
+            gene_names=self.data_loader.gene_names,
+            condition_columns=self.data_loader.condition_columns,
+            metadata={
+                "real_path": str(self.data_loader.real_path),
+                "generated_path": str(self.data_loader.generated_path),
+                "aggregate_method": self.aggregate_method,
+                "metric_names": [m.name for m in self.metrics],
+            }
+        )
+        
+        # Evaluate each split
+        for split in splits:
+            split_key = split if split != "all" else None
+            split_result = self._evaluate_split(split, split_key)
+            result.add_split(split_result)
+        
+        # Compute aggregate metrics
+        for split_result in result.splits.values():
+            split_result.compute_aggregates()
+        
+        # Print summary
+        if self.verbose:
+            self._print_summary(result)
+        
+        # Save if requested
+        if save_dir is not None:
+            result.save(save_dir)
+            self._log(f"Results saved to: {save_dir}")
+        
+        return result
+    
+    def _evaluate_split(
+        self,
+        split_name: str,
+        split_filter: Optional[str]
+    ) -> SplitResult:
+        """Evaluate a single split."""
+        split_result = SplitResult(split_name=split_name)
+        
+        conditions = list(self.data_loader.iterate_conditions(split_filter))
+        self._log(f"\n  Split '{split_name}': {len(conditions)} conditions")
+        
+        for i, (cond_key, real_data, gen_data, cond_info) in enumerate(conditions):
+            if self.verbose and (i + 1) % 10 == 0:
+                self._log(f"    Processing condition {i + 1}/{len(conditions)}")
+            
+            # Create condition result
+            cond_result = ConditionResult(
+                condition_key=cond_key,
+                split=split_name,
+                n_real_samples=real_data.shape[0],
+                n_generated_samples=gen_data.shape[0],
+                n_genes=real_data.shape[1],
+                gene_names=self.data_loader.gene_names,
+                perturbation=cond_info.get(self.data_loader.condition_columns[0]),
+                covariates=cond_info,
+            )
+            
+            # Store mean profiles
+            cond_result.real_mean = real_data.mean(axis=0)
+            cond_result.generated_mean = gen_data.mean(axis=0)
+            
+            # Compute all metrics
+            for metric in self.metrics:
+                try:
+                    metric_result = metric.compute(
+                        real=real_data,
+                        generated=gen_data,
+                        gene_names=self.data_loader.gene_names,
+                        aggregate_method=self.aggregate_method,
+                        condition=cond_key,
+                        split=split_name,
+                    )
+                    cond_result.add_metric(metric.name, metric_result)
+                except Exception as e:
+                    warnings.warn(
+                        f"Failed to compute {metric.name} for {cond_key}: {e}"
+                    )
+            
+            split_result.add_condition(cond_result)
+        
+        return split_result
+    
+    def _print_summary(self, result: EvaluationResult):
+        """Print summary of results."""
+        self._log("\n" + "=" * 60)
+        self._log("EVALUATION SUMMARY")
+        self._log("=" * 60)
+        
+        for split_name, split in result.splits.items():
+            self._log(f"\nSplit: {split_name} ({split.n_conditions} conditions)")
+            self._log("-" * 40)
+            
+            # Print aggregate metrics
+            for key, value in sorted(split.aggregate_metrics.items()):
+                if key.endswith("_mean"):
+                    metric_name = key[:-5]
+                    std_key = f"{metric_name}_std"
+                    std = split.aggregate_metrics.get(std_key, 0)
+                    self._log(f"  {metric_name}: {value:.4f} ± {std:.4f}")
+        
+        self._log("=" * 60)
+
+
+def evaluate(
+    real_path: Union[str, Path],
+    generated_path: Union[str, Path],
+    condition_columns: List[str],
+    split_column: Optional[str] = None,
+    output_dir: Optional[Union[str, Path]] = None,
+    metrics: Optional[List[Union[BaseMetric, Type[BaseMetric]]]] = None,
+    include_multivariate: bool = True,
+    verbose: bool = True,
+    **loader_kwargs
+) -> EvaluationResult:
+    """
+    Convenience function to run full evaluation.
+    
+    Parameters
+    ----------
+    real_path : str or Path
+        Path to real data h5ad file
+    generated_path : str or Path
+        Path to generated data h5ad file
+    condition_columns : List[str]
+        Columns to match between datasets
+    split_column : str, optional
+        Column indicating train/test split
+    output_dir : str or Path, optional
+        Directory to save results
+    metrics : List, optional
+        Metrics to compute
+    include_multivariate : bool
+        Whether to include multivariate metrics
+    verbose : bool
+        Print progress
+    **loader_kwargs
+        Additional arguments for data loader
+        
+    Returns
+    -------
+    EvaluationResult
+        Complete evaluation results
+        
+    Examples
+    --------
+    >>> results = evaluate(
+    ...     "real.h5ad",
+    ...     "generated.h5ad",
+    ...     condition_columns=["perturbation", "cell_type"],
+    ...     split_column="split",
+    ...     output_dir="evaluation_output/"
+    ... )
+    """
+    # Load data
+    loader = load_data(
+        real_path=real_path,
+        generated_path=generated_path,
+        condition_columns=condition_columns,
+        split_column=split_column,
+        **loader_kwargs
+    )
+    
+    # Create evaluator
+    evaluator = GeneEvalEvaluator(
+        data_loader=loader,
+        metrics=metrics,
+        include_multivariate=include_multivariate,
+        verbose=verbose,
+    )
+    
+    # Run evaluation
+    return evaluator.evaluate(save_dir=output_dir)
+
+
+class MetricRegistry:
+    """
+    Registry of available metrics.
+    
+    Allows registration of custom metrics and retrieval by name.
+    """
+    
+    _metrics: Dict[str, Type[BaseMetric]] = {}
+    
+    @classmethod
+    def register(cls, metric_class: Type[BaseMetric]):
+        """Register a metric class."""
+        instance = metric_class()
+        cls._metrics[instance.name] = metric_class
+    
+    @classmethod
+    def get(cls, name: str) -> Optional[Type[BaseMetric]]:
+        """Get metric class by name."""
+        return cls._metrics.get(name)
+    
+    @classmethod
+    def list_all(cls) -> List[str]:
+        """List all registered metric names."""
+        return list(cls._metrics.keys())
+    
+    @classmethod
+    def get_all(cls) -> List[Type[BaseMetric]]:
+        """Get all registered metric classes."""
+        return list(cls._metrics.values())
+
+
+# Register default metrics
+for metric_class in DEFAULT_METRICS:
+    MetricRegistry.register(metric_class)
+
+MetricRegistry.register(MultivariateWasserstein)
+MetricRegistry.register(MultivariateMMD)

geneval/evaluators/__init__.py

@@ -0,0 +1,4 @@
+from .base_evaluator import BaseEvaluator
+from .gene_expression_evaluator import GeneExpressionEvaluator
+
+__all__ = ["BaseEvaluator", "GeneExpressionEvaluator"]

geneval/evaluators/base_evaluator.py

@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+from abc import ABC
+from typing import Dict, Iterable, List, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+from anndata import AnnData
+from scipy import sparse
+
+from ..utils.preprocessing import to_dense
+
+
+class BaseEvaluator(ABC):
+    """
+    Base class for evaluation of generated data against real datasets.
+
+    Provides:
+    - Variable/gene alignment between real and generated AnnData objects
+    - Computation and application of control baselines per strata
+    """
+
+    def __init__(self, data, output: AnnData):
+        """
+        Parameters
+        ----------
+        data : object
+            An object providing at least:
+              - gene_expression_dataset.adata: AnnData
+              - perturbation_key: str
+              - split_key: str
+              - control: str
+              - condition_keys: Optional[List[str]]
+        output : AnnData
+            Generated data to evaluate.
+        """
+        self.data = data
+        self.output = output
+
+    # ---------- alignment utilities ----------
+
+    def _align_varnames_like(self, real: AnnData, generated: AnnData) -> Tuple[AnnData, AnnData]:
+        """
+        Align real and generated AnnData to the common set of var_names (genes),
+        preserving order based on the real AnnData.
+        """
+        real_genes = pd.Index(real.var_names.astype(str))
+        gen_genes = pd.Index(generated.var_names.astype(str))
+        common = real_genes.intersection(gen_genes)
+        if len(common) == 0:
+            raise ValueError("No overlapping genes between real and generated AnnData.")
+
+        # Reindex both adatas to the common genes in the order of real
+        real = real[:, real_genes.get_indexer(common)].copy()
+        generated = generated[:, generated.var_names.astype(str).isin(common)].copy()
+        # Reorder generated to match real
+        generated = generated[:, pd.Index(generated.var_names.astype(str)).get_indexer(common)].copy()
+
+        real.var_names = common
+        generated.var_names = common
+        return real, generated
+
+    # ---------- baseline utilities ----------
+
+    @staticmethod
+    def _key_from_values(values: Iterable[object]) -> str:
+        # stable string key for strata-tuples
+        return "####".join([str(v) for v in values])
+
+    def _compute_control_means(
+        self,
+        adata: AnnData,
+        perturbation_col: str,
+        control_value: str,
+        strata_cols: Optional[List[str]] = None,
+    ) -> Dict[str, np.ndarray]:
+        """
+        Compute per-strata control means across genes.
+
+        Returns a dict mapping a strata-key -> mean vector (n_genes,).
+        """
+        strata_cols = strata_cols or []
+        obs = adata.obs
+
+        if perturbation_col not in obs.columns:
+            raise KeyError(f"'{perturbation_col}' not found in adata.obs.")
+
+        is_control = (obs[perturbation_col].astype(str) == str(control_value)).to_numpy()
+        if not is_control.any():
+            # no controls; return empty means map
+            return {}
+
+        ctrl = adata[is_control]
+        if not strata_cols:
+            return {self._key_from_values([]): to_dense(ctrl.X).mean(axis=0)}
+
+        # group by strata columns (as strings to be robust)
+        df = ctrl.obs[strata_cols].astype(str)
+        means: Dict[str, np.ndarray] = {}
+        # compute mean per unique strata combination
+        for _, row in df.drop_duplicates().iterrows():
+            mask = np.ones(ctrl.n_obs, dtype=bool)
+            for c in strata_cols:
+                mask &= (df[c].to_numpy() == str(row[c]))
+            if not mask.any():
+                continue
+            key = self._key_from_values([row[c] for c in strata_cols])
+            means[key] = to_dense(ctrl.X[mask]).mean(axis=0)
+        return means
+
+    def _apply_baseline_per_strata(
+        self,
+        X,
+        obs: pd.DataFrame,
+        baseline: Dict[str, np.ndarray],
+        strata_cols: Optional[List[str]] = None,
+        mode: str = "subtract",
+    ):
+        """
+        Apply per-strata baseline vectors to rows in X based on obs[strata_cols].
+
+        mode: 'subtract' or 'add'
+        """
+        strata_cols = strata_cols or []
+        if mode not in ("subtract", "add"):
+            raise ValueError("mode must be 'subtract' or 'add'.")
+
+        if sparse.issparse(X):
+            X = X.tocsr(copy=True)
+            to_dense_first = False
+        else:
+            X = np.array(X, copy=True)
+            to_dense_first = True  # already dense
+
+        if not strata_cols:
+            key = self._key_from_values([])
+            b = baseline.get(key, None)
+            if b is None:
+                return X
+            if sparse.issparse(X):
+                # operate dense for simplicity
+                X = X.toarray()
+            if mode == "subtract":
+                X -= b
+            else:
+                X += b
+            return X
+
+        # Apply per group
+        df = obs[strata_cols].astype(str)
+        # iterate groups in baseline for efficiency
+        for key, b in baseline.items():
+            # decode key into tuple of values
+            parts = key.split("####") if key else []
+            if len(parts) != len(strata_cols):
+                # skip mismatched key
+                continue
+            mask = np.ones(df.shape[0], dtype=bool)
+            for col, val in zip(strata_cols, parts):
+                mask &= (df[col].to_numpy() == val)
+            if not mask.any():
+                continue
+
+            if sparse.issparse(X):
+                # operate in dense block then write back
+                block = X[mask].toarray()
+                if mode == "subtract":
+                    block -= b
+                else:
+                    block += b
+                X[mask] = sparse.csr_matrix(block)
+            else:
+                if mode == "subtract":
+                    X[mask] -= b
+                else:
+                    X[mask] += b
+
+        return X