diffbio 0.1.0__py3-none-any.whl

diffbio/core/soft_ops/straight_through.py

@@ -0,0 +1,261 @@
+"""Straight-through estimator decorators and _st variants.
+
+Straight-through estimators use the hard (exact, non-differentiable)
+function for the forward pass but route gradients through the soft
+(differentiable) version during backpropagation.
+
+This module provides:
+- :func:`st`: Decorator that creates a straight-through version of any
+  soft_ops function with a ``mode`` parameter.
+- :func:`grad_replace`: Lower-level decorator for custom forward/backward
+  split functions.
+- 27 pre-built ``_st`` variants (e.g., :func:`relu_st`, :func:`sort_st`).
+"""
+
+import functools
+import inspect
+from collections.abc import Callable
+
+import jax
+from jax import tree_util as jtu
+
+from diffbio.core.soft_ops import comparison, elementwise, quantile, sorting
+
+
+def grad_replace(fn: Callable) -> Callable:
+    """Decorator for custom forward/backward computation split.
+
+    The decorated function is called twice: once with ``forward=True``
+    (output used for forward pass) and once with ``forward=False``
+    (output used for gradient computation).
+
+    Args:
+        fn: Function accepting a ``forward: bool`` keyword argument.
+
+    Returns:
+        Wrapped function using hard forward, soft backward.
+    """
+
+    @functools.wraps(fn)
+    def wrapped(*args, **kwargs):
+        fw_y = fn(*args, **kwargs, forward=True)
+        bw_y = fn(*args, **kwargs, forward=False)
+        fw_leaves, fw_treedef = jtu.tree_flatten(
+            fw_y,
+            is_leaf=lambda x: x is None,
+        )
+        bw_leaves, _ = jtu.tree_flatten(
+            bw_y,
+            is_leaf=lambda x: x is None,
+        )
+        out_leaves = [
+            f if f is None or b is None else jax.lax.stop_gradient(f - b) + b
+            for f, b in zip(fw_leaves, bw_leaves)
+        ]
+        return jtu.tree_unflatten(fw_treedef, out_leaves)
+
+    return wrapped
+
+
+def st(fn: Callable) -> Callable:
+    """Decorator creating a straight-through estimator from a soft_ops function.
+
+    The decorated function is called twice: once with ``mode="hard"``
+    (forward pass) and once with the specified ``mode`` (backward pass).
+    The trick ``stop_gradient(hard - soft) + soft`` ensures the forward
+    output is hard but gradients flow through the soft version.
+
+    Args:
+        fn: Function with a ``mode`` parameter (e.g., any elementwise,
+            comparison, or sorting function).
+
+    Returns:
+        Wrapped straight-through estimator function.
+    """
+    sig = inspect.signature(fn)
+    mode_param = sig.parameters.get("mode")
+    if mode_param is not None:
+        mode_default = mode_param.default
+        mode_idx = list(sig.parameters.keys()).index("mode")
+    else:
+        mode_default = "smooth"
+        mode_idx = None
+
+    @functools.wraps(fn)
+    def wrapped(*args, **kwargs):
+        if mode_idx is not None and len(args) > mode_idx:
+            mode = args[mode_idx]
+            args = args[:mode_idx] + args[mode_idx + 1 :]
+        else:
+            mode = kwargs.pop("mode", mode_default)
+        fw_y = fn(*args, **kwargs, mode="hard")
+        bw_y = fn(*args, **kwargs, mode=mode)
+        fw_leaves, fw_treedef = jtu.tree_flatten(
+            fw_y,
+            is_leaf=lambda x: x is None,
+        )
+        bw_leaves, _ = jtu.tree_flatten(
+            bw_y,
+            is_leaf=lambda x: x is None,
+        )
+        out_leaves = [
+            f if f is None or b is None else jax.lax.stop_gradient(f - b) + b
+            for f, b in zip(fw_leaves, bw_leaves)
+        ]
+        return jtu.tree_unflatten(fw_treedef, out_leaves)
+
+    return wrapped
+
+
+# ---------------------------------------------------------------------------
+# Cached ST wrapper to avoid repeated inspect.signature() calls
+# ---------------------------------------------------------------------------
+
+_st_cache: dict[Callable, Callable] = {}
+
+
+def _cached_st(fn: Callable) -> Callable:
+    """Return a cached st() wrapper."""
+    if fn not in _st_cache:
+        _st_cache[fn] = st(fn)
+    return _st_cache[fn]
+
+
+# ---------------------------------------------------------------------------
+# 27 pre-built straight-through variants
+# ---------------------------------------------------------------------------
+
+
+def abs_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.elementwise.abs`."""
+    return _cached_st(elementwise.abs)(*args, **kwargs)
+
+
+def argmax_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.sorting.argmax`."""
+    return _cached_st(sorting.argmax)(*args, **kwargs)
+
+
+def argmedian_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.quantile.argmedian`."""
+    return _cached_st(quantile.argmedian)(*args, **kwargs)
+
+
+def argmin_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.sorting.argmin`."""
+    return _cached_st(sorting.argmin)(*args, **kwargs)
+
+
+def argpercentile_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.quantile.argpercentile`."""
+    return _cached_st(quantile.argpercentile)(*args, **kwargs)
+
+
+def argquantile_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.quantile.argquantile`."""
+    return _cached_st(quantile.argquantile)(*args, **kwargs)
+
+
+def argsort_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.sorting.argsort`."""
+    return _cached_st(sorting.argsort)(*args, **kwargs)
+
+
+def clip_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.elementwise.clip`."""
+    return _cached_st(elementwise.clip)(*args, **kwargs)
+
+
+def equal_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.comparison.equal`."""
+    return _cached_st(comparison.equal)(*args, **kwargs)
+
+
+def greater_equal_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.comparison.greater_equal`."""
+    return _cached_st(comparison.greater_equal)(*args, **kwargs)
+
+
+def greater_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.comparison.greater`."""
+    return _cached_st(comparison.greater)(*args, **kwargs)
+
+
+def heaviside_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.elementwise.heaviside`."""
+    return _cached_st(elementwise.heaviside)(*args, **kwargs)
+
+
+def isclose_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.comparison.isclose`."""
+    return _cached_st(comparison.isclose)(*args, **kwargs)
+
+
+def less_equal_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.comparison.less_equal`."""
+    return _cached_st(comparison.less_equal)(*args, **kwargs)
+
+
+def less_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.comparison.less`."""
+    return _cached_st(comparison.less)(*args, **kwargs)
+
+
+def max_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.sorting.max`."""
+    return _cached_st(sorting.max)(*args, **kwargs)
+
+
+def median_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.quantile.median`."""
+    return _cached_st(quantile.median)(*args, **kwargs)
+
+
+def min_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.sorting.min`."""
+    return _cached_st(sorting.min)(*args, **kwargs)
+
+
+def not_equal_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.comparison.not_equal`."""
+    return _cached_st(comparison.not_equal)(*args, **kwargs)
+
+
+def percentile_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.quantile.percentile`."""
+    return _cached_st(quantile.percentile)(*args, **kwargs)
+
+
+def quantile_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.quantile.quantile`."""
+    return _cached_st(quantile.quantile)(*args, **kwargs)
+
+
+def rank_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.sorting.rank`."""
+    return _cached_st(sorting.rank)(*args, **kwargs)
+
+
+def relu_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.elementwise.relu`."""
+    return _cached_st(elementwise.relu)(*args, **kwargs)
+
+
+def round_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.elementwise.round`."""
+    return _cached_st(elementwise.round)(*args, **kwargs)
+
+
+def sign_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.elementwise.sign`."""
+    return _cached_st(elementwise.sign)(*args, **kwargs)
+
+
+def sort_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.sorting.sort`."""
+    return _cached_st(sorting.sort)(*args, **kwargs)
+
+
+def top_k_st(*args, **kwargs):
+    """Straight-through :func:`~diffbio.core.soft_ops.sorting.top_k`."""
+    return _cached_st(sorting.top_k)(*args, **kwargs)

diffbio/core/uncertainty.py

@@ -0,0 +1,279 @@
+"""Uncertainty quantification wrappers for DiffBio operators.
+
+Provides ensemble and conformal prediction wrappers that add
+``uncertainty``, ``confidence_interval_lower``, and
+``confidence_interval_upper`` keys to any operator's output dict.
+
+Uses ``nnx.vmap`` with ``nnx.StateAxes`` to vectorize multiple forward
+passes across different RNG seeds (ensemble) or dropout samples
+(conformal), computing mean predictions and uncertainty estimates
+without Python for-loops.
+"""
+
+from __future__ import annotations
+
+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 PyTree
+
+from diffbio.core import soft_ops
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class EnsembleUQConfig(OperatorConfig):
+    """Configuration for ensemble-based uncertainty quantification.
+
+    Attributes:
+        n_members: Number of ensemble members (forward passes with
+            different random seeds).
+        confidence_level: Confidence level for intervals (0 to 1).
+    """
+
+    n_members: int = 5
+    confidence_level: float = 0.95
+
+
+@dataclass(frozen=True)
+class ConformalUQConfig(OperatorConfig):
+    """Configuration for conformal prediction-based UQ.
+
+    Attributes:
+        alpha: Significance level (1 - confidence). Smaller alpha
+            gives wider intervals.
+        num_samples: Number of Monte Carlo samples for interval estimation.
+    """
+
+    alpha: float = 0.1
+    num_samples: int = 20
+
+
+def _find_primary_output_key(result: dict[str, Any], data: dict[str, Any]) -> str | None:
+    """Find the main output key added by the operator.
+
+    Returns the first key in result that is not in data and holds a JAX array.
+    """
+    for key in result:
+        if key not in data and isinstance(result[key], jax.Array):
+            return key
+    return None
+
+
+def _run_n_samples(
+    operator: OperatorModule,
+    data: dict[str, Any],
+    state: dict[str, Any],
+    metadata: dict[str, Any] | None,
+    n_samples: int,
+    primary_key: str,
+) -> jax.Array:
+    """Run operator n_samples times and stack the primary output.
+
+    Uses jax.lax.scan to avoid Python for-loops. The operator is
+    deterministic per call (same params, same data), so all samples
+    are identical for non-stochastic operators. For stochastic operators
+    (dropout, noise), each call samples fresh randomness via NNX RNG state.
+
+    Args:
+        operator: The base operator to run.
+        data: Input data dict.
+        state: Element state.
+        metadata: Element metadata.
+        n_samples: Number of forward passes.
+        primary_key: Key to stack from the output dict.
+
+    Returns:
+        Stacked array of shape (n_samples, *output_shape).
+    """
+
+    def _step(_carry: None, _xs: None) -> tuple[None, jax.Array]:
+        result, _, _ = operator.apply(data, state, metadata)
+        return None, result[primary_key]
+
+    _, stacked = jax.lax.scan(_step, None, None, length=n_samples)
+    return stacked
+
+
+def _compute_uq_stats(
+    stacked: jax.Array,
+    confidence_level: float,
+) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
+    """Compute mean, std, and confidence interval from stacked samples.
+
+    Args:
+        stacked: Array of shape (n_samples, *output_shape).
+        confidence_level: Confidence level for quantile intervals.
+
+    Returns:
+        Tuple of (mean, std, lower, upper).
+    """
+    mean = jnp.mean(stacked, axis=0)
+    std = jnp.std(stacked, axis=0)
+    alpha = 1.0 - confidence_level
+    lower = soft_ops.quantile(stacked, jnp.array(alpha / 2), axis=0, softness=0.1)
+    upper = soft_ops.quantile(stacked, jnp.array(1 - alpha / 2), axis=0, softness=0.1)
+    return mean, std, lower, upper
+
+
+class EnsembleUQOperator(OperatorModule):
+    """Ensemble-based uncertainty quantification wrapper.
+
+    Runs the base operator multiple times and aggregates outputs to produce
+    mean predictions with uncertainty estimates (standard deviation) and
+    confidence intervals (quantile-based). Uses ``jax.lax.scan`` to avoid
+    Python for-loops.
+
+    Output adds:
+        - ``uncertainty``: Standard deviation across ensemble members.
+        - ``confidence_interval_lower``: Lower bound of confidence interval.
+        - ``confidence_interval_upper``: Upper bound of confidence interval.
+
+    Args:
+        config: EnsembleUQConfig with ensemble parameters.
+        base_operator: The operator to wrap.
+        rngs: Random number generators.
+        name: Optional operator name.
+    """
+
+    def __init__(
+        self,
+        config: EnsembleUQConfig,
+        *,
+        base_operator: OperatorModule,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ) -> None:
+        """Initialize ensemble UQ operator."""
+        super().__init__(config, rngs=rngs, name=name)
+        self.config: EnsembleUQConfig = config
+        self.base_operator = base_operator
+
+    def apply(
+        self,
+        data: PyTree,
+        state: PyTree,
+        metadata: dict[str, Any] | None,
+        random_params: Any = None,  # noqa: ARG002
+        stats: dict[str, Any] | None = None,  # noqa: ARG002
+    ) -> tuple[PyTree, PyTree, dict[str, Any] | None]:
+        """Run ensemble forward passes and aggregate with uncertainty.
+
+        Args:
+            data: Input data dict for the base operator.
+            state: Element state (passed through).
+            metadata: Element metadata (passed through).
+            random_params: Unused.
+            stats: Unused.
+
+        Returns:
+            Tuple of (output_with_uncertainty, state, metadata).
+        """
+        # Get a single result to find output keys and structure
+        base_result, _, _ = self.base_operator.apply(data, state, metadata)
+        primary_key = _find_primary_output_key(base_result, data)
+        if primary_key is None:
+            return base_result, state, metadata
+
+        stacked = _run_n_samples(
+            self.base_operator,
+            data,
+            state,
+            metadata,
+            self.config.n_members,
+            primary_key,
+        )
+        mean, std, lower, upper = _compute_uq_stats(
+            stacked,
+            self.config.confidence_level,
+        )
+
+        result = dict(base_result)
+        result[primary_key] = mean
+        result["uncertainty"] = std
+        result["confidence_interval_lower"] = lower
+        result["confidence_interval_upper"] = upper
+        return result, state, metadata
+
+
+class ConformalUQOperator(OperatorModule):
+    """Conformal prediction-based uncertainty quantification wrapper.
+
+    Uses Monte Carlo sampling to estimate empirical prediction intervals.
+    Runs the base operator multiple times via ``jax.lax.scan`` and computes
+    quantile-based intervals at the specified confidence level (1 - alpha).
+
+    Output adds:
+        - ``uncertainty``: Standard deviation across samples.
+        - ``confidence_interval_lower``: Lower quantile bound.
+        - ``confidence_interval_upper``: Upper quantile bound.
+
+    Args:
+        config: ConformalUQConfig with sampling parameters.
+        base_operator: The operator to wrap.
+        rngs: Random number generators.
+        name: Optional operator name.
+    """
+
+    def __init__(
+        self,
+        config: ConformalUQConfig,
+        *,
+        base_operator: OperatorModule,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ) -> None:
+        """Initialize conformal UQ operator."""
+        super().__init__(config, rngs=rngs, name=name)
+        self.config: ConformalUQConfig = config
+        self.base_operator = base_operator
+
+    def apply(
+        self,
+        data: PyTree,
+        state: PyTree,
+        metadata: dict[str, Any] | None,
+        random_params: Any = None,  # noqa: ARG002
+        stats: dict[str, Any] | None = None,  # noqa: ARG002
+    ) -> tuple[PyTree, PyTree, dict[str, Any] | None]:
+        """Run conformal prediction sampling.
+
+        Args:
+            data: Input data dict for the base operator.
+            state: Element state (passed through).
+            metadata: Element metadata (passed through).
+            random_params: Unused.
+            stats: Unused.
+
+        Returns:
+            Tuple of (output_with_intervals, state, metadata).
+        """
+        base_result, _, _ = self.base_operator.apply(data, state, metadata)
+        primary_key = _find_primary_output_key(base_result, data)
+        if primary_key is None:
+            return base_result, state, metadata
+
+        stacked = _run_n_samples(
+            self.base_operator,
+            data,
+            state,
+            metadata,
+            self.config.num_samples,
+            primary_key,
+        )
+        confidence_level = 1.0 - self.config.alpha
+        mean, std, lower, upper = _compute_uq_stats(stacked, confidence_level)
+
+        result = dict(base_result)
+        result[primary_key] = mean
+        result["uncertainty"] = std
+        result["confidence_interval_lower"] = lower
+        result["confidence_interval_upper"] = upper
+        return result, state, metadata

diffbio/evaluation/__init__.py

@@ -0,0 +1,42 @@
+"""Evaluation harness for scBench and SpatialBench benchmarks.
+
+Provides grading algorithms, benchmark problem definitions, task adapters,
+and a runner for evaluating DiffBio operators against real-world benchmark
+problems from single-cell and spatial transcriptomics domains.
+"""
+
+from diffbio.evaluation.adapters import TaskAdapter, compute_quality_metrics
+from diffbio.evaluation.graders import (
+    GradeResult,
+    grade_distribution_comparison,
+    grade_label_set_jaccard,
+    grade_marker_gene_precision_recall,
+    grade_multiple_choice,
+    grade_numeric_tolerance,
+)
+from diffbio.evaluation.problem import BenchmarkProblem, load_problems
+from diffbio.evaluation.runner import (
+    BenchmarkSummary,
+    EvalResult,
+    run_benchmark,
+    run_problem,
+    summarize,
+)
+
+__all__ = [
+    "BenchmarkProblem",
+    "BenchmarkSummary",
+    "EvalResult",
+    "GradeResult",
+    "TaskAdapter",
+    "compute_quality_metrics",
+    "grade_distribution_comparison",
+    "grade_label_set_jaccard",
+    "grade_marker_gene_precision_recall",
+    "grade_multiple_choice",
+    "grade_numeric_tolerance",
+    "load_problems",
+    "run_benchmark",
+    "run_problem",
+    "summarize",
+]