diffbio 0.1.0__py3-none-any.whl

diffbio/splitters/random.py

@@ -0,0 +1,217 @@
+"""Random and stratified splitters for DiffBio.
+
+This module provides random splitting utilities:
+- RandomSplitter: Simple random permutation-based splitting
+- StratifiedSplitter: Stratified splitting preserving class distribution
+"""
+
+import logging
+from dataclasses import dataclass
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+
+from datarax.core.data_source import DataSourceModule
+
+from diffbio.splitters.base import SplitResult, SplitterConfig, SplitterModule
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class RandomSplitterConfig(SplitterConfig):
+    """Configuration for random splitter.
+
+    Inherits all fields from SplitterConfig:
+        - train_frac: Fraction of data for training (default: 0.8)
+        - valid_frac: Fraction of data for validation (default: 0.1)
+        - test_frac: Fraction of data for testing (default: 0.1)
+        - seed: Random seed for reproducibility (optional)
+    """
+
+    pass
+
+
+class RandomSplitter(SplitterModule):
+    """Simple random splitting using JAX RNG.
+
+    Uses JAX random permutation for reproducible splits.
+    All data points are randomly assigned to train/valid/test sets
+    according to the configured fractions.
+
+    Example:
+        ```python
+        config = RandomSplitterConfig(train_frac=0.8, valid_frac=0.1, test_frac=0.1, seed=42)
+        splitter = RandomSplitter(config)
+        result = splitter.split(data_source)
+        print(f"Train size: {result.train_size}")
+        ```
+    """
+
+    def __init__(
+        self,
+        config: RandomSplitterConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize RandomSplitter.
+
+        Args:
+            config: Random splitter configuration
+            rngs: Random number generators
+            name: Optional module name
+        """
+        super().__init__(config, rngs=rngs, name=name)
+
+    def split(self, data_source: DataSourceModule) -> SplitResult:
+        """Split data source randomly.
+
+        Args:
+            data_source: Datarax DataSourceModule to split
+
+        Returns:
+            SplitResult with randomly assigned train/valid/test indices
+        """
+        n = len(data_source)
+        train_end = int(self.config.train_frac * n)
+        valid_end = int((self.config.train_frac + self.config.valid_frac) * n)
+
+        # Use JAX RNG for reproducibility
+        if self.config.seed is not None:
+            key = jax.random.key(self.config.seed)
+        elif self.rngs is not None and "split" in self.rngs:
+            key = self.rngs.split()
+        else:
+            key = jax.random.key(0)
+
+        indices = jax.random.permutation(key, jnp.arange(n))
+
+        return SplitResult(
+            train_indices=indices[:train_end],
+            valid_indices=indices[train_end:valid_end],
+            test_indices=indices[valid_end:],
+        )
+
+    def k_fold_split(
+        self, data_source: DataSourceModule, k: int = 5
+    ) -> list[tuple[jnp.ndarray, jnp.ndarray]]:
+        """K-fold cross-validation split.
+
+        Args:
+            data_source: Datarax DataSourceModule to split
+            k: Number of folds
+
+        Returns:
+            List of (train_indices, val_indices) tuples for each fold
+        """
+        n = len(data_source)
+
+        if self.config.seed is not None:
+            key = jax.random.key(self.config.seed)
+        else:
+            key = jax.random.key(0)
+
+        indices = jax.random.permutation(key, jnp.arange(n))
+        fold_size = n // k
+
+        folds = []
+        for i in range(k):
+            val_start = i * fold_size
+            val_end = (i + 1) * fold_size if i < k - 1 else n
+
+            val_indices = indices[val_start:val_end]
+            train_indices = jnp.concatenate([indices[:val_start], indices[val_end:]])
+            folds.append((train_indices, val_indices))
+
+        return folds
+
+
+@dataclass(frozen=True)
+class StratifiedSplitterConfig(SplitterConfig):
+    """Configuration for stratified splitter.
+
+    Attributes:
+        label_key: Key in data element containing labels (default: "y")
+    """
+
+    label_key: str = "y"
+
+
+class StratifiedSplitter(SplitterModule):
+    """Stratified splitting that preserves class distribution.
+
+    Ensures each split has approximately the same class distribution
+    as the original dataset. Useful for imbalanced classification tasks.
+
+    Example:
+        ```python
+        config = StratifiedSplitterConfig(seed=42, label_key="target")
+        splitter = StratifiedSplitter(config)
+        result = splitter.split(data_source)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: StratifiedSplitterConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize StratifiedSplitter.
+
+        Args:
+            config: Stratified splitter configuration
+            rngs: Random number generators
+            name: Optional module name
+        """
+        super().__init__(config, rngs=rngs, name=name)
+
+    def split(self, data_source: DataSourceModule) -> SplitResult:
+        """Split preserving class distribution.
+
+        Args:
+            data_source: Datarax DataSourceModule to split
+
+        Returns:
+            SplitResult with stratified train/valid/test indices
+        """
+        # Extract labels from data source
+        labels = jnp.array(
+            [data_source[i].data[self.config.label_key] for i in range(len(data_source))]
+        )
+
+        # Group indices by class
+        unique_labels = jnp.unique(labels)
+        class_indices = {int(label): jnp.where(labels == label)[0] for label in unique_labels}
+
+        # Use JAX RNG
+        if self.config.seed is not None:
+            key = jax.random.key(self.config.seed)
+        else:
+            key = jax.random.key(0)
+
+        train_inds: list[jnp.ndarray] = []
+        valid_inds: list[jnp.ndarray] = []
+        test_inds: list[jnp.ndarray] = []
+
+        for _label, indices in class_indices.items():
+            key, subkey = jax.random.split(key)
+            shuffled = jax.random.permutation(subkey, indices)
+
+            n_class = len(shuffled)
+            train_end = int(self.config.train_frac * n_class)
+            valid_end = int((self.config.train_frac + self.config.valid_frac) * n_class)
+
+            train_inds.append(shuffled[:train_end])
+            valid_inds.append(shuffled[train_end:valid_end])
+            test_inds.append(shuffled[valid_end:])
+
+        empty = jnp.array([], dtype=jnp.int32)
+        return SplitResult(
+            train_indices=jnp.concatenate(train_inds) if train_inds else empty,
+            valid_indices=jnp.concatenate(valid_inds) if valid_inds else empty,
+            test_indices=jnp.concatenate(test_inds) if test_inds else empty,
+        )

diffbio/splitters/sequence.py

@@ -0,0 +1,201 @@
+"""Sequence identity splitter for bioinformatics applications.
+
+This module provides sequence-aware splitting utilities:
+- SequenceIdentitySplitter: Split by sequence identity clustering
+
+For genomics/proteomics applications where similar sequences
+should not appear in both train and test sets.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Sequence
+
+from flax import nnx
+
+from datarax.core.data_source import DataSourceModule
+
+from diffbio.splitters.base import SplitResult, SplitterConfig, SplitterModule
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class SequenceIdentitySplitterConfig(SplitterConfig):
+    """Configuration for sequence identity splitter.
+
+    Attributes:
+        sequence_key: Key in data element containing sequence string (default: "sequence")
+        identity_threshold: Identity threshold for clustering (default: 0.3)
+            Sequences with identity > threshold are clustered together.
+        alignment_method: Method for identity computation ("simple" or "mmseqs2")
+    """
+
+    sequence_key: str = "sequence"
+    identity_threshold: float = 0.3
+    alignment_method: str = "simple"
+
+
+class SequenceIdentitySplitter(SplitterModule):
+    """Split sequences by identity threshold.
+
+    Groups similar sequences together using identity clustering,
+    then assigns clusters to train/valid/test to ensure structural
+    diversity between splits. This prevents data leakage from
+    similar sequences appearing in different splits.
+
+    Inherits from SplitterModule (StructuralModule) because:
+
+    - Non-parametric: clustering is deterministic
+    - Frozen config: splitting strategy doesn't change
+    - Domain-specific: requires sequence comparison
+
+    Similar to CD-HIT or MMseqs2 clustering approach.
+
+    Example:
+        ```python
+        config = SequenceIdentitySplitterConfig(identity_threshold=0.3)
+        splitter = SequenceIdentitySplitter(config)
+        result = splitter.split(sequence_source)
+        ```
+
+    References:
+        Li, Weizhong, and Adam Godzik. "Cd-hit: a fast program for clustering
+        and comparing large sets of protein or nucleotide sequences."
+        Bioinformatics 22.13 (2006): 1658-1659.
+    """
+
+    def __init__(
+        self,
+        config: SequenceIdentitySplitterConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize SequenceIdentitySplitter.
+
+        Args:
+            config: Sequence identity splitter configuration
+            rngs: Random number generators (unused for identity splitting)
+            name: Optional module name
+        """
+        super().__init__(config, rngs=rngs, name=name)
+
+    def _compute_identity(self, seq1: str, seq2: str) -> float:
+        """Compute sequence identity between two sequences.
+
+        Uses simple character matching. For unequal lengths,
+        compares up to the length of the shorter sequence.
+
+        Args:
+            seq1: First sequence
+            seq2: Second sequence
+
+        Returns:
+            Identity fraction between 0.0 and 1.0
+        """
+        if not seq1 or not seq2:
+            return 0.0
+
+        # Use shorter sequence for comparison
+        min_len = min(len(seq1), len(seq2))
+        seq1 = seq1[:min_len]
+        seq2 = seq2[:min_len]
+
+        if not seq1:
+            return 0.0
+
+        matches = sum(c1 == c2 for c1, c2 in zip(seq1, seq2))
+        return matches / len(seq1)
+
+    def _cluster_by_identity(self, sequences: Sequence[str]) -> list[list[int]]:
+        """Cluster sequences by identity threshold.
+
+        Uses greedy clustering: each sequence joins the first cluster
+        where it has identity > threshold with the representative.
+
+        Args:
+            sequences: List of sequence strings
+
+        Returns:
+            List of clusters, each cluster is a list of sequence indices
+        """
+        if self.config.alignment_method == "simple":
+            return self._simple_clustering(sequences)
+        elif self.config.alignment_method == "mmseqs2":
+            return self._mmseqs2_clustering(sequences)
+        else:
+            raise ValueError(f"Unknown alignment method: {self.config.alignment_method}")
+
+    def _simple_clustering(self, sequences: Sequence[str]) -> list[list[int]]:
+        """Simple greedy clustering by identity.
+
+        Each sequence joins the first cluster where it has
+        identity > threshold with the representative.
+
+        Args:
+            sequences: List of sequence strings
+
+        Returns:
+            List of clusters
+        """
+        clusters: list[list[int]] = []
+        representatives: list[str] = []
+
+        for idx, seq in enumerate(sequences):
+            assigned = False
+
+            for cluster_idx, rep in enumerate(representatives):
+                identity = self._compute_identity(seq, rep)
+                if identity > self.config.identity_threshold:
+                    clusters[cluster_idx].append(idx)
+                    assigned = True
+                    break
+
+            if not assigned:
+                clusters.append([idx])
+                representatives.append(seq)
+
+        return clusters
+
+    def _mmseqs2_clustering(self, sequences: Sequence[str]) -> list[list[int]]:
+        """Use MMseqs2 for clustering.
+
+        Requires MMseqs2 installation.
+
+        Args:
+            sequences: List of sequence strings
+
+        Returns:
+            List of clusters
+
+        Raises:
+            NotImplementedError: MMseqs2 integration not yet implemented
+        """
+        raise NotImplementedError(
+            "MMseqs2 clustering requires external tool installation. "
+            "Use alignment_method='simple' for built-in clustering."
+        )
+
+    def split(self, data_source: DataSourceModule) -> SplitResult:
+        """Split by sequence identity clustering.
+
+        Clusters sequences by identity, then assigns clusters
+        to train/valid/test splits. Largest clusters go to
+        train first.
+
+        Args:
+            data_source: Datarax DataSourceModule to split
+
+        Returns:
+            SplitResult with identity-based train/valid/test indices
+        """
+        # Extract sequences from data source
+        sequences = [data_source[i].data[self.config.sequence_key] for i in range(len(data_source))]
+
+        # Cluster by identity
+        clusters = self._cluster_by_identity(sequences)
+
+        # Sort clusters by size (largest first)
+        sorted_clusters = sorted(clusters, key=len, reverse=True)
+        return self.assign_groups_to_splits(sorted_clusters, len(data_source))

diffbio/utils/__init__.py

@@ -0,0 +1,55 @@
+"""Utility functions for DiffBio.
+
+This module provides utility functions for I/O, encoding, training,
+neural network building, and other common operations in bioinformatics pipelines.
+"""
+
+from diffbio.utils.dependency_runtime import (
+    ECOSYSTEM_PACKAGES,
+    DependencyRuntimeRecord,
+    FNOConstructorContract,
+    collect_dependency_runtime,
+    inspect_fno_constructor,
+    verify_canonical_dependency_runtime,
+)
+from diffbio.utils.quality import apply_quality_filter
+from diffbio.utils.nn_utils import (
+    ensure_rngs,
+    extract_windows_1d,
+    get_rng_key,
+    init_learnable_param,
+)
+from diffbio.utils.training import (
+    Trainer,
+    TrainingConfig,
+    TrainingState,
+    create_optax_optimizer,
+    create_synthetic_training_data,
+    cross_entropy_loss,
+    data_iterator,
+)
+
+__all__ = [
+    # Dependency runtime utilities
+    "ECOSYSTEM_PACKAGES",
+    "DependencyRuntimeRecord",
+    "FNOConstructorContract",
+    "collect_dependency_runtime",
+    "inspect_fno_constructor",
+    "verify_canonical_dependency_runtime",
+    # Training utilities
+    "Trainer",
+    "TrainingConfig",
+    "TrainingState",
+    "create_optax_optimizer",
+    "create_synthetic_training_data",
+    "cross_entropy_loss",
+    "data_iterator",
+    # Quality utilities
+    "apply_quality_filter",
+    # Neural network utilities
+    "ensure_rngs",
+    "extract_windows_1d",
+    "get_rng_key",
+    "init_learnable_param",
+]

diffbio/utils/dependency_runtime.py

@@ -0,0 +1,115 @@
+"""Helpers for verifying the installed ecosystem runtime."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+import importlib
+import inspect
+from pathlib import Path
+import site
+from types import ModuleType
+
+ECOSYSTEM_PACKAGES: tuple[str, ...] = ("datarax", "artifex", "opifex", "calibrax")
+
+
+@dataclass(frozen=True)
+class DependencyRuntimeRecord:
+    """Resolved runtime provenance for one ecosystem package."""
+
+    package: str
+    module_file: str
+    installed_from_site_packages: bool
+
+
+@dataclass(frozen=True)
+class FNOConstructorContract:
+    """Observed constructor contract for the live Opifex FNO surface."""
+
+    import_path: str
+    constructor_signature: str
+    supports_spatial_dims: bool
+
+
+def _site_packages_roots() -> tuple[Path, ...]:
+    """Return normalized site-packages roots for the active interpreter."""
+    return tuple(Path(root).resolve() for root in site.getsitepackages())
+
+
+def _resolve_module_file(module: ModuleType) -> Path:
+    """Return the concrete module file for an imported package.
+
+    Raises:
+        RuntimeError: If the imported module does not expose a file path.
+    """
+    module_file = getattr(module, "__file__", None)
+    if module_file is None:
+        msg = f"Imported module {module.__name__!r} does not expose __file__"
+        raise RuntimeError(msg)
+    return Path(module_file).resolve()
+
+
+def _is_in_site_packages(module_file: Path, site_roots: Sequence[Path]) -> bool:
+    """Return whether a module file resolves under one of the site-packages roots."""
+    return any(module_file.is_relative_to(root) for root in site_roots)
+
+
+def collect_dependency_runtime(
+    package_names: Sequence[str] = ECOSYSTEM_PACKAGES,
+) -> dict[str, DependencyRuntimeRecord]:
+    """Collect import provenance for the configured ecosystem packages."""
+    site_roots = _site_packages_roots()
+    runtime: dict[str, DependencyRuntimeRecord] = {}
+
+    for package_name in package_names:
+        module = importlib.import_module(package_name)
+        module_file = _resolve_module_file(module)
+        runtime[package_name] = DependencyRuntimeRecord(
+            package=package_name,
+            module_file=str(module_file),
+            installed_from_site_packages=_is_in_site_packages(module_file, site_roots),
+        )
+
+    return runtime
+
+
+def inspect_fno_constructor(
+    import_path: str = "opifex.neural.operators",
+) -> FNOConstructorContract:
+    """Inspect the live FourierNeuralOperator constructor contract."""
+    module = importlib.import_module(import_path)
+    constructor_signature = str(inspect.signature(module.FourierNeuralOperator.__init__))
+    return FNOConstructorContract(
+        import_path=import_path,
+        constructor_signature=constructor_signature,
+        supports_spatial_dims="spatial_dims" in constructor_signature,
+    )
+
+
+def verify_canonical_dependency_runtime(
+    package_names: Sequence[str] = ECOSYSTEM_PACKAGES,
+) -> tuple[dict[str, DependencyRuntimeRecord], FNOConstructorContract]:
+    """Validate the canonical installed runtime contract for ecosystem dependencies.
+
+    Raises:
+        RuntimeError: If any ecosystem package resolves outside site-packages or
+            if the live Opifex FNO constructor lacks ``spatial_dims`` support.
+    """
+    runtime = collect_dependency_runtime(package_names)
+    non_installed_packages = sorted(
+        package for package, record in runtime.items() if not record.installed_from_site_packages
+    )
+    if non_installed_packages:
+        package_list = ", ".join(non_installed_packages)
+        msg = (
+            "Canonical runtime must resolve ecosystem dependencies from installed "
+            f"site-packages; found non-installed imports for: {package_list}"
+        )
+        raise RuntimeError(msg)
+
+    fno_contract = inspect_fno_constructor()
+    if not fno_contract.supports_spatial_dims:
+        msg = "Live Opifex FourierNeuralOperator constructor does not expose spatial_dims"
+        raise RuntimeError(msg)
+
+    return runtime, fno_contract