diffbio 0.1.0__py3-none-any.whl

diffbio/operators/drug_discovery/property_predictor.py

@@ -0,0 +1,163 @@
+"""Molecular property prediction operator.
+
+This module implements a ChemProp-style molecular property predictor
+using message passing neural networks.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+from artifex.generative_models.core.base import MLP
+from datarax.core.config import OperatorConfig
+from datarax.core.operator import OperatorModule
+from flax import nnx
+
+from diffbio.operators.drug_discovery._graph_utils import (
+    build_optional_dropout,
+    graph_sum_readout,
+    initialize_graph_encoder_from_config,
+)
+from diffbio.utils.nn_utils import ARTIFEX_RELU_MLP_KWARGS
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class MolecularPropertyConfig(OperatorConfig):
+    """Configuration for molecular property predictor.
+
+    Attributes:
+        hidden_dim: Hidden dimension for message passing layers.
+        num_message_passing_steps: Number of message passing iterations.
+        num_output_tasks: Number of prediction tasks (multi-task learning).
+        dropout_rate: Dropout rate for regularization.
+        in_features: Number of input node features (default: DEFAULT_ATOM_FEATURES=34).
+        num_edge_features: Number of edge/bond features.
+    """
+
+    hidden_dim: int = 300
+    num_message_passing_steps: int = 3
+    num_output_tasks: int = 1
+    dropout_rate: float = 0.0
+    in_features: int = 4  # Default for tests; use DEFAULT_ATOM_FEATURES for real molecules
+    num_edge_features: int = 4
+
+
+class MolecularPropertyPredictor(OperatorModule):
+    """ChemProp-style molecular property predictor.
+
+    Implements a directed message passing neural network (D-MPNN) for
+    predicting molecular properties from graph representations.
+
+    The architecture consists of:
+    1. Message passing layers to compute atom representations
+    2. Graph-level readout via sum pooling
+    3. Feed-forward network for property prediction
+
+    Example:
+        ```python
+        config = MolecularPropertyConfig(hidden_dim=64, num_output_tasks=3)
+        predictor = MolecularPropertyPredictor(config, rngs=nnx.Rngs(42))
+        data = {
+            "node_features": node_features,
+            "adjacency": adjacency,
+            "node_mask": mask,
+        }
+        result, state, meta = predictor.apply(data, {}, None)
+        predictions = result["predictions"]  # shape: (3,)
+        ```
+    """
+
+    def __init__(self, config: MolecularPropertyConfig, *, rngs: nnx.Rngs | None = None):
+        """Initialize molecular property predictor.
+
+        Args:
+            config: Predictor configuration.
+            rngs: Flax NNX random number generators.
+        """
+        super().__init__(config, rngs=rngs)
+
+        rngs = initialize_graph_encoder_from_config(self, config, rngs=rngs)
+
+        self.ffn_backbone = MLP(
+            hidden_dims=[config.hidden_dim],
+            in_features=config.hidden_dim,
+            dropout_rate=config.dropout_rate,
+            rngs=rngs,
+            **ARTIFEX_RELU_MLP_KWARGS,
+        )
+        self.output_layer = nnx.Linear(config.hidden_dim, config.num_output_tasks, rngs=rngs)
+
+        self.dropout = build_optional_dropout(config.dropout_rate, rngs=rngs)
+
+    def apply(
+        self,
+        data: dict[str, Any],
+        state: dict[str, Any],
+        metadata: dict[str, Any] | None,
+        random_params: Any = None,
+        stats: dict[str, Any] | None = None,
+    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any] | None]:
+        """Predict molecular properties from graph representation.
+
+        Args:
+            data: Input data containing:
+                - node_features: (num_nodes, num_features) atom features
+                - adjacency: (num_nodes, num_nodes) adjacency matrix
+                - edge_features: Optional (num_nodes, num_nodes, num_edge_features)
+                - node_mask: (num_nodes,) mask for valid nodes
+            state: Per-element state (passed through).
+            metadata: Optional metadata.
+            random_params: Unused random parameters.
+            stats: Optional statistics dictionary.
+
+        Returns:
+            Tuple of:
+                - data with added "predictions" key
+                - unchanged state
+                - unchanged metadata
+        """
+        graph_repr = graph_sum_readout(data, self.encoder, dropout=self.dropout)
+
+        # Feed-forward prediction
+        ffn_output = self.ffn_backbone(graph_repr)
+        if isinstance(ffn_output, tuple):
+            raise TypeError("MolecularPropertyPredictor FFN must return a single tensor output.")
+        predictions = self.output_layer(ffn_output)
+
+        result = {
+            **data,
+            "predictions": predictions,
+            "graph_representation": graph_repr,
+        }
+
+        return result, state, metadata
+
+
+def create_property_predictor(
+    hidden_dim: int = 300,
+    num_layers: int = 3,
+    num_tasks: int = 1,
+    dropout_rate: float = 0.0,
+    seed: int = 42,
+) -> MolecularPropertyPredictor:
+    """Create a molecular property predictor.
+
+    Args:
+        hidden_dim: Hidden dimension for message passing.
+        num_layers: Number of message passing steps.
+        num_tasks: Number of prediction tasks.
+        dropout_rate: Dropout rate.
+        seed: Random seed.
+
+    Returns:
+        Configured MolecularPropertyPredictor.
+    """
+    config = MolecularPropertyConfig(
+        hidden_dim=hidden_dim,
+        num_message_passing_steps=num_layers,
+        num_output_tasks=num_tasks,
+        dropout_rate=dropout_rate,
+    )
+    return MolecularPropertyPredictor(config, rngs=nnx.Rngs(seed))

diffbio/operators/drug_discovery/similarity.py

@@ -0,0 +1,193 @@
+"""Differentiable molecular similarity operator.
+
+This module implements differentiable similarity metrics for comparing
+molecular fingerprints, enabling gradient-based optimization of similarity.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import jax.numpy as jnp
+from datarax.core.config import OperatorConfig
+from datarax.core.operator import OperatorModule
+from flax import nnx
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class MolecularSimilarityConfig(OperatorConfig):
+    """Configuration for molecular similarity operator.
+
+    Attributes:
+        similarity_type: Type of similarity metric ("tanimoto", "cosine", "dice").
+        temperature: Temperature for soft similarity (higher = sharper).
+    """
+
+    similarity_type: str = "tanimoto"
+    temperature: float = 1.0
+
+
+def tanimoto_similarity(a: jnp.ndarray, b: jnp.ndarray, eps: float = 1e-8) -> jnp.ndarray:
+    """Compute differentiable Tanimoto similarity.
+
+    For continuous vectors, uses the generalized Tanimoto formula:
+        T(a, b) = (a · b) / (|a|² + |b|² - a · b)
+
+    Args:
+        a: First fingerprint vector.
+        b: Second fingerprint vector.
+        eps: Small constant for numerical stability.
+
+    Returns:
+        Similarity score in [0, 1].
+    """
+    dot_product = jnp.sum(a * b)
+    norm_a_sq = jnp.sum(a * a)
+    norm_b_sq = jnp.sum(b * b)
+
+    # Generalized Tanimoto for continuous vectors
+    similarity = dot_product / (norm_a_sq + norm_b_sq - dot_product + eps)
+
+    # Clamp to [0, 1] for numerical stability
+    return jnp.clip(similarity, 0.0, 1.0)
+
+
+def cosine_similarity(a: jnp.ndarray, b: jnp.ndarray, eps: float = 1e-8) -> jnp.ndarray:
+    """Compute cosine similarity.
+
+    Args:
+        a: First vector.
+        b: Second vector.
+        eps: Small constant for numerical stability.
+
+    Returns:
+        Similarity score in [-1, 1].
+    """
+    dot_product = jnp.sum(a * b)
+    norm_a = jnp.linalg.norm(a)
+    norm_b = jnp.linalg.norm(b)
+
+    return dot_product / (norm_a * norm_b + eps)
+
+
+def dice_similarity(a: jnp.ndarray, b: jnp.ndarray, eps: float = 1e-8) -> jnp.ndarray:
+    """Compute Dice similarity coefficient.
+
+    For continuous vectors:
+        Dice(a, b) = 2 * (a · b) / (|a|² + |b|²)
+
+    Args:
+        a: First vector.
+        b: Second vector.
+        eps: Small constant for numerical stability.
+
+    Returns:
+        Similarity score in [0, 1].
+    """
+    dot_product = jnp.sum(a * b)
+    norm_a_sq = jnp.sum(a * a)
+    norm_b_sq = jnp.sum(b * b)
+
+    similarity = 2 * dot_product / (norm_a_sq + norm_b_sq + eps)
+
+    return jnp.clip(similarity, 0.0, 1.0)
+
+
+class MolecularSimilarityOperator(OperatorModule):
+    """Differentiable molecular similarity operator.
+
+    Computes similarity between molecular fingerprints using various
+    differentiable metrics. Supports Tanimoto, cosine, and Dice similarity.
+
+    Example:
+        ```python
+        config = MolecularSimilarityConfig(similarity_type="tanimoto")
+        sim_op = MolecularSimilarityOperator(config, rngs=nnx.Rngs(42))
+        data = {"fingerprint_a": fp1, "fingerprint_b": fp2}
+        result, _, _ = sim_op.apply(data, {}, None)
+        similarity = result["similarity"]  # scalar in [0, 1]
+        ```
+    """
+
+    def __init__(self, config: MolecularSimilarityConfig, *, rngs: nnx.Rngs | None = None):
+        """Initialize similarity operator.
+
+        Args:
+            config: Similarity configuration.
+            rngs: Flax NNX random number generators.
+        """
+        super().__init__(config, rngs=rngs)
+        self.config: MolecularSimilarityConfig = config
+
+        # Fix: wrap _unique_id as static for jax.grad compatibility
+        # (datarax stores it as plain int which causes gradient errors)
+        self._unique_id = nnx.static(self._unique_id)
+
+        # Select similarity function
+        if config.similarity_type == "tanimoto":
+            self._similarity_fn = tanimoto_similarity
+        elif config.similarity_type == "cosine":
+            self._similarity_fn = cosine_similarity
+        elif config.similarity_type == "dice":
+            self._similarity_fn = dice_similarity
+        else:
+            raise ValueError(f"Unknown similarity type: {config.similarity_type}")
+
+    def apply(
+        self,
+        data: dict[str, Any],
+        state: dict[str, Any],
+        metadata: dict[str, Any] | None,
+        random_params: Any = None,
+        stats: dict[str, Any] | None = None,
+    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any] | None]:
+        """Compute similarity between two fingerprints.
+
+        Args:
+            data: Input data containing:
+                - fingerprint_a: First fingerprint vector
+                - fingerprint_b: Second fingerprint vector
+            state: Per-element state (passed through).
+            metadata: Optional metadata.
+            random_params: Unused random parameters.
+            stats: Optional statistics dictionary.
+
+        Returns:
+            Tuple of:
+                - data with added "similarity" key
+                - unchanged state
+                - unchanged metadata
+        """
+        fp_a = data["fingerprint_a"]
+        fp_b = data["fingerprint_b"]
+
+        similarity = self._similarity_fn(fp_a, fp_b)
+
+        result = {
+            **data,
+            "similarity": similarity,
+        }
+
+        return result, state, metadata
+
+
+def create_similarity_operator(
+    similarity_type: str = "tanimoto",
+    temperature: float = 1.0,
+) -> MolecularSimilarityOperator:
+    """Create a molecular similarity operator.
+
+    Args:
+        similarity_type: Type of similarity ("tanimoto", "cosine", "dice").
+        temperature: Temperature parameter.
+
+    Returns:
+        Configured MolecularSimilarityOperator.
+    """
+    config = MolecularSimilarityConfig(
+        similarity_type=similarity_type,
+        temperature=temperature,
+    )
+    return MolecularSimilarityOperator(config, rngs=nnx.Rngs(42))

diffbio/operators/epigenomics/__init__.py

@@ -0,0 +1,35 @@
+"""Epigenomics operators for differentiable ChIP-seq and ATAC-seq analysis.
+
+This module provides differentiable operators for epigenomic data analysis,
+including peak calling and chromatin state annotation.
+"""
+
+from diffbio.operators.epigenomics.chromatin_state import (
+    ChromatinStateAnnotator,
+    ChromatinStateConfig,
+)
+from diffbio.operators.epigenomics.contextual import (
+    ContextualEpigenomicsConfig,
+    ContextualEpigenomicsOperator,
+    compute_contextual_epigenomics_loss,
+)
+from diffbio.operators.epigenomics.fno_peak_calling import (
+    FNOPeakCaller,
+    FNOPeakCallerConfig,
+)
+from diffbio.operators.epigenomics.peak_calling import (
+    DifferentiablePeakCaller,
+    PeakCallerConfig,
+)
+
+__all__ = [
+    "ChromatinStateAnnotator",
+    "ChromatinStateConfig",
+    "ContextualEpigenomicsConfig",
+    "ContextualEpigenomicsOperator",
+    "DifferentiablePeakCaller",
+    "FNOPeakCaller",
+    "FNOPeakCallerConfig",
+    "PeakCallerConfig",
+    "compute_contextual_epigenomics_loss",
+]