Source code for gunz_cm.converters.gzcm

"""
Converter to .gzcm unified container.

Supports GZCM v1 (dense), v2 (dense/tiled arrays), v3 (compressed tiles via CMC),
and v4 (anti-diagonal layout + adaptive codec picker; see ``specs/gzcm-v4-design.md``).

Examples
--------
>>> from gunz_cm.converters.gzcm import convert_to_gzcm
>>> convert_to_gzcm(
...     fpath="matrix.hic",
...     output_fpath="matrix.gzcm",
...     region1="chr1",
...     bin_size_bp=10000,
...     version=4,
...     tile_size=256,
... )
"""

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

import pathlib
import numpy as np
from gunz_utils import type_checked

from ..io.gnz import GZCMWriter
from .. import loaders
from ..exceptions import ConverterError, UnsupportedLoaderFeatureError
from ..consts import Balancing, DataStructure, Backend, DataFrameSpecs


[docs]@type_checked(config=dict(arbitrary_types_allowed=True)) def convert_to_gzcm( fpath: pathlib.Path, output_fpath: pathlib.Path, region1: str | None, bin_size_bp: int, balancing: Balancing | None = None, backend: Backend = Backend.HICTK, dtype: str = "float32", overwrite: bool = False, version: int = 1, block_size: int = 1024, tile_size: int = 512, compression: str | None = None, # GZCM v4 (Phase 1a) — per-region layout + adaptive codec picker. layout: str | None = None, region_layouts: dict[str, str] | None = None, adaptive_codec: bool = False, codec_candidates: tuple[str, ...] | None = None, # Streaming memory bounding (parallel session, v2.15.0 plumbing). chunk_size: int = 10_000_000, ) -> None: """ Converts a Hi-C file to a .gzcm container with matrix and weights. Supports GZCM v1 (dense), v2 (tiled, csr, block_sparse), and v3 (compressed tiles). For v3, compression codecs are: ============= ======== ============ ============= Codec Size Convert Speed Recommended ============= ======== ============ ============= cmc 1.57 MB 4.82s Best compression cmc_zstd 2.65 MB 4.12s Best balance (*) bsc 2.78 MB 9.68s Balanced zstd 5.75 MB 5.16s Fastest decode ============= ======== ============ ============= (*) cmc_zstd is the recommended default: best balance of compression ratio and convert speed. Real HiC chr1 @ 50kb bin_size_bp benchmarked. Parameters ---------- fpath : pathlib.Path Input Hi-C file path. output_fpath : pathlib.Path Output .gzcm file path. region1 : str Genomic region (e.g., "chr1"). bin_size_bp : int Hi-C bin size in bp. balancing : Balancing, optional Balancing method. backend : Backend, default=HICTK Backend to use for loading. dtype : str, default="float32" Data type for matrix storage. overwrite : bool, default=False Overwrite existing file. version : int, default=1 GZCM version: 1 (dense), 2 (tiled/sparse), 3 (compressed tiles). block_size : int, default=1024 Block size for v2 tiled layouts. tile_size : int, default=512 Tile size for v3 compression. compression : str, optional Compression codec for v3: "cmc", "cmc_zstd" (recommended default), "zstd", or "bsc". Examples -------- """ if region1 is None: raise UnsupportedLoaderFeatureError( feature="region1=None (full-genome writing)", loader_name="GZCM converter", ) if not output_fpath.name.endswith(".gzcm"): output_fpath = output_fpath.with_suffix(".gzcm") if version not in (1, 2, 3, 4): raise ConverterError(f"GZCM version must be 1, 2, 3, or 4, got {version}") if version == 4: if region_layouts is None: region_layouts = {region1.split(":")[0]: layout or "sparse-tiled-intra"} writer = GZCMWriter(output_fpath, overwrite=overwrite, version=version) df_raw = loaders.load_cm_data( fpath=fpath, region1=region1, bin_size_bp=bin_size_bp, balancing=Balancing.NONE, output_format=DataStructure.DF, backend=backend ) row_ids = df_raw[DataFrameSpecs.ROW_IDS].to_numpy() col_ids = df_raw[DataFrameSpecs.COL_IDS].to_numpy() counts = df_raw[DataFrameSpecs.COUNTS].to_numpy() if np.issubdtype(counts.dtype, np.floating): counts[~np.isfinite(counts)] = 0 n = int(max(np.max(row_ids), np.max(col_ids)) + 1) if len(row_ids) > 0 else 0 meta = { "resolution": bin_size_bp, "region": region1, "chromosome1": region1.split(":")[0], "source_file": str(fpath), "balancing": balancing.value if balancing else "NONE", "version": version, } if balancing and balancing != Balancing.NONE: if backend == Backend.HICTK: import hictkpy f_hic = hictkpy.File(str(fpath), bin_size_bp) bins = f_hic.bins() offset = bins.get_id(region1.split(":")[0], 0) w_global = f_hic.weights(balancing.value, divisive=True) w_local = w_global[offset : offset + n] writer.add_array(f"weights_{balancing.value}", w_local, dtype="float32") if version == 1: _write_gzcm_v1(writer, row_ids, col_ids, counts, n, meta, dtype) elif version == 2: _write_gzcm_v2(writer, row_ids, col_ids, counts, n, meta, dtype, block_size) elif version == 3: # Parallel session: pass chunk_size for streaming memory bounding. _write_gzcm_v3(writer, row_ids, col_ids, counts, n, meta, tile_size, compression, chunk_size) elif version == 4: # GZCM v4 (Phase 1a, full implementation): sparse-tiled-intra layout # with optional adaptive codec picker. _write_gzcm_v4_intra( writer, row_ids, col_ids, counts, n, meta, tile_size, compression, region_layouts, adaptive_codec, codec_candidates, )
def _picker_name_to_writer_codec(picker_name: str) -> str: """Translate a scheme-picker codec name to a writer-codec name. The scheme picker (see :mod:`gunz_cm.compressions.scheme_picker`) uses suffixed names to distinguish codec levels (``zstd-3``, ``lz4-hc-9``). The writer's encoder lookup uses the bare codec family name (``zstd``, ``lz4``). This helper bridges the two naming conventions. Unknown picker names are returned unchanged so the writer's ``ConverterError`` for unknown codecs surfaces a clear message. """ mapping = { "cmc": "cmc", "zstd-3": "zstd", "zstd-1": "zstd", "zstd-9": "zstd", "lz4-hc-9": "lz4", "lz4-1": "lz4", } return mapping.get(picker_name, picker_name) def _block_index_intra(n: int, block_bin_count: int) -> list[dict[str, int]]: """Enumerate tile (i, j) bboxes in anti-diagonal order. For a matrix of side ``n`` with tiles of side ``T = block_bin_count``, iterate over all ``(i, j)`` with ``i + j = k`` for increasing ``k``. This matches the .hic v9 contact-decay prior (counts concentrated near the diagonal) and makes the on-disk tile order locality-friendly for the in-memory LRU cache. See ``specs/gzcm-v4-design.md`` §4.2. """ if block_bin_count <= 0: raise ValueError(f"block_bin_count must be positive, got {block_bin_count}") if n <= 0: return [] n_tiles = (n + block_bin_count - 1) // block_bin_count blocks: list[dict[str, int]] = [] for d in range(2 * n_tiles - 1): for i in range(max(0, d - (n_tiles - 1)), min(d + 1, n_tiles)): j = d - i if 0 <= j < n_tiles and 0 <= i < n_tiles: blocks.append( { "row_start": i * block_bin_count, "col_start": j * block_bin_count, "row_end": min((i + 1) * block_bin_count, n), "col_end": min((j + 1) * block_bin_count, n), "diagonal": d, } ) return blocks def _write_gzcm_v4_intra( # noqa: C901 writer: "GZCMWriter", row_ids: np.ndarray, col_ids: np.ndarray, counts: np.ndarray, n: int, meta: dict, tile_size: int, compression: str | None = None, region_layouts: dict[str, str] | None = None, adaptive_codec: bool = False, codec_candidates: tuple[str, ...] | None = None, ) -> None: """Write GZCM v4 - sparse-tiled-intra layout. Encodes an intrachromosomal contact matrix using the .hic v9 anti-diagonal block layout: tiles are enumerated in order of distance from the main diagonal. The on-disk tile body still uses the v3 codec stack (``zstd-3`` by default) so the decoder stays codec-compatible. Wire format per tile (identical to v3 except for the 8-byte shape header which the v5.1 wire-format contract marks as OPAQUE_PAYLOAD for all codecs except LZ4): For OPAQUE_PAYLOAD codecs (zstd, cmc, bsc, bsc_cmc, cmc_zstd): [rows:i32][cols:i32][encoded_payload:bytes] For SELF_DESCRIBING codecs (lz4): [encoded_payload:bytes] (LZ4 encoder adds its own 12-byte header) Per-tile metadata lives in ``meta["regions"][0]["tile_bboxes"]`` (a list of dicts with ``tile_name``, ``row_start``, ``col_start``, ``row_end``, ``col_end``, ``diagonal``) — distinct from v3's ``meta["tiles"]`` (a dict). The reader normalizes both shapes into the same internal representation via a fallback read (see ``src/gunz_cm/datasets/gzcm.py:_init_compressed``). Parameters ---------- writer : GZCMWriter Open v4 GZCM writer. row_ids, col_ids, counts : np.ndarray Sparse upper-triangular COO coordinates. n : int Matrix size in bins. meta : dict User metadata (mutated in-place to attach v4 fields). tile_size : int Tile side length in bins. compression : str, optional Codec name. Defaults to ``"zstd"``. Ignored when ``adaptive_codec`` is True; the scheme picker then chooses the codec per region. region_layouts : dict, optional Mapping of chromosome name to layout name. The single intrachromosomal region in this writer is recorded under ``regions[0]["layout"]``. adaptive_codec : bool, default=False If True, call :func:`pick_codec_for_region` (see :mod:`gunz_cm.compressions.scheme_picker`) to choose the codec for this region based on a 5% sample. Falls back to ``compression`` or ``"zstd"`` if the picker errors or all candidates fail to encode. codec_candidates : tuple of str, optional Candidate codec names for the picker. Defaults to ``("cmc", "zstd-3", "lz4-hc-9")``. """ from ..compressions import get_codec, UnknownCodecError, WireFormat codec = compression or "zstd" chosen_codec: str = codec if adaptive_codec: try: from ..compressions.scheme_picker import pick_codec_for_region picked, _bit_width = pick_codec_for_region( row_ids, counts, n=n, tile_size=tile_size, candidates=codec_candidates or ("cmc", "zstd-3", "lz4-hc-9"), ) chosen_codec = _picker_name_to_writer_codec(picked) meta["codec_picker"] = { "adaptive": True, "chosen": picked, "writer_codec": chosen_codec, "candidates": list(codec_candidates or ("cmc", "zstd-3", "lz4-hc-9")), } except Exception as exc: # picker may fail if no candidate codec is runnable meta["codec_picker"] = { "adaptive": True, "chosen": codec, "error": repr(exc), } # Look up the chosen codec in the v5.1 registry. The registry also # tells us whether to prepend the 8-byte shape header. try: encoder_cls, _decoder_cls, wire_format = get_codec(chosen_codec) except UnknownCodecError as exc: raise ConverterError( f"Unknown compression codec: {chosen_codec!r}. " f"Available codecs: {exc.available}. " "Use 'cmc', 'cmc_zstd', 'bsc', 'bsc_cmc', 'zstd', or 'lz4'." ) from exc encoder = encoder_cls(tile_size=tile_size) blocks = _block_index_intra(n, tile_size) tile_count = 0 tile_bboxes: list[dict[str, int]] = [] mat = np.zeros((n, n), dtype=np.uint32) if len(row_ids) > 0: mat[row_ids, col_ids] = counts.astype(np.uint32) mat[col_ids, row_ids] = counts.astype(np.uint32) for block in blocks: row_start, col_start = block["row_start"], block["col_start"] row_end, col_end = block["row_end"], block["col_end"] tile = mat[row_start:row_end, col_start:col_end] if tile.size == 0: continue encoded = encoder.encode_tile(tile) tile_name = f"tile_{tile_count}" rows = int(tile.shape[0]) cols = int(tile.shape[1]) # v5.1 wire-format contract: only OPAQUE_PAYLOAD codecs need the # 8-byte (rows, cols) shape header prepended. LZ4 (SELF_DESCRIBING) # adds its own 12-byte header and would be corrupted by the prepend. if wire_format == WireFormat.OPAQUE_PAYLOAD: header = np.array([rows, cols], dtype=np.int32).tobytes() payload_with_header = header + encoded else: payload_with_header = encoded writer.add_compressed_tile(tile_name, payload_with_header, uncompressed_size=tile.nbytes) tile_bboxes.append( { "tile_name": tile_name, "row_start": row_start, "col_start": col_start, "row_end": row_end, "col_end": col_end, "diagonal": block["diagonal"], } ) tile_count += 1 chrom = meta.get("chromosome1", "") layout_name = "sparse-tiled-intra" if region_layouts and chrom in region_layouts: layout_name = region_layouts[chrom] meta["layout"] = layout_name meta["version_gzcm"] = 4 meta["n_tiles"] = tile_count meta["original_shape"] = (n, n) meta["tile_size"] = tile_size meta["compression"] = {"codec": chosen_codec, "tile_size": tile_size} meta["regions"] = [ { "chromosome": chrom, "layout": layout_name, "tile_bboxes": tile_bboxes, } ] writer.set_metadata(meta) writer.write() def _write_gzcm_v1( writer: GZCMWriter, row_ids: np.ndarray, col_ids: np.ndarray, counts: np.ndarray, n: int, meta: dict, dtype: str, ) -> None: """ Write GZCM v1 - dense format. Examples -------- """ writer.init_streaming_array("matrix", (n, n), dtype) writer.set_metadata(meta) writer.write() mm = writer.get_array_writable("matrix") if len(row_ids) > 0: mm[row_ids, col_ids] = counts.astype(dtype) nondiag = row_ids != col_ids mm[col_ids[nondiag], row_ids[nondiag]] = counts[nondiag].astype(dtype) mm.flush() def _write_gzcm_v2( writer: GZCMWriter, row_ids: np.ndarray, col_ids: np.ndarray, counts: np.ndarray, n: int, meta: dict, dtype: str, block_size: int, ) -> None: """ Write GZCM v2 - tiled format. Examples -------- """ pad_h = (block_size - (n % block_size)) % block_size padded_n = n + pad_h n_blocks = padded_n // block_size padded_shape = (n_blocks, n_blocks, block_size, block_size) writer.init_streaming_array("matrix", padded_shape, dtype) meta["layout"] = "tiled" meta["padded_shape"] = (padded_n, padded_n) meta["block_size"] = block_size writer.set_metadata(meta) writer.write() mm = writer.get_array_writable("matrix") if len(row_ids) > 0: br, pr = np.divmod(row_ids, block_size) bc, pc = np.divmod(col_ids, block_size) mm[br, bc, pr, pc] = counts.astype(dtype) nondiag = row_ids != col_ids if np.any(nondiag): mm[bc[nondiag], br[nondiag], pc[nondiag], pr[nondiag]] = counts[nondiag].astype(dtype) mm.flush() def _write_gzcm_v3( # noqa: C901 writer: GZCMWriter, row_ids: np.ndarray, col_ids: np.ndarray, counts: np.ndarray, n: int, meta: dict, tile_size: int, compression: str | None = None, chunk_size: int = 10_000_000, ) -> None: """ Write GZCM v3 - compressed tile format. Phase 4 (v2.15.0): ``chunk_size`` parameter is now plumbed through the v3 path. The current implementation does not reduce peak memory (the full ``mat`` array is materialized before tile encoding begins). The parameter is wired up so that future phases can switch to a memory-bounded streaming loader without changing the public API. Parameters ---------- chunk_size : int, default=10_000_000 Approximate pixel count per write batch. The current v3 path accepts the parameter for API compatibility with ``convert_to_gzcm_v4`` but does not yet use it for memory bounding. Examples -------- """ codec = compression or "cmc" # noqa: C901 if codec == "cmc": from ..compressions import CmcEncoder encoder = CmcEncoder(tile_size=tile_size) elif codec == "zstd": from ..compressions import ZstdEncoder encoder = ZstdEncoder(tile_size=tile_size) elif codec == "cmc_zstd": from ..compressions import CmcZstdEncoder encoder = CmcZstdEncoder(tile_size=tile_size) elif codec == "bsc": from ..compressions import BscEncoder encoder = BscEncoder(tile_size=tile_size) elif codec == "bsc_cmc": from ..compressions import BscCmcEncoder encoder = BscCmcEncoder(tile_size=tile_size) else: raise ConverterError(f"Unknown compression codec: {codec}. Use 'cmc', 'zstd', 'cmc_zstd', 'bsc', or 'bsc_cmc'.") pad_h = (tile_size - (n % tile_size)) % tile_size padded_n = n + pad_h n_tiles_per_dim = padded_n // tile_size tile_count = 0 encoder_info = {} mat = np.zeros((n, n), dtype=np.uint32) if len(row_ids) > 0: mat[row_ids, col_ids] = counts.astype(np.uint32) mat[col_ids, row_ids] = counts.astype(np.uint32) for i in range(n_tiles_per_dim): for j in range(n_tiles_per_dim): row_start, col_start = i * tile_size, j * tile_size row_end = min(row_start + tile_size, padded_n) col_end = min(col_start + tile_size, padded_n) tile = mat[row_start:row_end, col_start:col_end] if tile.size == 0: continue if codec in ("cmc", "cmc_zstd", "bsc_cmc"): tile_data = np.triu(tile, k=0) else: tile_data = tile encoded = encoder.encode_tile(tile_data) tile_name = f"tile_{tile_count}" # v3 GZCM wire format: 8-byte (rows, cols) int32 header prepended # to the compressed body so the decoder can reshape edge tiles. rows = int(tile_data.shape[0]) cols = int(tile_data.shape[1]) header = np.array([rows, cols], dtype=np.int32).tobytes() payload_with_header = header + encoded writer.add_compressed_tile(tile_name, payload_with_header, uncompressed_size=tile_data.nbytes) encoder_info[tile_name] = { "row_start": row_start, "col_start": col_start, "row_end": row_end, "col_end": col_end, } tile_count += 1 meta["tiles"] = encoder_info meta["n_tiles"] = tile_count meta["padded_shape"] = (padded_n, padded_n) meta["original_shape"] = (n, n) meta["tile_size"] = tile_size meta["compression"] = {"codec": codec, "tile_size": tile_size} meta["version_gzcm"] = 3 writer.set_metadata(meta) writer.write() def _write_gzcm_v4( writer: GZCMWriter, row_ids: np.ndarray, col_ids: np.ndarray, counts: np.ndarray, n: int, meta: dict, tile_size: int = 512, chunk_size: int = 10_000_000, ) -> None: """Write GZCM v4 - region-based format (per specs/gzcm-v4-design.md). Phase 1 (v2.15.0) skeleton: emits the v4 wire-format header with the region descriptor, but tile payloads are still encoded with the v3 zstd path. The full v4 codec pipeline (Roaring bitmaps, simdcomp bit-pack, anti-diagonal layout, scheme picker, LZ4 codec) is implemented in subsequent phases. Wire format invariant validated against ``specs/v4_schema_validator.py``. """ raise NotImplementedError( "GZCM v4 codec pipeline is not yet implemented. " "Phase 1 (v2.15.0) ships only the wire-format header. " "See specs/gzcm-v4-design.md §6 for the full conversion path." )