gengeneeval 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

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/lazy_evaluator.py

@@ -0,0 +1,424 @@
+"""
+Memory-efficient evaluator for large-scale gene expression datasets.
+
+Uses lazy loading and batched processing to minimize memory footprint.
+"""
+from __future__ import annotations
+
+from typing import Dict, List, Optional, Union, Type, Any, Generator
+from pathlib import Path
+import numpy as np
+import warnings
+from dataclasses import dataclass, field
+import gc
+
+from .data.lazy_loader import (
+    LazyGeneExpressionDataLoader,
+    load_data_lazy,
+    ConditionBatch,
+)
+from .metrics.base_metric import BaseMetric, MetricResult
+from .metrics.correlation import (
+    PearsonCorrelation,
+    SpearmanCorrelation,
+    MeanPearsonCorrelation,
+    MeanSpearmanCorrelation,
+)
+from .metrics.distances import (
+    Wasserstein1Distance,
+    Wasserstein2Distance,
+    MMDDistance,
+    EnergyDistance,
+)
+from .metrics.reconstruction import (
+    MSEDistance,
+)
+
+# These multivariate metrics don't support batched computation
+from .metrics.distances import MultivariateWasserstein, MultivariateMMD
+
+
+# Metrics that support incremental/batched computation
+BATCHABLE_METRICS = [
+    MSEDistance,
+    PearsonCorrelation,
+    SpearmanCorrelation,
+]
+
+# Metrics that require full data
+NON_BATCHABLE_METRICS = [
+    Wasserstein1Distance,
+    Wasserstein2Distance,
+    MMDDistance,
+    EnergyDistance,
+    MultivariateWasserstein,
+    MultivariateMMD,
+]
+
+
+@dataclass
+class StreamingMetricAccumulator:
+    """Accumulates values for streaming mean/std computation."""
+    n: int = 0
+    sum: float = 0.0
+    sum_sq: float = 0.0
+    
+    def add(self, value: float, count: int = 1):
+        """Add a value (or batch of values with same value)."""
+        self.n += count
+        self.sum += value * count
+        self.sum_sq += (value ** 2) * count
+    
+    def add_batch(self, values: np.ndarray):
+        """Add multiple values."""
+        self.n += len(values)
+        self.sum += np.sum(values)
+        self.sum_sq += np.sum(values ** 2)
+    
+    @property
+    def mean(self) -> float:
+        return self.sum / self.n if self.n > 0 else 0.0
+    
+    @property
+    def std(self) -> float:
+        if self.n <= 1:
+            return 0.0
+        variance = (self.sum_sq / self.n) - (self.mean ** 2)
+        return np.sqrt(max(0, variance))
+
+
+@dataclass
+class StreamingConditionResult:
+    """Lightweight result for a single condition."""
+    condition_key: str
+    n_real_samples: int = 0
+    n_generated_samples: int = 0
+    metrics: Dict[str, float] = field(default_factory=dict)
+    real_mean: Optional[np.ndarray] = None
+    generated_mean: Optional[np.ndarray] = None
+
+
+@dataclass
+class StreamingEvaluationResult:
+    """Memory-efficient evaluation result that streams to disk."""
+    output_dir: Path
+    n_conditions: int = 0
+    metric_accumulators: Dict[str, StreamingMetricAccumulator] = field(default_factory=dict)
+    condition_keys: List[str] = field(default_factory=list)
+    
+    def add_condition(self, result: StreamingConditionResult):
+        """Add a condition result and update accumulators."""
+        self.n_conditions += 1
+        self.condition_keys.append(result.condition_key)
+        
+        for metric_name, value in result.metrics.items():
+            if metric_name not in self.metric_accumulators:
+                self.metric_accumulators[metric_name] = StreamingMetricAccumulator()
+            self.metric_accumulators[metric_name].add(value)
+    
+    def get_summary(self) -> Dict[str, Dict[str, float]]:
+        """Get summary statistics."""
+        summary = {}
+        for name, acc in self.metric_accumulators.items():
+            summary[name] = {
+                "mean": acc.mean,
+                "std": acc.std,
+                "n": acc.n,
+            }
+        return summary
+    
+    def save_summary(self):
+        """Save summary to output directory."""
+        import json
+        
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        
+        summary = {
+            "n_conditions": self.n_conditions,
+            "metrics": self.get_summary(),
+            "condition_keys": self.condition_keys,
+        }
+        
+        with open(self.output_dir / "summary.json", "w") as f:
+            json.dump(summary, f, indent=2)
+
+
+class MemoryEfficientEvaluator:
+    """
+    Memory-efficient evaluator using lazy loading and batched processing.
+    
+    Features:
+    - Lazy data loading (one condition at a time)
+    - Batched processing within conditions
+    - Streaming metric accumulation
+    - Periodic garbage collection
+    - Progress streaming to disk
+    
+    Parameters
+    ----------
+    data_loader : LazyGeneExpressionDataLoader
+        Lazy data loader
+    metrics : List[BaseMetric], optional
+        Metrics to compute. Note: Some metrics (like MMD) may not support
+        batched computation and will use full condition data.
+    batch_size : int
+        Batch size for within-condition processing
+    gc_every_n_conditions : int
+        Run garbage collection every N conditions
+    verbose : bool
+        Print progress
+    """
+    
+    def __init__(
+        self,
+        data_loader: LazyGeneExpressionDataLoader,
+        metrics: Optional[List[Union[BaseMetric, Type[BaseMetric]]]] = None,
+        batch_size: int = 256,
+        gc_every_n_conditions: int = 10,
+        verbose: bool = True,
+    ):
+        self.data_loader = data_loader
+        self.batch_size = batch_size
+        self.gc_every_n_conditions = gc_every_n_conditions
+        self.verbose = verbose
+        
+        # Initialize metrics
+        self.metrics: List[BaseMetric] = []
+        metric_classes = metrics or [
+            MSEDistance,
+            PearsonCorrelation,
+            SpearmanCorrelation,
+            MeanPearsonCorrelation,
+            MeanSpearmanCorrelation,
+        ]
+        
+        for m in metric_classes:
+            if isinstance(m, type):
+                self.metrics.append(m())
+            else:
+                self.metrics.append(m)
+    
+    def _log(self, msg: str):
+        if self.verbose:
+            print(msg)
+    
+    def evaluate(
+        self,
+        split: Optional[str] = None,
+        output_dir: Optional[Union[str, Path]] = None,
+        save_per_condition: bool = False,
+    ) -> StreamingEvaluationResult:
+        """
+        Run memory-efficient evaluation.
+        
+        Parameters
+        ----------
+        split : str, optional
+            Split to evaluate
+        output_dir : str or Path, optional
+            Directory to save results. If provided, results are streamed to disk.
+        save_per_condition : bool
+            If True, save individual condition results to disk
+            
+        Returns
+        -------
+        StreamingEvaluationResult
+            Evaluation result with aggregated metrics
+        """
+        if output_dir is not None:
+            output_dir = Path(output_dir)
+            output_dir.mkdir(parents=True, exist_ok=True)
+        else:
+            output_dir = Path(".")
+        
+        result = StreamingEvaluationResult(output_dir=output_dir)
+        
+        # Get conditions
+        conditions = self.data_loader.get_common_conditions(split)
+        self._log(f"Evaluating {len(conditions)} conditions")
+        self._log(f"Memory estimate: {self.data_loader.estimate_memory_usage()}")
+        
+        # Iterate conditions (one at a time in memory)
+        for i, (cond_key, real_data, gen_data, cond_info) in enumerate(
+            self.data_loader.iterate_conditions(split)
+        ):
+            if self.verbose and (i + 1) % 10 == 0:
+                self._log(f"  Processing {i + 1}/{len(conditions)}: {cond_key}")
+            
+            # Compute metrics for this condition
+            cond_result = self._evaluate_condition(
+                cond_key, real_data, gen_data, cond_info
+            )
+            
+            # Add to streaming result
+            result.add_condition(cond_result)
+            
+            # Optionally save per-condition result
+            if save_per_condition and output_dir:
+                self._save_condition_result(cond_result, output_dir)
+            
+            # Periodic garbage collection
+            if (i + 1) % self.gc_every_n_conditions == 0:
+                gc.collect()
+        
+        # Final summary
+        result.save_summary()
+        
+        if self.verbose:
+            self._print_summary(result)
+        
+        return result
+    
+    def _evaluate_condition(
+        self,
+        cond_key: str,
+        real_data: np.ndarray,
+        gen_data: np.ndarray,
+        cond_info: Dict[str, str],
+    ) -> StreamingConditionResult:
+        """Evaluate a single condition."""
+        result = StreamingConditionResult(
+            condition_key=cond_key,
+            n_real_samples=real_data.shape[0],
+            n_generated_samples=gen_data.shape[0],
+        )
+        
+        # Compute means
+        result.real_mean = real_data.mean(axis=0)
+        result.generated_mean = gen_data.mean(axis=0)
+        
+        # Compute 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="mean",
+                    condition=cond_key,
+                )
+                result.metrics[metric.name] = metric_result.aggregate_value
+            except Exception as e:
+                warnings.warn(f"Failed to compute {metric.name} for {cond_key}: {e}")
+        
+        return result
+    
+    def _save_condition_result(
+        self,
+        result: StreamingConditionResult,
+        output_dir: Path,
+    ):
+        """Save a single condition result to disk."""
+        import json
+        
+        condition_dir = output_dir / "conditions"
+        condition_dir.mkdir(exist_ok=True)
+        
+        # Safe filename
+        safe_key = result.condition_key.replace("/", "_").replace("\\", "_")
+        
+        data = {
+            "condition_key": result.condition_key,
+            "n_real": result.n_real_samples,
+            "n_generated": result.n_generated_samples,
+            "metrics": result.metrics,
+        }
+        
+        with open(condition_dir / f"{safe_key}.json", "w") as f:
+            json.dump(data, f, indent=2)
+    
+    def _print_summary(self, result: StreamingEvaluationResult):
+        """Print summary."""
+        self._log("\n" + "=" * 60)
+        self._log("EVALUATION SUMMARY (Memory-Efficient)")
+        self._log("=" * 60)
+        self._log(f"Conditions evaluated: {result.n_conditions}")
+        self._log("-" * 40)
+        
+        for name, stats in result.get_summary().items():
+            self._log(f"  {name}: {stats['mean']:.4f} ± {stats['std']:.4f}")
+        
+        self._log("=" * 60)
+
+
+def evaluate_lazy(
+    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,
+    batch_size: int = 256,
+    use_backed: bool = False,
+    metrics: Optional[List[Union[BaseMetric, Type[BaseMetric]]]] = None,
+    verbose: bool = True,
+    save_per_condition: bool = False,
+    **kwargs
+) -> StreamingEvaluationResult:
+    """
+    Memory-efficient evaluation using lazy loading.
+    
+    Use this function for large datasets that don't fit in memory.
+    
+    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 for train/test split
+    output_dir : str or Path, optional
+        Directory to save results
+    batch_size : int
+        Batch size for processing
+    use_backed : bool
+        Use memory-mapped file access (for very large files)
+    metrics : List, optional
+        Metrics to compute
+    verbose : bool
+        Print progress
+    save_per_condition : bool
+        Save individual condition results
+        
+    Returns
+    -------
+    StreamingEvaluationResult
+        Aggregated evaluation results
+        
+    Examples
+    --------
+    >>> # For large datasets that don't fit in memory
+    >>> results = evaluate_lazy(
+    ...     "real.h5ad",
+    ...     "generated.h5ad",
+    ...     condition_columns=["perturbation"],
+    ...     output_dir="eval_output/",
+    ...     batch_size=256,
+    ...     use_backed=True,  # Memory-mapped for very large files
+    ... )
+    >>> print(results.get_summary())
+    """
+    # Create lazy loader
+    with load_data_lazy(
+        real_path=real_path,
+        generated_path=generated_path,
+        condition_columns=condition_columns,
+        split_column=split_column,
+        batch_size=batch_size,
+        use_backed=use_backed,
+    ) as loader:
+        # Create evaluator
+        evaluator = MemoryEfficientEvaluator(
+            data_loader=loader,
+            metrics=metrics,
+            batch_size=batch_size,
+            verbose=verbose,
+        )
+        
+        # Run evaluation
+        return evaluator.evaluate(
+            output_dir=output_dir,
+            save_per_condition=save_per_condition,
+        )

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",
 ]