Source code for gunz_cm.compressions.cmc_decoder

"""
CMC Decoder wrapper for GNZ v3 compression.

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

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

import typing as t

import numpy as np

from . import _paths

# Module-level registry of CMC decoder functions. Populated lazily on
# first access by ``_ensure_loaded()``. 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 decoder functions on first access.

    See ``cmc_encoder.py`` for the full rationale.
    """
    global _LOADED
    if _LOADED:
        return
    _cmc = _paths.load_cmc_decoder()
    _FUNCS["jbig_decode"] = _cmc["jbig_decode"]
    _FUNCS["debinarize_rc_bin_split_v2"] = _cmc["debinarize_rc_bin_split_v2"]
    _FUNCS["reverse_diag_transform_mode0"] = _cmc["reverse_diag_transform_mode0"]
    _LOADED = True


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


class CmcDecoder:
    """CMC 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.
    diag_transform : bool, default=True
        Reverse diagonal transform after decoding.

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

    def __init__(
        self,
        tile_size: int = 256,
        resolution: int = 50000,
        diag_transform: bool = True,
    ):
        """
        Examples
        --------
        """
        self.tile_size = tile_size
        self.resolution = resolution
        self.diag_transform = diag_transform

    def decode_tile(self, payload: bytes) -> np.ndarray:
        """Decode a single CMC-encoded tile.

        Parameters
        ----------
        payload : bytes
            CMC-encoded bitstream.

        Returns
        -------
        np.ndarray
            Decoded contact matrix tile.

        Examples
        --------
        """
        _ensure_loaded()
        bin_mat = _FUNCS["jbig_decode"](payload)
        debinarized = _FUNCS["debinarize_rc_bin_split_v2"](bin_mat, axis=0)
        if self.diag_transform:
            return _FUNCS["reverse_diag_transform_mode0"](debinarized)
        return debinarized

    def decode_tiles(self, payloads: list[bytes]) -> np.ndarray:
        """Decode multiple tiles into a 4D array.

        Parameters
        ----------
        payloads : list[bytes]
            List of encoded bitstreams.

        Returns
        -------
        np.ndarray
            4D array of decoded tiles (n_tile_rows, n_tile_cols, tile_size, tile_size).

        Examples
        --------
        """
        n_tiles = len(payloads)
        decoded = [self.decode_tile(p) for p in payloads]

        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=np.uint32)
        idx = 0
        for i in range(tile_rows):
            for j in range(tile_cols):
                result[i, j] = decoded[idx]
                idx += 1
        return result