diffbio 0.1.0__py3-none-any.whl

diffbio/operators/rnaseq/motif_discovery.py

@@ -0,0 +1,251 @@
+"""Differentiable motif discovery (MEME-style).
+
+This module implements a differentiable version of motif discovery with
+PWM (Position Weight Matrix) learning for end-to-end gradient flow.
+
+Inherits from TemperatureOperator to get:
+
+- _temperature property for temperature-controlled smoothing
+- soft_max() for logsumexp-based smooth maximum
+- soft_argmax() for soft position selection
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import flax.nnx as nnx
+import jax
+import jax.numpy as jnp
+from datarax.core.config import OperatorConfig
+
+from diffbio.core import soft_ops
+from diffbio.core.base_operators import TemperatureOperator
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class MotifDiscoveryConfig(OperatorConfig):
+    """Configuration for differentiable motif discovery.
+
+    Attributes:
+        motif_width: Width of the motif (number of positions).
+        num_motifs: Number of motifs to discover.
+        alphabet_size: Size of the sequence alphabet (4 for DNA).
+        temperature: Temperature for soft operations.
+        background_prior: Prior probability for background model.
+        stream_name: Name of the data stream to process.
+    """
+
+    motif_width: int = 12
+    num_motifs: int = 1
+    alphabet_size: int = 4
+    temperature: float = 1.0
+    learnable_temperature: bool = True
+    background_prior: float = 0.25  # Uniform for DNA
+
+
+class DifferentiableMotifDiscovery(TemperatureOperator):
+    """Differentiable motif discovery with PWM learning.
+
+    This operator implements a simplified differentiable version of MEME-style
+    motif discovery. It learns Position Weight Matrices (PWMs) that represent
+    sequence motifs and scans sequences to find motif occurrences.
+
+    The motif score at position i is computed as:
+        score(i) = sum_j PWM[j, seq[i+j]]
+
+    For one-hot encoded sequences, this is equivalent to:
+        score(i) = sum_j sum_k seq[i+j, k] * log(PWM[j, k])
+
+    Example:
+        ```python
+        config = MotifDiscoveryConfig(
+            motif_width=12,
+            num_motifs=3,
+        )
+        motif_op = DifferentiableMotifDiscovery(config, rngs=rngs)
+
+        data = {"sequence": one_hot_sequence}  # (length, alphabet_size)
+        result, state, metadata = motif_op.apply(data, {}, None)
+        motif_scores = result["motif_scores"]  # (num_positions, num_motifs)
+        pwm = result["pwm"]  # (num_motifs, motif_width, alphabet_size)
+        ```
+    """
+
+    def __init__(self, config: MotifDiscoveryConfig, *, rngs: nnx.Rngs | None = None):
+        """Initialize the motif discovery operator.
+
+        Args:
+            config: Configuration for the operator.
+            rngs: Random number generators for initialization.
+        """
+        super().__init__(config, rngs=rngs)
+        self.config = config
+
+        if rngs is None:
+            rngs = nnx.Rngs(0)
+
+        key = rngs.params() if hasattr(rngs, "params") else jax.random.key(0)
+
+        # Initialize PWM logits (before softmax normalization)
+        # Shape: (num_motifs, motif_width, alphabet_size)
+        # Initialize near uniform with small random noise
+        pwm_init = (
+            jax.random.normal(
+                key,
+                (config.num_motifs, config.motif_width, config.alphabet_size),
+            )
+            * 0.1
+        )
+
+        self.pwm_logits = nnx.Param(pwm_init)
+
+        # Temperature is managed by TemperatureOperator via self._temperature
+
+    def _get_pwm(self) -> jax.Array:
+        """Get normalized PWM from logits.
+
+        Returns:
+            PWM of shape (num_motifs, motif_width, alphabet_size) with
+            probabilities summing to 1 over the alphabet dimension.
+        """
+        temperature = jnp.abs(self._temperature) + 1e-6
+        return jax.nn.softmax(self.pwm_logits[...] / temperature, axis=-1)
+
+    def _scan_single_motif(self, sequence: jax.Array, pwm: jax.Array) -> jax.Array:
+        """Scan a sequence with a single PWM using convolution.
+
+        Args:
+            sequence: One-hot encoded sequence of shape (length, alphabet_size).
+            pwm: PWM of shape (motif_width, alphabet_size).
+
+        Returns:
+            Motif scores at each valid position, shape (num_positions,).
+        """
+        motif_width = pwm.shape[0]
+        seq_length = sequence.shape[0]
+        num_positions = seq_length - motif_width + 1
+
+        # Use log-odds scoring
+        # log_pwm = log(PWM) - log(background)
+        background = self.config.background_prior
+        log_pwm = jnp.log(pwm + 1e-8) - jnp.log(background)
+
+        # Compute score at each position using sliding window
+        def score_at_position(start_idx):
+            window = jax.lax.dynamic_slice(
+                sequence, (start_idx, 0), (motif_width, self.config.alphabet_size)
+            )
+            # Score = sum of log-odds weighted by sequence
+            score = jnp.sum(window * log_pwm)
+            return score
+
+        positions = jnp.arange(num_positions)
+        scores = jax.vmap(score_at_position)(positions)
+
+        return scores
+
+    def _scan_sequence(self, sequence: jax.Array) -> jax.Array:
+        """Scan a sequence with all motifs.
+
+        Args:
+            sequence: One-hot encoded sequence of shape (length, alphabet_size).
+
+        Returns:
+            Motif scores of shape (num_positions, num_motifs).
+        """
+        pwm = self._get_pwm()
+
+        # Scan with each motif
+        def scan_with_motif(pwm_single):
+            return self._scan_single_motif(sequence, pwm_single)
+
+        # Shape: (num_motifs, num_positions)
+        all_scores = jax.vmap(scan_with_motif)(pwm)
+
+        # Transpose to (num_positions, num_motifs)
+        return all_scores.T
+
+    def _find_motif_positions(self, scores: jax.Array, threshold: float = 0.0) -> jax.Array:
+        """Find soft motif positions based on scores.
+
+        Args:
+            scores: Motif scores of shape (num_positions, num_motifs).
+            threshold: Score threshold for calling a motif hit.
+
+        Returns:
+            Soft position indicators of shape (num_positions, num_motifs).
+        """
+        temperature = jnp.abs(self._temperature) + 1e-6
+        return soft_ops.greater(scores, threshold, softness=temperature)
+
+    def _apply_single(self, sequence: jax.Array) -> dict:
+        """Apply motif discovery to a single sequence.
+
+        Args:
+            sequence: One-hot encoded sequence of shape (length, alphabet_size).
+
+        Returns:
+            Dictionary with motif scores, positions, and PWM.
+        """
+        # Get current PWM
+        pwm = self._get_pwm()
+
+        # Scan sequence
+        motif_scores = self._scan_sequence(sequence)
+
+        # Find soft motif positions
+        motif_positions = self._find_motif_positions(motif_scores)
+
+        return {
+            "motif_scores": motif_scores,
+            "motif_positions": motif_positions,
+            "pwm": pwm,
+        }
+
+    def apply(
+        self,
+        data: dict[str, Any],
+        state: dict[str, Any],
+        metadata: dict | None,
+        random_params: dict | None = None,
+        stats: dict | None = None,
+    ) -> tuple[dict, dict, dict | None]:
+        """Apply motif discovery to sequence data.
+
+        Args:
+            data: Dictionary containing:
+                - 'sequence': One-hot encoded sequence(s) of shape
+                  (length, alphabet_size) or (batch, length, alphabet_size)
+            state: Operator state dictionary.
+            metadata: Optional metadata dictionary.
+            random_params: Optional random parameters (unused).
+            stats: Optional statistics dictionary (unused).
+
+        Returns:
+            Tuple of (output_data, state, metadata) where output_data contains:
+
+                - 'sequence': Original sequence data
+                - 'motif_scores': Log-odds scores at each position
+                - 'motif_positions': Soft motif occurrence indicators
+                - 'pwm': Current Position Weight Matrix
+        """
+        del random_params, stats  # Unused
+
+        sequence = data["sequence"]
+
+        # Handle single vs batched input
+        single_input = sequence.ndim == 2
+        if single_input:
+            result = self._apply_single(sequence)
+        else:
+            # Batched input - vmap over batch dimension
+            result = jax.vmap(self._apply_single)(sequence)
+            # PWM is shared, take from first (they're all the same)
+            result["pwm"] = self._get_pwm()
+
+        output_data = {**data, **result}
+
+        return output_data, state, metadata

diffbio/operators/rnaseq/splicing_psi.py

@@ -0,0 +1,216 @@
+"""Differentiable splicing PSI (Percent Spliced In) calculation.
+
+This module implements a differentiable PSI calculation operator for
+alternative splicing analysis with end-to-end gradient flow.
+
+Inherits from TemperatureOperator to get:
+
+- _temperature property for temperature-controlled smoothing
+- soft_max() for logsumexp-based smooth maximum
+- soft_argmax() for soft position selection
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import flax.nnx as nnx
+import jax
+import jax.numpy as jnp
+from datarax.core.config import OperatorConfig
+
+from diffbio.core import soft_ops
+from diffbio.core.base_operators import TemperatureOperator
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class SplicingPSIConfig(OperatorConfig):
+    """Configuration for differentiable PSI calculation.
+
+    Attributes:
+        pseudocount: Pseudocount added for numerical stability and regularization.
+        temperature: Temperature for confidence calculation.
+        min_total_reads: Minimum total reads for reliable PSI estimation.
+        stream_name: Name of the data stream to process.
+    """
+
+    pseudocount: float = 1.0
+    temperature: float = 1.0
+    learnable_temperature: bool = True
+    min_total_reads: int = 10
+
+
+class SplicingPSI(TemperatureOperator):
+    """Differentiable PSI calculation for alternative splicing analysis.
+
+    PSI (Percent Spliced In) quantifies alternative splicing by computing
+    the fraction of transcripts that include a specific exon or splice site.
+
+    The standard PSI formula is:
+        PSI = inclusion_reads / (inclusion_reads + exclusion_reads)
+
+    This operator adds:
+    - Learnable pseudocount for regularization
+    - Confidence estimation based on read coverage
+    - Full differentiability for end-to-end training
+
+    Example:
+        ```python
+        config = SplicingPSIConfig(
+            pseudocount=1.0,
+            min_total_reads=10,
+        )
+        psi_op = SplicingPSI(config, rngs=rngs)
+
+        data = {
+            "inclusion_counts": inclusion_reads,  # Junction reads supporting inclusion
+            "exclusion_counts": exclusion_reads,  # Junction reads supporting exclusion
+        }
+        result, state, metadata = psi_op.apply(data, {}, None)
+        psi_values = result["psi"]
+        confidence = result["psi_confidence"]
+        ```
+    """
+
+    def __init__(self, config: SplicingPSIConfig, *, rngs: nnx.Rngs | None = None):
+        """Initialize the PSI operator.
+
+        Args:
+            config: Configuration for the operator.
+            rngs: Random number generators for initialization.
+        """
+        super().__init__(config, rngs=rngs)
+        self.config = config
+
+        # Learnable pseudocount (must be positive)
+        self.pseudocount = nnx.Param(jnp.array(config.pseudocount))
+
+        # Temperature is managed by TemperatureOperator via self._temperature
+
+    def _compute_psi(
+        self, inclusion: jax.Array, exclusion: jax.Array, pseudocount: jax.Array | float
+    ) -> jax.Array:
+        """Compute PSI with pseudocount for numerical stability.
+
+        Args:
+            inclusion: Inclusion junction read counts.
+            exclusion: Exclusion junction read counts.
+            pseudocount: Pseudocount for regularization.
+
+        Returns:
+            PSI values in [0, 1].
+        """
+        # Add pseudocount to both numerator and denominator terms
+        inc_adj = inclusion + pseudocount
+        exc_adj = exclusion + pseudocount
+
+        # PSI = inclusion / (inclusion + exclusion)
+        psi = inc_adj / (inc_adj + exc_adj)
+
+        return psi
+
+    def _compute_confidence(
+        self, inclusion: jax.Array, exclusion: jax.Array, temperature: jax.Array | float
+    ) -> jax.Array:
+        """Compute confidence in PSI estimate based on read coverage.
+
+        Higher total reads = higher confidence in the PSI estimate.
+
+        Args:
+            inclusion: Inclusion junction read counts.
+            exclusion: Exclusion junction read counts.
+            temperature: Temperature for sigmoid scaling.
+
+        Returns:
+            Confidence values in [0, 1].
+        """
+        total_reads = inclusion + exclusion
+        min_reads = self.config.min_total_reads
+
+        # Sigmoid-based confidence: approaches 1 as reads increase
+        # Centered around min_total_reads
+        confidence = soft_ops.greater(
+            total_reads, min_reads, softness=temperature * min_reads + 1e-6
+        )
+
+        return confidence
+
+    def _compute_delta_psi_variance(
+        self, inclusion: jax.Array, exclusion: jax.Array, psi: jax.Array
+    ) -> jax.Array:
+        """Compute variance of PSI estimate using beta-binomial model.
+
+        This approximates the variance of PSI under a binomial model,
+        which can be used for significance testing.
+
+        Args:
+            inclusion: Inclusion junction read counts.
+            exclusion: Exclusion junction read counts.
+            psi: Computed PSI values.
+
+        Returns:
+            Variance estimates for each PSI value.
+        """
+        total = inclusion + exclusion + 2 * jnp.abs(self.pseudocount[...])
+
+        # Variance of beta distribution: psi * (1 - psi) / (n + 1)
+        variance = (psi * (1 - psi)) / (total + 1)
+
+        return variance
+
+    def apply(
+        self,
+        data: dict[str, Any],
+        state: dict[str, Any],
+        metadata: dict | None,
+        random_params: dict | None = None,
+        stats: dict | None = None,
+    ) -> tuple[dict, dict, dict | None]:
+        """Apply PSI calculation to junction read counts.
+
+        Args:
+            data: Dictionary containing:
+                - 'inclusion_counts': Reads supporting exon inclusion
+                - 'exclusion_counts': Reads supporting exon exclusion
+            state: Operator state dictionary.
+            metadata: Optional metadata dictionary.
+            random_params: Optional random parameters (unused).
+            stats: Optional statistics dictionary (unused).
+
+        Returns:
+            Tuple of (output_data, state, metadata) where output_data contains:
+
+                - 'inclusion_counts': Original inclusion counts
+                - 'exclusion_counts': Original exclusion counts
+                - 'psi': Computed PSI values
+                - 'psi_confidence': Confidence in PSI estimates
+                - 'psi_variance': Variance of PSI estimates
+        """
+        del random_params, stats  # Unused
+
+        inclusion = data["inclusion_counts"]
+        exclusion = data["exclusion_counts"]
+
+        # Get learnable parameters (ensure positive)
+        pseudocount = jnp.abs(self.pseudocount[...]) + 1e-6
+        temperature = jnp.abs(self._temperature) + 1e-6
+
+        # Compute PSI
+        psi = self._compute_psi(inclusion, exclusion, pseudocount)
+
+        # Compute confidence
+        confidence = self._compute_confidence(inclusion, exclusion, temperature)
+
+        # Compute variance for significance testing
+        variance = self._compute_delta_psi_variance(inclusion, exclusion, psi)
+
+        output_data = {
+            **data,
+            "psi": psi,
+            "psi_confidence": confidence,
+            "psi_variance": variance,
+        }
+
+        return output_data, state, metadata

diffbio/operators/singlecell/__init__.py

@@ -0,0 +1,193 @@
+"""Single-cell analysis operators for differentiable scRNA-seq processing.
+
+This module provides differentiable components for single-cell analysis:
+- DifferentiableAmbientRemoval: CellBender-style ambient RNA removal
+- DifferentiableCellAnnotator: Cell type annotation (celltypist/cellassign/scanvi)
+- DifferentiableDiffusionImputer: MAGIC-style diffusion imputation
+- DifferentiableDoubletScorer: Scrublet-style doublet detection
+- DifferentiableCellCommunication: GNN-based cell-cell communication analysis
+- DifferentiableGRN: GATv2-based gene regulatory network inference
+- DifferentiableLigandReceptor: Ligand-receptor co-expression scoring
+- SoftKMeansClustering: Differentiable soft k-means clustering
+- DifferentiableHarmony: Harmony-style batch correction
+- DifferentiableMMDBatchCorrection: MMD-regularised autoencoder batch correction
+- DifferentiableWGANBatchCorrection: Adversarial (WGAN) batch correction
+- DifferentiableSwitchDE: Sigmoidal switch differential expression
+- DifferentiableVelocity: RNA velocity via Neural ODEs
+- DifferentiablePseudotime: Diffusion-map pseudotime ordering
+- DifferentiableFateProbability: Absorption-based fate estimation
+- DifferentiableSpatialDomain: STAGATE-style spatial domain identification
+- DifferentiablePASTEAlignment: PASTE-style spatial slice alignment
+- DifferentiableDifferentialDistribution: scDD-style differential distribution testing
+- DifferentiableSimulator: Splatter-style single-cell count simulation
+- DifferentiableArchetypalAnalysis: PCHA-style archetypal analysis
+- DifferentiableOTTrajectory: Waddington-OT-style optimal transport trajectory
+"""
+
+from diffbio.operators.singlecell.archetypes import (
+    ArchetypalAnalysisConfig,
+    DifferentiableArchetypalAnalysis,
+)
+from diffbio.operators.singlecell.ambient_removal import (
+    AmbientRemovalConfig,
+    DifferentiableAmbientRemoval,
+)
+from diffbio.operators.singlecell.cell_annotation import (
+    CellAnnotatorConfig,
+    DifferentiableCellAnnotator,
+)
+from diffbio.operators.singlecell.communication import (
+    CellCommunicationConfig,
+    DifferentiableCellCommunication,
+    DifferentiableLigandReceptor,
+    LRScoringConfig,
+)
+from diffbio.operators.singlecell.batch_correction import (
+    BatchCorrectionConfig,
+    DifferentiableHarmony,
+)
+from diffbio.operators.singlecell.enhanced_batch_correction import (
+    DifferentiableMMDBatchCorrection,
+    DifferentiableWGANBatchCorrection,
+    MMDBatchCorrectionConfig,
+    WGANBatchCorrectionConfig,
+)
+from diffbio.operators.singlecell.grn_inference import (
+    DifferentiableGRN,
+    GRNInferenceConfig,
+)
+from diffbio.operators.singlecell.doublet_detection import (
+    DifferentiableDoubletScorer,
+    DifferentiableSoloDetector,
+    DoubletScorerConfig,
+    SoloDetectorConfig,
+)
+from diffbio.operators.singlecell.imputation import (
+    DifferentiableDiffusionImputer,
+    DifferentiableTransformerDenoiser,
+    DiffusionImputerConfig,
+    TransformerDenoiserConfig,
+)
+from diffbio.operators.singlecell.ot_trajectory import (
+    DifferentiableOTTrajectory,
+    OTTrajectoryConfig,
+)
+from diffbio.operators.singlecell.soft_clustering import (
+    SoftClusteringConfig,
+    SoftKMeansClustering,
+)
+from diffbio.operators.singlecell.switch_de import (
+    DifferentiableSwitchDE,
+    SwitchDEConfig,
+)
+from diffbio.operators.singlecell.trajectory import (
+    DifferentiableFateProbability,
+    DifferentiablePseudotime,
+    FateProbabilityConfig,
+    PseudotimeConfig,
+)
+from diffbio.operators.singlecell.spatial_domains import (
+    DifferentiablePASTEAlignment,
+    DifferentiableSpatialDomain,
+    PASTEAlignmentConfig,
+    SpatialDomainConfig,
+)
+from diffbio.operators.singlecell.differential_distribution import (
+    DifferentiableDifferentialDistribution,
+    DifferentialDistributionConfig,
+)
+from diffbio.operators.singlecell.simulation import (
+    DifferentiableSimulator,
+    SimulationConfig,
+)
+from diffbio.operators.singlecell.velocity import (
+    DifferentiableVelocity,
+    VelocityConfig,
+)
+from diffbio.operators.singlecell.downsampling import (
+    DownsamplingConfig,
+    ReadDownsampler,
+)
+from diffbio.operators.singlecell.knockdown_filter import (
+    KnockdownFilterConfig,
+    OnTargetKnockdownFilter,
+)
+from diffbio.operators.singlecell.sindy_grn import (
+    SINDyGRNConfig,
+    SINDyGRNOperator,
+)
+
+__all__ = [
+    # Archetypal Analysis
+    "ArchetypalAnalysisConfig",
+    "DifferentiableArchetypalAnalysis",
+    # Ambient Removal
+    "AmbientRemovalConfig",
+    "DifferentiableAmbientRemoval",
+    # Cell Annotation
+    "CellAnnotatorConfig",
+    "DifferentiableCellAnnotator",
+    # Clustering
+    "SoftClusteringConfig",
+    "SoftKMeansClustering",
+    # Communication (L-R scoring + GNN-based)
+    "CellCommunicationConfig",
+    "DifferentiableCellCommunication",
+    "DifferentiableLigandReceptor",
+    "LRScoringConfig",
+    # Batch Correction (Harmony)
+    "BatchCorrectionConfig",
+    "DifferentiableHarmony",
+    # Batch Correction (MMD + WGAN)
+    "DifferentiableMMDBatchCorrection",
+    "DifferentiableWGANBatchCorrection",
+    "MMDBatchCorrectionConfig",
+    "WGANBatchCorrectionConfig",
+    # GRN Inference
+    "DifferentiableGRN",
+    "GRNInferenceConfig",
+    # Doublet Detection
+    "DifferentiableDoubletScorer",
+    "DifferentiableSoloDetector",
+    "DoubletScorerConfig",
+    "SoloDetectorConfig",
+    # Imputation
+    "DifferentiableDiffusionImputer",
+    "DifferentiableTransformerDenoiser",
+    "DiffusionImputerConfig",
+    "TransformerDenoiserConfig",
+    # Switch DE
+    "DifferentiableSwitchDE",
+    "SwitchDEConfig",
+    # Trajectory Inference
+    "DifferentiableFateProbability",
+    "DifferentiablePseudotime",
+    "FateProbabilityConfig",
+    "PseudotimeConfig",
+    # Spatial Domain Identification
+    "DifferentiablePASTEAlignment",
+    "DifferentiableSpatialDomain",
+    "PASTEAlignmentConfig",
+    "SpatialDomainConfig",
+    # Differential Distribution
+    "DifferentiableDifferentialDistribution",
+    "DifferentialDistributionConfig",
+    # Simulation
+    "DifferentiableSimulator",
+    "SimulationConfig",
+    # OT Trajectory
+    "DifferentiableOTTrajectory",
+    "OTTrajectoryConfig",
+    # Velocity
+    "DifferentiableVelocity",
+    "VelocityConfig",
+    # Downsampling
+    "DownsamplingConfig",
+    "ReadDownsampler",
+    # Knockdown Filter
+    "KnockdownFilterConfig",
+    "OnTargetKnockdownFilter",
+    # SINDy GRN
+    "SINDyGRNConfig",
+    "SINDyGRNOperator",
+]