"""Collate function for sparse-COO dataset batches.
This module defines :func:`sparse_collate_fn`, which takes a list of
4-key patch dicts (as produced by :class:`SparseCODataset.__getitem__`
or any subclass) and returns a batched dict with MinkowskiEngine-style
batch indices prepended to the coordinate tensor.
The contract: every batch item MUST have keys ``coords``, ``features``,
``target``, ``info``. The output dict has keys ``coords``, ``features``,
``target``, ``infos`` (plural — a list of per-item info dicts, in input
order).
The ``coords`` output has an extra column prepended (the batch index).
This is the standard pattern for sparse-batch tensor operations.
Examples
--------
>>> from gunz_cm.datasets.sparse_collate import sparse_collate_fn
>>> import torch
>>> batch = [
... {
... "coords": torch.tensor([[0, 0], [0, 1], [1, 1]], dtype=torch.long),
... "features": torch.tensor([[1.], [2.], [3.]]),
... "target": torch.tensor([1., 2., 3.]),
... "info": {"chrom": "chr1", "start": 0, "end": 100},
... },
... {
... "coords": torch.tensor([[2, 0], [2, 2]], dtype=torch.long),
... "features": torch.tensor([[4.], [5.]]),
... "target": torch.tensor([4., 5.]),
... "info": {"chrom": "chr1", "start": 100, "end": 200},
... },
... ]
>>> out = sparse_collate_fn(batch)
>>> out["coords"].shape
torch.Size([5, 3])
>>> out["coords"][:, 0] # batch index column
tensor([0, 0, 0, 1, 1])
"""
from __future__ import annotations
__author__ = "Yeremia Gunawan Adhisantoso"
__email__ = "adhisant@tnt.uni-hannover.de"
__license__ = "Clear BSD"
from typing import Any, Dict, List
from ._torch_guard import require_torch
require_torch()
import torch # noqa: E402 (guarded by require_torch)
[docs]def sparse_collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
"""Collate sparse-COO batch into a MinkowskiEngine-style dict.
Every batch item MUST have keys ``coords``, ``features``, ``target``,
``info``. The output has keys ``coords`` (with batch index prepended),
``features``, ``target``, ``infos`` (a list of per-item info dicts).
Empty input (no items) returns:
- coords: shape (0, 3), int64
- features: shape (0, 1), float32
- target: shape (0,), float32
- infos: []
"""
batch_coords: List[torch.Tensor] = []
batch_feats: List[torch.Tensor] = []
batch_targets: List[torch.Tensor] = []
batch_infos: List[Dict[str, Any]] = []
for i, item in enumerate(batch):
coords = item["coords"]
# Validate shape: coords must be 2-D.
if coords.dim() != 2 or coords.size(1) != 2:
raise ValueError(
f"sparse_collate_fn: item {i} coords must be shape (N, 2); got {tuple(coords.shape)}"
)
# Prepend batch index as the first column.
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(item["features"])
batch_targets.append(item["target"])
batch_infos.append(item.get("info", {}))
if batch_coords:
out_coords = torch.cat(batch_coords, dim=0)
out_features = torch.cat(batch_feats, dim=0)
out_targets = torch.cat(batch_targets, dim=0)
else:
out_coords = torch.zeros((0, 3), dtype=torch.long)
out_features = torch.zeros((0, 1), dtype=torch.float32)
out_targets = torch.zeros((0,), dtype=torch.float32)
return {
"coords": out_coords,
"features": out_features,
"target": out_targets,
"infos": batch_infos,
}
__all__ = ["sparse_collate_fn"]