diffbio 0.1.0__py3-none-any.whl

diffbio/operators/crispr/guide_scoring.py

@@ -0,0 +1,269 @@
+"""Differentiable CRISPR Guide Scoring Operator.
+
+This module implements a DeepCRISPR-inspired differentiable guide RNA scoring
+operator using a CNN architecture to predict on-target efficiency.
+
+The architecture is inspired by DeepCRISPR which uses:
+1. A deep convolutional denoising neural network (DCDNN) autoencoder
+   for unsupervised representation learning
+2. A CNN classifier for efficiency prediction
+
+This implementation provides a simplified but differentiable version that:
+- Uses 1D convolutions over the one-hot encoded sequence
+- Supports optional epigenetic feature channels
+- Outputs efficiency scores in [0, 1]
+
+For SpCas9, the standard input is 20nt guide + 3nt PAM = 23nt context.
+
+References:
+    Chuai et al. (2018). "DeepCRISPR: Optimized CRISPR guide RNA design
+    by deep learning." Genome Biology.
+    https://github.com/bm2-lab/DeepCRISPR
+
+    Liu et al. (2021). "Enhancing CRISPR-Cas9 gRNA efficiency prediction
+    by data integration and deep learning." Nature Communications.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import jax.numpy as jnp
+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.utils.nn_utils import ARTIFEX_RELU_MLP_KWARGS
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class CRISPRScorerConfig(OperatorConfig):
+    """Configuration for DifferentiableCRISPRScorer.
+
+    Attributes:
+        guide_length: Length of guide RNA sequence (typically 20-23 nt).
+        alphabet_size: Size of nucleotide alphabet (4 for A/C/G/T).
+        hidden_channels: CNN hidden channel dimensions.
+        fc_dims: Fully connected layer dimensions.
+        dropout_rate: Dropout rate for regularization.
+    """
+
+    guide_length: int = 23
+    alphabet_size: int = 4
+    hidden_channels: tuple[int, ...] = (64, 128, 256)
+    fc_dims: tuple[int, ...] = (256, 128)
+    dropout_rate: float = 0.2
+
+    def __post_init__(self) -> None:
+        """Validate configuration."""
+        super().__post_init__()
+        if not self.hidden_channels:
+            raise ValueError(
+                "CRISPRScorerConfig.hidden_channels must contain at least one channel."
+            )
+        if not self.fc_dims:
+            raise ValueError("CRISPRScorerConfig.fc_dims must contain at least one hidden layer.")
+
+
+class DifferentiableCRISPRScorer(OperatorModule):
+    """DeepCRISPR-style differentiable guide RNA scoring.
+
+    This operator uses a 1D CNN architecture to predict CRISPR guide RNA
+    on-target efficiency from sequence features. The model learns sequence
+    patterns that correlate with efficient target cleavage.
+
+    The architecture consists of:
+    1. 1D convolutional layers for sequence feature extraction
+    2. Batch normalization and ReLU activations
+    3. Fully connected layers for efficiency prediction
+    4. Sigmoid output for efficiency score in [0, 1]
+
+    Attributes:
+        config: Operator configuration.
+        conv_layers: 1D convolutional layers.
+        conv_bn: Batch normalization layers for conv.
+        ffn_backbone: Shared Artifex MLP for score prediction.
+        output_head: Final output layer.
+
+    Example:
+        ```python
+        from diffbio.operators.crispr import (
+            DifferentiableCRISPRScorer,
+            CRISPRScorerConfig,
+        )
+        config = CRISPRScorerConfig(guide_length=23)
+        scorer = DifferentiableCRISPRScorer(config, rngs=nnx.Rngs(42))
+        data = {"guides": guide_sequences}  # (n_guides, length, 4)
+        result, _, _ = scorer.apply(data, {}, None)
+        scores = result["efficiency_scores"]  # (n_guides,)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: CRISPRScorerConfig,
+        *,
+        rngs: nnx.Rngs,
+    ) -> None:
+        """Initialize the CRISPR scorer.
+
+        Args:
+            config: Operator configuration.
+            rngs: Flax NNX random number generators.
+        """
+        super().__init__(config, rngs=rngs)
+
+        channel_pairs = zip((config.alphabet_size, *config.hidden_channels), config.hidden_channels)
+        self.conv_layers = nnx.List(
+            [
+                nnx.Conv(
+                    in_features=in_channels,
+                    out_features=out_channels,
+                    kernel_size=(3,),
+                    padding="SAME",
+                    rngs=rngs,
+                )
+                for in_channels, out_channels in channel_pairs
+            ]
+        )
+        self.conv_bn = nnx.List(
+            [nnx.BatchNorm(out_channels, rngs=rngs) for out_channels in config.hidden_channels]
+        )
+
+        # Calculate flattened size after convolutions
+        # With SAME padding, spatial size is preserved
+        flat_size = config.guide_length * config.hidden_channels[-1]
+
+        self.ffn_backbone = MLP(
+            hidden_dims=list(config.fc_dims),
+            in_features=flat_size,
+            dropout_rate=config.dropout_rate,
+            rngs=rngs,
+            **ARTIFEX_RELU_MLP_KWARGS,
+        )
+
+        # Output head for efficiency score
+        self.output_head = nnx.Linear(config.fc_dims[-1], 1, rngs=rngs)
+
+    def extract_features(self, guides: jnp.ndarray) -> jnp.ndarray:
+        """Extract features from guide sequences using CNN.
+
+        Args:
+            guides: One-hot encoded guides (n_guides, guide_length, 4).
+
+        Returns:
+            Feature vectors (n_guides, feature_dim).
+        """
+        # Input shape: (batch, length, channels)
+        x = guides
+
+        # Apply 1D convolutions with batch norm and ReLU
+        for conv, bn in zip(self.conv_layers, self.conv_bn):
+            x = conv(x)
+            x = bn(x)
+            x = nnx.relu(x)
+
+        # Flatten: (batch, length, channels) -> (batch, length * channels)
+        batch_size = x.shape[0]
+        x = x.reshape(batch_size, -1)
+
+        return x
+
+    def predict_efficiency(self, features: jnp.ndarray) -> jnp.ndarray:
+        """Predict efficiency score from features.
+
+        Args:
+            features: Feature vectors (n_guides, feature_dim).
+
+        Returns:
+            Efficiency scores (n_guides,) in range [0, 1].
+        """
+        backbone_output = self.ffn_backbone(features)
+        if isinstance(backbone_output, tuple):
+            raise TypeError("CRISPR scorer backbone must return a single tensor output.")
+
+        # Output layer with sigmoid for [0, 1] output
+        x = self.output_head(backbone_output)
+        scores = nnx.sigmoid(x).squeeze(-1)
+
+        return scores
+
+    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]:
+        """Apply CRISPR scoring to guide sequences.
+
+        Args:
+            data: Dictionary containing:
+                - "guides": One-hot encoded guides (n_guides, guide_length, 4).
+            state: Per-element state (passed through).
+            metadata: Optional metadata (passed through).
+            random_params: Random parameters for stochastic operations.
+            stats: Optional statistics dictionary.
+
+        Returns:
+            Tuple of (transformed_data, state, metadata) where transformed_data
+            contains:
+
+                - "guides": Original guide sequences.
+                - "efficiency_scores": Predicted efficiency (n_guides,).
+                - "features": Extracted feature vectors.
+        """
+        guides = data["guides"]
+
+        # Extract features using CNN
+        features = self.extract_features(guides)
+
+        # Predict efficiency scores
+        efficiency_scores = self.predict_efficiency(features)
+
+        # Build output
+        output = {
+            **data,
+            "efficiency_scores": efficiency_scores,
+            "features": features,
+        }
+
+        return output, state, metadata
+
+
+def create_crispr_scorer(
+    guide_length: int = 23,
+    hidden_channels: tuple[int, ...] = (64, 128, 256),
+    fc_dims: tuple[int, ...] = (256, 128),
+    dropout_rate: float = 0.2,
+    seed: int = 42,
+) -> DifferentiableCRISPRScorer:
+    """Factory function to create a CRISPR scorer.
+
+    Args:
+        guide_length: Length of guide RNA sequence.
+        hidden_channels: CNN hidden channel dimensions.
+        fc_dims: Fully connected layer dimensions.
+        dropout_rate: Dropout rate for regularization.
+        seed: Random seed for initialization.
+
+    Returns:
+        Configured DifferentiableCRISPRScorer instance.
+
+    Example:
+        ```python
+        scorer = create_crispr_scorer(guide_length=23)
+        result, _, _ = scorer.apply({"guides": data}, {}, None)
+        ```
+    """
+    config = CRISPRScorerConfig(
+        guide_length=guide_length,
+        hidden_channels=hidden_channels,
+        fc_dims=fc_dims,
+        dropout_rate=dropout_rate,
+    )
+
+    return DifferentiableCRISPRScorer(config, rngs=nnx.Rngs(seed))

diffbio/operators/drug_discovery/__init__.py

@@ -0,0 +1,133 @@
+"""Drug discovery operators for differentiable chemoinformatics.
+
+This module provides differentiable operators for molecular property prediction,
+fingerprint computation, and similarity scoring.
+
+Operators:
+    MolecularPropertyPredictor: ChemProp-style MPNN for property prediction
+    DifferentiableMolecularFingerprint: Neural graph fingerprints
+    CircularFingerprintOperator: Differentiable ECFP/Morgan fingerprints
+    MACCSKeysOperator: Differentiable MACCS 166 structural keys
+    AttentiveFP: Attention-based molecular fingerprint (Xiong et al. 2019)
+    ADMETPredictor: Multi-task ADMET property prediction
+    MolecularSimilarityOperator: Differentiable Tanimoto/cosine similarity
+    DifferentiableDTIPipeline: Integrated differentiable DTI scorer
+
+Utilities:
+    smiles_to_graph: Convert SMILES to molecular graph
+    batch_smiles_to_graphs: Batch conversion with padding
+"""
+
+from diffbio.operators.drug_discovery.admet_predictor import (
+    ADMET_TASK_NAMES,
+    ADMET_TASK_TYPES,
+    ADMETConfig,
+    ADMETPredictor,
+    create_admet_predictor,
+)
+from diffbio.operators.drug_discovery.attentive_fp import (
+    AttentiveFP,
+    AttentiveFPConfig,
+    create_attentive_fp,
+)
+from diffbio.operators.drug_discovery.dti import (
+    DTI_PIPELINE_INTEGRATION_LAYER,
+    PROTEIN_ONE_HOT_PREPROCESSING_VERSION,
+    DTIPipelineConfig,
+    DifferentiableDTIPipeline,
+    build_dti_pipeline_inputs,
+    encode_protein_sequences,
+)
+from diffbio.operators.drug_discovery.fingerprint import (
+    CircularFingerprintConfig,
+    CircularFingerprintOperator,
+    DifferentiableMolecularFingerprint,
+    MolecularFingerprintConfig,
+    create_ecfp4_operator,
+    create_ecfp6_operator,
+    create_fcfp4_operator,
+    create_fingerprint_operator,
+)
+from diffbio.operators.drug_discovery.maccs_keys import (
+    MACCSKeysConfig,
+    MACCSKeysOperator,
+    create_maccs_operator,
+)
+from diffbio.operators.drug_discovery.message_passing import (
+    MessagePassingLayer,
+    StackedMessagePassing,
+)
+from diffbio.operators.drug_discovery.primitives import (
+    AtomFeatureConfig,
+    DEFAULT_ATOM_CONFIG,
+    DEFAULT_ATOM_FEATURES,
+    batch_smiles_to_graphs,
+    smiles_to_graph,
+)
+from diffbio.operators.drug_discovery.property_predictor import (
+    MolecularPropertyConfig,
+    MolecularPropertyPredictor,
+    create_property_predictor,
+)
+from diffbio.operators.drug_discovery.similarity import (
+    MolecularSimilarityConfig,
+    MolecularSimilarityOperator,
+    cosine_similarity,
+    create_similarity_operator,
+    dice_similarity,
+    tanimoto_similarity,
+)
+
+__all__ = [
+    # Primitives
+    "smiles_to_graph",
+    "batch_smiles_to_graphs",
+    "DEFAULT_ATOM_FEATURES",
+    "DEFAULT_ATOM_CONFIG",
+    "AtomFeatureConfig",
+    # Message Passing
+    "MessagePassingLayer",
+    "StackedMessagePassing",
+    # Property Prediction
+    "MolecularPropertyConfig",
+    "MolecularPropertyPredictor",
+    "create_property_predictor",
+    # ADMET Prediction
+    "ADMETConfig",
+    "ADMETPredictor",
+    "create_admet_predictor",
+    "ADMET_TASK_NAMES",
+    "ADMET_TASK_TYPES",
+    # DTI Pipeline
+    "DTI_PIPELINE_INTEGRATION_LAYER",
+    "PROTEIN_ONE_HOT_PREPROCESSING_VERSION",
+    "DTIPipelineConfig",
+    "DifferentiableDTIPipeline",
+    "build_dti_pipeline_inputs",
+    "encode_protein_sequences",
+    # Neural Fingerprints
+    "MolecularFingerprintConfig",
+    "DifferentiableMolecularFingerprint",
+    "create_fingerprint_operator",
+    # Circular Fingerprints (ECFP/Morgan)
+    "CircularFingerprintConfig",
+    "CircularFingerprintOperator",
+    "create_ecfp4_operator",
+    "create_ecfp6_operator",
+    "create_fcfp4_operator",
+    # MACCS Keys
+    "MACCSKeysConfig",
+    "MACCSKeysOperator",
+    "create_maccs_operator",
+    # AttentiveFP
+    "AttentiveFPConfig",
+    "AttentiveFP",
+    "create_attentive_fp",
+    # Similarity
+    "MolecularSimilarityConfig",
+    "MolecularSimilarityOperator",
+    "create_similarity_operator",
+    "tanimoto_similarity",
+    "cosine_similarity",
+    "dice_similarity",
+]

diffbio/operators/drug_discovery/_graph_utils.py

@@ -0,0 +1,142 @@
+"""Shared helpers for graph-based drug-discovery operators."""
+
+import logging
+from typing import Any, Mapping
+
+import jax.numpy as jnp
+from flax import nnx
+
+from diffbio.operators.drug_discovery.message_passing import StackedMessagePassing
+from diffbio.utils.nn_utils import ensure_rngs
+
+logger = logging.getLogger(__name__)
+
+
+def stabilize_operator_id(module: Any) -> None:
+    """Mark operator unique ID as static for NNX/JAX transformations."""
+    module._unique_id = nnx.static(module._unique_id)
+
+
+def build_encoder(
+    *,
+    hidden_dim: int,
+    num_layers: int,
+    in_features: int,
+    rngs: nnx.Rngs,
+    num_edge_features: int | None = None,
+) -> StackedMessagePassing:
+    """Create a message-passing encoder with optional edge features."""
+    encoder_kwargs: dict[str, Any] = {
+        "hidden_dim": hidden_dim,
+        "num_layers": num_layers,
+        "in_features": in_features,
+        "rngs": rngs,
+    }
+    if num_edge_features is not None:
+        encoder_kwargs["num_edge_features"] = num_edge_features
+    return StackedMessagePassing(**encoder_kwargs)
+
+
+def _require_config_attr(config: Any, attr: str) -> Any:
+    """Read a required config attribute with a clear error message."""
+    if not hasattr(config, attr):
+        raise AttributeError(
+            f"{type(config).__name__} must define '{attr}' for graph encoder initialization."
+        )
+    return getattr(config, attr)
+
+
+def initialize_graph_encoder(
+    module: Any,
+    *,
+    rngs: nnx.Rngs | None,
+    hidden_dim: int,
+    num_layers: int,
+    in_features: int,
+    num_edge_features: int | None = None,
+    attr: str = "encoder",
+) -> nnx.Rngs:
+    """Stabilize ID, ensure RNGs, and attach a message-passing encoder."""
+    stabilize_operator_id(module)
+    resolved_rngs = ensure_rngs(rngs)
+    setattr(
+        module,
+        attr,
+        build_encoder(
+            hidden_dim=hidden_dim,
+            num_layers=num_layers,
+            in_features=in_features,
+            rngs=resolved_rngs,
+            num_edge_features=num_edge_features,
+        ),
+    )
+    return resolved_rngs
+
+
+def initialize_graph_encoder_from_config(
+    module: Any,
+    config: Any,
+    *,
+    rngs: nnx.Rngs | None,
+    num_layers_attr: str = "num_message_passing_steps",
+    hidden_dim_attr: str = "hidden_dim",
+    in_features_attr: str = "in_features",
+    num_edge_features_attr: str = "num_edge_features",
+    attr: str = "encoder",
+) -> nnx.Rngs:
+    """Initialize a standard graph encoder from a config object."""
+    hidden_dim = _require_config_attr(config, hidden_dim_attr)
+    num_layers = _require_config_attr(config, num_layers_attr)
+    in_features = _require_config_attr(config, in_features_attr)
+    num_edge_features = getattr(config, num_edge_features_attr, None)
+
+    return initialize_graph_encoder(
+        module,
+        rngs=rngs,
+        hidden_dim=hidden_dim,
+        num_layers=num_layers,
+        in_features=in_features,
+        num_edge_features=num_edge_features,
+        attr=attr,
+    )
+
+
+def build_optional_dropout(rate: float, *, rngs: nnx.Rngs) -> nnx.Dropout | None:
+    """Create a dropout module only when the configured rate is positive."""
+    if rate <= 0:
+        return None
+    return nnx.Dropout(rate=rate, rngs=rngs)
+
+
+def unpack_graph_inputs(data: Mapping[str, Any]) -> tuple[Any, Any, Any, Any]:
+    """Extract standard molecular graph tensors from an input dictionary."""
+    return (
+        data["node_features"],
+        data["adjacency"],
+        data.get("edge_features"),
+        data.get("node_mask"),
+    )
+
+
+def graph_sum_readout(
+    data: Mapping[str, Any],
+    encoder: StackedMessagePassing,
+    *,
+    dropout: nnx.Dropout | None = None,
+) -> jnp.ndarray:
+    """Encode a graph and sum-pool node states into a graph representation."""
+    node_features, adjacency, edge_features, node_mask = unpack_graph_inputs(data)
+
+    node_hidden = encoder(node_features, adjacency, edge_features)
+    if node_mask is not None:
+        node_hidden = node_hidden * node_mask[:, None]
+
+    graph_repr = jnp.sum(node_hidden, axis=0)
+    if dropout is not None:
+        graph_repr = dropout(graph_repr)
+    return graph_repr
+
+
+def attach_fingerprint(data: Mapping[str, Any], fingerprint: Any) -> dict[str, Any]:
+    """Return a shallow copy with a standardized fingerprint output field."""
+    return {**data, "fingerprint": fingerprint}