Machine Learning API¶

class SLAFDataLoader:
    """
    High-performance DataLoader for SLAF data optimized for ML training.

    SLAFDataLoader provides efficient streaming of single-cell data for machine learning
    applications with multiple loading strategies for different use cases. It uses async
    batch processing and provides device-agnostic CPU tensor output for maximum training flexibility.

    Key Features:
        - Multiple tokenization strategies (GeneFormer, scGPT)
        - Multiple loading modes for different entropy requirements:
            * Mixture of Scanners (MoS): Maximum entropy, best randomization (default)
            * Fragment-based loading: Higher entropy, moderate performance
            * Sequential loading: Fastest, lowest entropy
        - Pre-tokenized sequences for maximum performance (tokenized mode)
        - Raw data output for external processing (raw mode)
        - Device-agnostic CPU tensor output
        - Async batch processing with background prefetching
        - Memory-efficient streaming
        - Multi-epoch training support
        - Comprehensive error handling and validation

    Loading Modes:
        1. Mixture of Scanners (default): Randomly samples from multiple fragment generators
           for maximum entropy and randomization (88% of random entropy)
        2. Fragment-based: Loads complete Lance fragments for higher data entropy
        3. Sequential: Loads contiguous Lance batches for maximum throughput

    Examples:
        >>> # Basic usage with default settings (MoS loading)
        >>> slaf_array = SLAFArray("path/to/data.slaf")
        >>> dataloader = SLAFDataLoader(slaf_array)
        >>> for batch in dataloader:
        ...     print(f"Batch shape: {batch['input_ids'].shape}")
        ...     print(f"Cell IDs: {batch['cell_ids']}")
        ...     break
        Batch shape: torch.Size([32, 2048])
        Cell IDs: tensor([0, 1, 2, ..., 29, 30, 31])
        >>> print(f"MoS enabled: {dataloader.use_mixture_of_scanners}")
        MoS enabled: True

        >>> # Sequential loading for maximum throughput
        >>> dataloader = SLAFDataLoader(
        ...     slaf_array=slaf_array,
        ...     use_mixture_of_scanners=False,
        ...     by_fragment=False
        ... )
        >>> print(f"Sequential loading: {not dataloader.use_mixture_of_scanners}")
        Sequential loading: True

        >>> # Fragment-based loading for higher entropy
        >>> dataloader = SLAFDataLoader(
        ...     slaf_array=slaf_array,
        ...     use_mixture_of_scanners=False,
        ...     by_fragment=True
        ... )
        >>> print(f"Fragment-based loading: {dataloader.by_fragment}")
        Fragment-based loading: True

        >>> # Raw mode for external processing
        >>> dataloader = SLAFDataLoader(
        ...     slaf_array=slaf_array,
        ...     raw_mode=True
        ... )
        >>> for batch in dataloader:
        ...     print(f"Raw data type: {type(batch['x'])}")
        ...     break
        Raw data type: <class 'polars.dataframe.frame.DataFrame'>

        >>> # Multi-epoch training
        >>> dataloader = SLAFDataLoader(
        ...     slaf_array=slaf_array,
        ...     n_epochs=5
        ... )
        >>> print(f"Number of epochs: {dataloader.n_epochs}")
        Number of epochs: 5

        >>> # Custom configuration for training
        >>> dataloader = SLAFDataLoader(
        ...     slaf_array=slaf_array,
        ...     tokenizer_type="scgpt",
        ...     batch_size=64,
        ...     max_genes=1024
        ... )
        >>> print(f"Tokenizer type: {dataloader.tokenizer_type}")
        Tokenizer type: scgpt

        >>> # Training loop example
        >>> for batch_idx, batch in enumerate(dataloader):
        ...     input_ids = batch["input_ids"]
        ...     attention_mask = batch["attention_mask"]
        ...     cell_ids = batch["cell_ids"]
        ...     # Your training code here
        ...     if batch_idx >= 2:  # Just test first few batches
        ...         break
        >>> print("Training loop completed")
        Training loop completed

        >>> # Error handling for invalid tokenizer type
        >>> try:
        ...     dataloader = SLAFDataLoader(slaf_array, tokenizer_type="invalid")
        ... except ValueError as e:
        ...     print(f"Error: {e}")
        Error: Unsupported tokenizer type: invalid
    """

    device: Optional["torch.device"]  # type: ignore
    tokenizer: Optional["SLAFTokenizer"]  # type: ignore

    def __init__(
        self,
        slaf_array: SLAFArray,
        tokenizer_type: str = "geneformer",
        batch_size: int = 32,
        max_genes: int = 2048,
        vocab_size: int = 50000,
        n_expression_bins: int = 10,
        n_epochs: int = 1,  # Add n_epochs parameter
        raw_mode: bool = False,  # Add raw_mode parameter
        verbose: bool = True,  # Add verbose parameter
        batches_per_chunk: int = 1,  # Default to 1 for MoS (was 50 for sequential)
        by_fragment: bool = True,  # Default to True for MoS (was False for sequential)
        use_mixture_of_scanners: bool = True,  # Default to True for MoS (was False)
        n_scanners: int = 16,  # Add n_scanners parameter for MoS
        prefetch_batch_size: int = 4194304,  # Add prefetch_batch_size parameter for MoS
    ):
        """
        Initialize the SLAF DataLoader with training configuration.

        Args:
            slaf_array: SLAFArray instance containing the single-cell data.
                       Must be a valid SLAFArray with proper Lance dataset structure.

            # Tokenization Configuration
            tokenizer_type: Tokenization strategy to use. Options: "geneformer", "scgpt".
                          Geneformer uses ranked gene sequences, scGPT uses interleaved
                          gene-expression pairs. Ignored when raw_mode=True.
            max_genes: Maximum number of genes to include in each cell's tokenization.
                     For Geneformer: same as sequence length. For scGPT: number of
                     gene-expression pairs (sequence length = 2*max_genes+2).
            vocab_size: Size of the tokenizer vocabulary. Higher values allow more
                       genes but use more memory. Range: 1000-100000, default: 50000.
            n_expression_bins: Number of expression level bins for scGPT discretization.
                             Higher values provide finer expression resolution.
                             Range: 1-1000, default: 10.

            # Training Configuration
            batch_size: Number of cells per batch. Larger batches use more memory
                       but may improve training efficiency. Range: 1-512, default: 32.
            n_epochs: Number of epochs to run. The generator will automatically reset
                     after each epoch, enabling multi-epoch training on small datasets.
                     Default: 1.

            # Output Mode Configuration
            raw_mode: If True, return raw cell × gene data as Polars DataFrames
                     instead of pre-tokenized sequences. This bypasses tokenization
                     and windowing for maximum flexibility. Default: False.

            # Loading Strategy Configuration (MoS is now default)
            batches_per_chunk: Number of Lance batches to load per chunk for sequential loading.
                             Higher values use more memory but may improve throughput.
                             Range: 1-200, default: 1 (optimized for MoS). Only used when by_fragment=False.
            by_fragment: If True, use fragment-based loading instead of batch-based loading.
                        Fragment-based loading provides higher entropy but may be slightly slower.
                        Automatically enabled when use_mixture_of_scanners=True.
                        Default: True (enabled for MoS).
            use_mixture_of_scanners: If True, use mixture of scanners (MoS) approach for higher
                                   entropy by randomly sampling from multiple fragment generators.
                                   This provides the best randomization and is now the default
                                   for foundation model training. Default: True.
            n_scanners: Number of fragment generators to sample from simultaneously when using MoS.
                       Higher values provide better entropy but use more memory.
                       Range: 1-100, default: 16. Only used when use_mixture_of_scanners=True.
            prefetch_batch_size: Target number of rows to load per prefetch batch when using MoS.
                               Higher values improve throughput but use more memory.
                               Range: 1000-10000000, default: 4194304. Only used when
                               use_mixture_of_scanners=True.

            # System Configuration
            verbose: If True, print detailed timing and progress information.
                    If False, suppress all SLAF internal prints for clean output.
                    Default: True.

        Raises:
            ValueError: If tokenizer_type is not supported or parameters are invalid.
            RuntimeError: If PyTorch is not available or datasets module is missing.
            TypeError: If slaf_array is not a valid SLAFArray instance.
            ImportError: If required dependencies are not available.

        Loading Strategy Selection Guide:
            - For foundation model training: Use default settings (MoS provides 88% random entropy)
            - For maximum throughput: Set use_mixture_of_scanners=False, by_fragment=False
            - For external processing: Set raw_mode=True

        Examples:
            >>> # Basic initialization (MoS is now default)
            >>> slaf_array = SLAFArray("path/to/data.slaf")
            >>> dataloader = SLAFDataLoader(slaf_array)
            >>> print(f"Batch size: {dataloader.batch_size}")
            Batch size: 32
            >>> print(f"MoS enabled: {dataloader.use_mixture_of_scanners}")
            MoS enabled: True

            >>> # Sequential loading for maximum throughput
            >>> dataloader = SLAFDataLoader(
            ...     slaf_array=slaf_array,
            ...     use_mixture_of_scanners=False,
            ...     by_fragment=False
            ... )
            >>> print(f"Sequential loading: {not dataloader.use_mixture_of_scanners}")
            Sequential loading: True

            >>> # Fragment-based loading for higher entropy
            >>> dataloader = SLAFDataLoader(
            ...     slaf_array=slaf_array,
            ...     use_mixture_of_scanners=False,
            ...     by_fragment=True
            ... )
            >>> print(f"Fragment-based loading: {dataloader.by_fragment}")
            Fragment-based loading: True

            >>> # Raw mode for external processing
            >>> dataloader = SLAFDataLoader(
            ...     slaf_array=slaf_array,
            ...     raw_mode=True
            ... )
            >>> print(f"Raw mode: {dataloader.raw_mode}")
            Raw mode: True

            >>> # Error handling for invalid parameters
            >>> try:
            ...     dataloader = SLAFDataLoader(slaf_array, n_scanners=0)
            ... except ValueError as e:
            ...     print(f"Error: {e}")
            Error: n_scanners must be at least 1
        """
        self.slaf_array = slaf_array
        self.tokenizer_type = tokenizer_type
        self.batch_size = batch_size
        self.max_genes = max_genes
        self.n_epochs = n_epochs
        self.raw_mode = raw_mode  # Add raw_mode attribute
        self.verbose = verbose  # Add verbose attribute
        self.batches_per_chunk = batches_per_chunk  # Add batches_per_chunk attribute
        self.by_fragment = by_fragment  # Add by_fragment attribute
        self.use_mixture_of_scanners = use_mixture_of_scanners  # Add MoS attribute
        self.n_scanners = n_scanners  # Add n_scanners attribute
        self.prefetch_batch_size = (
            prefetch_batch_size  # Add prefetch_batch_size attribute
        )

        # Validate MoS parameters
        if self.use_mixture_of_scanners:
            if self.n_scanners < 1:
                raise ValueError("n_scanners must be at least 1")
            if self.n_scanners > 100:
                raise ValueError("n_scanners cannot exceed 100")
            if (
                self.prefetch_batch_size < 1000
            ):  # Allow smaller values for warm-up strategy
                raise ValueError("prefetch_batch_size must be at least 1,000")
            if self.prefetch_batch_size > 10000000:
                raise ValueError("prefetch_batch_size cannot exceed 10,000,000")

        # Device-agnostic: always return CPU tensors
        self.device = None

        # Check that required modules are available
        if not DATASETS_AVAILABLE:
            raise ImportError(
                "SLAFIterableDataset is required but not available. Please install required dependencies."
            )

        # Initialize tokenizer (only needed for non-raw mode)
        if not self.raw_mode:
            self.tokenizer = SLAFTokenizer(
                slaf_array=slaf_array,
                tokenizer_type=tokenizer_type,
                vocab_size=vocab_size,
                n_expression_bins=n_expression_bins,
            )

            # Get special tokens from tokenizer
            self.special_tokens = self.tokenizer.special_tokens
        else:
            # For raw mode, we don't need a tokenizer
            self.tokenizer = None
            self.special_tokens = None

        # Use IterableDataset
        self._dataset = SLAFIterableDataset(
            slaf_array=slaf_array,
            tokenizer=self.tokenizer,
            batch_size=batch_size,
            seed=42,  # TODO: make configurable
            max_queue_size=500,
            tokenizer_type=tokenizer_type,
            n_epochs=n_epochs,  # Pass n_epochs to dataset
            raw_mode=raw_mode,  # Pass raw_mode to dataset
            verbose=verbose,  # Pass verbose to dataset
            batches_per_chunk=batches_per_chunk,  # Pass batches_per_chunk to dataset
            by_fragment=by_fragment,  # Pass by_fragment to dataset
            use_mixture_of_scanners=use_mixture_of_scanners,  # Pass MoS to dataset
            n_scanners=n_scanners,  # Pass n_scanners to dataset
            prefetch_batch_size=prefetch_batch_size,  # Pass prefetch_batch_size to dataset
        )

    def __iter__(self):
        """
        Iterate through batches of single-cell data based on the configured mode.

        Yields batches of data suitable for machine learning training. The output format
        depends on the configuration:

        - **Tokenized mode** (default): Yields pre-tokenized sequences with attention masks
        - **Raw mode**: Yields raw Polars DataFrames for external processing
        - **Multi-epoch**: Automatically handles epoch transitions when n_epochs > 1

        The loading strategy (sequential, fragment-based, or Mixture of Scanners) affects
        data entropy and throughput but not the output format.

        Yields:
            dict: Batch dictionary containing:
                - **Tokenized mode** (raw_mode=False):
                    - input_ids: Pre-tokenized gene expression data (torch.Tensor)
                    - attention_mask: Boolean mask indicating valid tokens (torch.Tensor)
                    - cell_ids: Integer IDs of cells in the batch (torch.Tensor)
                - **Raw mode** (raw_mode=True):
                    - x: Raw cell × gene data as Polars DataFrame
                    - cell_ids: List of cell integer IDs in the batch
                - **Multi-epoch** (when n_epochs > 1):
                    - epoch: Current epoch number (int)

        Note:
            All tensors are returned on CPU for device-agnostic training.
            The training loop should handle device transfer as needed.

        Raises:
            ValueError: If the tokenizer type is not supported.
            RuntimeError: If batch processing fails.

        Examples:
            >>> # Basic iteration (tokenized mode)
            >>> slaf_array = SLAFArray("path/to/data.slaf")
            >>> dataloader = SLAFDataLoader(slaf_array, batch_size=16)
            >>> for batch in dataloader:
            ...     print(f"Batch keys: {list(batch.keys())}")
            ...     print(f"Input shape: {batch['input_ids'].shape}")
            ...     print(f"Cell IDs: {batch['cell_ids']}")
            ...     break
            Batch keys: ['input_ids', 'attention_mask', 'cell_ids']
            Input shape: (16, 2048)
            Cell IDs: tensor([0, 1, 2, ..., 13, 14, 15])

            >>> # Raw mode iteration
            >>> dataloader = SLAFDataLoader(slaf_array, raw_mode=True, batch_size=16)
            >>> for batch in dataloader:
            ...     print(f"Raw data type: {type(batch['x'])}")
            ...     print(f"Cell IDs: {batch['cell_ids']}")
            ...     break
            Raw data type: <class 'polars.dataframe.frame.DataFrame'>
            Cell IDs: [0, 1, 2, ..., 13, 14, 15]

            >>> # Multi-epoch training
            >>> dataloader = SLAFDataLoader(slaf_array, n_epochs=3)
            >>> epochs_seen = set()
            >>> for batch in dataloader:
            ...     if 'epoch' in batch:
            ...         epochs_seen.add(batch['epoch'])
            ...     if len(epochs_seen) >= 3:  # Stop after seeing all epochs
            ...         break
            >>> print(f"Epochs completed: {sorted(epochs_seen)}")
            Epochs completed: [0, 1, 2]

            >>> # Training loop with error handling
            >>> for batch_idx, batch in enumerate(dataloader):
            ...     try:
            ...         if 'input_ids' in batch:  # Tokenized mode
            ...             input_ids = batch["input_ids"]
            ...             attention_mask = batch["attention_mask"]
            ...             cell_ids = batch["cell_ids"]
            ...         else:  # Raw mode
            ...             x = batch["x"]
            ...             cell_ids = batch["cell_ids"]
            ...         # Your training code here
            ...         print(f"Processed batch {batch_idx}")
            ...     except Exception as e:
            ...         print(f"Error in batch {batch_idx}: {e}")
            ...         continue
            ...     if batch_idx >= 2:  # Just first few batches
            ...         break
            Processed batch 0
            Processed batch 1
            Processed batch 2

            >>> # Different tokenizer types
            >>> dataloader_geneformer = SLAFDataLoader(slaf_array, tokenizer_type="geneformer")
            >>> dataloader_scgpt = SLAFDataLoader(slaf_array, tokenizer_type="scgpt")
            >>>
            >>> # Compare batch shapes
            >>> for batch in dataloader_geneformer:
            ...     print(f"Geneformer input shape: {batch['input_ids'].shape}")
            ...     break
            Geneformer input shape: (32, 2048)
            >>> for batch in dataloader_scgpt:
            ...     print(f"scGPT input shape: {batch['input_ids'].shape}")
            ...     break
            scGPT input shape: (32, 1024)
        """
        yield from self._dataset

    def __len__(self):
        """
        Return the number of batches in the dataset.

        Note: Since SLAFDataLoader uses an IterableDataset that streams data,
        the exact number of batches is not known in advance. This method
        returns 0 to indicate an unknown length for streaming datasets.

        Returns:
            int: Always returns 0 to indicate unknown length for streaming datasets.

        Examples:
            >>> # Check dataset length
            >>> slaf_array = SLAFArray("path/to/data.slaf")
            >>> dataloader = SLAFDataLoader(slaf_array)
            >>> print(f"Dataset length: {len(dataloader)}")
            Dataset length: 0

            >>> # IterableDataset behavior
            >>> batch_count = 0
            >>> for batch in dataloader:
            ...     batch_count += 1
            ...     if batch_count >= 5:  # Just count first 5 batches
            ...         break
            >>> print(f"Actually processed {batch_count} batches")
            Actually processed 5 batches

            >>> # Length is consistent
            >>> print(f"Length check: {len(dataloader)}")
            Length check: 0
        """
        return 0  # Indicates unknown length

    def __del__(self):
        """
        Cleanup method to stop async prefetching.

        This method is called when the DataLoader object is garbage collected.
        It ensures that the underlying dataset's prefetcher is properly cleaned up
        to prevent resource leaks.

        Examples:
            >>> # DataLoader cleanup happens automatically
            >>> slaf_array = SLAFArray("path/to/data.slaf")
            >>> dataloader = SLAFDataLoader(slaf_array)
            >>> print("DataLoader created")
            DataLoader created
            >>> # When dataloader goes out of scope, __del__ is called automatically
            >>> del dataloader
            >>> print("DataLoader destroyed and cleaned up")
            DataLoader destroyed and cleaned up

            >>> # Manual cleanup (not usually needed)
            >>> dataloader = SLAFDataLoader(slaf_array)
            >>> dataloader.__del__()
            >>> print("Manual cleanup completed")
            Manual cleanup completed
        """
        if hasattr(self, "_dataset"):
            # The SLAFIterableDataset doesn't have a stop method,
            # so we just let it finish its current epoch.
            pass

@dataclass
class TileDBPrefetchBatch:
    """
    Container for a batch of TileDB data with metadata.

    This dataclass holds a processed batch of TileDB SOMA data along with
    associated metadata for tracking performance and debugging. It serves as
    the primary data structure passed between the batch processor and the
    async prefetcher.

    Attributes:
        batch_id: Unique identifier for this batch within the current epoch.
        batch_df: Polars DataFrame containing the cell-gene expression data
                 with columns: cell_integer_id, gene_integer_id, value.
        cell_integer_ids: List of unique cell IDs present in this batch.
        process_time: Time taken to process this batch (in seconds).
        memory_mb: Memory usage at the time of batch creation (in MB).

    Examples:
        >>> # Create a batch container
        >>> batch = TileDBPrefetchBatch(
        ...     batch_id=0,
        ...     batch_df=df,
        ...     cell_integer_ids=[0, 1, 2, 3],
        ...     process_time=0.1,
        ...     memory_mb=128.5
        ... )
        >>> print(f"Batch {batch.batch_id} has {len(batch.cell_integer_ids)} cells")
        Batch 0 has 4 cells
    """

    batch_id: int
    batch_df: (
        pl.DataFrame
    )  # Polars DataFrame with cell_integer_id, gene_integer_id, value
    cell_integer_ids: list[int]  # List of cell IDs in this batch
    process_time: float
    memory_mb: float

class TileDBBatchProcessor:
    """
    High-performance batch processor for TileDB SOMA data with multiple loading strategies.

    TileDBBatchProcessor provides efficient streaming and processing of single-cell data
    from TileDB SOMA format. It supports multiple loading strategies including Mixture
    of Scanners (MoS) for maximum entropy and sequential loading for maximum throughput.

    Key Features:
        - Multiple loading strategies:
            * Mixture of Scanners (MoS): Random sampling from multiple generators for
              maximum entropy and randomization (default)
            * Sequential loading: Contiguous data loading for maximum throughput
        - Streaming data processing with configurable batch sizes
        - Built-in shuffling strategies for data randomization
        - Multi-epoch training support with automatic epoch transitions
        - Comprehensive timing and memory monitoring
        - Error handling and recovery mechanisms
        - Configurable prefetch batch sizes for different dataset sizes

    Loading Strategies:
        1. Mixture of Scanners (default): Randomly samples from multiple fragment
           generators for maximum entropy and randomization
        2. Sequential: Loads contiguous data chunks for maximum throughput

    Examples:
        >>> # Basic usage with default MoS strategy
        >>> processor = TileDBBatchProcessor(
        ...     tiledb_path="path/to/experiment",
        ...     batch_size=32,
        ...     prefetch_batch_size=100
        ... )
        >>> batch = processor.load_prefetch_batch()
        >>> print(f"Loaded batch with {len(batch.cell_integer_ids)} cells")
        Loaded batch with 100 cells

        >>> # Sequential loading for maximum throughput
        >>> processor = TileDBBatchProcessor(
        ...     tiledb_path="path/to/experiment",
        ...     use_mixture_of_scanners=False,
        ...     batch_size=64
        ... )
        >>> print(f"MoS enabled: {processor.use_mixture_of_scanners}")
        MoS enabled: False

        >>> # Multi-epoch training
        >>> processor = TileDBBatchProcessor(
        ...     tiledb_path="path/to/experiment",
        ...     n_epochs=3
        ... )
        >>> print(f"Number of epochs: {processor.n_epochs}")
        Number of epochs: 3
    """

    def __init__(
        self,
        tiledb_path: str,
        batch_size: int = 32,
        prefetch_batch_size: int = 100,
        seed: int = 42,
        n_epochs: int = 1,
        verbose: bool = True,
        log_metrics: bool = False,
        use_mixture_of_scanners: bool = True,
        n_readers: int = 50,
        n_scanners: int = 8,
    ):
        """
        Initialize the TileDB batch processor with training configuration.

        Args:
            tiledb_path: Path to the TileDB SOMA experiment directory.
                         Must contain a valid SOMA experiment with RNA measurement data.
            batch_size: Number of cells per training batch. Larger batches use more
                       memory but may improve training efficiency. Range: 1-512, default: 32.
            prefetch_batch_size: Number of cells to load per prefetch batch from TileDB.
                               Higher values improve throughput but use more memory.
                               Range: 10-10000, default: 100.
            seed: Random seed for reproducible shuffling and MoS sampling.
                  Used for consistent data ordering across runs. Default: 42.
            n_epochs: Number of epochs to run. The processor will automatically reset
                     after each epoch, enabling multi-epoch training. Default: 1.
            verbose: If True, print detailed timing and progress information.
                    If False, suppress all internal prints for clean output. Default: True.
            log_metrics: If True, collect detailed timing metrics for performance analysis.
                        Metrics include loading time, shuffle time, and memory usage.
                        Default: False.
            use_mixture_of_scanners: If True, use MoS strategy for higher entropy by
                                   randomly sampling from multiple fragment generators.
                                   Provides better randomization for foundation model training.
                                   Default: True.
            n_readers: Total number of fragment generators to create when using MoS.
                      Higher values provide better entropy but use more memory.
                      Range: 1-1000, default: 50.
            n_scanners: Number of active scanners to sample from simultaneously when using MoS.
                       Higher values provide better entropy but use more memory.
                       Range: 1-100, default: 8.

        Raises:
            ImportError: If TileDB SOMA is not available.
            ValueError: If MoS parameters are invalid (n_readers < 1, n_scanners < 1,
                       or n_scanners > n_readers).
            RuntimeError: If the TileDB experiment cannot be opened or is invalid.

        Examples:
            >>> # Basic initialization with default MoS strategy
            >>> processor = TileDBBatchProcessor(
            ...     tiledb_path="path/to/experiment",
            ...     batch_size=32,
            ...     prefetch_batch_size=100
            ... )
            >>> print(f"Total cells: {processor.total_cells}")
            Total cells: 50000

            >>> # Sequential loading for maximum throughput
            >>> processor = TileDBBatchProcessor(
            ...     tiledb_path="path/to/experiment",
            ...     use_mixture_of_scanners=False,
            ...     batch_size=64
            ... )
            >>> print(f"MoS enabled: {processor.use_mixture_of_scanners}")
            MoS enabled: False

            >>> # High-entropy MoS configuration
            >>> processor = TileDBBatchProcessor(
            ...     tiledb_path="path/to/experiment",
            ...     n_readers=100,
            ...     n_scanners=16
            ... )
            >>> print(f"MoS readers: {processor.n_readers}, scanners: {processor.n_scanners}")
            MoS readers: 100, scanners: 16
        """
        if not TILEDB_AVAILABLE:
            raise ImportError("TileDB SOMA is required but not available")

        self.tiledb_path = tiledb_path
        self.batch_size = batch_size
        self.prefetch_batch_size = prefetch_batch_size
        self.seed = seed
        self.n_epochs = n_epochs
        self.verbose = verbose
        self.log_metrics = log_metrics
        self.use_mixture_of_scanners = use_mixture_of_scanners
        self.n_readers = n_readers
        self.n_scanners = n_scanners

        # Validate MoS parameters
        if self.use_mixture_of_scanners:
            if self.n_readers < 1:
                raise ValueError("n_readers must be at least 1")
            if self.n_scanners < 1:
                raise ValueError("n_scanners must be at least 1")
            if self.n_scanners > self.n_readers:
                raise ValueError("n_scanners cannot exceed n_readers")

        # Initialize state
        self.batch_id = 0
        self.current_epoch = 0
        self.total_cells = 0

        # Open TileDB experiment
        self.experiment = tiledbsoma.Experiment.open(tiledb_path)
        self.X = self.experiment.ms["RNA"].X["data"]

        # Get total number of cells
        self.total_cells = self.X.shape[0]

        # Initialize shuffling strategy (similar to SLAF)
        from slaf.ml.samplers import ShuffleType, create_shuffle

        self.shuffle = create_shuffle(ShuffleType.RANDOM)

        # Initialize MoS generators if enabled
        if self.use_mixture_of_scanners:
            self._initialize_mos_generators()

        # Initialize timing metrics for benchmarking
        self._timing_metrics: dict[str, list[float]] | None
        if self.log_metrics:
            self._timing_metrics = {
                "tiledb_loading": [],
                "shuffle": [],
                "total": [],
                "cells_processed": [],
            }
        else:
            self._timing_metrics = None

        # Initialize timing variables for consolidated reporting
        self._last_load_time = 0.0
        self._last_memory_mb = 0.0

    def _initialize_mos_generators(self):
        """Initialize MoS generators with evenly distributed scan ranges."""
        # Calculate scan ranges for each generator
        cells_per_reader = self.total_cells // self.n_readers
        remainder = self.total_cells % self.n_readers

        self.generators = []
        current_position = 0

        for i in range(self.n_readers):
            # Distribute remainder cells among first few readers
            reader_cell_count = cells_per_reader + (1 if i < remainder else 0)

            generator = {
                "generator_id": i,
                "start_position": current_position,
                "current_position": current_position,
                "end_position": current_position + reader_cell_count,
                "is_active": True,
            }

            self.generators.append(generator)
            current_position += reader_cell_count

        if self.verbose:
            print_prefetch(
                f"TileDB MoS initialized: {self.n_readers} generators, "
                f"{self.n_scanners} active scanners, "
                f"prefetch_batch_size={self.prefetch_batch_size}",
                self.verbose,
            )

    def reset_for_epoch(self, epoch: int) -> None:
        """
        Reset the processor for a new epoch.

        This method resets the batch processor state to start a new epoch,
        including resetting batch counters, MoS generator positions, and
        shuffling seeds. It is called automatically during multi-epoch training.

        Args:
            epoch: The epoch number to start (0-based indexing).
                  Must be 0 <= epoch < n_epochs.

        Raises:
            ValueError: If epoch is invalid (negative or >= n_epochs).

        Examples:
            >>> # Reset for epoch 1
            >>> processor = TileDBBatchProcessor("path/to/experiment", n_epochs=3)
            >>> processor.reset_for_epoch(1)
            >>> print(f"Current epoch: {processor.current_epoch}")
            Current epoch: 1

            >>> # Invalid epoch raises error
            >>> try:
            ...     processor.reset_for_epoch(5)  # n_epochs=3
            ... except ValueError as e:
            ...     print(f"Error: {e}")
            Error: Invalid epoch 5. Must be 0 <= epoch < 3
        """
        if epoch < 0 or epoch >= self.n_epochs:
            raise ValueError(
                f"Invalid epoch {epoch}. Must be 0 <= epoch < {self.n_epochs}"
            )

        self.current_epoch = epoch
        self.batch_id = 0

        # Reset MoS generators if enabled
        if self.use_mixture_of_scanners:
            for generator in self.generators:
                generator["current_position"] = generator["start_position"]
                generator["is_active"] = True

        if self.verbose:
            print(f"🔄 Reset TileDB processor for epoch {epoch}")

    def _record_timing(self, step: str, duration: float, cells_processed: int = 0):
        """Record timing for a processing step."""
        if not self.log_metrics or self._timing_metrics is None:
            return

        if step in self._timing_metrics:
            self._timing_metrics[step].append(duration)

        if cells_processed > 0:
            self._timing_metrics["cells_processed"].append(cells_processed)

    def load_prefetch_batch(self) -> TileDBPrefetchBatch:
        """
        Load and process a chunk of TileDB data into batches using configured strategy.

        This method loads a batch of data from TileDB SOMA format, applies shuffling,
        and returns a processed batch ready for training. It supports both MoS and
        sequential loading strategies and handles epoch transitions automatically.

        The method performs the following steps:
        1. Load data from TileDB using the configured strategy (MoS or sequential)
        2. Convert Arrow data to Polars DataFrame
        3. Apply shuffling strategy for data randomization
        4. Return processed batch with metadata

        Returns:
            TileDBPrefetchBatch: Processed batch containing:
                - batch_df: Polars DataFrame with cell-gene expression data
                - cell_integer_ids: List of unique cell IDs in the batch
                - process_time: Time taken to process the batch
                - memory_mb: Memory usage at batch creation time

        Raises:
            StopIteration: When all epochs are completed and no more data is available.
            RuntimeError: If TileDB data loading fails.

        Examples:
            >>> # Load a batch with MoS strategy
            >>> processor = TileDBBatchProcessor(
            ...     tiledb_path="path/to/experiment",
            ...     use_mixture_of_scanners=True
            ... )
            >>> batch = processor.load_prefetch_batch()
            >>> print(f"Batch {batch.batch_id} has {len(batch.cell_integer_ids)} cells")
            Batch 0 has 100 cells

            >>> # Load a batch with sequential strategy
            >>> processor = TileDBBatchProcessor(
            ...     tiledb_path="path/to/experiment",
            ...     use_mixture_of_scanners=False
            ... )
            >>> batch = processor.load_prefetch_batch()
            >>> print(f"Sequential batch shape: {batch.batch_df.shape}")
            Sequential batch shape: (100, 3)

            >>> # Handle epoch completion
            >>> processor = TileDBBatchProcessor("path/to/experiment", n_epochs=1)
            >>> try:
            ...     while True:
            ...         batch = processor.load_prefetch_batch()
            ...         print(f"Processed batch {batch.batch_id}")
            ... except StopIteration:
            ...     print("All epochs completed")
            All epochs completed
        """
        # Iterative approach to handle epoch transitions
        while True:
            start_time = time.time()

            if self.use_mixture_of_scanners:
                # MoS approach: randomly sample from active generators
                import numpy as np

                # Get indices of currently active generators
                active_generators = [g for g in self.generators if g["is_active"]]

                if not active_generators:
                    # Check if we should start a new epoch
                    if self.current_epoch + 1 < self.n_epochs:
                        if self.verbose:
                            print(
                                f"🔄 Epoch {self.current_epoch} complete, starting epoch {self.current_epoch + 1}"
                            )
                        self.reset_for_epoch(self.current_epoch + 1)
                        continue
                    else:
                        raise StopIteration("No more epochs available") from None

                # Randomly sample from active generators
                n_to_sample = min(self.n_scanners, len(active_generators))
                selected_generators = np.random.choice(
                    active_generators, size=n_to_sample, replace=False
                )

                if self.verbose and self.batch_id % 10 == 0:
                    print_prefetch(
                        f"TileDB MoS sampling: {len(active_generators)} active generators, "
                        f"sampling from {n_to_sample} generators",
                        self.verbose,
                    )

                # Load data from selected generators
                load_start = time.time()
                batch_dfs = []

                for generator in selected_generators:
                    try:
                        start_cell = generator["current_position"]
                        end_cell = min(
                            start_cell + self.prefetch_batch_size,
                            generator["end_position"],
                        )

                        if start_cell >= generator["end_position"]:
                            # Generator exhausted
                            generator["is_active"] = False
                            continue

                        # Read slice from TileDB
                        arrow_data = (
                            self.X.read((slice(start_cell, end_cell),))
                            .tables()
                            .concat()
                        )

                        # Convert Arrow table to Polars DataFrame
                        df = pl.from_arrow(arrow_data)  # type: ignore[assignment]
                        if not isinstance(df, pl.DataFrame):
                            raise TypeError("Expected DataFrame from Arrow table")

                        # Rename SOMA columns to expected names
                        df = df.rename(
                            {
                                "soma_dim_0": "cell_integer_id",
                                "soma_dim_1": "gene_integer_id",
                                "soma_data": "value",
                            }
                        )

                        batch_dfs.append(df)

                        # Update generator position
                        generator["current_position"] = end_cell

                        # Mark as inactive if exhausted
                        if generator["current_position"] >= generator["end_position"]:
                            generator["is_active"] = False

                    except Exception as e:
                        logger.error(
                            f"Error loading TileDB data from generator {generator['generator_id']}: {e}"
                        )
                        generator["is_active"] = False
                        continue

                if not batch_dfs:
                    # All selected generators are exhausted, continue to next iteration
                    continue

                # Combine all batches
                combined_df_mos = pl.concat(batch_dfs, how="vertical")
            else:
                # Sequential approach (original implementation)
                current_position = (
                    self.batch_id * self.prefetch_batch_size
                ) % self.total_cells

                # Only check for epoch transitions when we actually wrap around
                if self.batch_id > 0:
                    prev_position = (
                        (self.batch_id - 1) * self.prefetch_batch_size
                    ) % self.total_cells
                    if current_position < prev_position:  # We wrapped around
                        if self.current_epoch + 1 < self.n_epochs:
                            if self.verbose:
                                print(
                                    f"🔄 Epoch {self.current_epoch} complete, starting epoch {self.current_epoch + 1}"
                                )
                            self.reset_for_epoch(self.current_epoch + 1)
                            continue
                        else:
                            raise StopIteration("No more epochs available") from None

                # Load data from TileDB
                load_start = time.time()
                try:
                    # Read slice from TileDB as Arrow table
                    arrow_data = (
                        self.X.read(
                            (
                                slice(
                                    current_position,
                                    current_position + self.prefetch_batch_size,
                                ),
                            )
                        )
                        .tables()
                        .concat()
                    )

                    # Convert Arrow table to Polars DataFrame
                    combined_df = pl.from_arrow(arrow_data)  # type: ignore[assignment]
                    if not isinstance(combined_df, pl.DataFrame):
                        raise TypeError("Expected DataFrame from Arrow table")

                    # Rename SOMA columns to expected names
                    combined_df = combined_df.rename(
                        {
                            "soma_dim_0": "cell_integer_id",
                            "soma_dim_1": "gene_integer_id",
                            "soma_data": "value",
                        }
                    )

                except Exception as e:
                    logger.error(f"Error loading TileDB data: {e}")
                    raise StopIteration(f"Failed to load TileDB data: {e}") from e

            load_time = time.time() - load_start
            self._record_timing("tiledb_loading", load_time)

            # Print detailed loading breakdown every 10 batches
            if self.batch_id % 10 == 0:
                import psutil

                process = psutil.Process()
                memory_mb = process.memory_info().rss / 1024 / 1024

                # Store timing info for consolidated report
                self._last_load_time = load_time
                self._last_memory_mb = memory_mb

            # Apply shuffling strategy
            shuffle_start = time.time()

            # Apply shuffling with chunking
            if self.use_mixture_of_scanners:
                shuffled_chunks = self.shuffle.apply(
                    combined_df_mos,  # type: ignore[arg-type]
                    self.seed + self.batch_id + self.current_epoch * 10000,
                    batch_size=self.batch_size,
                )
            else:
                shuffled_chunks = self.shuffle.apply(
                    combined_df,  # type: ignore[arg-type]
                    self.seed + self.batch_id + self.current_epoch * 10000,
                    batch_size=self.batch_size,
                )

            shuffle_time = time.time() - shuffle_start
            total_time = time.time() - start_time

            # Record timing metrics
            self._record_timing("shuffle", shuffle_time)
            self._record_timing("total", total_time)

            # Count total cells across all chunks
            total_cells_in_chunks = sum(
                len(chunk.get_column("cell_integer_id").unique())
                for chunk in shuffled_chunks
                if isinstance(chunk, pl.DataFrame)
            )

            # Record cells processed
            self._record_timing("cells_processed", 0, total_cells_in_chunks)

            # Print consolidated prefetch batch reporting
            if self.batch_id % 10 == 0:
                strategy_name = "MoS" if self.use_mixture_of_scanners else "sequential"
                prefetch_report = f"TileDB {strategy_name} prefetch batch {self.batch_id} (epoch {self.current_epoch}):\n"
                prefetch_report += f"   TileDB loading: {self._last_load_time * 1000:.1f}ms ({self.prefetch_batch_size} cells)\n"
                prefetch_report += (
                    f"   Processing: {shuffle_time * 1000:.1f}ms shuffle\n"
                )
                prefetch_report += f"   Total: {total_time * 1000:.1f}ms, {len(shuffled_chunks)} chunks, {total_cells_in_chunks} cells, {self._last_memory_mb:.1f} MB"

                print_prefetch(prefetch_report, self.verbose)

            self.batch_id += 1

            # Return the first chunk as a batch (we'll handle multiple chunks in the dataloader)
            if shuffled_chunks:
                first_chunk = shuffled_chunks[0]
                return TileDBPrefetchBatch(
                    batch_id=self.batch_id - 1,
                    batch_df=first_chunk,
                    cell_integer_ids=first_chunk["cell_integer_id"].unique().to_list(),
                    process_time=shuffle_time,
                    memory_mb=self._last_memory_mb,
                )
            else:
                # No data in this chunk, continue to next iteration
                continue