"""MemmapTileDataset — TileDataset over an in-memory dense contact matrix.
This is the simplest of the three concrete TileDataset subclasses. It
loads a 2-D numpy array (or constructs one on the fly) and iterates
non-overlapping tiles of shape ``(window_size / bin_size_bp)`` square.
Examples
--------
>>> import numpy as np
>>> from gunz_cm.datasets.tile_memmap import MemmapTileDataset
>>> mat = np.eye(64, dtype=np.float32) + 0.1
>>> ds = MemmapTileDataset(
... matrix=mat, bin_size_bp=50_000, window_size=400_000,
... )
>>> len(ds)
1
>>> ds[0]["coords"].shape
(64, 2)
"""
from __future__ import annotations
from typing import Literal
import numpy as np
from ..consts import Balancing
from ._base import TileDataset, _TileIndex
from ._torch_guard import require_torch
require_torch() # noqa: E402
__all__ = ["MemmapTileDataset"]
[docs]class MemmapTileDataset(TileDataset):
"""TileDataset reading from a 2-D numpy array (memmap-friendly)."""
def __init__(
self,
matrix: np.ndarray,
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,
weights: np.ndarray | None = None,
) -> None:
self.matrix = np.asarray(matrix, dtype=np.float32)
if self.matrix.ndim != 2 or self.matrix.shape[0] != self.matrix.shape[1]:
raise ValueError(
f"MemmapTileDataset expects a square 2-D matrix; got shape {self.matrix.shape}"
)
self.n_bins = self.matrix.shape[0]
self._external_weights = weights
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 a non-overlapping tile index covering the full matrix."""
tile_bins = self.window_size // self.bin_size_bp
starts = np.arange(0, self.n_bins, tile_bins, dtype=np.int64)
ends = np.minimum(starts + tile_bins, self.n_bins).astype(np.int64)
df_dict = {
"start_bin": starts,
"end_bin": ends,
"start": starts * self.bin_size_bp,
"end": ends * self.bin_size_bp,
}
import pandas as pd
df = pd.DataFrame(df_dict)
return _TileIndex(df)
def _load_weights(self) -> None:
"""Honour the externally supplied weights array if any."""
if self._external_weights is not None:
self.weights = np.asarray(self._external_weights, dtype=np.float32)
def _fetch_patch(self, idx: int, s: int, e: int) -> np.ndarray:
"""Return the (h, w) dense slice for tile ``idx``."""
return self.matrix[s:e, s:e]
def _fetch_lr_patch(self, s: int, e: int) -> np.ndarray:
"""Downsample the HR patch in-place by ``lr_ds_ratio``."""
hr = self.matrix[s:e, s:e]
factor = self.lr_ds_ratio or 1
return hr[::factor, ::factor]