diffbio 0.1.0__py3-none-any.whl

diffbio/sources/bengrn_ground_truth.py

@@ -0,0 +1,306 @@
+"""BenGRN ground truth DataSource for GRN inference benchmarking.
+
+Loads mESC expression data and ChIP+Perturb ground truth edges from
+the benGRN repository (Stone & Sroy gold standards).
+
+Data source: /media/mahdi/ssd23/Works/benGRN/data/GroundTruth/
+
+References:
+    - benGRN: https://github.com/your-org/benGRN
+    - Stone & Sroy gold standards
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+
+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 datarax.core.data_source import DataSourceModule
+
+logger = logging.getLogger(__name__)
+
+_BASE_DIR = Path("/media/mahdi/ssd23/Works/benGRN/data/GroundTruth/stone_and_sroy")
+
+
+@dataclass(frozen=True, kw_only=True)
+class BenGRNConfig(StructuralConfig):
+    """Configuration for BenGRNSource.
+
+    Attributes:
+        data_dir: Root directory of benGRN ground truth data.
+        species: Species to load ('mouse' or 'human').
+        expression_dataset: Which expression dataset to use.
+        ground_truth: Which ground truth network to use.
+        max_genes: Maximum number of genes to include (for speed).
+    """
+
+    data_dir: str = str(_BASE_DIR)
+    species: str = "mouse"
+    expression_dataset: str = "duren"
+    ground_truth: str = "chipunion_KDUnion_intersect"
+    max_genes: int | None = None
+
+    def __post_init__(self) -> None:
+        """Validate configuration."""
+        super().__post_init__()
+        base = Path(self.data_dir)
+        if not base.exists():
+            raise FileNotFoundError(
+                f"benGRN data not found: {base}. Clone benGRN repo to ../benGRN/"
+            )
+
+
+class BenGRNSource(DataSourceModule):
+    """DataSource for benGRN ground truth GRN data.
+
+    Loads expression matrix and ground truth regulatory edges for
+    GRN inference benchmarking.
+    """
+
+    data: dict[str, Any] = nnx.data()
+
+    def __init__(
+        self,
+        config: BenGRNConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ) -> None:
+        """Load the benGRN ground truth data."""
+        super().__init__(config, rngs=rngs, name=name or "BenGRNSource")
+        self.data = self._load(config)
+        logger.info(
+            "Loaded benGRN: %d cells, %d genes, %d TFs, %d GT edges",
+            self.data["n_cells"],
+            self.data["n_genes"],
+            self.data["n_tfs"],
+            self.data["n_edges"],
+        )
+
+    def _load(self, config: BenGRNConfig) -> dict[str, Any]:
+        """Load expression data and ground truth edges.
+
+        Coordinates loading of expression, TF, and ground truth data,
+        then builds the final indexed dataset.
+        """
+        base = Path(config.data_dir)
+        prefix = "mESC" if config.species == "mouse" else "hESC"
+
+        gene_names, expression = self._load_expression(base, config.expression_dataset)
+        n_cells = expression.shape[0]
+
+        tf_names = self._load_tf_names(base, prefix, config.expression_dataset, gene_names)
+        gt_edges = self._load_ground_truth_edges(base, prefix, config.ground_truth)
+
+        common_genes = self._compute_common_genes(gene_names, gt_edges, config.max_genes)
+        gene_to_idx = {g: i for i, g in enumerate(common_genes)}
+
+        counts = self._build_expression_submatrix(expression, gene_names, common_genes)
+        gt_matrix, n_edges = self._build_adjacency_matrix(gt_edges, gene_to_idx)
+
+        tf_indices = np.array(
+            [gene_to_idx[tf] for tf in tf_names if tf in gene_to_idx],
+            dtype=np.int32,
+        )
+
+        return {
+            "counts": counts,
+            "gene_names": common_genes,
+            "tf_names": [tf for tf in tf_names if tf in gene_to_idx],
+            "tf_indices": tf_indices,
+            "ground_truth_matrix": gt_matrix,
+            "ground_truth_edges": gt_edges,
+            "n_cells": n_cells,
+            "n_genes": len(common_genes),
+            "n_tfs": len(tf_indices),
+            "n_edges": n_edges,
+        }
+
+    def _load_expression(self, base: Path, expression_dataset: str) -> tuple[list[str], np.ndarray]:
+        """Load and parse the gzipped expression TSV file.
+
+        Args:
+            base: Root data directory.
+            expression_dataset: Name of the expression dataset.
+
+        Returns:
+            Tuple of (gene_names, expression_matrix) where expression_matrix
+            is shape (n_cells, n_genes).
+
+        Raises:
+            FileNotFoundError: If the expression file does not exist.
+        """
+        import gzip  # noqa: PLC0415
+
+        expr_file = base / "scRNA" / f"{expression_dataset}_rna_filtered_log2.tsv.gz"
+        if not expr_file.exists():
+            raise FileNotFoundError(f"Expression data not found: {expr_file}")
+
+        with gzip.open(expr_file, "rt") as f:
+            lines = f.readlines()
+
+        # TSV layout: row 0 = cell/sample header, rows 1+ = gene_name \t values...
+        gene_names: list[str] = []
+        expr_cols: list[list[float]] = []
+        for line in lines[1:]:
+            parts = line.strip().split("\t")
+            gene_names.append(parts[0])
+            expr_cols.append([float(v) for v in parts[1:]])
+
+        # Transpose from (n_genes, n_cells) to (n_cells, n_genes)
+        expression = np.array(expr_cols, dtype=np.float32).T
+        return gene_names, expression
+
+    def _load_tf_names(
+        self,
+        base: Path,
+        prefix: str,
+        expression_dataset: str,
+        gene_names: list[str],
+    ) -> list[str]:
+        """Load transcription factor names from TSV file.
+
+        Falls back to the first 50 genes if no TF file is found.
+
+        Args:
+            base: Root data directory.
+            prefix: Species prefix (mESC or hESC).
+            expression_dataset: Name of the expression dataset.
+            gene_names: Gene names from expression data (fallback source).
+
+        Returns:
+            List of transcription factor names.
+        """
+        tf_file = base / f"{prefix}_{expression_dataset.capitalize()}_TFs.tsv"
+        if not tf_file.exists():
+            tf_file = base / f"{prefix}_Duren_TFs.tsv"
+
+        if tf_file.exists():
+            return [line.strip() for line in tf_file.read_text().splitlines() if line.strip()]
+
+        logger.warning("TF file not found, using first 50 genes")
+        return gene_names[:50]
+
+    def _load_ground_truth_edges(
+        self, base: Path, prefix: str, ground_truth: str
+    ) -> list[tuple[str, str]]:
+        """Load ground truth regulatory edges from file.
+
+        Args:
+            base: Root data directory.
+            prefix: Species prefix (mESC or hESC).
+            ground_truth: Name of the ground truth network.
+
+        Returns:
+            List of (source_gene, target_gene) edge tuples.
+
+        Raises:
+            FileNotFoundError: If the ground truth file does not exist.
+        """
+        gt_file = base / "gold_standards" / prefix / f"{prefix}_{ground_truth}.txt"
+        if not gt_file.exists():
+            raise FileNotFoundError(f"Ground truth not found: {gt_file}")
+
+        gt_edges: list[tuple[str, str]] = []
+        for line in gt_file.read_text().splitlines():
+            parts = line.strip().split("\t")
+            if len(parts) >= 2:
+                gt_edges.append((parts[0], parts[1]))
+        return gt_edges
+
+    def _compute_common_genes(
+        self,
+        gene_names: list[str],
+        gt_edges: list[tuple[str, str]],
+        max_genes: int | None,
+    ) -> list[str]:
+        """Compute sorted intersection of expression and ground truth genes.
+
+        Args:
+            gene_names: Gene names from expression data.
+            gt_edges: Ground truth edge list.
+            max_genes: Optional cap on the number of genes.
+
+        Returns:
+            Sorted list of common gene names.
+        """
+        expr_gene_set = set(gene_names)
+        gt_genes: set[str] = set()
+        for src, tgt in gt_edges:
+            gt_genes.add(src)
+            gt_genes.add(tgt)
+
+        common_genes = sorted(expr_gene_set & gt_genes)
+        if max_genes is not None:
+            common_genes = common_genes[:max_genes]
+        return common_genes
+
+    def _build_expression_submatrix(
+        self,
+        expression: np.ndarray,
+        gene_names: list[str],
+        common_genes: list[str],
+    ) -> jnp.ndarray:
+        """Extract expression columns for the common gene set.
+
+        Args:
+            expression: Full expression matrix (n_cells, n_genes).
+            gene_names: All gene names corresponding to expression columns.
+            common_genes: Subset of genes to keep.
+
+        Returns:
+            JAX array of shape (n_cells, len(common_genes)).
+        """
+        expr_col_idx = [gene_names.index(g) for g in common_genes if g in gene_names]
+        return jnp.array(expression[:, expr_col_idx])
+
+    def _build_adjacency_matrix(
+        self,
+        gt_edges: list[tuple[str, str]],
+        gene_to_idx: dict[str, int],
+    ) -> tuple[np.ndarray, int]:
+        """Build binary ground truth adjacency matrix.
+
+        Args:
+            gt_edges: Ground truth edge list.
+            gene_to_idx: Mapping from gene name to matrix index.
+
+        Returns:
+            Tuple of (adjacency_matrix, edge_count).
+        """
+        n_genes = len(gene_to_idx)
+        gt_matrix = np.zeros((n_genes, n_genes), dtype=np.float32)
+        n_edges = 0
+        for src, tgt in gt_edges:
+            if src in gene_to_idx and tgt in gene_to_idx:
+                gt_matrix[gene_to_idx[src], gene_to_idx[tgt]] = 1.0
+                n_edges += 1
+        return gt_matrix, n_edges
+
+    def load(self) -> dict[str, Any]:
+        """Return the full dataset."""
+        return self.data
+
+    def __len__(self) -> int:
+        """Return number of cells."""
+        return self.data["n_cells"]
+
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        """Iterate over cells."""
+        for i in range(len(self)):
+            yield {
+                k: v[i]
+                if hasattr(v, "__getitem__")
+                and k not in ("gene_names", "tf_names", "ground_truth_edges")
+                else v
+                for k, v in self.data.items()
+            }

diffbio/sources/contextual_epigenomics.py

@@ -0,0 +1,242 @@
+"""Shared source validation for contextual epigenomics workloads."""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+import jax.numpy as jnp
+import numpy as np
+
+from diffbio.sequences.dna import encode_dna_string
+
+CONTEXTUAL_EPIGENOMICS_DATASET_CONTRACT_KEYS = (
+    "sequence",
+    "tf_context",
+    "chromatin_contacts",
+    "targets",
+)
+CONTEXTUAL_TARGET_SEMANTICS = (
+    "binary_peak_mask",
+    "chromatin_state_id",
+)
+
+_GENERIC_MOTIF = "ACGTAC"
+
+
+def validate_contextual_epigenomics_dataset(
+    data: dict[str, Any],
+    *,
+    target_semantics: Literal["binary_peak_mask", "chromatin_state_id"] | None = None,
+    num_output_classes: int | None = None,
+) -> None:
+    """Validate the shared contextual epigenomics benchmark contract."""
+    missing_keys = [key for key in CONTEXTUAL_EPIGENOMICS_DATASET_CONTRACT_KEYS if key not in data]
+    if missing_keys:
+        raise ValueError(f"Contextual epigenomics data is missing required keys: {missing_keys}")
+
+    sequence = jnp.asarray(data["sequence"], dtype=jnp.float32)
+    tf_context = jnp.asarray(data["tf_context"], dtype=jnp.float32)
+    chromatin_contacts = jnp.asarray(data["chromatin_contacts"], dtype=jnp.float32)
+    targets = np.asarray(data["targets"])
+
+    if sequence.ndim != 3 or sequence.shape[-1] != 4:
+        raise ValueError("sequence must have shape (n_examples, sequence_length, 4).")
+    if tf_context.ndim != 2:
+        raise ValueError("tf_context must have shape (n_examples, n_tf_features).")
+    if chromatin_contacts.ndim != 3:
+        raise ValueError(
+            "chromatin_contacts must have shape (n_examples, sequence_length, sequence_length)."
+        )
+    if targets.ndim != 2:
+        raise ValueError("targets must have shape (n_examples, sequence_length).")
+
+    n_examples = int(sequence.shape[0])
+    sequence_length = int(sequence.shape[1])
+    if (
+        tf_context.shape[0] != n_examples
+        or chromatin_contacts.shape[0] != n_examples
+        or targets.shape[0] != n_examples
+    ):
+        raise ValueError(
+            "Contextual epigenomics data keys must all share the same leading dimension."
+        )
+
+    if (
+        chromatin_contacts.shape[1] != sequence_length
+        or chromatin_contacts.shape[2] != sequence_length
+    ):
+        raise ValueError(
+            "chromatin_contacts must align with sequence length and provide a square contact map."
+        )
+    if targets.shape[1] != sequence_length:
+        raise ValueError("targets must have shape (n_examples, sequence_length).")
+
+    sequence_mass = np.asarray(sequence.sum(axis=-1))
+    if not np.allclose(sequence_mass, 1.0, atol=1e-5):
+        raise ValueError(
+            "sequence must be one-hot or probability-normalized along the alphabet axis."
+        )
+
+    contacts = np.asarray(chromatin_contacts)
+    if not np.allclose(contacts, np.swapaxes(contacts, -1, -2), atol=1e-5):
+        raise ValueError("chromatin_contacts must be symmetric.")
+
+    if target_semantics is None:
+        return
+    if target_semantics not in CONTEXTUAL_TARGET_SEMANTICS:
+        raise ValueError(
+            "target_semantics must be one of "
+            f"{CONTEXTUAL_TARGET_SEMANTICS}, got {target_semantics!r}."
+        )
+
+    _validate_target_semantics(
+        targets=targets,
+        target_semantics=target_semantics,
+        num_output_classes=num_output_classes,
+    )
+
+
+def _validate_target_semantics(
+    *,
+    targets: np.ndarray,
+    target_semantics: Literal["binary_peak_mask", "chromatin_state_id"],
+    num_output_classes: int | None,
+) -> None:
+    """Validate target values against the declared contextual task semantics."""
+    if target_semantics == "binary_peak_mask":
+        unique_targets = set(np.unique(targets).tolist())
+        if not unique_targets.issubset({0, 1}):
+            raise ValueError("binary_peak_mask targets must contain only 0/1 labels.")
+        return
+
+    if num_output_classes is None or num_output_classes < 1:
+        raise ValueError("num_output_classes must be positive for chromatin_state_id targets.")
+
+    if not np.all(np.equal(targets, np.floor(targets))):
+        raise ValueError("chromatin_state_id targets must be integer class IDs.")
+
+    if np.any(targets < 0) or np.any(targets >= num_output_classes):
+        raise ValueError(
+            f"chromatin_state_id targets must be in the range [0, {num_output_classes})."
+        )
+
+
+def build_synthetic_contextual_epigenomics_dataset(
+    *,
+    n_examples: int,
+    sequence_length: int,
+    num_tf_features: int,
+    target_semantics: Literal["binary_peak_mask", "chromatin_state_id"],
+    num_output_classes: int,
+) -> dict[str, jnp.ndarray]:
+    """Build a deterministic contextual epigenomics dataset."""
+    if target_semantics not in CONTEXTUAL_TARGET_SEMANTICS:
+        raise ValueError(
+            "target_semantics must be one of "
+            f"{CONTEXTUAL_TARGET_SEMANTICS}, got {target_semantics!r}."
+        )
+
+    sequences: list[str] = []
+    tf_context_rows: list[np.ndarray] = []
+    chromatin_contacts: list[np.ndarray] = []
+    targets: list[np.ndarray] = []
+
+    motif_length = len(_GENERIC_MOTIF)
+    peak_width = max(motif_length, 4)
+    candidate_starts = _build_candidate_starts(
+        sequence_length=sequence_length,
+        num_tf_features=num_tf_features,
+        peak_width=peak_width,
+    )
+    template_sequence = _build_sequence_template(
+        sequence_length=sequence_length,
+        candidate_starts=candidate_starts,
+        motif=_GENERIC_MOTIF,
+    )
+
+    for example_index in range(n_examples):
+        tf_index = example_index % num_tf_features
+        region_start = candidate_starts[tf_index]
+        region_end = min(region_start + peak_width, sequence_length)
+        sequences.append(template_sequence)
+
+        tf_row = np.zeros(num_tf_features, dtype=np.float32)
+        tf_row[tf_index] = 1.0
+        tf_context_rows.append(tf_row)
+
+        contact_map = np.eye(sequence_length, dtype=np.float32) * 0.1
+        contact_map[region_start:region_end, region_start:region_end] = 1.0
+        flank_start = max(region_start - 1, 0)
+        flank_end = min(region_end + 1, sequence_length)
+        contact_map[flank_start:flank_end, flank_start:flank_end] = np.maximum(
+            contact_map[flank_start:flank_end, flank_start:flank_end],
+            0.5,
+        )
+        contact_map = np.maximum(contact_map, contact_map.T)
+        chromatin_contacts.append(contact_map)
+
+        target = np.zeros(sequence_length, dtype=np.int32)
+        if target_semantics == "binary_peak_mask":
+            target[region_start:region_end] = 1
+        else:
+            target[region_start:region_end] = 1
+            right_flank_end = min(region_end + peak_width // 2, sequence_length)
+            target[region_end:right_flank_end] = 2 % num_output_classes
+        targets.append(target)
+
+    one_hot_sequences = jnp.asarray(
+        np.stack(
+            [np.asarray(encode_dna_string(sequence), dtype=np.float32) for sequence in sequences]
+        ),
+        dtype=jnp.float32,
+    )
+    dataset = {
+        "sequence": one_hot_sequences,
+        "tf_context": jnp.asarray(np.stack(tf_context_rows), dtype=jnp.float32),
+        "chromatin_contacts": jnp.asarray(np.stack(chromatin_contacts), dtype=jnp.float32),
+        "targets": jnp.asarray(np.stack(targets)),
+    }
+    validate_contextual_epigenomics_dataset(
+        dataset,
+        target_semantics=target_semantics,
+        num_output_classes=num_output_classes,
+    )
+    return dataset
+
+
+def _build_candidate_starts(
+    *,
+    sequence_length: int,
+    num_tf_features: int,
+    peak_width: int,
+) -> list[int]:
+    """Build deterministic candidate windows selected by TF context."""
+    if sequence_length < peak_width + 2:
+        return [0 for _ in range(num_tf_features)]
+
+    max_start = max(sequence_length - peak_width - 1, 1)
+    raw_starts = np.linspace(
+        1,
+        max_start,
+        num=num_tf_features,
+        dtype=np.int32,
+    )
+    return [int(start) for start in raw_starts]
+
+
+def _build_sequence_template(
+    *,
+    sequence_length: int,
+    candidate_starts: list[int],
+    motif: str,
+) -> str:
+    """Build one shared sequence carrying several candidate regulatory windows."""
+    sequence = ["A"] * sequence_length
+    motif_chars = list(motif)
+    motif_length = len(motif_chars)
+
+    for start in candidate_starts:
+        end = min(start + motif_length, sequence_length)
+        sequence[start:end] = motif_chars[: end - start]
+
+    return "".join(sequence)