"""Base class for tile-based PyTorch Datasets over Hi-C contact matrices.
Provides the cross-format contract for all three subclass trees:
TileDataset (abstract)
HiCTileDataset — sparse .hic / .mcool via load_cm_data
MemmapTileDataset — dense memmap via load_memmap
GzcmTileDataset — GZCM v1/v2/v3 via GzcmReader
The base owns:
* Index construction over (chrom, start_bin, end_bin) tiles.
* Output contract: dense ``(1, H, W)`` tensor OR sparse
``{coords, features, target, info}`` dict.
* Weight normalization (KR / VC / ICE / etc.) for both dense and sparse
outputs, with NaN/Inf sanitization.
* Binomial downsampling augmentation.
* ``sparse_collate_fn`` (MinkowskiEngine-style batch-index prepending).
* Optional LR/HR pair mode for resolution-enhancement training.
Subclasses are responsible only for:
* ``_build_index()`` — return a pandas DataFrame of tile rows.
* ``_fetch_patch(s, e)`` — return a dense ``(h, w)`` numpy patch for the
window spanning bins ``[s, e)``.
* ``_load_weights()`` — populate ``self.weights`` (1-D float array aligned
to global bin ids) if balancing is requested.
Examples
--------
>>> from gunz_cm.datasets import HiCTileDataset, sparse_collate_fn
>>> from gunz_cm.consts import Balancing
>>> ds = HiCTileDataset(
... fpath="data.hic",
... bin_size_bp=1_000_000,
... window_size=5_000_000,
... balancing=Balancing.KR,
... downsample_ratio=(0.3, 0.7),
... )
>>> item = ds[0]
>>> item["coords"].shape # (N, 2)
>>> item["features"].dtype # torch.float32
"""
from __future__ import annotations
import typing as t
from typing import Literal
import numpy as np
import pandas as pd
from ..consts import Balancing
from ..exceptions import DatasetError
from ._torch_guard import require_torch
require_torch() # noqa: E402
import torch # noqa: E402 (guarded by require_torch)
from ._torch_guard import DatasetBase as DatasetType # noqa: E402
__all__ = ["TileDataset", "_TileIndex", "sparse_collate_fn"]
# ---------------------------------------------------------------------------
# Index type
# ---------------------------------------------------------------------------
class _TileIndex:
"""Lightweight wrapper around the per-format tile DataFrame.
Subclasses return this from ``_build_index``. The DataFrame must have
at least columns ``start_bin`` and ``end_bin`` (ints, half-open ``[s, e)``
bin intervals). Optional columns ``chrom``, ``start``, ``end`` carry the
genomic coordinates in bp; subclasses that don't have genomic context
(e.g. memmap of a single chromosome) may leave these as ``None``.
"""
__slots__ = ("df",)
def __init__(self, df: pd.DataFrame, bin_size_bp: int | None = None) -> None:
df = df.copy()
if "start_bin" not in df.columns or "end_bin" not in df.columns:
if "start" not in df.columns or "end" not in df.columns:
raise DatasetError(
"_TileIndex requires either ('start_bin', 'end_bin') or "
f"('start', 'end') columns. Got: {list(df.columns)}"
)
if bin_size_bp is None:
raise DatasetError(
"_TileIndex needs bin_size_bp to derive start_bin/end_bin "
"from start/end."
)
df["start_bin"] = (df["start"] // bin_size_bp).astype(np.int64)
df["end_bin"] = (df["end"] // bin_size_bp).astype(np.int64)
else:
df["start_bin"] = df["start_bin"].astype(np.int64)
df["end_bin"] = df["end_bin"].astype(np.int64)
self.df = df.reset_index(drop=True)
def __len__(self) -> int:
return len(self.df)
def row(self, idx: int) -> pd.Series:
return self.df.iloc[idx]
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _apply_downsampling(
r: np.ndarray,
c: np.ndarray,
counts: np.ndarray,
ratio: float | tuple[float, float] | None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Binomial subsampling for sparse coordinates.
Parameters
----------
r, c : ndarray of int
Tile-local row and column indices.
counts : ndarray of float
Per-binomial-trial counts (must be integer-typed for np.random.binomial).
ratio : float, tuple, or None
If ``None``, returns the input unchanged. If a tuple ``(lo, hi)``, a
random ratio is sampled uniformly per call.
Returns
-------
(r, c, counts) : ndarrays
Filtered to entries that survived subsampling.
"""
if ratio is None:
return r, c, counts
if isinstance(ratio, tuple):
alpha = np.random.uniform(float(ratio[0]), float(ratio[1]))
else:
alpha = float(ratio)
# Binomial subsampling: c' ~ Binomial(n=c, p=alpha).
# counts must be integer for np.random.binomial; coerce via int32.
sampled = np.random.binomial(counts.astype(np.int32), alpha)
mask = sampled > 0
return r[mask], c[mask], sampled[mask]
def _normalize_dense(
patch: np.ndarray,
weights: np.ndarray | None,
s: int,
e: int,
) -> np.ndarray:
"""Apply per-bin weight normalization to a dense patch.
Computes ``patch / np.outer(weights[s:e], weights[s:e])`` with NaN/Inf
sanitization. Output dtype is float32.
One-pass implementation: the patch is converted to float32 in place
(no copy when already float32) and divided by the normalization
denominator using ``np.divide(out=...)`` to avoid a second
temporary buffer.
"""
if weights is None:
return patch.astype(np.float32, copy=False)
w_slice = weights[s:e]
denom = np.outer(w_slice, w_slice)
out = patch.astype(np.float32, copy=False)
np.divide(out, denom, out=out)
np.nan_to_num(out, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
return out
def _normalize_sparse(
r: np.ndarray,
c: np.ndarray,
counts: np.ndarray,
weights: np.ndarray | None,
s: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Apply per-bin weight normalization to a sparse COO triplet.
Returns filtered (r, c, counts) with NaN/Inf entries removed.
"""
if weights is None:
return r, c, counts
w1 = weights[s + r]
w2 = weights[s + c]
out = counts.astype(np.float32) / (w1 * w2)
mask = np.isfinite(out)
if not np.all(mask):
return r[mask], c[mask], out[mask]
return r, c, out
def _build_dense_tensor(patch: np.ndarray) -> torch.Tensor:
"""Wrap a (H, W) array as a (1, H, W) float32 torch tensor."""
return torch.from_numpy(patch.astype(np.float32, copy=False)).unsqueeze(0)
# ---------------------------------------------------------------------------
# TileDataset base class
# ---------------------------------------------------------------------------
[docs]class TileDataset(DatasetType):
"""Abstract base class for tile-based Hi-C contact matrix datasets.
Subclasses MUST implement ``_build_index()`` and ``_fetch_patch(s, e)``.
Subclasses MAY override ``_load_weights()`` to populate ``self.weights``
from format-specific storage (e.g. ``.hic`` via ``load_cm_data``,
``.gzcm`` via the embedded ``weights_*`` array).
Parameters
----------
window_size : int
Tile width in base pairs.
bin_size_bp : int
Resolution in base pairs per axis bin.
output_type : {"sparse", "dense"}
Output format from ``__getitem__``. Default ``"sparse"``.
downsample_ratio : float, tuple, or None
Binomial subsampling ratio. A tuple ``(lo, hi)`` samples uniformly
per call. ``None`` disables downsampling.
balancing : Balancing or None
Normalization method applied at fetch time. ``None`` skips
normalization.
decompress : bool
If ``False``, GZCM v3 datasets return raw tile bytes instead of
decoded arrays. Ignored by non-GZCM subclasses. Default ``True``.
lr_fpath : str or None
Path to a low-resolution contact matrix for resolution-enhancement
training. Setting this enables LR/HR pair output mode.
lr_ds_ratio : int or None
Downscale factor for on-the-fly low-resolution downsampling when
``lr_fpath`` is ``None``. Required when ``lr_fpath`` is set;
ignored otherwise.
lr_balancing : Balancing or None
Balancing method for the LR data. Defaults to ``balancing`` if unset.
Examples
--------
"""
# Subclass-overridable
_index: _TileIndex
def __init__(
self,
window_size: int,
bin_size_bp: 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:
if output_type not in ("sparse", "dense"):
raise DatasetError(
f"output_type must be 'sparse' or 'dense'; got {output_type!r}"
)
if window_size <= 0 or bin_size_bp <= 0:
raise DatasetError(
f"window_size and bin_size_bp must be positive; "
f"got window_size={window_size}, bin_size_bp={bin_size_bp}"
)
if window_size % bin_size_bp != 0:
raise DatasetError(
f"window_size ({window_size}) must be a multiple of "
f"bin_size_bp ({bin_size_bp})"
)
self.window_size = window_size
self.bin_size_bp = bin_size_bp
self.output_type = output_type
self.downsample_ratio = downsample_ratio
self.balancing = balancing
self.decompress = decompress
self.weights: np.ndarray | None = None
self.lr_fpath = lr_fpath
self.lr_ds_ratio = lr_ds_ratio
self.lr_balancing = lr_balancing if lr_balancing is not None else balancing
if lr_fpath is not None and lr_ds_ratio is None:
raise DatasetError("lr_ds_ratio is required when lr_fpath is set")
if type(self) is TileDataset:
raise DatasetError(
"TileDataset is abstract; instantiate a concrete subclass."
)
self._load_weights()
# ------------------------------------------------------------------
# Abstract hooks for subclasses
# ------------------------------------------------------------------
def _build_index(self) -> _TileIndex:
"""Build the per-format tile index. Must be implemented by subclasses."""
raise NotImplementedError
def _fetch_patch(self, idx: int, s: int, e: int) -> np.ndarray:
"""Return a dense (h, w) numpy patch for the tile at ``idx``.
``s, e`` are the half-open bin interval ``[s, e)`` of the tile;
``h = e - s``. ``idx`` is the tile's index in ``self._index``,
exposed so subclasses that key region lookup off the index row
(e.g. ``HiCTileDataset``) can use the right entry instead of
reading a fixed row. Subclasses that key off (s, e) only
(e.g. ``MemmapTileDataset``, ``GzcmTileDataset``) may ignore
``idx``. Subclasses should not apply normalization or
downsampling here; the base class owns those transformations.
"""
raise NotImplementedError
def _load_weights(self) -> None:
"""Populate ``self.weights`` (1-D float array) if balancing is requested.
Default: no-op. Subclasses that support external weight fetching
(e.g. ``MemmapTileDataset`` reading from a sibling ``.hic``) override
this.
"""
return None
def _fetch_lr_patch(self, s: int, e: int) -> np.ndarray:
"""Return the LR patch for the same window. Required when lr_fpath is set.
Default: raise NotImplementedError. Subclasses that support LR/HR
pair mode override this.
"""
raise NotImplementedError(
f"{type(self).__name__} does not support LR/HR pair mode. "
"Use MemmapTileDataset with lr_fpath and lr_ds_ratio."
)
# ------------------------------------------------------------------
# Index helpers (subclasses call these from _build_index)
# ------------------------------------------------------------------
@staticmethod
def _make_genomic_index(
chrom: str,
start_bin: int,
end_bin: int,
step: int,
chrom_total_bins: int,
bin_size_bp: int,
chrom_offset: int = 0,
) -> _TileIndex:
"""Build a genomic chrom/start/end index over a single chromosome.
Used by HiCTileDataset (per-chrom) and MemmapTileDataset (single matrix).
"""
starts = np.arange(start_bin, end_bin, step, dtype=np.int64)
ends = np.minimum(starts + step, chrom_total_bins).astype(np.int64)
df = pd.DataFrame({
"chrom": chrom,
"start": (starts + chrom_offset) * bin_size_bp,
"end": (ends + chrom_offset) * bin_size_bp,
"start_bin": starts,
"end_bin": ends,
})
return _TileIndex(df)
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
def __len__(self) -> int:
return len(self._index)
@property
def index(self) -> pd.DataFrame:
"""Public DataFrame view of the tile index (chrom/start/end + start_bin/end_bin).
Kept for backward compatibility with code that introspects the
index directly (e.g. ``SpatialBatchSampler``, test fixtures).
"""
return self._index.df
def __getitem__(self, idx: int) -> dict[str, t.Any] | torch.Tensor:
row = self._index.row(idx)
s, e = int(row["start_bin"]), int(row["end_bin"])
chrom = row.get("chrom") if "chrom" in row.index else None
start_bp = int(row["start"]) if "start" in row.index else s * self.bin_size_bp
end_bp = int(row["end"]) if "end" in row.index else e * self.bin_size_bp
info = {"chrom": chrom, "start": start_bp, "end": end_bp}
patch = self._fetch_patch(idx, s, e)
if self.output_type == "dense":
normalized = _normalize_dense(patch, self.weights, s, e)
dense_out = _build_dense_tensor(normalized)
if self.lr_fpath is None:
return dense_out
return self._wrap_lr_hr(dense_out, info, s, e, patch)
r, c = np.nonzero(patch)
mask_ut = r <= c
r, c = r[mask_ut], c[mask_ut]
counts = patch[r, c]
r, c, normalized = _normalize_sparse(r, c, counts, self.weights, s)
target = normalized.copy() if self.lr_fpath is not None else normalized
r, c, features = _apply_downsampling(r, c, normalized, self.downsample_ratio)
coords = np.stack([r, c], axis=1) if len(r) > 0 else np.zeros((0, 2), dtype=np.int64)
sparse_out = {
"coords": torch.from_numpy(coords).long(),
"features": torch.from_numpy(features).unsqueeze(1)
if len(features) > 0
else torch.zeros((0, 1), dtype=torch.float32),
"target": torch.from_numpy(target)
if len(target) > 0
else torch.zeros((0,), dtype=torch.float32),
"info": info,
}
if self.lr_fpath is None:
return sparse_out
return self._wrap_lr_hr(sparse_out, info, s, e, patch)
# ------------------------------------------------------------------
# LR/HR pair output
# ------------------------------------------------------------------
def _wrap_lr_hr(
self,
hr_item: dict[str, t.Any] | torch.Tensor,
info: dict[str, t.Any],
s: int,
e: int,
hr_patch_dense: np.ndarray,
) -> dict[str, t.Any]:
"""Augment the HR item with LR pair data.
For dense output, the result is a dict ``{hr, lr, info}``.
For sparse output, the result is a dict with ``hr`` / ``lr`` sub-dicts
each containing the standard sparse triplet, plus a top-level ``info``.
"""
lr_patch = self._fetch_lr_patch(s, e)
lr_dense = _normalize_dense(lr_patch, self.weights, s, e)
lr_tensor = _build_dense_tensor(lr_dense)
if isinstance(hr_item, torch.Tensor):
return {
"hr": hr_item,
"lr": lr_tensor,
"info": info,
}
# Re-extract HR coords WITHOUT downsampling so the model can compare
# the full HR target against the downsampled LR input.
r_hr, c_hr = np.nonzero(hr_patch_dense)
mask_ut = r_hr <= c_hr
r_hr, c_hr = r_hr[mask_ut], c_hr[mask_ut]
counts_hr = hr_patch_dense[r_hr, c_hr]
r_hr, c_hr, counts_hr = _normalize_sparse(
r_hr, c_hr, counts_hr, self.weights, s
)
coords_hr = np.stack([r_hr, c_hr], axis=1) if len(r_hr) > 0 else np.zeros((0, 2), dtype=np.int64)
r_lr, c_lr = np.nonzero(lr_patch)
mask_ut_lr = r_lr <= c_lr
r_lr, c_lr = r_lr[mask_ut_lr], c_lr[mask_ut_lr]
counts_lr = lr_patch[r_lr, c_lr]
r_lr, c_lr, counts_lr = _normalize_sparse(
r_lr, c_lr, counts_lr, self.weights, s
)
coords_lr = np.stack([r_lr, c_lr], axis=1) if len(r_lr) > 0 else np.zeros((0, 2), dtype=np.int64)
return {
"hr": {
"coords": torch.from_numpy(coords_hr).long(),
"features": torch.from_numpy(counts_hr).unsqueeze(1)
if len(counts_hr) > 0
else torch.zeros((0, 1), dtype=torch.float32),
"target": torch.from_numpy(counts_hr)
if len(counts_hr) > 0
else torch.zeros((0,), dtype=torch.float32),
},
"lr": {
"coords": torch.from_numpy(coords_lr).long(),
"features": torch.from_numpy(counts_lr).unsqueeze(1)
if len(counts_lr) > 0
else torch.zeros((0, 1), dtype=torch.float32),
"target": torch.from_numpy(counts_lr)
if len(counts_lr) > 0
else torch.zeros((0,), dtype=torch.float32),
},
"info": info,
}
# ---------------------------------------------------------------------------
# Sparse collate function
# ---------------------------------------------------------------------------
def sparse_collate_fn(batch: list[t.Any]) -> dict[str, t.Any] | torch.Tensor:
"""Collate a batch of sparse tile items into a MinkowskiEngine-style dict.
Each input item is expected to have a ``"coords"`` key with shape
``(N_i, 2)`` and a ``"features"`` key with shape ``(N_i, 1)``.
The output ``coords`` has shape ``(sum_i N_i, 3)`` with columns
``[batch_idx, r, c]``. ``features`` is concatenated along dim 0.
``infos`` (if present) is preserved as a list in input order.
LR/HR pair mode is supported transparently: if the item has ``"hr"``
and ``"lr"`` sub-dicts, each is collated separately and the result is
``{"hr": ..., "lr": ..., "infos": [...]}``.
Examples
--------
"""
if not batch:
raise DatasetError("sparse_collate_fn received empty batch")
first = batch[0]
if isinstance(first, torch.Tensor):
tensors: list[torch.Tensor] = [t for t in batch if isinstance(t, torch.Tensor)] # type: ignore[misc]
return torch.stack(tensors, dim=0)
if "hr" in first and "lr" in first:
hr_items = [item["hr"] for item in batch]
lr_items = [item["lr"] for item in batch]
return {
"hr": sparse_collate_fn(hr_items),
"lr": sparse_collate_fn(lr_items),
"infos": [item.get("info", {}) for item in batch],
}
batch_coords: list[torch.Tensor] = []
batch_feats: list[torch.Tensor] = []
for i, item in enumerate(batch):
coords = item["coords"]
feats = item["features"]
batch_idx = torch.full((coords.shape[0], 1), i, dtype=torch.long)
batch_coords.append(torch.cat([batch_idx, coords], dim=1))
batch_feats.append(feats)
return {
"coords": torch.cat(batch_coords, dim=0),
"features": torch.cat(batch_feats, dim=0),
"infos": [item.get("info", {}) for item in batch],
}