diff --git a/AGENTS.md b/AGENTS.md index e4f0c40..c6ee758 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -93,7 +93,10 @@ Each layer calls only the layers below it: Content sniffing runs before extensions, because extensions lie (`.h5`). A format with no extension, such as an MTX folder (`10x-mtx`), is chosen as - an output only with `--to`. + an output only with `--to`. Order in `FORMATS` matters: the first sniffer + that matches wins, so a more specific format (`cellbender`, which is also a + 10x matrix) goes before the general one, and an output extension picks the + first format that can be written. - **`formats/`**: the per-file-type halves of `export` and `import`: npy, mtx, csv, json and images. `mtx.py` is a streaming Matrix Market codec (two passes in, block writes out) shared by `import`/`export sparse` and the diff --git a/README.md b/README.md index 3b307ef..c615916 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ A command-line tool for exploring huge AnnData stores (`.h5ad` and `.zarr`) with - Auto-detects `.h5ad` files vs `.zarr` directories - Chunked processing for dense and sparse matrices (CSR/CSC) - Reads every AnnData on-disk layout, from 0.7.x through the current spec, and always writes the current one -- Converts between HDF5 and Zarr (v2 and v3) in either direction with `adata convert`, and to and from 10x Genomics HDF5 matrices and 10x / STARsolo MTX folders +- Converts between HDF5 and Zarr (v2 and v3) in either direction with `adata convert`, and to and from 10x Genomics HDF5 matrices and 10x / STARsolo MTX folders, and from CellBender output - Rich terminal output with progress indicators, kept on stderr so results pipe cleanly **Documentation: [cellgeni.github.io/adata-cli](https://cellgeni.github.io/adata-cli/)** @@ -56,7 +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 store's format: h5ad ↔ Zarr v2/v3, 10x Genomics HDF5 (Cell Ranger's `*_feature_bc_matrix.h5`) and 10x / STARsolo MTX folders in and out, streaming, with the output written beside its destination and moved into place only once complete. +- `convert` – change a store's format: h5ad ↔ Zarr v2/v3, 10x Genomics HDF5 (Cell Ranger's `*_feature_bc_matrix.h5`) and 10x / STARsolo MTX folders in and out, CellBender output in, streaming, with the output written beside its destination and moved into place only once complete. - `recast` – change a matrix's dtype, layout (CSR/CSC/dense) or density, streaming; refuses a lossy cast or a large size increase unless forced. Called `convert` until 0.7.0. - `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). diff --git a/docs/COMMANDS.md b/docs/COMMANDS.md index 54bf7c5..50fca25 100644 --- a/docs/COMMANDS.md +++ b/docs/COMMANDS.md @@ -122,6 +122,7 @@ adata convert counts.h5ad -o counts_mtx --to 10x-mtx | `h5ad` | `.h5ad` | yes | yes | | `zarr` | `.zarr` (v2 or v3) | yes | yes | | `10x-h5` | `.h5` | Cell Ranger 2 and 3+, Space Ranger, Xenium | Cell Ranger 3+ layout | +| `cellbender` | `.h5` | CellBender `remove-background` 0.2 and later | no | | `10x-mtx` | a folder (`--to 10x-mtx`) | Cell Ranger 2 and 3+, STARsolo (velocyto too) | Cell Ranger 3+ layout, gzipped | Formats other than h5ad and Zarr convert to and from AnnData only: to turn a @@ -196,6 +197,42 @@ is refused, and `--force` writes the names as ids and `Gene Expression` as the type. A matrix of floats is written as it is, with a warning. A file imported by `convert` exports with its run metadata and tag fields intact. +### CellBender + +The output of CellBender's `remove-background` (`output.h5` and +`output_filtered.h5`, 0.2 and later) is recognised by its latent variables, +whatever its name. Its matrix is a 10x one and is read as above; the latents +are kept, under the names CellBender's own loaders give them, so code written +against those keeps working: + +| CellBender | AnnData | +|---|---| +| denoised counts | `X` | +| per-droplet latents (`cell_probability`, `cell_size`, `background_fraction`, `droplet_efficiency`; `latent_*` in 0.2) | obs columns, NaN for droplets CellBender did not analyse | +| which droplets were analysed | `obs["cellbender_analyzed"]` | +| `gene_expression_encoding` (`latent_gene_encoding` in 0.2) | obsm, same name | +| `ambient_expression` | `var["ambient_expression"]` | +| which features were analysed (`features_analyzed_inds`, 0.3+) | `var["cellbender_analyzed"]` | +| everything else (fitted constants, learning curve, settings) | `uns["cellbender"]` | + +```bash +adata convert output.h5 -o cellbender.h5ad +adata convert output.h5 -o cells.h5ad --min-cell-probability 0.5 +adata convert output_filtered.h5 -o cb.h5ad --raw raw_feature_bc_matrix.h5 +``` + +| Flag | Meaning | +|---|---| +| `--min-cell-probability` | Keep droplets whose `cell_probability` is at least this; CellBender's own cutoff is 0.5 | +| `--raw` | Add the raw Cell Ranger matrix CellBender ran on as `layers["cellranger"]`, and the denoised counts as `layers["cellbender"]`, as `load_anndata_from_input_and_output` does | + +The full `output.h5` keeps every droplet, as `read_10x_h5` would. `--raw` +matches barcodes, so it works for `output_filtered.h5` too, whose cells +CellBender orders by count; the raw file must list the same features in the +same order, and every output barcode, or it is refused. The 10x flags above +apply as well. `--from 10x-h5` reads a CellBender file as a plain 10x matrix, +without the latents. + ### 10x Genomics and STARsolo MTX folders A folder holding `matrix.mtx`, `barcodes.tsv` and `features.tsv` (or diff --git a/docs/TESTING.md b/docs/TESTING.md index ab9be43..d60358d 100644 --- a/docs/TESTING.md +++ b/docs/TESTING.md @@ -28,6 +28,7 @@ separate jobs. | `test_recast.py` | `recast` against scipy: casts, transposes, density changes and every refusal | | `test_tenx.py` | 10x HDF5 import and export against Cell Ranger-shaped fixtures (v2, v3, multimodal, probe-based): mapping, filters, round trips, refusals | | `test_mtx.py` | 10x / STARsolo MTX folders against hand-written fixtures (Cell Ranger 2 and 3, STARsolo, velocyto, GEO prefixes): mapping, entry orders, refusals that leave nothing, round trips | +| `test_cellbender.py` | CellBender 0.2 and 0.3+ output, full and filtered: each latent at its barcode, `--min-cell-probability`, `--raw` alignment, refusals | | `test_mtx_codec.py`, `test_csr_build.py` | The streaming Matrix Market reader and writer against scipy, and the bucket sort they and `recast` share | | `test_external_readers.py` | 10x conversions, HDF5 and MTX, checked by scanpy and PyTables in their own uv environments (`integration`) | | `test_performance.py` | Complexity guards -- what an operation costs, not how long it takes (see below) | @@ -245,6 +246,7 @@ should scale with, scale only that, and hold everything else fixed. | `convert` 10x import (whole, `--split-feature-types`), export | n_barcodes | linear | | `convert` 10x MTX import (in order, shuffled, `--split-feature-types`) | n_barcodes | linear; **flat** write calls | | `convert` to 10x MTX | n_obs | **flat** read calls | +| `convert` CellBender (plain, `--min-cell-probability`, `--raw` full and filtered) | n_barcodes | linear; **flat** write calls | | `recast --dtype`, `--layout csc`, `--layout dense` | nnz, n_obs | linear; the streaming transpose's **peak memory** is bounded | The `slow` marker is on the three guards that build a 65,536-row store; they still gate diff --git a/docs/index.md b/docs/index.md index 9ba4d2c..5a016a8 100644 --- a/docs/index.md +++ b/docs/index.md @@ -70,7 +70,7 @@ in.h5ad -o out.zarr` converts a whole store, and `subset`, `split` and `convert` also reads and writes 10x Genomics HDF5 matrices (Cell Ranger's `*_feature_bc_matrix.h5`) and 10x / STARsolo Matrix Market folders, velocyto's included, matching what scanpy's `read_10x_h5` and `read_10x_mtx` -read. +read, and CellBender's `remove-background` output with its latent variables. | Element | Read | Write | |---|---|---| diff --git a/src/adata/cli.py b/src/adata/cli.py index 67b1c5b..cbc4899 100644 --- a/src/adata/cli.py +++ b/src/adata/cli.py @@ -522,6 +522,8 @@ def concat( "gene_ids_column": "--gene-ids-column", "gene_names_column": "--gene-names-column", "feature_types_column": "--feature-types-column", + "raw": "--raw", + "min_cell_probability": "--min-cell-probability", } #: Flags that only make sense for format conversion; with recast's flags, @@ -645,6 +647,21 @@ def convert( "--feature-types-column", help="10x output: var column holding feature types (default feature_types)", ), + raw: Optional[Path] = typer.Option( + None, + "--raw", + help="CellBender input: the raw Cell Ranger .h5 it ran on, added as " + "layers['cellranger']", + exists=True, + dir_okay=False, + readable=True, + ), + min_cell_probability: Optional[float] = typer.Option( + None, + "--min-cell-probability", + help="CellBender input: keep droplets with cell_probability at least " + "this (CellBender's own cutoff is 0.5)", + ), # Deprecated: recast's flags, so `convert` keeps working as it did until # 1.0.0. Hidden from --help; the old matrix paths arrive in ctx.args. inplace: bool = typer.Option(False, "--inplace", hidden=True), @@ -662,7 +679,8 @@ def convert( Genomics matrix (Cell Ranger's *_feature_bc_matrix.h5); --to 10x-mtx writes a Matrix Market folder. Between h5ad and zarr every element is copied, streamed, keeping the source's chunking and compression where the - target can express them. A 10x matrix is read as scanpy's read_10x_h5 or + target can express them. CellBender output keeps its latent variables, as + CellBender's own loader names them. A 10x matrix is read as scanpy's read_10x_h5 or read_10x_mtx (gex_only=False) reads it, keeping integer counts, and written as Cell Ranger writes it. MTX folders from Cell Ranger and STARsolo, velocyto's included, are read in two streaming passes. @@ -684,6 +702,8 @@ def convert( adata convert Solo.out/Gene/filtered -o counts.h5ad adata convert Solo.out/Velocyto/raw -o velocyto.h5ad adata convert counts.h5ad -o counts_mtx --to 10x-mtx --layer counts + adata convert cellbender_output.h5 -o cb.h5ad --raw raw_feature_bc_matrix.h5 + adata convert cellbender_output.h5 -o cells.h5ad --min-cell-probability 0.5 """ entries = list(ctx.args) recast_flags = _given(ctx, _RECAST_FLAGS) @@ -767,6 +787,8 @@ def convert( gene_ids_column=gene_ids_column, gene_names_column=gene_names_column, feature_types_column=feature_types_column, + raw=raw, + min_cell_probability=min_cell_probability, ), ) except Exception as e: diff --git a/src/adata/conversion/__init__.py b/src/adata/conversion/__init__.py index 4cbb41d..9336448 100644 --- a/src/adata/conversion/__init__.py +++ b/src/adata/conversion/__init__.py @@ -17,7 +17,7 @@ import h5py -from adata.conversion import mtx, tenx +from adata.conversion import cellbender, mtx, tenx from adata.conversion.anndata_copy import AnnDataCopy from adata.conversion.base import ConversionPlan, Converter, Format, Options from adata.storage import has_valid_anndata_root_attrs, is_zarr_path @@ -66,6 +66,15 @@ def _sniff_h5ad(path: Path) -> bool: anndata=True, sniff=is_zarr_path, ), + # Before 10x-h5: a CellBender file is a 10x matrix too, and the first + # sniffer to match wins. `--from 10x-h5` still reads one as plain 10x. + "cellbender": Format( + name="cellbender", + description="CellBender remove-background output (0.2 and later)", + extensions=(".h5",), + sniff=cellbender.sniff, + importer=cellbender.CellbenderImport(), + ), "10x-h5": Format( name="10x-h5", description="10x Genomics feature-barcode matrix (Cell Ranger v2 or v3)", @@ -136,10 +145,21 @@ def output_format(path: Path, requested: Optional[str] = None) -> str: path = Path(path) if requested is not None and requested not in FORMATS: raise ValueError(f"Unknown format {requested!r}; use one of: {_names()}.") + if requested is not None and not ( + FORMATS[requested].anndata or FORMATS[requested].exporter is not None + ): + raise ValueError(f"{requested} can be read but not written.") suffix = path.suffix.lower() + # Only formats that can be written: `.h5` names CellBender output too, + # but an `.h5` output is a 10x matrix. by_extension = next( - (fmt.name for fmt in FORMATS.values() if suffix in fmt.extensions), None + ( + fmt.name + for fmt in FORMATS.values() + if suffix in fmt.extensions and (fmt.anndata or fmt.exporter is not None) + ), + None, ) if requested is None: if by_extension is None: diff --git a/src/adata/conversion/base.py b/src/adata/conversion/base.py index 5d2ccd2..55b42e8 100644 --- a/src/adata/conversion/base.py +++ b/src/adata/conversion/base.py @@ -35,6 +35,10 @@ class Options: gene_names_column: Optional[str] = None feature_types_column: Optional[str] = None + # Reading CellBender output. + raw: Optional[Path] = None + min_cell_probability: Optional[float] = None + @dataclass class ConversionPlan: diff --git a/src/adata/conversion/cellbender.py b/src/adata/conversion/cellbender.py new file mode 100644 index 0000000..06f905e --- /dev/null +++ b/src/adata/conversion/cellbender.py @@ -0,0 +1,654 @@ +"""CellBender `remove-background` output. + +CellBender writes the denoised counts as a 10x v3 matrix, readable as any +other (see `tenx`), plus its latent variables: + +* 0.3 and later: `/droplet_latents` (per analysed droplet: cell probability, + size, background fraction, efficiency, the gene-expression encoding), + `/global_latents` (the ambient profile and fitted constants) and + `/metadata` (learning curve, analysed barcodes, settings). +* 0.2: the same kinds of value inside the matrix group itself, named + `latent_cell_probability`, `latent_scale`, `latent_RT_efficiency`, + `latent_gene_encoding`, `ambient_expression`, and so on. The matrix group + is `/background_removed` when CellBender's input was Cell Ranger 2. + +The AnnData uses the names CellBender's own loaders give, so code written +against them keeps working: per-droplet latents become obs columns (NaN for +droplets CellBender did not analyse) plus a `cellbender_analyzed` flag, the +encoding `obsm`, per-feature globals `var`, and everything else +`uns["cellbender"]`. With `--raw`, the Cell Ranger counts CellBender started +from are added as `layers["cellranger"]`, and the output's own as +`layers["cellbender"]`, as `load_anndata_from_input_and_output` does. + +Which latent belongs to which barcode follows CellBender's own rule. In the +full `output.h5`, every input barcode is a row and `barcode_indices_for_latents` +places the analysed ones. In `output_filtered.h5` there is one latent row per +barcode -- but 0.3 and later still store every analysed droplet's index, so +the index is used only when its length matches the latents and it points +inside the file. +""" + +from __future__ import annotations + +from contextlib import ExitStack +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple + +import h5py +import numpy as np +from rich.console import Console + +from adata.conversion import tenx +from adata.conversion.base import ConversionPlan, Converter, Options +from adata.conversion.features import ( + STRING_STEP, + Target, + write_obs_index, + write_var, +) +from adata.core.csr_build import CsrAppender, count_keys, scratch_group, write_runs +from adata.core.recast import Matrix +from adata.core.subset import fan_out_sparse_matrix +from adata.elements import spec +from adata.elements.read import decode_str_array +from adata.elements.write import ( + ensure_anndata_skeleton, + set_column_order, + write_dense, + write_mapping, + write_scalar, + write_string_array, +) +from adata.formats.mtx import EntryMap, iter_matrix_entries +from adata.storage import create_dataset, is_zarr_group, open_store + +#: Datasets of a 0.2 matrix group that are the matrix, not a latent. +_MATRIX_KEYS = frozenset({ + "barcodes", "data", "indices", "indptr", "shape", "features", "genes", + "gene_names", "genome", "filtered_barcodes", +}) + +_INDEX = "barcode_indices_for_latents" + +#: The cell-probability column, by CellBender version. +_PROBABILITY = ("cell_probability", "latent_cell_probability") + +#: Rows of obsm written per block. +_ROW_STEP = 1 << 16 + + +def _matrix_group(handle: h5py.File) -> Optional[str]: + if "matrix" in handle: + return "matrix" + if "background_removed" in handle: + return "background_removed" + return None + + +def _has_latents(handle: h5py.File) -> bool: + if isinstance(handle.get("droplet_latents"), h5py.Group): + return True + name = _matrix_group(handle) + return name is not None and any( + key in handle[name] for key in ("latent_cell_probability", _INDEX) + ) + + +def sniff(path: Path) -> bool: + if not tenx.sniff(path): + return False + with h5py.File(path, "r") as handle: + return _has_latents(handle) + + +# --------------------------------------------------------------------------- +# the latents + + +@dataclass +class _Latents: + version: str + #: obs column -> dataset path, one value per latent row. + per_droplet: Dict[str, str] + #: obsm key -> dataset path, one row per latent row. + embeddings: Dict[str, str] + #: var column -> dataset path, one value per feature. + per_feature: Dict[str, str] + #: uns key -> dataset path. + other: Dict[str, str] + n_latent: int + #: The file row of each latent row; None when there is one per row. + rows: Optional[np.ndarray] + probability: Optional[str] = None + n_analyzed: int = 0 + + +def _datasets(group: h5py.Group, prefix: str) -> Iterator[Tuple[str, str, h5py.Dataset]]: + for key, item in group.items(): + if isinstance(item, h5py.Dataset): + yield key, f"{prefix}/{key}", item + + +def _read_latents(handle: h5py.File, group: str, n_barcodes: int, n_features: int) -> _Latents: + """Classify every latent dataset, and work out which row each belongs to.""" + per_droplet: Dict[str, str] = {} + embeddings: Dict[str, str] = {} + per_feature: Dict[str, str] = {} + other: Dict[str, str] = {} + lengths: Dict[str, int] = {} + index_path: Optional[str] = None + + def place_global(key: str, path: str, item: h5py.Dataset) -> None: + if item.ndim == 1 and item.shape[0] == n_features and item.dtype.kind in "iuf": + per_feature[key] = path + else: + other[key] = path + + if isinstance(handle.get("droplet_latents"), h5py.Group): + version = "0.3" + for key, path, item in _datasets(handle["droplet_latents"], "droplet_latents"): + if key == _INDEX: + index_path = path + elif item.ndim == 2: + embeddings[key] = path + lengths[key] = item.shape[0] + elif item.ndim == 1: + per_droplet[key] = path + lengths[key] = item.shape[0] + else: + other[key] = path + if isinstance(handle.get("global_latents"), h5py.Group): + for key, path, item in _datasets(handle["global_latents"], "global_latents"): + place_global(key, path, item) + # Metadata is uns whatever its length: `features_analyzed_inds` has + # one entry per feature when every feature was analysed, and is + # still not a property of each feature. + if isinstance(handle.get("metadata"), h5py.Group): + for key, path, _ in _datasets(handle["metadata"], "metadata"): + other[key if key not in other else f"metadata_{key}"] = path + else: + version = "0.2" + for key, path, item in _datasets(handle[group], group): + if key in _MATRIX_KEYS: + continue + if key == _INDEX: + index_path = path + elif key.startswith("latent_") and item.ndim == 2: + embeddings[key] = path + lengths[key] = item.shape[0] + elif key.startswith("latent_") and item.ndim == 1: + per_droplet[key] = path + lengths[key] = item.shape[0] + else: + place_global(key, path, item) + + sizes = set(lengths.values()) + if len(sizes) > 1: + raise ValueError( + "CellBender's per-droplet latents differ in length: " + + ", ".join(f"{k} {n}" for k, n in lengths.items()) + "." + ) + n_latent = sizes.pop() if sizes else 0 + + index = None if index_path is None else np.asarray(handle[index_path][...], dtype=np.int64) + n_analyzed = int(index.size) if index is not None else n_latent + rows: Optional[np.ndarray] + if index is not None and index.size == n_latent and ( + n_latent == 0 or (index.min() >= 0 and index.max() < n_barcodes) + ): + if np.unique(index).size != index.size: + raise ValueError(f"{_INDEX} repeats a barcode.") + rows = index + elif n_latent in (0, n_barcodes): + rows = None + else: + raise ValueError( + f"Cannot tell which barcode each of CellBender's {n_latent:,} latent " + f"rows belongs to: the matrix has {n_barcodes:,} barcodes, and " + f"{_INDEX} {'is missing' if index is None else f'has {index.size:,} entries'}." + ) + + probability = next((name for name in _PROBABILITY if name in per_droplet), None) + return _Latents( + version, per_droplet, embeddings, per_feature, other, n_latent, rows, + probability, n_analyzed, + ) + + +def _spread(values: np.ndarray, rows: Optional[np.ndarray], n: int) -> np.ndarray: + """Per-latent-row values as per-barcode ones: NaN where none was inferred.""" + if rows is None: + return values.astype(np.float64, copy=False) + full = np.full(n, np.nan) + full[rows] = values + return full + + +# --------------------------------------------------------------------------- +# --raw + + +@dataclass +class _Raw: + path: Path + group: str + n_barcodes: int + dtype: np.dtype + #: For each raw barcode, its row in the output, or -1. + to_output: np.ndarray + #: Whether output rows follow the raw file's order. + in_order: bool + + +def _barcode_blocks( + dataset: h5py.Dataset, rows: Optional[np.ndarray] = None +) -> Iterator[np.ndarray]: + """The barcodes, or those at `rows` (ascending), a block of the column at a time. + + Rows are picked out of contiguous reads. h5py's fancy indexing pays per + index selected: 65,536 scattered rows read some 500 times slower than a + slice of the whole column. + """ + n = dataset.shape[0] + for lo in range(0, n, STRING_STEP): + hi = min(lo + STRING_STEP, n) + if rows is None: + yield decode_str_array(np.asarray(dataset[lo:hi])) + continue + first, last = np.searchsorted(rows, (lo, hi)) + if first < last: + yield decode_str_array(np.asarray(dataset[lo:hi])[rows[first:last] - lo]) + + +def _plan_raw( + path: Path, options: Options, matrix: Any, barcodes: h5py.Dataset, out_rows: Optional[np.ndarray] +) -> _Raw: + """Check the raw file against the output, and map its barcodes to output rows.""" + with h5py.File(path, "r") as handle: + if not tenx.sniff(path): + raise ValueError(f"--raw {path.name} is not a 10x HDF5 matrix.") + raw = tenx._read_matrix(handle, options.genome) + same = ( + raw.n_features == matrix.n_features + and np.array_equal(raw.ids, matrix.ids) + and np.array_equal(raw.names, matrix.names) + ) + if same and raw.feature_types is not None and matrix.feature_types is not None: + same = np.array_equal(raw.feature_types, matrix.feature_types) + if not same: + raise ValueError( + f"--raw {path.name} does not list the same features, in the same " + "order, as the CellBender output; it must be the matrix CellBender " + "was run on." + ) + + n_out = barcodes.shape[0] if out_rows is None else out_rows.size + raw_barcodes = handle[raw.group]["barcodes"] + to_output = np.full(raw.n_barcodes, -1, dtype=np.int64) + + # The full output.h5 lists the raw file's barcodes in its own order: + # compare block by block before building a lookup of them all. + if out_rows is None and raw.n_barcodes == barcodes.shape[0] and all( + np.array_equal(a, b) + for a, b in zip(_barcode_blocks(raw_barcodes), _barcode_blocks(barcodes)) + ): + to_output[:] = np.arange(raw.n_barcodes) + elif n_out: + # The output's barcodes, sorted, to look each raw block up in: + # a numpy array, rather than a dict whose per-string overhead is + # several times the barcode itself. + wanted = np.concatenate(list(_barcode_blocks(barcodes, out_rows))) + order = np.argsort(wanted, kind="stable") + ranked = wanted[order] + start = 0 + for block in _barcode_blocks(raw_barcodes): + at = np.minimum(np.searchsorted(ranked, block), ranked.size - 1) + hit = ranked[at] == block + to_output[start : start + block.size][hit] = order[at[hit]] + start += block.size + found = int((to_output >= 0).sum()) + if found != n_out: + raise ValueError( + f"{n_out - found:,} of the output's {n_out:,} barcodes are not in " + f"--raw {path.name}; it must be the matrix CellBender was run on." + ) + mapped = to_output[to_output >= 0] + return _Raw( + path=path, + group=raw.group, + n_barcodes=raw.n_barcodes, + dtype=handle[raw.group]["data"].dtype, + to_output=to_output, + in_order=bool(np.all(np.diff(mapped) > 0)), + ) + + +def _write_raw( + raw: _Raw, roots: Sequence[Any], targets: Sequence[Target], n_out: int, + n_features: int, compression: Optional[str], +) -> None: + """layers["cellranger"]: the raw counts of each output's barcodes and features.""" + emap = _entry_map(n_out, n_features, targets) + destinations = [ + [CsrAppender( + root["layers"], "cellranger", shape=(n_out, emap.target_cols[t]), + data_dtype=raw.dtype, compression=compression, + )] + for t, root in enumerate(roots) + ] + with h5py.File(raw.path, "r") as handle: + group = handle[raw.group] + source = Matrix( + tenx._AsCsr(group, (raw.n_barcodes, n_features)), spec.CSR_MATRIX, + (raw.n_barcodes, n_features), group["data"].dtype, + ) + def blocks(): + for rows, cols, values in iter_matrix_entries(source, transpose=False): + out = raw.to_output[rows] + kept = out >= 0 + yield emap.apply(out[kept], cols[kept], values[kept]) + + grouped = raw.in_order and emap.n_targets == 1 + # Entries per output row, which size the buckets. Without a feature + # filter that is each raw row's length; with one, only a counting + # pass knows, and an estimate from widths alone can put a target's + # concentrated entries all in one bucket. + counts = np.zeros(emap.n_targets * n_out, dtype=np.int64) + if grouped: + pass # streamed straight through: no buckets to size + elif emap.owner is None: + indptr = np.asarray(group["indptr"][...], dtype=np.int64) + hit = raw.to_output >= 0 + counts[raw.to_output[hit]] = np.diff(indptr)[hit] + else: + for keys, _, _ in blocks(): + count_keys(counts, keys) + + write_runs( + blocks(), + n_rows=n_out, + n_targets=emap.n_targets, + destinations=destinations, + scratch_parent=roots[0], + grouped=grouped, + canonical=False, + counts=counts, + value_dtype=raw.dtype, + chunk=1 << 20, + ) + for appenders in destinations: + for appender in appenders: + appender.finish() + + +def _entry_map(n_rows: int, n_features: int, targets: Sequence[Target]) -> EntryMap: + if len(targets) == 1 and targets[0].keep is None: + return EntryMap(n_rows=n_rows, n_cols=n_features) + owner = np.full(n_features, -1, dtype=np.int64) + local = np.zeros(n_features, dtype=np.int64) + widths = [] + for t, target in enumerate(targets): + keep = np.arange(n_features) if target.keep is None else target.keep + owner[keep] = t + local[keep] = np.arange(keep.size) + widths.append(int(keep.size)) + return EntryMap( + n_rows=n_rows, n_cols=n_features, owner=owner, local=local, + n_targets=len(targets), target_cols=widths, + ) + + +# --------------------------------------------------------------------------- +# import + + +@dataclass +class _Plan: + matrix: Any + targets: List[Target] + latents: _Latents + out_rows: Optional[np.ndarray] + raw: Optional[_Raw] + zarr_format: Optional[int] + filtered_barcodes: bool = False + extra: Dict[str, Any] = field(default_factory=dict) + + +class CellbenderImport(Converter): + options = tenx.TenxImport.options | frozenset({"--raw", "--min-cell-probability"}) + + def plan( + self, + source: Path, + output: Path, + src_format: str, + dst_format: str, + options: Options, + ) -> ConversionPlan: + with h5py.File(source, "r") as handle: + matrix, outputs, targets, warnings = tenx.plan_matrix_import( + handle, output, options + ) + latents = _read_latents( + handle, matrix.group, matrix.n_barcodes, matrix.n_features + ) + + out_rows: Optional[np.ndarray] = None + if options.min_cell_probability is not None: + if not 0 <= options.min_cell_probability <= 1: + raise ValueError("--min-cell-probability must lie between 0 and 1.") + if latents.probability is None: + raise ValueError( + "This CellBender file records no cell probabilities to " + "filter on." + ) + values = np.asarray(handle[latents.per_droplet[latents.probability]][...]) + probability = _spread(values, latents.rows, matrix.n_barcodes) + with np.errstate(invalid="ignore"): + out_rows = np.flatnonzero(probability >= options.min_cell_probability) + + raw = None + if options.raw is not None: + raw = _plan_raw( + Path(options.raw), options, matrix, + handle[matrix.group]["barcodes"], out_rows, + ) + + n_out = matrix.n_barcodes if out_rows is None else out_rows.size + described = ( + f"{matrix.n_barcodes:,} droplets ({latents.n_analyzed:,} analysed) x " + f"{matrix.n_features:,} features" + ) + if out_rows is not None: + described += ( + f", keeping {n_out:,} with cell probability ≥ " + f"{options.min_cell_probability:g}" + ) + if len(outputs) > 1: + described += f", into {len(outputs)} outputs by feature type" + if raw is not None: + described += f", with the raw counts of {raw.path.name}" + return ConversionPlan( + source=source, + src_format=src_format, + dst_format=dst_format, + outputs=outputs, + summary=f"CellBender {latents.version} → {dst_format}, {described}", + warnings=warnings, + detail=_Plan( + matrix=matrix, + targets=targets, + latents=latents, + out_rows=out_rows, + raw=raw, + zarr_format=(options.zarr_format or 3) if dst_format == "zarr" else None, + ), + ) + + def write( + self, plan: ConversionPlan, staged: Sequence[Path], console: Console + ) -> None: + detail: _Plan = plan.detail + matrix, latents, out_rows = detail.matrix, detail.latents, detail.out_rows + n_out = matrix.n_barcodes if out_rows is None else out_rows.size + + with h5py.File(plan.source, "r") as handle, ExitStack() as stack: + group = handle[matrix.group] + roots = [ + stack.enter_context( + open_store(path, "w", zarr_format=detail.zarr_format) + ).root + for path in staged + ] + per_feature = { + key: np.asarray(handle[path][...]) for key, path in latents.per_feature.items() + } + if "features_analyzed_inds" in latents.other: + # As CellBender's anndata_from_h5 adds it. + analysed = np.zeros(matrix.n_features, dtype=bool) + analysed[np.asarray(handle[latents.other["features_analyzed_inds"]][...])] = True + per_feature["cellbender_analyzed"] = analysed + with console.status("[cyan]Writing obs, var and the latents...[/]"): + for root, target in zip(roots, detail.targets): + _write_obs(root, handle, group, latents, matrix.n_barcodes, out_rows) + columns = dict(target.columns) + for key, values in per_feature.items(): + columns[key] = values if target.keep is None else values[target.keep] + write_var(root, replace(target, columns=columns)) + ensure_anndata_skeleton(root) + _write_embeddings(root, handle, latents, matrix.n_barcodes, out_rows) + tenx._write_uns(root, matrix) + if matrix.group == "background_removed": + # CellBender's group, not a genome. + del root["uns"]["cellranger"]["genome"] + _write_uns(root, handle, latents) + + # One read of X, for X and, with --raw, layers["cellbender"]. + entries = [(root, out_rows, target.keep) for root, target in zip(roots, detail.targets)] + if detail.raw is not None: + entries += [ + (root["layers"], out_rows, target.keep, "cellbender") + for root, target in zip(roots, detail.targets) + ] + with console.status("[cyan]Copying X...[/]"): + fan_out_sparse_matrix( + tenx._AsCsr(group, (matrix.n_barcodes, matrix.n_features)), "X", entries + ) + + if detail.raw is not None: + with console.status(f"[cyan]Reading the raw counts of {detail.raw.path.name}...[/]"): + _write_raw( + detail.raw, roots, detail.targets, n_out, matrix.n_features, + None if is_zarr_group(roots[0]) else "gzip", + ) + + +def _write_obs( + root: Any, handle: h5py.File, group: h5py.Group, latents: _Latents, + n_barcodes: int, out_rows: Optional[np.ndarray], +) -> None: + n = n_barcodes if out_rows is None else out_rows.size + blocks = ([str(v) for v in block] for block in _barcode_blocks(group["barcodes"], out_rows)) + obs = write_obs_index(root, n, blocks) + + analyzed = np.zeros(n_barcodes, dtype=bool) + if latents.rows is None: + analyzed[:] = latents.n_latent == n_barcodes + else: + analyzed[latents.rows] = True + columns = ["cellbender_analyzed"] + write_dense(obs, "cellbender_analyzed", analyzed if out_rows is None else analyzed[out_rows]) + for key, path in latents.per_droplet.items(): + values = _spread(np.asarray(handle[path][...]), latents.rows, n_barcodes) + write_dense(obs, key, values if out_rows is None else values[out_rows]) + columns.append(key) + set_column_order(obs, columns) + + +def _write_embeddings( + root: Any, handle: h5py.File, latents: _Latents, n_barcodes: int, + out_rows: Optional[np.ndarray], +) -> None: + """Each embedding, a block of output rows at a time, NaN where none was inferred. + + Latent rows come in CellBender's order, analysed droplets by count, and + output rows in the file's, so one block of output draws on latents from + anywhere. Reading those rows by fancy indexing is some 500 times slower + than a slice, so the latents are read in slices and sorted, by the output + block they belong to, into a scratch copy that each block then reads in + one piece. + """ + if not latents.embeddings: + return + n_out = n_barcodes if out_rows is None else out_rows.size + position = np.arange(n_barcodes) if out_rows is None else np.full(n_barcodes, -1) + if out_rows is not None: + position[out_rows] = np.arange(n_out) + # Each latent row's output row, or -1. + target = ( + position[: latents.n_latent] if latents.rows is None else position[latents.rows] + ) + kept = np.flatnonzero(target >= 0) + block_of = target[kept] // _ROW_STEP + by_block = np.argsort(block_of, kind="stable") + n_blocks = -(-n_out // _ROW_STEP) + starts = np.concatenate(([0], np.cumsum(np.bincount(block_of, minlength=n_blocks)))) + # The output row of each scratch row. + placed = target[kept][by_block] + + with scratch_group(root, "__embedding_scratch__") as scratch: + for key, path in latents.embeddings.items(): + source = handle[path] + width = source.shape[1] + dtype = source.dtype if source.dtype.kind == "f" else np.float64 + sorted_rows = create_dataset( + scratch, key, shape=(kept.size, width), dtype=dtype, + chunks=(min(max(kept.size, 1), _ROW_STEP), width), + ) + cursor = starts[:-1].copy() + for lo in range(0, latents.n_latent, _ROW_STEP): + hi = min(lo + _ROW_STEP, latents.n_latent) + here = target[lo:hi] + have = np.flatnonzero(here >= 0) + if have.size == 0: + continue + values = np.asarray(source[lo:hi])[have] + blocks = here[have] // _ROW_STEP + order = np.argsort(blocks, kind="stable") + values, blocks = values[order], blocks[order] + edges = np.flatnonzero(np.diff(blocks)) + 1 + for run, b in zip(np.split(values, edges), blocks[np.r_[0, edges]]): + sorted_rows[cursor[b] : cursor[b] + len(run)] = run + cursor[b] += len(run) + + out = create_dataset( + root["obsm"], key, shape=(n_out, width), dtype=dtype, + chunks=(min(max(n_out, 1), _ROW_STEP), width), + ) + for b in range(n_blocks): + lo = b * _ROW_STEP + block = np.full((min(_ROW_STEP, n_out - lo), width), np.nan, dtype=dtype) + a, z = starts[b], starts[b + 1] + if a < z: + block[placed[a:z] - lo] = sorted_rows[a:z] + out[lo : lo + block.shape[0]] = block + spec.set_encoding(out, spec.ARRAY) + + +def _write_uns(root: Any, handle: h5py.File, latents: _Latents) -> None: + """Everything else CellBender recorded, as `uns["cellbender"]`.""" + meta = write_mapping(root["uns"], "cellbender") + write_scalar(meta, "version", latents.version) + for key, path in latents.other.items(): + value = np.asarray(handle[path][...]) + if value.ndim == 1 and value.size == 1: + # CellBender wraps each scalar setting in a one-element array. + value = value.reshape(()) + if value.ndim == 0: + item = value.item() + write_scalar(meta, key, item.decode("utf-8", "replace") if isinstance(item, bytes) else item) + elif value.dtype.kind in "SOU": + write_string_array(meta, key, decode_str_array(value)) + else: + write_dense(meta, key, value) diff --git a/src/adata/conversion/features.py b/src/adata/conversion/features.py index d7c33b6..a45bb78 100644 --- a/src/adata/conversion/features.py +++ b/src/adata/conversion/features.py @@ -215,7 +215,7 @@ def write_var(root: Any, target: Target) -> None: """ var = write_dataframe_header(root, "var", target.var_names, list(target.columns)) for key, values in target.columns.items(): - if values.dtype.kind == "b": + if values.dtype.kind in "biuf": write_dense(var, key, values) continue categories, codes = np.unique(values, return_inverse=True) diff --git a/src/adata/conversion/tenx.py b/src/adata/conversion/tenx.py index edd0f1d..7f370cc 100644 --- a/src/adata/conversion/tenx.py +++ b/src/adata/conversion/tenx.py @@ -261,6 +261,44 @@ def __getitem__(self, key: str) -> Any: # import +def plan_matrix_import( + handle: h5py.File, output: Path, options: Options +) -> Tuple[_Matrix, List[Path], List[Target], List[str]]: + """Read a 10x matrix's description and plan the AnnData outputs from it. + + Shared by every format built on a 10x matrix: the import flags mean the + same thing for all of them. + """ + check_import_options(options) + matrix = _read_matrix(handle, options.genome) + + if matrix.version == 2 and ( + options.feature_type is not None or options.split_feature_types + ): + raise ValueError( + "This is a Cell Ranger 2 file, which records no feature types; " + "--feature-type and --split-feature-types need a v3 file." + ) + + keep = np.ones(matrix.n_features, dtype=bool) + if options.genome is not None and matrix.version == 3: + genomes = matrix.columns.get("genome") + if genomes is None: + raise ValueError("This file records no genome per feature to filter on.") + if options.genome not in set(genomes): + raise ValueError( + f"No features of genome {options.genome!r}; this file has: " + f"{', '.join(sorted(set(genomes)))}." + ) + keep &= genomes == options.genome + + outputs, targets, warnings = plan_targets( + output, matrix.names, matrix.columns, matrix.feature_types, keep, options, + noun="file", + ) + return matrix, outputs, targets, warnings + + class TenxImport(Converter): options = frozenset( {"--genome", "--var-names", "--make-unique", "--feature-type", @@ -275,38 +313,13 @@ def plan( dst_format: str, options: Options, ) -> ConversionPlan: - check_import_options(options) - with h5py.File(source, "r") as handle: - matrix = _read_matrix(handle, options.genome) - has_filtered = "filtered_barcodes" in handle[matrix.group] - - if matrix.version == 2 and ( - options.feature_type is not None or options.split_feature_types - ): - raise ValueError( - "This is a Cell Ranger 2 file, which records no feature types; " - "--feature-type and --split-feature-types need a v3 file." + matrix, outputs, targets, warnings = plan_matrix_import( + handle, output, options ) - - keep = np.ones(matrix.n_features, dtype=bool) - if options.genome is not None and matrix.version == 3: - genomes = matrix.columns.get("genome") - if genomes is None: - raise ValueError("This file records no genome per feature to filter on.") - if options.genome not in set(genomes): - raise ValueError( - f"No features of genome {options.genome!r}; this file has: " - f"{', '.join(sorted(set(genomes)))}." - ) - keep &= genomes == options.genome - - types = matrix.feature_types - outputs, targets, warnings = plan_targets( - output, matrix.names, matrix.columns, types, keep, options, noun="file" - ) + has_filtered = "filtered_barcodes" in handle[matrix.group] described = describe_import( - matrix.n_barcodes, matrix.n_features, types, targets + matrix.n_barcodes, matrix.n_features, matrix.feature_types, targets ) return ConversionPlan( diff --git a/src/adata/core/csr_build.py b/src/adata/core/csr_build.py index c5a1c9b..36b406e 100644 --- a/src/adata/core/csr_build.py +++ b/src/adata/core/csr_build.py @@ -23,7 +23,7 @@ import tempfile from contextlib import contextmanager from pathlib import Path -from typing import Any, Iterator, List, Optional, Tuple +from typing import Any, Iterator, List, Optional, Sequence, Tuple import h5py import numpy as np @@ -326,3 +326,52 @@ def finish(self) -> int: ) dataset[:] = indptr return self.nnz + + +def write_runs( + blocks: Iterator[Triple], + *, + n_rows: int, + n_targets: int, + destinations: Sequence[Sequence[CsrAppender]], + scratch_parent: Any, + grouped: bool, + canonical: bool, + counts: np.ndarray, + value_dtype: Any, + chunk: int, +) -> int: + """Append entries, given in any order, to each output's CSR matrix. + + `blocks` yield `(key, column, value)` with `key = target * n_rows + row`, + so outputs share one sort. `counts` holds entries per key, from a first + pass, and `grouped` / `canonical` what that pass learned of the order: + grouped entries go straight through, others through on-disk buckets. + `destinations[t]` are the appenders for output `t`. Returns how many + duplicate entries were summed away. + """ + summer = DuplicateSummer() + + def consume(runs: Iterator[Triple]) -> None: + for keys, minors, values in summer(runs): + targets = keys // n_rows if n_targets > 1 else None + for t in np.unique(targets) if targets is not None else (0,): + if targets is not None: + pick = targets == t + rows, cols, vals = keys[pick] - t * n_rows, minors[pick], values[pick] + else: + rows, cols, vals = keys, minors, values + for appender in destinations[int(t)]: + appender.append(rows, cols, vals) + + if grouped: + consume(grouped_runs(blocks, sort=not canonical)) + else: + indptr = np.concatenate(([0], np.cumsum(counts))) + bounds = bucket_bounds(indptr, chunk=chunk) + with scratch_group(scratch_parent, "__csr_scratch__") as scratch: + sorter = BucketSorter(scratch, bounds, value_dtype) + for keys, minors, values in blocks: + sorter.add(keys, minors, values) + consume(sorter.runs()) + return summer.summed diff --git a/src/adata/core/subset.py b/src/adata/core/subset.py index 65ae03e..5f8971c 100644 --- a/src/adata/core/subset.py +++ b/src/adata/core/subset.py @@ -476,7 +476,11 @@ def fan_out_sparse_matrix( n_major, n_minor = (n_rows, n_cols) if csr else (n_cols, n_rows) outputs: List[_SparseOutput] = [] - for dst_parent, obs_idx, var_idx in targets: + for target in targets: + # A fourth element names this target's matrix, where it differs + # from `name`: one read can then fill X and a layer alike. + dst_parent, obs_idx, var_idx = target[:3] + out_name = target[3] if len(target) > 3 else name major_idx, minor_keep = (obs_idx, var_idx) if csr else (var_idx, obs_idx) _check_sorted(major_idx, "obs" if csr else "var") out_shape = ( @@ -487,7 +491,7 @@ def fan_out_sparse_matrix( _SparseOutput( src, dst_parent, - name, + out_name, enc, out_shape, major_idx, diff --git a/src/adata/formats/mtx.py b/src/adata/formats/mtx.py index fca7577..cd84549 100644 --- a/src/adata/formats/mtx.py +++ b/src/adata/formats/mtx.py @@ -29,14 +29,10 @@ import numpy as np from adata.core.csr_build import ( - BucketSorter, CsrAppender, - DuplicateSummer, - bucket_bounds, count_keys, - grouped_runs, - scratch_group, sum_keys, + write_runs, ) from adata.core.recast import DEFAULT_CHUNK, Matrix from adata.elements import spec @@ -331,7 +327,7 @@ def read_into( destinations: Sequence[Sequence[CsrAppender]], scratch_parent: Any, *, - chunk: int = DEFAULT_CHUNK, + chunk: Optional[int] = None, block_bytes: int = TEXT_BLOCK_BYTES, ) -> int: """Pass 2: append the file's entries to each output, in CSR order. @@ -340,36 +336,18 @@ def read_into( every entry of that output. Returns how many duplicate entries were summed away. """ - blocks = ( - emap.apply(*entry) - for entry in iter_entries(header, block_bytes=block_bytes) + return write_runs( + (emap.apply(*entry) for entry in iter_entries(header, block_bytes=block_bytes)), + n_rows=emap.n_rows, + n_targets=emap.n_targets, + destinations=destinations, + scratch_parent=scratch_parent, + grouped=scanned.grouped, + canonical=scanned.canonical, + counts=scanned.counts, + value_dtype=np.float64 if header.field == "real" else np.int64, + chunk=chunk or DEFAULT_CHUNK, ) - summer = DuplicateSummer() - - def consume(runs: Iterator[Tuple[np.ndarray, np.ndarray, np.ndarray]]) -> None: - for keys, minors, values in summer(runs): - targets = keys // emap.n_rows - for t in np.unique(targets) if emap.n_targets > 1 else (0,): - if emap.n_targets > 1: - pick = targets == t - rows, cols, vals = keys[pick] - t * emap.n_rows, minors[pick], values[pick] - else: - rows, cols, vals = keys, minors, values - for appender in destinations[int(t)]: - appender.append(rows, cols, vals) - - if scanned.grouped: - consume(grouped_runs(blocks, sort=not scanned.canonical)) - else: - indptr = np.concatenate(([0], np.cumsum(scanned.counts))) - bounds = bucket_bounds(indptr, chunk=chunk) - value_type = np.float64 if header.field == "real" else np.int64 - with scratch_group(scratch_parent, "__mtx_scratch__") as scratch: - sorter = BucketSorter(scratch, bounds, value_type) - for keys, minors, values in blocks: - sorter.add(keys, minors, values) - consume(sorter.runs()) - return summer.summed # --------------------------------------------------------------------------- diff --git a/tests/conftest.py b/tests/conftest.py index 98067bf..5c07521 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,7 @@ """Pytest configuration and fixtures for adata tests.""" from contextlib import contextmanager +from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, Optional, Sequence import tempfile @@ -496,3 +497,119 @@ def write_starsolo_velocyto( for name, counts in layers.items(): write_mtx_file(folder / f"{name}.mtx", counts) return folder + + +# --------------------------------------------------------------------------- +# CellBender remove-background output +# +# As CellBender writes it (checked against its source at 0.2.2, 0.3.2 and +# 0.4.0): a 10x matrix of the denoised counts plus latent variables. 0.3 and +# later keep the latents in their own groups; 0.2 keeps them inside the +# matrix group. Analysed droplets are ordered by descending UMI count, so a +# filtered file's barcodes are not in the full file's order. + + +@dataclass +class CellbenderTruth: + """What a fixture holds, for tests to check the import against.""" + + path: Path + rows: np.ndarray # which of the full barcodes the file holds, in order + analyzed: np.ndarray # barcode indices of the analysed droplets + latents: Dict[str, np.ndarray] # per analysed droplet, by stored name + encoding: np.ndarray + ambient: np.ndarray + + +def write_cellbender_h5( + path: Path, + counts: np.ndarray, + *, + barcodes: Sequence[str], + ids: Sequence[str], + names: Sequence[str], + types: Optional[Sequence[str]] = None, + version: str = "0.3", + filtered: bool = False, + n_analyzed: Optional[int] = None, + z_dim: int = 4, + cellranger_v2: bool = False, + seed: int = 0, +) -> CellbenderTruth: + """`counts` is every droplet x features; `filtered` writes cells only.""" + rng = np.random.default_rng(seed) + n_all, n_features = counts.shape + n_analyzed = n_analyzed or max(2, n_all // 2) + analyzed = np.argsort(counts.sum(axis=1), kind="stable")[::-1][:n_analyzed] + p = rng.uniform(0, 1, n_analyzed) + p[0], p[-1] = 0.99, 0.01 # at least one cell and one empty droplet + latents = { + "cell_probability": p, + "cell_size": rng.uniform(100, 1000, n_analyzed), + "droplet_efficiency": rng.uniform(0.5, 1.5, n_analyzed), + "background_fraction": rng.uniform(0, 0.5, n_analyzed), + } + if version == "0.2": + latents = { + "latent_cell_probability": latents["cell_probability"], + "latent_scale": latents["cell_size"], + "latent_RT_efficiency": latents["droplet_efficiency"], + } + encoding = rng.normal(size=(n_analyzed, z_dim)).astype(np.float32) + ambient = rng.dirichlet(np.ones(n_features)) + + if filtered: + keep = p > 0.5 + rows = analyzed[keep] + latents = {k: v[keep] for k, v in latents.items()} + encoding = encoding[keep] + else: + rows = np.arange(n_all) + + file_barcodes = [barcodes[i] for i in rows] + file_counts = counts[rows] + if cellranger_v2: + with h5py.File(path, "w") as f: + group = f.create_group("background_removed") + _write_10x_arrays(group, file_counts, file_barcodes) + group.create_dataset("genes", data=_fixed(ids)) + group.create_dataset("gene_names", data=_fixed(names)) + group.create_dataset("genome", data=_fixed(["NA"] * n_features)) + else: + write_10x_h5( + path, file_counts, barcodes=file_barcodes, ids=ids, names=names, types=types + ) + + with h5py.File(path, "a") as f: + if version == "0.2": + group = f["background_removed" if cellranger_v2 else "matrix"] + if not filtered: + group.create_dataset("barcode_indices_for_latents", data=analyzed) + group.create_dataset("latent_gene_encoding", data=encoding) + for key, values in latents.items(): + group.create_dataset(key, data=values) + group.create_dataset("ambient_expression", data=ambient) + group.create_dataset("target_false_positive_rate", data=0.01) + group.create_dataset("lambda_multiplier", data=0.5) + group.create_dataset("training_elbo_per_epoch", data=rng.normal(size=20)) + group.create_dataset("overdispersion_mean_and_scale", data=[0.1, 0.2]) + else: + droplets = f.create_group("droplet_latents") + # In 0.3+ every analysed droplet, even in the filtered file. + droplets.create_dataset("barcode_indices_for_latents", data=analyzed) + droplets.create_dataset("gene_expression_encoding", data=encoding) + for key, values in latents.items(): + droplets.create_dataset(key, data=values) + glob = f.create_group("global_latents") + glob.create_dataset("ambient_expression", data=ambient) + glob.create_dataset("empty_droplet_size_lognormal_loc", data=5.5) + glob.create_dataset("swapping_fraction_dist_params", data=[1.0, 20.0]) + meta = f.create_group("metadata") + meta.create_dataset("learning_curve_train_elbo", data=rng.normal(size=20)) + meta.create_dataset("barcodes_analyzed", data=_fixed([barcodes[i] for i in analyzed])) + meta.create_dataset("barcodes_analyzed_inds", data=analyzed) + meta.create_dataset("features_analyzed_inds", data=np.arange(n_features)) + meta.create_dataset("target_false_positive_rate", data=[0.01]) + meta.create_dataset("estimator", data=_fixed(["mckp"])) + + return CellbenderTruth(path, rows, analyzed, latents, encoding, ambient) diff --git a/tests/test_cellbender.py b/tests/test_cellbender.py new file mode 100644 index 0000000..71e9df4 --- /dev/null +++ b/tests/test_cellbender.py @@ -0,0 +1,352 @@ +"""Tests for `adata convert` of CellBender `remove-background` output. + +Fixtures are written as CellBender writes its files (see conftest), in both +the 0.2 and 0.3+ layouts, full and filtered. What anndata reads back is +checked against what each fixture was built from: the denoised counts, and +each latent at the barcode it belongs to. +""" + +from __future__ import annotations + +import re +from pathlib import Path +from typing import List + +import h5py +import numpy as np +import pytest +from typer.testing import CliRunner + +from adata.cli import app +from tests.conftest import write_10x_h5, write_cellbender_h5 + +ad = pytest.importorskip("anndata") + +runner = CliRunner() +_ANSI = re.compile(r"\x1b\[[0-9;]*[a-zA-Z]") + +N_DROPLETS = 40 +IDS = [f"ENSG{i:05d}" for i in range(6)] + ["CD3", "CD4"] +NAMES = ["A", "B", "C", "D", "E", "F", "CD3_TotalSeqB", "CD4_TotalSeqB"] +TYPES = ["Gene Expression"] * 6 + ["Antibody Capture"] * 2 + + +def _out(result) -> str: + text = result.stdout + (result.stderr or "") + return " ".join(_ANSI.sub("", text).split()) + + +def _barcodes(n: int = N_DROPLETS) -> List[str]: + return [f"AAACCCAAG{i:07d}-1" for i in range(n)] + + +def _read(path: Path): + return ad.read_zarr(path) if path.suffix == ".zarr" else ad.read_h5ad(path) + + +def _convert(*args): + return runner.invoke(app, ["convert", *map(str, args)]) + + +def _ok(*args): + result = _convert(*args) + assert result.exit_code == 0, _out(result) + return result + + +def _raw_counts(seed: int = 0) -> np.ndarray: + counts = np.random.default_rng(seed).poisson(1.5, (N_DROPLETS, len(IDS))) + counts[-3:] = 0 # empty droplets + return counts + + +def _denoised(raw: np.ndarray) -> np.ndarray: + return np.maximum(raw - 1, 0) + + +def _cellbender(tmp_path: Path, **kwargs): + raw = _raw_counts() + name = f"cb{kwargs.get('version', '0.3')}{'f' if kwargs.get('filtered') else ''}.h5" + return write_cellbender_h5( + tmp_path / name, _denoised(raw), barcodes=_barcodes(), ids=IDS, + names=NAMES, types=TYPES, **kwargs, + ) + + +def _raw_file(tmp_path: Path, **kwargs) -> Path: + return write_10x_h5( + tmp_path / "raw_feature_bc_matrix.h5", _raw_counts(), + barcodes=kwargs.pop("barcodes", _barcodes()), ids=kwargs.pop("ids", IDS), + names=NAMES, types=TYPES, + ) + + +def _expected(truth, values: np.ndarray) -> np.ndarray: + """Per-latent-row values, at the barcodes the file holds.""" + if len(values) == len(truth.rows): + return values + full = np.full(N_DROPLETS, np.nan) + full[truth.analyzed] = values + return full[truth.rows] + + +# --------------------------------------------------------------------------- +# the mapping + + +@pytest.mark.parametrize("suffix", [".h5ad", ".zarr"]) +@pytest.mark.parametrize("filtered", [False, True], ids=["full", "filtered"]) +@pytest.mark.parametrize("version", ["0.3", "0.2"]) +def test_each_latent_lands_on_its_own_barcode(tmp_path, version, filtered, suffix): + truth = _cellbender(tmp_path, version=version, filtered=filtered) + out = tmp_path / f"out{suffix}" + result = _ok(truth.path, "-o", out) + assert f"CellBender {version}" in _out(result) + + got = _read(out) + assert np.array_equal(got.X.toarray(), _denoised(_raw_counts())[truth.rows]) + assert list(got.obs_names) == [_barcodes()[i] for i in truth.rows] + assert list(got.var_names) == NAMES + + for key, values in truth.latents.items(): + np.testing.assert_array_equal(got.obs[key].to_numpy(), _expected(truth, values), key) + analyzed = np.isin(truth.rows, truth.analyzed) + assert got.obs["cellbender_analyzed"].tolist() == analyzed.tolist() + + encoding_key = "latent_gene_encoding" if version == "0.2" else "gene_expression_encoding" + encoding = got.obsm[encoding_key] + if filtered: + np.testing.assert_array_equal(encoding, truth.encoding) + else: + position = {b: i for i, b in enumerate(truth.analyzed)} + for row, barcode in enumerate(truth.rows): + if barcode in position: + np.testing.assert_array_equal(encoding[row], truth.encoding[position[barcode]]) + else: + assert np.isnan(encoding[row]).all() + + np.testing.assert_allclose(got.var["ambient_expression"], truth.ambient) + uns = got.uns["cellbender"] + assert uns["version"] == version + assert float(uns["target_false_positive_rate"]) == 0.01 + if version == "0.3": + assert "features_analyzed_inds" in uns, "metadata stays in uns" + assert "features_analyzed_inds" not in got.var + assert got.var["cellbender_analyzed"].all(), "as CellBender's anndata_from_h5 adds it" + + +@pytest.mark.parametrize("filtered", [False, True], ids=["full", "filtered"]) +def test_latents_land_right_when_read_in_many_blocks(tmp_path, filtered, monkeypatch): + """The embedding is sorted into output blocks through scratch, and the + barcodes picked out of contiguous reads; tiny blocks make every row of + the fixture cross a block boundary.""" + from adata.conversion import cellbender + + monkeypatch.setattr(cellbender, "_ROW_STEP", 3) + monkeypatch.setattr(cellbender, "STRING_STEP", 7) + truth = _cellbender(tmp_path, filtered=filtered) + raw = _raw_file(tmp_path) + for extra in ([], ["--min-cell-probability", "0.5"]): + out = tmp_path / f"out{len(extra)}.h5ad" + _ok(truth.path, "-o", out, "--raw", raw, *extra) + got = _read(out) + position = {b: i for i, b in enumerate(truth.analyzed)} + latent_row = ( + np.arange(len(truth.rows)) if filtered + else np.array([position.get(b, -1) for b in truth.rows]) + ) + index = {barcode: i for i, barcode in enumerate(_barcodes()[r] for r in truth.rows)} + for row, barcode in enumerate(got.obs_names): + source = latent_row[index[barcode]] + expected = truth.encoding[source] if source >= 0 else np.full(truth.encoding.shape[1], np.nan) + np.testing.assert_array_equal(got.obsm["gene_expression_encoding"][row], expected) + rows = [_barcodes().index(b) for b in got.obs_names] + assert np.array_equal(got.layers["cellranger"].toarray(), _raw_counts()[rows]) + + +def test_cellbender_0_2_on_cell_ranger_2_input(tmp_path): + truth = _cellbender(tmp_path, version="0.2", cellranger_v2=True) + out = tmp_path / "out.h5ad" + _ok(truth.path, "-o", out) + got = _read(out) + assert "latent_cell_probability" in got.obs + assert "genome" not in got.uns["cellranger"], "background_removed is not a genome" + + +# --------------------------------------------------------------------------- +# sniffing + + +def test_a_cellbender_file_is_recognised_and_plain_10x_still_is_10x(tmp_path): + truth = _cellbender(tmp_path) + plain = _raw_file(tmp_path) + out = tmp_path / "plain.h5ad" + result = _ok(plain, "-o", out) + assert "10x h5 v3" in _out(result) + assert "cellbender" not in _read(out).uns + + out = tmp_path / "as10x.h5ad" + _ok(truth.path, "-o", out, "--from", "10x-h5") + got = _read(out) + assert "cell_probability" not in got.obs + assert np.array_equal(got.X.toarray(), _denoised(_raw_counts())) + + +# --------------------------------------------------------------------------- +# --min-cell-probability + + +@pytest.mark.parametrize("version", ["0.3", "0.2"]) +def test_min_cell_probability_keeps_likely_cells_in_file_order(tmp_path, version): + truth = _cellbender(tmp_path, version=version) + key = "latent_cell_probability" if version == "0.2" else "cell_probability" + expected = sorted(truth.analyzed[truth.latents[key] >= 0.5]) + + out = tmp_path / "cells.h5ad" + _ok(truth.path, "-o", out, "--min-cell-probability", "0.5") + got = _read(out) + assert list(got.obs_names) == [_barcodes()[i] for i in expected] + assert (got.obs[key] >= 0.5).all() + assert np.array_equal(got.X.toarray(), _denoised(_raw_counts())[expected]) + + +def test_min_cell_probability_is_checked(tmp_path): + truth = _cellbender(tmp_path) + result = _convert(truth.path, "-o", tmp_path / "o.h5ad", "--min-cell-probability", "2") + assert result.exit_code == 1 + assert "between 0 and 1" in _out(result) + + with h5py.File(truth.path, "a") as handle: + del handle["droplet_latents/cell_probability"] + result = _convert(truth.path, "-o", tmp_path / "o.h5ad", "--min-cell-probability", "0.5") + assert result.exit_code == 1 + assert "no cell probabilities" in _out(result) + assert not (tmp_path / "o.h5ad").exists() + + +# --------------------------------------------------------------------------- +# --raw + + +@pytest.mark.parametrize("filtered", [False, True], ids=["full", "filtered"]) +def test_raw_adds_the_cell_ranger_counts_of_the_same_barcodes( + tmp_path, filtered, monkeypatch +): + """The full output lists the raw barcodes in their own order, so the raw + counts stream straight across; the filtered one orders cells by count, + so they go through the buckets.""" + import adata.core.csr_build as csr_build + + sorters = [] + original = csr_build.BucketSorter + + def watching(*args, **kwargs): + sorters.append(1) + return original(*args, **kwargs) + + monkeypatch.setattr(csr_build, "BucketSorter", watching) + truth = _cellbender(tmp_path, filtered=filtered) + raw = _raw_file(tmp_path) + out = tmp_path / "out.h5ad" + _ok(truth.path, "-o", out, "--raw", raw) + assert bool(sorters) == filtered + + got = _read(out) + assert np.array_equal(got.layers["cellranger"].toarray(), _raw_counts()[truth.rows]) + assert np.array_equal(got.layers["cellbender"].toarray(), got.X.toarray()) + assert got.layers["cellranger"].has_sorted_indices + + +def test_raw_follows_feature_filters_and_splits(tmp_path): + truth = _cellbender(tmp_path, filtered=True) + raw = _raw_file(tmp_path) + + out = tmp_path / "gex.h5ad" + _ok(truth.path, "-o", out, "--raw", raw, "--feature-type", "Gene Expression") + got = _read(out) + assert np.array_equal(got.layers["cellranger"].toarray(), _raw_counts()[truth.rows][:, :6]) + np.testing.assert_allclose(got.var["ambient_expression"], truth.ambient[:6]) + + out = tmp_path / "split.h5ad" + _ok(truth.path, "-o", out, "--raw", raw, "--split-feature-types") + adt = _read(tmp_path / "split.Antibody_Capture.h5ad") + assert np.array_equal(adt.layers["cellranger"].toarray(), _raw_counts()[truth.rows][:, 6:]) + assert np.array_equal(adt.X.toarray(), _denoised(_raw_counts())[truth.rows][:, 6:]) + + +def test_raw_buckets_are_sized_by_exact_counts_under_a_feature_filter(tmp_path, monkeypatch): + """Bucket sizes used to be estimated as each row's length scaled by the + kept share of features, which rounds to zero for a narrow filter: every + entry then landed in one bucket, held whole in memory.""" + from adata.conversion import cellbender + + seen = [] + original = cellbender.write_runs + + def watching(blocks, **kwargs): + seen.append(kwargs["counts"].copy()) + return original(blocks, **kwargs) + + monkeypatch.setattr(cellbender, "write_runs", watching) + truth = _cellbender(tmp_path, filtered=True) + raw = _raw_file(tmp_path) + _ok(truth.path, "-o", tmp_path / "adt.h5ad", "--raw", raw, "--feature-type", "Antibody Capture") + expected = (_raw_counts()[truth.rows][:, 6:] != 0).sum(axis=1) + assert np.array_equal(seen[-1], expected) + + +@pytest.mark.parametrize( + "spoil,message", + [ + ("features", "same features"), + ("barcodes", "are not in --raw"), + ], +) +def test_a_raw_file_that_does_not_match_is_refused(tmp_path, spoil, message): + truth = _cellbender(tmp_path, filtered=True) + if spoil == "features": + raw = _raw_file(tmp_path, ids=[i + "x" for i in IDS]) + else: + raw = _raw_file(tmp_path, barcodes=[b.replace("AAA", "TTT") for b in _barcodes()]) + out = tmp_path / "out.h5ad" + result = _convert(truth.path, "-o", out, "--raw", raw) + assert result.exit_code == 1 + assert message in _out(result) + assert not out.exists() + + +def test_cellbender_flags_are_refused_for_other_formats(tmp_path): + plain = _raw_file(tmp_path) + result = _convert(plain, "-o", tmp_path / "o.h5ad", "--min-cell-probability", "0.5") + assert result.exit_code == 1 + assert "does not apply" in _out(result) + + +def test_latents_of_mismatched_lengths_are_refused(tmp_path): + truth = _cellbender(tmp_path) + with h5py.File(truth.path, "a") as handle: + del handle["droplet_latents/cell_size"] + handle["droplet_latents"].create_dataset("cell_size", data=np.ones(3)) + result = _convert(truth.path, "-o", tmp_path / "o.h5ad") + assert result.exit_code == 1 + assert "differ in length" in _out(result) + + +def test_latents_that_fit_no_barcode_are_refused(tmp_path): + truth = _cellbender(tmp_path, version="0.2") + with h5py.File(truth.path, "a") as handle: + del handle["matrix/barcode_indices_for_latents"] + result = _convert(truth.path, "-o", tmp_path / "o.h5ad") + assert result.exit_code == 1 + assert "Cannot tell which barcode" in _out(result) + + +def test_cellbender_output_cannot_be_written(tmp_path): + truth = _cellbender(tmp_path) + store = tmp_path / "s.h5ad" + _ok(truth.path, "-o", store) + result = _convert(store, "-o", tmp_path / "x.h5", "--to", "cellbender") + assert result.exit_code == 1 + assert "can be read but not written" in _out(result) + # An .h5 output is still a 10x matrix. + _ok(store, "-o", tmp_path / "x.h5") diff --git a/tests/test_external_readers.py b/tests/test_external_readers.py index 2ead4bb..23f8d93 100644 --- a/tests/test_external_readers.py +++ b/tests/test_external_readers.py @@ -26,6 +26,7 @@ from adata.cli import app from tests.conftest import ( write_10x_h5, + write_cellbender_h5, write_10x_mtx, write_10x_v2_h5, write_starsolo_velocyto, @@ -313,6 +314,48 @@ def test_scanpy_reads_the_mtx_folder_we_export(tmp_path, layout): assert problems == [], problems +@pytest.mark.parametrize("filtered", [False, True], ids=["full", "filtered"]) +def test_a_cellbender_file_imports_as_scanpy_reads_its_matrix(tmp_path, filtered): + """scanpy reads CellBender output as a 10x matrix and ignores the latents; + the matrix, barcodes and features must agree. The latents themselves are + checked in test_cellbender.py.""" + truth = write_cellbender_h5( + tmp_path / "cellbender.h5", + _counts(40, 6), + barcodes=_barcodes(40), + ids=[f"ENSG{i}" for i in range(6)], + names=[f"G{i}" for i in range(6)], + filtered=filtered, + ) + out = tmp_path / "cellbender.h5ad" + _convert(str(truth.path), "-o", str(out)) + problems = _run(["scanpy"], _COMPARE_CELLBENDER, str(truth.path), str(out)) + assert problems == [], problems + + +_COMPARE_CELLBENDER = r""" +import json, sys, warnings +warnings.filterwarnings("ignore") +import numpy as np, scanpy as sc, anndata as ad + +ref, got = sc.read_10x_h5(sys.argv[1], gex_only=False), ad.read_h5ad(sys.argv[2]) +problems = [] +if ref.shape != got.shape: + problems.append(f"shape {ref.shape} != {got.shape}") +else: + if not np.array_equal(ref.X.toarray(), got.X.toarray()): + problems.append("X differs") + if list(ref.obs_names) != list(got.obs_names): + problems.append("obs_names differ") + if list(ref.var_names) != list(got.var_names): + problems.append("var_names differ") + for column in ref.var.columns: + if list(ref.var[column].astype(str)) != list(got.var[column].astype(str)): + problems.append(f"var[{column!r}] differs") +print("RESULT " + json.dumps(problems)) +""" + + _READ_EXPORT_WITH_SCANPY = r""" import json, sys, warnings warnings.filterwarnings("ignore") diff --git a/tests/test_mtx.py b/tests/test_mtx.py index 8844d4a..185fe40 100644 --- a/tests/test_mtx.py +++ b/tests/test_mtx.py @@ -112,6 +112,7 @@ def test_each_layout_imports_as_scanpy_reads_it(tmp_path, counts, layout, suffix @pytest.mark.parametrize("order", ["cell", "gene", "shuffled"]) def test_entry_order_does_not_change_the_store(tmp_path, counts, order, monkeypatch): + import adata.core.csr_build as csr_build import adata.formats.mtx as mtx_module if order == "cell": @@ -119,7 +120,7 @@ def test_entry_order_does_not_change_the_store(tmp_path, counts, order, monkeypa def refuse(*args, **kwargs): raise AssertionError("an ordered file must not go through the buckets") - monkeypatch.setattr(mtx_module, "BucketSorter", refuse) + monkeypatch.setattr(csr_build, "BucketSorter", refuse) else: monkeypatch.setattr(mtx_module, "DEFAULT_CHUNK", 8) diff --git a/tests/test_performance.py b/tests/test_performance.py index b652e72..50926fb 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -592,6 +592,68 @@ def measure(n: int) -> int: assert_independent_of(measure, what="convert h5ad to 10x mtx read calls", axis="n_obs") +def _cellbender_pair(path: Path, n: int, filtered: bool): + from tests.conftest import write_10x_h5, write_cellbender_h5 + + counts = np.random.default_rng(0).poisson(0.5, (n, 8)) + kwargs = dict( + barcodes=[f"BC{i:08d}-1" for i in range(n)], + ids=[f"ENSG{i}" for i in range(8)], + names=[f"G{i}" for i in range(8)], + ) + truth = write_cellbender_h5(path / "cb.h5", counts, filtered=filtered, **kwargs) + raw = write_10x_h5(path / "raw.h5", counts, **kwargs) + return truth.path, raw + + +@pytest.mark.parametrize( + "filtered,extra", + [ + (False, {}), + (False, {"min_cell_probability": 0.5}), + (False, {"raw": True}), + (True, {"raw": True}), + ], +) +def test_cellbender_import_work_grows_linearly(tmp_path, filtered, extra): + """Latents, obs and the embedding stream with X; --raw reads the raw + matrix once, through the buckets when the rows come in another order.""" + from adata.commands.convert import convert_store + from adata.conversion import Options + + tag = f"{filtered}{sorted(extra)}".replace(" ", "") + + def measure(n: int) -> int: + where = tmp_path / f"cb{tag}{n}" + where.mkdir() + source, raw = _cellbender_pair(where, n, filtered) + options = dict(extra) + if options.pop("raw", False): + options["raw"] = raw + with count_io() as io: + convert_store(source, where / "out.h5ad", QUIET, options=Options(**options)) + return io.work + + assert_grows_linearly( + measure, what=f"convert CellBender ({'filtered' if filtered else 'full'}, {extra})", + axis="n_barcodes", + ) + + +def test_cellbender_embedding_is_written_in_blocks(tmp_path): + from adata.commands.convert import convert_store + + def measure(n: int) -> int: + where = tmp_path / f"ce{n}" + where.mkdir() + source, _ = _cellbender_pair(where, n, False) + with count_io() as io: + convert_store(source, where / "out.h5ad", QUIET) + return io.h5_write_calls + + assert_independent_of(measure, what="convert CellBender write calls", axis="n_barcodes") + + def test_split_reads_grow_linearly_in_group_count(tmp_path, monkeypatch): """Splitting into k groups must stay linear in k at fixed store size.