diffbio 0.1.0__py3-none-any.whl

diffbio/core/soft_ops/_types.py

@@ -0,0 +1,15 @@
+"""Type aliases for soft differentiable operations.
+
+SoftBool represents a probability in [0, 1] -- a soft relaxation of a
+boolean value. SoftIndex represents a probability distribution over
+discrete indices -- a soft relaxation of an integer index.
+"""
+
+from jax import Array
+from jaxtyping import Float
+
+SoftBool = Float[Array, "..."]
+"""Soft boolean: probability in [0, 1]."""
+
+SoftIndex = Float[Array, "..."]
+"""Soft index: probabilities summing to 1 along the last axis."""

diffbio/core/soft_ops/_utils.py

@@ -0,0 +1,342 @@
+"""Internal utility functions for soft_ops.
+
+These are shared helpers used across the soft_ops submodules.
+Not part of the public API.
+"""
+
+from typing import Literal, overload
+
+import jax
+import jax.numpy as jnp
+from jax import Array
+
+
+def validate_softness(softness: float | Array) -> None:
+    """Raise ``ValueError`` if softness is not positive.
+
+    Validation is skipped inside JAX-traced contexts (jit, grad, vmap)
+    since concrete values are not available during tracing.
+
+    Args:
+        softness: The softness (temperature) parameter. Must be > 0.
+    """
+    if isinstance(softness, jax.core.Tracer):
+        return
+    if float(softness) <= 0:
+        msg = f"softness must be positive, got {softness}"
+        raise ValueError(msg)
+
+
+def ensure_float(x: float | Array | jnp.ndarray) -> jnp.ndarray:
+    """Cast to default float dtype if not already floating point.
+
+    Args:
+        x: Input array or scalar.
+
+    Returns:
+        Array with floating-point dtype.
+    """
+    x = jnp.asarray(x)
+    if jnp.issubdtype(x.dtype, jnp.floating):
+        return x
+    return x.astype(jnp.result_type(float))
+
+
+def normalize_axis_argument(
+    x: jnp.ndarray,
+    axis: int | None,
+) -> tuple[jnp.ndarray, int]:
+    """Flatten when ``axis`` is None, otherwise canonicalize the axis."""
+    if axis is None:
+        return jnp.ravel(x), 0
+    return x, canonicalize_axis(axis, x.ndim)
+
+
+@overload
+def standardize_and_squash(
+    x: jnp.ndarray,
+    axis: int = ...,
+    eps: float = ...,
+    temperature: float = ...,
+    return_mean_std: Literal[False] = ...,
+) -> jnp.ndarray: ...
+
+
+@overload
+def standardize_and_squash(
+    x: jnp.ndarray,
+    axis: int = ...,
+    eps: float = ...,
+    temperature: float = ...,
+    *,
+    return_mean_std: Literal[True],
+) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: ...
+
+
+def standardize_and_squash(
+    x: jnp.ndarray,
+    axis: int = -1,
+    eps: float = 1e-6,
+    temperature: float = 1.0,
+    return_mean_std: bool = False,
+) -> jnp.ndarray | tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
+    """Standardize along axis then squash to (0, 1) via sigmoid.
+
+    Steps:
+        1. Standardize: ``(x - mean) / std`` along ``axis``
+        2. Scale by ``1 / temperature``
+        3. Squash: apply sigmoid to map to (0, 1)
+
+    Args:
+        x: Input array.
+        axis: Axis along which to standardize.
+        eps: Epsilon for numerical stability in std computation.
+        temperature: Controls sharpness. Lower = sharper.
+        return_mean_std: If True, also return mean and std.
+
+    Returns:
+        Squashed array in (0, 1). If ``return_mean_std``, returns
+        ``(squashed, mean, std)`` tuple.
+    """
+    mean = jnp.mean(x, axis=axis, keepdims=True)
+    var = jnp.mean((x - mean) ** 2, axis=axis, keepdims=True)
+    std = jnp.sqrt(var + eps)
+    z = (x - mean) / std
+    z = z / temperature
+    z = jax.nn.sigmoid(z)
+    if return_mean_std:
+        return z, mean, std
+    return z
+
+
+def unsquash_and_destandardize(
+    y: jnp.ndarray,
+    mean: jnp.ndarray,
+    std: jnp.ndarray,
+    eps: float = 1e-10,
+    temperature: float = 1.0,
+) -> jnp.ndarray:
+    """Inverse of :func:`standardize_and_squash`.
+
+    Steps:
+        1. Unsquash: logit to map from (0, 1) back to R
+        2. Scale by ``temperature``
+        3. Destandardize: ``z * std + mean``
+
+    Args:
+        y: Squashed array in (0, 1).
+        mean: Mean from standardization.
+        std: Std from standardization.
+        eps: Clipping epsilon for logit stability.
+        temperature: Must match the value used in squash.
+
+    Returns:
+        Recovered array in original scale.
+    """
+    safe_eps = jnp.maximum(eps, 10 * jnp.finfo(y.dtype).eps)
+    y = jnp.clip(y, safe_eps, 1.0 - safe_eps)
+    z = jnp.log(y / (1.0 - y))
+    z = z * temperature
+    return z * std + mean
+
+
+def quantile_interpolation_params(
+    q: jnp.ndarray,
+    n: int,
+    method: Literal["linear", "lower", "higher", "nearest", "midpoint"],
+) -> tuple[jnp.ndarray, jnp.ndarray, bool]:
+    """Compute interpolation parameters for quantile computation.
+
+    Maps a quantile ``q`` in [0, 1] to index ``k``, interpolation
+    weight ``a``, and whether to take the next element.
+
+    Args:
+        q: Quantile value(s) in [0, 1].
+        n: Number of elements in the sorted array.
+        method: Interpolation method.
+
+    Returns:
+        Tuple of (k, a, take_next) where k is the integer index,
+        a is the interpolation weight, and take_next indicates
+        whether to interpolate with the next element.
+    """
+    p = q * (n - 1)
+
+    if method == "linear":
+        k = jnp.floor(p).astype(jnp.int32)
+        a = p - k
+        take_next = True
+    elif method == "lower":
+        k = jnp.floor(p).astype(jnp.int32)
+        k = jnp.clip(k, 0, n - 1)
+        a = jnp.zeros_like(p)
+        take_next = False
+    elif method == "higher":
+        k = jnp.ceil(p).astype(jnp.int32)
+        k = jnp.clip(k, 0, n - 1)
+        a = jnp.zeros_like(p)
+        take_next = False
+    elif method == "nearest":
+        flag = jnp.less_equal(p - jnp.floor(p), 0.5)
+        k = jnp.where(flag, jnp.floor(p), jnp.ceil(p)).astype(jnp.int32)
+        a = jnp.zeros_like(p)
+        take_next = False
+    elif method == "midpoint":
+        k = jnp.floor(p).astype(jnp.int32)
+        a = jnp.full_like(p, 0.5)
+        is_int = jnp.isclose(p, jnp.round(p))
+        a = jnp.where(is_int, 0.0, a)
+        take_next = True
+    else:
+        msg = f"Unknown quantile method: {method!r}"
+        raise ValueError(msg)
+    return k, a, take_next
+
+
+def map_in_chunks(
+    f,
+    xs: jnp.ndarray,
+    chunk_size: int,
+) -> jnp.ndarray:
+    """Map ``f`` row-wise over axis 0 using checkpointed ``lax.scan``.
+
+    ``f`` receives a chunk of shape ``(chunk_size, *rest)`` and must
+    return ``(chunk_size, *out_rest)``. Uses ``jax.checkpoint`` for
+    O(n) backward memory.
+
+    Args:
+        f: Function to apply to each chunk.
+        xs: Input array with shape ``(n, *rest)``.
+        chunk_size: Number of rows per chunk.
+
+    Returns:
+        Output array with shape ``(n, *out_rest)``.
+    """
+    n = xs.shape[0]
+    if chunk_size >= n:
+        return f(xs)
+    remainder = n % chunk_size
+    if remainder:
+        pad_size = chunk_size - remainder
+        padding = jnp.zeros((pad_size, *xs.shape[1:]), dtype=xs.dtype)
+        xs_padded = jnp.concatenate([xs, padding], axis=0)
+    else:
+        xs_padded = xs
+    n_padded = xs_padded.shape[0]
+    xs_chunked = xs_padded.reshape(
+        n_padded // chunk_size,
+        chunk_size,
+        *xs.shape[1:],
+    )
+    f_remat = jax.checkpoint(f)
+    _, ys = jax.lax.scan(
+        lambda _, chunk: (None, f_remat(chunk)),
+        None,
+        xs_chunked,
+    )
+    ys = ys.reshape(n_padded, *ys.shape[2:])
+    return ys[:n]
+
+
+def reduce_in_chunks(
+    f,
+    xs: jnp.ndarray,
+    chunk_size: int,
+) -> jnp.ndarray:
+    """Apply ``f`` to chunks of ``xs`` along axis 0 and sum results.
+
+    ``f`` receives a chunk of shape ``(chunk_size, *rest)`` and must
+    return a result whose shape does **not** include the chunk
+    dimension. Remainder rows are zero-padded; the correction
+    ``f(zeros)`` is subtracted.
+
+    Uses ``jax.checkpoint`` for memory-efficient backpropagation.
+
+    Args:
+        f: Reduction function applied per chunk.
+        xs: Input array with shape ``(n, *rest)``.
+        chunk_size: Number of rows per chunk.
+
+    Returns:
+        Summed result across all chunks.
+    """
+    n = xs.shape[0]
+    if chunk_size >= n:
+        return f(xs)
+    remainder = n % chunk_size
+    if remainder:
+        pad_size = chunk_size - remainder
+        padding = jnp.zeros((pad_size, *xs.shape[1:]), dtype=xs.dtype)
+        xs_padded = jnp.concatenate([xs, padding], axis=0)
+    else:
+        pad_size = 0
+        xs_padded = xs
+    n_padded = xs_padded.shape[0]
+    xs_chunked = xs_padded.reshape(
+        n_padded // chunk_size,
+        chunk_size,
+        *xs.shape[1:],
+    )
+
+    out_struct = jax.eval_shape(f, xs_chunked[0])
+    init = jnp.zeros(out_struct.shape, dtype=out_struct.dtype)
+
+    f_remat = jax.checkpoint(f)
+
+    def body(acc: jnp.ndarray, chunk: jnp.ndarray):
+        """Accumulate reduction result for one chunk."""
+        return acc + f_remat(chunk), None
+
+    result, _ = jax.lax.scan(body, init, xs_chunked)
+
+    if remainder:
+        zero_pad = jnp.zeros((pad_size, *xs.shape[1:]), dtype=xs.dtype)
+        result = result - f(zero_pad)
+    return result
+
+
+def canonicalize_axis(axis: int | None, num_dims: int) -> int:
+    """Normalize axis to a positive integer.
+
+    Args:
+        axis: Axis index (positive or negative). Must not be None.
+        num_dims: Number of dimensions in the array.
+
+    Returns:
+        Normalized positive axis index.
+
+    Raises:
+        ValueError: If axis is None or out of bounds.
+    """
+    if axis is None:
+        msg = "axis must be specified"
+        raise ValueError(msg)
+    if not -num_dims <= axis < num_dims:
+        msg = f"axis {axis} is out of bounds for array of dimension {num_dims}"
+        raise ValueError(msg)
+    if axis < 0:
+        axis += num_dims
+    return axis
+
+
+def flatten_or_canonicalize_axis(
+    x: jnp.ndarray,
+    axis: int | None,
+) -> tuple[jnp.ndarray, int, int | None]:
+    """Flatten ``x`` if ``axis`` is None, otherwise canonicalize.
+
+    Returns ``(x_prepared, resolved_axis, original_ndim_if_flattened)``.
+    The third element is None when axis was not None.
+
+    Args:
+        x: Input array.
+        axis: Axis index or None (flatten).
+
+    Returns:
+        Tuple of prepared array, resolved axis, and original ndim
+        (only set when the array was flattened).
+    """
+    if axis is None:
+        num_dims = x.ndim
+        return jnp.ravel(x), 0, num_dims
+    return x, canonicalize_axis(axis, x.ndim), None

diffbio/core/soft_ops/autograd_safe.py

@@ -0,0 +1,120 @@
+"""Autograd-safe math operations.
+
+Provides NaN-free alternatives to standard JAX math functions by using
+the double-where trick: the forward pass computes the correct value
+even at domain boundaries, and the backward pass produces finite (zero)
+gradients instead of NaN/Inf.
+
+The double-where trick works by:
+    1. Replacing problematic inputs with safe values (e.g., 0 -> 1 for sqrt)
+    2. Computing the function on the safe input
+    3. Using ``jnp.where`` to select the safe output or a fallback (e.g., 0)
+
+Because JAX traces through both branches of ``jnp.where``, step 1
+ensures the "unused" branch never produces NaN in its gradient.
+"""
+
+import jax.numpy as jnp
+from jax import Array
+
+
+def sqrt(x: Array) -> Array:
+    """Autograd-safe square root.
+
+    Returns ``sqrt(x)`` for ``x > 0`` and ``0`` otherwise, without
+    producing NaN gradients at ``x = 0``.
+
+    Args:
+        x: Input array.
+
+    Returns:
+        Elementwise square root, safe for autodiff.
+    """
+    safe_x = jnp.where(x > 0, x, 1.0)
+    return jnp.where(x > 0, jnp.sqrt(safe_x), 0.0)
+
+
+def arcsin(x: Array) -> Array:
+    """Autograd-safe arcsine.
+
+    Returns ``arcsin(x)`` for ``|x| < 1`` and ``+/-pi/2`` at the
+    boundary, without producing NaN gradients at ``x = +/-1``.
+
+    Args:
+        x: Input array with values in [-1, 1].
+
+    Returns:
+        Elementwise arcsine, safe for autodiff.
+    """
+    interior = jnp.abs(x) < 1
+    safe_x = jnp.where(interior, x, 0.0)
+    return jnp.where(interior, jnp.arcsin(safe_x), jnp.sign(x) * (jnp.pi / 2))
+
+
+def arccos(x: Array) -> Array:
+    """Autograd-safe arccosine.
+
+    Returns ``arccos(x)`` for ``|x| < 1``, ``0`` at ``x = 1``, and
+    ``pi`` at ``x = -1``, without producing NaN gradients at the
+    boundary.
+
+    Args:
+        x: Input array with values in [-1, 1].
+
+    Returns:
+        Elementwise arccosine, safe for autodiff.
+    """
+    interior = jnp.abs(x) < 1
+    safe_x = jnp.where(interior, x, 0.0)
+    return jnp.where(interior, jnp.arccos(safe_x), jnp.where(x >= 1, 0.0, jnp.pi))
+
+
+def div(x: Array, y: Array) -> Array:
+    """Autograd-safe division.
+
+    Returns ``x / y`` when ``y != 0`` and ``0`` otherwise, without
+    producing NaN gradients at ``y = 0``.
+
+    Args:
+        x: Numerator array.
+        y: Denominator array.
+
+    Returns:
+        Elementwise safe division.
+    """
+    nonzero = y != 0
+    safe_y = jnp.where(nonzero, y, 1.0)
+    return jnp.where(nonzero, x / safe_y, 0.0)
+
+
+def log(x: Array) -> Array:
+    """Autograd-safe natural logarithm.
+
+    Returns ``log(x)`` for ``x > 0`` and ``0`` otherwise, without
+    producing NaN gradients at ``x = 0``.
+
+    Args:
+        x: Input array.
+
+    Returns:
+        Elementwise natural logarithm, safe for autodiff.
+    """
+    safe_x = jnp.where(x > 0, x, 1.0)
+    return jnp.where(x > 0, jnp.log(safe_x), 0.0)
+
+
+def norm(x: Array, axis: int | None = None, keepdims: bool = False) -> Array:
+    """Autograd-safe L2 norm.
+
+    Computes ``sqrt(sum(x**2))`` using :func:`sqrt`, avoiding NaN
+    gradients when the norm is zero.
+
+    Args:
+        x: Input array.
+        axis: Axis or axes along which to compute the norm.
+        keepdims: If True, retains reduced axes with size 1.
+
+    Returns:
+        L2 norm along the given axis, safe for autodiff.
+    """
+    return sqrt(jnp.sum(x * x, axis=axis, keepdims=keepdims))

diffbio/core/soft_ops/comparison.py

@@ -0,0 +1,235 @@
+"""Soft comparison operators.
+
+Provides differentiable relaxations of elementwise comparison operations
+(greater, less, equal, etc.) returning SoftBool values in [0, 1].
+
+Each function uses :func:`~diffbio.core.soft_ops.elementwise.sigmoidal`
+as the underlying smooth step function, inheriting the multi-mode
+smoothness options.
+"""
+
+import jax.numpy as jnp
+from jax import Array
+
+from diffbio.core.soft_ops._types import SoftBool
+from diffbio.core.soft_ops._utils import ensure_float
+from diffbio.core.soft_ops.elementwise import Mode, abs, sigmoidal
+from diffbio.core.soft_ops.logical import logical_not
+
+
+def greater(
+    x: Array,
+    y: float | Array,
+    softness: float | Array = 0.1,
+    mode: Mode = "smooth",
+    epsilon: float = 1e-10,
+) -> SoftBool:
+    """Soft ``x > y``.
+
+    Uses sigmoidal on ``x - y - epsilon`` so the output approaches 0
+    at equality as softness -> 0.
+
+    Args:
+        x: First input array.
+        y: Second input array (broadcastable with x).
+        softness: Width of transition (> 0).
+        mode: ``"hard"`` or sigmoidal mode.
+        epsilon: Small offset for strict inequality at the limit.
+
+    Returns:
+        SoftBool in [0, 1].
+    """
+    x = ensure_float(x)
+    if mode == "hard":
+        return jnp.greater(x, y).astype(x.dtype)
+    return sigmoidal(x - y - epsilon, softness=softness, mode=mode)
+
+
+def greater_equal(
+    x: Array,
+    y: float | Array,
+    softness: float | Array = 0.1,
+    mode: Mode = "smooth",
+    epsilon: float = 1e-10,
+) -> SoftBool:
+    """Soft ``x >= y``.
+
+    Uses sigmoidal on ``x - y + epsilon`` so the output approaches 1
+    at equality as softness -> 0.
+
+    Args:
+        x: First input array.
+        y: Second input array.
+        softness: Width of transition (> 0).
+        mode: ``"hard"`` or sigmoidal mode.
+        epsilon: Small offset for non-strict inequality at the limit.
+
+    Returns:
+        SoftBool in [0, 1].
+    """
+    x = ensure_float(x)
+    if mode == "hard":
+        return jnp.greater_equal(x, y).astype(x.dtype)
+    return sigmoidal(x - y + epsilon, softness=softness, mode=mode)
+
+
+def less(
+    x: Array,
+    y: float | Array,
+    softness: float | Array = 0.1,
+    mode: Mode = "smooth",
+    epsilon: float = 1e-10,
+) -> SoftBool:
+    """Soft ``x < y``. Complement of :func:`greater_equal`.
+
+    Args:
+        x: First input array.
+        y: Second input array.
+        softness: Width of transition (> 0).
+        mode: ``"hard"`` or sigmoidal mode.
+        epsilon: Small offset.
+
+    Returns:
+        SoftBool in [0, 1].
+    """
+    x = ensure_float(x)
+    if mode == "hard":
+        return jnp.less(x, y).astype(x.dtype)
+    return logical_not(
+        greater_equal(x, y, softness=softness, mode=mode, epsilon=epsilon),
+    )
+
+
+def less_equal(
+    x: Array,
+    y: float | Array,
+    softness: float | Array = 0.1,
+    mode: Mode = "smooth",
+    epsilon: float = 1e-10,
+) -> SoftBool:
+    """Soft ``x <= y``. Complement of :func:`greater`.
+
+    Args:
+        x: First input array.
+        y: Second input array.
+        softness: Width of transition (> 0).
+        mode: ``"hard"`` or sigmoidal mode.
+        epsilon: Small offset.
+
+    Returns:
+        SoftBool in [0, 1].
+    """
+    x = ensure_float(x)
+    if mode == "hard":
+        return jnp.less_equal(x, y).astype(x.dtype)
+    return logical_not(
+        greater(x, y, softness=softness, mode=mode, epsilon=epsilon),
+    )
+
+
+def equal(
+    x: Array,
+    y: Array,
+    softness: float | Array = 0.1,
+    mode: Mode = "smooth",
+    epsilon: float = 1e-10,
+) -> SoftBool:
+    """Soft ``x == y``.
+
+    Implemented as soft ``abs(x - y) <= 0``, scaled to [0, 1].
+
+    Args:
+        x: First input array.
+        y: Second input array.
+        softness: Width of transition (> 0).
+        mode: ``"hard"`` or sigmoidal mode.
+        epsilon: Small offset.
+
+    Returns:
+        SoftBool in [0, 1].
+    """
+    x = ensure_float(x)
+    if mode == "hard":
+        return jnp.equal(x, y).astype(x.dtype)
+    diff = abs(x - y, softness=softness, mode=mode)
+    return 2.0 * less_equal(
+        diff,
+        jnp.zeros_like(diff),
+        mode=mode,
+        softness=softness,
+        epsilon=epsilon,
+    )
+
+
+def not_equal(
+    x: Array,
+    y: Array,
+    softness: float | Array = 0.1,
+    mode: Mode = "smooth",
+    epsilon: float = 1e-10,
+) -> SoftBool:
+    """Soft ``x != y``.
+
+    Implemented as soft ``abs(x - y) > 0``, scaled to [0, 1].
+
+    Args:
+        x: First input array.
+        y: Second input array.
+        softness: Width of transition (> 0).
+        mode: ``"hard"`` or sigmoidal mode.
+        epsilon: Small offset.
+
+    Returns:
+        SoftBool in [0, 1].
+    """
+    x = ensure_float(x)
+    if mode == "hard":
+        return jnp.not_equal(x, y).astype(x.dtype)
+    diff = abs(x - y, softness=softness, mode=mode)
+    tmp = greater(
+        diff,
+        jnp.zeros_like(diff),
+        mode=mode,
+        softness=softness,
+        epsilon=epsilon,
+    )
+    return 2.0 * tmp - 1.0
+
+
+def isclose(
+    x: Array,
+    y: Array,
+    softness: float | Array = 0.1,
+    rtol: float = 1e-05,
+    atol: float = 1e-08,
+    mode: Mode = "smooth",
+    epsilon: float = 1e-10,
+) -> SoftBool:
+    """Soft approximate equality.
+
+    Implements soft ``abs(x - y) <= atol + rtol * abs(y)``.
+
+    Args:
+        x: First input array.
+        y: Second input array.
+        softness: Width of transition (> 0).
+        rtol: Relative tolerance.
+        atol: Absolute tolerance.
+        mode: ``"hard"`` or sigmoidal mode.
+        epsilon: Small offset.
+
+    Returns:
+        SoftBool in [0, 1].
+    """
+    x = ensure_float(x)
+    if mode == "hard":
+        return jnp.isclose(x, y, atol=atol, rtol=rtol).astype(x.dtype)
+    diff = abs(x - y, softness=softness, mode=mode)
+    y_abs = abs(y, softness=softness, mode=mode)
+    return 2.0 * less_equal(
+        diff,
+        atol + rtol * y_abs,
+        mode=mode,
+        softness=softness,
+        epsilon=epsilon,
+    )