diff --git a/pyproject.toml b/pyproject.toml index 793ddac..9b475db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,7 @@ classifiers = [ ] dependencies = [ "adjusttext", - "anndata", + "anndata>=0.12", "bed-reader", "dask", "limix-core", @@ -41,21 +41,16 @@ dependencies = [ "requests", "rich", "scanpy", - "xarray", - "zarr>=3", # for debug logging (referenced from the issue template) "session-info", + "xarray", + "zarr>=3", ] optional-dependencies.datasets = [ - "liftover", - "sgkit[plink]", # TODO cbgen is a pain to install on mac. I think cmake issues. "bio2zarr[all]", -] - -optional-dependencies.pgen = [ - "pgenlib", - "zarr>3", + "liftover", + "sgkit[plink]", # TODO cbgen is a pain to install on mac. I think cmake issues. ] optional-dependencies.dev = [ @@ -64,9 +59,9 @@ optional-dependencies.dev = [ ] optional-dependencies.doc = [ "docutils>=0.8,!=0.18.*,!=0.19.*", - "linkify-it-py", "ipykernel", "ipython", + "linkify-it-py", "myst-nb>=1.1", # Until pybtex >0.23.0 releases: https://bitbucket.org/pybtex-devs/pybtex/issues/169/ "setuptools", @@ -93,12 +88,21 @@ optional-dependencies.mixmil = [ optional-dependencies.ml = [ "pytorch-lightning", ] +optional-dependencies.pgen = [ + "pgenlib", + "zarr>3", +] + optional-dependencies.rvat = [ # installation wit pip doesn't work, install via conda install -c conda-forge chiscore # "chiscore", "limix-core", "pybiomart", ] +optional-dependencies.scdrs = [ + "numpy<2", + "scdrs", +] optional-dependencies.snpeff = [ # not available via pip, has to be installed via conda install bioconda::snpeff # "snpeff", @@ -107,10 +111,6 @@ optional-dependencies.tensorqtl = [ "tensorqtl", #"plink2" ] -optional-dependencies.scdrs = [ - "scdrs", - "numpy<2.0.0", -] optional-dependencies.test = [ "coverage", "pytest", @@ -120,6 +120,8 @@ urls.Documentation = "https://cellink-docs.readthedocs.io/" urls.Homepage = "https://github.com/theislab/cellink" urls.Source = "https://github.com/theislab/cellink" +scripts.cellink-pgen = "cellink.cli.pgen:main" + [tool.hatch.build.targets.wheel] packages = [ "src/cellink" ] @@ -203,6 +205,3 @@ skip = [ "docs/references.md", "docs/notebooks/example.ipynb", ] - -[project.scripts] -cellink-pgen = "cellink.cli.pgen:main" diff --git a/src/cellink/_core/donordata.py b/src/cellink/_core/donordata.py index 8f35d15..0654856 100644 --- a/src/cellink/_core/donordata.py +++ b/src/cellink/_core/donordata.py @@ -1,10 +1,12 @@ from __future__ import annotations import logging +import warnings from collections.abc import Callable from dataclasses import dataclass import h5py +import numpy as np import pandas as pd import scanpy as sc import zarr @@ -22,6 +24,33 @@ logger = logging.getLogger(__name__) + +def _has_lazy_X(adata) -> bool: + """Check if an AnnData has a lazy-backed X matrix (dask or on-disk sparse). + + Safe for MuData / non-AnnData inputs (returns False). + """ + lazy_types = [] + try: + import dask.array as da + + lazy_types.append(da.Array) + except ImportError: + pass + try: + from anndata.abc import CSCDataset, CSRDataset + + lazy_types.extend([CSRDataset, CSCDataset]) + except ImportError: + try: + from anndata._core.sparse_dataset import BaseCompressedSparseDataset + + lazy_types.append(BaseCompressedSparseDataset) + except ImportError: + pass + return bool(lazy_types) and isinstance(getattr(adata, "X", None), tuple(lazy_types)) + + HIGHLIGHT_COLOR = "bold deep_pink2" @@ -54,19 +83,38 @@ def __init__( ): if donor_id not in C.obs.columns: raise ValueError(f"'{donor_id}' not found in C.obs") - if donor_id not in G.obs.columns and donor_id != G.obs.index.name: - raise ValueError(f"'{donor_id}' must be in gdata.obs or set as index") - if donor_id != G.obs.index.name: - G.obs = G.obs.set_index(donor_id) + + if isinstance(G.obs, pd.DataFrame): + if donor_id not in G.obs.columns and donor_id != G.obs.index.name: + raise ValueError(f"'{donor_id}' must be in gdata.obs or set as index") + if donor_id != G.obs.index.name: + G.obs = G.obs.set_index(donor_id) + else: + # Lazy AnnData (e.g. from read_lazy) — obs is xarray-backed. + # donor_id must already be the obs index (obs_names). + logger.debug("G.obs is not a pandas DataFrame; assuming donor_id is the obs index.") self._var_dims_to_sync = [] if var_dims_to_sync is None else var_dims_to_sync self.donor_id = donor_id + # Deferred obs filter for lazy G: applying fancy row-indexing on a lazy array + # before var-slicing forces materialising a large intermediate. We store the + # desired donor subset here and apply it in-memory inside to_memory(), AFTER + # the cheap var-slice has already narrowed the columns. + self._lazy_G_obs_filter: pd.Index | None = None + # Same deferral idea for lazy C — but stores the *donors to keep*; the + # cell-level boolean mask & sort are computed at to_memory() time, after + # the var-slice has narrowed C. + self._lazy_C_obs_filter: pd.Index | None = None self._match_donors(G, C) self.uns = uns def _match_donors(self, G: AnnData | MuData, C: AnnData | MuData) -> None: - G_idx = G.obs.index - C_idx = pd.Index(C.obs[self.donor_id].unique()) + G_idx = pd.Index(G.obs_names) + # C.obs[donor_id] may be a pandas.Series (eager) or an xarray.DataArray + # (when C comes from anndata.experimental.read_lazy). Coerce to numpy. + donor_col = C.obs[self.donor_id] + donor_arr = np.asarray(donor_col.values if hasattr(donor_col, "values") else donor_col) + C_idx = pd.Index(np.unique(donor_arr)) keep_donors = G_idx.intersection(C_idx) if len(keep_donors) == 0: @@ -77,29 +125,163 @@ def _match_donors(self, G: AnnData | MuData, C: AnnData | MuData) -> None: ), self.donor_id, ) - if not keep_donors.equals(G.obs_names): - G = G[keep_donors] - keep_cells = C.obs[self.donor_id].isin(keep_donors) - if not keep_cells.all(): - C = C[keep_cells] + # --- G obs filter --- + if not keep_donors.equals(G_idx): + if _has_lazy_X(G): + # Defer obs-filtering for lazy G. Applying a fancy row-index on a + # lazy array before a column-slice forces it to materialise the + # full row-filtered intermediate (all columns) before selecting the + # target columns — potentially reading gigabytes of data. + # We record keep_donors and apply the filter cheaply in to_memory() + # once only the needed columns have been loaded. + self._lazy_G_obs_filter = keep_donors + else: + G = G[keep_donors] - # Sort cells by donor order - sorted_cells = C.obs.iloc[ - pd.Categorical(C.obs[self.donor_id], categories=keep_donors, ordered=True).argsort() - ].index - if not sorted_cells.equals(C.obs_names): - C = C[sorted_cells] + # --- C cell filter & sort --- + if _has_lazy_X(C): + # Mirror the G deferral. Store keep_donors; the cell mask & donor-order + # sort are computed in to_memory() once C.X has been narrowed (var-slice). + self._lazy_C_obs_filter = keep_donors + else: + keep_cells = C.obs[self.donor_id].isin(keep_donors) + if not keep_cells.all(): + C = C[keep_cells] + # Sort cells by donor order (only C — G order is kept as-is). + # Stable sort so within-donor cell order is preserved deterministically. + sorted_cells = C.obs.iloc[ + pd.Categorical(C.obs[self.donor_id], categories=keep_donors, ordered=True).argsort(kind="stable") + ].index + if not sorted_cells.equals(C.obs_names): + C = C[sorted_cells] self._C = C self._G = G + @classmethod + def _from_parts( + cls, + *, + G: AnnData | MuData, + C: AnnData | MuData, + donor_id: str, + var_dims_to_sync: list[str], + uns: dict, + lazy_G_obs_filter: pd.Index | None = None, + lazy_C_obs_filter: pd.Index | None = None, + ) -> DonorData: + """Construct a DonorData from already-matched G and C, skipping _match_donors.""" + obj = object.__new__(cls) + obj._G = G + obj._C = C + obj.donor_id = donor_id + obj._var_dims_to_sync = var_dims_to_sync + obj.uns = uns + obj._lazy_G_obs_filter = lazy_G_obs_filter + obj._lazy_C_obs_filter = lazy_C_obs_filter + return obj + + @property + def _lazy_G(self) -> bool: + """Whether G has a lazy-backed X (lazy loading).""" + return _has_lazy_X(self._G) + + @property + def _lazy_C(self) -> bool: + """Whether C has a lazy-backed X (lazy loading).""" + return _has_lazy_X(self._C) + + @property + def is_lazy(self) -> bool: + """Whether G or C has a lazy-backed X matrix (lazy loading).""" + return self._lazy_G or self._lazy_C + + def to_memory(self, *, inplace: bool = False) -> DonorData: + """Return a DonorData with both G and C materialised into memory. + + If neither side is lazy, returns ``self`` unchanged. For lazy G or C, + calls ``.to_memory()`` on the respective side. + + When obs filters were deferred (to avoid expensive lazy fancy row-indexing + before column-slicing), they are applied here in-memory after the var-slice + has already narrowed the columns. + + Parameters + ---------- + inplace: + If True, materialise G and C into ``self`` and return ``self``. + If False (default), return a new DonorData leaving ``self`` unchanged. + """ + if not self.is_lazy: + return self + + # --- G side --- + G = self._G + if self._lazy_G: + G = G.to_memory() if hasattr(G, "to_memory") else G.copy() + # read_lazy → to_memory() can lose the obs index name; restore it + if G.obs.index.name is None and self.donor_id is not None: + G.obs.index.name = self.donor_id + if self._lazy_G_obs_filter is not None: + G = G[self._lazy_G_obs_filter] + + # --- C side --- + C = self._C + if self._lazy_C: + C = C.to_memory() if hasattr(C, "to_memory") else C.copy() + if self._lazy_C_obs_filter is not None: + keep_donors = self._lazy_C_obs_filter + # Cell-level filter & donor-order sort, mirroring the eager path + # in _match_donors(). Now that obs is a pandas DataFrame, .isin() + # and Categorical-argsort behave normally. + if self.donor_id in C.obs.columns: + keep_cells = C.obs[self.donor_id].isin(keep_donors) + if not keep_cells.all(): + C = C[keep_cells] + sorted_cells = C.obs.iloc[ + pd.Categorical(C.obs[self.donor_id], categories=keep_donors, ordered=True).argsort( + kind="stable" + ) + ].index + if not sorted_cells.equals(C.obs_names): + C = C[sorted_cells] + if C.is_view: + C = C.copy() + + if inplace: + self._G = G + self._C = C + self._lazy_G_obs_filter = None + self._lazy_C_obs_filter = None + return self + + return DonorData._from_parts( + G=G, + C=C, + donor_id=self.donor_id, + var_dims_to_sync=self._var_dims_to_sync, + uns=self.uns, + ) + + def compute(self, *, inplace: bool = False) -> DonorData: + """Alias for :meth:`to_memory`, following dask naming conventions.""" + return self.to_memory(inplace=inplace) + def copy(self) -> DonorData: - if self._G.is_view: - self._G = self._G.copy() - if self._C.is_view: - self._C = self._C.copy() - return self + # For lazy sides, don't materialise — share the reference. + # For eager sides, copy if it's currently a view to detach it. + G = self._G.copy() if not self._lazy_G and self._G.is_view else self._G + C = self._C.copy() if not self._lazy_C and self._C.is_view else self._C + return DonorData._from_parts( + G=G, + C=C, + donor_id=self.donor_id, + var_dims_to_sync=list(self._var_dims_to_sync), + uns=self.uns, + lazy_G_obs_filter=self._lazy_G_obs_filter, + lazy_C_obs_filter=self._lazy_C_obs_filter, + ) def _write_dd(self, f: h5py.File): if isinstance(self.G, MuData): @@ -196,6 +378,7 @@ def C(self) -> AnnData: def C(self, value: AnnData | MuData) -> None: if not isinstance(value, AnnData | MuData): raise ValueError("C must be an AnnData or MuData object") + self._lazy_C_obs_filter = None # reset — new C takes over self._C = value self._match_donors(self._G, self._C) @@ -207,6 +390,7 @@ def G(self) -> AnnData: def G(self, value: AnnData | MuData) -> None: if not isinstance(value, AnnData | MuData): raise ValueError("G must be an AnnData or MuData object") + self._lazy_G_obs_filter = None # reset — new G takes over self._G = value self._match_donors(self._G, self._C) @@ -231,9 +415,9 @@ def sel( C_obs: slice = slice(None), C_var: slice = slice(None), ): - _G = self.G[G_obs] + _G = self._G[G_obs] _G = _G[:, G_var] - _C = self.C[C_obs] + _C = self._C[C_obs] _C = _C[:, C_var] _G = self._sync_var_dims(_G, _C) @@ -251,10 +435,11 @@ def __getitem__(self, key): (idx,) if isinstance(idx, str) else idx for idx in key ) # needed because Mudata[str] looks up modalities key = key + (slice(None),) * (4 - len(key)) - # Only slice if key is not slice(None) - _G = self.G[key[0]] if key[0] is not slice(None) else self.G + # Slice self._G directly (not self.G) so the deferred obs-filter view is never + # embedded into a new lazy dask chain — that would recreate the bottleneck. + _G = self._G[key[0]] if key[0] is not slice(None) else self._G _G = _G[:, key[1]] if key[1] is not slice(None) else _G - _C = self.C[key[2]] if key[2] is not slice(None) else self.C + _C = self._C[key[2]] if key[2] is not slice(None) else self._C _C = _C[:, key[3]] if key[3] is not slice(None) else _C _G = self._sync_var_dims(_G, _C) @@ -319,6 +504,12 @@ def aggregate( verbose: Whether to print verbose output. Defaults to False. """ + if self._lazy_C: + raise RuntimeError( + "Cannot aggregate on a lazy C. " + "Call dd.to_memory() first or subset with dd[..., :, var_subset].to_memory()." + ) + if filter_key is not None: assert filter_value is not None, "filter_value must be provided if filter_key is provided" @@ -374,9 +565,10 @@ def aggregate( def prep_repr(self) -> str: """String representation of DonorData showing side-by-side dd.G and dd.C views.""" + n_G_obs, n_G_vars, n_C_obs, n_C_vars = self.shape # Split the representations into lines - G_lines, G_highlight = _anndata_repr(self.G, self.G.n_obs, self.G.n_vars, self._var_dims_to_sync) - C_lines, C_highlight = _anndata_repr(self.C, self.C.n_obs, self.C.n_vars, [self.donor_id]) + G_lines, G_highlight = _anndata_repr(self.G, n_G_obs, n_G_vars, self._var_dims_to_sync) + C_lines, C_highlight = _anndata_repr(self.C, n_C_obs, n_C_vars, [self.donor_id]) def pad_lists(l1, l2): max_lines = max(len(l1), len(l2)) @@ -412,7 +604,7 @@ def highlight_lines(lines, highlights): n_cells_per_donor = self.C.obs[self.donor_id].value_counts() min_n_cells, max_n_cells = n_cells_per_donor.min(), n_cells_per_donor.max() header_line = Text( - f"DonorData(n_donors={self.G.shape[0]:,}, " + f"DonorData(n_donors={n_G_obs:,}, " f"n_cells_per_donor=[{min_n_cells:,}-{max_n_cells:,}], " f"donor_id='{self.donor_id}')", style=HIGHLIGHT_COLOR, @@ -441,7 +633,21 @@ def __str__(self) -> str: @property def shape(self) -> tuple[tuple[int, int], tuple[int, int]]: - return *self.G.shape, *self.C.shape + # When obs filters are deferred for lazy G/C, self._G/self._C still hold all + # rows; report the filtered counts so dd.shape stays user-meaningful. + n_G_obs = len(self._lazy_G_obs_filter) if self._lazy_G_obs_filter is not None else self._G.shape[0] + # For lazy C, _lazy_C_obs_filter stores donors (not cells), so we can't + # report exact cell counts without reading C.obs[donor_id]; fall back to + # the unfiltered count and warn. + if self._lazy_C_obs_filter is not None: + warnings.warn( + "dd.shape reports the unfiltered cell count for lazy C because the donor " + "filter is still pending. Call dd.to_memory() to get the exact count.", + UserWarning, + stacklevel=2, + ) + n_C_obs = self._C.shape[0] + return n_G_obs, self._G.shape[1], n_C_obs, self._C.shape[1] def _anndata_repr(adata, n_obs, n_vars, highlight_keys=None) -> str: diff --git a/src/cellink/io/__init__.py b/src/cellink/io/__init__.py index 4502c3d..31aa701 100644 --- a/src/cellink/io/__init__.py +++ b/src/cellink/io/__init__.py @@ -1,4 +1,4 @@ from ._export import to_plink, write_variants_to_vcf -from ._readwrite import read_dd, read_h5_dd, read_zarr_dd +from ._readwrite import read_dd, read_h5_dd, read_lazy_dd, read_zarr_dd from ._sgkit import from_sgkit_dataset, read_bgen, read_plink, read_sgkit_zarr from ._pgen import stream_pgen_to_zarr, read_pgen_zarr diff --git a/src/cellink/io/_readwrite.py b/src/cellink/io/_readwrite.py index a489459..3241602 100644 --- a/src/cellink/io/_readwrite.py +++ b/src/cellink/io/_readwrite.py @@ -1,3 +1,6 @@ +from __future__ import annotations + +import logging import warnings import h5py @@ -13,6 +16,8 @@ from .._core import DonorData +logger = logging.getLogger(__name__) + warnings.filterwarnings( "ignore", message="The return type of `Dataset.dims` will be changed", @@ -46,7 +51,6 @@ def _read_mudata(group: StorageType, backed: bool = True) -> MuData: - Adapted from `mudata._core.io.read_h5mu`. - Preserves modality ordering if the `mod-order` attribute is present in the group. """ - d = {} for k in group.keys(): if k in ["obs", "var"]: @@ -100,7 +104,6 @@ def _read_dd(f: h5py.File) -> DonorData: ValueError If the encoding type of `G` or `C` is not recognized. """ - if f["G"].attrs.get("encoding-type") == "MuData": G = _read_mudata(group=f["G"]) elif f["G"].attrs.get("encoding-type") == "anndata": @@ -135,7 +138,7 @@ def read_h5_dd(path: str) -> DonorData: Parameters ---------- - path : str + path : str | Path Path to the HDF5 file containing serialized DonorData. Returns @@ -143,8 +146,7 @@ def read_h5_dd(path: str) -> DonorData: DonorData A DonorData object with genotype (`G`), cell expression (`C`), and metadata. """ - - with h5py.File(path, "r") as f: + with h5py.File(str(path), "r") as f: return _read_dd(f) @@ -154,7 +156,7 @@ def read_zarr_dd(path: str) -> DonorData: Parameters ---------- - path : str + path : str | Path Path to the Zarr store containing serialized DonorData. Returns @@ -162,8 +164,7 @@ def read_zarr_dd(path: str) -> DonorData: DonorData A DonorData object with genotype (`G`), cell expression (`C`), and metadata. """ - - f = zarr.open(path, mode="r") + f = zarr.open(str(path), mode="r") return _read_dd(f) @@ -205,7 +206,8 @@ def read_dd(path: str, fmt: str = None) -> DonorData: >>> dd = read_dd("donor_data.dd.zarr") >>> dd = read_dd("custom_data.h5", fmt="h5") """ - + path = str(path) + if fmt is None: if path.endswith(".h5") or path.endswith(".dd.h5"): fmt = "h5" @@ -220,3 +222,299 @@ def read_dd(path: str, fmt: str = None) -> DonorData: return read_zarr_dd(path) else: raise ValueError("Unknown format: use 'h5' or 'zarr'.") + + +def _eagerize_obs_var(adata): + """Materialise obs/var (and obsm/varm/uns) into in-memory pandas/numpy. + + Keeps ``X`` and ``layers`` as-is — for an AnnData returned by + :func:`anndata.experimental.read_lazy`, that means dask-backed. + + The result is an AnnData whose ``.obs`` / ``.var`` are pandas DataFrames + (so all the usual ``.loc``, ``.isin``, ``.merge`` operations work), while + ``.X`` reads chunks on demand. + """ + import anndata as ad + import numpy as np + import pandas as pd + + def _to_dataframe(x): + # xarray-backed (read_lazy): exposes .ds attribute on the obs/var proxy + if hasattr(x, "ds"): + return x.ds.load().to_dataframe() + if isinstance(x, pd.DataFrame): + return x.copy() + if hasattr(x, "to_dataframe"): + return x.to_dataframe() + return pd.DataFrame(x) + + def _materialise(x): + if isinstance(x, (pd.DataFrame, pd.Series, np.ndarray)): + return x + if hasattr(x, "to_pandas"): + return x.to_pandas() + if hasattr(x, "compute"): + return np.asarray(x.compute()) + if hasattr(x, "ds"): + ds = x.ds.load() + return ds.to_dataframe() if hasattr(ds, "to_dataframe") else ds + return x + + obs_df = _to_dataframe(adata.obs) + var_df = _to_dataframe(adata.var) + + obsm = {k: _materialise(adata.obsm[k]) for k in adata.obsm} if adata.obsm else None + varm = {k: _materialise(adata.varm[k]) for k in adata.varm} if adata.varm else None + uns = {k: _materialise(adata.uns[k]) for k in (adata.uns or {})} or None + layers = dict(adata.layers) if adata.layers else None # keep lazy + + return ad.AnnData( + X=adata.X, + obs=obs_df, + var=var_df, + obsm=obsm, + varm=varm, + layers=layers, + uns=uns, + ) + + +def _load_lazy_anndata_zarr(path: str, *, G_reader: str = "auto", load_annotation_index: bool = True): + """Load an AnnData zarr lazily (dask-backed X). Used for both G and C subgroups.""" + a = None + if G_reader in ("auto", "read_lazy"): + try: + from anndata.experimental import read_lazy + + a = read_lazy(path, load_annotation_index=load_annotation_index) + logger.info("Loaded %s lazily via anndata.experimental.read_lazy", path) + except Exception as e: + if G_reader == "read_lazy": + raise + logger.debug("read_lazy failed (%s), trying read_pgen_zarr", e) + + if a is None and G_reader in ("auto", "read_pgen_zarr"): + from cellink.io._pgen import read_pgen_zarr + + a = read_pgen_zarr(path) + logger.info("Loaded %s lazily via read_pgen_zarr", path) + + if a is None: + raise ValueError(f"Could not load {path} lazily with reader={G_reader!r}") + return a + + +def _read_lazy_dd_zarr( + path: str, + *, + lazy_G: bool, + lazy_C: bool, + eager_obs_var: bool, + G_reader: str, + load_annotation_index: bool, +) -> DonorData: + import anndata as ad + + f = zarr.open(path, mode="r") + donor_id = f.attrs.get("donor_id", "donor_id") + var_dims_to_sync = list(f.attrs.get("var_dims_to_sync", [])) + + uns = {} + uns_group = f.get("uns") + if uns_group: + for key in uns_group: + try: + uns[key] = uns_group[key][()] + except (TypeError, KeyError, OSError) as e: + warnings.warn( + f"Skipped loading uns['{key}']: {e}", + UserWarning, + stacklevel=2, + ) + + # --- G --- + if lazy_G: + G = _load_lazy_anndata_zarr(f"{path}/G", G_reader=G_reader, load_annotation_index=load_annotation_index) + if eager_obs_var: + G = _eagerize_obs_var(G) + else: + G = ad.read_zarr(f"{path}/G") if "G" in f else read_elem(f["G"]) + + # --- C --- + # G_reader is intentionally NOT applied to C; C is always read via the + # standard zarr reader (read_lazy when lazy, ad.read_zarr otherwise). + if lazy_C: + C = _load_lazy_anndata_zarr(f"{path}/C", G_reader="read_lazy", load_annotation_index=load_annotation_index) + if eager_obs_var: + C = _eagerize_obs_var(C) + else: + C = ad.read_zarr(f"{path}/C") if "C" in f else read_elem(f["C"]) + + return DonorData( + G=G, + C=C, + donor_id=donor_id, + var_dims_to_sync=var_dims_to_sync, + uns=uns, + ) + + +def _read_lazy_dd_h5( + path: str, + *, + lazy_G: bool, + lazy_C: bool, + eager_obs_var: bool, + load_annotation_index: bool, +) -> DonorData: + f = h5py.File(path, "r") # NOT a context manager — handle must outlive dd + + donor_id = f.attrs.get("donor_id", "donor_id") + var_dims_to_sync = list(f.attrs.get("var_dims_to_sync", [])) + + uns = {} + uns_group = f.get("uns") + if uns_group is not None: + for key in uns_group: + try: + uns[key] = uns_group[key][()] + except (TypeError, KeyError, OSError) as e: + warnings.warn( + f"Skipped loading uns['{key}']: {e}", + UserWarning, + stacklevel=2, + ) + + def _open_lazy(group): + from anndata.experimental import read_lazy + + return read_lazy(group, load_annotation_index=load_annotation_index) + + # --- G --- + if lazy_G: + G = _open_lazy(f["G"]) + if eager_obs_var: + G = _eagerize_obs_var(G) + else: + G = read_elem(f["G"]) + + # --- C --- + if lazy_C: + C = _open_lazy(f["C"]) + if eager_obs_var: + C = _eagerize_obs_var(C) + else: + C = read_elem(f["C"]) + + dd = DonorData( + G=G, + C=C, + donor_id=donor_id, + var_dims_to_sync=var_dims_to_sync, + uns=uns, + ) + # Pin the h5py file handle so it isn't gc'd while dask still references its datasets. + # Always set — at least one side is lazy here (the both-eager case is handled + # upstream by delegating to read_dd, so the file is fully read and closed). + dd._h5_handle = f + return dd + + +def read_lazy_dd( + path: str, + *, + lazy_G: bool = True, + lazy_C: bool = False, + eager_obs_var: bool = False, + G_reader: str = "auto", + load_annotation_index: bool = True, +) -> DonorData: + """Read a DonorData store with per-side control over lazy vs eager loading. + + Supports both zarr (``.dd.zarr``) and HDF5 (``.dd.h5``) stores. + + The two boolean flags ``lazy_G`` and ``lazy_C`` are independent. Any side + flagged ``True`` is loaded with a dask-backed ``X`` (and, by default, + xarray-backed ``obs`` / ``var``). Any side flagged ``False`` is loaded + fully into memory using the standard non-lazy reader. + + When **both** flags are ``False`` this function delegates to + :func:`read_dd`, i.e. it behaves identically to a regular non-lazy read. + + For HDF5 with at least one lazy side, the underlying ``h5py.File`` handle + is kept alive on the returned DonorData (``dd._h5_handle``) — closing the + file or dropping the DonorData invalidates lazy reads. + + Parameters + ---------- + path : str + Path to a ``.dd.zarr`` or ``.dd.h5`` DonorData store. + lazy_G : bool, default True + If True, load G with a dask-backed ``X``. If False, load G eagerly. + lazy_C : bool, default False + If True, load C with a dask-backed ``X``. If False, load C eagerly. + (Cells × genes is typically small enough that eager is fine.) + eager_obs_var : bool, default False + Only meaningful for sides that are lazy. If True, materialise + ``obs`` / ``var`` (and ``obsm`` / ``varm`` / ``uns``) as in-memory + pandas / numpy while keeping ``X`` (and ``layers``) lazy. Useful when + you want fast pandas-based filtering on obs/var without paying for X + materialisation. No effect on sides that are already eager. + G_reader : {'auto', 'read_lazy', 'read_pgen_zarr'}, default 'auto' + How to open **G** lazily (only used when ``lazy_G=True`` and the store + is zarr; ignored otherwise — ``C`` is always read via the standard + zarr reader). + + - ``'read_lazy'``: :func:`anndata.experimental.read_lazy` + (requires zarr-backed AnnData written by anndata >=0.11). + - ``'read_pgen_zarr'``: :func:`cellink.io.read_pgen_zarr` + (for stores written by :func:`stream_pgen_to_zarr`). + - ``'auto'``: tries ``read_lazy`` first, falls back to ``read_pgen_zarr``. + load_annotation_index : bool, default True + Passed to ``anndata.experimental.read_lazy`` for faster var indexing. + + Returns + ------- + DonorData + A DonorData object. ``dd.is_lazy`` is True if either side is lazy. + Call ``dd.to_memory()`` or ``dd[..., :, var_slice].to_memory()`` to + materialise. + + Examples + -------- + >>> # Default: lazy G, eager C + >>> dd = read_lazy_dd("donor_data.dd.zarr") + + >>> # Both lazy + pandas obs/var on both sides + >>> dd = read_lazy_dd("donor_data.dd.zarr", lazy_C=True, eager_obs_var=True) + + >>> # Eager G, lazy C (uncommon but supported) + >>> dd = read_lazy_dd("donor_data.dd.zarr", lazy_G=False, lazy_C=True) + + >>> # Both eager — equivalent to read_dd(path) + >>> dd = read_lazy_dd("donor_data.dd.zarr", lazy_G=False, lazy_C=False) + """ + path = str(path) + + # Both eager: delegate to the regular non-lazy reader (no dask, no file pinning). + if not lazy_G and not lazy_C: + return read_dd(path) + + if path.endswith((".zarr", ".dd.zarr")): + return _read_lazy_dd_zarr( + path, + lazy_G=lazy_G, + lazy_C=lazy_C, + eager_obs_var=eager_obs_var, + G_reader=G_reader, + load_annotation_index=load_annotation_index, + ) + if path.endswith((".h5", ".dd.h5")): + return _read_lazy_dd_h5( + path, + lazy_G=lazy_G, + lazy_C=lazy_C, + eager_obs_var=eager_obs_var, + load_annotation_index=load_annotation_index, + ) + raise ValueError(f"Cannot infer format from {path!r}; expected .dd.zarr / .zarr or .dd.h5 / .h5") diff --git a/src/cellink/pl/__init__.py b/src/cellink/pl/__init__.py index 2a2182f..ac76f95 100644 --- a/src/cellink/pl/__init__.py +++ b/src/cellink/pl/__init__.py @@ -1,2 +1,2 @@ from ._at import locus, manhattan, qq -from ._expression import expression_by_genotype, volcano +from ._expression import expression_by_genotype, scatter, volcano diff --git a/tests/test_lazy_donordata.py b/tests/test_lazy_donordata.py new file mode 100644 index 0000000..5e2761a --- /dev/null +++ b/tests/test_lazy_donordata.py @@ -0,0 +1,497 @@ +"""Tests for lazy DonorData support (dask-backed G).""" + +import dask.array as da +import numpy as np +import pandas as pd +import pytest +from anndata import AnnData + +from cellink._core.data_fields import DAnn +from cellink._core.donordata import DonorData, _has_dask_X +from cellink.io import read_lazy_dd, read_zarr_dd + +# ── helpers ────────────────────────────────────────────────────────────────── + + +def _make_lazy_gdata(gdata: AnnData) -> AnnData: + """Convert an in-memory AnnData to one with a dask-backed X.""" + lazy_X = da.from_array(np.asarray(gdata.X), chunks=(gdata.n_obs, 2)) + return AnnData(X=lazy_X, obs=gdata.obs.copy(), var=gdata.var.copy()) + + +def _make_lazy_adata(adata: AnnData) -> AnnData: + """Convert an in-memory single-cell AnnData to one with a dask-backed X. + + obs is preserved as a pandas DataFrame (so .isin / Categorical-argsort work); + only X becomes a dask.Array. _has_dask_X(adata) returns True. + """ + lazy_X = da.from_array(np.asarray(adata.X), chunks=(adata.n_obs, 2)) + return AnnData(X=lazy_X, obs=adata.obs.copy(), var=adata.var.copy()) + + +# ── _has_dask_X ───────────────────────────────────────────────────────────── + + +def test_has_dask_X_false(gdata): + assert _has_dask_X(gdata) is False + + +def test_has_dask_X_true(gdata): + lazy = _make_lazy_gdata(gdata) + assert _has_dask_X(lazy) is True + + +# ── DonorData with lazy G ──────────────────────────────────────────────────── + + +def test_lazy_init(adata, gdata): + lazy_g = _make_lazy_gdata(gdata) + dd = DonorData(G=lazy_g, C=adata) + assert dd.is_lazy is True + + +def test_eager_init(adata, gdata): + dd = DonorData(G=gdata, C=adata) + assert dd.is_lazy is False + + +def test_lazy_shape(adata, gdata): + lazy_g = _make_lazy_gdata(gdata) + dd = DonorData(G=lazy_g, C=adata) + assert dd.shape == (*lazy_g.shape, *adata.shape) + + +def test_lazy_donor_matching(adata, gdata): + """Lazy G with extra donors: filter is deferred (avoids fancy row-indexing on + dask before var-slicing) and applied at to_memory().""" + extra_obs = pd.DataFrame(index=pd.Index(["EXTRA_DONOR"], name=DAnn.donor)) + extra_X = da.zeros((1, gdata.n_vars), chunks=(1, gdata.n_vars)) + extra = AnnData( + X=da.concatenate([da.from_array(np.asarray(gdata.X)), extra_X], axis=0), + obs=pd.concat([gdata.obs, extra_obs]), + var=gdata.var.copy(), + ) + dd = DonorData(G=extra, C=adata) + # Filter deferred — _G still holds EXTRA_DONOR + assert "EXTRA_DONOR" in dd.G.obs_names + assert dd._lazy_G_obs_filter is not None and "EXTRA_DONOR" not in dd._lazy_G_obs_filter + # After to_memory(), the deferred filter is applied + dd_mem = dd.to_memory() + assert "EXTRA_DONOR" not in dd_mem.G.obs_names + + +def test_lazy_slice_variants(adata, gdata): + lazy_g = _make_lazy_gdata(gdata) + dd = DonorData(G=lazy_g, C=adata) + keep = ["SNP0", "SNP1"] + dd_sub = dd[:, keep] + assert dd_sub.shape[1] == len(keep) + # Slicing lazy should remain lazy + assert dd_sub.is_lazy is True + + +# ── to_memory ──────────────────────────────────────────────────────────────── + + +def test_to_memory(adata, gdata): + lazy_g = _make_lazy_gdata(gdata) + dd = DonorData(G=lazy_g, C=adata) + dd_mem = dd.to_memory() + assert dd_mem.is_lazy is False + assert isinstance(dd_mem.G.X, np.ndarray) + np.testing.assert_array_equal(dd_mem.G.X, np.asarray(gdata.X)) + + +def test_to_memory_idempotent(adata, gdata): + dd = DonorData(G=gdata, C=adata) + dd_mem = dd.to_memory() + assert dd_mem is dd # no-op for in-memory + + +# ── aggregate guard ────────────────────────────────────────────────────────── + + +def test_aggregate_works_with_lazy_G_only(adata, gdata): + """With per-side laziness, aggregate operates on C — lazy G alone should not block it.""" + lazy_g = _make_lazy_gdata(gdata) + dd = DonorData(G=lazy_g, C=adata) + dd.aggregate(key_added="Gex") + assert "Gex" in dd.G.obsm + + +def test_aggregate_raises_on_lazy_C(adata, gdata): + """Aggregation reads C — lazy C must still block until to_memory().""" + lazy_c = _make_lazy_adata(adata) + dd = DonorData(G=gdata, C=lazy_c) + with pytest.raises(RuntimeError, match="lazy"): + dd.aggregate(key_added="Gex") + + +def test_aggregate_works_after_to_memory(adata, gdata): + lazy_g = _make_lazy_gdata(gdata) + dd = DonorData(G=lazy_g, C=adata) + dd_mem = dd.to_memory() + dd_mem.aggregate(key_added="Gex") + assert "Gex" in dd_mem.G.obsm + + +# ── copy ───────────────────────────────────────────────────────────────────── + + +def test_copy_lazy(adata, gdata): + lazy_g = _make_lazy_gdata(gdata) + dd = DonorData(G=lazy_g, C=adata) + dd_copy = dd.copy() + # copy should not materialise + assert dd_copy.is_lazy is True + assert isinstance(dd_copy.G.X, da.Array) + + +# ── read_lazy_dd roundtrip ─────────────────────────────────────────────────── + + +@pytest.mark.slow +def test_read_lazy_dd_roundtrip(tmp_path, adata, gdata): + """Write a DonorData as zarr, read it back lazily, verify shapes.""" + out = tmp_path / "test.dd.zarr" + dd = DonorData(G=gdata, C=adata) + dd.write_zarr_dd(str(out)) + + dd_lazy = read_lazy_dd(str(out)) + assert dd_lazy.is_lazy is True + assert dd_lazy.G.shape == dd.G.shape + assert dd_lazy.C.shape == dd.C.shape + + # Materialise and check values + dd_mem = dd_lazy.to_memory() + assert dd_mem.is_lazy is False + np.testing.assert_array_almost_equal(dd_mem.G.X, np.asarray(dd.G.X)) + + +@pytest.mark.slow +def test_read_lazy_dd_matches_eager(tmp_path, adata, gdata): + """Lazy and eager reads of the same zarr should produce identical data.""" + out = tmp_path / "test.dd.zarr" + dd = DonorData(G=gdata, C=adata) + dd.write_zarr_dd(str(out)) + + dd_eager = read_zarr_dd(str(out)) + dd_lazy = read_lazy_dd(str(out)) + dd_lazy_mem = dd_lazy.to_memory() + + assert dd_lazy_mem.G.shape == dd_eager.G.shape + assert dd_lazy_mem.C.shape == dd_eager.C.shape + np.testing.assert_array_almost_equal(dd_lazy_mem.G.X, np.asarray(dd_eager.G.X)) + + +# ── per-side laziness flags ────────────────────────────────────────────────── + + +def test_lazy_G_only(adata, gdata): + lazy_g = _make_lazy_gdata(gdata) + dd = DonorData(G=lazy_g, C=adata) + assert dd._lazy_G is True + assert dd._lazy_C is False + assert dd.is_lazy is True + + +def test_lazy_C_only(adata, gdata): + lazy_c = _make_lazy_adata(adata) + dd = DonorData(G=gdata, C=lazy_c) + assert dd._lazy_G is False + assert dd._lazy_C is True + assert dd.is_lazy is True + + +def test_lazy_both(adata, gdata): + lazy_g = _make_lazy_gdata(gdata) + lazy_c = _make_lazy_adata(adata) + dd = DonorData(G=lazy_g, C=lazy_c) + assert dd._lazy_G is True + assert dd._lazy_C is True + assert dd.is_lazy is True + + +def test_eager_both(adata, gdata): + dd = DonorData(G=gdata, C=adata) + assert dd._lazy_G is False + assert dd._lazy_C is False + assert dd.is_lazy is False + + +def test_lazy_flags_track_setter(adata, gdata): + """_lazy_G / _lazy_C are computed from current G/C, not cached at init.""" + lazy_g = _make_lazy_gdata(gdata) + dd = DonorData(G=lazy_g, C=adata) + assert dd._lazy_G is True + # Replace G with eager — flag must flip + dd.G = gdata + assert dd._lazy_G is False + assert dd.is_lazy is False + + +# ── to_memory: lazy C side ─────────────────────────────────────────────────── + + +def test_to_memory_lazy_C(adata, gdata): + # Eager reference: cells get sorted by donor order during _match_donors, + # so we compare against an eager DonorData's C.X (same sort order). + eager_dd = DonorData(G=gdata, C=adata.copy()) + expected = np.asarray(eager_dd.C.X) + + lazy_c = _make_lazy_adata(adata) + dd = DonorData(G=gdata, C=lazy_c) + dd_mem = dd.to_memory() + assert dd_mem.is_lazy is False + assert isinstance(dd_mem.C.X, np.ndarray) + np.testing.assert_array_equal(dd_mem.C.X, expected) + + +def test_to_memory_lazy_both(adata, gdata): + eager_dd = DonorData(G=gdata, C=adata.copy()) + expected_C = np.asarray(eager_dd.C.X) + + lazy_g = _make_lazy_gdata(gdata) + lazy_c = _make_lazy_adata(adata) + dd = DonorData(G=lazy_g, C=lazy_c) + dd_mem = dd.to_memory() + assert dd_mem.is_lazy is False + assert isinstance(dd_mem.G.X, np.ndarray) + assert isinstance(dd_mem.C.X, np.ndarray) + np.testing.assert_array_equal(dd_mem.G.X, np.asarray(gdata.X)) + np.testing.assert_array_equal(dd_mem.C.X, expected_C) + + +def test_to_memory_lazy_C_preserves_donor_intersection(adata, gdata): + """When C has cells from a donor not in G, lazy-C donors filter is deferred + and applied at to_memory.""" + # Drop one donor from G to force a mismatch + g_subset = gdata[gdata.obs_names[:-1]].copy() + lazy_c = _make_lazy_adata(adata) + dd = DonorData(G=g_subset, C=lazy_c) + # Filter is deferred for lazy C + assert dd._lazy_C_obs_filter is not None + dd_mem = dd.to_memory() + # After materialisation, no cells from the dropped donor remain + dropped = set(adata.obs[DAnn.donor]) - set(g_subset.obs_names) + assert dropped, "test setup should produce a non-trivial drop" + assert not set(dd_mem.C.obs[DAnn.donor]).intersection(dropped) + + +# ── _eagerize_obs_var helper ───────────────────────────────────────────────── + + +def test_eagerize_obs_var_keeps_X_lazy(gdata): + from cellink.io._readwrite import _eagerize_obs_var + + lazy = _make_lazy_gdata(gdata) + eager_ov = _eagerize_obs_var(lazy) + assert isinstance(eager_ov.X, da.Array) # X stays dask + assert isinstance(eager_ov.obs, pd.DataFrame) + assert isinstance(eager_ov.var, pd.DataFrame) + + +def test_eagerize_obs_var_preserves_columns_and_index(gdata): + from cellink.io._readwrite import _eagerize_obs_var + + lazy = _make_lazy_gdata(gdata) + eager_ov = _eagerize_obs_var(lazy) + assert list(eager_ov.obs.columns) == list(gdata.obs.columns) + assert list(eager_ov.var.columns) == list(gdata.var.columns) + assert list(eager_ov.obs.index) == list(gdata.obs.index) + assert list(eager_ov.var.index) == list(gdata.var.index) + + +# ── read_lazy_dd: zarr + new flags ─────────────────────────────────────────── + + +@pytest.mark.slow +def test_read_lazy_dd_zarr_eager_obs_var(tmp_path, adata, gdata): + out = tmp_path / "test.dd.zarr" + DonorData(G=gdata, C=adata).write_zarr_dd(str(out)) + + dd = read_lazy_dd(str(out), eager_obs_var=True) + # G: X dask, obs/var pandas + assert isinstance(dd.G.X, da.Array) + assert isinstance(dd.G.obs, pd.DataFrame) + assert isinstance(dd.G.var, pd.DataFrame) + # C is still eager (default lazy_C=False) — obs/var also pandas + assert isinstance(dd.C.obs, pd.DataFrame) + + +@pytest.mark.slow +def test_read_lazy_dd_zarr_lazy_C(tmp_path, adata, gdata): + out = tmp_path / "test.dd.zarr" + DonorData(G=gdata, C=adata).write_zarr_dd(str(out)) + + dd = read_lazy_dd(str(out), lazy_C=True) + assert dd._lazy_G is True + assert dd._lazy_C is True + assert isinstance(dd.G.X, da.Array) + assert isinstance(dd.C.X, da.Array) + + +@pytest.mark.slow +def test_read_lazy_dd_zarr_lazy_C_eager_obs_var(tmp_path, adata, gdata): + out = tmp_path / "test.dd.zarr" + eager_dd = DonorData(G=gdata, C=adata) + eager_dd.write_zarr_dd(str(out)) + expected_C = np.asarray(eager_dd.C.X) + + dd = read_lazy_dd(str(out), lazy_C=True, eager_obs_var=True) + # X dask on both sides + assert isinstance(dd.G.X, da.Array) + assert isinstance(dd.C.X, da.Array) + # obs/var pandas on both sides + assert isinstance(dd.G.obs, pd.DataFrame) + assert isinstance(dd.C.obs, pd.DataFrame) + # Round-trip values still correct after to_memory + dd_mem = dd.to_memory() + np.testing.assert_array_almost_equal(dd_mem.G.X, np.asarray(gdata.X)) + np.testing.assert_array_almost_equal(dd_mem.C.X, expected_C) + + +# ── read_lazy_dd: HDF5 ─────────────────────────────────────────────────────── + + +@pytest.mark.slow +def test_read_lazy_dd_h5_basic(tmp_path, adata, gdata): + out = tmp_path / "test.dd.h5" + DonorData(G=gdata, C=adata).write_h5_dd(str(out)) + + dd = read_lazy_dd(str(out)) + assert dd._lazy_G is True + assert dd._lazy_C is False # default + assert isinstance(dd.G.X, da.Array) + + dd_mem = dd.to_memory() + np.testing.assert_array_almost_equal(dd_mem.G.X, np.asarray(gdata.X)) + + +@pytest.mark.slow +def test_read_lazy_dd_h5_eager_obs_var(tmp_path, adata, gdata): + out = tmp_path / "test.dd.h5" + DonorData(G=gdata, C=adata).write_h5_dd(str(out)) + + dd = read_lazy_dd(str(out), eager_obs_var=True) + assert isinstance(dd.G.X, da.Array) + assert isinstance(dd.G.obs, pd.DataFrame) + assert isinstance(dd.G.var, pd.DataFrame) + + +@pytest.mark.slow +def test_read_lazy_dd_h5_lazy_C(tmp_path, adata, gdata): + out = tmp_path / "test.dd.h5" + eager_dd = DonorData(G=gdata, C=adata) + eager_dd.write_h5_dd(str(out)) + expected_C = np.asarray(eager_dd.C.X) + + dd = read_lazy_dd(str(out), lazy_C=True, eager_obs_var=True) + assert dd._lazy_G is True + assert dd._lazy_C is True + assert isinstance(dd.G.X, da.Array) + assert isinstance(dd.C.X, da.Array) + + dd_mem = dd.to_memory() + np.testing.assert_array_almost_equal(dd_mem.G.X, np.asarray(gdata.X)) + np.testing.assert_array_almost_equal(dd_mem.C.X, expected_C) + + +@pytest.mark.slow +def test_read_lazy_dd_h5_handle_pinned(tmp_path, adata, gdata): + """The h5py file handle must be pinned on dd._h5_handle so dask reads stay valid.""" + import h5py + + out = tmp_path / "test.dd.h5" + DonorData(G=gdata, C=adata).write_h5_dd(str(out)) + + dd = read_lazy_dd(str(out)) + assert hasattr(dd, "_h5_handle") + assert isinstance(dd._h5_handle, h5py.File) + assert dd._h5_handle.id.valid # file is open + + +@pytest.mark.slow +def test_read_lazy_dd_h5_compute_after_slice(tmp_path, adata, gdata): + """Slicing then computing on a lazy h5-backed G should return the correct values.""" + out = tmp_path / "test.dd.h5" + DonorData(G=gdata, C=adata).write_h5_dd(str(out)) + + dd = read_lazy_dd(str(out)) + keep = list(dd.G.var_names[:2]) + sub = dd.G[:, keep] + arr = np.asarray(sub.X) if not hasattr(sub.X, "compute") else sub.X.compute() + assert arr.shape[1] == len(keep) + + +# ── read_lazy_dd: per-side lazy flags + non-lazy fallback ──────────────────── + + +@pytest.mark.slow +def test_read_lazy_dd_eager_G_lazy_C_zarr(tmp_path, adata, gdata): + """Uncommon combo: G eager, C lazy.""" + out = tmp_path / "test.dd.zarr" + DonorData(G=gdata, C=adata).write_zarr_dd(str(out)) + + dd = read_lazy_dd(str(out), lazy_G=False, lazy_C=True) + assert dd._lazy_G is False + assert dd._lazy_C is True + assert isinstance(dd.G.X, np.ndarray) + assert isinstance(dd.C.X, da.Array) + + +@pytest.mark.slow +def test_read_lazy_dd_both_eager_delegates_zarr(tmp_path, adata, gdata): + """lazy_G=False, lazy_C=False → behaves as a regular non-lazy read_dd.""" + from cellink.io import read_dd + + out = tmp_path / "test.dd.zarr" + DonorData(G=gdata, C=adata).write_zarr_dd(str(out)) + + dd_eager = read_dd(str(out)) + dd = read_lazy_dd(str(out), lazy_G=False, lazy_C=False) + assert dd.is_lazy is False + assert isinstance(dd.G.X, np.ndarray) + assert isinstance(dd.C.X, np.ndarray) + assert dd.G.shape == dd_eager.G.shape + assert dd.C.shape == dd_eager.C.shape + np.testing.assert_array_almost_equal(dd.G.X, np.asarray(dd_eager.G.X)) + + +@pytest.mark.slow +def test_read_lazy_dd_eager_G_lazy_C_h5(tmp_path, adata, gdata): + out = tmp_path / "test.dd.h5" + DonorData(G=gdata, C=adata).write_h5_dd(str(out)) + + dd = read_lazy_dd(str(out), lazy_G=False, lazy_C=True) + assert dd._lazy_G is False + assert dd._lazy_C is True + assert isinstance(dd.G.X, np.ndarray) + assert isinstance(dd.C.X, da.Array) + + +@pytest.mark.slow +def test_read_lazy_dd_both_eager_delegates_h5(tmp_path, adata, gdata): + from cellink.io import read_dd + + out = tmp_path / "test.dd.h5" + DonorData(G=gdata, C=adata).write_h5_dd(str(out)) + + dd_eager = read_dd(str(out)) + dd = read_lazy_dd(str(out), lazy_G=False, lazy_C=False) + assert dd.is_lazy is False + assert isinstance(dd.G.X, np.ndarray) + assert isinstance(dd.C.X, np.ndarray) + np.testing.assert_array_almost_equal(dd.G.X, np.asarray(dd_eager.G.X)) + # No h5 handle pinned — both-eager path goes through read_dd which closes the file. + assert not hasattr(dd, "_h5_handle") + + +# ── read_lazy_dd dispatch by extension ─────────────────────────────────────── + + +def test_read_lazy_dd_unknown_extension(tmp_path): + out = tmp_path / "test.unknown" + out.touch() + with pytest.raises(ValueError, match="Cannot infer format"): + read_lazy_dd(str(out))