diffbio 0.1.0__py3-none-any.whl

diffbio/sources/multiomics.py

@@ -0,0 +1,279 @@
+"""Shared multi-omics provenance and embedding-artifact sources."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Mapping
+
+import jax.numpy as jnp
+from flax import nnx
+
+from diffbio.sources.indexed_embeddings import IndexedEmbeddingSource, IndexedEmbeddingSourceConfig
+
+MULTIOMICS_DATASET_PROVENANCE_KEYS = (
+    "dataset_name",
+    "source_type",
+    "modalities",
+    "curation_status",
+    "biological_validation",
+    "promotion_eligible",
+    "source_path",
+)
+MULTIOMICS_ARTIFACT_METADATA_KEYS = (
+    "artifact_id",
+    "artifact_type",
+    "modalities",
+    "embedding_source",
+    "foundation_source_name",
+    "promotion_eligible",
+)
+_SUPPORTED_MULTIOMICS_MODALITIES = frozenset(
+    {
+        "rna",
+        "atac",
+        "protein",
+        "spatial",
+        "metabolomics",
+        "mass_spectrometry",
+    }
+)
+
+
+@dataclass(frozen=True)
+class MultiOmicsEmbeddingSourceConfig(IndexedEmbeddingSourceConfig):
+    """Configuration for sample-indexed multi-omics embedding artifacts."""
+
+    row_id_key: str = "sample_ids"
+
+
+class MultiOmicsEmbeddingSource(IndexedEmbeddingSource):
+    """Indexed embedding source specialized for multi-omics sample artifacts."""
+
+    config: MultiOmicsEmbeddingSourceConfig  # pyright: ignore[reportIncompatibleVariableOverride]
+
+    @classmethod
+    def from_path(
+        cls,
+        path: Path | str,
+        *,
+        rngs: nnx.Rngs | None = None,
+    ) -> MultiOmicsEmbeddingSource:
+        """Build the canonical sample-indexed source for a multi-omics artifact."""
+        return cls(MultiOmicsEmbeddingSourceConfig(file_path=str(path)), rngs=rngs)
+
+    @property
+    def sample_ids(self) -> tuple[str, ...] | None:
+        """Tuple of persisted sample identifiers, if present."""
+        return self.row_ids
+
+    def align_to_reference_sample_ids(
+        self,
+        *,
+        reference_sample_ids: Sequence[str],
+        require_sample_ids: bool = True,
+    ) -> jnp.ndarray:
+        """Align imported embeddings to a reference multi-omics sample order."""
+        return self.align_to_reference_ids(
+            reference_ids=tuple(reference_sample_ids),
+            require_row_ids=require_sample_ids,
+            artifact_label="Multi-omics",
+            id_display_name="Sample ID",
+        )
+
+
+@dataclass(frozen=True)
+class MetabolomicsEmbeddingSourceConfig(IndexedEmbeddingSourceConfig):
+    """Configuration for spectrum-indexed metabolomics embedding artifacts."""
+
+    row_id_key: str = "spectrum_ids"
+
+
+class MetabolomicsEmbeddingSource(IndexedEmbeddingSource):
+    """Indexed embedding source specialized for metabolomics spectrum artifacts."""
+
+    config: MetabolomicsEmbeddingSourceConfig  # pyright: ignore[reportIncompatibleVariableOverride]
+
+    @classmethod
+    def from_path(
+        cls,
+        path: Path | str,
+        *,
+        rngs: nnx.Rngs | None = None,
+    ) -> MetabolomicsEmbeddingSource:
+        """Build the canonical spectrum-indexed source for a metabolomics artifact."""
+        return cls(MetabolomicsEmbeddingSourceConfig(file_path=str(path)), rngs=rngs)
+
+    @property
+    def spectrum_ids(self) -> tuple[str, ...] | None:
+        """Tuple of persisted spectrum identifiers, if present."""
+        return self.row_ids
+
+    def align_to_reference_spectrum_ids(
+        self,
+        *,
+        reference_spectrum_ids: Sequence[str],
+        require_spectrum_ids: bool = True,
+    ) -> jnp.ndarray:
+        """Align imported embeddings to a reference spectrum order."""
+        return self.align_to_reference_ids(
+            reference_ids=tuple(reference_spectrum_ids),
+            require_row_ids=require_spectrum_ids,
+            artifact_label="Metabolomics",
+            id_display_name="Spectrum ID",
+        )
+
+
+def build_multiomics_dataset_provenance(
+    *,
+    dataset_name: str,
+    source_type: str,
+    modalities: Sequence[str],
+    curation_status: str,
+    biological_validation: str,
+    promotion_eligible: bool,
+    source_path: str | None = None,
+) -> dict[str, Any]:
+    """Build canonical provenance for benchmarked multi-omics datasets."""
+    provenance: dict[str, Any] = {
+        "dataset_name": dataset_name,
+        "source_type": source_type,
+        "modalities": list(modalities),
+        "curation_status": curation_status,
+        "biological_validation": biological_validation,
+        "promotion_eligible": promotion_eligible,
+        "source_path": source_path,
+    }
+    return validate_multiomics_dataset_provenance(provenance)
+
+
+def validate_multiomics_dataset_provenance(
+    provenance: Mapping[str, Any],
+) -> dict[str, Any]:
+    """Validate and normalize a multi-omics dataset provenance payload."""
+    missing = [key for key in MULTIOMICS_DATASET_PROVENANCE_KEYS if key not in provenance]
+    if missing:
+        raise ValueError(f"multi-omics dataset_provenance is missing required keys: {missing}")
+
+    normalized = {key: provenance[key] for key in MULTIOMICS_DATASET_PROVENANCE_KEYS}
+    for key in (
+        "dataset_name",
+        "source_type",
+        "curation_status",
+        "biological_validation",
+    ):
+        _require_non_empty_string(normalized[key], field_name=f"dataset_provenance.{key}")
+
+    normalized["modalities"] = _normalize_modalities(normalized["modalities"])
+    if not isinstance(normalized["promotion_eligible"], bool):
+        raise TypeError("multi-omics dataset_provenance.promotion_eligible must be a bool.")
+    if str(normalized["source_type"]).startswith("synthetic") and normalized["promotion_eligible"]:
+        raise ValueError("Synthetic multi-omics provenance cannot be promotion_eligible.")
+    if normalized["source_path"] is not None:
+        _require_non_empty_string(
+            normalized["source_path"],
+            field_name="dataset_provenance.source_path",
+        )
+    return normalized
+
+
+def build_multiomics_artifact_metadata(
+    *,
+    artifact_id: str,
+    artifact_type: str,
+    modalities: Sequence[str],
+    embedding_source: str,
+    foundation_source_name: str,
+    promotion_eligible: bool,
+) -> dict[str, Any]:
+    """Build canonical metadata for imported or benchmark-produced omics artifacts."""
+    metadata: dict[str, Any] = {
+        "artifact_id": artifact_id,
+        "artifact_type": artifact_type,
+        "modalities": list(modalities),
+        "embedding_source": embedding_source,
+        "foundation_source_name": foundation_source_name,
+        "promotion_eligible": promotion_eligible,
+    }
+    return validate_multiomics_artifact_metadata(metadata)
+
+
+def validate_multiomics_artifact_metadata(
+    metadata: Mapping[str, Any],
+) -> dict[str, Any]:
+    """Validate and normalize multi-omics artifact metadata."""
+    missing = [key for key in MULTIOMICS_ARTIFACT_METADATA_KEYS if key not in metadata]
+    if missing:
+        raise ValueError(f"multi-omics artifact metadata is missing required keys: {missing}")
+
+    normalized = {key: metadata[key] for key in MULTIOMICS_ARTIFACT_METADATA_KEYS}
+    for key in ("artifact_id", "artifact_type", "embedding_source", "foundation_source_name"):
+        _require_non_empty_string(normalized[key], field_name=f"artifact_metadata.{key}")
+    normalized["modalities"] = _normalize_modalities(normalized["modalities"])
+    if not isinstance(normalized["promotion_eligible"], bool):
+        raise TypeError("multi-omics artifact metadata promotion_eligible must be a bool.")
+    return normalized
+
+
+def load_multiomics_embedding_source(
+    path: Path | str,
+    *,
+    rngs: nnx.Rngs | None = None,
+) -> MultiOmicsEmbeddingSource:
+    """Build the canonical sample-indexed multi-omics embedding source."""
+    return MultiOmicsEmbeddingSource.from_path(path, rngs=rngs)
+
+
+def align_multiomics_embeddings(
+    *,
+    reference_sample_ids: Sequence[str],
+    artifact_path: Path | str,
+    require_sample_ids: bool = True,
+) -> jnp.ndarray:
+    """Align imported multi-omics embeddings to a reference sample order."""
+    return load_multiomics_embedding_source(artifact_path).align_to_reference_sample_ids(
+        reference_sample_ids=reference_sample_ids,
+        require_sample_ids=require_sample_ids,
+    )
+
+
+def load_metabolomics_embedding_source(
+    path: Path | str,
+    *,
+    rngs: nnx.Rngs | None = None,
+) -> MetabolomicsEmbeddingSource:
+    """Build the canonical spectrum-indexed metabolomics embedding source."""
+    return MetabolomicsEmbeddingSource.from_path(path, rngs=rngs)
+
+
+def align_metabolomics_embeddings(
+    *,
+    reference_spectrum_ids: Sequence[str],
+    artifact_path: Path | str,
+    require_spectrum_ids: bool = True,
+) -> jnp.ndarray:
+    """Align imported metabolomics embeddings to a reference spectrum order."""
+    return load_metabolomics_embedding_source(artifact_path).align_to_reference_spectrum_ids(
+        reference_spectrum_ids=reference_spectrum_ids,
+        require_spectrum_ids=require_spectrum_ids,
+    )
+
+
+def _normalize_modalities(raw_modalities: Any) -> list[str]:
+    """Normalize and validate modality identifiers."""
+    if not isinstance(raw_modalities, Sequence) or isinstance(raw_modalities, str):
+        raise TypeError("modalities must be a non-empty sequence of strings.")
+    modalities = [str(modality) for modality in raw_modalities]
+    if not modalities:
+        raise ValueError("modalities must contain at least one modality.")
+    invalid = sorted(set(modalities) - _SUPPORTED_MULTIOMICS_MODALITIES)
+    if invalid:
+        raise ValueError(f"Unsupported multi-omics modalities: {invalid}")
+    return modalities
+
+
+def _require_non_empty_string(value: Any, *, field_name: str) -> None:
+    """Require one non-empty string field."""
+    if not isinstance(value, str) or not value:
+        raise TypeError(f"{field_name} must be a non-empty string.")

diffbio/sources/pancreas.py

@@ -0,0 +1,108 @@
+"""Pancreas endocrinogenesis DataSource for trajectory benchmarks.
+
+Loads the scVelo pancreas dataset (Bastidas-Ponce et al. 2019,
+Bergen et al. 2020) with spliced/unspliced layers for RNA velocity
+and precomputed PCA embeddings for pseudotime inference.
+
+Dataset: 3,696 cells, 27,998 genes, 8 cell types, 5 coarse types.
+Source: scVelo tutorial dataset.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import jax.numpy as jnp
+import numpy as np
+from flax import nnx
+
+from datarax.core.config import StructuralConfig
+
+from diffbio.sources._benchmark_source import (
+    BenchmarkDataSource,
+    encode_label_column,
+)
+from diffbio.sources._utils import to_dense_float32 as _to_dense
+
+logger = logging.getLogger(__name__)
+
+_FILENAME = "endocrinogenesis_day15.h5ad"
+
+
+@dataclass(frozen=True, kw_only=True)
+class PancreasConfig(StructuralConfig):
+    """Configuration for PancreasSource.
+
+    Attributes:
+        data_dir: Directory containing the downloaded h5ad file.
+        subsample: If set, randomly subsample this many cells.
+        cluster_key: Column in obs for cell type labels.
+    """
+
+    data_dir: str = "/media/mahdi/ssd23/Data/scvelo"
+    subsample: int | None = None
+    cluster_key: str = "clusters"
+
+    def __post_init__(self) -> None:
+        """Validate configuration."""
+        super().__post_init__()
+        path = Path(self.data_dir) / _FILENAME
+        if not path.exists():
+            raise FileNotFoundError(
+                f"Dataset not found: {path}. Download the scVelo "
+                f"pancreas dataset to {self.data_dir}/"
+            )
+
+
+class PancreasSource(BenchmarkDataSource):
+    """DataSource for the scVelo pancreas endocrinogenesis dataset.
+
+    Provides counts, spliced/unspliced layers, PCA embeddings,
+    and cell type labels for trajectory and velocity benchmarks.
+    """
+
+    def __init__(
+        self,
+        config: PancreasConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ) -> None:
+        """Load the pancreas dataset."""
+        super().__init__(config, rngs=rngs, name=name or "PancreasSource")
+        self.data = self._load(config)
+        self._log_loaded_summary(logger, "pancreas", ("n_cells", "n_genes", "n_types"))
+
+    def _load(self, config: PancreasConfig) -> dict[str, Any]:
+        """Load and preprocess the h5ad file."""
+        adata, counts = self._load_benchmark_counts(config, _FILENAME, _to_dense)
+
+        # Spliced/unspliced for velocity
+        spliced = jnp.array(_to_dense(adata.layers["spliced"]))
+        unspliced = jnp.array(_to_dense(adata.layers["unspliced"]))
+
+        # Cell type labels
+        labels = encode_label_column(adata.obs[config.cluster_key])
+
+        # Embeddings
+        embeddings = jnp.array(
+            np.asarray(
+                adata.obsm.get("X_pca", np.zeros((adata.n_obs, 50))),
+                dtype=np.float32,
+            )
+        )
+
+        return {
+            "counts": counts,
+            "spliced": spliced,
+            "unspliced": unspliced,
+            "cell_type_labels": labels,
+            "embeddings": embeddings,
+            "gene_names": list(adata.var_names),
+            "n_cells": adata.n_obs,
+            "n_genes": adata.n_vars,
+            "n_types": int(len(np.unique(labels))),
+        }

diffbio/sources/perturbation/__init__.py

@@ -0,0 +1,69 @@
+"""Perturbation-aware data loading for single-cell experiments.
+
+This sub-package provides data sources, configuration, and utilities for
+loading single-cell perturbation experiment data (e.g., CRISPR screens),
+porting and adapting features from the cell-load library to the
+JAX/datarax ecosystem.
+
+Sources:
+    PerturbationAnnDataSource: Single-file perturbation-aware AnnData source
+    PerturbationConcatSource: Multi-dataset concatenation source
+
+Configuration:
+    ExperimentConfig: TOML-based experiment configuration
+    PerturbationSourceConfig: Source configuration
+
+Control Mapping:
+    RandomControlMapping: Random control cell mapping within cell type
+    BatchControlMapping: Batch-aware control cell mapping
+
+Utilities:
+    H5MetadataCache: Singleton cache for H5 metadata
+    GlobalH5MetadataCache: Process-global cache manager
+"""
+
+from diffbio.sources.perturbation._types import MappingStrategy, OutputSpaceMode
+from diffbio.sources.perturbation.concat_source import PerturbationConcatSource
+from diffbio.sources.perturbation.control_mapping import (
+    BatchControlMapping,
+    ControlMappingConfig,
+    RandomControlMapping,
+)
+from diffbio.sources.perturbation.experiment_config import (
+    DatasetEntry,
+    ExperimentConfig,
+    FewshotEntry,
+    ZeroshotEntry,
+    load_experiment_config,
+)
+from diffbio.sources.perturbation.h5_metadata_cache import (
+    GlobalH5MetadataCache,
+    H5MetadataCache,
+)
+from diffbio.sources.perturbation.perturbation_source import (
+    PerturbationAnnDataSource,
+    PerturbationSourceConfig,
+)
+
+__all__ = [
+    # Types
+    "MappingStrategy",
+    "OutputSpaceMode",
+    # Sources
+    "PerturbationAnnDataSource",
+    "PerturbationSourceConfig",
+    "PerturbationConcatSource",
+    # Configuration
+    "ExperimentConfig",
+    "DatasetEntry",
+    "ZeroshotEntry",
+    "FewshotEntry",
+    "load_experiment_config",
+    # Control Mapping
+    "ControlMappingConfig",
+    "RandomControlMapping",
+    "BatchControlMapping",
+    # Cache
+    "H5MetadataCache",
+    "GlobalH5MetadataCache",
+]

diffbio/sources/perturbation/_types.py

@@ -0,0 +1,51 @@
+"""Shared type aliases and enums for perturbation data loading.
+
+Provides canonical string enums and type aliases used across the perturbation
+sub-package: output space modes, mapping strategies, and cell/perturbation
+label types.
+"""
+
+from __future__ import annotations
+
+from enum import StrEnum
+from typing import TypeAlias
+
+# ---------------------------------------------------------------------------
+# Type aliases
+# ---------------------------------------------------------------------------
+
+CellIndex: TypeAlias = int
+PerturbationLabel: TypeAlias = str
+CellTypeLabel: TypeAlias = str
+BatchLabel: TypeAlias = str
+
+
+# ---------------------------------------------------------------------------
+# Enums
+# ---------------------------------------------------------------------------
+
+
+class OutputSpaceMode(StrEnum):
+    """Output representation mode for perturbation data.
+
+    Attributes:
+        GENE: Highly variable gene (HVG) subset of the count matrix.
+        ALL: Full gene expression matrix.
+        EMBEDDING: Pre-computed embedding only (no raw counts).
+    """
+
+    GENE = "gene"
+    ALL = "all"
+    EMBEDDING = "embedding"
+
+
+class MappingStrategy(StrEnum):
+    """Strategy for mapping perturbed cells to control cells.
+
+    Attributes:
+        BATCH: Map within same batch and cell type.
+        RANDOM: Map to random control of same cell type.
+    """
+
+    BATCH = "batch"
+    RANDOM = "random"

diffbio/sources/perturbation/_utils.py

@@ -0,0 +1,125 @@
+"""Shared utility functions for the perturbation sub-package.
+
+Ports and adapts utility functions from cell-load's data_utils module to work
+with JAX arrays and numpy instead of PyTorch tensors.
+
+References:
+    - cell-load/src/cell_load/utils/data_utils.py
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
+import jax.numpy as jnp
+import numpy as np
+
+
+def safe_decode_array(arr: Any) -> np.ndarray:
+    """Decode byte-string arrays to UTF-8 and cast all entries to Python str.
+
+    Args:
+        arr: Array-like of bytes or other objects.
+
+    Returns:
+        Numpy string array with all elements decoded to str.
+    """
+    decoded: list[str] = []
+    for x in arr:
+        if isinstance(x, (bytes, bytearray)):
+            decoded.append(x.decode("utf-8", errors="ignore"))
+        else:
+            decoded.append(str(x))
+    return np.array(decoded, dtype=str)
+
+
+def generate_onehot_map(keys: Iterable[str]) -> dict[str, jnp.ndarray]:
+    """Build a map from each unique key to a one-hot JAX array.
+
+    Keys are sorted to ensure deterministic ordering across runs.
+
+    Args:
+        keys: Iterable of hashable string items.
+
+    Returns:
+        Dict mapping each unique key to a float32 one-hot JAX array
+        of length equal to the number of unique keys.
+    """
+    unique_keys = sorted(set(keys))
+    n = len(unique_keys)
+    identity = jnp.eye(n, dtype=jnp.float32)
+    return {k: identity[i] for i, k in enumerate(unique_keys)}
+
+
+def is_discrete_counts(x: jnp.ndarray, n_cells: int = 100) -> bool:
+    """Detect if data appears to be raw integer counts.
+
+    Checks whether the row sums of the first ``n_cells`` rows are integers
+    (fractional part approximately zero).
+
+    Args:
+        x: Array of shape ``(n_cells, n_genes)``.
+        n_cells: Number of cells to sample for detection.
+
+    Returns:
+        True if data appears to be discrete/raw counts.
+    """
+    top_n = min(x.shape[0], n_cells)
+    row_sums = x[:top_n].sum(axis=1)
+    frac_part = row_sums - jnp.floor(row_sums)
+    return bool(jnp.all(jnp.abs(frac_part) < 1e-7))
+
+
+def is_log_transformed(x: jnp.ndarray) -> bool:
+    """Detect if data is log-transformed by checking the global maximum.
+
+    Log1p-transformed data typically has a maximum below 15.
+
+    Args:
+        x: Array of expression values.
+
+    Returns:
+        True if data appears to be log-transformed.
+    """
+    return bool(x.max() < 15.0)
+
+
+def split_perturbations_by_cell_fraction(
+    pert_groups: dict[str, np.ndarray],
+    val_fraction: float,
+    rng: np.random.Generator,
+) -> tuple[list[str], list[str]]:
+    """Partition perturbations so the val subset approximates a target cell fraction.
+
+    Uses a greedy algorithm: shuffles perturbations, then greedily assigns each
+    to the val subset if doing so brings the val cell count closer to the target.
+
+    Args:
+        pert_groups: Dict mapping perturbation names to arrays of cell indices.
+        val_fraction: Target fraction of total cells to assign to validation.
+        rng: Numpy random generator for shuffling.
+
+    Returns:
+        Tuple of (train_perturbation_names, val_perturbation_names).
+    """
+    total_cells = sum(len(indices) for indices in pert_groups.values())
+    target_val_cells = val_fraction * total_cells
+
+    pert_size_list = [(p, len(pert_groups[p])) for p in pert_groups]
+    rng.shuffle(pert_size_list)
+
+    val_perts: list[str] = []
+    current_val_cells = 0
+
+    for pert, size in pert_size_list:
+        new_val_cells = current_val_cells + size
+        diff_if_add = abs(new_val_cells - target_val_cells)
+        diff_if_skip = abs(current_val_cells - target_val_cells)
+
+        if diff_if_add < diff_if_skip:
+            val_perts.append(pert)
+            current_val_cells = new_val_cells
+
+    train_perts = [p for p, _ in pert_size_list if p not in set(val_perts)]
+    return train_perts, val_perts