diffbio 0.1.0__py3-none-any.whl

diffbio/sources/balifam.py

@@ -0,0 +1,207 @@
+"""BAliBASE reference alignment DataSource (balifam).
+
+Loads protein families from the balifam repository, which provides
+curated subsets of BAliBASE reference alignments at three tiers
+(100, 1000, 10000 sequences per family).
+
+Each family contains unaligned input sequences and a reference
+alignment with mixed-case convention:
+- Uppercase residues: core (scored) positions
+- Lowercase residues: insert (unscored) positions
+- Dots (.): gap characters
+
+The dataset must be available locally. See:
+https://github.com/steineggerlab/balifam
+
+Data directory layout::
+
+    balifam100/
+        in/       # Unaligned FASTA files (one per family)
+        ref/      # Reference alignments (FASTA with gaps)
+        info/     # Family ID lists
+"""
+
+from __future__ import annotations
+
+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/balifam"
+
+
+def _parse_fasta(path: Path) -> list[tuple[str, str]]:
+    """Parse a FASTA file into (name, sequence) tuples.
+
+    Handles multi-line sequences by concatenating continuation lines.
+
+    Args:
+        path: Path to FASTA file.
+
+    Returns:
+        List of (sequence_name, sequence_string) tuples.
+    """
+    entries: list[tuple[str, str]] = []
+    current_name: str | None = None
+    current_seq_parts: list[str] = []
+
+    with path.open() as fh:
+        for line in fh:
+            line = line.rstrip("\n")
+            if line.startswith(">"):
+                if current_name is not None:
+                    entries.append((current_name, "".join(current_seq_parts)))
+                current_name = line[1:].strip()
+                current_seq_parts = []
+            elif current_name is not None:
+                current_seq_parts.append(line)
+
+    if current_name is not None:
+        entries.append((current_name, "".join(current_seq_parts)))
+
+    return entries
+
+
+@dataclass(frozen=True, kw_only=True)
+class BalifamConfig(StructuralConfig):
+    """Configuration for BalifamSource.
+
+    Attributes:
+        data_dir: Root directory of the balifam repository.
+        tier: Family size tier (100, 1000, or 10000 sequences).
+        max_families: Maximum number of families to load.
+            None loads all available families.
+    """
+
+    data_dir: str = _DEFAULT_DATA_DIR
+    tier: int = 100
+    max_families: int | None = None
+
+    def __post_init__(self) -> None:
+        """Validate configuration."""
+        super().__post_init__()
+        valid_tiers = {100, 1000, 10000}
+        if self.tier not in valid_tiers:
+            raise ValueError(f"tier must be one of {valid_tiers}, got {self.tier}")
+        tier_dir = Path(self.data_dir) / f"balifam{self.tier}"
+        if not tier_dir.exists():
+            raise FileNotFoundError(
+                f"Balifam tier directory not found: {tier_dir}. "
+                f"Clone from: https://github.com/steineggerlab/balifam"
+            )
+
+
+class BalifamSource(DataSourceModule):
+    """DataSource for BAliBASE reference alignments (balifam).
+
+    Loads protein family alignments from balifam for evaluating
+    multiple sequence alignment methods. Each family provides
+    unaligned input sequences and a curated reference alignment.
+
+    The reference alignment contains a subset of the input sequences
+    with known structural alignment, using mixed-case annotation
+    (uppercase = core/scored, lowercase = insert/unscored).
+
+    Example:
+        ```python
+        config = BalifamConfig(tier=100, max_families=5)
+        source = BalifamSource(config)
+        families = source.load()
+        print(families[0]["family_id"])
+        ```
+    """
+
+    families: list[dict[str, Any]] = nnx.data()
+
+    def __init__(
+        self,
+        config: BalifamConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ) -> None:
+        """Load balifam families.
+
+        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 "BalifamSource")
+        self.families = self._load_families(config)
+        logger.info(
+            "Loaded %d balifam%d families",
+            len(self.families),
+            config.tier,
+        )
+
+    def _load_families(self, config: BalifamConfig) -> list[dict[str, Any]]:
+        """Load families from disk.
+
+        Args:
+            config: Source configuration.
+
+        Returns:
+            List of family dicts with keys: family_id, sequences,
+            reference, n_sequences, n_reference.
+        """
+        tier_dir = Path(config.data_dir) / f"balifam{config.tier}"
+        in_dir = tier_dir / "in"
+        ref_dir = tier_dir / "ref"
+
+        # Discover families from reference directory (authoritative)
+        family_files = sorted(ref_dir.iterdir())
+        if config.max_families is not None:
+            family_files = family_files[: config.max_families]
+
+        families: list[dict[str, Any]] = []
+        for ref_path in family_files:
+            family_id = ref_path.name
+            in_path = in_dir / family_id
+
+            if not in_path.exists():
+                logger.warning(
+                    "Input file missing for family %s, skipping",
+                    family_id,
+                )
+                continue
+
+            sequences = _parse_fasta(in_path)
+            reference = _parse_fasta(ref_path)
+
+            families.append(
+                {
+                    "family_id": family_id,
+                    "sequences": sequences,
+                    "reference": reference,
+                    "n_sequences": len(sequences),
+                    "n_reference": len(reference),
+                }
+            )
+
+        return families
+
+    def load(self) -> list[dict[str, Any]]:
+        """Return all loaded families.
+
+        Returns:
+            List of family dicts, each with keys: family_id,
+            sequences, reference, n_sequences, n_reference.
+        """
+        return self.families
+
+    def __len__(self) -> int:
+        """Return the number of loaded families."""
+        return len(self.families)
+
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        """Iterate over families."""
+        return iter(self.families)

diffbio/sources/bam.py

@@ -0,0 +1,265 @@
+"""BAM/CRAM file data source for genomics workflows.
+
+This module provides BAMSource for reading aligned sequencing reads
+from BAM/CRAM files with lazy loading and efficient indexed access.
+
+Based on best practices from:
+- pysam (HTSlib Python wrapper)
+- Google Nucleus genomics library
+- DeepVariant BAM handling patterns
+
+References:
+    - https://pysam.readthedocs.io/
+    - https://github.com/pysam-developers/pysam
+"""
+
+import logging
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+import jax.numpy as jnp
+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 BAMSourceConfig(StructuralConfig):
+    """Configuration for BAM/CRAM data source.
+
+    Attributes:
+        file_path: Path to BAM/CRAM file
+        reference_path: Optional path to reference FASTA (required for CRAM)
+        include_unmapped: Whether to include unmapped reads (default: False)
+        min_mapping_quality: Minimum mapping quality to include (default: None)
+        region: Optional genomic region to query (e.g., "chr1:1000-2000")
+        handle_n: How to handle N nucleotides in sequences
+    """
+
+    file_path: Path = None  # type: ignore[assignment]  # Required, validated in post_init
+    reference_path: Path | None = None
+    include_unmapped: bool = False
+    min_mapping_quality: int | None = None
+    region: str | None = None
+    handle_n: Literal["uniform", "zero"] = "uniform"
+
+    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 BAMSource(IndexedBatchSourceMixin, DataSourceModule):
+    """BAM/CRAM file data source extending Datarax DataSourceModule.
+
+    Provides efficient access to aligned sequencing reads with:
+
+    - Lazy loading using pysam iterators
+    - Indexed random access via BAI/CRAI files
+    - Quality filtering at load time
+    - One-hot encoded sequence output
+
+    Inherits from DataSourceModule (StructuralModule) because:
+
+    - Non-parametric: BAM reading is deterministic
+    - Frozen config: file parameters don't change
+    - Domain-specific: requires genomics-specific handling
+
+    Example:
+        ```python
+        config = BAMSourceConfig(file_path=Path("sample.bam"))
+        source = BAMSource(config)
+        for element in source:
+            print(element.data["read_name"], element.data["sequence"].shape)
+        ```
+
+    Performance Tips (from pysam best practices):
+
+    - Use indexed BAM files for random access
+    - Filter by region to reduce data loading
+    - Set min_mapping_quality to filter at read time
+    """
+
+    # Annotate data storage for Flax NNX
+    _reads: list = nnx.data()
+
+    def __init__(
+        self,
+        config: BAMSourceConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize BAMSource.
+
+        Args:
+            config: BAM source configuration
+            rngs: Random number generators (unused for data loading)
+            name: Optional module name
+
+        Raises:
+            FileNotFoundError: If BAM file not found
+            ImportError: If pysam is not installed
+        """
+        super().__init__(config, rngs=rngs, name=name)
+
+        # Import pysam lazily to allow installation without it
+        try:
+            import pysam
+
+            self._pysam = pysam
+        except ImportError as err:
+            raise ImportError(
+                "pysam is required for BAMSource. Install with: pip install pysam"
+            ) from err
+
+        # Validate file exists
+        if not config.file_path.exists():
+            raise FileNotFoundError(f"BAM file not found: {config.file_path}")
+
+        # Load read index (metadata only, not full sequences)
+        self._reads = self._index_reads()
+        self._current_idx = 0
+
+    def _index_reads(self) -> list[dict]:
+        """Build an index of reads for random access.
+
+        This loads read metadata without fully parsing sequences,
+        enabling lazy loading on access.
+
+        Returns:
+            List of read metadata dictionaries
+        """
+        config = self.config
+        reads: list[dict] = []
+
+        # Open BAM file
+        mode = "rb" if str(config.file_path).endswith(".bam") else "rc"
+        reference = str(config.reference_path) if config.reference_path else None
+
+        with self._pysam.AlignmentFile(
+            str(config.file_path), mode, reference_filename=reference
+        ) as bam:
+            for read in self._iter_reads(bam):
+                if self._should_skip_read(read):
+                    continue
+                reads.append(self._read_info(read))
+
+        return reads
+
+    def _iter_reads(self, bam) -> object:
+        """Return the configured BAM iterator."""
+        if self.config.region:
+            return bam.fetch(region=self.config.region)
+        return bam.fetch(until_eof=True)
+
+    def _should_skip_read(self, read) -> bool:
+        """Check whether a read should be excluded by source filters."""
+        if read.is_unmapped and not self.config.include_unmapped:
+            return True
+        min_mapq = self.config.min_mapping_quality
+        return min_mapq is not None and read.mapping_quality < min_mapq
+
+    @staticmethod
+    def _read_info(read) -> dict[str, object]:
+        """Extract index metadata from a pysam read."""
+        return {
+            "query_name": read.query_name,
+            "reference_id": read.reference_id,
+            "reference_start": read.reference_start,
+            "is_unmapped": read.is_unmapped,
+            "query_sequence": read.query_sequence,
+            "query_qualities": list(read.query_qualities)
+            if read.query_qualities is not None
+            else None,
+            "mapping_quality": read.mapping_quality,
+            "reference_name": read.reference_name if not read.is_unmapped else None,
+        }
+
+    def _read_to_element(self, idx: int, read_info: dict) -> Element:
+        """Convert read info to Element with one-hot encoded sequence.
+
+        Args:
+            idx: Index of the read
+            read_info: Read metadata dictionary
+
+        Returns:
+            Element with sequence, quality scores, and metadata
+        """
+        # Encode sequence as one-hot
+        sequence_str = read_info["query_sequence"] or ""
+        if sequence_str:
+            sequence = encode_dna_string(sequence_str, handle_n=self.config.handle_n)
+        else:
+            sequence = jnp.zeros((0, 4), dtype=jnp.float32)
+
+        # Get quality scores
+        if read_info["query_qualities"] is not None:
+            quality_scores = jnp.array(read_info["query_qualities"], dtype=jnp.float32)
+        else:
+            # Default quality if not available
+            quality_scores = jnp.ones(len(sequence_str), dtype=jnp.float32) * 30.0
+
+        data = {
+            "sequence": sequence,
+            "quality_scores": quality_scores,
+            "read_name": read_info["query_name"],
+        }
+
+        metadata = {
+            "idx": idx,
+            "reference_name": read_info["reference_name"],
+            "reference_start": read_info["reference_start"],
+            "mapping_quality": read_info["mapping_quality"],
+            "unmapped": read_info["is_unmapped"],
+        }
+
+        return Element(data=data, state={}, metadata=metadata)  # pyright: ignore[reportArgumentType]
+
+    def __len__(self) -> int:
+        """Return the number of reads in the source."""
+        return len(self._reads)
+
+    def __getitem__(self, idx: int) -> Element | None:
+        """Get read by index.
+
+        Args:
+            idx: Index of the read
+
+        Returns:
+            Element at the given index, or None if out of bounds
+        """
+        if idx < 0 or idx >= len(self._reads):
+            return None
+        return self._read_to_element(idx, self._reads[idx])
+
+    def __iter__(self) -> Iterator[Element]:  # type: ignore[override]
+        """Return iterator over reads."""
+        self._current_idx = 0
+        return self
+
+    def __next__(self) -> Element:
+        """Get next read in iteration."""
+        if self._current_idx >= len(self._reads):
+            raise StopIteration
+        elem = self._read_to_element(self._current_idx, self._reads[self._current_idx])
+        self._current_idx += 1
+        return elem
+
+    def _batch_total_size(self) -> int:
+        """Return number of indexed reads for mixin batch iteration."""
+        return len(self._reads)
+
+    def _batch_element(self, idx: int) -> Element:
+        """Build the indexed read element for mixin batch iteration."""
+        return self._read_to_element(idx, self._reads[idx])