diff --git a/README.md b/README.md index 898fb39..13a01cf 100644 --- a/README.md +++ b/README.md @@ -143,6 +143,15 @@ va = vsparse.load_packed("archived.h5ad") norm = va.X.normalized("parafac2") # lazy view; norm @ B works without densifying ``` +### 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. + +```python +va = vsparse.VCSCAnnData.read_h5ad("archived.h5ad") +keep = va.X.stats.row_sums >= 500 # QC filter from the sidecar alone +``` + ### CUDA (`to_gpu`) A normalized view moves onto an NVIDIA GPU with `.to_gpu()`, keeping its value-compressed layout there rather than expanding to one float per nonzero — so a matrix that only fits in device memory compressed still fits. Everything floating-point on the device is float32. diff --git a/docs/api.md b/docs/api.md index d9a9d84..f6e7e77 100644 --- a/docs/api.md +++ b/docs/api.md @@ -13,6 +13,9 @@ :show-inheritance: :inherited-members: +.. autoclass:: vsparse.VCSStats + :members: + .. autofunction:: vsparse.from_anndata .. autofunction:: vsparse.to_layer diff --git a/docs/usage.md b/docs/usage.md index affbf92..0f82ea7 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -127,6 +127,23 @@ va = vsparse.load_packed("archived.ivcsr.h5ad") norm = va.X.normalized("parafac2") ``` +### Statistics sidecar + +Each `VCSCArray`/`VCSRArray` carries a {class}`vsparse.VCSStats` as `arr.stats`: +per-row `row_sums`/`row_nnz` (total counts, n-genes-by-counts), per-column +`col_sums`/`col_nnz` (total counts, n-cells-by-counts), and `col_depth_sums`, the +per-column sum of depth-normalized values `sum_i x[i, j] / row_sums[i]`. They are +computed on first access (one pass over the data) and kept. + +`write_h5ad`/`write_zarr` always store them, as a versioned `stats` subgroup of +`X`/`raw_X`, and readers (including `load_packed`) attach them, so QC filters on +these columns need no read of `X`. Normalization takes every recipe's +per-cell and per-gene scales (`a`/`b` in `y = (g(x * a * b) - c) * s`) from the +sidecar; the centering/scaling terms `c`/`s` depend on the nonlinear `g`, so they +still take one pass per recipe. A file without a usable sidecar -- written by an +earlier version, or with an unrecognized sidecar version -- reads as before, and the +statistics are recomputed when first needed. + ## CUDA A normalized view can be moved onto an NVIDIA GPU with `.to_gpu()`, which needs diff --git a/src/vsparse/__init__.py b/src/vsparse/__init__.py index 37c10a2..b3f2b49 100644 --- a/src/vsparse/__init__.py +++ b/src/vsparse/__init__.py @@ -12,6 +12,7 @@ from vsparse._cuda import CudaNormalizedView, cuda_is_available from vsparse._norm_common import RECIPES, Recipe from vsparse._rapid_load import load_packed +from vsparse._stats import VCSStats from vsparse._vcs_norm import VCSCArrayNormalized, VCSRArrayNormalized __all__ = [ @@ -23,6 +24,7 @@ "VCSCArrayNormalized", "VCSRArray", "VCSRArrayNormalized", + "VCSStats", "cuda_is_available", "from_anndata", "load_packed", diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py index 6a116e1..b8d2a31 100644 --- a/src/vsparse/_base.py +++ b/src/vsparse/_base.py @@ -17,7 +17,7 @@ import numpy as np import scipy.sparse as sp -from vsparse import _construct, _ops +from vsparse import _construct, _ops, _stats from vsparse._indexutils import is_full_slice as _is_full_slice from vsparse._indexutils import normalize_major_idx as _normalize_major_idx from vsparse._indexutils import smallest_index_dtype as _smallest_index_dtype @@ -42,7 +42,7 @@ class _VCSBase: # __rmatmul__/__rmul__ instead of trying to broadcast us as an ndarray. __array_ufunc__ = None - __slots__ = ("_norm_cache", "indices", "major_ptr", "shape", "value_ptr", "values") + __slots__ = ("_norm_cache", "_stats", "indices", "major_ptr", "shape", "value_ptr", "values") def __init__( self, @@ -81,6 +81,7 @@ def __init__( self.value_ptr = value_ptr self.indices = indices self._norm_cache = _NormCache() + self._stats: _stats.VCSStats | None = None # -- axis bookkeeping ------------------------------------------------ @@ -117,14 +118,29 @@ def __repr__(self) -> str: # pragma: no cover - cosmetic f"<{cls} shape={self.shape} dtype={self.dtype} nnz={self.nnz} n_unique={self.n_unique}>" ) + @property + def stats(self) -> _stats.VCSStats: + """Per-row/per-column totals and nonzero counts -- see :mod:`vsparse._stats`. + + Computed on first access (``O(nnz)``) and kept, or restored from the + sidecar of a file written by :meth:`~vsparse.VCSCAnnData.write_h5ad`/ + :meth:`~vsparse.VCSCAnnData.write_zarr`. Normalization reads its + per-cell/per-gene scales (``a``/``b``) from here. + """ + if self._stats is None: + self._stats = _stats.compute_stats(self) + return self._stats + def copy(self): - return type(self)( + out = type(self)( self.shape, self.major_ptr.copy(), self.values.copy(), self.value_ptr.copy(), self.indices.copy(), ) + out._stats = self._stats # same data, so the same (read-only) statistics + return out # -- construction / conversion --------------------------------------- diff --git a/src/vsparse/_construct.py b/src/vsparse/_construct.py index 3ca22bc..6e4d8f7 100644 --- a/src/vsparse/_construct.py +++ b/src/vsparse/_construct.py @@ -53,7 +53,9 @@ def _compress(major_ptr_in, minor_indices, data, n_major): seg_vals = data[start:end] seg_minor = minor_indices[start:end] - order = np.argsort(seg_vals) + # Stable, so each value group keeps its minor indices in input order + # (sorted, for canonical scipy input) -- see vsparse._stats. + order = np.argsort(seg_vals, kind="mergesort") sorted_vals = seg_vals[order] sorted_minor = seg_minor[order] diff --git a/src/vsparse/_io.py b/src/vsparse/_io.py index e8f5fc9..a8db421 100644 --- a/src/vsparse/_io.py +++ b/src/vsparse/_io.py @@ -6,6 +6,12 @@ :class:`~vsparse.VCSCAnnData` -- and survive a round trip through HDF5 or zarr without any manual (de)serialization. +Both formats store the array's statistics sidecar (see :mod:`vsparse._stats`) +as a ``stats`` subgroup next to the array data. Readers attach it when it's +present and valid, and otherwise leave the array to recompute it on demand, so +files without one -- written by earlier versions, or by other tools -- still +read normally. + This relies on ``anndata._io.specs.registry``, an internal (not yet public) extension mechanism -- the same one anndata itself uses to support dask/awkward/cupy arrays. It is not guaranteed stable across anndata @@ -23,7 +29,7 @@ import numpy as np import zarr -from vsparse import _ivcsc +from vsparse import _ivcsc, _stats from vsparse._base import VCSCArray, VCSRArray, _VCSBase from vsparse._compression import read_array from vsparse._indexutils import smallest_index_dtype @@ -55,20 +61,31 @@ def _write_vcs( g.attrs["shape"] = v.shape for name in _ARRAY_KEYS: _writer.write_elem(g, name, getattr(v, name), dataset_kwargs=dataset_kwargs) + _stats.write_stats(g, v.stats, write_elem=_writer.write_elem, dataset_kwargs=dataset_kwargs) + + +def _read_shape(elem: GroupStorageType) -> tuple[int, int]: + shape_vals = [int(s) for s in np.asarray(elem.attrs["shape"]).tolist()] + return (shape_vals[0], shape_vals[1]) + + +def _with_stats(arr: _VCSBase, elem: GroupStorageType) -> _VCSBase: + """Attach the sidecar stored under ``elem``, if there's a usable one -- see :mod:`vsparse._stats`.""" + arr._stats = _stats.read_stats(elem, arr.shape) + return arr def _make_read_vcs(cls: type[_VCSBase]): def _read(elem: GroupStorageType, *, _reader: Reader) -> _VCSBase: - shape_vals = [int(s) for s in np.asarray(elem.attrs["shape"]).tolist()] - shape = (shape_vals[0], shape_vals[1]) arrays = {name: read_array(elem[name]) for name in _ARRAY_KEYS} - return cls( - shape, + arr = cls( + _read_shape(elem), arrays["major_ptr"], arrays["values"], arrays["value_ptr"], arrays["indices"], ) + return _with_stats(arr, elem) return _read @@ -88,6 +105,9 @@ def write_ivcs_elem( the registry, e.g. via ``anndata.io.read_elem``) always reconstructs a normal VCSCArray/VCSRArray with a plain ``indices`` array. + ``v``'s statistics sidecar (:attr:`~vsparse.VCSCArray.stats`) is written + alongside, computing it first if ``v`` doesn't carry it yet. + Unlike the VCSC/VCSR codec registered below, this is not reachable through anndata's generic ``write_elem`` type dispatch (which can only route by the Python type of ``v``, and both formats share the same @@ -106,21 +126,21 @@ def write_ivcs_elem( ad.io.write_elem(g, name, getattr(v, name), dataset_kwargs=dataset_kwargs) packed = _ivcsc.pack_indices(v.value_ptr, v.indices) ad.io.write_elem(g, "packed_indices", packed, dataset_kwargs=dataset_kwargs) + _stats.write_stats(g, v.stats, write_elem=ad.io.write_elem, dataset_kwargs=dataset_kwargs) g.attrs["encoding-type"] = "ivcsc" if v._format == "csc" else "ivcsr" g.attrs["encoding-version"] = _IVCS_SPEC_VERSION def _make_read_ivcs(cls: type[_VCSBase]): def _read(elem: GroupStorageType, *, _reader: Reader) -> _VCSBase: - shape_vals = [int(s) for s in np.asarray(elem.attrs["shape"]).tolist()] - shape = (shape_vals[0], shape_vals[1]) major_ptr = read_array(elem["major_ptr"]) values = read_array(elem["values"]) value_ptr = read_array(elem["value_ptr"]) packed = read_array(elem["packed_indices"]) dtype = np.dtype(cast(str, elem.attrs["indices_dtype"])) indices = _ivcsc.unpack_indices(value_ptr, packed, dtype) - return cls(shape, major_ptr, values, value_ptr, indices) + arr = cls(_read_shape(elem), major_ptr, values, value_ptr, indices) + return _with_stats(arr, elem) return _read diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index 85ce485..d169143 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -29,19 +29,18 @@ Statistics (``a``, ``b``, ``c``, ``s``) are computed once, at construction, over the whole wrapped array, and cached on the view -- see :meth:`~NormalizedViewBase.select`/:meth:`_VCSBase.normalized` for how a -subset or an alternate recipe is (re)computed. Computing them still requires -touching every nonzero (once for ``b``, once more for ``c``/``s``, since -those need ``b`` first), done with parallel numba kernels below, specialized +subset or an alternate recipe is (re)computed. ``a`` and ``b`` come from the +array's recipe-independent :attr:`~vsparse.VCSCArray.stats` sidecar (see +:mod:`vsparse._stats`) without touching the data. ``c``/``s`` still require a +pass over every nonzero, done with parallel numba kernels below, specialized per storage format: -- major=columns (VCSC): both passes collapse into one fused kernel, fully - parallel over columns with no cross-thread writes -- each column's own - elements carry everything needed to compute both its ``b`` and its +- major=columns (VCSC): fully parallel over columns with no cross-thread + writes -- each column's own elements carry everything needed for its ``c``/``s`` (:func:`_column_stats_major_is_col`). -- major=rows (VCSR): each pass is a scatter-add across columns from - many rows, so it's parallelized row-chunked with thread-local partial - column arrays, reduced by summing across threads - (:func:`_scaled_col_sums_vcs`, :func:`_gstats_col_sums_vcs`). +- major=rows (VCSR): a scatter-add across columns from many rows, so it's + parallelized row-chunked with thread-local partial column arrays, reduced + by summing across threads (:func:`_gstats_col_sums_vcs`). These kernels only need ``major_ptr``/``values``/``value_ptr``/``indices`` arrays, from the plain (:mod:`vsparse._base`) array types. The matmul kernels @@ -65,6 +64,8 @@ import numpy as np import numpy.typing as npt +from vsparse._stats import VCSStats + __all__ = [ "DEFAULT_RECIPE", "NORM_CACHE_MAXSIZE", @@ -238,30 +239,22 @@ def _g_np(x: np.ndarray, g_code: int) -> np.ndarray: @numba.njit(cache=True, parallel=True) def _column_stats_major_is_col( - major_ptr, values, value_ptr, indices, row_scale, need_b, need_gstats, g_code, n_rows + major_ptr, values, value_ptr, indices, row_scale, gene_scale, g_code, n_rows ): - """Per-column ``gsum`` (raw material for ``b``) and the ``g(scaled)`` variance inputs. + """Per-column ``g(scaled)`` variance inputs. Sums cover the stored entries only; :func:`_finish_variance` folds in the implicit zeros. A VCSC column's statistics depend only on its own - nonzeros, so all three loops fuse into one pass per column. + nonzeros, so both loops run within one parallel pass per column. """ n_major = major_ptr.shape[0] - 1 - gsum = np.ones(n_major, dtype=np.float64) col_sum = np.zeros(n_major, dtype=np.float64) col_m2 = np.zeros(n_major, dtype=np.float64) col_corr = np.zeros(n_major, dtype=np.float64) col_nnz = np.zeros(n_major, dtype=np.float64) for j in numba.prange(n_major): # ty: ignore[not-iterable] - gs = 1.0 - if need_b: - gs = 0.0 - for u in range(major_ptr[j], major_ptr[j + 1]): - v = values[u] - for k in range(value_ptr[u], value_ptr[u + 1]): - gs += v / row_scale[indices[k]] - gsum[j] = gs - if need_gstats and gs > 0.0: + gs = gene_scale[j] + if gs > 0.0: s0 = 0.0 count = 0.0 for u in range(major_ptr[j], major_ptr[j + 1]): @@ -286,30 +279,12 @@ def _column_stats_major_is_col( corr += d col_m2[j] = m2 col_corr[j] = corr - return gsum, col_sum, col_m2, col_corr, col_nnz + return col_sum, col_m2, col_corr, col_nnz # -- statistics: major=rows -- scatter-add passes ---------------------------- -@numba.njit(cache=True, parallel=True) -def _scaled_col_sums_vcs(major_ptr, values, value_ptr, indices, row_scale, n_cols, nthreads): - n_major = major_ptr.shape[0] - 1 - chunk = (n_major + nthreads - 1) // nthreads - partial = np.zeros((nthreads, n_cols), dtype=np.float64) - for t in numba.prange(nthreads): # ty: ignore[not-iterable] - start = t * chunk - end = min(n_major, start + chunk) - local = partial[t] - for i in range(start, end): - rs = row_scale[i] - for u in range(major_ptr[i], major_ptr[i + 1]): - v = values[u] - for k in range(value_ptr[u], value_ptr[u + 1]): - local[indices[k]] += v / rs - return partial.sum(axis=0) - - @numba.njit(cache=True, parallel=True) def _gstats_col_sums_vcs( major_ptr, values, value_ptr, indices, row_scale, gene_scale, g_code, n_cols, nthreads @@ -655,21 +630,33 @@ def _prep_key(key: Any) -> Any: return key -def _compute_row_scale(arr: Any, recipe: Recipe) -> np.ndarray: - n_rows, _n_cols = arr.shape +def _depth_target(stats: VCSStats, recipe: Recipe) -> float | None: + """The depth every row is scaled to, or ``None`` for no depth normalization (``a = 1``).""" if recipe.depth_target is None: - return np.ones(n_rows, dtype=np.float64) - - row_totals = np.asarray(arr.sum(axis=1), dtype=np.float64) + return None if recipe.depth_target == "median": - target = float(np.median(row_totals)) if row_totals.shape[0] else 0.0 + row_sums = stats.row_sums + target = float(np.median(row_sums)) if row_sums.shape[0] else 0.0 else: target = float(recipe.depth_target) - if target <= 0.0: - return np.ones(n_rows, dtype=np.float64) - row_scale = row_totals / target - row_scale[row_scale == 0.0] = 1.0 # rows with no counts: avoid div-by-zero (unused otherwise) - return row_scale + return target if target > 0.0 else None + + +def _row_gene_scales(stats: VCSStats, recipe: Recipe) -> tuple[np.ndarray, np.ndarray]: + """``row_scale``/``gene_scale`` (``1 / a``/``1 / b``), from the array's stats sidecar alone.""" + n_rows, n_cols = stats.shape + target = _depth_target(stats, recipe) + if target is None: + row_scale = np.ones(n_rows, dtype=np.float64) + col_scaled_sums = stats.col_sums + else: + row_scale = stats.row_sums / target + # Rows with no counts: avoid div-by-zero (unused otherwise). + row_scale[row_scale == 0.0] = 1.0 + col_scaled_sums = target * stats.col_depth_sums + if not recipe.gene_scale: + return row_scale, np.ones(n_cols, dtype=np.float64) + return row_scale, np.array(col_scaled_sums, dtype=np.float64) def _is_constant_column(variance: np.ndarray, mean: np.ndarray, n_rows: int) -> np.ndarray: @@ -683,58 +670,36 @@ def _is_constant_column(variance: np.ndarray, mean: np.ndarray, n_rows: int) -> return variance <= n_rows * eps * variance + (n_rows * mean * eps) ** 2 -def _column_gstats(arr, row_scale, need_b, need_gstats, g_code, n_rows, n_cols): - """``gene_scale`` plus the stored-only sum/m2/corr/nnz feeding :func:`_finish_variance`. +def _column_gstats(arr, row_scale, gene_scale, g_code, n_rows, n_cols): + """The stored-only sum/m2/corr/nnz of ``g(scaled)`` feeding :func:`_finish_variance`. - Dispatches on ``arr._format``: VCSC fuses both statistics into one pass - per column (see :func:`_column_stats_major_is_col`); VCSR needs two - genuinely separate passes, since a column's ``gene_scale`` isn't final - until every row has scattered into it. + Dispatches on ``arr._format``: VCSC runs one pass per column (see + :func:`_column_stats_major_is_col`); VCSR needs two separate scatter + passes, since a column's mean isn't final until every row has scattered + into it. """ - indices = arr.indices # decode once; shared by both statistics passes below - zeros = np.zeros(n_cols, dtype=np.float64) - ones = np.ones(n_cols, dtype=np.float64) - + kernel_args = (arr.major_ptr, arr.values, arr.value_ptr, arr.indices, row_scale, gene_scale) if arr._format == "csc": - if not (need_b or need_gstats): - return ones, zeros, zeros, zeros, zeros - return _column_stats_major_is_col( - arr.major_ptr, - arr.values, - arr.value_ptr, - indices, - row_scale, - need_b, - need_gstats, - g_code, - n_rows, - ) - - gene_scale = ones - if need_b: - nthreads = numba.get_num_threads() - gene_scale = _scaled_col_sums_vcs( - arr.major_ptr, arr.values, arr.value_ptr, indices, row_scale, n_cols, nthreads - ) - if not need_gstats: - return gene_scale, zeros, zeros, zeros, zeros + return _column_stats_major_is_col(*kernel_args, g_code, n_rows) nthreads = numba.get_num_threads() - kernel_args = (arr.major_ptr, arr.values, arr.value_ptr, indices, row_scale, gene_scale, g_code) - col_sum, col_nnz = _gstats_col_sums_vcs(*kernel_args, n_cols, nthreads) - col_m2, col_corr = _gstats_col_deviations_vcs( - *kernel_args, col_sum / n_rows if n_rows > 0 else zeros, n_cols, nthreads - ) - return gene_scale, col_sum, col_m2, col_corr, col_nnz + col_sum, col_nnz = _gstats_col_sums_vcs(*kernel_args, g_code, n_cols, nthreads) + col_mean = col_sum / n_rows if n_rows > 0 else np.zeros(n_cols, dtype=np.float64) + col_m2, col_corr = _gstats_col_deviations_vcs(*kernel_args, g_code, col_mean, n_cols, nthreads) + return col_sum, col_m2, col_corr, col_nnz -def _column_mean_scale(recipe, col_sum, col_m2, col_corr, col_nnz, n_rows, n_cols, *, need_gstats): +def _column_mean_scale(arr, recipe, row_scale, gene_scale): """Per-column ``col_mean``/``col_post_scale``, or the recipe's no-op defaults.""" + n_rows, n_cols = arr.shape zeros = np.zeros(n_cols, dtype=np.float64) ones = np.ones(n_cols, dtype=np.float64) - if not need_gstats: + if not (recipe.center or recipe.post_scale): return zeros, ones + col_sum, col_m2, col_corr, col_nnz = _column_gstats( + arr, row_scale, gene_scale, recipe.g_code, n_rows, n_cols + ) mean, variance = _finish_variance(col_sum, col_m2, col_corr, col_nnz, n_rows) col_mean = mean if recipe.center else zeros if not recipe.post_scale: @@ -785,18 +750,9 @@ def __init__( self.recipe = resolve_recipe(recipe) self.stale = stale - n_rows, n_cols = arr.shape - row_scale = _compute_row_scale(arr, self.recipe) - self.row_scale = row_scale - - need_b = self.recipe.gene_scale - need_gstats = self.recipe.center or self.recipe.post_scale - gene_scale, col_sum, col_m2, col_corr, col_nnz = _column_gstats( - arr, row_scale, need_b, need_gstats, self.recipe.g_code, n_rows, n_cols - ) - self.gene_scale = gene_scale + self.row_scale, self.gene_scale = _row_gene_scales(arr.stats, self.recipe) self.col_mean, self.col_post_scale = _column_mean_scale( - self.recipe, col_sum, col_m2, col_corr, col_nnz, n_rows, n_cols, need_gstats=need_gstats + arr, self.recipe, self.row_scale, self.gene_scale ) @classmethod diff --git a/src/vsparse/_stats.py b/src/vsparse/_stats.py new file mode 100644 index 0000000..a6d3f04 --- /dev/null +++ b/src/vsparse/_stats.py @@ -0,0 +1,212 @@ +"""Per-row/per-column summary statistics carried by a VCS array -- the "stats sidecar". + +Every :class:`~vsparse.VCSCArray`/:class:`~vsparse.VCSRArray` can carry a +:class:`VCSStats` (see :attr:`~vsparse.VCSCArray.stats`): the per-cell and +per-gene totals QC filtering and normalization keep needing, held at +``O(n_rows + n_cols)`` instead of being re-derived from the ``O(nnz)`` data +each time (https://github.com/meyer-lab/vsparse/issues/26). + +The fields are chosen so that the per-cell ``a`` and per-gene ``b`` of every +normalization recipe (``y = (g(x * a * b) - c) * s``, see +:mod:`vsparse._norm_common`) follow from them without touching the data: + +- ``a = target / row_sums`` (or ``1``), where ``target`` is a constant or + ``median(row_sums)``; +- ``b = 1 / sum_i(x[i, j] * a[i])``, which is ``1 / col_sums`` when ``a = 1`` + and ``1 / (target * col_depth_sums)`` otherwise, since ``target`` factors + out of the sum. + +``c``/``s`` are moments of ``g(x * a * b)``, a nonlinear function of the +recipe, so there's no recipe-independent summary of them; they still take +one pass over the data per recipe. + +The statistics are computed lazily, on first access, and persisted alongside +the array by :meth:`~vsparse.VCSCAnnData.write_h5ad`/:meth:`~vsparse.VCSCAnnData.write_zarr` +as a ``stats`` subgroup of the array's own group (:func:`write_stats`). On read +(:func:`read_stats`) a missing, unrecognized or malformed sidecar is ignored +rather than trusted: the array then simply recomputes its statistics on first +use, so files written before the sidecar existed load exactly as before. +""" + +from __future__ import annotations + +import warnings +from dataclasses import dataclass, fields +from types import MappingProxyType +from typing import TYPE_CHECKING, Any + +import numba +import numpy as np + +from vsparse import _ops +from vsparse._compression import read_array + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping + +__all__ = ["VCSStats", "compute_stats", "read_stats", "write_stats"] + +STATS_KEY = "stats" +STATS_ENCODING_TYPE = "vsparse-stats" +STATS_ENCODING_VERSION = "0.1.0" + + +@dataclass(frozen=True, slots=True, eq=False) +class VCSStats: + """Per-row/per-column statistics of one VCS array -- see :mod:`vsparse._stats`. + + The arrays are read-only, since a :class:`VCSStats` is shared between an + array and its copies. + """ + + #: Per-row sum of stored values (scanpy's ``obs["total_counts"]``). + row_sums: np.ndarray + #: Per-row count of stored values that are nonzero (``obs["n_genes_by_counts"]``). + row_nnz: np.ndarray + #: Per-column sum of stored values (``var["total_counts"]``). + col_sums: np.ndarray + #: Per-column count of stored values that are nonzero (``var["n_cells_by_counts"]``). + col_nnz: np.ndarray + #: Per-column sum of depth-normalized values, ``sum_i(x[i, j] / row_sums[i])`` + #: (rows summing to zero divide by ``1`` instead). + col_depth_sums: np.ndarray + + def __post_init__(self) -> None: + for f in fields(self): + dtype = np.int64 if f.name.endswith("_nnz") else np.float64 + arr = np.array(getattr(self, f.name), dtype=dtype) + arr.flags.writeable = False + object.__setattr__(self, f.name, arr) + + @property + def shape(self) -> tuple[int, int]: + return (int(self.row_sums.shape[0]), int(self.col_sums.shape[0])) + + def _is_consistent(self) -> bool: + n_rows, n_cols = self.shape + return ( + self.row_nnz.shape == (n_rows,) + and self.col_nnz.shape == (n_cols,) + and self.col_depth_sums.shape == (n_cols,) + ) + + +def depth_divisors(row_sums: np.ndarray) -> np.ndarray: + """``row_sums``, with rows summing to zero mapped to ``1`` so they can be divided by.""" + return np.where(row_sums != 0.0, row_sums, 1.0) + + +def _nonzero_counts(arr: Any) -> tuple[np.ndarray, np.ndarray]: + """Per-major and per-minor counts of stored values that are nonzero.""" + nonzero = arr.values != 0 + nz_group_sizes = np.diff(arr.value_ptr) * nonzero + group_of_major = np.repeat(np.arange(arr.n_major, dtype=np.int64), np.diff(arr.major_ptr)) + major = np.bincount(group_of_major, weights=nz_group_sizes, minlength=arr.n_major) + minor = _ops.minor_sums( + nonzero.astype(np.float64), + arr.value_ptr, + arr.indices, + arr.n_minor, + _ops.accumulator_threads(arr.n_minor), + ) + return major.astype(np.int64), np.rint(minor).astype(np.int64) + + +@numba.njit(cache=True, inline="always") +def _is_sorted(seg: np.ndarray) -> bool: + for k in range(1, seg.shape[0]): # noqa: SIM110 -- numba has no all(genexpr) + if seg[k] < seg[k - 1]: + return False + return True + + +@numba.njit(cache=True, parallel=True) +def _depth_sums_major_is_col(major_ptr, values, value_ptr, indices, inv_depth): + """``inv_depth @ arr`` for a VCSC ``arr``, summing each value group in row order. + + A plain gather (``_ops.minor_matvec``) would add a group's rows up in + whatever order they're stored, and the IVCSC on-disk format re-sorts them + (:mod:`vsparse._ivcsc`), so a sidecar computed before writing could miss + one recomputed after reading by an ulp. Summing in sorted order makes the + result independent of that order; groups are usually sorted already. + """ + n_major = major_ptr.shape[0] - 1 + out = np.zeros(n_major, dtype=np.float64) + for j in numba.prange(n_major): # ty: ignore[not-iterable] + total = 0.0 + for u in range(major_ptr[j], major_ptr[j + 1]): + seg = indices[value_ptr[u] : value_ptr[u + 1]] + if not _is_sorted(seg): + seg = np.sort(seg) + acc = 0.0 + for r in seg: + acc += inv_depth[r] + total += np.float64(values[u]) * acc + out[j] = total + return out + + +def _col_depth_sums(arr: Any, row_sums: np.ndarray) -> np.ndarray: + inv_depth = 1.0 / depth_divisors(row_sums) + if arr._format == "csc": + return _depth_sums_major_is_col( + arr.major_ptr, arr.values, arr.value_ptr, arr.indices, inv_depth + ) + # Major=rows scatters each row once into each of its columns, in row + # order, so storage order within a value group doesn't matter here. + return inv_depth @ arr + + +def compute_stats(arr: Any) -> VCSStats: + """Compute ``arr``'s :class:`VCSStats` from its data, in ``O(nnz)``.""" + row_sums = np.asarray(arr.sum(axis=1), dtype=np.float64) + col_sums = np.asarray(arr.sum(axis=0), dtype=np.float64) + major_nnz, minor_nnz = _nonzero_counts(arr) + row_nnz, col_nnz = (minor_nnz, major_nnz) if arr._format == "csc" else (major_nnz, minor_nnz) + col_depth_sums = _col_depth_sums(arr, row_sums) + return VCSStats(row_sums, row_nnz, col_sums, col_nnz, col_depth_sums) + + +def write_stats( + g: Any, + stats: VCSStats, + *, + write_elem: Callable[..., None], + dataset_kwargs: Mapping[str, Any] = MappingProxyType({}), +) -> None: + """Write ``stats`` as the ``stats`` subgroup of an array's group ``g``.""" + sg = g.require_group(STATS_KEY) + for f in fields(stats): + write_elem(sg, f.name, getattr(stats, f.name), dataset_kwargs=dataset_kwargs) + sg.attrs["encoding-type"] = STATS_ENCODING_TYPE + sg.attrs["encoding-version"] = STATS_ENCODING_VERSION + + +def read_stats(g: Any, shape: tuple[int, int]) -> VCSStats | None: + """The sidecar stored under an array's group ``g``, or ``None`` if there isn't a usable one. + + ``None`` means the caller should recompute: the sidecar is absent (a file + written before it existed), of an encoding this version doesn't know, or + missing a field. One whose lengths don't match ``shape`` is also + discarded, with a warning, since that can only mean it's corrupt. + """ + if STATS_KEY not in g: + return None + sg = g[STATS_KEY] + if ( + sg.attrs.get("encoding-type") != STATS_ENCODING_TYPE + or sg.attrs.get("encoding-version") != STATS_ENCODING_VERSION + ): + return None + names = [f.name for f in fields(VCSStats)] + if not all(name in sg for name in names): + return None + stats = VCSStats(**{name: read_array(sg[name]) for name in names}) + if not (stats.shape == tuple(shape) and stats._is_consistent()): + warnings.warn( + f"ignoring a stats sidecar of shape {stats.shape} stored with an array " + f"of shape {tuple(shape)}; it will be recomputed", + stacklevel=2, + ) + return None + return stats diff --git a/tests/test_stats_sidecar.py b/tests/test_stats_sidecar.py new file mode 100644 index 0000000..8dc642b --- /dev/null +++ b/tests/test_stats_sidecar.py @@ -0,0 +1,186 @@ +"""The per-row/per-column statistics sidecar (``arr.stats``, :mod:`vsparse._stats`).""" + +from __future__ import annotations + +import warnings +from typing import Any + +import h5py +import numpy as np +import pytest +import scipy.sparse as sp +from hypothesis import given + +from _hypothesis_strategies import dense_matrices, slow_first_call +from vsparse import RECIPES, VCSCAnnData, VCSCArray, VCSRArray, load_packed +from vsparse import _stats as stats_mod + +CLASSES = [VCSCArray, VCSRArray] +_FIELDS = ("row_sums", "row_nnz", "col_sums", "col_nnz", "col_depth_sums") + + +def _counts(seed: int = 0, shape: tuple[int, int] = (40, 25)) -> np.ndarray: + rng = np.random.default_rng(seed) + dense = rng.poisson(1.0, size=shape).astype(np.float64) + dense[rng.random(shape) < 0.5] = 0.0 + return dense + + +def _arrays(adata: VCSCAnnData) -> tuple[Any, Any]: + """``adata``'s ``X``/``raw_X``, typed loosely: both are set in every test here.""" + return adata.X, adata.raw_X + + +def _assert_stats_equal(left: Any, right: Any) -> None: + for name in _FIELDS: + np.testing.assert_array_equal(getattr(left, name), getattr(right, name), err_msg=name) + + +def _assert_views_equal(left, right) -> None: + for name in ("row_scale", "gene_scale", "col_mean", "col_post_scale"): + np.testing.assert_array_equal(getattr(left, name), getattr(right, name), err_msg=name) + + +@slow_first_call +@given(dense=dense_matrices()) +@pytest.mark.parametrize("cls", CLASSES) +def test_stats_match_dense(cls, dense): + stats = cls.from_scipy(sp.csr_array(dense)).stats + row_sums = dense.sum(axis=1) + divisors = np.where(row_sums != 0, row_sums, 1.0) + np.testing.assert_allclose(stats.row_sums, row_sums) + np.testing.assert_allclose(stats.col_sums, dense.sum(axis=0)) + np.testing.assert_array_equal(stats.row_nnz, (dense != 0).sum(axis=1)) + np.testing.assert_array_equal(stats.col_nnz, (dense != 0).sum(axis=0)) + np.testing.assert_allclose(stats.col_depth_sums, (dense / divisors[:, None]).sum(axis=0)) + assert stats.shape == dense.shape + + +@pytest.mark.parametrize("cls", CLASSES) +def test_explicit_zeros_are_not_counted(cls): + mat = sp.csr_array((np.array([0.0, 2.0]), np.array([0, 1]), np.array([0, 2, 2])), shape=(2, 2)) + stats = cls.from_scipy(mat).stats + np.testing.assert_array_equal(stats.row_nnz, [1, 0]) + np.testing.assert_array_equal(stats.col_nnz, [0, 1]) + + +@pytest.mark.parametrize("cls", CLASSES) +def test_stats_are_cached_read_only_and_shared_by_copies(cls): + arr = cls.from_scipy(sp.csr_array(_counts())) + assert arr._stats is None + stats = arr.stats + assert arr.stats is stats + assert arr.copy()._stats is stats + assert arr[:, :]._stats is stats + with pytest.raises(ValueError, match="read-only"): + stats.row_sums[0] = 1.0 + # A subset is different data, so it gets its own. + assert arr[:5]._stats is None + + +@pytest.mark.parametrize("cls", CLASSES) +@pytest.mark.parametrize("recipe", sorted(RECIPES)) +def test_normalization_scales_come_from_the_sidecar(cls, recipe, monkeypatch): + arr = cls.from_scipy(sp.csr_array(_counts())) + expected = arr.normalized(recipe) + + # Restored statistics are used as-is: no recomputation, identical results. + restored = cls.from_scipy(sp.csr_array(_counts())) + restored._stats = arr.stats + monkeypatch.setattr(stats_mod, "compute_stats", lambda _arr: pytest.fail("stats recomputed")) + _assert_views_equal(restored.normalized(recipe), expected) + + +@pytest.mark.parametrize("cls", CLASSES) +@pytest.mark.parametrize("fmt", ["h5ad", "zarr"]) +def test_write_read_round_trips_the_sidecar(cls, fmt, tmp_path): + adata = VCSCAnnData( + X=cls.from_scipy(sp.csr_array(_counts(0))), + raw_X=cls.from_scipy(sp.csr_array(_counts(1))), + # Nested arrays go through the registered (unpacked) VCSC/VCSR codec. + uns={"counts": cls.from_scipy(sp.csr_array(_counts(2)))}, + ) + x, raw_x = _arrays(adata) + assert x._stats is None # never computed before writing ... + path = tmp_path / f"out.{fmt}" + getattr(adata, f"write_{fmt}")(path) + back = getattr(VCSCAnnData, f"read_{fmt}")(path) + back_x, back_raw_x = _arrays(back) + + # ... but written anyway, and read back exactly. + for name, original, read in ( + ("X", x, back_x), + ("raw_X", raw_x, back_raw_x), + ("uns", adata.uns["counts"], back.uns["counts"]), + ): + assert read._stats is not None, name + _assert_stats_equal(read._stats, stats_mod.compute_stats(read)) + _assert_stats_equal(read._stats, original.stats) + + for recipe in sorted(RECIPES): + _assert_views_equal(back_x.normalized(recipe), x.normalized(recipe)) + + +def test_load_packed_attaches_the_sidecar(tmp_path): + adata = VCSCAnnData(X=VCSRArray.from_scipy(sp.csr_array(_counts()))) + path = tmp_path / "out.h5ad" + adata.write_h5ad(path) + _assert_stats_equal(_arrays(load_packed(path))[0]._stats, _arrays(adata)[0].stats) + + +def _write_then_edit(tmp_path, edit) -> Any: + """Write a file, apply ``edit`` to its ``X`` group, and read back ``X``.""" + adata = VCSCAnnData(X=VCSCArray.from_scipy(sp.csr_array(_counts()))) + path = tmp_path / "out.h5ad" + adata.write_h5ad(path) + with h5py.File(path, "r+") as f: + edit(f["X"]) + return _arrays(VCSCAnnData.read_h5ad(path))[0] + + +def _drop_group(g): + del g["stats"] + + +def _bump_version(g): + g["stats"].attrs["encoding-version"] = "99.0.0" + + +def _drop_field(g): + del g["stats/col_depth_sums"] + + +@pytest.mark.parametrize("edit", [_drop_group, _bump_version, _drop_field]) +def test_unusable_sidecar_is_ignored_and_recomputed(tmp_path, edit): + with warnings.catch_warnings(): + warnings.simplefilter("error") + back_x = _write_then_edit(tmp_path, edit) + assert back_x._stats is None + expected = VCSCArray.from_scipy(sp.csr_array(_counts())) + _assert_stats_equal(back_x.stats, expected.stats) + _assert_views_equal(back_x.normalized(), expected.normalized()) + + +def test_mismatched_sidecar_warns_and_is_recomputed(tmp_path): + def truncate(g): + del g["stats/row_sums"] + g["stats"].create_dataset("row_sums", data=np.ones(3)) + + with pytest.warns(UserWarning, match="stats sidecar"): + back_x = _write_then_edit(tmp_path, truncate) + assert back_x._stats is None + assert back_x.stats.shape == back_x.shape + + +@pytest.mark.parametrize("cls", CLASSES) +def test_stats_do_not_depend_on_order_within_value_groups(cls): + # IVCS storage re-sorts each value group's indices, so the sidecar + # written before a round trip has to match what's recomputed after it. + arr = cls.from_scipy(sp.csr_array(_counts(shape=(300, 40)))) + rng = np.random.default_rng(0) + shuffled = arr.indices.copy() + for u in range(arr.n_unique): + seg = shuffled[arr.value_ptr[u] : arr.value_ptr[u + 1]] + seg[:] = rng.permutation(seg) + other = cls(arr.shape, arr.major_ptr, arr.values, arr.value_ptr, shuffled) + _assert_stats_equal(other.stats, arr.stats)