auroraomics 0.1.0.dev0__py3-none-any.whl

auroraomics/h5ad.py

@@ -0,0 +1,398 @@
+"""Write the standard result file with ``h5py`` alone.
+
+The result of a prediction is an ``.h5ad``: the expression matrix in ``X``, one
+row per spot and one column per gene, with the per-spot table in ``obs``, the
+per-gene table in ``var``, the spot coordinates in ``obsm["spatial"]``, the run
+provenance in ``uns`` and any derived matrix in ``layers``. It reads back as an
+ordinary ``AnnData``.
+
+**Why this module does not import anndata.** Two reasons, and the second is the
+load-bearing one:
+
+* anndata pulls pandas and its stack. This code runs inside GPU images whose
+  compiled dependencies are ABI-pinned, and beside a model that wants the
+  memory; the format is a handful of HDF5 attributes, so paying for that stack
+  to write them is a poor trade.
+* a matrix bigger than memory cannot be handed to a library that takes an
+  in-memory array. Accumulating a (300k x 19k) float32 matrix to write it in
+  one call needs ~23 GB, which is simply not available — and the model that
+  produced it is holding several GB of its own. So ``X`` is created empty,
+  chunked and compressed, and filled batch by batch as the batches arrive:
+  peak memory is one batch, whatever the slide.
+
+The format is the on-disk contract anndata 0.10 and 0.11 read, expressed as
+``encoding-type``/``encoding-version`` attributes on every element. The
+:func:`write_result` docstring names each one; the test suite holds the actual
+guarantee, by opening what this module writes with anndata itself.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping, Sequence
+from pathlib import Path
+from typing import Any
+
+import h5py
+import numpy as np
+
+from . import contracts
+
+# ── the on-disk encoding vocabulary ─────────────────────────────────────────
+_ANNDATA = {"encoding-type": "anndata", "encoding-version": "0.1.0"}
+_DICT = {"encoding-type": "dict", "encoding-version": "0.1.0"}
+_DATAFRAME = {"encoding-type": "dataframe", "encoding-version": "0.2.0"}
+_ARRAY = {"encoding-type": "array", "encoding-version": "0.2.0"}
+_STRING_ARRAY = {"encoding-type": "string-array", "encoding-version": "0.2.0"}
+_STRING = {"encoding-type": "string", "encoding-version": "0.2.0"}
+_SCALAR = {"encoding-type": "numeric-scalar", "encoding-version": "0.2.0"}
+
+#: The name a dataframe group gives its index when the caller has none. This is
+#: anndata's own convention for an unnamed index, and it round-trips as one.
+_DEFAULT_INDEX_NAME = "_index"
+
+#: Rows per chunk of a matrix. Chunks are whole rows across every column, so a
+#: chunk is one compression unit and one disk write; 1024 rows of 19k float32
+#: genes is ~78 MB, which is why the cache below matters.
+_CHUNK_ROWS = 1024
+
+#: Floor and ceiling for the HDF5 raw-chunk cache. See :func:`_cache_bytes`.
+_CACHE_FLOOR = 32 * 1024 * 1024
+_CACHE_CEILING = 256 * 1024 * 1024
+
+_STRING_DTYPE = h5py.string_dtype(encoding="utf-8")
+
+
+class ResultError(ValueError):
+    """The pieces handed to :func:`write_result` cannot form a valid result."""
+
+
+def _cache_bytes(n_cols: int, itemsize: int, n_rows: int) -> int:
+    """How much raw-chunk cache to give the file, and why it is not the default.
+
+    Batches advance monotonically through the matrix and are far smaller than a
+    chunk — a model predicts a few dozen rows at a time while a chunk is
+    :data:`_CHUNK_ROWS` of them. With HDF5's default 1 MB cache the active
+    chunk does not fit, so every batch evicts it, and every eviction
+    re-compresses and rewrites the *whole* chunk: a read-modify-write storm
+    that stalls the producer whenever the target is a slow disk. Sized to hold
+    several whole chunks, each chunk is compressed and written exactly once, as
+    it fills.
+
+    The ceiling keeps the process modest, and the floor keeps a narrow matrix
+    (few genes, so small chunks) from inheriting a cache too small to help.
+    """
+    chunk_rows = max(1, min(_CHUNK_ROWS, n_rows))
+    chunk_bytes = chunk_rows * max(1, n_cols) * itemsize
+    return int(min(max(chunk_bytes * 4, _CACHE_FLOOR), _CACHE_CEILING))
+
+
+def _is_string_data(values: Any) -> bool:
+    """Does this column hold text rather than numbers?"""
+    arr = np.asarray(values)
+    if arr.dtype.kind in "US":
+        return True
+    if arr.dtype.kind == "O":
+        return all(isinstance(v, (str, bytes)) for v in arr.ravel())
+    return False
+
+
+def _write_string_array(
+    parent: h5py.Group, name: str, values: Any, compression: str | None
+) -> h5py.Dataset:
+    data = np.asarray(
+        [v.decode() if isinstance(v, bytes) else str(v) for v in np.asarray(values)],
+        dtype=object,
+    )
+    dataset = parent.create_dataset(
+        name, data=data, dtype=_STRING_DTYPE, compression=compression
+    )
+    dataset.attrs.update(_STRING_ARRAY)
+    return dataset
+
+
+def _write_column(
+    parent: h5py.Group, name: str, values: Any, compression: str | None
+) -> None:
+    if _is_string_data(values):
+        _write_string_array(parent, name, values, compression)
+        return
+    arr = np.asarray(values)
+    if arr.ndim != 1:
+        raise ResultError(
+            f"column {name!r} must be one-dimensional, got shape {arr.shape}"
+        )
+    dataset = parent.create_dataset(name, data=arr, compression=compression)
+    dataset.attrs.update(_ARRAY)
+
+
+def _write_dataframe(
+    parent: h5py.Group,
+    name: str,
+    columns: Mapping[str, Any],
+    index: Sequence[Any],
+    index_name: str,
+    compression: str | None,
+) -> None:
+    """Write one table as a ``dataframe`` group.
+
+    The group carries the column order and the index's name as attributes: the
+    index is a dataset like any other, and ``_index`` is what says which one it
+    is. Order is an attribute rather than the file's own key order because HDF5
+    sorts its keys alphabetically, so ``column-order`` is the only place the
+    caller's column order survives.
+    """
+    if index_name in columns:
+        raise ResultError(
+            f"{name}: {index_name!r} is the index name, so it cannot also be a "
+            "column — rename one of them"
+        )
+    group = parent.create_group(name)
+    group.attrs.update(_DATAFRAME)
+    group.attrs["_index"] = index_name
+    group.attrs.create(
+        "column-order", data=list(columns), shape=(len(columns),), dtype=_STRING_DTYPE
+    )
+    _write_string_array(group, index_name, index, compression)
+    for column, values in columns.items():
+        arr = np.asarray(values)
+        if arr.shape[0] != len(index):
+            raise ResultError(
+                f"{name}: column {column!r} has {arr.shape[0]} values for "
+                f"{len(index)} rows"
+            )
+        _write_column(group, column, values, compression)
+
+
+def _write_mapping(
+    parent: h5py.Group, name: str, values: Mapping[str, Any], compression: str | None
+) -> h5py.Group:
+    """Write a nested mapping as a ``dict`` group, recursing into sub-mappings.
+
+    Scalars, strings, numeric arrays and string arrays each get the encoding
+    anndata expects for them; anything else is refused by name rather than
+    written as something a reader will mis-parse.
+    """
+    group = parent.create_group(name)
+    group.attrs.update(_DICT)
+    for key, value in values.items():
+        if not isinstance(key, str):
+            raise ResultError(f"{name}: keys must be strings, got {key!r}")
+        if isinstance(value, Mapping):
+            _write_mapping(group, key, value, compression)
+        elif isinstance(value, (str, bytes)):
+            dataset = group.create_dataset(
+                key,
+                data=value.decode() if isinstance(value, bytes) else value,
+                dtype=_STRING_DTYPE,
+            )
+            dataset.attrs.update(_STRING)
+        elif isinstance(value, (bool, int, float, np.bool_, np.number)):
+            dataset = group.create_dataset(key, data=value)
+            dataset.attrs.update(_SCALAR)
+        elif isinstance(value, (Sequence, np.ndarray)):
+            if _is_string_data(value):
+                _write_string_array(group, key, value, compression)
+            else:
+                dataset = group.create_dataset(
+                    key, data=np.asarray(value), compression=compression
+                )
+                dataset.attrs.update(_ARRAY)
+        else:
+            raise ResultError(
+                f"{name}.{key} is a {type(value).__name__}, which has no "
+                "representation in this file format; convert it to a mapping, "
+                "a string, a number or an array first"
+            )
+    return group
+
+
+def _write_matrix(
+    parent: h5py.Group,
+    name: str,
+    source: Any,
+    n_rows: int,
+    n_cols: int,
+    dtype: Any,
+    compression: str | None,
+) -> None:
+    """Create a chunked matrix and fill it from an array or a stream of batches.
+
+    ``source`` is either a two-dimensional array, or an iterable of row batches
+    that together cover exactly ``n_rows``. Batches are written where they
+    land and then dropped, so nothing here ever holds more than one of them —
+    that is the whole point of taking an iterable, and a test pins it by
+    checking that earlier batches have been collected before the last arrives.
+    """
+    dataset = parent.create_dataset(
+        name,
+        shape=(n_rows, n_cols),
+        dtype=dtype,
+        chunks=(max(1, min(_CHUNK_ROWS, n_rows)), max(1, n_cols)),
+        compression=compression,
+    )
+    dataset.attrs.update(_ARRAY)
+
+    # A whole matrix, or a stream of row batches. Anything that already knows
+    # it is two-dimensional (a numpy array, an open dataset from another file)
+    # is the whole matrix; everything else is iterated.
+    batches: Iterable[Any]
+    if getattr(source, "ndim", None) == 2:
+        batches = (source,)
+    else:
+        batches = source
+
+    written = 0
+    for index, batch in enumerate(batches):
+        arr = np.asarray(batch)
+        if arr.ndim != 2 or arr.shape[1] != n_cols:
+            raise ResultError(
+                f"{name}: batch {index} has shape {arr.shape}, expected "
+                f"(rows, {n_cols})"
+            )
+        end = written + arr.shape[0]
+        if end > n_rows:
+            raise ResultError(
+                f"{name}: batches carry more than {n_rows} rows (batch {index} "
+                f"would reach row {end})"
+            )
+        dataset[written:end, :] = arr.astype(dtype, copy=False)
+        written = end
+    if written != n_rows:
+        raise ResultError(
+            f"{name}: batches carried {written} rows for {n_rows} observations "
+            "— a result whose matrix is part zero-filled is worse than no file"
+        )
+
+
+def _row_count(columns: Mapping[str, Any] | None, index: Sequence[Any] | None) -> int:
+    if index is not None:
+        return len(index)
+    if columns:
+        first = next(iter(columns.values()))
+        return int(np.asarray(first).shape[0])
+    raise ResultError(
+        "cannot tell how many rows this result has: pass an index, or at least "
+        "one column. A streamed matrix is created before its first batch "
+        "arrives, so the shape has to be known in advance"
+    )
+
+
+def write_result(
+    path: str | Path,
+    *,
+    obs: Mapping[str, Any] | None = None,
+    var: Mapping[str, Any] | None = None,
+    x: Any = None,
+    obs_index: Sequence[Any] | None = None,
+    var_index: Sequence[Any] | None = None,
+    spatial: Any = None,
+    uns: Mapping[str, Any] | None = None,
+    layers: Mapping[str, Any] | None = None,
+    obs_index_name: str = _DEFAULT_INDEX_NAME,
+    var_index_name: str | None = None,
+    dtype: Any = np.float32,
+    compression: str | None = "gzip",
+) -> Path:
+    """Write one result file and return its path.
+
+    Parameters
+    ----------
+    path:
+        Where to write. An existing file is replaced.
+    obs, var:
+        The per-spot and per-gene tables, as mappings of column name to a
+        one-dimensional sequence. Text columns are stored as text; everything
+        else keeps its numpy dtype.
+    x:
+        The expression matrix: either a two-dimensional array, or an **iterable
+        of row batches** covering every observation in order. The iterable form
+        is what makes a matrix larger than memory writable — see
+        :func:`_write_matrix`.
+    obs_index, var_index:
+        Row and column names. Default to the positions as strings.
+    var_index_name:
+        The name ``var``'s index carries in the file. Defaults to the gene
+        identifier the shared result contract declares, so a reader knows what
+        kind of identifier it is looking at without guessing.
+    spatial:
+        An ``(n_obs, 2)`` array of spot coordinates, written to
+        ``obsm["spatial"]`` — where every viewer looks for them.
+    uns:
+        Free-form provenance, nested mappings allowed.
+    layers:
+        Extra matrices of the same shape as ``x``, each an array or an iterable
+        of batches, written under ``layers/<name>``.
+    dtype:
+        The matrix dtype. float32 by default: float16 halves the file but is
+        not enough precision for the statistics these matrices get fed to, and
+        several tools refuse it outright.
+    compression:
+        Passed to HDF5. ``"gzip"`` by default; ``None`` writes uncompressed.
+    """
+    obs = dict(obs or {})
+    var = dict(var or {})
+    n_obs = _row_count(obs, obs_index)
+    n_vars = _row_count(var, var_index)
+    if obs_index is None:
+        obs_index = [str(i) for i in range(n_obs)]
+    if var_index is None:
+        var_index = [str(i) for i in range(n_vars)]
+    if var_index_name is None:
+        var_index_name = str(contracts.result_layout()["var_index"])
+    if x is None:
+        raise ResultError("a result needs a matrix; pass x=")
+
+    if layers and "X" in layers:
+        raise ResultError(
+            "a layer cannot be called 'X' — that is the matrix itself, and one "
+            "of the two would silently win"
+        )
+    matrices: dict[str, Any] = {"X": x}
+    matrices.update(layers or {})
+    itemsize = max(np.dtype(dtype).itemsize, 1)
+
+    path = Path(path)
+    # The cache is a property of the OPEN FILE, so it has to be sized before
+    # anything is created in it — which is why the shape is settled above.
+    with h5py.File(
+        path,
+        "w",
+        rdcc_nbytes=_cache_bytes(n_vars, itemsize, n_obs),
+        rdcc_nslots=1009,
+        # Prefer evicting chunks that are fully written, which is every chunk
+        # behind the advancing write position.
+        rdcc_w0=1.0,
+    ) as handle:
+        handle.attrs.update(_ANNDATA)
+        _write_dataframe(handle, "obs", obs, obs_index, obs_index_name, compression)
+        _write_dataframe(handle, "var", var, var_index, var_index_name, compression)
+
+        obsm = handle.create_group("obsm")
+        obsm.attrs.update(_DICT)
+        if spatial is not None:
+            coords = np.asarray(spatial)
+            if coords.ndim != 2 or coords.shape[0] != n_obs:
+                raise ResultError(
+                    f"spatial must be (n_obs, dims) = ({n_obs}, ...), got shape "
+                    f"{coords.shape}"
+                )
+            dataset = obsm.create_dataset(
+                "spatial", data=coords, compression=compression
+            )
+            dataset.attrs.update(_ARRAY)
+
+        _write_mapping(handle, "uns", dict(uns or {}), compression)
+
+        layers_group = handle.create_group("layers")
+        layers_group.attrs.update(_DICT)
+        for name, source in matrices.items():
+            parent = handle if name == "X" else layers_group
+            _write_matrix(
+                parent, name, source, n_obs, n_vars, dtype, compression
+            )
+
+        # anndata writes these three even when empty, and a reader that expects
+        # them should not have to special-case a file this module produced.
+        for name in ("varm", "obsp", "varp"):
+            handle.create_group(name).attrs.update(_DICT)
+    return path