diffbio 0.1.0__py3-none-any.whl

diffbio/sequences/dna.py

@@ -0,0 +1,239 @@
+"""DNA sequence data types and encoding utilities for DiffBio.
+
+This module provides functions and utilities for working with DNA sequences
+in a JAX-compatible, differentiable manner.
+"""
+
+from typing import Literal
+
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array, Float
+
+
+# DNA nucleotide alphabet
+DNA_ALPHABET = "ACGT"
+DNA_ALPHABET_SIZE = 4
+
+# Mapping from nucleotide to index
+_NUC_TO_IDX = {"A": 0, "C": 1, "G": 2, "T": 3, "N": -1}
+_IDX_TO_NUC = {0: "A", 1: "C", 2: "G", 3: "T"}
+
+
+def encode_dna_string(sequence: str, handle_n: Literal["uniform", "zero"] = "uniform") -> Array:
+    """Encode a DNA string as a one-hot JAX array.
+
+    Args:
+        sequence: DNA string containing only A, C, G, T, N characters.
+        handle_n: How to handle N (unknown) nucleotides:
+            - "uniform": Encode as uniform distribution [0.25, 0.25, 0.25, 0.25]
+            - "zero": Encode as zeros [0, 0, 0, 0]
+
+    Returns:
+        One-hot encoded array of shape (len(sequence), 4).
+        Columns represent A, C, G, T in that order.
+
+    Example:
+        ```python
+        encode_dna_string("ACGT")
+        ```
+        Array([[1, 0, 0, 0],
+               [0, 1, 0, 0],
+               [0, 0, 1, 0],
+               [0, 0, 0, 1]], dtype=float32)
+    """
+    sequence = sequence.upper()
+
+    # Convert to indices
+    indices = []
+    for nuc in sequence:
+        if nuc in _NUC_TO_IDX:
+            indices.append(_NUC_TO_IDX[nuc])
+        else:
+            raise ValueError(f"Invalid nucleotide: {nuc}. Expected A, C, G, T, or N.")
+
+    indices_array = jnp.array(indices, dtype=jnp.int32)
+
+    # One-hot encode valid nucleotides (indices 0-3)
+    # N nucleotides have index -1 and will be handled separately
+    valid_mask = indices_array >= 0
+    safe_indices = jnp.where(valid_mask, indices_array, 0)  # Use 0 for safe indexing
+    one_hot = jax.nn.one_hot(safe_indices, DNA_ALPHABET_SIZE)
+
+    # Handle N nucleotides
+    if handle_n == "uniform":
+        n_encoding = jnp.ones(DNA_ALPHABET_SIZE) / DNA_ALPHABET_SIZE
+    else:  # "zero"
+        n_encoding = jnp.zeros(DNA_ALPHABET_SIZE)
+
+    # Apply N encoding where needed
+    n_mask = ~valid_mask
+    one_hot = jnp.where(n_mask[:, None], n_encoding, one_hot)
+
+    return one_hot
+
+
+def decode_dna_onehot(encoded: Array, threshold: float = 0.5) -> str:
+    """Decode a one-hot encoded DNA array back to a string.
+
+    Args:
+        encoded: One-hot encoded array of shape (length, 4).
+        threshold: Minimum confidence to assign a nucleotide.
+            Below threshold, returns 'N'.
+
+    Returns:
+        DNA string.
+    """
+    # Get argmax for each position
+    indices = jnp.argmax(encoded, axis=-1)
+    max_vals = jnp.max(encoded, axis=-1)
+
+    # Build string
+    result = []
+    for i in range(len(indices)):
+        idx = int(indices[i])
+        max_val = float(max_vals[i])
+        if max_val >= threshold and idx in _IDX_TO_NUC:
+            result.append(_IDX_TO_NUC[idx])
+        else:
+            result.append("N")
+
+    return "".join(result)
+
+
+def phred_to_probability(phred_scores: Array) -> Array:
+    """Convert Phred quality scores to error probabilities.
+
+    Args:
+        phred_scores: Array of Phred scores (typically 0-40).
+
+    Returns:
+        Array of error probabilities in range [0, 1].
+
+    Note:
+        Error probability = 10^(-Q/10) where Q is Phred score.
+    """
+    return jnp.power(10.0, -phred_scores / 10.0)
+
+
+def probability_to_phred(error_prob: Array, max_phred: float = 60.0) -> Array:
+    """Convert error probabilities to Phred quality scores.
+
+    Args:
+        error_prob: Array of error probabilities in range (0, 1].
+        max_phred: Maximum Phred score to return (for numerical stability).
+
+    Returns:
+        Array of Phred scores.
+    """
+    # Clip to avoid log(0)
+    safe_prob = jnp.clip(error_prob, 1e-10, 1.0)
+    phred = -10.0 * jnp.log10(safe_prob)
+    return jnp.clip(phred, 0.0, max_phred)
+
+
+def soft_encode_dna(
+    sequence: Array,
+    quality_scores: Array,
+    temperature: float = 1.0,
+) -> Array:
+    """Create soft one-hot encoding weighted by quality scores.
+
+    This creates a differentiable encoding where positions with low quality
+    have more uniform distributions (higher entropy).
+
+    Args:
+        sequence: One-hot encoded sequence of shape (length, 4).
+        quality_scores: Phred quality scores of shape (length,).
+        temperature: Temperature parameter controlling softness.
+            Lower = sharper (more like hard one-hot)
+            Higher = softer (more uniform)
+
+    Returns:
+        Soft-encoded array of shape (length, 4) where each row
+        sums to 1 but may not be exactly one-hot.
+    """
+    # Convert quality to confidence (higher quality = higher confidence)
+    confidence = 1.0 - phred_to_probability(quality_scores)
+
+    # Scale one-hot by confidence and add uniform noise for uncertainty
+    uniform = jnp.ones(DNA_ALPHABET_SIZE) / DNA_ALPHABET_SIZE
+
+    # Weighted combination: high confidence -> one-hot, low confidence -> uniform
+    soft_encoded = confidence[:, None] * sequence + (1 - confidence[:, None]) * uniform
+
+    # Apply temperature scaling and softmax for normalization
+    logits = jnp.log(soft_encoded + 1e-10) / temperature
+    return jax.nn.softmax(logits, axis=-1)
+
+
+def complement_dna(encoded: Array) -> Array:
+    """Get the complement of a DNA sequence.
+
+    Complement mapping: A<->T, C<->G
+
+    Args:
+        encoded: One-hot encoded DNA of shape (..., 4).
+
+    Returns:
+        Complemented sequence of same shape.
+    """
+    # Complement swaps A<->T (indices 0<->3) and C<->G (indices 1<->2)
+    # Permutation: [0,1,2,3] -> [3,2,1,0]
+    return encoded[..., ::-1]
+
+
+def reverse_complement_dna(encoded: Array) -> Array:
+    """Get the reverse complement of a DNA sequence.
+
+    Args:
+        encoded: One-hot encoded DNA of shape (length, 4).
+
+    Returns:
+        Reverse complemented sequence of same shape.
+    """
+    return complement_dna(encoded[::-1])
+
+
+def gc_content(encoded: Array) -> Float[Array, ""]:
+    """Calculate GC content of a DNA sequence.
+
+    Args:
+        encoded: One-hot encoded DNA of shape (length, 4).
+
+    Returns:
+        Scalar GC content as fraction in [0, 1].
+    """
+    # G is index 2, C is index 1
+    gc_sum = jnp.sum(encoded[:, 1] + encoded[:, 2])
+    total = jnp.sum(encoded)
+    return gc_sum / (total + 1e-10)
+
+
+def create_dna_element_data(
+    sequence: str | Array,
+    quality_scores: Array | None = None,
+) -> dict:
+    """Create data dictionary for a DNA sequence Element.
+
+    This creates a data structure compatible with Datarax's Element.
+
+    Args:
+        sequence: DNA string or pre-encoded one-hot array.
+        quality_scores: Optional Phred quality scores.
+
+    Returns:
+        Dictionary suitable for Element(data=...).
+    """
+    # Encode sequence if string
+    if isinstance(sequence, str):
+        encoded_seq = encode_dna_string(sequence)
+    else:
+        encoded_seq = sequence
+
+    data = {"sequence": encoded_seq}
+
+    if quality_scores is not None:
+        data["quality_scores"] = quality_scores
+
+    return data

diffbio/sources/__init__.py

@@ -0,0 +1,149 @@
+"""DiffBio data sources module.
+
+This module provides data source implementations extending Datarax's DataSourceModule
+for bioinformatics and drug discovery applications.
+
+Sources:
+    AnnDataSource: AnnData (.h5ad) file reading for single-cell data
+    ENCODEPeakSource: ENCODE narrowPeak BED file reading for ChIP-seq peaks
+    IndexedViewSource: Lazy-loading view into a data source using index mapping
+    MolNetSource: MoleculeNet benchmark datasets for drug discovery
+    BAMSource: BAM/CRAM file reading for aligned sequencing reads
+    BioSNAPDTISource: Deterministic binary DTI benchmark source
+    DavisDTISource: Deterministic affinity DTI benchmark source
+    FastaSource: FASTA file reading for DNA/RNA sequences
+
+Interop:
+    to_anndata: Convert DiffBio data dict to AnnData object
+    from_anndata: Convert AnnData object to DiffBio data dict
+"""
+
+from diffbio.sources.anndata_interop import from_anndata, to_anndata
+from diffbio.sources.anndata_source import AnnDataSource, AnnDataSourceConfig
+from diffbio.sources.bam import BAMSource, BAMSourceConfig
+from diffbio.sources.contextual_epigenomics import (
+    CONTEXTUAL_EPIGENOMICS_DATASET_CONTRACT_KEYS,
+    CONTEXTUAL_TARGET_SEMANTICS,
+    build_synthetic_contextual_epigenomics_dataset,
+    validate_contextual_epigenomics_dataset,
+)
+from diffbio.sources.dti import (
+    DTI_DATASET_CONTRACT_KEYS,
+    BioSNAPDTISource,
+    DTISourceConfig,
+    DavisDTISource,
+    build_paired_dti_batch,
+    deterministic_dti_split,
+    validate_dti_dataset,
+)
+from diffbio.sources.embeddings import (
+    EmbeddingArtifactSource,
+    EmbeddingArtifactSourceConfig,
+)
+from diffbio.sources.encode_peaks import ENCODEPeakConfig, ENCODEPeakSource
+from diffbio.sources.fasta import FastaSource, FastaSourceConfig
+from diffbio.sources.indexed_view import IndexedViewSource, IndexedViewSourceConfig
+from diffbio.sources.molnet import MolNetSource, MolNetSourceConfig
+from diffbio.sources.multiomics import (
+    MULTIOMICS_ARTIFACT_METADATA_KEYS,
+    MULTIOMICS_DATASET_PROVENANCE_KEYS,
+    MetabolomicsEmbeddingSource,
+    MetabolomicsEmbeddingSourceConfig,
+    MultiOmicsEmbeddingSource,
+    MultiOmicsEmbeddingSourceConfig,
+    align_metabolomics_embeddings,
+    align_multiomics_embeddings,
+    build_multiomics_artifact_metadata,
+    build_multiomics_dataset_provenance,
+    load_metabolomics_embedding_source,
+    load_multiomics_embedding_source,
+    validate_multiomics_artifact_metadata,
+    validate_multiomics_dataset_provenance,
+)
+from diffbio.sources.perturbation import (
+    BatchControlMapping,
+    ControlMappingConfig,
+    ExperimentConfig,
+    GlobalH5MetadataCache,
+    H5MetadataCache,
+    PerturbationAnnDataSource,
+    PerturbationConcatSource,
+    PerturbationSourceConfig,
+    RandomControlMapping,
+    load_experiment_config,
+)
+from diffbio.sources.singlecell_foundation import (
+    SingleCellEmbeddingSource,
+    SingleCellEmbeddingSourceConfig,
+    align_singlecell_embeddings,
+    load_singlecell_embedding_source,
+)
+from diffbio.sources.sequence_foundation import (
+    SequenceEmbeddingSource,
+    SequenceEmbeddingSourceConfig,
+    align_sequence_embeddings,
+    load_sequence_embedding_source,
+)
+
+__all__ = [
+    "AnnDataSource",
+    "AnnDataSourceConfig",
+    "BAMSource",
+    "BAMSourceConfig",
+    "BioSNAPDTISource",
+    "CONTEXTUAL_EPIGENOMICS_DATASET_CONTRACT_KEYS",
+    "CONTEXTUAL_TARGET_SEMANTICS",
+    "DTI_DATASET_CONTRACT_KEYS",
+    "DTISourceConfig",
+    "DavisDTISource",
+    "EmbeddingArtifactSource",
+    "EmbeddingArtifactSourceConfig",
+    "ENCODEPeakConfig",
+    "ENCODEPeakSource",
+    "FastaSource",
+    "FastaSourceConfig",
+    "IndexedViewSource",
+    "IndexedViewSourceConfig",
+    "MolNetSource",
+    "MolNetSourceConfig",
+    "MULTIOMICS_ARTIFACT_METADATA_KEYS",
+    "MULTIOMICS_DATASET_PROVENANCE_KEYS",
+    "MetabolomicsEmbeddingSource",
+    "MetabolomicsEmbeddingSourceConfig",
+    "MultiOmicsEmbeddingSource",
+    "MultiOmicsEmbeddingSourceConfig",
+    "from_anndata",
+    "to_anndata",
+    # Perturbation
+    "BatchControlMapping",
+    "ControlMappingConfig",
+    "ExperimentConfig",
+    "GlobalH5MetadataCache",
+    "H5MetadataCache",
+    "PerturbationAnnDataSource",
+    "PerturbationConcatSource",
+    "PerturbationSourceConfig",
+    "RandomControlMapping",
+    "build_synthetic_contextual_epigenomics_dataset",
+    "build_paired_dti_batch",
+    "deterministic_dti_split",
+    "validate_dti_dataset",
+    "validate_contextual_epigenomics_dataset",
+    "load_experiment_config",
+    "SequenceEmbeddingSource",
+    "SequenceEmbeddingSourceConfig",
+    "align_sequence_embeddings",
+    "align_metabolomics_embeddings",
+    "align_multiomics_embeddings",
+    "build_multiomics_artifact_metadata",
+    "build_multiomics_dataset_provenance",
+    "load_sequence_embedding_source",
+    "load_metabolomics_embedding_source",
+    "load_multiomics_embedding_source",
+    "SingleCellEmbeddingSource",
+    "SingleCellEmbeddingSourceConfig",
+    "align_singlecell_embeddings",
+    "load_singlecell_embedding_source",
+    "validate_multiomics_artifact_metadata",
+    "validate_multiomics_dataset_provenance",
+]

diffbio/sources/_anndata_shared.py

@@ -0,0 +1,89 @@
+"""Shared helpers for AnnData-backed DiffBio sources."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import jax.numpy as jnp
+import numpy as np
+from flax import nnx
+
+from diffbio.sources._utils import _require_anndata
+
+
+def to_dense_array(matrix: Any) -> np.ndarray:
+    """Convert a dense or sparse matrix into a float32 NumPy array."""
+    import scipy.sparse  # noqa: PLC0415
+
+    if scipy.sparse.issparse(matrix):
+        return np.asarray(matrix.toarray(), dtype=np.float32)
+    return np.asarray(matrix, dtype=np.float32)
+
+
+def read_h5ad(config: Any) -> Any:
+    """Load an AnnData object after validating the configured file path."""
+    anndata_mod = _require_anndata()
+    file_path = Path(str(config.file_path))
+    if not file_path.exists():
+        raise FileNotFoundError(f"AnnData file not found: {file_path}")
+    return anndata_mod.read_h5ad(file_path, backed="r" if config.backed else None)
+
+
+def load_obsm(adata: Any) -> dict[str, jnp.ndarray]:
+    """Load AnnData embedding matrices as float32 JAX arrays."""
+    if adata.obsm is None or len(adata.obsm) == 0:
+        return {}
+
+    return {
+        key: jnp.array(np.asarray(adata.obsm[key], dtype=np.float32)) for key in adata.obsm.keys()
+    }
+
+
+def extract_anndata_annotations(
+    adata: Any,
+) -> tuple[dict[str, Any], dict[str, Any], dict[str, jnp.ndarray]]:
+    """Extract obs, var, and obsm tables into the standard in-memory layout."""
+    obs = {col: np.asarray(adata.obs[col]) for col in adata.obs.columns}
+    var = {col: np.asarray(adata.var[col]) for col in adata.var.columns}
+    obsm = load_obsm(adata)
+    return obs, var, obsm
+
+
+def build_anndata_data(
+    *,
+    counts: jnp.ndarray,
+    obs: dict[str, Any],
+    var: dict[str, Any],
+    obsm: dict[str, jnp.ndarray],
+) -> dict[str, Any]:
+    """Assemble the canonical in-memory AnnData payload for DiffBio sources."""
+    return {
+        "counts": counts,
+        "obs": obs,
+        "var": var,
+        "obsm": obsm,
+    }
+
+
+def initialize_eager_source_state(
+    source: Any,
+    *,
+    data: dict[str, Any],
+    length: int,
+    seed: int,
+    shuffle: bool,
+    dataset_name: str | None,
+    split_name: str | None,
+    dataset_info: dict[str, int],
+) -> None:
+    """Populate the common eager-source bookkeeping fields."""
+    source.data = data
+    source.length = length
+    source.index = nnx.Variable(0)
+    source.epoch = nnx.Variable(0)
+    source._seed = seed
+    source.shuffle = shuffle
+    source.dataset_name = dataset_name
+    source.split_name = split_name
+    source._dataset_info = dataset_info

diffbio/sources/_batch_iteration.py

@@ -0,0 +1,37 @@
+"""Utilities for simple stateful batch iteration in data sources."""
+
+import logging
+from collections.abc import Callable
+
+import jax
+
+from datarax.typing import Element
+
+logger = logging.getLogger(__name__)
+
+
+def reset_iteration_state(source: object, seed: int | None = None) -> None:
+    """Reset a source object exposing `_current_idx` iteration state."""
+    del seed  # API compatibility with DataSourceModule reset signature
+    source._current_idx = 0
+
+
+def next_batch(
+    *,
+    batch_size: int,
+    key: jax.Array | None,
+    current_idx: int,
+    total_size: int,
+    get_element: Callable[[int], Element],
+) -> tuple[list[Element], int]:
+    """Collect up to `batch_size` elements starting from `current_idx`."""
+    del key  # API compatibility with DataSourceModule.get_batch signature
+
+    batch: list[Element] = []
+    idx = current_idx
+    for _ in range(batch_size):
+        if idx >= total_size:
+            break
+        batch.append(get_element(idx))
+        idx += 1
+    return batch, idx

diffbio/sources/_benchmark_source.py

@@ -0,0 +1,152 @@
+"""Shared substrate for benchmark-oriented eager data sources."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+import logging
+from pathlib import Path
+from typing import Any, Callable, Literal, overload
+
+import jax.numpy as jnp
+import numpy as np
+from flax import nnx
+
+from datarax.core.data_source import DataSourceModule
+
+from diffbio.sources._utils import _require_anndata
+
+
+def load_benchmark_adata(
+    *,
+    data_dir: str,
+    filename: str,
+    subsample: int | None = None,
+    seed: int = 42,
+) -> Any:
+    """Load a benchmark h5ad file with optional deterministic subsampling."""
+    anndata_mod = _require_anndata()
+    path = Path(data_dir) / filename
+    adata = anndata_mod.read_h5ad(path)
+
+    if subsample is not None and subsample < adata.n_obs:
+        rng = np.random.default_rng(seed)
+        indices = rng.choice(adata.n_obs, size=subsample, replace=False)
+        indices.sort()
+        adata = adata[indices].copy()
+
+    return adata
+
+
+def load_benchmark_counts(
+    *,
+    data_dir: str,
+    filename: str,
+    to_dense: Callable[[Any], np.ndarray],
+    subsample: int | None = None,
+    seed: int = 42,
+) -> tuple[Any, jnp.ndarray]:
+    """Load a benchmark h5ad file and convert its count matrix to a JAX array."""
+    adata = load_benchmark_adata(
+        data_dir=data_dir,
+        filename=filename,
+        subsample=subsample,
+        seed=seed,
+    )
+    counts = jnp.array(to_dense(adata.X))
+    return adata, counts
+
+
+@overload
+def encode_label_column(
+    column: Any,
+    *,
+    include_names: Literal[False] = False,
+) -> np.ndarray: ...
+
+
+@overload
+def encode_label_column(
+    column: Any,
+    *,
+    include_names: Literal[True],
+) -> tuple[np.ndarray, list[str]]: ...
+
+
+def encode_label_column(
+    column: Any,
+    *,
+    include_names: bool = False,
+) -> np.ndarray | tuple[np.ndarray, list[str]]:
+    """Encode an observation column into stable int32 label codes."""
+    if hasattr(column, "cat"):
+        codes = np.asarray(column.cat.codes, dtype=np.int32)
+        names = [str(value) for value in column.cat.categories]
+    else:
+        unique_labels, codes = np.unique(np.asarray(column), return_inverse=True)
+        codes = codes.astype(np.int32)
+        names = [str(value) for value in unique_labels]
+
+    if include_names:
+        return codes, names
+    return codes
+
+
+def iter_loaded_rows(
+    data: dict[str, Any],
+    *,
+    static_keys: tuple[str, ...] = ("gene_names",),
+) -> Iterator[dict[str, Any]]:
+    """Iterate row-wise over an eager benchmark payload."""
+    static_key_set = set(static_keys)
+    for i in range(int(data["n_cells"])):
+        yield {
+            key: value[i] if hasattr(value, "__getitem__") and key not in static_key_set else value
+            for key, value in data.items()
+        }
+
+
+class BenchmarkDataSource(DataSourceModule):
+    """Shared base class for eager benchmark data sources backed by a data dict."""
+
+    data: dict[str, Any] = nnx.data()
+    iter_static_keys: tuple[str, ...] = ("gene_names",)
+
+    def load(self) -> dict[str, Any]:
+        """Return the eagerly loaded dataset payload."""
+        return self.data
+
+    def __len__(self) -> int:
+        """Return the number of rows in the eagerly loaded dataset."""
+        return int(self.data["n_cells"])
+
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        """Iterate row-wise through the loaded dataset payload."""
+        return iter_loaded_rows(self.data, static_keys=self.iter_static_keys)
+
+    def _load_benchmark_counts(
+        self,
+        config: Any,
+        filename: str,
+        to_dense: Callable[[Any], np.ndarray],
+    ) -> tuple[Any, jnp.ndarray]:
+        """Load benchmark AnnData and convert its count matrix to a JAX array."""
+        return load_benchmark_counts(
+            data_dir=config.data_dir,
+            filename=filename,
+            to_dense=to_dense,
+            subsample=getattr(config, "subsample", None),
+        )
+
+    def _log_loaded_summary(
+        self,
+        logger: logging.Logger,
+        dataset_name: str,
+        metric_keys: tuple[str, ...],
+    ) -> None:
+        """Log a standard loaded-dataset summary using the requested metric keys."""
+        labels = [key[2:] if key.startswith("n_") else key for key in metric_keys]
+        metric_summary = ", ".join(f"%d {label}" for label in labels)
+        logger.info(
+            f"Loaded {dataset_name}: {metric_summary}",
+            *(int(self.data[key]) for key in metric_keys),
+        )

diffbio/sources/_indexed_batch_source.py

@@ -0,0 +1,38 @@
+"""Mixin for stateful index-based batching in data sources."""
+
+import logging
+
+import jax
+
+from datarax.typing import Element
+
+from diffbio.sources._batch_iteration import next_batch, reset_iteration_state
+
+logger = logging.getLogger(__name__)
+
+
+class IndexedBatchSourceMixin:
+    """Reusable `reset` and `get_batch` logic for indexable data sources."""
+
+    def _batch_total_size(self) -> int:
+        """Return the total number of elements available for batching."""
+        raise NotImplementedError
+
+    def _batch_element(self, idx: int) -> Element:
+        """Return the element at the given index."""
+        raise NotImplementedError
+
+    def reset(self, seed: int | None = None) -> None:
+        """Reset iteration state, optionally with a new seed."""
+        reset_iteration_state(self, seed)
+
+    def get_batch(self, batch_size: int, key: jax.Array | None = None) -> list[Element]:
+        """Return the next batch of elements, advancing the internal index."""
+        batch, self._current_idx = next_batch(
+            batch_size=batch_size,
+            key=key,
+            current_idx=self._current_idx,
+            total_size=self._batch_total_size(),
+            get_element=self._batch_element,
+        )
+        return batch