"""
Dataset for .gzcm unified container.
Supports GZCM v1 (dense), v2 (tiled/csr/block_sparse), v3 (compressed tiles), and
v4 (sparse-tiled-intra: anti-diagonal block order, adaptive codec picker).
Examples
--------
>>> from gunz_cm.datasets.gzcm import GzcmDataset
>>> ds = GzcmDataset("matrix.gzcm", window_size=1000000)
>>> item = ds[0]
"""
__author__ = "Yeremia Gunawan Adhisantoso"
__email__ = "adhisant@tnt.uni-hannover.de"
__license__ = "Clear BSD"
import pathlib
import numpy as np
import pandas as pd
import threading
import typing as t
from cachetools import LRUCache
# v2.15.0 (parallel session, GZCM v4): keep HEAD's post-rename
# GZCMReader (Phase 1a) AND add the new v4 symbols so the dataset
# dispatcher can use them.
from ..exceptions import DatasetError, GzcmV4FormatError
from ..io.gnz import GZCMReader, GzcmV4Reader
from ._torch_guard import require_torch
require_torch()
import torch # noqa: E402 (guarded by require_torch)
from ._torch_guard import DatasetBase as DatasetType
# Backward-compatible default tile-cache size. v2.26.0 raised this from 32
# to 256 because the v2.15.0 cache was the documented bottleneck for "very
# large windows or high concurrency" (docs/source/user_guide/datasets.md:120).
_DEFAULT_TILE_CACHE_SIZE = 256
[docs]class GZCMDataset(DatasetType):
"""
Dataset for .gzcm unified container format.
Supports GZCM v1 (dense), v2 (tiled/csr/block_sparse), and v3 (compressed tiles).
Parameters
----------
fpath : str
Path to .gzcm file.
window_size : int
Window size in bp.
output_type : str, default="sparse"
Output type: "sparse" or "dense".
downsample_ratio : float or tuple, optional
Downsampling ratio.
decompress : bool, default=True
If True, decode compressed tiles on access. If False, return raw bytes.
tile_cache_size : int, default=256
Maximum number of decoded tiles to keep in the in-memory LRU cache.
The cache is thread-safe (``threading.RLock``) so it is safe under
PyTorch ``DataLoader(num_workers > 0)``. Set to 0 to disable caching.
Examples
--------
"""
def __init__(
self,
fpath: str,
window_size: int,
output_type: str = "sparse",
downsample_ratio: float | tuple[float, float] | None = None,
decompress: bool = True,
tile_cache_size: int = _DEFAULT_TILE_CACHE_SIZE,
):
# v2.15.0 (parallel session, GZCM v4): use the dispatcher so the
# dataset transparently opens a v3 or v4 reader based on the
# file's header.version. See `_open_reader` below.
self.reader = GzcmDataset._open_reader(fpath)
self.metadata = self.reader.get_metadata()
self.version = self.reader.version
self.bin_size_bp = self.metadata["resolution"]
self.tile_cache_size = tile_cache_size
self.window_size = window_size
self.output_type = output_type
self.downsample_ratio = downsample_ratio
self.decompress = decompress
self.layout = self.metadata.get("layout", "dense")
self.block_size = self.metadata.get("block_size")
self.compression = self.metadata.get("compression")
if self.version == 3 and self.compression:
self._init_compressed()
elif self.version == 4 and self.layout == "sparse-tiled-intra" and self.compression:
# v4 sparse-tiled-intra: same tile-codec stack as v3 but
# tiles written in anti-diagonal order with per-tile bboxes
# in meta["regions"][0]["tile_bboxes"]. The codec stack is
# identical to v3, so we route through the same _init_compressed
# path. The metadata shape mismatch is normalized below in
# _init_compressed.
self._init_compressed()
elif self.layout == "dense":
self._init_dense()
elif self.layout == "tiled":
self._init_tiled()
elif self.layout == "csr":
self._init_csr()
elif self.layout == "block_sparse":
self._init_block_sparse()
else:
raise DatasetError(f"Unknown layout: {self.layout}")
self._init_weights()
step = window_size // self.bin_size_bp
starts = np.arange(0, self.n_bins, step)
ends = np.clip(starts + step, 0, self.n_bins)
self.index = pd.DataFrame({"start_bin": starts, "end_bin": ends})
def _init_compressed(self):
"""Initialize GZCM v3 compressed tile layout."""
self.tile_size = self.compression.get("tile_size", 256)
self.codec = self.compression.get("codec", "cmc")
self.n_bins = self.metadata.get("original_shape", [0, 0])[0]
# Bug 0.1 fix: v3 stores tile metadata under meta["tiles"] (a dict
# keyed by tile_name); v4 stores it under
# meta["regions"][0]["tile_bboxes"] (a list of dicts with the
# same fields plus a "diagonal" index). Without the v4 fallback,
# GzcmDataset's _build_tile_index would be empty for every v4
# file and __getitem__ would return an empty sparse dict.
self.tiles_meta = self.metadata.get("tiles", {})
if not self.tiles_meta:
v4_regions = self.metadata.get("regions", [])
if v4_regions:
v4_bboxes = v4_regions[0].get("tile_bboxes", [])
if v4_bboxes:
self.tiles_meta = {
bb["tile_name"]: {
"row_start": bb["row_start"],
"col_start": bb["col_start"],
"row_end": bb["row_end"],
"col_end": bb["col_end"],
"diagonal": bb.get("diagonal", 0),
}
for bb in v4_bboxes
if "tile_name" in bb
}
self._tile_index = None
self._decoder = None
# v2.26.0 tile cache: cachetools.LRUCache + threading.RLock. Replaces
# the pre-v2.26.0 bare OrderedDict (maxsize=32, no locking) which
# was the documented bottleneck for high-concurrency NN training.
# cachetools promotes re-accessed tiles on hit and evicts the
# least-recently-used entry on insert; the lock makes the cache
# safe under PyTorch DataLoader(num_workers > 0).
self._tile_cache_lock = threading.RLock()
self._init_cache()
@staticmethod
def _open_reader(fpath: str | pathlib.Path) -> "GZCMReader | GzcmV4Reader":
"""Open a v1/v2/v3 or v4 reader based on the file's header.version.
Phase 3 of the v4 implementation: the public Dataset API transparently
dispatches between ``GZCMReader`` (v1/v2/v3, post-Phase-1a rename)
and ``GzcmV4Reader`` (v4) so that ``GZCMDataset(...)`` accepts both.
v1/v2/v3 files continue to use the existing ``GZCMReader`` path; v4
files use ``GzcmV4Reader``. Tile-payload decoding is the
responsibility of the subclass.
"""
probe = GZCMReader(fpath)
if probe.version == 4:
return GzcmV4Reader(fpath)
return probe
def _build_tile_index(self):
"""Build O(1) lookup index for tiles."""
if self._tile_index is None:
self._tile_index = {}
for name, meta in self.tiles_meta.items():
key = (meta["row_start"], meta["col_start"])
self._tile_index[key] = name
return self._tile_index
def _init_dense(self):
"""Initialize dense layout."""
self.matrix = self.reader.get_array("matrix", mode="r")
self.n_bins = self.matrix.shape[0]
self.dtype = self.matrix.dtype
def _init_tiled(self):
"""Initialize tiled layout."""
self.matrix = self.reader.get_array("matrix", mode="r")
self.n_bins = self.matrix.shape[0] * self.block_size
self.dtype = self.matrix.dtype
def _init_csr(self):
"""Initialize CSR layout."""
self.indptr = self.reader.get_array("indptr", mode="r")
self.indices = self.reader.get_array("indices", mode="r")
self.data = self.reader.get_array("data", mode="r")
self.n_bins = len(self.indptr) - 1
self.dtype = self.data.dtype
def _init_block_sparse(self):
"""Initialize block sparse layout."""
self.block_index = self.reader.get_array("block_index", mode="r")
self.block_data = self.reader.get_array("block_data", mode="r")
self.n_bins = self.block_index.shape[0] * self.block_size
self.dtype = self.block_data.dtype
def _init_weights(self):
"""Initialize weights if present."""
self.weights = None
for key in self.reader.keys():
if key.startswith("weights_"):
self.weights = self.reader.get_array(key, mode="r")
break
def _get_decoder(self):
"""Lazy-load decoder based on codec, via the v5.1 codec registry.
Replaces the prior hardcoded if/elif chain over cmc/zstd/cmc_zstd/
bsc/bsc_cmc/lz4. The registry (see
:mod:`gunz_cm.compressions`) is the single source of truth for the
codec -> decoder class mapping; adding a new codec is one
``register_codec(...)`` call.
"""
if self._decoder is None:
from ..compressions import get_codec, UnknownCodecError as _UC
try:
_enc_cls, dec_cls, _wire_format = get_codec(self.codec)
except _UC as exc:
raise DatasetError(
f"Unknown codec: {self.codec!r}. "
f"Available codecs: {exc.available}. "
"Did the file use a newer codec than this version of gunz-cm supports?"
) from exc
self._decoder = dec_cls(tile_size=self.tile_size)
return self._decoder
def _init_cache(self) -> None:
"""Allocate the per-dataset tile cache (v2.26.0+).
Uses ``cachetools.LRUCache`` which provides O(1) get/put with
automatic LRU eviction. The cache is thread-safe under
``self._tile_cache_lock`` (a ``threading.RLock``) so it can be
shared across ``DataLoader(num_workers > 0)`` worker threads.
Setting ``tile_cache_size=0`` disables caching entirely.
"""
if self.tile_cache_size > 0:
self._tile_cache: LRUCache[str, np.ndarray] = LRUCache(maxsize=self.tile_cache_size)
else:
self._tile_cache = None
self._cache_maxsize = self.tile_cache_size
def _decode_tile(self, tile_name: str) -> np.ndarray:
"""Decode a compressed tile with LRU cache eviction (v2.26.0+).
Thread-safe under ``DataLoader(num_workers > 0)`` via the
``_tile_cache_lock``. cachetools' LRUCache promotes a hit to the
most-recently-used position and evicts the LRU entry on insert
when at maxsize.
"""
with self._tile_cache_lock:
cache = self._tile_cache
if cache is not None:
try:
return cache[tile_name]
except KeyError:
pass
payload, shape = self.reader.get_compressed_tile(tile_name, return_shape=True)
decoder = self._get_decoder()
decoded = decoder.decode_tile(payload, shape=shape)
if self.decompress:
with self._tile_cache_lock:
if self._tile_cache is not None:
self._tile_cache[tile_name] = decoded
return decoded
def _get_compressed_patch(self, s: int, e: int) -> np.ndarray:
"""Get a patch from compressed tiles with O(1) tile lookup."""
h = e - s
patch = np.zeros((h, h), dtype=np.uint32)
tile_size = self.tile_size
n_bins = self.n_bins
t_start = s // tile_size
t_end = (e - 1) // tile_size + 1
tile_index = self._build_tile_index()
# Vectorize the per-tile bbox intersection math. The dictionary lookup
# and the decoder call (which performs I/O) remain per-tile, but the
# Python-level min/max arithmetic is hoisted into numpy.
ti_arr = np.arange(t_start, t_end)
tj_arr = np.arange(t_start, t_end)
ti_grid, tj_grid = np.meshgrid(ti_arr, tj_arr, indexing="ij")
tile_origin_r = ti_grid * tile_size
tile_origin_c = tj_grid * tile_size
# Destination (patch) row/col start/end: clipped to [s, e).
dr_s_grid = np.maximum(s, tile_origin_r)
dr_e_grid = np.minimum(e, tile_origin_r + tile_size)
dc_s_grid = np.maximum(s, tile_origin_c)
dc_e_grid = np.minimum(e, tile_origin_c + tile_size)
# Source (decoded) row/col start/end within each tile.
pr_s_grid = dr_s_grid - tile_origin_r
pc_s_grid = dc_s_grid - tile_origin_c
for ii in range(ti_arr.size):
for jj in range(tj_arr.size):
ti = int(ti_arr[ii])
tj = int(tj_arr[jj])
key = (ti * tile_size, tj * tile_size)
tile_name = tile_index.get(key)
if tile_name is None:
continue
decoded = self._decode_tile(tile_name)
dr_s = int(dr_s_grid[ii, jj])
dr_e = int(dr_e_grid[ii, jj])
dc_s = int(dc_s_grid[ii, jj])
dc_e = int(dc_e_grid[ii, jj])
pr_s = int(pr_s_grid[ii, jj])
pc_s = int(pc_s_grid[ii, jj])
pr_e = pr_s + (dr_e - dr_s)
pc_e = pc_s + (dc_e - dc_s)
if dr_s < n_bins and dc_s < n_bins:
pr_e_c = min(pr_e, decoded.shape[0])
pc_e_c = min(pc_e, decoded.shape[1])
patch_rows = pr_e_c - pr_s
patch_cols = pc_e_c - pc_s
if patch_rows > 0 and patch_cols > 0 and pr_s < decoded.shape[0] and pc_s < decoded.shape[1]:
patch[dr_s - s : dr_s - s + patch_rows, dc_s - s : dc_s - s + patch_cols] = decoded[pr_s:pr_e_c, pc_s:pc_e_c]
return patch
def __len__(self):
return len(self.index)
def _get_dense_patch_tiled(self, s, e):
B = self.block_size
b_start = s // B
b_end = (e - 1) // B + 1
patch_h = e - s
patch = np.zeros((patch_h, patch_h), dtype=self.dtype)
for br in range(b_start, b_end):
for bc in range(b_start, b_end):
gs_r, ge_r = br * B, (br + 1) * B
gs_c, ge_c = bc * B, (bc + 1) * B
is_r, ie_r = max(s, gs_r), min(e, ge_r)
is_c, ie_c = max(s, gs_c), min(e, ge_c)
if is_r < ie_r and is_c < ie_c:
if self.layout == "tiled":
if br < self.matrix.shape[0] and bc < self.matrix.shape[1]:
block = self.matrix[br, bc]
else:
continue
else:
if br < self.block_index.shape[0] and bc < self.block_index.shape[1]:
idx = self.block_index[br, bc]
if idx == -1:
continue
block = self.block_data[idx]
else:
continue
patch[is_r - s : ie_r - s, is_c - s : ie_c - s] = block[is_r - gs_r : ie_r - gs_r, is_c - gs_c : ie_c - gs_c]
return patch
def _get_csr_coo(self, s, e):
r_list, c_list, v_list = [], [], []
for r in range(s, e):
p0 = self.indptr[r]
p1 = self.indptr[r + 1]
if p1 > p0:
cols = self.indices[p0:p1]
vals = self.data[p0:p1]
mask = (cols >= s) & (cols < e)
if np.any(mask):
r_list.append(np.full(mask.sum(), r - s, dtype=np.int64))
c_list.append(cols[mask] - s)
v_list.append(vals[mask])
if not r_list:
return np.array([], dtype=np.int64), np.array([], dtype=np.int64), np.array([], dtype=self.dtype)
return np.concatenate(r_list), np.concatenate(c_list), np.concatenate(v_list)
def __getitem__(self, idx):
row = self.index.iloc[idx]
s, e = int(row["start_bin"]), int(row["end_bin"])
if self.version == 3 and self.compression:
return self._get_item_compressed(s, e)
if self.version == 4 and self.layout == "sparse-tiled-intra" and self.compression:
# v4 uses the same compressed-tile patch path as v3; only the
# metadata shape (regions[0].tile_bboxes vs tiles) and the
# write-side layout differ.
return self._get_item_compressed(s, e)
return self._get_item_uncompressed(s, e)
def _get_item_uncompressed(self, s: int, e: int):
"""Handle uncompressed (v1/v2) item retrieval."""
if self.output_type == "dense":
return self._get_item_dense(s, e)
return self._get_item_sparse(s, e)
def _get_item_compressed(self, s: int, e: int):
"""Handle compressed (v3 and v4 sparse-tiled-intra) item retrieval."""
patch = self._get_compressed_patch(s, e)
if self.output_type == "dense":
return self._apply_weights_dense(patch, s, e)
r, c = np.nonzero(patch)
mask = r <= c
r, c = r[mask], c[mask]
counts = patch[r, c]
return self._build_sparse_output(r, c, counts, s)
def _get_item_dense(self, s: int, e: int):
"""Dense output path for uncompressed layouts."""
if self.layout == "dense":
patch = self.matrix[s:e, s:e]
elif self.layout in ["tiled", "block_sparse"]:
patch = self._get_dense_patch_tiled(s, e)
elif self.layout == "csr":
r, c, v = self._get_csr_coo(s, e)
h = e - s
patch = np.zeros((h, h), dtype=self.dtype)
if len(r) > 0:
patch[r, c] = v
else:
raise DatasetError(f"Unknown layout: {self.layout}")
return self._apply_weights_dense(patch, s, e)
def _get_item_sparse(self, s: int, e: int):
"""Sparse output path for uncompressed layouts."""
if self.layout == "dense":
patch = self.matrix[s:e, s:e]
r, c = np.nonzero(patch)
mask = r <= c
r, c = r[mask], c[mask]
counts = patch[r, c]
elif self.layout in ["tiled", "block_sparse"]:
patch = self._get_dense_patch_tiled(s, e)
r, c = np.nonzero(patch)
mask = r <= c
r, c = r[mask], c[mask]
counts = patch[r, c]
elif self.layout == "csr":
r, c, counts = self._get_csr_coo(s, e)
mask = r <= c
r, c, counts = r[mask], c[mask], counts[mask]
else:
raise DatasetError(f"Unknown layout: {self.layout}")
return self._build_sparse_output(r, c, counts, s)
def _apply_weights_dense(self, patch, s, e):
"""Apply weights to dense patch."""
if self.weights is not None:
w_slice = self.weights[s:e]
denom = np.outer(w_slice, w_slice)
patch = patch.astype(np.float32) / denom
np.nan_to_num(patch, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
return torch.from_numpy(patch).float().unsqueeze(0)
def _build_sparse_output(self, r, c, counts, s):
"""Build sparse dict output from coordinates and counts."""
if self.weights is not None:
w1 = self.weights[s + r]
w2 = self.weights[s + c]
counts = counts.astype(np.float32) / (w1 * w2)
mask_v = np.isfinite(counts)
if not np.all(mask_v):
r, c, counts = r[mask_v], c[mask_v], counts[mask_v]
target_counts = counts.copy()
if self.downsample_ratio is not None:
if isinstance(self.downsample_ratio, tuple):
alpha = np.random.uniform(*self.downsample_ratio)
else:
alpha = self.downsample_ratio
counts = np.random.binomial(counts.astype(np.int32), alpha)
mask = counts > 0
r, c, counts = r[mask], c[mask], counts[mask]
coords = np.stack([r, c], axis=1)
return {
"coords": torch.from_numpy(coords).long(),
"features": torch.from_numpy(counts).float().unsqueeze(1),
"target": torch.from_numpy(target_counts).float(),
"info": {"start": s * self.bin_size_bp},
}
GnzSparseDataset = GZCMDataset # backward-compat alias for 1 release
# 1-release deprecation alias; remove in 2.30.0
GzcmDataset = GZCMDataset