from __future__ import annotations from dataclasses import dataclass from typing import Any, Dict, List, Optional, Tuple from collections import OrderedDict import os import torch import anndata as ad import numpy as np import scipy.sparse as sp from torch.utils.data import IterableDataset, get_worker_info from .bin_cache import load_value_bin_cache # ============================================================================= # Debug / Logging # ============================================================================= DEBUG_PRINT: bool = False def _log(msg: str) -> None: if DEBUG_PRINT: print(f"[sc-token] {msg}") def _safe_close_anndata(adata_obj: Any) -> None: f = getattr(adata_obj, "file", None) if f is not None: try: f.close() except Exception: pass # ============================================================================= # Per-file materialized reader # ============================================================================= class H5ADFileReader: """ Per-file reader. Matrix semantics: - use_raw=True : adata.raw.X counts -> log1p(CP10K) - use_raw=False : adata.X is assumed to already be log1p(CP10K) Design notes: - This reader always returns ALL observed nonzero genes for a cell. - No top-k truncation happens here anymore. - Top-k is now applied later inside the collator. - Gene order is preserved here; any shuffling happens in the collator. """ def __init__( self, path: str, dataset_id: int, *, gene_name_to_id: Dict[str, int], meta_maps: Dict[str, Dict[str, int]], meta_keys: List[str], use_raw: bool, cp_target_sum: float = 1e4, backed: Optional[str] = "r", prefer_csr: bool = True, strict_gene: bool = True, filter_missing_genes: bool = False, value_bin_cache_dir: str = "", require_value_bin_cache: bool = False, ) -> None: self.path = path self.dataset_id = int(dataset_id) self.gene_name_to_id = gene_name_to_id self.meta_maps = meta_maps self.meta_keys = list(meta_keys) self.use_raw = bool(use_raw) self.cp_target_sum = float(cp_target_sum) self.prefer_csr = bool(prefer_csr) self.filter_missing_genes = bool(filter_missing_genes) self.value_bin_cache_dir = str(value_bin_cache_dir or "") self.require_value_bin_cache = bool(require_value_bin_cache) self._bin_cache_indptr = None self._bin_cache_ids = None _log(f"open: {path} (use_raw={self.use_raw}, backed={backed})") self.adata = ad.read_h5ad(path, backed=backed) self.n_cells = int(self.adata.n_obs) if self.use_raw: if (self.adata.raw is None) or (self.adata.raw.X is None): raise ValueError(f"[{path}] use_raw=True but adata.raw.X missing.") self.X = self.adata.raw.X self.var_names = self.adata.raw.var_names else: if self.adata.X is None: raise ValueError(f"[{path}] use_raw=False but adata.X missing.") self.X = self.adata.X self.var_names = self.adata.var_names self.local_to_global = self._build_local_to_global(strict=strict_gene) self.meta_int_cols = self._build_meta_int_cols() self._load_value_bin_cache() self._is_csr = sp.isspmatrix_csr(self.X) if self._is_csr and self.prefer_csr: self._csr_indptr = self.X.indptr self._csr_indices = self.X.indices self._csr_data = self.X.data _log(f"[{path}] n_cells={self.n_cells}") def close(self) -> None: _safe_close_anndata(self.adata) def _load_value_bin_cache(self) -> None: if not self.value_bin_cache_dir: return cache = load_value_bin_cache( self.value_bin_cache_dir, self.path, expected_n_cells=self.n_cells, mmap_mode="r", ) if cache is None: if self.require_value_bin_cache: raise FileNotFoundError( f"Missing value bin cache for {self.path} under {self.value_bin_cache_dir}" ) return self._bin_cache_indptr = cache["indptr"] self._bin_cache_ids = cache["bin_ids"] def has_value_bin_cache(self) -> bool: return self._bin_cache_indptr is not None and self._bin_cache_ids is not None def get_cached_bin_ids(self, cell_idx: int) -> Optional[np.ndarray]: if not self.has_value_bin_cache(): return None i = int(cell_idx) start = int(self._bin_cache_indptr[i]) end = int(self._bin_cache_indptr[i + 1]) return np.asarray(self._bin_cache_ids[start:end], dtype=np.int64) def _build_local_to_global(self, *, strict: bool) -> np.ndarray: ids = np.fromiter( (self.gene_name_to_id.get(str(g), 0) for g in list(self.var_names)), dtype=np.int64, count=len(self.var_names), ) if strict and int(np.sum(ids == 0)) > 0: missing = int(np.sum(ids == 0)) raise ValueError( f"[{self.path}] {missing} genes missing in global gene vocab. " "Rebuild vocab with full manifest or check gene naming consistency." ) return ids def _build_meta_int_cols(self) -> Dict[str, np.ndarray]: out: Dict[str, np.ndarray] = {} for k in self.meta_keys: if k not in self.adata.obs.columns: raise ValueError(f"[{self.path}] obs has no column '{k}'") if k not in self.meta_maps: raise ValueError(f"meta_maps missing key='{k}'. Build meta vocabs first.") gmap = self.meta_maps[k] col = self.adata.obs[k].astype("category") cats = col.cat.categories.astype(str).tolist() codes = col.cat.codes.to_numpy() cat_ids = np.array([gmap.get(c, 0) for c in cats], dtype=np.int32) ids = np.zeros((self.n_cells,), dtype=np.int32) valid = codes >= 0 ids[valid] = cat_ids[codes[valid]] out[k] = ids return out def _csr_row_indices_data(self, cell_idx: int) -> Tuple[np.ndarray, np.ndarray]: i = int(cell_idx) start = int(self._csr_indptr[i]) end = int(self._csr_indptr[i + 1]) return self._csr_indices[start:end], self._csr_data[start:end] def _row_indices_data(self, cell_idx: int) -> Tuple[np.ndarray, np.ndarray]: if self._is_csr and self.prefer_csr: idx, val = self._csr_row_indices_data(cell_idx) else: row = self.X[int(cell_idx)] if sp.issparse(row): if not sp.isspmatrix_csr(row): row = row.tocsr() idx, val = row.indices, row.data else: dense = np.asarray(row).ravel() idx = np.flatnonzero(dense) val = dense[idx] if val.size > 0: keep = val != 0 if not np.all(keep): idx = idx[keep] val = val[keep] return idx, val def get_cell_tokens(self, cell_idx: int) -> Tuple[List[int], List[float]]: idx, val = self._row_indices_data(cell_idx) if idx.size == 0: return [], [] row_sum = float(val.sum()) if self.use_raw: if row_sum <= 0: return [], [] expr = np.log1p(val.astype(np.float32) * (self.cp_target_sum / row_sum)) else: expr = val.astype(np.float32) gene_ids = self.local_to_global[idx.astype(np.int64)].astype(np.int64) if self.filter_missing_genes: keep = gene_ids > 0 if not np.all(keep): gene_ids = gene_ids[keep] expr = expr[keep] if gene_ids.size == 0: return [], [] return gene_ids.tolist(), expr.tolist() def get_cell_observation_stats(self, cell_idx: int) -> Tuple[float, int]: idx, val = self._row_indices_data(cell_idx) detected_genes = int(idx.size) if self.use_raw: return float(val.sum()), detected_genes obs = self.adata.obs row = obs.iloc[int(cell_idx)] total_counts = None for key in ("total_counts", "n_counts"): if key not in obs.columns: continue candidate = float(row[key]) if np.isfinite(candidate): total_counts = candidate break if total_counts is None: total_counts = float(detected_genes) for key in ("n_genes_by_counts", "n_genes"): if key not in obs.columns: continue candidate = float(row[key]) if np.isfinite(candidate): detected_genes = int(candidate) break return total_counts, detected_genes def get_meta_batch(self, cell_ids: np.ndarray, *, include_dataset_id: bool = False) -> Dict[str, List[int]]: ids = cell_ids.astype(np.int64) out: Dict[str, List[int]] = {} if include_dataset_id: out["dataset_id"] = [self.dataset_id + 1] * int(ids.size) for k in self.meta_keys: out[k] = self.meta_int_cols[k][ids].astype(np.int64).tolist() return out # ============================================================================= # Global pseudo-batch style index metadata # ============================================================================= @dataclass class _PseudoBatchGroupIndex: """ One synthetic group inside one file. For scRep we keep one group containing all cells in the file. """ raw_cell_ids: np.ndarray # int32 n_batches: int usable_cells: int @dataclass class _FileIndex: path: str dataset_id: int groups: List[_PseudoBatchGroupIndex] # ============================================================================= # Small LRU reader cache # ============================================================================= class _ReaderCache: """ Small LRU cache of open H5ADFileReader objects. Each worker / rank gets its own cache inside __iter__. """ def __init__( self, *, max_open_files: int, gene_name_to_id: Dict[str, int], meta_maps: Dict[str, Dict[str, int]], meta_keys: List[str], use_raw: bool, cp_target_sum: float, backed: Optional[str], prefer_csr: bool, strict_gene: bool, filter_missing_genes: bool, value_bin_cache_dir: str, require_value_bin_cache: bool, ) -> None: self.max_open_files = int(max_open_files) self.gene_name_to_id = gene_name_to_id self.meta_maps = meta_maps self.meta_keys = meta_keys self.use_raw = bool(use_raw) self.cp_target_sum = float(cp_target_sum) self.backed = backed self.prefer_csr = bool(prefer_csr) self.strict_gene = bool(strict_gene) self.filter_missing_genes = bool(filter_missing_genes) self.value_bin_cache_dir = str(value_bin_cache_dir or "") self.require_value_bin_cache = bool(require_value_bin_cache) self._cache: "OrderedDict[int, H5ADFileReader]" = OrderedDict() def get(self, file_id: int, file_meta: _FileIndex) -> H5ADFileReader: if file_id in self._cache: reader = self._cache.pop(file_id) self._cache[file_id] = reader return reader reader = H5ADFileReader( path=file_meta.path, dataset_id=file_meta.dataset_id, gene_name_to_id=self.gene_name_to_id, meta_maps=self.meta_maps, meta_keys=self.meta_keys, use_raw=self.use_raw, cp_target_sum=self.cp_target_sum, backed=self.backed, prefer_csr=self.prefer_csr, strict_gene=self.strict_gene, filter_missing_genes=self.filter_missing_genes, value_bin_cache_dir=self.value_bin_cache_dir, require_value_bin_cache=self.require_value_bin_cache, ) self._cache[file_id] = reader while len(self._cache) > self.max_open_files: _, old_reader = self._cache.popitem(last=False) try: old_reader.close() except Exception: pass return reader def close_all(self) -> None: for _, reader in self._cache.items(): try: reader.close() except Exception: pass self._cache.clear() # ============================================================================= # Multi-file single-cell stream # ============================================================================= class MultiH5ADCellStream(IterableDataset): """ Globally shuffled single-cell stream. Changes from the old implementation: - Build a static global cell-slot index once. - Sample slots by a global permutation, so cells are drawn with weight proportional to file size (number of cells), not by equal-file coin flips. - Still materializes one cell per yielded example. """ def __init__( self, manifest: List[Dict[str, Any]], *, gene_name_to_id: Dict[str, int], meta_maps: Dict[str, Dict[str, int]], meta_keys: Optional[List[str]] = None, buffer_size: int = 16, seed: int = 42, use_raw: bool = True, cp_target_sum: float = 1e4, backed: Optional[str] = "r", prefer_csr: bool = True, sort_cell_ids_for_io: bool = False, include_dataset_id: bool = False, skip_bad_files: bool = False, filter_missing_genes: bool = False, value_bin_cache_dir: str = "", require_value_bin_cache: bool = False, ) -> None: self.manifest = list(manifest) self.gene_name_to_id = gene_name_to_id self.meta_maps = meta_maps self.meta_keys = list(meta_keys or []) self.full_meta_keys = list(self.meta_keys) # Kept for compatibility with existing call sites / configs. self.buffer_size = int(buffer_size) self.seed = int(seed) # Equivalent to sampling one-cell pseudo-batches. self.N = 1 self.use_raw = bool(use_raw) self.cp_target_sum = float(cp_target_sum) self.backed = backed self.prefer_csr = bool(prefer_csr) self.filter_missing_genes = bool(filter_missing_genes) self.strict_gene = not self.filter_missing_genes self.sort_cell_ids_for_io = bool(sort_cell_ids_for_io) self.include_dataset_id = bool(include_dataset_id) self.skip_bad_files = bool(skip_bad_files) self.value_bin_cache_dir = str(value_bin_cache_dir or "") self.require_value_bin_cache = bool(require_value_bin_cache) # reader cache self.max_open_files = 1024 self.cycle = True self._iter_id = 0 if len(self.manifest) == 0: raise RuntimeError("Empty manifest.") # Built once self.file_index: List[_FileIndex] = [] # shape: (num_global_pseudo_batches, 3), columns = [file_id, group_id, batch_id] self.batch_refs: np.ndarray = np.empty((0, 3), dtype=np.int32) self._build_global_index() # ------------------------------------------------------------------------- # Global index construction # ------------------------------------------------------------------------- def _scan_one_file(self, item: Dict[str, Any], file_id: int) -> Optional[_FileIndex]: path = item["path"] dataset_id = int(item.get("dataset_id", 0)) adata = None try: adata = ad.read_h5ad(path, backed=self.backed) n_cells = int(adata.n_obs) if n_cells < self.N: _log(f"skip file (n_cells < N): {path}") return None # Keep failure mode aligned with old stream: invalid/missing columns fail at scan # rather than mid-iteration. for k in self.full_meta_keys: if k not in adata.obs.columns: raise ValueError(f"[{path}] obs has no required column '{k}'.") raw_cells = np.arange(n_cells, dtype=np.int32) n_batches = int(raw_cells.size // self.N) if n_batches <= 0: _log(f"skip file (no full pseudo-batches): {path}") return None usable_cells = int(n_batches * self.N) group = _PseudoBatchGroupIndex( raw_cell_ids=raw_cells, n_batches=n_batches, usable_cells=usable_cells, ) return _FileIndex(path=path, dataset_id=dataset_id, groups=[group]) finally: if adata is not None: _safe_close_anndata(adata) def _build_global_index(self) -> None: file_index: List[_FileIndex] = [] batch_refs: List[Tuple[int, int, int]] = [] for src_file_id, item in enumerate(self.manifest): try: fidx = self._scan_one_file(item, src_file_id) except Exception as e: if self.skip_bad_files: _log(f"skip bad file: {item.get('path')} | {e}") continue raise if fidx is None: continue new_file_id = len(file_index) file_index.append(fidx) for group_id, g in enumerate(fidx.groups): for batch_id in range(int(g.n_batches)): batch_refs.append((new_file_id, group_id, batch_id)) if len(file_index) == 0 or len(batch_refs) == 0: raise RuntimeError( "No eligible cells found. " "Check manifest, requested meta keys, and skip_bad_files setting." ) self.file_index = file_index self.batch_refs = np.asarray(batch_refs, dtype=np.int32) total_files = len(self.file_index) total_groups = int(sum(len(f.groups) for f in self.file_index)) total_batches = int(self.batch_refs.shape[0]) total_cells_usable = int(sum(g.usable_cells for f in self.file_index for g in f.groups)) _log( f"global index built: files={total_files}, groups={total_groups}, " f"slots={total_batches}, usable_cells={total_cells_usable}" ) # ------------------------------------------------------------------------- # Introspection helpers # ------------------------------------------------------------------------- @property def num_global_pseudo_batches(self) -> int: return int(self.batch_refs.shape[0]) def estimate_batches_per_shard(self, num_shards: int) -> int: if num_shards <= 0: return self.num_global_pseudo_batches return (self.num_global_pseudo_batches + num_shards - 1) // num_shards # ------------------------------------------------------------------------- # Per-epoch group reshuffle # ------------------------------------------------------------------------- def _build_epoch_group_batches(self, epoch_seed: int) -> List[List[np.ndarray]]: """ For the current epoch/pass, rebuild each file/group by: - reshuffling cell ids - dropping tail - reshaping to (n_batches, N) Returns: epoch_group_batches[file_id][group_id] = np.ndarray of shape (n_batches, N) """ epoch_group_batches: List[List[np.ndarray]] = [] for file_id, fidx in enumerate(self.file_index): file_batches: List[np.ndarray] = [] for group_id, gidx in enumerate(fidx.groups): group_seed = ( int(epoch_seed) + 1000003 * int(file_id + 1) + 9176 * int(group_id + 1) ) % (2**63 - 1) rng = np.random.default_rng(group_seed) raw = gidx.raw_cell_ids perm = rng.permutation(raw.size) shuffled = raw[perm] usable = shuffled[: gidx.usable_cells] batch_cells = usable.reshape(gidx.n_batches, self.N).astype(np.int32, copy=False) file_batches.append(batch_cells) epoch_group_batches.append(file_batches) return epoch_group_batches # ------------------------------------------------------------------------- # Materialization # ------------------------------------------------------------------------- def _materialize_one_batch( self, *, file_id: int, group_id: int, batch_id: int, epoch_group_batches: List[List[np.ndarray]], cache: _ReaderCache, ) -> Dict[str, Any]: fmeta = self.file_index[file_id] cell_ids = epoch_group_batches[file_id][group_id][batch_id] if self.sort_cell_ids_for_io: cell_ids = np.sort(cell_ids) reader = cache.get(file_id, fmeta) meta_acc = reader.get_meta_batch( cell_ids, include_dataset_id=self.include_dataset_id, ) cid = int(cell_ids[0]) gene_ids, expr = reader.get_cell_tokens(cid) total_counts, n_genes_by_counts = reader.get_cell_observation_stats(cid) out = { "gene_ids": gene_ids, "expr": expr, "meta": meta_acc, "total_counts": float(total_counts), "n_genes_by_counts": int(n_genes_by_counts), } cached_bin_ids = reader.get_cached_bin_ids(cid) if cached_bin_ids is not None: if int(cached_bin_ids.shape[0]) != len(gene_ids): raise ValueError( f"Cached bin length mismatch for {reader.path} cell={cid}: " f"bin_ids={cached_bin_ids.shape[0]} genes={len(gene_ids)}" ) out["bin_ids"] = cached_bin_ids.astype(np.int64, copy=False).tolist() return out # ------------------------------------------------------------------------- # Iteration # ------------------------------------------------------------------------- def __iter__(self): # ------------------------- # rank / world # ------------------------- if torch.distributed.is_available() and torch.distributed.is_initialized(): rank = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() else: rank = int(os.environ.get("RANK", "0")) world_size = int(os.environ.get("WORLD_SIZE", "1")) # ------------------------- # worker info # ------------------------- info = get_worker_info() if info is None: worker_id, num_workers = 0, 1 else: worker_id, num_workers = info.id, info.num_workers shard_id = rank * num_workers + worker_id num_shards = world_size * num_workers iter_id = getattr(self, "_iter_id", 0) self._iter_id = iter_id + 1 base_seed = int(self.seed) mix = (base_seed + 10000019 * iter_id) % (2**63 - 1) rng = np.random.default_rng(mix) if DEBUG_PRINT: _log( f"rank={rank}/{world_size} worker={worker_id}/{num_workers} " f"iter_id={iter_id} seed={mix} shard={shard_id}/{num_shards} " f"global_pseudo_batches={self.num_global_pseudo_batches}" ) cache = _ReaderCache( max_open_files=self.max_open_files, gene_name_to_id=self.gene_name_to_id, meta_maps=self.meta_maps, meta_keys=self.full_meta_keys, use_raw=self.use_raw, cp_target_sum=self.cp_target_sum, backed=self.backed, prefer_csr=self.prefer_csr, strict_gene=self.strict_gene, filter_missing_genes=self.filter_missing_genes, value_bin_cache_dir=self.value_bin_cache_dir, require_value_bin_cache=self.require_value_bin_cache, ) try: epoch_idx = 0 while True: epoch_seed = (mix + 1000003 * epoch_idx) % (2**63 - 1) # 1) Re-shuffle every group for this epoch epoch_group_batches = self._build_epoch_group_batches(epoch_seed) # 2) Global shuffle over static slot table global_order = rng.permutation(self.num_global_pseudo_batches) # 3) Shard by rank * worker local_order = global_order[shard_id::num_shards] if DEBUG_PRINT: _log( f"new epoch/pass: epoch_idx={epoch_idx}, total={self.num_global_pseudo_batches}, " f"local={len(local_order)}, shard={shard_id}/{num_shards}" ) for ref_idx in local_order.tolist(): file_id, group_id, batch_id = self.batch_refs[int(ref_idx)].tolist() yield self._materialize_one_batch( file_id=int(file_id), group_id=int(group_id), batch_id=int(batch_id), epoch_group_batches=epoch_group_batches, cache=cache, ) epoch_idx += 1 if not self.cycle: break finally: cache.close_all()