From 80ff7aa0d3f7d30f4825fc09a9ecfffbce8fa035 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Tue, 20 Oct 2026 14:31:30 +0200 Subject: [PATCH] reps only where selected --- src/scfit/data/_io.py | 51 ++++++++++++++++++++++++++++++++++----- src/scfit/data/_loader.py | 9 ++++--- src/scfit/data/_source.py | 36 +++++++++++++++++++++------ tests/data/test_source.py | 44 ++++++++++++++++++++++++++++++++- 4 files changed, 122 insertions(+), 18 deletions(-) diff --git a/src/scfit/data/_io.py b/src/scfit/data/_io.py index 914a0c5..75c6d78 100644 --- a/src/scfit/data/_io.py +++ b/src/scfit/data/_io.py @@ -50,6 +50,36 @@ def get_from_container(source: ad.AnnData | list[ad.AnnData], loc: str) -> list[ return [getattr(source, field)[sub]] +def has_rep(adata: ad.AnnData, loc: str) -> bool: + """Whether ``adata`` carries rep ``loc`` (``"X"`` | ``"obsm/"`` | ``"layers/"``).""" + if loc == "X": + return adata.X is not None + field, sub = loc.split("/", 1) + return sub in getattr(adata, field) + + +def rep_placeholder(template, n_obs: int): + """A never-read stand-in for a rep a file lacks: ``n_obs`` rows shaped and typed like ``template``, costing no memory. + + Stands in for a file whose rows a stream never selects, so annbatch still gets one backing per file (its global + row indexing spans all of them). Backed templates get an unwritten in-memory zarr array (all fill value) or an + empty CSR group; in-memory ones a broadcast zero view or an empty CSR matrix. + """ + import scipy.sparse as sp + + shape, dtype = (n_obs, int(template.shape[1])), template.dtype + sparse = isinstance(template, ad.abc.CSRDataset) or sp.issparse(template) + if not is_backed_array(template): + return sp.csr_matrix(shape, dtype=dtype) if sparse else np.broadcast_to(np.zeros((), dtype=dtype), shape) + group = zarr.open_group(zarr.storage.MemoryStore(), mode="w") + if sparse: + ad.io.write_elem(group, "x", sp.csr_matrix(shape, dtype=dtype)) + return ad.io.sparse_dataset(group["x"]) + return group.create_array( + "x", shape=shape, chunks=(max(1, min(n_obs, 65_536)), shape[1]), dtype=dtype, fill_value=0 + ) + + def _read_rows(source: Container, loc: str, row_idx: np.ndarray) -> np.ndarray: r"""Densely read the global rows ``row_idx`` (ascending, unique) of rep ``loc`` from a source's backings. @@ -58,17 +88,24 @@ def _read_rows(source: Container, loc: str, row_idx: np.ndarray) -> np.ndarray: scattered rows degrade to one slice per row (still correct). Slicing reads uniformly across dense zarr, backed ``CSRDataset``\\s and in-memory arrays, so no per-backend indexing branch is needed. """ - backings = get_from_container(source, loc) # each has ``.shape`` (sparse groups already wrapped by read_backing) - offs = np.concatenate([[0], np.cumsum([b.shape[0] for b in backings])]).astype(np.int64) - parts = [] - for d, b in enumerate(backings): + adatas = source if isinstance(source, list) else [source] + offs = np.concatenate([[0], np.cumsum([a.n_obs for a in adatas])]).astype(np.int64) + parts, width = [], None + for d, a in enumerate(adatas): lb = row_idx[(row_idx >= offs[d]) & (row_idx < offs[d + 1])] - offs[d] # ascending within this dataset - runs = np.split(lb, np.flatnonzero(np.diff(lb) != 1) + 1) if lb.size else () # maximal consecutive runs + if has_rep(a, loc): + width = width if width is not None else int(get_from_container(a, loc)[0].shape[1]) + if not lb.size: + continue # only files with selected rows need to carry the rep + if not has_rep(a, loc): + raise KeyError(f"rep {loc!r} is missing from AnnData {d}, which has {lb.size} of the selected rows.") + (b,) = get_from_container(a, loc) # has ``.shape`` (sparse groups already wrapped by read_backing) + runs = np.split(lb, np.flatnonzero(np.diff(lb) != 1) + 1) # maximal consecutive runs for run in runs: sub = b[int(run[0]) : int(run[-1]) + 1] # contiguous slice read dense = getattr(sub, "todense", None) parts.append(np.asarray(dense()) if dense is not None else np.asarray(sub)) - return np.concatenate(parts, axis=0) if parts else np.empty((0, int(backings[0].shape[1])), dtype=np.float32) + return np.concatenate(parts, axis=0) if parts else np.empty((0, width or 0), dtype=np.float32) def materialize_node( @@ -232,6 +269,8 @@ def load_backed_adata(g, *, keys: Sequence[str], cols: Sequence[str] = (), uns: if key == "X": continue field, sub = key.split("/", 1) # "obsm/" | "layers/" + if field not in g or sub not in g[field]: + continue # a file may lack a rep its streams never select from; `Source.rep` decides getattr(adata, field)[sub] = read_backing(g[field][sub]) for key in uns: adata.uns[key] = ad.io.read_elem(g["uns"][key]) diff --git a/src/scfit/data/_loader.py b/src/scfit/data/_loader.py index aa6f5cb..7f93974 100644 --- a/src/scfit/data/_loader.py +++ b/src/scfit/data/_loader.py @@ -241,14 +241,17 @@ def _build_per_rep_loaders(self, name: str, sampler: ClassSampler | BoundClassSa # One annbatch Loader per rep; each deepcopies the sampler so the reps stay aligned. Each file carries # its rows' leaf as obs, so a batch reports the group it was drawn from (`_LEAF` in its obs). src, s, to = self._resolved[name], self._streams[name], self._to - cats = self._st[name]["cats"] + cats, w = self._st[name]["cats"], self._st[name]["w"] + # Files the stream selects rows from; the others may lack a rep (a placeholder stands in, never read). + offsets = np.cumsum([0, *(a.n_obs for a in src.adatas)]) + selected = w[cats.codes] > 0 + needed = [bool(selected[a:b].any()) for a, b in itertools.pairwise(offsets)] loaders: dict[str, AnnbatchLoader] = {} for loc in s.reps: base = AnnbatchLoader( batch_sampler=deepcopy(sampler), return_index=False, to=to, preload_to_gpu=self._preload_to_gpu ) - backings = src.rep(loc) - offsets = np.cumsum([0, *(b.shape[0] for b in backings)]) + backings = src.rep(loc, needed) obs = [ pd.DataFrame({_LEAF: cats[a:b]}, index=np.arange(a, b).astype(str)) for a, b in itertools.pairwise(offsets) diff --git a/src/scfit/data/_source.py b/src/scfit/data/_source.py index 343e0c5..15f5290 100644 --- a/src/scfit/data/_source.py +++ b/src/scfit/data/_source.py @@ -6,14 +6,14 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, NamedTuple import anndata as ad import numpy as np import pandas as pd -from scfit.data._io import get_from_container, leaf_codes, materialize_node, obs_columns +from scfit.data._io import get_from_container, has_rep, leaf_codes, materialize_node, obs_columns, rep_placeholder if TYPE_CHECKING: from scfit.data._schema import Container, Stream @@ -71,21 +71,41 @@ def factorize(self, group_by: tuple[str, ...]) -> _Factorized: self._leaf_cache[gb] = _Factorized(codes, leaves, cats) return self._leaf_cache[gb] - def rep(self, loc: str) -> list: + def rep(self, loc: str, needed: Sequence[bool] | None = None) -> list: """The array(s) backing rep ``loc``, one per AnnData in list order — with a shared-width check. - annbatch stacks a streamed rep into one array, so its ``shape[1]`` must agree across the dataset's - files; a mismatch raises here rather than deep inside annbatch. + A file may lack the rep when ``needed`` says its rows are never selected (e.g. a target-side rep in the + controls' file): it gets a never-read placeholder (:func:`~scfit.data._io.rep_placeholder`), so annbatch + still sees one backing per file. A missing rep in a needed file raises. annbatch stacks a streamed rep + into one array, so its ``shape[1]`` must agree across the files that carry it; a mismatch raises here + rather than deep inside annbatch. """ - backings = get_from_container(self._adatas, loc) - widths = {int(b.shape[1]) for b in backings} + needed = [True] * len(self._adatas) if needed is None else list(needed) + backings = [] + for i, (a, need) in enumerate(zip(self._adatas, needed, strict=True)): + if has_rep(a, loc): + backings.append(get_from_container(a, loc)[0]) + elif not need: + backings.append(None) + else: + raise KeyError( + f"rep {loc!r} is missing from AnnData {i} of this source, and the stream selects rows from it. " + "Add the rep to that file, or give its rows zero weight in this stream." + ) + real = [b for b in backings if b is not None] + if not real: + raise KeyError(f"rep {loc!r} is in none of this source's AnnData.") + widths = {int(b.shape[1]) for b in real} if len(widths) > 1: raise ValueError( f"rep {loc!r} has inconsistent shape[1] across this source's AnnData ({sorted(widths)}); a " "streamed representation must share its feature dimension across files (align to a common " "space — e.g. a shared obsm embedding or a common gene panel — or use separate source_keys)." ) - return backings + return [ + b if b is not None else rep_placeholder(real[0], a.n_obs) + for a, b in zip(self._adatas, backings, strict=True) + ] def materialize(self, stream: Stream) -> Source: """A new :class:`Source` holding this stream's selected (positive-weight) cells in RAM. diff --git a/tests/data/test_source.py b/tests/data/test_source.py index fbaa1b3..6bfad63 100644 --- a/tests/data/test_source.py +++ b/tests/data/test_source.py @@ -14,7 +14,8 @@ import pytest from scheme_helpers import KEY, encoded_adata, perturbation_streams -from scfit.data import Loader, Stream +from scfit.data import Loader, Stream, open_source +from scfit.data._io import _read_rows from scfit.data._source import Source READ = {"sampler": {"batch_size": 4, "chunk_size": 1, "preload_nchunks": 4}, "to": None} @@ -62,6 +63,47 @@ def test_shared_obsm_rep_ok_despite_differing_x_width(): assert [x.shape[1] for x in src.rep("obsm/emb")] == [4, 4] # equal-width rep → no error +# ── reps a file may lack ──────────────────────────────────────────────────────────────────────── +def test_streamed_rep_only_in_the_selected_file(): + # e.g. a target-side rep that only the perturbed file carries: the primary never selects file A's rows. + a, b = _file(["T"], file_id=0), _file(["NK"], file_id=1) + b.obsm["emb"] = np.full((b.n_obs, 3), 7, "float32") + loader = Loader( + {KEY: [a, b]}, + primary=Stream(KEY, group_by=("cell_type",), weights={("NK",): 1.0}, reps=("X", "obsm/emb")), + **READ, + ) + for _ in range(4): + batch = next(iter(loader))["primary"] + assert (np.asarray(batch["X"])[:, 0] == 1).all() and (np.asarray(batch["obsm/emb"]) == 7).all() + + +def test_rep_missing_from_a_selected_file_raises(): + a, b = _file(["T"], 0), _file(["NK"], 1) + b.obsm["emb"] = np.zeros((b.n_obs, 3), "float32") + with pytest.raises(KeyError, match="missing from AnnData 0"): + Source([a, b]).rep("obsm/emb") + assert Source([a, b]).rep("obsm/emb", needed=[False, True])[0].shape == (a.n_obs, 3) + + +def test_read_rows_touches_only_files_with_selected_rows(): + a, b = _file(["T"], 0), _file(["NK"], 1) + a.obsm["emb"] = np.full((a.n_obs, 2), 3, "float32") + np.testing.assert_array_equal(_read_rows([a, b], "obsm/emb", np.arange(2, 5)), np.full((3, 2), 3)) + with pytest.raises(KeyError, match="missing from AnnData 1"): + _read_rows([a, b], "obsm/emb", np.array([0, a.n_obs])) + + +def test_backed_placeholder_for_a_file_lacking_the_rep(tmp_path): + a, b = _file(["T"], 0), _file(["NK"], 1) + b.obsm["emb"] = np.ones((b.n_obs, 3), "float32") + a.write_zarr(tmp_path / "a.zarr") + b.write_zarr(tmp_path / "b.zarr") + src = Source(open_source([tmp_path / "a.zarr", tmp_path / "b.zarr"], keys=["X", "obsm/emb"], cols=["cell_type"])) + placeholder, real = src.rep("obsm/emb", needed=[False, True]) + assert placeholder.shape == (a.n_obs, 3) and type(placeholder) is type(real) + + # ── shared factorization ──────────────────────────────────────────────────────────────────────── def test_shared_source_factorized_once(): # a primary and its matched control name the same source_key + group_by → the Source factorizes it once.