diffbio 0.1.0__py3-none-any.whl

diffbio/sources/encode_peaks.py

@@ -0,0 +1,223 @@
+"""ENCODE narrowPeak BED data source.
+
+Loads peak calls from ENCODE narrowPeak BED files (gzipped or plain),
+such as those produced by the ENCODE ChIP-seq pipeline. Each row in
+the BED file encodes one called peak with:
+
+    chr  start  end  name  score  strand  signalValue  pValue  qValue  peak
+
+Only the first ten columns are required. The ``peak`` column (column 10)
+gives the offset within the peak region to the summit position.
+
+Peaks can be filtered to a single chromosome for speed and optionally
+capped at a maximum count.
+
+Reference:
+    ENCODE Project Consortium. "An integrated encyclopedia of DNA
+    elements in the human genome." Nature 489, 57-74 (2012).
+"""
+
+from __future__ import annotations
+
+import gzip
+import logging
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+from datarax.core.config import StructuralConfig
+from datarax.core.data_source import DataSourceModule
+from flax import nnx
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_DATA_PATH = "/media/mahdi/ssd23/Data/encode/CTCF_K562_narrowPeak.bed.gz"
+
+
+@dataclass(frozen=True, kw_only=True)
+class ENCODEPeakConfig(StructuralConfig):
+    """Configuration for ENCODEPeakSource.
+
+    Attributes:
+        data_path: Path to the narrowPeak BED file (gzipped or plain).
+        chromosome: Chromosome to filter peaks to. None loads all.
+        max_peaks: Maximum number of peaks to load. None loads all.
+    """
+
+    data_path: str = _DEFAULT_DATA_PATH
+    chromosome: str | None = "chr22"
+    max_peaks: int | None = None
+
+    def __post_init__(self) -> None:
+        """Validate that the data file exists."""
+        super().__post_init__()
+        path = Path(self.data_path)
+        if not path.exists():
+            raise FileNotFoundError(
+                f"ENCODE narrowPeak file not found: {path}. "
+                f"Download from https://www.encodeproject.org/"
+            )
+
+
+@dataclass(frozen=True, kw_only=True)
+class ENCODEPeak:
+    """A single ENCODE narrowPeak record.
+
+    Attributes:
+        chromosome: Chromosome name (e.g. ``"chr22"``).
+        start: 0-based start coordinate.
+        end: 0-based end coordinate (exclusive).
+        signal_value: Fold-enrichment signal value.
+        p_value: -log10 p-value (or -1 if unavailable).
+        q_value: -log10 q-value (or -1 if unavailable).
+        summit_offset: Offset from ``start`` to the summit position.
+    """
+
+    chromosome: str
+    start: int
+    end: int
+    signal_value: float
+    p_value: float
+    q_value: float
+    summit_offset: int
+
+
+def _parse_narrowpeak(
+    file_path: Path,
+    chromosome: str | None,
+    max_peaks: int | None,
+) -> list[ENCODEPeak]:
+    """Parse a narrowPeak BED file into ENCODEPeak records.
+
+    Args:
+        file_path: Path to gzipped or plain BED file.
+        chromosome: Filter to this chromosome. None keeps all.
+        max_peaks: Maximum peaks to return. None returns all.
+
+    Returns:
+        Sorted list of ENCODEPeak records (by start position).
+    """
+    peaks: list[ENCODEPeak] = []
+    opener = gzip.open if file_path.suffix == ".gz" else open
+
+    with opener(file_path, "rt", encoding="utf-8") as fh:
+        for line in fh:
+            line = line.strip()
+            if not line or line.startswith("#"):
+                continue
+            fields = line.split("\t")
+            if len(fields) < 10:
+                continue
+
+            chrom = fields[0]
+            if chromosome is not None and chrom != chromosome:
+                continue
+
+            peak = ENCODEPeak(
+                chromosome=chrom,
+                start=int(fields[1]),
+                end=int(fields[2]),
+                signal_value=float(fields[6]),
+                p_value=float(fields[7]),
+                q_value=float(fields[8]),
+                summit_offset=int(fields[9]),
+            )
+            peaks.append(peak)
+
+            if max_peaks is not None and len(peaks) >= max_peaks:
+                break
+
+    # Sort by genomic position for deterministic ordering
+    peaks.sort(key=lambda p: (p.chromosome, p.start))
+    return peaks
+
+
+class ENCODEPeakSource(DataSourceModule):
+    """DataSource for ENCODE narrowPeak ChIP-seq peak calls.
+
+    Loads peak positions and signal values from an ENCODE narrowPeak
+    BED file. Peaks are optionally filtered to a single chromosome
+    and/or capped at a maximum count.
+
+    Each loaded peak provides genomic coordinates, signal enrichment,
+    statistical significance, and summit position.
+
+    Example:
+        ```python
+        config = ENCODEPeakConfig(chromosome="chr22", max_peaks=500)
+        source = ENCODEPeakSource(config)
+        data = source.load()
+        print(data["n_peaks"])       # Number of peaks loaded
+        print(data["starts"][:5])    # First 5 start positions
+        ```
+    """
+
+    data: dict[str, Any] = nnx.data()
+
+    def __init__(
+        self,
+        config: ENCODEPeakConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ) -> None:
+        """Load ENCODE narrowPeak data from BED file.
+
+        Args:
+            config: Configuration with file path and filters.
+            rngs: Optional RNG state (unused, for interface compat).
+            name: Optional module name.
+        """
+        super().__init__(config, rngs=rngs, name=name or "ENCODEPeakSource")
+        file_path = Path(config.data_path)
+        peaks = _parse_narrowpeak(file_path, config.chromosome, config.max_peaks)
+
+        if not peaks:
+            chrom_msg = f" on {config.chromosome}" if config.chromosome else ""
+            raise ValueError(
+                f"No peaks found in {file_path}{chrom_msg}. "
+                f"Check the file format and chromosome filter."
+            )
+
+        starts = np.array([p.start for p in peaks], dtype=np.int64)
+        ends = np.array([p.end for p in peaks], dtype=np.int64)
+        signals = np.array([p.signal_value for p in peaks], dtype=np.float64)
+        summits = np.array(
+            [p.start + p.summit_offset for p in peaks],
+            dtype=np.int64,
+        )
+
+        self.data = {
+            "peaks": peaks,
+            "starts": starts,
+            "ends": ends,
+            "signal_values": signals,
+            "summit_positions": summits,
+            "n_peaks": len(peaks),
+            "chromosome": config.chromosome,
+        }
+        logger.info(
+            "Loaded ENCODE peaks: %d peaks from %s%s",
+            len(peaks),
+            file_path.name,
+            f" ({config.chromosome})" if config.chromosome else "",
+        )
+
+    def load(self) -> dict[str, Any]:
+        """Return the full dataset as a dictionary.
+
+        Returns:
+            Dict with keys: peaks, starts, ends, signal_values,
+            summit_positions, n_peaks, chromosome.
+        """
+        return self.data
+
+    def __len__(self) -> int:
+        """Return the number of loaded peaks."""
+        return self.data["n_peaks"]
+
+    def __iter__(self) -> Iterator[ENCODEPeak]:
+        """Iterate over individual peak records."""
+        yield from self.data["peaks"]

diffbio/sources/fasta.py

@@ -0,0 +1,226 @@
+"""FASTA file data source for genomics workflows.
+
+This module provides FastaSource for reading DNA/RNA sequences
+from FASTA files with lazy loading and efficient indexed access.
+
+Based on best practices from:
+- pyfaidx (samtools-compatible FASTA indexing)
+- BioPython SeqIO.index patterns
+- Google Nucleus FASTA handling
+
+References:
+    - https://github.com/mdshw5/pyfaidx
+    - https://pythonhosted.org/pyfaidx/
+"""
+
+import logging
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+from flax import nnx
+
+from datarax.core.config import StructuralConfig
+from datarax.core.data_source import DataSourceModule
+from datarax.typing import Element
+
+from diffbio.sequences.dna import encode_dna_string
+from diffbio.sources._indexed_batch_source import IndexedBatchSourceMixin
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class FastaSourceConfig(StructuralConfig):
+    """Configuration for FASTA data source.
+
+    Attributes:
+        file_path: Path to FASTA file
+        handle_n: How to handle N nucleotides ("uniform" or "zero")
+        create_index: Whether to create .fai index if not exists (default: True)
+    """
+
+    file_path: Path = None  # type: ignore[assignment]  # Required, validated in post_init
+    handle_n: Literal["uniform", "zero"] = "uniform"
+    create_index: bool = True
+
+    def __post_init__(self) -> None:
+        """Validate configuration after initialization."""
+        super().__post_init__()
+        if self.file_path is None:
+            raise ValueError("file_path is required")
+
+
+class FastaSource(IndexedBatchSourceMixin, DataSourceModule):
+    """FASTA file data source extending Datarax DataSourceModule.
+
+    Provides efficient access to DNA/RNA sequences with:
+
+    - Lazy loading using samtools-compatible .fai index
+    - Dictionary-like access by sequence name
+    - One-hot encoded sequence output
+    - Support for compressed BGZF files
+
+    Inherits from DataSourceModule (StructuralModule) because:
+
+    - Non-parametric: FASTA reading is deterministic
+    - Frozen config: file parameters don't change
+    - Domain-specific: requires genomics-specific handling
+
+    Example:
+        ```python
+        config = FastaSourceConfig(file_path=Path("genome.fasta"))
+        source = FastaSource(config)
+        elem = source.get_by_name("chr1")
+        print(elem.data["sequence"].shape)
+        ```
+
+    Performance Tips (from pyfaidx best practices):
+
+    - Use indexed FASTA files (.fai) for random access
+    - Access regions with slicing for large chromosomes
+    - BGZF compression reduces disk space while maintaining random access
+    """
+
+    # Annotate data storage for Flax NNX
+    _sequence_names: list = nnx.data()
+    _name_to_idx: dict = nnx.data()
+
+    def __init__(
+        self,
+        config: FastaSourceConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize FastaSource.
+
+        Args:
+            config: FASTA source configuration
+            rngs: Random number generators (unused for data loading)
+            name: Optional module name
+
+        Raises:
+            FileNotFoundError: If FASTA file not found
+            ImportError: If pyfaidx is not installed
+        """
+        super().__init__(config, rngs=rngs, name=name)
+
+        # Import pyfaidx lazily to allow installation without it
+        try:
+            import pyfaidx
+
+            self._pyfaidx = pyfaidx
+        except ImportError as err:
+            raise ImportError(
+                "pyfaidx is required for FastaSource. Install with: pip install pyfaidx"
+            ) from err
+
+        # Validate file exists
+        if not config.file_path.exists():
+            raise FileNotFoundError(f"FASTA file not found: {config.file_path}")
+
+        # Open FASTA file with pyfaidx (creates index if needed)
+        self._fasta = self._pyfaidx.Fasta(
+            str(config.file_path),
+            build_index=config.create_index,
+        )
+
+        # Build sequence name index
+        self._sequence_names = list(self._fasta.keys())
+        self._name_to_idx = {name: idx for idx, name in enumerate(self._sequence_names)}
+        self._current_idx = 0
+
+    def _sequence_to_element(self, idx: int, seq_name: str) -> Element:
+        """Convert FASTA sequence to Element with one-hot encoding.
+
+        Args:
+            idx: Index of the sequence
+            seq_name: Name/ID of the sequence
+
+        Returns:
+            Element with one-hot encoded sequence and metadata
+        """
+        # Get sequence from pyfaidx (lazy loaded)
+        fasta_seq = self._fasta[seq_name]
+        sequence_str = str(fasta_seq).upper()
+
+        # Encode sequence as one-hot
+        sequence = encode_dna_string(sequence_str, handle_n=self.config.handle_n)
+
+        # Get description if available
+        description = getattr(fasta_seq, "long_name", seq_name)
+
+        data = {
+            "sequence": sequence,
+            "sequence_id": seq_name,
+            "description": description,
+        }
+
+        metadata = {
+            "idx": idx,
+            "length": len(sequence_str),
+            "file_path": str(self.config.file_path),
+        }
+
+        return Element(data=data, state={}, metadata=metadata)  # pyright: ignore[reportArgumentType]
+
+    def __len__(self) -> int:
+        """Return the number of sequences in the source."""
+        return len(self._sequence_names)
+
+    def __getitem__(self, idx: int) -> Element | None:
+        """Get sequence by index.
+
+        Args:
+            idx: Index of the sequence
+
+        Returns:
+            Element at the given index, or None if out of bounds
+        """
+        if idx < 0 or idx >= len(self._sequence_names):
+            return None
+        seq_name = self._sequence_names[idx]
+        return self._sequence_to_element(idx, seq_name)
+
+    def __iter__(self) -> Iterator[Element]:  # type: ignore[override]
+        """Return iterator over sequences."""
+        self._current_idx = 0
+        return self
+
+    def __next__(self) -> Element:
+        """Get next sequence in iteration."""
+        if self._current_idx >= len(self._sequence_names):
+            raise StopIteration
+        seq_name = self._sequence_names[self._current_idx]
+        elem = self._sequence_to_element(self._current_idx, seq_name)
+        self._current_idx += 1
+        return elem
+
+    def _batch_total_size(self) -> int:
+        """Return number of indexed sequences for mixin batch iteration."""
+        return len(self._sequence_names)
+
+    def _batch_element(self, idx: int) -> Element:
+        """Build the indexed sequence element for mixin batch iteration."""
+        return self._sequence_to_element(idx, self._sequence_names[idx])
+
+    def get_by_name(self, name: str) -> Element | None:
+        """Get sequence by name/ID.
+
+        Args:
+            name: Sequence identifier (e.g., "chr1", "seq1")
+
+        Returns:
+            Element for the sequence, or None if not found
+        """
+        if name not in self._name_to_idx:
+            return None
+        idx = self._name_to_idx[name]
+        return self._sequence_to_element(idx, name)
+
+    @property
+    def sequence_names(self) -> list[str]:
+        """Get list of all sequence names in the FASTA file."""
+        return list(self._sequence_names)

diffbio/sources/immune_human.py

@@ -0,0 +1,172 @@
+"""Immune human integration benchmark DataSource.
+
+Loads the human immune cell atlas dataset from the scib benchmark
+(Luecken et al., Nature Methods 2022). This is the standard dataset
+for evaluating single-cell batch integration methods.
+
+Dataset: 33,506 cells, 12,303 genes, 10 batches, 16 cell types.
+Source: https://figshare.com/articles/dataset/12420968
+
+The dataset must be pre-downloaded to a local directory. See
+``benchmarks/README.md`` for download instructions.
+"""
+
+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 = "Immune_ALL_human.h5ad"
+
+
+@dataclass(frozen=True, kw_only=True)
+class ImmuneHumanConfig(StructuralConfig):
+    """Configuration for ImmuneHumanSource.
+
+    Attributes:
+        data_dir: Directory containing the downloaded h5ad file.
+        subsample: If set, randomly subsample this many cells.
+            Use for quick/CI benchmark runs.
+        batch_key: Column name in obs for batch labels.
+        label_key: Column name in obs for cell type labels.
+        embedding_key: Key in obsm for precomputed embeddings.
+    """
+
+    data_dir: str = "/media/mahdi/ssd23/Data/scib"
+    subsample: int | None = None
+    batch_key: str = "batch"
+    label_key: str = "final_annotation"
+    embedding_key: str = "X_pca"
+
+    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}. "
+                f"Download from: https://ndownloader.figshare.com/"
+                f"files/25717328"
+            )
+
+
+class ImmuneHumanSource(BenchmarkDataSource):
+    """DataSource for the scib immune human integration benchmark.
+
+    Loads the human immune cell atlas (33,506 cells, 12,303 genes,
+    10 batches, 16 cell types) from a pre-downloaded h5ad file.
+
+    Follows the datarax DataSourceModule pattern: eager loading at
+    init, dict-based access via ``load()``, length via ``__len__``.
+
+    Example:
+        ```python
+        config = ImmuneHumanConfig(data_dir="/path/to/data")
+        source = ImmuneHumanSource(config)
+        data = source.load()
+        print(data["counts"].shape)  # (33506, 12303)
+        ```
+    """
+
+    def __init__(
+        self,
+        config: ImmuneHumanConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ) -> None:
+        """Load the immune human dataset.
+
+        Args:
+            config: Configuration with data directory and options.
+            rngs: Optional RNG state (unused, for interface compat).
+            name: Optional module name.
+        """
+        super().__init__(config, rngs=rngs, name=name or "ImmuneHumanSource")
+        self.data = self._load(config)
+        self._log_loaded_summary(
+            logger,
+            "immune_human",
+            ("n_cells", "n_genes", "n_batches", "n_types"),
+        )
+
+    def _load(self, config: ImmuneHumanConfig) -> dict[str, Any]:
+        """Load and preprocess the h5ad file."""
+        adata, counts = self._load_benchmark_counts(config, _FILENAME, _to_dense)
+
+        # Encode categorical labels as integer codes
+        batch_labels = encode_label_column(adata.obs[config.batch_key])
+        cell_type_labels = encode_label_column(adata.obs[config.label_key])
+
+        # Get embeddings (PCA from obsm, or compute on the fly)
+        if config.embedding_key in adata.obsm:
+            embeddings = jnp.array(np.asarray(adata.obsm[config.embedding_key], dtype=np.float32))
+        else:
+            logger.info(
+                "Embedding key '%s' not in obsm (%s). Computing PCA (50 components).",
+                config.embedding_key,
+                list(adata.obsm.keys()),
+            )
+            embeddings = self._compute_pca(counts, n_components=50)
+
+        gene_names = list(adata.var_names)
+        cell_ids = [str(cell_id) for cell_id in adata.obs_names]
+
+        return {
+            "counts": counts,
+            "batch_labels": batch_labels,
+            "cell_type_labels": cell_type_labels,
+            "cell_ids": cell_ids,
+            "embeddings": embeddings,
+            "gene_names": gene_names,
+            "n_cells": adata.n_obs,
+            "n_genes": adata.n_vars,
+            "n_batches": int(len(np.unique(batch_labels))),
+            "n_types": int(len(np.unique(cell_type_labels))),
+        }
+
+    @staticmethod
+    def _compute_pca(counts: jnp.ndarray, n_components: int = 50) -> jnp.ndarray:
+        """Compute PCA embeddings from count matrix.
+
+        Log-normalizes, then computes truncated SVD for PCA.
+
+        Args:
+            counts: Dense count matrix (n_cells, n_genes).
+            n_components: Number of principal components.
+
+        Returns:
+            PCA embeddings of shape (n_cells, n_components).
+        """
+        # Log-normalize: log1p(counts / total * 10000)
+        totals = jnp.sum(counts, axis=1, keepdims=True)
+        totals = jnp.maximum(totals, 1.0)
+        normalized = jnp.log1p(counts / totals * 10000.0)
+
+        # Center
+        mean = jnp.mean(normalized, axis=0)
+        centered = normalized - mean
+
+        # Truncated SVD via numpy (JAX full SVD on 33K x 12K is expensive)
+        centered_np = np.asarray(centered)
+        from sklearn.decomposition import TruncatedSVD  # noqa: PLC0415
+
+        svd = TruncatedSVD(n_components=n_components, random_state=42)
+        embeddings = svd.fit_transform(centered_np)
+        return jnp.array(embeddings.astype(np.float32))