diffbio 0.1.0__py3-none-any.whl

diffbio/pipelines/perturbation.py

@@ -0,0 +1,407 @@
+"""End-to-end perturbation experiment data pipeline.
+
+Orchestrates the full data setup workflow for single-cell perturbation
+experiments: loading, QC filtering, train/val/test splitting, batch sampling,
+and control cell mapping.
+
+This pipeline produces ready-to-train data sources with paired
+(perturbed, control) cell elements. It is a structural (non-differentiable)
+setup pipeline — the differentiable training loop operates on its outputs.
+
+References:
+    - cell-load PerturbationDataModule
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+from flax import nnx
+
+from datarax.core.config import StructuralConfig
+from datarax.core.data_source import DataSourceModule
+
+from diffbio.operators.singlecell.knockdown_filter import (
+    KnockdownFilterConfig,
+    OnTargetKnockdownFilter,
+)
+from diffbio.samplers.perturbation_sampler import (
+    PerturbationBatchSampler,
+    PerturbationSamplerConfig,
+)
+from diffbio.sources.perturbation.concat_source import PerturbationConcatSource
+from diffbio.sources.perturbation.control_mapping import (
+    BatchControlMapping,
+    ControlMappingConfig,
+    RandomControlMapping,
+)
+from diffbio.sources.perturbation.perturbation_source import (
+    PerturbationAnnDataSource,
+    PerturbationSourceConfig,
+)
+from diffbio.splitters.perturbation import (
+    FewShotSplitter,
+    FewShotSplitterConfig,
+    ZeroShotSplitter,
+    ZeroShotSplitterConfig,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class _PipelineSourceConfig:
+    """Perturbation source selection and output-space configuration."""
+
+    pert_col: str = "perturbation"
+    cell_type_col: str = "cell_type"
+    batch_col: str = "batch"
+    control_pert: str = "non-targeting"
+    output_space: str = "all"
+    embedding_key: str | None = None
+    hvg_col: str | None = None
+
+
+@dataclass(frozen=True)
+class _PipelineBatchingAndSplitConfig:
+    """Control mapping, batching, and split selection configuration."""
+
+    mapping_strategy: str = "random"
+    n_basal_samples: int = 1
+    sentence_size: int = 512
+    sentences_per_batch: int = 1
+    split_mode: str = "random"
+    held_out_cell_types: tuple[str, ...] = ()
+    held_out_perturbations: tuple[str, ...] = ()
+
+
+@dataclass(frozen=True)
+class _PipelineFilterConfig:
+    """Knockdown-filter configuration."""
+
+    enable_knockdown_filter: bool = False
+    residual_expression: float = 0.30
+    cell_residual_expression: float = 0.50
+    min_cells: int = 30
+    var_gene_col: str | None = None
+
+
+@dataclass(frozen=True)
+class _PipelineFractionConfig:
+    """Split fraction and seed configuration."""
+
+    train_frac: float = 0.8
+    valid_frac: float = 0.1
+    seed: int = 42
+
+
+@dataclass(frozen=True)
+class PerturbationPipelineConfig(
+    _PipelineSourceConfig,
+    _PipelineBatchingAndSplitConfig,
+    _PipelineFilterConfig,
+    _PipelineFractionConfig,
+    StructuralConfig,
+):
+    """Configuration for PerturbationPipeline."""
+
+    def __post_init__(self) -> None:
+        """Validate pipeline configuration."""
+        super().__post_init__()
+
+        if self.output_space == "embedding" and self.embedding_key is None:
+            raise ValueError("embedding_key is required when output_space='embedding'")
+
+        if self.mapping_strategy not in {"batch", "random"}:
+            raise ValueError(
+                "mapping_strategy must be either 'batch' or 'random', "
+                f"got {self.mapping_strategy!r}"
+            )
+
+        if self.split_mode not in {"random", "zeroshot", "fewshot"}:
+            raise ValueError(
+                "split_mode must be one of 'random', 'zeroshot', or 'fewshot', "
+                f"got {self.split_mode!r}"
+            )
+
+        if self.n_basal_samples <= 0:
+            raise ValueError("n_basal_samples must be positive")
+        if self.sentence_size <= 0:
+            raise ValueError("sentence_size must be positive")
+        if self.sentences_per_batch <= 0:
+            raise ValueError("sentences_per_batch must be positive")
+        if self.min_cells <= 0:
+            raise ValueError("min_cells must be positive")
+
+        if not 0.0 <= self.train_frac <= 1.0:
+            raise ValueError("train_frac must be between 0.0 and 1.0")
+        if not 0.0 <= self.valid_frac <= 1.0:
+            raise ValueError("valid_frac must be between 0.0 and 1.0")
+        if self.train_frac + self.valid_frac > 1.0:
+            raise ValueError("train_frac + valid_frac must be <= 1.0")
+
+        if not 0.0 < self.residual_expression <= 1.0:
+            raise ValueError("residual_expression must be in (0.0, 1.0]")
+        if not 0.0 < self.cell_residual_expression <= 1.0:
+            raise ValueError("cell_residual_expression must be in (0.0, 1.0]")
+
+        if self.split_mode == "random":
+            if self.held_out_cell_types or self.held_out_perturbations:
+                raise ValueError(
+                    "held_out_cell_types and held_out_perturbations are only valid for "
+                    "zeroshot/fewshot split modes"
+                )
+        elif self.split_mode == "zeroshot":
+            if not self.held_out_cell_types:
+                raise ValueError("held_out_cell_types is required when split_mode='zeroshot'")
+            if self.held_out_perturbations:
+                raise ValueError("held_out_perturbations is not used when split_mode='zeroshot'")
+        elif not self.held_out_perturbations:
+            raise ValueError("held_out_perturbations is required when split_mode='fewshot'")
+        elif self.held_out_cell_types:
+            raise ValueError("held_out_cell_types is not used when split_mode='fewshot'")
+
+
+class PerturbationPipeline:
+    """End-to-end data setup pipeline for perturbation experiments.
+
+    Orchestrates the full workflow from raw H5AD files to ready-to-train
+    data sources with paired (perturbed, control) cell output.
+
+    Workflow:
+        1. Load H5AD file(s) into PerturbationAnnDataSource(s)
+        2. (Optional) Apply on-target knockdown QC filter
+        3. Split into train/val/test via zero-shot, few-shot, or random
+        4. Build control cell mapping for each split
+        5. Create batch sampler for training
+
+    Example::
+
+        config = PerturbationPipelineConfig(
+            split_mode="zeroshot",
+            held_out_cell_types=("TypeA",),
+            mapping_strategy="batch",
+        )
+        pipeline = PerturbationPipeline(config)
+        result = pipeline.setup(["path/to/data.h5ad"])
+
+        for batch_indices in result.train_sampler:
+            elements = [result.train_source[i] for i in batch_indices]
+            paired = result.get_paired_batch(batch_indices, split="train")
+    """
+
+    def __init__(self, config: PerturbationPipelineConfig) -> None:
+        self._config = config
+
+    def setup(self, file_paths: list[str | Path]) -> PerturbationPipelineResult:
+        """Execute the full data setup workflow.
+
+        Args:
+            file_paths: Paths to H5AD files to load.
+
+        Returns:
+            PerturbationPipelineResult with sources, splits, samplers,
+            and control mappings.
+        """
+        config = self._config
+
+        # 1. Load sources
+        logger.info("Loading %d H5AD file(s)...", len(file_paths))
+        sources = []
+        for path in file_paths:
+            src_config = PerturbationSourceConfig(
+                file_path=str(path),
+                pert_col=config.pert_col,
+                cell_type_col=config.cell_type_col,
+                batch_col=config.batch_col,
+                control_pert=config.control_pert,
+                output_space=config.output_space,
+                embedding_key=config.embedding_key,
+                hvg_col=config.hvg_col,
+                seed=config.seed,
+            )
+            sources.append(PerturbationAnnDataSource(src_config))
+
+        source: DataSourceModule
+        if len(sources) == 1:
+            source = sources[0]
+        else:
+            source = PerturbationConcatSource(sources=sources)
+
+        # 2. Knockdown filter (optional)
+        filter_mask: np.ndarray | None = None
+        if config.enable_knockdown_filter:
+            logger.info("Applying knockdown QC filter...")
+            filt = OnTargetKnockdownFilter(
+                KnockdownFilterConfig(
+                    pert_col=config.pert_col,
+                    control_pert=config.control_pert,
+                    residual_expression=config.residual_expression,
+                    cell_residual_expression=config.cell_residual_expression,
+                    min_cells=config.min_cells,
+                    var_gene_col=config.var_gene_col,
+                )
+            )
+            filter_mask = filt.process(source)
+            n_kept = int(filter_mask.sum())
+            logger.info(
+                "Knockdown filter: %d / %d cells pass",
+                n_kept,
+                len(filter_mask),
+            )
+
+        # 3. Split
+        logger.info("Splitting data (mode=%s)...", config.split_mode)
+        if config.split_mode == "zeroshot":
+            splitter = ZeroShotSplitter(
+                ZeroShotSplitterConfig(
+                    held_out_cell_types=config.held_out_cell_types,
+                    pert_col=config.pert_col,
+                    cell_type_col=config.cell_type_col,
+                    train_frac=config.train_frac,
+                    valid_frac=config.valid_frac,
+                    test_frac=1.0 - config.train_frac - config.valid_frac,
+                    seed=config.seed,
+                ),
+                rngs=nnx.Rngs(config.seed),
+            )
+        elif config.split_mode == "fewshot":
+            splitter = FewShotSplitter(
+                FewShotSplitterConfig(
+                    held_out_perturbations=config.held_out_perturbations,
+                    pert_col=config.pert_col,
+                    cell_type_col=config.cell_type_col,
+                    control_pert=config.control_pert,
+                    train_frac=config.train_frac,
+                    valid_frac=config.valid_frac,
+                    test_frac=1.0 - config.train_frac - config.valid_frac,
+                    seed=config.seed,
+                ),
+                rngs=nnx.Rngs(config.seed),
+            )
+        else:
+            from diffbio.splitters.random import (  # noqa: PLC0415
+                RandomSplitter,
+                RandomSplitterConfig,
+            )
+
+            splitter = RandomSplitter(
+                RandomSplitterConfig(
+                    train_frac=config.train_frac,
+                    valid_frac=config.valid_frac,
+                    test_frac=1.0 - config.train_frac - config.valid_frac,
+                    seed=config.seed,
+                ),
+                rngs=nnx.Rngs(config.seed),
+            )
+
+        split_result = splitter.split(source)
+
+        # Apply filter mask to split indices if knockdown filter was used
+        if filter_mask is not None:
+            passing = set(np.where(filter_mask)[0])
+            train_idx = np.array([i for i in split_result.train_indices if int(i) in passing])
+            valid_idx = np.array([i for i in split_result.valid_indices if int(i) in passing])
+            test_idx = np.array([i for i in split_result.test_indices if int(i) in passing])
+        else:
+            train_idx = np.array(split_result.train_indices)
+            valid_idx = np.array(split_result.valid_indices)
+            test_idx = np.array(split_result.test_indices)
+
+        logger.info(
+            "Split sizes: train=%d, val=%d, test=%d",
+            len(train_idx),
+            len(valid_idx),
+            len(test_idx),
+        )
+
+        # 4. Control mapping
+        logger.info(
+            "Building control mapping (strategy=%s)...",
+            config.mapping_strategy,
+        )
+        mapping_config = ControlMappingConfig(
+            strategy=config.mapping_strategy,
+            n_basal_samples=config.n_basal_samples,
+            seed=config.seed,
+        )
+        if config.mapping_strategy == "batch":
+            mapper = BatchControlMapping(mapping_config)
+        else:
+            mapper = RandomControlMapping(mapping_config)
+
+        control_mapping = mapper.build_mapping(source)
+
+        # 5. Train sampler
+        group_codes = source.get_group_codes()
+        train_group_codes = group_codes[train_idx]
+        sampler_config = PerturbationSamplerConfig(
+            sentence_size=config.sentence_size,
+            sentences_per_batch=config.sentences_per_batch,
+            seed=config.seed,
+        )
+        train_sampler = PerturbationBatchSampler(sampler_config, train_group_codes)
+
+        return PerturbationPipelineResult(
+            source=source,
+            train_indices=train_idx,
+            valid_indices=valid_idx,
+            test_indices=test_idx,
+            control_mapping=control_mapping,
+            train_sampler=train_sampler,
+            filter_mask=filter_mask,
+        )
+
+
+class PerturbationPipelineResult:
+    """Result of PerturbationPipeline.setup().
+
+    Holds all artifacts needed for training: the data source, split indices,
+    control mapping, and train sampler.
+
+    Attributes:
+        source: The underlying data source (single or concatenated).
+        train_indices: Cell indices for training.
+        valid_indices: Cell indices for validation.
+        test_indices: Cell indices for testing.
+        control_mapping: Array mapping perturbed cell index to control indices.
+        train_sampler: Batch sampler for training iteration.
+        filter_mask: Boolean QC mask (None if filtering was disabled).
+    """
+
+    def __init__(
+        self,
+        source: Any,
+        train_indices: np.ndarray,
+        valid_indices: np.ndarray,
+        test_indices: np.ndarray,
+        control_mapping: np.ndarray,
+        train_sampler: PerturbationBatchSampler,
+        filter_mask: np.ndarray | None = None,
+    ) -> None:
+        self.source = source
+        self.train_indices = train_indices
+        self.valid_indices = valid_indices
+        self.test_indices = test_indices
+        self.control_mapping = control_mapping
+        self.train_sampler = train_sampler
+        self.filter_mask = filter_mask
+
+    def get_element(self, global_idx: int) -> dict[str, Any]:
+        """Get a single cell element by global index.
+
+        Args:
+            global_idx: Index into the full source.
+
+        Returns:
+            Per-cell dictionary with counts and metadata.
+        """
+        return self.source[global_idx]
+
+    def get_var_dims(self) -> dict[str, int]:
+        """Return dimensionality info from the source."""
+        return self.source.get_var_dims()

diffbio/pipelines/preprocessing.py

@@ -0,0 +1,267 @@
+"""End-to-end differentiable preprocessing pipeline.
+
+This module provides a complete read preprocessing pipeline that composes:
+1. Quality filtering - Filter low-quality bases
+2. Adapter removal - Soft trim adapter sequences
+3. Duplicate weighting - Assign probabilistic weights based on uniqueness
+4. Error correction - Neural network-based base correction
+
+The pipeline is fully differentiable, enabling gradient-based optimization
+of all preprocessing components jointly.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+from datarax.core.config import OperatorConfig
+from datarax.core.operator import OperatorModule
+from flax import nnx
+from jaxtyping import Array
+
+from diffbio.operators.preprocessing import (
+    AdapterRemovalConfig,
+    DifferentiableDuplicateWeighting,
+    DuplicateWeightingConfig,
+    ErrorCorrectionConfig,
+    SoftAdapterRemoval,
+    SoftErrorCorrection,
+)
+from diffbio.operators.quality_filter import (
+    DifferentiableQualityFilter,
+    QualityFilterConfig,
+)
+from diffbio.utils.quality import apply_quality_filter
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class PreprocessingPipelineConfig(OperatorConfig):
+    # pylint: disable=too-many-instance-attributes
+    """Configuration for the preprocessing pipeline.
+
+    Attributes:
+        read_length: Expected read length for initialization.
+        adapter_sequence: Adapter sequence to remove (Illumina universal default).
+        quality_threshold: Initial quality score threshold for filtering.
+        adapter_match_threshold: Threshold for adapter matching.
+        adapter_temperature: Temperature for soft adapter trimming.
+        duplicate_similarity_threshold: Similarity threshold for duplicate detection.
+        error_correction_window: Window size for error correction.
+        error_correction_hidden_dim: Hidden dimension for error correction network.
+        enable_adapter_removal: Whether to enable adapter removal step.
+        enable_duplicate_weighting: Whether to enable duplicate weighting step.
+        enable_error_correction: Whether to enable error correction step.
+    """
+
+    read_length: int = 150
+    adapter_sequence: str = "AGATCGGAAGAG"
+    quality_threshold: float = 20.0
+    adapter_match_threshold: float = 0.8
+    adapter_temperature: float = 1.0
+    duplicate_similarity_threshold: float = 0.95
+    error_correction_window: int = 11
+    error_correction_hidden_dim: int = 64
+    enable_adapter_removal: bool = True
+    enable_duplicate_weighting: bool = True
+    enable_error_correction: bool = True
+
+
+class PreprocessingPipeline(OperatorModule):
+    """End-to-end differentiable preprocessing pipeline.
+
+    This pipeline processes sequencing reads through multiple preprocessing steps:
+
+    Input data structure:
+        - reads: Float[Array, "num_reads read_length 4"] - One-hot encoded reads
+        - quality: Float[Array, "num_reads read_length"] - Base quality scores
+
+    Output data structure (adds):
+        - preprocessed_reads: Float[Array, "num_reads read_length 4"] - Processed reads
+        - preprocessed_quality: Float[Array, "num_reads read_length"] - Processed quality
+        - read_weights: Float[Array, "num_reads"] - Read uniqueness weights
+
+    The pipeline is fully differentiable, supporting gradient-based training
+    to optimize all preprocessing components jointly.
+
+    Example:
+        ```python
+        config = PreprocessingPipelineConfig(read_length=150)
+        pipeline = PreprocessingPipeline(config, rngs=nnx.Rngs(42))
+        result, state, meta = pipeline.apply(data, {}, None)
+        processed = result["preprocessed_reads"]
+        ```
+    """
+
+    def __init__(
+        self,
+        config: PreprocessingPipelineConfig,
+        *,
+        rngs: nnx.Rngs,
+        name: str | None = None,
+    ):
+        """Initialize the preprocessing pipeline.
+
+        Args:
+            config: Pipeline configuration.
+            rngs: Random number generators for parameter initialization.
+            name: Optional name for the pipeline.
+        """
+        super().__init__(config, rngs=rngs, name=name)
+
+        # 1. Quality filter (always enabled)
+        self.quality_filter = DifferentiableQualityFilter(
+            QualityFilterConfig(initial_threshold=config.quality_threshold),
+            rngs=rngs,
+        )
+
+        # 2. Adapter removal (optional)
+        self.adapter_removal = (
+            SoftAdapterRemoval(
+                AdapterRemovalConfig(
+                    adapter_sequence=config.adapter_sequence,
+                    match_threshold=config.adapter_match_threshold,
+                    temperature=config.adapter_temperature,
+                ),
+                rngs=rngs,
+            )
+            if config.enable_adapter_removal
+            else None
+        )
+
+        # 3. Duplicate weighting (optional)
+        self.duplicate_weighting = (
+            DifferentiableDuplicateWeighting(
+                DuplicateWeightingConfig(
+                    similarity_threshold=config.duplicate_similarity_threshold,
+                ),
+                rngs=rngs,
+            )
+            if config.enable_duplicate_weighting
+            else None
+        )
+
+        # 4. Error correction (optional)
+        self.error_correction = (
+            SoftErrorCorrection(
+                ErrorCorrectionConfig(
+                    window_size=config.error_correction_window,
+                    hidden_dim=config.error_correction_hidden_dim,
+                ),
+                rngs=rngs,
+            )
+            if config.enable_error_correction
+            else None
+        )
+
+    def apply(
+        self,
+        data: dict[str, Array],
+        state: dict[str, Any],
+        metadata: dict[str, Any] | None,
+        random_params: Any = None,
+        stats: dict[str, Any] | None = None,
+    ) -> tuple[dict[str, Array], dict[str, Any], dict[str, Any] | None]:
+        """Apply the full preprocessing pipeline to reads.
+
+        Args:
+            data: Input data containing:
+                - reads: Float[Array, "num_reads read_length 4"]
+                - quality: Float[Array, "num_reads read_length"]
+            state: Element state (passed through).
+            metadata: Element metadata (passed through).
+            random_params: Not used (deterministic pipeline).
+            stats: Optional statistics dict.
+
+        Returns:
+            Tuple of (output_data, state, metadata) where output_data contains
+            all input keys plus preprocessed outputs.
+        """
+        reads = data["reads"]
+        quality = data["quality"]
+        num_reads = reads.shape[0]
+
+        # Initialize read weights to 1.0 (all reads equally weighted)
+        read_weights = jnp.ones((num_reads,))
+
+        # Step 1: Quality filtering (per-base)
+        filtered_reads, filtered_quality = apply_quality_filter(self.quality_filter, reads, quality)
+
+        # Step 2: Adapter removal (optional) - apply per-read using vmap
+        if self.adapter_removal is not None:
+
+            def apply_adapter_removal(read, quality):
+                """Apply adapter removal to a single read."""
+                adapter_data = {"sequence": read, "quality_scores": quality}
+                adapter_result, _, _ = self.adapter_removal.apply(adapter_data, {}, None)
+                return adapter_result["sequence"], adapter_result["quality_scores"]
+
+            filtered_reads, filtered_quality = jax.vmap(apply_adapter_removal)(
+                filtered_reads, filtered_quality
+            )
+
+        # Step 3: Duplicate weighting (optional) - operates on full batch for cross-read comparison
+        # Use apply_batch which returns weights for all reads (not just first)
+        if self.duplicate_weighting is not None:
+            raw_weights, _ = self.duplicate_weighting.apply_batch(filtered_reads, filtered_quality)
+            # Normalize weights to [0, 1] range for use as probabilities
+            read_weights = raw_weights / jnp.max(raw_weights)
+
+        # Step 4: Error correction (optional) - apply per-read using vmap
+        if self.error_correction is not None:
+
+            def apply_error_correction(read, quality):
+                """Apply error correction to a single read."""
+                ec_data = {"sequence": read, "quality_scores": quality}
+                ec_result, _, _ = self.error_correction.apply(ec_data, {}, None)
+                return ec_result["sequence"]
+
+            filtered_reads = jax.vmap(apply_error_correction)(filtered_reads, filtered_quality)
+
+        # Build output preserving input keys
+        output_data = {
+            **data,
+            "preprocessed_reads": filtered_reads,
+            "preprocessed_quality": filtered_quality,
+            "read_weights": read_weights,
+        }
+
+        return output_data, state, metadata
+
+
+def create_preprocessing_pipeline(
+    read_length: int = 150,
+    quality_threshold: float = 20.0,
+    adapter_sequence: str = "AGATCGGAAGAG",
+    enable_adapter_removal: bool = True,
+    enable_duplicate_weighting: bool = True,
+    enable_error_correction: bool = True,
+    seed: int = 42,
+) -> PreprocessingPipeline:
+    """Factory function to create a preprocessing pipeline.
+
+    Args:
+        read_length: Expected read length.
+        quality_threshold: Quality score threshold.
+        adapter_sequence: Adapter sequence to remove.
+        enable_adapter_removal: Whether to enable adapter removal.
+        enable_duplicate_weighting: Whether to enable duplicate weighting.
+        enable_error_correction: Whether to enable error correction.
+        seed: Random seed.
+
+    Returns:
+        Configured PreprocessingPipeline instance.
+    """
+    config = PreprocessingPipelineConfig(
+        read_length=read_length,
+        quality_threshold=quality_threshold,
+        adapter_sequence=adapter_sequence,
+        enable_adapter_removal=enable_adapter_removal,
+        enable_duplicate_weighting=enable_duplicate_weighting,
+        enable_error_correction=enable_error_correction,
+    )
+    rngs = nnx.Rngs(seed)
+    return PreprocessingPipeline(config, rngs=rngs)