gsMap3D 0.1.0a1__py3-none-any.whl

gsMap/latent2gene/marker_scores.py

@@ -0,0 +1,1265 @@
+"""
+Marker score calculation using homogeneous neighbors
+Implements weighted geometric mean calculation in log space with JAX acceleration
+"""
+
+import gc
+import json
+import logging
+import queue
+import threading
+import time
+import traceback
+from dataclasses import dataclass
+from functools import partial
+from pathlib import Path
+
+import anndata as ad
+import jax.numpy as jnp
+import numpy as np
+import pandas as pd
+import scanpy as sc
+from jax import jit
+
+# Progress bar imports
+from rich.console import Console, Group
+from rich.panel import Panel
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    SpinnerColumn,
+    TaskProgressColumn,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
+from rich.table import Table
+
+from gsMap.config import DatasetType, MarkerScoreCrossSliceStrategy
+
+from .connectivity import ConnectivityMatrixBuilder
+from .memmap_io import MemMapDense
+from .row_ordering_jax import optimize_row_order_jax
+
+logger = logging.getLogger(__name__)
+
+from .memmap_io import ComponentThroughput, ParallelMarkerScoreWriter, ParallelRankReader
+
+
+class ParallelMarkerScoreComputer:
+    """Multi-threaded computer pool for marker score calculation (reusable across cell types)"""
+
+    def __init__(
+        self,
+        global_log_gmean: np.ndarray,
+        global_expr_frac: np.ndarray,
+        homogeneous_neighbors: int,
+        num_workers: int = 4,
+        input_queue: queue.Queue = None,
+        output_queue: queue.Queue = None,
+        cross_slice_strategy: str = None,
+        n_slices: int = 1,
+        num_homogeneous_per_slice: int = None,
+        no_expression_fraction: bool = False
+    ):
+        """
+        Initialize computer pool
+
+        Args:
+            global_log_gmean: Global log geometric mean
+            global_expr_frac: Global expression fraction
+            homogeneous_neighbors: Number of homogeneous neighbors
+            num_workers: Number of compute workers
+            input_queue: Optional input queue (from reader)
+            output_queue: Optional output queue (to writer)
+            cross_slice_strategy: Strategy for 3D ('per_slice_pool' or 'hierarchical_pool')
+            n_slices: Number of slices for 3D data
+            num_homogeneous_per_slice: Neighbors per slice for 3D strategies
+            no_expression_fraction: Skip expression fraction filtering if True
+        """
+        self.num_workers = num_workers
+        self.homogeneous_neighbors = homogeneous_neighbors
+        self.cross_slice_strategy = cross_slice_strategy
+        self.n_slices = n_slices
+        self.num_homogeneous_per_slice = num_homogeneous_per_slice or homogeneous_neighbors
+        self.no_expression_fraction = no_expression_fraction
+
+        # Store global statistics as JAX arrays
+        self.global_log_gmean = jnp.array(global_log_gmean)
+        self.global_expr_frac = jnp.array(global_expr_frac)
+
+        # Queues for communication
+        self.compute_queue = input_queue if input_queue else queue.Queue(maxsize=num_workers * 2)
+        self.result_queue = output_queue if output_queue else queue.Queue(maxsize=num_workers * 2)
+
+        # Throughput tracking
+        self.throughput = ComponentThroughput()
+        self.throughput_lock = threading.Lock()
+
+        # Current processing context
+        self.neighbor_weights = None
+        self.cell_indices_sorted = None
+        self.batch_size = None
+        self.n_cells = None
+
+        # Exception handling
+        self.exception_queue = queue.Queue()
+        self.has_error = threading.Event()
+
+        # Start worker threads
+        self.workers = []
+        self.stop_workers = threading.Event()
+        self.active_cell_type = None  # Track current cell type being processed
+        self._start_workers()
+
+    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.info(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
+
+                # Unpack data from reader
+                batch_idx, rank_data, rank_indices, original_shape, batch_metadata = item
+
+                # Track timing
+                start_time = time.time()
+
+                # Extract batch parameters from metadata
+                # These should always be provided, but check for safety
+                if self.batch_size is None or self.neighbor_weights is None:
+                    logger.error(f"Compute worker {worker_id}: batch context not set")
+                    raise RuntimeError("Batch context must be set before processing")
+
+                batch_start = batch_metadata['batch_start']
+                batch_end = batch_metadata['batch_end']
+                actual_batch_size = batch_end - batch_start
+
+                # Verify shape
+                assert original_shape == (actual_batch_size, self.homogeneous_neighbors * self.n_slices), \
+                    f"Unexpected rank data shape: {original_shape}, expected ({actual_batch_size}, {self.homogeneous_neighbors * self.n_slices})"
+
+                # Get batch-specific data
+                batch_weights = self.neighbor_weights[batch_start:batch_end]
+                batch_cell_indices = self.cell_indices_sorted[batch_start:batch_end]
+
+                # Convert to JAX for efficient computation
+                rank_data_jax = jnp.array(rank_data)
+                rank_indices_jax = jnp.array(rank_indices)
+
+                # Use JAX fancy indexing
+                batch_ranks = rank_data_jax[rank_indices_jax]
+
+                # Compute marker scores using appropriate strategy
+                if self.cross_slice_strategy == 'hierarchical_pool':
+                    # Use hierarchical pooling (per-slice marker score average) for 3D data
+                    marker_scores = compute_marker_scores_3d_hierarchical_pool_jax(
+                        batch_ranks,
+                        batch_weights,
+                        actual_batch_size,
+                        self.n_slices,
+                        self.num_homogeneous_per_slice,
+                        self.global_log_gmean,
+                        self.global_expr_frac,
+                        self.no_expression_fraction
+                    )
+                else:
+                    # Use standard computation (includes mean pooling via weights)
+                    marker_scores = compute_marker_scores_jax(
+                        batch_ranks,
+                        batch_weights,
+                        actual_batch_size,
+                        self.homogeneous_neighbors * self.n_slices,
+                        self.global_log_gmean,
+                        self.global_expr_frac,
+                        self.no_expression_fraction
+                    )
+
+                # Convert back to numpy as float16 for memory efficiency
+                marker_scores_np = np.array(marker_scores, dtype=np.float16)
+
+                # Track throughput
+                elapsed = time.time() - start_time
+                with self.throughput_lock:
+                    self.throughput.record_batch(elapsed)
+
+                # Put result directly to writer queue
+                self.result_queue.put((batch_idx, marker_scores_np, batch_cell_indices))
+                self.compute_queue.task_done()
+
+            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()
+                self.stop_workers.set()  # Signal all workers to stop
+                break
+
+        logger.info(f"Compute worker {worker_id} stopped")
+
+    def set_batch_context(self, neighbor_weights: np.ndarray, cell_indices_sorted: np.ndarray,
+                         batch_size: int, n_cells: int):
+        """Set context for processing batches of current cell type"""
+        self.neighbor_weights = jnp.asarray(neighbor_weights) # transfer to jax array only once
+        self.cell_indices_sorted = cell_indices_sorted
+        self.batch_size = batch_size
+        self.n_cells = n_cells
+
+    def reset_for_cell_type(self, cell_type: str):
+        """Reset for processing a new cell type"""
+        self.active_cell_type = cell_type
+        self.neighbor_weights = None
+        self.cell_indices_sorted = None
+        self.batch_size = None
+        self.n_cells = None
+        with self.throughput_lock:
+            self.throughput = ComponentThroughput()
+        logger.debug(f"Reset computer throughput for {cell_type}")
+
+    def get_queue_sizes(self):
+        """Get current queue sizes for progress tracking"""
+        return self.compute_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"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")
+
+
+
+
+
+@dataclass
+class PipelineStats:
+    """Statistics for pipeline monitoring"""
+    total_batches: int
+    completed_reads: int = 0
+    completed_computes: int = 0
+    completed_writes: int = 0
+    pending_compute: int = 0
+    pending_write: int = 0
+    pending_read: int = 0
+    start_time: float = 0
+
+    # Component throughput tracking
+    reader_throughput: ComponentThroughput = None
+    computer_throughput: ComponentThroughput = None
+    writer_throughput: ComponentThroughput = None
+
+    def __post_init__(self):
+        self.start_time = time.time()
+        self.reader_throughput = ComponentThroughput()
+        self.computer_throughput = ComponentThroughput()
+        self.writer_throughput = ComponentThroughput()
+
+    @property
+    def elapsed_time(self):
+        return time.time() - self.start_time
+
+    @property
+    def throughput(self):
+        if self.elapsed_time > 0:
+            return self.completed_writes / self.elapsed_time
+        return 0
+
+
+class MarkerScoreMessageQueue:
+    """Streamlined pipeline for marker score calculation (reusable across cell types)"""
+
+    def __init__(
+        self,
+        reader: ParallelRankReader,
+        computer: ParallelMarkerScoreComputer,
+        writer: ParallelMarkerScoreWriter,
+        batch_size: int,
+    ):
+        """
+        Initialize pipeline with shared pools
+
+        Args:
+            reader: Rank reader pool
+            computer: Marker score computer pool
+            writer: Marker score writer pool
+            batch_size: Size of each batch
+        """
+        self.reader = reader
+        self.computer = computer
+        self.writer = writer
+        self.batch_size = batch_size
+
+        # Cell type specific parameters (set via reset_for_cell_type)
+        self.n_batches = None
+        self.n_cells = None
+        self.active_cell_type = None
+        self.stats = None
+
+    def reset_for_cell_type(self, cell_type: str, n_cells: int):
+        """Reset the queue for processing a new cell type
+
+        Args:
+            cell_type: Name of the cell type
+            n_cells: Total number of cells for this cell type
+        """
+        # Reset all components for new cell type
+        self.reader.reset_for_cell_type(cell_type)
+        self.computer.reset_for_cell_type(cell_type)
+        self.writer.reset_for_cell_type(cell_type)
+
+        self.active_cell_type = cell_type
+        self.n_cells = n_cells
+        self.n_batches = (n_cells + self.batch_size - 1) // self.batch_size
+        self.stats = PipelineStats(total_batches=self.n_batches)
+        logger.debug(f"Reset MarkerScoreMessageQueue for {cell_type}: {self.n_batches} batches, {n_cells} cells")
+
+    def _submit_batches(
+        self,
+        neighbor_indices: np.ndarray,
+        neighbor_weights: np.ndarray,
+        cell_indices_sorted: np.ndarray
+    ):
+        """Submit all batches to the reader"""
+        # Set context for computer
+        self.computer.set_batch_context(
+            neighbor_weights=neighbor_weights,
+            cell_indices_sorted=cell_indices_sorted,
+            batch_size=self.batch_size,
+            n_cells=self.n_cells
+        )
+
+        # Submit all batches with metadata
+        for batch_idx in range(self.n_batches):
+            batch_start = batch_idx * self.batch_size
+            batch_end = min(batch_start + self.batch_size, self.n_cells)
+            batch_neighbors = neighbor_indices[batch_start:batch_end]
+
+            # Include metadata for computer
+            batch_metadata = {
+                'batch_start': batch_start,
+                'batch_end': batch_end,
+                'batch_idx': batch_idx
+            }
+
+            self.reader.submit_batch(batch_idx, batch_neighbors, batch_metadata)
+
+    def _update_stats(self):
+        """Update pipeline statistics"""
+        # Update completed counts
+        self.stats.completed_writes = self.writer.get_completed_count()
+
+        # Update queue sizes
+        read_pending, read_ready = self.reader.get_queue_sizes()
+        self.stats.pending_read = read_pending
+
+        compute_pending, compute_ready = self.computer.get_queue_sizes()
+        self.stats.pending_compute = compute_pending
+
+        self.stats.pending_write = self.writer.get_queue_size()
+
+        # Update throughput stats
+        self.stats.reader_throughput = self.reader.throughput
+        self.stats.computer_throughput = self.computer.throughput
+        self.stats.writer_throughput = self.writer.throughput
+
+        # Estimate completed reads and computes based on queue states
+        self.stats.completed_reads = self.n_batches - read_pending
+        self.stats.completed_computes = self.stats.completed_writes + self.stats.pending_write
+
+    def start(
+        self,
+        neighbor_indices: np.ndarray,
+        neighbor_weights: np.ndarray,
+        cell_indices_sorted: np.ndarray,
+        enable_profiling: bool = False
+    ):
+        """Run pipeline with rich progress display
+
+        Args:
+            neighbor_indices: Neighbor indices for each cell
+            neighbor_weights: Weights for each neighbor
+            cell_indices_sorted: Sorted cell indices
+            enable_profiling: Whether to enable profiling
+        """
+
+        # Ensure the queue has been reset for this cell type
+        if self.active_cell_type is None or self.stats is None:
+            raise RuntimeError("MarkerScoreMessageQueue must be reset before starting. Call reset_for_cell_type first.")
+
+        # Optional profiling
+        tracer = None
+        if enable_profiling:
+            import viztracer
+            tracer = viztracer.VizTracer(
+                output_file=f"marker_score_{self.active_cell_type}_{self.n_batches}.json",
+                max_stack_depth=10,
+            )
+            tracer.start()
+
+        console = Console()
+
+        # Define queue color mapping based on queue size
+        def get_queue_color(size: int) -> str:
+            """Get color based on queue size"""
+            if size == 0:
+                return "dim"
+            elif size < 5:
+                return "green"
+            elif size < 10:
+                return "yellow"
+            elif size < 20:
+                return "bright_yellow"
+            else:
+                return "red"
+
+        try:
+            # Create progress bars
+            with Progress(
+                SpinnerColumn(),
+                TextColumn("[bold blue]{task.description}"),
+                BarColumn(),
+                MofNCompleteColumn(),
+                TaskProgressColumn(),
+                TimeElapsedColumn(),
+                TimeRemainingColumn(),
+                console=console,
+                refresh_per_second=10
+            ) as progress:
+
+                # Add single task for pipeline
+                pipeline_task = progress.add_task(
+                    f"[bold]{self.active_cell_type}[/bold]", total=self.n_batches
+                )
+
+                # Submit all batches to start the pipeline
+                self._submit_batches(neighbor_indices, neighbor_weights, cell_indices_sorted)
+
+                # Monitor progress
+                while self.stats.completed_writes < self.n_batches:
+                    # Check for errors in any component
+                    self.reader.check_errors()
+                    self.computer.check_errors()
+                    self.writer.check_errors()
+
+                    self._update_stats()
+
+                    # Update progress based on completed writes (final stage)
+                    progress.update(pipeline_task, completed=self.stats.completed_writes)
+
+                    # Color code queue sizes based on fullness
+                    compute_color = get_queue_color(self.stats.pending_compute)
+                    write_color = get_queue_color(self.stats.pending_write)
+
+                    # Update description with queue information
+                    progress.update(
+                        pipeline_task,
+                        description=(
+                            f"[bold]{self.active_cell_type}[/bold] | "
+                            f"Queues: [{compute_color}]R→C:{self.stats.pending_compute}[/{compute_color}] "
+                            f"[{write_color}]C→W:{self.stats.pending_write}[/{write_color}]"
+                        )
+                    )
+
+                    time.sleep(0.1)
+
+                # Final update
+                progress.update(pipeline_task, completed=self.n_batches)
+
+        except Exception as e:
+            logger.error(f"Pipeline failed for {self.active_cell_type}: {e}")
+            # Stop all workers to prevent hanging
+            self.stop()
+            raise
+
+        finally:
+            # Stop profiling if enabled
+            if tracer is not None:
+                tracer.stop()
+                tracer.save()
+                logger.info(f"Profiling saved to marker_score_{self.active_cell_type}_{self.n_batches}.json")
+
+        # No threads to wait for - processing happens in component worker threads
+
+        # Print summary with component throughputs
+        # Calculate effective pipeline throughput (limited by bottleneck)
+        min(
+            self.stats.reader_throughput.throughput * self.reader.num_workers if self.stats.reader_throughput.throughput > 0 else float('inf'),
+            self.stats.computer_throughput.throughput * self.computer.num_workers if self.stats.computer_throughput.throughput > 0 else float('inf'),
+            self.stats.writer_throughput.throughput * self.writer.num_workers if self.stats.writer_throughput.throughput > 0 else float('inf')
+        )
+
+        # Calculate cells per second for each component and pipeline
+        pipeline_cells_per_sec = self.stats.throughput * self.batch_size if self.stats.throughput > 0 else 0
+        reader_cells_per_sec = self.stats.reader_throughput.throughput * self.batch_size * self.reader.num_workers if self.stats.reader_throughput.throughput > 0 else 0
+        computer_cells_per_sec = self.stats.computer_throughput.throughput * self.batch_size * self.computer.num_workers if self.stats.computer_throughput.throughput > 0 else 0
+        writer_cells_per_sec = self.stats.writer_throughput.throughput * self.batch_size * self.writer.num_workers if self.stats.writer_throughput.throughput > 0 else 0
+
+        # Create summary table
+        summary_table = Table(title=f"[bold green]✓ Completed {self.active_cell_type}[/bold green]", box=None)
+        summary_table.add_column("Property", style="dim")
+        summary_table.add_column("Value", style="green", justify="right")
+
+        summary_table.add_row("Total Batches", str(self.n_batches))
+        summary_table.add_row("Time Elapsed", f"{self.stats.elapsed_time:.2f}s")
+        summary_table.add_row("Pipeline Throughput", f"{self.stats.throughput:.2f} batches/s ({pipeline_cells_per_sec:.0f} cells/s)")
+
+        perf_table = Table(title="[bold]Component Performance (per worker)[/bold]", show_header=True, header_style="bold blue")
+        perf_table.add_column("Component")
+        perf_table.add_column("Throughput", justify="right")
+        perf_table.add_column("Workers", justify="right")
+        perf_table.add_column("Total Throughput", justify="right", style="green")
+
+        perf_table.add_row("Reader", f"{self.stats.reader_throughput.throughput:.2f} b/s", str(self.reader.num_workers), f"{reader_cells_per_sec:.0f} c/s")
+        perf_table.add_row("Computer", f"{self.stats.computer_throughput.throughput:.2f} b/s", str(self.computer.num_workers), f"{computer_cells_per_sec:.0f} c/s")
+        perf_table.add_row("Writer", f"{self.stats.writer_throughput.throughput:.2f} b/s", str(self.writer.num_workers), f"{writer_cells_per_sec:.0f} c/s")
+
+        console.print(Panel(
+            Group(summary_table, perf_table),
+            title="Pipeline Summary",
+            border_style="green"
+        ))
+
+        # Final garbage collection
+        import gc
+        gc.collect()
+
+        logger.info(f"✓ Completed processing {self.active_cell_type}")
+
+    def stop(self):
+        """Stop all workers in the pipeline components"""
+        logger.info("Stopping MarkerScoreMessageQueue workers...")
+        self.reader.stop_workers.set()
+        self.computer.stop_workers.set()
+        self.writer.stop_workers.set()
+        logger.info("MarkerScoreMessageQueue workers stopped")
+
+
+
+@partial(jit, static_argnums=(2, 3, 6))
+def compute_marker_scores_jax(
+    log_ranks: jnp.ndarray,  # (B*N) × G matrix
+    weights: jnp.ndarray,  # B × N weight matrix
+    batch_size: int,
+    num_neighbors: int,
+    global_log_gmean: jnp.ndarray,  # G-dimensional vector
+    global_expr_frac: jnp.ndarray,  # G-dimensional vector
+    no_expression_fraction: bool = False  # Skip expression fraction filtering if True
+) -> jnp.ndarray:
+    """
+    JAX-accelerated marker score computation
+
+    Returns:
+        B × G marker scores
+    """
+    n_genes = log_ranks.shape[1]
+
+    # Reshape to batch format
+    log_ranks_3d = log_ranks.reshape(batch_size, num_neighbors, n_genes)
+
+    # Compute weighted geometric mean in log space
+    weights = weights / weights.sum(axis=1, keepdims=True)  # Normalize weights
+    weighted_log_mean = jnp.einsum('bn,bng->bg', weights, log_ranks_3d)
+
+    # Calculate marker score
+    marker_score = jnp.exp(weighted_log_mean - global_log_gmean)
+    marker_score = jnp.where(marker_score < 1.0, 0.0, marker_score)
+
+    # Apply expression fraction filter (only if not disabled)
+    if not no_expression_fraction:
+        # Compute expression fraction (mean of is_expressed across neighbors)
+        # Treat min log rank as non-expressed
+        is_expressed = (log_ranks_3d != log_ranks_3d.min(axis=-1, keepdims=True))
+
+        # Create mask for valid neighbors (where weights > 0)
+        valid_mask = weights > 0  # Shape: (batch_size, num_neighbors)
+
+        # Apply mask and compute mean only for valid neighbors
+        is_expressed_masked = jnp.where(valid_mask[:, :, None], is_expressed, 0)
+        valid_counts = valid_mask.sum(axis=1, keepdims=True)  # Count of valid neighbors per cell
+
+        # Compute mean only over valid neighbors (avoid division by zero)
+        expr_frac = jnp.where(
+            valid_counts > 0,
+            is_expressed_masked.astype(jnp.float16).sum(axis=1) / valid_counts,
+            0.0
+        )
+
+        frac_mask = expr_frac > global_expr_frac
+        marker_score = jnp.where(frac_mask, marker_score, 0.0)
+
+    marker_score = jnp.exp(marker_score ** 1.5) - 1.0
+
+    # Return as float16 for memory efficiency
+    return marker_score.astype(jnp.float16)
+
+
+@partial(jit, static_argnums=(2, 3, 4, 7))
+def compute_marker_scores_3d_hierarchical_pool_jax(
+    log_ranks: jnp.ndarray,  # (B*N) × G matrix where N = n_slices * num_homogeneous_per_slice
+    weights: jnp.ndarray,  # B × N weight matrix
+    batch_size: int,
+    n_slices: int,
+    num_homogeneous_per_slice: int,
+    global_log_gmean: jnp.ndarray,  # G-dimensional vector
+    global_expr_frac: jnp.ndarray,  # G-dimensional vector
+    no_expression_fraction: bool = False  # Skip expression fraction filtering if True
+) -> jnp.ndarray:
+    """
+    JAX-accelerated marker score computation with hierarchical pooling for 3D spatial data.
+    Computes marker scores independently for each slice and takes the average.
+
+    Args:
+        log_ranks: Flattened log ranks (batch_size * total_neighbors, n_genes)
+        weights: Flattened weights (batch_size, total_neighbors)
+        batch_size: Number of cells in batch
+        n_slices: Number of slices (1 + 2 * n_adjacent_slices)
+        num_homogeneous_per_slice: Number of homogeneous neighbors per slice
+        global_log_gmean: Global log geometric mean
+        global_expr_frac: Global expression fraction
+
+    Returns:
+        (batch_size, n_genes) marker scores using average cross slices
+    """
+    n_genes = log_ranks.shape[1]
+    n_slices * num_homogeneous_per_slice
+
+    # Reshape to separate slices: (batch_size, n_slices, num_homogeneous_per_slice, n_genes)
+    log_ranks_4d = log_ranks.reshape(batch_size, n_slices, num_homogeneous_per_slice, n_genes)
+
+    # Reshape weights: (batch_size, n_slices, num_homogeneous_per_slice)
+    weights_3d = weights.reshape(batch_size, n_slices, num_homogeneous_per_slice)
+
+    # Normalize weights within each slice (sum to 1 along num_homogeneous_per_slice axis)
+    weights_sum = weights_3d.sum(axis=2, keepdims=True) # Shape: (batch_size, n_slices, 1)
+    weights_normalized = weights_3d / jnp.where(weights_sum > 0, weights_sum, 1.0)
+
+    # Compute weighted geometric mean in log space for each slice
+    # Result: (batch_size, n_slices, n_genes)
+    weighted_log_mean = jnp.einsum('bsn,bsng->bsg', weights_normalized, log_ranks_4d)
+
+    # Calculate marker score for each slice
+    marker_score_per_slice = jnp.exp(weighted_log_mean - global_log_gmean[None, None, :])
+    marker_score_per_slice = jnp.where(marker_score_per_slice < 1.0, 0.0, marker_score_per_slice)
+
+    # Apply expression fraction filter for each slice (only if not disabled)
+    if not no_expression_fraction:
+        # Compute expression fraction for each slice
+        # Treat min log rank as non-expressed
+        min_log_rank = log_ranks_4d.min(axis=-1, keepdims=True)
+        is_expressed = (log_ranks_4d != min_log_rank)
+
+        # Create mask for valid neighbors within each slice (where weights > 0)
+        valid_mask = weights_3d > 0  # Shape: (batch_size, n_slices, num_homogeneous_per_slice)
+
+        # Apply mask and compute mean only for valid neighbors within each slice
+        is_expressed_masked = jnp.where(valid_mask[:, :, :, None], is_expressed, 0)
+        valid_counts = valid_mask.sum(axis=2, keepdims=True)  # Count of valid neighbors per slice
+
+        # Compute mean only over valid neighbors (avoid division by zero)
+        # Result: (batch_size, n_slices, n_genes)
+        # valid_counts has shape (batch_size, n_slices, 1), need to broadcast properly
+        expr_frac = jnp.where(
+            valid_counts > 0,
+            is_expressed_masked.astype(jnp.float16).sum(axis=2) / valid_counts,
+            0.0
+        )
+
+        frac_mask = expr_frac > global_expr_frac[None, None, :]
+        marker_score_per_slice = jnp.where(frac_mask, marker_score_per_slice, 0.0)
+
+    marker_score_per_slice = jnp.exp(marker_score_per_slice ** 1.5) - 1.0
+
+    # Calculate median across slices instead of max pooling to reduce noise
+    # Result: (batch_size, n_genes)
+
+    # Get the invalid slice mask before calculating median.
+    # These slices are invalid could be:
+    # 1. no neighbors from this slice for the outermost slices
+    # 2. the adjacent slices that do not have enough high quality cells which has been filtered out
+    invalid_slice_mask = (weights_sum == 0).squeeze(-1)  # Shape: (batch_size, n_slices)
+
+    # Set invalid slice scores to NaN, then calculate the median
+    # NaN values will be ignored by jnp.nanmedian
+    marker_score_per_slice = jnp.where(
+        invalid_slice_mask[:, :, None],  # Broadcast to (batch_size, n_slices, n_genes)
+        jnp.nan,
+        marker_score_per_slice
+    )
+
+    # Calculate median across slices (axis=1), ignoring NaN values
+    marker_score = jnp.nanmean(marker_score_per_slice, axis=1)
+    # marker_score = jnp.nanmedian(marker_score_per_slice, axis=1)
+
+    # Handle cases where all slices are invalid (all NaN) - set to 0
+    marker_score = jnp.where(jnp.isnan(marker_score), 0.0, marker_score)
+
+    # Return as float16 for memory efficiency
+    return marker_score.astype(jnp.float16)
+
+
+class MarkerScoreCalculator:
+    """Main class for calculating marker scores"""
+
+    def __init__(self, config):
+        """
+        Initialize with configuration
+
+        Args:
+            config: LatentToGeneConfig object
+        """
+        self.config = config
+        self.connectivity_builder = ConnectivityMatrixBuilder(config)
+
+    def load_global_stats(self, mean_frac_path: str) -> tuple[np.ndarray, np.ndarray]:
+        """Load pre-calculated global geometric mean and expression fraction from parquet"""
+
+        logger.info("Loading global statistics from parquet...")
+        parquet_path = Path(mean_frac_path)
+
+        if not parquet_path.exists():
+            raise FileNotFoundError(f"Global stats file not found: {parquet_path}")
+
+        # Load the dataframe
+        mean_frac_df = pd.read_parquet(parquet_path)
+
+        # Extract global log geometric mean and expression fraction
+        global_log_gmean = mean_frac_df['G_Mean'].values.astype(np.float32)
+        global_expr_frac = mean_frac_df['frac'].values.astype(np.float32)
+
+        logger.info(f"Loaded global stats for {len(global_log_gmean)} genes")
+
+        return global_log_gmean, global_expr_frac
+
+    def _display_input_summary(self, adata, cell_types, n_cells, n_genes):
+        """Display summary of input data and cell types"""
+        table = Table(title="[bold cyan]Marker Score Input Summary[/bold cyan]", show_header=True, header_style="bold magenta")
+        table.add_column("Property", style="dim")
+        table.add_column("Value", style="green")
+
+        table.add_row("Total Cells", str(n_cells))
+        table.add_row("Total Genes", str(n_genes))
+        table.add_row("Cell Types", str(len(cell_types)))
+        table.add_row("Dataset Type", str(self.config.dataset_type.value))
+        table.add_row("Annotation Key", str(self.config.annotation))
+
+        self.console.print(table)
+
+        # Display cell type breakdown
+        ct_table = Table(title="[bold cyan]Cell Type Breakdown[/bold cyan]", show_header=True, header_style="bold blue")
+        ct_table.add_column("Cell Type", style="dim")
+        ct_table.add_column("Count", justify="right")
+
+        ct_counts = adata.obs[self.config.annotation].value_counts()
+        for ct in cell_types:
+            count = ct_counts.get(ct, 0)
+            style = "green" if count >= self.config.min_cells_per_type else "red"
+            ct_table.add_row(ct, f"[{style}]{count}[/{style}]")
+
+        self.console.print(ct_table)
+
+    def _load_input_data(
+        self,
+        adata_path: str,
+        rank_memmap_path: str,
+        mean_frac_path: str
+    ) -> tuple[ad.AnnData, MemMapDense, np.ndarray, np.ndarray, int, int, np.ndarray]:
+        """Load input data: AnnData, rank memory map, global statistics, and high quality mask
+
+        Returns:
+            Tuple of (adata, rank_memmap, global_log_gmean, global_expr_frac, n_cells, n_genes, high_quality_mask)
+        """
+        # Load concatenated AnnData
+        logger.info(f"Loading concatenated AnnData from {adata_path}")
+        if not Path(adata_path).exists():
+            raise FileNotFoundError(f"Concatenated AnnData not found: {adata_path}")
+        adata = sc.read_h5ad(adata_path)
+
+        # Load pre-calculated global statistics
+        global_log_gmean, global_expr_frac = self.load_global_stats(mean_frac_path)
+
+        # Open rank memory map and get dimensions
+        rank_memmap_path = Path(rank_memmap_path)
+        meta_path = rank_memmap_path.with_suffix('.meta.json')
+        with open(meta_path) as f:
+            meta = json.load(f)
+
+        rank_memmap = MemMapDense(
+            path=rank_memmap_path,
+            shape=tuple(meta['shape']),
+            dtype=np.dtype(meta['dtype']),
+            mode='r',
+            tmp_dir=self.config.memmap_tmp_dir
+        )
+
+        logger.info(f"Opened rank memory map from {rank_memmap_path}")
+        n_cells = adata.n_obs
+        n_cells_rank = rank_memmap.shape[0]
+        n_genes = rank_memmap.shape[1]
+
+        logger.info(f"AnnData dimensions: {n_cells} cells × {adata.n_vars} genes")
+        logger.info(f"Rank MemMap dimensions: {n_cells_rank} cells × {n_genes} genes")
+
+        # Cells should match exactly since filtering is done before rank memmap creation
+        assert n_cells == n_cells_rank, \
+            f"Cell count mismatch: AnnData has {n_cells} cells, Rank MemMap has {n_cells_rank} cells. " \
+            f"This indicates the filtering was not applied consistently during rank calculation."
+
+        # Load high quality mask based on configuration
+        if self.config.high_quality_neighbor_filter:
+            if 'High_quality' not in adata.obs.columns:
+                raise ValueError("High_quality column not found in AnnData obs. Please ensure QC was applied during find_latent_representation step.")
+            high_quality_mask = adata.obs['High_quality'].values.astype(bool)
+            logger.info(f"Loaded high quality mask: {high_quality_mask.sum()}/{len(high_quality_mask)} cells marked as high quality")
+        else:
+            # Create all-True mask when high quality filtering is disabled
+            high_quality_mask = np.ones(n_cells, dtype=bool)
+            logger.info("High quality filtering disabled - using all cells")
+
+        return adata, rank_memmap, global_log_gmean, global_expr_frac, n_cells, n_genes, high_quality_mask
+
+    def _prepare_embeddings(
+        self,
+        adata: ad.AnnData
+    ) -> tuple[np.ndarray | None, np.ndarray, np.ndarray, np.ndarray | None]:
+        """Prepare and normalize embeddings based on dataset type
+
+        Returns:
+            Tuple of (coords, emb_niche, emb_indv, slice_ids)
+        """
+        logger.info("Loading shared data structures...")
+
+        coords = None
+        emb_niche = None
+        slice_ids = None
+
+        if self.config.dataset_type in ['spatial2D', 'spatial3D']:
+            # Load spatial coordinates for spatial datasets
+            coords = adata.obsm[self.config.spatial_key]
+
+            # Load slice IDs if provided (for both spatial2D and spatial3D)
+            assert 'slice_id' in adata.obs.columns
+            slice_ids = adata.obs['slice_id'].values.astype(np.int32)
+
+            # Try to load niche embeddings if they exist
+            if self.config.latent_representation_niche in adata.obsm:
+                emb_niche = adata.obsm[self.config.latent_representation_niche]
+
+        # Load cell embeddings for all dataset types
+        emb_indv = adata.obsm[self.config.latent_representation_cell].astype(np.float16)
+
+        # --- ELEGANT FIX START ---
+        # If emb_niche is missing (scRNA-seq or spatial without niche),
+        # create a dummy (N, 1) array of ones.
+        # Normalizing a vector of 1s results in 1.0, so cosine similarity will be 1.0 everywhere.
+        if emb_niche is None:
+            logger.info("No niche embeddings found. Using dummy embeddings (all ones).")
+            emb_niche = np.ones((emb_indv.shape[0], 1), dtype=np.float32)
+        # --- ELEGANT FIX END ---
+
+        # Normalize embeddings
+        logger.info("Normalizing embeddings...")
+
+        # L2 normalize niche embeddings (always exists now)
+        emb_niche_norm = np.linalg.norm(emb_niche, axis=1, keepdims=True)
+        emb_niche = emb_niche / (emb_niche_norm + 1e-8)
+
+        # L2 normalize individual embeddings
+        emb_indv_norm = np.linalg.norm(emb_indv, axis=1, keepdims=True)
+        emb_indv = emb_indv / (emb_indv_norm + 1e-8)
+
+        return coords, emb_niche, emb_indv, slice_ids
+
+    def _get_cell_types(self, adata: ad.AnnData) -> np.ndarray:
+        """Get cell types from annotation key
+
+        Returns:
+            Array of unique cell types
+        """
+        annotation_key = self.config.annotation
+
+        if annotation_key is not None:
+            assert annotation_key in adata.obs.columns, f"Annotation key '{annotation_key}' not found in adata.obs"
+            # Get unique cell types, excluding NaN values
+            cell_types = adata.obs[annotation_key].dropna().unique()
+
+            # Check if there are any NaN values and handle them
+            nan_count = adata.obs[annotation_key].isna().sum()
+            if nan_count > 0:
+                logger.warning(f"Found {nan_count} cells with NaN annotation in '{annotation_key}', these will be skipped")
+        else:
+            logger.warning(f"Annotation {annotation_key} not found, processing all cells as one type")
+            cell_types = ["all"]
+            adata.obs[annotation_key] = "all"
+
+        logger.info(f"Processing {len(cell_types)} cell types")
+        return cell_types
+
+    def _initialize_pipeline(
+        self,
+        rank_memmap: MemMapDense,
+        output_memmap: MemMapDense,
+        global_log_gmean: np.ndarray,
+        global_expr_frac: np.ndarray
+    ):
+        """Initialize the processing pipeline with shared pools and queues"""
+        logger.info("Initializing shared processing pools with direct queue connections...")
+
+        # Create shared queues to connect components with configured sizes
+        reader_to_computer_queue = queue.Queue(maxsize=self.config.mkscore_compute_workers * self.config.compute_input_queue_size)
+        computer_to_writer_queue = queue.Queue(maxsize=self.config.writer_queue_size)
+
+        self.reader = ParallelRankReader(
+            rank_memmap,
+            num_workers=self.config.rank_read_workers,
+            output_queue=reader_to_computer_queue  # Direct connection to computer
+        )
+
+        # Determine 3D strategy parameters
+        cross_slice_strategy = None
+        n_slices = 1
+        num_homogeneous_per_slice = self.config.homogeneous_neighbors
+
+        if (self.config.dataset_type == DatasetType.SPATIAL_3D and
+            self.config.cross_slice_marker_score_strategy in [
+                MarkerScoreCrossSliceStrategy.PER_SLICE_POOL,
+                MarkerScoreCrossSliceStrategy.HIERARCHICAL_POOL
+            ]):
+
+            cross_slice_strategy = self.config.cross_slice_marker_score_strategy
+            n_slices = 1 + 2 * self.config.n_adjacent_slices
+            # For pooling strategies, num_homogeneous is per slice
+            num_homogeneous_per_slice = self.config.homogeneous_neighbors
+
+        self.computer = ParallelMarkerScoreComputer(
+            global_log_gmean,
+            global_expr_frac,
+            self.config.homogeneous_neighbors,
+            num_workers=self.config.mkscore_compute_workers,
+            input_queue=reader_to_computer_queue,  # Input from reader
+            output_queue=computer_to_writer_queue,  # Output to writer
+            cross_slice_strategy=cross_slice_strategy,
+            n_slices=n_slices,
+            num_homogeneous_per_slice=num_homogeneous_per_slice,
+            no_expression_fraction=self.config.no_expression_fraction
+        )
+
+        self.writer = ParallelMarkerScoreWriter(
+            output_memmap,
+            num_workers=self.config.mkscore_write_workers,
+            input_queue=computer_to_writer_queue  # Input from computer
+        )
+
+        logger.info(f"Processing pools initialized: {self.config.rank_read_workers} readers, "
+                   f"{self.config.mkscore_compute_workers} computers, "
+                   f"{self.config.mkscore_write_workers} writers")
+
+        self.marker_score_queue = MarkerScoreMessageQueue(
+            reader=self.reader,
+            computer=self.computer,
+            writer=self.writer,
+            batch_size=self.config.mkscore_batch_size
+        )
+
+    def _find_homogeneous_spots(
+        self,
+        adata: ad.AnnData,
+        cell_type: str,
+        annotation_key: str,
+        coords: np.ndarray | None,
+        emb_niche: np.ndarray,
+        emb_indv: np.ndarray,
+        slice_ids: np.ndarray | None,
+        rank_shape: tuple[int, int],
+        high_quality_mask: np.ndarray
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray | None, np.ndarray, int] | None:
+        """
+        Prepare batch data for a cell type
+
+        Returns:
+            Tuple of (neighbor_indices, cell_sims, niche_sims, cell_indices_sorted, n_cells)
+            or None if cell type should be skipped
+        """
+        # Get cells of this type, excluding NaN values
+        cell_mask = (adata.obs[annotation_key] == cell_type) & (adata.obs[annotation_key].notna())
+        cell_indices = np.where(cell_mask)[0]
+        n_cells = len(cell_indices)
+
+        # Check minimum cells
+        min_cells = self.config.min_cells_per_type
+        if n_cells < min_cells:
+            logger.warning(f"Skipping {cell_type}: only {n_cells} cells (min: {min_cells})")
+
+        logger.info(f"Processing {cell_type}: {n_cells} cells")
+
+        # Build connectivity matrix
+        logger.info("Building connectivity matrix...")
+        neighbor_indices, cell_sims, niche_sims = self.connectivity_builder.build_connectivity_matrix(
+            coords=coords,
+            emb_niche=emb_niche,
+            emb_indv=emb_indv,
+            cell_mask=cell_mask,
+            high_quality_mask=high_quality_mask,
+            slice_ids=slice_ids,
+            k_central=self.config.spatial_neighbors,
+            k_adjacent=self.config.adjacent_slice_spatial_neighbors,
+            n_adjacent_slices=self.config.n_adjacent_slices
+        )
+        gc.collect()
+
+        # Validate neighbor indices
+        max_valid_idx = rank_shape[0] - 1
+        assert neighbor_indices.max() <= max_valid_idx, \
+            f"Neighbor indices exceed bounds (max: {neighbor_indices.max()}, limit: {max_valid_idx})"
+
+        # Optimize row order using JAX implementation
+        logger.info("Optimizing row order for cache efficiency...")
+        row_order = optimize_row_order_jax(
+            neighbor_indices=neighbor_indices[:,:self.config.homogeneous_neighbors],
+            cell_indices=cell_indices,
+            neighbor_weights=cell_sims[:,:self.config.homogeneous_neighbors],
+        )
+
+        neighbor_indices = neighbor_indices[row_order]
+        cell_sims = cell_sims[row_order]
+        if niche_sims is not None:
+            niche_sims = niche_sims[row_order]
+        cell_indices_sorted = cell_indices[row_order]
+
+        # Save homogeneous neighbor data to adata
+        has_real_niche_embedding = self.config.latent_representation_niche is not None
+        self._save_homogeneous_data_to_adata(
+            adata=adata,
+            neighbor_indices=neighbor_indices,
+            cell_sims=cell_sims,
+            niche_sims=niche_sims,
+            cell_indices_sorted=cell_indices_sorted,
+            has_real_niche_embedding=has_real_niche_embedding
+        )
+
+        # warning for cells not find homo neighbors
+        homo_neighbor_count = np.count_nonzero(cell_sims > 0, axis=1)
+        zero_homo_neighbor_mask = (homo_neighbor_count <= 5)
+        if np.any(zero_homo_neighbor_mask):
+            logger.warning(f"Cell type {cell_type}: {zero_homo_neighbor_mask.sum()} cells can't find enough homogeneous neighbors")
+
+        return neighbor_indices, cell_sims, niche_sims, cell_indices_sorted, n_cells
+
+    def _save_homogeneous_data_to_adata(
+        self,
+        adata: ad.AnnData,
+        neighbor_indices: np.ndarray,
+        cell_sims: np.ndarray,
+        niche_sims: np.ndarray | None,
+        cell_indices_sorted: np.ndarray,
+        has_real_niche_embedding: bool
+    ):
+        """
+        Save homogeneous neighbor data to adata obsm and obs.
+
+        Args:
+            adata: AnnData object to save data to
+            neighbor_indices: (n_cells, k) array of neighbor indices
+            cell_sims: (n_cells, k) array of cell similarity scores
+            niche_sims: (n_cells, k) array of niche similarity scores or None
+            cell_indices_sorted: (n_cells,) array of cell indices in sorted order
+            has_real_niche_embedding: Whether real niche embedding was provided
+        """
+        # Initialize obsm matrices if they don't exist
+        if 'gsMap_homo_indices' not in adata.obsm.keys() or (adata.obsm['gsMap_homo_indices'].shape[1] != neighbor_indices.shape[1]):
+            adata.obsm['gsMap_homo_indices'] = np.zeros((adata.n_obs, neighbor_indices.shape[1]), dtype=neighbor_indices.dtype)
+        if 'gsMap_homo_cell_sims' not in adata.obsm.keys() or (adata.obsm['gsMap_homo_cell_sims'].shape[1] != neighbor_indices.shape[1]):
+            adata.obsm['gsMap_homo_cell_sims'] = np.zeros((adata.n_obs, cell_sims.shape[1]), dtype=cell_sims.dtype)
+
+        # Store the neighbor indices and cell_sims for this cell type
+        adata.obsm['gsMap_homo_indices'][cell_indices_sorted] = neighbor_indices
+        adata.obsm['gsMap_homo_cell_sims'][cell_indices_sorted] = cell_sims
+
+        # Only store niche_sims if niche embedding was provided (not dummy)
+        if has_real_niche_embedding and niche_sims is not None:
+            if 'gsMap_homo_niche_sims' not in adata.obsm.keys() or (adata.obsm['gsMap_homo_niche_sims'].shape[1] != neighbor_indices.shape[1]):
+                adata.obsm['gsMap_homo_niche_sims'] = np.zeros((adata.n_obs, niche_sims.shape[1]), dtype=niche_sims.dtype)
+            adata.obsm['gsMap_homo_niche_sims'][cell_indices_sorted] = niche_sims
+
+        # Initialize obs columns if they don't exist
+        if 'gsMap_homo_neighbor_count' not in adata.obs.columns:
+            adata.obs['gsMap_homo_neighbor_count'] = 0
+        if 'gsMap_homo_cell_sims_median' not in adata.obs.columns:
+            adata.obs['gsMap_homo_cell_sims_median'] = np.nan
+        if has_real_niche_embedding and 'gsMap_homo_niche_sims_median' not in adata.obs.columns:
+            adata.obs['gsMap_homo_niche_sims_median'] = np.nan
+
+        # Calculate statistics for each cell (only consider valid neighbors where cell_sims > 0)
+        valid_mask = cell_sims > 0  # (n_cells, k) boolean mask
+
+        # Count of valid homogeneous neighbors
+        homo_neighbor_count = valid_mask.sum(axis=1)
+        adata.obs.loc[adata.obs.index[cell_indices_sorted], 'gsMap_homo_neighbor_count'] = homo_neighbor_count
+
+        # Median of cell similarity (only for valid neighbors)
+        cell_sims_masked = np.where(valid_mask, cell_sims, np.nan)
+        cell_sims_median = np.nanmedian(cell_sims_masked, axis=1)
+        adata.obs.loc[adata.obs.index[cell_indices_sorted], 'gsMap_homo_cell_sims_median'] = cell_sims_median
+
+        # Median of niche similarity (only for valid neighbors, only if niche embedding provided)
+        if has_real_niche_embedding and niche_sims is not None:
+            niche_sims_masked = np.where(valid_mask, niche_sims, np.nan)
+            niche_sims_median = np.nanmedian(niche_sims_masked, axis=1)
+            adata.obs.loc[adata.obs.index[cell_indices_sorted], 'gsMap_homo_niche_sims_median'] = niche_sims_median
+
+    def _calculate_marker_scores_by_cell_type(
+        self,
+        adata: ad.AnnData,
+        cell_type: str,
+        coords: np.ndarray | None,
+        emb_niche: np.ndarray,
+        emb_indv: np.ndarray,
+        annotation_key: str,
+        slice_ids: np.ndarray | None = None,
+        high_quality_mask: np.ndarray = None
+    ):
+        """Process a single cell type with shared pools"""
+
+        # Find homogeneous spots
+        neighbor_indices, cell_sims, niche_sims, cell_indices_sorted, n_cells = self._find_homogeneous_spots(
+            adata, cell_type, annotation_key, coords, emb_niche,
+            emb_indv, slice_ids, self.reader.shape, high_quality_mask
+        )
+
+        # Reset the message queue for this cell type
+        self.marker_score_queue.reset_for_cell_type(cell_type, n_cells)
+
+        # Run the marker_score_queue to compute marker score
+        self.marker_score_queue.start(
+            neighbor_indices=neighbor_indices,
+            neighbor_weights=cell_sims,
+            cell_indices_sorted=cell_indices_sorted,
+            enable_profiling=self.config.enable_profiling
+        )
+
+    def calculate_marker_scores(
+        self,
+        adata_path: str,
+        rank_memmap_path: str,
+        mean_frac_path: str,
+        output_path: str | Path | None = None
+    ) -> str | Path:
+        """
+        Main execution function for marker score calculation
+
+        Args:
+            adata_path: Path to concatenated latent adata
+            rank_memmap_path: Path to rank memory map
+            mean_frac_path: Path to mean expression fraction parquet
+            output_path: Optional output path for marker scores
+
+        Returns:
+            Path to output marker score memory map file
+        """
+        logger.info("Starting marker score calculation...")
+        self.console = Console()
+
+        self.console.print(Panel.fit(
+            "[bold cyan]Marker Score Calculation Stage[/bold cyan]",
+            subtitle="gsMap Stage 2",
+            border_style="cyan"
+        ))
+
+        # Use config path if not specified
+        if output_path is None:
+            output_path = Path(self.config.marker_scores_memmap_path)
+        else:
+            output_path = Path(output_path)
+
+        # Load all input data
+        adata, rank_memmap, global_log_gmean, global_expr_frac, n_cells, n_genes, high_quality_mask = self._load_input_data(
+            adata_path, rank_memmap_path, mean_frac_path
+        )
+
+        # Initialize output memory map
+        output_memmap = MemMapDense(
+            output_path,
+            shape=(n_cells, n_genes),
+            dtype=np.float16,  # Use float16 to save memory
+            mode='w',
+            num_write_workers=self.config.mkscore_write_workers,
+            tmp_dir=self.config.memmap_tmp_dir
+        )
+
+        # Get cell types to process
+        cell_types = self._get_cell_types(adata)
+        annotation_key = self.config.annotation
+
+        # Display summary
+        self._display_input_summary(adata, cell_types, n_cells, n_genes)
+
+        # Prepare embeddings
+        coords, emb_niche, emb_indv, slice_ids = self._prepare_embeddings(adata)
+
+        # Initialize processing pipeline
+        self._initialize_pipeline(rank_memmap, output_memmap, global_log_gmean, global_expr_frac)
+
+        # Process each cell type
+        for cell_type in cell_types:
+            self._calculate_marker_scores_by_cell_type(
+                adata,
+                cell_type,
+                coords,
+                emb_niche,
+                emb_indv,
+                annotation_key,
+                slice_ids,
+                high_quality_mask
+            )
+
+        # Save updated AnnData with neighbor matrices
+        logger.info(f"Saving updated AnnData with homogeneous spot matrices to {adata_path}")
+        adata.write_h5ad(adata_path)
+
+        # Close all shared pools after all cell types are processed
+        logger.info("Closing shared processing pools...")
+        self.reader.close()
+        self.computer.close()
+        self.writer.close()
+
+        # Close memory maps
+        rank_memmap.close()
+        output_memmap.close()
+
+        self.console.print(Panel.fit(
+            "[bold green]✓ Marker score calculation complete![/bold green]",
+            border_style="green"
+        ))
+
+        logger.info(f"Results saved to {output_path}")
+
+        return str(output_path)