A file may lack a rep its streams never select from

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/<k>"`` | ``"layers/<k>"``)."""
+    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/<k>" | "layers/<k>"
+        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])

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)

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.

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.