optional ExpressionPreprocessor for COO values before rank/bin

slaf/ml/__init__.py

@@ -11,6 +11,7 @@
     PrefetchBatchProcessor,
     SLAFIterableDataset,
 )
+from .expression_preprocessor import ExpressionPreprocessor
 from .samplers import (
     RandomShuffle,
     Shuffle,
@@ -24,6 +25,7 @@
 __all__ = [
     # Core DataLoaders
     "SLAFDataLoader",
+    "ExpressionPreprocessor",
     "TileDBDataLoader",
     "TileDBIterableDataset",
     # Dataset and Processing

slaf/ml/aggregators.py

@@ -12,6 +12,10 @@
 import polars as pl
 
 from slaf.core.tabular_schema import DataSchema
+from slaf.ml.expression_preprocessor import (
+    ExpressionPreprocessor,
+    apply_expression_preprocessor,
+)
 
 
 class Window(ABC):
@@ -84,6 +88,12 @@ def apply(
             "use_binned_expressions", True
         )  # Default to True for scGPT
 
+        preprocessor = kwargs.get("expression_preprocessor")
+        fragment_df = apply_expression_preprocessor(fragment_df, schema, preprocessor)
+        already_log1p = (
+            isinstance(preprocessor, ExpressionPreprocessor) and preprocessor.log1p
+        )
+
         if use_binned_expressions:
             grouped = (
                 fragment_df.with_columns(
@@ -93,7 +103,11 @@ def apply(
                     .alias("gene_rank")
                 )
                 .filter(pl.col("gene_rank") <= max_items)
-                .with_columns(pl.col(vk).log1p().alias("log_value"))
+                .with_columns(
+                    (pl.col(vk) if already_log1p else pl.col(vk).log1p()).alias(
+                        "log_value"
+                    )
+                )
                 .with_columns(
                     pl.when(pl.col("log_value") > 0)
                     .then(
@@ -164,6 +178,9 @@ def apply(
 
         min_percentile = kwargs.get("min_percentile", None)
 
+        preprocessor = kwargs.get("expression_preprocessor")
+        fragment_df = apply_expression_preprocessor(fragment_df, schema, preprocessor)
+
         if min_percentile is not None:
             grouped = (
                 fragment_df.with_columns(
@@ -237,6 +254,9 @@ def apply(
         item_out = schema.item_list_key
         value_out = schema.value_list_key or "expr_sequence"
 
+        preprocessor = kwargs.get("expression_preprocessor")
+        fragment_df = apply_expression_preprocessor(fragment_df, schema, preprocessor)
+
         result = fragment_df.with_columns(
             [
                 pl.col(vk)

slaf/ml/dataloaders.py

@@ -4,6 +4,7 @@
 
 from slaf.core.slaf import SLAFArray
 
+from .expression_preprocessor import ExpressionPreprocessor
 from .tokenizers import GeneformerTokenizer, ScGPTTokenizer, SLAFTokenizer
 
 # Try to import torch, but make it optional
@@ -284,6 +285,7 @@ def __init__(
         max_queue_size: int = 5000,  # Add max_queue_size parameter
         parallelize_fragment_reads: bool = False,  # Parallelize fragment reads in MoS (cloud optimization)
         prefetcher_ready_timeout: float = 10.0,  # Timeout for waiting for prefetcher to be ready
+        expression_preprocessor: ExpressionPreprocessor | None = None,
     ):
         """
         Initialize the SLAF DataLoader with training configuration.
@@ -339,6 +341,8 @@ def __init__(
                                Higher values improve throughput but use more memory.
                                Range: 1000-10000000, default: 4194304. Only used when
                                use_mixture_of_scanners=True.
+            expression_preprocessor: Optional per-cell normalize_total and element-wise log1p on
+                COO values before rank/bin. Default None preserves the historical window behavior.
             parallelize_fragment_reads: Whether to parallelize fragment reads in MoS mode.
                                        Critical for cloud scenarios where network latency
                                        dominates (can improve throughput 10-30x). For local
@@ -501,6 +505,7 @@ def __init__(
             prefetch_batch_size=prefetch_batch_size,  # Pass prefetch_batch_size to dataset
             parallelize_fragment_reads=parallelize_fragment_reads,  # Pass parallelize_fragment_reads
             prefetcher_ready_timeout=prefetcher_ready_timeout,  # Pass prefetcher_ready_timeout
+            expression_preprocessor=expression_preprocessor,
         )
 
     def __iter__(self):

slaf/ml/datasets.py

@@ -55,6 +55,7 @@
 
 from slaf.core.slaf import SLAFArray
 from slaf.core.tabular_schema import SLAF_LANCE_COO_SCHEMA
+from slaf.ml.expression_preprocessor import ExpressionPreprocessor
 from slaf.ml.samplers import Shuffle
 from slaf.ml.tokenizers import SLAFTokenizer
 
@@ -328,6 +329,7 @@ def __init__(
         n_scanners: int = 16,  # Add n_scanners parameter for MoS
         prefetch_batch_size: int = 4194304,  # Add prefetch_batch_size parameter for MoS
         parallelize_fragment_reads: bool = False,  # Parallelize fragment reads in MoS (cloud optimization)
+        expression_preprocessor: ExpressionPreprocessor | None = None,
     ):
         """Initialize the PrefetchBatchProcessor."""
         self.slaf_array = slaf_array
@@ -347,6 +349,7 @@ def __init__(
         self.n_scanners = n_scanners
         self.prefetch_batch_size = prefetch_batch_size
         self.parallelize_fragment_reads = parallelize_fragment_reads
+        self.expression_preprocessor = expression_preprocessor
 
         # Validate MoS parameters
         if self.use_mixture_of_scanners:
@@ -1023,13 +1026,17 @@ def load_prefetch_batch(self) -> PrefetchBatch:
 
                     shuffle_time = time.time() - shuffle_start
                     window_start = time.time()
-                    window_params = {
+                    window_params: dict[str, Any] = {
                         "n_expression_bins": self.n_expression_bins,
                         "use_binned_expressions": self.use_binned_expressions,
                     }
                     window_params.update(
                         self.window_kwargs
                     )  # Add any additional kwargs
+                    if self.expression_preprocessor is not None:
+                        window_params["expression_preprocessor"] = (
+                            self.expression_preprocessor
+                        )
 
                     tokenizer = self.tokenizer
                     if tokenizer is None:
@@ -1539,6 +1546,7 @@ def __init__(
         prefetch_batch_size: int = 4194304,  # Add prefetch_batch_size parameter for MoS
         parallelize_fragment_reads: bool = False,  # Parallelize fragment reads in MoS (cloud optimization)
         prefetcher_ready_timeout: float = 10.0,  # Timeout for waiting for prefetcher to be ready
+        expression_preprocessor: ExpressionPreprocessor | None = None,
     ):
         super().__init__()
         self.slaf_array = slaf_array
@@ -1595,6 +1603,7 @@ def __init__(
             n_scanners=n_scanners,  # Pass MoS parameters
             prefetch_batch_size=prefetch_batch_size,  # Pass MoS parameters
             parallelize_fragment_reads=parallelize_fragment_reads,  # Pass parallelize_fragment_reads
+            expression_preprocessor=expression_preprocessor,
         )
         self.prefetcher = AsyncPrefetcher(
             batch_processor=self.batch_processor,

slaf/ml/distributed.py

@@ -20,6 +20,7 @@
 from slaf.distributed.dataloader import DecompressingQueueWrapper, DistributedDataLoader
 
 # Import SLAF-specific components (for type hints and adapters)
+from slaf.ml.expression_preprocessor import ExpressionPreprocessor
 from slaf.ml.tokenizers import GeneformerTokenizer, ScGPTTokenizer
 
 # Configure Modal image for SLAF workers.
@@ -232,6 +233,7 @@ def __init__(
         seed: int = 42,
         queue_name: str | None = None,
         modal_queue_environment: str | None = None,
+        expression_preprocessor: ExpressionPreprocessor | None = None,
         **window_kwargs: Any,
     ):
         """
@@ -271,6 +273,8 @@ def __init__(
             queue_name: Name of Modal Queue (auto-generated if None)
             modal_queue_environment: Optional Modal Environment name for Queue/Dict
                 (match consumer’s ``modal.Queue.from_name(..., environment_name=…)``, e.g. ``"main"`` in some apps).
+            expression_preprocessor: Optional COO value preprocessing before rank/bin (same as
+                ``SLAFDataLoader``). Merged into ``window_kwargs`` for workers.
             **window_kwargs: Additional window function parameters
         """
         self.slaf_array = slaf_array
@@ -303,6 +307,8 @@ def __init__(
         window_kwargs = dict(window_kwargs)
         window_kwargs.setdefault("n_expression_bins", n_expression_bins)
         window_kwargs.setdefault("use_binned_expressions", True)
+        if expression_preprocessor is not None:
+            window_kwargs["expression_preprocessor"] = expression_preprocessor
 
         tokenizer_factory_kwargs: dict[str, Any] | None = None
         tokenizer_cls: type[GeneformerTokenizer] | type[ScGPTTokenizer] | None = None

slaf/ml/expression_preprocessor.py

@@ -0,0 +1,52 @@
+"""Optional COO value transforms before window rank / bin (training pipeline)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import polars as pl
+
+from slaf.core.tabular_schema import DataSchema
+
+
+@dataclass(frozen=True)
+class ExpressionPreprocessor:
+    """Per-cell and element-wise transforms on the COO ``value`` column."""
+
+    normalize_total_target: float | None = None
+    """If set, scale each cell so its values sum to this target (Scanpy-style)."""
+
+    log1p: bool = False
+    """If True, apply ``log1p`` element-wise after optional normalization."""
+
+
+def apply_expression_preprocessor(
+    fragment_df: pl.DataFrame,
+    schema: DataSchema,
+    preprocessor: ExpressionPreprocessor | None,
+) -> pl.DataFrame:
+    """Update ``schema.value_key`` when ``preprocessor`` applies ops; no-op if ``None``."""
+    if preprocessor is None:
+        return fragment_df
+    if not isinstance(preprocessor, ExpressionPreprocessor):
+        raise TypeError(
+            "expression_preprocessor must be an ExpressionPreprocessor instance or None, "
+            f"got {type(preprocessor)!r}"
+        )
+    if preprocessor.normalize_total_target is None and not preprocessor.log1p:
+        return fragment_df
+
+    vk = schema.value_key
+    out = fragment_df
+    if preprocessor.normalize_total_target is not None:
+        t = float(preprocessor.normalize_total_target)
+        cell_sum = pl.col(vk).sum().over(schema.group_key)
+        out = out.with_columns(
+            pl.when(cell_sum > 0)
+            .then(pl.col(vk) * (t / cell_sum))
+            .otherwise(0.0)
+            .alias(vk)
+        )
+    if preprocessor.log1p:
+        out = out.with_columns(pl.col(vk).log1p().alias(vk))
+    return out

tests/test_aggregators.py

@@ -4,6 +4,7 @@
 This module tests the window function implementations for different tokenization strategies.
 """
 
+import numpy as np
 import polars as pl
 import pytest
 
@@ -13,6 +14,7 @@
     ScGPTWindow,
     Window,
 )
+from slaf.ml.expression_preprocessor import ExpressionPreprocessor
 
 
 class TestWindow:
@@ -185,6 +187,36 @@ def test_apply_single_cell(self):
         # Should contain the actual expression values
         assert 3.0 in expr_seq_raw or 5.0 in expr_seq_raw
 
+    def test_expression_preprocessor_normalize_log1p_raw(self):
+        """normalize_total + log1p before rank/raw aggregate matches numpy reference."""
+        result = self.window.apply(
+            self.test_data,
+            self.schema,
+            2,
+            use_binned_expressions=False,
+            expression_preprocessor=ExpressionPreprocessor(
+                normalize_total_target=10.0,
+                log1p=True,
+            ),
+        )
+        cell_0 = result.filter(pl.col("cell_integer_id") == 0)
+        gene_seq = cell_0["gene_sequence"][0]
+        expr_seq = cell_0["expr_sequence"][0]
+        assert list(gene_seq) == [10, 20]
+        raw = np.array([5.0, 3.0, 1.0], dtype=np.float64)
+        scaled = raw * (10.0 / raw.sum())
+        expected = np.log1p(scaled[:2])
+        np.testing.assert_allclose(np.array(expr_seq), expected, rtol=1e-5)
+
+    def test_expression_preprocessor_invalid_type(self):
+        with pytest.raises(TypeError, match="ExpressionPreprocessor"):
+            self.window.apply(
+                self.test_data,
+                self.schema,
+                2,
+                expression_preprocessor="not_a_preprocess",  # type: ignore[arg-type]
+            )
+
 
 class TestGeneformerWindow:
     """Test GeneformerWindow implementation"""