diffbio 0.1.0__py3-none-any.whl

diffbio/core/soft_ops/quantile.py

@@ -0,0 +1,376 @@
+"""Soft quantile, median, and percentile operators.
+
+Provides differentiable relaxations of quantile-based statistics.
+Quantiles are computed via :func:`~diffbio.core.soft_ops.sorting.argsort`
+or :func:`~diffbio.core.soft_ops.sorting.sort`, with interpolation
+following the same methods as ``jax.numpy.quantile``.
+"""
+
+from typing import Literal, cast
+
+import jax
+import jax.numpy as jnp
+from jax import Array
+
+from diffbio.core.soft_ops._projections_simplex import SimplexMode, proj_simplex
+from diffbio.core.soft_ops._types import SoftIndex
+from diffbio.core.soft_ops._utils import (
+    canonicalize_axis,
+    ensure_float,
+    quantile_interpolation_params,
+    standardize_and_squash,
+)
+from diffbio.core.soft_ops.selection import take_along_axis
+from diffbio.core.soft_ops.sorting import (
+    _neuralsort_a_sum,
+    _sorting_network_permutation,
+)
+
+Mode = Literal["hard", "smooth", "c0", "c1", "c2"]
+ArgMethod = Literal["softsort", "neuralsort", "sorting_network"]
+
+
+def _sorting_network_quantile_index(
+    x_last: Array,
+    softness: float | Array,
+    mode: Mode,
+    *,
+    standardize: bool,
+    take_next: bool,
+    a_b: Array,
+    k: Array,
+    kp1: Array,
+) -> Array:
+    """Return the soft quantile index from the sorting-network backend."""
+    perm = _sorting_network_permutation(
+        x_last,
+        softness,
+        cast(SimplexMode, mode),
+        descending=False,
+        standardized=standardize,
+    )
+    if take_next:
+        return (1.0 - a_b) * perm[..., k, :] + a_b * perm[..., kp1, :]
+    return perm[..., k, :]
+
+
+def argquantile(
+    x: Array,
+    q: Array,
+    axis: int | None = None,
+    keepdims: bool = False,
+    softness: float | Array = 0.1,
+    mode: Mode = "smooth",
+    method: ArgMethod = "neuralsort",
+    quantile_method: Literal[
+        "linear",
+        "lower",
+        "higher",
+        "nearest",
+        "midpoint",
+    ] = "linear",
+    standardize: bool = True,
+) -> SoftIndex:
+    """Soft argquantile returning SoftIndex.
+
+    Args:
+        x: Input array.
+        q: Quantile(s) in [0, 1]. Scalar or 1-D array.
+        axis: Axis along which to compute. None flattens.
+        keepdims: If True, keep reduced dimension.
+        softness: Controls sharpness (> 0).
+        mode: Smoothness mode.
+        method: Algorithm.
+        quantile_method: Interpolation method.
+        standardize: If True, standardize input.
+
+    Returns:
+        SoftIndex probability distribution over quantile position(s).
+    """
+    q_arr = jnp.asarray(q)
+    if q_arr.ndim > 1:
+        msg = f"q must be scalar or 1-D, got shape {q_arr.shape}"
+        raise ValueError(msg)
+    if q_arr.ndim == 1:
+
+        def _single(qi: Array) -> SoftIndex:
+            return argquantile(
+                x,
+                q=qi,
+                axis=axis,
+                keepdims=keepdims,
+                softness=softness,
+                mode=mode,
+                method=method,
+                quantile_method=quantile_method,
+                standardize=standardize,
+            )
+
+        return jax.vmap(_single)(q_arr)
+
+    orig_axis_is_none = axis is None
+    if axis is None:
+        num_dims = x.ndim
+        x = jnp.ravel(x)
+        axis = 0
+    else:
+        axis = canonicalize_axis(axis, x.ndim)
+        num_dims = None
+
+    if mode != "hard":
+        x = ensure_float(x)
+    if standardize and mode != "hard":
+        x = standardize_and_squash(x, axis=axis)
+
+    x_last = jnp.moveaxis(x, axis, -1)
+    *batch_dims, n = x_last.shape
+
+    q_val = jnp.clip(q, 0.0, 1.0)
+    k, a, take_next = quantile_interpolation_params(q_val, n, quantile_method)
+    a_b = jnp.expand_dims(a, axis=-1)
+    kp1 = jnp.minimum(k + 1, n - 1)
+
+    if mode == "hard":
+        indices = jnp.argsort(x_last, axis=-1, descending=False)
+        if take_next:
+            idx_pair = jnp.stack(
+                [indices[..., k], indices[..., kp1]],
+                axis=-1,
+            )
+            oh = jax.nn.one_hot(idx_pair, num_classes=n, axis=-1)
+            soft_index = (1.0 - a_b) * oh[..., 0, :] + a_b * oh[..., 1, :]
+        else:
+            soft_index = jax.nn.one_hot(
+                indices[..., k],
+                num_classes=n,
+                axis=-1,
+            )
+    elif method == "softsort":
+        x_sorted = jnp.sort(x_last, axis=-1, descending=False)
+        if take_next:
+            anchors = jnp.stack(
+                [x_sorted[..., k], x_sorted[..., kp1]],
+                axis=-1,
+            )
+            abs_diff = jnp.abs(
+                anchors[..., :, None] - x_last[..., None, :],
+            )
+            proj = proj_simplex(-abs_diff, axis=-1, softness=softness, mode=mode)
+            soft_index = (1.0 - a_b) * proj[..., 0, :] + a_b * proj[..., 1, :]
+        else:
+            anchors = x_sorted[..., k, None]
+            abs_diff = jnp.abs(
+                anchors[..., :, None] - x_last[..., None, :],
+            )
+            soft_index = proj_simplex(
+                -abs_diff,
+                axis=-1,
+                softness=softness,
+                mode=mode,
+            )[..., 0, :]
+    elif method == "neuralsort":
+        a_sum = _neuralsort_a_sum(x_last, mode=mode, softness=softness)
+        if take_next:
+            i = jnp.array([k + 1, k + 2])
+            coef = n + 1 - 2 * i
+            coef = jnp.broadcast_to(coef, (*batch_dims, 2))
+            z = -(coef[..., :, None] * x_last[..., None, :] + a_sum[..., None, :])
+            proj = proj_simplex(z, axis=-1, softness=softness, mode=mode)
+            soft_index = (1.0 - a_b) * proj[..., 0, :] + a_b * proj[..., 1, :]
+        else:
+            coef = jnp.array([n + 1 - 2 * (k + 1)])
+            coef = jnp.broadcast_to(coef, (*batch_dims, 1))
+            z = -(coef[..., :, None] * x_last[..., None, :] + a_sum[..., None, :])
+            soft_index = proj_simplex(
+                z,
+                axis=-1,
+                softness=softness,
+                mode=mode,
+            )[..., 0, :]
+    elif method == "sorting_network":
+        soft_index = _sorting_network_quantile_index(
+            x_last,
+            softness,
+            mode,
+            standardize=standardize,
+            take_next=take_next,
+            a_b=a_b,
+            k=k,
+            kp1=kp1,
+        )
+    else:
+        msg = f"Invalid method: {method!r}"
+        raise ValueError(msg)
+
+    if keepdims:
+        if orig_axis_is_none:
+            soft_index = soft_index.reshape(*(1,) * num_dims, n)
+        else:
+            soft_index = jnp.expand_dims(soft_index, axis=axis)
+
+    return soft_index
+
+
+def quantile(
+    x: Array,
+    q: Array,
+    axis: int | None = None,
+    keepdims: bool = False,
+    softness: float | Array = 0.1,
+    mode: Mode = "smooth",
+    method: ArgMethod = "neuralsort",
+    quantile_method: Literal[
+        "linear",
+        "lower",
+        "higher",
+        "nearest",
+        "midpoint",
+    ] = "linear",
+    standardize: bool = True,
+    gated_grad: bool = True,
+) -> Array:
+    """Soft quantile returning value.
+
+    Implemented as :func:`argquantile` + :func:`take_along_axis`
+    for most methods.
+
+    Args:
+        x: Input array.
+        q: Quantile(s) in [0, 1].
+        axis: Axis along which to compute.
+        keepdims: If True, keep reduced dimension.
+        softness: Controls sharpness (> 0).
+        mode: Smoothness mode.
+        method: Algorithm.
+        quantile_method: Interpolation method.
+        standardize: If True, standardize input.
+        gated_grad: If False, stop gradient through soft index.
+
+    Returns:
+        Quantile value(s).
+    """
+    if mode == "hard":
+        return jnp.quantile(x, q, axis=axis, keepdims=keepdims, method=quantile_method)
+
+    soft_idx = argquantile(
+        x,
+        q,
+        axis=axis,
+        keepdims=True,
+        softness=softness,
+        mode=mode,
+        method=method,
+        quantile_method=quantile_method,
+        standardize=standardize,
+    )
+    if not gated_grad:
+        soft_idx = jax.lax.stop_gradient(soft_idx)
+
+    _axis = 0 if axis is None else canonicalize_axis(axis, x.ndim)
+    if axis is None:
+        x = jnp.ravel(x)
+
+    q_arr = jnp.asarray(q)
+    if q_arr.ndim == 1:
+        result = jax.vmap(lambda si: take_along_axis(x, si, axis=_axis))(soft_idx)
+    else:
+        result = take_along_axis(x, soft_idx, axis=_axis)
+
+    if not keepdims:
+        if q_arr.ndim == 0:
+            result = jnp.squeeze(result, axis=_axis if axis is not None else 0)
+    return result
+
+
+def argmedian(
+    x: Array,
+    axis: int | None = None,
+    keepdims: bool = False,
+    softness: float | Array = 0.1,
+    mode: Mode = "smooth",
+    method: ArgMethod = "neuralsort",
+    standardize: bool = True,
+) -> SoftIndex:
+    """Soft argmedian: :func:`argquantile` with ``q=0.5``."""
+    return argquantile(
+        x,
+        q=jnp.array(0.5),
+        axis=axis,
+        keepdims=keepdims,
+        softness=softness,
+        mode=mode,
+        method=method,
+        standardize=standardize,
+    )
+
+
+def median(
+    x: Array,
+    axis: int | None = None,
+    keepdims: bool = False,
+    softness: float | Array = 0.1,
+    mode: Mode = "smooth",
+    method: ArgMethod = "neuralsort",
+    standardize: bool = True,
+    gated_grad: bool = True,
+) -> Array:
+    """Soft median: :func:`quantile` with ``q=0.5``."""
+    return quantile(
+        x,
+        q=jnp.array(0.5),
+        axis=axis,
+        keepdims=keepdims,
+        softness=softness,
+        mode=mode,
+        method=method,
+        standardize=standardize,
+        gated_grad=gated_grad,
+    )
+
+
+def argpercentile(
+    x: Array,
+    p: Array,
+    axis: int | None = None,
+    keepdims: bool = False,
+    softness: float | Array = 0.1,
+    mode: Mode = "smooth",
+    method: ArgMethod = "neuralsort",
+    standardize: bool = True,
+) -> SoftIndex:
+    """Soft argpercentile: :func:`argquantile` with ``q = p / 100``."""
+    return argquantile(
+        x,
+        q=jnp.asarray(p) / 100.0,
+        axis=axis,
+        keepdims=keepdims,
+        softness=softness,
+        mode=mode,
+        method=method,
+        standardize=standardize,
+    )
+
+
+def percentile(
+    x: Array,
+    p: Array,
+    axis: int | None = None,
+    keepdims: bool = False,
+    softness: float | Array = 0.1,
+    mode: Mode = "smooth",
+    method: ArgMethod = "neuralsort",
+    standardize: bool = True,
+    gated_grad: bool = True,
+) -> Array:
+    """Soft percentile: :func:`quantile` with ``q = p / 100``."""
+    return quantile(
+        x,
+        q=jnp.asarray(p) / 100.0,
+        axis=axis,
+        keepdims=keepdims,
+        softness=softness,
+        mode=mode,
+        method=method,
+        standardize=standardize,
+        gated_grad=gated_grad,
+    )

diffbio/core/soft_ops/selection.py

@@ -0,0 +1,236 @@
+"""Soft selection and indexing operators.
+
+Provides differentiable relaxations of array selection operations using
+SoftBool conditions and SoftIndex probability distributions instead of
+discrete boolean masks and integer indices.
+
+Key types:
+- **SoftBool**: Probability in [0, 1], used in :func:`where`.
+- **SoftIndex**: Probability distribution over indices (sums to 1),
+  used in :func:`take_along_axis`, :func:`take`, :func:`choose`, etc.
+"""
+
+from collections.abc import Sequence
+
+import jax
+import jax.numpy as jnp
+from jax import Array
+
+from diffbio.core.soft_ops._types import SoftBool, SoftIndex
+from diffbio.core.soft_ops._utils import canonicalize_axis, normalize_axis_argument
+
+
+def where(condition: SoftBool, x: Array, y: Array) -> Array:
+    """Soft where: ``x * condition + y * (1 - condition)``.
+
+    Unlike ``jnp.where``, this smoothly interpolates between ``x`` and
+    ``y`` based on the continuous condition value.
+
+    Args:
+        condition: SoftBool in [0, 1], same shape as x and y.
+        x: Values selected when condition is 1.
+        y: Values selected when condition is 0.
+
+    Returns:
+        Interpolated array.
+    """
+    return x * condition + y * (1.0 - condition)
+
+
+def take_along_axis(
+    x: Array,
+    soft_index: SoftIndex,
+    axis: int | None = -1,
+) -> Array:
+    """Soft take_along_axis via weighted dot product.
+
+    ``soft_index`` must have one more dimension than ``x``: the extra
+    (last) dimension contains the probability distribution over the
+    elements along ``axis``.
+
+    Args:
+        x: Input array of shape ``(..., n, ...)``.
+        soft_index: SoftIndex of shape ``(..., k, ..., [n])`` where
+            ``[n]`` is the probability distribution dimension.
+        axis: Axis in ``x`` to select from. If None, ``x`` is flattened.
+
+    Returns:
+        Array of shape ``(..., k, ...)``.
+    """
+    x, axis = normalize_axis_argument(x, axis)
+    if x.ndim + 1 != soft_index.ndim:
+        msg = (
+            f"x.ndim + 1 == soft_index.ndim required, "
+            f"got x.ndim={x.ndim}, soft_index.ndim={soft_index.ndim}"
+        )
+        raise ValueError(msg)
+    x = jnp.moveaxis(x, axis, -1)
+    soft_index = jnp.moveaxis(soft_index, axis, -2)
+    dotprod = jnp.einsum("...n,...kn->...k", x, soft_index)
+    return jnp.moveaxis(dotprod, -1, axis)
+
+
+def take(
+    x: Array,
+    soft_index: SoftIndex,
+    axis: int | None = None,
+) -> Array:
+    """Soft take via weighted dot product.
+
+    Unlike :func:`take_along_axis`, ``soft_index`` is a 2-D matrix
+    of shape ``(k, [n])`` applied uniformly across batch dimensions.
+
+    Args:
+        x: Input array of shape ``(..., n, ...)``.
+        soft_index: SoftIndex of shape ``(k, [n])``.
+        axis: Axis to select from. If None, ``x`` is flattened.
+
+    Returns:
+        Array of shape ``(..., k, ...)``.
+    """
+    if soft_index.ndim != 2:
+        msg = f"soft_index must be (k, [n]), got shape {soft_index.shape}"
+        raise ValueError(msg)
+    x, axis = normalize_axis_argument(x, axis)
+    if axis != x.ndim - 1:
+        x = jnp.moveaxis(x, axis, -1)
+    soft_index = jnp.reshape(
+        soft_index,
+        (1,) * (x.ndim - 1) + soft_index.shape,
+    )
+    x = jnp.expand_dims(x, axis)
+    soft_index = jnp.moveaxis(soft_index, -2, axis)
+    return jnp.sum(x * soft_index, axis=-1)
+
+
+def choose(
+    soft_index: SoftIndex,
+    choices: Array,
+) -> Array:
+    """Soft choose among multiple arrays.
+
+    Softly selects among ``choices`` using ``soft_index`` weights.
+
+    Args:
+        soft_index: SoftIndex of shape ``(..., [n])``.
+        choices: Array of shape ``(n, ...)``.
+
+    Returns:
+        Weighted combination of choices.
+    """
+    if soft_index.ndim != choices.ndim or soft_index.shape[-1] != choices.shape[0]:
+        msg = (
+            f"Incompatible shapes: soft_index={soft_index.shape}, "
+            f"choices={choices.shape}. Need soft_index.shape=(..., [n]) "
+            f"and choices.shape=(n, ...)"
+        )
+        raise ValueError(msg)
+    tgt_shape = jnp.broadcast_shapes(choices.shape[1:], soft_index.shape[:-1])
+    choices_bcast = jnp.broadcast_to(choices, (choices.shape[0], *tgt_shape))
+    choices_bcast = jnp.moveaxis(choices_bcast, 0, -1)
+    return jnp.sum(choices_bcast * soft_index, axis=-1)
+
+
+def dynamic_index_in_dim(
+    x: Array,
+    soft_index: SoftIndex,
+    axis: int = 0,
+    keepdims: bool = True,
+) -> Array:
+    """Soft dynamic indexing along a dimension.
+
+    Selects a single element (weighted combination) along ``axis``
+    using the probability distribution ``soft_index``.
+
+    Args:
+        x: Input array of shape ``(..., n, ...)``.
+        soft_index: SoftIndex of shape ``([n],)``.
+        axis: Axis to index.
+        keepdims: If True, retains the indexed dimension as size 1.
+
+    Returns:
+        Indexed array.
+    """
+    axis = canonicalize_axis(axis, x.ndim)
+    if x.shape[axis] != soft_index.shape[0]:
+        msg = (
+            f"Dimension mismatch: x.shape[{axis}]={x.shape[axis]} "
+            f"vs soft_index.shape[0]={soft_index.shape[0]}"
+        )
+        raise ValueError(msg)
+    x = jnp.moveaxis(x, axis, -1)
+    x_reshaped = jnp.reshape(x, (-1, x.shape[-1]))
+    dotprod = jnp.sum(x_reshaped * soft_index[None, :], axis=-1)
+    y = jnp.reshape(dotprod, x.shape[:-1])
+    if keepdims:
+        y = jnp.expand_dims(y, axis=axis)
+    return y
+
+
+def dynamic_slice_in_dim(
+    x: Array,
+    soft_start_index: SoftIndex,
+    slice_size: int,
+    axis: int = 0,
+) -> Array:
+    """Soft dynamic slicing along a dimension.
+
+    Extracts a soft slice of ``slice_size`` elements starting at the
+    position defined by ``soft_start_index``.
+
+    Args:
+        x: Input array of shape ``(..., n, ...)``.
+        soft_start_index: SoftIndex of shape ``([n],)``.
+        slice_size: Number of elements to extract.
+        axis: Axis to slice.
+
+    Returns:
+        Array of shape ``(..., slice_size, ...)``.
+    """
+    axis = canonicalize_axis(axis, x.ndim)
+    if not (0 < slice_size <= x.shape[axis]):
+        msg = (
+            f"slice_size must satisfy 0 < slice_size <= x.shape[axis], "
+            f"got slice_size={slice_size}, x.shape[axis]={x.shape[axis]}"
+        )
+        raise ValueError(msg)
+
+    x_last = jnp.moveaxis(x, axis, -1)
+    t_idx = jnp.arange(slice_size)
+
+    def one_step(t: Array) -> Array:
+        rolled = jnp.roll(x_last, shift=-t, axis=-1)
+        return jnp.einsum("...n,n->...", rolled, soft_start_index)
+
+    y_stack = jax.vmap(one_step)(t_idx)
+    y_last = jnp.moveaxis(y_stack, 0, -1)
+    return jnp.moveaxis(y_last, -1, axis)
+
+
+def dynamic_slice(
+    x: Array,
+    soft_start_indices: Sequence[SoftIndex],
+    slice_sizes: Sequence[int],
+) -> Array:
+    """Soft dynamic slicing across multiple dimensions.
+
+    Applies :func:`dynamic_slice_in_dim` sequentially along each axis.
+
+    Args:
+        x: Input array of shape ``(n_1, n_2, ..., n_k)``.
+        soft_start_indices: One SoftIndex per dimension.
+        slice_sizes: One slice length per dimension.
+
+    Returns:
+        Array of shape ``(l_1, l_2, ..., l_k)``.
+    """
+    if not (len(soft_start_indices) == len(slice_sizes) == x.ndim):
+        msg = (
+            f"len(soft_start_indices) == len(slice_sizes) == x.ndim required, "
+            f"got {len(soft_start_indices)}, {len(slice_sizes)}, {x.ndim}"
+        )
+        raise ValueError(msg)
+    y = x
+    for axis, (start, size) in enumerate(zip(soft_start_indices, slice_sizes)):
+        y = dynamic_slice_in_dim(y, start, size, axis=axis)
+    return y