diff --git a/AGENTS.md b/AGENTS.md index c45acfe..fa84d22 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -61,6 +61,9 @@ Each layer calls only the layers below it: - `concat.py` concatenates on disk; - `recast.py` handles dtype, layout and density changes, including an out-of-core transpose; + - `csr_build.py` builds a CSR matrix from entries in any order, through + row-range buckets on disk. The transpose uses it, and so does anything + that reads entries out of order; - `query.py` and `select.py` hold the expression language behind `--obs-query`. - **`elements/`**: the AnnData on-disk spec. @@ -82,7 +85,9 @@ Each layer calls only the layers below it: other formats. To add a format: - add one `Format` entry in `conversion/__init__.py::FORMATS` with a content sniffer; - - write a converter module (`tenx.py` is the model); + - write a converter module (`tenx.py` is the model). The feature table that + every 10x-style format shares (var names, feature-type splits, var to + id/name/type on export) is in `features.py`; reuse it; - add Cell Ranger-style h5py fixtures in `tests/conftest.py`; - add a scanpy comparison in `tests/test_external_readers.py`. diff --git a/src/adata/conversion/features.py b/src/adata/conversion/features.py new file mode 100644 index 0000000..d7c33b6 --- /dev/null +++ b/src/adata/conversion/features.py @@ -0,0 +1,380 @@ +"""The feature table shared by every 10x-style format. + +A 10x HDF5 matrix and a 10x MTX folder record the same things about each +feature -- an id, a name, a type -- and import and export them the same way. +What is common lives here, so `tenx` and `mtx` differ only in how they read +and write their files: + +* import: which features each output keeps (`plan_targets`), how they are + named (`annotate`), and how obs and var are written (`write_obs_index`, + `write_var`); +* export: which matrix to write (`plan_matrix_source`) and each feature's id, + name and type, worked out from var (`export_features`). + +The AnnData produced matches what scanpy's `read_10x_h5` and `read_10x_mtx` +give with `gex_only=False`. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple + +import numpy as np + +from adata.conversion.base import Options +from adata.core.recast import describe +from adata.elements import spec +from adata.elements.read import ( + dataframe_columns, + element_len, + read_str_all, + resolve_index, +) +from adata.elements.strings import string_dtype_for +from adata.elements.write import ( + write_categorical, + write_dataframe_header, + write_dense, + write_string_array, +) +from adata.storage import create_dataset, is_group, is_zarr_group, open_store +from adata.util.names import count_duplicates, make_index_unique, unique_names + +VAR_NAMES = ("gene_symbols", "gene_ids") + +#: Rows of a string column handled per read when streaming one. +STRING_STEP = 1 << 18 + + +# --------------------------------------------------------------------------- +# import + + +@dataclass +class Target: + """One AnnData to write: which features it keeps, and their annotation.""" + + label: Optional[str] + keep: Optional[np.ndarray] + var_names: List[str] + columns: Dict[str, np.ndarray] + + +def type_counts(types: np.ndarray) -> List[Tuple[str, int]]: + """Feature types in order of first appearance, with how many of each.""" + labels, first, counts = np.unique(types, return_index=True, return_counts=True) + order = np.argsort(first) + return [(str(labels[i]), int(counts[i])) for i in order] + + +def check_import_options(options: Options) -> None: + if options.var_names not in VAR_NAMES: + raise ValueError(f"--var-names must be one of: {', '.join(VAR_NAMES)}.") + if options.feature_type is not None and options.split_feature_types: + raise ValueError( + "--feature-type keeps one type and --split-feature-types writes " + "every type separately; pass one or the other." + ) + + +def annotate( + names: np.ndarray, + columns: Dict[str, np.ndarray], + keep: Optional[np.ndarray], + options: Options, + label: Optional[str], +) -> Target: + """var for the features in `keep`: their names, and their columns.""" + + def pick(values: np.ndarray) -> np.ndarray: + return values if keep is None else values[keep] + + picked = {key: pick(values) for key, values in columns.items()} + if options.var_names == "gene_ids": + # As scanpy's read_10x_mtx does it: ids become the names, and the + # symbols move into a column where the ids were. + index = list(picked["gene_ids"]) + picked = { + ("gene_symbols" if key == "gene_ids" else key): ( + pick(names) if key == "gene_ids" else values + ) + for key, values in picked.items() + } + else: + index = list(pick(names)) + if options.make_unique: + index = make_index_unique(index) + return Target(label=label, keep=keep, var_names=index, columns=picked) + + +def plan_targets( + output: Path, + names: np.ndarray, + columns: Dict[str, np.ndarray], + types: Optional[np.ndarray], + keep: np.ndarray, + options: Options, + *, + noun: str, +) -> Tuple[List[Path], List[Target], List[str]]: + """The outputs to write, the features each keeps, and any warnings. + + One output, unless `--split-feature-types` asks for one per type. `keep` + is what survives filters the caller applied first, such as `--genome`; + `noun` names the input in messages ("file", "folder"). + """ + groups: List[Tuple[Optional[str], np.ndarray]] = [] + if options.feature_type is not None: + available = [t for t, _ in type_counts(types)] + if options.feature_type not in available: + raise ValueError( + f"No features of type {options.feature_type!r}; this {noun} " + f"has: {', '.join(available)}." + ) + groups.append((options.feature_type, keep & (types == options.feature_type))) + elif options.split_feature_types: + for label, _ in type_counts(types[keep]): + groups.append((label, keep & (types == label))) + else: + groups.append((None, keep)) + + stems = unique_names([label for label, _ in groups if label is not None]) + outputs: List[Path] = [] + targets: List[Target] = [] + warnings: List[str] = [] + for label, mask in groups: + path = output + if options.split_feature_types: + path = output.with_name(f"{output.stem}.{stems[label]}{output.suffix}") + selection = None if mask.all() else np.flatnonzero(mask) + target = annotate(names, columns, selection, options, label) + duplicated = count_duplicates(target.var_names) + if duplicated: + where = f" in {path.name}" if options.split_feature_types else "" + warnings.append( + f"{duplicated:,} var names{where} repeat an earlier one. " + "Pass --make-unique to suffix them as anndata's " + "var_names_make_unique does." + ) + outputs.append(path) + targets.append(target) + return outputs, targets, warnings + + +def describe_import( + n_barcodes: int, + n_features: int, + types: Optional[np.ndarray], + targets: Sequence[Target], +) -> str: + """`1,222 barcodes x 33,538 features (Gene Expression 33,538)`, and so on.""" + described = f"{n_barcodes:,} barcodes x {n_features:,} features" + if types is not None: + described += " (" + ", ".join(f"{t} {n:,}" for t, n in type_counts(types)) + ")" + if len(targets) > 1: + described += f", into {len(targets)} outputs by feature type" + elif targets[0].keep is not None: + described += f", keeping {len(targets[0].keep):,} features" + return described + + +def write_obs_index(root: Any, n: int, blocks: Iterable[Sequence[str]]) -> Any: + """obs, holding only its index: the barcodes, written a block at a time. + + Returns the group, for the caller to add columns to. A source that gives + more or fewer than `n` barcodes is refused. + """ + obs = root.create_group("obs") + spec.set_encoding(obs, spec.DATAFRAME) + obs.attrs["_index"] = "_index" + + backend = "zarr" if is_zarr_group(root) else "hdf5" + index = create_dataset(obs, "_index", shape=(n,), dtype=string_dtype_for(backend)) + written = 0 + for block in blocks: + if written + len(block) > n: + raise ValueError(f"There are more than {n:,} barcodes.") + index[written : written + len(block)] = np.asarray(block, dtype=object) + written += len(block) + if written != n: + raise ValueError(f"There are {written:,} barcodes, not {n:,}.") + spec.set_encoding(index, spec.STRING_ARRAY) + return obs + + +def write_var(root: Any, target: Target) -> None: + """var, with string columns encoded as anndata's writer would encode them. + + anndata turns any string column with a repeated value into a categorical + when it writes (`strings_to_categoricals`), so an h5ad scanpy wrote from + the same file holds `feature_types` and `genome` as categoricals. Doing + the same keeps the two interchangeable, and those columns a few bytes + instead of a string per feature. + """ + var = write_dataframe_header(root, "var", target.var_names, list(target.columns)) + for key, values in target.columns.items(): + if values.dtype.kind == "b": + write_dense(var, key, values) + continue + categories, codes = np.unique(values, return_inverse=True) + if len(categories) < len(values): + write_categorical(var, key, codes, list(categories)) + else: + write_string_array(var, key, values) + + +# --------------------------------------------------------------------------- +# export + + +def matrix_path(layer: Optional[str]) -> str: + """`--layer` as a path: X, raw/X, a layer's name, or layers/NAME.""" + if layer is None or layer in ("X", "raw/X") or layer.startswith("layers/"): + return layer or "X" + return f"layers/{layer}" + + +def resolve(root: Any, path: str) -> Any: + node = root + for part in path.split("/"): + if not is_group(node) or part not in node: + raise KeyError(f"{path!r} not found in the store.") + node = node[part] + return node + + +@dataclass +class ExportFeatures: + ids: List[str] + names: List[str] + types: List[str] + genomes: List[str] + tags: Dict[str, List[str]] + + +@dataclass +class MatrixSource: + """The matrix an export writes, and what is known about it before writing.""" + + path: str + n_obs: int + n_var: int + dtype: np.dtype + features: ExportFeatures + + +def plan_matrix_source( + source: Path, options: Options, warnings: List[str] +) -> MatrixSource: + """Check the matrix `--layer` names against obs and var; read var's features. + + Warns when the matrix holds floats, which 10x formats are not expected to. + """ + path = matrix_path(options.layer) + with open_store(source, "r", require_anndata=False) as store: + root = store.root + matrix = describe(resolve(root, path)) + var_path = "raw/var" if path.startswith("raw/") else "var" + var, obs = resolve(root, var_path), resolve(root, "obs") + obs_index, _ = resolve_index(obs, "obs") + var_index, _ = resolve_index(var, "var") + n_obs, n_var = element_len(obs_index), element_len(var_index) + if matrix.shape != (n_obs, n_var): + raise ValueError( + f"{path} is {matrix.shape[0]} x {matrix.shape[1]}, but obs has " + f"{n_obs} and {var_path} has {n_var} entries." + ) + features = export_features(root, var, var_index, options, warnings) + + if matrix.dtype.kind == "f": + warnings.append( + f"{path} holds {matrix.dtype} values. Cell Ranger writes integer " + "counts, and tools that read its output may expect them: export " + "a counts layer with --layer, or recast to int32 first." + ) + return MatrixSource( + path=path, n_obs=n_obs, n_var=n_var, dtype=matrix.dtype, features=features + ) + + +def export_features( + root: Any, + var: Any, + var_index: Any, + options: Options, + warnings: List[str], +) -> ExportFeatures: + """Work out each feature's id, name, type and genome from var.""" + columns = dataframe_columns(var) + n = element_len(var_index) + + def column(name: str) -> List[str]: + return read_str_all(var[name], parent_group=var) + + def requested(flag: str, name: Optional[str]) -> Optional[str]: + if name is not None and name not in columns: + raise ValueError( + f"{flag} {name}: no such var column. It has: " + f"{', '.join(columns) or 'none'}." + ) + return name + + ids_col = requested("--gene-ids-column", options.gene_ids_column) + names_col = requested("--gene-names-column", options.gene_names_column) + types_col = requested("--feature-types-column", options.feature_types_column) + index = read_str_all(var_index) + + missing = [] + if ids_col is None and names_col is None and "gene_ids" not in columns and "gene_symbols" in columns: + # var_names are ids already, as `--var-names gene_ids` leaves them. + ids, names = index, column("gene_symbols") + else: + names = column(names_col) if names_col else index + if ids_col is None and "gene_ids" in columns: + ids_col = "gene_ids" + if ids_col is not None: + ids = column(ids_col) + else: + missing.append("gene ids (a gene_ids column, or --gene-ids-column)") + ids = names + if types_col is None and "feature_types" in columns: + types_col = "feature_types" + if types_col is not None: + types = column(types_col) + else: + missing.append("feature types (a feature_types column, or --feature-types-column)") + types = ["Gene Expression"] * n + + if missing: + if not options.force: + raise ValueError( + f"var has no {' and no '.join(missing)}, which a 10x matrix " + "records for every feature. Pass --force to write placeholders: " + "the names as ids, and 'Gene Expression' as the type." + ) + warnings.append(f"Writing placeholders for {' and '.join(missing)}.") + + genomes = column("genome") if "genome" in columns else [""] * n + tags: Dict[str, List[str]] = {} + for key in recorded_tag_keys(root): + if key != "genome" and key in columns: + tags[key] = column(key) + return ExportFeatures(ids=ids, names=names, types=types, genomes=genomes, tags=tags) + + +def cellranger_meta(root: Any) -> Optional[Any]: + """`uns/cellranger`, as an import recorded it, if there is one.""" + if "uns" not in root or "cellranger" not in root["uns"]: + return None + meta = root["uns"]["cellranger"] + return meta if is_group(meta) else None + + +def recorded_tag_keys(root: Any) -> List[str]: + meta = cellranger_meta(root) + if meta is None or "feature_tag_keys" not in meta: + return [] + return read_str_all(meta["feature_tag_keys"]) + diff --git a/src/adata/conversion/tenx.py b/src/adata/conversion/tenx.py index 40c5dc2..edd0f1d 100644 --- a/src/adata/conversion/tenx.py +++ b/src/adata/conversion/tenx.py @@ -32,23 +32,32 @@ from rich.console import Console from adata.conversion.base import ConversionPlan, Converter, Options +from adata.conversion.features import ( + STRING_STEP, + ExportFeatures, + Target, + cellranger_meta, + check_import_options, + describe_import, + plan_matrix_source, + plan_targets, + resolve, + write_obs_index, + write_var, +) from adata.core.recast import describe, sparsify, transpose_sparse_streaming from adata.core.subset import fan_out_sparse_matrix from adata.elements import spec from adata.elements.read import ( - dataframe_columns, decode_str_array, element_len, read_str_all, read_str_chunk, resolve_index, ) -from adata.elements.strings import string_dtype_for from adata.elements.write import ( ensure_anndata_skeleton, set_column_order, - write_categorical, - write_dataframe_header, write_dense, write_mapping, write_scalar, @@ -57,14 +66,10 @@ from adata.storage import ( GZIP_LEVEL, copy_dataset, - create_dataset, has_valid_anndata_root_attrs, is_dataset, - is_group, - is_zarr_group, open_store, ) -from adata.util.names import count_duplicates, make_index_unique, unique_names #: What a group must hold to be a matrix, in each layout. _V3_KEYS = frozenset({"barcodes", "data", "indices", "indptr", "shape", "features"}) @@ -87,10 +92,6 @@ "version", ) -#: Rows of a string column handled per read when streaming one. -_STRING_STEP = 1 << 18 - -VAR_NAMES = ("gene_symbols", "gene_ids") # --------------------------------------------------------------------------- @@ -260,23 +261,6 @@ def __getitem__(self, key: str) -> Any: # import -@dataclass -class _Output: - """One AnnData to write: which features it keeps, and their annotation.""" - - label: Optional[str] - keep: Optional[np.ndarray] - var_names: List[str] - columns: Dict[str, np.ndarray] - - -def _type_counts(types: np.ndarray) -> List[Tuple[str, int]]: - """Feature types in order of first appearance, with how many of each.""" - labels, first, counts = np.unique(types, return_index=True, return_counts=True) - order = np.argsort(first) - return [(str(labels[i]), int(counts[i])) for i in order] - - class TenxImport(Converter): options = frozenset( {"--genome", "--var-names", "--make-unique", "--feature-type", @@ -291,13 +275,7 @@ def plan( dst_format: str, options: Options, ) -> ConversionPlan: - if options.var_names not in VAR_NAMES: - raise ValueError(f"--var-names must be one of: {', '.join(VAR_NAMES)}.") - if options.feature_type is not None and options.split_feature_types: - raise ValueError( - "--feature-type keeps one type and --split-feature-types writes " - "every type separately; pass one or the other." - ) + check_import_options(options) with h5py.File(source, "r") as handle: matrix = _read_matrix(handle, options.genome) @@ -324,49 +302,12 @@ def plan( keep &= genomes == options.genome types = matrix.feature_types - groups: List[Tuple[Optional[str], np.ndarray]] = [] - if options.feature_type is not None: - available = [t for t, _ in _type_counts(types)] - if options.feature_type not in available: - raise ValueError( - f"No features of type {options.feature_type!r}; this file " - f"has: {', '.join(available)}." - ) - groups.append((options.feature_type, keep & (types == options.feature_type))) - elif options.split_feature_types: - for label, _ in _type_counts(types[keep]): - groups.append((label, keep & (types == label))) - else: - groups.append((None, keep)) - - stems = unique_names([label for label, _ in groups if label is not None]) - outputs: List[Path] = [] - targets: List[_Output] = [] - warnings: List[str] = [] - for label, mask in groups: - path = output - if options.split_feature_types: - path = output.with_name(f"{output.stem}.{stems[label]}{output.suffix}") - selection = None if mask.all() else np.flatnonzero(mask) - target = self._annotate(matrix, selection, options, label) - duplicated = count_duplicates(target.var_names) - if duplicated: - where = f" in {path.name}" if options.split_feature_types else "" - warnings.append( - f"{duplicated:,} var names{where} repeat an earlier one. " - "Pass --make-unique to suffix them as anndata's " - "var_names_make_unique does." - ) - outputs.append(path) - targets.append(target) - - described = f"{matrix.n_barcodes:,} barcodes x {matrix.n_features:,} features" - if types is not None: - described += " (" + ", ".join(f"{t} {n:,}" for t, n in _type_counts(types)) + ")" - if len(outputs) > 1: - described += f", into {len(outputs)} outputs by feature type" - elif targets[0].keep is not None: - described += f", keeping {len(targets[0].keep):,} features" + outputs, targets, warnings = plan_targets( + output, matrix.names, matrix.columns, types, keep, options, noun="file" + ) + described = describe_import( + matrix.n_barcodes, matrix.n_features, types, targets + ) return ConversionPlan( source=source, @@ -383,39 +324,12 @@ def plan( }, ) - @staticmethod - def _annotate( - matrix: _Matrix, - keep: Optional[np.ndarray], - options: Options, - label: Optional[str], - ) -> _Output: - def pick(values: np.ndarray) -> np.ndarray: - return values if keep is None else values[keep] - - columns = {key: pick(values) for key, values in matrix.columns.items()} - if options.var_names == "gene_ids": - # As scanpy's read_10x_mtx does it: ids become the names, and the - # symbols move into a column where the ids were. - index = list(columns["gene_ids"]) - columns = { - ("gene_symbols" if key == "gene_ids" else key): ( - pick(matrix.names) if key == "gene_ids" else values - ) - for key, values in columns.items() - } - else: - index = list(pick(matrix.names)) - if options.make_unique: - index = make_index_unique(index) - return _Output(label=label, keep=keep, var_names=index, columns=columns) - def write( self, plan: ConversionPlan, staged: Sequence[Path], console: Console ) -> None: detail = plan.detail matrix: _Matrix = detail["matrix"] - targets: List[_Output] = detail["targets"] + targets: List[Target] = detail["targets"] with h5py.File(plan.source, "r") as handle, ExitStack() as stack: group = handle[matrix.group] @@ -428,7 +342,7 @@ def write( for root, target in zip(roots, targets): with console.status("[cyan]Writing obs and var...[/]"): _write_obs(group, root, detail["filtered_barcodes"]) - _write_var(root, target) + write_var(root, target) ensure_anndata_skeleton(root) _write_uns(root, matrix) @@ -443,18 +357,16 @@ def write( def _write_obs(group: h5py.Group, root: Any, filtered_barcodes: bool) -> None: """obs: the barcodes, streamed, and Space Ranger's `filtered_barcodes`.""" - obs = root.create_group("obs") - spec.set_encoding(obs, spec.DATAFRAME) - obs.attrs["_index"] = "_index" - barcodes = group["barcodes"] n = barcodes.shape[0] - backend = "zarr" if is_zarr_group(root) else "hdf5" - index = create_dataset(obs, "_index", shape=(n,), dtype=string_dtype_for(backend)) - for lo in range(0, n, _STRING_STEP): - hi = min(lo + _STRING_STEP, n) - index[lo:hi] = np.asarray(_strings_slice(barcodes, lo, hi), dtype=object) - spec.set_encoding(index, spec.STRING_ARRAY) + obs = write_obs_index( + root, + n, + ( + _strings_slice(barcodes, lo, min(lo + STRING_STEP, n)) + for lo in range(0, n, STRING_STEP) + ), + ) columns = [] if filtered_barcodes: @@ -467,27 +379,6 @@ def _strings_slice(dataset: Any, lo: int, hi: int) -> List[str]: return [str(v) for v in decode_str_array(np.asarray(dataset[lo:hi]))] -def _write_var(root: Any, target: _Output) -> None: - """var, with string columns encoded as anndata's writer would encode them. - - anndata turns any string column with a repeated value into a categorical - when it writes (`strings_to_categoricals`), so an h5ad scanpy wrote from - the same file holds `feature_types` and `genome` as categoricals. Doing - the same keeps the two interchangeable, and those columns a few bytes - instead of a string per feature. - """ - var = write_dataframe_header(root, "var", target.var_names, list(target.columns)) - for key, values in target.columns.items(): - if values.dtype.kind == "b": - write_dense(var, key, values) - continue - categories, codes = np.unique(values, return_inverse=True) - if len(categories) < len(values): - write_categorical(var, key, codes, list(categories)) - else: - write_string_array(var, key, values) - - def _write_uns(root: Any, matrix: _Matrix) -> None: """The run's metadata, which scanpy drops and an export puts back.""" meta = write_mapping(root["uns"], "cellranger") @@ -521,22 +412,6 @@ def _write_value(parent: Any, key: str, value: Any) -> None: # export -def _matrix_path(layer: Optional[str]) -> str: - """`--layer` as a path: X, raw/X, a layer's name, or layers/NAME.""" - if layer is None or layer in ("X", "raw/X") or layer.startswith("layers/"): - return layer or "X" - return f"layers/{layer}" - - -def _resolve(root: Any, path: str) -> Any: - node = root - for part in path.split("/"): - if not is_group(node) or part not in node: - raise KeyError(f"{path!r} not found in the store.") - node = node[part] - return node - - def _read_value(obj: Any) -> Any: """An AnnData element read back as the attribute value it came from.""" enc = spec.encoding_type(obj) @@ -552,15 +427,6 @@ def _read_value(obj: Any) -> Any: return value -@dataclass -class _Features: - ids: List[str] - names: List[str] - types: List[str] - genomes: List[str] - tags: Dict[str, List[str]] - - class TenxExport(Converter): options = frozenset( {"--layer", "--gene-ids-column", "--gene-names-column", "--feature-types-column"} @@ -574,31 +440,11 @@ def plan( dst_format: str, options: Options, ) -> ConversionPlan: - path = _matrix_path(options.layer) warnings: List[str] = [] + matrix = plan_matrix_source(source, options, warnings) with open_store(source, "r", require_anndata=False) as store: - root = store.root - matrix = describe(_resolve(root, path)) - var_path = "raw/var" if path.startswith("raw/") else "var" - var, obs = _resolve(root, var_path), _resolve(root, "obs") - obs_index, _ = resolve_index(obs, "obs") - var_index, index_name = resolve_index(var, "var") - n_obs, n_var = element_len(obs_index), element_len(var_index) - if matrix.shape != (n_obs, n_var): - raise ValueError( - f"{path} is {matrix.shape[0]} x {matrix.shape[1]}, but obs has " - f"{n_obs} and {var_path} has {n_var} entries." - ) - - features = _features(root, var, var_index, options, warnings) - attrs = _run_attrs(root) - - if matrix.dtype.kind == "f": - warnings.append( - f"{path} holds {matrix.dtype} values. Cell Ranger writes integer " - "counts, and tools that read its output may expect them: export " - "a counts layer with --layer, or recast to int32 first." - ) + attrs = _run_attrs(store.root) + path, n_obs, n_var = matrix.path, matrix.n_obs, matrix.n_var return ConversionPlan( source=source, @@ -610,19 +456,19 @@ def plan( f"from {path}" ), warnings=warnings, - detail={"path": path, "features": features, "attrs": attrs}, + detail={"path": path, "features": matrix.features, "attrs": attrs}, ) def write( self, plan: ConversionPlan, staged: Sequence[Path], console: Console ) -> None: detail = plan.detail - features: _Features = detail["features"] + features: ExportFeatures = detail["features"] with open_store(plan.source, "r", require_anndata=False) as store, h5py.File( staged[0], "w" ) as out: root = store.root - matrix = describe(_resolve(root, detail["path"])) + matrix = describe(resolve(root, detail["path"])) with console.status(f"[cyan]Copying {detail['path']}...[/]"): group = _write_matrix(matrix, out) n_obs, n_var = matrix.shape @@ -647,90 +493,10 @@ def write( out.attrs[key] = value -def _features( - root: Any, - var: Any, - var_index: Any, - options: Options, - warnings: List[str], -) -> _Features: - """Work out each feature's id, name, type and genome from var.""" - columns = dataframe_columns(var) - n = element_len(var_index) - - def column(name: str) -> List[str]: - return read_str_all(var[name], parent_group=var) - - def requested(flag: str, name: Optional[str]) -> Optional[str]: - if name is not None and name not in columns: - raise ValueError( - f"{flag} {name}: no such var column. It has: " - f"{', '.join(columns) or 'none'}." - ) - return name - - ids_col = requested("--gene-ids-column", options.gene_ids_column) - names_col = requested("--gene-names-column", options.gene_names_column) - types_col = requested("--feature-types-column", options.feature_types_column) - index = read_str_all(var_index) - - missing = [] - if ids_col is None and names_col is None and "gene_ids" not in columns and "gene_symbols" in columns: - # var_names are ids already, as `--var-names gene_ids` leaves them. - ids, names = index, column("gene_symbols") - else: - names = column(names_col) if names_col else index - if ids_col is None and "gene_ids" in columns: - ids_col = "gene_ids" - if ids_col is not None: - ids = column(ids_col) - else: - missing.append("gene ids (a gene_ids column, or --gene-ids-column)") - ids = names - if types_col is None and "feature_types" in columns: - types_col = "feature_types" - if types_col is not None: - types = column(types_col) - else: - missing.append("feature types (a feature_types column, or --feature-types-column)") - types = ["Gene Expression"] * n - - if missing: - if not options.force: - raise ValueError( - f"var has no {' and no '.join(missing)}, which a 10x matrix " - "records for every feature. Pass --force to write placeholders: " - "the names as ids, and 'Gene Expression' as the type." - ) - warnings.append(f"Writing placeholders for {' and '.join(missing)}.") - - genomes = column("genome") if "genome" in columns else [""] * n - tags: Dict[str, List[str]] = {} - for key in _recorded_tag_keys(root): - if key != "genome" and key in columns: - tags[key] = column(key) - return _Features(ids=ids, names=names, types=types, genomes=genomes, tags=tags) - - -def _cellranger_meta(root: Any) -> Optional[Any]: - """`uns/cellranger`, as an import recorded it, if there is one.""" - if "uns" not in root or "cellranger" not in root["uns"]: - return None - meta = root["uns"]["cellranger"] - return meta if is_group(meta) else None - - -def _recorded_tag_keys(root: Any) -> List[str]: - meta = _cellranger_meta(root) - if meta is None or "feature_tag_keys" not in meta: - return [] - return read_str_all(meta["feature_tag_keys"]) - - def _run_attrs(root: Any) -> Dict[str, Any]: """Root attributes: the recorded run's, with Cell Ranger's format markers.""" attrs: Dict[str, Any] = {} - meta = _cellranger_meta(root) + meta = cellranger_meta(root) if meta is not None and "attrs" in meta: for key in _RUN_ATTRS: if key in meta["attrs"] and is_dataset(meta["attrs"][key]): @@ -802,11 +568,11 @@ def _write_fixed_strings_streamed( n = element_len(column) cache: Dict[str, np.ndarray] = {} width = 1 - for lo in range(0, n, _STRING_STEP): - block = read_str_chunk(column, lo, min(lo + _STRING_STEP, n), cache, frame) + for lo in range(0, n, STRING_STEP): + block = read_str_chunk(column, lo, min(lo + STRING_STEP, n), cache, frame) width = max(width, _fixed_width(block)) dataset = _create_fixed(parent, name, n, width) - for lo in range(0, n, _STRING_STEP): - hi = min(lo + _STRING_STEP, n) + for lo in range(0, n, STRING_STEP): + hi = min(lo + STRING_STEP, n) block = read_str_chunk(column, lo, hi, cache, frame) dataset[lo:hi] = np.array([v.encode("utf-8") for v in block], dtype=f"S{width}") diff --git a/src/adata/core/csr_build.py b/src/adata/core/csr_build.py new file mode 100644 index 0000000..3c38cb7 --- /dev/null +++ b/src/adata/core/csr_build.py @@ -0,0 +1,178 @@ +"""Building a compressed sparse matrix from entries that arrive out of order. + +A streaming transpose and an MTX import have the same problem: entries come +in an order unrelated to the rows they belong to, and the output cannot be +written until each row is complete. Both solve it the same way, with memory +set by the bucket size rather than by nnz: + +1. count the entries per output row (`count_keys`), which gives `indptr`; +2. scatter the entries into buckets, each covering a slice of the output rows + and holding about as many entries as one read (`bucket_bounds`, + `BucketSorter.add`), kept on disk in a scratch group (`scratch_group`); +3. read one bucket at a time, sort it, and append it to the output + (`BucketSorter.runs`). +""" + +from __future__ import annotations + +import os +import tempfile +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Iterator, List, Optional, Tuple + +import h5py +import numpy as np + +from adata.storage import is_hdf5_group + +#: Smallest bucket worth making. Below this the per-bucket overhead +#: dominates and the extra passes buy nothing. +MIN_BUCKET_ENTRIES = 1 << 12 + + +def count_keys(counts: np.ndarray, keys: np.ndarray) -> None: + """Add one to `counts[k]` for every `k` in `keys`. + + `np.bincount(keys, minlength=len(counts))` costs the length of `counts` + on every call, which for a block of a few thousand entries against + millions of rows is almost all waste -- and quadratic over a matrix read + in many blocks. Count over the span the block actually covers instead, + or by sorting when that span is sparse. + """ + if keys.size == 0: + return + lo, hi = int(keys.min()), int(keys.max()) + span = hi - lo + 1 + if span <= 4 * keys.size: + counts[lo : hi + 1] += np.bincount(keys - lo, minlength=span) + else: + unique, seen = np.unique(keys, return_counts=True) + counts[unique] += seen + + +def bucket_bounds( + indptr: np.ndarray, *, chunk: int, bucket_entries: Optional[int] = None +) -> np.ndarray: + """Row bounds splitting the output into buckets of about equal nnz. + + One bucket per slice of the output rows, holding about as many nonzeros + as one read. Tying this to `chunk` rather than a constant is what makes + the peak follow the setting the caller chose: with a fixed bucket size, + anything below it went into a single bucket and the "streaming" path + quietly held the whole matrix. An explicit `bucket_entries` is honoured + as given; the floor applies only to the value derived from `chunk`, + where it stops a tiny chunk producing thousands of buckets whose + overhead outweighs the saving. + + Split by nonzero count, not by coordinate. Equal-width bounds put nearly + everything in one bucket whenever the matrix is skewed -- and single-cell + matrices are: a handful of genes carry most of the counts. Measured on + one such matrix, the largest of three equal-width buckets held 88% of + the entries, so the bucket, not the chunk, set the peak. + """ + n_rows = int(indptr.size - 1) + nnz = int(indptr[-1]) if indptr.size else 0 + per_bucket = ( + max(1, int(bucket_entries)) + if bucket_entries + else max(MIN_BUCKET_ENTRIES, int(chunk)) + ) + n_buckets = max(1, int(np.ceil(nnz / per_bucket))) if nnz else 1 + n_buckets = min(n_buckets, n_rows) or 1 + + targets = np.linspace(0, nnz, n_buckets + 1)[1:-1] + bounds = np.concatenate(( + [0], + np.searchsorted(indptr, targets, side="left").astype(np.int64), + [n_rows], + )) + return np.unique(bounds) + + +@contextmanager +def scratch_group(dst_parent: Any, name: str) -> Iterator[Any]: + """A group for intermediate buckets, gone afterwards. + + HDF5 never returns the space of a deleted dataset to the file, so scratch + written into an HDF5 output stayed in it as dead space: a streamed + transpose came out three times the size of an in-memory one. On HDF5 the + scratch is a separate file beside the output instead. Deleting from a + Zarr store frees the space, so there it stays in place. + """ + if not is_hdf5_group(dst_parent): + group = dst_parent.create_group(name) + try: + yield group + finally: + del dst_parent[name] + return + + # Beside the output rather than in the system temp directory, which on + # a cluster node is often far too small for a matrix's worth of entries. + where = Path(str(dst_parent.file.filename)).parent + fd, path = tempfile.mkstemp( + prefix=".adata-", suffix=".transpose-scratch.h5", + dir=str(where) if where.is_dir() else None, + ) + os.close(fd) + try: + with h5py.File(path, "w") as handle: + yield handle.create_group(name) + finally: + os.unlink(path) + + +class BucketSorter: + """Scatter (row, column, value) entries into row-range buckets on disk. + + `add` takes entries in any order; `runs` gives them back one bucket at a + time, sorted by row and then column, buckets in row order. Concatenated, + the runs are the matrix in CSR order. + """ + + def __init__(self, scratch: Any, bounds: np.ndarray, value_dtype: Any) -> None: + # Looked up through the module rather than imported by name, so a + # test that watches `_append` sees every bucket write. + from adata.core import subset + + self._subset = subset + self.bounds = np.asarray(bounds, dtype=np.int64) + self.n_buckets = len(self.bounds) - 1 + self.buckets: List[Tuple[Any, Any, Any]] = [ + ( + subset._growable(scratch, f"major{b}", np.int64), + subset._growable(scratch, f"minor{b}", np.int64), + subset._growable(scratch, f"value{b}", value_dtype), + ) + for b in range(self.n_buckets) + ] + + def add(self, major: np.ndarray, minor: np.ndarray, values: np.ndarray) -> None: + if major.size == 0: + return + which = np.clip( + np.searchsorted(self.bounds, major, side="right") - 1, + 0, + self.n_buckets - 1, + ) + # One stable sort by bucket, then a slice per bucket the block + # touches, rather than a full-block comparison for every bucket. + order = np.argsort(which, kind="stable") + starts = np.searchsorted(which[order], np.arange(self.n_buckets + 1)) + for b in np.flatnonzero(np.diff(starts)): + pick = order[starts[b] : starts[b + 1]] + majors, minors, vals = self.buckets[b] + self._subset._append(majors, major[pick]) + self._subset._append(minors, minor[pick]) + self._subset._append(vals, values[pick]) + + def runs(self) -> Iterator[Tuple[np.ndarray, np.ndarray, np.ndarray]]: + for majors, minors, vals in self.buckets: + if majors.shape[0] == 0: + continue + major = np.asarray(majors[...], dtype=np.int64) + minor = np.asarray(minors[...], dtype=np.int64) + values = np.asarray(vals[...]) + order = np.lexsort((minor, major)) + yield major[order], minor[order], values[order] diff --git a/src/adata/core/recast.py b/src/adata/core/recast.py index 850dab8..7c227fa 100644 --- a/src/adata/core/recast.py +++ b/src/adata/core/recast.py @@ -23,19 +23,20 @@ from __future__ import annotations -import os -import tempfile -from contextlib import contextmanager from dataclasses import dataclass -from pathlib import Path -from typing import Any, Dict, Iterator, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple -import h5py import numpy as np from rich.console import Console from adata.elements import spec from adata.elements.write import set_shape_attr +from adata.core.csr_build import ( + BucketSorter, + bucket_bounds, + count_keys, + scratch_group, +) from adata.core.subset import _growable_like, _sparse_dataset from adata.storage import ( copy_attrs, @@ -44,7 +45,6 @@ fit_chunks, is_dataset, is_group, - is_hdf5_group, is_zarr_group, with_compression, zarr_format_of, @@ -74,9 +74,6 @@ #: Values per read while streaming `data`/`indices`. DEFAULT_CHUNK = 1 << 20 -#: Smallest bucket worth making during a streaming transpose. Below this -#: the per-bucket overhead dominates and the extra passes buy nothing. -MIN_BUCKET_ENTRIES = 1 << 12 def parse_dtype(text: str, *, allowed: Tuple[str, ...] = DTYPES) -> np.dtype: @@ -414,37 +411,8 @@ def transpose_sparse_in_memory( ) -@contextmanager -def _scratch_group(dst_parent: Any, name: str) -> Iterator[Any]: - """A group for the transpose's intermediate buckets, gone afterwards. - - HDF5 never returns the space of a deleted dataset to the file, so scratch - written into an HDF5 output stayed in it as dead space: a streamed - transpose came out three times the size of an in-memory one. On HDF5 the - scratch is a separate file beside the output instead. Deleting from a - Zarr store frees the space, so there it stays in place. - """ - if not is_hdf5_group(dst_parent): - group = dst_parent.create_group(name) - try: - yield group - finally: - del dst_parent[name] - return - - # Beside the output rather than in the system temp directory, which on - # a cluster node is often far too small for a matrix's worth of entries. - where = Path(str(dst_parent.file.filename)).parent - fd, path = tempfile.mkstemp( - prefix=".adata-", suffix=".transpose-scratch.h5", - dir=str(where) if where.is_dir() else None, - ) - os.close(fd) - try: - with h5py.File(path, "w") as handle: - yield handle.create_group(name) - finally: - os.unlink(path) +#: Kept under its old name: the scratch group moved to `core.csr_build`. +_scratch_group = scratch_group def transpose_sparse_streaming( @@ -476,8 +444,6 @@ def transpose_sparse_streaming( Costs about two extra passes over nnz, which is the price of not holding the matrix. `--in-memory` is there for when you would rather pay in RAM. """ - from adata.core.subset import _append, _growable - target = ( spec.CSC_MATRIX if src.kind == spec.CSR_MATRIX else spec.CSR_MATRIX ) @@ -492,39 +458,10 @@ def transpose_sparse_streaming( counts = np.zeros(n_major_out, dtype=np.int64) for start in range(0, nnz, chunk): block = np.asarray(source_indices[start : min(start + chunk, nnz)], dtype=np.int64) - counts += np.bincount(block, minlength=n_major_out) + count_keys(counts, block) out_indptr = np.concatenate(([0], np.cumsum(counts))) - # One bucket per slice of the output major axis, holding about as many - # nonzeros as one read. Tying this to `chunk` rather than a constant is - # what makes the peak follow the setting the caller chose: with a fixed - # bucket size, anything below it went into a single bucket and the - # "streaming" path quietly held the whole matrix. - # An explicit request is honoured as given; the floor applies only to - # the value derived from `chunk`, where it stops a tiny chunk producing - # thousands of buckets whose overhead outweighs the saving. - per_bucket = ( - max(1, int(bucket_entries)) - if bucket_entries - else max(MIN_BUCKET_ENTRIES, int(chunk)) - ) - n_buckets = max(1, int(np.ceil(nnz / per_bucket))) if nnz else 1 - n_buckets = min(n_buckets, n_major_out) or 1 - - # Split by nonzero count, not by coordinate. Equal-width bounds put - # nearly everything in one bucket whenever the matrix is skewed -- and - # single-cell matrices are: a handful of genes carry most of the - # counts. Measured on one such matrix, the largest of three equal-width - # buckets held 88% of the entries, so the bucket, not the chunk, set - # the peak. `counts` is already to hand from pass 1. - cumulative = out_indptr - targets = np.linspace(0, nnz, n_buckets + 1)[1:-1] - bounds = np.concatenate(( - [0], - np.searchsorted(cumulative, targets, side="left").astype(np.int64), - [n_major_out], - )) - bounds = np.unique(bounds) + bounds = bucket_bounds(out_indptr, chunk=chunk, bucket_entries=bucket_entries) n_buckets = len(bounds) - 1 if console is not None and n_buckets > 1: console.print( @@ -542,15 +479,8 @@ def transpose_sparse_streaming( ) written = 0 - with _scratch_group(dst_parent, f"__{name}_transpose_scratch__") as scratch: - buckets = [ - ( - _growable(scratch, f"major{b}", np.int64), - _growable(scratch, f"minor{b}", np.int64), - _growable(scratch, f"value{b}", src.dtype), - ) - for b in range(n_buckets) - ] + with scratch_group(dst_parent, f"__{name}_transpose_scratch__") as scratch: + sorter = BucketSorter(scratch, bounds, src.dtype) # Pass 2: scatter the input into buckets, one input block at a time. # Step over the major axis in blocks holding about `chunk` nonzeros, @@ -568,29 +498,16 @@ def transpose_sparse_streaming( major = np.repeat( np.arange(lo, hi, dtype=np.int64), np.diff(indptr[lo : hi + 1]) ) - which = np.clip(np.searchsorted(bounds, minor, side="right") - 1, 0, n_buckets - 1) - for b in range(n_buckets): - pick = which == b - if not pick.any(): - continue - _append(buckets[b][0], minor[pick]) - _append(buckets[b][1], major[pick]) - _append(buckets[b][2], values[pick]) + # The input's minor axis is the output's major one. + sorter.add(minor, major, values) # Pass 3: each bucket in turn, sorted into output order. - for b in range(n_buckets): - new_major, new_minor, value = buckets[b] - if new_major.shape[0] == 0: - continue - majors = np.asarray(new_major[...], dtype=np.int64) - minors = np.asarray(new_minor[...], dtype=np.int64) - values = np.asarray(value[...]) - order = np.lexsort((minors, majors)) - count = order.size - out_indices[written : written + count] = minors[order].astype( + for _, new_minor, values in sorter.runs(): + count = new_minor.size + out_indices[written : written + count] = new_minor.astype( index_dtype, copy=False ) - out_data[written : written + count] = values[order].astype( + out_data[written : written + count] = values.astype( data_dtype, copy=False ) written += count diff --git a/src/adata/storage/__init__.py b/src/adata/storage/__init__.py index 148bd9f..c2c0ee8 100644 --- a/src/adata/storage/__init__.py +++ b/src/adata/storage/__init__.py @@ -15,6 +15,8 @@ import numpy as np +from adata.util.path import is_zarr_path # noqa: F401 - re-exported + ROOT_ENCODING_TYPE = "anndata" ROOT_ENCODING_VERSION = "0.1.0" @@ -92,16 +94,6 @@ def is_dataset(obj: Any) -> bool: return is_hdf5_dataset(obj) or is_zarr_array(obj) -def is_zarr_path(path: Path) -> bool: - if not path.exists() or not path.is_dir(): - return False - if (path / "zarr.json").exists(): - return True - if (path / ".zgroup").exists() or (path / ".zattrs").exists(): - return True - return False - - def detect_backend(path: Path) -> str: if path.exists(): if path.is_dir(): diff --git a/src/adata/util/path.py b/src/adata/util/path.py index 24d3832..7cfb103 100644 --- a/src/adata/util/path.py +++ b/src/adata/util/path.py @@ -17,3 +17,15 @@ def same_path(left: Path, right: Path) -> bool: return Path(left).resolve() == Path(right).resolve() except OSError: # pragma: no cover - unresolvable path return Path(left).absolute() == Path(right).absolute() + + +def is_zarr_path(path: Path) -> bool: + """Is this an existing Zarr store (v2 or v3), judged by its metadata files?""" + path = Path(path) + if not path.exists() or not path.is_dir(): + return False + if (path / "zarr.json").exists(): + return True + if (path / ".zgroup").exists() or (path / ".zattrs").exists(): + return True + return False diff --git a/src/adata/util/staging.py b/src/adata/util/staging.py index df2f7c7..779c3ac 100644 --- a/src/adata/util/staging.py +++ b/src/adata/util/staging.py @@ -19,6 +19,8 @@ from pathlib import Path from typing import Iterator, List, Sequence, Tuple +from adata.util.path import is_zarr_path + def _sibling(path: Path, tag: str) -> Path: """`out.zarr` -> `out..zarr`, keeping the suffix that picks the backend.""" @@ -27,9 +29,10 @@ def _sibling(path: Path, tag: str) -> Path: def _temp_for(path: Path) -> Path: temp = _sibling(path, "adata-tmp") - if path.is_dir() and path.suffix != ".zarr": + if path.suffix != ".zarr" and is_zarr_path(path): # A Zarr store named without the suffix. A temporary named the same - # way would not exist yet, and so be opened as HDF5. + # way would not exist yet, and so be opened as HDF5. Only a Zarr + # store: any other directory, such as an MTX folder, keeps its name. temp = temp.with_name(f"{temp.name}.zarr") return temp diff --git a/tests/test_csr_build.py b/tests/test_csr_build.py new file mode 100644 index 0000000..47b3731 --- /dev/null +++ b/tests/test_csr_build.py @@ -0,0 +1,96 @@ +"""Tests for `adata.core.csr_build`, the bucket sort behind streaming transposes. + +The streaming transpose is tested end to end in test_recast.py; these pin the +pieces directly, because the MTX import builds on them with inputs a +transpose never produces: blocks whose keys cover a narrow span, and entries +in no particular order. +""" + +from __future__ import annotations + +import h5py +import numpy as np +import pytest + +from adata.core.csr_build import BucketSorter, bucket_bounds, count_keys, scratch_group + + +@pytest.mark.parametrize("span", [10, 900_000], ids=["dense span", "sparse span"]) +def test_count_keys_matches_bincount(span): + rng = np.random.default_rng(0) + counts = np.zeros(1_000_000, dtype=np.int64) + expected = np.zeros_like(counts) + for _ in range(5): + keys = rng.integers(500, 500 + span, size=1000) + count_keys(counts, keys) + expected += np.bincount(keys, minlength=counts.size) + assert np.array_equal(counts, expected) + + +def test_count_keys_ignores_an_empty_block(): + counts = np.zeros(4, dtype=np.int64) + count_keys(counts, np.array([], dtype=np.int64)) + assert not counts.any() + + +def test_bucket_bounds_cover_every_row_and_balance_nonzeros(): + # One heavy row among many light ones: bounds follow the counts. + counts = np.ones(1000, dtype=np.int64) + counts[10] = 5000 + indptr = np.concatenate(([0], np.cumsum(counts))) + + bounds = bucket_bounds(indptr, chunk=1000, bucket_entries=1000) + + assert bounds[0] == 0 and bounds[-1] == 1000 + assert np.all(np.diff(bounds) > 0) + per_bucket = np.diff(indptr[bounds]) + assert per_bucket.max() <= 5000 + 1000 + + +@pytest.mark.parametrize("order", ["sorted", "reversed", "shuffled"]) +def test_the_runs_are_the_entries_in_csr_order_whatever_order_they_came_in( + tmp_path, order +): + rng = np.random.default_rng(1) + n_rows, n_cols, nnz = 300, 50, 4000 + flat = rng.choice(n_rows * n_cols, size=nnz, replace=False) + rows, cols = flat // n_cols, flat % n_cols + values = rng.integers(1, 100, size=nnz).astype(np.int32) + + arrival = { + "sorted": np.lexsort((cols, rows)), + "reversed": np.lexsort((cols, rows))[::-1], + "shuffled": rng.permutation(nnz), + }[order] + + counts = np.bincount(rows, minlength=n_rows) + indptr = np.concatenate(([0], np.cumsum(counts))) + bounds = bucket_bounds(indptr, chunk=500, bucket_entries=500) + + with h5py.File(tmp_path / "out.h5", "w") as handle: + with scratch_group(handle, "scratch") as scratch: + sorter = BucketSorter(scratch, bounds, np.int32) + for block in np.array_split(arrival, 7): + sorter.add(rows[block], cols[block], values[block]) + runs = list(sorter.runs()) + + assert len(runs) > 1, "expected several buckets" + got_rows = np.concatenate([r for r, _, _ in runs]) + got_cols = np.concatenate([c for _, c, _ in runs]) + got_values = np.concatenate([v for _, _, v in runs]) + + expected = np.lexsort((cols, rows)) + assert np.array_equal(got_rows, rows[expected]) + assert np.array_equal(got_cols, cols[expected]) + assert np.array_equal(got_values, values[expected]) + + +def test_hdf5_scratch_is_a_separate_file_removed_afterwards(tmp_path): + with h5py.File(tmp_path / "out.h5", "w") as handle: + with scratch_group(handle, "scratch") as scratch: + assert scratch.file.filename != handle.filename + during = sorted(p.name for p in tmp_path.iterdir()) + after = sorted(p.name for p in tmp_path.iterdir()) + + assert any("scratch" in name for name in during) + assert after == ["out.h5"] diff --git a/tests/test_staging.py b/tests/test_staging.py index a781822..c45d863 100644 --- a/tests/test_staging.py +++ b/tests/test_staging.py @@ -111,3 +111,18 @@ def test_a_zarr_store_named_without_the_suffix_stays_zarr(tmp_path): assert (out / "zarr.json").read_text() == '{"new": true}' assert _names(tmp_path) == ["store"] + + +def test_a_directory_that_is_not_zarr_keeps_its_name(tmp_path): + """An MTX output folder is a directory too, and must not be staged as Zarr.""" + out = tmp_path / "filtered_feature_bc_matrix" + out.mkdir() + (out / "matrix.mtx.gz").write_bytes(b"old") + + with staged_outputs([out]) as (temp,): + assert temp.name == "filtered_feature_bc_matrix.adata-tmp" + temp.mkdir() + (temp / "matrix.mtx.gz").write_bytes(b"new") + + assert (out / "matrix.mtx.gz").read_bytes() == b"new" + assert _names(tmp_path) == ["filtered_feature_bc_matrix"]