diffbio 0.1.0__py3-none-any.whl

diffbio/operators/multiomics/hic_contact.py

@@ -0,0 +1,377 @@
+"""Hi-C chromatin contact analysis operator.
+
+This module provides differentiable analysis of Hi-C contact matrices
+for chromatin structure inference.
+
+Key technique: Uses neural network to learn bin embeddings from contact
+patterns, then predicts compartments and TAD boundaries using attention
+over neighboring bins.
+
+Applications: Chromatin compartment identification, TAD boundary detection,
+3D genome structure prediction.
+
+Inherits from TemperatureOperator to get:
+
+- _temperature property for temperature-controlled smoothing
+- soft_max() for logsumexp-based smooth maximum
+- soft_argmax() for soft position selection
+"""
+
+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 datarax.core.config import OperatorConfig
+from flax import nnx
+from jaxtyping import Array, Float, PyTree
+
+from diffbio.core.base_operators import TemperatureOperator
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class HiCContactAnalysisConfig(OperatorConfig):
+    """Configuration for HiCContactAnalysis.
+
+    Attributes:
+        n_bins: Number of genomic bins.
+        hidden_dim: Hidden dimension for neural networks.
+        num_layers: Number of encoder layers.
+        num_heads: Number of attention heads.
+        bin_features: Dimension of input bin features.
+        dropout_rate: Dropout rate for regularization.
+        temperature: Temperature for softmax operations.
+    """
+
+    n_bins: int = 1000
+    hidden_dim: int = 128
+    num_layers: int = 3
+    num_heads: int = 4
+    bin_features: int = 16
+    dropout_rate: float = 0.1
+    temperature: float = 1.0
+
+    def __post_init__(self) -> None:
+        """Validate Hi-C contact analysis configuration."""
+        super().__post_init__()
+        if self.num_layers < 1:
+            raise ValueError("HiCContactAnalysisConfig.num_layers must be at least 1.")
+
+
+class ContactEncoder(nnx.Module):
+    """Encoder for Hi-C contact patterns."""
+
+    def __init__(
+        self,
+        n_bins: int,
+        hidden_dim: int,
+        num_layers: int,
+        *,
+        rngs: nnx.Rngs,
+    ):
+        """Initialize the contact encoder.
+
+        Args:
+            n_bins: Number of bins.
+            hidden_dim: Hidden dimension.
+            num_layers: Number of layers.
+            rngs: Random number generators.
+        """
+        super().__init__()
+        self.backbone = MLP(
+            hidden_dims=[hidden_dim] * num_layers,
+            in_features=n_bins,
+            activation="gelu",
+            output_activation="gelu",
+            use_batch_norm=False,
+            rngs=rngs,
+        )
+
+    def __call__(
+        self,
+        contact_matrix: Float[Array, "n_bins n_bins"],
+    ) -> Float[Array, "n_bins hidden_dim"]:
+        """Encode contact patterns.
+
+        Args:
+            contact_matrix: Hi-C contact matrix.
+
+        Returns:
+            Bin embeddings from contact patterns.
+        """
+        backbone_output = self.backbone(contact_matrix)
+        if isinstance(backbone_output, tuple):
+            raise TypeError("Hi-C contact encoder backbone must return a single tensor.")
+        return backbone_output
+
+
+class BinFeatureEncoder(nnx.Module):
+    """Encoder for genomic bin features."""
+
+    def __init__(
+        self,
+        bin_features: int,
+        hidden_dim: int,
+        *,
+        rngs: nnx.Rngs,
+    ):
+        """Initialize the bin feature encoder.
+
+        Args:
+            bin_features: Input feature dimension.
+            hidden_dim: Hidden dimension.
+            rngs: Random number generators.
+        """
+        super().__init__()
+        self.backbone = MLP(
+            hidden_dims=[hidden_dim, hidden_dim],
+            in_features=bin_features,
+            activation="gelu",
+            output_activation=None,
+            use_batch_norm=False,
+            rngs=rngs,
+        )
+
+    def __call__(
+        self,
+        bin_features: Float[Array, "n_bins bin_features"],
+    ) -> Float[Array, "n_bins hidden_dim"]:
+        """Encode bin features.
+
+        Args:
+            bin_features: Genomic bin features.
+
+        Returns:
+            Bin feature embeddings.
+        """
+        backbone_output = self.backbone(bin_features)
+        if isinstance(backbone_output, tuple):
+            raise TypeError("Hi-C feature encoder backbone must return a single tensor.")
+        return backbone_output
+
+
+class LocalAttention(nnx.Module):
+    """Local attention for TAD boundary detection."""
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        num_heads: int,
+        *,
+        rngs: nnx.Rngs,
+    ):
+        """Initialize local attention.
+
+        Args:
+            hidden_dim: Hidden dimension.
+            num_heads: Number of attention heads.
+            rngs: Random number generators.
+        """
+        super().__init__()
+
+        self.num_heads = num_heads
+        self.head_dim = hidden_dim // num_heads
+        self.scale = self.head_dim**-0.5
+
+        self.query = nnx.Linear(in_features=hidden_dim, out_features=hidden_dim, rngs=rngs)
+        self.key = nnx.Linear(in_features=hidden_dim, out_features=hidden_dim, rngs=rngs)
+        self.value = nnx.Linear(in_features=hidden_dim, out_features=hidden_dim, rngs=rngs)
+        self.output = nnx.Linear(in_features=hidden_dim, out_features=hidden_dim, rngs=rngs)
+
+    def __call__(
+        self,
+        x: Float[Array, "n_bins hidden_dim"],
+    ) -> Float[Array, "n_bins hidden_dim"]:
+        """Apply local attention.
+
+        Args:
+            x: Input embeddings.
+
+        Returns:
+            Attended embeddings.
+        """
+        n_bins = x.shape[0]
+
+        Q = self.query(x).reshape(n_bins, self.num_heads, self.head_dim)
+        K = self.key(x).reshape(n_bins, self.num_heads, self.head_dim)
+        V = self.value(x).reshape(n_bins, self.num_heads, self.head_dim)
+
+        # Full attention (could be windowed for efficiency)
+        attn = jnp.einsum("qhd,khd->hqk", Q, K) * self.scale
+        attn = jax.nn.softmax(attn, axis=-1)
+
+        out = jnp.einsum("hqk,khd->qhd", attn, V)
+        out = out.reshape(n_bins, -1)
+        out = self.output(out)
+
+        return out
+
+
+class HiCContactAnalysis(TemperatureOperator):
+    """Differentiable Hi-C contact analysis.
+
+    This operator analyzes Hi-C contact matrices to identify
+    chromatin compartments and TAD boundaries using neural networks.
+
+    Algorithm:
+    1. Encode contact patterns per bin
+    2. Encode genomic bin features
+    3. Combine contact and feature embeddings
+    4. Apply attention for context
+    5. Predict compartment scores
+    6. Detect TAD boundaries
+    7. Reconstruct contacts from embeddings
+
+    Args:
+        config: HiCContactAnalysisConfig with model parameters.
+        rngs: Flax NNX random number generators.
+        name: Optional operator name.
+
+    Example:
+        ```python
+        config = HiCContactAnalysisConfig(n_bins=1000)
+        analyzer = HiCContactAnalysis(config, rngs=nnx.Rngs(42))
+        data = {"contact_matrix": contacts, "bin_features": features}
+        result, state, meta = analyzer.apply(data, {}, None)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: HiCContactAnalysisConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize the Hi-C contact analyzer.
+
+        Args:
+            config: Analysis 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 is managed by TemperatureOperator via self._temperature
+
+        # Contact pattern encoder
+        self.contact_encoder = ContactEncoder(
+            n_bins=config.n_bins,
+            hidden_dim=config.hidden_dim,
+            num_layers=config.num_layers,
+            rngs=rngs,
+        )
+
+        # Bin feature encoder
+        self.feature_encoder = BinFeatureEncoder(
+            bin_features=config.bin_features,
+            hidden_dim=config.hidden_dim,
+            rngs=rngs,
+        )
+
+        # Combine contact and feature embeddings
+        self.combine = nnx.Linear(
+            in_features=config.hidden_dim * 2,
+            out_features=config.hidden_dim,
+            rngs=rngs,
+        )
+
+        # Attention for context
+        self.attention = LocalAttention(
+            hidden_dim=config.hidden_dim,
+            num_heads=config.num_heads,
+            rngs=rngs,
+        )
+
+        # Output heads
+        self.compartment_head = nnx.Linear(
+            in_features=config.hidden_dim,
+            out_features=1,
+            rngs=rngs,
+        )
+
+        self.boundary_head = nnx.Linear(
+            in_features=config.hidden_dim,
+            out_features=1,
+            rngs=rngs,
+        )
+
+    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 Hi-C contact analysis.
+
+        Args:
+            data: Dictionary containing:
+                - "contact_matrix": Hi-C contact matrix (n_bins, n_bins)
+                - "bin_features": Bin genomic features (n_bins, bin_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:
+
+                    - "contact_matrix": Original contacts
+                    - "bin_features": Original features
+                    - "bin_embeddings": Learned bin embeddings
+                    - "compartment_scores": A/B compartment scores
+                    - "tad_boundary_scores": TAD boundary probabilities
+                    - "predicted_contacts": Reconstructed contacts
+                - state is passed through unchanged
+                - metadata is passed through unchanged
+        """
+        contact_matrix = data["contact_matrix"]
+        bin_features = data["bin_features"]
+
+        # Encode contact patterns
+        contact_emb = self.contact_encoder(contact_matrix)  # (n_bins, hidden_dim)
+
+        # Encode bin features
+        feature_emb = self.feature_encoder(bin_features)  # (n_bins, hidden_dim)
+
+        # Combine embeddings
+        combined = jnp.concatenate([contact_emb, feature_emb], axis=-1)
+        combined = nnx.gelu(self.combine(combined))  # (n_bins, hidden_dim)
+
+        # Apply attention for context
+        attended = self.attention(combined) + combined  # Residual
+
+        # Predict compartment scores (continuous A/B)
+        compartment_scores = self.compartment_head(attended).squeeze(-1)  # (n_bins,)
+
+        # Predict TAD boundary scores
+        boundary_logits = self.boundary_head(attended).squeeze(-1)  # (n_bins,)
+        tad_boundary_scores = jax.nn.sigmoid(boundary_logits)
+
+        # Reconstruct contacts from embeddings (dot product)
+        # (n_bins, hidden_dim) @ (hidden_dim, n_bins) -> (n_bins, n_bins)
+        predicted_contacts = jnp.einsum("ih,jh->ij", attended, attended)
+        predicted_contacts = jax.nn.softplus(predicted_contacts)  # Non-negative
+
+        # Build output
+        transformed_data = {
+            "contact_matrix": contact_matrix,
+            "bin_features": bin_features,
+            "bin_embeddings": attended,
+            "compartment_scores": compartment_scores,
+            "tad_boundary_scores": tad_boundary_scores,
+            "predicted_contacts": predicted_contacts,
+        }
+
+        return transformed_data, state, metadata

diffbio/operators/multiomics/multiomics_vae.py

@@ -0,0 +1,325 @@
+"""Multi-omics VAE with Product-of-Experts fusion.
+
+This module implements a differentiable multi-omics variational autoencoder
+that jointly integrates data from multiple modalities (e.g., RNA-seq, ATAC-seq,
+protein). The model learns a shared latent space via Product-of-Experts (PoE)
+fusion of per-modality posterior distributions, following the approach
+described in MULTIVI (Ashuach et al., 2023).
+
+Key algorithm:
+    1. Per-modality encoders map counts to (mu_m, logvar_m).
+    2. PoE fuses posteriors:
+       precision_joint = sum(1 / sigma_m^2)
+       mu_joint = (sum mu_m / sigma_m^2) / precision_joint
+    3. Reparameterised sample z ~ N(mu_joint, 1/precision_joint).
+    4. Per-modality decoders reconstruct counts from z.
+    5. ELBO = sum_m w_m * recon_loss_m + KL(q(z) || N(0,I)).
+"""
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+from artifex.generative_models.core.base import MLP
+from artifex.generative_models.core.losses.base import reduce_loss
+from artifex.generative_models.core.losses.divergence import gaussian_kl_divergence
+from datarax.core.config import OperatorConfig
+from flax import nnx
+from jaxtyping import Array, Float, PyTree
+
+from diffbio.constants import EPSILON
+from diffbio.core.base_operators import EncoderDecoderOperator
+from diffbio.operators._loss_balancing import LossBalancingMixin
+from diffbio.utils.nn_utils import ensure_rngs
+
+logger = logging.getLogger(__name__)
+
+# Canonical key names for the two most common modalities.
+_DEFAULT_MODALITY_KEYS = ("rna", "atac")
+
+
+@dataclass(frozen=True)
+class MultiOmicsVAEConfig(OperatorConfig):
+    """Configuration for DifferentiableMultiOmicsVAE.
+
+    Attributes:
+        modality_dims: Feature dimension for each modality.
+        latent_dim: Shared latent space dimension.
+        hidden_dim: Hidden layer width for all encoders / decoders.
+        modality_weight_mode: How reconstruction losses are weighted.
+            'equal' gives uniform weight; 'learnable' uses softmax over
+            a learnable log-weight vector.
+    """
+
+    modality_dims: list[int] = field(default_factory=lambda: [2000, 500])
+    latent_dim: int = 10
+    hidden_dim: int = 64
+    modality_weight_mode: str = "equal"
+    use_gradnorm: bool = False
+
+    def __post_init__(self) -> None:
+        """Set stochastic defaults for VAE sampling and validate."""
+        object.__setattr__(self, "stochastic", True)
+        if self.stream_name is None:
+            object.__setattr__(self, "stream_name", "sample")
+        super().__post_init__()
+
+
+class DifferentiableMultiOmicsVAE(LossBalancingMixin, EncoderDecoderOperator):
+    """Multi-omics VAE with Product-of-Experts latent fusion.
+
+    For each modality a dedicated encoder produces (mu_m, logvar_m).
+    These are combined via PoE into a joint posterior from which z is
+    sampled.  Per-modality decoders then reconstruct counts from z.
+
+    The ELBO objective uses MSE reconstruction loss per modality,
+    optionally weighted by learnable per-modality weights, plus
+    a KL divergence term against a standard-normal prior.
+
+    Data keys follow the convention ``<name>_counts`` for input and
+    ``<name>_reconstructed`` for output.  When exactly two modalities
+    are used the canonical names ``rna`` and ``atac`` are applied;
+    otherwise ``modality_<i>`` is used.
+
+    Attributes:
+        encoders: Per-modality encoder modules.
+        decoders: Per-modality decoder modules.
+        mu_heads: Per-modality linear projection for latent mean.
+        logvar_heads: Per-modality linear projection for latent logvar.
+        log_modality_weights: Learnable log-weights (only in 'learnable' mode).
+    """
+
+    def __init__(
+        self,
+        config: MultiOmicsVAEConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ) -> None:
+        """Initialise the multi-omics VAE.
+
+        Args:
+            config: Operator configuration.
+            rngs: Flax NNX random number generators.
+            name: Optional operator name.
+        """
+        super().__init__(config, rngs=rngs, name=name)
+
+        rngs = ensure_rngs(rngs)
+        n_modalities = len(config.modality_dims)
+
+        # Per-modality encoders ------------------------------------------
+        encoders: list[MLP] = []
+        mu_heads: list[nnx.Linear] = []
+        logvar_heads: list[nnx.Linear] = []
+
+        for dim in config.modality_dims:
+            encoders.append(
+                MLP(
+                    hidden_dims=[config.hidden_dim, config.hidden_dim],
+                    in_features=dim,
+                    activation="relu",
+                    output_activation="relu",
+                    use_batch_norm=False,
+                    rngs=rngs,
+                )
+            )
+            mu_heads.append(nnx.Linear(config.hidden_dim, config.latent_dim, rngs=rngs))
+            logvar_heads.append(nnx.Linear(config.hidden_dim, config.latent_dim, rngs=rngs))
+
+        self.encoders = nnx.List(encoders)
+        self.mu_heads = nnx.List(mu_heads)
+        self.logvar_heads = nnx.List(logvar_heads)
+
+        # Per-modality decoders ------------------------------------------
+        decoders: list[MLP] = []
+        for dim in config.modality_dims:
+            decoders.append(
+                MLP(
+                    hidden_dims=[config.hidden_dim, config.hidden_dim, dim],
+                    in_features=config.latent_dim,
+                    activation="relu",
+                    use_batch_norm=False,
+                    rngs=rngs,
+                )
+            )
+        self.decoders = nnx.List(decoders)
+
+        # Modality weight mode -------------------------------------------
+        self._weight_mode = nnx.static(config.modality_weight_mode)
+        if config.modality_weight_mode == "learnable":
+            self.log_modality_weights = nnx.Param(jnp.zeros(n_modalities))
+
+        # Assign canonical key names -------------------------------------
+        if n_modalities == 2:
+            self._modality_keys: tuple[str, ...] = _DEFAULT_MODALITY_KEYS
+        else:
+            self._modality_keys = tuple(f"modality_{i}" for i in range(n_modalities))
+
+    # ------------------------------------------------------------------
+    # PoE fusion
+    # ------------------------------------------------------------------
+
+    def product_of_experts(
+        self,
+        mu_list: list[Float[Array, "batch latent"]],
+        logvar_list: list[Float[Array, "batch latent"]],
+    ) -> tuple[Float[Array, "batch latent"], Float[Array, "batch latent"]]:
+        """Fuse per-modality posteriors via Product-of-Experts.
+
+        For M modalities the PoE joint posterior is Gaussian with:
+            precision_joint = sum_m precision_m
+            mu_joint = (sum_m mu_m * precision_m) / precision_joint
+        where precision_m = 1 / sigma_m^2 = exp(-logvar_m).
+
+        Args:
+            mu_list: Per-modality means.
+            logvar_list: Per-modality log-variances.
+
+        Returns:
+            (mu_joint, logvar_joint) tuple for the fused posterior.
+        """
+        # Stack precisions and weighted means across modalities
+        precision_sum = jnp.zeros_like(mu_list[0])
+        weighted_mu_sum = jnp.zeros_like(mu_list[0])
+
+        for mu_m, logvar_m in zip(mu_list, logvar_list):
+            precision_m = jnp.exp(-logvar_m)
+            precision_sum = precision_sum + precision_m
+            weighted_mu_sum = weighted_mu_sum + mu_m * precision_m
+
+        mu_joint = weighted_mu_sum / (precision_sum + EPSILON)
+        logvar_joint = -jnp.log(precision_sum + EPSILON)
+
+        return mu_joint, logvar_joint
+
+    # ------------------------------------------------------------------
+    # Modality weights
+    # ------------------------------------------------------------------
+
+    def _get_modality_weights(self) -> Float[Array, "n_modalities"]:
+        """Return normalised modality weights.
+
+        Returns:
+            Weight vector summing to 1 of length n_modalities.
+        """
+        n_modalities = len(self.config.modality_dims)
+        if self._weight_mode == "learnable":
+            return jax.nn.softmax(self.log_modality_weights[...])
+        return jnp.ones(n_modalities) / n_modalities
+
+    # ------------------------------------------------------------------
+    # Data key helpers
+    # ------------------------------------------------------------------
+
+    def _input_key(self, idx: int) -> str:
+        """Return the data-dict key for modality *idx* input counts.
+
+        Args:
+            idx: Modality index.
+
+        Returns:
+            String key such as ``rna_counts`` or ``modality_0_counts``.
+        """
+        return f"{self._modality_keys[idx]}_counts"
+
+    def _output_key(self, idx: int) -> str:
+        """Return the data-dict key for modality *idx* reconstruction.
+
+        Args:
+            idx: Modality index.
+
+        Returns:
+            String key such as ``rna_reconstructed``.
+        """
+        return f"{self._modality_keys[idx]}_reconstructed"
+
+    # ------------------------------------------------------------------
+    # apply
+    # ------------------------------------------------------------------
+
+    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]:
+        """Run the multi-omics VAE forward pass.
+
+        Steps:
+            1. Encode each modality to (mu_m, logvar_m).
+            2. PoE fusion -> (mu_joint, logvar_joint).
+            3. Reparameterise -> z.
+            4. Decode each modality from z.
+            5. Compute ELBO = weighted recon + KL.
+
+        Args:
+            data: Dictionary with ``<modality>_counts`` keys, each of shape
+                (n_cells, modality_dim).
+            state: Operator state (passed through unchanged).
+            metadata: Operator metadata (passed through unchanged).
+            random_params: Not used.
+            stats: Not used.
+
+        Returns:
+            Tuple of (result_data, state, metadata) where result_data
+            contains the original inputs plus ``joint_latent``,
+            ``<modality>_reconstructed``, and ``elbo_loss``.
+        """
+        n_modalities = len(self.config.modality_dims)
+
+        # 1. Encode each modality ----------------------------------------
+        mu_list: list[jax.Array] = []
+        logvar_list: list[jax.Array] = []
+
+        for i in range(n_modalities):
+            counts = data[self._input_key(i)]
+            h: jax.Array = self.encoders[i](jnp.log1p(counts))
+            mu = self.mu_heads[i](h)
+            logvar = jnp.clip(self.logvar_heads[i](h), -10.0, 10.0)
+            mu_list.append(mu)
+            logvar_list.append(logvar)
+
+        # 2. PoE fusion --------------------------------------------------
+        mu_joint, logvar_joint = self.product_of_experts(mu_list, logvar_list)
+
+        # 3. Sample z via reparameterisation (inherited) -----------------
+        z = self.reparameterize(mu_joint, logvar_joint)
+
+        # 4. Decode each modality ----------------------------------------
+        reconstructions: list[jax.Array] = []
+        for i in range(n_modalities):
+            reconstruction: jax.Array = self.decoders[i](z)
+            reconstructions.append(reconstruction)
+
+        # 5. Compute ELBO ------------------------------------------------
+        weights = self._get_modality_weights()
+
+        # Weighted reconstruction loss (MSE per modality, summed over features)
+        total_recon = jnp.array(0.0)
+        for i in range(n_modalities):
+            counts = data[self._input_key(i)]
+            per_sample = jnp.sum((counts - reconstructions[i]) ** 2, axis=-1)
+            mean_recon = reduce_loss(per_sample, reduction="mean")
+            total_recon = total_recon + weights[i] * mean_recon
+
+        # KL divergence (batch_sum: sum over latent, mean over batch)
+        kl = gaussian_kl_divergence(mu_joint, logvar_joint, reduction="batch_sum")
+
+        elbo_loss = total_recon + kl
+
+        # Build output dict -----------------------------------------------
+        result: dict[str, Any] = dict(data)
+        result["joint_latent"] = z
+        result["joint_mu"] = mu_joint
+        result["joint_logvar"] = logvar_joint
+        result["elbo_loss"] = elbo_loss
+
+        for i in range(n_modalities):
+            result[self._output_key(i)] = reconstructions[i]
+
+        return result, state, metadata