"""GzcmTileDataset — TileDataset over a .gzcm container.
Reads GZCM v1/v2/v3 files via :class:`gunz_cm.io.gnz.GZCMReader` and
v4 files via :class:`gunz_cm.io.gnz.GzcmV4Reader`. The dispatch is
performed by the static ``GZCMDataset._open_reader`` method on the
file-backed Dataset side (see ``src/gunz_cm/datasets/gzcm.py``), which
inspects the header version and returns the appropriate reader. This
class reuses that dispatcher so the two file-backed Dataset classes
share the same v3/v4 routing logic.
Tile decoding: for v3, we use the existing GZCMDataset path. For v4,
each region exposes its ``tile_bboxes`` and ``codec_per_tile`` lists;
this implementation currently returns the full dense matrix for the
region (a v4 sparse-tile decode path is tracked as a follow-up).
"""
from __future__ import annotations
from typing import Literal
import numpy as np
from ..consts import Balancing
from ..exceptions import DatasetError
from ._base import TileDataset, _TileIndex
from ._torch_guard import require_torch
require_torch() # noqa: E402
__all__ = ["GzcmTileDataset"]
[docs]class GzcmTileDataset(TileDataset):
"""TileDataset reading from a ``.gzcm`` container file.
Parameters
----------
fpath : str
Path to the ``.gzcm`` file. May be v1, v2, v3, or v4.
bin_size_bp : int
Resolution in base pairs per axis bin.
window_size : int
Tile width in base pairs.
"""
def __init__(
self,
fpath: str,
bin_size_bp: int,
window_size: int,
output_type: Literal["sparse", "dense"] = "sparse",
downsample_ratio: float | tuple[float, float] | None = None,
balancing: Balancing | None = None,
decompress: bool = True,
lr_fpath: str | None = None,
lr_ds_ratio: int | None = None,
lr_balancing: Balancing | None = None,
) -> None:
self.fpath = fpath
super().__init__(
window_size=window_size,
bin_size_bp=bin_size_bp,
output_type=output_type,
downsample_ratio=downsample_ratio,
balancing=balancing,
decompress=decompress,
lr_fpath=lr_fpath,
lr_ds_ratio=lr_ds_ratio,
lr_balancing=lr_balancing,
)
self._reader = self._open_reader()
self._index = self._build_index()
def _open_reader(self):
"""Use the v3/v4 dispatcher from ``src/gunz_cm/datasets/gzcm.py``."""
from .gzcm import GZCMDataset
return GZCMDataset._open_reader(self.fpath)
def _build_index(self) -> _TileIndex:
"""Build a single-region tile index from the .gzcm metadata."""
import pandas as pd
n_bins = self._infer_n_bins()
tile_bins = self.window_size // self.bin_size_bp
starts = np.arange(0, n_bins, tile_bins, dtype=np.int64)
ends = np.minimum(starts + tile_bins, n_bins).astype(np.int64)
df = pd.DataFrame(
{
"chrom": getattr(self._reader, "chrom", "chr1"),
"start": starts * self.bin_size_bp,
"end": ends * self.bin_size_bp,
"start_bin": starts,
"end_bin": ends,
}
)
return _TileIndex(df)
def _infer_n_bins(self) -> int:
"""Infer the matrix size from the reader's metadata."""
meta = getattr(self._reader, "metadata", {})
if "original_shape" in meta:
return int(meta["original_shape"][0])
arrays_info = getattr(self._reader, "arrays_info", {}) or {}
matrix = arrays_info.get("matrix")
if matrix is not None:
shape = matrix.get("shape")
if shape and len(shape) >= 2:
return int(shape[0])
raise DatasetError(
f"GzcmTileDataset: cannot infer matrix size from {self.fpath}; "
f"metadata.original_shape missing"
)
def _fetch_patch(self, idx: int, s: int, e: int) -> np.ndarray:
"""Fetch a (h, w) dense patch via the existing tile-cache path."""
from .gzcm import GZCMDataset
if isinstance(self._reader, GZCMDataset):
return self._reader._get_compressed_patch(s, e)
dense = self._fetch_dense_matrix()
return dense[s:e, s:e]
def _fetch_dense_matrix(self) -> np.ndarray:
"""Materialise the full dense matrix for v3/v4 readers."""
arrays_info = getattr(self._reader, "arrays_info", {}) or {}
matrix_info = arrays_info.get("matrix")
if matrix_info is None:
raise DatasetError(
f"GzcmTileDataset: no 'matrix' array in {self.fpath}; v4 sparse-tiled "
f"intracellular decode not yet implemented"
)
return np.asarray(
np.memmap(
self.fpath,
dtype=np.dtype(matrix_info["dtype"]),
mode="r",
offset=matrix_info["offset"],
shape=tuple(matrix_info["shape"]),
)
)
def _fetch_lr_patch(self, s: int, e: int) -> np.ndarray:
"""LR/HR pair mode is not supported for GzcmTileDataset."""
raise NotImplementedError(
"GzcmTileDataset does not support LR/HR pair mode; use "
"MemmapTileDataset with lr_fpath=... and lr_ds_ratio=... instead."
)