scRep_pretrain/bin_cache.py

from __future__ import annotations

import hashlib
import json
from pathlib import Path
from typing import Any, Dict, Optional

import numpy as np


def cache_key_for_h5ad(path: str) -> str:
    resolved = str(Path(path).expanduser().resolve())
    digest = hashlib.sha1(resolved.encode("utf-8")).hexdigest()[:12]
    return f"{Path(path).stem}.{digest}"


def value_bin_cache_paths(cache_dir: str, h5ad_path: str) -> Dict[str, Path]:
    root = Path(cache_dir).expanduser().resolve()
    key = cache_key_for_h5ad(h5ad_path)
    return {
        "indptr": root / f"{key}.indptr.npy",
        "bin_ids": root / f"{key}.bin_ids.npy",
        "meta": root / f"{key}.meta.json",
    }


def load_value_bin_cache(
    cache_dir: str,
    h5ad_path: str,
    *,
    expected_n_cells: Optional[int] = None,
    mmap_mode: str = "r",
) -> Optional[Dict[str, Any]]:
    if not cache_dir:
        return None
    paths = value_bin_cache_paths(cache_dir, h5ad_path)
    if not paths["indptr"].exists() or not paths["bin_ids"].exists():
        return None

    indptr = np.load(paths["indptr"], mmap_mode=mmap_mode)
    bin_ids = np.load(paths["bin_ids"], mmap_mode=mmap_mode)
    if expected_n_cells is not None and int(indptr.shape[0]) != int(expected_n_cells) + 1:
        raise ValueError(
            f"bin cache n_cells mismatch for {h5ad_path}: "
            f"indptr_len={indptr.shape[0]} expected={int(expected_n_cells) + 1}"
        )

    meta: Dict[str, Any] = {}
    if paths["meta"].exists():
        with open(paths["meta"], "r", encoding="utf-8") as f:
            meta = json.load(f)
    return {"indptr": indptr, "bin_ids": bin_ids, "meta": meta, "paths": paths}