diff --git a/README.md b/README.md index 13a01cf..4776e58 100644 --- a/README.md +++ b/README.md @@ -143,6 +143,16 @@ va = vsparse.load_packed("archived.h5ad") norm = va.X.normalized("parafac2") # lazy view; norm @ B works without densifying ``` +`load_packed` can also subset as it decodes: `obs_idx`/`var_idx` take a boolean mask, an array of integer positions, or a slice, and give the same result as indexing a full load, without ever building the full matrix. With IVCSR, only the selected cells are decoded and unselected genes are dropped in the same pass (IVCSC is the same with the axes swapped). Every packed byte is still read, so reading the stream sets the floor: on a 1.26M-cell file, a 1% cell subset decodes `X` in 0.21 s against 0.50 s for all of it. + +```python +import anndata, h5py + +with h5py.File("archived.h5ad") as f: # read only obs to build the mask + obs = anndata.io.read_elem(f["obs"]) +va_blood = vsparse.load_packed("archived.h5ad", obs_idx=obs["tissue"] == "blood") +``` + ### Statistics sidecar Every `VCSCArray`/`VCSRArray` carries per-row and per-column summary statistics as `arr.stats` (a `vsparse.VCSStats`): `row_sums`/`col_sums` (total counts), `row_nnz`/`col_nnz` (n-genes/n-cells by counts), and `col_depth_sums` (per-gene sums of depth-normalized counts). They are computed on first access and written next to `X`/`raw_X` by `write_h5ad`/`write_zarr` whether or not they were ever computed or a normalized view was ever made, so a file read back has them without touching the data. Every recipe's per-cell and per-gene scales (`a`/`b`) are derived from them, so only centering/scaling recipes still need a pass over the data. Files without a usable sidecar (older versions, other tools) load as before, and the statistics are recomputed when first needed. diff --git a/benchmarks/cases.py b/benchmarks/cases.py index b5bb151..642269b 100644 --- a/benchmarks/cases.py +++ b/benchmarks/cases.py @@ -352,6 +352,62 @@ def large_layout_and_matmul() -> dict[str, float]: } +@slow +def subset_load_vs_full_load() -> dict[str, float]: + """Reading ``X`` subset as it's unpacked, against reading all of it, from a zstd IVCSR file. + + Random selections of 0.1/1/10% of cells, a "clustered" 1% (one contiguous + run, as selecting a single sample tends to be), and 1% of cells with 10% + of genes. Each ratio is that read's time over a full read of ``X``. Every + packed byte is still read either way, so reading the stream is the floor. + Whole ``load_packed`` calls also read ``obs``/``var``, which is recorded + separately since it costs the same either way (and, with a string index of + this many cells, more than ``X`` does). + """ + import tempfile + from pathlib import Path + + import anndata as ad + import h5py + import pandas as pd + + from vsparse import VCSCAnnData, VCSRArray, load_packed + from vsparse._io import read_vcs_subset + + n_obs, n_vars = 400_000, 3_000 + mat = integer_counts_csr(n_obs, n_vars, density=0.02) + adata = VCSCAnnData( + X=VCSRArray.from_scipy(mat), + obs=pd.DataFrame(index=np.arange(n_obs).astype(str)), + var=pd.DataFrame(index=np.arange(n_vars).astype(str)), + ) + del mat + rng = np.random.default_rng(0) + start = n_obs // 3 + selections = { + f"random_{frac:g}": (np.flatnonzero(rng.random(n_obs) < frac), None) + for frac in (0.001, 0.01, 0.1) + } + selections["clustered_0.01"] = (np.arange(start, start + n_obs // 100), None) + selections["cells_0.01_genes_0.1"] = ( + np.flatnonzero(rng.random(n_obs) < 0.01), + np.flatnonzero(rng.random(n_vars) < 0.1), + ) + out: dict[str, float] = {} + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "bench.ivcsr.h5ad" + adata.write_h5ad(path) + del adata + with h5py.File(path, "r") as f: + full = best_time(lambda: ad.io.read_elem(f["X"]), repeat=5) + out["full_x_seconds"] = full + for name, (rows, cols) in selections.items(): + t = best_time(lambda rows=rows, cols=cols: read_vcs_subset(f["X"], rows, cols), 5) + out[f"time_ratio_{name}"] = t / full + out["full_load_packed_seconds"] = best_time(lambda: load_packed(path), repeat=3) + return out + + # -- CUDA kernels vs materializing to a sparse matrix ------------------------ # # The question `vsparse._cuda` exists to answer: is walking the diff --git a/docs/usage.md b/docs/usage.md index 0f82ea7..5de94cf 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -127,6 +127,31 @@ va = vsparse.load_packed("archived.ivcsr.h5ad") norm = va.X.normalized("parafac2") ``` +### Subsetting while loading + +`load_packed(path, obs_idx=..., var_idx=...)` returns exactly +`load_packed(path)[obs_idx, var_idx]`, but subsets `X` as it decodes rather than +after. Each selection is a boolean mask over the axis, an array of integer positions +(negative positions count from the end; unsorted or repeated positions come back in +that order), or a slice; select by label by converting first, e.g. with +`obs.index.get_indexer(labels)`. + +For an IVCSR file only the selected cells' indices are decoded, and a gene selection +drops and renumbers indices in that same pass (IVCSC: the same, with cells and genes +swapped). The full matrix is never allocated. Every packed byte is still read from +disk, though: with no stored per-row byte offsets, finding where each selected cell +starts takes a parallel scan of the stream. The subset doesn't carry the file's +statistics sidecar (see below), which describes the whole matrix; its own `stats` are +computed on first use. + +```python +import anndata, h5py + +with h5py.File("archived.ivcsr.h5ad") as f: # read only obs to build the mask + mask = anndata.io.read_elem(f["obs"])["condition"] == "control" +controls = vsparse.load_packed("archived.ivcsr.h5ad", obs_idx=mask, var_idx=hvg_positions) +``` + ### Statistics sidecar Each `VCSCArray`/`VCSRArray` carries a {class}`vsparse.VCSStats` as `arr.stats`: diff --git a/src/vsparse/_anndata_class.py b/src/vsparse/_anndata_class.py index 69402e7..6dec78f 100644 --- a/src/vsparse/_anndata_class.py +++ b/src/vsparse/_anndata_class.py @@ -212,7 +212,18 @@ def __getitem__(self, index: Any) -> VCSCAnnData: # ty: ignore[invalid-method-o oidx, vidx = self._normalize_indices(index) oidx = _as_slice_index(oidx, self.n_obs) vidx = _as_slice_index(vidx, self.n_vars) + return type(self)( + X=_subset_2d(self._vcs_X, oidx, vidx), + raw_X=_subset_2d(self._vcs_raw_X, oidx, vidx), + **self._subset_fields(oidx, vidx), + ) + + def _subset_fields(self, oidx: Any, vidx: Any) -> dict[str, Any]: + """Every field but ``X``/``raw_X``, subset to positions ``oidx``/``vidx``. + Shared by :meth:`__getitem__` and :func:`~vsparse.load_packed`'s + ``obs_idx``/``var_idx``, which decodes ``X`` already subset. + """ obs = cast(pd.DataFrame, self.obs).iloc[oidx].copy() var = cast(pd.DataFrame, self.var).iloc[vidx].copy() @@ -229,18 +240,16 @@ def __getitem__(self, index: Any) -> VCSCAnnData: # ty: ignore[invalid-method-o if _VSPARSE_UNS_KEY in uns: uns = {**uns, _VSPARSE_UNS_KEY: {**uns[_VSPARSE_UNS_KEY], "stale": True}} - return type(self)( - X=_subset_2d(self._vcs_X, oidx, vidx), - raw_X=_subset_2d(self._vcs_raw_X, oidx, vidx), - obs=obs, - var=var, - uns=uns, - obsm=_map_subset_1d(self.obsm, oidx), - varm=_map_subset_1d(self.varm, vidx), - obsp=_map_subset_2d(self.obsp, oidx, oidx), - varp=_map_subset_2d(self.varp, vidx, vidx), - layers=_map_subset_2d(self.layers, oidx, vidx), - ) + return { + "obs": obs, + "var": var, + "uns": uns, + "obsm": _map_subset_1d(self.obsm, oidx), + "varm": _map_subset_1d(self.varm, vidx), + "obsp": _map_subset_2d(self.obsp, oidx, oidx), + "varp": _map_subset_2d(self.varp, vidx, vidx), + "layers": _map_subset_2d(self.layers, oidx, vidx), + } def copy(self) -> VCSCAnnData: # ty: ignore[invalid-method-override] """A deep copy, preserving the VCSC/VCSR-backed ``X``/``raw_X``. diff --git a/src/vsparse/_io.py b/src/vsparse/_io.py index a8db421..3ed3ca1 100644 --- a/src/vsparse/_io.py +++ b/src/vsparse/_io.py @@ -41,7 +41,7 @@ GroupStorageType = h5py.Group | zarr.Group -__all__ = ["write_ivcs_elem"] +__all__ = ["read_vcs_subset", "write_ivcs_elem"] _ARRAY_KEYS = ("major_ptr", "values", "value_ptr", "indices") _SPEC_VERSION = "0.1.0" @@ -145,6 +145,85 @@ def _read(elem: GroupStorageType, *, _reader: Reader) -> _VCSBase: return _read +def _sorted_unique(idx: np.ndarray) -> tuple[np.ndarray, np.ndarray | None]: + """``idx`` as sorted unique positions, plus how to recover ``idx`` from them (``None``: no need).""" + if idx.shape[0] < 2 or bool(np.all(idx[1:] > idx[:-1])): + return idx, None + uniq, inverse = np.unique(idx, return_inverse=True) + return uniq, inverse + + +def read_vcs_subset( + elem: GroupStorageType, rows: np.ndarray | None, cols: np.ndarray | None +) -> _VCSBase: + """Read the stored VCS array ``elem``, keeping only ``rows``/``cols`` (int positions, or ``None`` for all). + + Equal to reading the whole array and indexing it with ``[rows, cols]``. An + IVCSC/IVCSR array is subset as it's decoded (see + :func:`vsparse._ivcsc.unpack_selected`): every packed byte is still read, + but only the selected major slices are decoded, and a minor-axis selection + is applied in the same pass, so the full matrix is never built. A + selection that isn't sorted and unique is decoded as one that is, then + reordered in memory. + + A subset doesn't carry the stored statistics sidecar, which describes the + whole array; its own :attr:`~vsparse.VCSCArray.stats` are computed on first use. + """ + spec = cast(str, elem.attrs.get("encoding-type", "")) + if spec in _IVCS_SPEC_NAMES.values(): + return _read_ivcs_subset(elem, VCSCArray if spec == "ivcsc" else VCSRArray, rows, cols) + arr = cast(_VCSBase, ad.io.read_elem(elem)) + return arr[slice(None) if rows is None else rows, slice(None) if cols is None else cols] + + +def _major_selection(key: np.ndarray | None, n_major: int) -> tuple[np.ndarray, np.ndarray | None]: + """The sorted unique major slices to decode, and how to reorder them into ``key``.""" + if key is None: + return np.arange(n_major, dtype=np.int64), None + return _sorted_unique(key) + + +def _minor_remap( + key: np.ndarray | None, n_minor: int +) -> tuple[np.ndarray | None, np.ndarray | None, int]: + """Each minor index's new one (``-1``: dropped), how to reorder into ``key``, and the new length.""" + if key is None: + return None, None, n_minor + kept, inverse = _sorted_unique(key) + remap = np.full(n_minor, -1, dtype=np.int64) + remap[kept] = np.arange(kept.shape[0]) + return remap, inverse, kept.shape[0] + + +def _read_ivcs_subset( + elem: GroupStorageType, + cls: type[_VCSBase], + rows: np.ndarray | None, + cols: np.ndarray | None, +) -> _VCSBase: + n_major, n_minor = cls._swap(_read_shape(elem)) + major_key, minor_key = (cols, rows) if cls is VCSCArray else (rows, cols) + sel, major_inverse = _major_selection(major_key, n_major) + remap, minor_inverse, n_minor_new = _minor_remap(minor_key, n_minor) + + values = read_array(elem["values"]) + major_ptr, groups, value_ptr, indices = _ivcsc.unpack_selected( + read_array(elem["major_ptr"]), + read_array(elem["value_ptr"]), + read_array(elem["packed_indices"]), + sel, + remap, + n_minor_new, + ) + n_rows, n_cols = cls._swap((sel.shape[0], n_minor_new)) + arr = cls((n_rows, n_cols), major_ptr, values[groups], value_ptr, indices) + if major_inverse is not None: + arr = arr._select_major(major_inverse) + if minor_inverse is not None: + arr = arr._select_minor(minor_inverse) + return arr + + def _register() -> None: from anndata._io.specs.registry import _REGISTRY, IOSpec diff --git a/src/vsparse/_ivcsc.py b/src/vsparse/_ivcsc.py index a7b3a2e..028f944 100644 --- a/src/vsparse/_ivcsc.py +++ b/src/vsparse/_ivcsc.py @@ -42,7 +42,9 @@ import numba import numpy as np -__all__ = ["pack_indices", "unpack_indices"] +from vsparse._indexutils import smallest_index_dtype + +__all__ = ["pack_indices", "unpack_indices", "unpack_selected"] # Below this many packed bytes, decode serially -- the fixed overhead of the # counting pass plus a second (parallel) pass isn't worth it. @@ -50,6 +52,8 @@ # Lower bound on how small a per-thread chunk is allowed to get, so we don't # over-split modestly-sized arrays into slivers. _MIN_CHUNK_BYTES = 1 << 16 # 64 KiB +# Bytes :func:`_locate_ranks` counts at a time when skipping ahead. +_LOCATE_BLOCK = 256 @numba.njit(cache=True) @@ -140,14 +144,14 @@ def _realign_forward(buf: np.ndarray, guess: np.int64) -> np.int64: @numba.njit(cache=True) def _count_varints_range(buf: np.ndarray, start: np.int64, end: np.int64) -> np.int64: - """Count complete varints in buf[start:end); start/end must be true boundaries.""" - pos = start + """Count varints ending in buf[start:end) -- i.e. bytes with the continuation bit clear. + + Between true boundaries that's the number of complete varints. Counted + branch-free, one byte at a time, so the loop vectorizes. + """ count = np.int64(0) - while pos < end: - while buf[pos] & 0x80: - pos += 1 - pos += 1 - count += 1 + for pos in range(start, end): + count += 1 - (buf[pos] >> 7) return count @@ -189,20 +193,14 @@ def _decode_chunks( out[kk] = prev -def _group_chunk_boundaries( - value_ptr: np.ndarray, buf: np.ndarray, n_chunks: int -) -> tuple[np.ndarray, np.ndarray]: - """Split ``buf`` into ``n_chunks`` pieces landing on exact group boundaries. +def _varint_chunks(buf: np.ndarray, n_chunks: int) -> tuple[np.ndarray, np.ndarray]: + """Split ``buf`` into ``n_chunks`` spans starting on varint boundaries. - Returns ``(chunk_group, chunk_byte)``, each of length ``n_chunks + 1``: - chunk ``t`` owns groups ``[chunk_group[t], chunk_group[t + 1])``, whose - packed bytes start at ``chunk_byte[t]``. Used by :func:`_unpack_parallel` - to parallelize decoding safely. + Returns ``(real_starts, cum)``, each of length ``n_chunks + 1``: span ``t`` + covers bytes ``[real_starts[t], real_starts[t + 1])``, and its first varint + is varint number ``cum[t]`` of the whole stream. """ n_bytes = buf.shape[0] - n_groups = value_ptr.shape[0] - 1 - nnz = int(value_ptr[-1]) - guesses = np.linspace(0, n_bytes, n_chunks + 1).astype(np.int64) real_starts = np.empty(n_chunks + 1, dtype=np.int64) real_starts[0] = 0 @@ -215,7 +213,23 @@ def _group_chunk_boundaries( cum = np.empty(n_chunks + 1, dtype=np.int64) cum[0] = 0 cum[1:] = np.cumsum(counts) - assert cum[n_chunks] == nnz + return real_starts, cum + + +def _group_chunk_boundaries( + value_ptr: np.ndarray, buf: np.ndarray, n_chunks: int +) -> tuple[np.ndarray, np.ndarray]: + """Split ``buf`` into ``n_chunks`` pieces landing on exact group boundaries. + + Returns ``(chunk_group, chunk_byte)``, each of length ``n_chunks + 1``: + chunk ``t`` owns groups ``[chunk_group[t], chunk_group[t + 1])``, whose + packed bytes start at ``chunk_byte[t]``. Used by :func:`_unpack_parallel` + to parallelize decoding safely. + """ + n_bytes = buf.shape[0] + n_groups = value_ptr.shape[0] - 1 + real_starts, cum = _varint_chunks(buf, n_chunks) + assert cum[n_chunks] == int(value_ptr[-1]) chunk_group = np.empty(n_chunks + 1, dtype=np.int64) chunk_byte = np.empty(n_chunks + 1, dtype=np.int64) @@ -264,3 +278,178 @@ def unpack_indices(value_ptr: np.ndarray, packed: np.ndarray, dtype: np.dtype) - else: _unpack_parallel(value_ptr, packed, out, n_chunks) return out + + +# -- selective decode ----------------------------------------------------------- +# +# Decoding only some major slices (rows of an IVCSR array) still needs every +# byte of the stream in memory, since nothing records where a slice's bytes +# start. But it only has to *count* varints up to each selected slice, not +# decode them: :func:`_varint_chunks` splits the stream on varint boundaries +# with a parallel count, each span then walks forward to the selected slices +# that start inside it (:func:`_locate_ranks`), and only those slices are +# decoded (:func:`_count_kept`/:func:`_fill_selected`), in parallel. A minor-axis +# selection is applied during that same decode by renumbering each index +# through a lookup table, dropping the ones mapped to ``-1``. + + +@numba.njit(cache=True, parallel=True) +def _locate_ranks( + buf: np.ndarray, + real_starts: np.ndarray, + cum: np.ndarray, + targets: np.ndarray, + span_ptr: np.ndarray, + out: np.ndarray, +) -> None: + """Byte offset of varint number ``targets[i]``, for sorted ``targets``. + + Span ``t`` (from :func:`_varint_chunks`) owns targets + ``[span_ptr[t], span_ptr[t + 1])`` and walks forward to each in turn, so + the whole pass touches each byte at most once. + """ + for t in numba.prange(real_starts.shape[0] - 1): # ty: ignore[not-iterable] + pos = real_starts[t] + end = real_starts[t + 1] + rank = cum[t] + for i in range(span_ptr[t], span_ptr[t + 1]): + need = targets[i] - rank + # Skip whole blocks the target isn't in, counting with the + # vectorized kernel. That can stop partway through a varint; + # _advance_n_varints then finishes it as the first of `need`. + while need > 0 and pos + _LOCATE_BLOCK <= end: + ended = _count_varints_range(buf, pos, pos + _LOCATE_BLOCK) + if ended >= need: + break + need -= ended + pos += _LOCATE_BLOCK + pos = _advance_n_varints(buf, pos, need) + rank = targets[i] + out[i] = pos + + +@numba.njit(cache=True, parallel=True) +def _selected_groups(major_ptr, value_ptr, sel, group_offset, groups, sizes): + """Every stored group of the selected slices, in order, and how many indices each holds.""" + for i in numba.prange(sel.shape[0]): # ty: ignore[not-iterable] + out_g = group_offset[i] + for g in range(major_ptr[sel[i]], major_ptr[sel[i] + 1]): + groups[out_g] = g + sizes[out_g] = value_ptr[g + 1] - value_ptr[g] + out_g += 1 + + +@numba.njit(cache=True, parallel=True) +def _count_kept(major_ptr, value_ptr, buf, sel, starts, group_offset, remap, out_counts): + """How many indices of each selected group ``remap`` keeps (maps to ``>= 0``).""" + for i in numba.prange(sel.shape[0]): # ty: ignore[not-iterable] + pos = starts[i] + out_g = group_offset[i] + for g in range(major_ptr[sel[i]], major_ptr[sel[i] + 1]): + prev = np.int64(-1) + kept = 0 + for _ in range(value_ptr[g], value_ptr[g + 1]): + result, pos = _decode_varint(buf, pos) + prev = prev + 1 + np.int64(result) + if remap[prev] >= 0: + kept += 1 + out_counts[out_g] = kept + out_g += 1 + + +@numba.njit(cache=True, parallel=True) +def _fill_selected(major_ptr, value_ptr, buf, sel, starts, group_offset, remap, out_ptr, out): + """Decode the selected groups, writing each kept index's ``remap`` from ``out_ptr``.""" + for i in numba.prange(sel.shape[0]): # ty: ignore[not-iterable] + pos = starts[i] + out_g = group_offset[i] + for g in range(major_ptr[sel[i]], major_ptr[sel[i] + 1]): + prev = np.int64(-1) + k = out_ptr[out_g] + for _ in range(value_ptr[g], value_ptr[g + 1]): + result, pos = _decode_varint(buf, pos) + prev = prev + 1 + np.int64(result) + new = remap[prev] + if new >= 0: + out[k] = new + k += 1 + out_g += 1 + + +def _slice_starts( + major_ptr: np.ndarray, value_ptr: np.ndarray, buf: np.ndarray, sel: np.ndarray +) -> np.ndarray: + """Byte offset in ``buf`` where each selected major slice's indices start.""" + targets = value_ptr[major_ptr[sel]] + n_chunks = _num_chunks(buf.shape[0]) + real_starts, cum = _varint_chunks(buf, n_chunks) + # A target at the very end of the stream (an empty last slice) belongs to + # the last span, which walks it up to the final byte. + span_of = np.clip(np.searchsorted(cum, targets, side="right") - 1, 0, n_chunks - 1) + span_ptr = np.searchsorted(span_of, np.arange(n_chunks + 1)).astype(np.int64) + starts = np.empty(sel.shape[0], dtype=np.int64) + _locate_ranks(buf, real_starts, cum, targets, span_ptr, starts) + return starts + + +def unpack_selected( + major_ptr: np.ndarray, + value_ptr: np.ndarray, + packed: np.ndarray, + sel: np.ndarray, + remap: np.ndarray | None, + n_minor: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Decode only major slices ``sel``, and only the minor indices ``remap`` keeps. + + Parameters + ---------- + major_ptr, value_ptr, packed + The stored IVCS arrays. + sel + Sorted, unique major indices to decode. + remap + ``None`` to keep every minor index, or an ``int64`` array giving each + old minor index its new one, ``-1`` to drop it. + n_minor + Length of the minor axis afterwards. + + Returns + ------- + ``(major_ptr, groups, value_ptr, indices)`` for the selection, where + ``groups`` are the stored groups it keeps (so ``values[groups]`` are its + values). Groups left empty by ``remap`` are dropped, as indexing an + in-memory array would. + """ + major_ptr = np.ascontiguousarray(major_ptr, dtype=np.int64) + value_ptr = np.ascontiguousarray(value_ptr, dtype=np.int64) + buf = np.ascontiguousarray(packed, dtype=np.uint8) + sel = np.ascontiguousarray(sel, dtype=np.int64) + + # Every group of the selected slices, in order: group_offset is the + # selection's major_ptr before any group is emptied by `remap`. + group_offset = np.zeros(sel.shape[0] + 1, dtype=np.int64) + np.cumsum(major_ptr[sel + 1] - major_ptr[sel], out=group_offset[1:]) + groups = np.empty(int(group_offset[-1]), dtype=np.int64) + counts = np.empty(groups.shape[0], dtype=np.int64) + _selected_groups(major_ptr, value_ptr, sel, group_offset, groups, counts) + + starts = _slice_starts(major_ptr, value_ptr, buf, sel) + filtering = remap is not None + if remap is None: + remap = np.arange(n_minor, dtype=np.int64) + else: + _count_kept(major_ptr, value_ptr, buf, sel, starts, group_offset, remap, counts) + + out_ptr = np.zeros(groups.shape[0] + 1, dtype=np.int64) + np.cumsum(counts, out=out_ptr[1:]) + indices = np.empty(int(out_ptr[-1]), dtype=smallest_index_dtype(n_minor)) + _fill_selected(major_ptr, value_ptr, buf, sel, starts, group_offset, remap, out_ptr, indices) + if not filtering: + return group_offset, groups, out_ptr, indices + + nonempty = counts > 0 + nonempty_before = np.zeros(groups.shape[0] + 1, dtype=np.int64) + np.cumsum(nonempty, out=nonempty_before[1:]) + new_value_ptr = np.append(out_ptr[:-1][nonempty], out_ptr[-1]) + return nonempty_before[group_offset], groups[nonempty], new_value_ptr, indices diff --git a/src/vsparse/_rapid_load.py b/src/vsparse/_rapid_load.py index a1b4768..e2da7b7 100644 --- a/src/vsparse/_rapid_load.py +++ b/src/vsparse/_rapid_load.py @@ -2,18 +2,22 @@ :func:`load_packed` reads the packed ``X`` and decodes it immediately into an ordinary :class:`~vsparse.VCSCArray`/:class:`~vsparse.VCSRArray` inside a -:class:`~vsparse.VCSCAnnData`. Filtering and normalization are separate steps -(subsetting, then the ``.normalized()`` views), not part of loading. +:class:`~vsparse.VCSCAnnData`, optionally keeping only some cells and genes +as it decodes (``obs_idx``/``var_idx``). Normalization is a separate step (the +``.normalized()`` views), not part of loading. """ from __future__ import annotations from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import anndata as ad +import numpy as np +from vsparse import _io from vsparse._anndata_class import VCSCAnnData +from vsparse._indexutils import normalize_major_idx if TYPE_CHECKING: from os import PathLike @@ -23,12 +27,46 @@ _FIELD_KEYS = ("obs", "var", "obsm", "varm", "obsp", "varp", "layers", "uns") -def load_packed(path: str | PathLike[str], *, x_key: str = "X") -> VCSCAnnData: +def _positions(key: Any, n: int, name: str) -> np.ndarray | None: + """``key`` (a boolean mask, integer positions, or a slice) as int positions on an axis of ``n``.""" + if key is None: + return None + if isinstance(key, slice): + return np.arange(*key.indices(n), dtype=np.int64) + arr = np.asarray(key) + if arr.size == 0: + return np.empty(0, dtype=np.int64) + _check_selection(arr, n, name) + return normalize_major_idx(arr, n).astype(np.int64, copy=False) + + +def _check_selection(arr: np.ndarray, n: int, name: str) -> None: + if arr.ndim != 1 or not (arr.dtype == bool or np.issubdtype(arr.dtype, np.integer)): + raise TypeError( + f"{name} must be a 1-D boolean mask, an array of integer positions, " + f"or a slice; got {arr.dtype} with shape {arr.shape}" + ) + if arr.dtype == bool and arr.shape[0] != n: + raise IndexError(f"{name} is a boolean mask of length {arr.shape[0]}, expected {n}") + + +def load_packed( + path: str | PathLike[str], + *, + x_key: str = "X", + obs_idx: Any = None, + var_idx: Any = None, +) -> VCSCAnnData: """Load an IVCSR/IVCSC-backed ``.h5ad`` file, decoding ``X`` immediately. ``X`` comes back as an ordinary VCSCArray/VCSRArray with plain - ``indices``, decoded from the packed bytes read from disk. Subset and - normalize afterwards with the usual indexing and ``.normalized()`` views. + ``indices``, decoded from the packed bytes read from disk. Normalize + afterwards with the ``.normalized()`` views. + + ``load_packed(path, obs_idx=o, var_idx=v)`` equals ``load_packed(path)[o, v]``, + but decodes only the selected cells (for IVCSR; genes for IVCSC) and drops + unselected genes (cells) during that same decode, so the full matrix is + never built. Every packed byte is still read from disk. Parameters ---------- @@ -38,12 +76,33 @@ def load_packed(path: str | PathLike[str], *, x_key: str = "X") -> VCSCAnnData: :meth:`~vsparse.VCSCAnnData.write_h5ad`). x_key Top-level h5ad group holding the IVCSR/IVCSC array (``"X"`` by default). + obs_idx, var_idx + Cells/genes to keep: a boolean mask over the axis, an array of integer + positions (negative ones count from the end; unsorted or repeated + positions give that order), or a slice. ``None`` (the default) keeps + all. A mask built from metadata, e.g. ``obs["tissue"] == "blood"``, + works as is; to select by label, convert first with + ``obs.index.get_indexer(labels)``. """ import h5py import hdf5plugin # noqa: F401 -- registers the zstd/Blosc2 HDF5 filters with h5py.File(Path(path), "r") as f: - X = ad.io.read_elem(f[x_key]) kwargs = {k: ad.io.read_elem(f[k]) for k in _FIELD_KEYS if k in f} + if obs_idx is None and var_idx is None: + X = ad.io.read_elem(f[x_key]) + return VCSCAnnData(X=X, **kwargs) # ty: ignore[invalid-argument-type] + return _load_subset(f[x_key], kwargs, obs_idx, var_idx) + - return VCSCAnnData(X=X, **kwargs) # ty: ignore[invalid-argument-type] +def _load_subset(x_elem: Any, kwargs: dict[str, Any], obs_idx: Any, var_idx: Any) -> VCSCAnnData: + """:func:`load_packed` with a selection: ``X`` subset as it's decoded, the rest after reading.""" + n_obs, n_vars = _io._read_shape(x_elem) + rows = _positions(obs_idx, n_obs, "obs_idx") + cols = _positions(var_idx, n_vars, "var_idx") + X = _io.read_vcs_subset(x_elem, rows, cols) + meta = VCSCAnnData(**kwargs) + fields = meta._subset_fields( + slice(None) if rows is None else rows, slice(None) if cols is None else cols + ) + return VCSCAnnData(X=X, **fields) diff --git a/tests/test_load_subset.py b/tests/test_load_subset.py new file mode 100644 index 0000000..0e4e9eb --- /dev/null +++ b/tests/test_load_subset.py @@ -0,0 +1,158 @@ +"""``load_packed(obs_idx=..., var_idx=...)``: subsetting while unpacking IVCSR/IVCSC.""" + +from __future__ import annotations + +import warnings +from typing import Any + +import numpy as np +import pandas as pd +import pytest +import scipy.sparse as sp + +from vsparse import VCSCAnnData, VCSCArray, VCSRArray, _ivcsc, load_packed + +N_OBS, N_VARS = 120, 3000 # wide, so many index deltas need multi-byte varints + + +def _counts(rng: np.random.Generator) -> sp.csr_array: + mat = sp.random_array((N_OBS, N_VARS), density=0.02, format="lil", random_state=0) + mat[:7] = 0 # empty rows, at both ends and in the middle + mat[-5:] = 0 + mat[50:53] = 0 + mat[:, 100:300] = 0 # and empty columns + out = sp.csr_array(mat) + out.data = rng.integers(1, 40, size=out.data.shape[0]).astype(np.float64) + return out + + +@pytest.fixture(params=[VCSRArray, VCSCArray], ids=["ivcsr", "ivcsc"]) +def packed_path(request, tmp_path): + rng = np.random.default_rng(0) + adata = VCSCAnnData( + X=request.param.from_scipy(_counts(rng)), + obs=pd.DataFrame( + {"group": rng.integers(0, 4, N_OBS)}, index=[f"c{i}" for i in range(N_OBS)] + ), + var=pd.DataFrame(index=[f"g{i}" for i in range(N_VARS)]), + obsm={"pca": rng.normal(size=(N_OBS, 3))}, + varm={"loadings": rng.normal(size=(N_VARS, 2))}, + obsp={"dist": sp.random_array((N_OBS, N_OBS), density=0.1, format="csr")}, + uns={"note": "kept"}, + ) + path = tmp_path / "packed.h5ad" + adata.write_h5ad(path) + return path + + +@pytest.fixture(params=[False, True], ids=["serial", "many_chunks"]) +def chunking(request, monkeypatch): + """Exercise both a single-span scan and many spans (normally only past 1 MiB).""" + if request.param: + monkeypatch.setattr(_ivcsc, "_num_chunks", lambda n_bytes: max(1, min(13, n_bytes))) + + +def _dense(mat: Any) -> np.ndarray: + return mat.toarray() + + +def _assert_same_x(left, right) -> None: + assert left.shape == right.shape + assert type(left) is type(right) + for name in ("major_ptr", "values", "value_ptr", "indices"): + np.testing.assert_array_equal(getattr(left, name), getattr(right, name), err_msg=name) + + +def _assert_matches_indexing(path, obs_idx, var_idx) -> None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) # repeated positions -> duplicate names + subset = load_packed(path, obs_idx=obs_idx, var_idx=var_idx) + expected = load_packed(path)[ + slice(None) if obs_idx is None else obs_idx, + slice(None) if var_idx is None else var_idx, + ] + _assert_same_x(subset.X, expected.X) + pd.testing.assert_frame_equal(subset.obs, expected.obs) + pd.testing.assert_frame_equal(subset.var, expected.var) + np.testing.assert_array_equal(subset.obsm["pca"], expected.obsm["pca"]) + np.testing.assert_array_equal(subset.varm["loadings"], expected.varm["loadings"]) + np.testing.assert_array_equal(_dense(subset.obsp["dist"]), _dense(expected.obsp["dist"])) + assert subset.uns["note"] == "kept" + + +_RNG = np.random.default_rng(1) +SELECTIONS = { + "obs_mask": (_RNG.random(N_OBS) < 0.2, None), + "var_mask": (None, _RNG.random(N_VARS) < 0.3), + "both_masks": (_RNG.random(N_OBS) < 0.5, _RNG.random(N_VARS) < 0.1), + "sorted_positions": (np.array([0, 6, 7, 40, N_OBS - 1]), np.arange(0, N_VARS, 7)), + "unsorted_repeated_negative": (np.array([9, 3, 3, -1, 60]), np.array([2999, 5, 5, -3, 0])), + "slices": (slice(5, 100, 3), slice(None, None, -2)), + "empty_rows_only": (np.array([0, 1, 50, 51, N_OBS - 1]), None), + "nothing": (np.array([], dtype=np.int64), np.zeros(N_VARS, dtype=bool)), + "everything": (np.ones(N_OBS, dtype=bool), np.arange(N_VARS)), + "list_of_ints": ([4, 8, 15, 16, 23, 42], None), +} + + +@pytest.mark.parametrize("selection", SELECTIONS.values(), ids=SELECTIONS.keys()) +def test_subset_load_matches_indexing_a_full_load(packed_path, chunking, selection): + _assert_matches_indexing(packed_path, *selection) + + +def test_subset_recomputes_rather_than_inherits_the_stats_sidecar(packed_path): + mask = np.zeros(N_OBS, dtype=bool) + mask[::3] = True + subset: Any = load_packed(packed_path, obs_idx=mask).X + full: Any = load_packed(packed_path).X + assert subset._stats is None + expected = full[mask] + np.testing.assert_array_equal(subset.stats.row_sums, expected.stats.row_sums) + np.testing.assert_array_equal(subset.stats.col_depth_sums, expected.stats.col_depth_sums) + + +def test_metadata_masks_work_directly(packed_path): + full = load_packed(packed_path) + mask = full.obs["group"] == 2 # a pandas Series + _assert_matches_indexing(packed_path, mask, None) + + +@pytest.mark.parametrize( + ("kwargs", "error"), + [ + ({"obs_idx": np.array(["c1", "c2"])}, TypeError), + ({"var_idx": np.array([0.5, 1.0])}, TypeError), + ({"obs_idx": np.ones((2, 2), dtype=bool)}, TypeError), + ({"obs_idx": np.ones(N_OBS - 1, dtype=bool)}, IndexError), + ({"var_idx": np.array([N_VARS])}, IndexError), + ({"obs_idx": np.array([-N_OBS - 1])}, IndexError), + ], +) +def test_invalid_selections_are_rejected(packed_path, kwargs, error): + with pytest.raises(error): + load_packed(packed_path, **kwargs) + + +@pytest.mark.parametrize("seed", range(5)) +def test_unpack_selected_matches_a_full_decode(chunking, seed): + rng = np.random.default_rng(seed) + arr = VCSRArray.from_scipy(_counts(rng)) + packed = _ivcsc.pack_indices(arr.value_ptr, arr.indices) + full = VCSRArray( + arr.shape, + arr.major_ptr, + arr.values, + arr.value_ptr, + _ivcsc.unpack_indices(arr.value_ptr, packed, arr.indices.dtype), + ) + sel = np.flatnonzero(rng.random(N_OBS) < 0.3) + cols = np.flatnonzero(rng.random(N_VARS) < 0.4) + remap = np.full(N_VARS, -1, dtype=np.int64) + remap[cols] = np.arange(cols.shape[0]) + + major_ptr, groups, value_ptr, indices = _ivcsc.unpack_selected( + arr.major_ptr, arr.value_ptr, packed, sel, remap, cols.shape[0] + ) + values = arr.values[groups] # noqa: PD011 -- a VCSRArray, not pandas + got = VCSRArray((sel.shape[0], cols.shape[0]), major_ptr, values, value_ptr, indices) + _assert_same_x(got, full[sel][:, cols])