gsMap3D 0.1.0a1__py3-none-any.whl

gsMap/spatial_ldsc/ldscore_quick_mode.py

@@ -0,0 +1,912 @@
+"""
+Unified spatial LDSC processor combining chunk production, parallel loading, and result accumulation.
+"""
+
+import gc
+import logging
+import queue
+import sys
+import threading
+import time
+import traceback
+from collections import deque
+from dataclasses import dataclass
+from pathlib import Path
+
+import anndata as ad
+import jax.numpy as jnp
+import numpy as np
+import pandas as pd
+from rich.console import Console
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    SpinnerColumn,
+    TaskProgressColumn,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
+from scipy.stats import norm
+from statsmodels.stats.multitest import multipletests
+
+from gsMap.config import SpatialLDSCConfig
+
+from .io import prepare_trait_data
+
+logger = logging.getLogger("gsMap.spatial_ldsc_processor")
+
+
+def prepare_snp_data_for_blocks(data: dict, n_blocks: int) -> dict:
+    """Prepare SNP-related data arrays for equal-sized blocks."""
+    if 'chisq' in data:
+        n_snps = len(data['chisq'])
+    elif 'N' in data:
+        n_snps = len(data['N'])
+    else:
+        raise ValueError("Cannot determine number of SNPs from data")
+
+    block_size = n_snps // n_blocks
+    n_snps_used = block_size * n_blocks
+    n_dropped = n_snps - n_snps_used
+
+    if n_dropped > 0:
+        logger.info(f"Truncating SNP data: dropping {n_dropped} SNPs "
+                   f"({n_dropped/n_snps*100:.3f}%) for {n_blocks} blocks of size {block_size}")
+
+    truncated = {}
+    snp_keys = ['baseline_ld_sum', 'w_ld', 'chisq', 'N', 'snp_positions']
+
+    for key, value in data.items():
+        if key in snp_keys and isinstance(value, np.ndarray | jnp.ndarray):
+            truncated[key] = value[:n_snps_used]
+        elif key == 'baseline_ld':
+            truncated[key] = value.iloc[:n_snps_used]
+        else:
+            truncated[key] = value
+
+    truncated['block_size'] = block_size
+    truncated['n_blocks'] = n_blocks
+    truncated['n_snps_used'] = n_snps_used
+    truncated['n_snps_original'] = n_snps
+
+    return truncated
+
+
+@dataclass
+class ComponentThroughput:
+    """Track throughput for individual pipeline components"""
+    total_batches: int = 0
+    total_time: float = 0.0
+    last_batch_time: float = 0.0
+
+    def record_batch(self, elapsed_time: float):
+        """Record a batch completion"""
+        self.total_batches += 1
+        self.total_time += elapsed_time
+        self.last_batch_time = elapsed_time
+
+    @property
+    def average_time(self) -> float:
+        """Average time per batch"""
+        if self.total_batches > 0:
+            return self.total_time / self.total_batches
+        return 0.0
+
+    @property
+    def throughput(self) -> float:
+        """Batches per second"""
+        if self.average_time > 0:
+            return 1.0 / self.average_time
+        return 0.0
+
+
+class ParallelLDScoreReader:
+    """Multi-threaded reader for fetching LD score chunks from memory-mapped marker scores"""
+
+    def __init__(
+        self,
+        processor,  # Reference to SpatialLDSCProcessor for data access
+        num_workers: int = 4,
+        output_queue: queue.Queue = None
+    ):
+        """Initialize reader pool"""
+        self.processor = processor
+        self.num_workers = num_workers
+
+        # Queues for communication
+        self.read_queue = queue.Queue()
+        self.result_queue = output_queue if output_queue else queue.Queue(maxsize=num_workers * 8)
+
+        # Throughput tracking
+        self.throughput = ComponentThroughput()
+        self.throughput_lock = threading.Lock()
+
+        # Exception handling
+        self.exception_queue = queue.Queue()
+        self.has_error = threading.Event()
+
+        # Start worker threads
+        self.workers = []
+        self.stop_workers = threading.Event()
+        self._start_workers()
+
+    def _start_workers(self):
+        """Start worker threads"""
+        for i in range(self.num_workers):
+            worker = threading.Thread(
+                target=self._worker,
+                args=(i,),
+                daemon=True
+            )
+            worker.start()
+            self.workers.append(worker)
+        logger.info(f"Started {self.num_workers} reader threads")
+
+    def _worker(self, worker_id: int):
+        """Worker thread for reading LD score chunks"""
+        logger.debug(f"Reader worker {worker_id} started")
+
+        while not self.stop_workers.is_set():
+            try:
+                # Get chunk request
+                item = self.read_queue.get(timeout=1)
+                if item is None:
+                    break
+
+                chunk_idx = item
+
+                # Track timing
+                start_time = time.time()
+
+                # Fetch the chunk using processor's method
+                ldscore, spot_names, abs_start, abs_end = self.processor._fetch_ldscore_chunk(chunk_idx)
+
+
+                # Track throughput
+                elapsed = time.time() - start_time
+                with self.throughput_lock:
+                    self.throughput.record_batch(elapsed)
+
+                # Put result for computer
+                self.result_queue.put({
+                    'chunk_idx': chunk_idx,
+                    'ldscore': ldscore,
+                    'spot_names': spot_names,
+                    'abs_start': abs_start,
+                    'abs_end': abs_end,
+                    'worker_id': worker_id,
+                    'success': True
+                })
+                self.read_queue.task_done()
+
+            except queue.Empty:
+                continue
+            except Exception as e:
+                error_trace = traceback.format_exc()
+                logger.error(f"Reader worker {worker_id} error on chunk {chunk_idx}: {e}\nTraceback:\n{error_trace}")
+                self.exception_queue.put((worker_id, e, error_trace))
+                self.has_error.set()
+                self.result_queue.put({
+                    'chunk_idx': chunk_idx if 'chunk_idx' in locals() else -1,
+                    'worker_id': worker_id,
+                    'success': False,
+                    'error': str(e)
+                })
+                break
+
+        logger.debug(f"Reader worker {worker_id} stopped")
+
+    def submit_chunk(self, chunk_idx: int):
+        """Submit chunk for reading"""
+        self.read_queue.put(chunk_idx)
+
+    def get_result(self):
+        """Get next completed chunk"""
+        return self.result_queue.get()
+
+    def get_queue_sizes(self):
+        """Get current queue sizes for monitoring"""
+        return self.read_queue.qsize(), self.result_queue.qsize()
+
+    def check_errors(self):
+        """Check if any worker encountered an error"""
+        if self.has_error.is_set():
+            try:
+                worker_id, exception, error_trace = self.exception_queue.get_nowait()
+                raise RuntimeError(f"Reader worker {worker_id} failed: {exception}\nOriginal traceback:\n{error_trace}") from exception
+            except queue.Empty:
+                raise RuntimeError("Reader worker failed with unknown error")
+
+    def close(self):
+        """Clean up resources"""
+        self.stop_workers.set()
+        for _ in range(self.num_workers):
+            self.read_queue.put(None)
+        for worker in self.workers:
+            worker.join(timeout=5)
+        logger.info("Reader pool closed")
+
+
+class ParallelLDScoreComputer:
+    """Multi-threaded computer for processing LD scores with JAX"""
+
+    def __init__(
+        self,
+        processor,  # Reference to SpatialLDSCProcessor
+        process_chunk_jit_fn,  # JIT-compiled processing function
+        num_workers: int = 4,
+        input_queue: queue.Queue = None
+    ):
+        """Initialize computer pool"""
+        self.processor = processor
+        self.process_chunk_jit_fn = process_chunk_jit_fn
+        self.num_workers = num_workers
+
+        # Queues for communication
+        self.compute_queue = input_queue if input_queue else queue.Queue(maxsize=num_workers * 2)
+
+        # Throughput tracking
+        self.throughput = ComponentThroughput()
+        self.throughput_lock = threading.Lock()
+
+        # Processing statistics
+        self.total_cells_processed = 0
+        self.total_chunks_processed = 0
+        self.stats_lock = threading.Lock()
+
+        # Exception handling
+        self.exception_queue = queue.Queue()
+        self.has_error = threading.Event()
+
+        # Results storage
+        self.results = []
+        self.results_lock = threading.Lock()
+
+        # Prepare static JAX arrays
+        self._prepare_jax_arrays()
+
+        # Start worker threads
+        self.workers = []
+        self.stop_workers = threading.Event()
+        self._start_workers()
+
+    def _prepare_jax_arrays(self):
+        """Prepare static JAX arrays from data_truncated"""
+        n_snps_used = self.processor.data_truncated['n_snps_used']
+
+        baseline_ann = (self.processor.data_truncated['baseline_ld'].values.astype(np.float32) *
+                       self.processor.data_truncated['N'].reshape(-1, 1).astype(np.float32) /
+                       self.processor.data_truncated['Nbar'])
+        baseline_ann = np.concatenate([baseline_ann,
+                                      np.ones((n_snps_used, 1), dtype=np.float32)], axis=1)
+
+        # Convert to JAX arrays
+        self.baseline_ld_sum_jax = jnp.asarray(self.processor.data_truncated['baseline_ld_sum'], dtype=jnp.float32)
+        self.chisq_jax = jnp.asarray(self.processor.data_truncated['chisq'], dtype=jnp.float32)
+        self.N_jax = jnp.asarray(self.processor.data_truncated['N'], dtype=jnp.float32)
+        self.baseline_ann_jax = jnp.asarray(baseline_ann, dtype=jnp.float32)
+        self.w_ld_jax = jnp.asarray(self.processor.data_truncated['w_ld'], dtype=jnp.float32)
+        self.Nbar = self.processor.data_truncated['Nbar']
+
+        del baseline_ann
+        gc.collect()
+
+    def _start_workers(self):
+        """Start compute worker threads"""
+        for i in range(self.num_workers):
+            worker = threading.Thread(
+                target=self._compute_worker,
+                args=(i,),
+                daemon=True
+            )
+            worker.start()
+            self.workers.append(worker)
+        logger.info(f"Started {self.num_workers} compute workers")
+
+    def _compute_worker(self, worker_id: int):
+        """Compute worker thread"""
+        logger.debug(f"Compute worker {worker_id} started")
+
+        while not self.stop_workers.is_set():
+            try:
+                # Get data from reader
+                item = self.compute_queue.get(timeout=1)
+                if item is None:
+                    break
+
+                # Skip failed chunks
+                if not item.get('success', False):
+                    logger.error(f"Skipping chunk {item.get('chunk_idx')} due to read error")
+                    continue
+
+                # Unpack data
+                chunk_idx = item['chunk_idx']
+                ldscore = item['ldscore']
+                spot_names = item['spot_names']
+                abs_start = item['abs_start']
+                abs_end = item['abs_end']
+
+                # Track timing
+                start_time = time.time()
+
+                # Convert to JAX and process
+                spatial_ld_jax = jnp.asarray(ldscore, dtype=jnp.float32)
+
+                # Process with batched JIT function (no batch_size needed - processes all spots at once)
+                betas, ses = self.process_chunk_jit_fn(
+                    self.processor.config.n_blocks,
+                    spatial_ld_jax,
+                    self.baseline_ld_sum_jax,
+                    self.chisq_jax,
+                    self.N_jax,
+                    self.baseline_ann_jax,
+                    self.w_ld_jax,
+                    self.Nbar
+                )
+
+                # Ensure computation completes
+                betas.block_until_ready()
+                ses.block_until_ready()
+
+                # Convert to numpy
+                betas_np = np.array(betas)
+                ses_np = np.array(ses)
+
+                # Track throughput and statistics
+                elapsed = time.time() - start_time
+                n_cells_in_chunk = abs_end - abs_start
+
+                with self.throughput_lock:
+                    self.throughput.record_batch(elapsed)
+
+                with self.stats_lock:
+                    self.total_cells_processed += n_cells_in_chunk
+                    self.total_chunks_processed += 1
+
+                # Store result
+                with self.results_lock:
+                    self.processor._add_chunk_result(
+                        chunk_idx, betas_np, ses_np, spot_names,
+                        abs_start, abs_end
+                    )
+
+                # Clean up
+                del spatial_ld_jax, betas, ses
+
+            except queue.Empty:
+                continue
+            except Exception as e:
+                error_trace = traceback.format_exc()
+                logger.error(f"Compute worker {worker_id} error: {e}\nTraceback:\n{error_trace}")
+                self.exception_queue.put((worker_id, e, error_trace))
+                self.has_error.set()
+                break
+
+        logger.debug(f"Compute worker {worker_id} stopped")
+
+    def get_queue_size(self):
+        """Get compute queue size"""
+        return self.compute_queue.qsize()
+
+    def get_stats(self):
+        """Get processing statistics"""
+        with self.stats_lock:
+            return self.total_cells_processed, self.total_chunks_processed
+
+    def check_errors(self):
+        """Check if any worker encountered an error"""
+        if self.has_error.is_set():
+            try:
+                worker_id, exception, error_trace = self.exception_queue.get_nowait()
+                raise RuntimeError(f"Compute worker {worker_id} failed: {exception}\nOriginal traceback:\n{error_trace}") from exception
+            except queue.Empty:
+                raise RuntimeError("Compute worker failed with unknown error")
+
+    def close(self):
+        """Close compute pool"""
+        self.stop_workers.set()
+        for _ in range(self.num_workers):
+            self.compute_queue.put(None)
+        for worker in self.workers:
+            worker.join(timeout=5)
+        logger.info("Compute pool closed")
+
+
+class SpatialLDSCProcessor:
+    """
+    Unified processor for spatial LDSC that combines:
+    - ChunkProducer: Loading spatial LD chunks
+    - ParallelChunkLoader: Managing parallel chunk loading with adjacent fetching
+    - QuickModeLDScore: Handling memory-mapped marker scores
+    - ResultAccumulator: Validating, merging and saving results
+    """
+
+    def __init__(self,
+                 config: SpatialLDSCConfig,
+                 output_dir: Path,
+                 marker_score_adata: ad.AnnData,
+                 snp_gene_weight_adata: ad.AnnData,
+                 baseline_ld: pd.DataFrame,
+                 w_ld: pd.DataFrame,
+                 n_loader_threads: int = 10):
+        """
+        Initialize the unified processor.
+
+        Args:
+            config: Configuration object
+            output_dir: Output directory for results
+            marker_score_adata: AnnData object containing marker scores in .X
+            snp_gene_weight_adata: AnnData object containing SNP-gene weights
+            baseline_ld: Baseline LD scores common to all traits
+            w_ld: Weights common to all traits
+            n_loader_threads: Number of parallel loader threads
+        """
+        self.config = config
+        self.output_dir = output_dir
+        self.marker_score_adata = marker_score_adata
+        self.snp_gene_weight_adata = snp_gene_weight_adata
+        self.n_spots = marker_score_adata.n_obs
+        self.baseline_ld = baseline_ld
+        self.w_ld = w_ld
+        self.n_loader_threads = n_loader_threads
+
+        # Trait specific state
+        self.trait_name = None
+        self.data_truncated = None
+
+        self.results = []
+        self.processed_chunks = set()
+        self.min_spot_start = float('inf')
+        self.max_spot_end = 0
+
+        logger.info(f"Detected Marker scores shape: (n_spots={self.n_spots}, n_genes={self.marker_score_adata.n_vars})")
+
+        # Find common genes and initialize indices
+        gene_names_from_adata = np.array(self.marker_score_adata.var_names)
+        self.spot_names_all = np.array(self.marker_score_adata.obs_names)
+
+        if self.snp_gene_weight_adata is None:
+             raise ValueError("snp_gene_weight_adata must be provided")
+
+        marker_score_genes_series = pd.Series(gene_names_from_adata)
+        common_genes_mask = marker_score_genes_series.isin(self.snp_gene_weight_adata.var.index)
+        common_genes = gene_names_from_adata[common_genes_mask]
+
+        self.marker_score_gene_indices = np.where(common_genes_mask)[0]
+        self.weight_gene_indices = [self.snp_gene_weight_adata.var.index.get_loc(g) for g in common_genes]
+
+        logger.info(f"Found {len(common_genes)} common genes")
+
+        # Filter by sample if specified
+        if self.config.sample_filter:
+            logger.info(f"Filtering spots by sample: {self.config.sample_filter}")
+            sample_info = self.marker_score_adata.obs.get('sample', self.marker_score_adata.obs.get('sample_name', None))
+
+            if sample_info is None:
+                raise ValueError("No 'sample' or 'sample_name' column found in obs")
+
+            sample_info = sample_info.to_numpy()
+            self.spot_indices = np.where(sample_info == self.config.sample_filter)[0]
+
+            # Verify spots are contiguous for efficient slicing
+            expected_range = list(range(self.spot_indices[0], self.spot_indices[-1] + 1))
+            if self.spot_indices.tolist() != expected_range:
+                raise ValueError("Spot indices for sample must be contiguous")
+
+            self.sample_start_offset = self.spot_indices[0]
+            self.spot_names_filtered = self.spot_names_all[self.spot_indices]
+            logger.info(f"Found {len(self.spot_indices)} spots for sample '{self.config.sample_filter}'")
+        else:
+            self.spot_indices = np.arange(self.n_spots)
+            self.spot_names_filtered = self.spot_names_all
+            self.sample_start_offset = 0
+
+        self.n_spots_filtered = len(self.spot_indices)
+
+        # Set up chunking
+        self.chunk_size = self.config.spots_per_chunk_quick_mode
+
+        # Handle cell indices range if specified
+        if self.config.cell_indices_range:
+            start_cell, end_cell = self.config.cell_indices_range
+            # Adjust for filtered spots
+            start_cell = max(0, start_cell)
+            end_cell = min(end_cell, self.n_spots_filtered)
+            self.chunk_starts = list(range(start_cell, end_cell, self.chunk_size))
+            logger.info(f"Processing cell range [{start_cell}, {end_cell})")
+            self.total_cells_to_process = end_cell - start_cell
+        else:
+            self.chunk_starts = list(range(0, self.n_spots_filtered, self.chunk_size))
+            self.total_cells_to_process = self.n_spots_filtered
+
+        self.total_chunks = len(self.chunk_starts)
+        logger.info(f"Total chunks to process: {self.total_chunks}")
+
+
+    def setup_trait(self, trait_name: str, sumstats_file: str):
+        """
+        Setup processor for a new trait.
+        Loads sumstats, prepares data, and initializes trait-specific components.
+        """
+        self.trait_name = trait_name
+        logger.info(f"Setting up processor for trait: {trait_name}")
+
+        # Prepare trait data
+        data, common_snps = prepare_trait_data(
+            self.config,
+            trait_name,
+            sumstats_file,
+            self.baseline_ld,
+            self.w_ld,
+            self.snp_gene_weight_adata
+        )
+
+        # Prepare blocks
+        self.data_truncated = prepare_snp_data_for_blocks(data, self.config.n_blocks)
+
+        # Extract SNP-gene weight matrix for this trait's SNPs
+        logger.info(f"Initializing trait-specific components for trait: {self.trait_name}")
+        snp_positions = self.data_truncated.get('snp_positions', None)
+        if snp_positions is None:
+            raise ValueError("snp_positions not found in data_truncated")
+
+        # Extract SNP-gene weight matrix for this trait's SNPs
+        self.snp_gene_weight_sparse = self.snp_gene_weight_adata.X[snp_positions, :][:, self.weight_gene_indices]
+
+        if hasattr(self.snp_gene_weight_sparse, 'tocsr'):
+            self.snp_gene_weight_sparse = self.snp_gene_weight_sparse.tocsr()
+
+        logger.info(f"SNP-gene weight matrix shape: {self.snp_gene_weight_sparse.shape}")
+
+        # Reset results
+        self.results = []
+        self.processed_chunks = set()
+        self.min_spot_start = float('inf')
+        self.max_spot_end = 0
+
+
+
+
+    def _fetch_ldscore_chunk(self, chunk_index: int) -> tuple[np.ndarray, pd.Index, int, int]:
+        """
+        Fetch LD score chunk for given index.
+
+        Returns:
+            Tuple of (ldscore_array, spot_names, absolute_start, absolute_end)
+        """
+        if chunk_index >= len(self.chunk_starts):
+            raise ValueError(f"Invalid chunk index {chunk_index}")
+
+        start = self.chunk_starts[chunk_index]
+        end = min(start + self.chunk_size, self.n_spots_filtered)
+
+        # Calculate absolute positions in memmap
+        memmap_start = self.sample_start_offset + start
+        memmap_end = self.sample_start_offset + end
+
+        # Load chunk from marker_score_adata
+        mk_score_chunk = self.marker_score_adata.X[memmap_start:memmap_end, self.marker_score_gene_indices]
+        mk_score_chunk = mk_score_chunk.T.astype(np.float32)
+
+        # Compute LD scores via sparse matrix multiplication
+        ldscore_chunk = self.snp_gene_weight_sparse @ mk_score_chunk
+
+        if hasattr(ldscore_chunk, 'toarray'):
+            ldscore_chunk = ldscore_chunk.toarray()
+
+        # Get spot names
+        spot_names = pd.Index(self.spot_names_filtered[start:end])
+
+        # Calculate absolute positions in original data
+        absolute_start = self.spot_indices[start] if start < len(self.spot_indices) else start
+        absolute_end = self.spot_indices[end - 1] + 1 if end > 0 else absolute_start
+
+        return ldscore_chunk.astype(np.float32, copy=False), spot_names, absolute_start, absolute_end
+
+
+    def process_all_chunks(self, process_chunk_jit_fn) -> pd.DataFrame:
+        """
+        Process all chunks using parallel reader-computer pipeline.
+
+        Args:
+            process_chunk_jit_fn: JIT-compiled function for processing chunks
+
+        Returns:
+            Merged DataFrame with all results
+        """
+        # Create the reader-computer pipeline
+        reader = ParallelLDScoreReader(
+            processor=self,
+            num_workers=self.config.ldsc_read_workers,
+        )
+
+        computer = ParallelLDScoreComputer(
+            processor=self,
+            process_chunk_jit_fn=process_chunk_jit_fn,
+            num_workers=self.config.ldsc_compute_workers,
+            input_queue=reader.result_queue  # Connect reader output to computer input
+        )
+
+        try:
+            # Submit all chunks to reader
+            for chunk_idx in range(self.total_chunks):
+                reader.submit_chunk(chunk_idx)
+
+            # Build description with sample name and range info
+            desc_parts = [f"Processing {self.total_chunks:,} chunks ({self.total_cells_to_process:,} cells)"]
+
+            if hasattr(self.config, 'sample_filter') and self.config.sample_filter:
+                desc_parts.append(f"Sample: {self.config.sample_filter}")
+
+            if self.config.cell_indices_range:
+                start_cell, end_cell = self.config.cell_indices_range
+                desc_parts.append(f"Range: [{start_cell:,}-{end_cell:,})")
+
+            description = " | ".join(desc_parts)
+
+            # # Start JAX profiling if needed
+            # # if hasattr(self.config, 'enable_jax_profiling') and self.config.enable_jax_profiling:
+            # print('starting jax profiler...')
+            # print('starting jax profiler...')
+            # print('starting jax profiler...')
+            # print('starting jax profiler...')
+            # print('starting jax profiler...')
+            # jax.profiler.start_trace("/tmp/jax-trace-ldsc")
+            # Auto-detect terminal vs redirected output
+            # When output is redirected, Console will disable progress bar animations
+            is_terminal = sys.stdout.isatty()
+            console = Console(soft_wrap=True)
+
+            # Log interval for non-terminal mode (every 10% or at least every 60 seconds)
+            log_interval_pct = 10
+            last_log_pct = 0
+            last_log_time = 0
+
+            with Progress(
+                SpinnerColumn(),
+                TextColumn("[bold blue]{task.description}"),
+                BarColumn(),
+                MofNCompleteColumn(),
+                TaskProgressColumn(),
+                TextColumn("[bold green]{task.fields[speed]} cells/s"),
+                TextColumn("[dim]R→C: {task.fields[r_to_c_queue]}"),
+                TimeElapsedColumn(),
+                TimeRemainingColumn(),
+                console=console,
+                refresh_per_second=2,
+                disable=not is_terminal  # Disable rich progress bar when not in terminal
+            ) as progress:
+                task = progress.add_task(
+                    description,
+                    total=self.total_cells_to_process,
+                    speed="0",
+                    r_to_c_queue="0"
+                )
+
+                start_time = time.time()
+                last_update_time = start_time
+                last_chunks_processed = 0
+
+                # For 10s moving average
+                speed_window = deque() # Stores (timestamp, n_cells_processed)
+
+                while last_chunks_processed < self.total_chunks:
+                    # Check for errors
+                    reader.check_errors()
+                    computer.check_errors()
+
+                    # Get current stats from computer
+                    n_cells_processed, n_chunks_processed = computer.get_stats()
+
+                    # Update progress periodically
+                    current_time = time.time()
+                    if current_time - last_update_time > 0.5:  # Update every 0.5 seconds
+                        # Get queue sizes
+                        r_pending, r_to_c = reader.get_queue_sizes()
+
+                        # Calculate speed in cells/s (10s moving average)
+                        speed_window.append((current_time, n_cells_processed))
+
+                        # Remove entries older than 10s
+                        while speed_window and current_time - speed_window[0][0] > 10.0:
+                            speed_window.popleft()
+
+                        if len(speed_window) > 1:
+                            t_diff = speed_window[-1][0] - speed_window[0][0]
+                            c_diff = speed_window[-1][1] - speed_window[0][1]
+                            speed_10s = c_diff / t_diff if t_diff > 0 else 0
+                        else:
+                            # Fallback to cumulative speed if not enough data for window
+                            elapsed_total = current_time - start_time
+                            speed_10s = n_cells_processed / elapsed_total if elapsed_total > 0 else 0
+
+                        # Update progress bar (only effective in terminal mode)
+                        progress.update(
+                            task,
+                            completed=n_cells_processed,
+                            speed=f"{speed_10s:,.0f}",
+                            r_to_c_queue=f"{r_to_c}"
+                        )
+
+                        # Log progress when output is redirected
+                        if not is_terminal:
+                            current_pct = (n_cells_processed / self.total_cells_to_process) * 100 if self.total_cells_to_process > 0 else 0
+                            time_since_last_log = current_time - last_log_time
+
+                            # Log every 10% or every 60 seconds
+                            if current_pct >= last_log_pct + log_interval_pct or time_since_last_log >= 60:
+                                elapsed = current_time - start_time
+                                logger.info(
+                                    f"Progress: {n_cells_processed:,}/{self.total_cells_to_process:,} cells "
+                                    f"({current_pct:.1f}%) | {n_chunks_processed}/{self.total_chunks} chunks | "
+                                    f"Speed: {speed_10s:,.0f} cells/s | Elapsed: {elapsed:.1f}s"
+                                )
+                                last_log_pct = (current_pct // log_interval_pct) * log_interval_pct
+                                last_log_time = current_time
+
+                        last_update_time = current_time
+
+                    # Update last processed count
+                    if n_chunks_processed > last_chunks_processed:
+                        last_chunks_processed = n_chunks_processed
+
+                        # This would block all threads, so we avoid it
+                        # # Periodic memory check
+                        # if n_chunks_processed % 100 == 0:
+                        #     gc.collect()
+                    # Small sleep to prevent busy waiting
+                    time.sleep(0.5)
+
+            # # if hasattr(self.config, 'enable_jax_profiling') and self.config.enable_jax_profiling:
+            # jax.profiler.stop_trace()
+            # print(f'stopped jax profiler, trace saved to /tmp/jax-trace-ldsc')
+            # print(f'stopped jax profiler, trace saved to /tmp/jax-trace-ldsc')
+            # print(f'stopped jax profiler, trace saved to /tmp/jax-trace-ldsc')
+            #     # logger.info("JAX profiling trace saved to /tmp/jax-trace-ldsc")
+
+        finally:
+            # Clean up resources
+            reader.close()
+            computer.close()
+
+            # Log overall speed
+            total_elapsed = time.time() - start_time
+            n_cells_total, _ = computer.get_stats()
+            overall_speed = n_cells_total / total_elapsed if total_elapsed > 0 else 0
+            logger.info(f"Processing complete. Overall speed: {overall_speed:,.0f} cells/s")
+
+            # Note: mkscore_memmap is NOT closed here to allow reuse across multiple traits.
+            # It will be closed by the caller after all traits are processed.
+
+        # Validate and merge results
+        return self._validate_merge_and_save()
+
+    def _add_chunk_result(self, chunk_idx: int, betas: np.ndarray, ses: np.ndarray,
+                         spot_names: pd.Index, abs_start: int, abs_end: int):
+        """Add processed chunk result to accumulator."""
+        # Update coverage tracking
+        self.min_spot_start = min(self.min_spot_start, abs_start)
+        self.max_spot_end = max(self.max_spot_end, abs_end)
+
+        # Store result
+        self.results.append({
+            'chunk_idx': chunk_idx,
+            'betas': betas,
+            'ses': ses,
+            'spot_names': spot_names,
+            'abs_start': abs_start,
+            'abs_end': abs_end
+        })
+        self.processed_chunks.add(chunk_idx)
+
+    def _validate_merge_and_save(self) -> pd.DataFrame:
+        """
+        Validate completeness, merge results, and save with appropriate filename.
+
+        Returns:
+            Merged DataFrame with all results
+        """
+        if not self.results:
+            raise ValueError("No results to merge")
+
+        # Check completeness
+        expected_chunks = set(range(self.total_chunks))
+        missing_chunks = expected_chunks - self.processed_chunks
+
+        if missing_chunks:
+            logger.warning(f"Missing chunks: {sorted(missing_chunks)}")
+            logger.warning(f"Processed {len(self.processed_chunks)}/{self.total_chunks} chunks")
+
+        # Sort results by chunk index
+        sorted_results = sorted(self.results, key=lambda x: x['chunk_idx'])
+
+        # Merge all results
+        dfs = []
+        for result in sorted_results:
+            betas = result['betas'].astype(np.float64)
+            ses = result['ses'].astype(np.float64)
+
+            # Calculate statistics
+            z_scores = betas / ses
+            p_values = norm.sf(z_scores)
+            log10_p = -np.log10(np.maximum(p_values, 1e-300))
+
+            chunk_df = pd.DataFrame({
+                'spot': result['spot_names'],
+                'beta': result['betas'],
+                'se': result['ses'],
+                'z': z_scores.astype(np.float32),
+                'p': p_values,
+                'neg_log10_p': log10_p
+            })
+            dfs.append(chunk_df)
+
+        merged_df = pd.concat(dfs, ignore_index=True)
+
+        # Generate filename with cell range information
+        filename = self._generate_output_filename()
+        output_path = self.output_dir / filename
+
+        # Save results
+        logger.info(f"Saving results to {output_path}")
+        merged_df.to_csv(output_path, index=False, compression='gzip')
+
+        # Log statistics
+        self._log_statistics(merged_df, output_path)
+
+        return merged_df
+
+    def _generate_output_filename(self) -> str:
+        """Generate output filename including cell range information."""
+        base_name = f"{self.config.project_name}_{self.trait_name}"
+
+        # If we have cell indices range, include it in filename
+        if self.config.cell_indices_range:
+            start_cell, end_cell = self.config.cell_indices_range
+            # Adjust for actual processed range
+            actual_start = max(self.min_spot_start, start_cell)
+            actual_end = min(self.max_spot_end, end_cell)
+            return f"{base_name}_cells_{actual_start}_{actual_end}.csv.gz"
+
+        # Check if we have complete coverage
+        if self.min_spot_start == 0 and self.max_spot_end == self.n_spots:
+            return f"{base_name}.csv.gz"
+
+        if self.config.sample_filter:
+            return f"{base_name}_{self.config.sample_filter}_start{self.min_spot_start}_end{self.max_spot_end}_total{self.n_spots}.csv.gz"
+
+        # Partial coverage without explicit range
+        return f"{base_name}_start{self.min_spot_start}_end{self.max_spot_end}_total{self.n_spots}.csv.gz"
+
+    def _log_statistics(self, df: pd.DataFrame, output_path: Path):
+        """Log statistical summary of results."""
+        n_spots = len(df)
+        bonferroni_threshold = 0.05 / n_spots
+        n_bonferroni_sig = (df['p'] < bonferroni_threshold).sum()
+
+        # FDR correction
+        reject, _, _, _ = multipletests(
+            df['p'], alpha=0.001, method='fdr_bh'
+        )
+        n_fdr_sig = reject.sum()
+
+        logger.info("=" * 70)
+        logger.info("STATISTICAL SUMMARY")
+        logger.info("=" * 70)
+        logger.info(f"Total spots: {n_spots:,}")
+        logger.info(f"Cell range processed: [{self.min_spot_start}, {self.max_spot_end})")
+        logger.info(f"Max -log10(p): {df['neg_log10_p'].max():.2f}")
+        logger.info("-" * 70)
+        logger.info(f"Nominally significant (p < 0.05): {(df['p'] < 0.05).sum():,}")
+        logger.info(f"Bonferroni threshold: {bonferroni_threshold:.2e}")
+        logger.info(f"Bonferroni significant: {n_bonferroni_sig:,}")
+        logger.info(f"FDR significant (alpha=0.001): {n_fdr_sig:,}")
+        logger.info("=" * 70)
+        logger.info(f"Results saved to: {output_path}")
+
+        # Warn if incomplete
+        if len(self.processed_chunks) < self.total_chunks:
+            logger.warning(f"WARNING: Only processed {len(self.processed_chunks)}/{self.total_chunks} chunks")
+