diff --git a/CHANGELOG.md b/CHANGELOG.md index 437917c..cd9589b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,38 @@ Notable changes to `adata-cli`. Versions are `MAJOR.MINOR.PATCH`; tags carry no ## Unreleased +### Added + +- **`adata convert` changes a matrix's dtype, layout or density on disk** + ([#13](https://github.com/cellgeni/adata-cli/issues/13)). Counts stored as + float64 halve with `--dtype float32`; `--indices-dtype int32` halves the + index arrays of a matrix small enough to address that way; `--layout + csr|csc` transposes between the sparse encodings; `--layout dense|sparse` + changes the density. Works on `X`, any layer, `raw/X` or any 2-D array by + path, and on all of them at once with `--all`. + + A cast that would not round-trip is refused before anything is written -- + every value is cast and cast back, because whether float64 counts survive + float32 depends on the counts, not on the dtypes. So is a densification + that would inflate the store beyond four times its size. `--force` + overrides either, and both messages say what they measured. + + Transposing streams by default, holding one bucket of nonzeros rather than + the matrix, so it works on files too large to load; `--in-memory` is + faster when the matrix fits. Buckets are balanced by nonzero count rather + than by coordinate range, because a single-cell matrix is skewed -- a few + genes carry most of the counts -- and equal-width bounds put most of one + in a single bucket. + + `--indices-dtype` is checked against both what `indices` must address and + what `indptr` must reach, which differ: a narrow matrix with more than + 2^31 nonzeros needs int64 offsets over int32 coordinates. Both are + preserved from the source when not specified. + +- **`concat` now names the command to run** when inputs disagree about a + matrix encoding. The check itself is not new, but nothing tested it and it + could not suggest a fix, because there was none. + ### Fixed - **`concat --merge` never finished on a real store.** Aligning a var column diff --git a/README.md b/README.md index b435b73..0e75618 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,7 @@ Run help at any level (e.g. `adata --help`, `adata export --help`). - `subset` – stream and write a filtered copy, selected by obs/var name lists (`--obs`/`--var`) or by expression (`--obs-query`/`--var-query`). - `split` – write one store per distinct value of an annotation column, with a CSV manifest. - `concat` – concatenate stores along the obs axis, with `--join inner|outer` and merge strategies for var and uns. +- `convert` – change a matrix's dtype, layout (CSR/CSC/dense) or density, streaming; refuses a lossy cast or a large size increase unless forced. - `export` – extract data from a store; subcommands: `dataframe` (any dataframe group to CSV), `array` (dense to `.npy`), `sparse` (CSR/CSC to `.mtx`), `dict` (JSON), `image` (PNG). Results go to stdout when no `--output` is given. - `import` – write new data into a store at any path; subcommands: `dataframe` (CSV), `array` (`.npy`), `sparse` (`.mtx`), `dict` (JSON), `image` (PNG/JPEG/TIFF). @@ -78,6 +79,19 @@ adata split data.h5ad --by sample -o per_sample/ adata concat per_sample/*.h5ad -o merged.h5ad --join outer --label sample ``` +### Shrinking a store, and making encodings agree + +```bash +adata convert data.h5ad X --inplace --dtype float32 +adata convert data.h5ad --all -o small.h5ad --dtype float32 --indices-dtype int32 +adata convert data.h5ad X -o csc.h5ad --layout csc +``` + +Counts written as float64 halve with no loss — and `convert` proves that +before it writes, by casting every value and casting it back. `concat` +refuses inputs whose matrices disagree about CSR versus CSC; `--layout` is +how you make them agree. + ## Documentation - [Get started](docs/GET_STARTED.md) — a short tutorial diff --git a/benchmarks/cases.py b/benchmarks/cases.py index 7269b51..6f752b3 100644 --- a/benchmarks/cases.py +++ b/benchmarks/cases.py @@ -485,6 +485,61 @@ def _py(body: str) -> str: ), ], ), + # -- convert ------------------------------------------------------------ + Case( + name="convert-dtype", + question="Rewrite X as float32.", + contenders=[ + Contender( + "adata-cli", + argv=[ + "adata", "convert", "{input0}", "X", "-o", "{output}", + "--dtype", "float32", "--force", + ], + ), + Contender( + "anndata (in memory)", + script=_py( + "obj = ad.read_h5ad(IN[0])\n" + "obj.X = obj.X.astype('float32')\n" + "obj.write_h5ad(OUT, compression=COMPRESSION)\n" + ), + ), + Contender( + "scanpy", + unsupported="no dtype rewrite; scanpy defers to anndata", + ), + ], + ), + Case( + name="convert-layout", + question="Transpose X from CSR to CSC.", + contenders=[ + Contender( + "adata-cli (streaming)", + argv=[ + "adata", "convert", "{input0}", "X", "-o", "{output}", + "--layout", "csc", + ], + ), + Contender( + "adata-cli (--in-memory)", + argv=[ + "adata", "convert", "{input0}", "X", "-o", "{output}", + "--layout", "csc", "--in-memory", + ], + ), + Contender( + "anndata (in memory)", + script=_py( + "obj = ad.read_h5ad(IN[0])\n" + "obj.X = obj.X.tocsc()\n" + "obj.write_h5ad(OUT, compression=COMPRESSION)\n" + ), + ), + ], + tags=["headline"], + ), # -- the claim itself ------------------------------------------------- Case( name="rss-vs-size", diff --git a/docs/COMMANDS.md b/docs/COMMANDS.md index 77e4540..4cac9b8 100644 --- a/docs/COMMANDS.md +++ b/docs/COMMANDS.md @@ -81,6 +81,52 @@ adata subset data.h5ad --inplace --obs barcodes.txt `raw/` is carried over and matched against its **own** var axis, which usually holds more genes than the main object. +## `convert` + +Change a matrix's dtype, layout or density. Counts held as float64 cost twice +the disk and twice the read for no information; a tool that wants CSC cannot +use a CSR store; and `concat` refuses inputs whose encodings disagree. + +```bash +adata convert data.h5ad X -o out.h5ad --dtype float32 +adata convert data.h5ad X -o out.h5ad --layout csc +adata convert data.h5ad X --inplace --dtype float32 --indices-dtype int32 +adata convert data.h5ad --all -o out.h5ad --dtype float32 +adata convert data.h5ad layers/counts -o out.h5ad --layout dense --force +``` + +| Flag | Meaning | +|---|---| +| `--output`, `-o` | Output path. Required unless `--inplace` | +| `--inplace` | Replace the source (written to a temporary path first) | +| `--all` | Convert `X`, every layer and `raw/X` | +| `--dtype` | New dtype for the values, e.g. `float32` | +| `--indices-dtype` | New dtype for sparse indices: `int32` or `int64` | +| `--layout` | `csr`, `csc`, `dense` or `sparse` | +| `--force` | Convert despite a lossy cast or a large size increase | +| `--in-memory` | Transpose in memory rather than streaming | +| `--chunk`, `-C` | Row chunk size for dense matrices | +| `--zarr-format` | Zarr version to write; defaults to the source's | + +Two things are refused before anything is written. A cast that would not +round-trip -- checked by casting every value and casting it back, not by +comparing dtypes -- and a densification that would inflate the store beyond +four times its size. `--force` overrides either, and the message says how +many values would change or how large the result would be. + +The index dtype is **preserved** unless `--indices-dtype` asks otherwise, so +narrowing the values does not silently widen the indices and leave the file +bigger than it started. `indices` and `indptr` keep their own widths, which +can differ: a narrow matrix with more than 2^31 nonzeros needs int64 offsets +over int32 coordinates, and both are range-checked before writing. + +The output path may not name the input; use `--inplace`, which writes to a +temporary file and swaps it in only once the conversion has finished. + +Transposing streams by default and works on matrices too large to load, at +the cost of two extra passes over the nonzeros. `--in-memory` is faster when +the matrix fits. + ## `split` One store per distinct value of a column. diff --git a/docs/index.md b/docs/index.md index 669fd73..b90e217 100644 --- a/docs/index.md +++ b/docs/index.md @@ -51,6 +51,8 @@ adata subset data.h5ad -o cortex.h5ad --obs-query "cluster == Cortex_2" adata split data.h5ad --by sample -o per_sample/ adata concat per_sample/*.h5ad -o merged.h5ad --join outer --label sample +adata convert data.h5ad X -o small.h5ad --dtype float32 + adata create new.h5ad --n-obs 5000 --n-var 2000 adata import sparse new.h5ad X counts.mtx --inplace ``` diff --git a/src/adata/cli.py b/src/adata/cli.py index 45f1f30..cd496a5 100644 --- a/src/adata/cli.py +++ b/src/adata/cli.py @@ -6,9 +6,11 @@ from rich.console import Console import typer +from adata.core.convert import LAYOUTS from adata.commands import ( MERGE_CHOICES, concat_stores, + convert_store, create_store, split_store, list_store, @@ -476,6 +478,134 @@ def concat( raise typer.Exit(code=1) +# ============================================================================ +# CONVERT command +# ============================================================================ +@app.command("convert") +def convert( + file: Path = typer.Argument( + ..., + help="Input .h5ad/.zarr", + exists=True, + readable=True, + dir_okay=True, + file_okay=True, + ), + entries: Optional[List[str]] = typer.Argument( + None, + help="Matrix paths to convert, e.g. 'X', 'layers/counts', 'raw/X'", + ), + output: Optional[Path] = typer.Option( + None, + "--output", + "-o", + help="Output .h5ad/.zarr path. Required unless --inplace.", + dir_okay=True, + file_okay=True, + ), + inplace: bool = typer.Option( + False, + "--inplace", + help="Modify source file directly.", + ), + convert_all: bool = typer.Option( + False, + "--all", + help="Convert X, every layer, and raw/X", + ), + dtype: Optional[str] = typer.Option( + None, + "--dtype", + help="New dtype for the values, e.g. float32", + ), + indices_dtype: Optional[str] = typer.Option( + None, + "--indices-dtype", + help="New dtype for sparse indices: int32 or int64", + ), + layout: Optional[str] = typer.Option( + None, + "--layout", + help="Target layout: csr, csc, dense or sparse", + ), + force: bool = typer.Option( + False, + "--force", + help="Convert despite a lossy cast or a large size increase", + ), + in_memory: bool = typer.Option( + False, + "--in-memory", + help="Transpose in memory instead of streaming (faster if it fits)", + ), + chunk_rows: int = typer.Option( + 1024, + "--chunk", + "-C", + help="Row chunk size for dense matrices", + ), + zarr_format: Optional[int] = typer.Option( + None, + "--zarr-format", + help="Zarr spec version to write (defaults to the source store's)", + ), +) -> None: + """ + Change a matrix's dtype, layout or density. + + Counts held as float64 cost twice the disk and twice the read for no + information; a tool that wants CSC cannot use a CSR store; and concat + refuses inputs whose encodings disagree. All three are this command. + + A cast that would not round-trip is refused before anything is written, + as is a densification that would inflate the store; --force overrides + both. Transposing streams by default, so it works on matrices too large + to load. + + Examples: + adata convert data.h5ad X -o out.h5ad --dtype float32 + adata convert data.h5ad X -o out.h5ad --layout csc + adata convert data.h5ad X --inplace --dtype float32 --indices-dtype int32 + adata convert data.h5ad --all -o out.h5ad --dtype float32 + """ + if not inplace and output is None: + console.print( + "[bold red]Error:[/] Output file is required. " + "Use --output/-o or --inplace.", + ) + raise typer.Exit(code=1) + + if layout is not None and layout not in LAYOUTS: + console.print( + f"[bold red]Error:[/] --layout must be one of: {', '.join(LAYOUTS)}." + ) + raise typer.Exit(code=1) + + if zarr_format is not None and zarr_format not in (2, 3): + console.print("[bold red]Error:[/] --zarr-format must be 2 or 3.") + raise typer.Exit(code=1) + + try: + convert_store( + file=file, + entries=list(entries) if entries else None, + output=output, + console=console, + dtype=dtype, + indices_dtype=indices_dtype, + layout=layout, + convert_all=convert_all, + inplace=inplace, + chunk_rows=chunk_rows, + in_memory=in_memory, + force=force, + zarr_format=zarr_format, + ) + except Exception as e: + console.print(f"[bold red]Error:[/] {e}") + raise typer.Exit(code=1) + + # ============================================================================ # SPLIT command # ============================================================================ diff --git a/src/adata/commands/__init__.py b/src/adata/commands/__init__.py index 3cc464b..5ed2e4c 100644 --- a/src/adata/commands/__init__.py +++ b/src/adata/commands/__init__.py @@ -6,3 +6,4 @@ from adata.commands.create import create_store from adata.commands.split import split_store from adata.commands.concat import MERGE_CHOICES, MERGE_STRATEGIES, concat_stores +from adata.commands.convert import convert_store diff --git a/src/adata/commands/convert.py b/src/adata/commands/convert.py new file mode 100644 index 0000000..54844ff --- /dev/null +++ b/src/adata/commands/convert.py @@ -0,0 +1,218 @@ +"""The `convert` command: rewrite matrices, copy everything else. + +Writes a whole new store rather than editing one in place, for the same +reason `subset` does: a conversion that fails half way leaves the original +untouched, and `--inplace` becomes an atomic swap of a finished file rather +than a partial edit of a live one. +""" + +from __future__ import annotations + +import shutil +from pathlib import Path +from typing import Any, List, Optional + +from rich.console import Console + +from adata.core.convert import ( + DEFAULT_CHUNK, + LAYOUTS, + Plan, + convert_matrix, + parse_dtype, + plan_conversion, +) +from adata.elements import spec +from adata.elements.write import ensure_anndata_skeleton +from adata.storage import copy_tree, detect_backend, is_group, open_store + +#: What `--all` means. The matrices an AnnData object is built from, not +#: every 2-D array in the file: obsm/varm hold embeddings, whose dtype is +#: rarely what anyone is trying to shrink. Those still convert by path. +ALL_PREFIXES = ("X", "layers", "raw/X") + + +def discover_matrices(root: Any) -> List[str]: + """Entry paths `--all` resolves to, in a stable order.""" + found: List[str] = [] + if "X" in root: + found.append("X") + if "layers" in root and is_group(root["layers"]): + found.extend(f"layers/{key}" for key in sorted(root["layers"].keys())) + if "raw" in root and is_group(root["raw"]) and "X" in root["raw"]: + found.append("raw/X") + return found + + +def _same_path(left: Path, right: Path) -> bool: + """Do these name the same store, through symlinks and `..` alike?""" + try: + return left.resolve() == right.resolve() + except OSError: # pragma: no cover - unresolvable path + return left.absolute() == right.absolute() + + +def _resolve(root: Any, path: str) -> Any: + obj = root + for part in path.split("/"): + if part not in obj: + raise KeyError(f"{path!r} not found in the store.") + obj = obj[part] + return obj + + +def convert_store( + file: Path, + entries: Optional[List[str]], + output: Optional[Path], + console: Console, + *, + dtype: Optional[str] = None, + indices_dtype: Optional[str] = None, + layout: Optional[str] = None, + convert_all: bool = False, + inplace: bool = False, + chunk: int = DEFAULT_CHUNK, + chunk_rows: int = 1024, + in_memory: bool = False, + force: bool = False, + zarr_format: Optional[int] = None, +) -> None: + """Write a copy of `file` with the named matrices converted.""" + if not inplace and output is None: + raise ValueError("Output file is required unless --inplace is specified.") + if dtype is None and indices_dtype is None and layout is None: + raise ValueError( + "Nothing to do: pass at least one of --dtype, --indices-dtype " + "or --layout." + ) + if layout is not None and layout not in LAYOUTS: + raise ValueError(f"--layout must be one of: {', '.join(LAYOUTS)}.") + + data_dtype = parse_dtype(dtype) if dtype else None + index_dtype = ( + parse_dtype(indices_dtype, allowed=("int32", "int64")) + if indices_dtype + else None + ) + + if output is not None and not inplace and _same_path(file, output): + # Opening the destination "w" clears it while the source is still + # being read from it. HDF5 refuses; Zarr does not, and quietly + # produced an empty store where the data used to be. + raise ValueError( + f"Output path is the input: {output}. Use --inplace to replace " + "it, which writes to a temporary file first." + ) + + if inplace: + backend = detect_backend(file) + if backend == "zarr": + base = file.stem if file.suffix else file.name + dst_path = file.with_name(f"{base}.convert-tmp.zarr") + else: + dst_path = file.with_name(f"{file.name}.convert-tmp") + if dst_path.exists(): + raise FileExistsError(f"Temporary path already exists: {dst_path}") + else: + dst_path = output + + if zarr_format is None and detect_backend(file) == "zarr": + with open_store(file, "r") as probe: + zarr_format = probe.zarr_format + + with open_store(file, "r") as src_store: + src = src_store.root + targets = discover_matrices(src) if convert_all else list(entries or []) + if not targets: + raise ValueError( + "No matrices selected. Name one (e.g. `X`) or pass --all." + ) + # Every check first, before the destination exists. A refusal must + # not leave behind a store holding a copy of obs and var and nothing + # else -- the caller cannot tell that from a finished conversion. + plans = { + path: plan_conversion( + _resolve(src, path), + path, + dtype=data_dtype, + index_dtype=index_dtype, + layout=layout, + chunk=chunk, + force=force, + console=console, + ) + for path in targets + } + + console.print( + f"[cyan]Converting {len(targets)} matrix/matrices:[/] " + + ", ".join(targets) + ) + + with open_store(dst_path, "w", zarr_format=zarr_format) as dst_store: + dst = dst_store.root + _write( + src, dst, plans, + chunk=chunk, chunk_rows=chunk_rows, in_memory=in_memory, + console=console, + ) + ensure_anndata_skeleton(dst) + + if inplace: + if file.is_dir(): + shutil.rmtree(file) + elif file.exists(): + file.unlink() + if dst_path.is_dir(): + shutil.move(str(dst_path), str(file)) + else: + dst_path.replace(file) + console.print(f"[green]Converted[/] {file}") + else: + console.print(f"[green]Wrote[/] {dst_path}") + + +def _write(src: Any, dst: Any, plans: dict, **options: Any) -> None: + """Copy the store across, converting the targeted entries as they pass.""" + by_parent: dict = {} + for path in plans: + parent, _, leaf = path.rpartition("/") + by_parent.setdefault(parent, {})[leaf] = plans[path] + + for key in src.keys(): + if key in by_parent.get("", {}): + convert_matrix( + src[key], dst, key, plan=by_parent[""][key], **options + ) + elif key in by_parent: + # Any parent, not just layers and raw. Restricting it to those + # two meant an explicitly named `obsm/X_pca` was copied + # unconverted and the command still reported success. + _write_group(src[key], dst, key, by_parent, **options) + else: + copy_tree(src[key], dst, key) + + +def _write_group( + group: Any, dst: Any, name: str, by_parent: dict, **options: Any +) -> None: + """Recreate one container, converting the members that were named.""" + from adata.storage import copy_attrs, is_zarr_group + + out = dst.create_group(name) + copy_attrs( + group.attrs, + out.attrs, + target_backend="zarr" if is_zarr_group(dst) else "hdf5", + ) + if not spec.encoding_type(out): + spec.set_encoding(out, spec.RAW if name == "raw" else spec.DICT) + + + wanted = by_parent.get(name, {}) + for key in group.keys(): + if key in wanted: + convert_matrix(group[key], out, key, plan=wanted[key], **options) + else: + copy_tree(group[key], out, key) diff --git a/src/adata/core/concat.py b/src/adata/core/concat.py index f1c4ecc..39aa7ea 100644 --- a/src/adata/core/concat.py +++ b/src/adata/core/concat.py @@ -561,10 +561,13 @@ def _check(label: str, sources: List[Any]) -> None: kinds = {_matrix_kind(s) for s in sources} if kinds in ({spec.CSR_MATRIX}, {spec.CSC_MATRIX}, {"dense"}): return + wanted = sorted(kinds)[0] raise ValueError( f"Cannot concatenate {label!r}: inputs use " f"{', '.join(sorted(kinds))}. Every input must use the same " - "encoding -- convert them to match first." + f"encoding. Convert them to match first, e.g. " + f"`adata convert INPUT {label} -o converted.h5ad --layout " + f"{'dense' if wanted == 'dense' else wanted.removesuffix('_matrix')}`." ) if all("X" in r for r in roots): @@ -624,6 +627,9 @@ def _concat_matrix( ) return True + # A backstop. `check_matrix_encodings` runs before the output store is + # created and should have raised already; this catches an element it + # does not cover, where failing late still beats writing nonsense. raise ValueError( f"Cannot concatenate {name!r}: inputs use {', '.join(sorted(kinds))}." ) diff --git a/src/adata/core/convert.py b/src/adata/core/convert.py new file mode 100644 index 0000000..fb3cfd3 --- /dev/null +++ b/src/adata/core/convert.py @@ -0,0 +1,906 @@ +"""Rewriting a matrix's dtype, layout or density, on disk. + +Three conversions, one command. Counts stored as float64 cost twice the disk +and twice the read for no information (issue #13); a CSR store handed to a +tool that wants CSC has to be transposed somewhere; and `concat` refuses +inputs whose encodings disagree, which until now left nothing to do about it. + +Everything here streams. The matrix these conversions matter for is the one +too large to load, so a converter that loads it would only work on the files +that did not need converting. + +Two safety rules, both of which fail before anything is written -- the same +principle `check_matrix_encodings` states for concat, because a half-written +store is worse than a refusal: + +* a cast that does not round-trip is refused unless forced, so `int32` + indices that would overflow, or a float downcast that would lose real + precision, are reported rather than silently written; +* densifying is refused when it would inflate the store beyond + `MAX_GROWTH_FACTOR`, because a 5% dense matrix becomes ten times its size + and that should not be a surprise. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +from rich.console import Console + +from adata.elements import spec +from adata.elements.write import set_shape_attr +from adata.storage import ( + copy_attrs, + create_dataset, + dataset_create_kwargs, + is_dataset, + is_group, + is_zarr_group, +) + +#: Layouts the user can ask for. "sparse" means "whichever sparse encoding +#: keeps the major axis it already has", so densify/sparsify round-trips. +LAYOUTS = ("csr", "csc", "dense", "sparse") + +#: What `--dtype` accepts. An allowlist rather than `np.dtype(text)`: that +#: would cheerfully accept "S10" or "datetime64[ns]" and produce a store no +#: reader expects from a matrix. +DTYPES = ( + "float16", "float32", "float64", + "int8", "int16", "int32", "int64", + "uint8", "uint16", "uint32", "uint64", + "bool", +) + +#: Indices and indptr must be integers, and signed: anndata and scipy both +#: expect a signed index type. +INDEX_DTYPES = ("int32", "int64") + +#: How much larger a densified store may get before it needs --force. +MAX_GROWTH_FACTOR = 4.0 + +#: 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: + """Resolve a user-supplied dtype name, or say what is accepted.""" + name = text.strip().lower() + if name not in allowed: + raise ValueError( + f"Unknown dtype {text!r}. Choose from: {', '.join(allowed)}." + ) + return np.dtype(name) + + +# --------------------------------------------------------------------------- +# safety + + +@dataclass +class CastReport: + """What a cast would do to the values, measured rather than assumed.""" + + total: int = 0 + changed: int = 0 + worst_absolute: float = 0.0 + overflowed: bool = False + + @property + def lossless(self) -> bool: + return self.changed == 0 and not self.overflowed + + def describe(self, source: Any, target: Any) -> str: + if self.lossless: + return f"{source} -> {target}: every value round-trips" + share = self.changed / self.total if self.total else 0.0 + detail = ( + "values exceed its range" + if self.overflowed + else f"largest change {self.worst_absolute:g}" + ) + return ( + f"{source} -> {target} is lossy: {self.changed:,} of " + f"{self.total:,} values ({share:.2%}) do not round-trip, {detail}" + ) + + +def check_cast(dataset: Any, target: np.dtype, *, chunk: int = DEFAULT_CHUNK) -> CastReport: + """Would casting `dataset` to `target` lose anything? + + Streams the values, casts each block and casts it back. An exact + round-trip is the only honest test: whether float64 counts survive + float32 depends on the counts, not on the dtypes, and that is precisely + the question issue #13 is asking. + """ + report = CastReport() + source = np.dtype(getattr(dataset, "dtype", "float64")) + n = int(dataset.shape[0]) if getattr(dataset, "shape", None) else 0 + + info = np.finfo(target) if target.kind == "f" else ( + np.iinfo(target) if target.kind in "iu" else None + ) + + for start in range(0, n, chunk): + block = np.asarray(dataset[start : min(start + chunk, n)]) + report.total += block.size + if block.size == 0: + continue + + with np.errstate(invalid="ignore", over="ignore"): + cast = block.astype(target) + back = cast.astype(source) + + if info is not None: + finite = block[np.isfinite(block)] if source.kind == "f" else block + if finite.size and ( + float(finite.max()) > float(info.max) + or float(finite.min()) < float(info.min) + ): + report.overflowed = True + + # NaN never equals itself, so compare those separately rather than + # counting every missing value as a loss. + differs = back != block + if source.kind == "f": + both_nan = np.isnan(block) & np.isnan(back) + differs &= ~both_nan + count = int(differs.sum()) + if count: + report.changed += count + delta = np.abs( + block[differs].astype("float64") - back[differs].astype("float64") + ) + finite_delta = delta[np.isfinite(delta)] + if finite_delta.size: + report.worst_absolute = max( + report.worst_absolute, float(finite_delta.max()) + ) + + return report + + +def check_index_dtype( + group: Any, + index_dtype: np.dtype, + pointer_dtype: np.dtype, + shape: Tuple[int, int], +) -> None: + """Refuse index dtypes that cannot address this matrix. + + Cheaper than `check_cast`: the largest value each array must hold is + bounded by the dimensions and the nonzero count, so no pass over the + data is needed. `indices` holds coordinates, bounded by the larger + dimension; `indptr` holds offsets, bounded by nnz. They are checked + separately because they can legitimately need different widths. + """ + nnz = int(group["indices"].shape[0]) + for name, dtype, largest, what in ( + ("indices", index_dtype, max(int(shape[0]), int(shape[1])), "coordinates"), + ("indptr", pointer_dtype, nnz, "offsets"), + ): + limit = int(np.iinfo(dtype).max) + if largest > limit: + raise ValueError( + f"{dtype} cannot hold this matrix's {what}: {name} must reach " + f"{largest:,} for a {shape[0]:,} x {shape[1]:,} matrix with " + f"{nnz:,} nonzeros, and {dtype} tops out at {limit:,}. " + "Use int64." + ) + + +def _stored_bytes(obj: Any) -> int: + """Bytes this element occupies, as best the backend will say.""" + total = 0 + targets = [obj] if is_dataset(obj) else [ + obj[k] for k in ("data", "indices", "indptr") if k in obj + ] + for item in targets: + try: + total += int(item.nbytes) + except Exception: # pragma: no cover - backend without nbytes + shape = getattr(item, "shape", ()) or () + size = int(np.prod(shape)) if shape else 0 + total += size * int(getattr(item.dtype, "itemsize", 8) or 8) + return total + + +def check_growth( + current: int, projected: int, *, force: bool, what: str +) -> None: + """Refuse a conversion that inflates the store, unless asked twice.""" + if force or current <= 0 or projected <= current * MAX_GROWTH_FACTOR: + return + raise ValueError( + f"{what} would grow from {current / 1e6:,.0f} MB to " + f"{projected / 1e6:,.0f} MB ({projected / current:.1f}x). That is " + f"above the {MAX_GROWTH_FACTOR:g}x limit; pass --force if it is what " + "you want." + ) + + +# --------------------------------------------------------------------------- +# reading a source + + +@dataclass +class Matrix: + """A matrix on disk, described enough to convert it.""" + + obj: Any + kind: str # "csr_matrix" | "csc_matrix" | "dense" + shape: Tuple[int, int] + dtype: np.dtype + + @property + def sparse(self) -> bool: + return self.kind in spec.SPARSE_TYPES + + @property + def nnz(self) -> int: + return int(self.obj["indices"].shape[0]) if self.sparse else 0 + + +def describe(obj: Any) -> Matrix: + """Classify a matrix element, or say why it is not one.""" + enc = spec.encoding_type(obj) + if is_group(obj) and enc in spec.SPARSE_TYPES: + shape = obj.attrs.get("shape", None) + if shape is None: + raise ValueError("Sparse matrix group is missing its 'shape' attribute.") + return Matrix(obj, enc, (int(shape[0]), int(shape[1])), obj["data"].dtype) + + if is_dataset(obj): + if getattr(obj, "ndim", 0) != 2: + raise ValueError( + f"Only 2-D matrices can be converted; this one has " + f"{getattr(obj, 'ndim', '?')} dimension(s)." + ) + return Matrix(obj, "dense", (int(obj.shape[0]), int(obj.shape[1])), obj.dtype) + + raise ValueError( + f"Not a matrix: encoding {enc!r}. `convert` handles X, layers, raw/X " + "and any 2-D dense array." + ) + + +def resolve_layout(source: Matrix, requested: Optional[str]) -> str: + """The concrete target encoding for a possibly-vague request.""" + if requested is None: + return source.kind + if requested == "dense": + return "dense" + if requested == "sparse": + # Keep the major axis it already had, so sparse -> dense -> sparse is + # the identity rather than a silent transpose. + return source.kind if source.sparse else spec.CSR_MATRIX + return spec.CSR_MATRIX if requested == "csr" else spec.CSC_MATRIX + + +# --------------------------------------------------------------------------- +# writers + + +def _new_sparse_group( + src: Matrix, dst_parent: Any, name: str, enc: str, shape: Tuple[int, int] +) -> Any: + group = dst_parent.create_group(name) + backend = "zarr" if is_zarr_group(dst_parent) else "hdf5" + copy_attrs(src.obj.attrs, group.attrs, target_backend=backend) + spec.set_encoding(group, enc) + set_shape_attr(group, shape) + return group + + +def _sparse_dataset( + group: Any, name: str, dtype: np.dtype, n: int, template: Any +) -> Any: + """A 1-D dataset of `n` elements laid out like `template`. + + Forwarding compression and chunking matters more here than anywhere + else: the point of a dtype change is usually to make the file smaller, + and creating the output with a fixed 65,536-element chunk made a + 2,400-nonzero matrix allocate 786 KB of mostly empty chunk -- eight + times the source, from a conversion asked for to halve it. + """ + from adata.core.subset import _clamp_chunks + + backend = "zarr" if is_zarr_group(group) else "hdf5" + kw = dataset_create_kwargs(template, target_backend=backend, dst_parent=group) + kw = _clamp_chunks(kw, max(1, n)) + if "chunks" not in kw and n: + kw["chunks"] = (min(n, 1 << 16),) + return create_dataset(group, name, shape=(n,), dtype=dtype, **kw) + + +def _growable_like(group: Any, name: str, dtype: np.dtype, template: Any) -> Any: + """Like `_sparse_dataset`, but extensible for a size not yet known.""" + backend = "zarr" if is_zarr_group(group) else "hdf5" + kw = dataset_create_kwargs(template, target_backend=backend, dst_parent=group) + kw.pop("shards", None) + chunks = kw.pop("chunks", None) + step = int(chunks[0]) if chunks else 1 << 16 + if is_zarr_group(group): + return group.create_array(name, shape=(0,), dtype=dtype, chunks=(step,), **kw) + return group.create_dataset( + name, shape=(0,), maxshape=(None,), dtype=dtype, chunks=(step,), **kw + ) + + +def _write_sparse_arrays( + group: Any, + data: np.ndarray, + indices: np.ndarray, + indptr: np.ndarray, + *, + data_dtype: np.dtype, + index_dtype: np.dtype, + pointer_dtype: np.dtype, + source: Any, +) -> None: + for name, values, dtype, template in ( + ("data", data, data_dtype, source["data"]), + ("indices", indices, index_dtype, source["indices"]), + ("indptr", indptr, pointer_dtype, source["indptr"]), + ): + cast = values.astype(dtype, copy=False) + dataset = _sparse_dataset(group, name, dtype, cast.size, template) + if cast.size: + dataset[:] = cast + + +def cast_sparse( + src: Matrix, + dst_parent: Any, + name: str, + *, + data_dtype: np.dtype, + index_dtype: np.dtype, + pointer_dtype: np.dtype, + chunk: int = DEFAULT_CHUNK, +) -> None: + """Rewrite a sparse matrix with new dtypes, keeping its layout. + + The cheap case, and the one issue #13 asks for: the structure is + untouched, so this is a straight streamed copy of `data` and `indices` + into differently typed datasets. Sized up front, because nnz is known. + """ + group = _new_sparse_group(src, dst_parent, name, src.kind, src.shape) + source_data, source_indices = src.obj["data"], src.obj["indices"] + nnz = int(source_data.shape[0]) + + out_data = _sparse_dataset(group, "data", data_dtype, nnz, source_data) + out_indices = _sparse_dataset(group, "indices", index_dtype, nnz, source_indices) + + for start in range(0, nnz, chunk): + end = min(start + chunk, nnz) + out_data[start:end] = np.asarray(source_data[start:end]).astype( + data_dtype, copy=False + ) + out_indices[start:end] = np.asarray(source_indices[start:end]).astype( + index_dtype, copy=False + ) + + indptr = np.asarray(src.obj["indptr"][...]).astype(pointer_dtype, copy=False) + out_indptr = _sparse_dataset( + group, "indptr", pointer_dtype, indptr.size, src.obj["indptr"] + ) + out_indptr[:] = indptr + + +def transpose_sparse_in_memory( + src: Matrix, + dst_parent: Any, + name: str, + *, + data_dtype: np.dtype, + index_dtype: np.dtype, + pointer_dtype: np.dtype = np.dtype("int64"), +) -> None: + """Swap CSR<->CSC by loading the matrix and sorting it once. + + Faster than streaming whenever the matrix fits, and the whole matrix is + what it needs -- so it is opt-in, never the default. + """ + target = ( + spec.CSC_MATRIX if src.kind == spec.CSR_MATRIX else spec.CSR_MATRIX + ) + n_major_in = src.shape[0] if src.kind == spec.CSR_MATRIX else src.shape[1] + n_major_out = src.shape[1] if src.kind == spec.CSR_MATRIX else src.shape[0] + + indptr = np.asarray(src.obj["indptr"][...], dtype=np.int64) + minor = np.asarray(src.obj["indices"][...], dtype=np.int64) + values = np.asarray(src.obj["data"][...]) + major = np.repeat(np.arange(n_major_in, dtype=np.int64), np.diff(indptr)) + + # Sort by the new major axis, then the new minor, which is what both + # encodings require of `indices` within a row. + order = np.lexsort((major, minor)) + out_indptr = np.concatenate( + ([0], np.cumsum(np.bincount(minor, minlength=n_major_out))) + ) + + group = _new_sparse_group(src, dst_parent, name, target, src.shape) + _write_sparse_arrays( + group, + values[order], + major[order], + out_indptr, + data_dtype=data_dtype, + index_dtype=index_dtype, + pointer_dtype=pointer_dtype, + source=src.obj, + ) + + +def transpose_sparse_streaming( + src: Matrix, + dst_parent: Any, + name: str, + *, + data_dtype: np.dtype, + index_dtype: np.dtype, + pointer_dtype: np.dtype = np.dtype("int64"), + chunk: int = DEFAULT_CHUNK, + bucket_entries: Optional[int] = None, + console: Optional[Console] = None, +) -> None: + """Swap CSR<->CSC without loading the matrix. + + A transpose cannot be done in one pass: the first entry of the output + may come from the last row of the input. Three passes instead, with + memory set by `bucket_entries` rather than by nnz: + + 1. count nonzeros per output major, by streaming `indices` alone. That + gives the output `indptr` by cumulative sum. + 2. stream the input again, splitting each block's entries into buckets by + which slice of the output they land in, and append each bucket to + scratch datasets. + 3. read one bucket at a time, sort it, and append to the output in order. + + 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 + ) + n_major_in = src.shape[0] if src.kind == spec.CSR_MATRIX else src.shape[1] + n_major_out = src.shape[1] if src.kind == spec.CSR_MATRIX else src.shape[0] + + indptr = np.asarray(src.obj["indptr"][...], dtype=np.int64) + source_indices, source_data = src.obj["indices"], src.obj["data"] + nnz = int(source_indices.shape[0]) + + # Pass 1: how many entries land in each output major. + 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) + 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) + n_buckets = len(bounds) - 1 + if console is not None and n_buckets > 1: + console.print( + f"[dim]Transposing {nnz:,} nonzeros through {n_buckets} buckets[/]" + ) + + # nnz is invariant under a transpose, so the output can be sized now and + # written by slice rather than grown block by block. + group = _new_sparse_group(src, dst_parent, name, target, src.shape) + out_data = _sparse_dataset(group, "data", data_dtype, nnz, source_data) + out_indices = _sparse_dataset(group, "indices", index_dtype, nnz, source_indices) + written = 0 + + scratch_name = f"__{name}_transpose_scratch__" + scratch = dst_parent.create_group(scratch_name) + try: + 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) + ] + + # Pass 2: scatter the input into buckets, one input block at a time. + # Step over the major axis in blocks holding about `chunk` nonzeros, + # so the read size is set by the data rather than by how many rows + # happen to be empty. + average = max(1, nnz // max(1, n_major_in)) + major_step = max(1, chunk // average) + for lo in range(0, n_major_in, major_step): + hi = min(lo + major_step, n_major_in) + start, end = int(indptr[lo]), int(indptr[hi]) + if end <= start: + continue + minor = np.asarray(source_indices[start:end], dtype=np.int64) + values = np.asarray(source_data[start:end]) + 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]) + + # 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( + index_dtype, copy=False + ) + out_data[written : written + count] = values[order].astype( + data_dtype, copy=False + ) + written += count + finally: + del dst_parent[scratch_name] + + cast_indptr = out_indptr.astype(pointer_dtype, copy=False) + dataset = _sparse_dataset( + group, "indptr", pointer_dtype, cast_indptr.size, src.obj["indptr"] + ) + dataset[:] = cast_indptr + + +def densify( + src: Matrix, + dst_parent: Any, + name: str, + *, + data_dtype: np.dtype, + chunk_rows: int = 1024, +) -> None: + """Write a sparse matrix out as a dense array, a block of rows at a time. + + The zeros are the point: nothing here materialises the whole grid, so a + matrix too large to densify in memory still converts -- it just produces + a file that is honestly much larger, which `check_growth` warns about + before any of it is written. + """ + n_rows, n_cols = src.shape + backend = "zarr" if is_zarr_group(dst_parent) else "hdf5" + dst = create_dataset( + dst_parent, name, shape=(n_rows, n_cols), dtype=data_dtype + ) + copy_attrs(src.obj.attrs, dst.attrs, target_backend=backend) + spec.set_encoding(dst, spec.ARRAY) + # `shape` belongs to the sparse encoding and would contradict the array's + # own shape if it were carried over. + if "shape" in dst.attrs: + del dst.attrs["shape"] + + indptr = np.asarray(src.obj["indptr"][...], dtype=np.int64) + indices, values = src.obj["indices"], src.obj["data"] + csr = src.kind == spec.CSR_MATRIX + n_major = n_rows if csr else n_cols + + for lo in range(0, n_major, chunk_rows): + hi = min(lo + chunk_rows, n_major) + start, end = int(indptr[lo]), int(indptr[hi]) + block = np.zeros( + (hi - lo, n_cols) if csr else (n_rows, hi - lo), dtype=data_dtype + ) + if end > start: + minor = np.asarray(indices[start:end], dtype=np.int64) + data = np.asarray(values[start:end]) + major = np.repeat( + np.arange(hi - lo, dtype=np.int64), np.diff(indptr[lo : hi + 1]) + ) + # Repeated coordinates are legal in a CSR/CSC store and mean + # their sum, which is what scipy's own `toarray` produces. + # Plain assignment keeps whichever came last, so a + # non-canonical input silently changed value on densifying. + if csr: + np.add.at(block, (major, minor), data) + else: + np.add.at(block, (minor, major), data) + if csr: + dst[lo:hi, :] = block + else: + dst[:, lo:hi] = block + + +def sparsify( + src: Matrix, + dst_parent: Any, + name: str, + *, + enc: str, + data_dtype: np.dtype, + index_dtype: np.dtype, + pointer_dtype: np.dtype = np.dtype("int64"), + chunk_rows: int = 1024, + console: Optional[Console] = None, +) -> Tuple[int, float]: + """Write a dense array out as CSR or CSC, a block at a time. + + Returns `(nnz, density)` so the caller can say whether it was worth it -- + a matrix that is half nonzero gets bigger, not smaller, and the user + should hear that from us rather than from `du`. + """ + from adata.core.subset import _append, _growable + + n_rows, n_cols = src.shape + csr = enc == spec.CSR_MATRIX + n_major = n_rows if csr else n_cols + + group = dst_parent.create_group(name) + backend = "zarr" if is_zarr_group(dst_parent) else "hdf5" + copy_attrs(src.obj.attrs, group.attrs, target_backend=backend) + spec.set_encoding(group, enc) + set_shape_attr(group, src.shape) + + out_data = _growable_like(group, "data", data_dtype, src.obj) + out_indices = _growable_like(group, "indices", index_dtype, src.obj) + counts: List[int] = [] + + for lo in range(0, n_major, chunk_rows): + hi = min(lo + chunk_rows, n_major) + block = np.asarray(src.obj[lo:hi, :] if csr else src.obj[:, lo:hi]) + if not csr: + block = block.T # iterate majors as rows either way + nonzero_major, nonzero_minor = np.nonzero(block) + counts.extend(np.bincount(nonzero_major, minlength=hi - lo).tolist()) + _append(out_indices, nonzero_minor.astype(index_dtype, copy=False)) + _append(out_data, block[nonzero_major, nonzero_minor].astype( + data_dtype, copy=False + )) + + indptr = np.concatenate(([0], np.cumsum(counts, dtype=np.int64))) + create_dataset(group, "indptr", data=indptr.astype(pointer_dtype, copy=False)) + + + nnz = int(indptr[-1]) + density = nnz / max(1, n_rows * n_cols) + if console is not None and density > 0.5: + console.print( + f"[yellow]{name} is {density:.0%} nonzero; the sparse form is " + f"larger than the dense one below about 33%.[/]" + ) + return nnz, density + + +def cast_dense( + src: Matrix, + dst_parent: Any, + name: str, + *, + data_dtype: np.dtype, + chunk_rows: int = 1024, +) -> None: + """Rewrite a dense matrix with a new dtype, a block of rows at a time.""" + n_rows, n_cols = src.shape + backend = "zarr" if is_zarr_group(dst_parent) else "hdf5" + kw = dataset_create_kwargs( + src.obj, target_backend=backend, dst_parent=dst_parent + ) + from adata.core.subset import _clamp_chunks + + dst = create_dataset( + dst_parent, + name, + shape=(n_rows, n_cols), + dtype=data_dtype, + **_clamp_chunks(kw, n_rows, n_cols), + ) + copy_attrs(src.obj.attrs, dst.attrs, target_backend=backend) + + for lo in range(0, n_rows, chunk_rows): + hi = min(lo + chunk_rows, n_rows) + dst[lo:hi, :] = np.asarray(src.obj[lo:hi, :]).astype( + data_dtype, copy=False + ) + + +# --------------------------------------------------------------------------- +# dispatch + + +@dataclass +class Plan: + """A checked conversion, ready to run. + + Separating the decision from the writing is what makes "fails before + anything is written" true rather than aspirational: the caller resolves + every plan first, and only then creates the destination store. An + earlier version checked inside the writer, so a refusal still left an + output file holding a copy of obs and var. + """ + + source: Matrix + layout: str + data_dtype: np.dtype + index_dtype: np.dtype + #: `indptr` is tracked apart from `indices` because the two can + #: legitimately differ: a narrow matrix with more than 2^31 nonzeros + #: needs int64 offsets over int32 column indices. Inferring one from + #: the other silently overflowed the offsets and corrupted the matrix. + pointer_dtype: np.dtype = np.dtype("int64") + report: Optional[CastReport] = None + + +def plan_conversion( + obj: Any, + name: str, + *, + dtype: Optional[np.dtype] = None, + index_dtype: Optional[np.dtype] = None, + layout: Optional[str] = None, + chunk: int = DEFAULT_CHUNK, + force: bool = False, + console: Optional[Console] = None, +) -> Plan: + """Decide what to do, and refuse here if it should not be done.""" + src = describe(obj) + target_layout = resolve_layout(src, layout) + data_dtype = np.dtype(dtype) if dtype is not None else src.dtype + + if index_dtype is not None: + idx_dtype = ptr_dtype = np.dtype(index_dtype) + elif src.sparse: + # Keep what the source used, each independently. Defaulting to + # int64 doubled the index arrays of every int32 store; inferring + # indptr from indices narrowed the offsets of every store that + # needed them wider. + idx_dtype = np.dtype(src.obj["indices"].dtype) + ptr_dtype = np.dtype(src.obj["indptr"].dtype) + else: + idx_dtype = ptr_dtype = np.dtype("int64") + + report: Optional[CastReport] = None + if dtype is not None and data_dtype != src.dtype: + values = src.obj["data"] if src.sparse else src.obj + report = check_cast(values, data_dtype, chunk=chunk) + message = report.describe(src.dtype, data_dtype) + if not report.lossless and not force: + raise ValueError(f"{name}: {message}. Pass --force to convert anyway.") + if console is not None: + colour = "dim" if report.lossless else "yellow" + console.print(f"[{colour}]{name}: {message}[/]") + + # Always, not only when asked: an inferred dtype can be too narrow too, + # and a silently overflowed offset is indistinguishable from corruption. + if src.sparse and not force: + check_index_dtype(src.obj, idx_dtype, ptr_dtype, src.shape) + + if target_layout == "dense" and src.sparse: + projected = src.shape[0] * src.shape[1] * data_dtype.itemsize + check_growth( + _stored_bytes(src.obj), projected, force=force, what=f"{name} as dense" + ) + + return Plan(src, target_layout, data_dtype, idx_dtype, ptr_dtype, report) + + +def convert_matrix( + obj: Any, + dst_parent: Any, + name: str, + *, + plan: Optional[Plan] = None, + dtype: Optional[np.dtype] = None, + index_dtype: Optional[np.dtype] = None, + layout: Optional[str] = None, + chunk: int = DEFAULT_CHUNK, + chunk_rows: int = 1024, + in_memory: bool = False, + force: bool = False, + console: Optional[Console] = None, +) -> None: + """Write `obj` into `dst_parent` under `name`, converted. + + Pass a `plan` from `plan_conversion` to have the checks already done; + without one they run here, which is convenient for a direct caller but + means the destination already exists by the time a refusal is raised. + """ + if plan is None: + plan = plan_conversion( + obj, name, dtype=dtype, index_dtype=index_dtype, layout=layout, + chunk=chunk, force=force, console=console, + ) + src = plan.source + target_layout = plan.layout + data_dtype = plan.data_dtype + idx_dtype = plan.index_dtype + ptr_dtype = plan.pointer_dtype + + # --- then write ------------------------------------------------------- + if target_layout == "dense": + if src.sparse: + densify(src, dst_parent, name, data_dtype=data_dtype, chunk_rows=chunk_rows) + else: + cast_dense(src, dst_parent, name, data_dtype=data_dtype, chunk_rows=chunk_rows) + return + + if not src.sparse: + sparsify( + src, + dst_parent, + name, + enc=target_layout, + data_dtype=data_dtype, + index_dtype=idx_dtype, + pointer_dtype=ptr_dtype, + chunk_rows=chunk_rows, + console=console, + ) + return + + if target_layout == src.kind: + cast_sparse( + src, + dst_parent, + name, + data_dtype=data_dtype, + index_dtype=idx_dtype, + pointer_dtype=ptr_dtype, + chunk=chunk, + ) + return + + transpose = ( + transpose_sparse_in_memory if in_memory else transpose_sparse_streaming + ) + extra: Dict[str, Any] = ( + {} if in_memory else {"chunk": chunk, "console": console} + ) + transpose( + src, + dst_parent, + name, + data_dtype=data_dtype, + index_dtype=idx_dtype, + pointer_dtype=ptr_dtype, + **extra, + ) diff --git a/tests/test_commands_phase2.py b/tests/test_commands_phase2.py index 7f1b367..3ea6b99 100644 --- a/tests/test_commands_phase2.py +++ b/tests/test_commands_phase2.py @@ -859,3 +859,78 @@ def test_concat_merge_drop_is_accepted_and_keeps_no_var_columns(tmp_path): got = ad.read_h5ad(out) assert list(got.var.columns) == [] assert dict(got.uns) == {} + + +# --------------------------------------------------------------------------- +# concat refuses mismatched matrix encodings +# +# The check has been there since concat was written and nothing exercised it. +# It is the error a user is most likely to meet with real per-sample files, +# since whether a matrix lands as CSR or CSC depends on how it was made. + + +def _store_with_layout(path, layout, *, cells, layer=None): + matrix = sparse.csr_matrix(np.ones((len(cells), 3), dtype="float32")) + if layout == "csc": + matrix = matrix.tocsc() + elif layout == "dense": + matrix = matrix.toarray() + obj = ad.AnnData( + X=matrix, + obs=pd.DataFrame(index=cells), + var=pd.DataFrame(index=["g1", "g2", "g3"]), + ) + if layer is not None: + obj.layers["counts"] = ( + sparse.csr_matrix(np.ones((len(cells), 3), dtype="float32")) + if layer == "csr" + else np.ones((len(cells), 3), dtype="float32") + ) + obj.write_h5ad(path) + return path + + +@pytest.mark.parametrize( + "left,right", [("csr", "csc"), ("csr", "dense"), ("csc", "dense")] +) +def test_concat_refuses_mismatched_x_encodings(tmp_path, left, right): + a = _store_with_layout(tmp_path / "a.h5ad", left, cells=["c1", "c2"]) + b = _store_with_layout(tmp_path / "b.h5ad", right, cells=["c3", "c4"]) + out = tmp_path / "m.h5ad" + + result = runner.invoke(app, ["concat", str(a), str(b), "-o", str(out)]) + assert result.exit_code == 1 + text = _out(result) + assert "Cannot concatenate 'X'" in text + # The error has to name the way out, which until `convert` existed it + # could not do. + assert "adata convert" in text and "--layout" in text + assert not out.exists(), "nothing may be written when the check fails" + + +def test_concat_refuses_mismatched_layer_encodings(tmp_path): + a = _store_with_layout(tmp_path / "a.h5ad", "csr", cells=["c1", "c2"], + layer="csr") + b = _store_with_layout(tmp_path / "b.h5ad", "csr", cells=["c3", "c4"], + layer="dense") + out = tmp_path / "m.h5ad" + + result = runner.invoke(app, ["concat", str(a), str(b), "-o", str(out)]) + assert result.exit_code == 1 + assert "layers/counts" in _out(result) + + +def test_concat_succeeds_once_the_encodings_are_converted(tmp_path): + """The suggested fix has to actually work, so the test follows it.""" + a = _store_with_layout(tmp_path / "a.h5ad", "csr", cells=["c1", "c2"]) + b = _store_with_layout(tmp_path / "b.h5ad", "csc", cells=["c3", "c4"]) + + converted = tmp_path / "b-csr.h5ad" + assert runner.invoke( + app, ["convert", str(b), "X", "-o", str(converted), "--layout", "csr"] + ).exit_code == 0 + + out = tmp_path / "m.h5ad" + result = runner.invoke(app, ["concat", str(a), str(converted), "-o", str(out)]) + assert result.exit_code == 0, _out(result) + assert ad.read_h5ad(out).shape == (4, 3) diff --git a/tests/test_convert.py b/tests/test_convert.py new file mode 100644 index 0000000..b53e0d6 --- /dev/null +++ b/tests/test_convert.py @@ -0,0 +1,693 @@ +"""Tests for `adata convert`. + +scipy is the oracle throughout. A transpose or a cast is only correct if it +agrees exactly with what scipy would have produced from the same matrix, and +"looks plausible" is not a standard worth having for something that rewrites +people's data. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import numpy as np +import pytest +from typer.testing import CliRunner + +from adata.cli import app + +ad = pytest.importorskip("anndata", reason="anndata is required for these tests") +pd = pytest.importorskip("pandas") +sparse = pytest.importorskip("scipy.sparse") + +runner = CliRunner() +_ANSI = re.compile(r"\x1b\[[0-9;]*[a-zA-Z]") + + +def _out(result) -> str: + text = result.stdout + (result.stderr or "") + return " ".join(_ANSI.sub("", text).split()) + + +def _matrix(n_obs=40, n_var=25, density=0.2, layout="csr", seed=0, integral=True): + rng = np.random.default_rng(seed) + matrix = sparse.random( + n_obs, n_var, density=density, format="csr", dtype="float64", + random_state=rng, + ) + if integral: + # Counts, which is what issue #13 is about: float32 holds them exactly. + matrix.data = np.round(matrix.data * 100) + return matrix.tocsc() if layout == "csc" else matrix + + +def _store(path: Path, matrix, *, layers=None, raw=False) -> Path: + obj = ad.AnnData( + X=matrix, + obs=pd.DataFrame(index=[f"c{i}" for i in range(matrix.shape[0])]), + var=pd.DataFrame(index=[f"g{i}" for i in range(matrix.shape[1])]), + ) + for name, value in (layers or {}).items(): + obj.layers[name] = value + if raw: + obj.raw = obj + obj.write_h5ad(path) + return path + + +def _dense(matrix): + return matrix.toarray() if sparse.issparse(matrix) else np.asarray(matrix) + + +# --------------------------------------------------------------------------- +# layout + + +@pytest.mark.parametrize("in_memory", [False, True], ids=["streaming", "in-memory"]) +@pytest.mark.parametrize("source,target", [("csr", "csc"), ("csc", "csr")]) +def test_transpose_matches_scipy_exactly(tmp_path, source, target, in_memory): + """Both paths, both directions, against scipy's own conversion. + + Having two implementations is only worth it if they agree; that is the + whole reason the streaming one is not checked merely against itself. + """ + matrix = _matrix(layout=source) + store = _store(tmp_path / "in.h5ad", matrix) + + out = tmp_path / f"{source}-{target}-{in_memory}.h5ad" + argv = ["convert", str(store), "X", "-o", str(out), "--layout", target] + if in_memory: + argv.append("--in-memory") + result = runner.invoke(app, argv) + assert result.exit_code == 0, _out(result) + + got = ad.read_h5ad(out).X + expected = matrix.tocsc() if target == "csc" else matrix.tocsr() + assert got.format == target + assert np.array_equal(got.indptr, expected.indptr) + assert np.array_equal(got.indices, expected.indices) + assert np.array_equal(got.data, expected.data) + + +def test_transposing_twice_is_the_identity(tmp_path): + matrix = _matrix() + store = _store(tmp_path / "in.h5ad", matrix) + + mid, back = tmp_path / "mid.h5ad", tmp_path / "back.h5ad" + assert runner.invoke( + app, ["convert", str(store), "X", "-o", str(mid), "--layout", "csc"] + ).exit_code == 0 + assert runner.invoke( + app, ["convert", str(mid), "X", "-o", str(back), "--layout", "csr"] + ).exit_code == 0 + + result = ad.read_h5ad(back).X + assert result.format == "csr" + assert np.array_equal(result.indptr, matrix.indptr) + assert np.array_equal(result.indices, matrix.indices) + assert np.array_equal(result.data, matrix.data) + + +def test_the_two_transpose_paths_agree_on_an_awkward_matrix(tmp_path): + """Empty rows, a full row, and a single-entry column all in one.""" + dense = np.zeros((6, 5)) + dense[0, :] = [1, 2, 3, 4, 5] # full row + dense[3, 2] = 7 # lone entry + # rows 1, 2, 4, 5 stay empty + matrix = sparse.csr_matrix(dense) + store = _store(tmp_path / "in.h5ad", matrix) + + outputs = {} + for tag, extra in (("stream", []), ("memory", ["--in-memory"])): + out = tmp_path / f"{tag}.h5ad" + assert runner.invoke( + app, + ["convert", str(store), "X", "-o", str(out), "--layout", "csc", *extra], + ).exit_code == 0 + outputs[tag] = ad.read_h5ad(out).X + + expected = matrix.tocsc() + for tag, got in outputs.items(): + assert np.array_equal(got.indptr, expected.indptr), tag + assert np.array_equal(got.indices, expected.indices), tag + assert np.array_equal(got.data, expected.data), tag + + +# --------------------------------------------------------------------------- +# dtype -- what issue #13 asked for + + +def test_counts_stored_as_float64_convert_to_float32_and_shrink(tmp_path): + """The reported case, end to end, including that the file gets smaller. + + An early version produced an output eight times the input, because it + created the arrays with a fixed 65,536-element chunk and promoted int32 + indices to int64. A conversion asked for to halve a file must not + enlarge it, so the size is part of the test. + """ + import h5py + + matrix = _matrix(n_obs=200, n_var=120, density=0.1) + store = _store(tmp_path / "in.h5ad", matrix) + out = tmp_path / "f32.h5ad" + + result = runner.invoke( + app, ["convert", str(store), "X", "-o", str(out), "--dtype", "float32"] + ) + assert result.exit_code == 0, _out(result) + + got = ad.read_h5ad(out).X + assert got.dtype == np.dtype("float32") + assert np.array_equal(got.toarray(), matrix.toarray().astype("float32")) + + def stored(path): + with h5py.File(path) as handle: + return sum( + handle["X"][key].id.get_storage_size() + for key in ("data", "indices", "indptr") + ) + + assert stored(out) < stored(store), ( + f"X grew from {stored(store):,} to {stored(out):,} bytes on a " + "float64 -> float32 conversion" + ) + + +def test_the_index_dtype_is_preserved_unless_asked(tmp_path): + """Defaulting to int64 silently doubled every int32 store's indices.""" + import h5py + + store = _store(tmp_path / "in.h5ad", _matrix()) + with h5py.File(store) as handle: + source_dtype = handle["X/indices"].dtype + + out = tmp_path / "out.h5ad" + assert runner.invoke( + app, ["convert", str(store), "X", "-o", str(out), "--dtype", "float32"] + ).exit_code == 0 + + with h5py.File(out) as handle: + assert handle["X/indices"].dtype == source_dtype + + +def test_indices_can_be_narrowed_when_asked(tmp_path): + import h5py + + store = _store(tmp_path / "in.h5ad", _matrix()) + out = tmp_path / "out.h5ad" + result = runner.invoke( + app, + ["convert", str(store), "X", "-o", str(out), "--indices-dtype", "int32"], + ) + assert result.exit_code == 0, _out(result) + + with h5py.File(out) as handle: + assert handle["X/indices"].dtype == np.dtype("int32") + assert handle["X/indptr"].dtype == np.dtype("int32") + assert np.array_equal(ad.read_h5ad(out).X.toarray(), _matrix().toarray()) + + +# --------------------------------------------------------------------------- +# refusing before writing + + +def test_a_lossy_cast_is_refused_and_nothing_is_written(tmp_path): + """The store must be left alone, not half converted.""" + matrix = _matrix(integral=False) + matrix.data = matrix.data + 0.123456789012345 + store = _store(tmp_path / "in.h5ad", matrix) + out = tmp_path / "out.h5ad" + + result = runner.invoke( + app, ["convert", str(store), "X", "-o", str(out), "--dtype", "float32"] + ) + assert result.exit_code == 1 + text = _out(result) + assert "lossy" in text and "--force" in text + assert "do not round-trip" in text + assert not out.exists(), "a refused conversion must leave no output behind" + + +def test_force_converts_anyway(tmp_path): + matrix = _matrix(integral=False) + matrix.data = matrix.data + 0.123456789012345 + store = _store(tmp_path / "in.h5ad", matrix) + out = tmp_path / "out.h5ad" + + result = runner.invoke( + app, + ["convert", str(store), "X", "-o", str(out), "--dtype", "float32", "--force"], + ) + assert result.exit_code == 0, _out(result) + assert ad.read_h5ad(out).X.dtype == np.dtype("float32") + + +def test_a_lossless_cast_says_so(tmp_path): + """Worth telling the user: it is the question they were asking.""" + store = _store(tmp_path / "in.h5ad", _matrix()) + result = runner.invoke( + app, + ["convert", str(store), "X", "-o", str(tmp_path / "o.h5ad"), + "--dtype", "float32"], + ) + assert result.exit_code == 0 + assert "round-trips" in _out(result) + + +class _FakeGroup: + """Just enough of a sparse group for the index-range check.""" + + def __init__(self, nnz: int): + self._nnz = nnz + + def __getitem__(self, key): + assert key == "indices" + return type("D", (), {"shape": (self._nnz,)})() + + +def test_an_index_dtype_too_small_to_address_the_matrix_is_refused(): + """indices and indptr are bounded by different things, so both are checked. + + `indices` holds coordinates, bounded by the larger dimension; `indptr` + holds offsets, bounded by nnz. A matrix can legitimately need int64 for + one and not the other. + """ + from adata.core.convert import check_index_dtype + + wide = _FakeGroup(nnz=10) + with pytest.raises(ValueError, match="cannot hold this matrix's coordinates"): + check_index_dtype( + wide, np.dtype("int32"), np.dtype("int64"), (10, 3_000_000_000) + ) + + many = _FakeGroup(nnz=3_000_000_000) + with pytest.raises(ValueError, match="cannot hold this matrix's offsets"): + check_index_dtype(many, np.dtype("int64"), np.dtype("int32"), (10, 10)) + + # Narrow matrix, huge nnz: int32 coordinates are fine, offsets are not. + check_index_dtype( + _FakeGroup(nnz=10), np.dtype("int32"), np.dtype("int64"), (10, 10) + ) + + +def test_indptr_keeps_its_own_width(tmp_path): + """A store with int32 indices and int64 indptr must keep both. + + Inferring indptr's dtype from indices narrowed the offsets of any + matrix with more than 2^31 nonzeros, which corrupts it silently. + """ + import h5py + + store = _store(tmp_path / "in.h5ad", _matrix()) + with h5py.File(store, "a") as handle: + pointers = handle["X/indptr"][...] + coordinates = handle["X/indices"][...] + del handle["X/indptr"], handle["X/indices"] + handle["X"].create_dataset("indptr", data=pointers.astype("int64")) + handle["X"].create_dataset("indices", data=coordinates.astype("int32")) + + out = tmp_path / "out.h5ad" + assert runner.invoke( + app, ["convert", str(store), "X", "-o", str(out), "--dtype", "float32"] + ).exit_code == 0 + + with h5py.File(out) as handle: + assert handle["X/indptr"].dtype == np.dtype("int64"), "offsets narrowed" + assert handle["X/indices"].dtype == np.dtype("int32") + + +# --------------------------------------------------------------------------- +# density + + +def test_densify_then_sparsify_round_trips(tmp_path): + matrix = _matrix(density=0.3) + store = _store(tmp_path / "in.h5ad", matrix) + + dense_path, sparse_path = tmp_path / "dense.h5ad", tmp_path / "sparse.h5ad" + assert runner.invoke( + app, + ["convert", str(store), "X", "-o", str(dense_path), "--layout", "dense", + "--force"], + ).exit_code == 0 + got_dense = ad.read_h5ad(dense_path).X + assert not sparse.issparse(got_dense) + assert np.array_equal(_dense(got_dense), matrix.toarray()) + + assert runner.invoke( + app, + ["convert", str(dense_path), "X", "-o", str(sparse_path), + "--layout", "csr"], + ).exit_code == 0 + got = ad.read_h5ad(sparse_path).X + assert got.format == "csr" + assert np.array_equal(got.toarray(), matrix.toarray()) + + +def test_densifying_a_sparse_matrix_is_refused_when_it_would_explode(tmp_path): + """0.5% dense over a wide matrix is exactly the case that hurts.""" + matrix = _matrix(n_obs=200, n_var=2000, density=0.005) + store = _store(tmp_path / "in.h5ad", matrix) + out = tmp_path / "out.h5ad" + + result = runner.invoke( + app, ["convert", str(store), "X", "-o", str(out), "--layout", "dense"] + ) + assert result.exit_code == 1 + text = _out(result) + assert "would grow" in text and "--force" in text + assert not out.exists() + + +def test_sparsifying_a_mostly_dense_matrix_warns(tmp_path): + dense = np.ones((20, 10)) + store = _store(tmp_path / "in.h5ad", sparse.csr_matrix(dense)) + dense_path = tmp_path / "dense.h5ad" + assert runner.invoke( + app, + ["convert", str(store), "X", "-o", str(dense_path), "--layout", "dense", + "--force"], + ).exit_code == 0 + + result = runner.invoke( + app, + ["convert", str(dense_path), "X", "-o", str(tmp_path / "s.h5ad"), + "--layout", "csr"], + ) + assert result.exit_code == 0, _out(result) + assert "nonzero" in _out(result) + + +# --------------------------------------------------------------------------- +# selection and plumbing + + +def test_all_converts_x_layers_and_raw(tmp_path): + matrix = _matrix() + store = _store( + tmp_path / "in.h5ad", matrix, layers={"counts": matrix.copy()}, raw=True + ) + out = tmp_path / "out.h5ad" + + result = runner.invoke( + app, ["convert", str(store), "--all", "-o", str(out), "--layout", "csc"] + ) + assert result.exit_code == 0, _out(result) + + got = ad.read_h5ad(out) + assert got.X.format == "csc" + assert got.layers["counts"].format == "csc" + assert got.raw is not None and got.raw.X.format == "csc" + assert list(got.obs_names) == [f"c{i}" for i in range(matrix.shape[0])] + + +def test_untargeted_elements_are_carried_over_untouched(tmp_path): + matrix = _matrix() + obj = ad.AnnData( + X=matrix, + obs=pd.DataFrame( + {"group": pd.Categorical(["a", "b"] * (matrix.shape[0] // 2))}, + index=[f"c{i}" for i in range(matrix.shape[0])], + ), + var=pd.DataFrame(index=[f"g{i}" for i in range(matrix.shape[1])]), + ) + obj.layers["untouched"] = matrix.copy() + obj.obsm["X_pca"] = np.arange(matrix.shape[0] * 3, dtype="float32").reshape(-1, 3) + obj.uns["note"] = "keep me" + store = tmp_path / "in.h5ad" + obj.write_h5ad(store) + + out = tmp_path / "out.h5ad" + assert runner.invoke( + app, ["convert", str(store), "X", "-o", str(out), "--layout", "csc"] + ).exit_code == 0 + + got = ad.read_h5ad(out) + assert got.X.format == "csc" + assert got.layers["untouched"].format == "csr", "an untargeted layer changed" + assert got.uns["note"] == "keep me" + assert np.array_equal(got.obsm["X_pca"], obj.obsm["X_pca"]) + assert list(got.obs["group"]) == list(obj.obs["group"]) + + +def test_inplace_replaces_the_source(tmp_path): + matrix = _matrix() + store = _store(tmp_path / "in.h5ad", matrix) + + result = runner.invoke( + app, ["convert", str(store), "X", "--inplace", "--dtype", "float32"] + ) + assert result.exit_code == 0, _out(result) + + got = ad.read_h5ad(store) + assert got.X.dtype == np.dtype("float32") + assert np.array_equal(got.X.toarray(), matrix.toarray().astype("float32")) + assert not list(tmp_path.glob("*convert-tmp*")), "temp file left behind" + + +def test_converting_to_zarr_keeps_the_values(tmp_path): + matrix = _matrix() + store = _store(tmp_path / "in.h5ad", matrix) + out = tmp_path / "out.zarr" + + result = runner.invoke( + app, ["convert", str(store), "X", "-o", str(out), "--layout", "csc"] + ) + assert result.exit_code == 0, _out(result) + got = ad.read_zarr(out).X + assert got.format == "csc" + assert np.array_equal(got.toarray(), matrix.toarray()) + + +# --------------------------------------------------------------------------- +# argument handling + + +def test_convert_needs_something_to_do(tmp_path): + store = _store(tmp_path / "in.h5ad", _matrix()) + result = runner.invoke( + app, ["convert", str(store), "X", "-o", str(tmp_path / "o.h5ad")] + ) + assert result.exit_code == 1 + assert "at least one of --dtype" in _out(result) + + +def test_convert_needs_an_output(tmp_path): + store = _store(tmp_path / "in.h5ad", _matrix()) + result = runner.invoke(app, ["convert", str(store), "X", "--dtype", "float32"]) + assert result.exit_code == 1 + assert "Output file is required" in _out(result) + + +@pytest.mark.parametrize( + "flag,value,expected", + [ + ("--dtype", "complex128", "Unknown dtype"), + ("--indices-dtype", "float32", "Unknown dtype"), + ("--layout", "coo", "--layout must be one of"), + ], +) +def test_bad_values_are_rejected_by_name(tmp_path, flag, value, expected): + store = _store(tmp_path / "in.h5ad", _matrix()) + result = runner.invoke( + app, + ["convert", str(store), "X", "-o", str(tmp_path / "o.h5ad"), flag, value], + ) + assert result.exit_code == 1 + assert expected in _out(result) + + +def test_converting_something_that_is_not_a_matrix_says_so(tmp_path): + store = _store(tmp_path / "in.h5ad", _matrix()) + result = runner.invoke( + app, + ["convert", str(store), "obs", "-o", str(tmp_path / "o.h5ad"), + "--dtype", "float32"], + ) + assert result.exit_code == 1 + assert "Not a matrix" in _out(result) + + +def test_a_missing_entry_names_the_path(tmp_path): + store = _store(tmp_path / "in.h5ad", _matrix()) + result = runner.invoke( + app, + ["convert", str(store), "layers/nope", "-o", str(tmp_path / "o.h5ad"), + "--dtype", "float32"], + ) + assert result.exit_code == 1 + assert "layers/nope" in _out(result) + + +# --------------------------------------------------------------------------- +# what review found +# +# Five findings on the first version of this command, four of them able to +# change or destroy data while reporting success. Each gets a test. + + +def test_an_output_that_names_the_input_is_refused(tmp_path): + """Writing over the store being read from destroyed it. + + HDF5 happens to refuse the second open; Zarr does not, and the command + completed successfully leaving a store with zero nonzeros where the + data had been. + """ + matrix = _matrix() + store = _store(tmp_path / "in.h5ad", matrix) + + result = runner.invoke( + app, ["convert", str(store), "X", "-o", str(store), "--dtype", "float32"] + ) + assert result.exit_code == 1 + assert "Output path is the input" in _out(result) + assert np.array_equal(ad.read_h5ad(store).X.toarray(), matrix.toarray()) + + +def test_an_output_that_aliases_the_input_through_a_relative_path_is_refused( + tmp_path, +): + matrix = _matrix() + store = _store(tmp_path / "in.h5ad", matrix) + alias = tmp_path / "sub" / ".." / "in.h5ad" + (tmp_path / "sub").mkdir() + + result = runner.invoke( + app, ["convert", str(store), "X", "-o", str(alias), "--dtype", "float32"] + ) + assert result.exit_code == 1 + assert np.array_equal(ad.read_h5ad(store).X.toarray(), matrix.toarray()) + + +def test_an_explicitly_named_obsm_matrix_is_actually_converted(tmp_path): + """Only `layers` and `raw` were descended into, so this silently no-opped. + + The command reported success and copied the matrix over unchanged, + which is the worst way to get this wrong: the user has no signal. + """ + matrix = _matrix() + obj = ad.AnnData( + X=matrix, + obs=pd.DataFrame(index=[f"c{i}" for i in range(matrix.shape[0])]), + var=pd.DataFrame(index=[f"g{i}" for i in range(matrix.shape[1])]), + ) + obj.obsm["X_pca"] = np.ones((matrix.shape[0], 4), dtype="float64") + store = tmp_path / "in.h5ad" + obj.write_h5ad(store) + + out = tmp_path / "out.h5ad" + result = runner.invoke( + app, + ["convert", str(store), "obsm/X_pca", "-o", str(out), "--dtype", "float32"], + ) + assert result.exit_code == 0, _out(result) + + got = ad.read_h5ad(out) + assert got.obsm["X_pca"].dtype == np.dtype("float32") + assert np.array_equal(got.obsm["X_pca"], obj.obsm["X_pca"].astype("float32")) + assert got.X.dtype == matrix.dtype, "X should not have been touched" + + +def test_densifying_sums_duplicate_coordinates(tmp_path): + """A repeated coordinate means the sum, which is what scipy produces. + + Legal in a CSR store and not what anndata writes, so it takes a + hand-built file to reach -- but plain assignment kept whichever entry + came last and changed the matrix's values on the way to dense. + """ + import h5py + + store = _store(tmp_path / "in.h5ad", sparse.csr_matrix(np.zeros((2, 3)))) + with h5py.File(store, "a") as handle: + for key in ("data", "indices", "indptr"): + del handle["X"][key] + handle["X"].create_dataset("data", data=np.array([1.0, 2.0])) + handle["X"].create_dataset("indices", data=np.array([1, 1])) + handle["X"].create_dataset("indptr", data=np.array([0, 2, 2])) + + assert ad.read_h5ad(store).X.toarray()[0, 1] == 3.0, "scipy sums them" + + out = tmp_path / "dense.h5ad" + assert runner.invoke( + app, + ["convert", str(store), "X", "-o", str(out), "--layout", "dense", + "--force"], + ).exit_code == 0 + assert np.asarray(ad.read_h5ad(out).X)[0, 1] == 3.0 + + +def test_transpose_buckets_are_balanced_by_nonzeros_not_by_coordinate( + tmp_path, monkeypatch +): + """A skewed matrix must not land in one bucket. + + Single-cell matrices are skewed -- a few genes carry most of the counts + -- so equal-width coordinate bounds defeat the streaming guarantee + exactly where it matters. Measured on the matrix below, the largest of + three equal-width buckets held 88% of the entries, so the bucket rather + than the chunk set the peak. + """ + import adata.core.subset as subset_module + from adata.core.convert import describe, transpose_sparse_streaming + from adata.storage import open_store + + rng = np.random.default_rng(0) + n = 400 + rows, cols = [], [] + for row in range(n): + for _ in range(20): + rows.append(row) + cols.append( + int(rng.integers(0, 5)) if rng.random() < 0.95 + else int(rng.integers(0, n)) + ) + skewed = sparse.csr_matrix((np.ones(len(rows)), (rows, cols)), shape=(n, n)) + store = _store(tmp_path / "skew.h5ad", skewed) + + # `transpose_sparse_streaming` imports `_append` from this module when it + # runs, so this is the binding it will pick up. + per_bucket: dict = {} + original = subset_module._append + + def watching_append(dataset, values): + name = str(getattr(dataset, "name", "") or getattr(dataset, "path", "")) + if "major" in name: + per_bucket[name] = per_bucket.get(name, 0) + int(values.size) + return original(dataset, values) + + monkeypatch.setattr(subset_module, "_append", watching_append) + + out = tmp_path / "out.h5ad" + with open_store(store, "r") as src, open_store(out, "w") as dst: + transpose_sparse_streaming( + describe(src.root["X"]), + dst.root, + "X", + data_dtype=np.dtype("float64"), + index_dtype=np.dtype("int64"), + chunk=1000, + bucket_entries=1000, + ) + + sizes = sorted(per_bucket.values(), reverse=True) + assert len(sizes) > 1, f"expected several buckets, saw {per_bucket}" + assert sizes[0] <= skewed.nnz * 0.5, ( + f"the largest bucket held {sizes[0]} of {skewed.nnz} nonzeros " + f"({sizes[0] / skewed.nnz:.0%}); buckets must be balanced by count, " + "not by coordinate range" + ) + + # Read X back directly: this wrote only the matrix, not a whole store. + import h5py + + expected = skewed.tocsc() + with h5py.File(out) as handle: + got = sparse.csc_matrix( + (handle["X/data"][...], handle["X/indices"][...], + handle["X/indptr"][...]), + shape=tuple(handle["X"].attrs["shape"]), + ) + assert np.array_equal(got.toarray(), expected.toarray()), ( + "balancing the buckets changed the result" + ) diff --git a/tests/test_docs_are_accurate.py b/tests/test_docs_are_accurate.py index d474a34..f6f4608 100644 --- a/tests/test_docs_are_accurate.py +++ b/tests/test_docs_are_accurate.py @@ -100,7 +100,8 @@ def _command_path(tokens: List[str]) -> List[str]: _GROUPS = {"export", "import"} _TOP_LEVEL = { - "view", "ls", "subset", "split", "concat", "create", "export", "import", + "view", "ls", "subset", "split", "concat", "convert", "create", + "export", "import", } _SUBCOMMANDS = { "export": {"dataframe", "array", "sparse", "dict", "image"}, diff --git a/tests/test_performance.py b/tests/test_performance.py index 87c0023..f272055 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -1138,3 +1138,126 @@ def test_copying_wide_strings_does_not_read_the_whole_array(tmp_path): f"{TARGET_READ_BYTES / 1e6:.0f} MB read budget. A step larger than " "the array means the whole array is read in one go." ) + + +# --------------------------------------------------------------------------- +# convert +# +# The streaming transpose exists because the in-memory one does not scale. +# That is a claim about peak memory, so it is the peak that is asserted -- +# a correctness test cannot tell the two implementations apart, which is +# exactly why they are both allowed to exist. + + +def _sparse_store(path: Path, n_obs: int, n_var: int = 32, density: float = 0.2): + rng = np.random.default_rng(0) + matrix = sparse.random( + n_obs, n_var, density=density, format="csr", dtype="float32", + random_state=rng, + ) + ad.AnnData( + X=matrix, + obs=pd.DataFrame(index=[f"c{i}" for i in range(n_obs)]), + var=pd.DataFrame(index=[f"g{i}" for i in range(n_var)]), + ).write_h5ad(path) + return path + + +def _convert(source: Path, out: Path, **kwargs): + from adata.commands.convert import convert_store + + convert_store(source, ["X"], out, QUIET, **kwargs) + + +def test_convert_dtype_reads_grow_linearly_in_nonzeros(tmp_path): + def measure(n: int) -> int: + source = _sparse_store(tmp_path / f"cd{n}.h5ad", n) + with count_io() as io: + _convert(source, tmp_path / f"od{n}.h5ad", dtype="float64") + return io.work + + assert_grows_linearly(measure, what="convert --dtype", axis="nnz") + + +@pytest.mark.parametrize("layout", ["csc", "dense"]) +def test_convert_layout_reads_grow_linearly(tmp_path, layout): + def measure(n: int) -> int: + source = _sparse_store(tmp_path / f"cl{layout}{n}.h5ad", n) + with count_io() as io: + _convert( + source, tmp_path / f"ol{layout}{n}.h5ad", layout=layout, force=True + ) + return io.work + + assert_grows_linearly(measure, what=f"convert --layout {layout}", axis="n_obs") + + +@pytest.mark.slow +def test_the_streaming_transpose_does_not_hold_the_matrix(tmp_path): + """Peak allocation must not track nnz. This is the whole point of it. + + Measured against the in-memory path in the same run, which does hold the + matrix and therefore does grow -- so the comparison shows the difference + is real rather than an artefact of how the fixture is built. + """ + + def peak(n: int, in_memory: bool) -> int: + source = _sparse_store(tmp_path / f"tp{in_memory}{n}.h5ad", n, n_var=64) + with count_allocations() as measured: + _convert( + source, + tmp_path / f"otp{in_memory}{n}.h5ad", + layout="csc", + in_memory=in_memory, + chunk=4096, + ) + return measured[0] + + small, large = 256, 8192 + streaming = peak(large, False) / max(1, peak(small, False)) + loaded = peak(large, True) / max(1, peak(small, True)) + + assert streaming < loaded, ( + f"the streaming transpose grew {streaming:.1f}x over a 32x larger " + f"matrix and the in-memory one grew {loaded:.1f}x -- if streaming is " + "not the cheaper of the two it has no reason to exist" + ) + assert streaming <= 8.0, ( + f"streaming transpose peak grew {streaming:.1f}x for 32x the " + "nonzeros; it is supposed to be bounded by the chunk" + ) + + +def test_the_transpose_streams_unless_asked_not_to(tmp_path): + """The safe path is the default; --in-memory is opt-in. + + Asserted by watching which function runs, because the two produce + identical output and no result can distinguish them. + """ + import adata.core.convert as convert_module + + called: List[str] = [] + for name in ("transpose_sparse_streaming", "transpose_sparse_in_memory"): + original = getattr(convert_module, name) + + def record(*args, _name=name, _original=original, **kwargs): + called.append(_name) + return _original(*args, **kwargs) + + setattr(convert_module, name, record) + + try: + source = _sparse_store(tmp_path / "d.h5ad", 64) + _convert(source, tmp_path / "default.h5ad", layout="csc") + assert called == ["transpose_sparse_streaming"], called + + called.clear() + _convert(source, tmp_path / "asked.h5ad", layout="csc", in_memory=True) + assert called == ["transpose_sparse_in_memory"], called + finally: + for name in ("transpose_sparse_streaming", "transpose_sparse_in_memory"): + setattr( + convert_module, name, getattr(convert_module, name).__wrapped__ + if hasattr(getattr(convert_module, name), "__wrapped__") + else getattr(convert_module, name) + )