"""
Zstd decoder wrapper for GZCM v3 compression.
Uses zlib as fallback since zstandard may not be installed.
Examples
--------
"""
__author__ = "Yeremia Gunawan Adhisantoso"
__email__ = "adhisant@tnt.uni-hannover.de"
__license__ = "Clear BSD"
import zlib
import numpy as np
try:
import zstandard as zstd
HAS_ZSTD = True
except ImportError:
HAS_ZSTD = False
[docs]class ZstdDecoder:
"""Zstd decoder for contact matrix tiles.
Parameters
----------
tile_size : int, default=256
Tile size for block processing.
resolution : int, default=50000
Hi-C resolution in bp.
dtype : np.dtype, default=np.uint32
Data type for decoded tiles.
use_zstd : bool, default=True
Use zstd if available, otherwise zlib fallback.
Examples
--------
"""
def __init__(
self,
tile_size: int = 256,
resolution: int = 50000,
dtype: np.dtype = np.uint32,
use_zstd: bool = True,
):
"""
Examples
--------
"""
self.tile_size = tile_size
self.resolution = resolution
self.dtype = np.dtype(dtype)
self.use_zstd = use_zstd and HAS_ZSTD
[docs] def decode_tile(
self,
payload: bytes,
shape: tuple[int, int] | None = None,
) -> np.ndarray:
"""Decode a single compressed tile.
Parameters
----------
payload : bytes
Compressed bitstream. When ``shape`` is provided (v3 GZCM path),
the first 8 bytes encode ``(rows, cols)`` as ``np.int32`` and are
stripped before decompression. When ``shape`` is ``None`` the
payload is treated as raw compressed bytes (legacy / codec tests).
shape : tuple of (int, int), optional
Actual tile shape ``(rows, cols)``. Required for edge tiles where
``rows != self.tile_size`` or ``cols != self.tile_size``. If
``None``, falls back to ``(self.tile_size, self.tile_size)`` and
assumes no 8-byte header is present.
Returns
-------
np.ndarray
Decoded contact matrix tile with the requested shape.
Examples
--------
"""
if shape is not None:
# v3 GZCM path: payload = 8-byte (rows, cols) int32 header + body
rows, cols = shape
body = payload[8:]
else:
# Legacy / direct codec API: raw compressed bytes, square tile.
rows = cols = self.tile_size
body = payload
if self.use_zstd:
ctx = zstd.ZstdDecompressor()
data = ctx.decompress(body)
else:
data = zlib.decompress(body)
return np.frombuffer(data, dtype=self.dtype).reshape(rows, cols)
[docs] def decode_tiles(
self,
payloads: list[bytes],
shapes: list[tuple[int, int]] | None = None,
) -> np.ndarray:
"""Decode multiple tiles into a 4D array.
Parameters
----------
payloads : list[bytes]
List of encoded bitstreams.
shapes : list of tuple of (int, int), optional
Per-tile shapes matching ``payloads``. When provided, each payload
is assumed to carry the 8-byte header. Defaults to ``None``
(legacy: square tiles, no header).
Returns
-------
np.ndarray
4D array of decoded tiles (n_tile_rows, n_tile_cols, tile_rows, tile_cols).
Examples
--------
"""
n_tiles = len(payloads)
if shapes is None:
decoded = [self.decode_tile(p) for p in payloads]
else:
decoded = [self.decode_tile(p, shape=s) for p, s in zip(payloads, shapes, strict=True)]
tile_shape = decoded[0].shape
tile_rows = int(np.sqrt(n_tiles))
tile_cols = n_tiles // tile_rows if tile_rows > 0 else 1
result = np.empty((tile_rows, tile_cols, *tile_shape), dtype=self.dtype)
idx = 0
for i in range(tile_rows):
for j in range(tile_cols):
result[i, j] = decoded[idx]
idx += 1
return result