B8c: Pickles, spatial and images

docs/guides/inspect-export.md

@@ -64,7 +64,10 @@ missing, and text as fixed-width unicode.
 ## Metadata and images
 
 `export dict` writes a group or scalar as JSON, and `export image` writes a
-2-D or 3-D array as PNG, JPEG or TIFF, such as a Visium tissue image:
+2-D or 3-D array as PNG, JPEG or TIFF, such as a Visium tissue image. It
+keeps the values, a 16-bit image as a 16-bit PNG or TIFF, and refuses ones
+an image cannot hold unless `--scale` stretches them onto 0-255 (see
+[`export`](../reference/commands.md#export)):
 
 ```bash
 adata export dict data.h5ad uns/neighbors -o neighbors.json

docs/reference/commands.md

@@ -460,6 +460,20 @@ with `export dataframe`, which writes them as empty fields.
 `export dict` writes NaN and infinite values as `null`, since JSON has no
 number for them, and says how many it wrote.
 
+`export image` keeps the values as they are. Whole numbers from 0 to 255
+are an 8-bit image, whatever their dtype. One channel of `uint16`, or of
+whole numbers up to 65,535, is a 16-bit image as `.png` or `.tiff`, the
+formats that keep 16 bits; a colour image keeps 8 bits a channel. Floats
+from 0 to 1 are a normalised image, scaled up to 0-255; booleans become 0
+and 255. Values an image cannot keep (negative, too large for the format,
+not finite) are refused, naming their range, rather than clipped.
+`--scale` stretches them onto 0-255 instead, the smallest to 0 and the
+largest to 255, with anything not finite at 0:
+
+```bash
+adata export image data.h5ad uns/spatial/morphology -o morphology.png --scale
+```
+
 A file output is written beside its destination and renamed into place once
 complete, so an export that fails or is refused leaves an existing file as it
 was. An output that names the input store, lies inside a Zarr input, or names

docs/reference/formats/cistopic.md

@@ -29,20 +29,51 @@ adata convert imputed_acc_obj.pkl -o imputed.h5ad --dense
 | `projections["cell"]` and `["region"]`, such as UMAP | `obsm["X_<name>"]`, `varm["X_<name>"]` |
 | `project`, `path_to_fragments`, the topic names, the model's metrics, coherence, parameters and other tables | `uns["pycistopic"]` |
 
-A `CistopicImputedFeatures` becomes X, cells x features, CSR (`--dense`
-keeps it dense), with `feature_names` and `cell_names` as var and obs names.
-
-The pickle is read without running anything from it: an unpickler that
-resolves only the names a pycisTopic object holds (numpy arrays, scipy
-sparse matrices, pandas frames, a few builtins) and refuses any other by
-name, so a pickle naming `os.system` is refused rather than run.
-pycisTopic's own classes are stood in for by plain holders of their
-attributes. A pickle written by pycisTopic's pandas 1.5 reads with pandas 2
-or 3, as CI checks. Columns are written as anndata writes them: Categorical
-as categorical, nullable integers and booleans as nullable, numbers held as
+A `CistopicImputedFeatures` becomes X, cells x features, CSR, or dense with
+`--dense`, with `feature_names` and `cell_names` as var and obs names. Its
+matrix is dense for imputed accessibility and sparse for gene activity;
+`--dense` makes X dense either way. A matrix that is not numbers is refused.
+
+## Pickles are trusted input
+
+Read a pickle only from a source you would trust with the data itself. It
+is read without running anything from it: an unpickler that resolves only
+the names a pycisTopic object holds (numpy arrays, scipy sparse matrices,
+pandas frames, a few builtins) and refuses any other by name, so a pickle
+naming `os.system` is refused rather than run. pycisTopic's own classes are
+stood in for by plain holders of their attributes.
+
+What those names are given is checked too, before numpy, scipy or pandas
+act on it, and a pickle that fails a check is refused before anything is
+written:
+
+- every numpy dtype and array is one numpy itself would write: a plain
+  dtype (no structured ones), and data that fills the array's shape
+  exactly;
+- pandas' and scipy's classes are created empty and given their state, as
+  every pickle of them does, never constructed from the file's arguments;
+- every sparse matrix has its indices within its shape and an `indptr`
+  that rises from 0 to its number of entries, every categorical's codes
+  lie within its categories, and every table's columns hold as many rows
+  as its index.
+
+Dates and times are not read; pycisTopic's objects hold none.
+
+A pickle written by pycisTopic's pandas 1.5 reads with pandas 2 or 3, as CI
+checks. Columns are written as anndata writes them: Categorical as
+categorical, nullable integers and booleans as nullable, numbers held as
 `object` as numbers, repeated text as categorical. Anything else on the
 object is named in a warning and left out.
 
+## Memory
+
 A pickle cannot be read in part, so it is held in memory whole, as it was
-when pycisTopic saved it; the fragment matrix is turned into CSR on disk by
-the streaming transpose rather than copied in memory.
+when pycisTopic saved it. Reading it checks each value, which takes about
+ten times as long as Python's own `pickle.load`: under a second for a
+200 MB pickle. The
+fragment matrix is turned into CSR on disk by the streaming transpose
+rather than copied in memory. An imputed matrix is written a block of about
+a million entries at a time (a block of cells, or of one cell's features
+when a cell has more), whatever the number of features, and a dense X is
+chunked the same way. A sparse matrix written with `--dense` goes through
+CSR on disk first, then a block of cells at a time.

docs/reference/formats/xenium.md

@@ -30,7 +30,9 @@ adata convert xenium_bundle -o region1.h5ad --feature-type "Gene Expression"
 
 Xenium ships `analysis/` packed as `analysis.tar.gz`. When it has not been
 unpacked, the three tables above are copied out of it, by name, into a
-temporary folder beside the output, removed when the conversion ends. A cell an
+temporary folder beside the output (in the nearest of its folders that
+exists, when the output goes into a new one), removed when the conversion
+ends. A cell an
 analysis table lacks is left empty (no cluster, NaN), with a warning
 giving the count.
 

src/adata/cli.py

@@ -1445,21 +1445,30 @@ def export_image(
         "-o",
         help="Output image file (.png, .jpg, .tiff); defaults to PNG on stdout",
     ),
+    scale: bool = typer.Option(
+        False,
+        "--scale",
+        help="Stretch the values onto 0-255 (smallest to 0, largest to 255), "
+        "for values an image cannot keep as they are",
+    ),
 ) -> None:
     """
     Export an image-like array to PNG/JPG/TIFF format.
 
     The array should be 2D (H,W) or 3D (H,W,C) with C in {1,3,4}. The
     format follows the output's extension; without --output the image goes
-    to stdout as PNG.
+    to stdout as PNG. Values are kept as they are: one channel of 16 bits
+    as a 16-bit PNG or TIFF. Values an image cannot keep are refused, not
+    clipped; --scale stretches them onto 0-255.
 
     Examples:
         adata export image data.h5ad uns/spatial/image -o tissue.png
         adata export image data.h5ad uns/spatial/image > tissue.png
+        adata export image data.h5ad uns/spatial/morphology -o morphology.png --scale
     """
 
     try:
-        export_image_cmd(file=file, obj=entry, out=output, console=console)
+        export_image_cmd(file=file, obj=entry, out=output, console=console, scale=scale)
     except Exception as e:
         console.print(f"[bold red]Error:[/] {escape(str(e))}")
         raise typer.Exit(code=1)

src/adata/commands/export.py

@@ -128,9 +128,11 @@ def export_json(
         )
 
 
-def export_image(file: Path, obj: str, out: Optional[Path], console: Console) -> None:
+def export_image(
+    file: Path, obj: str, out: Optional[Path], console: Console, scale: bool = False
+) -> None:
     with _output(file, out, console) as target, open_store(file, "r") as store:
-        export_image_format(store.root, obj=obj, out=target, console=console)
+        export_image_format(store.root, obj=obj, out=target, console=console, scale=scale)
 
 
 __all__ = [

src/adata/conversion/cistopic.py

@@ -24,14 +24,17 @@
 A pickle cannot be read in part, so the object is held in memory, as it was
 when pycisTopic saved it. What is written from it is not copied again: the
 fragment matrix, regions x cells CSR, is the same three arrays as a CSC cells
-x regions, and is turned into CSR on disk by the streaming transpose.
+x regions, and is turned into CSR on disk by the streaming transpose. A
+dense imputed matrix is written `BLOCK_ENTRIES` at a time, which is also
+the dense output's chunk, so neither grows with the number of features; a
+sparse one asked for dense goes through that CSR on disk first.
 """
 
 from __future__ import annotations
 
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Sequence, Tuple
+from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple
 
 import numpy as np
 from rich.console import Console
@@ -69,8 +72,11 @@
     "topic_region",
 })
 
-#: Cells of a dense imputed matrix written at once.
-DENSE_CELLS = 4096
+#: Entries of an imputed matrix written at once: a block of cells, every
+#: feature of each, or a run of one cell's features when a cell has more.
+#: Also the chunk of a dense output. Sparse output from a dense matrix takes
+#: about 40 bytes an entry while a block is sorted out, so 40 MiB at most.
+BLOCK_ENTRIES = 1 << 20
 
 
 def require_extra() -> None:
@@ -130,8 +136,10 @@ def plan(
         name = Path(source).name
         try:
             obj = pickled.load(source)
-        except pickled.RefusedGlobal as exc:
+        except pickled.RefusedPickle as exc:
             raise ValueError(f"{name}: {exc}") from None
+        except Exception as exc:  # numpy's, scipy's or pandas' own refusal, or a file cut short
+            raise ValueError(f"{name} could not be read as a pycisTopic pickle: {exc}") from None
         kind = getattr(obj, "kind", type(obj).__name__)
         warnings: List[str] = []
         if kind == "CistopicObject":
@@ -294,6 +302,11 @@ def _plan_imputed(obj: Any, name: str) -> _Import:
             f"{name}: its matrix is {shape}, but it names {len(features):,} features "
             f"and {len(cells):,} cells."
         )
+    if not _is_sparse(matrix):
+        if not isinstance(matrix, np.ndarray) or matrix.ndim != 2:
+            raise ValueError(f"{name}: its matrix is a {type(matrix).__name__}, not an array.")
+        if matrix.dtype.kind not in "biuf":
+            raise ValueError(f"{name}: its matrix holds {matrix.dtype} values, not numbers.")
     return _Import("CistopicImputedFeatures", obj, len(cells), len(features), cells, features)
 
 
@@ -350,35 +363,90 @@ def _write_csr(parent: Any, name: str, matrix: Any, shape: Tuple[int, int], deta
         create_dataset(group, key, data=values, **kw)
 
 
+def _tiles(n_cells: int, n_features: int) -> Tuple[int, int]:
+    """Cells and features a block holds, `BLOCK_ENTRIES` at most: whole
+    cells when one fits, else a run of one cell's features."""
+    budget = max(1, BLOCK_ENTRIES)
+    if n_features <= budget:
+        return max(1, budget // max(1, n_features)), max(1, n_features)
+    return 1, budget
+
+
+def _blocks(n_cells: int, n_features: int) -> Iterator[Tuple[slice, slice]]:
+    """(cells, features) of each block, in row order: a cell is split only
+    when it is a block of its own."""
+    rows, cols = _tiles(n_cells, n_features)
+    for lo in range(0, n_cells, rows):
+        for start in range(0, n_features, cols):
+            yield slice(lo, min(lo + rows, n_cells)), slice(start, min(start + cols, n_features))
+
+
+def _dense_target(root: Any, detail: _Import, dtype: np.dtype) -> Any:
+    """X, dense, chunked as the blocks it is written in."""
+    shape = (detail.n_obs, detail.n_var)
+    kw: Dict[str, Any] = {}
+    if all(shape):  # no chunk is larger than the array, nor compressed without one
+        rows, cols = _tiles(*shape)
+        kw = {"chunks": (min(rows, shape[0]), min(cols, shape[1]))}
+        if detail.compression is not None:
+            kw = with_compression(
+                kw, detail.compression, shape=shape, dtype=dtype,
+                backend="zarr" if is_zarr_group(root) else "hdf5", zarr_format=detail.zarr_format,
+            )
+    target = create_dataset(root, "X", shape=shape, dtype=dtype, **kw)
+    spec.set_encoding(target, spec.ARRAY)
+    return target
+
+
 def _write_imputed(root: Any, detail: _Import) -> None:
-    """The features x cells matrix as cells x features, a block of cells at a time."""
+    """The features x cells matrix as cells x features, a block at a time."""
     matrix = detail.obj.mtx
     if _is_sparse(matrix):
-        _write_fragments(root, "X", matrix, detail)
+        if detail.dense:
+            _write_dense_from_sparse(root, matrix, detail)
+        else:
+            _write_fragments(root, "X", matrix, detail)
         return
-    matrix = np.asarray(matrix)
     shape = (detail.n_obs, detail.n_var)
     dtype = np.dtype(matrix.dtype.str)
     if detail.dense:
-        kw: Dict[str, Any] = {"chunks": (min(max(shape[0], 1), DENSE_CELLS), max(shape[1], 1))}
-        if detail.compression is not None:
-            kw = with_compression(
-                kw, detail.compression, shape=shape, dtype=dtype,
-                backend="zarr" if is_zarr_group(root) else "hdf5", zarr_format=detail.zarr_format,
-            )
-        target = create_dataset(root, "X", shape=shape, dtype=dtype, **kw)
-        spec.set_encoding(target, spec.ARRAY)
-        for lo in range(0, shape[0], DENSE_CELLS):
-            target[lo : lo + DENSE_CELLS] = matrix[:, lo : lo + DENSE_CELLS].T
+        target = _dense_target(root, detail, dtype)
+        for cells, features in _blocks(*shape):
+            target[cells, features] = matrix[features, cells].T
         return
     appender = CsrAppender(root, "X", shape=shape, data_dtype=dtype, compression=detail.compression)
-    for lo in range(0, shape[0], DENSE_CELLS):
-        block = matrix[:, lo : lo + DENSE_CELLS].T
+    for cells, features in _blocks(*shape):
+        block = matrix[features, cells].T
         rows, cols = np.nonzero(block)
-        appender.append(rows + lo, cols, block[rows, cols])
+        appender.append(rows + cells.start, cols + features.start, block[rows, cols])
     appender.finish()
 
 
+def _write_dense_from_sparse(root: Any, matrix: Any, detail: _Import) -> None:
+    """A sparse features x cells matrix as dense cells x features: as CSR
+    first, on disk through the streaming transpose, then a block of cells at
+    a time from that, entries repeated at a place summed as scipy sums
+    them."""
+    dtype = np.dtype(matrix.dtype.str)
+    with scratch_group(root, "__cistopic_dense__") as scratch:
+        _write_fragments(scratch, "csr", matrix, detail)
+        csr = scratch["csr"]
+        indptr = np.asarray(csr["indptr"][...], dtype=np.int64)
+        target = _dense_target(root, detail, dtype)
+        read = None  # the cells whose entries are held: one read for all their runs of features
+        for cells, features in _blocks(detail.n_obs, detail.n_var):
+            if read != cells:
+                read = cells
+                start, end = int(indptr[cells.start]), int(indptr[cells.stop])
+                cols = np.asarray(csr["indices"][start:end], dtype=np.int64)
+                values = np.asarray(csr["data"][start:end])
+                rows = np.repeat(np.arange(cells.stop - cells.start), np.diff(indptr[cells.start : cells.stop + 1]))
+            block = np.zeros((cells.stop - cells.start, features.stop - features.start), dtype=dtype)
+            keep = (cols >= features.start) & (cols < features.stop)
+            np.add.at(block, (rows[keep], cols[keep] - features.start), values[keep])
+            target[cells, features] = block
+
+
 def _write_uns(uns: Any, detail: _Import, warnings: List[str]) -> None:
     from adata.formats.pandas_frame import write_frame
 

src/adata/conversion/xenium.py

@@ -109,15 +109,23 @@ def _bundle(source: Path) -> Path:
     return bundle
 
 
-def _unpack_analysis(archive: Path, beside: Path) -> tempfile.TemporaryDirectory:
+def _unpack_analysis(archive: Path, output: Path) -> tempfile.TemporaryDirectory:
     """The analysis tables this reads, out of Xenium's analysis.tar.gz.
 
     Only those members, each copied to a path built here rather than the
     one the archive names, so the archive cannot write anywhere else. Into
     a folder beside the output rather than the system's temporary one,
-    which on a cluster node is often too small for a large run's PCA.
+    which on a cluster node is often too small for a large run's PCA: the
+    nearest of the output's folders that exists, since this runs while
+    planning and the output's own is made only once every check has passed.
     """
-    holder = tempfile.TemporaryDirectory(prefix=".adata-xenium-", dir=beside)
+    beside = Path(output).resolve().parent
+    while not beside.exists() and beside != beside.parent:
+        beside = beside.parent
+    try:
+        holder = tempfile.TemporaryDirectory(prefix=".adata-xenium-", dir=beside)
+    except OSError as exc:
+        raise ValueError(f"Cannot unpack {archive.name} in {beside}: {exc.strerror or exc}.") from None
     wanted = {str(Path("analysis") / p): p for p in analysis_paths()}
     with tarfile.open(archive, "r:*") as tar:
         for member in tar:
@@ -156,7 +164,7 @@ def plan(
         unpacked = None
         packed_but_empty = False
         if not analysis.is_dir() and (bundle / "analysis.tar.gz").is_file():
-            unpacked = _unpack_analysis(bundle / "analysis.tar.gz", Path(output).resolve().parent)
+            unpacked = _unpack_analysis(bundle / "analysis.tar.gz", Path(output))
             analysis = Path(unpacked.name)
             packed_but_empty = not any(analysis.rglob("*.csv"))
         warnings: List[str] = []