diffbio 0.1.0__py3-none-any.whl

diffbio/operators/singlecell/velocity.py

@@ -0,0 +1,361 @@
+"""RNA velocity estimation via Neural ODEs.
+
+This module provides differentiable RNA velocity estimation using neural
+networks to learn splicing kinetics and integrate ODEs.
+
+Key technique: Uses neural networks to learn per-gene kinetics parameters
+(transcription, splicing, degradation rates) and integrates the splicing
+ODE differentiably using Euler method.
+
+Applications: Inferring cell state transitions and developmental trajectories
+from single-cell RNA-seq data with spliced/unspliced counts.
+"""
+
+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, Float, PyTree
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class VelocityConfig(OperatorConfig):
+    """Configuration for DifferentiableVelocity.
+
+    Attributes:
+        n_genes: Number of genes.
+        hidden_dim: Hidden dimension for neural networks.
+        dt: Time step for ODE integration.
+        n_steps: Number of integration steps.
+        kinetics_model: Type of kinetics model ("standard" or "dynamical").
+    """
+
+    n_genes: int = 2000
+    hidden_dim: int = 64
+    dt: float = 0.1
+    n_steps: int = 10
+    kinetics_model: str = "standard"
+
+
+class TimeEncoder(nnx.Module):
+    """Encoder for estimating latent time from expression."""
+
+    def __init__(
+        self,
+        n_genes: int,
+        hidden_dim: int,
+        *,
+        rngs: nnx.Rngs,
+    ):
+        """Initialize the time encoder.
+
+        Args:
+            n_genes: Number of input genes.
+            hidden_dim: Hidden dimension.
+            rngs: Random number generators.
+        """
+        super().__init__()
+
+        # Encode spliced + unspliced to latent time
+        self.linear1 = nnx.Linear(
+            in_features=n_genes * 2,
+            out_features=hidden_dim,
+            rngs=rngs,
+        )
+        self.linear2 = nnx.Linear(
+            in_features=hidden_dim,
+            out_features=hidden_dim,
+            rngs=rngs,
+        )
+        self.time_proj = nnx.Linear(
+            in_features=hidden_dim,
+            out_features=1,
+            rngs=rngs,
+        )
+        self.norm = nnx.LayerNorm(num_features=hidden_dim, rngs=rngs)
+
+    def __call__(
+        self,
+        spliced: Float[Array, "n_cells n_genes"],
+        unspliced: Float[Array, "n_cells n_genes"],
+    ) -> Float[Array, "n_cells"]:
+        """Estimate latent time for each cell.
+
+        Args:
+            spliced: Spliced (mature) mRNA counts.
+            unspliced: Unspliced (nascent) mRNA counts.
+
+        Returns:
+            Latent time estimates per cell (bounded 0-1).
+        """
+        # Concatenate and log-transform
+        x = jnp.concatenate([jnp.log1p(spliced), jnp.log1p(unspliced)], axis=-1)
+
+        # Encode
+        x = nnx.gelu(self.linear1(x))
+        x = self.norm(nnx.gelu(self.linear2(x)))
+
+        # Project to time (sigmoid for 0-1 range)
+        time = jax.nn.sigmoid(self.time_proj(x)).squeeze(-1)
+
+        return time
+
+
+class KineticsEncoder(nnx.Module):
+    """Encoder for learning per-gene kinetics parameters."""
+
+    def __init__(
+        self,
+        n_genes: int,
+        *,
+        rngs: nnx.Rngs,
+    ):
+        """Initialize the kinetics encoder.
+
+        Args:
+            n_genes: Number of genes.
+            rngs: Random number generators.
+        """
+        super().__init__()
+
+        # Learnable per-gene kinetics (alpha, beta, gamma)
+        # Initialize with reasonable defaults
+        key = rngs.params()
+        k1, k2, k3 = jax.random.split(key, 3)
+
+        # Transcription rate (alpha): typically 0.1-10
+        self.log_alpha = nnx.Param(jax.random.normal(k1, (n_genes,)) * 0.5)
+
+        # Splicing rate (beta): typically 0.1-1
+        self.log_beta = nnx.Param(jax.random.normal(k2, (n_genes,)) * 0.5 - 1.0)
+
+        # Degradation rate (gamma): typically 0.01-0.5
+        self.log_gamma = nnx.Param(jax.random.normal(k3, (n_genes,)) * 0.5 - 2.0)
+
+    def __call__(
+        self,
+    ) -> tuple[
+        Float[Array, "n_genes"],
+        Float[Array, "n_genes"],
+        Float[Array, "n_genes"],
+    ]:
+        """Get kinetics parameters.
+
+        Returns:
+            Tuple of (alpha, beta, gamma) - all positive via softplus.
+        """
+        alpha = jax.nn.softplus(self.log_alpha[...])
+        beta = jax.nn.softplus(self.log_beta[...])
+        gamma = jax.nn.softplus(self.log_gamma[...])
+
+        return alpha, beta, gamma
+
+
+class DifferentiableVelocity(OperatorModule):
+    """Differentiable RNA velocity estimation via Neural ODEs.
+
+    This operator estimates RNA velocity from spliced and unspliced
+    counts using learned kinetics parameters and differentiable ODE
+    integration.
+
+    Algorithm:
+    1. Encode expression to latent time per cell
+    2. Learn per-gene kinetics (alpha, beta, gamma)
+    3. Compute velocity from splicing ODE:
+       ds/dt = beta * u - gamma * s
+       du/dt = alpha - beta * u
+    4. Integrate ODE using Euler method
+
+    Args:
+        config: VelocityConfig with model parameters.
+        rngs: Flax NNX random number generators.
+        name: Optional operator name.
+
+    Example:
+        ```python
+        config = VelocityConfig(n_genes=2000)
+        velocity = DifferentiableVelocity(config, rngs=nnx.Rngs(42))
+        data = {"spliced": spliced, "unspliced": unspliced}
+        result, state, meta = velocity.apply(data, {}, None)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: VelocityConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize the velocity operator.
+
+        Args:
+            config: Velocity configuration.
+            rngs: Random number generators for initialization.
+            name: Optional operator name.
+        """
+        super().__init__(config, rngs=rngs, name=name)
+
+        if rngs is None:
+            rngs = nnx.Rngs(0)
+
+        self.n_genes = config.n_genes
+        self.dt = config.dt
+        self.n_steps = config.n_steps
+
+        # Time encoder
+        self.time_encoder = TimeEncoder(
+            n_genes=config.n_genes,
+            hidden_dim=config.hidden_dim,
+            rngs=rngs,
+        )
+
+        # Kinetics encoder
+        self.kinetics_encoder = KineticsEncoder(
+            n_genes=config.n_genes,
+            rngs=rngs,
+        )
+
+    def _compute_velocity(
+        self,
+        spliced: Float[Array, "n_cells n_genes"],
+        unspliced: Float[Array, "n_cells n_genes"],
+        beta: Float[Array, "n_genes"],
+        gamma: Float[Array, "n_genes"],
+    ) -> Float[Array, "n_cells n_genes"]:
+        """Compute RNA velocity from splicing dynamics.
+
+        The standard RNA velocity model:
+        ds/dt = beta * u - gamma * s
+
+        Args:
+            spliced: Spliced counts.
+            unspliced: Unspliced counts.
+            beta: Splicing rate.
+            gamma: Degradation rate.
+
+        Returns:
+            Velocity (ds/dt) for each cell-gene pair.
+        """
+        # Velocity is the rate of change of spliced mRNA
+        # ds/dt = splicing_in - degradation_out
+        # ds/dt = beta * unspliced - gamma * spliced
+        velocity = beta[None, :] * unspliced - gamma[None, :] * spliced
+
+        return velocity
+
+    def _euler_step(
+        self,
+        s: Float[Array, "n_cells n_genes"],
+        u: Float[Array, "n_cells n_genes"],
+        alpha: Float[Array, "n_genes"],
+        beta: Float[Array, "n_genes"],
+        gamma: Float[Array, "n_genes"],
+        dt: float,
+    ) -> tuple[Float[Array, "n_cells n_genes"], Float[Array, "n_cells n_genes"]]:
+        """Single Euler integration step.
+
+        Args:
+            s: Current spliced counts.
+            u: Current unspliced counts.
+            alpha: Transcription rate.
+            beta: Splicing rate.
+            gamma: Degradation rate.
+            dt: Time step.
+
+        Returns:
+            Tuple of (new_spliced, new_unspliced).
+        """
+        # ODE system:
+        # du/dt = alpha - beta * u
+        # ds/dt = beta * u - gamma * s
+        du_dt = alpha[None, :] - beta[None, :] * u
+        ds_dt = beta[None, :] * u - gamma[None, :] * s
+
+        # Euler update
+        u_new = u + dt * du_dt
+        s_new = s + dt * ds_dt
+
+        # Ensure non-negative
+        u_new = jnp.maximum(u_new, 0.0)
+        s_new = jnp.maximum(s_new, 0.0)
+
+        return s_new, u_new
+
+    def apply(
+        self,
+        data: PyTree,
+        state: PyTree,
+        metadata: dict[str, Any] | None,
+        random_params: Any = None,
+        stats: dict[str, Any] | None = None,
+    ) -> tuple[PyTree, PyTree, dict[str, Any] | None]:
+        """Apply RNA velocity estimation.
+
+        Args:
+            data: Dictionary containing:
+                - "spliced": Spliced mRNA counts (n_cells, n_genes)
+                - "unspliced": Unspliced mRNA counts (n_cells, n_genes)
+            state: Element state (passed through unchanged)
+            metadata: Element metadata (passed through unchanged)
+            random_params: Not used
+            stats: Not used
+
+        Returns:
+            Tuple of (transformed_data, state, metadata):
+                - transformed_data contains:
+
+                    - "spliced": Original spliced counts
+                    - "unspliced": Original unspliced counts
+                    - "velocity": RNA velocity estimates
+                    - "latent_time": Estimated latent time per cell
+                    - "alpha": Transcription rate per gene
+                    - "beta": Splicing rate per gene
+                    - "gamma": Degradation rate per gene
+                    - "projected_spliced": Projected future spliced
+                - state is passed through unchanged
+                - metadata is passed through unchanged
+        """
+        spliced = data["spliced"]
+        unspliced = data["unspliced"]
+
+        # Estimate latent time per cell
+        latent_time = self.time_encoder(spliced, unspliced)
+
+        # Get kinetics parameters
+        alpha, beta, gamma = self.kinetics_encoder()
+
+        # Compute velocity
+        velocity = self._compute_velocity(spliced, unspliced, beta, gamma)
+
+        # Project forward using learned dynamics via jax.lax.scan
+        def _euler_body(
+            carry: tuple[jax.Array, jax.Array], _: None
+        ) -> tuple[tuple[jax.Array, jax.Array], None]:
+            s, u = carry
+            s_new, u_new = self._euler_step(s, u, alpha, beta, gamma, self.dt)
+            return (s_new, u_new), None
+
+        (s_proj, u_proj), _ = jax.lax.scan(
+            _euler_body, (spliced, unspliced), None, length=self.n_steps
+        )
+
+        transformed_data = {
+            "spliced": spliced,
+            "unspliced": unspliced,
+            "velocity": velocity,
+            "latent_time": latent_time,
+            "alpha": alpha,
+            "beta": beta,
+            "gamma": gamma,
+            "projected_spliced": s_proj,
+        }
+
+        return transformed_data, state, metadata

diffbio/operators/statistical/__init__.py

@@ -0,0 +1,35 @@
+"""Differentiable statistical model operators.
+
+This module provides operators for:
+
+- Hidden Markov Models with differentiable forward algorithm
+- Negative Binomial GLM for differential expression
+- Unrolled EM for transcript quantification
+
+All operators maintain gradient flow for end-to-end training.
+"""
+
+from diffbio.operators.statistical.em_quantification import (
+    DifferentiableEMQuantifier,
+    EMQuantifierConfig,
+)
+from diffbio.operators.statistical.hmm import (
+    DifferentiableHMM,
+    HMMConfig,
+)
+from diffbio.operators.statistical.nb_glm import (
+    DifferentiableNBGLM,
+    NBGLMConfig,
+)
+
+__all__ = [
+    # HMM
+    "HMMConfig",
+    "DifferentiableHMM",
+    # Negative Binomial GLM
+    "NBGLMConfig",
+    "DifferentiableNBGLM",
+    # EM Quantification
+    "EMQuantifierConfig",
+    "DifferentiableEMQuantifier",
+]

diffbio/operators/statistical/em_quantification.py

@@ -0,0 +1,260 @@
+"""Differentiable EM-based transcript quantification operator.
+
+This module provides an unrolled EM algorithm for transcript
+quantification, inspired by Salmon and Kallisto.
+
+Key technique: Fixed number of EM iterations enables gradient flow
+through all steps of the algorithm.
+
+Applications: RNA-seq transcript quantification, isoform abundance estimation.
+
+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 jax
+import jax.numpy as jnp
+from datarax.core.config import OperatorConfig
+from flax import nnx
+from jaxtyping import Array, Float, PyTree
+
+from diffbio.core.base_operators import TemperatureOperator
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class EMQuantifierConfig(OperatorConfig):
+    """Configuration for DifferentiableEMQuantifier.
+
+    Attributes:
+        n_transcripts: Number of transcripts to quantify.
+        n_iterations: Fixed number of EM iterations (for unrolling).
+        temperature: Temperature for softmax in E-step.
+    """
+
+    n_transcripts: int = 1000
+    n_iterations: int = 10
+    temperature: float = 1.0
+
+
+class DifferentiableEMQuantifier(TemperatureOperator):
+    """Differentiable EM for transcript quantification.
+
+    This operator implements the EM algorithm for estimating transcript
+    abundances from read-to-transcript compatibility data. The fixed
+    number of iterations enables gradient flow through the entire
+    quantification process.
+
+    Algorithm:
+    1. Initialize abundances (learnable prior)
+    2. E-step: Probabilistic assignment of reads to transcripts
+       weights = softmax(compatibility * abundances / temperature)
+    3. M-step: Update abundances from weighted counts
+       abundances = sum(weights) / effective_lengths
+       abundances = abundances / sum(abundances)
+    4. Repeat for n_iterations
+
+    Inherits from TemperatureOperator to get:
+
+    - _temperature property for temperature-controlled smoothing
+    - soft_max() for logsumexp-based smooth maximum
+    - soft_argmax() for soft position selection
+
+    Args:
+        config: EMQuantifierConfig with model parameters.
+        rngs: Flax NNX random number generators.
+        name: Optional operator name.
+
+    Example:
+        ```python
+        config = EMQuantifierConfig(n_transcripts=1000, n_iterations=10)
+        quantifier = DifferentiableEMQuantifier(config, rngs=nnx.Rngs(42))
+        data = {"compatibility": compat_matrix, "effective_lengths": eff_lens}
+        result, state, meta = quantifier.apply(data, {}, None)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: EMQuantifierConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize the EM quantifier operator.
+
+        Args:
+            config: EM quantifier configuration.
+            rngs: Random number generators for initialization.
+            name: Optional operator name.
+        """
+        super().__init__(config, rngs=rngs, name=name)
+
+        if rngs is None:
+            rngs = nnx.Rngs(0)
+
+        self.n_transcripts = config.n_transcripts
+        self.n_iterations = config.n_iterations
+        # Temperature is now managed by TemperatureOperator via self._temperature
+
+        # Initialize log abundances (will be normalized via softmax)
+        # Shape: (n_transcripts,)
+        key = rngs.params()
+        init_log_abundances = jax.random.normal(key, (config.n_transcripts,)) * 0.01
+        self.log_initial_abundances = nnx.Param(init_log_abundances)
+
+    def get_initial_abundances(self) -> Float[Array, "n_transcripts"]:
+        """Get normalized initial abundances.
+
+        Returns:
+            Initial abundance distribution (n_transcripts,), sums to 1.
+        """
+        return jax.nn.softmax(self.log_initial_abundances[...])
+
+    def em_step(
+        self,
+        abundances: Float[Array, "n_transcripts"],
+        compatibility: Float[Array, "n_reads n_transcripts"],
+        effective_lengths: Float[Array, "n_transcripts"],
+    ) -> Float[Array, "n_transcripts"]:
+        """Perform one EM iteration.
+
+        Args:
+            abundances: Current abundance estimates.
+            compatibility: Read-transcript compatibility matrix.
+            effective_lengths: Effective transcript lengths.
+
+        Returns:
+            Updated abundance estimates.
+        """
+        # E-step: Compute read assignment probabilities
+        # P(transcript | read) propto compatibility * abundance / length
+        rate = abundances / (effective_lengths + 1e-8)  # (n_transcripts,)
+
+        # Score for each read-transcript pair
+        scores = compatibility * rate  # (n_reads, n_transcripts)
+
+        # Normalize per read (softmax with temperature)
+        # Use inherited _temperature property from TemperatureOperator
+        weights = jax.nn.softmax(
+            jnp.log(scores + 1e-10) / self._temperature, axis=1
+        )  # (n_reads, n_transcripts)
+
+        # M-step: Update abundances
+        # Expected count for each transcript
+        expected_counts = jnp.sum(weights, axis=0)  # (n_transcripts,)
+
+        # Normalize by effective length
+        new_abundances = expected_counts / (effective_lengths + 1e-8)
+
+        # Normalize to sum to 1
+        new_abundances = new_abundances / (jnp.sum(new_abundances) + 1e-8)
+
+        return new_abundances
+
+    def quantify(
+        self,
+        compatibility: Float[Array, "n_reads n_transcripts"],
+        effective_lengths: Float[Array, "n_transcripts"],
+    ) -> Float[Array, "n_transcripts"]:
+        """Run EM algorithm for quantification.
+
+        Args:
+            compatibility: Read-transcript compatibility matrix.
+            effective_lengths: Effective transcript lengths.
+
+        Returns:
+            Final transcript abundance estimates.
+        """
+        # Initialize abundances
+        abundances = self.get_initial_abundances()
+
+        # Run fixed number of EM iterations
+        def em_iteration(abundances, _):
+            new_abundances = self.em_step(abundances, compatibility, effective_lengths)
+            return new_abundances, None
+
+        final_abundances, _ = jax.lax.scan(em_iteration, abundances, None, length=self.n_iterations)
+
+        return final_abundances
+
+    def compute_tpm(
+        self,
+        abundances: Float[Array, "n_transcripts"],
+        effective_lengths: Float[Array, "n_transcripts"],
+    ) -> Float[Array, "n_transcripts"]:
+        """Convert abundances to TPM (Transcripts Per Million).
+
+        Args:
+            abundances: Normalized abundance estimates.
+            effective_lengths: Effective transcript lengths.
+
+        Returns:
+            TPM values (sum to 1 million).
+        """
+        # TPM = (abundance / length) / sum(abundance / length) * 1e6
+        rate = abundances / (effective_lengths + 1e-8)
+        tpm = rate / (jnp.sum(rate) + 1e-8) * 1e6
+        return tpm
+
+    def apply(
+        self,
+        data: PyTree,
+        state: PyTree,
+        metadata: dict[str, Any] | None,
+        random_params: Any = None,
+        stats: dict[str, Any] | None = None,
+    ) -> tuple[PyTree, PyTree, dict[str, Any] | None]:
+        """Apply EM quantification to read assignment data.
+
+        This method runs the EM algorithm to estimate transcript
+        abundances from read-transcript compatibility data.
+
+        Args:
+            data: Dictionary containing:
+                - "compatibility": Read-transcript compatibility matrix
+                  (n_reads, n_transcripts)
+                - "effective_lengths": Effective transcript lengths
+                  (n_transcripts,)
+            state: Element state (passed through unchanged)
+            metadata: Element metadata (passed through unchanged)
+            random_params: Not used (deterministic operator)
+            stats: Not used
+
+        Returns:
+            Tuple of (transformed_data, state, metadata):
+                - transformed_data contains:
+
+                    - "compatibility": Original compatibility matrix
+                    - "effective_lengths": Original effective lengths
+                    - "abundances": Estimated transcript abundances
+                    - "tpm": TPM (Transcripts Per Million) values
+                - state is passed through unchanged
+                - metadata is passed through unchanged
+        """
+        compatibility = data["compatibility"]
+        effective_lengths = data["effective_lengths"]
+
+        # Run EM quantification
+        abundances = self.quantify(compatibility, effective_lengths)
+
+        # Compute TPM
+        tpm = self.compute_tpm(abundances, effective_lengths)
+
+        # Build output data
+        transformed_data = {
+            "compatibility": compatibility,
+            "effective_lengths": effective_lengths,
+            "abundances": abundances,
+            "tpm": tpm,
+        }
+
+        return transformed_data, state, metadata