diffbio 0.1.0__py3-none-any.whl

diffbio/core/graph_utils.py

@@ -0,0 +1,149 @@
+"""Graph utility functions for k-NN graph construction and similarity computation.
+
+Provides reusable, differentiable functions for pairwise distance computation,
+k-nearest-neighbor graph construction, fuzzy set membership, and graph
+symmetrization. These primitives underpin UMAP, trajectory inference,
+imputation, and other graph-based operators.
+"""
+
+import jax
+import jax.numpy as jnp
+
+from diffbio.core import soft_ops
+
+__all__ = [
+    "compute_pairwise_distances",
+    "compute_knn_graph",
+    "compute_fuzzy_membership",
+    "symmetrize_graph",
+]
+
+
+def compute_pairwise_distances(
+    features: jax.Array,
+    metric: str = "euclidean",
+) -> jax.Array:
+    """Compute pairwise distance matrix between all samples.
+
+    Args:
+        features: Input feature matrix of shape ``(n_samples, n_features)``.
+        metric: Distance metric, either ``"euclidean"`` or ``"cosine"``.
+
+    Returns:
+        Distance matrix of shape ``(n_samples, n_samples)`` where entry
+        ``(i, j)`` is the distance from sample *i* to sample *j*.
+    """
+    if metric == "cosine":
+        norms = jnp.linalg.norm(features, axis=-1, keepdims=True)
+        features_norm = features / jnp.maximum(norms, jnp.finfo(features.dtype).eps)
+        similarity = jnp.dot(features_norm, features_norm.T)
+        distances = jnp.clip(1.0 - similarity, 0.0, 2.0)
+    else:
+        diff = features[:, None, :] - features[None, :, :]
+        distances = jnp.sqrt(jnp.sum(diff**2, axis=-1) + 1e-16)
+
+    # Self-distance is zero by definition; float32 matmul cannot guarantee this
+    n = features.shape[0]
+    distances = distances.at[jnp.diag_indices(n)].set(0.0)
+
+    return distances
+
+
+def compute_knn_graph(
+    distances: jax.Array,
+    k: int,
+) -> tuple[jax.Array, jax.Array]:
+    """Build a k-nearest-neighbor graph from a dense distance matrix.
+
+    For each node the *k* closest neighbours (by distance) are selected.
+    Self-connections are assumed to already be masked out by setting the
+    diagonal to a large value before calling this function.
+
+    Args:
+        distances: Dense distance matrix of shape ``(n, n)``.  The diagonal
+            should contain large sentinel values (e.g. ``DISTANCE_MASK_SENTINEL``) so that
+            self-loops are never selected.
+        k: Number of nearest neighbours per node.  Clipped to ``n - 1``
+            when larger than the number of samples minus one.
+
+    Returns:
+        A tuple ``(edge_indices, edge_weights)`` where:
+
+        - ``edge_indices`` has shape ``(n * k_eff, 2)`` with each row
+            ``[source, target]``.
+        - ``edge_weights`` has shape ``(n * k_eff,)`` containing the
+            corresponding distances.
+
+        ``k_eff = min(k, n - 1)``.
+    """
+    n = distances.shape[0]
+    k_eff = min(k, n - 1)
+
+    # Argsort each row; first k_eff entries are the nearest neighbours
+    sorted_indices = jnp.argsort(distances, axis=-1)
+    knn_indices = sorted_indices[:, :k_eff]  # (n, k_eff)
+
+    # Source indices: each node repeated k_eff times
+    sources = jnp.repeat(jnp.arange(n), k_eff)  # (n * k_eff,)
+    targets = knn_indices.reshape(-1)  # (n * k_eff,)
+
+    edge_indices = jnp.stack([sources, targets], axis=-1)  # (n * k_eff, 2)
+    edge_weights = distances[sources, targets]  # (n * k_eff,)
+
+    return edge_indices, edge_weights
+
+
+def compute_fuzzy_membership(
+    distances: jax.Array,
+    k: int,
+    softness: float = 0.1,
+) -> jax.Array:
+    """Compute fuzzy set membership using a Gaussian kernel with local bandwidth.
+
+    The bandwidth (sigma) for each sample is set to the distance to its *k*-th
+    nearest neighbour, making the kernel adapt to local density.  The diagonal
+    of the output is forced to zero (no self-similarity).
+
+    Args:
+        distances: Dense distance matrix of shape ``(n, n)``.  The diagonal
+            should contain large sentinel values so that self-distances are
+            excluded from the bandwidth computation.
+        k: Number of neighbours used to determine the local bandwidth.
+            Clipped to ``n - 1`` when larger.
+
+    Returns:
+        Fuzzy membership matrix of shape ``(n, n)`` with values in ``[0, 1]``.
+    """
+    n = distances.shape[0]
+    k_eff = min(k, n - 1)
+
+    # Local bandwidth: distance to the k-th nearest neighbour
+    sorted_dists = soft_ops.sort(distances, axis=-1, softness=softness)
+    sigma = sorted_dists[:, k_eff - 1 : k_eff].squeeze(-1)
+    sigma = jnp.maximum(sigma, 1e-8)
+
+    # Gaussian kernel with local bandwidth
+    p_ij = jnp.exp(-distances / sigma[:, None])
+
+    # Zero out self-similarity
+    p_ij = p_ij * (1.0 - jnp.eye(n))
+
+    return p_ij
+
+
+def symmetrize_graph(adjacency: jax.Array) -> jax.Array:
+    """Symmetrize a directed adjacency matrix via fuzzy set union.
+
+    Applies the probabilistic (fuzzy) union:
+    ``p_sym = p + p^T - p * p^T``
+
+    This ensures the output is symmetric and, when inputs are in ``[0, 1]``,
+    the outputs remain in ``[0, 1]``.
+
+    Args:
+        adjacency: Directed adjacency / membership matrix of shape ``(n, n)``.
+
+    Returns:
+        Symmetric adjacency matrix of shape ``(n, n)``.
+    """
+    return adjacency + adjacency.T - adjacency * adjacency.T

diffbio/core/neural_components.py

@@ -0,0 +1,270 @@
+"""Reusable neural network components for DiffBio.
+
+This module provides neural network building blocks that are specific to
+bioinformatics applications and not available in Flax NNX built-ins.
+
+IMPORTANT: For standard components, use Flax NNX built-ins:
+- nnx.MultiHeadAttention for attention
+- nnx.Linear for dense layers
+- nnx.Conv for convolutions
+- nnx.LayerNorm for normalization
+- nnx.Dropout for dropout
+- nnx.Sequential for layer composition
+
+For reusable components, import from artifex:
+- PositionalEncoding (sinusoidal, learned, RoPE)
+- ResidualBlock1D, ResidualBlock2D
+- TransformerBlock
+- Various loss functions
+
+This module only provides DiffBio-specific components:
+- GumbelSoftmaxModule: Differentiable discrete sampling
+- GraphMessagePassing: GNN message passing for graph-structured data
+"""
+
+from typing import Literal
+
+import jax
+import jax.numpy as jnp
+from artifex.generative_models.core.base import MLP
+from flax import nnx
+from jaxtyping import Array, Float, Int
+
+from diffbio.constants import DEFAULT_TEMPERATURE, EPSILON
+from diffbio.utils.nn_utils import get_rng_key
+
+# =============================================================================
+# Re-export from artifex (import when available, provide stubs otherwise)
+# =============================================================================
+
+from artifex.generative_models.core.layers.positional import (
+    PositionalEncoding,
+    RotaryPositionalEncoding as RoPE,
+    SinusoidalPositionalEncoding,
+)
+from artifex.generative_models.core.layers.residual import (
+    Conv1DResidualBlock as ResidualBlock1D,
+    Conv2DResidualBlock as ResidualBlock2D,
+)
+
+
+__all__ = [
+    # DiffBio-specific components
+    "GumbelSoftmaxModule",
+    "GraphMessagePassing",
+    # Re-exported from artifex
+    "PositionalEncoding",
+    "SinusoidalPositionalEncoding",
+    "RoPE",
+    "ResidualBlock1D",
+    "ResidualBlock2D",
+]
+
+
+class GumbelSoftmaxModule(nnx.Module):
+    """Neural network module for Gumbel-softmax sampling.
+
+    This module wraps the gumbel_softmax function as an nnx.Module,
+    providing differentiable categorical sampling during forward pass.
+
+    Useful for:
+
+    - Discrete latent variable models (VQ-VAE variants)
+    - Hard attention mechanisms
+    - Discrete sequence generation
+
+    Args:
+        temperature: Initial temperature for sampling.
+        hard: If True, use straight-through estimator for discrete samples.
+        rngs: Flax NNX random number generators.
+
+    Example:
+        ```python
+        module = GumbelSoftmaxModule(temperature=0.5, rngs=nnx.Rngs(42))
+        logits = jnp.array([[1.0, 2.0, 3.0]])
+        samples = module(logits)
+        ```
+    """
+
+    def __init__(
+        self,
+        temperature: float = DEFAULT_TEMPERATURE,
+        hard: bool = False,
+        *,
+        rngs: nnx.Rngs,
+    ):
+        """Initialize GumbelSoftmaxModule.
+
+        Args:
+            temperature: Temperature for Gumbel-softmax.
+            hard: Whether to use hard (one-hot) samples.
+            rngs: Random number generators.
+        """
+        super().__init__()
+        self.temperature = temperature
+        self.hard = hard
+        self.rngs = rngs
+
+    def __call__(self, logits: Float[Array, "... n"]) -> Float[Array, "... n"]:
+        """Apply Gumbel-softmax sampling.
+
+        Args:
+            logits: Unnormalized log-probabilities of shape (..., n).
+
+        Returns:
+            Samples of same shape as logits.
+        """
+        key = get_rng_key(self.rngs, "dropout", fallback_seed=0)
+        gumbel_noise = jax.random.gumbel(key, logits.shape)
+        perturbed = (logits + gumbel_noise) / self.temperature
+        soft_sample = jax.nn.softmax(perturbed, axis=-1)
+        if self.hard:
+            hard_sample = jax.nn.one_hot(
+                jnp.argmax(soft_sample, axis=-1),
+                logits.shape[-1],
+            )
+            return hard_sample - jax.lax.stop_gradient(soft_sample) + soft_sample
+        return soft_sample
+
+
+class GraphMessagePassing(nnx.Module):
+    """Graph neural network message passing layer.
+
+    Implements a standard message passing scheme:
+    1. Compute messages from source nodes and edge features
+    2. Aggregate messages at destination nodes
+    3. Update node features with aggregated messages
+
+    Supports different aggregation functions (sum, mean, max).
+
+    Args:
+        node_features: Input node feature dimension.
+        edge_features: Edge feature dimension.
+        hidden_dim: Output hidden dimension.
+        aggregation: Aggregation function ("sum", "mean", "max").
+        rngs: Flax NNX random number generators.
+
+    Example:
+        ```python
+        layer = GraphMessagePassing(
+            node_features=32, edge_features=8, hidden_dim=64,
+            rngs=nnx.Rngs(42)
+        )
+        node_feat = jnp.ones((5, 32))  # 5 nodes
+        edge_feat = jnp.ones((8, 8))   # 8 edges
+        edge_index = jnp.array([[0,0,1,1,2,2,3,4], [1,2,2,3,3,4,4,0]])
+        output = layer(node_feat, edge_feat, edge_index)
+        ```
+    """
+
+    def __init__(
+        self,
+        node_features: int,
+        edge_features: int,
+        hidden_dim: int,
+        aggregation: Literal["sum", "mean", "max"] = "sum",
+        *,
+        rngs: nnx.Rngs,
+    ):
+        """Initialize GraphMessagePassing layer.
+
+        Args:
+            node_features: Input node feature dimension.
+            edge_features: Edge feature dimension.
+            hidden_dim: Output dimension.
+            aggregation: Aggregation method.
+            rngs: Random number generators.
+        """
+        super().__init__()
+        if node_features <= 0:
+            raise ValueError(f"node_features must be positive, got {node_features}")
+        if edge_features <= 0:
+            raise ValueError(f"edge_features must be positive, got {edge_features}")
+        if hidden_dim <= 0:
+            raise ValueError(f"hidden_dim must be positive, got {hidden_dim}")
+        if aggregation not in {"sum", "mean", "max"}:
+            raise ValueError(f"Unknown aggregation: {aggregation}")
+
+        self.hidden_dim = nnx.static(hidden_dim)
+        self.aggregation = nnx.static(aggregation)
+        self.message_mlp = MLP(
+            [hidden_dim, hidden_dim],
+            in_features=node_features + edge_features,
+            activation=nnx.relu,
+            rngs=rngs,
+        )
+        self.update_mlp = MLP(
+            [hidden_dim, hidden_dim],
+            in_features=node_features + hidden_dim,
+            activation=nnx.relu,
+            rngs=rngs,
+        )
+
+    def __call__(
+        self,
+        node_features: Float[Array, "num_nodes node_feat"],
+        edge_features: Float[Array, "num_edges edge_feat"],
+        edge_index: Int[Array, "2 num_edges"],
+    ) -> Float[Array, "num_nodes hidden_dim"]:
+        """Apply message passing.
+
+        Args:
+            node_features: Node feature matrix (num_nodes, node_features).
+            edge_features: Edge feature matrix (num_edges, edge_features).
+            edge_index: Edge indices [source, dest] of shape (2, num_edges).
+
+        Returns:
+            Updated node features (num_nodes, hidden_dim).
+        """
+        num_nodes = node_features.shape[0]
+        num_edges = edge_index.shape[1]
+
+        # Handle empty graph case
+        if num_edges == 0:
+            # No messages, just transform node features
+            update_output = self.update_mlp(
+                jnp.concatenate([node_features, jnp.zeros((num_nodes, self.hidden_dim))], axis=-1)
+            )
+            if isinstance(update_output, tuple):
+                return update_output[0]
+            return update_output
+
+        # Extract source and destination node indices
+        source_idx = edge_index[0]
+        dest_idx = edge_index[1]
+
+        # Get source node features for each edge
+        source_features = node_features[source_idx]  # (num_edges, node_feat)
+
+        # Compute messages: MLP(concat(source_features, edge_features))
+        message_input = jnp.concatenate([source_features, edge_features], axis=-1)
+        messages = self.message_mlp(message_input)
+        if isinstance(messages, tuple):
+            messages = messages[0]
+
+        # Aggregate messages at destination nodes
+        if self.aggregation == "sum":
+            aggregated = jax.ops.segment_sum(messages, dest_idx, num_segments=num_nodes)
+        elif self.aggregation == "mean":
+            sum_messages = jax.ops.segment_sum(messages, dest_idx, num_segments=num_nodes)
+            counts = jax.ops.segment_sum(jnp.ones(num_edges), dest_idx, num_segments=num_nodes)
+            aggregated = sum_messages / (counts[:, None] + EPSILON)
+        elif self.aggregation == "max":
+            # segment_max with default of -inf for empty segments
+            aggregated = jax.ops.segment_max(
+                messages,
+                dest_idx,
+                num_segments=num_nodes,
+                indices_are_sorted=False,
+            )
+            # Replace -inf with 0 for nodes with no incoming edges
+            aggregated = jnp.where(jnp.isinf(aggregated), jnp.zeros_like(aggregated), aggregated)
+        else:
+            raise ValueError(f"Unknown aggregation: {self.aggregation}")
+
+        # Update node features: MLP(concat(node_features, aggregated))
+        update_input = jnp.concatenate([node_features, aggregated], axis=-1)
+        updated = self.update_mlp(update_input)
+        if isinstance(updated, tuple):
+            return updated[0]
+        return updated

diffbio/core/optimal_transport.py

@@ -0,0 +1,133 @@
+"""Optimal transport layers for differentiable assignment and matching.
+
+Ownership note: DiffBio retains this transport-plan layer because downstream
+operators need a differentiable marginal-constrained plan module. Calibrax
+currently exposes scalar Sinkhorn-style metrics, not this operator contract.
+
+This module provides differentiable optimal transport solvers using the
+Sinkhorn algorithm in log-domain for numerical stability.
+
+Components:
+
+- **SinkhornLayer**: Computes the entropy-regularised optimal transport plan
+  between two discrete distributions given a cost matrix, using the
+  Sinkhorn-Knopp algorithm in log-domain.
+"""
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from jaxtyping import Array, Float
+
+__all__ = [
+    "SinkhornLayer",
+]
+
+
+class SinkhornLayer(nnx.Module):
+    """Sinkhorn optimal transport layer (log-domain).
+
+    Computes the entropy-regularised optimal transport plan between two
+    discrete marginal distributions ``a`` and ``b`` given a cost matrix ``C``,
+    by solving::
+
+        min_{P >= 0}  <P, C> - epsilon * H(P)
+        s.t.  P @ 1 = a,   P^T @ 1 = b
+
+    The algorithm runs in log-domain for numerical stability::
+
+        f, g = 0, 0
+        for _ in range(num_iters):
+            f = epsilon * log(a) - epsilon * logsumexp((-C + g) / epsilon, axis=1)
+            g = epsilon * log(b) - epsilon * logsumexp((-C + f) / epsilon, axis=0)
+        P = exp((f[:, None] + g[None, :] - C) / epsilon)
+
+    Args:
+        epsilon: Entropy regularisation strength (larger = smoother plan).
+        num_iters: Number of Sinkhorn iterations.
+        rngs: Flax NNX random number generators (unused, kept for API consistency).
+    """
+
+    def __init__(
+        self,
+        epsilon: float,
+        num_iters: int,
+        *,
+        rngs: nnx.Rngs,
+    ) -> None:
+        """Initialize the Sinkhorn layer.
+
+        Args:
+            epsilon: Regularisation strength.
+            num_iters: Number of Sinkhorn iterations.
+            rngs: Random number generators (for API consistency).
+        """
+        super().__init__()
+        self.epsilon = nnx.static(epsilon)
+        self.num_iters = nnx.static(num_iters)
+
+    def __call__(
+        self,
+        cost: Float[Array, "n m"],
+        a: Float[Array, " n"],
+        b: Float[Array, " m"],
+    ) -> Float[Array, "n m"]:
+        """Compute the optimal transport plan.
+
+        Args:
+            cost: Cost matrix of shape ``(n, m)``.
+            a: Source marginal distribution of shape ``(n,)``, must sum to 1.
+            b: Target marginal distribution of shape ``(m,)``, must sum to 1.
+
+        Returns:
+            Transport plan of shape ``(n, m)`` satisfying (approximately)
+            ``P @ 1 = a`` and ``P^T @ 1 = b``.
+        """
+        return _sinkhorn_log_domain(cost, a, b, self.epsilon, self.num_iters)
+
+
+def _sinkhorn_log_domain(
+    cost: Float[Array, "n m"],
+    a: Float[Array, " n"],
+    b: Float[Array, " m"],
+    epsilon: float,
+    num_iters: int,
+) -> Float[Array, "n m"]:
+    """Run the Sinkhorn algorithm in log-domain.
+
+    Args:
+        cost: Cost matrix ``(n, m)``.
+        a: Source marginal ``(n,)``.
+        b: Target marginal ``(m,)``.
+        epsilon: Regularisation parameter.
+        num_iters: Number of iterations.
+
+    Returns:
+        Transport plan ``(n, m)``.
+    """
+    log_a = jnp.log(a + 1e-30)
+    log_b = jnp.log(b + 1e-30)
+
+    f = jnp.zeros_like(a)
+    g = jnp.zeros_like(b)
+
+    def _step(
+        carry: tuple[Float[Array, " n"], Float[Array, " m"]], _: None
+    ) -> tuple[tuple[Float[Array, " n"], Float[Array, " m"]], None]:
+        """Perform one Sinkhorn iteration updating dual variables f and g."""
+        f_prev, g_prev = carry
+        # f update: epsilon * log(a) - epsilon * logsumexp((-C + g) / epsilon, axis=1)
+        f_new = epsilon * log_a - epsilon * jax.scipy.special.logsumexp(
+            (-cost + g_prev[None, :]) / epsilon, axis=1
+        )
+        # g update: epsilon * log(b) - epsilon * logsumexp((-C + f) / epsilon, axis=0)
+        g_new = epsilon * log_b - epsilon * jax.scipy.special.logsumexp(
+            (-cost + f_new[:, None]) / epsilon, axis=0
+        )
+        return (f_new, g_new), None
+
+    (f, g), _ = jax.lax.scan(_step, (f, g), None, length=num_iters)
+
+    # Recover the transport plan
+    log_plan = (f[:, None] + g[None, :] - cost) / epsilon
+    return jnp.exp(log_plan)