diffbio 0.1.0__py3-none-any.whl

diffbio/sources/anndata_source.py

@@ -0,0 +1,361 @@
+"""AnnData (.h5ad) data source for single-cell genomics.
+
+This module provides AnnDataSource for loading single-cell RNA-seq data from
+.h5ad files (AnnData format) and converting them to JAX-compatible data dicts
+suitable for DiffBio operators.
+
+Follows the datarax eager-loading pattern (same as HFEagerSource): all data is
+loaded to JAX arrays at init, then iteration/batching uses pure JAX operations
+with O(1) memory shuffling via Grain's index_shuffle.
+
+Handles both dense and sparse count matrices, cell/gene metadata, and
+optional embeddings (PCA, UMAP, etc.).
+
+References:
+    - https://anndata.readthedocs.io/
+    - Wolf et al. "SCANPY: large-scale single-cell gene expression data analysis"
+      Genome Biology, 2018.
+"""
+
+# Ownership note: DiffBio owns the AnnData biological schema adapter; batching
+# and iteration stay on Datarax eager-source primitives.
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass
+from typing import Any
+
+import jax
+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.sources._eager_source_ops import eager_get_batch, eager_iter, eager_reset
+
+from diffbio.sources._anndata_shared import (
+    build_anndata_data,
+    extract_anndata_annotations,
+    initialize_eager_source_state,
+    read_h5ad,
+    to_dense_array,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class AnnDataSourceConfig(StructuralConfig):
+    """Configuration for AnnDataSource.
+
+    Attributes:
+        file_path: Path to the .h5ad file (string or Path object).
+        backed: Whether to open in backed mode (memory-mapped).
+        shuffle: Whether to shuffle during iteration.
+        seed: Integer seed for Grain's index_shuffle.
+        split: Optional split name for pipeline integration.
+    """
+
+    file_path: str | None = None
+    backed: bool = False
+    shuffle: bool = False
+    seed: int = 42
+    split: str | None = None
+
+    def __post_init__(self) -> None:
+        """Validate configuration after initialization."""
+        if self.shuffle:
+            object.__setattr__(self, "stochastic", True)
+            if self.stream_name is None:
+                object.__setattr__(self, "stream_name", "shuffle")
+        else:
+            object.__setattr__(self, "stochastic", False)
+
+        super().__post_init__()
+
+        if self.file_path is None:
+            raise ValueError("file_path is required for AnnDataSourceConfig")
+
+
+class AnnDataSource(DataSourceModule):
+    """Eager-loading AnnData source for single-cell RNA-seq data.
+
+    Loads all data from .h5ad files to JAX arrays at initialization, then
+    provides pure JAX iteration, batching, and indexed access. Follows the
+    same eager-loading pattern as datarax's HFEagerSource.
+
+    Provides:
+        - Full dataset loading via ``load()``
+        - Per-cell indexed access via ``__getitem__``
+        - Iteration via ``__iter__`` with optional O(1) memory shuffling
+        - Batch retrieval via ``get_batch(batch_size)``
+        - Automatic sparse-to-dense conversion
+        - JAX array output for count matrices and embeddings
+
+    Output dictionary keys:
+        - ``counts``: Dense JAX array of shape (n_cells, n_genes) from ``.X``
+        - ``obs``: Dict of cell metadata columns from ``.obs``
+        - ``var``: Dict of gene metadata columns from ``.var``
+        - ``obsm``: Dict of embedding JAX arrays from ``.obsm`` (empty if absent)
+
+    Example:
+        ```python
+        config = AnnDataSourceConfig(file_path="pbmc3k.h5ad")
+        source = AnnDataSource(config)
+        print(len(source))                # 2700
+        print(source.load()["counts"].shape)  # (2700, 32738)
+
+        for cell in source:
+            print(cell["counts"].shape)   # (32738,)
+            break
+
+        batch = source.get_batch(32)
+        print(batch["counts"].shape)      # (32, 32738)
+        ```
+    """
+
+    # Annotate data storage for Flax NNX (prevents parameter tracking)
+    data: dict[str, Any] = nnx.data()
+
+    def __init__(
+        self,
+        config: AnnDataSourceConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ) -> None:
+        """Initialize AnnDataSource from a .h5ad file.
+
+        Loads all data to JAX arrays at construction time.
+
+        Args:
+            config: AnnDataSourceConfig with file path and options.
+            rngs: Optional RNG state for shuffling.
+            name: Optional module name.
+
+        Raises:
+            FileNotFoundError: If the file does not exist.
+            ImportError: If anndata is not installed.
+        """
+        if name is None:
+            name = f"AnnDataSource({config.file_path})"
+        super().__init__(config, rngs=rngs, name=name)
+
+        adata = read_h5ad(config)
+
+        # Convert count matrix to JAX array
+        counts = jnp.array(to_dense_array(adata.X))
+        obs, var, obsm = extract_anndata_annotations(adata)
+
+        self._initialize_loaded_source(
+            config=config,
+            adata=adata,
+            data=build_anndata_data(counts=counts, obs=obs, var=var, obsm=obsm),
+            length=adata.n_obs,
+        )
+
+    # =================================================================
+    # Public API: load / info
+    # =================================================================
+
+    def load(self) -> dict[str, Any]:
+        """Return the full dataset as a dictionary of JAX arrays and metadata.
+
+        Returns:
+            Dictionary with keys ``counts``, ``obs``, ``var``, ``obsm``.
+        """
+        return dict(self.data)
+
+    def get_dataset_info(self) -> dict[str, int]:
+        """Return cached dataset metadata.
+
+        Returns:
+            Dict with ``n_genes`` and ``n_cells``.
+        """
+        return self._dataset_info
+
+    def _initialize_loaded_source(
+        self,
+        *,
+        config: AnnDataSourceConfig,
+        adata: Any,
+        data: dict[str, Any],
+        length: int,
+    ) -> None:
+        """Initialize common eager-source state after AnnData loading."""
+        initialize_eager_source_state(
+            self,
+            data=data,
+            length=length,
+            seed=config.seed,
+            shuffle=config.shuffle,
+            dataset_name=str(config.file_path),
+            split_name=config.split,
+            dataset_info={
+                "n_genes": adata.n_vars,
+                "n_cells": adata.n_obs,
+            },
+        )
+
+    def _iter_with_builder(
+        self,
+        build_element: Callable[[dict[str, Any], int], dict[str, Any]],
+    ) -> Iterator[dict[str, Any]]:
+        """Iterate using the standard eager-source bookkeeping state."""
+        return eager_iter(
+            self.data,
+            self.length,
+            self.index,
+            self.epoch,
+            self.shuffle,
+            self._seed,
+            build_element,
+        )
+
+    def _get_batch_with_gather(
+        self,
+        batch_size: int,
+        key: jax.Array | None,
+        gather_fn: Callable[[dict[str, Any], Any], dict[str, Any]],
+    ) -> dict[str, Any]:
+        """Collect a batch using the standard eager-source bookkeeping state."""
+        return eager_get_batch(
+            self.data,
+            self.length,
+            self.index,
+            self.epoch,
+            self.shuffle,
+            self._seed,
+            batch_size,
+            key,
+            gather_fn,
+        )
+
+    # =================================================================
+    # DataSourceModule protocol: __len__, __iter__, __next__, __getitem__
+    # =================================================================
+
+    def __len__(self) -> int:
+        """Return the number of cells in the dataset."""
+        return self.length
+
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        """Iterate over cells with optional O(1) memory shuffling.
+
+        Yields:
+            Per-cell dictionaries with ``counts``, ``obs``, ``obsm`` keys.
+        """
+        return self._iter_with_builder(_build_cell_element)
+
+    def __next__(self) -> dict[str, Any]:
+        """Get the next cell element (required by DataSourceModule).
+
+        Raises:
+            StopIteration: When iteration is exhausted.
+        """
+        raise StopIteration
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        """Get data for a single cell by index.
+
+        Supports negative indexing.
+
+        Args:
+            idx: Cell index (supports negative indexing).
+
+        Returns:
+            Dictionary with ``counts``, ``obs``, ``obsm`` keys.
+
+        Raises:
+            IndexError: If idx is out of bounds.
+        """
+        if idx < 0:
+            idx = self.length + idx
+        if idx < 0 or idx >= self.length:
+            raise IndexError(f"Cell index {idx} out of range for dataset with {self.length} cells")
+        return _build_cell_element(self.data, idx)
+
+    # =================================================================
+    # Batch retrieval
+    # =================================================================
+
+    def get_batch(self, batch_size: int, key: jax.Array | None = None) -> dict[str, Any]:
+        """Get a batch of cells.
+
+        Stateful (advances internal index) when called without ``key``.
+        Stateless (random sampling) when called with ``key``.
+
+        Args:
+            batch_size: Number of cells per batch.
+            key: Optional RNG key for stateless random sampling.
+
+        Returns:
+            Dictionary with batched arrays.
+        """
+
+        def _gather(data: dict[str, Any], indices: jax.Array) -> dict[str, Any]:
+            counts = data["counts"][indices]
+            obs = {col: np.asarray(arr)[np.array(indices)] for col, arr in data["obs"].items()}
+            obsm: dict[str, jnp.ndarray] = {}
+            for emb_name, emb_arr in data["obsm"].items():
+                obsm[emb_name] = emb_arr[indices]
+            return {"counts": counts, "obs": obs, "obsm": obsm}
+
+        return self._get_batch_with_gather(batch_size, key, _gather)
+
+    # =================================================================
+    # State management
+    # =================================================================
+
+    def reset(self, seed: int | None = None) -> None:
+        """Reset source to the beginning.
+
+        Args:
+            seed: Unused (uses config seed).
+        """
+        del seed
+        eager_reset(self.index, self.epoch, self._cache)
+
+    def set_shuffle(self, shuffle: bool) -> None:
+        """Enable or disable shuffling.
+
+        Args:
+            shuffle: Whether to shuffle data during iteration.
+        """
+        self.shuffle = shuffle
+
+    def __repr__(self) -> str:
+        """Return string representation."""
+        return (
+            f"AnnDataSource("
+            f"dataset={self.dataset_name}, "
+            f"length={self.length}, "
+            f"shuffle={self.shuffle}, "
+            f"epoch={self.epoch.get_value()})"
+        )
+
+
+# =====================================================================
+# Module-level helpers (kept outside the class)
+# =====================================================================
+
+
+def _build_cell_element(data: dict[str, Any], idx: int) -> dict[str, Any]:
+    """Build a per-cell dictionary from the full dataset at a given index.
+
+    Args:
+        data: The full data dict with ``counts``, ``obs``, ``obsm`` keys.
+        idx: Cell index.
+
+    Returns:
+        Per-cell dictionary with scalar obs values and 1D arrays.
+    """
+    cell_counts = data["counts"][idx]
+    cell_obs = {col: arr[idx] for col, arr in data["obs"].items()}
+    cell_obsm: dict[str, jnp.ndarray] = {}
+    for emb_name, emb_arr in data["obsm"].items():
+        cell_obsm[emb_name] = emb_arr[idx]
+    return {"counts": cell_counts, "obs": cell_obs, "obsm": cell_obsm}

diffbio/sources/archive_ii.py

@@ -0,0 +1,174 @@
+"""ArchiveII RNA secondary structure DataSource.
+
+Loads RNA sequences with known secondary structures from the ArchiveII
+benchmark dataset (Sloma & Mathews, 2016). Structures are provided in
+dot-bracket notation (DBN) and represent experimentally validated RNA
+secondary structures from crystal structures and NMR.
+
+The dataset is read from CSV files produced by RNAFoldAssess, with
+columns: name, sequence, ground_truth_type, ground_truth_data.
+
+Only rows with ``ground_truth_type == "DBN"`` are loaded; rows without
+structure annotations are silently skipped.
+
+References:
+    Sloma, M. F. & Mathews, D. H. (2016). Exact calculation of loop
+    formation probability identifies folding motifs in RNA secondary
+    structures. RNA 22, 1808-1818.
+"""
+
+from __future__ import annotations
+
+import csv
+import logging
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from datarax.core.config import StructuralConfig
+from datarax.core.data_source import DataSourceModule
+from flax import nnx
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_DATA_DIR = "/media/mahdi/ssd23/Works/RNAFoldAssess/tutorial/processed_data"
+_STRUCTURE_FILENAME = "example_data_structure.csv"
+
+
+@dataclass(frozen=True, kw_only=True)
+class ArchiveIIConfig(StructuralConfig):
+    """Configuration for ArchiveIISource.
+
+    Attributes:
+        data_dir: Directory containing the ArchiveII CSV file.
+        filename: Name of the CSV file with structure annotations.
+        max_sequences: Maximum number of sequences to load.
+            None means load all available sequences.
+    """
+
+    data_dir: str = _DEFAULT_DATA_DIR
+    filename: str = _STRUCTURE_FILENAME
+    max_sequences: int | None = None
+
+    def __post_init__(self) -> None:
+        """Validate that the data file exists."""
+        super().__post_init__()
+        path = Path(self.data_dir) / self.filename
+        if not path.exists():
+            raise FileNotFoundError(
+                f"ArchiveII data not found: {path}. "
+                f"Expected a CSV with columns: name, sequence, "
+                f"ground_truth_type, ground_truth_data."
+            )
+
+
+def _parse_csv(
+    csv_path: Path,
+    max_sequences: int | None,
+) -> list[dict[str, str]]:
+    """Parse ArchiveII CSV and return DBN-annotated entries.
+
+    Args:
+        csv_path: Path to the CSV file.
+        max_sequences: Maximum entries to return. None for all.
+
+    Returns:
+        List of dicts with keys: name, sequence, structure.
+    """
+    entries: list[dict[str, str]] = []
+    with csv_path.open(encoding="utf-8") as fh:
+        reader = csv.DictReader(fh)
+        for row in reader:
+            gt_type = row.get("ground_truth_type", "").strip()
+            gt_data = row.get("ground_truth_data", "").strip()
+            if gt_type != "DBN" or not gt_data:
+                continue
+            sequence = row["sequence"].strip().upper()
+            name = row["name"].strip()
+            entries.append(
+                {
+                    "name": name,
+                    "sequence": sequence,
+                    "structure": gt_data,
+                }
+            )
+            if max_sequences is not None and len(entries) >= max_sequences:
+                break
+    return entries
+
+
+class ArchiveIISource(DataSourceModule):
+    """DataSource for ArchiveII RNA secondary structures.
+
+    Loads RNA sequences and their known dot-bracket notation (DBN)
+    structures from a CSV file. Only entries with ground_truth_type
+    ``"DBN"`` are included.
+
+    Each entry is a dict with keys:
+        - ``name``: Sequence identifier (e.g. ``"1KXK_chain_0"``)
+        - ``sequence``: RNA sequence string (e.g. ``"GUCUACC..."``)
+        - ``structure``: DBN string (e.g. ``"....(((...)))..."``)
+
+    Example:
+        ```python
+        config = ArchiveIIConfig(max_sequences=10)
+        source = ArchiveIISource(config)
+        data = source.load()
+        print(data["n_sequences"])
+        print(data["entries"][0]["name"])
+        ```
+    """
+
+    data: dict[str, Any] = nnx.data()
+
+    def __init__(
+        self,
+        config: ArchiveIIConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ) -> None:
+        """Load ArchiveII RNA structures from CSV.
+
+        Args:
+            config: Configuration with data directory and limits.
+            rngs: Optional RNG state (unused, for interface compat).
+            name: Optional module name.
+        """
+        super().__init__(config, rngs=rngs, name=name or "ArchiveIISource")
+        csv_path = Path(config.data_dir) / config.filename
+        entries = _parse_csv(csv_path, config.max_sequences)
+
+        if not entries:
+            raise ValueError(
+                f"No DBN-annotated sequences found in {csv_path}. "
+                f"Ensure the CSV has rows with "
+                f"ground_truth_type='DBN'."
+            )
+
+        self.data = {
+            "entries": entries,
+            "n_sequences": len(entries),
+        }
+        logger.info(
+            "Loaded ArchiveII: %d sequences from %s",
+            len(entries),
+            csv_path,
+        )
+
+    def load(self) -> dict[str, Any]:
+        """Return the full dataset as a dictionary.
+
+        Returns:
+            Dict with keys: entries (list of dicts), n_sequences.
+        """
+        return self.data
+
+    def __len__(self) -> int:
+        """Return the number of loaded sequences."""
+        return self.data["n_sequences"]
+
+    def __iter__(self) -> Iterator[dict[str, str]]:
+        """Iterate over individual RNA entries."""
+        yield from self.data["entries"]