diffbio 0.1.0__py3-none-any.whl

diffbio/operators/assembly/gnn_assembly.py

@@ -0,0 +1,252 @@
+"""Graph Neural Network-based assembly navigator.
+
+This module provides a differentiable approach to assembly graph
+traversal using message passing neural networks.
+
+Key technique: Uses graph attention for message passing between nodes,
+then scores edges for soft path selection, enabling gradient flow
+through the assembly process.
+
+Applications: Differentiable genome assembly, assembly polishing,
+scaffolding optimization.
+
+Inherits from GraphOperator to get:
+
+- scatter_aggregate() for message aggregation
+- global_pool() for graph-level pooling
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import jax.numpy as jnp
+from flax import nnx
+from jaxtyping import Array, Float, Int, PyTree
+
+from diffbio.configs import TemperatureConfig
+from diffbio.core import soft_ops
+from diffbio.core.base_operators import GraphOperator
+from diffbio.core.gnn_components import GraphAttentionBlock
+from diffbio.utils.nn_utils import init_learnable_param
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class GNNAssemblyNavigatorConfig(TemperatureConfig):
+    """Configuration for GNNAssemblyNavigator.
+
+    Attributes:
+        node_features: Dimension of input node features.
+        hidden_dim: Hidden dimension for GNN layers.
+        num_layers: Number of GNN layers.
+        num_heads: Number of attention heads.
+        edge_features: Dimension of edge features.
+        dropout_rate: Dropout rate for regularization.
+        temperature: Temperature for softmax operations.
+    """
+
+    node_features: int = 64
+    hidden_dim: int = 128
+    num_layers: int = 3
+    num_heads: int = 4
+    edge_features: int = 8
+    dropout_rate: float = 0.1
+    temperature: float = 1.0
+
+
+class GNNAssemblyNavigator(GraphOperator):
+    """Graph Neural Network for assembly graph traversal.
+
+    This operator uses message passing with graph attention to update
+    node embeddings and predict edge traversal probabilities for
+    differentiable assembly.
+
+    Algorithm:
+    1. Project input node features to hidden dimension
+    2. Apply multiple GNN layers with graph attention
+    3. Compute edge scores from source/target node embeddings
+    4. Apply sigmoid for traversal probabilities
+    5. Compute path confidence from edge scores
+
+    Inherits from GraphOperator to get:
+
+    - scatter_aggregate() for message aggregation utilities
+    - global_pool() for graph-level pooling
+
+    Uses temperature-controlled smoothing:
+    - _temperature property for temperature-controlled sigmoid
+
+    Args:
+        config: GNNAssemblyNavigatorConfig with model parameters.
+        rngs: Flax NNX random number generators.
+        name: Optional operator name.
+
+    Example:
+        ```python
+        config = GNNAssemblyNavigatorConfig(hidden_dim=128)
+        navigator = GNNAssemblyNavigator(config, rngs=nnx.Rngs(42))
+        data = {"node_features": nodes, "edge_index": edges, "edge_features": edge_attr}
+        result, state, meta = navigator.apply(data, {}, None)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: GNNAssemblyNavigatorConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize the GNN assembly navigator.
+
+        Args:
+            config: Navigator 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.hidden_dim = config.hidden_dim
+
+        # Temperature management (similar to TemperatureOperator pattern)
+        if config.learnable_temperature:
+            self._temperature_param = init_learnable_param(config.temperature)
+        else:
+            self._temperature_param = None
+            self._fixed_temperature = config.temperature
+
+        # Input projection
+        self.input_projection = nnx.Linear(
+            in_features=config.node_features,
+            out_features=config.hidden_dim,
+            rngs=rngs,
+        )
+
+        # GNN layers
+        self.gnn_layers = nnx.List(
+            [
+                GraphAttentionBlock(
+                    hidden_dim=config.hidden_dim,
+                    num_heads=config.num_heads,
+                    edge_features=config.edge_features,
+                    dropout_rate=config.dropout_rate,
+                    rngs=rngs,
+                )
+                for _ in range(config.num_layers)
+            ]
+        )
+
+        # Edge scoring MLP
+        self.edge_mlp = nnx.Linear(
+            in_features=config.hidden_dim * 2,
+            out_features=1,
+            rngs=rngs,
+        )
+
+    @property
+    def _temperature(self) -> Array | float:
+        """Get current temperature value."""
+        if self._temperature_param is not None:
+            return jnp.abs(self._temperature_param[...]) + 1e-6
+        return self._fixed_temperature
+
+    def compute_edge_scores(
+        self,
+        node_embeddings: Float[Array, "n_nodes hidden_dim"],
+        edge_index: Int[Array, "2 n_edges"],
+    ) -> Float[Array, "n_edges"]:
+        """Compute edge scores from node embeddings.
+
+        Args:
+            node_embeddings: Node embedding matrix.
+            edge_index: Edge indices (source, target).
+
+        Returns:
+            Score for each edge.
+        """
+        sources = edge_index[0]
+        targets = edge_index[1]
+
+        # Get source and target embeddings
+        source_emb = node_embeddings[sources]  # (n_edges, hidden_dim)
+        target_emb = node_embeddings[targets]  # (n_edges, hidden_dim)
+
+        # Concatenate and score
+        edge_repr = jnp.concatenate([source_emb, target_emb], axis=-1)
+        scores = self.edge_mlp(edge_repr).squeeze(-1)  # (n_edges,)
+
+        return scores
+
+    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 GNN assembly navigation.
+
+        Args:
+            data: Dictionary containing:
+                - "node_features": Node features (n_nodes, node_features)
+                - "edge_index": Edge indices (2, n_edges)
+                - "edge_features": Edge features (n_edges, edge_features)
+            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:
+
+                    - "node_features": Original node features
+                    - "edge_index": Original edge indices
+                    - "edge_features": Original edge features
+                    - "node_embeddings": Updated node embeddings
+                    - "edge_scores": Scores for each edge
+                    - "traversal_probs": Sigmoid probabilities for traversal
+                    - "path_confidence": Confidence score for paths
+                - state is passed through unchanged
+                - metadata is passed through unchanged
+        """
+        node_features = data["node_features"]
+        edge_index = data["edge_index"]
+        edge_features = data["edge_features"]
+
+        # Project to hidden dimension
+        node_emb = self.input_projection(node_features)
+
+        # Apply GNN layers
+        deterministic = not self.config.stochastic
+        for layer in self.gnn_layers:
+            node_emb = layer(node_emb, edge_index, edge_features, deterministic=deterministic)
+
+        # Compute edge scores
+        edge_scores = self.compute_edge_scores(node_emb, edge_index)
+
+        # Traversal probabilities via sigmoid
+        # Use inherited _temperature property for temperature-controlled sigmoid
+        traversal_probs = soft_ops.greater(edge_scores, 0.0, softness=self._temperature)
+
+        # Path confidence: mean probability weighted by score magnitude
+        path_confidence = jnp.mean(traversal_probs * jnp.abs(edge_scores))
+
+        # Build output
+        transformed_data = {
+            "node_features": node_features,
+            "edge_index": edge_index,
+            "edge_features": edge_features,
+            "node_embeddings": node_emb,
+            "edge_scores": edge_scores,
+            "traversal_probs": traversal_probs,
+            "path_confidence": path_confidence,
+        }
+
+        return transformed_data, state, metadata

diffbio/operators/assembly/metagenomic_binning.py

@@ -0,0 +1,296 @@
+"""Differentiable metagenomic binning operators.
+
+This module provides VAE-based approaches to metagenomic binning
+inspired by VAMB (Variational Autoencoders for Metagenomic Binning).
+
+The approach encodes tetranucleotide frequencies (TNF) and abundance
+profiles into a latent space where contigs from the same genome cluster together.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+from artifex.generative_models.core.base import MLP
+from flax import nnx
+from jaxtyping import Array, Float
+
+from diffbio.configs import TemperatureConfig, apply_stochastic_sampling_defaults
+from diffbio.core.base_operators import EncoderDecoderOperator, TemperatureOperator
+from diffbio.utils.nn_utils import ARTIFEX_RELU_BATCH_NORM_MLP_KWARGS
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class MetagenomicBinnerConfig(TemperatureConfig):
+    """Configuration for metagenomic binning VAE.
+
+    Attributes:
+        n_tnf_features: Number of tetranucleotide frequency features (default 136).
+        n_abundance_features: Number of sample abundance features.
+        latent_dim: Dimension of the latent space.
+        hidden_dims: tuple of hidden layer dimensions for encoder/decoder.
+        dropout_rate: Dropout rate for regularization.
+        beta: KL divergence weight (beta-VAE).
+        n_clusters: Number of clusters for soft binning.
+    """
+
+    n_tnf_features: int = 136  # 4^4 / 2 for canonical k-mers
+    n_abundance_features: int = 10
+    latent_dim: int = 32
+    hidden_dims: tuple[int, ...] = (512, 256)
+    dropout_rate: float = 0.2
+    beta: float = 1.0
+    n_clusters: int = 100
+
+    def __post_init__(self) -> None:
+        """Set stochastic config and validate."""
+        apply_stochastic_sampling_defaults(self)
+        super().__post_init__()
+        if not self.hidden_dims:
+            raise ValueError(
+                "MetagenomicBinnerConfig.hidden_dims must contain at least one hidden layer."
+            )
+
+
+class DifferentiableMetagenomicBinner(TemperatureOperator, EncoderDecoderOperator):
+    """VAMB-style differentiable metagenomic binning.
+
+    This operator implements a Variational Autoencoder for metagenomic binning,
+    encoding tetranucleotide frequencies (TNF) and abundance profiles into a
+    shared latent space where contigs from the same genome cluster together.
+
+    The approach is fully differentiable, enabling:
+    - End-to-end optimization with downstream tasks
+    - Soft cluster assignments via temperature-controlled softmax
+    - Integration with neural abundance estimation
+
+    Input data structure:
+        - tnf: Float[Array, "n_contigs n_tnf"] - Tetranucleotide frequencies
+        - abundance: Float[Array, "n_contigs n_samples"] - Sample abundances
+
+    Output data structure (adds):
+        - latent_z: Float[Array, "n_contigs latent_dim"] - Latent representations
+        - latent_mu: Float[Array, "n_contigs latent_dim"] - Latent means
+        - latent_logvar: Float[Array, "n_contigs latent_dim"] - Latent log variance
+        - cluster_assignments: Float[Array, "n_contigs n_clusters"] - Soft bins
+        - reconstructed_tnf: Float[Array, "n_contigs n_tnf"] - Reconstructed TNF
+        - reconstructed_abundance: Float[Array, "n_contigs n_samples"] - Recon. abundance
+
+    Example:
+        ```python
+        config = MetagenomicBinnerConfig(n_abundance_features=5, n_clusters=50)
+        binner = DifferentiableMetagenomicBinner(config, rngs=nnx.Rngs(42))
+        result, state, meta = binner.apply(data, {}, None)
+        bins = result["cluster_assignments"].argmax(axis=-1)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: MetagenomicBinnerConfig,
+        *,
+        rngs: nnx.Rngs,
+        name: str | None = None,
+    ):
+        """Initialize the metagenomic binner.
+
+        Args:
+            config: Binner configuration.
+            rngs: Random number generators.
+            name: Optional name for the operator.
+        """
+        super().__init__(config, rngs=rngs, name=name)
+
+        input_dim = config.n_tnf_features + config.n_abundance_features
+
+        self.encoder_backbone = MLP(
+            hidden_dims=list(config.hidden_dims),
+            in_features=input_dim,
+            dropout_rate=config.dropout_rate,
+            rngs=rngs,
+            **ARTIFEX_RELU_BATCH_NORM_MLP_KWARGS,
+        )
+        self.fc_latent = nnx.List(
+            [
+                nnx.Linear(config.hidden_dims[-1], config.latent_dim, rngs=rngs),
+                nnx.Linear(config.hidden_dims[-1], config.latent_dim, rngs=rngs),
+            ]
+        )
+
+        decoder_hidden_dims = list(reversed(config.hidden_dims))
+        self.decoder_backbone = MLP(
+            hidden_dims=decoder_hidden_dims,
+            in_features=config.latent_dim,
+            dropout_rate=config.dropout_rate,
+            rngs=rngs,
+            **ARTIFEX_RELU_BATCH_NORM_MLP_KWARGS,
+        )
+        self.decoder_heads = nnx.List(
+            [
+                nnx.Linear(decoder_hidden_dims[-1], config.n_tnf_features, rngs=rngs),
+                nnx.Linear(decoder_hidden_dims[-1], config.n_abundance_features, rngs=rngs),
+            ]
+        )
+
+        # Tracks train/eval mode for latent sampling even when no stochastic
+        # submodule on the encoder is active for a given configuration.
+        self.latent_sampling_mode = nnx.Dropout(rate=0.0, rngs=rngs)
+
+        # Learnable cluster centroids
+        self.centroids = nnx.Param(
+            jax.random.normal(rngs.params(), (config.n_clusters, config.latent_dim)) * 0.1
+        )
+
+    def encode(
+        self, x: Float[Array, "batch input_dim"]
+    ) -> tuple[Float[Array, "batch latent"], Float[Array, "batch latent"]]:
+        """Encode input to latent distribution.
+
+        Args:
+            x: Concatenated TNF and abundance features.
+
+        Returns:
+            Tuple of (mu, logvar).
+        """
+        encoded = self.encoder_backbone(x)
+        if isinstance(encoded, tuple):
+            raise TypeError("Metagenomic binner encoder backbone must return a single tensor.")
+
+        fc_mu, fc_logvar = self.fc_latent
+        mu = fc_mu(encoded)
+        logvar = fc_logvar(encoded)
+        return mu, logvar
+
+    def decode(
+        self, z: Float[Array, "batch latent"]
+    ) -> tuple[Float[Array, "batch n_tnf"], Float[Array, "batch n_abundance"]]:
+        """Decode latent to reconstructed features.
+
+        Args:
+            z: Latent representation.
+
+        Returns:
+            Tuple of (tnf, abundance).
+        """
+        decoded = self.decoder_backbone(z)
+        if isinstance(decoded, tuple):
+            raise TypeError("Metagenomic binner decoder backbone must return a single tensor.")
+
+        fc_tnf, fc_abundance = self.decoder_heads
+        # TNF uses softmax (frequencies sum to 1)
+        tnf_recon = nnx.softmax(fc_tnf(decoded), axis=-1)
+        # Abundance uses softplus (positive values)
+        abundance_recon = nnx.softplus(fc_abundance(decoded))
+
+        return tnf_recon, abundance_recon
+
+    def soft_cluster(self, z: Float[Array, "batch latent"]) -> Float[Array, "batch n_clusters"]:
+        """Compute soft cluster assignments.
+
+        Args:
+            z: Latent representations.
+
+        Returns:
+            Soft cluster assignment probabilities.
+        """
+        # Compute squared distances to centroids
+        z_expanded = z[:, None, :]  # (batch, 1, latent)
+        centroids_expanded = self.centroids[...][None, :, :]  # (1, n_clusters, latent)
+        sq_distances = jnp.sum((z_expanded - centroids_expanded) ** 2, axis=-1)
+
+        # Soft assignment via softmax
+        temperature = jnp.maximum(self._temperature, 1e-6)
+        assignments = nnx.softmax(-sq_distances / temperature, axis=-1)
+        return assignments
+
+    def apply(
+        self,
+        data: dict[str, Array],
+        state: dict[str, Any],
+        metadata: dict[str, Any] | None,
+        random_params: Any = None,  # noqa: ARG002
+        stats: dict[str, Any] | None = None,  # noqa: ARG002
+    ) -> tuple[dict[str, Array], dict[str, Any], dict[str, Any] | None]:
+        """Apply metagenomic binning.
+
+        Args:
+            data: Input data containing:
+                - tnf: Float[Array, "n_contigs n_tnf"]
+                - abundance: Float[Array, "n_contigs n_samples"]
+            state: Element state (passed through).
+            metadata: Element metadata (passed through).
+            random_params: Random parameters.
+            stats: Optional statistics dict.
+
+        Returns:
+            Tuple of (output_data, state, metadata).
+        """
+        tnf = data["tnf"]
+        abundance = data["abundance"]
+
+        # Concatenate features
+        x = jnp.concatenate([tnf, abundance], axis=-1)
+
+        # Encode
+        mu, logvar = self.encode(x)
+
+        # Sample latent (use mu during eval for determinism)
+        if self.latent_sampling_mode.deterministic:
+            z = mu  # Eval mode: deterministic
+        else:
+            z = self.reparameterize(mu, logvar)  # Train mode: stochastic
+
+        # Decode
+        tnf_recon, abundance_recon = self.decode(z)
+
+        # Soft cluster assignments
+        cluster_assignments = self.soft_cluster(z)
+
+        # Build output
+        output_data = {
+            **data,
+            "latent_z": z,
+            "latent_mu": mu,
+            "latent_logvar": logvar,
+            "cluster_assignments": cluster_assignments,
+            "reconstructed_tnf": tnf_recon,
+            "reconstructed_abundance": abundance_recon,
+        }
+
+        return output_data, state, metadata
+
+
+def create_metagenomic_binner(
+    n_abundance_features: int = 10,
+    n_clusters: int = 100,
+    latent_dim: int = 32,
+    hidden_dims: tuple[int, ...] | None = None,
+    seed: int = 42,
+) -> DifferentiableMetagenomicBinner:
+    """Factory function to create a metagenomic binner.
+
+    Args:
+        n_abundance_features: Number of sample abundance features.
+        n_clusters: Number of clusters/bins.
+        latent_dim: Dimension of latent space.
+        hidden_dims: Hidden layer dimensions.
+        seed: Random seed.
+
+    Returns:
+        Configured DifferentiableMetagenomicBinner instance.
+    """
+    if hidden_dims is None:
+        hidden_dims = (512, 256)
+
+    config = MetagenomicBinnerConfig(
+        n_abundance_features=n_abundance_features,
+        n_clusters=n_clusters,
+        latent_dim=latent_dim,
+        hidden_dims=hidden_dims,
+    )
+    rngs = nnx.Rngs(seed)
+    return DifferentiableMetagenomicBinner(config, rngs=rngs)

diffbio/operators/crispr/__init__.py

@@ -0,0 +1,17 @@
+"""CRISPR guide design operators.
+
+This module provides differentiable operators for CRISPR guide RNA
+design and scoring, including on-target efficiency prediction.
+"""
+
+from diffbio.operators.crispr.guide_scoring import (
+    CRISPRScorerConfig,
+    DifferentiableCRISPRScorer,
+    create_crispr_scorer,
+)
+
+__all__ = [
+    "CRISPRScorerConfig",
+    "DifferentiableCRISPRScorer",
+    "create_crispr_scorer",
+]