diffbio 0.1.0__py3-none-any.whl

diffbio/sources/sequence_foundation.py

@@ -0,0 +1,68 @@
+"""Sequence embedding sources following the Datarax source model."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import jax.numpy as jnp
+from flax import nnx
+
+from diffbio.sources.indexed_embeddings import IndexedEmbeddingSource, IndexedEmbeddingSourceConfig
+
+
+@dataclass(frozen=True)
+class SequenceEmbeddingSourceConfig(IndexedEmbeddingSourceConfig):
+    """Configuration for sequence embedding artifacts."""
+
+    row_id_key: str = "sequence_ids"
+
+
+class SequenceEmbeddingSource(IndexedEmbeddingSource):
+    """Indexed embedding source specialized for sequence artifacts."""
+
+    config: SequenceEmbeddingSourceConfig  # pyright: ignore[reportIncompatibleVariableOverride]
+
+    @property
+    def sequence_ids(self) -> tuple[str, ...] | None:
+        """Tuple of persisted sequence identifiers, if present."""
+        return self.row_ids
+
+    def align_to_reference_sequence_ids(
+        self,
+        *,
+        reference_sequence_ids: list[str] | tuple[str, ...],
+        require_sequence_ids: bool = True,
+    ) -> jnp.ndarray:
+        """Align external embeddings to the benchmark sequence order."""
+        return self.align_to_reference_ids(
+            reference_ids=reference_sequence_ids,
+            require_row_ids=require_sequence_ids,
+            artifact_label="Sequence",
+            id_display_name="Sequence ID",
+        )
+
+
+def load_sequence_embedding_source(
+    path: Path | str,
+    *,
+    rngs: nnx.Rngs | None = None,
+) -> SequenceEmbeddingSource:
+    """Build the canonical sequence embedding source for an artifact."""
+    return SequenceEmbeddingSource(
+        SequenceEmbeddingSourceConfig(file_path=str(path)),
+        rngs=rngs,
+    )
+
+
+def align_sequence_embeddings(
+    *,
+    reference_sequence_ids: list[str] | tuple[str, ...],
+    artifact_path: Path | str,
+    require_sequence_ids: bool = True,
+) -> jnp.ndarray:
+    """Align external embeddings to the benchmark sequence order."""
+    return load_sequence_embedding_source(artifact_path).align_to_reference_sequence_ids(
+        reference_sequence_ids=reference_sequence_ids,
+        require_sequence_ids=require_sequence_ids,
+    )

diffbio/sources/singlecell_foundation.py

@@ -0,0 +1,68 @@
+"""Single-cell embedding sources following the Datarax source model."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import jax.numpy as jnp
+from flax import nnx
+
+from diffbio.sources.indexed_embeddings import IndexedEmbeddingSource, IndexedEmbeddingSourceConfig
+
+
+@dataclass(frozen=True)
+class SingleCellEmbeddingSourceConfig(IndexedEmbeddingSourceConfig):
+    """Configuration for single-cell embedding artifacts."""
+
+    row_id_key: str = "cell_ids"
+
+
+class SingleCellEmbeddingSource(IndexedEmbeddingSource):
+    """Indexed embedding source specialized for single-cell artifacts."""
+
+    config: SingleCellEmbeddingSourceConfig  # pyright: ignore[reportIncompatibleVariableOverride]
+
+    @property
+    def cell_ids(self) -> tuple[str, ...] | None:
+        """Tuple of persisted cell identifiers, if present."""
+        return self.row_ids
+
+    def align_to_reference_cell_ids(
+        self,
+        *,
+        reference_cell_ids: list[str] | tuple[str, ...],
+        require_cell_ids: bool = True,
+    ) -> jnp.ndarray:
+        """Align external embeddings to the benchmark cell order."""
+        return self.align_to_reference_ids(
+            reference_ids=reference_cell_ids,
+            require_row_ids=require_cell_ids,
+            artifact_label="Single-cell",
+            id_display_name="Cell ID",
+        )
+
+
+def load_singlecell_embedding_source(
+    path: Path | str,
+    *,
+    rngs: nnx.Rngs | None = None,
+) -> SingleCellEmbeddingSource:
+    """Build the canonical single-cell embedding source for an artifact."""
+    return SingleCellEmbeddingSource(
+        SingleCellEmbeddingSourceConfig(file_path=str(path)),
+        rngs=rngs,
+    )
+
+
+def align_singlecell_embeddings(
+    *,
+    reference_cell_ids: list[str] | tuple[str, ...],
+    artifact_path: Path | str,
+    require_cell_ids: bool = True,
+) -> jnp.ndarray:
+    """Align external embeddings to the benchmark cell order."""
+    return load_singlecell_embedding_source(artifact_path).align_to_reference_cell_ids(
+        reference_cell_ids=reference_cell_ids,
+        require_cell_ids=require_cell_ids,
+    )

diffbio/splitters/__init__.py

@@ -0,0 +1,63 @@
+"""DiffBio splitters module.
+
+This module provides dataset splitting utilities extending Datarax's StructuralModule
+for train/validation/test splitting in bioinformatics and drug discovery applications.
+
+Splitters:
+    SplitterModule: Base class for all splitters
+    SplitResult: NamedTuple containing split indices
+    RandomSplitter: Simple random splitting
+    StratifiedSplitter: Stratified splitting preserving class distribution
+    ScaffoldSplitter: Molecular scaffold-based splitting for drug discovery
+    TanimotoClusterSplitter: Fingerprint similarity clustering for drug discovery
+    SequenceIdentitySplitter: Sequence identity clustering for bioinformatics
+"""
+
+from diffbio.splitters.base import SplitResult, SplitterConfig, SplitterModule
+from diffbio.splitters.molecular import (
+    ScaffoldSplitter,
+    ScaffoldSplitterConfig,
+    TanimotoClusterSplitter,
+    TanimotoClusterSplitterConfig,
+)
+from diffbio.splitters.random import (
+    RandomSplitter,
+    RandomSplitterConfig,
+    StratifiedSplitter,
+    StratifiedSplitterConfig,
+)
+from diffbio.splitters.perturbation import (
+    FewShotSplitter,
+    FewShotSplitterConfig,
+    ZeroShotSplitter,
+    ZeroShotSplitterConfig,
+)
+from diffbio.splitters.sequence import (
+    SequenceIdentitySplitter,
+    SequenceIdentitySplitterConfig,
+)
+
+__all__ = [
+    # Base classes
+    "SplitterModule",
+    "SplitterConfig",
+    "SplitResult",
+    # Random splitters
+    "RandomSplitter",
+    "RandomSplitterConfig",
+    "StratifiedSplitter",
+    "StratifiedSplitterConfig",
+    # Molecular splitters
+    "ScaffoldSplitter",
+    "ScaffoldSplitterConfig",
+    "TanimotoClusterSplitter",
+    "TanimotoClusterSplitterConfig",
+    # Sequence splitters
+    "SequenceIdentitySplitter",
+    "SequenceIdentitySplitterConfig",
+    # Perturbation splitters
+    "FewShotSplitter",
+    "FewShotSplitterConfig",
+    "ZeroShotSplitter",
+    "ZeroShotSplitterConfig",
+]

diffbio/splitters/base.py

@@ -0,0 +1,251 @@
+"""Base splitter classes for DiffBio.
+
+This module provides the base classes for dataset splitting:
+- SplitResult: NamedTuple containing train/valid/test indices
+- SplitterConfig: Base configuration for splitters
+- SplitterModule: Base class for all splitters
+"""
+
+import logging
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import NamedTuple
+
+import jax.numpy as jnp
+import numpy as np
+from flax import nnx
+
+from datarax.core.config import StructuralConfig
+from datarax.core.data_source import DataSourceModule
+from datarax.core.structural import StructuralModule
+
+logger = logging.getLogger(__name__)
+
+
+class SplitResult(NamedTuple):
+    """Result of a dataset split operation.
+
+    Attributes:
+        train_indices: Array of indices for training set
+        valid_indices: Array of indices for validation set
+        test_indices: Array of indices for test set
+    """
+
+    train_indices: jnp.ndarray
+    valid_indices: jnp.ndarray
+    test_indices: jnp.ndarray
+
+    @property
+    def train_size(self) -> int:
+        """Return number of training samples."""
+        return len(self.train_indices)
+
+    @property
+    def valid_size(self) -> int:
+        """Return number of validation samples."""
+        return len(self.valid_indices)
+
+    @property
+    def test_size(self) -> int:
+        """Return number of test samples."""
+        return len(self.test_indices)
+
+
+@dataclass(frozen=True)
+class SplitterConfig(StructuralConfig):
+    """Base configuration for splitters.
+
+    Frozen because splitters are non-parametric (StructuralModule).
+
+    Attributes:
+        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)
+    """
+
+    train_frac: float = 0.8
+    valid_frac: float = 0.1
+    test_frac: float = 0.1
+    seed: int | None = None
+
+    def __post_init__(self):
+        """Validate configuration after initialization."""
+        super().__post_init__()
+        total = self.train_frac + self.valid_frac + self.test_frac
+        if not np.isclose(total, 1.0):
+            raise ValueError(f"Split fractions must sum to 1.0, got {total}")
+
+
+class SplitterModule(StructuralModule):
+    """Base class for dataset splitters.
+
+    Inherits from StructuralModule because:
+
+    - Non-parametric (no learnable parameters)
+    - Frozen config (splitting strategy is fixed)
+    - Uses process() method pattern
+    - Integrates with Datarax data sources
+
+    Splitters divide data into train/valid/test sets, while Datarax
+    SamplerModule controls iteration ORDER within those sets.
+
+    Args:
+        config: Splitter configuration
+        rngs: Random number generators for stochastic splitting
+        name: Optional name for the module
+    """
+
+    def __init__(
+        self,
+        config: SplitterConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize SplitterModule.
+
+        Args:
+            config: 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 a data source into train/valid/test indices.
+
+        Subclasses must implement this method.
+
+        Args:
+            data_source: Datarax DataSourceModule to split
+
+        Returns:
+            SplitResult with train/valid/test indices
+        """
+        raise NotImplementedError("Subclasses must implement split()")
+
+    def process(self, data_source: DataSourceModule) -> SplitResult:
+        """Process data source using the split method.
+
+        Args:
+            data_source: Datarax DataSourceModule to split
+
+        Returns:
+            SplitResult with train/valid/test indices
+        """
+        return self.split(data_source)
+
+    def k_fold_split(
+        self, data_source: DataSourceModule, k: int = 5
+    ) -> list[tuple[jnp.ndarray, jnp.ndarray]]:
+        """K-fold cross-validation split.
+
+        Subclasses may implement this method.
+
+        Args:
+            data_source: Datarax DataSourceModule to split
+            k: Number of folds
+
+        Returns:
+            List of (train_indices, val_indices) tuples for each fold
+        """
+        raise NotImplementedError("Subclasses may implement k_fold_split()")
+
+    def assign_groups_to_splits(
+        self,
+        groups: Iterable[Iterable[int]],
+        total_size: int,
+    ) -> SplitResult:
+        """Assign grouped indices to train/valid/test by configured fractions.
+
+        Each input group is placed wholly into one split, preserving group boundaries
+        (for example scaffold clusters or sequence-identity clusters).
+        """
+        train_cutoff = self.config.train_frac * total_size
+        valid_cutoff = (self.config.train_frac + self.config.valid_frac) * total_size
+
+        train_inds: list[int] = []
+        valid_inds: list[int] = []
+        test_inds: list[int] = []
+
+        for group in groups:
+            group_indices = list(group)
+            if len(train_inds) < train_cutoff:
+                train_inds.extend(group_indices)
+            elif len(train_inds) + len(valid_inds) < valid_cutoff:
+                valid_inds.extend(group_indices)
+            else:
+                test_inds.extend(group_indices)
+
+        return SplitResult(
+            train_indices=jnp.array(train_inds, dtype=jnp.int32),
+            valid_indices=jnp.array(valid_inds, dtype=jnp.int32),
+            test_indices=jnp.array(test_inds, dtype=jnp.int32),
+        )
+
+    def create_split_sources(
+        self,
+        data_source: DataSourceModule,
+        split_result: SplitResult | None = None,
+        lazy: bool = True,
+    ) -> tuple[DataSourceModule, DataSourceModule, DataSourceModule]:
+        """Create separate data sources for each split.
+
+        This creates views into the original data source using the split indices.
+        Each returned source can be used with Datarax samplers independently.
+
+        Args:
+            data_source: Original data source
+            split_result: Pre-computed split (or compute if None)
+            lazy: If True, use lazy loading (IndexedViewSource). If False,
+                  eagerly load into MemorySource (faster iteration but uses memory).
+
+        Returns:
+            Tuple of (train_source, valid_source, test_source)
+        """
+        if split_result is None:
+            split_result = self.split(data_source)
+
+        if lazy:
+            # LAZY LOADING: Create view sources that delegate to original
+            from diffbio.sources.indexed_view import (
+                IndexedViewSource,
+                IndexedViewSourceConfig,
+            )
+
+            train_config = IndexedViewSourceConfig(shuffle=True, seed=self.config.seed)
+            valid_config = IndexedViewSourceConfig(shuffle=False)
+            test_config = IndexedViewSourceConfig(shuffle=False)
+
+            return (
+                IndexedViewSource(
+                    train_config, data_source, split_result.train_indices, rngs=self.rngs
+                ),
+                IndexedViewSource(
+                    valid_config, data_source, split_result.valid_indices, rngs=self.rngs
+                ),
+                IndexedViewSource(
+                    test_config, data_source, split_result.test_indices, rngs=self.rngs
+                ),
+            )
+        else:
+            # EAGER LOADING: Load all elements into memory (faster iteration)
+            from datarax.sources import MemorySource, MemorySourceConfig
+
+            train_elements = [data_source[int(i)] for i in split_result.train_indices]
+            valid_elements = [data_source[int(i)] for i in split_result.valid_indices]
+            test_elements = [data_source[int(i)] for i in split_result.test_indices]
+
+            train_config = MemorySourceConfig(
+                shuffle=True,
+                seed=self.config.seed,  # pyright: ignore[reportCallIssue]
+            )
+            valid_config = MemorySourceConfig(shuffle=False)
+            test_config = MemorySourceConfig(shuffle=False)
+
+            return (
+                MemorySource(train_config, data=train_elements, rngs=self.rngs),
+                MemorySource(valid_config, data=valid_elements, rngs=self.rngs),
+                MemorySource(test_config, data=test_elements, rngs=self.rngs),
+            )