"""
PyTorch Dataset implementation for Fully Sparse Hi-C data loading.
Supports on-the-fly binomial downsampling and genomic window indexing.
Examples
--------
"""
__author__ = "Yeremia Gunawan Adhisantoso"
__email__ = "adhisant@tnt.uni-hannover.de"
__license__ = "Clear BSD"
import numpy as np
import pandas as pd
import typing as t
import warnings
from pydantic import validate_call
from ..loaders import load_cm_data, get_bins, DataStructure, Balancing
from ..consts import Backend
from ..utils import intervals
from ._torch_guard import require_torch
from .sparse_coo import SparseCODataset
from .sparse_collate import sparse_collate_fn
require_torch()
import torch # noqa: E402 (guarded by require_torch)
from ._torch_guard import DatasetBase as DatasetType
[docs]class HiCDataset(SparseCODataset):
"""
A PyTorch Dataset for on-the-fly loading of Hi-C patches from sparse files.
Inherits from :class:`SparseCODataset`; subclasses only need to implement
:meth:`_load_patch` (the RCV fetch from the file) and the genomic-index
mapping via :meth:`_patch_boundaries`. The 4-key output dict
(``coords``, ``features``, ``target``, ``info``), the downsampling
logic, and the dense output path all live in the base class.
"""
def __init__(
self,
fpath: str,
bin_size_bp: int,
window_size: int,
blacklist: pd.DataFrame | None = None,
downsample_ratio: float | tuple[float, float] | None = None,
balancing: Balancing | None = Balancing.NONE,
output_type: str = "sparse",
**kwargs,
):
# The base class owns the post-processing: downsample, output_type,
# and the 4-key contract. The subclass owns only the file-specific
# bits (path, kwargs, index) and overrides _load_patch.
super().__init__(
bin_size_bp=bin_size_bp,
window_size=window_size,
downsample_ratio=downsample_ratio,
output_type=output_type,
)
self.fpath = fpath
self.balancing = balancing
self.kwargs = kwargs
# 1. Generate Index
# We use window_size as the binning step for the training windows
self.index = get_bins(fpath, window_size)
# 2. Filter Index
if blacklist is not None:
self.index = intervals.subtract(self.index, blacklist)
def __len__(self) -> int:
return len(self.index)
def _patch_boundaries(self, idx: int):
row = self.index.iloc[idx]
return row['chrom'], int(row['start']), int(row['end'])
def _load_patch(self, idx: int):
chrom, start, end = self._patch_boundaries(idx)
data = load_cm_data(
self.fpath,
bin_size_bp=self.bin_size_bp,
region1=f"{chrom}:{start}-{end}",
balancing=self.balancing,
output_format=DataStructure.RCV,
backend=self.kwargs.pop("backend", Backend.HICTK),
**self.kwargs,
)
r_ids, c_ids, counts = data
return (
np.asarray(r_ids, dtype=np.int64),
np.asarray(c_ids, dtype=np.int64),
np.asarray(counts, dtype=np.float64),
{"chrom": chrom, "start": start, "end": end},
)
def _global_to_local(
self, row_ids: np.ndarray, col_ids: np.ndarray, start_bin: int
):
# The base class subtracts start_bin; here we just normalize.
return row_ids - start_bin, col_ids - start_bin
# 1-release deprecation alias; remove in 2.30.0
HiCDataset = HiCDataset
# 1-release deprecation alias; remove in 2.30.0
HiCSparseDataset = HiCDataset