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

geneval/__init__.py

@@ -8,6 +8,7 @@ Features:
 - Multiple distance and correlation metrics (per-gene and aggregate)
 - Condition-based matching (perturbation, cell type, etc.)
 - Train/test split support
+- Memory-efficient lazy loading for large datasets
 - Publication-quality visualizations
 - Command-line interface
 
@@ -20,12 +21,22 @@ Quick Start:
     ...     output_dir="output/"
     ... )
 
+Memory-Efficient Mode (for large datasets):
+    >>> from geneval import evaluate_lazy
+    >>> results = evaluate_lazy(
+    ...     real_path="real.h5ad",
+    ...     generated_path="generated.h5ad",
+    ...     condition_columns=["perturbation"],
+    ...     batch_size=256,
+    ...     use_backed=True,  # Memory-mapped access
+    ... )
+
 CLI Usage:
     $ geneval --real real.h5ad --generated generated.h5ad \\
               --conditions perturbation cell_type --output results/
 """
 
-__version__ = "0.2.0"
+__version__ = "0.3.0"
 __author__ = "GenEval Team"
 
 # Main evaluation interface
@@ -35,12 +46,26 @@ from .evaluator import (
     MetricRegistry,
 )
 
+# Memory-efficient evaluation
+from .lazy_evaluator import (
+    evaluate_lazy,
+    MemoryEfficientEvaluator,
+    StreamingEvaluationResult,
+)
+
 # Data loading
 from .data.loader import (
     GeneExpressionDataLoader,
     load_data,
 )
 
+# Memory-efficient data loading
+from .data.lazy_loader import (
+    LazyGeneExpressionDataLoader,
+    load_data_lazy,
+    ConditionBatch,
+)
+
 # Results
 from .results import (
     EvaluationResult,
@@ -76,6 +101,14 @@ from .metrics.reconstruction import (
     R2Score,
 )
 
+# Accelerated computation
+from .metrics.accelerated import (
+    AccelerationConfig,
+    ParallelMetricComputer,
+    get_available_backends,
+    compute_metrics_accelerated,
+)
+
 # Visualization
 from .visualization.visualizer import (
     EvaluationVisualizer,
@@ -99,9 +132,17 @@ __all__ = [
     "evaluate",
     "GeneEvalEvaluator",
     "MetricRegistry",
+    # Memory-efficient evaluation
+    "evaluate_lazy",
+    "MemoryEfficientEvaluator",
+    "StreamingEvaluationResult",
     # Data loading
     "GeneExpressionDataLoader",
     "load_data",
+    # Memory-efficient data loading
+    "LazyGeneExpressionDataLoader",
+    "load_data_lazy",
+    "ConditionBatch",
     # Results
     "EvaluationResult",
     "SplitResult", 
@@ -123,6 +164,16 @@ __all__ = [
     "EnergyDistance",
     "MultivariateWasserstein",
     "MultivariateMMD",
+    # Reconstruction metrics
+    "MSEDistance",
+    "RMSEDistance",
+    "MAEDistance",
+    "R2Score",
+    # Acceleration
+    "AccelerationConfig",
+    "ParallelMetricComputer",
+    "get_available_backends",
+    "compute_metrics_accelerated",
     # Visualization
     "EvaluationVisualizer",
     "visualize",

geneval/data/__init__.py

@@ -2,6 +2,7 @@
 Data loading module for gene expression evaluation.
 
 Provides data loaders for paired real and generated datasets.
+Includes both standard and memory-efficient lazy loading options.
 """
 
 from .loader import (
@@ -9,15 +10,28 @@ from .loader import (
     load_data,
     DataLoaderError,
 )
+from .lazy_loader import (
+    LazyGeneExpressionDataLoader,
+    load_data_lazy,
+    LazyDataLoaderError,
+    ConditionBatch,
+)
 from .gene_expression_datamodule import (
     GeneExpressionDataModule,
     DataModuleError,
 )
 
 __all__ = [
+    # Standard loader
     "GeneExpressionDataLoader",
     "load_data",
     "DataLoaderError",
+    # Lazy loader (memory-efficient)
+    "LazyGeneExpressionDataLoader",
+    "load_data_lazy",
+    "LazyDataLoaderError",
+    "ConditionBatch",
+    # DataModule
     "GeneExpressionDataModule",
     "DataModuleError",
 ]

geneval/data/lazy_loader.py

@@ -0,0 +1,562 @@
+"""
+Memory-efficient lazy data loader for large-scale gene expression datasets.
+
+Provides lazy loading and batched iteration over AnnData h5ad files without
+loading entire datasets into memory. Supports backed mode for very large files.
+"""
+from __future__ import annotations
+
+from typing import Optional, List, Union, Dict, Tuple, Iterator, Generator
+from pathlib import Path
+import warnings
+import numpy as np
+import pandas as pd
+from scipy import sparse
+from dataclasses import dataclass
+
+try:
+    import anndata as ad
+    import scanpy as sc
+except ImportError:
+    raise ImportError("anndata and scanpy are required. Install with: pip install anndata scanpy")
+
+
+class LazyDataLoaderError(Exception):
+    """Custom exception for lazy data loading errors."""
+    pass
+
+
+@dataclass
+class ConditionBatch:
+    """Container for a batch of samples from a condition."""
+    condition_key: str
+    condition_info: Dict[str, str]
+    real_data: np.ndarray
+    generated_data: np.ndarray
+    batch_idx: int
+    n_batches: int
+    is_last_batch: bool
+    
+
+class LazyGeneExpressionDataLoader:
+    """
+    Memory-efficient lazy data loader for paired gene expression datasets.
+    
+    Unlike GeneExpressionDataLoader, this class:
+    - Uses backed mode for h5ad files to avoid loading entire datasets
+    - Supports batched iteration over conditions
+    - Only loads data into memory when explicitly requested
+    - Provides memory usage estimates
+    
+    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
+    batch_size : int
+        Maximum number of samples per batch when iterating
+    use_backed : bool
+        If True, use backed mode (memory-mapped). May be slower but uses minimal memory.
+        If False, loads full file but processes in batches.
+    min_samples_per_condition : int
+        Minimum samples required per condition to include
+    """
+    
+    def __init__(
+        self,
+        real_path: Union[str, Path],
+        generated_path: Union[str, Path],
+        condition_columns: List[str],
+        split_column: Optional[str] = None,
+        batch_size: int = 256,
+        use_backed: bool = False,
+        min_samples_per_condition: int = 2,
+    ):
+        self.real_path = Path(real_path)
+        self.generated_path = Path(generated_path)
+        self.condition_columns = condition_columns
+        self.split_column = split_column
+        self.batch_size = batch_size
+        self.use_backed = use_backed
+        self.min_samples_per_condition = min_samples_per_condition
+        
+        # Lazy-loaded references (backed or full)
+        self._real: Optional[ad.AnnData] = None
+        self._generated: Optional[ad.AnnData] = None
+        
+        # Metadata (always loaded - lightweight)
+        self._real_obs: Optional[pd.DataFrame] = None
+        self._generated_obs: Optional[pd.DataFrame] = None
+        self._real_var_names: Optional[pd.Index] = None
+        self._generated_var_names: Optional[pd.Index] = None
+        
+        # Gene alignment info
+        self._common_genes: Optional[List[str]] = None
+        self._real_gene_idx: Optional[np.ndarray] = None
+        self._gen_gene_idx: Optional[np.ndarray] = None
+        
+        # Pre-computed condition indices for fast access
+        self._condition_indices: Optional[Dict[str, Dict[str, np.ndarray]]] = None
+        
+        # State
+        self._is_initialized = False
+    
+    def initialize(self) -> "LazyGeneExpressionDataLoader":
+        """
+        Initialize loader by reading metadata only (not expression data).
+        
+        This loads obs DataFrames and var_names to prepare for iteration,
+        but does not load the expression matrices.
+        
+        Returns
+        -------
+        self
+            For method chaining
+        """
+        if self._is_initialized:
+            return self
+        
+        # Validate paths
+        if not self.real_path.exists():
+            raise LazyDataLoaderError(f"Real data file not found: {self.real_path}")
+        if not self.generated_path.exists():
+            raise LazyDataLoaderError(f"Generated data file not found: {self.generated_path}")
+        
+        # Load metadata only (obs and var_names)
+        # This is much faster and lighter than loading full data
+        self._load_metadata()
+        
+        # Validate columns
+        self._validate_columns()
+        
+        # Compute gene alignment indices
+        self._compute_gene_alignment()
+        
+        # Pre-compute condition indices
+        self._precompute_condition_indices()
+        
+        self._is_initialized = True
+        return self
+    
+    def _load_metadata(self):
+        """Load only metadata (obs, var_names) without expression data."""
+        # For backed mode, we open but don't load X
+        # For non-backed, we still only read metadata initially
+        
+        # Try backed mode if requested
+        if self.use_backed:
+            try:
+                self._real = sc.read_h5ad(self.real_path, backed='r')
+                self._generated = sc.read_h5ad(self.generated_path, backed='r')
+            except Exception as e:
+                warnings.warn(f"Backed mode failed, falling back to standard loading: {e}")
+                self._real = None
+                self._generated = None
+        
+        if self._real is None:
+            # Load only what we need: obs and var
+            # For very large files, read in low-memory mode
+            self._real = sc.read_h5ad(self.real_path)
+            self._generated = sc.read_h5ad(self.generated_path)
+        
+        # Cache lightweight metadata
+        self._real_obs = self._real.obs.copy()
+        self._generated_obs = self._generated.obs.copy()
+        self._real_var_names = pd.Index(self._real.var_names.astype(str))
+        self._generated_var_names = pd.Index(self._generated.var_names.astype(str))
+    
+    def _validate_columns(self):
+        """Validate that required columns exist."""
+        for col in self.condition_columns:
+            if col not in self._real_obs.columns:
+                raise LazyDataLoaderError(
+                    f"Condition column '{col}' not found in real data. "
+                    f"Available: {list(self._real_obs.columns)}"
+                )
+            if col not in self._generated_obs.columns:
+                raise LazyDataLoaderError(
+                    f"Condition column '{col}' not found in generated data. "
+                    f"Available: {list(self._generated_obs.columns)}"
+                )
+    
+    def _compute_gene_alignment(self):
+        """Pre-compute gene alignment indices."""
+        common = self._real_var_names.intersection(self._generated_var_names)
+        
+        if len(common) == 0:
+            raise LazyDataLoaderError(
+                "No overlapping genes between real and generated data."
+            )
+        
+        self._common_genes = common.tolist()
+        self._real_gene_idx = self._real_var_names.get_indexer(common)
+        self._gen_gene_idx = self._generated_var_names.get_indexer(common)
+        
+        n_real_only = len(self._real_var_names) - len(common)
+        n_gen_only = len(self._generated_var_names) - len(common)
+        
+        if n_real_only > 0 or n_gen_only > 0:
+            warnings.warn(
+                f"Gene alignment: {len(common)} common genes. "
+                f"Dropped {n_real_only} from real, {n_gen_only} from generated."
+            )
+    
+    def _get_condition_key(self, row: pd.Series) -> str:
+        """Generate unique key for a condition combination."""
+        return "####".join([str(row[c]) for c in self.condition_columns])
+    
+    def _precompute_condition_indices(self):
+        """Pre-compute sample indices for each condition (lightweight)."""
+        self._condition_indices = {"real": {}, "generated": {}}
+        
+        # Real data conditions
+        real_conditions = self._real_obs[self.condition_columns].astype(str).drop_duplicates()
+        for _, row in real_conditions.iterrows():
+            key = self._get_condition_key(row)
+            mask = np.ones(len(self._real_obs), dtype=bool)
+            for col in self.condition_columns:
+                mask &= (self._real_obs[col].astype(str) == str(row[col])).values
+            
+            indices = np.where(mask)[0]
+            if len(indices) >= self.min_samples_per_condition:
+                self._condition_indices["real"][key] = indices
+        
+        # Generated data conditions
+        gen_conditions = self._generated_obs[self.condition_columns].astype(str).drop_duplicates()
+        for _, row in gen_conditions.iterrows():
+            key = self._get_condition_key(row)
+            mask = np.ones(len(self._generated_obs), dtype=bool)
+            for col in self.condition_columns:
+                mask &= (self._generated_obs[col].astype(str) == str(row[col])).values
+            
+            indices = np.where(mask)[0]
+            if len(indices) >= self.min_samples_per_condition:
+                self._condition_indices["generated"][key] = indices
+    
+    def get_splits(self) -> List[str]:
+        """Get available splits."""
+        if not self._is_initialized:
+            self.initialize()
+        
+        if self.split_column is None or self.split_column not in self._real_obs.columns:
+            return ["all"]
+        
+        return list(self._real_obs[self.split_column].astype(str).unique())
+    
+    def get_common_conditions(self, split: Optional[str] = None) -> List[str]:
+        """Get conditions present in both real and generated data."""
+        if not self._is_initialized:
+            self.initialize()
+        
+        real_keys = set(self._condition_indices["real"].keys())
+        gen_keys = set(self._condition_indices["generated"].keys())
+        
+        common = real_keys & gen_keys
+        
+        # Filter by split if specified
+        if split is not None and split != "all" and self.split_column is not None:
+            filtered = set()
+            for key in common:
+                real_idx = self._condition_indices["real"][key]
+                split_vals = self._real_obs.iloc[real_idx][self.split_column].astype(str)
+                if (split_vals == split).any():
+                    filtered.add(key)
+            common = filtered
+        
+        return sorted(common)
+    
+    def _extract_data_subset(
+        self,
+        adata: ad.AnnData,
+        indices: np.ndarray,
+        gene_idx: np.ndarray,
+    ) -> np.ndarray:
+        """Extract and align a subset of data."""
+        # Handle backed vs loaded data
+        if hasattr(adata, 'isbacked') and adata.isbacked:
+            # Backed mode: read only what we need
+            X = adata.X[indices][:, gene_idx]
+        else:
+            # Standard mode
+            X = adata.X[indices][:, gene_idx]
+        
+        # Convert to dense if sparse
+        if sparse.issparse(X):
+            X = X.toarray()
+        
+        return np.asarray(X, dtype=np.float32)
+    
+    def iterate_conditions(
+        self,
+        split: Optional[str] = None,
+    ) -> Generator[Tuple[str, np.ndarray, np.ndarray, Dict[str, str]], None, None]:
+        """
+        Iterate over conditions, loading one condition at a time.
+        
+        This loads data for one condition, yields it, then releases memory
+        before loading the next condition.
+        
+        Parameters
+        ----------
+        split : str, optional
+            Filter to this split only
+            
+        Yields
+        ------
+        Tuple[str, np.ndarray, np.ndarray, Dict[str, str]]
+            (condition_key, real_data, generated_data, condition_info)
+        """
+        if not self._is_initialized:
+            self.initialize()
+        
+        common_conditions = self.get_common_conditions(split)
+        
+        for key in common_conditions:
+            real_indices = self._condition_indices["real"][key]
+            gen_indices = self._condition_indices["generated"][key]
+            
+            # Filter by split if needed
+            if split is not None and split != "all" and self.split_column is not None:
+                split_mask = self._real_obs.iloc[real_indices][self.split_column].astype(str) == split
+                real_indices = real_indices[split_mask.values]
+            
+            if len(real_indices) < self.min_samples_per_condition:
+                continue
+            
+            # Load data for this condition only
+            real_data = self._extract_data_subset(
+                self._real, real_indices, self._real_gene_idx
+            )
+            gen_data = self._extract_data_subset(
+                self._generated, gen_indices, self._gen_gene_idx
+            )
+            
+            # Parse condition info
+            parts = key.split("####")
+            condition_info = dict(zip(self.condition_columns, parts))
+            
+            yield key, real_data, gen_data, condition_info
+    
+    def iterate_conditions_batched(
+        self,
+        split: Optional[str] = None,
+        batch_size: Optional[int] = None,
+    ) -> Generator[ConditionBatch, None, None]:
+        """
+        Iterate over conditions in batches for memory efficiency.
+        
+        Useful when even a single condition is too large to fit in memory.
+        
+        Parameters
+        ----------
+        split : str, optional
+            Filter to this split only
+        batch_size : int, optional
+            Override default batch size
+            
+        Yields
+        ------
+        ConditionBatch
+            Batch of samples from a condition
+        """
+        if not self._is_initialized:
+            self.initialize()
+        
+        batch_size = batch_size or self.batch_size
+        common_conditions = self.get_common_conditions(split)
+        
+        for key in common_conditions:
+            real_indices = self._condition_indices["real"][key]
+            gen_indices = self._condition_indices["generated"][key]
+            
+            # Filter by split
+            if split is not None and split != "all" and self.split_column is not None:
+                split_mask = self._real_obs.iloc[real_indices][self.split_column].astype(str) == split
+                real_indices = real_indices[split_mask.values]
+            
+            if len(real_indices) < self.min_samples_per_condition:
+                continue
+            
+            # Parse condition info
+            parts = key.split("####")
+            condition_info = dict(zip(self.condition_columns, parts))
+            
+            # Calculate number of batches (use max of real/gen for alignment)
+            n_real = len(real_indices)
+            n_gen = len(gen_indices)
+            n_batches = max(
+                (n_real + batch_size - 1) // batch_size,
+                (n_gen + batch_size - 1) // batch_size
+            )
+            
+            for batch_idx in range(n_batches):
+                start_real = batch_idx * batch_size
+                end_real = min(start_real + batch_size, n_real)
+                start_gen = batch_idx * batch_size
+                end_gen = min(start_gen + batch_size, n_gen)
+                
+                # Handle case where one dataset is smaller
+                if start_real >= n_real:
+                    # Wrap around for real data
+                    batch_real_idx = real_indices[start_real % n_real:end_real % n_real + 1]
+                else:
+                    batch_real_idx = real_indices[start_real:end_real]
+                
+                if start_gen >= n_gen:
+                    batch_gen_idx = gen_indices[start_gen % n_gen:end_gen % n_gen + 1]
+                else:
+                    batch_gen_idx = gen_indices[start_gen:end_gen]
+                
+                if len(batch_real_idx) == 0 or len(batch_gen_idx) == 0:
+                    continue
+                
+                real_data = self._extract_data_subset(
+                    self._real, batch_real_idx, self._real_gene_idx
+                )
+                gen_data = self._extract_data_subset(
+                    self._generated, batch_gen_idx, self._gen_gene_idx
+                )
+                
+                yield ConditionBatch(
+                    condition_key=key,
+                    condition_info=condition_info,
+                    real_data=real_data,
+                    generated_data=gen_data,
+                    batch_idx=batch_idx,
+                    n_batches=n_batches,
+                    is_last_batch=(batch_idx == n_batches - 1),
+                )
+    
+    @property
+    def gene_names(self) -> List[str]:
+        """Get common gene names."""
+        if not self._is_initialized:
+            self.initialize()
+        return self._common_genes
+    
+    @property
+    def n_genes(self) -> int:
+        """Number of common genes."""
+        return len(self.gene_names)
+    
+    def estimate_memory_usage(self) -> Dict[str, float]:
+        """
+        Estimate memory usage in MB for different loading strategies.
+        
+        Returns
+        -------
+        Dict[str, float]
+            Memory estimates in MB
+        """
+        if not self._is_initialized:
+            self.initialize()
+        
+        n_real = len(self._real_obs)
+        n_gen = len(self._generated_obs)
+        n_genes = self.n_genes
+        
+        # 4 bytes per float32
+        bytes_per_element = 4
+        
+        full_real = n_real * n_genes * bytes_per_element / 1e6
+        full_gen = n_gen * n_genes * bytes_per_element / 1e6
+        
+        # Average condition size
+        n_conditions = len(self.get_common_conditions())
+        avg_per_condition = (n_real + n_gen) / max(n_conditions, 1)
+        per_condition = avg_per_condition * n_genes * bytes_per_element / 1e6
+        
+        per_batch = self.batch_size * 2 * n_genes * bytes_per_element / 1e6
+        
+        return {
+            "full_load_mb": full_real + full_gen,
+            "per_condition_mb": per_condition,
+            "per_batch_mb": per_batch,
+            "metadata_mb": (n_real + n_gen) * 100 / 1e6,  # rough obs estimate
+        }
+    
+    def close(self):
+        """Close backed file handles if any."""
+        if self._real is not None and hasattr(self._real, 'file'):
+            try:
+                self._real.file.close()
+            except:
+                pass
+        if self._generated is not None and hasattr(self._generated, 'file'):
+            try:
+                self._generated.file.close()
+            except:
+                pass
+    
+    def __enter__(self):
+        """Context manager entry."""
+        self.initialize()
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        self.close()
+    
+    def __repr__(self) -> str:
+        if not self._is_initialized:
+            return f"LazyGeneExpressionDataLoader(not initialized, backed={self.use_backed})"
+        
+        return (
+            f"LazyGeneExpressionDataLoader("
+            f"real={len(self._real_obs)}x{len(self._real_var_names)}, "
+            f"gen={len(self._generated_obs)}x{len(self._generated_var_names)}, "
+            f"common_genes={self.n_genes}, "
+            f"batch_size={self.batch_size})"
+        )
+
+
+def load_data_lazy(
+    real_path: Union[str, Path],
+    generated_path: Union[str, Path],
+    condition_columns: List[str],
+    split_column: Optional[str] = None,
+    batch_size: int = 256,
+    use_backed: bool = False,
+    **kwargs
+) -> LazyGeneExpressionDataLoader:
+    """
+    Convenience function to create a lazy data loader.
+    
+    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
+    batch_size : int
+        Maximum samples per batch
+    use_backed : bool
+        Use memory-mapped access for very large files
+    **kwargs
+        Additional arguments for LazyGeneExpressionDataLoader
+        
+    Returns
+    -------
+    LazyGeneExpressionDataLoader
+        Initialized lazy data loader
+    """
+    loader = LazyGeneExpressionDataLoader(
+        real_path=real_path,
+        generated_path=generated_path,
+        condition_columns=condition_columns,
+        split_column=split_column,
+        batch_size=batch_size,
+        use_backed=use_backed,
+        **kwargs
+    )
+    loader.initialize()
+    return loader