Source code for gunz_cm.compressions.cmc_zstd_encoder

"""
CMC Transforms + Zstd Encoder for GZCM v3 compression.

Combines CMC's domain-specific transforms (diagonal transform, binarization)
with Zstd entropy coding for faster decode than pure CMC.

Examples
--------
"""

__author__ = "Yeremia Gunawan Adhisantoso"
__email__ = "adhisant@tnt.uni-hannover.de"
__license__ = "Clear BSD"

import typing as t
import zlib
import numpy as np

try:
    import zstandard as zstd

    HAS_ZSTD = True
except ImportError:
    HAS_ZSTD = False

from . import _paths

# Module-level registry of CMC+Zstd encoder functions. See ``cmc_encoder.py``
# for the rationale (DOP Phase 3, v2.20.0).
_FUNCS: dict[str, t.Any] = {}
_LOADED: bool = False


def _ensure_loaded() -> None:
    """Lazily resolve CMC+Zstd encoder functions on first access.

    See ``cmc_encoder.py`` for the full rationale.
    """
    global _LOADED
    if _LOADED:
        return
    _cmc = _paths.load_cmc_zstd_encoder()
    _FUNCS["binarize_rc_bin_split_v2"] = _cmc["binarize_rc_bin_split_v2"]
    _FUNCS["diag_transform"] = _cmc["diag_transform"]
    _LOADED = True


def __getattr__(name: str) -> t.Any:
    """PEP 562 module-level lazy attribute access for CMC+Zstd encoder functions."""
    if name in {"binarize_rc_bin_split_v2", "diag_transform"}:
        _ensure_loaded()
        return _FUNCS[name]
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")


[docs]class CmcZstdEncoder: """CMC Transforms + Zstd encoder for contact matrix tiles. Uses CMC's domain-specific transforms (diagonal transform, binarization) with Zstd entropy coding for better compression and faster decode. Parameters ---------- tile_size : int, default=256 Tile size for block processing. resolution : int, default=50000 Hi-C resolution in bp. level : int, default=3 Compression level (1-22 for zstd, 1-9 for zlib fallback). Examples -------- """ def __init__( self, tile_size: int = 256, resolution: int = 50000, level: int = 3, ): """ Examples -------- """ self.tile_size = tile_size self.resolution = resolution self.level = level
[docs] def encode_tile(self, mat: np.ndarray) -> bytes: """Encode a single contact matrix tile. Parameters ---------- mat : np.ndarray 2D contact matrix tile (upper triangular). Returns ------- bytes Compressed bitstream (shape info + encoded data). Examples -------- """ _ensure_loaded() mat = _FUNCS["diag_transform"](mat, mode=0) bin_mat = _FUNCS["binarize_rc_bin_split_v2"](mat, axis=0) data = bin_mat.tobytes() shape = np.array(bin_mat.shape, dtype=np.int32).tobytes() if HAS_ZSTD: ctx = zstd.ZstdCompressor(level=self.level) compressed = ctx.compress(data) else: compressed = zlib.compress(data, level=min(self.level, 9)) return shape + compressed
[docs] def encode_tiles(self, tiles: np.ndarray) -> list[bytes]: """Encode multiple tiles. Parameters ---------- tiles : np.ndarray 4D array of shape (n_tile_rows, n_tile_cols, tile_size, tile_size). Returns ------- list[bytes] List of encoded bitstreams, one per tile. Examples -------- """ n_tile_rows, n_tile_cols = tiles.shape[0], tiles.shape[1] results = [] for i in range(n_tile_rows): for j in range(n_tile_cols): results.append(self.encode_tile(tiles[i, j])) return results
[docs] def get_compression_info(self) -> dict: """Return compression metadata. Returns ------- dict Compression parameters for header. Examples -------- """ return { "codec": "cmc_zstd", "version": "1.0", "tile_size": self.tile_size, "resolution": self.resolution, "level": self.level, }