From a724d9adcf27095dac543dfb14a522ad8b7d0992 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Wed, 14 Oct 2026 07:22:09 -0700 Subject: [PATCH] Remove load_and_normalize in favor of load, subset, then normalize The fused loader and its helpers are replaced by load_packed plus the .normalized() views. Drops the tests and docs that only covered it. Claude-Session: https://claude.ai/code/session_01QBeTZy6BKKGKGJ7fwvkWgM --- README.md | 14 +- docs/api.md | 1 - docs/usage.md | 21 +- src/vsparse/__init__.py | 3 +- src/vsparse/_rapid_load.py | 580 +------------------------------------ src/vsparse/_vcs_norm.py | 3 +- tests/test_index_dtypes.py | 45 --- tests/test_rapid_load.py | 388 +------------------------ 8 files changed, 25 insertions(+), 1030 deletions(-) diff --git a/README.md b/README.md index d0dc2ae..125ee26 100644 --- a/README.md +++ b/README.md @@ -133,19 +133,13 @@ va.write_h5ad("archived.h5ad", format="ivcsc") va_loaded = vsparse.VCSCAnnData.read_h5ad("archived.h5ad") ``` -### Fast Filtering & Depth Normalization (`load_and_normalize`) +### Loading, subsetting, and normalizing -For IVCSR-backed datasets, `vsparse.load_and_normalize` bypasses full array decompression for rapid preprocessing, reproducing the filtering and depth normalization from `parafac2.normalize.prepare_dataset`: - -- **Cell filtering without decoding**: Computes cell totals in $O(n_{\text{unique}})$ time directly from group sizes without unpacking varint indices. -- **Fused filter & normalization**: Performs gene filtering, cell/gene depth-scaling, and $\log_{10}(1000x + 1)$ transform in parallel passes over the compact CSR representation. +Loading, subsetting, and normalization are separate steps, so they can be mixed and matched. `vsparse.load_packed` decodes an IVCSR/IVCSC-stored `.h5ad` into a `VCSCAnnData`; normalization is then a view over the stored array (`.normalized("parafac2")` reproduces `parafac2.normalize.prepare_dataset`'s depth normalization and log10 transform) that is never materialized. ```python -adata_norm = vsparse.load_and_normalize( - "archived.h5ad", - min_cell_counts=10.0, # Filter cells with counts <= 10 - gene_threshold=0.05, # Filter genes with counts <= 0.05 * n_cells -) +va = vsparse.load_packed("archived.h5ad") +norm = va.X.normalized("parafac2") # lazy view; norm @ B works without densifying ``` ### CUDA (`to_gpu`) diff --git a/docs/api.md b/docs/api.md index afec1d9..d9a9d84 100644 --- a/docs/api.md +++ b/docs/api.md @@ -17,7 +17,6 @@ .. autofunction:: vsparse.to_layer -.. autofunction:: vsparse.load_and_normalize .. autoclass:: vsparse.CudaNormalizedView :members: diff --git a/docs/usage.md b/docs/usage.md index b26039d..a523e21 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -111,27 +111,20 @@ views, concatenation, most of the scanpy/anndata ecosystem -- are **not** suppor directly, so they round-trip through `write_elem`/`read_elem` even when nested inside a plain `AnnData`'s `.uns`, independent of `VCSCAnnData`. -## Fast IVCSR loading and normalization +## Loading and normalizing IVCSR files -When working with large IVCSR-stored `.h5ad` files, `vsparse.load_and_normalize` provides -a fused loader and depth-normalization kernel reproducing the preprocessing in -`parafac2.normalize.prepare_dataset`: +`vsparse.load_packed` decodes an IVCSR/IVCSC-stored `.h5ad` into a `VCSCAnnData`. +Loading, subsetting, and normalization are separate steps; normalization is a +lazy view (`.normalized("parafac2")`) over the loaded array, so the steps can be +mixed and matched rather than fused into one loader: ```python import vsparse -adata = vsparse.load_and_normalize( - "archived.ivcsr.h5ad", - min_cell_counts=10.0, - gene_threshold=0.05, -) +va = vsparse.load_packed("archived.ivcsr.h5ad") +norm = va.X.normalized("parafac2") ``` -This bypasses generic `AnnData` indexing and decompression overhead: -1. Cell counts are computed in $O(n_{\text{unique}})$ time from unique-value group sizes without touching or decoding the packed `indices` byte stream. -2. The delta+varint indices are unpacked in parallel across CPU cores. -3. Row/gene filtering, compaction, and depth normalization ($(\text{cell\_scale}) \times (\text{gene\_sums})$ followed by $\log_{10}(1000x + 1)$) are executed in fused parallel passes directly into the output CSR representation. - ## 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 5474f48..a1f749f 100644 --- a/src/vsparse/__init__.py +++ b/src/vsparse/__init__.py @@ -11,7 +11,7 @@ from vsparse._base import VCSCArray, VCSRArray from vsparse._cuda import CudaNormalizedView, cuda_is_available from vsparse._norm_common import RECIPES, Recipe -from vsparse._rapid_load import load_and_normalize, load_packed +from vsparse._rapid_load import load_packed from vsparse._vcs_norm import VCSCArrayNormalized, VCSRArrayNormalized __all__ = [ @@ -25,7 +25,6 @@ "VCSRArrayNormalized", "cuda_is_available", "from_anndata", - "load_and_normalize", "load_packed", "to_layer", ] diff --git a/src/vsparse/_rapid_load.py b/src/vsparse/_rapid_load.py index 19d8a3e..d63ba13 100644 --- a/src/vsparse/_rapid_load.py +++ b/src/vsparse/_rapid_load.py @@ -1,594 +1,34 @@ -"""Fast filter/normalize loader for IVCSR-stored ``.h5ad`` files. +"""Loader for IVCSR/IVCSC-stored ``.h5ad`` files. -Mirrors the preprocessing in `parafac2.normalize.prepare_dataset -`_ -(minimum-count cell filtering, minimum-expression gene filtering, per-cell -read-depth normalization, and a log10 transform) but is written directly -against the on-disk VCSR/IVCSR layout instead of going through -:meth:`~vsparse.VCSCAnnData.read_h5ad` + generic ``AnnData`` indexing. - -Why this needs its own code path ---------------------------------- -A VCSR major (row/cell) slice stores each *unique* value once, alongside the -list of minor-axis (gene) indices that share it -- ``values``/``value_ptr`` -are indexed per unique-value group, not per nonzero. Two consequences drive -the design here: - -- **Cell filtering is (almost) free.** A cell's total count is - ``sum(value * group_size for each unique-value group in that row)``, a sum - over ``n_unique`` (~46x fewer entries than ``nnz`` on the bundled example - dataset) that touches only ``major_ptr``/``values``/``value_ptr``. It needs - none of the (compressed, delta+varint-packed) ``indices`` array, so the - cell mask can be computed without ever decoding it. -- **Gene filtering and normalization must touch every retained nonzero.** A - gene's total count is a scatter-add keyed by minor-axis index, so there is - no shortcut analogous to the cell case. Without an ``obs_filter``, every - index is decoded. With an ``obs_filter``, the packed stream must still be - scanned because row byte offsets are not stored, but indices and values - are materialized only for selected rows. Once decoded, the per-cell scale - factor differs per row while the per-gene scale factor differs per column, - so after both are applied a row's values are no longer mostly-repeated -- - the whole point of VCSC/VCSR's deduplication is gone. Normalization is - therefore done on a plain CSR array, built once, right after the two masks - are known -- via a fused numba pass rather than scipy's generic per-step - (row sum, row scale, column sum, column scale, multiply, add, log10) - elementwise pipeline, each step of which is another single-threaded - full-``nnz`` pass with its own temporary. - -The one thing this implementation does *not* do is what -:func:`parafac2.normalize.prepare_dataset` calls "indexing for subsetting the -data": ``X[cell_mask, gene_mask]``. That line re-derives a filtered matrix via -scipy's generic fancy indexing (row selection, then a column selection that -internally goes through a CSC conversion). Here, row/gene filtering, -compaction, and column-index remapping happen in one fused parallel pass -(:func:`_count_kept`/:func:`_fill_kept`) directly on the decoded arrays, so -the filtered CSR is built once instead of assembled and then re-sliced. -Gene means (``X.var["means"]``) aren't computed either -- nothing here needs -them, and computing them is a simple ``X.mean(axis=0)`` for callers that do. +: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. """ from __future__ import annotations -from collections.abc import Callable from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING import anndata as ad -import numba -import numpy as np -import pandas as pd -from scipy.sparse import csr_array -from vsparse import _ivcsc from vsparse._anndata_class import VCSCAnnData -from vsparse._indexutils import smallest_index_dtype -from vsparse._ops import accumulator_threads if TYPE_CHECKING: from os import PathLike -__all__ = ["load_and_normalize", "load_packed"] - - -# -- group-level (no-decode) cell totals ------------------------------------- - - -def _cell_totals(major_ptr: np.ndarray, values: np.ndarray, value_ptr: np.ndarray) -> np.ndarray: - """Per-row total counts, computed without touching ``indices``. - - Each unique-value group of size ``value_ptr[k + 1] - value_ptr[k]`` - contributes ``values[k] * group_size`` to its row's total -- so this - only costs ``O(n_unique)``, not ``O(nnz)``. - """ - n_major = major_ptr.shape[0] - 1 - group_sizes = np.diff(value_ptr) - weighted = values.astype(np.float64) * group_sizes - row_of_group = np.repeat(np.arange(n_major, dtype=np.int64), np.diff(major_ptr)) - return np.bincount(row_of_group, weights=weighted, minlength=n_major) - - -# -- parallel value expansion (values/value_ptr -> per-nonzero data) -------- - - -@numba.njit(cache=True, parallel=True) -def _expand_values(values: np.ndarray, value_ptr: np.ndarray, out: np.ndarray) -> None: - n_unique = values.shape[0] - for k in numba.prange(n_unique): # ty: ignore[not-iterable] - v = values[k] - for p in range(value_ptr[k], value_ptr[k + 1]): - out[p] = v - - -def _build_data(values: np.ndarray, value_ptr: np.ndarray, nnz: int) -> np.ndarray: - out = np.empty(nnz, dtype=values.dtype) - _expand_values(values, value_ptr, out) - return out - - -# -- selective packed decode (obs-filtered IVCSR rows only) ----------------- - - -@numba.njit(cache=True) -def _decode_selected_rows( - major_ptr: np.ndarray, - values: np.ndarray, - value_ptr: np.ndarray, - packed: np.ndarray, - row_mask: np.ndarray, - out_indices: np.ndarray, - out_data: np.ndarray, -) -> None: - """Decode only selected IVCSR rows while scanning past excluded rows.""" - pos = np.int64(0) - out_pos = 0 - n_rows = major_ptr.shape[0] - 1 - - for r in range(n_rows): - keep = row_mask[r] - for g in range(major_ptr[r], major_ptr[r + 1]): - if keep: - prev = np.int64(-1) - value = values[g] - for _ in range(value_ptr[g], value_ptr[g + 1]): - result, pos = _ivcsc._decode_varint(packed, pos) - prev = prev + 1 + np.int64(result) - out_indices[out_pos] = prev - out_data[out_pos] = value - out_pos += 1 - else: - for _ in range(value_ptr[g], value_ptr[g + 1]): - while packed[pos] & 0x80: - pos += 1 - pos += 1 - - -def _build_selected_rows( - major_ptr: np.ndarray, - values: np.ndarray, - value_ptr: np.ndarray, - packed: np.ndarray, - indices_dtype: np.dtype, - row_mask: np.ndarray, -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - full_indptr = value_ptr[major_ptr] - row_nnz = np.diff(full_indptr)[row_mask] - nnz = int(row_nnz.sum()) - ptr_dtype = smallest_index_dtype(nnz) - - indptr = np.zeros(row_nnz.shape[0] + 1, dtype=ptr_dtype) - np.cumsum(row_nnz, out=indptr[1:]) - indices = np.empty(nnz, dtype=indices_dtype) - data = np.empty(nnz, dtype=values.dtype) - _decode_selected_rows(major_ptr, values, value_ptr, packed, row_mask, indices, data) - return indptr, indices, data - - -# -- parallel weighted bincount (raw per-gene totals, for gene_mask) -------- - - -@numba.njit(cache=True, parallel=True) -def _weighted_bincount( - indices: np.ndarray, data: np.ndarray, n_bins: int, nthreads: int -) -> np.ndarray: - n = indices.shape[0] - chunk = (n + nthreads - 1) // nthreads - partial = np.zeros((nthreads, n_bins), dtype=np.float64) - for t in numba.prange(nthreads): # ty: ignore[not-iterable] - start = t * chunk - end = min(n, start + chunk) - local = partial[t] - for k in range(start, end): - local[indices[k]] += data[k] - return partial.sum(axis=0) - - -@numba.njit(cache=True, parallel=True) -def _gene_detection_counts( - indices: np.ndarray, data: np.ndarray, n_genes: int, nthreads: int -) -> np.ndarray: - """Number of cells with positive expression for each gene.""" - n = indices.shape[0] - chunk = (n + nthreads - 1) // nthreads - partial = np.zeros((nthreads, n_genes), dtype=np.int64) - for t in numba.prange(nthreads): # ty: ignore[not-iterable] - start = t * chunk - end = min(n, start + chunk) - local = partial[t] - for k in range(start, end): - if data[k] > 0: - local[indices[k]] += 1 - return partial.sum(axis=0) - - -# -- fused row+column filter/compaction -------------------------------------- - - -@numba.njit(cache=True, parallel=True) -def _count_kept( - row_indptr: np.ndarray, - indices: np.ndarray, - gene_mask: np.ndarray, - kept_rows: np.ndarray, - out_counts: np.ndarray, -) -> None: - for j in numba.prange(kept_rows.shape[0]): # ty: ignore[not-iterable] - r = kept_rows[j] - c = 0 - for k in range(row_indptr[r], row_indptr[r + 1]): - if gene_mask[indices[k]]: - c += 1 - out_counts[j] = c - - -@numba.njit(cache=True, parallel=True) -def _fill_kept( - row_indptr: np.ndarray, - indices: np.ndarray, - data: np.ndarray, - gene_remap: np.ndarray, - kept_rows: np.ndarray, - new_indptr: np.ndarray, - out_indices: np.ndarray, - out_data: np.ndarray, -) -> None: - for j in numba.prange(kept_rows.shape[0]): # ty: ignore[not-iterable] - r = kept_rows[j] - pos = new_indptr[j] - for k in range(row_indptr[r], row_indptr[r + 1]): - g = gene_remap[indices[k]] - if g >= 0: - out_indices[pos] = g - out_data[pos] = data[k] - pos += 1 - - -def _filter_and_compact( - row_indptr: np.ndarray, - indices: np.ndarray, - data: np.ndarray, - cell_mask: np.ndarray, - gene_mask: np.ndarray, -) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int]: - kept_rows = np.nonzero(cell_mask)[0] - n_kept_genes = int(gene_mask.sum()) - gene_remap = (np.cumsum(gene_mask) - 1).astype(smallest_index_dtype(n_kept_genes)) - gene_remap[~gene_mask] = -1 - - # ``new_indptr`` is indexed by nonzero count and needs int64 once nnz - # passes INT32_MAX; ``out_indices`` holds gene indices and never does. - ptr_dtype = smallest_index_dtype(int(indices.shape[0])) - col_dtype = smallest_index_dtype(n_kept_genes) - - counts = np.empty(kept_rows.shape[0], dtype=ptr_dtype) - _count_kept(row_indptr, indices, gene_mask, kept_rows, counts) - new_indptr = np.zeros(kept_rows.shape[0] + 1, dtype=ptr_dtype) - np.cumsum(counts, out=new_indptr[1:]) - - nnz_filtered = int(new_indptr[-1]) - out_indices = np.empty(nnz_filtered, dtype=col_dtype) - out_data = np.empty(nnz_filtered, dtype=np.float32) - _fill_kept(row_indptr, indices, data, gene_remap, kept_rows, new_indptr, out_indices, out_data) - - return new_indptr, out_indices, out_data, kept_rows, n_kept_genes - - -# -- fused row-scale / gene-scale / log10 transform -------------------------- -# -# scipy's route to the same result -- X.sum(axis=1), a np.repeat of the -# per-row scale over nnz, an in-place divide, X.sum(axis=0), a fancy-index -# gather of the per-column scale over nnz, another in-place divide, then -# three more elementwise passes (*1000, +1, log10) -- is ~9 passes over the -# (still nnz-scale, filtering rarely drops much on real data) array plus two -# extra full-nnz temporaries (the repeat and the gather). Every one of those -# passes and temporaries is single-threaded generic numpy code. Fusing the -# whole tail into three parallel passes over the already-filtered CSR (row -# sums -> per-column sums of the row-scaled data -> the final transform, -# each row-parallel with no nnz-sized temporary beyond the output) removes -# both the redundant passes and the memory pressure from those temporaries. - - -@numba.njit(cache=True, parallel=True) -def _row_sums(indptr: np.ndarray, data: np.ndarray, out: np.ndarray) -> None: - for r in numba.prange(out.shape[0]): # ty: ignore[not-iterable] - s = 0.0 - for k in range(indptr[r], indptr[r + 1]): - s += data[k] - out[r] = s - - -@numba.njit(cache=True, parallel=True) -def _scaled_col_sums( - indptr: np.ndarray, - indices: np.ndarray, - data: np.ndarray, - row_scale: np.ndarray, - n_bins: int, - nthreads: int, -) -> np.ndarray: - n_rows = indptr.shape[0] - 1 - chunk = (n_rows + nthreads - 1) // nthreads - partial = np.zeros((nthreads, n_bins), dtype=np.float64) - for t in numba.prange(nthreads): # ty: ignore[not-iterable] - start = t * chunk - end = min(n_rows, start + chunk) - local = partial[t] - for r in range(start, end): - sc = row_scale[r] - for k in range(indptr[r], indptr[r + 1]): - local[indices[k]] += data[k] / sc - return partial.sum(axis=0) - - -@numba.njit(cache=True, parallel=True) -def _fused_normalize_transform( - indptr: np.ndarray, - indices: np.ndarray, - data: np.ndarray, - row_scale: np.ndarray, - gene_sums: np.ndarray, -) -> None: - """Overwrites ``data`` in place -- each output only depends on its own input.""" - n_rows = indptr.shape[0] - 1 - for r in numba.prange(n_rows): # ty: ignore[not-iterable] - sc = row_scale[r] - for k in range(indptr[r], indptr[r + 1]): - v = (data[k] / sc) / gene_sums[indices[k]] - data[k] = np.log10(np.float32(1.0) + np.float32(1000.0) * v) - - -def _normalize_and_transform( - indptr: np.ndarray, indices: np.ndarray, data: np.ndarray, n_genes: int -) -> np.ndarray: - """Depth normalization + log10(1000x + 1), fused, overwriting ``data`` in place.""" - n_rows = indptr.shape[0] - 1 - nthreads = numba.get_num_threads() - - counts_per_cell = np.empty(n_rows, dtype=np.float64) - _row_sums(indptr, data, counts_per_cell) - counts_per_cell /= np.median(counts_per_cell) - - gene_sums = _scaled_col_sums(indptr, indices, data, counts_per_cell, n_genes, nthreads) - - _fused_normalize_transform(indptr, indices, data, counts_per_cell, gene_sums) - return data - - -# -- top-level entry point --------------------------------------------------- - - -def _compute_gene_mask( - indices: np.ndarray, - data: np.ndarray, - n_genes: int, - denom: int, - gene_threshold: float, - min_cells: int | None, -) -> np.ndarray: - """Genes with total raw counts above ``gene_threshold * denom``, and detected in >= ``min_cells``.""" - gene_totals_raw = _weighted_bincount(indices, data, n_genes, accumulator_threads(n_genes)) - gene_mask = gene_totals_raw > (gene_threshold * denom) - if min_cells is not None: - gene_detection_counts = _gene_detection_counts( - indices, data, n_genes, accumulator_threads(n_genes) - ) - gene_mask &= gene_detection_counts >= min_cells - return gene_mask - - -def _rows_without_obs_filter( - major_ptr: np.ndarray, - values: np.ndarray, - value_ptr: np.ndarray, - packed: np.ndarray, - indices_dtype: np.dtype, - cell_totals: np.ndarray, - min_cell_counts: float, -) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - """Decode every row; ``metadata_cell_mask`` is just ``cell_mask`` since none were dropped upfront.""" - cell_mask = cell_totals > min_cell_counts - indices = _ivcsc.unpack_indices(value_ptr, packed, indices_dtype) - data = _build_data(values, value_ptr, indices.shape[0]) - row_indptr = value_ptr[major_ptr] - return cell_mask, row_indptr, indices, data, cell_mask - - -def _rows_with_obs_filter( - major_ptr: np.ndarray, - values: np.ndarray, - value_ptr: np.ndarray, - packed: np.ndarray, - indices_dtype: np.dtype, - cell_totals: np.ndarray, - obs: Any, - obs_filter: Callable[[pd.DataFrame], object], - min_cell_counts: float, - n_cells: int, -) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, int]: - """Validate ``obs_filter``, then decode only the rows it selects.""" - if not isinstance(obs, pd.DataFrame): - raise ValueError("obs_filter requires an obs table in the h5ad file") - if not callable(obs_filter): - raise TypeError("obs_filter must be callable or None") - - obs_mask = np.asarray(obs_filter(obs)) - if obs_mask.ndim != 1 or obs_mask.shape[0] != n_cells: - raise ValueError(f"obs_filter must return a one-dimensional mask of length {n_cells}") - if obs_mask.dtype != np.bool_: - raise ValueError("obs_filter must return a boolean mask") - if not np.any(obs_mask): - raise ValueError("obs_filter selected no cells") - obs_mask = np.ascontiguousarray(obs_mask) - - selected_rows = np.nonzero(obs_mask)[0] - cell_mask = cell_totals[obs_mask] > min_cell_counts - row_indptr, indices, data = _build_selected_rows( - major_ptr, values, value_ptr, packed, indices_dtype, obs_mask - ) - - metadata_cell_mask = np.zeros(n_cells, dtype=np.bool_) - metadata_cell_mask[selected_rows[cell_mask]] = True - return cell_mask, row_indptr, indices, data, metadata_cell_mask, selected_rows.shape[0] - +__all__ = ["load_packed"] _FIELD_KEYS = ("obs", "var", "obsm", "varm", "obsp", "varp", "layers", "uns") -def _validate_min_cells(min_cells: int | None) -> None: - if min_cells is None: - return - if isinstance(min_cells, bool) or not isinstance(min_cells, int): - raise TypeError("min_cells must be an integer or None") - if min_cells < 0: - raise ValueError("min_cells must be non-negative") - - -def _read_ivcsr_group( - f: Any, x_key: str -) -> tuple[ - tuple[int, int], np.dtype, np.ndarray, np.ndarray, np.ndarray, np.ndarray, dict[str, Any] -]: - """Read the packed IVCSR array at ``f[x_key]`` plus the surrounding obs/var/etc. fields.""" - g = f[x_key] - shape = (int(g.attrs["shape"][0]), int(g.attrs["shape"][1])) - indices_dtype = np.dtype(g.attrs["indices_dtype"]) - major_ptr = g["major_ptr"][...] - values = g["values"][...] - value_ptr = g["value_ptr"][...] - packed = g["packed_indices"][...] - - kwargs = {k: ad.io.read_elem(f[k]) for k in _FIELD_KEYS if k in f} - if "raw" in f: - kwargs["raw"] = ad.io.read_elem(f["raw"]) - return shape, indices_dtype, major_ptr, values, value_ptr, packed, kwargs - - -def load_and_normalize( - path: str | PathLike[str], - *, - min_cell_counts: float = 10.0, - gene_threshold: float = 0.0, - min_cells: int | None = None, - obs_filter: Callable[[pd.DataFrame], object] | None = None, - x_key: str = "X", -) -> ad.AnnData: - """Load, filter, and depth-normalize a VCSR/IVCSR-backed ``.h5ad`` file. - - Reproduces ``parafac2.normalize.prepare_dataset``: cells with total - counts <= ``min_cell_counts`` and genes with total counts <= - ``gene_threshold * n_cells`` are dropped. When ``min_cells`` is given, - genes expressed in fewer than ``min_cells`` cells are also dropped. Gene - filters are measured on the raw counts after any ``obs_filter``. - The remaining matrix is row-normalized to the median per-cell depth, then - column-normalized by gene sum, then transformed as ``log10(1000x + 1)``. - Surrounding metadata (``obs``, ``var``, ``obsm``, etc.) is sliced to - match the retained cells and genes. - - Parameters - ---------- - path - Path to an ``.h5ad`` file whose ``X`` (or ``layers[x_key]``) was - written with ``format="ivcsc"``/``"ivcsr"`` (see - :meth:`~vsparse.VCSCAnnData.write_h5ad`). - min_cell_counts - Cells with total raw counts <= this are dropped. - gene_threshold - Minimum threshold fraction for gene inclusion, as in - ``parafac2.normalize.prepare_dataset``: genes with total raw counts - <= ``gene_threshold * n_cells`` are dropped. - min_cells - Optional gene filter. Genes expressed in fewer than this - many cells are dropped. Expression is defined as a raw count > 0. - obs_filter - Optional callable receiving ``obs`` and returning a one-dimensional - boolean mask. When provided, rows are subset before cell filtering, - gene filtering, and normalization, so gene totals and - ``gene_threshold`` are computed using only the selected cells. The - packed IVCSR stream is still read in full, but indices and values are - materialized only for selected rows. - x_key - Top-level h5ad group holding the IVCSR array (``"X"`` by default). - - Returns - ------- - ad.AnnData - Filtered, depth-normalized AnnData object with ``X`` as a CSR array - and sliced metadata. - - Examples - -------- - Select cells using multiple ``obs`` columns and multiple accepted values:: - - load_and_normalize( - path, - obs_filter=lambda obs: ( - obs["condition"].isin(["control", "vehicle"]) - & (obs["timepoint"] == "T3") - ), - ) - """ - import h5py - import hdf5plugin # noqa: F401 -- registers the Blosc2 HDF5 filter - - _validate_min_cells(min_cells) - - with h5py.File(Path(path), "r") as f: - shape, indices_dtype, major_ptr, values, value_ptr, packed, kwargs = _read_ivcsr_group( - f, x_key - ) - - n_cells, n_genes = shape - - # Cell mask: no decode needed. - cell_totals = _cell_totals(major_ptr, values, value_ptr) - - if obs_filter is None: - cell_mask, row_indptr, indices, data, metadata_cell_mask = _rows_without_obs_filter( - major_ptr, values, value_ptr, packed, indices_dtype, cell_totals, min_cell_counts - ) - gene_denom = n_cells - else: - cell_mask, row_indptr, indices, data, metadata_cell_mask, gene_denom = ( - _rows_with_obs_filter( - major_ptr, - values, - value_ptr, - packed, - indices_dtype, - cell_totals, - kwargs.get("obs"), - obs_filter, - min_cell_counts, - n_cells, - ) - ) - del packed - - gene_mask = _compute_gene_mask(indices, data, n_genes, gene_denom, gene_threshold, min_cells) - - new_indptr, out_indices, out_data, kept_rows, n_kept_genes = _filter_and_compact( - row_indptr, indices, data, cell_mask, gene_mask - ) - del indices, data, row_indptr - - normalized = _normalize_and_transform(new_indptr, out_indices, out_data, n_kept_genes) - X = csr_array((normalized, out_indices, new_indptr), shape=(kept_rows.shape[0], n_kept_genes)) - - if "obs" in kwargs and "var" in kwargs: - adata = ad.AnnData(**kwargs) - else: - adata = ad.AnnData(shape=shape, **kwargs) - adata = adata[metadata_cell_mask, gene_mask].copy() - adata.X = X - - return adata - - def load_packed(path: str | PathLike[str], *, x_key: str = "X") -> VCSCAnnData: """Load an IVCSR/IVCSC-backed ``.h5ad`` file, decoding ``X`` immediately. - Unlike :func:`load_and_normalize`, this does no filtering or - normalization -- ``X`` comes back as an ordinary VCSCArray/VCSRArray with - plain ``indices``, decoded from the packed bytes read from disk. + ``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. Parameters ---------- diff --git a/src/vsparse/_vcs_norm.py b/src/vsparse/_vcs_norm.py index eaab3cd..17ef87a 100644 --- a/src/vsparse/_vcs_norm.py +++ b/src/vsparse/_vcs_norm.py @@ -3,8 +3,7 @@ :class:`VCSCArrayNormalized`/:class:`VCSRArrayNormalized` wrap a plain (unpacked) VCSC/VCSR array and behave like a normalized matrix -- read-depth normalized, transformed, optionally centered/scaled, per one of -:data:`vsparse._norm_common.RECIPES` (the default, ``"parafac2"``, matches -what :func:`vsparse._rapid_load.load_and_normalize` builds) -- without ever +:data:`vsparse._norm_common.RECIPES` (the default is ``"parafac2"``) -- without ever materializing it. See :mod:`vsparse._norm_common` for the shared recipe/ statistics/materialization logic (:class:`~vsparse._norm_common.NormalizedViewBase`) and :mod:`vsparse._vcs_matmul` for the matmul kernels these use (a direct diff --git a/tests/test_index_dtypes.py b/tests/test_index_dtypes.py index 8291244..63b3521 100644 --- a/tests/test_index_dtypes.py +++ b/tests/test_index_dtypes.py @@ -6,7 +6,6 @@ from vsparse import VCSCAnnData, VCSCArray, VCSRArray from vsparse._indexutils import smallest_index_dtype -from vsparse._rapid_load import _filter_and_compact INT32_MAX = np.iinfo(np.int32).max @@ -127,47 +126,3 @@ def test_packed_write_narrows_a_wide_in_memory_array(tmp_path, dense): assert isinstance(back.X, VCSRArray) assert back.X.indices.dtype == np.int32 np.testing.assert_allclose(back.X.toarray(), dense) - - -def _small_filter_inputs(): - dense = np.array( - [[1.0, 0.0, 2.0, 0.0], [0.0, 3.0, 0.0, 4.0], [5.0, 0.0, 6.0, 0.0]], - dtype=np.float32, - ) - return sp.csr_array(dense), np.array([True, False, True]), np.array([True, False, True, False]) - - -def test_filter_and_compact_uses_int32_when_both_bounds_fit(): - X, cell_mask, gene_mask = _small_filter_inputs() - new_indptr, out_indices, out_data, kept_rows, n_kept = _filter_and_compact( - X.indptr, X.indices, X.data, cell_mask, gene_mask - ) - - assert new_indptr.dtype == np.int32 - assert out_indices.dtype == np.int32 - assert n_kept == 2 - np.testing.assert_array_equal(kept_rows, [0, 2]) - np.testing.assert_allclose(out_data, [1.0, 2.0, 5.0, 6.0]) - np.testing.assert_array_equal(out_indices, [0, 1, 0, 1]) - - -def test_gene_indices_stay_int32_when_pointers_need_int64(monkeypatch): - """Forced, since a >INT32_MAX-nonzero matrix cannot be allocated in a test.""" - import vsparse._rapid_load as rapid_load - - X, cell_mask, gene_mask = _small_filter_inputs() - nnz_in = int(X.indices.shape[0]) - real = rapid_load.smallest_index_dtype - - monkeypatch.setattr( - rapid_load, - "smallest_index_dtype", - lambda n: np.dtype(np.int64) if n == nnz_in else real(n), - ) - new_indptr, out_indices, _, _, _ = _filter_and_compact( - X.indptr, X.indices, X.data, cell_mask, gene_mask - ) - - assert new_indptr.dtype == np.int64 - assert out_indices.dtype == np.int32 - np.testing.assert_array_equal(out_indices, [0, 1, 0, 1]) diff --git a/tests/test_rapid_load.py b/tests/test_rapid_load.py index 3f5ee22..d0c8c71 100644 --- a/tests/test_rapid_load.py +++ b/tests/test_rapid_load.py @@ -2,40 +2,10 @@ import anndata as ad import numpy as np -import pytest import scipy.sparse as sp from vsparse import VCSCAnnData, VCSRArray -from vsparse._rapid_load import load_and_normalize, load_packed - - -def _reference_prepare( - dense: np.ndarray, - min_cell_counts: float, - gene_threshold: float, - min_cells: int | None = None, -): - """Direct translation of parafac2.normalize.prepare_dataset's math, on a plain array.""" - X = sp.csr_array(dense) - cell_mask = np.ravel(X.sum(axis=1)) > min_cell_counts - gene_mask = np.ravel(X.sum(axis=0)) > (gene_threshold * X.shape[0]) - if min_cells is not None: - gene_mask &= np.ravel((X > 0).sum(axis=0)) >= min_cells - Xf = sp.csr_array(X[cell_mask][:, gene_mask]) - Xf.data = Xf.data.astype(np.float32) - - counts_per_cell = np.ravel(Xf.sum(axis=1)).astype(np.float32) - counts_per_cell /= np.median(counts_per_cell) - Xf.data /= np.repeat(counts_per_cell, np.diff(Xf.indptr)) - - gene_sums = np.ravel(Xf.sum(axis=0)).astype(np.float32) - Xf.data /= gene_sums[Xf.indices] - - Xf.data *= np.float32(1000.0) - Xf.data += np.float32(1.0) - np.log10(Xf.data, out=Xf.data) - - return Xf, np.nonzero(cell_mask)[0], np.nonzero(gene_mask)[0] +from vsparse._rapid_load import load_packed def _write_ivcsr(tmp_path, dense: np.ndarray, name="data.ivcsr.h5ad"): @@ -46,362 +16,8 @@ def _write_ivcsr(tmp_path, dense: np.ndarray, name="data.ivcsr.h5ad"): return path -@pytest.mark.parametrize( - ("min_cell_counts", "gene_threshold"), - [(-1.0, 0.0), (10.0, 0.05), (3.0, 0.2)], -) -def test_matches_reference_normalization(tmp_path, rng, min_cell_counts, gene_threshold): - dense = rng.integers(0, 8, size=(120, 40)).astype(np.float64) - dense[rng.random(dense.shape) < 0.5] = 0.0 - path = _write_ivcsr(tmp_path, dense) - - result = load_and_normalize( - path, min_cell_counts=min_cell_counts, gene_threshold=gene_threshold - ) - ref_X, _, _ = _reference_prepare(dense, min_cell_counts, gene_threshold) - - assert isinstance(result, ad.AnnData) - assert result.shape == ref_X.shape - assert isinstance(result.X, sp.csr_array) - np.testing.assert_allclose(result.X.toarray(), ref_X.toarray(), rtol=1e-5, atol=1e-5) - - -def test_no_filtering_keeps_everything(tmp_path, rng): - dense = rng.integers(0, 5, size=(30, 10)).astype(np.float64) - path = _write_ivcsr(tmp_path, dense) - - result = load_and_normalize(path, min_cell_counts=-1.0, gene_threshold=0.0) - - assert isinstance(result, ad.AnnData) - assert result.shape == dense.shape - assert result.n_obs == dense.shape[0] - assert result.n_vars == dense.shape[1] - - -def test_min_cells_filters_genes_by_detection_count(tmp_path): - dense = np.array( - [ - [5, 1, 1, 1], - [0, 1, 1, 1], - [0, 0, 1, 1], - [0, 0, 0, 1], - [0, 0, 0, 1], - [0, 0, 0, 1], - ], - dtype=np.float64, - ) - path = _write_ivcsr(tmp_path, dense) - - result = load_and_normalize( - path, - min_cell_counts=-1.0, - gene_threshold=0.0, - min_cells=3, - ) - ref_X, _, ref_genes = _reference_prepare(dense, -1.0, 0.0, min_cells=3) - - assert list(result.var_names) == [str(i) for i in ref_genes] - assert list(result.var_names) == ["2", "3"] - assert isinstance(result.X, sp.csr_array) - np.testing.assert_allclose(result.X.toarray(), ref_X.toarray(), rtol=1e-5, atol=1e-5) - - -def test_min_cells_uses_obs_filtered_cohort(tmp_path): - import pandas as pd - - dense = np.array( - [ - [1, 1, 1], - [1, 0, 1], - [0, 1, 1], - [1, 0, 1], - [0, 0, 1], - [1, 0, 1], - ], - dtype=np.float64, - ) - obs = pd.DataFrame( - {"condition": ["control", "treated"] * 3}, - index=[f"cell_{i}" for i in range(dense.shape[0])], - ) - adata = ad.AnnData(X=sp.csr_array(dense), obs=obs) - path = tmp_path / "min_cells_obs_filter.ivcsr.h5ad" - VCSCAnnData.from_anndata(adata, format="csr").write_h5ad(path, format="ivcsc") - - obs_mask = np.asarray(obs["condition"] == "control") - result = load_and_normalize( - path, - min_cell_counts=-1.0, - gene_threshold=0.0, - min_cells=2, - obs_filter=lambda obs: obs["condition"] == "control", - ) - ref_X, _, ref_genes = _reference_prepare(dense[obs_mask], -1.0, 0.0, min_cells=2) - - assert list(result.var_names) == [str(i) for i in ref_genes] - assert list(result.var_names) == ["1", "2"] - assert isinstance(result.X, sp.csr_array) - np.testing.assert_allclose(result.X.toarray(), ref_X.toarray(), rtol=1e-5, atol=1e-5) - - -@pytest.mark.parametrize( - ("min_cells", "error", "match"), - [ - (-1, ValueError, "non-negative"), - (1.5, TypeError, "integer or None"), - (True, TypeError, "integer or None"), - ], -) -def test_min_cells_validation(tmp_path, rng, min_cells, error, match): - dense = rng.integers(0, 5, size=(10, 4)).astype(np.float64) - path = _write_ivcsr(tmp_path, dense) - - with pytest.raises(error, match=match): - load_and_normalize(path, min_cells=min_cells) - - -def test_obs_filter_matches_subset_reference(tmp_path): - import pandas as pd - - dense = np.array( - [ - [10, 0, 2], - [0, 100, 0], - [10, 0, 2], - [0, 100, 0], - [10, 0, 2], - [0, 100, 0], - ], - dtype=np.float64, - ) - obs = pd.DataFrame( - {"condition": ["control", "treated"] * 3}, - index=[f"cell_{i}" for i in range(dense.shape[0])], - ) - adata = ad.AnnData(X=sp.csr_array(dense), obs=obs) - path = tmp_path / "obs_filter.ivcsr.h5ad" - VCSCAnnData.from_anndata(adata, format="csr").write_h5ad(path, format="ivcsc") - - obs_mask = np.asarray(obs["condition"] == "control") - result = load_and_normalize( - path, - min_cell_counts=-1.0, - gene_threshold=1.0, - obs_filter=lambda obs: obs["condition"] == "control", - ) - ref_X, _, ref_genes = _reference_prepare(dense[obs_mask], -1.0, 1.0) - - assert list(result.obs_names) == ["cell_0", "cell_2", "cell_4"] - assert list(result.var_names) == [str(i) for i in ref_genes] - assert isinstance(result.X, sp.csr_array) - np.testing.assert_allclose(result.X.toarray(), ref_X.toarray(), rtol=1e-5, atol=1e-5) - - -def test_obs_filter_applied_before_cell_filtering_and_slices_metadata(tmp_path, rng): - import pandas as pd - - n_cells, n_genes = 12, 6 - dense = rng.integers(0, 8, size=(n_cells, n_genes)).astype(np.float64) - dense[rng.random(dense.shape) < 0.4] = 0.0 - dense[3] = 0.0 - - obs = pd.DataFrame( - { - "condition": ["control", "treated", "control"] * 4, - "batch": np.arange(n_cells), - }, - index=[f"cell_{i}" for i in range(n_cells)], - ) - var = pd.DataFrame(index=[f"gene_{j}" for j in range(n_genes)]) - obsm = {"embedding": rng.random((n_cells, 2))} - obsp = {"distances": rng.random((n_cells, n_cells))} - adata = ad.AnnData(X=sp.csr_array(dense), obs=obs, var=var, obsm=obsm, obsp=obsp) - path = tmp_path / "obs_filter_metadata.ivcsr.h5ad" - VCSCAnnData.from_anndata(adata, format="csr").write_h5ad(path, format="ivcsc") - - obs_mask = np.asarray(obs["condition"] == "control") - selected_rows = np.nonzero(obs_mask)[0] - ref_X, ref_cells, ref_genes = _reference_prepare(dense[obs_mask], 5.0, 0.2) - retained_rows = selected_rows[ref_cells] - - result = load_and_normalize( - path, - min_cell_counts=5.0, - gene_threshold=0.2, - obs_filter=lambda obs: obs["condition"] == "control", - ) - - assert list(result.obs_names) == [f"cell_{i}" for i in retained_rows] - assert list(result.var_names) == [f"gene_{j}" for j in ref_genes] - assert list(result.obs["batch"]) == list(retained_rows) - np.testing.assert_allclose( - np.asarray(result.obsm["embedding"]), np.asarray(obsm["embedding"])[retained_rows] - ) - np.testing.assert_allclose( - np.asarray(result.obsp["distances"]), - np.asarray(obsp["distances"])[retained_rows][:, retained_rows], - ) - assert isinstance(result.X, sp.csr_array) - np.testing.assert_allclose(result.X.toarray(), ref_X.toarray(), rtol=1e-5, atol=1e-5) - - -@pytest.mark.parametrize( - ("obs_filter", "match"), - [ - (lambda obs: np.ones(len(obs) - 1, dtype=bool), "mask of length"), - (lambda obs: np.ones(len(obs), dtype=np.int8), "boolean mask"), - (lambda obs: np.zeros(len(obs), dtype=bool), "selected no cells"), - ], -) -def test_obs_filter_validation(tmp_path, rng, obs_filter, match): - dense = rng.integers(0, 5, size=(10, 4)).astype(np.float64) - path = _write_ivcsr(tmp_path, dense) - - with pytest.raises(ValueError, match=match): - load_and_normalize(path, obs_filter=obs_filter) - - -def test_obs_filter_must_be_callable(tmp_path, rng): - dense = rng.integers(0, 5, size=(10, 4)).astype(np.float64) - path = _write_ivcsr(tmp_path, dense) - - with pytest.raises(TypeError, match="must be callable"): - load_and_normalize(path, obs_filter=True) # ty: ignore[invalid-argument-type] - - -def test_obs_filter_requires_obs(tmp_path): - import h5py - - from vsparse import _io - - dense = np.array([[1.0, 0.0], [0.0, 1.0]]) - path = tmp_path / "no_obs.ivcsr.h5ad" - with h5py.File(path, "w") as f: - _io.write_ivcs_elem(f, "X", VCSRArray.from_scipy(sp.csr_array(dense))) - - with pytest.raises(ValueError, match="requires an obs table"): - load_and_normalize(path, obs_filter=lambda obs: np.ones(len(obs), dtype=bool)) - - -def test_metadata_sliced_correctly(tmp_path, rng): - import pandas as pd - - n_cells, n_genes = 60, 20 - dense = rng.integers(0, 8, size=(n_cells, n_genes)).astype(np.float64) - dense[rng.random(dense.shape) < 0.5] = 0.0 - - obs = pd.DataFrame( - {"cell_type": [f"type_{i % 3}" for i in range(n_cells)], "batch": range(n_cells)}, - index=pd.Index([f"cell_{i}" for i in range(n_cells)]), - ) - var = pd.DataFrame( - {"gene_symbol": [f"SYM_{j}" for j in range(n_genes)]}, - index=pd.Index([f"gene_{j}" for j in range(n_genes)]), - ) - obsm = {"pca": rng.random((n_cells, 5))} - varm = {"loadings": rng.random((n_genes, 3))} - obsp = {"distances": rng.random((n_cells, n_cells))} - varp = {"correlations": rng.random((n_genes, n_genes))} - uns = {"dataset_info": "test_metadata", "version": 1} - - adata = ad.AnnData( - X=sp.csr_array(dense), - obs=obs, - var=var, - obsm=obsm, - varm=varm, - obsp=obsp, - varp=varp, - uns=uns, - ) - - path = tmp_path / "metadata_test.ivcsr.h5ad" - vad = VCSCAnnData.from_anndata(adata, format="csr") - vad.write_h5ad(path, format="ivcsc") - - min_cell_counts = 10.0 - gene_threshold = 0.05 - ref_X, ref_cells, ref_genes = _reference_prepare(dense, min_cell_counts, gene_threshold) - - result = load_and_normalize( - path, min_cell_counts=min_cell_counts, gene_threshold=gene_threshold - ) - - assert isinstance(result, ad.AnnData) - assert result.shape == ref_X.shape - - # Check obs/var indices and contents - assert list(result.obs_names) == [f"cell_{i}" for i in ref_cells] - assert list(result.var_names) == [f"gene_{j}" for j in ref_genes] - assert list(result.obs["cell_type"]) == [f"type_{i % 3}" for i in ref_cells] - assert list(result.var["gene_symbol"]) == [f"SYM_{j}" for j in ref_genes] - - # Check obsm/varm - np.testing.assert_allclose(np.asarray(result.obsm["pca"]), np.asarray(obsm["pca"])[ref_cells]) - np.testing.assert_allclose( - np.asarray(result.varm["loadings"]), np.asarray(varm["loadings"])[ref_genes] - ) - - # Check obsp/varp - np.testing.assert_allclose( - np.asarray(result.obsp["distances"]), np.asarray(obsp["distances"])[ref_cells][:, ref_cells] - ) - np.testing.assert_allclose( - np.asarray(result.varp["correlations"]), - np.asarray(varp["correlations"])[ref_genes][:, ref_genes], - ) - - # Check uns - assert result.uns == uns - - # Check X - assert isinstance(result.X, sp.csr_array) - np.testing.assert_allclose(result.X.toarray(), ref_X.toarray(), rtol=1e-5, atol=1e-5) - - -def test_rapid_load_preserves_raw(tmp_path, rng): - """Verify that load_and_normalize restores adata.raw when present in the h5ad file.""" - import h5py - - from vsparse import VCSRArray, _io - - dense = rng.integers(0, 8, size=(40, 20)).astype(np.float64) - adata = ad.AnnData(X=sp.csr_array(dense)) - adata.raw = adata - - path = tmp_path / "raw_test.ivcsr.h5ad" - with h5py.File(path, "w") as f: - _io.write_ivcs_elem(f, "X", VCSRArray.from_scipy(sp.csr_array(dense))) - ad.io.write_elem(f, "obs", adata.obs) - ad.io.write_elem(f, "var", adata.var) - ad.io.write_elem(f, "raw", {"X": sp.csr_array(dense), "var": adata.var}) - - result = load_and_normalize(path, min_cell_counts=-1.0, gene_threshold=0.0) - assert result.raw is not None - - -def test_rapid_load_custom_x_key(tmp_path, rng): - """Verify load_and_normalize reads from a custom top-level group key.""" - import h5py - - from vsparse import VCSRArray, _io - - dense = rng.integers(0, 8, size=(30, 15)).astype(np.float64) - vcsr = VCSRArray.from_scipy(sp.csr_array(dense)) - - path = tmp_path / "custom_key.h5ad" - with h5py.File(path, "w") as f: - _io.write_ivcs_elem(f, "custom_matrix", vcsr) - - result = load_and_normalize( - path, x_key="custom_matrix", min_cell_counts=-1.0, gene_threshold=0.0 - ) - assert isinstance(result, ad.AnnData) - assert result.shape == dense.shape - - def test_load_packed_decodes_x(tmp_path, rng): - """load_packed is the alternative route: no filtering/normalization, X decoded eagerly.""" + """X is decoded eagerly into a plain VCSRArray.""" dense = rng.integers(0, 8, size=(20, 12)).astype(np.float64) dense[rng.random(dense.shape) < 0.4] = 0.0 path = _write_ivcsr(tmp_path, dense)