"""HiCTileDataset — TileDataset over .hic / .mcool via load_cm_data.
This is one of three concrete subclasses of :class:`TileDataset` that
ship with the gunz-cm dataset-zoo (see ``specs/tile-dataset-consolidation.md``).
The other two are :class:`MemmapTileDataset` (memmap input) and
:class:`GzcmTileDataset` (GZCM v1-v3 + v4 input).
The class owns only the file-specific bits: how to build the per-chrom
tile index and how to fetch a (chrom, start_bin, end_bin) patch via
``load_cm_data``. Everything else (output contract, downsample,
balancing, LR/HR pair mode, sparse collate) lives in the base class.
"""
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__ = ["HiCTileDataset"]
[docs]class HiCTileDataset(TileDataset):
"""TileDataset reading from ``.hic`` or ``.mcool`` via ``load_cm_data``.
Parameters
----------
fpath : str
Path to the ``.hic`` / ``.mcool`` / ``.cool`` source file.
bin_size_bp : int
Resolution in base pairs per axis bin.
window_size : int
Tile width in base pairs. Must be a multiple of ``bin_size_bp``.
chrom : str, optional
Restrict iteration to this chromosome. If ``None``, iterate over
every chromosome reported by the loader.
downsample_ratio, output_type, balancing, decompress, lr_fpath,
lr_ds_ratio, lr_balancing : forwarded to :class:`TileDataset`.
Examples
--------
>>> ds = HiCTileDataset(
... fpath="data.hic",
... bin_size_bp=50_000,
... window_size=500_000,
... chrom="chr1",
... )
>>> len(ds) > 0
True
>>> item = ds[0]
>>> item["coords"].dtype
torch.int64
"""
def __init__(
self,
fpath: str,
bin_size_bp: int,
window_size: int,
chrom: str | None = None,
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
self._chrom = chrom
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._index = self._build_index()
def _build_index(self) -> _TileIndex:
"""Build the per-chrom tile index from ``load_cm_data``."""
from ..loaders import get_bins
bins = get_bins(self.fpath, self.window_size)
if self._chrom is not None:
bins = bins[bins["chrom"] == self._chrom]
return _TileIndex(bins)
def _fetch_patch(self, idx: int, s: int, e: int) -> np.ndarray:
"""Fetch a (h, w) dense patch for tile ``idx`` via ``load_cm_data``."""
from ..loaders import load_cm_data, DataStructure
row = self._index.row(idx)
chrom = row["chrom"]
start_bp = int(row["start"])
end_bp = int(row["end"])
result = load_cm_data(
self.fpath,
chrom=chrom,
start=start_bp,
end=end_bp,
resolution=self.bin_size_bp,
balancing=self.balancing,
output_format=DataStructure.NUMPY,
)
if result is None or not hasattr(result, "data"):
raise DatasetError(
f"HiCTileDataset: load_cm_data returned no data for "
f"{chrom}:{start_bp}-{end_bp} at {self.bin_size_bp}"
)
return np.asarray(result.data, dtype=np.float32)
def _fetch_lr_patch(self, s: int, e: int) -> np.ndarray:
"""Return the LR patch via the dedicated LR file or on-the-fly downsample.
If ``lr_fpath`` is set, fetch from that file at the same (s, e) window.
If unset, downsample the HR patch by ``lr_ds_ratio``.
"""
if self.lr_fpath is not None:
from ..loaders import load_cm_data, DataStructure
row = self._index.row(0)
chrom = row["chrom"]
start_bp = int(row["start"])
end_bp = int(row["end"])
lr_bin = self.bin_size_bp * (self.lr_ds_ratio or 1)
lr_result = load_cm_data(
self.lr_fpath,
chrom=chrom,
start=start_bp,
end=end_bp,
resolution=lr_bin,
balancing=self.lr_balancing,
output_format=DataStructure.NUMPY,
)
if lr_result is None or not hasattr(lr_result, "data"):
raise DatasetError(
f"HiCTileDataset: LR load_cm_data returned no data for "
f"{chrom}:{start_bp}-{end_bp}"
)
return np.asarray(lr_result.data, dtype=np.float32)
hr = self._fetch_patch(0, s, e)
factor = self.lr_ds_ratio or 1
h, w = hr.shape
return hr[::factor, ::factor]