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

geneval/metrics/accelerated.py

@@ -0,0 +1,857 @@
+"""
+Accelerated metric computation with CPU parallelization and GPU support.
+
+This module provides performance optimizations for metric computation:
+- CPU parallelization via joblib for multi-core speedup
+- GPU acceleration via PyTorch/geomloss for batch computation
+- Vectorized operations for improved NumPy performance
+
+Example usage:
+    >>> from geneval.metrics.accelerated import ParallelMetricComputer
+    >>> computer = ParallelMetricComputer(n_jobs=8, device="cuda")
+    >>> results = computer.compute_all(real, generated, metrics)
+"""
+from __future__ import annotations
+
+import warnings
+from typing import List, Optional, Dict, Any, Union, Literal
+from dataclasses import dataclass
+import numpy as np
+
+from .base_metric import BaseMetric, MetricResult
+
+
+# Check for optional dependencies
+try:
+    from joblib import Parallel, delayed
+    HAS_JOBLIB = True
+except ImportError:
+    HAS_JOBLIB = False
+
+try:
+    import torch
+    HAS_TORCH = True
+except ImportError:
+    HAS_TORCH = False
+
+try:
+    from geomloss import SamplesLoss
+    HAS_GEOMLOSS = True
+except ImportError:
+    HAS_GEOMLOSS = False
+
+
+@dataclass
+class AccelerationConfig:
+    """Configuration for accelerated metric computation.
+    
+    Attributes
+    ----------
+    n_jobs : int
+        Number of CPU jobs for parallel computation. 
+        -1 uses all available cores. Default is 1 (no parallelization).
+    device : str
+        Device for computation: "cpu", "cuda", "cuda:0", etc.
+        Default is "cpu".
+    batch_genes : bool
+        If True, batch all genes for GPU computation. Default is True.
+    gene_batch_size : int or None
+        If set, process genes in batches of this size to manage memory.
+        None means process all genes at once.
+    prefer_gpu : bool
+        If True and GPU is available, prefer GPU implementations.
+        Default is True.
+    verbose : bool
+        Print acceleration info. Default is False.
+    """
+    n_jobs: int = 1
+    device: str = "cpu"
+    batch_genes: bool = True
+    gene_batch_size: Optional[int] = None
+    prefer_gpu: bool = True
+    verbose: bool = False
+
+
+def get_available_backends() -> Dict[str, bool]:
+    """Check which acceleration backends are available.
+    
+    Returns
+    -------
+    Dict[str, bool]
+        Dictionary with backend availability.
+    """
+    backends = {
+        "joblib": HAS_JOBLIB,
+        "torch": HAS_TORCH,
+        "geomloss": HAS_GEOMLOSS,
+        "cuda": HAS_TORCH and torch.cuda.is_available(),
+        "mps": HAS_TORCH and hasattr(torch.backends, "mps") and torch.backends.mps.is_available(),
+    }
+    return backends
+
+
+def _get_device(device: str) -> "torch.device":
+    """Get PyTorch device, handling availability checks.
+    
+    Parameters
+    ----------
+    device : str
+        Device string ("cpu", "cuda", "cuda:0", "mps", "auto")
+        
+    Returns
+    -------
+    torch.device
+        PyTorch device object
+    """
+    if not HAS_TORCH:
+        raise ImportError("PyTorch is required for GPU acceleration")
+    
+    if device == "auto":
+        if torch.cuda.is_available():
+            return torch.device("cuda")
+        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            return torch.device("mps")
+        else:
+            return torch.device("cpu")
+    
+    return torch.device(device)
+
+
+class ParallelMetricComputer:
+    """Parallel and GPU-accelerated metric computation.
+    
+    This class wraps metric computation with parallelization and GPU
+    acceleration options for significant speedups on large datasets.
+    
+    Parameters
+    ----------
+    n_jobs : int
+        Number of parallel jobs. -1 for all cores.
+    device : str
+        Compute device ("cpu", "cuda", "auto")
+    batch_genes : bool
+        Whether to batch genes for GPU computation.
+    gene_batch_size : int, optional
+        Process genes in chunks of this size.
+    verbose : bool
+        Print progress information.
+        
+    Examples
+    --------
+    >>> computer = ParallelMetricComputer(n_jobs=8)
+    >>> results = computer.compute_metric(metric, real, generated)
+    
+    >>> # GPU acceleration
+    >>> computer = ParallelMetricComputer(device="cuda")
+    >>> results = computer.compute_metric(metric, real, generated)
+    """
+    
+    def __init__(
+        self,
+        n_jobs: int = 1,
+        device: str = "cpu",
+        batch_genes: bool = True,
+        gene_batch_size: Optional[int] = None,
+        verbose: bool = False,
+    ):
+        self.n_jobs = n_jobs
+        self.device = device
+        self.batch_genes = batch_genes
+        self.gene_batch_size = gene_batch_size
+        self.verbose = verbose
+        
+        # Validate configuration
+        if n_jobs != 1 and not HAS_JOBLIB:
+            warnings.warn("joblib not available, falling back to sequential processing")
+            self.n_jobs = 1
+        
+        if device != "cpu" and not HAS_TORCH:
+            warnings.warn("PyTorch not available, falling back to CPU")
+            self.device = "cpu"
+        
+        if self.verbose:
+            backends = get_available_backends()
+            print(f"Acceleration backends: {backends}")
+            print(f"Using n_jobs={self.n_jobs}, device={self.device}")
+    
+    def compute_metric_parallel(
+        self,
+        metric: BaseMetric,
+        real: np.ndarray,
+        generated: np.ndarray,
+        gene_names: Optional[List[str]] = None,
+    ) -> MetricResult:
+        """Compute a metric with CPU parallelization.
+        
+        Splits genes across multiple CPU cores for parallel computation.
+        
+        Parameters
+        ----------
+        metric : BaseMetric
+            Metric to compute
+        real : np.ndarray
+            Real data, shape (n_samples, n_genes)
+        generated : np.ndarray
+            Generated data, shape (n_samples, n_genes)
+        gene_names : List[str], optional
+            Gene names
+            
+        Returns
+        -------
+        MetricResult
+            Computed metric result
+        """
+        n_genes = real.shape[1]
+        if gene_names is None:
+            gene_names = [f"gene_{i}" for i in range(n_genes)]
+        
+        if self.n_jobs == 1 or not HAS_JOBLIB:
+            # Sequential computation
+            per_gene = metric.compute_per_gene(real, generated)
+        else:
+            # Parallel computation across genes
+            if self.gene_batch_size:
+                # Process in batches
+                batches = [
+                    (i, min(i + self.gene_batch_size, n_genes))
+                    for i in range(0, n_genes, self.gene_batch_size)
+                ]
+            else:
+                # Split evenly across jobs
+                n_effective_jobs = min(self.n_jobs if self.n_jobs > 0 else 8, n_genes)
+                batch_size = max(1, n_genes // n_effective_jobs)
+                batches = [
+                    (i, min(i + batch_size, n_genes))
+                    for i in range(0, n_genes, batch_size)
+                ]
+            
+            def compute_batch(start: int, end: int) -> np.ndarray:
+                return metric.compute_per_gene(
+                    real[:, start:end],
+                    generated[:, start:end]
+                )
+            
+            results = Parallel(n_jobs=self.n_jobs, prefer="threads")(
+                delayed(compute_batch)(start, end) for start, end in batches
+            )
+            
+            per_gene = np.concatenate(results)
+        
+        aggregate = metric.compute_aggregate(per_gene, method="mean")
+        
+        return MetricResult(
+            name=metric.name,
+            per_gene_values=per_gene,
+            gene_names=gene_names,
+            aggregate_value=aggregate,
+            aggregate_method="mean",
+            metadata={
+                "higher_is_better": metric.higher_is_better,
+                "accelerated": True,
+                "n_jobs": self.n_jobs,
+            }
+        )
+
+
+# =============================================================================
+# GPU-Accelerated Distance Metrics
+# =============================================================================
+
+class GPUWasserstein1:
+    """GPU-accelerated Wasserstein-1 distance computation.
+    
+    Computes W1 distance for all genes in parallel on GPU using
+    vectorized sorting and quantile interpolation.
+    """
+    
+    def __init__(self, device: str = "cuda"):
+        if not HAS_TORCH:
+            raise ImportError("PyTorch required for GPU acceleration")
+        self.device = _get_device(device)
+    
+    def compute_batch(
+        self,
+        real: np.ndarray,
+        generated: np.ndarray,
+    ) -> np.ndarray:
+        """Compute W1 for all genes in batch on GPU.
+        
+        Parameters
+        ----------
+        real : np.ndarray
+            Real data, shape (n_samples_real, n_genes)
+        generated : np.ndarray
+            Generated data, shape (n_samples_gen, n_genes)
+            
+        Returns
+        -------
+        np.ndarray
+            W1 distance per gene
+        """
+        # Move to GPU
+        real_t = torch.tensor(real, dtype=torch.float32, device=self.device)
+        gen_t = torch.tensor(generated, dtype=torch.float32, device=self.device)
+        
+        n_genes = real_t.shape[1]
+        n_quantiles = max(real_t.shape[0], gen_t.shape[0])
+        
+        # Sort each gene column
+        real_sorted, _ = torch.sort(real_t, dim=0)
+        gen_sorted, _ = torch.sort(gen_t, dim=0)
+        
+        # Interpolate to same number of quantiles
+        quantile_positions = torch.linspace(0, 1, n_quantiles, device=self.device)
+        
+        # Interpolate real
+        real_indices = quantile_positions * (real_sorted.shape[0] - 1)
+        real_floor = real_indices.long().clamp(0, real_sorted.shape[0] - 2)
+        real_frac = (real_indices - real_floor.float()).unsqueeze(1)
+        real_interp = (
+            real_sorted[real_floor] * (1 - real_frac) + 
+            real_sorted[real_floor + 1] * real_frac
+        )
+        
+        # Interpolate generated
+        gen_indices = quantile_positions * (gen_sorted.shape[0] - 1)
+        gen_floor = gen_indices.long().clamp(0, gen_sorted.shape[0] - 2)
+        gen_frac = (gen_indices - gen_floor.float()).unsqueeze(1)
+        gen_interp = (
+            gen_sorted[gen_floor] * (1 - gen_frac) + 
+            gen_sorted[gen_floor + 1] * gen_frac
+        )
+        
+        # W1 = mean absolute difference
+        w1 = torch.mean(torch.abs(real_interp - gen_interp), dim=0)
+        
+        return w1.cpu().numpy()
+
+
+class GPUWasserstein2:
+    """GPU-accelerated Wasserstein-2 distance using geomloss.
+    
+    Batches all genes together for efficient GPU computation.
+    """
+    
+    def __init__(self, device: str = "cuda", blur: float = 0.01):
+        if not HAS_TORCH:
+            raise ImportError("PyTorch required for GPU acceleration")
+        if not HAS_GEOMLOSS:
+            raise ImportError("geomloss required for Wasserstein-2 GPU acceleration")
+        
+        self.device = _get_device(device)
+        self.blur = blur
+        self.loss_fn = SamplesLoss(loss="sinkhorn", p=2, blur=blur, backend="tensorized")
+    
+    def compute_batch(
+        self,
+        real: np.ndarray,
+        generated: np.ndarray,
+    ) -> np.ndarray:
+        """Compute W2 for all genes in batch on GPU.
+        
+        Parameters
+        ----------
+        real : np.ndarray
+            Real data, shape (n_samples_real, n_genes)
+        generated : np.ndarray
+            Generated data, shape (n_samples_gen, n_genes)
+            
+        Returns
+        -------
+        np.ndarray
+            W2 distance per gene
+        """
+        n_genes = real.shape[1]
+        
+        # Move to GPU
+        real_t = torch.tensor(real, dtype=torch.float32, device=self.device)
+        gen_t = torch.tensor(generated, dtype=torch.float32, device=self.device)
+        
+        distances = torch.zeros(n_genes, device=self.device)
+        
+        # Process each gene (geomloss requires separate calls per distribution pair)
+        # But we can batch by treating genes as batch dimension
+        for i in range(n_genes):
+            r = real_t[:, i:i+1]  # Keep 2D
+            g = gen_t[:, i:i+1]
+            distances[i] = self.loss_fn(r, g)
+        
+        return distances.cpu().numpy()
+
+
+class GPUMMD:
+    """GPU-accelerated MMD computation with RBF kernel.
+    
+    Uses PyTorch for vectorized kernel computation across all genes.
+    """
+    
+    def __init__(self, device: str = "cuda", sigma: Optional[float] = None):
+        if not HAS_TORCH:
+            raise ImportError("PyTorch required for GPU acceleration")
+        
+        self.device = _get_device(device)
+        self.sigma = sigma
+    
+    def compute_batch(
+        self,
+        real: np.ndarray,
+        generated: np.ndarray,
+    ) -> np.ndarray:
+        """Compute MMD for all genes in batch on GPU.
+        
+        Parameters
+        ----------
+        real : np.ndarray
+            Real data, shape (n_samples_real, n_genes)
+        generated : np.ndarray
+            Generated data, shape (n_samples_gen, n_genes)
+            
+        Returns
+        -------
+        np.ndarray
+            MMD per gene
+        """
+        real_t = torch.tensor(real, dtype=torch.float32, device=self.device)
+        gen_t = torch.tensor(generated, dtype=torch.float32, device=self.device)
+        
+        n_genes = real_t.shape[1]
+        n_x, n_y = real_t.shape[0], gen_t.shape[0]
+        
+        mmd_values = torch.zeros(n_genes, device=self.device)
+        
+        for g in range(n_genes):
+            x = real_t[:, g:g+1]
+            y = gen_t[:, g:g+1]
+            
+            # Median heuristic for sigma
+            if self.sigma is None:
+                combined = torch.cat([x, y], dim=0)
+                pairwise = torch.abs(combined - combined.T)
+                sigma = torch.median(pairwise[pairwise > 0]).item()
+                if sigma == 0:
+                    sigma = 1.0
+            else:
+                sigma = self.sigma
+            
+            # RBF kernel
+            def rbf(a, b, s):
+                sq_dist = (a - b.T) ** 2
+                return torch.exp(-sq_dist / (2 * s ** 2))
+            
+            K_xx = rbf(x, x, sigma)
+            K_yy = rbf(y, y, sigma)
+            K_xy = rbf(x, y, sigma)
+            
+            # Unbiased MMD
+            mmd = (
+                (K_xx.sum() - K_xx.trace()) / (n_x * (n_x - 1)) +
+                (K_yy.sum() - K_yy.trace()) / (n_y * (n_y - 1)) -
+                2 * K_xy.sum() / (n_x * n_y)
+            )
+            
+            mmd_values[g] = torch.clamp(mmd, min=0)
+        
+        return mmd_values.cpu().numpy()
+
+
+class GPUEnergyDistance:
+    """GPU-accelerated Energy distance computation."""
+    
+    def __init__(self, device: str = "cuda"):
+        if not HAS_TORCH:
+            raise ImportError("PyTorch required for GPU acceleration")
+        
+        self.device = _get_device(device)
+    
+    def compute_batch(
+        self,
+        real: np.ndarray,
+        generated: np.ndarray,
+    ) -> np.ndarray:
+        """Compute Energy distance for all genes in batch on GPU.
+        
+        Parameters
+        ----------
+        real : np.ndarray
+            Real data, shape (n_samples_real, n_genes)
+        generated : np.ndarray
+            Generated data, shape (n_samples_gen, n_genes)
+            
+        Returns
+        -------
+        np.ndarray
+            Energy distance per gene
+        """
+        real_t = torch.tensor(real, dtype=torch.float32, device=self.device)
+        gen_t = torch.tensor(generated, dtype=torch.float32, device=self.device)
+        
+        n_genes = real_t.shape[1]
+        
+        energy_values = torch.zeros(n_genes, device=self.device)
+        
+        for g in range(n_genes):
+            x = real_t[:, g]
+            y = gen_t[:, g]
+            
+            # E[|X - Y|]
+            xy_dist = torch.mean(torch.abs(x.unsqueeze(1) - y.unsqueeze(0)))
+            
+            # E[|X - X'|]
+            xx_dist = torch.mean(torch.abs(x.unsqueeze(1) - x.unsqueeze(0)))
+            
+            # E[|Y - Y'|]
+            yy_dist = torch.mean(torch.abs(y.unsqueeze(1) - y.unsqueeze(0)))
+            
+            energy = 2 * xy_dist - xx_dist - yy_dist
+            energy_values[g] = torch.clamp(energy, min=0)
+        
+        return energy_values.cpu().numpy()
+
+
+# =============================================================================
+# Vectorized NumPy Implementations (for CPU speedup without joblib)
+# =============================================================================
+
+def vectorized_wasserstein1(
+    real: np.ndarray,
+    generated: np.ndarray,
+) -> np.ndarray:
+    """Compute W1 for all genes using vectorized NumPy.
+    
+    This is faster than the loop-based scipy implementation.
+    
+    Parameters
+    ----------
+    real : np.ndarray
+        Real data, shape (n_samples_real, n_genes)
+    generated : np.ndarray
+        Generated data, shape (n_samples_gen, n_genes)
+        
+    Returns
+    -------
+    np.ndarray
+        W1 distance per gene
+    """
+    n_genes = real.shape[1]
+    n_quantiles = max(real.shape[0], generated.shape[0])
+    
+    # Sort each column
+    real_sorted = np.sort(real, axis=0)
+    gen_sorted = np.sort(generated, axis=0)
+    
+    # Interpolate to same number of quantiles
+    real_positions = np.linspace(0, 1, real_sorted.shape[0])
+    gen_positions = np.linspace(0, 1, gen_sorted.shape[0])
+    target_positions = np.linspace(0, 1, n_quantiles)
+    
+    # Interpolate each gene column
+    real_interp = np.zeros((n_quantiles, n_genes))
+    gen_interp = np.zeros((n_quantiles, n_genes))
+    
+    for g in range(n_genes):
+        real_interp[:, g] = np.interp(target_positions, real_positions, real_sorted[:, g])
+        gen_interp[:, g] = np.interp(target_positions, gen_positions, gen_sorted[:, g])
+    
+    # W1 = mean absolute difference
+    return np.mean(np.abs(real_interp - gen_interp), axis=0)
+
+
+def vectorized_mmd(
+    real: np.ndarray,
+    generated: np.ndarray,
+    sigma: Optional[float] = None,
+) -> np.ndarray:
+    """Compute MMD for all genes using vectorized NumPy.
+    
+    Parameters
+    ----------
+    real : np.ndarray
+        Real data, shape (n_samples_real, n_genes)
+    generated : np.ndarray
+        Generated data, shape (n_samples_gen, n_genes)
+    sigma : float, optional
+        Kernel bandwidth. Uses median heuristic if None.
+        
+    Returns
+    -------
+    np.ndarray
+        MMD per gene
+    """
+    n_genes = real.shape[1]
+    n_x, n_y = real.shape[0], generated.shape[0]
+    
+    mmd_values = np.zeros(n_genes)
+    
+    for g in range(n_genes):
+        x = real[:, g:g+1]
+        y = generated[:, g:g+1]
+        
+        # Median heuristic
+        if sigma is None:
+            combined = np.vstack([x, y])
+            pairwise = np.abs(combined - combined.T)
+            s = float(np.median(pairwise[pairwise > 0]))
+            if s == 0:
+                s = 1.0
+        else:
+            s = sigma
+        
+        # RBF kernel
+        K_xx = np.exp(-(x - x.T) ** 2 / (2 * s ** 2))
+        K_yy = np.exp(-(y - y.T) ** 2 / (2 * s ** 2))
+        K_xy = np.exp(-(x - y.T) ** 2 / (2 * s ** 2))
+        
+        # Unbiased MMD
+        mmd = (
+            (np.sum(K_xx) - np.trace(K_xx)) / (n_x * (n_x - 1)) +
+            (np.sum(K_yy) - np.trace(K_yy)) / (n_y * (n_y - 1)) -
+            2 * np.sum(K_xy) / (n_x * n_y)
+        )
+        
+        mmd_values[g] = max(0, mmd)
+    
+    return mmd_values
+
+
+# =============================================================================
+# High-Level Accelerated Evaluation Interface
+# =============================================================================
+
+def compute_metrics_accelerated(
+    real: np.ndarray,
+    generated: np.ndarray,
+    metrics: List[str] = ["wasserstein_1", "wasserstein_2", "mmd", "energy"],
+    n_jobs: int = 1,
+    device: str = "cpu",
+    gene_names: Optional[List[str]] = None,
+    verbose: bool = False,
+) -> Dict[str, MetricResult]:
+    """Compute multiple metrics with acceleration.
+    
+    This is the main entry point for accelerated metric computation.
+    Automatically selects the best available backend.
+    
+    Parameters
+    ----------
+    real : np.ndarray
+        Real data, shape (n_samples_real, n_genes)
+    generated : np.ndarray
+        Generated data, shape (n_samples_gen, n_genes)
+    metrics : List[str]
+        Metrics to compute: "wasserstein_1", "wasserstein_2", "mmd", "energy"
+    n_jobs : int
+        Number of CPU jobs (-1 for all cores)
+    device : str
+        Compute device ("cpu", "cuda", "auto")
+    gene_names : List[str], optional
+        Gene names
+    verbose : bool
+        Print progress
+        
+    Returns
+    -------
+    Dict[str, MetricResult]
+        Dictionary of metric results
+    """
+    backends = get_available_backends()
+    
+    if device == "auto":
+        if backends["cuda"]:
+            device = "cuda"
+        elif backends["mps"]:
+            device = "mps"
+        else:
+            device = "cpu"
+    
+    if verbose:
+        print(f"Using device: {device}, n_jobs: {n_jobs}")
+        print(f"Available backends: {backends}")
+    
+    n_genes = real.shape[1]
+    if gene_names is None:
+        gene_names = [f"gene_{i}" for i in range(n_genes)]
+    
+    results = {}
+    
+    for metric_name in metrics:
+        if verbose:
+            print(f"Computing {metric_name}...")
+        
+        if device != "cpu" and backends["torch"]:
+            # GPU path
+            if metric_name == "wasserstein_1":
+                gpu_metric = GPUWasserstein1(device=device)
+                per_gene = gpu_metric.compute_batch(real, generated)
+            elif metric_name == "wasserstein_2" and backends["geomloss"]:
+                gpu_metric = GPUWasserstein2(device=device)
+                per_gene = gpu_metric.compute_batch(real, generated)
+            elif metric_name == "mmd":
+                gpu_metric = GPUMMD(device=device)
+                per_gene = gpu_metric.compute_batch(real, generated)
+            elif metric_name == "energy":
+                gpu_metric = GPUEnergyDistance(device=device)
+                per_gene = gpu_metric.compute_batch(real, generated)
+            else:
+                # Fallback to vectorized CPU
+                per_gene = _compute_cpu_metric(metric_name, real, generated, n_jobs)
+        else:
+            # CPU path
+            per_gene = _compute_cpu_metric(metric_name, real, generated, n_jobs)
+        
+        results[metric_name] = MetricResult(
+            name=metric_name,
+            per_gene_values=per_gene,
+            gene_names=gene_names,
+            aggregate_value=float(np.nanmean(per_gene)),
+            aggregate_method="mean",
+            metadata={
+                "device": device,
+                "n_jobs": n_jobs,
+                "accelerated": True,
+            }
+        )
+    
+    return results
+
+
+def _compute_cpu_metric(
+    metric_name: str,
+    real: np.ndarray,
+    generated: np.ndarray,
+    n_jobs: int,
+) -> np.ndarray:
+    """Compute metric on CPU with optional parallelization."""
+    if metric_name == "wasserstein_1":
+        if n_jobs != 1 and HAS_JOBLIB:
+            return _parallel_w1(real, generated, n_jobs)
+        else:
+            return vectorized_wasserstein1(real, generated)
+    elif metric_name == "wasserstein_2":
+        return _compute_w2_cpu(real, generated, n_jobs)
+    elif metric_name == "mmd":
+        if n_jobs != 1 and HAS_JOBLIB:
+            return _parallel_mmd(real, generated, n_jobs)
+        else:
+            return vectorized_mmd(real, generated)
+    elif metric_name == "energy":
+        return _compute_energy_cpu(real, generated, n_jobs)
+    else:
+        raise ValueError(f"Unknown metric: {metric_name}")
+
+
+def _parallel_w1(real: np.ndarray, generated: np.ndarray, n_jobs: int) -> np.ndarray:
+    """Parallel W1 computation."""
+    from scipy.stats import wasserstein_distance
+    
+    n_genes = real.shape[1]
+    
+    def compute_single(g):
+        r = real[:, g]
+        gen = generated[:, g]
+        r = r[~np.isnan(r)]
+        gen = gen[~np.isnan(gen)]
+        if len(r) == 0 or len(gen) == 0:
+            return np.nan
+        return wasserstein_distance(r, gen)
+    
+    results = Parallel(n_jobs=n_jobs)(
+        delayed(compute_single)(g) for g in range(n_genes)
+    )
+    
+    return np.array(results)
+
+
+def _parallel_mmd(real: np.ndarray, generated: np.ndarray, n_jobs: int) -> np.ndarray:
+    """Parallel MMD computation."""
+    n_genes = real.shape[1]
+    
+    def compute_single(g):
+        x = real[:, g:g+1]
+        y = generated[:, g:g+1]
+        
+        combined = np.vstack([x, y])
+        pairwise = np.abs(combined - combined.T)
+        sigma = float(np.median(pairwise[pairwise > 0]))
+        if sigma == 0:
+            sigma = 1.0
+        
+        n_x, n_y = len(x), len(y)
+        
+        K_xx = np.exp(-(x - x.T) ** 2 / (2 * sigma ** 2))
+        K_yy = np.exp(-(y - y.T) ** 2 / (2 * sigma ** 2))
+        K_xy = np.exp(-(x - y.T) ** 2 / (2 * sigma ** 2))
+        
+        mmd = (
+            (np.sum(K_xx) - np.trace(K_xx)) / (n_x * (n_x - 1)) +
+            (np.sum(K_yy) - np.trace(K_yy)) / (n_y * (n_y - 1)) -
+            2 * np.sum(K_xy) / (n_x * n_y)
+        )
+        
+        return max(0, mmd)
+    
+    results = Parallel(n_jobs=n_jobs)(
+        delayed(compute_single)(g) for g in range(n_genes)
+    )
+    
+    return np.array(results)
+
+
+def _compute_w2_cpu(real: np.ndarray, generated: np.ndarray, n_jobs: int) -> np.ndarray:
+    """CPU W2 computation (quantile-based)."""
+    n_genes = real.shape[1]
+    
+    def compute_single(g):
+        r = real[:, g]
+        gen = generated[:, g]
+        
+        r = r[~np.isnan(r)]
+        gen = gen[~np.isnan(gen)]
+        
+        if len(r) == 0 or len(gen) == 0:
+            return np.nan
+        
+        r_sorted = np.sort(r)
+        g_sorted = np.sort(gen)
+        
+        n = max(len(r_sorted), len(g_sorted))
+        r_q = np.interp(np.linspace(0, 1, n), np.linspace(0, 1, len(r_sorted)), r_sorted)
+        g_q = np.interp(np.linspace(0, 1, n), np.linspace(0, 1, len(g_sorted)), g_sorted)
+        
+        return np.sqrt(np.mean((r_q - g_q) ** 2))
+    
+    if n_jobs != 1 and HAS_JOBLIB:
+        results = Parallel(n_jobs=n_jobs)(
+            delayed(compute_single)(g) for g in range(n_genes)
+        )
+        return np.array(results)
+    else:
+        return np.array([compute_single(g) for g in range(n_genes)])
+
+
+def _compute_energy_cpu(real: np.ndarray, generated: np.ndarray, n_jobs: int) -> np.ndarray:
+    """CPU Energy distance computation."""
+    n_genes = real.shape[1]
+    
+    def compute_single(g):
+        x = real[:, g]
+        y = generated[:, g]
+        
+        x = x[~np.isnan(x)]
+        y = y[~np.isnan(y)]
+        
+        if len(x) < 2 or len(y) < 2:
+            return np.nan
+        
+        xy_dist = np.mean(np.abs(x[:, np.newaxis] - y[np.newaxis, :]))
+        xx_dist = np.mean(np.abs(x[:, np.newaxis] - x[np.newaxis, :]))
+        yy_dist = np.mean(np.abs(y[:, np.newaxis] - y[np.newaxis, :]))
+        
+        return max(0, 2 * xy_dist - xx_dist - yy_dist)
+    
+    if n_jobs != 1 and HAS_JOBLIB:
+        results = Parallel(n_jobs=n_jobs)(
+            delayed(compute_single)(g) for g in range(n_genes)
+        )
+        return np.array(results)
+    else:
+        return np.array([compute_single(g) for g in range(n_genes)])