diffbio 0.1.0__py3-none-any.whl

diffbio/sources/perturbation/concat_source.py

@@ -0,0 +1,115 @@
+"""Multi-dataset concatenation source for perturbation experiments.
+
+Combines multiple PerturbationAnnDataSource instances into a single unified
+source with global indexing. Validates metadata consistency across sources.
+
+References:
+    - cell-load/src/cell_load/dataset/_metadata.py (MetadataConcatDataset)
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator
+from typing import Any
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+class PerturbationConcatSource:
+    """Concatenation of multiple PerturbationAnnDataSource instances.
+
+    Provides unified indexing across multiple sources. Global index ``i``
+    is mapped to the correct underlying source and local index.
+
+    Validates that all sources share consistent metadata column names
+    (control_pert, pert_col, etc.).
+
+    Args:
+        sources: List of PerturbationAnnDataSource instances.
+
+    Raises:
+        ValueError: If no sources are provided.
+    """
+
+    def __init__(
+        self,
+        sources: list[Any],
+    ) -> None:
+        if not sources:
+            raise ValueError("PerturbationConcatSource requires at least one source.")
+
+        self._sources = sources
+
+        # Build cumulative length offsets for global -> local index mapping
+        self._offsets: list[int] = []
+        cumulative = 0
+        for s in sources:
+            self._offsets.append(cumulative)
+            cumulative += len(s)
+        self._total_length = cumulative
+
+    def __len__(self) -> int:
+        """Return total number of cells across all sources."""
+        return self._total_length
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        """Get element by global index.
+
+        Args:
+            idx: Global cell index.
+
+        Returns:
+            Per-cell dictionary from the appropriate source.
+
+        Raises:
+            IndexError: If index is out of range.
+        """
+        if idx < 0:
+            idx = self._total_length + idx
+        if idx < 0 or idx >= self._total_length:
+            raise IndexError(
+                f"Index {idx} out of range for concat source with {self._total_length} cells"
+            )
+
+        source_idx, local_idx = self._global_to_local(idx)
+        return self._sources[source_idx][local_idx]
+
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        """Iterate over all cells across all sources in order."""
+        for source in self._sources:
+            yield from source
+
+    def get_control_mask(self) -> np.ndarray:
+        """Return concatenated boolean control mask."""
+        return np.concatenate([s.get_control_mask() for s in self._sources])
+
+    def get_group_codes(self) -> np.ndarray:
+        """Return concatenated group codes."""
+        return np.concatenate([s.get_group_codes() for s in self._sources])
+
+    def get_pert_codes(self) -> np.ndarray:
+        """Return concatenated perturbation codes."""
+        return np.concatenate([s.get_pert_codes() for s in self._sources])
+
+    def get_cell_type_codes(self) -> np.ndarray:
+        """Return concatenated cell type codes."""
+        return np.concatenate([s.get_cell_type_codes() for s in self._sources])
+
+    def get_batch_codes(self) -> np.ndarray:
+        """Return concatenated batch codes."""
+        return np.concatenate([s.get_batch_codes() for s in self._sources])
+
+    @property
+    def sources(self) -> list[Any]:
+        """Return the underlying source list."""
+        return self._sources
+
+    def _global_to_local(self, global_idx: int) -> tuple[int, int]:
+        """Map a global index to (source_index, local_index)."""
+        for i in range(len(self._sources) - 1, -1, -1):
+            if global_idx >= self._offsets[i]:
+                return i, global_idx - self._offsets[i]
+        return 0, global_idx  # pragma: no cover

diffbio/sources/perturbation/control_mapping.py

@@ -0,0 +1,215 @@
+"""Control cell mapping strategies for perturbation experiments.
+
+Maps perturbed cells to control cells using batch-based or random strategies.
+Mappings are precomputed at setup time as numpy index arrays.
+
+References:
+    - cell-load/src/cell_load/mapping_strategies/batch.py
+    - cell-load/src/cell_load/mapping_strategies/random.py
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+from datarax.core.config import StructuralConfig
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class ControlMappingConfig(StructuralConfig):
+    """Configuration for control cell mapping.
+
+    Attributes:
+        strategy: Mapping strategy (``"batch"`` or ``"random"``).
+        n_basal_samples: Number of control cells per perturbed cell.
+        seed: Random seed for reproducibility.
+        map_controls: Whether to also map control cells to other controls.
+        cache_pairs: Whether to cache the mapping after first computation.
+    """
+
+    strategy: str = "random"
+    n_basal_samples: int = 1
+    seed: int = 42
+    map_controls: bool = False
+    cache_pairs: bool = False
+
+
+def _build_ctrl_pools_by_ct(ctrl_mask: np.ndarray, ct_codes: np.ndarray) -> dict[int, np.ndarray]:
+    """Build control index pools grouped by cell type."""
+    pools: dict[int, list[int]] = {}
+    for idx in np.where(ctrl_mask)[0]:
+        ct = int(ct_codes[idx])
+        if ct not in pools:
+            pools[ct] = []
+        pools[ct].append(idx)
+    return {k: np.array(v) for k, v in pools.items()}
+
+
+def _map_cells_to_controls(
+    cell_indices: np.ndarray,
+    ct_codes: np.ndarray,
+    ctrl_by_ct: dict[int, np.ndarray],
+    n_basal: int,
+    rng: np.random.Generator,
+) -> np.ndarray:
+    """Map a set of cell indices to control cells from the same cell type."""
+    mapping = np.empty((len(cell_indices), n_basal), dtype=np.int64)
+    for i, cidx in enumerate(cell_indices):
+        ct = int(ct_codes[cidx])
+        pool = ctrl_by_ct.get(ct, np.array([], dtype=np.int64))
+        if len(pool) == 0:
+            mapping[i] = -1
+            continue
+        chosen = rng.choice(pool, size=n_basal, replace=len(pool) < n_basal)
+        mapping[i] = chosen
+    return mapping
+
+
+class RandomControlMapping:
+    """Map perturbed cells to random controls of the same cell type.
+
+    For each perturbed cell, randomly selects ``n_basal_samples`` control cells
+    from the same cell type, pooled across all batches.
+
+    When ``map_controls=True``, also maps each control cell to another random
+    control of the same cell type.
+
+    When ``cache_pairs=True``, the mapping is computed once and cached for
+    subsequent calls.
+
+    Args:
+        config: Mapping configuration.
+    """
+
+    def __init__(self, config: ControlMappingConfig) -> None:
+        self._config = config
+        self._cached_mapping: np.ndarray | None = None
+
+    def build_mapping(self, source: Any) -> np.ndarray:
+        """Build mapping from cells to control cells.
+
+        Args:
+            source: A PerturbationAnnDataSource or PerturbationConcatSource.
+
+        Returns:
+            Array of shape ``(n_mapped, n_basal_samples)`` with control
+            cell indices. When ``map_controls=False``, ``n_mapped`` equals
+            the number of perturbed cells. When ``True``, ``n_mapped``
+            equals the total number of cells.
+        """
+        if self._config.cache_pairs and self._cached_mapping is not None:
+            return self._cached_mapping
+
+        ctrl_mask = source.get_control_mask()
+        ct_codes = source.get_cell_type_codes()
+        n_basal = self._config.n_basal_samples
+        rng = np.random.default_rng(self._config.seed)
+
+        ctrl_by_ct = _build_ctrl_pools_by_ct(ctrl_mask, ct_codes)
+
+        if self._config.map_controls:
+            # Map ALL cells (perturbed + controls) to controls
+            all_indices = np.arange(len(ctrl_mask))
+            mapping = _map_cells_to_controls(all_indices, ct_codes, ctrl_by_ct, n_basal, rng)
+        else:
+            # Map only perturbed cells
+            pert_indices = np.where(~ctrl_mask)[0]
+            mapping = _map_cells_to_controls(pert_indices, ct_codes, ctrl_by_ct, n_basal, rng)
+
+        if self._config.cache_pairs:
+            self._cached_mapping = mapping
+
+        return mapping
+
+
+class BatchControlMapping:
+    """Map perturbed cells to controls within the same batch and cell type.
+
+    Prefers controls from the same (batch, cell_type) group. Falls back to
+    all controls from the same cell type if the batch group is empty.
+
+    When ``map_controls=True``, also maps control cells.
+    When ``cache_pairs=True``, caches after first computation.
+
+    Args:
+        config: Mapping configuration.
+    """
+
+    def __init__(self, config: ControlMappingConfig) -> None:
+        self._config = config
+        self._cached_mapping: np.ndarray | None = None
+
+    def build_mapping(self, source: Any) -> np.ndarray:
+        """Build mapping from cells to control cells.
+
+        Args:
+            source: A PerturbationAnnDataSource or PerturbationConcatSource.
+
+        Returns:
+            Array of shape ``(n_mapped, n_basal_samples)`` with control
+            cell indices.
+        """
+        if self._config.cache_pairs and self._cached_mapping is not None:
+            return self._cached_mapping
+
+        ctrl_mask = source.get_control_mask()
+        ct_codes = source.get_cell_type_codes()
+        batch_codes = source.get_batch_codes()
+        n_basal = self._config.n_basal_samples
+        rng = np.random.default_rng(self._config.seed)
+
+        # Build control pools by (batch, cell_type) and by cell_type
+        ctrl_indices = np.where(ctrl_mask)[0]
+
+        ctrl_by_batch_ct: dict[tuple[int, int], list[int]] = {}
+        ctrl_by_ct: dict[int, list[int]] = {}
+
+        for idx in ctrl_indices:
+            ct = int(ct_codes[idx])
+            batch = int(batch_codes[idx])
+            key = (batch, ct)
+            if key not in ctrl_by_batch_ct:
+                ctrl_by_batch_ct[key] = []
+            ctrl_by_batch_ct[key].append(idx)
+
+            if ct not in ctrl_by_ct:
+                ctrl_by_ct[ct] = []
+            ctrl_by_ct[ct].append(idx)
+
+        ctrl_by_batch_ct_arr = {k: np.array(v) for k, v in ctrl_by_batch_ct.items()}
+        ctrl_by_ct_arr = {k: np.array(v) for k, v in ctrl_by_ct.items()}
+
+        # Determine which cells to map
+        if self._config.map_controls:
+            cells_to_map = np.arange(len(ctrl_mask))
+        else:
+            cells_to_map = np.where(~ctrl_mask)[0]
+
+        mapping = np.empty((len(cells_to_map), n_basal), dtype=np.int64)
+
+        for i, cidx in enumerate(cells_to_map):
+            ct = int(ct_codes[cidx])
+            batch = int(batch_codes[cidx])
+            key = (batch, ct)
+
+            pool = ctrl_by_batch_ct_arr.get(key)
+            if pool is None or len(pool) == 0:
+                pool = ctrl_by_ct_arr.get(ct, np.array([], dtype=np.int64))
+
+            if len(pool) == 0:
+                mapping[i] = -1
+                continue
+
+            chosen = rng.choice(pool, size=n_basal, replace=len(pool) < n_basal)
+            mapping[i] = chosen
+
+        if self._config.cache_pairs:
+            self._cached_mapping = mapping
+
+        return mapping

diffbio/sources/perturbation/experiment_config.py

@@ -0,0 +1,261 @@
+"""Experiment configuration for perturbation experiments.
+
+Parses TOML configuration files following the cell-load schema and provides
+frozen dataclass representations of experiment settings: dataset paths,
+training assignments, zero-shot cell types, and few-shot perturbation splits.
+
+References:
+    - cell-load/src/cell_load/config.py (ExperimentConfig)
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from datarax.core.config import StructuralConfig
+
+logger = logging.getLogger(__name__)
+
+
+def _require_toml() -> Any:
+    """Import tomllib from the standard library (Python 3.11+)."""
+    import tomllib  # noqa: PLC0415
+
+    return tomllib
+
+
+# ---------------------------------------------------------------------------
+# Data classes
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class DatasetEntry:
+    """A single dataset in an experiment.
+
+    Attributes:
+        name: Identifier for the dataset.
+        path: Filesystem path to the dataset directory or file.
+    """
+
+    name: str
+    path: str
+
+
+@dataclass(frozen=True)
+class ZeroshotEntry:
+    """A cell type held out entirely for zero-shot evaluation.
+
+    Attributes:
+        dataset: Name of the dataset containing this cell type.
+        cell_type: Cell type to hold out.
+        split: Target split (``"val"`` or ``"test"``).
+    """
+
+    dataset: str
+    cell_type: str
+    split: str
+
+
+@dataclass(frozen=True)
+class FewshotEntry:
+    """Perturbations held out within a cell type for few-shot evaluation.
+
+    Attributes:
+        dataset: Name of the dataset containing this cell type.
+        cell_type: Cell type within which perturbations are split.
+        val_perturbations: Perturbation names assigned to validation.
+        test_perturbations: Perturbation names assigned to testing.
+    """
+
+    dataset: str
+    cell_type: str
+    val_perturbations: tuple[str, ...] = ()
+    test_perturbations: tuple[str, ...] = ()
+
+
+@dataclass(frozen=True)
+class ExperimentConfig(StructuralConfig):
+    """Top-level experiment configuration.
+
+    Encapsulates all settings parsed from a TOML file: dataset paths,
+    training assignments, zero-shot cell type holdouts, and few-shot
+    perturbation splits.
+
+    Attributes:
+        datasets: Registered dataset entries.
+        training_datasets: Names of datasets used for training.
+        zeroshot: Cell types held out entirely.
+        fewshot: Perturbation-level splits within cell types.
+    """
+
+    datasets: tuple[DatasetEntry, ...] = ()
+    training_datasets: tuple[str, ...] = ()
+    zeroshot: tuple[ZeroshotEntry, ...] = ()
+    fewshot: tuple[FewshotEntry, ...] = ()
+
+    def get_all_datasets(self) -> set[str]:
+        """Return the set of all dataset names referenced in the config."""
+        names = set(self.training_datasets)
+        for z in self.zeroshot:
+            names.add(z.dataset)
+        for f in self.fewshot:
+            names.add(f.dataset)
+        return names
+
+    def get_zeroshot_celltypes(self, dataset: str) -> dict[str, str]:
+        """Get zero-shot cell types and their target splits for a dataset.
+
+        Args:
+            dataset: Dataset name to filter by.
+
+        Returns:
+            Dict mapping cell type names to split names.
+        """
+        return {z.cell_type: z.split for z in self.zeroshot if z.dataset == dataset}
+
+    def get_fewshot_celltypes(self, dataset: str) -> dict[str, FewshotEntry]:
+        """Get few-shot cell type entries for a dataset.
+
+        Args:
+            dataset: Dataset name to filter by.
+
+        Returns:
+            Dict mapping cell type names to FewshotEntry objects.
+        """
+        return {f.cell_type: f for f in self.fewshot if f.dataset == dataset}
+
+    def validate(self) -> None:
+        """Validate configuration consistency.
+
+        Raises:
+            ValueError: If referenced datasets lack paths or splits are invalid.
+        """
+        all_referenced = self.get_all_datasets()
+        dataset_names = {d.name for d in self.datasets}
+        missing = all_referenced - dataset_names
+        if missing:
+            raise ValueError(f"Missing dataset paths for: {missing}")
+
+        valid_splits = {"train", "val", "test"}
+        for z in self.zeroshot:
+            if z.split not in valid_splits:
+                raise ValueError(
+                    f"Invalid split '{z.split}' for zeroshot entry "
+                    f"{z.dataset}.{z.cell_type}. Must be one of {valid_splits}"
+                )
+
+        logger.info("Configuration validation passed")
+
+    def save_config(self, path: Path) -> None:
+        """Save configuration to a TOML file.
+
+        Args:
+            path: Output file path.
+        """
+        lines: list[str] = ["[datasets]"]
+        for d in self.datasets:
+            lines.append(f'{d.name} = "{d.path}"')
+
+        lines.append("")
+        lines.append("[training]")
+        for name in self.training_datasets:
+            lines.append(f'{name} = "train"')
+
+        if self.zeroshot:
+            lines.append("")
+            lines.append("[zeroshot]")
+            for z in self.zeroshot:
+                lines.append(f'"{z.dataset}.{z.cell_type}" = "{z.split}"')
+
+        for f in self.fewshot:
+            lines.append("")
+            lines.append(f'[fewshot."{f.dataset}.{f.cell_type}"]')
+            if f.val_perturbations:
+                val_list = ", ".join(f'"{p}"' for p in f.val_perturbations)
+                lines.append(f"val = [{val_list}]")
+            if f.test_perturbations:
+                test_list = ", ".join(f'"{p}"' for p in f.test_perturbations)
+                lines.append(f"test = [{test_list}]")
+
+        path.write_text("\n".join(lines) + "\n")
+
+
+# ---------------------------------------------------------------------------
+# TOML loading
+# ---------------------------------------------------------------------------
+
+
+def load_experiment_config(path: Path) -> ExperimentConfig:
+    """Load an experiment configuration from a TOML file.
+
+    Supports the cell-load TOML schema::
+
+        [datasets]
+        dataset_name = "/path/to/data"
+
+        [training]
+        dataset_name = "train"
+
+        [zeroshot]
+        "dataset.celltype" = "test"
+
+        [fewshot."dataset.celltype"]
+        val = ["Gene1"]
+        test = ["Gene2"]
+
+    Args:
+        path: Path to the TOML file.
+
+    Returns:
+        Parsed ExperimentConfig.
+
+    Raises:
+        FileNotFoundError: If the TOML file does not exist.
+    """
+    path = Path(path)
+    if not path.exists():
+        raise FileNotFoundError(f"TOML config file not found: {path}")
+
+    tomllib = _require_toml()
+    with open(path, "rb") as f:
+        raw = tomllib.load(f)
+
+    # Parse [datasets]
+    datasets_raw: dict[str, str] = raw.get("datasets", {})
+    datasets = tuple(DatasetEntry(name=name, path=p) for name, p in datasets_raw.items())
+
+    # Parse [training]
+    training_raw: dict[str, str] = raw.get("training", {})
+    training_datasets = tuple(training_raw.keys())
+
+    # Parse [zeroshot] -- keys are "dataset.celltype"
+    zeroshot_raw: dict[str, str] = raw.get("zeroshot", {})
+    zeroshot_entries: list[ZeroshotEntry] = []
+    for key, split in zeroshot_raw.items():
+        dataset, cell_type = key.split(".", 1)
+        zeroshot_entries.append(ZeroshotEntry(dataset=dataset, cell_type=cell_type, split=split))
+
+    # Parse [fewshot] -- keys are "dataset.celltype", values are {split: [perts]}
+    fewshot_raw: dict[str, dict[str, list[str]]] = raw.get("fewshot", {})
+    fewshot_entries: list[FewshotEntry] = []
+    for key, pert_config in fewshot_raw.items():
+        dataset, cell_type = key.split(".", 1)
+        fewshot_entries.append(
+            FewshotEntry(
+                dataset=dataset,
+                cell_type=cell_type,
+                val_perturbations=tuple(pert_config.get("val", [])),
+                test_perturbations=tuple(pert_config.get("test", [])),
+            )
+        )
+
+    return ExperimentConfig(
+        datasets=datasets,
+        training_datasets=training_datasets,
+        zeroshot=tuple(zeroshot_entries),
+        fewshot=tuple(fewshot_entries),
+    )