diffbio 0.1.0__py3-none-any.whl

diffbio/__init__.py

@@ -0,0 +1,39 @@
+"""DiffBio: End-to-end differentiable bioinformatics pipelines built on the wider JAX/NNX ecosystem.
+
+This package provides differentiable bioinformatics pipeline components that
+integrate with Datarax, Artifex, Opifex, and Calibrax for gradient-based
+optimization of genomics workflows.
+
+Key components:
+- sequences: Biological sequence data types (DNA, RNA, Protein)
+- operators: Differentiable bioinformatics operators (alignment, quality filtering)
+- losses: Loss functions and biological regularization
+- pipelines: Pre-built differentiable pipeline templates
+- configs: Base configuration classes for operators
+- constants: Centralized constants for the library
+"""
+
+from diffbio import (
+    configs,
+    constants,
+    evaluation,
+    losses,
+    operators,
+    pipelines,
+    sequences,
+    utils,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "__version__",
+    "configs",
+    "constants",
+    "evaluation",
+    "losses",
+    "operators",
+    "pipelines",
+    "sequences",
+    "utils",
+]

diffbio/configs.py

@@ -0,0 +1,75 @@
+"""Base configuration classes for DiffBio operators.
+
+This module provides base configuration classes that reduce duplication
+across operator configs by providing common fields with sensible defaults.
+
+Note: The `stochastic` and `stream_name` fields are already defined in
+datarax.core.config.OperatorConfig, so we don't re-declare them here.
+DiffBio configs inherit these fields automatically.
+"""
+
+from dataclasses import dataclass
+
+from datarax.core.config import OperatorConfig
+
+from diffbio.constants import (
+    DEFAULT_DROPOUT_RATE,
+    DEFAULT_HIDDEN_DIM,
+    DEFAULT_NUM_CLASSES,
+    DEFAULT_NUM_LAYERS,
+    DEFAULT_TEMPERATURE,
+)
+
+
+def apply_stochastic_sampling_defaults(
+    config: OperatorConfig,
+    *,
+    stream_name: str = "sample",
+) -> None:
+    """Apply DiffBio's default stochastic sampling contract to a config.
+
+    Args:
+        config: Operator config instance being finalized in ``__post_init__``.
+        stream_name: Default RNG stream name to use when none is supplied.
+    """
+    object.__setattr__(config, "stochastic", True)
+    if config.stream_name is None:
+        object.__setattr__(config, "stream_name", stream_name)
+
+
+@dataclass(frozen=True)
+class TemperatureConfig(OperatorConfig):
+    """Configuration for operators with temperature parameter.
+
+    Use this base class for operators that use temperature-based
+    smoothing (logsumexp relaxation, soft thresholding, etc.).
+
+    Attributes:
+        temperature: Temperature for smooth operations.
+            Lower = sharper (closer to hard operations).
+            Higher = smoother (more gradient flow).
+        learnable_temperature: Whether temperature is a learnable parameter.
+            If True, temperature will be an nnx.Param that receives gradients.
+    """
+
+    temperature: float = DEFAULT_TEMPERATURE
+    learnable_temperature: bool = False
+
+
+@dataclass(frozen=True)
+class ClassifierConfig(OperatorConfig):
+    """Base configuration for classifier operators.
+
+    Provides common fields for neural network classifiers.
+
+    Attributes:
+        num_classes: Number of output classes.
+        hidden_dim: Hidden layer dimension.
+        num_layers: Number of hidden layers.
+        dropout_rate: Dropout rate for regularization.
+    """
+
+    num_classes: int = DEFAULT_NUM_CLASSES
+    hidden_dim: int = DEFAULT_HIDDEN_DIM
+    num_layers: int = DEFAULT_NUM_LAYERS
+    dropout_rate: float = DEFAULT_DROPOUT_RATE

diffbio/constants.py

@@ -0,0 +1,204 @@
+"""Centralized constants for DiffBio.
+
+This module provides centralized constants used across operators, pipelines,
+and utilities. Using these constants ensures consistency and makes maintenance
+easier.
+"""
+
+from enum import StrEnum
+
+# =============================================================================
+# Neural Network Architecture Defaults
+# =============================================================================
+
+DEFAULT_HIDDEN_DIM = 64
+"""Default hidden layer dimension for MLP architectures."""
+
+DEFAULT_HIDDEN_DIM_LARGE = 128
+"""Larger hidden dimension for complex models."""
+
+DEFAULT_EMBEDDING_DIM = 64
+"""Default embedding dimension for sequence embeddings."""
+
+DEFAULT_DROPOUT_RATE = 0.1
+"""Default dropout rate for regularization."""
+
+DEFAULT_NUM_LAYERS = 2
+"""Default number of hidden layers in MLPs."""
+
+DEFAULT_TEMPERATURE = 1.0
+"""Default temperature for logsumexp smoothing and softmax."""
+
+
+# =============================================================================
+# Quality Score Constants (Phred Scale)
+# =============================================================================
+
+PHRED_QUALITY_MAX = 40.0
+"""Maximum Phred quality score (99.99% accuracy)."""
+
+PHRED_QUALITY_THRESHOLD = 20.0
+"""Default Phred quality threshold (99% accuracy, 1% error rate)."""
+
+PHRED_QUALITY_MIN = 0.0
+"""Minimum Phred quality score."""
+
+
+# =============================================================================
+# Alignment Constants
+# =============================================================================
+
+DEFAULT_GAP_OPEN = -10.0
+"""Default gap opening penalty for alignment."""
+
+DEFAULT_GAP_EXTEND = -1.0
+"""Default gap extension penalty for alignment."""
+
+DNA_MATCH_SCORE = 2.0
+"""Default match score for DNA alignment."""
+
+DNA_MISMATCH_SCORE = -1.0
+"""Default mismatch penalty for DNA alignment."""
+
+
+# =============================================================================
+# Numerical Stability Constants
+# =============================================================================
+
+EPSILON = 1e-8
+"""Small value for numerical stability in divisions and log operations."""
+
+EPSILON_LOG = 1e-10
+"""Smaller epsilon specifically for log operations to prevent -inf."""
+
+DISTANCE_MASK_SENTINEL = 1e10
+"""Large value used to mask out self-distances in pairwise distance matrices."""
+
+
+# =============================================================================
+# Pileup and Coverage Constants
+# =============================================================================
+
+DEFAULT_PILEUP_WINDOW_SIZE = 21
+"""Default window size for pileup context."""
+
+DEFAULT_MIN_COVERAGE = 1
+"""Default minimum coverage threshold."""
+
+DEFAULT_MAX_COVERAGE = 100
+"""Default maximum coverage for normalization."""
+
+
+# =============================================================================
+# Variant Calling Constants
+# =============================================================================
+
+DEFAULT_NUM_CLASSES = 3
+"""Default number of variant classes (REF, SNV, INDEL)."""
+
+
+class ClassifierType(StrEnum):
+    """Classifier type for variant calling pipelines."""
+
+    MLP = "mlp"
+    """Multi-layer perceptron classifier."""
+
+    CNN = "cnn"
+    """Convolutional neural network classifier."""
+
+
+# =============================================================================
+# DNA/RNA Constants
+# =============================================================================
+
+DNA_ALPHABET_SIZE = 4
+"""Size of DNA alphabet (A, C, G, T)."""
+
+RNA_ALPHABET_SIZE = 4
+"""Size of RNA alphabet (A, C, G, U)."""
+
+PROTEIN_ALPHABET_SIZE = 20
+"""Size of standard protein alphabet."""
+
+
+# =============================================================================
+# VAE / Latent Space Constants
+# =============================================================================
+
+DEFAULT_LATENT_DIM = 10
+"""Default latent dimension for VAE models."""
+
+DEFAULT_BETA_VAE = 1.0
+"""Default beta parameter for beta-VAE (KL weight)."""
+
+
+# =============================================================================
+# Graph Neural Network Constants
+# =============================================================================
+
+DEFAULT_NODE_FEATURES = 32
+"""Default node feature dimension for GNN."""
+
+DEFAULT_EDGE_FEATURES = 8
+"""Default edge feature dimension for GNN."""
+
+DEFAULT_NUM_HEADS = 4
+"""Default number of attention heads."""
+
+
+# =============================================================================
+# HMM Constants
+# =============================================================================
+
+DEFAULT_HMM_STATES = 3
+"""Default number of hidden states for HMM."""
+
+DEFAULT_HMM_EMISSIONS = 4
+"""Default number of emissions for HMM (matches DNA alphabet)."""
+
+
+# =============================================================================
+# Sequence Length Limits
+# =============================================================================
+
+DEFAULT_MAX_SEQ_LENGTH = 1000
+"""Default maximum sequence length."""
+
+DEFAULT_MAX_ALIGNMENT_LENGTH = 500
+"""Default maximum alignment length."""
+
+
+# =============================================================================
+# Data Dictionary Keys
+# =============================================================================
+# These constants ensure consistent key naming across operators
+
+
+class DataKeys:
+    """Standard keys for data dictionaries in DiffBio operators."""
+
+    # Sequence data
+    SEQUENCE = "sequence"
+    SEQ1 = "seq1"
+    SEQ2 = "seq2"
+    READS = "reads"
+
+    # Quality and positions
+    QUALITY = "quality"
+    QUALITY_SCORES = "quality_scores"
+    POSITIONS = "positions"
+
+    # Pileup related
+    PILEUP = "pileup"
+    PILEUP_WINDOW = "pileup_window"
+    COVERAGE = "coverage"
+    MEAN_QUALITY = "mean_quality"
+
+    # Classification outputs
+    LOGITS = "logits"
+    PROBABILITIES = "probabilities"
+
+    # Alignment outputs
+    SCORE = "score"
+    ALIGNMENT_MATRIX = "alignment_matrix"
+    SOFT_ALIGNMENT = "soft_alignment"

diffbio/core/__init__.py

@@ -0,0 +1,127 @@
+"""Core module for DiffBio.
+
+This module provides foundational components for building differentiable
+bioinformatics operators:
+
+- **soft_ops**: Differentiable soft operations (soft sorting, argmax,
+  comparisons, logical ops, selection, quantile, straight-through estimators)
+- **base_operators**: Domain-specific base classes (TemperatureOperator,
+  SequenceOperator, EncoderDecoderOperator, GraphOperator, HMMOperator)
+- **neural_components**: DiffBio-specific neural network modules
+  (GumbelSoftmaxModule, GraphMessagePassing)
+- **gnn_components**: Graph attention modules
+  (GraphAttentionLayer, GraphAttentionBlock, GATv2Layer, GATv2Block)
+- **optimal_transport**: Optimal transport solvers (SinkhornLayer)
+- **data_types**: Type aliases and protocols for type safety
+
+Usage::
+
+    from diffbio.core import soft_ops
+    from diffbio.core import TemperatureOperator
+    from diffbio.core import GraphAttentionLayer, GATv2Layer
+    from diffbio.core import SinkhornLayer
+    from diffbio.core.data_types import SequenceData
+    from diffbio.core.soft_ops import SoftBool, SoftIndex
+"""
+
+import diffbio.core.base_operators as _base_operators
+import diffbio.core.gnn_components as _gnn_components
+import diffbio.core.graph_utils as _graph_utils
+import diffbio.core.neural_components as _neural_components
+import diffbio.core.optimal_transport as _optimal_transport
+import diffbio.core.soft_ops as soft_ops  # noqa: F401 -- public submodule
+
+# Soft operation types (re-exported for convenience)
+SoftBool = soft_ops.SoftBool
+SoftIndex = soft_ops.SoftIndex
+
+# Base operators
+TemperatureOperator = _base_operators.TemperatureOperator
+SequenceOperator = _base_operators.SequenceOperator
+EncoderDecoderOperator = _base_operators.EncoderDecoderOperator
+GraphOperator = _base_operators.GraphOperator
+HMMOperator = _base_operators.HMMOperator
+
+# Neural components
+GumbelSoftmaxModule = _neural_components.GumbelSoftmaxModule
+GraphMessagePassing = _neural_components.GraphMessagePassing
+PositionalEncoding = _neural_components.PositionalEncoding
+SinusoidalPositionalEncoding = _neural_components.SinusoidalPositionalEncoding
+RoPE = _neural_components.RoPE
+ResidualBlock1D = _neural_components.ResidualBlock1D
+ResidualBlock2D = _neural_components.ResidualBlock2D
+
+# GNN components
+GraphAttentionLayer = _gnn_components.GraphAttentionLayer
+GraphAttentionBlock = _gnn_components.GraphAttentionBlock
+GATv2Layer = _gnn_components.GATv2Layer
+GATv2Block = _gnn_components.GATv2Block
+
+# Optimal transport
+SinkhornLayer = _optimal_transport.SinkhornLayer
+
+# Graph utilities
+compute_pairwise_distances = _graph_utils.compute_pairwise_distances
+compute_knn_graph = _graph_utils.compute_knn_graph
+compute_fuzzy_membership = _graph_utils.compute_fuzzy_membership
+symmetrize_graph = _graph_utils.symmetrize_graph
+
+# Data types (including SoftBool and SoftIndex from soft_ops)
+from diffbio.core.data_types import (  # noqa: E402
+    AlignmentResultData,
+    BatchArray,
+    DifferentiableOperator,
+    GraphData,
+    LatentData,
+    LossFunction,
+    MetadataDict,
+    OperatorOutput,
+    PositionWeightMatrix,
+    Probability,
+    ProbabilityArray,
+    Regularizer,
+    ScoreMatrix,
+    SequenceArray,
+    SequenceData,
+    SequenceEncoder,
+    StateDict,
+    Temperature,
+    VariantData,
+)
+
+__all__ = [  # pyright: ignore[reportUnsupportedDunderAll]
+    # Soft operations (full module)
+    "soft_ops",
+    "SoftBool",
+    "SoftIndex",
+    # Base operators
+    *_base_operators.__all__,
+    # GNN components
+    *_gnn_components.__all__,
+    # Graph utilities
+    *_graph_utils.__all__,
+    # Neural components
+    *_neural_components.__all__,
+    # Optimal transport
+    *_optimal_transport.__all__,
+    # Data types
+    "SequenceData",
+    "AlignmentResultData",
+    "VariantData",
+    "LatentData",
+    "GraphData",
+    "StateDict",
+    "MetadataDict",
+    "OperatorOutput",
+    "DifferentiableOperator",
+    "SequenceEncoder",
+    "LossFunction",
+    "Regularizer",
+    "Temperature",
+    "Probability",
+    "SequenceArray",
+    "BatchArray",
+    "ProbabilityArray",
+    "ScoreMatrix",
+    "PositionWeightMatrix",
+]