diff --git a/.gitignore b/.gitignore index 1fd1f144..2962c3d1 100644 --- a/.gitignore +++ b/.gitignore @@ -93,6 +93,9 @@ ipython_config.py # install all needed dependencies. #Pipfile.lock +# uv (lock file is not tracked in this repo) +uv.lock + # PEP 582; used by e.g. github.com/David-OConnor/pyflow __pypackages__/ diff --git a/pyproject.toml b/pyproject.toml index 13e8e696..12b09e61 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,7 @@ classifiers = [ dynamic = [ "version" ] dependencies = [ "adjusttext", - "anndata", + "anndata>=0.13", # lazy read accessors (read_lazy) — the loader reads reps via AnnData, not key_backings "cloudpickle", "coverage", "dask", @@ -47,6 +47,15 @@ dependencies = [ "session-info", ] +optional-dependencies.annbatch = [ + # Pinned to the branch adding `SequentialClassSampler` (streaming eval path); revert to a released + # `annbatch>=0.2.x` once it lands upstream. Direct URL so pip/uv/tox all resolve the branch. + "annbatch @ git+https://github.com/selmanozleyen/annbatch.git@feat/bound-class-wip", + # Rust-backed zarr v3 codec pipeline. Always installed with the streaming path: without it, zarr + # decode runs on a single GIL-bound Python thread (profiled bottleneck when building/reading large + # sparse Tahoe zarrs); `zarrs` decodes multithreaded in Rust and actually uses the allocated cores. + "zarrs", +] optional-dependencies.dev = [ "furo", "myst-nb", @@ -89,6 +98,7 @@ optional-dependencies.pp = [ "rdkit", ] optional-dependencies.test = [ + "cellflow-tools[annbatch]", "cellflow-tools[embedding]", "cellflow-tools[external]", "cellflow-tools[pp]", @@ -105,12 +115,19 @@ urls.Home-page = "https://github.com/theislab/cellflow" urls.Source = "https://github.com/theislab/cellflow" [tool.hatch.build.targets.wheel] -packages = [ 'src/cellflow' ] +# `dagloader` is vendored for now but kept as a standalone top-level package (import path +# `dagloader`) so it can later be extracted into an external dependency with no code changes. +packages = [ 'src/cellflow', 'src/dagloader' ] [tool.hatch.version] source = "vcs" fallback-version = "0.1.0" +[tool.hatch.metadata] +# the `annbatch` extra pins the fork by git URL (a direct reference); hatchling rejects +# direct references in metadata unless this is set. Drop once annbatch lands upstream. +allow-direct-references = true + [tool.ruff] line-length = 120 src = [ "src" ] @@ -172,6 +189,12 @@ markers = [ "slow: marks tests as slow (deselect with '-m \"not slow\"')", "internet: marks tests that require internet access (deselect with '-m \"not internet\"')", ] +filterwarnings = [ + # jaxopt is unmaintained and warns on import; it is pulled in transitively by ott-jax + # (ott.geometry.costs imports it on every version, incl. main) — not by cellflow. Nothing to fix + # on our side, so silence the noise rather than fail on it. + "ignore:JAXopt is no longer maintained:DeprecationWarning", +] [tool.coverage.run] branch = true diff --git a/src/cellflow/data/__init__.py b/src/cellflow/data/__init__.py index e6f6f2de..71205c14 100644 --- a/src/cellflow/data/__init__.py +++ b/src/cellflow/data/__init__.py @@ -1,6 +1,6 @@ -from cellflow.data._data import BaseDataMixin, ConditionData, PredictionData, TrainingData, ValidationData -from cellflow.data._dataloader import PredictionSampler, TrainSampler, ValidationSampler +from cellflow.data._data import BaseDataMixin, ConditionData, PredictionData from cellflow.data._datamanager import DataManager +from cellflow.data._legacy import PredictionSampler, TrainingData, TrainSampler, ValidationData, ValidationSampler __all__ = [ "DataManager", diff --git a/src/cellflow/data/_annbatch.py b/src/cellflow/data/_annbatch.py new file mode 100644 index 00000000..5503b408 --- /dev/null +++ b/src/cellflow/data/_annbatch.py @@ -0,0 +1,207 @@ +"""Build the annbatch/dagloader streaming training path from a CellFlow covariate spec. + +Turns the ``prepare_data`` covariate arguments into a :class:`dagloader.Scheme` (perturbed root, +matched-control child) and a ``condition_fn`` mapping each sampled leaf to its condition embedding. The +embeddings reuse the in-memory machinery — a cell-free ``AnnData`` shell (``obs`` + ``uns``) drives a +:class:`~cellflow.data._datamanager.DataManager` and +:func:`~cellflow.data._condition.build_condition_data` — so they match the in-memory path exactly. Only +``obs`` (and the embedding tables) are read here; cells are streamed later by ``dagloader``. +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import anndata as ad +import numpy as np + +from cellflow.data._condition import _key_layout, build_condition_data, enumerate_perturbations +from cellflow.data._datamanager import DataManager + +if TYPE_CHECKING: + from cellflow._types import ArrayLike + from dagloader import Container, Scheme + +Leaf = tuple[object, ...] # a scheme leaf: one value per grouping column + +__all__ = [ + "AnnbatchTraining", + "assert_source_chunkable", + "build_annbatch_training", + "sample_rep_to_key", +] + + +def assert_source_chunkable(source: Container, cols: Sequence[str], chunk_size: int) -> None: + """Raise unless ``source`` satisfies annbatch's run-length rule for ``chunk_size`` (reads ``obs`` only). + + With ``chunk_size > 1`` annbatch reads contiguous slices, so every contiguous run of each category + (a ``cols`` combination) must be at least ``chunk_size`` cells; a category may span several runs. + On a shorter run, raises pointing at ``DatasetCollection.add_adatas(groupby=...)``. In-memory sources + are grouped automatically by :func:`build_annbatch_training`, so this only bites out-of-core inputs. + """ + from dagloader._io import leaf_codes, obs_columns + + codes, _ = leaf_codes(obs_columns(source, list(cols)), list(cols)) + if len(codes) == 0: + return + run_starts = np.concatenate([[0], np.flatnonzero(np.diff(codes) != 0) + 1]) + run_lengths = np.diff(np.concatenate([run_starts, [len(codes)]])) + shortest = int(run_lengths.min()) + if shortest < chunk_size: + raise ValueError( + f"chunk_size={chunk_size} requires every contiguous run of each category (a {list(cols)} " + f"combination) to be at least chunk_size cells, but the source has a run of only {shortest}. " + f"Group the source so each category forms long runs — e.g. create the DatasetCollection with " + f"`add_adatas(..., groupby={list(cols)})` — or use chunk_size=1. (A category may span several " + f"runs; only each run's length matters.)" + ) + + +def sample_rep_to_key(sample_rep: str) -> str: + """CellFlow ``sample_rep`` → dagloader representation key (``"X"`` or ``"obsm/"``).""" + return "X" if sample_rep == "X" else f"obsm/{sample_rep}" + + +@dataclass(frozen=True) +class AnnbatchTraining: + """Everything the model needs to stream-train, assembled from a covariate spec (cells untouched).""" + + scheme: Scheme + condition_fn: Callable[[Leaf], dict[str, np.ndarray]] + condition_data: dict[str, np.ndarray] + data_manager: DataManager + data_dim: int + max_combination_length: int + + +def build_annbatch_training( + source: Container, + *, + sample_rep: str, + control_key: str, + perturbation_covariates: Mapping[str, Sequence[str]], + perturbation_covariate_reps: Mapping[str, str] | None = None, + sample_covariates: Sequence[str] | None = None, + sample_covariate_reps: Mapping[str, str] | None = None, + split_covariates: Sequence[str] | None = None, + max_combination_length: int | None = None, + null_value: float = 0.0, + rep_dict: Mapping[str, Mapping[str, ArrayLike]] | None = None, + seed: int = 0, + control_in_memory: bool = False, +) -> AnnbatchTraining: + """Assemble the :class:`dagloader.Scheme` + ``condition_fn`` for the streaming path (obs only). + + ``source`` is an out-of-core :class:`annbatch.DatasetCollection` or an in-memory ``AnnData``. + ``rep_dict`` holds the covariate embedding tables (as ``adata.uns`` would); pass :obj:`None` when the + primary covariate is categorical (one-hot). + + ``control_in_memory`` tells dagloader to materialize the control (child) node into RAM (sets + :attr:`~dagloader.Node.in_memory`; dagloader owns the read via :func:`~dagloader._io.materialize_node`), + so the matched control is served from memory while the perturbed target keeps streaming out of core — a + large dataloader speedup, since controls are re-drawn every batch. Only enable it when the controls fit + in host RAM (the small population by design). + """ + from dagloader import Bind, Node, Scheme, uniform + from dagloader._io import key_backings, obs_columns + + context = tuple(split_covariates or ()) + pert_cols = tuple(c for grp in perturbation_covariates.values() for c in grp) + samp_cols = tuple(sample_covariates or ()) + cols = tuple(dict.fromkeys((*context, *pert_cols, *samp_cols))) # grouping cols (deduped, ordered) + key = sample_rep_to_key(sample_rep) + + obs = obs_columns(source, [*cols, control_key]) + + # In-memory source: stable-sort by the grouping columns so `chunk_size > 1` reads contiguous slices + # (cheap, and cell order is irrelevant). Out-of-core sources aren't reordered (expensive zarr re-sort) + # — they must be built grouped; `assert_source_chunkable` enforces that. + if isinstance(source, ad.AnnData): + order = obs[list(cols)].reset_index(drop=True).sort_values(list(cols), kind="stable").index.to_numpy() + source = source[order].copy() + obs = obs_columns(source, [*cols, control_key]) + + # DataManager as a covariate-encoder factory: reads only obs + uns, so a cell-free shell suffices. + # `sample_rep` is stored for validation's `_get_cell_data` (verification is type-only, so it's safe here). + shell = ad.AnnData(obs=obs.copy()) + shell.uns = dict(rep_dict or {}) + dm = DataManager( + shell, + sample_rep=sample_rep, + control_key=control_key, + perturbation_covariates=dict(perturbation_covariates), + perturbation_covariate_reps=dict(perturbation_covariate_reps) if perturbation_covariate_reps else None, + sample_covariates=list(samp_cols), + sample_covariate_reps=dict(sample_covariate_reps) if sample_covariate_reps else None, + split_covariates=list(context), + max_combination_length=max_combination_length, + null_value=null_value, + ) + + # Per-condition embeddings — the shared helper → identical to the in-memory path (parity-tested). + condition_data = build_condition_data( + obs, + shell.uns, + control_key=control_key, + perturb_covar_keys=dm._perturb_covar_keys, + split_covariates=list(context), + sample_covariates=list(samp_cols), + perturbation_covariates=dict(perturbation_covariates), + covariate_reps=dm._covariate_reps, + covar_to_idx=dm.covar_to_idx, + is_categorical=dm.is_categorical, + primary_one_hot_encoder=dm.primary_one_hot_encoder, + primary_group=dm.primary_group, + linked_perturb_covars=dm.linked_perturb_covars, + max_combination_length=dm.max_combination_length, + null_value=null_value, + ) + + # leaf → perturbation index → embedding. `enumerate_perturbations` lays tuples out in `tuple_keys` + # order (differs from `cols`), so re-project the leaf; string-normalize both sides so dtype quirks + # (categorical / numpy scalars) don't break the match. + idx_to_cov = enumerate_perturbations( + obs, + control_key=control_key, + perturb_covar_keys=dm._perturb_covar_keys, + split_covariates=list(context), + sample_covariates=list(samp_cols), + ) + _, tuple_keys = _key_layout(dm._perturb_covar_keys, list(context), list(samp_cols)) + cov_to_idx = {tuple(map(str, cov)): i for i, cov in idx_to_cov.items()} + reorder = [cols.index(c) for c in tuple_keys] + + def condition_fn(leaf: Leaf) -> dict[str, np.ndarray]: + idx = cov_to_idx[tuple(str(leaf[i]) for i in reorder)] + return {group: condition_data[group][[idx]] for group in condition_data} + + # The Scheme: root = perturbed combos, child = matched-control combos (bound on the context columns). + ctrl_flag = obs[control_key].to_numpy().astype(bool) + pert = [tuple(r) for r in obs.loc[~ctrl_flag, list(cols)].drop_duplicates().to_numpy()] + ctrl = [tuple(r) for r in obs.loc[ctrl_flag, list(cols)].drop_duplicates().to_numpy()] + # `control_in_memory` just tells dagloader to materialize the control (child) node into RAM — the + # perturbed target keeps streaming out of core. dagloader owns the materialization (Node.in_memory); + # the bind still matches control↔target by the context columns. + scheme = Scheme( + sources={"data": source}, + nodes={ + "pert": Node("data", cols, key, uniform(pert)), + "ctrl": Node("data", cols, key, uniform(ctrl), in_memory=control_in_memory), + }, + root="pert", + binds=(Bind("pert", "ctrl", common=context),), + seed=seed, + ) + + data_dim = int(key_backings(source, key)[0].shape[1]) # key_backings wraps sparse groups → has .shape + return AnnbatchTraining( + scheme=scheme, + condition_fn=condition_fn, + condition_data=condition_data, + data_manager=dm, + data_dim=data_dim, + max_combination_length=dm.max_combination_length, + ) diff --git a/src/cellflow/data/_condition.py b/src/cellflow/data/_condition.py new file mode 100644 index 00000000..bd43a599 --- /dev/null +++ b/src/cellflow/data/_condition.py @@ -0,0 +1,152 @@ +"""Pure helpers for building condition data from ``obs`` + covariate representations. + +Extracted from :class:`~cellflow.data._datamanager.DataManager` so the in-memory path and the +annbatch/dagloader streaming path share **one** implementation — they must produce identical +condition embeddings. Nothing here touches the cell matrix (``X``); only the covariate spec, the +``obs`` table, and the representation dict (``uns``) are used. ``DataManager`` delegates to these. +""" + +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Mapping, Sequence +from typing import Any + +import numpy as np +import pandas as pd + +from cellflow._logging import logger +from cellflow.data._utils import _to_list + +__all__ = ["build_condition_data", "enumerate_perturbations", "get_max_combination_length"] + + +def _key_layout( + perturb_covar_keys: Sequence[str], + split_covariates: Sequence[str], + sample_covariates: Sequence[str], +) -> tuple[list[str], list[str]]: + """Column layouts used by DataManager's condition enumeration. + + Returns ``(all_combs_keys, tuple_keys)``: the sort order for assigning perturbation indices, and + the order each ``perturbation_idx_to_covariates`` tuple is laid out in. Shared by + :func:`enumerate_perturbations` and :func:`build_condition_data` so both agree with DataManager. + """ + uniq_sample_keys = list(split_covariates) if len(split_covariates) else list(sample_covariates) + perturbation_covariates_keys = [k for k in perturb_covar_keys if k not in uniq_sample_keys] + if len(split_covariates): + all_combs_keys = uniq_sample_keys + perturbation_covariates_keys + else: + all_combs_keys = perturbation_covariates_keys + uniq_sample_keys + tuple_keys = perturbation_covariates_keys + uniq_sample_keys + return all_combs_keys, tuple_keys + + +def get_max_combination_length( + perturbation_covariates: dict[str, Sequence[str]], + max_combination_length: int | None = None, +) -> int: + """Maximum number of perturbations in a combination (a pure function of the spec). + + This is the largest covariate-group size in ``perturbation_covariates`` (e.g. ``{"drug": + ("drug_1", "drug_2")}`` ⇒ ``2``). ``max_combination_length`` acts only as a floor: a value below + the observed maximum is raised to it (with a warning); a larger value is kept as-is. + """ + obs_max_combination_length = max(len(comb) for comb in perturbation_covariates.values()) + if max_combination_length is None: + return obs_max_combination_length + elif max_combination_length < obs_max_combination_length: + logger.warning( + f"Provided `max_combination_length` is smaller than the observed maximum combination length of the perturbation covariates. Setting maximum combination length to {obs_max_combination_length}.", + stacklevel=2, + ) + return obs_max_combination_length + else: + return max_combination_length + + +def enumerate_perturbations( + obs: pd.DataFrame, + *, + control_key: str, + perturb_covar_keys: Sequence[str], + split_covariates: Sequence[str], + sample_covariates: Sequence[str], +) -> dict[int, tuple]: + """Map each perturbation index to its covariate tuple, matching ``DataManager`` (no dask/masks). + + Reproduces the target-combination enumeration inside + :meth:`~cellflow.data._datamanager.DataManager._get_condition_data`: the unique non-control + combinations, indexed by ``arange`` after sorting by ``all_combs_keys``, each tuple laid out in + ``perturbation_covariates_keys + uniq_sample_keys`` order. Uses plain pandas (no per-cell masks, + no dask), so it is cheap enough to run off a streamed ``obs`` table. + """ + all_combs_keys, tuple_keys = _key_layout(perturb_covar_keys, split_covariates, sample_covariates) + + df = obs[[*all_combs_keys, control_key]].copy() + for col in all_combs_keys: # mirror DataManager: categorical sort order (only cast if needed) + if df[col].dtype != "category": + df[col] = df[col].astype("category") + df = df[~df[control_key].astype(bool)] + combos = df[all_combs_keys].drop_duplicates().sort_values(by=all_combs_keys).reset_index(drop=True) + return {i: tuple(combos.loc[i, tuple_keys]) for i in range(len(combos))} + + +def build_condition_data( + obs: pd.DataFrame, + rep_dict: Mapping[str, Any], + *, + control_key: str, + perturb_covar_keys: Sequence[str], + split_covariates: Sequence[str], + sample_covariates: Sequence[str], + perturbation_covariates: Mapping[str, Sequence[str]], + covariate_reps: Mapping[str, str], + covar_to_idx: Mapping[str, int], + is_categorical: bool, + primary_one_hot_encoder: Any, + primary_group: str, + linked_perturb_covars: Mapping[str, Any], + max_combination_length: int, + null_value: float, +) -> dict[str, np.ndarray]: + """Assemble the per-condition embeddings (``condition_data``), matching ``DataManager`` (no dask). + + Enumerates the perturbations (:func:`enumerate_perturbations`), then reuses the same per-condition + embedding function as the in-memory path — :meth:`DataManager._get_embeddings` — over a plain + serial loop instead of DataManager's ``dask.delayed`` fan-out. The returned dict maps each covariate + group to an ``(n_perturbations, max_combination_length, dim)`` array, aligned to the perturbation + index from :func:`enumerate_perturbations`. Values are identical to + ``DataManager._get_condition_data(...).condition_data`` (parity-tested); only the orchestration + differs, so it is cheap enough to run off a streamed ``obs`` table. + """ + from cellflow.data._datamanager import DataManager # lazy: reuse the shared per-condition embedding + + _, tuple_keys = _key_layout(perturb_covar_keys, split_covariates, sample_covariates) + idx_to_covariates = enumerate_perturbations( + obs, + control_key=control_key, + perturb_covar_keys=perturb_covar_keys, + split_covariates=split_covariates, + sample_covariates=sample_covariates, + ) + perturb_covariates = OrderedDict({k: sorted(_to_list(v)) for k, v in perturbation_covariates.items()}) + condition_data: dict[str, list[np.ndarray]] = {k: [] for k in covar_to_idx} + for idx in sorted(idx_to_covariates): + tgt_cond = dict(zip(tuple_keys, idx_to_covariates[idx], strict=True)) + embeddings = DataManager._get_embeddings( + condition_data=tgt_cond, + rep_dict=rep_dict, + perturb_covariates=perturb_covariates, + covariate_reps=covariate_reps, + is_categorical=is_categorical, + primary_one_hot_encoder=primary_one_hot_encoder, + null_value=null_value, + max_combination_length=max_combination_length, + linked_perturb_covars=linked_perturb_covars, + sample_covariates=sample_covariates, + primary_group=primary_group, + ) + for pert_cov, emb in embeddings.items(): + condition_data[pert_cov].append(emb) + return {pert_cov: np.array(emb) for pert_cov, emb in condition_data.items()} diff --git a/src/cellflow/data/_data.py b/src/cellflow/data/_data.py index 0f51d304..ee82897f 100644 --- a/src/cellflow/data/_data.py +++ b/src/cellflow/data/_data.py @@ -11,8 +11,6 @@ "BaseDataMixin", "ConditionData", "PredictionData", - "TrainingData", - "ValidationData", ] @@ -81,96 +79,6 @@ class ConditionData(BaseDataMixin): data_manager: Any -@dataclass -class TrainingData(BaseDataMixin): - """Training data. - - Parameters - ---------- - cell_data - The representation of cell data, e.g. PCA of gene expression data. - split_covariates_mask - Mask of the split covariates. - split_idx_to_covariates - Dictionary explaining values in ``split_covariates_mask``. - perturbation_covariates_mask - Mask of the perturbation covariates. - perturbation_idx_to_covariates - Dictionary explaining values in ``perturbation_covariates_mask``. - condition_data - Dictionary with embeddings for conditions. - control_to_perturbation - Mapping from control index to target distribution indices. - max_combination_length - Maximum number of covariates in a combination. - data_manager - The data manager - """ - - cell_data: np.ndarray # (n_cells, n_features) - split_covariates_mask: np.ndarray # (n_cells,), which cell assigned to which source distribution - split_idx_to_covariates: dict[int, tuple[Any, ...]] # (n_sources,) dictionary explaining split_covariates_mask - perturbation_covariates_mask: np.ndarray # (n_cells,), which cell assigned to which target distribution - perturbation_idx_to_covariates: dict[ - int, tuple[str, ...] - ] # (n_targets,), dictionary explaining perturbation_covariates_mask - perturbation_idx_to_id: dict[int, Any] - condition_data: dict[str, np.ndarray] # (n_targets,) all embeddings for conditions - control_to_perturbation: dict[int, np.ndarray] # mapping from control idx to target distribution idcs - max_combination_length: int - null_value: Any - data_manager: Any - - -@dataclass -class ValidationData(BaseDataMixin): - """Data container for the validation data. - - Parameters - ---------- - cell_data - The representation of cell data, e.g. PCA of gene expression data. - split_covariates_mask - Mask of the split covariates. - split_idx_to_covariates - Dictionary explaining values in ``split_covariates_mask``. - perturbation_covariates_mask - Mask of the perturbation covariates. - perturbation_idx_to_covariates - Dictionary explaining values in ``perturbation_covariates_mask``. - condition_data - Dictionary with embeddings for conditions. - control_to_perturbation - Mapping from control index to target distribution indices. - max_combination_length - Maximum number of covariates in a combination. - data_manager - The data manager - n_conditions_on_log_iteration - Number of conditions to use for computation callbacks at each logged iteration. - If :obj:`None`, use all conditions. - n_conditions_on_train_end - Number of conditions to use for computation callbacks at the end of training. - If :obj:`None`, use all conditions. - """ - - cell_data: np.ndarray # (n_cells, n_features) - split_covariates_mask: np.ndarray # (n_cells,), which cell assigned to which source distribution - split_idx_to_covariates: dict[int, tuple[Any, ...]] # (n_sources,) dictionary explaining split_covariates_mask - perturbation_covariates_mask: np.ndarray # (n_cells,), which cell assigned to which target distribution - perturbation_idx_to_covariates: dict[ - int, tuple[str, ...] - ] # (n_targets,), dictionary explaining perturbation_covariates_mask - perturbation_idx_to_id: dict[int, Any] - condition_data: dict[str, np.ndarray] # (n_targets,) all embeddings for conditions - control_to_perturbation: dict[int, np.ndarray] # mapping from control idx to target distribution idcs - max_combination_length: int - null_value: Any - data_manager: Any - n_conditions_on_log_iteration: int | None = None - n_conditions_on_train_end: int | None = None - - @dataclass class PredictionData(BaseDataMixin): """Data container to perform prediction. diff --git a/src/cellflow/data/_dataloader.py b/src/cellflow/data/_dataloader.py index 70bd91ef..3c4d1fbc 100644 --- a/src/cellflow/data/_dataloader.py +++ b/src/cellflow/data/_dataloader.py @@ -1,303 +1,97 @@ -import abc -import queue -import threading -from collections.abc import Generator -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal -import jax import numpy as np -from cellflow.data._data import PredictionData, TrainingData, ValidationData +from cellflow._types import ArrayLike -__all__ = ["TrainSampler", "ValidationSampler", "PredictionSampler", "OOCTrainSampler"] +if TYPE_CHECKING: + from dagloader import DAGEvalLoader, DAGLoader +__all__ = [ + "DAGEvalAdapter", + "DAGTrainAdapter", +] -class TrainSampler: - """Data sampler for :class:`~cellflow.data.TrainingData`. - Parameters - ---------- - data - The training data. - batch_size - The batch size. +def _densify(x: ArrayLike) -> ArrayLike: + """Densify a possibly-sparse streamed cell batch, staying on the array's native backend. + ``dagloader`` yields native jax arrays (``to="jax"``): dense reps pass straight through, and a + sparse rep arrives as a ``jax.experimental.sparse`` CSR — densified here via ``.todense()`` (still + on-device). No numpy cast and no host round-trip, so a GPU-resident batch reaches the solver as-is. """ + todense = getattr(x, "todense", None) # (jax/scipy) sparse → dense; dense arrays lack this + return todense() if todense is not None else x - def __init__(self, data: TrainingData, batch_size: int = 1024): - self._data = data - self._data_idcs = np.arange(data.cell_data.shape[0]) - self.batch_size = batch_size - self.n_source_dists = data.n_controls - self.n_target_dists = data.n_perturbations - - self._control_to_perturbation_keys = sorted(data.control_to_perturbation.keys()) - self._has_condition_data = data.condition_data is not None - - def _sample_target_dist_idx(self, source_dist_idx, rng): - """Sample a target distribution index given the source distribution index.""" - return rng.choice(self._data.control_to_perturbation[source_dist_idx]) - - def _get_embeddings(self, idx, condition_data) -> dict[str, np.ndarray]: - """Get embeddings for a given index.""" - result = {} - for key, arr in condition_data.items(): - result[key] = np.expand_dims(arr[idx], 0) - return result - - def _sample_from_mask(self, rng, mask) -> np.ndarray: - """Sample indices according to a mask.""" - # Convert mask to probability distribution - valid_indices = np.where(mask)[0] - - # Handle case with no valid indices (should not happen in practice) - if len(valid_indices) == 0: - raise ValueError("No valid indices found in the mask") - - # Sample from valid indices with equal probability - batch_idcs = rng.choice(valid_indices, self.batch_size, replace=True) - return batch_idcs - - def sample(self, rng) -> dict[str, Any]: - """Sample a batch of data. - - Parameters - ---------- - seed : int, optional - Random seed - - Returns - ------- - Dictionary with source and target data - """ - # Sample source distribution index - source_dist_idx = rng.integers(0, self.n_source_dists) - - # Get source cells - source_cells_mask = self._data.split_covariates_mask == source_dist_idx - source_batch_idcs = self._sample_from_mask(rng, source_cells_mask) - source_batch = self._data.cell_data[source_batch_idcs] - - target_dist_idx = self._sample_target_dist_idx(source_dist_idx, rng) - target_cells_mask = self._data.perturbation_covariates_mask == target_dist_idx - target_batch_idcs = self._sample_from_mask(rng, target_cells_mask) - target_batch = self._data.cell_data[target_batch_idcs] - - if not self._has_condition_data: - return {"src_cell_data": source_batch, "tgt_cell_data": target_batch} - else: - condition_batch = self._get_embeddings(target_dist_idx, self._data.condition_data) - return { - "src_cell_data": source_batch, - "tgt_cell_data": target_batch, - "condition": condition_batch, - } - - @property - def data(self): - """The training data.""" - return self._data - - -class BaseValidSampler(abc.ABC): - @abc.abstractmethod - def sample(*args, **kwargs): - pass - def _get_key(self, cond_idx: int) -> tuple[str, ...]: - if len(self._data.perturbation_idx_to_id): # type: ignore[attr-defined] - return self._data.perturbation_idx_to_id[cond_idx] # type: ignore[attr-defined] - cov_combination = self._data.perturbation_idx_to_covariates[cond_idx] # type: ignore[attr-defined] - return tuple(cov_combination[i] for i in range(len(cov_combination))) +class DAGTrainAdapter: + """Adapt a ``dagloader.DAGLoader`` stream to the trainer's ``sample(rng)`` batch contract. - def _get_perturbation_to_control(self, data: ValidationData | PredictionData) -> dict[int, np.ndarray]: - d = {} - for k, v in data.control_to_perturbation.items(): - for el in v: - d[el] = k - return d - - def _get_condition_data(self, cond_idx: int) -> dict[str, np.ndarray]: - return {k: v[[cond_idx], ...] for k, v in self._data.condition_data.items()} # type: ignore[attr-defined] - - -class ValidationSampler(BaseValidSampler): - """Data sampler for :class:`~cellflow.data.ValidationData`. - - Parameters - ---------- - val_data - The validation data. - seed - Random seed. + Renames the loader's ``{"target", "source", "condition"}`` batch to the model's + ``{"src_cell_data", "tgt_cell_data", "condition"}`` and densifies any sparse cell rep, keeping the + arrays on their native (jax) backend, so the in-memory and streaming paths reach the solver + identically. This is the *training* adapter; validation uses :class:`DAGEvalAdapter`. """ - def __init__(self, val_data: ValidationData, seed: int = 0) -> None: - self._data = val_data - self.perturbation_to_control = self._get_perturbation_to_control(val_data) - self.n_conditions_on_log_iteration = ( - val_data.n_conditions_on_log_iteration - if val_data.n_conditions_on_log_iteration is not None - else val_data.n_perturbations - ) - self.n_conditions_on_train_end = ( - val_data.n_conditions_on_train_end - if val_data.n_conditions_on_train_end is not None - else val_data.n_perturbations - ) - self.rng = np.random.default_rng(seed) - if self._data.condition_data is None: - raise NotImplementedError("Validation data must have condition data.") + def __init__(self, loader: "DAGLoader"): + self._loader = loader + self._iter = iter(loader) - def sample(self, mode: Literal["on_log_iteration", "on_train_end"]) -> Any: - """Sample data for validation. + def sample(self, rng: np.random.Generator | None = None) -> dict[str, ArrayLike | dict[str, ArrayLike]]: + """Return the next streamed batch as a trainer batch dict. - Parameters - ---------- - mode - Sampling mode. Either ``"on_log_iteration"`` or ``"on_train_end"``. - - Returns - ------- - Dictionary with source, condition, and target data from the validation data. + ``rng`` is unused — the ``DAGLoader`` owns its own reproducible RNG. """ - size = self.n_conditions_on_log_iteration if mode == "on_log_iteration" else self.n_conditions_on_train_end - condition_idcs = self.rng.choice(self._data.n_perturbations, size=(size,), replace=False) - - source_idcs = [self.perturbation_to_control[cond_idx] for cond_idx in condition_idcs] - source_cells_mask = [self._data.split_covariates_mask == source_idx for source_idx in source_idcs] - source_cells = [self._data.cell_data[mask] for mask in source_cells_mask] - target_cells_mask = [cond_idx == self._data.perturbation_covariates_mask for cond_idx in condition_idcs] - target_cells = [self._data.cell_data[mask] for mask in target_cells_mask] - conditions = [self._get_condition_data(cond_idx) for cond_idx in condition_idcs] - cell_rep_dict = {} - cond_dict = {} - true_dict = {} - for i in range(len(condition_idcs)): - k = self._get_key(condition_idcs[i]) - cell_rep_dict[k] = source_cells[i] - cond_dict[k] = conditions[i] - true_dict[k] = target_cells[i] - - return {"source": cell_rep_dict, "condition": cond_dict, "target": true_dict} + batch = next(self._iter) + out: dict[str, ArrayLike | dict[str, ArrayLike]] = { + "src_cell_data": _densify(batch["source"]), + "tgt_cell_data": _densify(batch["target"]), + } + if "condition" in batch: + out["condition"] = batch["condition"] + return out - @property - def data(self) -> ValidationData: - """The validation data.""" - return self._data +class DAGEvalAdapter: + """Adapt a ``dagloader.DAGEvalLoader`` to the trainer's validation ``sample(mode)`` contract. -class PredictionSampler(BaseValidSampler): - """Data sampler for :class:`~cellflow.data.PredictionData`. + Yields per-condition ``{"source", "condition", "target"}`` dicts (keyed by the perturbed leaf), as the + trainer's validation step expects. The control-rooted ``DAGEvalLoader`` reads each control population's + cells in full for the source and samples a matched perturbed batch for the target — via annbatch, no + boolean masking. ``n_conditions_*`` sets how many (control-population, drug) batches to draw per mode. + This mirrors the in-memory :class:`~cellflow.data._legacy.ValidationSampler` contract for the streaming + path, so the trainer's validation step is identical either way. Parameters ---------- - pred_data - The prediction data. - + eval_loader + A :class:`dagloader.DAGEvalLoader` built over the validation source. + n_conditions_on_log_iteration, n_conditions_on_train_end + How many conditions to draw per mode; :obj:`None` visits each control population once. """ - def __init__(self, pred_data: PredictionData) -> None: - self._data = pred_data - self.perturbation_to_control = self._get_perturbation_to_control(pred_data) - if self._data.condition_data is None: - raise NotImplementedError("Validation data must have condition data.") - - def sample(self) -> Any: - """Sample data for prediction. - - Returns - ------- - Dictionary with source and condition data from the prediction data. - """ - condition_idcs = range(self._data.n_perturbations) - - source_idcs = [self.perturbation_to_control[cond_idx] for cond_idx in condition_idcs] - source_cells_mask = [self._data.split_covariates_mask == source_idx for source_idx in source_idcs] - source_cells = [self._data.cell_data[mask] for mask in source_cells_mask] - conditions = [self._get_condition_data(cond_idx) for cond_idx in condition_idcs] - cell_rep_dict = {} - cond_dict = {} - for i in range(len(condition_idcs)): - k = self._get_key(condition_idcs[i]) - cell_rep_dict[k] = source_cells[i] - cond_dict[k] = conditions[i] - - return { - "source": cell_rep_dict, - "condition": cond_dict, - } - - @property - def data(self) -> PredictionData: - """The training data.""" - return self._data - - -def prefetch_to_device( - sampler: TrainSampler, seed: int, num_iterations: int, prefetch_factor: int = 2, num_workers: int = 4 -) -> Generator[dict[str, Any], None, None]: - seq = np.random.SeedSequence(seed) - random_generators = [np.random.default_rng(s) for s in seq.spawn(num_workers)] - - q: queue.Queue[dict[str, Any]] = queue.Queue(maxsize=prefetch_factor * num_workers) - sem = threading.Semaphore(num_iterations) - stop_event = threading.Event() - - def worker(rng: np.random.Generator): - while not stop_event.is_set() and sem.acquire(blocking=False): - batch = sampler.sample(rng) - batch = jax.device_put(batch, jax.devices()[0], donate=True) - jax.block_until_ready(batch) - while not stop_event.is_set(): - try: - q.put(batch, timeout=1.0) - break # Batch successfully put into the queue; break out of retry loop - except queue.Full: - continue - - return - - # Start multiple worker threads - ts = [] - for i in range(num_workers): - t = threading.Thread(target=worker, daemon=True, name=f"worker-{i}", args=(random_generators[i],)) - t.start() - ts.append(t) - - try: - for _ in range(num_iterations): - # Yield batches from the queue; will block waiting for available batch - yield q.get() - finally: - # When the generator is closed or garbage collected, clean up the worker threads - stop_event.set() # Signal all workers to exit - for t in ts: - t.join() # Wait for all worker threads to finish - - -class OOCTrainSampler: def __init__( - self, data: TrainingData, seed: int, batch_size: int = 1024, num_workers: int = 4, prefetch_factor: int = 2 - ): - self.inner = TrainSampler(data=data, batch_size=batch_size) - self.num_workers = num_workers - self.prefetch_factor = prefetch_factor - self.seed = seed - self._iterator = None - - def set_sampler(self, num_iterations: int) -> None: - self._iterator = prefetch_to_device( - sampler=self.inner, seed=self.seed, num_iterations=num_iterations, prefetch_factor=self.prefetch_factor - ) - - def sample(self, rng=None) -> dict[str, Any]: - if self._iterator is None: - raise ValueError( - "Sampler not set. Use `set_sampler` to set the sampler with" - "the number of iterations. Without the number of iterations," - " the sampler will not be able to sample the data." - ) - if rng is not None: - del rng - return next(self._iterator) + self, + eval_loader: "DAGEvalLoader", + *, + n_conditions_on_log_iteration: int | None = None, + n_conditions_on_train_end: int | None = None, + ) -> None: + self._eval = eval_loader + self.n_conditions_on_log_iteration = n_conditions_on_log_iteration + self.n_conditions_on_train_end = n_conditions_on_train_end + + def sample(self, mode: Literal["on_log_iteration", "on_train_end"]) -> dict[str, dict[Any, Any]]: + """Sample a validation batch: per-condition source/condition/target dicts (keyed by perturbed leaf).""" + # Densify exactly as the training adapter does: with the GPU read path (cupy) annbatch yields a + # jax *sparse* CSR, which the solver's `predict` can't index — so eval source/target must be dense. + n = self.n_conditions_on_log_iteration if mode == "on_log_iteration" else self.n_conditions_on_train_end + source: dict[Any, Any] = {} + condition: dict[Any, Any] = {} + target: dict[Any, Any] = {} + for out in self._eval.iter_conditions(n_conditions=n): + key = tuple(out["leaf"]) + source[key] = _densify(out["source"]) + condition[key] = out["condition"] + target[key] = _densify(out["target"]) + return {"source": source, "condition": condition, "target": target} diff --git a/src/cellflow/data/_datamanager.py b/src/cellflow/data/_datamanager.py index 1b085d39..63845348 100644 --- a/src/cellflow/data/_datamanager.py +++ b/src/cellflow/data/_datamanager.py @@ -13,9 +13,10 @@ from dask.diagnostics import ProgressBar from pandas.api.types import is_numeric_dtype -from cellflow._logging import logger from cellflow._types import ArrayLike -from cellflow.data._data import ConditionData, PredictionData, ReturnData, TrainingData, ValidationData +from cellflow.data._condition import get_max_combination_length +from cellflow.data._data import ConditionData, PredictionData, ReturnData +from cellflow.data._legacy import TrainingData, ValidationData from ._utils import _flatten_list, _to_list @@ -998,17 +999,8 @@ def _get_max_combination_length( perturbation_covariates: dict[str, list[str]], max_combination_length: int | None, ) -> int: - obs_max_combination_length = max(len(comb) for comb in perturbation_covariates.values()) - if max_combination_length is None: - return obs_max_combination_length - elif max_combination_length < obs_max_combination_length: - logger.warning( - f"Provided `max_combination_length` is smaller than the observed maximum combination length of the perturbation covariates. Setting maximum combination length to {obs_max_combination_length}.", - stacklevel=2, - ) - return obs_max_combination_length - else: - return max_combination_length + # Delegates to the shared implementation so the in-memory and streaming paths stay identical. + return get_max_combination_length(perturbation_covariates, max_combination_length) def _get_primary_covar_encoder( self, diff --git a/src/cellflow/data/_legacy.py b/src/cellflow/data/_legacy.py new file mode 100644 index 00000000..3abbafbd --- /dev/null +++ b/src/cellflow/data/_legacy.py @@ -0,0 +1,413 @@ +"""Legacy in-memory data samplers. + +These samplers operate on materialized :mod:`cellflow.data` containers +(:class:`~cellflow.data.TrainingData`, :class:`~cellflow.data.ValidationData`, +:class:`~cellflow.data.PredictionData`) held fully in memory. They are superseded by the +annbatch/``dagloader`` streaming path (see :class:`cellflow.data._dataloader.DAGTrainAdapter` +and :meth:`cellflow.model.CellFlowAnnbatch.prepare_data`), which streams cells out of core and also +accepts an in-memory ``AnnData``. Kept here for backward compatibility. +""" + +import abc +import queue +import threading +from collections.abc import Generator +from dataclasses import dataclass +from typing import Any, Literal + +import jax +import numpy as np + +from cellflow.data._data import BaseDataMixin, PredictionData + +__all__ = [ + "TrainingData", + "ValidationData", + "TrainSampler", + "BaseValidSampler", + "ValidationSampler", + "PredictionSampler", + "OOCTrainSampler", + "prefetch_to_device", +] + + +@dataclass +class TrainingData(BaseDataMixin): + """Training data (in-memory path). + + Parameters + ---------- + cell_data + The representation of cell data, e.g. PCA of gene expression data. + split_covariates_mask + Mask of the split covariates. + split_idx_to_covariates + Dictionary explaining values in ``split_covariates_mask``. + perturbation_covariates_mask + Mask of the perturbation covariates. + perturbation_idx_to_covariates + Dictionary explaining values in ``perturbation_covariates_mask``. + condition_data + Dictionary with embeddings for conditions. + control_to_perturbation + Mapping from control index to target distribution indices. + max_combination_length + Maximum number of covariates in a combination. + data_manager + The data manager + """ + + cell_data: np.ndarray # (n_cells, n_features) + split_covariates_mask: np.ndarray # (n_cells,), which cell assigned to which source distribution + split_idx_to_covariates: dict[int, tuple[Any, ...]] # (n_sources,) dictionary explaining split_covariates_mask + perturbation_covariates_mask: np.ndarray # (n_cells,), which cell assigned to which target distribution + perturbation_idx_to_covariates: dict[ + int, tuple[str, ...] + ] # (n_targets,), dictionary explaining perturbation_covariates_mask + perturbation_idx_to_id: dict[int, Any] + condition_data: dict[str, np.ndarray] # (n_targets,) all embeddings for conditions + control_to_perturbation: dict[int, np.ndarray] # mapping from control idx to target distribution idcs + max_combination_length: int + null_value: Any + data_manager: Any + + +@dataclass +class ValidationData(BaseDataMixin): + """Validation data (in-memory path). + + Parameters + ---------- + cell_data + The representation of cell data, e.g. PCA of gene expression data. + split_covariates_mask + Mask of the split covariates. + split_idx_to_covariates + Dictionary explaining values in ``split_covariates_mask``. + perturbation_covariates_mask + Mask of the perturbation covariates. + perturbation_idx_to_covariates + Dictionary explaining values in ``perturbation_covariates_mask``. + condition_data + Dictionary with embeddings for conditions. + control_to_perturbation + Mapping from control index to target distribution indices. + max_combination_length + Maximum number of covariates in a combination. + data_manager + The data manager + n_conditions_on_log_iteration + Number of conditions to use for computation callbacks at each logged iteration. + If :obj:`None`, use all conditions. + n_conditions_on_train_end + Number of conditions to use for computation callbacks at the end of training. + If :obj:`None`, use all conditions. + """ + + cell_data: np.ndarray # (n_cells, n_features) + split_covariates_mask: np.ndarray # (n_cells,), which cell assigned to which source distribution + split_idx_to_covariates: dict[int, tuple[Any, ...]] # (n_sources,) dictionary explaining split_covariates_mask + perturbation_covariates_mask: np.ndarray # (n_cells,), which cell assigned to which target distribution + perturbation_idx_to_covariates: dict[ + int, tuple[str, ...] + ] # (n_targets,), dictionary explaining perturbation_covariates_mask + perturbation_idx_to_id: dict[int, Any] + condition_data: dict[str, np.ndarray] # (n_targets,) all embeddings for conditions + control_to_perturbation: dict[int, np.ndarray] # mapping from control idx to target distribution idcs + max_combination_length: int + null_value: Any + data_manager: Any + n_conditions_on_log_iteration: int | None = None + n_conditions_on_train_end: int | None = None + + +class TrainSampler: + """Data sampler for :class:`~cellflow.data.TrainingData`. + + Parameters + ---------- + data + The training data. + batch_size + The batch size. + + """ + + def __init__(self, data: TrainingData, batch_size: int = 1024): + self._data = data + self._data_idcs = np.arange(data.cell_data.shape[0]) + self.batch_size = batch_size + self.n_source_dists = data.n_controls + self.n_target_dists = data.n_perturbations + + self._control_to_perturbation_keys = sorted(data.control_to_perturbation.keys()) + self._has_condition_data = data.condition_data is not None + + def _sample_target_dist_idx(self, source_dist_idx, rng): + """Sample a target distribution index given the source distribution index.""" + return rng.choice(self._data.control_to_perturbation[source_dist_idx]) + + def _get_embeddings(self, idx, condition_data) -> dict[str, np.ndarray]: + """Get embeddings for a given index.""" + result = {} + for key, arr in condition_data.items(): + result[key] = np.expand_dims(arr[idx], 0) + return result + + def _sample_from_mask(self, rng, mask) -> np.ndarray: + """Sample indices according to a mask.""" + # Convert mask to probability distribution + valid_indices = np.where(mask)[0] + + # Handle case with no valid indices (should not happen in practice) + if len(valid_indices) == 0: + raise ValueError("No valid indices found in the mask") + + # Sample from valid indices with equal probability + batch_idcs = rng.choice(valid_indices, self.batch_size, replace=True) + return batch_idcs + + def sample(self, rng) -> dict[str, Any]: + """Sample a batch of data. + + Parameters + ---------- + seed : int, optional + Random seed + + Returns + ------- + Dictionary with source and target data + """ + # Sample source distribution index + source_dist_idx = rng.integers(0, self.n_source_dists) + + # Get source cells + source_cells_mask = self._data.split_covariates_mask == source_dist_idx + source_batch_idcs = self._sample_from_mask(rng, source_cells_mask) + source_batch = self._data.cell_data[source_batch_idcs] + + target_dist_idx = self._sample_target_dist_idx(source_dist_idx, rng) + target_cells_mask = self._data.perturbation_covariates_mask == target_dist_idx + target_batch_idcs = self._sample_from_mask(rng, target_cells_mask) + target_batch = self._data.cell_data[target_batch_idcs] + + if not self._has_condition_data: + return {"src_cell_data": source_batch, "tgt_cell_data": target_batch} + else: + condition_batch = self._get_embeddings(target_dist_idx, self._data.condition_data) + return { + "src_cell_data": source_batch, + "tgt_cell_data": target_batch, + "condition": condition_batch, + } + + @property + def data(self): + """The training data.""" + return self._data + + +class BaseValidSampler(abc.ABC): + @abc.abstractmethod + def sample(*args, **kwargs): + pass + + def _get_key(self, cond_idx: int) -> tuple[str, ...]: + if len(self._data.perturbation_idx_to_id): # type: ignore[attr-defined] + return self._data.perturbation_idx_to_id[cond_idx] # type: ignore[attr-defined] + cov_combination = self._data.perturbation_idx_to_covariates[cond_idx] # type: ignore[attr-defined] + return tuple(cov_combination[i] for i in range(len(cov_combination))) + + def _get_perturbation_to_control(self, data: ValidationData | PredictionData) -> dict[int, np.ndarray]: + d = {} + for k, v in data.control_to_perturbation.items(): + for el in v: + d[el] = k + return d + + def _get_condition_data(self, cond_idx: int) -> dict[str, np.ndarray]: + return {k: v[[cond_idx], ...] for k, v in self._data.condition_data.items()} # type: ignore[attr-defined] + + +class ValidationSampler(BaseValidSampler): + """Data sampler for :class:`~cellflow.data.ValidationData`. + + Parameters + ---------- + val_data + The validation data. + seed + Random seed. + """ + + def __init__(self, val_data: ValidationData, seed: int = 0) -> None: + self._data = val_data + self.perturbation_to_control = self._get_perturbation_to_control(val_data) + self.n_conditions_on_log_iteration = ( + val_data.n_conditions_on_log_iteration + if val_data.n_conditions_on_log_iteration is not None + else val_data.n_perturbations + ) + self.n_conditions_on_train_end = ( + val_data.n_conditions_on_train_end + if val_data.n_conditions_on_train_end is not None + else val_data.n_perturbations + ) + self.rng = np.random.default_rng(seed) + if self._data.condition_data is None: + raise NotImplementedError("Validation data must have condition data.") + + def sample(self, mode: Literal["on_log_iteration", "on_train_end"]) -> Any: + """Sample data for validation. + + Parameters + ---------- + mode + Sampling mode. Either ``"on_log_iteration"`` or ``"on_train_end"``. + + Returns + ------- + Dictionary with source, condition, and target data from the validation data. + """ + size = self.n_conditions_on_log_iteration if mode == "on_log_iteration" else self.n_conditions_on_train_end + condition_idcs = self.rng.choice(self._data.n_perturbations, size=(size,), replace=False) + + source_idcs = [self.perturbation_to_control[cond_idx] for cond_idx in condition_idcs] + source_cells_mask = [self._data.split_covariates_mask == source_idx for source_idx in source_idcs] + source_cells = [self._data.cell_data[mask] for mask in source_cells_mask] + target_cells_mask = [cond_idx == self._data.perturbation_covariates_mask for cond_idx in condition_idcs] + target_cells = [self._data.cell_data[mask] for mask in target_cells_mask] + conditions = [self._get_condition_data(cond_idx) for cond_idx in condition_idcs] + cell_rep_dict = {} + cond_dict = {} + true_dict = {} + for i in range(len(condition_idcs)): + k = self._get_key(condition_idcs[i]) + cell_rep_dict[k] = source_cells[i] + cond_dict[k] = conditions[i] + true_dict[k] = target_cells[i] + + return {"source": cell_rep_dict, "condition": cond_dict, "target": true_dict} + + @property + def data(self) -> ValidationData: + """The validation data.""" + return self._data + + +class PredictionSampler(BaseValidSampler): + """Data sampler for :class:`~cellflow.data.PredictionData`. + + Parameters + ---------- + pred_data + The prediction data. + + """ + + def __init__(self, pred_data: PredictionData) -> None: + self._data = pred_data + self.perturbation_to_control = self._get_perturbation_to_control(pred_data) + if self._data.condition_data is None: + raise NotImplementedError("Validation data must have condition data.") + + def sample(self) -> Any: + """Sample data for prediction. + + Returns + ------- + Dictionary with source and condition data from the prediction data. + """ + condition_idcs = range(self._data.n_perturbations) + + source_idcs = [self.perturbation_to_control[cond_idx] for cond_idx in condition_idcs] + source_cells_mask = [self._data.split_covariates_mask == source_idx for source_idx in source_idcs] + source_cells = [self._data.cell_data[mask] for mask in source_cells_mask] + conditions = [self._get_condition_data(cond_idx) for cond_idx in condition_idcs] + cell_rep_dict = {} + cond_dict = {} + for i in range(len(condition_idcs)): + k = self._get_key(condition_idcs[i]) + cell_rep_dict[k] = source_cells[i] + cond_dict[k] = conditions[i] + + return { + "source": cell_rep_dict, + "condition": cond_dict, + } + + @property + def data(self) -> PredictionData: + """The training data.""" + return self._data + + +def prefetch_to_device( + sampler: TrainSampler, seed: int, num_iterations: int, prefetch_factor: int = 2, num_workers: int = 4 +) -> Generator[dict[str, Any], None, None]: + seq = np.random.SeedSequence(seed) + random_generators = [np.random.default_rng(s) for s in seq.spawn(num_workers)] + + q: queue.Queue[dict[str, Any]] = queue.Queue(maxsize=prefetch_factor * num_workers) + sem = threading.Semaphore(num_iterations) + stop_event = threading.Event() + + def worker(rng: np.random.Generator): + while not stop_event.is_set() and sem.acquire(blocking=False): + batch = sampler.sample(rng) + batch = jax.device_put(batch, jax.devices()[0], donate=True) + jax.block_until_ready(batch) + while not stop_event.is_set(): + try: + q.put(batch, timeout=1.0) + break # Batch successfully put into the queue; break out of retry loop + except queue.Full: + continue + + return + + # Start multiple worker threads + ts = [] + for i in range(num_workers): + t = threading.Thread(target=worker, daemon=True, name=f"worker-{i}", args=(random_generators[i],)) + t.start() + ts.append(t) + + try: + for _ in range(num_iterations): + # Yield batches from the queue; will block waiting for available batch + yield q.get() + finally: + # When the generator is closed or garbage collected, clean up the worker threads + stop_event.set() # Signal all workers to exit + for t in ts: + t.join() # Wait for all worker threads to finish + + +class OOCTrainSampler: + def __init__( + self, data: TrainingData, seed: int, batch_size: int = 1024, num_workers: int = 4, prefetch_factor: int = 2 + ): + self.inner = TrainSampler(data=data, batch_size=batch_size) + self.num_workers = num_workers + self.prefetch_factor = prefetch_factor + self.seed = seed + self._iterator = None + + def set_sampler(self, num_iterations: int) -> None: + self._iterator = prefetch_to_device( + sampler=self.inner, seed=self.seed, num_iterations=num_iterations, prefetch_factor=self.prefetch_factor + ) + + def sample(self, rng=None) -> dict[str, Any]: + if self._iterator is None: + raise ValueError( + "Sampler not set. Use `set_sampler` to set the sampler with" + "the number of iterations. Without the number of iterations," + " the sampler will not be able to sample the data." + ) + if rng is not None: + del rng + return next(self._iterator) diff --git a/src/cellflow/model/__init__.py b/src/cellflow/model/__init__.py index 8731f241..8c912082 100644 --- a/src/cellflow/model/__init__.py +++ b/src/cellflow/model/__init__.py @@ -1,3 +1,4 @@ from cellflow.model._cellflow import CellFlow +from cellflow.model._cellflow_annbatch import CellFlowAnnbatch -__all__ = ["CellFlow"] +__all__ = ["CellFlow", "CellFlowAnnbatch"] diff --git a/src/cellflow/model/_base.py b/src/cellflow/model/_base.py new file mode 100644 index 00000000..463f7051 --- /dev/null +++ b/src/cellflow/model/_base.py @@ -0,0 +1,668 @@ +"""Shared, data-path-agnostic CellFlow base: model setup, training, prediction, save/load.""" + +import abc +import functools +import os +import types +import warnings +from collections.abc import Callable, Sequence +from dataclasses import field as dc_field +from typing import Any, Literal + +import anndata as ad +import cloudpickle +import flax.linen as nn +import jax +import jax.numpy as jnp +import numpy as np +import optax +import pandas as pd + +from cellflow import _constants +from cellflow._compat import BrownianBridge, ConstantNoiseFlow +from cellflow._types import ArrayLike, Layers_separate_input_t, Layers_t +from cellflow.data._data import ConditionData +from cellflow.data._datamanager import DataManager +from cellflow.data._legacy import PredictionSampler, ValidationData +from cellflow.model._utils import _write_predictions +from cellflow.networks import _velocity_field +from cellflow.plotting import _utils +from cellflow.solvers import SOLVER_REGISTRY, _genot, _otfm +from cellflow.training._callbacks import BaseCallback +from cellflow.training._trainer import CellFlowTrainer +from cellflow.utils import match_linear + +__all__ = ["BaseCellFlow"] + + +class BaseCellFlow(abc.ABC): + """Base class holding everything independent of how training cells are sourced. + + Concrete subclasses (:class:`~cellflow.model.CellFlow` in-memory, ``CellFlowAnnbatch`` streaming) + implement the data-preparation methods and the path-specific hooks. + """ + + def __init__(self, solver: Literal["otfm", "genot"] = "otfm"): + if solver not in SOLVER_REGISTRY: + raise ValueError(f"Unknown solver {solver!r}. Registered solvers: {sorted(SOLVER_REGISTRY)}.") + self._solver_class, self._vf_class = SOLVER_REGISTRY[solver] + self._dm: DataManager | None = None # set by a subclass `prepare_*` (both paths) + self._data_dim: int | None = None + self._dataloader: Any | None = None + self._trainer: CellFlowTrainer | None = None + self._validation_data: dict[str, Any] = {"predict_kwargs": {}} + self._solver: _otfm.OTFlowMatching | _genot.GENOT | None = None + self._condition_dim: int | None = None + self._vf: _velocity_field.ConditionalVelocityField | _velocity_field.GENOTConditionalVelocityField | None = None + + def prepare_model( + self, + condition_mode: Literal["deterministic", "stochastic"] = "deterministic", + regularization: float = 0.0, + pooling: Literal["mean", "attention_token", "attention_seed"] = "attention_token", + pooling_kwargs: dict[str, Any] = types.MappingProxyType({}), + layers_before_pool: Layers_separate_input_t | Layers_t = dc_field(default_factory=lambda: []), + layers_after_pool: Layers_t = dc_field(default_factory=lambda: []), + condition_embedding_dim: int = 256, + cond_output_dropout: float = 0.9, + condition_encoder_kwargs: dict[str, Any] | None = None, + pool_sample_covariates: bool = True, + time_freqs: int = 1024, + time_max_period: int | None = 10000, + time_encoder_dims: Sequence[int] = (2048, 2048, 2048), + time_encoder_dropout: float = 0.0, + hidden_dims: Sequence[int] = (2048, 2048, 2048), + hidden_dropout: float = 0.0, + conditioning: Literal["concatenation", "film", "resnet"] = "concatenation", + conditioning_kwargs: dict[str, Any] = dc_field(default_factory=lambda: {}), + decoder_dims: Sequence[int] = (4096, 4096, 4096), + decoder_dropout: float = 0.0, + vf_act_fn: Callable[[jnp.ndarray], jnp.ndarray] = nn.silu, + vf_kwargs: dict[str, Any] | None = None, + probability_path: dict[Literal["constant_noise", "bridge"], float] | None = None, + match_fn: Callable[[ArrayLike, ArrayLike], ArrayLike] = match_linear, + optimizer: optax.GradientTransformation = optax.MultiSteps(optax.adam(5e-5), 20), + solver_kwargs: dict[str, Any] | None = None, + layer_norm_before_concatenation: bool = False, + linear_projection_before_concatenation: bool = False, + seed=0, + ) -> None: + """Prepare the model for training. + + This function sets up the neural network architecture and specificities of the + :attr:`solver`. When :attr:`solver` is an instance of :class:`cellflow.solvers._genot.GENOT`, + the following arguments have to be passed to ``'condition_encoder_kwargs'``: + + + Parameters + ---------- + condition_mode + Mode of the encoder, should be one of: + + - ``'deterministic'``: Learns condition encoding point-wise. + - ``'stochastic'``: Learns a Gaussian distribution for representing conditions. + + regularization + Regularization strength in the latent space: + + - For deterministic mode, it is the strength of the L2 regularization. + - For stochastic mode, it is the strength of the VAE regularization. + + pooling + Pooling method, should be one of: + + - ``'mean'``: Aggregates combinations of covariates by the mean of their + learned embeddings. + - ``'attention_token'``: Aggregates combinations of covariates by an attention + mechanism with a class token. + - ``'attention_seed'``: Aggregates combinations of covariates by seed attention. + + pooling_kwargs + Keyword arguments for the pooling method corresponding to: + + - :class:`cellflow.networks.TokenAttentionPooling` if ``'pooling'`` is + ``'attention_token'``. + - :class:`cellflow.networks.SeedAttentionPooling` if ``'pooling'`` is ``'attention_seed'``. + + layers_before_pool + Layers applied to the condition embeddings before pooling. Can be of type + + - :class:`tuple` with elements corresponding to dictionaries with keys: + + - ``'layer_type'`` of type :class:`str` indicating the type of the layer, can be + ``'mlp'`` or ``'self_attention'``. + - Further keyword arguments for the layer type :class:`cellflow.networks.MLPBlock` or + :class:`cellflow.networks.SelfAttentionBlock`. + + - :class:`dict` with keys corresponding to perturbation covariate keys, and values + correspondinng to the above mentioned tuples. + + layers_after_pool + Layers applied to the condition embeddings after pooling, and before applying the last + layer of size ``'condition_embedding_dim'``. Should be of type :class:`tuple` with + elements corresponding to dictionaries with keys: + + - ``'layer_type'`` of type :class:`str` indicating the type of the layer, can be + ``'mlp'`` or ``'self_attention'``. + - Further keys depend on the layer type, either for :class:`cellflow.networks.MLPBlock` or + for :class:`cellflow.networks.SelfAttentionBlock`. + + condition_embedding_dim + Dimensions of the condition embedding, i.e. the last layer of the + :class:`cellflow.networks.ConditionEncoder`. + cond_output_dropout + Dropout rate for the last layer of the :class:`cellflow.networks.ConditionEncoder`. + condition_encoder_kwargs + Keyword arguments for the :class:`cellflow.networks.ConditionEncoder`. + pool_sample_covariates + Whether to include sample covariates in the pooling. + time_freqs + Frequency of the sinusoidal time encoding + (:func:`ott.neural.networks.layers.sinusoidal_time_encoder`). + time_max_period + Controls the frequency of the time embeddings, see + :func:`cellflow.networks.utils.sinusoidal_time_encoder`. + time_encoder_dims + Dimensions of the layers processing the time embedding in + :attr:`cellflow.networks.ConditionalVelocityField.time_encoder`. + time_encoder_dropout + Dropout rate for the :attr:`cellflow.networks.ConditionalVelocityField.time_encoder`. + hidden_dims + Dimensions of the layers processing the input to the velocity field + via :attr:`cellflow.networks.ConditionalVelocityField.x_encoder`. + hidden_dropout + Dropout rate for :attr:`cellflow.networks.ConditionalVelocityField.x_encoder`. + conditioning + Conditioning method, should be one of: + + - ``'concatenation'``: Concatenate the time, data, and condition embeddings. + - ``'film'``: Use FiLM conditioning, i.e. learn FiLM weights from time and condition embedding + to scale the data embeddings. + - ``'resnet'``: Use residual conditioning. + + conditioning_kwargs + Keyword arguments for the conditioning method. + decoder_dims + Dimensions of the output layers in + :attr:`cellflow.networks.ConditionalVelocityField.decoder`. + decoder_dropout + Dropout rate for the output layer + :attr:`cellflow.networks.ConditionalVelocityField.decoder`. + vf_act_fn + Activation function of the :class:`cellflow.networks.ConditionalVelocityField`. + vf_kwargs + Additional keyword arguments for the solver-specific vector field. + For instance, when ``'solver==genot'``, the following keyword argument can be passed: + + - ``'genot_source_dims'`` of type :class:`tuple` with the dimensions + of the :class:`cellflow.networks.MLPBlock` processing the source cell. + - ``'genot_source_dropout'`` of type :class:`float` indicating the dropout rate + for the source cell processing. + probability_path + Probability path to use for training. Should be a :class:`dict` of the form + + - ``'{"constant_noise": noise_val'`` + - ``'{"bridge": noise_val}'`` + + If :obj:`None`, defaults to ``'{"constant_noise": 0.0}'``. + match_fn + Matching function between unperturbed and perturbed cells. Should take as input source + and target data and return the optimal transport matrix, see e.g. + :func:`cellflow.utils.match_linear`. + optimizer + Optimizer used for training. + solver_kwargs + Keyword arguments for the solver :class:`cellflow.solvers.OTFlowMatching` or + :class:`cellflow.solvers.GENOT`. + layer_norm_before_concatenation + If :obj:`True`, applies layer normalization before concatenating + the embedded time, embedded data, and condition embeddings. + linear_projection_before_concatenation + If :obj:`True`, applies a linear projection before concatenating + the embedded time, embedded data, and embedded condition. + seed + Random seed. + + Returns + ------- + Updates the following fields: + + - :attr:`cellflow.model.CellFlow.velocity_field` - an instance of the + :class:`cellflow.networks.ConditionalVelocityField`. + - :attr:`cellflow.model.CellFlow.solver` - an instance of :class:`cellflow.solvers.OTFlowMatching` + or :class:`cellflow.solvers.GENOT`. + - :attr:`cellflow.model.CellFlow.trainer` - an instance of the + :class:`cellflow.training.CellFlowTrainer`. + """ + # Condition embeddings + max combination length are path-specific (in-memory vs streaming). + condition_data, max_combination_length = self._encoder_conditions() + + if condition_mode == "stochastic": + if regularization == 0.0: + raise ValueError("Stochastic condition embeddings require `regularization`>0.") + + condition_encoder_kwargs = condition_encoder_kwargs or {} + # Each velocity field owns which solver-specific `vf_kwargs` it accepts (validated/defaulted here). + vf_kwargs = self._vf_class._normalize_vf_kwargs(vf_kwargs) + covariates_not_pooled = [] if pool_sample_covariates else self._dm.sample_covariates + solver_kwargs = solver_kwargs or {} + probability_path = probability_path or {"constant_noise": 0.0} + + self.vf = self._vf_class( + output_dim=self._data_dim, + max_combination_length=max_combination_length, + condition_mode=condition_mode, + regularization=regularization, + condition_embedding_dim=condition_embedding_dim, + covariates_not_pooled=covariates_not_pooled, + pooling=pooling, + pooling_kwargs=pooling_kwargs, + layers_before_pool=layers_before_pool, + layers_after_pool=layers_after_pool, + cond_output_dropout=cond_output_dropout, + condition_encoder_kwargs=condition_encoder_kwargs, + act_fn=vf_act_fn, + time_freqs=time_freqs, + time_max_period=time_max_period, + time_encoder_dims=time_encoder_dims, + time_encoder_dropout=time_encoder_dropout, + hidden_dims=hidden_dims, + hidden_dropout=hidden_dropout, + conditioning=conditioning, + conditioning_kwargs=conditioning_kwargs, + decoder_dims=decoder_dims, + decoder_dropout=decoder_dropout, + layer_norm_before_concatenation=layer_norm_before_concatenation, + linear_projection_before_concatenation=linear_projection_before_concatenation, + **vf_kwargs, + ) + + probability_path, noise = next(iter(probability_path.items())) + if probability_path == "constant_noise": + probability_path = ConstantNoiseFlow(noise) + elif probability_path == "bridge": + probability_path = BrownianBridge(noise) + else: + raise NotImplementedError( + f"The key of `probability_path` must be `'constant_noise'` or `'bridge'` but found {probability_path}." + ) + + # Each solver owns how it names its match function / whether it needs source-target dims. + self._solver = self._solver_class( + vf=self.vf, + probability_path=probability_path, + optimizer=optimizer, + conditions=condition_data, + rng=jax.random.PRNGKey(seed), + **self._solver_class._match_kwargs(match_fn=match_fn, data_dim=self._data_dim), + **solver_kwargs, + ) + + self._trainer = CellFlowTrainer(solver=self.solver, predict_kwargs=self.validation_data["predict_kwargs"]) # type: ignore[arg-type] + + def train( + self, + num_iterations: int, + batch_size: int = 1024, + valid_freq: int = 1000, + callbacks: Sequence[BaseCallback] = [], + monitor_metrics: Sequence[str] = [], + out_of_core_dataloading: bool = False, + ) -> None: + """Train the model. + + Note + ---- + A low value of ``'valid_freq'`` results in long training + because predictions are time-consuming compared to training steps. + + Parameters + ---------- + num_iterations + Number of iterations to train the model. + batch_size + Batch size. + valid_freq + Frequency of validation. + callbacks + Callbacks to perform at each validation step. There are two types of callbacks: + - Callbacks for computations should inherit from + :class:`~cellflow.training.ComputationCallback` see e.g. :class:`cellflow.training.Metrics`. + - Callbacks for logging should inherit from :class:`~cellflow.training.LoggingCallback` see + e.g. :class:`~cellflow.training.WandbLogger`. + monitor_metrics + Metrics to monitor. + out_of_core_dataloading + If :obj:`True`, use out-of-core dataloading. Uses the :class:`cellflow.data._legacy.OOCTrainSampler` + to load data that does not fit into GPU memory. + + Returns + ------- + Updates the following fields: + + - :attr:`cellflow.model.CellFlow.dataloader` - the training dataloader. + - :attr:`cellflow.model.CellFlow.solver` - the trained solver. + """ + if self._dm is None: + raise ValueError("Data not initialized. Please call `prepare_data` first.") + + if self.trainer is None: + raise ValueError("Model not initialized. Please call `prepare_model` first.") + + self._bind_train_dataloader(batch_size, out_of_core_dataloading) # in-memory binds; streaming is a no-op + validation_loaders = self._build_validation_loaders() + self._trainer.predict_kwargs = self.validation_data.get("predict_kwargs", {}) + + self._solver = self.trainer.train( + dataloader=self._dataloader, + num_iterations=num_iterations, + valid_freq=valid_freq, + valid_loaders=validation_loaders, + callbacks=callbacks, + monitor_metrics=monitor_metrics, + ) + + def predict( + self, + adata: ad.AnnData, + covariate_data: pd.DataFrame, + sample_rep: str | None = None, + condition_id_key: str | None = None, + key_added_prefix: str | None = None, + rng: ArrayLike | None = None, + **kwargs: Any, + ) -> dict[str, ArrayLike] | None: + """Predict perturbation responses. + + Parameters + ---------- + adata + An :class:`~anndata.AnnData` object with the source representation. + covariate_data + Covariate data defining the condition to predict. This :class:`~pandas.DataFrame` + should have the same columns as :attr:`~anndata.AnnData.obs` of + :attr:`cellflow.model.CellFlow.adata`, and as registered in + :attr:`cellflow.model.CellFlow.data_manager`. + sample_rep + Key in :attr:`~anndata.AnnData.obsm` where the sample representation is stored or + ``'X'`` to use :attr:`~anndata.AnnData.X`. If :obj:`None`, the key is assumed to be + the same as for the training data. + condition_id_key + Key in ``'covariate_data'`` defining the condition name. + key_added_prefix + If not :obj:`None`, prefix to store the prediction in :attr:`~anndata.AnnData.obsm`. + If :obj:`None`, the predictions are not stored, and the predictions are returned as a + :class:`dict`. + rng + Random number generator. If :obj:`None` and :attr:`cellflow.model.CellFlow.conditino_mode` + is ``'stochastic'``, the condition vector will be the mean of the learnt distributions, + otherwise samples from the distribution. + kwargs + Keyword arguments for the predict function, i.e. + :meth:`cellflow.solvers.OTFlowMatching.predict` or :meth:`cellflow.solvers.GENOT.predict`. + + Returns + ------- + If ``'key_added_prefix'`` is :obj:`None`, a :class:`dict` with the predicted sample + representation for each perturbation, otherwise stores the predictions in + :attr:`~anndata.AnnData.obsm` and returns :obj:`None`. + """ + if self.solver is None or not self.solver.is_trained: + raise ValueError("Model not trained. Please call `train` first.") + + if sample_rep is None: + sample_rep = self._dm.sample_rep + + if adata is not None and covariate_data is not None: + if covariate_data.empty: + raise ValueError("`covariate_data` is empty.") + if self._dm.control_key not in adata.obs.columns: + raise ValueError( + f"If both `adata` and `covariate_data` are given, the control key `{self._dm.control_key}` must be in `adata.obs`." + ) + if not adata.obs[self._dm.control_key].all(): + raise ValueError( + f"If both `adata` and `covariate_data` are given, all samples in `adata` must be control samples, and thus `adata.obs[`{self._dm.control_key}`] must be set to `True` everywhere." + ) + pred_data = self._dm.get_prediction_data( + adata, + sample_rep=sample_rep, # type: ignore[arg-type] + covariate_data=covariate_data, + condition_id_key=condition_id_key, + ) + pred_loader = PredictionSampler(pred_data) + batch = pred_loader.sample() + src = batch["source"] + condition = batch.get("condition", None) + # using jax.tree.map to batch the prediction + # because PredictionSampler can return a different number of cells for each condition + out = jax.tree.map( + functools.partial(self.solver.predict, rng=rng, **kwargs), + src, + condition, # type: ignore[attr-defined] + ) + if key_added_prefix is None: + return out + if len(pred_data.control_to_perturbation) > 1: + raise ValueError( + f"When saving predictions to `adata`, all control cells must be from the same control \ + population, but found {len(pred_data.control_to_perturbation)} control populations." + ) + out_np = {k: np.array(v) for k, v in out.items()} + _write_predictions( + adata=adata, + predictions=out_np, + key_added_prefix=key_added_prefix, + ) + + def get_condition_embedding( + self, + covariate_data: pd.DataFrame | ConditionData, + rep_dict: dict[str, str] | None = None, + condition_id_key: str | None = None, + key_added: str | None = _constants.CONDITION_EMBEDDING, + ) -> tuple[pd.DataFrame, pd.DataFrame]: + """Get the embedding of the conditions. + + Outputs the mean and variance of the learnt embeddings + generated by the :class:`~cellflow.networks.ConditionEncoder`. + + Parameters + ---------- + covariate_data + Can be one of + + - a :class:`~pandas.DataFrame` defining the conditions with the same columns as the + :class:`~anndata.AnnData` used for the initialisation of :class:`~cellflow.model.CellFlow`. + - an instance of :class:`~cellflow.data.ConditionData`. + + rep_dict + Dictionary containing the representations of the perturbation covariates. Will be considered an + empty dictionary if :obj:`None`. + condition_id_key + Key defining the name of the condition. Only available + if ``'covariate_data'`` is a :class:`~pandas.DataFrame`. + key_added + Key to store the condition embedding in :attr:`~anndata.AnnData.uns`. The mean is + stored under ``key_added`` and the variance under ``f"{key_added}_var"``. If + :obj:`None`, the embeddings are not stored. + + Returns + ------- + A :class:`tuple` of :class:`~pandas.DataFrame` with the mean and variance of the condition embeddings. + """ + if self.solver is None or not self.solver.is_trained: + raise ValueError("Model not trained. Please call `train` first.") + + if hasattr(covariate_data, "condition_data"): + cond_data = covariate_data + elif isinstance(covariate_data, pd.DataFrame): + cond_data = self._dm.get_condition_data( + covariate_data=covariate_data, + rep_dict=rep_dict, + condition_id_key=condition_id_key, + ) + else: + raise ValueError("Covariate data must be a `pandas.DataFrame` or an instance of `BaseData`.") + + condition_embeddings_mean: dict[str, ArrayLike] = {} + condition_embeddings_var: dict[str, ArrayLike] = {} + n_conditions = len(next(iter(cond_data.condition_data.values()))) + for i in range(n_conditions): + condition = {k: v[[i], :] for k, v in cond_data.condition_data.items()} + if condition_id_key: + c_key = cond_data.perturbation_idx_to_id[i] + else: + cov_combination = cond_data.perturbation_idx_to_covariates[i] + c_key = tuple(cov_combination[i] for i in range(len(cov_combination))) + condition_embeddings_mean[c_key], condition_embeddings_var[c_key] = self.solver.get_condition_embedding( + condition + ) + + df_mean = pd.DataFrame.from_dict({k: v[0] for k, v in condition_embeddings_mean.items()}).T + df_var = pd.DataFrame.from_dict({k: v[0] for k, v in condition_embeddings_var.items()}).T + + if condition_id_key: + df_mean.index.set_names([condition_id_key], inplace=True) + df_var.index.set_names([condition_id_key], inplace=True) + else: + df_mean.index.set_names(list(self._dm.perturb_covar_keys), inplace=True) + df_var.index.set_names(list(self._dm.perturb_covar_keys), inplace=True) + + if key_added is not None: + if self.adata is None: # streaming/annbatch path: no `adata` to store into + warnings.warn( + "No `adata` is attached (streaming path); returning the condition embeddings without " + "storing them under `adata.uns`. Pass `key_added=None` to silence this.", + stacklevel=2, + ) + else: # mean under `key_added`, variance under `f"{key_added}_var"` (distinct keys, #295) + _utils.set_plotting_vars(self.adata, key=key_added, value=df_mean) + _utils.set_plotting_vars(self.adata, key=f"{key_added}_var", value=df_var) + + return df_mean, df_var + + def save( + self, + dir_path: str, + file_prefix: str | None = None, + overwrite: bool = False, + ) -> None: + """ + Save the model. + + Pickles the :class:`~cellflow.model.CellFlow` object. + + Parameters + ---------- + dir_path + Path to a directory, defaults to current directory + file_prefix + Prefix to prepend to the file name. + overwrite + Overwrite existing data or not. + + Returns + ------- + :obj:`None` + """ + file_name = ( + f"{file_prefix}_{self.__class__.__name__}.pkl" + if file_prefix is not None + else f"{self.__class__.__name__}.pkl" + ) + file_dir = os.path.join(dir_path, file_name) if dir_path is not None else file_name + + if not overwrite and os.path.exists(file_dir): + raise RuntimeError(f"Unable to save to an existing file `{file_dir}` use `overwrite=True` to overwrite it.") + with open(file_dir, "wb") as f: + cloudpickle.dump(self, f) + + @classmethod + def load( + cls, + filename: str, + ) -> "BaseCellFlow": + """ + Load a :class:`~cellflow.model.CellFlow` model from a saved instance. + + Parameters + ---------- + filename + Path to the saved file. + + Returns + ------- + Loaded instance of the model. + """ + # Check if filename is a directory + file_name = os.path.join(filename, f"{cls.__name__}.pkl") if os.path.isdir(filename) else filename + + with open(file_name, "rb") as f: + model = cloudpickle.load(f) + + if type(model) is not cls: + raise TypeError(f"Expected the model to be type of `{cls}`, found `{type(model)}`.") + return model + + @property + def adata(self) -> ad.AnnData | None: + """The :class:`~anndata.AnnData` used for training, or :obj:`None` in the streaming path.""" + return getattr(self, "_adata", None) + + @property + def solver(self) -> _otfm.OTFlowMatching | _genot.GENOT | None: + """The solver.""" + return self._solver + + @property + def dataloader(self) -> Any | None: + """The dataloader used for training (in-memory sampler or a streaming adapter).""" + return self._dataloader + + @property + def trainer(self) -> CellFlowTrainer | None: + """The trainer used for training.""" + return self._trainer + + @property + def validation_data(self) -> dict[str, ValidationData]: + """The validation data.""" + return self._validation_data + + @property + def data_manager(self) -> DataManager: + """The data manager, initialised with :attr:`cellflow.model.CellFlow.adata`.""" + return self._dm + + @property + def velocity_field( + self, + ) -> _velocity_field.ConditionalVelocityField | _velocity_field.GENOTConditionalVelocityField | None: + """The conditional velocity field.""" + return self._vf + + @velocity_field.setter # type: ignore[attr-defined,no-redef] + def velocity_field(self, vf: _velocity_field.ConditionalVelocityField) -> None: + """Set the velocity field.""" + if not isinstance(vf, _velocity_field.ConditionalVelocityField): + raise ValueError(f"Expected `vf` to be an instance of `ConditionalVelocityField`, found `{type(vf)}`.") + self._vf = vf + + @property + def condition_mode(self) -> Literal["deterministic", "stochastic"]: + """The mode of the encoder.""" + return self.velocity_field.condition_mode + + # ── path-specific hooks (implemented by CellFlow / CellFlowAnnbatch) ────────────────────────── + @abc.abstractmethod + def _encoder_conditions(self) -> tuple[dict[str, np.ndarray], int]: + """The condition embeddings + ``max_combination_length`` used by :meth:`prepare_model`.""" + + @abc.abstractmethod + def _bind_train_dataloader(self, batch_size: int, out_of_core_dataloading: bool) -> None: + """Bind :attr:`_dataloader` for :meth:`train` (in-memory builds one; streaming already has it).""" + + @abc.abstractmethod + def _build_validation_loaders(self) -> dict[str, Any]: + """The validation samplers (keyed by name) that :meth:`train` feeds to the trainer.""" diff --git a/src/cellflow/model/_cellflow.py b/src/cellflow/model/_cellflow.py index f1596077..9bf047c3 100644 --- a/src/cellflow/model/_cellflow.py +++ b/src/cellflow/model/_cellflow.py @@ -1,64 +1,42 @@ -import functools -import os -import types -from collections.abc import Callable, Sequence -from dataclasses import field as dc_field +"""In-memory (legacy) CellFlow: training data materialized from an ``AnnData`` via :meth:`prepare_data`. + +The streaming path (:class:`~cellflow.model.CellFlowAnnbatch`) is the default; this in-memory path is +kept for backward compatibility. +""" + +import warnings +from collections.abc import Sequence from typing import Any, Literal import anndata as ad -import cloudpickle -import flax.linen as nn -import jax -import jax.numpy as jnp import numpy as np -import optax -import pandas as pd - -from cellflow import _constants -from cellflow._compat import BrownianBridge, ConstantNoiseFlow -from cellflow._types import ArrayLike, Layers_separate_input_t, Layers_t -from cellflow.data._data import ConditionData, TrainingData, ValidationData -from cellflow.data._dataloader import OOCTrainSampler, PredictionSampler, TrainSampler, ValidationSampler + from cellflow.data._datamanager import DataManager -from cellflow.model._utils import _write_predictions -from cellflow.networks import _velocity_field -from cellflow.plotting import _utils -from cellflow.solvers import SOLVER_REGISTRY, _genot, _otfm -from cellflow.training._callbacks import BaseCallback -from cellflow.training._trainer import CellFlowTrainer -from cellflow.utils import match_linear +from cellflow.data._legacy import OOCTrainSampler, TrainingData, TrainSampler, ValidationSampler +from cellflow.model._base import BaseCellFlow __all__ = ["CellFlow"] -class CellFlow: - """CellFlow model for perturbation prediction using Flow Matching and Optimal Transport. - - CellFlow builds upon neural optimal transport estimators extending :cite:`tong:23`, - :cite:`pooladian:23`, :cite:`eyring:24`, :cite:`klein:23` which are all based on - Flow Matching :cite:`lipman:22`. +class CellFlow(BaseCellFlow): + """CellFlow with in-memory training data extracted from an :class:`~anndata.AnnData`. - Parameters - ---------- - adata - An :class:`~anndata.AnnData` object to extract the training data from. - solver - Solver to use for training. Any name registered in - :data:`cellflow.solvers.SOLVER_REGISTRY` (``'otfm'`` or ``'genot'`` by default, extendable - via :func:`cellflow.solvers.register_solver`). + Prepare training data with :meth:`prepare_data`; model setup, training, prediction and IO are + inherited from :class:`~cellflow.model._base.BaseCellFlow`. For large or out-of-core datasets use + :class:`~cellflow.model.CellFlowAnnbatch` (the default streaming path). """ - def __init__(self, adata: ad.AnnData, solver: str = "otfm"): + def __init__(self, adata: ad.AnnData | None = None, solver: Literal["otfm", "genot"] = "otfm"): + super().__init__(solver) + if adata is not None: + warnings.warn( + "Passing `adata` to `CellFlow(...)` is deprecated and will be removed in a future " + "release; pass it to `prepare_data(adata=...)` instead.", + FutureWarning, + stacklevel=2, + ) self._adata = adata - if solver not in SOLVER_REGISTRY: - raise ValueError(f"Unknown solver {solver!r}. Registered solvers: {sorted(SOLVER_REGISTRY)}.") - self._solver_class, self._vf_class = SOLVER_REGISTRY[solver] - self._dataloader: TrainSampler | OOCTrainSampler | None = None - self._trainer: CellFlowTrainer | None = None - self._validation_data: dict[str, ValidationData] = {"predict_kwargs": {}} - self._solver: _otfm.OTFlowMatching | _genot.GENOT | None = None - self._condition_dim: int | None = None - self._vf: _velocity_field.ConditionalVelocityField | _velocity_field.GENOTConditionalVelocityField | None = None + self._train_data: TrainingData | None = None def prepare_data( self, @@ -71,6 +49,7 @@ def prepare_data( split_covariates: Sequence[str] | None = None, max_combination_length: int | None = None, null_value: float = 0.0, + adata: ad.AnnData | None = None, ) -> None: """Prepare the dataloader for training from :attr:`~cellflow.model.CellFlow.adata`. @@ -120,6 +99,9 @@ def prepare_data( as the maximal number of perturbations a cell has been treated with. null_value Value to use for padding to ``'max_combination_length'``. + adata + The :class:`~anndata.AnnData` object to extract the training data from. If :obj:`None`, + the object passed to the (deprecated) constructor argument is used instead. Returns ------- @@ -170,6 +152,14 @@ def prepare_data( split_covariates=split_covariates, ) """ + adata = adata if adata is not None else self._adata + if adata is None: + raise ValueError( + "No `adata` provided. Pass it as `prepare_data(adata=...)` (recommended) or " + "construct `CellFlow(adata=...)`." + ) + self._adata = adata # kept for downstream predict / validation / plotting + self._dm = DataManager( self.adata, sample_rep=sample_rep, @@ -196,6 +186,13 @@ def prepare_validation_data( ) -> None: """Prepare the validation data. + Validation is always in-memory (metrics need materialized cells): pass a held-out + :class:`~anndata.AnnData` and its cells (at ``sample_rep``) + condition embeddings become a + ``ValidationData``. Works for both the in-memory and streaming training paths; in the streaming + path the ``adata`` must carry the same ``sample_rep`` and the covariate embeddings in ``.uns``. + (Unrelated to the ``val`` split from :meth:`split_annbatch_data`, a streaming loader — not a + validation set.) + Parameters ---------- adata @@ -222,9 +219,7 @@ def prepare_validation_data( """ if self.train_data is None: - raise ValueError( - "Dataloader not initialized. Training data needs to be set up before preparing validation data. Please call prepare_data first." - ) + raise ValueError("Model data not initialized. Call `prepare_data(...)` before preparing validation data.") val_data = self._dm.get_validation_data( adata, n_conditions_on_log_iteration=n_conditions_on_log_iteration, @@ -241,589 +236,6 @@ def prepare_validation_data( predict_kwargs = self._validation_data["predict_kwargs"] self._validation_data["predict_kwargs"] = predict_kwargs - def prepare_model( - self, - condition_mode: Literal["deterministic", "stochastic"] = "deterministic", - regularization: float = 0.0, - pooling: Literal["mean", "attention_token", "attention_seed"] = "attention_token", - pooling_kwargs: dict[str, Any] = types.MappingProxyType({}), - layers_before_pool: Layers_separate_input_t | Layers_t = dc_field(default_factory=lambda: []), - layers_after_pool: Layers_t = dc_field(default_factory=lambda: []), - condition_embedding_dim: int = 256, - cond_output_dropout: float = 0.9, - condition_encoder_kwargs: dict[str, Any] | None = None, - pool_sample_covariates: bool = True, - time_freqs: int = 1024, - time_max_period: int | None = 10000, - time_encoder_dims: Sequence[int] = (2048, 2048, 2048), - time_encoder_dropout: float = 0.0, - hidden_dims: Sequence[int] = (2048, 2048, 2048), - hidden_dropout: float = 0.0, - conditioning: Literal["concatenation", "film", "resnet"] = "concatenation", - conditioning_kwargs: dict[str, Any] = dc_field(default_factory=lambda: {}), - decoder_dims: Sequence[int] = (4096, 4096, 4096), - decoder_dropout: float = 0.0, - vf_act_fn: Callable[[jnp.ndarray], jnp.ndarray] = nn.silu, - vf_kwargs: dict[str, Any] | None = None, - probability_path: dict[Literal["constant_noise", "bridge"], float] | None = None, - match_fn: Callable[[ArrayLike, ArrayLike], ArrayLike] = match_linear, - optimizer: optax.GradientTransformation = optax.MultiSteps(optax.adam(5e-5), 20), - solver_kwargs: dict[str, Any] | None = None, - layer_norm_before_concatenation: bool = False, - linear_projection_before_concatenation: bool = False, - seed=0, - ) -> None: - """Prepare the model for training. - - This function sets up the neural network architecture and specificities of the - :attr:`solver`. When :attr:`solver` is an instance of :class:`cellflow.solvers._genot.GENOT`, - the following arguments have to be passed to ``'condition_encoder_kwargs'``: - - - Parameters - ---------- - condition_mode - Mode of the encoder, should be one of: - - - ``'deterministic'``: Learns condition encoding point-wise. - - ``'stochastic'``: Learns a Gaussian distribution for representing conditions. - - regularization - Regularization strength in the latent space: - - - For deterministic mode, it is the strength of the L2 regularization. - - For stochastic mode, it is the strength of the VAE regularization. - - pooling - Pooling method, should be one of: - - - ``'mean'``: Aggregates combinations of covariates by the mean of their - learned embeddings. - - ``'attention_token'``: Aggregates combinations of covariates by an attention - mechanism with a class token. - - ``'attention_seed'``: Aggregates combinations of covariates by seed attention. - - pooling_kwargs - Keyword arguments for the pooling method corresponding to: - - - :class:`cellflow.networks.TokenAttentionPooling` if ``'pooling'`` is - ``'attention_token'``. - - :class:`cellflow.networks.SeedAttentionPooling` if ``'pooling'`` is ``'attention_seed'``. - - layers_before_pool - Layers applied to the condition embeddings before pooling. Can be of type - - - :class:`tuple` with elements corresponding to dictionaries with keys: - - - ``'layer_type'`` of type :class:`str` indicating the type of the layer, can be - ``'mlp'`` or ``'self_attention'``. - - Further keyword arguments for the layer type :class:`cellflow.networks.MLPBlock` or - :class:`cellflow.networks.SelfAttentionBlock`. - - - :class:`dict` with keys corresponding to perturbation covariate keys, and values - correspondinng to the above mentioned tuples. - - layers_after_pool - Layers applied to the condition embeddings after pooling, and before applying the last - layer of size ``'condition_embedding_dim'``. Should be of type :class:`tuple` with - elements corresponding to dictionaries with keys: - - - ``'layer_type'`` of type :class:`str` indicating the type of the layer, can be - ``'mlp'`` or ``'self_attention'``. - - Further keys depend on the layer type, either for :class:`cellflow.networks.MLPBlock` or - for :class:`cellflow.networks.SelfAttentionBlock`. - - condition_embedding_dim - Dimensions of the condition embedding, i.e. the last layer of the - :class:`cellflow.networks.ConditionEncoder`. - cond_output_dropout - Dropout rate for the last layer of the :class:`cellflow.networks.ConditionEncoder`. - condition_encoder_kwargs - Keyword arguments for the :class:`cellflow.networks.ConditionEncoder`. - pool_sample_covariates - Whether to include sample covariates in the pooling. - time_freqs - Frequency of the sinusoidal time encoding - (:func:`ott.neural.networks.layers.sinusoidal_time_encoder`). - time_max_period - Controls the frequency of the time embeddings, see - :func:`cellflow.networks.utils.sinusoidal_time_encoder`. - time_encoder_dims - Dimensions of the layers processing the time embedding in - :attr:`cellflow.networks.ConditionalVelocityField.time_encoder`. - time_encoder_dropout - Dropout rate for the :attr:`cellflow.networks.ConditionalVelocityField.time_encoder`. - hidden_dims - Dimensions of the layers processing the input to the velocity field - via :attr:`cellflow.networks.ConditionalVelocityField.x_encoder`. - hidden_dropout - Dropout rate for :attr:`cellflow.networks.ConditionalVelocityField.x_encoder`. - conditioning - Conditioning method, should be one of: - - - ``'concatenation'``: Concatenate the time, data, and condition embeddings. - - ``'film'``: Use FiLM conditioning, i.e. learn FiLM weights from time and condition embedding - to scale the data embeddings. - - ``'resnet'``: Use residual conditioning. - - conditioning_kwargs - Keyword arguments for the conditioning method. - decoder_dims - Dimensions of the output layers in - :attr:`cellflow.networks.ConditionalVelocityField.decoder`. - decoder_dropout - Dropout rate for the output layer - :attr:`cellflow.networks.ConditionalVelocityField.decoder`. - vf_act_fn - Activation function of the :class:`cellflow.networks.ConditionalVelocityField`. - vf_kwargs - Additional keyword arguments for the solver-specific vector field. - For instance, when ``'solver==genot'``, the following keyword argument can be passed: - - - ``'genot_source_dims'`` of type :class:`tuple` with the dimensions - of the :class:`cellflow.networks.MLPBlock` processing the source cell. - - ``'genot_source_dropout'`` of type :class:`float` indicating the dropout rate - for the source cell processing. - probability_path - Probability path to use for training. Should be a :class:`dict` of the form - - - ``'{"constant_noise": noise_val'`` - - ``'{"bridge": noise_val}'`` - - If :obj:`None`, defaults to ``'{"constant_noise": 0.0}'``. - match_fn - Matching function between unperturbed and perturbed cells. Should take as input source - and target data and return the optimal transport matrix, see e.g. - :func:`cellflow.utils.match_linear`. - optimizer - Optimizer used for training. - solver_kwargs - Keyword arguments for the solver :class:`cellflow.solvers.OTFlowMatching` or - :class:`cellflow.solvers.GENOT`. - layer_norm_before_concatenation - If :obj:`True`, applies layer normalization before concatenating - the embedded time, embedded data, and condition embeddings. - linear_projection_before_concatenation - If :obj:`True`, applies a linear projection before concatenating - the embedded time, embedded data, and embedded condition. - seed - Random seed. - - Returns - ------- - Updates the following fields: - - - :attr:`cellflow.model.CellFlow.velocity_field` - an instance of the - :class:`cellflow.networks.ConditionalVelocityField`. - - :attr:`cellflow.model.CellFlow.solver` - an instance of :class:`cellflow.solvers.OTFlowMatching` - or :class:`cellflow.solvers.GENOT`. - - :attr:`cellflow.model.CellFlow.trainer` - an instance of the - :class:`cellflow.training.CellFlowTrainer`. - """ - if self.train_data is None: - raise ValueError("Dataloader not initialized. Please call `prepare_data` first.") - - if condition_mode == "stochastic": - if regularization == 0.0: - raise ValueError("Stochastic condition embeddings require `regularization`>0.") - - condition_encoder_kwargs = condition_encoder_kwargs or {} - # Each velocity field owns which solver-specific `vf_kwargs` it accepts (validated/defaulted here). - vf_kwargs = self._vf_class._normalize_vf_kwargs(vf_kwargs) - covariates_not_pooled = [] if pool_sample_covariates else self._dm.sample_covariates - solver_kwargs = solver_kwargs or {} - probability_path = probability_path or {"constant_noise": 0.0} - - self.vf = self._vf_class( - output_dim=self._data_dim, - max_combination_length=self.train_data.max_combination_length, - condition_mode=condition_mode, - regularization=regularization, - condition_embedding_dim=condition_embedding_dim, - covariates_not_pooled=covariates_not_pooled, - pooling=pooling, - pooling_kwargs=pooling_kwargs, - layers_before_pool=layers_before_pool, - layers_after_pool=layers_after_pool, - cond_output_dropout=cond_output_dropout, - condition_encoder_kwargs=condition_encoder_kwargs, - act_fn=vf_act_fn, - time_freqs=time_freqs, - time_max_period=time_max_period, - time_encoder_dims=time_encoder_dims, - time_encoder_dropout=time_encoder_dropout, - hidden_dims=hidden_dims, - hidden_dropout=hidden_dropout, - conditioning=conditioning, - conditioning_kwargs=conditioning_kwargs, - decoder_dims=decoder_dims, - decoder_dropout=decoder_dropout, - layer_norm_before_concatenation=layer_norm_before_concatenation, - linear_projection_before_concatenation=linear_projection_before_concatenation, - **vf_kwargs, - ) - - probability_path, noise = next(iter(probability_path.items())) - if probability_path == "constant_noise": - probability_path = ConstantNoiseFlow(noise) - elif probability_path == "bridge": - probability_path = BrownianBridge(noise) - else: - raise NotImplementedError( - f"The key of `probability_path` must be `'constant_noise'` or `'bridge'` but found {probability_path}." - ) - - # Each solver owns how it names its match function / whether it needs source-target dims. - self._solver = self._solver_class( - vf=self.vf, - probability_path=probability_path, - optimizer=optimizer, - conditions=self.train_data.condition_data, - rng=jax.random.PRNGKey(seed), - **self._solver_class._match_kwargs(match_fn=match_fn, data_dim=self._data_dim), - **solver_kwargs, - ) - - self._trainer = CellFlowTrainer(solver=self.solver, predict_kwargs=self.validation_data["predict_kwargs"]) # type: ignore[arg-type] - - def train( - self, - num_iterations: int, - batch_size: int = 1024, - valid_freq: int = 1000, - callbacks: Sequence[BaseCallback] = [], - monitor_metrics: Sequence[str] = [], - out_of_core_dataloading: bool = False, - ) -> None: - """Train the model. - - Note - ---- - A low value of ``'valid_freq'`` results in long training - because predictions are time-consuming compared to training steps. - - Parameters - ---------- - num_iterations - Number of iterations to train the model. - batch_size - Batch size. - valid_freq - Frequency of validation. - callbacks - Callbacks to perform at each validation step. There are two types of callbacks: - - Callbacks for computations should inherit from - :class:`~cellflow.training.ComputationCallback` see e.g. :class:`cellflow.training.Metrics`. - - Callbacks for logging should inherit from :class:`~cellflow.training.LoggingCallback` see - e.g. :class:`~cellflow.training.WandbLogger`. - monitor_metrics - Metrics to monitor. - out_of_core_dataloading - If :obj:`True`, use out-of-core dataloading. Uses the :class:`cellflow.data._dataloader.OOCTrainSampler` - to load data that does not fit into GPU memory. - - Returns - ------- - Updates the following fields: - - - :attr:`cellflow.model.CellFlow.dataloader` - the training dataloader. - - :attr:`cellflow.model.CellFlow.solver` - the trained solver. - """ - if self.train_data is None: - raise ValueError("Data not initialized. Please call `prepare_data` first.") - - if self.trainer is None: - raise ValueError("Model not initialized. Please call `prepare_model` first.") - - if out_of_core_dataloading: - self._dataloader = OOCTrainSampler(data=self.train_data, batch_size=batch_size) - else: - self._dataloader = TrainSampler(data=self.train_data, batch_size=batch_size) - validation_loaders = {k: ValidationSampler(v) for k, v in self.validation_data.items() if k != "predict_kwargs"} - self._trainer.predict_kwargs = self.validation_data.get("predict_kwargs", {}) - - self._solver = self.trainer.train( - dataloader=self._dataloader, - num_iterations=num_iterations, - valid_freq=valid_freq, - valid_loaders=validation_loaders, - callbacks=callbacks, - monitor_metrics=monitor_metrics, - ) - - def predict( - self, - adata: ad.AnnData, - covariate_data: pd.DataFrame, - sample_rep: str | None = None, - condition_id_key: str | None = None, - key_added_prefix: str | None = None, - rng: ArrayLike | None = None, - **kwargs: Any, - ) -> dict[str, ArrayLike] | None: - """Predict perturbation responses. - - Parameters - ---------- - adata - An :class:`~anndata.AnnData` object with the source representation. - covariate_data - Covariate data defining the condition to predict. This :class:`~pandas.DataFrame` - should have the same columns as :attr:`~anndata.AnnData.obs` of - :attr:`cellflow.model.CellFlow.adata`, and as registered in - :attr:`cellflow.model.CellFlow.data_manager`. - sample_rep - Key in :attr:`~anndata.AnnData.obsm` where the sample representation is stored or - ``'X'`` to use :attr:`~anndata.AnnData.X`. If :obj:`None`, the key is assumed to be - the same as for the training data. - condition_id_key - Key in ``'covariate_data'`` defining the condition name. - key_added_prefix - If not :obj:`None`, prefix to store the prediction in :attr:`~anndata.AnnData.obsm`. - If :obj:`None`, the predictions are not stored, and the predictions are returned as a - :class:`dict`. - rng - Random number generator. If :obj:`None` and :attr:`cellflow.model.CellFlow.conditino_mode` - is ``'stochastic'``, the condition vector will be the mean of the learnt distributions, - otherwise samples from the distribution. - kwargs - Keyword arguments for the predict function, i.e. - :meth:`cellflow.solvers.OTFlowMatching.predict` or :meth:`cellflow.solvers.GENOT.predict`. - - Returns - ------- - If ``'key_added_prefix'`` is :obj:`None`, a :class:`dict` with the predicted sample - representation for each perturbation, otherwise stores the predictions in - :attr:`~anndata.AnnData.obsm` and returns :obj:`None`. - """ - if self.solver is None or not self.solver.is_trained: - raise ValueError("Model not trained. Please call `train` first.") - - if sample_rep is None: - sample_rep = self._dm.sample_rep - - if adata is not None and covariate_data is not None: - if covariate_data.empty: - raise ValueError("`covariate_data` is empty.") - if self._dm.control_key not in adata.obs.columns: - raise ValueError( - f"If both `adata` and `covariate_data` are given, the control key `{self._dm.control_key}` must be in `adata.obs`." - ) - if not adata.obs[self._dm.control_key].all(): - raise ValueError( - f"If both `adata` and `covariate_data` are given, all samples in `adata` must be control samples, and thus `adata.obs[`{self._dm.control_key}`] must be set to `True` everywhere." - ) - pred_data = self._dm.get_prediction_data( - adata, - sample_rep=sample_rep, # type: ignore[arg-type] - covariate_data=covariate_data, - condition_id_key=condition_id_key, - ) - pred_loader = PredictionSampler(pred_data) - batch = pred_loader.sample() - src = batch["source"] - condition = batch.get("condition", None) - # using jax.tree.map to batch the prediction - # because PredictionSampler can return a different number of cells for each condition - out = jax.tree.map( - functools.partial(self.solver.predict, rng=rng, **kwargs), - src, - condition, # type: ignore[attr-defined] - ) - if key_added_prefix is None: - return out - if len(pred_data.control_to_perturbation) > 1: - raise ValueError( - f"When saving predictions to `adata`, all control cells must be from the same control \ - population, but found {len(pred_data.control_to_perturbation)} control populations." - ) - out_np = {k: np.array(v) for k, v in out.items()} - _write_predictions( - adata=adata, - predictions=out_np, - key_added_prefix=key_added_prefix, - ) - - def get_condition_embedding( - self, - covariate_data: pd.DataFrame | ConditionData, - rep_dict: dict[str, str] | None = None, - condition_id_key: str | None = None, - key_added: str | None = _constants.CONDITION_EMBEDDING, - ) -> tuple[pd.DataFrame, pd.DataFrame]: - """Get the embedding of the conditions. - - Outputs the mean and variance of the learnt embeddings - generated by the :class:`~cellflow.networks.ConditionEncoder`. - - Parameters - ---------- - covariate_data - Can be one of - - - a :class:`~pandas.DataFrame` defining the conditions with the same columns as the - :class:`~anndata.AnnData` used for the initialisation of :class:`~cellflow.model.CellFlow`. - - an instance of :class:`~cellflow.data.ConditionData`. - - rep_dict - Dictionary containing the representations of the perturbation covariates. Will be considered an - empty dictionary if :obj:`None`. - condition_id_key - Key defining the name of the condition. Only available - if ``'covariate_data'`` is a :class:`~pandas.DataFrame`. - key_added - Key to store the condition embedding in :attr:`~anndata.AnnData.uns`. The mean is - stored under ``key_added`` and the variance under ``f"{key_added}_var"``. If - :obj:`None`, the embeddings are not stored. - - Returns - ------- - A :class:`tuple` of :class:`~pandas.DataFrame` with the mean and variance of the condition embeddings. - """ - if self.solver is None or not self.solver.is_trained: - raise ValueError("Model not trained. Please call `train` first.") - - if hasattr(covariate_data, "condition_data"): - cond_data = covariate_data - elif isinstance(covariate_data, pd.DataFrame): - cond_data = self._dm.get_condition_data( - covariate_data=covariate_data, - rep_dict=rep_dict, - condition_id_key=condition_id_key, - ) - else: - raise ValueError("Covariate data must be a `pandas.DataFrame` or an instance of `BaseData`.") - - condition_embeddings_mean: dict[str, ArrayLike] = {} - condition_embeddings_var: dict[str, ArrayLike] = {} - n_conditions = len(next(iter(cond_data.condition_data.values()))) - for i in range(n_conditions): - condition = {k: v[[i], :] for k, v in cond_data.condition_data.items()} - if condition_id_key: - c_key = cond_data.perturbation_idx_to_id[i] - else: - cov_combination = cond_data.perturbation_idx_to_covariates[i] - c_key = tuple(cov_combination[i] for i in range(len(cov_combination))) - condition_embeddings_mean[c_key], condition_embeddings_var[c_key] = self.solver.get_condition_embedding( - condition - ) - - df_mean = pd.DataFrame.from_dict({k: v[0] for k, v in condition_embeddings_mean.items()}).T - df_var = pd.DataFrame.from_dict({k: v[0] for k, v in condition_embeddings_var.items()}).T - - if condition_id_key: - df_mean.index.set_names([condition_id_key], inplace=True) - df_var.index.set_names([condition_id_key], inplace=True) - else: - df_mean.index.set_names(list(self._dm.perturb_covar_keys), inplace=True) - df_var.index.set_names(list(self._dm.perturb_covar_keys), inplace=True) - - if key_added is not None: - _utils.set_plotting_vars(self.adata, key=key_added, value=df_mean) - _utils.set_plotting_vars(self.adata, key=f"{key_added}_var", value=df_var) - - return df_mean, df_var - - def save( - self, - dir_path: str, - file_prefix: str | None = None, - overwrite: bool = False, - ) -> None: - """ - Save the model. - - Pickles the :class:`~cellflow.model.CellFlow` object. - - Parameters - ---------- - dir_path - Path to a directory, defaults to current directory - file_prefix - Prefix to prepend to the file name. - overwrite - Overwrite existing data or not. - - Returns - ------- - :obj:`None` - """ - file_name = ( - f"{file_prefix}_{self.__class__.__name__}.pkl" - if file_prefix is not None - else f"{self.__class__.__name__}.pkl" - ) - file_dir = os.path.join(dir_path, file_name) if dir_path is not None else file_name - - if not overwrite and os.path.exists(file_dir): - raise RuntimeError(f"Unable to save to an existing file `{file_dir}` use `overwrite=True` to overwrite it.") - with open(file_dir, "wb") as f: - cloudpickle.dump(self, f) - - @classmethod - def load( - cls, - filename: str, - ) -> "CellFlow": - """ - Load a :class:`~cellflow.model.CellFlow` model from a saved instance. - - Parameters - ---------- - filename - Path to the saved file. - - Returns - ------- - Loaded instance of the model. - """ - # Check if filename is a directory - file_name = os.path.join(filename, f"{cls.__name__}.pkl") if os.path.isdir(filename) else filename - - with open(file_name, "rb") as f: - model = cloudpickle.load(f) - - if type(model) is not cls: - raise TypeError(f"Expected the model to be type of `{cls}`, found `{type(model)}`.") - return model - - @property - def adata(self) -> ad.AnnData: - """The :class:`~anndata.AnnData` object used for training.""" - return self._adata - - @property - def solver(self) -> _otfm.OTFlowMatching | _genot.GENOT | None: - """The solver.""" - return self._solver - - @property - def dataloader(self) -> TrainSampler | OOCTrainSampler | None: - """The dataloader used for training.""" - return self._dataloader - - @property - def trainer(self) -> CellFlowTrainer | None: - """The trainer used for training.""" - return self._trainer - - @property - def validation_data(self) -> dict[str, ValidationData]: - """The validation data.""" - return self._validation_data - - @property - def data_manager(self) -> DataManager: - """The data manager, initialised with :attr:`cellflow.model.CellFlow.adata`.""" - return self._dm - - @property - def velocity_field( - self, - ) -> _velocity_field.ConditionalVelocityField | _velocity_field.GENOTConditionalVelocityField | None: - """The conditional velocity field.""" - return self._vf - @property def train_data(self) -> TrainingData | None: """The training data.""" @@ -836,14 +248,18 @@ def train_data(self, data: TrainingData) -> None: raise ValueError(f"Expected `data` to be an instance of `TrainingData`, found `{type(data)}`.") self._train_data = data - @velocity_field.setter # type: ignore[attr-defined,no-redef] - def velocity_field(self, vf: _velocity_field.ConditionalVelocityField) -> None: - """Set the velocity field.""" - if not isinstance(vf, _velocity_field.ConditionalVelocityField): - raise ValueError(f"Expected `vf` to be an instance of `ConditionalVelocityField`, found `{type(vf)}`.") - self._vf = vf + # ── path-specific hook implementations ─────────────────────────────────────────────────────── + def _encoder_conditions(self) -> tuple[dict[str, np.ndarray], int]: + if self.train_data is None: + raise ValueError("Data not initialized. Please call `prepare_data(...)` first.") + return self.train_data.condition_data, self.train_data.max_combination_length + + def _bind_train_dataloader(self, batch_size: int, out_of_core_dataloading: bool) -> None: + self._dataloader = ( + OOCTrainSampler(data=self.train_data, batch_size=batch_size) + if out_of_core_dataloading + else TrainSampler(data=self.train_data, batch_size=batch_size) + ) - @property - def condition_mode(self) -> Literal["deterministic", "stochastic"]: - """The mode of the encoder.""" - return self.velocity_field.condition_mode + def _build_validation_loaders(self) -> dict[str, Any]: + return {k: ValidationSampler(v) for k, v in self.validation_data.items() if k != "predict_kwargs"} diff --git a/src/cellflow/model/_cellflow_annbatch.py b/src/cellflow/model/_cellflow_annbatch.py new file mode 100644 index 00000000..9a8271ea --- /dev/null +++ b/src/cellflow/model/_cellflow_annbatch.py @@ -0,0 +1,325 @@ +"""Streaming (annbatch/dagloader) CellFlow: the default path. + +Trains, validates and predicts over an out-of-core :class:`~annbatch.DatasetCollection` or an +in-memory ``AnnData`` alike. +""" + +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Literal + +import anndata as ad +import numpy as np +import pandas as pd + +from cellflow._types import ArrayLike +from cellflow.data._dataloader import DAGEvalAdapter, DAGTrainAdapter +from cellflow.model._base import BaseCellFlow + +if TYPE_CHECKING: + from annbatch import DatasetCollection # optional dep — only imported for typing + + from dagloader import DAGEvalLoader, SamplerConfig, Scheme # optional deps — typing only + +__all__ = ["CellFlowAnnbatch"] + + +class CellFlowAnnbatch(BaseCellFlow): + """CellFlow over the annbatch/dagloader streaming path (cells sourced out-of-core or in-memory). + + Cells are streamed from an :class:`~annbatch.DatasetCollection` (out-of-core) or an in-memory + ``AnnData`` via :meth:`prepare_data`; validation and prediction read each condition's full cell + set through :class:`~dagloader.DAGEvalLoader` (no boolean masking). Model setup, training, prediction + and IO are inherited from :class:`~cellflow.model._base.BaseCellFlow`. + """ + + def __init__(self, solver: Literal["otfm", "genot"] = "otfm"): + super().__init__(solver) + self._scheme: Scheme | None = None + self._split_schemes: dict[str, Scheme] | None = None + self._split_assignment: pd.DataFrame | None = None + self._annbatch_sampler_configs: dict[str, SamplerConfig] | None = None + self._eval_cfg: SamplerConfig | None = None # target read params for DAGEvalLoader + self._condition_data: dict[str, np.ndarray] | None = None # condition embeddings + self._max_combination_length: int | None = None + self._condition_fn = None # leaf -> embedding (set by `prepare_data`) + self._prep_kwargs: dict[str, Any] | None = None # covariate spec, reused for validation sources + self._seed: int = 0 + self._split_eval_loaders: dict[str, DAGEvalLoader] = {} + + def prepare_data( + self, + source: "ad.AnnData | DatasetCollection", + sample_rep: str, + control_key: str, + perturbation_covariates: dict[str, Sequence[str]], + perturbation_covariate_reps: dict[str, str] | None = None, + sample_covariates: Sequence[str] | None = None, + sample_covariate_reps: dict[str, str] | None = None, + split_covariates: Sequence[str] | None = None, + max_combination_length: int | None = None, + null_value: float = 0.0, + rep_dict: Mapping[str, Mapping[str, ArrayLike]] | None = None, + sampler_config: "SamplerConfig | Mapping[str, SamplerConfig] | None" = None, + seed: int = 0, + control_in_memory: bool = False, + split_by: Sequence[str] | None = None, + split_ratios: Mapping[str, float] | None = None, + split_force_training_values: Mapping[str, object] | None = None, + split_random_state: int = 42, + ) -> None: + """Prepare the annbatch/dagloader streaming training path (the default, recommended path). + + The covariate arguments (``sample_rep``, ``control_key``, ``perturbation_covariates``, + ``perturbation_covariate_reps``, ``sample_covariates``, ``sample_covariate_reps``, + ``split_covariates``, ``max_combination_length``, ``null_value``) mean exactly what they do in + :meth:`prepare_data`; only the cells differ — streamed from an out-of-core + :class:`annbatch.DatasetCollection` instead of materialized. Requires the ``annbatch`` extra and + a model constructed as ``CellFlowAnnbatch()``. + + Parameters + ---------- + source + An out-of-core :class:`annbatch.DatasetCollection` to stream cells from (an in-memory + ``AnnData`` also works — the ``dagloader`` is container-agnostic). Its ``obs`` supplies the + grouping / condition columns; ``sample_rep`` selects the streamed representation. + rep_dict + The covariate embedding tables that ``adata.uns`` would hold in the in-memory path (keys + match the values of ``perturbation_covariate_reps`` / ``sample_covariate_reps``). Required + when a covariate group is embedded; may be :obj:`None` when the primary covariate is + categorical (one-hot encoded). + sampler_config + Read parameters for the streamed loader(s), **one per split**: a single + :class:`dagloader.SamplerConfig` for all splits, or a ``{split_name: SamplerConfig}`` mapping + covering every split (see :func:`dagloader.resolve_split_configs`). With no split the only + split is ``"train"``. ``chunk_size > 1`` reads contiguous slices, so every run of each + category must be at least ``chunk_size`` cells — in-memory sources are grouped automatically, + an out-of-core :class:`~annbatch.DatasetCollection` must be built grouped + (``add_adatas(..., groupby=[...])``) or a clear error is raised. + seed + Reproducibility seed for the ``dagloader`` per-node RNG streams. + split_by + If given, split the prepared ``Scheme``'s target combinations into train/val/test in the + same call (delegates to :meth:`split_annbatch_data`). Columns whose unique combinations are + held out (⊆ the target columns). If :obj:`None`, no split is made (the model would train on + all combinations). + split_ratios + ``{split_name: fraction}`` summing to 1.0 for the split. Defaults to + ``{"train": 0.6, "val": 0.2, "test": 0.2}``. Only used when ``split_by`` is given. + split_force_training_values + ``{column: value}`` (keys ⊆ ``split_by``) forced into the training split. Only used when + ``split_by`` is given. + split_random_state + Seed for the split's combination shuffle. Only used when ``split_by`` is given. + + Returns + ------- + :obj:`None`, and sets up the streaming training data used by :meth:`train`. + + Notes + ----- + The ``"train"`` split feeds :meth:`train`; when a split is made, the ``val`` / ``test`` splits are + read via :class:`~dagloader.DAGEvalLoader` (see :attr:`split_eval_loaders`), not streamed. + """ + from cellflow.data._annbatch import build_annbatch_training + from dagloader import DAGLoader, resolve_split_configs + + # Build the Scheme + condition_fn + condition embeddings from the covariate spec (obs only). + built = build_annbatch_training( + source, + sample_rep=sample_rep, + control_key=control_key, + perturbation_covariates=perturbation_covariates, + perturbation_covariate_reps=perturbation_covariate_reps, + sample_covariates=sample_covariates, + sample_covariate_reps=sample_covariate_reps, + split_covariates=split_covariates, + max_combination_length=max_combination_length, + null_value=null_value, + rep_dict=rep_dict, + seed=seed, + control_in_memory=control_in_memory, + ) + self._scheme = built.scheme + self._dm = built.data_manager + self._condition_data = built.condition_data + self._data_dim = built.data_dim + self._max_combination_length = built.max_combination_length + condition_fn = built.condition_fn + # kept so `prepare_validation_data` / the auto-wired split loaders can build DAGEvalLoaders. + self._condition_fn = condition_fn + self._seed = seed + self._prep_kwargs = { + "sample_rep": sample_rep, + "control_key": control_key, + "perturbation_covariates": perturbation_covariates, + "perturbation_covariate_reps": perturbation_covariate_reps, + "sample_covariates": sample_covariates, + "sample_covariate_reps": sample_covariate_reps, + "split_covariates": split_covariates, + "max_combination_length": max_combination_length, + "null_value": null_value, + "rep_dict": rep_dict, + "seed": seed, + "control_in_memory": control_in_memory, + } + + # Splitting step — kept in `prepare_data` so preparing and splitting are one call. + if split_by is not None: + self._split_assignment = self.split_annbatch_data( + split_by=split_by, + ratios=split_ratios, + force_training_values=split_force_training_values, + random_state=split_random_state, + ) + + # One `SamplerConfig` per split (a single spec ⇒ all splits; a per-split dict ⇒ all specified). + # The splits are the split schemes when a split was made, else the single ``"train"`` scheme. + schemes = self._split_schemes if self._split_schemes is not None else {"train": self._scheme} + if sampler_config is None: + raise ValueError("`sampler_config` is required: give a SamplerConfig or a {split: SamplerConfig} mapping.") + self._annbatch_sampler_configs = resolve_split_configs(sampler_config, list(schemes)) + self._eval_cfg = self._annbatch_sampler_configs["train"] # target-batch read params for eval loaders + + # `chunk_size > 1` reads contiguous slices → every category's runs must be >= chunk_size + # (in-memory sources are auto-grouped above; out-of-core collections must be built grouped). + max_chunk = max(cfg.chunk_size for cfg in self._annbatch_sampler_configs.values()) + if max_chunk > 1: + from cellflow.data._annbatch import assert_source_chunkable + + root = self._scheme.nodes[self._scheme.root] + assert_source_chunkable(self._scheme.sources["data"], root.cols, max_chunk) + + # Only the "train" split is streamed (feeds `train()`); val/test are read via DAGEvalLoader + # below, so we don't build unused per-split streaming loaders. + self._dataloader = DAGTrainAdapter( + DAGLoader(schemes["train"], self._annbatch_sampler_configs["train"], condition_fn=condition_fn) + ) + + # Auto-wire the val/test split combinations as evaluation sources over the same cells: a + # `DAGEvalLoader` reads each held-out condition's full cell set + matched controls. The "val" + # split (if any) feeds training-time metrics; every non-train split is kept on + # `split_eval_loaders` for post-hoc evaluation. Override with `prepare_validation_data(...)`. + from dagloader import DAGEvalLoader + + self._split_eval_loaders = {} + if self._split_schemes is not None: + for split_name, sch in self._split_schemes.items(): + if split_name == "train": + continue + self._split_eval_loaders[split_name] = DAGEvalLoader(sch, self._eval_cfg, condition_fn, seed=seed) + if "val" in self._split_eval_loaders: + self._validation_data["val"] = DAGEvalAdapter(self._split_eval_loaders["val"]) + + def split_annbatch_data( + self, + *, + split_by: Sequence[str], + ratios: Mapping[str, float] | None = None, + force_training_values: Mapping[str, object] | None = None, + random_state: int = 42, + ) -> pd.DataFrame: + """Partition the prepared annbatch ``Scheme``'s target combinations into named splits. + + Call after :meth:`prepare_data`. Splits hold out whole *combinations* of ``split_by`` + (a subset of the perturbation / split-covariate columns), not cells; controls are carried through + every split. See :func:`dagloader.split_scheme` for the mechanics. + + Parameters + ---------- + split_by + Columns whose unique combinations are partitioned across splits (⊆ the scheme's target + columns). + ratios + ``{split_name: fraction}`` summing to 1.0. Defaults to + ``{"train": 0.6, "val": 0.2, "test": 0.2}``. The first split is the training split. + force_training_values + ``{column: value}`` (keys ⊆ ``split_by``): any combination matching a value is forced into + the training (first) split. + random_state + Seed for the combination shuffle. + + Returns + ------- + A :class:`~pandas.DataFrame` of the target combinations and their assigned split. Also stores + the per-split schemes on the model. + """ + from dagloader import split_assignment, split_scheme + + if self._scheme is None: + raise ValueError( + "No annbatch `Scheme` to split. Call `prepare_data(...)` first " + "(the out-of-core streaming path)." + ) + self._split_schemes = split_scheme( + self._scheme, + split_by=split_by, + ratios=ratios, + force_training_values=force_training_values, + random_state=random_state, + ) + # `prepare_data` wires the "train" split into the streaming loader and the val/test splits into + # `DAGEvalLoader`s (`split_eval_loaders`); here we only produce the schemes + table. + return split_assignment(self._split_schemes) + + def prepare_validation_data( + self, + source: "ad.AnnData | DatasetCollection", + name: str, + n_conditions_on_log_iteration: int | None = None, + n_conditions_on_train_end: int | None = None, + predict_kwargs: dict[str, Any] | None = None, + ) -> None: + """Register a validation set read via :class:`~dagloader.DAGEvalLoader` (no boolean masking). + + The ``source`` (an :class:`~anndata.AnnData` or an :class:`~annbatch.DatasetCollection`) is + grouped by the same covariate spec passed to :meth:`prepare_data`; each of its conditions is + read in full (all its cells + all matched controls) at validation time. Overrides any same-named + auto-wired split loader. Preserves the legacy per-condition metric semantics. + + Parameters + ---------- + source + The held-out validation source (out-of-core or in-memory). + name + Key under which the validation set is stored in :attr:`validation_data`. + n_conditions_on_log_iteration, n_conditions_on_train_end + Conditions to sample per validation step; :obj:`None` uses all. + predict_kwargs + Keyword arguments for the solver's ``predict`` used during validation. + """ + if self._scheme is None or self._prep_kwargs is None: + raise ValueError("Set up training first via `prepare_data(...)` before preparing validation data.") + + from cellflow.data._annbatch import build_annbatch_training + from dagloader import DAGEvalLoader + + built = build_annbatch_training(source, **self._prep_kwargs) + eval_loader = DAGEvalLoader(built.scheme, self._eval_cfg, built.condition_fn, seed=self._seed) + self._validation_data[name] = DAGEvalAdapter( + eval_loader, + n_conditions_on_log_iteration=n_conditions_on_log_iteration, + n_conditions_on_train_end=n_conditions_on_train_end, + ) + predict_kwargs = predict_kwargs or {} + if len(self._validation_data.get("predict_kwargs", {})) > 0 and len(predict_kwargs) > 0: + self._validation_data["predict_kwargs"].update(predict_kwargs) + predict_kwargs = self._validation_data["predict_kwargs"] + self._validation_data["predict_kwargs"] = predict_kwargs + + @property + def split_eval_loaders(self) -> dict[str, "DAGEvalLoader"]: + """Per-split :class:`~dagloader.DAGEvalLoader` objects for non-train splits (e.g. ``val``, ``test``).""" + return self._split_eval_loaders + + # ── path-specific hook implementations ─────────────────────────────────────────────────────── + def _encoder_conditions(self) -> tuple[dict[str, np.ndarray], int]: + if self._condition_data is None or self._max_combination_length is None: + raise ValueError("Data not initialized. Please call `prepare_data(...)` first.") + return self._condition_data, self._max_combination_length + + def _bind_train_dataloader(self, batch_size: int, out_of_core_dataloading: bool) -> None: + # the streaming `_dataloader` (DAGTrainAdapter) was already set in `prepare_data`. + return None + + def _build_validation_loaders(self) -> dict[str, Any]: + return {k: v for k, v in self.validation_data.items() if k != "predict_kwargs"} diff --git a/src/cellflow/networks/_velocity_field.py b/src/cellflow/networks/_velocity_field.py index b1a45db9..6f04147f 100644 --- a/src/cellflow/networks/_velocity_field.py +++ b/src/cellflow/networks/_velocity_field.py @@ -13,7 +13,71 @@ from cellflow.networks._set_encoders import ConditionEncoder from cellflow.networks._utils import FilmBlock, MLPBlock, ResNetBlock, sinusoidal_time_encoder -__all__ = ["ConditionalVelocityField", "GENOTConditionalVelocityField"] +__all__ = [ + "ConditionalVelocityField", + "GENOTConditionalVelocityField", + "null_condition_embedding", + "null_condition_input", +] + + +def null_condition_embedding( + cond_embedding: jnp.ndarray, + *, + condition_dropout_prob: float, + make_rng: Callable[[str], jax.Array], + train: bool, + force_uncond: bool, +) -> jnp.ndarray: + """Null the condition *embedding* for classifier-free guidance (``condition_null='zero_embedding'``). + + - If ``force_uncond`` is set, the condition is always dropped, yielding the + unconditional velocity field (used at inference time). + - Otherwise, during training the condition is dropped independently per set + element with probability ``condition_dropout_prob``. + + With ``force_uncond=False`` and ``condition_dropout_prob == 0.0`` this is a no-op + that reproduces the standard conditional behavior byte-for-byte: no random number + is drawn and the embedding is returned as-is. ``make_rng`` is the owning module's + :meth:`flax.linen.Module.make_rng` (drawn only when a dropout mask is needed). + """ + if force_uncond: + return jnp.zeros_like(cond_embedding) + if train and condition_dropout_prob > 0.0: + keep = jax.random.bernoulli( + make_rng("dropout"), + p=1.0 - condition_dropout_prob, + shape=(cond_embedding.shape[0], 1), + ) + return jnp.where(keep, cond_embedding, jnp.zeros_like(cond_embedding)) + return cond_embedding + + +def null_condition_input( + cond: dict[str, jnp.ndarray], + *, + condition_dropout_prob: float, + mask_value: float, + make_rng: Callable[[str], jax.Array], + train: bool, + force_uncond: bool, +) -> dict[str, jnp.ndarray]: + """Null the *raw* condition by filling it with ``mask_value`` (``condition_null='mask_value'``). + + Routes an all-masked condition set through the condition encoder, so the + unconditional representation is whatever the encoder maps a fully-masked set to + (matching how padded conditions are handled). Same drop policy as + :func:`null_condition_embedding`: always when ``force_uncond`` is set, otherwise per + set element with probability ``condition_dropout_prob`` during training. With the + defaults it returns ``cond`` unchanged and draws no random number. + """ + if force_uncond: + return jax.tree_util.tree_map(lambda c: jnp.full_like(c, mask_value), cond) + if train and condition_dropout_prob > 0.0: + n = next(iter(cond.values())).shape[0] + keep = jax.random.bernoulli(make_rng("dropout"), p=1.0 - condition_dropout_prob, shape=(n, 1, 1)) + return jax.tree_util.tree_map(lambda c: jnp.where(keep, c, jnp.full_like(c, mask_value)), cond) + return cond class ConditionalVelocityField(nn.Module): @@ -172,13 +236,7 @@ def setup(self): self.layer_cond_output_dropout = nn.Dropout(rate=self.cond_output_dropout) self.layer_norm_condition = nn.LayerNorm() if self.layer_norm_before_concatenation else lambda x: x - self.time_encoder = MLPBlock( - dims=self.time_encoder_dims, - act_fn=self.act_fn, - dropout_rate=self.time_encoder_dropout, - act_last_layer=False, - ) - self.layer_norm_time = nn.LayerNorm() if self.layer_norm_before_concatenation else lambda x: x + self._setup_time() self.x_encoder = MLPBlock( dims=self.hidden_dims, @@ -197,6 +255,28 @@ def setup(self): self.output_layer = nn.Dense(self.output_dim) + self._setup_conditioning(conditioning_kwargs) + + def _setup_time(self) -> None: + """Build the time encoder and its optional pre-concatenation LayerNorm. + + Override to a no-op in a time-less velocity field (e.g. Equilibrium Matching). + """ + self.time_encoder = MLPBlock( + dims=self.time_encoder_dims, + act_fn=self.act_fn, + dropout_rate=self.time_encoder_dropout, + act_last_layer=False, + ) + self.layer_norm_time = nn.LayerNorm() if self.layer_norm_before_concatenation else lambda x: x + + def _setup_conditioning(self, conditioning_kwargs: dict[str, Any]) -> None: + """Build the ``conditioning``-mode-specific submodules. + + Called at the end of :meth:`setup` (so all shared submodules already exist on + ``self``). Override in a subclass to add new ``conditioning`` modes, delegating to + ``super()._setup_conditioning`` for the built-in ones. + """ if self.conditioning == "film": self.film_block = FilmBlock( input_dim=self.hidden_dims[-1], @@ -223,7 +303,30 @@ def __call__( train: bool = True, force_uncond: bool = False, ) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: - squeeze = x_t.ndim == 1 + # Split into "encode the condition" and "velocity given the embedding" so the predict path can + # encode once and reuse the embedding across every ODE step (the condition is constant along the + # trajectory). `__call__` composes them, so the training path and any single-shot call are + # unchanged — see `encode_condition` / `velocity_from_embedding`. + cond_embedding, cond_mean, cond_logvar = self.encode_condition( + cond, encoder_noise, train=train, force_uncond=force_uncond + ) + out = self.velocity_from_embedding(t, x_t, cond_embedding, train=train) + return out, cond_mean, cond_logvar + + def encode_condition( + self, + cond: dict[str, jnp.ndarray], + encoder_noise: jnp.ndarray, + train: bool = True, + force_uncond: bool = False, + ) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Encode the condition into its (pre-concatenation) embedding, returning also the raw + Gaussian parameters used by the stochastic-encoder regularization. + + This is the part of the forward pass that only depends on ``cond`` (not on ``t`` or ``x``), so + the predict path evaluates it once per condition and passes the result to + :meth:`velocity_from_embedding` at every ODE step. + """ if self.condition_null == "mask_value": cond = self._maybe_null_input(cond, train=train, force_uncond=force_uncond) cond_mean, cond_logvar = self.condition_encoder(cond, training=train) @@ -235,31 +338,89 @@ def __call__( cond_embedding = self.layer_cond_output_dropout(cond_embedding, deterministic=not train) if self.condition_null == "zero_embedding": cond_embedding = self._maybe_null_embedding(cond_embedding, train=train, force_uncond=force_uncond) + cond_embedding = self.layer_norm_condition(cond_embedding) + return cond_embedding, cond_mean, cond_logvar + + def velocity_from_embedding( + self, + t: jnp.ndarray, + x_t: jnp.ndarray, + cond_embedding: jnp.ndarray, + train: bool = True, + ) -> jnp.ndarray: + """Velocity for ``x_t`` at time ``t`` given an already-encoded (and normed) ``cond_embedding``. + This is the ODE right-hand side; it contains no condition-encoder work, so integrating with a + precomputed embedding avoids re-encoding the condition on every step. + """ + squeeze = x_t.ndim == 1 t_encoded = sinusoidal_time_encoder(t, time_freqs=self.time_freqs, time_max_period=self.time_max_period) t_encoded = self.time_encoder(t_encoded, training=train) - x_encoded = self.x_encoder(x_t, training=train) + x_encoded = self._encode_x(x_t, squeeze, train) t_encoded = self.layer_norm_time(t_encoded) x_encoded = self.layer_norm_x(x_encoded) - cond_embedding = self.layer_norm_condition(cond_embedding) if squeeze: cond_embedding = jnp.squeeze(cond_embedding) # , 0) elif cond_embedding.shape[0] != x_t.shape[0]: # type: ignore[attr-defined] cond_embedding = jnp.tile(cond_embedding, (x_t.shape[0], 1)) + return self._combine_and_decode(t_encoded, x_encoded, cond_embedding, squeeze, train) + + def _encode_x(self, x_t: jnp.ndarray, squeeze: bool, train: bool) -> jnp.ndarray: + """Encode ``x_t`` before conditioning. Override to insert pre-conditioning processing.""" + return self.x_encoder(x_t, training=train) + + def _conditioning_signals( + self, + t_encoded: jnp.ndarray, + x_encoded: jnp.ndarray, + cond_embedding: jnp.ndarray, + x_0_encoded: jnp.ndarray | None = None, + ) -> tuple[tuple[jnp.ndarray, ...], jnp.ndarray]: + """Return ``(concat_inputs, conditioning_vec)`` for the conditioning step. + + ``concat_inputs`` is what ``'concatenation'`` concatenates (order matters for + checkpoint compatibility); ``conditioning_vec`` is what modulates ``x`` for + ``'film'``/``'resnet'`` (and ``'adaln_zero'`` in subclasses). ``GENOT`` folds the + encoded source ``x_0`` into both. + """ + if x_0_encoded is None: + concat_inputs = (t_encoded, x_encoded, cond_embedding) + conditioning_vec = jnp.concatenate((t_encoded, cond_embedding), axis=-1) + else: + concat_inputs = (t_encoded, x_encoded, x_0_encoded, cond_embedding) + conditioning_vec = jnp.concatenate((t_encoded, x_0_encoded, cond_embedding), axis=-1) + return concat_inputs, conditioning_vec + + def _combine_and_decode( + self, + t_encoded: jnp.ndarray, + x_encoded: jnp.ndarray, + cond_embedding: jnp.ndarray, + squeeze: bool, + train: bool, + x_0_encoded: jnp.ndarray | None = None, + ) -> jnp.ndarray: + """Combine ``(t, x, [x_0,] condition)`` per the conditioning mode, decode, and project. + + Override in a subclass to add new ``conditioning`` modes, delegating to + ``super()._combine_and_decode`` for the built-in ones. ``x_0_encoded`` is the encoded + GENOT source (folded into the conditioning by :meth:`_conditioning_signals`). + """ + concat_inputs, conditioning_vec = self._conditioning_signals(t_encoded, x_encoded, cond_embedding, x_0_encoded) if self.conditioning == "concatenation": - out = jnp.concatenate((t_encoded, x_encoded, cond_embedding), axis=-1) + out = jnp.concatenate(concat_inputs, axis=-1) elif self.conditioning == "film": - out = self.film_block(x_encoded, jnp.concatenate((t_encoded, cond_embedding), axis=-1)) + out = self.film_block(x_encoded, conditioning_vec) elif self.conditioning == "resnet": - out = self.resnet_block(x_encoded, jnp.concatenate((t_encoded, cond_embedding), axis=-1)) + out = self.resnet_block(x_encoded, conditioning_vec) else: raise ValueError(f"Unknown conditioning mode: {self.conditioning}.") out = self.decoder(out, training=train) - return self.output_layer(out), cond_mean, cond_logvar + return self.output_layer(out) def _maybe_null_embedding( self, @@ -278,17 +439,13 @@ def _maybe_null_embedding( defaults) this is a no-op that reproduces the standard conditional behavior byte-for-byte: no random number is drawn and the embedding is returned as-is. """ - if force_uncond: - return jnp.zeros_like(cond_embedding) - if train and self.condition_dropout_prob > 0.0: - drop_rng = self.make_rng("dropout") - keep = jax.random.bernoulli( - drop_rng, - p=1.0 - self.condition_dropout_prob, - shape=(cond_embedding.shape[0], 1), - ) - return jnp.where(keep, cond_embedding, jnp.zeros_like(cond_embedding)) - return cond_embedding + return null_condition_embedding( + cond_embedding, + condition_dropout_prob=self.condition_dropout_prob, + make_rng=self.make_rng, + train=train, + force_uncond=force_uncond, + ) def _maybe_null_input( self, @@ -306,14 +463,14 @@ def _maybe_null_input( :attr:`condition_dropout_prob` during training. With the defaults it returns ``cond`` unchanged and draws no random number. """ - if force_uncond: - return jax.tree_util.tree_map(lambda c: jnp.full_like(c, self.mask_value), cond) - if train and self.condition_dropout_prob > 0.0: - drop_rng = self.make_rng("dropout") - n = next(iter(cond.values())).shape[0] - keep = jax.random.bernoulli(drop_rng, p=1.0 - self.condition_dropout_prob, shape=(n, 1, 1)) - return jax.tree_util.tree_map(lambda c: jnp.where(keep, c, jnp.full_like(c, self.mask_value)), cond) - return cond + return null_condition_input( + cond, + condition_dropout_prob=self.condition_dropout_prob, + mask_value=self.mask_value, + make_rng=self.make_rng, + train=train, + force_uncond=force_uncond, + ) def get_condition_embedding(self, condition: dict[str, jnp.ndarray]) -> tuple[jnp.ndarray, jnp.ndarray]: """Get the embedding of the condition. @@ -529,9 +686,15 @@ def _normalize_vf_kwargs(vf_kwargs: dict[str, Any] | None) -> dict[str, Any]: """ if vf_kwargs is None: return {"genot_source_dims": [1024, 1024, 1024], "genot_source_dropout": 0.0} - assert isinstance(vf_kwargs, dict) - assert "genot_source_dims" in vf_kwargs - assert "genot_source_dropout" in vf_kwargs + if not isinstance(vf_kwargs, dict): + raise TypeError(f"`vf_kwargs` must be a dict or None, got {type(vf_kwargs).__name__}.") + allowed = {"genot_source_dims", "genot_source_dropout"} + unknown = set(vf_kwargs) - allowed + if unknown: + raise ValueError(f"Unexpected `vf_kwargs` keys {sorted(unknown)}; allowed: {sorted(allowed)}.") + missing = allowed - set(vf_kwargs) + if missing: + raise ValueError(f"Missing `vf_kwargs` keys {sorted(missing)}; required: {sorted(allowed)}.") return vf_kwargs def setup(self): @@ -556,13 +719,7 @@ def setup(self): self.layer_cond_output_dropout = nn.Dropout(rate=self.cond_output_dropout) self.layer_norm_condition = nn.LayerNorm() if self.layer_norm_before_concatenation else lambda x: x - self.time_encoder = MLPBlock( - dims=self.time_encoder_dims, - act_fn=self.act_fn, - dropout_rate=self.time_encoder_dropout, - act_last_layer=False, - ) - self.layer_norm_time = nn.LayerNorm() if self.layer_norm_before_concatenation else lambda x: x + self._setup_time() self.x_encoder = MLPBlock( dims=self.hidden_dims, @@ -588,22 +745,7 @@ def setup(self): self.output_layer = nn.Dense(self.output_dim) - if self.conditioning == "film": - self.film_block = FilmBlock( - input_dim=self.hidden_dims[-1], - cond_dim=self.time_encoder_dims[-1] + self.condition_embedding_dim, - **conditioning_kwargs, - ) - elif self.conditioning == "resnet": - self.resnet_block = ResNetBlock( - input_dim=self.hidden_dims[-1], - **self.conditioning_kwargs, - ) - elif self.conditioning == "concatenation": - if len(conditioning_kwargs) > 0: - raise ValueError("If `conditioning=='concatenation' mode, no conditioning kwargs can be passed.") - else: - raise ValueError(f"Unknown conditioning mode: {self.conditioning}") + self._setup_conditioning(conditioning_kwargs) def __call__( self, @@ -613,14 +755,19 @@ def __call__( cond: dict[str, jnp.ndarray], encoder_noise: jnp.ndarray, train: bool = True, + force_uncond: bool = False, ): squeeze = x_t.ndim == 1 + if self.condition_null == "mask_value": + cond = self._maybe_null_input(cond, train=train, force_uncond=force_uncond) cond_mean, cond_logvar = self.condition_encoder(cond, training=train) if self.condition_mode == "deterministic": cond_embedding = cond_mean else: cond_embedding = cond_mean + encoder_noise * jnp.exp(cond_logvar / 2.0) cond_embedding = self.layer_cond_output_dropout(cond_embedding, deterministic=not train) + if self.condition_null == "zero_embedding": + cond_embedding = self._maybe_null_embedding(cond_embedding, train=train, force_uncond=force_uncond) t_encoded = sinusoidal_time_encoder(t, time_freqs=self.time_freqs, time_max_period=self.time_max_period) t_encoded = self.time_encoder(t_encoded, training=train) x_encoded = self.x_encoder(x_t, training=train) @@ -636,17 +783,8 @@ def __call__( elif cond_embedding.shape[0] != x_t.shape[0]: # type: ignore[attr-defined] cond_embedding = jnp.tile(cond_embedding, (x_t.shape[0], 1)) - if self.conditioning == "concatenation": - out = jnp.concatenate((t_encoded, x_encoded, x_0_encoded, cond_embedding), axis=-1) - elif self.conditioning == "film": - out = self.film_block(x_encoded, jnp.concatenate((t_encoded, x_0_encoded, cond_embedding), axis=-1)) - elif self.conditioning == "resnet": - out = self.resnet_block(x_encoded, jnp.concatenate((t_encoded, x_0_encoded, cond_embedding), axis=-1)) - else: - raise ValueError(f"Unknown conditioning mode: {self.conditioning}.") - - out = self.decoder(out, training=train) - return self.output_layer(out), cond_mean, cond_logvar + out = self._combine_and_decode(t_encoded, x_encoded, cond_embedding, squeeze, train, x_0_encoded=x_0_encoded) + return out, cond_mean, cond_logvar def create_train_state( self, diff --git a/src/cellflow/solvers/_base.py b/src/cellflow/solvers/_base.py index 9d0dc8ed..7d911313 100644 --- a/src/cellflow/solvers/_base.py +++ b/src/cellflow/solvers/_base.py @@ -35,6 +35,9 @@ def __init__( self.probability_path = probability_path self.time_sampler = time_sampler self._predict_fn_cache: dict[Any, Any] = {} + # Separate cache for the flat (condition-batched) predict fn — takes a per-cell embedding and + # vmaps over the concatenated cells of all conditions at once. See `OTFlowMatching.predict`. + self._flat_predict_fn_cache: dict[Any, Any] = {} @property def _inference_state(self) -> train_state.TrainState: diff --git a/src/cellflow/solvers/_genot.py b/src/cellflow/solvers/_genot.py index 93b162bc..11049ae8 100644 --- a/src/cellflow/solvers/_genot.py +++ b/src/cellflow/solvers/_genot.py @@ -17,6 +17,7 @@ from cellflow._types import ArrayLike from cellflow.model._utils import _multivariate_normal from cellflow.solvers._base import BaseSolver +from cellflow.solvers._otfm import ClassifierFreeGuidance, Guidance, VelocityFn __all__ = ["GENOT"] @@ -79,11 +80,13 @@ def __init__( target_dim: int, time_sampler: Callable[[jax.Array, int], jnp.ndarray] = solver_utils.uniform_sampler, latent_noise_fn: (Callable[[jax.Array, tuple[int, ...]], jnp.ndarray] | None) = None, + guidance: Guidance | None = None, **kwargs: Any, ): super().__init__(vf, probability_path, time_sampler) self.data_match_fn = jax.jit(data_match_fn) self.source_dim = source_dim + self.guidance = guidance if latent_noise_fn is None: latent_noise_fn = functools.partial(_multivariate_normal, dim=target_dim) self.latent_noise_fn = latent_noise_fn @@ -265,22 +268,70 @@ def predict( x_pred = self._predict_jit(x, condition, rng, rng_genot, **kwargs) return np.array(x_pred) + @property + def cfg_enabled(self) -> bool: + """Whether classifier-free guidance is available at predict time. + + Guidance needs a meaningful unconditional velocity ``v_null``, which only exists + when the velocity field was trained with condition dropout + (``condition_dropout_prob > 0``). When ``False``, a per-call ``guidance_scale`` is + ignored (with a warning) and the plain conditional velocity is used. + """ + return float(getattr(self.vf, "condition_dropout_prob", 0.0)) > 0.0 + + def _base_velocity(self) -> VelocityFn: + """Base (conditional) velocity closure for the predict path. + + Signature ``(t, x, args, force_uncond=False) -> velocity`` with ``args`` being + ``(params, x_0, condition, encoder_noise)``. Nulling only the condition (``x_0`` is + kept) gives the unconditional source→target velocity used by guidance. + """ + + def vf( + t: float, + x: jnp.ndarray, + args: tuple[Any, jnp.ndarray, dict[str, jnp.ndarray], jnp.ndarray], + force_uncond: bool = False, + ) -> jnp.ndarray: + params, x_0, condition, encoder_noise = args + return self.vf_state.apply_fn( + {"params": params}, t, x, x_0, condition, encoder_noise, train=False, force_uncond=force_uncond + )[0] + + return vf + def _get_predict_fn(self, kwargs_frozen: frozen_dict.FrozenDict) -> Callable: """Build and cache a jit+vmap predict function for the given diffrax kwargs. - The returned function is created once per unique set of diffrax kwargs, - then reused on subsequent calls. + The base velocity from :meth:`_base_velocity` is wrapped by guidance when it + applies: a per-call ``guidance_scale != 1.0`` (popped here, requires + :attr:`cfg_enabled`) builds a :class:`~cellflow.solvers.ClassifierFreeGuidance` + for this call, overriding the construction-time ``guidance``; otherwise the + construction-time ``guidance`` is used (``None`` = plain conditional velocity). + The returned function is created once per unique set of kwargs, then reused. """ if kwargs_frozen in self._predict_fn_cache: return self._predict_fn_cache[kwargs_frozen] kwargs = dict(kwargs_frozen) + guidance_scale = float(kwargs.pop("guidance_scale", 1.0)) + guidance = self.guidance + if guidance_scale != 1.0: + if self.cfg_enabled: + # v = v_null + scale·(v_cond − v_null); overrides construction-time guidance. + guidance = ClassifierFreeGuidance(scale=guidance_scale) + else: + warnings.warn( + f"guidance_scale={guidance_scale} ignored: the velocity field was not trained " + "with classifier-free guidance (condition_dropout_prob == 0), so the " + "unconditional velocity is undefined. Using the plain conditional velocity.", + stacklevel=2, + ) + guidance = None - def vf( - t: float, x: jnp.ndarray, args: tuple[Any, jnp.ndarray, dict[str, jnp.ndarray], jnp.ndarray] - ) -> jnp.ndarray: - params, x_0, condition, encoder_noise = args - return self.vf_state.apply_fn({"params": params}, t, x, x_0, condition, encoder_noise, train=False)[0] + vf = self._base_velocity() + if guidance is not None: + vf = guidance.wrap(vf) def solve_ode( params: Any, diff --git a/src/cellflow/solvers/_otfm.py b/src/cellflow/solvers/_otfm.py index 3886dff2..2b1eacf9 100644 --- a/src/cellflow/solvers/_otfm.py +++ b/src/cellflow/solvers/_otfm.py @@ -28,12 +28,15 @@ class Guidance(Protocol): """Pluggable transform applied to the base velocity field on the predict path. - A guidance strategy receives the base (conditional) velocity closure and the - inference train state, and returns a new velocity closure with the same - ``(t, x, args) -> velocity`` signature. + A guidance strategy receives the base velocity closure ``vf(t, x, args, + force_uncond=False)`` — which owns the field's call signature and returns the + conditional (``force_uncond=False``) or unconditional (``True``) velocity — and + returns a plain ``(t, x, args) -> velocity`` closure. Taking the closure (rather + than the train state) keeps guidance agnostic to solver-specific ``args`` such as + GENOT's source ``x_0``. """ - def wrap(self, vf: VelocityFn, inference_state: train_state.TrainState) -> VelocityFn: + def wrap(self, vf: Callable) -> VelocityFn: """Wrap the base velocity ``vf`` and return the guided velocity.""" ... @@ -80,22 +83,18 @@ def from_ode_weight(cls, cfg_ode_weight: float) -> "ClassifierFreeGuidance": raise ValueError("cfg_ode_weight must be non-negative.") return cls(scale=1.0 + cfg_ode_weight) - def wrap(self, vf: VelocityFn, inference_state: train_state.TrainState) -> VelocityFn: - """Return a velocity closure computing ``v_null + scale * (v_cond - v_null)``.""" + def wrap(self, vf: Callable) -> VelocityFn: + """Return a velocity closure computing ``v_null + scale * (v_cond - v_null)``. + + ``vf`` is the base velocity ``vf(t, x, args, force_uncond=False)``; it owns the + field's call signature, so this blend is agnostic to solver-specific ``args`` + (e.g. GENOT threads its source ``x_0`` through ``args``). + """ scale = self.scale def guided_vf(t: jnp.ndarray, x: jnp.ndarray, args: tuple[Any, ...]) -> jnp.ndarray: - params, condition, encoder_noise = args - v_cond = vf(t, x, args) - v_null = inference_state.apply_fn( - {"params": params}, - t, - x, - condition, - encoder_noise, - train=False, - force_uncond=True, - )[0] + v_cond = vf(t, x, args, force_uncond=False) + v_null = vf(t, x, args, force_uncond=True) return v_null + scale * (v_cond - v_null) return guided_vf @@ -282,9 +281,33 @@ def _base_velocity(self) -> VelocityFn: inference velocity field conditionally (``force_uncond=False``). """ - def vf(t: jnp.ndarray, x: jnp.ndarray, args: tuple[Any, dict[str, jnp.ndarray], jnp.ndarray]) -> jnp.ndarray: + def vf( + t: jnp.ndarray, + x: jnp.ndarray, + args: tuple[Any, dict[str, jnp.ndarray], jnp.ndarray], + force_uncond: bool = False, + ) -> jnp.ndarray: params, condition, encoder_noise = args - return self.vf_state_inference.apply_fn({"params": params}, t, x, condition, encoder_noise, train=False)[0] + return self.vf_state_inference.apply_fn( + {"params": params}, t, x, condition, encoder_noise, train=False, force_uncond=force_uncond + )[0] + + return vf + + def _base_velocity_from_embedding(self) -> Callable: + """Velocity closure taking a *precomputed* condition embedding in ``args``. + + ``args`` is ``(params, cond_embedding)``. Used on the predict path when no guidance is active, so + the condition encoder runs once (outside the ODE, in :meth:`_get_predict_fn`) rather than on + every integration step; the embedding is constant along the trajectory, so the result is + identical to :meth:`_base_velocity`. + """ + + def vf(t: jnp.ndarray, x: jnp.ndarray, args: tuple[Any, jnp.ndarray]) -> jnp.ndarray: + params, cond_embedding = args + return self.vf_state_inference.apply_fn( + {"params": params}, t, x, cond_embedding, train=False, method="velocity_from_embedding" + ) return vf @@ -327,28 +350,118 @@ def _get_predict_fn(self, kwargs_frozen: frozen_dict.FrozenDict) -> Callable: ) guidance = None - vf = self._base_velocity() - if guidance is not None: - vf = guidance.wrap(vf, self.vf_state_inference) + if guidance is None: + # No guidance: encode the condition once, then integrate the embedding-only rhs — the + # condition encoder never runs inside the ODE. Identical output to the per-step path. + vf_emb = self._base_velocity_from_embedding() - def solve_ode( - params: Any, x: jnp.ndarray, condition: dict[str, jnp.ndarray], encoder_noise: jnp.ndarray - ) -> jnp.ndarray: - ode_term = diffrax.ODETerm(vf) + def solve_ode_emb(params: Any, x: jnp.ndarray, cond_embedding: jnp.ndarray) -> jnp.ndarray: + result = diffrax.diffeqsolve( + diffrax.ODETerm(vf_emb), t0=0.0, t1=1.0, y0=x, args=(params, cond_embedding), **kwargs + ) + return result.ys[0] + + vmapped = jax.vmap(solve_ode_emb, in_axes=[None, 0, None]) + + def fn( + params: Any, x: jnp.ndarray, condition: dict[str, jnp.ndarray], encoder_noise: jnp.ndarray + ) -> jnp.ndarray: + cond_embedding = self.vf_state_inference.apply_fn( + {"params": params}, condition, encoder_noise, train=False, method="encode_condition" + )[0] + return vmapped(params, x, cond_embedding) + + fn = jax.jit(fn) + else: + # Guidance wraps the full velocity (needs both conditional and null embeddings per step), so + # keep encoding inside the ODE. + vf = guidance.wrap(self._base_velocity()) + + def solve_ode( + params: Any, x: jnp.ndarray, condition: dict[str, jnp.ndarray], encoder_noise: jnp.ndarray + ) -> jnp.ndarray: + result = diffrax.diffeqsolve( + diffrax.ODETerm(vf), t0=0.0, t1=1.0, y0=x, args=(params, condition, encoder_noise), **kwargs + ) + return result.ys[0] + + fn = jax.jit(jax.vmap(solve_ode, in_axes=[None, 0, None, None])) + + self._predict_fn_cache[kwargs_frozen] = fn + return fn + + def _get_flat_predict_fn(self, kwargs_frozen: frozen_dict.FrozenDict) -> Callable: + """Build/cache the condition-batched predict fn. + + One ``jit(vmap)`` solve over the *concatenated* cells of all conditions, each cell carrying its + own precomputed embedding (``in_axes=[None, 0, 0]``). Only valid with no guidance active (guidance + needs the encoder inside the ODE), so it integrates the embedding-only rhs. + """ + if kwargs_frozen in self._flat_predict_fn_cache: + return self._flat_predict_fn_cache[kwargs_frozen] + + kwargs = dict(kwargs_frozen) + vf_emb = self._base_velocity_from_embedding() + + def solve_ode_emb(params: Any, x: jnp.ndarray, cond_embedding: jnp.ndarray) -> jnp.ndarray: result = diffrax.diffeqsolve( - ode_term, - t0=0.0, - t1=1.0, - y0=x, - args=(params, condition, encoder_noise), - **kwargs, + diffrax.ODETerm(vf_emb), t0=0.0, t1=1.0, y0=x, args=(params, cond_embedding), **kwargs ) return result.ys[0] - fn = jax.jit(jax.vmap(solve_ode, in_axes=[None, 0, None, None])) - self._predict_fn_cache[kwargs_frozen] = fn + fn = jax.jit(jax.vmap(solve_ode_emb, in_axes=[None, 0, 0])) + self._flat_predict_fn_cache[kwargs_frozen] = fn return fn + def _predict_batched( + self, + x: dict[str, ArrayLike], + condition: dict[str, dict[str, ArrayLike]], + rng: jax.Array | None = None, + **kwargs: Any, + ) -> dict[str, ArrayLike]: + """Predict every condition in one condition-batched ODE solve (no-guidance path). + + Encodes each condition once, concatenates all conditions' source cells into a single batch (each + cell tagged with its condition's embedding), integrates them in one vmapped solve, then splits the + result back per condition. Differing per-condition cell counts are handled by concatenation (no + padding). Numerically identical to the per-condition loop: each cell's ODE is independent. + """ + kwargs.setdefault("dt0", None) + kwargs.setdefault("solver", diffrax.Tsit5()) + kwargs.setdefault("stepsize_controller", diffrax.PIDController(rtol=1e-5, atol=1e-5)) + kwargs.pop("guidance_scale", None) # == 1.0 on this path (guaranteed by the caller) + kwargs_frozen = frozen_dict.freeze(kwargs) + + keys = list(x) + # One encoder-noise draw shared across conditions (matches the per-condition loop under one rng). + noise_dim = (1, self.vf.condition_embedding_dim) + use_mean = rng is None or self.condition_encoder_mode == "deterministic" + rng = utils.default_prng_key(rng) + encoder_noise = jnp.zeros(noise_dim) if use_mean else jax.random.normal(rng, noise_dim) + + params = self.vf_state_inference.params + groups = list(condition[keys[0]]) + cond_stacked = {g: jnp.concatenate([jnp.asarray(condition[k][g]) for k in keys], axis=0) for g in groups} + cond_embedding = self.vf_state_inference.apply_fn( + {"params": params}, cond_stacked, encoder_noise, train=False, method="encode_condition" + )[0] # (n_conditions, embedding_dim) + + sizes = [int(jnp.asarray(x[k]).shape[0]) for k in keys] + x_flat = jnp.concatenate([jnp.asarray(x[k]) for k in keys], axis=0) + emb_flat = jnp.concatenate( + [jnp.broadcast_to(cond_embedding[i : i + 1], (sizes[i], cond_embedding.shape[-1])) for i in range(len(keys))], + axis=0, + ) + out_flat = self._get_flat_predict_fn(kwargs_frozen)(params, x_flat, emb_flat) + + out: dict[str, ArrayLike] = {} + offset = 0 + for k, n in zip(keys, sizes, strict=True): + out[k] = out_flat[offset : offset + n] + offset += n + return out + def _predict_jit( self, x: ArrayLike, @@ -418,7 +531,14 @@ def predict( return {} if isinstance(x, dict): - jax_results = {k: self._predict_jit(x[k], condition[k], rng, **kwargs) for k in x} + # With no guidance, batch every condition into one condition-batched solve (encode once per + # condition, one vmapped kernel over all cells). Guidance stays on the per-condition loop + # since it needs the conditional + null embeddings inside the ODE. Same result either way. + guidance_scale = float(kwargs.get("guidance_scale", 1.0)) + if self.guidance is None and guidance_scale == 1.0: + jax_results = self._predict_batched(x, condition, rng, **kwargs) + else: + jax_results = {k: self._predict_jit(x[k], condition[k], rng, **kwargs) for k in x} return {k: np.array(v) for k, v in jax_results.items()} else: x_pred = self._predict_jit(x, condition, rng, **kwargs) diff --git a/src/cellflow/training/_trainer.py b/src/cellflow/training/_trainer.py index f0130690..a2c4ef1d 100644 --- a/src/cellflow/training/_trainer.py +++ b/src/cellflow/training/_trainer.py @@ -6,7 +6,7 @@ from numpy.typing import ArrayLike from tqdm import tqdm -from cellflow.data._dataloader import OOCTrainSampler, TrainSampler, ValidationSampler +from cellflow.data._legacy import OOCTrainSampler, TrainSampler, ValidationSampler from cellflow.solvers import _genot, _otfm from cellflow.training._callbacks import BaseCallback, CallbackRunner diff --git a/src/dagloader/README.md b/src/dagloader/README.md new file mode 100644 index 00000000..32b3daed --- /dev/null +++ b/src/dagloader/README.md @@ -0,0 +1,252 @@ +# `dagloader` — declarative, index-free sampling over annbatch + +A small, declarative layer that turns a **scheme of columns + weights** into a stream of matched +`{source, target, condition}` training batches, read (in- or out-of-core) through +[`annbatch`](https://github.com/laminlabs/annbatch). It is the streaming sampler that both +**cellflow** (this repo's model) and **sc-flow-tools** (`theislab/sc-flow-tools`) can share: one +condition per batch, source↔target matching, categorical + continuous condition payloads, and +reproducible weighted sampling — expressed as data, not as materialized row-index tables. + +``` +Scheme (sources + nodes + binds + weights) + SamplerConfig (batch/chunk/preload) + │ │ + └───────────────────► DAGLoader ◄────────────┘ + │ + {"target": X[perturbed], "source": X[matched control], "condition": emb} + │ + ├─ every node streams through its own ScheduledClassSampler + annbatch.Loader + └─ root category ~ weights; each bound child's category derived from the parent's +``` + +Structure (`Scheme`) and read parameters (`SamplerConfig`) are separate objects, both handed to the +loader — so the same scheme can run with different batch/chunk settings, and neither lives on the +`Node`. + +--- + +## 1. Mental model + +- **Source** — a named cell store: an in-memory `AnnData` **or** an out-of-core + `annbatch.DatasetCollection`. The loader never branches on which; obs is read for grouping, cells + are read only when a batch is materialized. +- **Node** — a partition of *one source's* cells into **leaves** = the unique combinations of its + `cols`, plus a `{combo: weight}` mapping. `cols` are the grouping / condition columns; `key` is the + single representation to stream (`"X"`, `"obsm/X_pca"`, `"layers/…"`) — add another node for another + rep. The `Node` carries no read parameters; those live in `SamplerConfig` (below). +- **Weight = the selection.** A combination with weight `0` (or absent from the mapping) is simply + not sampled. There is **no separate `select` mask** — inclusion *is* a positive weight. This is + native to annbatch's `ClassSampler`. +- **Bind(parent, child, common)** — condition the child on the parent: each batch, the child's leaf + is derived from the parent's leaf by matching the `common` column values. With `common` = the + context (cell line, …), the child (control/source) is drawn from the **same context** as the parent + (perturbed/target). This is the source↔target matching. +- **Scheme** — sources + nodes + a rooted tree of binds + the reproducibility `seed`. Pure structure. + The **root** is the streamed target; bound children are the sources. Batches-per-pass is *not* set + here — the loader derives it from the root (a natural epoch, `root_n_obs // batch_size`). +- **SamplerConfig** — the annbatch read parameters (`batch_size`, `chunk_size`, `preload_nchunks`), + kept off the `Node` and the `Scheme`. Passed to the loader as one config (all nodes) or a + `{node_name: SamplerConfig}` mapping (per-node). Nodes may use **different** `batch_size`\s (source + and target row counts need not match); the root's drives the shared batch count. +- **DAGLoader** — iterates batches. Each node streams through its own `ScheduledClassSampler`. + +### Batch contract + +```python +{ + "target": ndarray (B_t, d), # root node's streamed rep (Node.key); B_t = root batch_size + "source": ndarray (B_s, d), # bound child's streamed rep — matched to target's context + "condition": ndarray (B_t, e), # condition_fn(leaf) broadcast over the target batch (optional) +} +``` +Each node's row count is its own `SamplerConfig.batch_size` (`B_t` for the target, `B_s` for the +source) — they need not be equal. Sparse `X` stays sparse. + +--- + +## 2. Sampling schemes + +The "sampling scheme" is entirely the **weights** plus the per-node sampler cadence — nothing is +hardcoded. + +**Weight builders** (plain functions returning `{combo: weight}`; use them or write the dict yourself): + +| helper | meaning | +|---|---| +| `uniform(combos)` | every combination equally likely | +| `frequency(counts)` | sample ∝ cell count (favor abundant conditions) | +| `inverse_frequency(counts)` | sample ∝ 1 / cell count (balance rare vs abundant) | + +**`ScheduledClassSampler`** — the one genuinely new piece. It is an annbatch `ClassSampler` whose +per-batch category can be *supplied* (`set_schedule`) instead of drawn internally. It does **not copy** +annbatch's chunk math: `ClassSampler._iter_requests` makes exactly one *weighted* draw +(`rng.choice(n_classes, size=n_groups, p=…)`) to choose each group's class, so `ScheduledClassSampler` +runs that method unchanged but temporarily wraps `self._rng` so that one call returns the scheduled +positions. Every other draw (run selection, slice starts, shuffling) uses the real generator, so a +scheduled pass is bit-identical to what annbatch would produce for those group classes. That is what +makes the bind work: + +- the **root** schedule is drawn from the root's weights (one category per batch, `_n_batches` total); +- each **bound child**'s schedule is *derived* from the root's (parent leaf → `common` value → + matching child leaf; the child's RNG breaks ties / picks a fallback); +- the loaders then zip batch-for-batch with **no per-step reconfiguration**. + +With `schedule=None`, `ScheduledClassSampler` is exactly `ClassSampler` (draws ∝ weights) — a strict +superset. (If annbatch exposed a `_group_positions(n_groups)` hook, even the rng-wrapping would go +away and this would be a one-method override — an upstream candidate.) + +**Read parameters — `SamplerConfig` (batch / chunk / preload):** + +- `batch_size` (`B`): rows per emitted batch for that node. Nodes may differ (target rows need not + equal source rows) — every node still draws the same *number* of batches (the root's), so the + schedules zip; only the per-batch row count differs. +- `chunk_size` (default `1`): annbatch read-slice size. `1` = per-row reads (any layout). `>1` = + contiguous chunked reads for higher on-disk throughput; it assumes each sampled leaf is a contiguous + run ≥ `chunk_size` (which a condition-sorted collection provides) and must divide `batch_size`. +- `preload_nchunks` (default `batch_size // chunk_size`): chunks per annbatch read window; a larger + multiple packs more classes into one read. + +Pass one `SamplerConfig` (applied to every node) or a `{node_name: SamplerConfig}` mapping (per-node — +e.g. a chunked root and a per-row control). + +**Cadence / reproducibility (on the `Scheme`):** + +- batches per pass: **derived**, not configured — the loader uses a natural epoch over the root + (`root_n_obs // root batch_size`) and restarts each pass → effectively infinite with a fixed, + reproducible restart cadence. All nodes restart together (the root drives the count). +- `seed`: per-node RNG streams are spawned from one `SeedSequence(seed)` (one independent stream per + node, by sorted name) so nodes never correlate and the full `(source, target)` sequence is + reproducible from the seed. + +> **On layout / sortedness:** the default (`chunk_size=1`) reads per-row and is indifferent to on-disk +> order. `chunk_size>1` is purely a throughput opt-in for a condition-sorted collection; the loader +> does not police layout — it forwards annbatch's own behavior. Ordering never affects *correctness* +> (matching is recovered from the schedule / columns, not from row order). + +--- + +## 3. The cellflow case + +cellflow trains a conditional flow from a **control** population to a **perturbed** population, +matched within a context, one perturbation condition per batch. That maps directly onto a two-node +scheme built by `perturbation_scheme` (`_schemes.py`): + +| cellflow concept | dagloader | +|---|---| +| `sample_rep` (`X` / an obsm key) | `Node.key` | +| `split_covariates` (context / grouping) | `context` → the `Bind.common` columns | +| `perturbation_covariates` columns | `perturbation` → the extra `cols` on both nodes | +| `control_key` / which cells are control | `control_values` → which combos land in the `ctrl` node | +| control = same group as target | `Bind("pert", "ctrl", common=context)` | +| one perturbation per batch, weighted | root `pert` node weights + `ScheduledClassSampler` | +| `batch_size` | `SamplerConfig.batch_size` | +| `perturbation_covariate_reps` (embeddings) | `condition_fn(leaf) -> embedding` | + +```python +from dagloader import DAGLoader, SamplerConfig, perturbation_scheme + +scheme = perturbation_scheme( + adata_or_collection, + context=["cell_line"], # split_covariates + perturbation=["drug"], # perturbation_covariates + control_values={"drug": "control"}, + key="X", # sample_rep + seed=0, +) +loader = DAGLoader( + scheme, + SamplerConfig(batch_size=256, chunk_size=1, preload_nchunks=256), # all read params explicit (no hidden defaults) + condition_fn=lambda leaf: DRUG_EMB[leaf[-1]], +) +batch = next(loader) # {"target", "source", "condition"} — source is a matched-cell-line control +``` + +**How cellflow uses it.** The collection/streaming path is **additive** to cellflow's in-memory +`adata` + `perturbation_covariates` API — add `perturbation_scheme` + `DAGLoader` as a new +dataloader alongside the existing ones (see `docs/perturbation_combinations.md` → "minimal churn"). +The condition encoding (categorical embeddings, combinations, pooling) stays in the model; the loader +only needs a `condition_fn` that turns a leaf tuple into the per-condition embedding. + +--- + +## 4. The sc-flow-tools cases + +sc-flow-tools already owns grouping (`HierarchicalIndexer`), matching (`control_values_dict`, +`matched_keys`), representations, and an in-memory multi-node `TrainSampler`. The one thing it lacks +is **out-of-core streaming**. `DAGLoader` is exactly that streaming sampler: its scheme +reproduces what a `DataManager` config already describes, reading cells from a `DatasetCollection` +instead of holding them in memory. + +| sc-flow-tools concept | dagloader | +|---|---| +| `HierarchicalIndexer` grouping (obs Categorical codes) | `Node.cols` → leaves (obs-only, computed at build) | +| `control_values_dict` (which values are control) | the `ctrl` node's weighted combos (zero-weight = excluded) | +| default same-context coupling (one-to-many) | `Bind(..., common=context)` | +| `matched_keys` (explicit fixed source→target pairing) | a `Bind` whose `common` covers the pairing columns, so parent leaf → child leaf is forced | +| `sample_rep` in PCA space (an obsm key) | `Node.key = "obsm/X_pca"` | +| `*_reps` (uns embedding tables) | `condition_fn(leaf)` (concatenate per-covariate embeddings) | +| weighted / multi-node `TrainSampler` | root weights + a rooted tree of binds | + +**Direct scheme construction** (the sc-flow-tools style — assemble nodes from a `DataManager` config): + +```python +from dagloader import Bind, DAGLoader, Node, SamplerConfig, Scheme, uniform + +# grouping columns: context (cell_line) + perturbation (drug); state read from a PCA obsm key +cols = ("cell_line", "drug") +scheme = Scheme( + sources={"data": collection}, + nodes={ + "target": Node("data", cols, key="obsm/X_pca", weights=uniform(perturbed_combos)), + "control": Node("data", cols, key="obsm/X_pca", weights=uniform(control_combos)), + }, + root="target", + binds=(Bind("target", "control", common=("cell_line",)),), # control_values_dict + same-context + seed=0, +) +loader = DAGLoader( + scheme, + SamplerConfig(batch_size=256, chunk_size=1, preload_nchunks=256), # per-node: {"target": SamplerConfig(batch_size=256, chunk_size=32, preload_nchunks=8), "control": SamplerConfig(batch_size=256, chunk_size=1, preload_nchunks=256)} + condition_fn=lambda leaf: concat(cell_emb[leaf[0]], drug_emb[leaf[1]]), # *_reps combination +) +batch = next(loader) # {"target", "source", "condition"} +``` + +**How sc-flow-tools uses it.** Per the shared-data-layer plan (`docs/shared-data-layer-status.md`), +sc-flow-tools hosts the shared data layer and CellFlow2 depends on it. In that plan the tree nodes +carry **row indices** (obs-only), and the annbatch loader (this package's `DAGLoader`) streams +the sampled indices from a `DatasetCollection`. The `TrainSampler`'s node dispatch becomes a scheme: +each node → a `Node`, the matched-subpopulation tree → the binds, `control_values_dict` → weights, +`matched_keys` → binds with pairing columns in `common`. + +### What is and isn't in this prototype + +- **In:** single-level bind (root + one bound child) matched on shared columns; a single streamed rep + per node (`X` / obsm / layer); categorical condition via `condition_fn`; in-memory and out-of-core + sources; weighted per-condition sampling; per-node `SamplerConfig`; reproducible per-node RNG. +- **Documented mapping, not yet wired:** multi-child trees emit a single `source` today (the batch + contract would grow a `sources[child_name]` map); **per-cell continuous covariates / Gromov + `state_lin`+`state_quad`** (multiple aligned reps for the same streamed rows) would need annbatch to + return several reps per batch, or a second aligned read — deliberately left out rather than + reintroduce a bespoke row-gather sampler. + +--- + +## 5. Files + +| file | contents | +|---|---| +| `_schema.py` | `Node`, `Bind`, `Scheme`, `SamplerConfig`, `Weights`, `uniform` / `frequency` / `inverse_frequency` — structural, data-free | +| `_scheduled_sampler.py` | `ScheduledClassSampler` (reuses annbatch `_iter_requests` via an rng wrapper; the upstream candidate) | +| `_io.py` | container-agnostic obs / leaf-code / rep-backing helpers (AnnData or DatasetCollection; X / obsm / layers) | +| `_loader.py` | `DAGLoader` (config resolution, per-pass scheduling, bind derivation, batch assembly) | +| `_schemes.py` | `perturbation_scheme` factory (the cellflow-shaped case) | + +Tests live in `tests/dagloader/`, organized by the cases above: + +| test file | case | +|---|---| +| `test_scheduled_sampler.py` | `ScheduledClassSampler`: schedule adherence, chunk coherence, `schedule=None` ≡ `ClassSampler`, validation | +| `test_sampling_schemes.py` | weights as the selection; `uniform` / `frequency` / `inverse_frequency` empirical distributions; per-node independent RNG; reproducibility; `Node` / `Scheme` / `SamplerConfig` validation | +| `test_cellflow_case.py` | `perturbation_scheme`: matched control↔perturbed batches, condition, obsm rep, on-disk collection, `chunk_size>1` | +| `test_scflow_cases.py` | direct scheme construction: multi-covariate condition, explicit matched pairing via `common`, obsm rep as state, out-of-core | +| `test_train.py` | end-to-end: scheme → loader → a real flow-matching training step (loss decreases), in-memory + DatasetCollection + obsm + `chunk_size>1` | diff --git a/src/dagloader/__init__.py b/src/dagloader/__init__.py new file mode 100644 index 00000000..dd2c514f --- /dev/null +++ b/src/dagloader/__init__.py @@ -0,0 +1,47 @@ +r"""Declarative, index-free class-mapping sampler over annbatch. + +A :class:`Scheme` is a rooted tree of :class:`Node`\\s over named cell sources; each node partitions +its source's cells into leaves (unique column-combinations) with a per-combination weight mapping. +:class:`DAGLoader` streams matched ``{source, target, condition}`` batches — one condition per +batch — where the root's per-batch category is drawn from its weights (annbatch ``ClassSampler``) and +each bound child replays that schedule via an annbatch ``BoundClassSampler`` (matched on the bind's +shared columns) so the loaders zip batch-for-batch. No row indices are exposed: the scheme is +columns / keys / weights. + +See ``README.md`` (next to this file) for the model, the sampling schemes, and the mapping to +cellflow and sc-flow-tools use-cases. +""" + +from dagloader._eval_loader import DAGEvalLoader +from dagloader._loader import DAGLoader +from dagloader._schema import ( + Bind, + Container, + Node, + SamplerConfig, + Scheme, + Weights, + frequency, + inverse_frequency, + uniform, +) +from dagloader._schemes import perturbation_scheme +from dagloader._split import resolve_split_configs, split_assignment, split_scheme + +__all__ = [ + "Bind", + "Container", + "DAGEvalLoader", + "DAGLoader", + "Node", + "SamplerConfig", + "Scheme", + "Weights", + "frequency", + "inverse_frequency", + "perturbation_scheme", + "resolve_split_configs", + "split_assignment", + "split_scheme", + "uniform", +] diff --git a/src/dagloader/_eval_loader.py b/src/dagloader/_eval_loader.py new file mode 100644 index 00000000..b6345aec --- /dev/null +++ b/src/dagloader/_eval_loader.py @@ -0,0 +1,165 @@ +"""``DAGEvalLoader`` — control-rooted eval reader: a Sequential control *inner* + a bound perturbed target. + +Anchors on the source (control) — "there is no perturbed without a source". Two loaders, both driven by +one deterministic :class:`~annbatch.samplers.SequentialClassSampler` over the control populations: + +* the **source** loader *is* that inner — it reads each scheduled control population **in full** (all its + controls), and +* the **target** loader is an :class:`~annbatch.samplers.BoundClassSampler` on the *same* inner, matched on + the bind's ``common`` (context) columns, that **samples** a perturbed leaf (drug) within each matched + context. + +annbatch does all the class matching; nothing is derived or updated per pass here. The condition of each +batch is the perturbed leaf the bound drew (read from an identically-seeded oracle, so it lines up with +the target loader's own draw). Which control populations are visited — and how many batches — is set via +the schedule (``iter_conditions``). Works over an in-memory ``AnnData`` and a ``DatasetCollection`` alike. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator, Mapping + +import numpy as np +from annbatch import Loader +from annbatch.samplers import BoundClassSampler, SequentialClassSampler + +from dagloader._io import key_backings, leaf_codes, materialize_node, obs_columns +from dagloader._loader import _HAS_CUPY, _flat_categorical +from dagloader._schema import SamplerConfig, Scheme, _weight_vector + +__all__ = ["DAGEvalLoader"] + + +class DAGEvalLoader: + """Yield ``{"source", "target", "condition", "leaf"}`` per control population (source-anchored). + + Parameters + ---------- + scheme + A prepared :class:`~dagloader.Scheme` (root = perturbed/target node, one bound control child). + Its bind's ``common`` columns are the matching context (e.g. ``cell_line``). + sampler_config + Read parameters for the bound perturbed-target sampler (``batch_size`` = target cells per + condition; controls are read in full regardless). + condition_fn + Maps a perturbed leaf (over the target node's ``cols``) to its condition embedding. + seed + Seed for the target's drug sampling (reproducible across ``iter_conditions`` calls). + """ + + def __init__( + self, + scheme: Scheme, + sampler_config: SamplerConfig, + condition_fn: Callable[[tuple], np.ndarray | Mapping[str, np.ndarray]] | None = None, + *, + seed: int = 0, + ) -> None: + self.s = scheme + self._cond_fn = condition_fn + self._cfg = sampler_config + self._seed = seed + binds = [b for b in scheme.binds if b.parent == scheme.root] + if len(binds) != 1: + raise ValueError("DAGEvalLoader expects exactly one bound source of the root.") + b = binds[0] + self._pert = scheme.nodes[scheme.root] # perturbed / target + self._ctrl = scheme.nodes[b.child] # control / source + self._context = b.common + # control source honors Node.in_memory (materialize the control cells into RAM once) + ctrl_src = scheme.sources[self._ctrl.source] + self._src = materialize_node(ctrl_src, self._ctrl) if self._ctrl.in_memory else ctrl_src + self._src_p = scheme.sources[self._pert.source] + + # per-node tuple-labelled categorical + weight vector (obs only). Control/perturbed cells live in + # the same source; the weights mark which leaves belong to each node. + cc, cl = leaf_codes(obs_columns(self._src, self._ctrl.cols), self._ctrl.cols) + pc, pl = leaf_codes(obs_columns(self._src_p, self._pert.cols), self._pert.cols) + self._ctrl_cats = _flat_categorical(cc, cl) + self._pert_cats = _flat_categorical(pc, pl) + self._ctrl_w = _weight_vector(self._ctrl.weights, cl) + self._pert_w = _weight_vector(self._pert.weights, pl) + self._ctrl_leaves = cl + # inner (control) tuple position -> target tuple position, for each shared context column + self._on = {self._ctrl.cols.index(c): self._pert.cols.index(c) for c in self._context} + # the control populations to visit (positive-weight control leaves) + self._ctrl_codes = np.array([i for i in range(len(cl)) if self._ctrl_w[i] > 0], dtype=np.int64) + if self._ctrl_codes.size == 0: + raise ValueError("no control population (positive-weight control leaf) to evaluate.") + + @property + def control_populations(self) -> list[tuple]: + """The control leaves (contexts) this loader iterates over.""" + return [self._ctrl_leaves[i] for i in self._ctrl_codes] + + def _inner(self, schedule: np.ndarray) -> SequentialClassSampler: + # deterministic: same schedule ⇒ same control-population order across the source loader and every + # bound inner, so source[j] and target[j] refer to the same matched context. + return SequentialClassSampler(self._ctrl_cats, schedule=schedule) + + def _bound(self, schedule: np.ndarray) -> BoundClassSampler: + cfg = self._cfg + return BoundClassSampler( + self._inner(schedule), + cfg.chunk_size, + cfg.preload_nchunks, + cfg.batch_size, + classes_to_bind_on=self._pert_cats, + on=self._on, + classes=self._pert_cats, # secondary = the perturbed leaf; weights pick positive-weight drugs + class_weights=self._pert_w, + rng=np.random.default_rng(self._seed), + ) + + + # TODO(selmanozleyen): rename to _loaders_from_node + def _node_loaders(self, src, node, make_sampler) -> dict: + cfg = self._cfg + preload_to_gpu = cfg.preload_to_gpu if cfg.preload_to_gpu is not None else _HAS_CUPY # None ⇒ auto + loaders = {} + for key in node.keys: + loader = Loader( + batch_sampler=make_sampler(), return_index=False, to=cfg.to, preload_to_gpu=preload_to_gpu + ).add_datasets(key_backings(src, key)) + loaders[key] = loader + return loaders + + def iter_conditions(self, n_conditions: int | None = None) -> Iterator[dict]: + """Yield one batch per scheduled control population. + + With ``n_conditions`` set, the control populations are cycled to that many batches (each re-reads + the population's controls and the bound samples a fresh drug); otherwise every control population + is visited once. + """ + if n_conditions is None: + schedule = self._ctrl_codes.copy() + else: + reps = int(np.ceil(n_conditions / self._ctrl_codes.size)) + schedule = np.tile(self._ctrl_codes, reps)[:n_conditions].astype(np.int64) + + # condition oracle: identically-seeded bound → its per-batch drawn perturbed leaf lines up with + # the target loader's own draw (annbatch reproduces the same class sequence from the same seed). + oracle = self._bound(schedule) + vocab = oracle.vocab + ctx = len(self._context) + cond_leaves = [tuple(vocab[int(c)])[ctx:] for c in oracle.batch_codes()] # strip the shared context prefix + + # TODO(selmanozleyen): inline x_loaders since there is no need + src_loaders = self._node_loaders(self._src, self._ctrl, lambda: self._inner(schedule)) + tgt_loaders = self._node_loaders(self._src_p, self._pert, lambda: self._bound(schedule)) + src_iters = {k: iter(ld) for k, ld in src_loaders.items()} + tgt_iters = {k: iter(ld) for k, ld in tgt_loaders.items()} + skeys, tkeys = list(src_loaders), list(tgt_loaders) + + for j in range(len(schedule)): + src = {k: next(src_iters[k])["X"] for k in skeys} + tgt = {k: next(tgt_iters[k])["X"] for k in tkeys} + leaf = cond_leaves[j] + out: dict = {"leaf": leaf, "source": src[skeys[0]], "target": tgt[tkeys[0]]} + if len(skeys) > 1: + out["source_reps"] = src + if len(tkeys) > 1: + out["target_reps"] = tgt + if self._cond_fn is not None: + out["condition"] = self._cond_fn(leaf) + yield out diff --git a/src/dagloader/_io.py b/src/dagloader/_io.py new file mode 100644 index 00000000..a659cddd --- /dev/null +++ b/src/dagloader/_io.py @@ -0,0 +1,163 @@ +"""Container-agnostic data access for the loader: obs columns, leaf codes, and rep backings. + +Every helper works uniformly over an in-memory ``AnnData`` and an out-of-core annbatch +``DatasetCollection``, and over ``X`` / ``obsm`` / ``layers`` representations — so the loader never +branches on the source kind. No cell matrices are touched for grouping (obs only); cells are read +only when a batch is materialized (by annbatch's own loader over the returned backings). +""" + +from __future__ import annotations + +import os +from collections.abc import Sequence + +import anndata as ad +import numpy as np +import pandas as pd + +from dagloader._schema import Container + +__all__ = ["key_backings", "leaf_codes", "load_backed_adata", "materialize_node", "obs_columns", "open_source"] + + +def _readable(x): + """A rep backing annbatch can read (dense passes through; a sparse zarr group gets wrapped). + + A **sparse** rep in a ``DatasetCollection`` is a zarr *group* (CSR: data/indices/indptr) with no + ``.shape``, so wrap it as an anndata ``CSRDataset`` (which exposes ``.shape`` + row indexing). + In-memory ``AnnData`` reps (scipy/numpy) and dense zarr arrays already qualify and pass through. + """ + if hasattr(x, "shape"): + return x + from anndata.io import sparse_dataset + + return sparse_dataset(x) + + +def key_backings(source: Container, loc: str) -> list: + """The array(s) backing rep ``loc`` for a source, ready to feed one annbatch ``add_datasets``. + + annbatch's ``add_datasets`` concatenates on the obs axis and needs equal feature dims, so each rep + gets its own loader over its own array(s). For a ``DatasetCollection`` (or a ``list`` of AnnData — + one backing per adata) the per-dataset arrays are gathered in order (matching the global row + layout); a sparse rep's zarr group is wrapped so it is readable (see :func:`_readable`). + """ + if isinstance(source, list): # list of (backed) AnnData: one backing per adata, in list order + return [b for a in source for b in key_backings(a, loc)] + if loc == "X": + return [source.X] if isinstance(source, ad.AnnData) else [_readable(g["X"]) for g in source] + field, sub = loc.split("/", 1) # "obsm/X_pca" | "layers/log1p" + if isinstance(source, ad.AnnData): + return [getattr(source, field)[sub]] + return [_readable(g[field][sub]) for g in source] # DatasetCollection: one backing per dataset + + +def _read_rows(source: Container, loc: str, row_idx: np.ndarray) -> np.ndarray: + """Densely read the global rows ``row_idx`` (ascending) of rep ``loc`` from a source's backings.""" + backings = key_backings(source, loc) # each has ``.shape`` (sparse groups already wrapped by _readable) + offs = np.concatenate([[0], np.cumsum([b.shape[0] for b in backings])]).astype(np.int64) + parts = [] + for d, b in enumerate(backings): + lb = row_idx[(row_idx >= offs[d]) & (row_idx < offs[d + 1])] - offs[d] # ascending within this dataset + if lb.size: + sub = b.oindex[lb] if hasattr(b, "oindex") else b[lb] # zarr array: orthogonal; CSRDataset/scipy: [] + dense = getattr(sub, "todense", None) + parts.append(np.asarray(dense()) if dense is not None else np.asarray(sub)) + return np.concatenate(parts, axis=0) if parts else np.empty((0, int(backings[0].shape[1])), dtype=np.float32) + + +def materialize_node(source: Container, node) -> ad.AnnData: + """Materialize a node's selected (positive-weight) cells into an in-memory (dense) ``AnnData``. + + The general "read this node's cells into RAM" op behind :attr:`~dagloader.Node.in_memory`: reads only + the rows whose leaf has positive weight — for each of the node's reps (``keys``) — and sorts them by + ``cols`` so ``chunk_size > 1`` still reads contiguous runs. Handles dense and sparse (CSR-group) + backings alike (via :func:`key_backings`). Cells must fit host RAM (the intended use is a small, + frequently re-drawn population such as matched controls). + """ + from dagloader._schema import _weight_vector + + obs = obs_columns(source, node.cols) + codes, leaves = leaf_codes(obs, node.cols) + selected = np.flatnonzero(_weight_vector(node.weights, leaves) > 0) # leaf codes with positive weight + row_idx = np.flatnonzero(np.isin(codes, selected)) # ascending global rows of the selected leaves + reps = {key: _read_rows(source, key, row_idx) for key in node.keys} + sub_obs = obs.iloc[row_idx].reset_index(drop=True) + order = sub_obs.sort_values(list(node.cols), kind="stable").index.to_numpy() # contiguous runs for chunk>1 + sub_obs, reps = sub_obs.iloc[order].reset_index(drop=True), {k: v[order] for k, v in reps.items()} + adata = ad.AnnData(X=reps["X"], obs=sub_obs) if "X" in reps else ad.AnnData(obs=sub_obs) # X-> n_var inferred + for key, v in reps.items(): + if key != "X": + adata.obsm[key.split("/", 1)[1]] = v + return adata + + +def leaf_codes(obs: pd.DataFrame, cols: Sequence[str]) -> tuple[np.ndarray, list[tuple]]: + """Per-cell leaf code + the ordered leaf combinations (the grouping over ``cols``).""" + tuples = [tuple(row) for row in obs[list(cols)].to_numpy()] + leaves = sorted(set(tuples), key=lambda t: tuple(map(str, t))) + code_of = {lf: i for i, lf in enumerate(leaves)} + return np.array([code_of[t] for t in tuples], dtype=np.int64), leaves + + +def obs_columns(source: Container, cols: Sequence[str]) -> pd.DataFrame: + """Obs columns from any container (AnnData attr vs DatasetCollection reader vs list) — no cell matrices.""" + if isinstance(source, list): # list of AnnData: concatenate obs in list order (= key_backings order) + return pd.concat([obs_columns(a, cols) for a in source], ignore_index=True) + if isinstance(source, ad.AnnData): + return source.obs[list(cols)] + return source.obs(columns=list(cols)) # DatasetCollection + + +def _backed(x): + """A zarr rep as a readable backing: dense zarr array passes through; a sparse group is wrapped.""" + import zarr + + return x if isinstance(x, zarr.Array) else ad.io.sparse_dataset(x) + + +def load_backed_adata(g, *, keys: Sequence[str], cols: Sequence[str] = ()) -> ad.AnnData: + """Open a zarr adata group as a (backed) AnnData, reading only the reps in ``keys`` and obs ``cols``. + + ``X`` is read as a backed sparse dataset or lazy dense zarr array; ``obsm/`` / ``layers/`` reps + likewise; ``var`` is reduced to its index and ``obs`` to ``cols`` (all of obs if ``cols`` is empty). + Only the reps a node actually streams are materialized, so unused representations are never touched. + """ + var = g["var"] + obs = ad.io.read_elem(g["obs"]) + kw: dict = { + "obs": obs[list(cols)] if cols else obs, + "var": pd.DataFrame(index=pd.Index(ad.io.read_elem(var[var.attrs.get("_index")]))), + } + if "X" in keys: + kw["X"] = _backed(g["X"]) + adata = ad.AnnData(**kw) + for key in keys: + if key == "X": + continue + field, sub = key.split("/", 1) # "obsm/X_pca" | "layers/log1p" + getattr(adata, field)[sub] = _backed(g[field][sub]) + return adata + + +def open_source(src, *, keys: Sequence[str], cols: Sequence[str] = ()) -> Container: + """Resolve one :class:`~dagloader.Scheme` source value to a :data:`~dagloader.Container`. + + An in-memory ``AnnData`` or ``DatasetCollection`` passes through unchanged. A single path is opened as + zarr and auto-detected: an annbatch collection root (``encoding-type`` ``"annbatch-preshuffled"``) + becomes a ``DatasetCollection``; otherwise it is a single adata read backed via :func:`load_backed_adata`. + A list of paths becomes a list of backed AnnData (the loader feeds one backing per adata to + ``add_datasets``). Only the reps in ``keys`` and obs ``cols`` are read (see :func:`load_backed_adata`). + """ + import zarr + from annbatch import DatasetCollection + + if isinstance(src, (ad.AnnData, DatasetCollection)): + return src + if isinstance(src, (str, os.PathLike)): + g = zarr.open_group(src, mode="r") + if g.attrs.get("encoding-type") == "annbatch-preshuffled": + return DatasetCollection(src, mode="r") + return load_backed_adata(g, keys=keys, cols=cols) + # list/sequence of adata zarr paths (or already-open AnnData) → list of backed AnnData + return [a if isinstance(a, ad.AnnData) else load_backed_adata(zarr.open_group(a, mode="r"), keys=keys, cols=cols) for a in src] diff --git a/src/dagloader/_loader.py b/src/dagloader/_loader.py new file mode 100644 index 00000000..9c4bda8f --- /dev/null +++ b/src/dagloader/_loader.py @@ -0,0 +1,270 @@ +"""``DAGLoader`` — streams matched ``{source, target, condition}`` batches from a :class:`Scheme`. + +Each pass is a fresh epoch. The root (target) node draws a per-batch class schedule ∝ its weights via an +annbatch :class:`~annbatch.samplers.ClassSampler`; every bound child replays that schedule onto its own +cells via an annbatch :class:`~annbatch.samplers.BoundClassSampler` — matched by *label* on the bind's +``common`` columns (select via child weights + project via ``common``). The loader +never wraps annbatch's RNG: every sampler that must agree within a pass (the schedule oracle, the target +reps, and each child's inner) is **reseeded from one per-pass seed** ``(node seed, pass index)``, so +target/condition/source stay aligned, a node's reps read the same rows, and a pickled loader resumes the +exact same stream. See ``README.md``. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator, Mapping +from importlib.util import find_spec + +import numpy as np +import pandas as pd +from annbatch import Loader +from annbatch.samplers import BoundClassSampler, ClassSampler + +from dagloader._io import key_backings, leaf_codes, materialize_node, obs_columns +from dagloader._schema import Bind, Container, SamplerConfig, Scheme, _weight_vector + +__all__ = ["DAGLoader"] + +# annbatch's GPU path (cupy vstack/indexing → `to="jax"` yields a GPU-resident array via dlpack) needs +# cupy. When it's absent (CPU-only envs — Mac, CI), fall back so `to="jax"` still yields a CPU jax array. +_HAS_CUPY = find_spec("cupy") is not None + + +def _flat_categorical(codes: np.ndarray, leaves: list[tuple]) -> pd.Categorical: + """A tuple-labelled categorical: per-cell leaf code over ``leaves`` (categories are the leaf tuples). + + Tuple labels (not opaque integer codes) are what let :class:`~annbatch.samplers.BoundClassSampler` + match a child to its parent by the bind's ``common`` columns — it projects the label by position. + """ + categories = pd.MultiIndex.from_tuples(leaves).to_flat_index() + return pd.Categorical.from_codes(codes, categories=categories) + + +class DAGLoader: + """Yields ``{"source", "target", "condition"}`` batches; every node streams through its own loader.""" + + def __init__( + self, + scheme: Scheme, + sampler_config: SamplerConfig | Mapping[str, SamplerConfig], + condition_fn: Callable[[tuple], np.ndarray | Mapping[str, np.ndarray]] | None = None, + ) -> None: + self.s = scheme + self._cond_fn = condition_fn + self._cfg = self._resolve_configs(sampler_config) + self._B = self._cfg[self.s.root].batch_size # root/target batch size — drives the pass length + + # root's direct children are the bound sources (parity with the previous loader: depth-1 sources) + self._child_binds: list[Bind] = [b for b in scheme.binds if b.parent == self.s.root] + + # per-node stable sub-seed from one SeedSequence, so nodes don't correlate; a pass's seed is + # (sub-seed, pass index) → a pass is fully reproducible from its index. + self._node_seeds: dict[str, int] = { + name: int(seq.generate_state(1)[0]) + for name, seq in zip( + sorted(scheme.nodes), np.random.SeedSequence(scheme.seed).spawn(len(scheme.nodes)), strict=True + ) + } + + # resolve each node's source; a `Node.in_memory` node is materialized into RAM once (see _io). + self._nodes: dict[str, Container] = { + name: (materialize_node(scheme.sources[node.source], node) if node.in_memory else scheme.sources[node.source]) + for name, node in scheme.nodes.items() + } + + # per-node leaf partition + weights + tuple-labelled categorical (obs only — no cell matrices) + self._st: dict[str, dict] = {} + for name, node in scheme.nodes.items(): + obs = obs_columns(self._nodes[name], node.cols) + codes, leaves = leaf_codes(obs, node.cols) + self._st[name] = { + "node": node, + "leaves": leaves, + "w": _weight_vector(node.weights, leaves), + "cats": _flat_categorical(codes, leaves), + } + + # A natural epoch over the root (target) node: its cell count // the root's batch_size. The root + # drives the zip, so every node draws the same number of batches; each node's num_samples == + # _n_batches * that node's batch_size (source rows need not equal target rows). + n_root_obs = len(self._st[self.s.root]["cats"]) + self._n_batches = max(1, n_root_obs // self._B) + + self._build_samplers_and_loaders() + + self._iters: dict[str, dict[str, Iterator[dict]]] | None = None + self._schedule: np.ndarray | None = None # per-batch root leaf code for the current pass + self._pos = 0 + + def _resolve_configs(self, cfg: SamplerConfig | Mapping[str, SamplerConfig]) -> dict[str, SamplerConfig]: + """Normalize to one config per node (nodes may use different batch sizes).""" + if isinstance(cfg, SamplerConfig): + return dict.fromkeys(self.s.nodes, cfg) + missing = set(self.s.nodes) - set(cfg) + if missing: + raise ValueError(f"sampler_config mapping is missing node(s): {sorted(missing)}.") + return {name: cfg[name] for name in self.s.nodes} + + # ── build ──────────────────────────────────────────────────────────── + def _new_class_sampler(self, name: str) -> ClassSampler: + cfg = self._cfg[name] + try: # annbatch enforces its own run-length rule for chunk>1; forward with node context + return ClassSampler( + chunk_size=cfg.chunk_size, + preload_nchunks=cfg.preload_nchunks, + batch_size=cfg.batch_size, + classes=self._st[name]["cats"], + num_samples=self._n_batches * cfg.batch_size, + class_weights=self._st[name]["w"], + drop_last=True, + rng=np.random.default_rng(self._node_seeds[name]), + ) + except ValueError as e: + raise ValueError(f"node {name!r}: {e}") from e + + def _bound_on(self, b: Bind) -> dict[int, int]: + """Map root tuple positions → child tuple positions for the bind's ``common`` columns.""" + rcols, ccols = self._st[b.parent]["node"].cols, self._st[b.child]["node"].cols + return {rcols.index(c): ccols.index(c) for c in b.common} + + def _new_bound_sampler(self, b: Bind) -> BoundClassSampler: + inner = self._new_class_sampler(self.s.root) + # Match on the bind's shared columns; the child's leaf weights (0 for excluded leaves, e.g. + # perturbed cells in a control node) go in as the *secondary* class so only positive-weight + # child leaves are drawn within each matched context — the exclusion `classes_to_bind_on` alone + # can't express (it groups all cells sharing the context). + return self._make_bound( + b, + inner, + on=self._bound_on(b), + classes=self._st[b.child]["cats"], + class_weights=self._st[b.child]["w"], + ) + + def _make_bound( + self, + b: Bind, + inner: ClassSampler, + *, + on: dict[int, int] | None, + classes: pd.Categorical | None = None, + class_weights: np.ndarray | None = None, + ) -> BoundClassSampler: + cfg = self._cfg[b.child] + try: + return BoundClassSampler( + inner, + cfg.chunk_size, + cfg.preload_nchunks, + cfg.batch_size, + classes_to_bind_on=self._st[b.child]["cats"], + on=on, + classes=classes, + class_weights=class_weights, + rng=np.random.default_rng(self._node_seeds[b.child]), + ) + except ValueError as e: + raise ValueError(f"node {b.child!r}: {e}") from e + + def _build_samplers_and_loaders(self) -> None: + """Per node/key: a ClassSampler (root) or BoundClassSampler (child) + annbatch Loader. + + A per-child schedule *oracle* and each bound's inner are root-seeded so their class draws agree + with the target's; all of a node's keys share the node seed, so the (identical) samplers select + the same rows every batch — every rep of a node is the same cells. Reps need separate Loaders + (annbatch can't mix feature dims in one loader), each with native chunked reads. + """ + self._oracle = self._new_class_sampler(self.s.root) # supplies the per-batch condition schedule + + self._samplers: dict[str, dict[str, ClassSampler | BoundClassSampler]] = {} + self._loaders: dict[str, dict[str, Loader]] = {} + self._add_node_loaders(self.s.root, lambda: self._new_class_sampler(self.s.root)) + for b in self._child_binds: + self._add_node_loaders(b.child, lambda b=b: self._new_bound_sampler(b)) + + def _add_node_loaders(self, name: str, make_sampler: Callable[[], ClassSampler | BoundClassSampler]) -> None: + node = self._st[name]["node"] + src = self._nodes[name] + cfg = self._cfg[name] + preload_to_gpu = cfg.preload_to_gpu if cfg.preload_to_gpu is not None else _HAS_CUPY # None ⇒ auto + self._samplers[name], self._loaders[name] = {}, {} + for ki, key in enumerate(node.keys): + sampler = make_sampler() + return_index = name == self.s.root and ki == 0 # only for the schedule↔row alignment check + # `to` (default "jax") + `preload_to_gpu` are user-set via SamplerConfig. `to="jax"` yields + # native jax arrays (no host round-trip); `preload_to_gpu` keeps the read window on-GPU (needs + # cupy), else it defers the device copy to the step. Auto-selects cupy when unset. + loader = Loader( + batch_sampler=sampler, return_index=return_index, to=cfg.to, preload_to_gpu=preload_to_gpu + ).add_datasets(key_backings(src, key)) + self._samplers[name][key], self._loaders[name][key] = sampler, loader + + # ── per-pass scheduling ──────────────────────────────────────────────── + def _start_pass(self) -> None: + """Draw the schedule and rebuild iterators for a fresh epoch (advancing every sampler's RNG once). + + The oracle's ``batch_codes()`` and each target/child ``iter()`` each consume one class draw, so — + all root-referencing samplers having started from the root seed — the oracle, the target reps and + every bound child's inner stay in lockstep and draw the *same* per-batch class each pass. The RNG + advances across passes (a real epoch stream), so a pickled loader — whose sampler RNG state is + kept — resumes the next pass rather than replaying. + """ + self._schedule = self._oracle.batch_codes() + self._iters = {name: {key: iter(ld) for key, ld in loaders.items()} for name, loaders in self._loaders.items()} + self._pos = 0 + + # ── pickling ───────────────────────────────────────────────────────────── + def __getstate__(self) -> dict[str, object]: + """Pickle without the live annbatch iterators (generators aren't picklable). + + Every sampler's RNG state is kept, so a reloaded loader resumes the same reproducible stream + (the next pass) on the next ``__next__``; only ``_iters`` is dropped and rebuilt. + """ + state = self.__dict__.copy() + state["_iters"] = None + return state + + def __setstate__(self, state: dict[str, object]) -> None: + self.__dict__.update(state) + self._iters = None # force `_start_pass` on the next `__next__`, using the restored sampler RNG + + # ── iteration ────────────────────────────────────────────────────────── + def __iter__(self) -> DAGLoader: + return self + + def _nodes_next(self, name: str) -> dict[str, np.ndarray]: + """One batch per key of a node — identical samplers pick the same rows, so the reps are aligned.""" + node = self._st[name]["node"] + return {key: next(self._iters[name][key])["X"] for key in node.keys} + + def __next__(self) -> dict[str, np.ndarray]: + if self._iters is None or self._pos >= self._n_batches: + self._start_pass() # first pass, next epoch, or resume after unpickling + j = self._pos + self._pos += 1 + + st = self._st[self.s.root] + node = st["node"] + leaf = st["leaves"][int(self._schedule[j])] # per-batch category — from the schedule oracle + reps = self._nodes_next(self.s.root) + target = reps[node.keys[0]] # primary streamed rep + B = target.shape[0] + + out: dict = {"target": target} + if len(node.keys) > 1: # aligned reps of the target cells (state, per-cell condition, …) + out["target_reps"] = reps + if self._cond_fn is not None: + cond = self._cond_fn(leaf) + if isinstance(cond, Mapping): # per-leaf structured condition (e.g. set-encoded) — emit as-is + out["condition"] = {k: np.asarray(v, dtype=np.float32) for k, v in cond.items()} + else: # a per-leaf vector broadcast across the (class-coherent) batch + cond = np.asarray(cond, dtype=np.float32) + out["condition"] = np.broadcast_to(cond, (B, cond.shape[-1])).copy() + + for b in self._child_binds: # bound child source, replayed via its BoundClassSampler + cnode = self._st[b.child]["node"] + creps = self._nodes_next(b.child) + out["source"] = creps[cnode.keys[0]] + if len(cnode.keys) > 1: + out["source_reps"] = creps + return out diff --git a/src/dagloader/_schema.py b/src/dagloader/_schema.py new file mode 100644 index 00000000..0a558edd --- /dev/null +++ b/src/dagloader/_schema.py @@ -0,0 +1,273 @@ +r"""Declarative schema for :class:`~dagloader.DAGLoader`. + +A :class:`Scheme` is a rooted tree of :class:`Node`\\s over named cell *sources* — pure structure +(sources, grouping columns, weights, binds). How each node is *read* (chunk / preload / batch sizes) +lives in a separate :class:`SamplerConfig` passed to the loader, deliberately kept off the ``Node`` so +the same structure can be run with different sampler settings. + +Each node partitions its source's cells into **leaves** (unique combinations of ``cols``) with a +per-combination :data:`Weights` mapping. A weight of 0 (or a combination absent from the mapping) is +*excluded* — that IS the selection, native to annbatch's ``ClassSampler``. :class:`Bind` links a +parent to a child on shared columns, so the child is sampled *conditioned* on the parent's values. +See ``README.md`` for the model and the cellflow / sc-flow-tools mapping. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from os import PathLike + +import anndata as ad +import numpy as np +from annbatch import DatasetCollection + +# A cell source: an in-memory/backed AnnData, an out-of-core annbatch DatasetCollection, or a list of +# AnnData (streamed as one logical source — one annbatch backing per adata, in list order). +Container = ad.AnnData | DatasetCollection | list[ad.AnnData] + +# A sampling scheme is just a mapping {combination -> weight}. A combination absent from the mapping +# (or with weight 0) is excluded — that IS the selection. ``uniform`` / ``frequency`` / +# ``inverse_frequency`` are plain helpers that build such a dict; nothing about them is privileged. +Weights = Mapping[tuple, float] + +__all__ = [ + "Bind", + "Container", + "Node", + "SamplerConfig", + "Scheme", + "Weights", + "frequency", + "inverse_frequency", + "uniform", +] + + +def uniform(combos: Sequence[tuple]) -> dict[tuple, float]: + """Every combination equally likely.""" + return {tuple(c): 1.0 for c in combos} + + +def frequency(counts: Mapping[tuple, int]) -> dict[tuple, float]: + """Sample each combination ∝ its cell count (favor abundant conditions).""" + return {tuple(k): float(c) for k, c in counts.items()} + + +def inverse_frequency(counts: Mapping[tuple, int]) -> dict[tuple, float]: + """Sample each combination ∝ 1 / cell count (balance rare vs abundant conditions).""" + return {tuple(k): 1.0 / c for k, c in counts.items()} + + +def _weight_vector(weights: Weights, leaves: Sequence[tuple]) -> np.ndarray: + """Resolve ``{combo: weight}`` to normalized per-leaf weights (→ ``ClassSampler.class_weights``).""" + v = np.array([float(weights.get(tuple(lf), 0.0)) for lf in leaves], dtype=float) + s = v.sum() + if s <= 0: + raise ValueError("weights resolve to all-zero over these leaves — nothing to sample.") + return v / s + + +@dataclass(frozen=True) +class Node: + """A partition of one source's cells into leaves, with a per-leaf sampling weight. + + Parameters + ---------- + source + Key into :attr:`Scheme.sources`. + cols + Tree levels → leaves are the unique combinations of these columns (over ALL the source's + cells). These are the grouping/condition columns (cellflow's ``split_covariates`` + + ``perturbation_covariates`` columns; sc-flow-tools' grouping keys). + keys + Representation location(s) to stream: ``"X"`` | ``"obsm/"`` | ``"layers/"`` + (cellflow's ``sample_rep``). A single string streams one rep; a tuple streams SEVERAL + **aligned** reps of the *same* sampled cells — e.g. the state plus a per-cell continuous + condition. The first key drives sampling (via annbatch's ``ClassSampler``); the rest are + read back for the exact same rows, so every rep of a batch is the same cells. + weights + ``{combo: weight}``; a combination absent or with weight 0 is excluded (= the selection). + in_memory + If :obj:`True`, the loader materializes this node's selected (positive-weight) cells into an + in-memory ``AnnData`` once and streams them from RAM instead of re-reading the source every batch + (see :func:`~dagloader._io.materialize_node`). Use for a small, frequently re-drawn population + (e.g. matched controls). Requires those cells to fit in host RAM. + """ + + source: str + cols: tuple[str, ...] + keys: str | tuple[str, ...] = "X" # one rep, or several aligned reps of the same cells + weights: Weights = field(default_factory=dict) + in_memory: bool = False # materialize this node's selected cells into RAM (see dagloader._io) + + def __post_init__(self) -> None: # structural checks (data-free) + object.__setattr__(self, "keys", (self.keys,) if isinstance(self.keys, str) else tuple(self.keys)) + if not self.cols: + raise ValueError("Node.cols must be non-empty.") + if not self.keys or any(not k for k in self.keys): + raise ValueError("Node.keys must be one or more non-empty representation locations.") + for k in self.weights: + if len(k) != len(self.cols): + raise ValueError(f"weight key {k!r} arity != cols {self.cols}.") + if any(w < 0 for w in self.weights.values()): + raise ValueError("weights must be non-negative.") + + +@dataclass(frozen=True) +class Bind: + """Condition ``child`` on ``parent``: match on the ``common`` columns (⊆ their shared cols). + + Each batch, the child's sampled leaf is derived from the parent's leaf via the ``common`` values + (parent leaf → shared-column values → matching child leaf). This is the source↔target matching: + with ``common`` = the context (e.g. cell line), the child (control) is drawn from the *same* + context as the parent (perturbed) — cellflow's "control = same group", sc-flow-tools' + ``control_values_dict`` + default same-context coupling. + + Conditioning is **required**: if a parent value has no matching positive-weight child leaf the + loader raises (no silent fallback). When several child leaves share the bound value — the child + partitions on columns beyond ``common`` (e.g. child cols ``(a, x)`` bound on ``a``) — one is drawn + ∝ the child's leaf weights, so ``P(child extra cols | common)`` is weight-controlled. Pass + ``common=()`` to opt into unconditional child sampling explicitly. + + Matching is thus *select* (the child's per-leaf weights) + *project* (``common``) — nothing more is + needed. An arbitrary parent-leaf → child-leaf pairing (sc-flow-tools' ``matched_keys``) is not a + separate mechanism: because it is a function, tag each side with a shared key column (parent cells + with the child leaf they map to, child cells with their own leaf) and bind on that column. + """ + + parent: str + child: str + common: tuple[str, ...] = () # ⊆ parent.cols ∩ child.cols; () ⇒ unconditional + + +@dataclass(frozen=True, kw_only=True) +class SamplerConfig: + r"""annbatch read parameters for a node's sampler — kept separate from the structural :class:`Node`. + + Passed to :class:`~dagloader.DAGLoader` as either one config (applied to every node) + or a ``{node_name: SamplerConfig}`` mapping (per-node). Nodes may use **different** ``batch_size``\\s: + every node draws the same number of batches (the root's, derived from its cell count), but a node's + batch carries its own row count — so source and target row counts need not match. + + Parameters + ---------- + batch_size + Rows per emitted batch for this node (target rows for the root; source rows for a bound child). + chunk_size + annbatch read-slice size. **Required** and explicit — there is no hidden default; set it + deliberately. ``1`` ⇒ per-row reads (any on-disk layout); ``>1`` ⇒ contiguous chunked reads + (higher throughput on disk), assuming each sampled leaf sits in a contiguous run ≥ + ``chunk_size``. Must divide ``batch_size`` (one category per batch). + preload_nchunks + Chunks per annbatch read window. **Required** and explicit — there is no hidden default; set it + deliberately (e.g. ``batch_size // chunk_size`` for one batch per window). Must be a positive + multiple of ``batch_size // chunk_size``. + to + annbatch ``Loader`` output backend for the yielded batches — ``"jax"`` (default), ``"torch"``, or + :obj:`None` (annbatch's own default). Forwarded verbatim to :class:`annbatch.Loader`. + preload_to_gpu + Whether annbatch keeps the read window on-GPU (needs ``cupy``). :obj:`None` (default) auto-selects + it from cupy availability; pass ``True``/``False`` to force it. + """ + + batch_size: int + chunk_size: int + preload_nchunks: int + to: str = "jax" + preload_to_gpu: bool | None = None + + +@dataclass(frozen=True) +class Scheme: + """The structural sampling spec: sources, a rooted tree of nodes, and the reproducibility cadence. + + Read parameters (chunk / preload / batch sizes) are NOT here — they are a separate + :class:`SamplerConfig` given to the loader. + + Parameters + ---------- + sources + ``{name: AnnData | DatasetCollection}`` — the cell sources the nodes reference. + nodes + ``{name: Node}``. Exactly one is the ``root`` (the streamed target); the rest are bound + children (sources/controls) via ``binds``. + root + Name of the root node (must have no parent). + seed + Reproducibility seed. Per-node RNG streams are spawned from one ``SeedSequence(seed)`` so nodes + do not correlate and the whole stream is reproducible. + binds + Parent→child links (see :class:`Bind`). Must form a rooted tree over ``nodes``. + + Notes + ----- + Batches per with-replacement pass is *not* configured here: the loader derives it from the root + (target) node — a natural epoch of ``root_n_obs // batch_size`` — and restarts each pass. The root + drives the zip, so every node's sampler draws the same number of batches. + """ + + sources: Mapping[str, Container] + nodes: Mapping[str, Node] + root: str + seed: int + binds: tuple[Bind, ...] = () + + @classmethod + def from_paths( + cls, + sources: Mapping[str, Container | str | PathLike | Sequence[str | PathLike]], + nodes: Mapping[str, Node], + root: str, + seed: int, + binds: tuple[Bind, ...] = (), + ) -> Scheme: + """Build a :class:`Scheme` where a source may be given as a zarr path (or list of paths). + + Same signature as the constructor, but each ``sources`` value may additionally be: + + * a **path** to a single zarr adata → read backed (only the reps the referencing nodes use); + * a **path** to an annbatch collection root → opened as a :class:`~annbatch.DatasetCollection` + (auto-detected from its ``encoding-type``); + * a **list of paths** to zarr adatas → a list of backed AnnData streamed as one source (one + annbatch backing per adata, in list order). + + An already-constructed :data:`Container` (AnnData / DatasetCollection / list of AnnData) passes + through unchanged. Everything on disk is expected in **zarr**. Only the ``keys`` and ``cols`` the + nodes referencing a source actually need are read (see :func:`~dagloader._io.open_source`). + """ + from dagloader._io import open_source + + keys_by_src: dict[str, set[str]] = {name: set() for name in sources} + cols_by_src: dict[str, set[str]] = {name: set() for name in sources} + for node in nodes.values(): + if node.source in keys_by_src: # unknown sources are reported by __post_init__ + keys_by_src[node.source].update(node.keys) + cols_by_src[node.source].update(node.cols) + resolved = { + name: open_source(src, keys=keys_by_src[name], cols=cols_by_src[name]) + for name, src in sources.items() + } + return cls(sources=resolved, nodes=nodes, root=root, seed=seed, binds=binds) + + def __post_init__(self) -> None: # structural: rooted tree + references + if self.root not in self.nodes: + raise ValueError(f"root {self.root!r} not in nodes.") + for name, n in self.nodes.items(): + if n.source not in self.sources: + raise ValueError(f"node {name!r} references unknown source {n.source!r}.") + parents: dict[str, str] = {} + for b in self.binds: + if b.parent not in self.nodes or b.child not in self.nodes: + raise ValueError("bind references unknown node.") + if b.child in parents: + raise ValueError(f"node {b.child!r} has multiple parents — must be a rooted tree.") + parents[b.child] = b.parent + shared = set(self.nodes[b.parent].cols) & set(self.nodes[b.child].cols) + if not set(b.common) <= shared: + raise ValueError(f"bind.common {b.common} must be ⊆ shared cols of {b.parent}&{b.child} ({shared}).") + if self.root in parents: + raise ValueError("root must have no parent.") + for name in self.nodes: + if name != self.root and name not in parents: + raise ValueError(f"non-root node {name!r} is not bound to the tree.") diff --git a/src/dagloader/_schemes.py b/src/dagloader/_schemes.py new file mode 100644 index 00000000..74898c96 --- /dev/null +++ b/src/dagloader/_schemes.py @@ -0,0 +1,58 @@ +"""Factory helpers ("the above layer") that fill a :class:`Scheme` from an obs table. + +These are conveniences, not privileged: they build the same ``Node`` / ``Bind`` / ``Scheme`` a caller +could assemble by hand. :func:`perturbation_scheme` is the cellflow-shaped case (control → perturbed, +matched on context). Read parameters (batch / chunk / preload) are a separate :class:`SamplerConfig` +given to the loader, so the factory only builds structure. sc-flow-tools-shaped schemes are usually +assembled directly from a ``DataManager`` config; see ``README.md``. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence + +from dagloader._io import obs_columns +from dagloader._schema import Bind, Container, Node, Scheme, uniform + +__all__ = ["perturbation_scheme"] + + +def perturbation_scheme( + source: Container, + *, + context: Sequence[str], + perturbation: Sequence[str], + control_values: Mapping[str, object], + key: str = "X", + seed: int = 0, +) -> Scheme: + """Fill a perturbation Scheme from the obs table: root = perturbed combos, child = control combos. + + ``source`` is an in-memory AnnData or an on-disk DatasetCollection. There is no ``select`` step — + control vs perturbed is encoded purely by which combinations carry weight. The control node is + bound to the perturbed root on ``context``, so each batch's control cells come from the same + context (cell line, …) as the perturbed cells — the source↔target matching. + + Parameters mirror cellflow: ``context`` = ``split_covariates`` (grouping/context), ``perturbation`` + = the perturbation columns, ``control_values`` = which value marks control per column, ``key`` = + ``sample_rep``. Read parameters (batch/chunk/preload) go to the loader's ``SamplerConfig``. + """ + cols = (*context, *perturbation) + combos = [tuple(r) for r in obs_columns(source, cols).drop_duplicates().to_numpy()] + + def is_control(combo: tuple) -> bool: + return all(combo[cols.index(c)] == v for c, v in control_values.items()) + + pert = [c for c in combos if not is_control(c)] + ctrl = [c for c in combos if is_control(c)] + return Scheme( + sources={"data": source}, + nodes={ + # non-control combos weighted (rest excluded = the selection); control combos weighted. + "pert": Node("data", cols, key, uniform(pert)), + "ctrl": Node("data", cols, key, uniform(ctrl)), + }, + root="pert", + binds=(Bind("pert", "ctrl", common=tuple(context)),), + seed=seed, + ) diff --git a/src/dagloader/_split.py b/src/dagloader/_split.py new file mode 100644 index 00000000..e8f07e2f --- /dev/null +++ b/src/dagloader/_split.py @@ -0,0 +1,203 @@ +"""Split a :class:`~dagloader.Scheme`'s target combinations into named splits — weights-only. + +A leaf's weight *is* its selection (weight 0 / absent ⇒ not sampled), so a split just restricts the +root node's weights to a subset of combinations. :func:`split_scheme` returns one :class:`Scheme` per +split, identical to the input but with the root weights restricted; bound children (controls) are +carried through unchanged, since a matched control must stay available in every split. Like CellFlow2's +combination-level splitter (hold out whole conditions, not cells), but on the ``Scheme`` — no cells, no +``DatasetCollection``, no loader. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import replace + +import numpy as np +import pandas as pd + +from dagloader._schema import SamplerConfig, Scheme + +__all__ = ["resolve_split_configs", "split_assignment", "split_scheme"] + +_DEFAULT_RATIOS = {"train": 0.6, "val": 0.2, "test": 0.2} + + +def _allocate(n: int, ratios: Mapping[str, float], names: Sequence[str]) -> dict[str, int]: + """Largest-remainder apportionment of ``n`` groups to splits by ratio (the sizes sum to ``n``).""" + exact = {name: ratios[name] * n for name in names} + sizes = {name: int(exact[name]) for name in names} # floor + remainder = n - sum(sizes.values()) + order = sorted(names, key=lambda nm: exact[nm] - sizes[nm], reverse=True) # largest fractional part first + for nm in order[:remainder]: + sizes[nm] += 1 + return sizes + + +def split_scheme( + scheme: Scheme, + *, + split_by: Sequence[str], + ratios: Mapping[str, float] | None = None, + force_training_values: Mapping[str, object] | None = None, + random_state: int = 42, +) -> dict[str, Scheme]: + """Partition the root (target) combinations of ``scheme`` into named splits. + + Splits are over whole *combinations* of ``split_by`` — combos sharing the same ``split_by`` values + land in the same split — so entire conditions are held out rather than random cells. + + Parameters + ---------- + scheme + The prepared scheme; its ``root`` node's positive-weight combinations are the split universe. + split_by + Columns whose unique combinations are partitioned across splits — a subset of the root's + ``cols``. + ratios + ``{split_name: fraction}`` summing to 1.0 (all > 0). Defaults to + ``{"train": 0.6, "val": 0.2, "test": 0.2}``. Insertion order defines the split order; the first + split is the "training" split that :paramref:`force_training_values` forces into. + force_training_values + ``{column: value}`` (keys ⊆ ``split_by``): any group matching *any* of these (column, value) + pairs is forced into the first split, regardless of the shuffle. + random_state + Seed for the group shuffle (reproducible). + + Returns + ------- + ``{split_name: Scheme}`` — each a copy of ``scheme`` with the root node's ``weights`` restricted to + that split's combinations; every other node (bound children / controls) is carried through + unchanged. + """ + ratios = dict(_DEFAULT_RATIOS if ratios is None else ratios) + split_by = tuple(split_by) + force = dict(force_training_values or {}) + + root = scheme.nodes[scheme.root] + cols = root.cols + + # --- validation ----------------------------------------------------------------------------- + if not split_by: + raise ValueError("split_by must be non-empty.") + missing = [c for c in split_by if c not in cols] + if missing: + raise ValueError(f"split_by columns {missing} are not in the root node's cols {cols}.") + if not ratios: + raise ValueError("ratios must be a non-empty {split_name: fraction} mapping.") + if any(r <= 0 for r in ratios.values()): + raise ValueError(f"ratios must all be > 0; got {ratios}.") + if not np.isclose(sum(ratios.values()), 1.0): + raise ValueError(f"ratios must sum to 1.0; got {ratios} (sum={sum(ratios.values())}).") + bad_force = [k for k in force if k not in split_by] + if bad_force: + raise ValueError(f"force_training_values keys {bad_force} must be a subset of split_by {split_by}.") + + names = list(ratios) # deterministic split order (insertion order) + train_name = names[0] # forced groups land here + + # --- the split universe: the root's positive-weight combinations (weight = the selection) ----- + combos = [combo for combo, w in root.weights.items() if w > 0] + if not combos: + raise ValueError("root node has no positive-weight combinations to split.") + + proj_idx = [cols.index(c) for c in split_by] + + def project(combo: tuple) -> tuple: # a combination → its split_by group + return tuple(combo[i] for i in proj_idx) + + groups = sorted({project(c) for c in combos}, key=lambda g: tuple(map(str, g))) # deterministic + + def is_forced(group: tuple) -> bool: # OR across keys, like CellFlow2._contains_value + gd = dict(zip(split_by, group, strict=True)) + return any(gd.get(k) == v for k, v in force.items()) + + forced = [g for g in groups if is_forced(g)] + free = [g for g in groups if not is_forced(g)] + + rng = np.random.default_rng(random_state) + free = [free[i] for i in rng.permutation(len(free))] + + sizes = _allocate(len(free), ratios, names) + # a split that must hold groups but got 0 is an error (mirrors CellFlow2). The train split is + # exempt when it will still receive forced groups. + empty = [n for n in names if sizes[n] == 0 and not (n == train_name and forced)] + if empty: + raise ValueError( + f"splits {empty} received 0 of {len(groups)} `split_by` group(s); increase the number of " + f"groups or adjust ratios {ratios}." + ) + + split_of: dict[tuple, str] = {} + pos = 0 + for name in names: + for g in free[pos : pos + sizes[name]]: + split_of[g] = name + pos += sizes[name] + for g in forced: # forced groups always go to the first (training) split + split_of[g] = train_name + + # --- build one Scheme per split: restrict the root weights, carry everything else ------------ + out: dict[str, Scheme] = {} + for name in names: + kept = {combo: root.weights[combo] for combo in combos if split_of[project(combo)] == name} + new_root = replace(root, weights=kept) + out[name] = replace(scheme, nodes={**scheme.nodes, scheme.root: new_root}) + return out + + +def split_assignment(splits: Mapping[str, Scheme]) -> pd.DataFrame: + """A tidy ``{*root cols, "split"}`` table of which target combination went to which split. + + Reconstructed from the split schemes' root weights (positive-weight combinations); useful for + inspecting a :func:`split_scheme` result without touching any cells. + """ + if not splits: + raise ValueError("splits mapping is empty.") + any_scheme = next(iter(splits.values())) + cols = any_scheme.nodes[any_scheme.root].cols + rows = [(*combo, name) for name, sch in splits.items() for combo, w in sch.nodes[sch.root].weights.items() if w > 0] + df = pd.DataFrame(rows, columns=[*cols, "split"]) + return df.sort_values(["split", *cols]).reset_index(drop=True) + + +def resolve_split_configs( + config: SamplerConfig | Mapping[str, SamplerConfig], + split_names: Sequence[str], +) -> dict[str, SamplerConfig]: + """Resolve a read-parameter spec into exactly one :class:`SamplerConfig` per split. + + ``config`` is either + + * **a single** :class:`SamplerConfig` — applied to every split; or + * **a per-split mapping** ``{split_name: SamplerConfig}`` — in which case **every** split in + ``split_names`` must be present (no partial specs, no unknown split names). + + Returns ``{split_name: SamplerConfig}`` for exactly ``split_names``. + """ + names = list(split_names) + if not names: + raise ValueError("split_names must be non-empty.") + + if isinstance(config, SamplerConfig): # one config → all splits + return dict.fromkeys(names, config) + + if isinstance(config, Mapping): # per-split mapping: every split specified, no extras + missing = [n for n in names if n not in config] + if missing: + raise ValueError( + f"sampler_config is per-split but is missing config(s) for split(s) {missing}; specify all of {names}." + ) + extra = [k for k in config if k not in names] + if extra: + raise ValueError(f"sampler_config has config(s) for unknown split(s) {extra}; splits are {names}.") + bad = [k for k, v in config.items() if not isinstance(v, SamplerConfig)] + if bad: + raise ValueError( + f"sampler_config values must be SamplerConfig instances; got a non-SamplerConfig for {bad}." + ) + return {name: config[name] for name in names} + + raise ValueError( + f"sampler_config must be a SamplerConfig or a {{split: SamplerConfig}} mapping; got {type(config).__name__}." + ) diff --git a/tests/conftest.py b/tests/conftest.py index e8ef3428..a6ca73de 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,7 +4,7 @@ import pandas as pd import pytest -from cellflow.data._dataloader import ValidationSampler +from cellflow.data._legacy import ValidationSampler @pytest.fixture diff --git a/tests/dagloader/test_coherence.py b/tests/dagloader/test_coherence.py new file mode 100644 index 00000000..44c62c6e --- /dev/null +++ b/tests/dagloader/test_coherence.py @@ -0,0 +1,150 @@ +"""Condition-coherence of the BoundClassSampler-based DAGLoader — the guarantee the old suite missed. + +Each cell's ``(cell_line, drug, control)`` is encoded into ``X`` so every yielded row can be decoded and +checked: the target batch is one perturbed condition, the ``condition`` vector matches that condition, +and the source batch is control cells of the *matched* context (``common=`` → same cell line). This is +what silently broke when dagloader's rng-wrapping scheduler met annbatch's refactored class draw. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("annbatch") + +import anndata as ad + +from dagloader import Bind, DAGLoader, Node, SamplerConfig, Scheme, uniform + +LINES = ("A", "B") +DRUGS = ("control", "d1", "d2", "d3") +_LINE = {c: i for i, c in enumerate(LINES)} +_DRUG = {d: i for i, d in enumerate(DRUGS)} + + +def _adata(n_per_combo: int = 16) -> ad.AnnData: + rows = [(cl, dr) for cl in LINES for dr in DRUGS for _ in range(n_per_combo)] + obs = pd.DataFrame(rows, columns=["cell_line", "drug"]) + obs["control"] = obs["drug"] == "control" + obs.index = obs.index.astype(str) + # X encodes the cell's identity so a batch row can be decoded: [cell_line, drug, is_control] + x = np.stack( + [ + obs["cell_line"].map(_LINE).to_numpy(), + obs["drug"].map(_DRUG).to_numpy(), + obs["control"].to_numpy().astype(float), + ], + axis=1, + ).astype("float32") + return ad.AnnData(X=x, obs=obs) + + +def _condition_fn(leaf: tuple) -> np.ndarray: + return np.array([_LINE[leaf[0]], _DRUG[leaf[1]]], dtype=float) + + +def _scheme(adata: ad.AnnData, *, bind: Bind, seed: int = 0) -> Scheme: + cols = ("cell_line", "drug") + combos = {tuple(r) for r in adata.obs[list(cols)].to_numpy()} + pert = [c for c in combos if c[1] != "control"] + ctrl = [c for c in combos if c[1] == "control"] + return Scheme( + sources={"data": adata}, + nodes={"pert": Node("data", cols, "X", uniform(pert)), "ctrl": Node("data", cols, "X", uniform(ctrl))}, + root="pert", + binds=(bind,), + seed=seed, + ) + + +def test_common_bind_target_condition_and_context_coherent(): + adata = _adata() + loader = DAGLoader( + _scheme(adata, bind=Bind("pert", "ctrl", common=("cell_line",))), + SamplerConfig(batch_size=8, chunk_size=1, preload_nchunks=8), + condition_fn=_condition_fn, + ) + it = iter(loader) + for _ in range(12): + batch = next(it) + tgt = np.asarray(batch["target"]) + src = np.asarray(batch["source"]) + cond = np.asarray(batch["condition"]) + + # target batch is one perturbed condition + assert len(np.unique(tgt[:, 0])) == 1, "target batch mixes cell lines" + assert len(np.unique(tgt[:, 1])) == 1, "target batch mixes drugs" + assert tgt[0, 2] == 0.0, "target must be perturbed (not control)" + + # condition vector matches that exact (cell_line, drug) + assert np.all(cond[:, 0] == tgt[0, 0]) and np.all(cond[:, 1] == tgt[0, 1]), "condition ≠ target condition" + + # source batch is control cells of the SAME context (cell line) + assert np.all(src[:, 2] == 1.0), "source must be control cells" + assert len(np.unique(src[:, 0])) == 1 and src[0, 0] == tgt[0, 0], "source context ≠ target context" + + +def test_reps_are_aligned_same_cells(): + # two aligned reps of the target: X and an obsm copy of X. Same sampled rows → identical values. + adata = _adata() + adata.obsm["rep"] = adata.X.copy() + cols = ("cell_line", "drug") + combos = {tuple(r) for r in adata.obs[list(cols)].to_numpy()} + pert = [c for c in combos if c[1] != "control"] + ctrl = [c for c in combos if c[1] == "control"] + scheme = Scheme( + sources={"data": adata}, + nodes={ + "pert": Node("data", cols, ("X", "obsm/rep"), uniform(pert)), + "ctrl": Node("data", cols, "X", uniform(ctrl)), + }, + root="pert", + binds=(Bind("pert", "ctrl", common=("cell_line",)),), + seed=0, + ) + loader = DAGLoader(scheme, SamplerConfig(batch_size=8, chunk_size=1, preload_nchunks=8)) + batch = next(iter(loader)) + x = np.asarray(batch["target_reps"]["X"]) + rep = np.asarray(batch["target_reps"]["obsm/rep"]) + np.testing.assert_array_equal(x, rep) # aligned reps must be the same cells + + +def test_materialize_node_selects_positive_weight_rows(): + # `materialize_node` reads only a node's positive-weight (here: control) cells into an in-memory AnnData + from dagloader._io import materialize_node + + adata = _adata() + cols = ("cell_line", "drug") + ctrl = [c for c in {tuple(r) for r in adata.obs[list(cols)].to_numpy()} if c[1] == "control"] + mem = materialize_node(adata, Node("data", cols, "X", uniform(ctrl))) + assert isinstance(mem, ad.AnnData) + assert (mem.obs["drug"] == "control").all() # only positive-weight (control) rows materialized + assert mem.n_obs == int((adata.obs["drug"] == "control").sum()) + assert (np.asarray(mem.X)[:, 2] == 1.0).all() # X[:,2] encodes is_control + + +def test_in_memory_node_materialized(): + # Node.in_memory → the loader materializes that node into RAM (served from memory, not re-read each + # batch). SamplerConfig also carries the user-set `to` / `preload_to_gpu` (exercised here). + adata = _adata() + cols = ("cell_line", "drug") + combos = {tuple(r) for r in adata.obs[list(cols)].to_numpy()} + pert = [c for c in combos if c[1] != "control"] + ctrl = [c for c in combos if c[1] == "control"] + scheme = Scheme( + sources={"data": adata}, + nodes={ + "pert": Node("data", cols, "X", uniform(pert)), + "ctrl": Node("data", cols, "X", uniform(ctrl), in_memory=True), + }, + root="pert", + binds=(Bind("pert", "ctrl", common=("cell_line",)),), + seed=0, + ) + cfg = SamplerConfig(batch_size=8, chunk_size=1, preload_nchunks=8, to="jax", preload_to_gpu=False) + dl = DAGLoader(scheme, cfg, condition_fn=_condition_fn) + assert isinstance(dl._nodes["ctrl"], ad.AnnData) # ctrl node materialized into RAM + batch = next(iter(dl)) + assert (np.asarray(batch["source"])[:, 2] == 1.0).all() # source = matched control cells (from RAM) diff --git a/tests/dagloader/test_eval_loader.py b/tests/dagloader/test_eval_loader.py new file mode 100644 index 00000000..787d7b4e --- /dev/null +++ b/tests/dagloader/test_eval_loader.py @@ -0,0 +1,125 @@ +"""``DAGEvalLoader``: control-rooted eval reader (Sequential control inner + BoundClassSampler target). + +For each control population (context, e.g. ``cell_line``) the source is **all** its control cells (read in +full via the Sequential inner); the target is a matched perturbed batch (annbatch samples a drug within +the context). The condition is the perturbed leaf the target drew. Cells carry a unique index in ``X`` so +we can assert the source is exactly the context's controls and the target is that context's perturbed cells. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("annbatch") + +import anndata as ad + +from dagloader import Bind, DAGEvalLoader, Node, SamplerConfig, Scheme, uniform + +LINES = ("A", "B") +DRUGS = ("control", "d1", "d2", "d3") +_LINE = {c: i for i, c in enumerate(LINES)} +_DRUG = {d: i for i, d in enumerate(DRUGS)} +_CFG = SamplerConfig(batch_size=4, chunk_size=1, preload_nchunks=4) + + +def _adata(n_per_combo: int = 8) -> ad.AnnData: + rows = [(cl, dr) for cl in LINES for dr in DRUGS for _ in range(n_per_combo)] + obs = pd.DataFrame(rows, columns=["cell_line", "drug"]) + obs["control"] = obs["drug"] == "control" + obs.index = obs.index.astype(str) + x = np.stack( + [ + obs["cell_line"].map(_LINE).to_numpy(), + obs["drug"].map(_DRUG).to_numpy(), + obs["control"].to_numpy().astype(float), + np.arange(len(obs)), # unique cell index + ], + axis=1, + ).astype("float32") + return ad.AnnData(x, obs=obs) + + +def _condition_fn(leaf: tuple) -> dict[str, np.ndarray]: + return {"drug": np.array([[_LINE[leaf[0]], _DRUG[leaf[1]]]], dtype=float)} + + +def _scheme(adata: ad.AnnData) -> Scheme: + cols = ("cell_line", "drug") + combos = {tuple(r) for r in adata.obs[list(cols)].to_numpy()} + pert = [c for c in combos if c[1] != "control"] + ctrl = [c for c in combos if c[1] == "control"] + return Scheme( + sources={"data": adata}, + nodes={"pert": Node("data", cols, "X", uniform(pert)), "ctrl": Node("data", cols, "X", uniform(ctrl))}, + root="pert", + binds=(Bind("pert", "ctrl", common=("cell_line",)),), + seed=0, + ) + + +def _idx(batch_x) -> set[int]: + return set(np.asarray(batch_x)[:, 3].astype(int).tolist()) + + +def test_source_is_all_controls_of_context_target_matches_condition(): + adata = _adata() + obs = adata.obs + loader = DAGEvalLoader(_scheme(adata), _CFG, _condition_fn) + assert set(loader.control_populations) == {("A", "control"), ("B", "control")} + + seen_ctx = [] + for out in loader.iter_conditions(): # one batch per control population + cl, dr = out["leaf"] + seen_ctx.append(cl) + assert dr != "control" # target is a perturbed condition + # source = ALL control cells of this cell line (read in full) + src_truth = set(np.flatnonzero((obs["cell_line"] == cl).to_numpy() & obs["control"].to_numpy()).tolist()) + assert _idx(out["source"]) == src_truth + # target cells are perturbed cells of this (cell_line, drug) + rows = np.asarray(out["target"]) + assert np.all(rows[:, 0].astype(int) == _LINE[cl]) # same cell line + assert np.all(rows[:, 1].astype(int) == _DRUG[dr]) # the drawn drug + assert np.all(rows[:, 2] == 0.0) # not control + # condition embedding is the drawn perturbed leaf + np.testing.assert_array_equal(out["condition"]["drug"], [[_LINE[cl], _DRUG[dr]]]) + + assert set(seen_ctx) == set(LINES) # each control population once + + +def test_n_conditions_cycles_control_populations(): + adata = _adata() + loader = DAGEvalLoader(_scheme(adata), _CFG, _condition_fn) + outs = list(loader.iter_conditions(n_conditions=5)) + assert len(outs) == 5 # 2 control populations cycled to 5 batches + for out in outs: + assert out["leaf"][1] != "control" + + +def test_deterministic_across_calls(): + adata = _adata() + loader = DAGEvalLoader(_scheme(adata), _CFG, _condition_fn) + a = [out["leaf"] for out in loader.iter_conditions(n_conditions=6)] + b = [out["leaf"] for out in loader.iter_conditions(n_conditions=6)] + assert a == b # same seed ⇒ same drawn conditions each call + + +def test_reps_aligned_same_cells(): + adata = _adata() + adata.obsm["rep"] = adata.X.copy() + cols = ("cell_line", "drug") + combos = {tuple(r) for r in adata.obs[list(cols)].to_numpy()} + scheme = Scheme( + sources={"data": adata}, + nodes={ + "pert": Node("data", cols, ("X", "obsm/rep"), uniform([c for c in combos if c[1] != "control"])), + "ctrl": Node("data", cols, "X", uniform([c for c in combos if c[1] == "control"])), + }, + root="pert", + binds=(Bind("pert", "ctrl", common=("cell_line",)),), + seed=0, + ) + out = next(DAGEvalLoader(scheme, _CFG, _condition_fn).iter_conditions()) + np.testing.assert_array_equal(np.asarray(out["target_reps"]["X"]), np.asarray(out["target_reps"]["obsm/rep"])) diff --git a/tests/dagloader/test_from_paths.py b/tests/dagloader/test_from_paths.py new file mode 100644 index 00000000..0478a916 --- /dev/null +++ b/tests/dagloader/test_from_paths.py @@ -0,0 +1,156 @@ +"""``Scheme.from_paths`` — resolving zarr paths (single adata / collection root / list) to sources. + +Covers the three source-value shapes the classmethod adds on top of the constructor, the auto-detection +of a single path (adata vs annbatch collection root), the "load only the reps the nodes use" contract, +and that a resolved scheme streams matched batches end-to-end just like a constructed one. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("annbatch") + +import warnings + +import anndata as ad +import scipy.sparse as sp +from annbatch import DatasetCollection + +from dagloader import Bind, DAGLoader, Node, SamplerConfig, Scheme, uniform +from dagloader._io import load_backed_adata, open_source + +LINES = ("A", "B") +DRUGS = ("control", "d1", "d2") +COLS = ("cell_line", "drug") + + +def _adata(n_per_combo: int = 16, seed: int = 0) -> ad.AnnData: + rng = np.random.default_rng(seed) + rows = [(cl, dr) for cl in LINES for dr in DRUGS for _ in range(n_per_combo)] + obs = pd.DataFrame(rows, columns=list(COLS)) + obs["extra"] = "unused" # an obs column no node references — must not be required + obs.index = obs.index.astype(str) + a = ad.AnnData(X=sp.csr_matrix(rng.random((len(obs), 5), dtype="float32")), obs=obs) + a.obsm["emb"] = rng.random((len(obs), 3), dtype="float32") + a.layers["log1p"] = sp.csr_matrix(rng.random((len(obs), 5), dtype="float32")) + return a + + +def _write_zarr(adata: ad.AnnData, path) -> str: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") # zarr v2/v3 default-format UserWarning + adata.write_zarr(str(path)) + return str(path) + + +def _weights(): + combos = {(cl, dr) for cl in LINES for dr in DRUGS} + pert = uniform([c for c in combos if c[1] != "control"]) + ctrl = uniform([c for c in combos if c[1] == "control"]) + return pert, ctrl + + +def _cfg() -> SamplerConfig: + return SamplerConfig(batch_size=16, chunk_size=16, preload_nchunks=1, to=None) + + +# ── load_backed_adata: only the requested reps/cols are materialized ─────────────────────────── + + +def test_load_backed_adata_reads_only_requested_keys(tmp_path): + g = _open_group(_write_zarr(_adata(), tmp_path / "a.zarr")) + backed = load_backed_adata(g, keys=("X", "obsm/emb"), cols=COLS) + + assert isinstance(backed.X, ad.abc.CSRDataset) # sparse rep stays backed (not read into RAM) + assert list(backed.obsm) == ["emb"] # obsm/emb requested → present + assert "log1p" not in backed.layers # layers/log1p not requested → never touched + assert list(backed.obs.columns) == list(COLS) # obs reduced to the requested cols + np.testing.assert_array_equal(backed.obsm["emb"].shape, (backed.n_obs, 3)) + + +def test_load_backed_adata_layers_and_no_x(tmp_path): + g = _open_group(_write_zarr(_adata(), tmp_path / "a.zarr")) + backed = load_backed_adata(g, keys=("layers/log1p",), cols=("cell_line",)) + + assert backed.X is None # X not requested + assert "log1p" in backed.layers # requested layer present + assert list(backed.obs.columns) == ["cell_line"] + + +def _open_group(path): + import zarr + + return zarr.open_group(path, mode="r") + + +# ── open_source / from_paths: the three source-value shapes ──────────────────────────────────── + + +def test_from_paths_single_adata_path_autodetects_adata(tmp_path): + pert, ctrl = _weights() + p = _write_zarr(_adata(), tmp_path / "a.zarr") + nodes = {"pert": Node("data", COLS, "X", pert), "ctrl": Node("data", COLS, ("X", "obsm/emb"), ctrl)} + s = Scheme.from_paths(sources={"data": p}, nodes=nodes, root="pert", seed=0, binds=(Bind("pert", "ctrl", ("cell_line",)),)) + + src = s.sources["data"] + assert isinstance(src, ad.AnnData) + assert isinstance(src.X, ad.abc.CSRDataset) # backed, not in-memory + assert list(src.obsm) == ["emb"] # union of the nodes' keys → only obsm/emb + assert "log1p" not in src.layers + + +def test_from_paths_collection_root_autodetects_dataset_collection(tmp_path): + pert, ctrl = _weights() + ap = _write_zarr(_adata(), tmp_path / "a.zarr") + cp = str(tmp_path / "coll.zarr") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + DatasetCollection(cp, mode="a").add_adatas([ap], groupby=list(COLS), shuffle=False) + + nodes = {"pert": Node("data", COLS, "X", pert), "ctrl": Node("data", COLS, "X", ctrl)} + s = Scheme.from_paths(sources={"data": cp}, nodes=nodes, root="pert", seed=0, binds=(Bind("pert", "ctrl", ("cell_line",)),)) + assert isinstance(s.sources["data"], DatasetCollection) + + +def test_from_paths_list_of_paths_gives_list_of_backed_adata(tmp_path): + pert, ctrl = _weights() + paths = [_write_zarr(_adata(seed=i), tmp_path / f"a{i}.zarr") for i in range(2)] + nodes = {"pert": Node("data", COLS, "X", pert), "ctrl": Node("data", COLS, "X", ctrl)} + s = Scheme.from_paths(sources={"data": paths}, nodes=nodes, root="pert", seed=0, binds=(Bind("pert", "ctrl", ("cell_line",)),)) + + src = s.sources["data"] + assert isinstance(src, list) and len(src) == 2 + assert all(isinstance(a, ad.AnnData) for a in src) + + +def test_from_paths_passes_through_constructed_containers(tmp_path): + pert, ctrl = _weights() + adata = _adata() # already in-memory AnnData + nodes = {"pert": Node("data", COLS, "X", pert), "ctrl": Node("data", COLS, "X", ctrl)} + s = Scheme.from_paths(sources={"data": adata}, nodes=nodes, root="pert", seed=0, binds=(Bind("pert", "ctrl", ("cell_line",)),)) + assert s.sources["data"] is adata # unchanged + + +# ── end-to-end: a from_paths scheme streams matched batches ──────────────────────────────────── + + +@pytest.mark.parametrize( + "make_source", + [ + pytest.param(lambda tp: _write_zarr(_adata(), tp / "a.zarr"), id="single-path"), + pytest.param(lambda tp: [_write_zarr(_adata(seed=i), tp / f"a{i}.zarr") for i in range(2)], id="list-of-paths"), + ], +) +def test_from_paths_streams_matched_batches(tmp_path, make_source): + pert, ctrl = _weights() + src = make_source(tmp_path) + nodes = {"pert": Node("data", COLS, "X", pert), "ctrl": Node("data", COLS, ("X", "obsm/emb"), ctrl)} + s = Scheme.from_paths(sources={"data": src}, nodes=nodes, root="pert", seed=0, binds=(Bind("pert", "ctrl", ("cell_line",)),)) + + batch = next(iter(DAGLoader(s, _cfg()))) + assert batch["target"].shape == (16, 5) + assert batch["source"].shape == (16, 5) # matched control rows + assert batch["source_reps"]["obsm/emb"].shape == (16, 3) # aligned obsm rep of the same cells diff --git a/tests/dagloader/test_pickle.py b/tests/dagloader/test_pickle.py new file mode 100644 index 00000000..5e680d5d --- /dev/null +++ b/tests/dagloader/test_pickle.py @@ -0,0 +1,51 @@ +"""DAGLoader is picklable mid-stream: the live annbatch iterators are dropped, RNG/state is kept.""" + +from __future__ import annotations + +import pickle + +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("annbatch") + +import anndata as ad + +from dagloader import Bind, DAGLoader, Node, SamplerConfig, Scheme, uniform + + +def _loader(seed=0): + rows = [(cl, dr) for cl in ("A", "B") for dr in ("control", "d1", "d2") for _ in range(16)] + obs = pd.DataFrame(rows, columns=["cell_line", "drug"]) + obs.index = obs.index.astype(str) + obs["control"] = obs["drug"] == "control" + adata = ad.AnnData(X=np.random.default_rng(0).normal(size=(len(obs), 4)).astype("float32"), obs=obs) + cols = ("cell_line", "drug") + combos = {tuple(r) for r in obs[list(cols)].to_numpy()} + pert = [c for c in combos if c[1] != "control"] + ctrl = [c for c in combos if c[1] == "control"] + scheme = Scheme( + sources={"data": adata}, + nodes={"pert": Node("data", cols, "X", uniform(pert)), "ctrl": Node("data", cols, "X", uniform(ctrl))}, + root="pert", + binds=(Bind("pert", "ctrl", common=("cell_line",)),), + seed=seed, + ) + # condition_fn omitted so the loader is plain-picklable (no closure); state/RNG is what we test here + return DAGLoader(scheme, SamplerConfig(batch_size=8, chunk_size=1, preload_nchunks=8)) + + +def test_pickle_mid_stream_and_resume_deterministic(): + loader = _loader() + it = iter(loader) + for _ in range(2): # advance into a pass → live annbatch generators + a wrapped sampler RNG exist + next(it) + blob = pickle.dumps(loader) # would raise "cannot pickle 'generator'" without __getstate__/__reduce__ + + # two loaders restored from the same checkpoint must produce identical streams + la, lb = pickle.loads(blob), pickle.loads(blob) + a = [next(la)["target"] for _ in range(4)] + b = [next(lb)["target"] for _ in range(4)] + assert all(np.array_equal(x, y) for x, y in zip(a, b, strict=True)) + assert a[0].shape == (8, 4) diff --git a/tests/dagloader/test_split.py b/tests/dagloader/test_split.py new file mode 100644 index 00000000..e8203cd7 --- /dev/null +++ b/tests/dagloader/test_split.py @@ -0,0 +1,195 @@ +"""Tests for :func:`dagloader.split_scheme` / :func:`dagloader.split_assignment`. + +The split is a weights-only transform on a :class:`~dagloader.Scheme`, so these exercise it on a scheme +built from a tiny in-memory ``AnnData`` via :func:`~dagloader.perturbation_scheme` — no +``DatasetCollection`` and no ``DAGLoader`` are constructed. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("annbatch") # dagloader imports annbatch; skip the module if it is not installed + +import anndata as ad + +from dagloader import ( + SamplerConfig, + Scheme, + perturbation_scheme, + resolve_split_configs, + split_assignment, + split_scheme, +) + + +def _toy_scheme(n_drugs: int = 5, cell_lines: tuple[str, ...] = ("A", "B"), seed: int = 0) -> Scheme: + """A perturbation scheme with ``n_drugs`` perturbations across ``cell_lines`` (+ a control each).""" + drugs = ["control", *[f"d{i}" for i in range(1, n_drugs + 1)]] + rows = [(cl, dr) for cl in cell_lines for dr in drugs for _ in range(2)] # 2 cells per combo + obs = pd.DataFrame(rows, columns=["cell_line", "drug"]) + obs["control"] = obs["drug"] == "control" + obs.index = obs.index.astype(str) + adata = ad.AnnData(X=np.zeros((len(obs), 3), dtype="float32"), obs=obs) + return perturbation_scheme( + adata, + context=["cell_line"], + perturbation=["drug"], + control_values={"drug": "control"}, + key="X", + seed=seed, + ) + + +def _combos(scheme: Scheme, node: str) -> set[tuple]: + n = scheme.nodes[node] + return {combo for combo, w in n.weights.items() if w > 0} + + +def _drugs(scheme: Scheme) -> set[str]: + root = scheme.nodes[scheme.root] + i = root.cols.index("drug") + return {combo[i] for combo, w in root.weights.items() if w > 0} + + +class TestSplitScheme: + def test_returns_named_schemes(self): + splits = split_scheme(_toy_scheme(), split_by=["drug"]) + assert set(splits) == {"train", "val", "test"} + for sch in splits.values(): + assert isinstance(sch, Scheme) + assert sch.root == "pert" # root node name unchanged, only its weights are restricted + + def test_partition_covers_and_is_disjoint(self): + scheme = _toy_scheme(n_drugs=5) + splits = split_scheme(scheme, split_by=["drug"], random_state=0) + drug_sets = [_drugs(s) for s in splits.values()] + # disjoint + for a in range(len(drug_sets)): + for b in range(a + 1, len(drug_sets)): + assert drug_sets[a].isdisjoint(drug_sets[b]) + # cover exactly the original perturbation drugs + assert set().union(*drug_sets) == _drugs(scheme) + # and the combinations partition the original root combinations + combo_union = set().union(*(_combos(s, "pert") for s in splits.values())) + assert combo_union == _combos(scheme, "pert") + + def test_default_ratio_sizes(self): + # 5 drugs, 60/20/20 -> 3/1/1 drug-groups; each drug spans 2 cell lines -> 6/2/2 combos + splits = split_scheme(_toy_scheme(n_drugs=5), split_by=["drug"], random_state=1) + assert {k: len(_drugs(v)) for k, v in splits.items()} == {"train": 3, "val": 1, "test": 1} + assert {k: len(_combos(v, "pert")) for k, v in splits.items()} == {"train": 6, "val": 2, "test": 2} + + def test_groups_by_projection_when_split_by_is_subset(self): + # split_by=["drug"] but cols=(cell_line, drug): both cell lines of a drug share a split + splits = split_scheme(_toy_scheme(cell_lines=("A", "B")), split_by=["drug"], random_state=2) + for sch in splits.values(): + root = sch.nodes[sch.root] + by_drug: dict[str, set[str]] = {} + di, ci = root.cols.index("drug"), root.cols.index("cell_line") + for combo, w in root.weights.items(): + if w > 0: + by_drug.setdefault(combo[di], set()).add(combo[ci]) + for drug, lines in by_drug.items(): + assert lines == {"A", "B"}, f"{drug} split across cell lines: {lines}" + + def test_controls_carried_through_unchanged(self): + scheme = _toy_scheme() + splits = split_scheme(scheme, split_by=["drug"]) + for sch in splits.values(): + assert sch.nodes["ctrl"].weights == scheme.nodes["ctrl"].weights + assert sch.sources is scheme.sources # same source objects, no copy + + def test_reproducible_for_same_seed(self): + a = split_scheme(_toy_scheme(), split_by=["drug"], random_state=7) + b = split_scheme(_toy_scheme(), split_by=["drug"], random_state=7) + assert {k: _drugs(v) for k, v in a.items()} == {k: _drugs(v) for k, v in b.items()} + + def test_force_training_values(self): + # d1 must land in train regardless of the shuffle seed + for seed in range(5): + splits = split_scheme( + _toy_scheme(), split_by=["drug"], force_training_values={"drug": "d1"}, random_state=seed + ) + assert "d1" in _drugs(splits["train"]) + assert "d1" not in _drugs(splits["val"]) + assert "d1" not in _drugs(splits["test"]) + + def test_split_by_context_column_with_custom_ratios(self): + # "or whatever": hold out a whole cell line, two-way split + scheme = _toy_scheme(cell_lines=("A", "B")) + splits = split_scheme(scheme, split_by=["cell_line"], ratios={"train": 0.5, "test": 0.5}) + assert set(splits) == {"train", "test"} + lines = [{c[scheme.nodes["pert"].cols.index("cell_line")] for c in _combos(s, "pert")} for s in splits.values()] + assert lines[0].isdisjoint(lines[1]) + assert set().union(*lines) == {"A", "B"} + for sch in splits.values(): # controls for BOTH lines stay available in each split + assert sch.nodes["ctrl"].weights == scheme.nodes["ctrl"].weights + + +class TestSplitAssignment: + def test_assignment_table(self): + scheme = _toy_scheme(n_drugs=5) + splits = split_scheme(scheme, split_by=["drug"], random_state=0) + df = split_assignment(splits) + assert list(df.columns) == ["cell_line", "drug", "split"] + assert len(df) == len(_combos(scheme, "pert")) # one row per target combination + assert set(df["split"]) == {"train", "val", "test"} + # every (cell_line, drug) appears exactly once + assert not df.duplicated(subset=["cell_line", "drug"]).any() + + +class TestSplitValidation: + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"split_by": []}, "split_by must be non-empty"), + ({"split_by": ["nope"]}, "not in the root node's cols"), + ({"split_by": ["drug"], "ratios": {"train": 0.5, "val": 0.2, "test": 0.2}}, "sum to 1.0"), + ({"split_by": ["drug"], "ratios": {"train": 1.2, "val": -0.2}}, "must all be > 0"), + ({"split_by": ["drug"], "force_training_values": {"cell_line": "A"}}, "subset of split_by"), + ], + ) + def test_invalid_arguments_raise(self, kwargs, match): + with pytest.raises(ValueError, match=match): + split_scheme(_toy_scheme(), **kwargs) + + def test_too_few_groups_for_ratios_raises(self): + # 2 cell lines cannot fill 3 non-empty splits + with pytest.raises(ValueError, match="received 0 of"): + split_scheme(_toy_scheme(cell_lines=("A", "B")), split_by=["cell_line"]) + + +class TestResolveSplitConfigs: + NAMES = ["train", "val", "test"] + + def test_single_config_object_applies_to_all(self): + cfg = SamplerConfig(batch_size=256, chunk_size=1, preload_nchunks=256) + out = resolve_split_configs(cfg, self.NAMES) + assert set(out) == set(self.NAMES) + assert all(out[n] is cfg for n in self.NAMES) # same object shared + + def test_per_split_mapping_ok(self): + train, val, test = ( + SamplerConfig(batch_size=256, chunk_size=1, preload_nchunks=256), + SamplerConfig(batch_size=64, chunk_size=1, preload_nchunks=64), + SamplerConfig(batch_size=64, chunk_size=1, preload_nchunks=64), + ) + out = resolve_split_configs({"train": train, "val": val, "test": test}, self.NAMES) + assert out == {"train": train, "val": val, "test": test} + + def test_per_split_missing_split_raises(self): + with pytest.raises(ValueError, match="missing config"): + resolve_split_configs({"train": SamplerConfig(batch_size=1, chunk_size=1, preload_nchunks=1)}, self.NAMES) + + def test_per_split_unknown_split_raises(self): + cfg = SamplerConfig(batch_size=1, chunk_size=1, preload_nchunks=1) + with pytest.raises(ValueError, match="unknown split"): + resolve_split_configs({"train": cfg, "val": cfg, "test": cfg, "bogus": cfg}, self.NAMES) + + def test_per_split_non_sampler_config_value_raises(self): + cfg = SamplerConfig(batch_size=1, chunk_size=1, preload_nchunks=1) + with pytest.raises(ValueError, match="must be SamplerConfig"): + resolve_split_configs({"train": cfg, "val": cfg, "test": {"batch_size": 1}}, self.NAMES) diff --git a/tests/data/test_cfsampler.py b/tests/data/test_cfsampler.py index 0ae5405a..40b458df 100644 --- a/tests/data/test_cfsampler.py +++ b/tests/data/test_cfsampler.py @@ -1,7 +1,7 @@ import numpy as np import pytest -from cellflow.data._dataloader import OOCTrainSampler, PredictionSampler, TrainSampler +from cellflow.data._legacy import OOCTrainSampler, PredictionSampler, TrainSampler from cellflow.data._datamanager import DataManager @@ -86,7 +86,7 @@ def test_sampling_no_combinations(self, adata_perturbation, batch_size: int): class TestValidationSampler: @pytest.mark.parametrize("n_conditions_on_log_iteration", [None, 1, 3]) def test_valid_sampler(self, adata_perturbation, n_conditions_on_log_iteration): - from cellflow.data._dataloader import ValidationSampler + from cellflow.data._legacy import ValidationSampler from cellflow.data._datamanager import DataManager control_key = "control" diff --git a/tests/data/test_condition.py b/tests/data/test_condition.py new file mode 100644 index 00000000..f01aed36 --- /dev/null +++ b/tests/data/test_condition.py @@ -0,0 +1,132 @@ +"""Tests for the extracted condition helpers in :mod:`cellflow.data._condition`. + +These guard the "act the same" contract: the standalone functions must agree with the +:class:`~cellflow.data._datamanager.DataManager` methods they were extracted from (DataManager +delegates to them, so any divergence is a regression). +""" + +import anndata as ad +import numpy as np +import pytest + +from cellflow.data._condition import build_condition_data, enumerate_perturbations, get_max_combination_length +from cellflow.data._datamanager import DataManager + +SPECS = [ + {"drug": ["drug1"]}, + {"drug": ["drug1", "drug2"]}, + {"drug": ["drug1", "drug2"], "dosage": ["dosage_a", "dosage_b"]}, + {"drug": ["drug_a", "drug_b", "drug_c"], "dosage": ["dosage_a", "dosage_b", "dosage_c"]}, +] + + +class TestGetMaxCombinationLength: + @pytest.mark.parametrize("spec", SPECS) + @pytest.mark.parametrize("override", [None, 1, 2, 5, 100]) + def test_matches_datamanager(self, spec, override): + # single source of truth: standalone == DataManager (which delegates to it) + assert get_max_combination_length(spec, override) == DataManager._get_max_combination_length(spec, override) + + def test_no_override_is_observed_max(self): + assert get_max_combination_length({"drug": ["d1", "d2"], "dose": ["x1"]}) == 2 + + def test_larger_override_kept(self): + assert get_max_combination_length({"drug": ["d1", "d2"]}, 5) == 5 + + def test_smaller_override_raised_to_observed_with_warning(self, caplog): + assert get_max_combination_length({"drug": ["d1", "d2"]}, 1) == 2 + + +ENUM_CASES = [ + ({"drug": ["drug1"]}, ["cell_type"], []), + ({"drug": ["drug1", "drug2"]}, ["cell_type"], []), + ({"drug": ["drug1"]}, [], []), + ({"drug": ["drug1"]}, [], ["dosage_c"]), + ({"drug": ["drug1"], "dosage": ["dosage_a"]}, ["cell_type"], []), +] + + +class TestEnumeratePerturbations: + @pytest.mark.parametrize(("perturbation_covariates", "split_covariates", "sample_covariates"), ENUM_CASES) + def test_matches_datamanager_idx_to_covariates( + self, + adata_perturbation: ad.AnnData, + perturbation_covariates, + split_covariates, + sample_covariates, + ): + dm = DataManager( + adata_perturbation, + sample_rep="X", + control_key="control", + perturbation_covariates=perturbation_covariates, + perturbation_covariate_reps={"drug": "drug"}, + split_covariates=split_covariates, + sample_covariates=sample_covariates, + ) + expected = { + int(k): tuple(v) + for k, v in dm._get_condition_data(adata=adata_perturbation).perturbation_idx_to_covariates.items() + } + got = enumerate_perturbations( + adata_perturbation.obs, + control_key="control", + perturb_covar_keys=dm._perturb_covar_keys, + split_covariates=dm._split_covariates, + sample_covariates=dm._sample_covariates, + ) + assert got == expected + + +BUILD_CASES = [ + ({"drug": ["drug1"]}, {"drug": "drug"}, ["cell_type"], []), + ({"drug": ["drug1", "drug2"]}, {"drug": "drug"}, ["cell_type"], []), + ({"drug": ["drug1"]}, {"drug": "drug"}, [], ["dosage_c"]), + ({"drug": ["drug1"], "dosage": ["dosage_a"]}, {"drug": "drug"}, ["cell_type"], []), + ({"drug": ["drug1"]}, {}, ["cell_type"], []), # no rep -> primary one-hot-encoder path +] + + +class TestBuildConditionData: + @pytest.mark.parametrize( + ("perturbation_covariates", "perturbation_covariate_reps", "split_covariates", "sample_covariates"), + BUILD_CASES, + ) + def test_matches_datamanager_condition_data( + self, + adata_perturbation: ad.AnnData, + perturbation_covariates, + perturbation_covariate_reps, + split_covariates, + sample_covariates, + ): + dm = DataManager( + adata_perturbation, + sample_rep="X", + control_key="control", + perturbation_covariates=perturbation_covariates, + perturbation_covariate_reps=perturbation_covariate_reps, + split_covariates=split_covariates, + sample_covariates=sample_covariates, + ) + ref = dm._get_condition_data(adata=adata_perturbation).condition_data + got = build_condition_data( + adata_perturbation.obs, + adata_perturbation.uns, + control_key="control", + perturb_covar_keys=dm.perturb_covar_keys, + split_covariates=dm.split_covariates, + sample_covariates=dm.sample_covariates, + perturbation_covariates=dm.perturbation_covariates, + covariate_reps=dm.covariate_reps, + covar_to_idx=dm._covar_to_idx, + is_categorical=dm.is_categorical, + primary_one_hot_encoder=dm.primary_one_hot_encoder, + primary_group=dm.primary_group, + linked_perturb_covars=dm.linked_perturb_covars, + max_combination_length=dm.max_combination_length, + null_value=dm.null_value, + ) + assert set(got) == set(ref) + for key in ref: + np.testing.assert_array_equal(got[key], ref[key]) diff --git a/tests/model/test_annbatch_ooc.py b/tests/model/test_annbatch_ooc.py new file mode 100644 index 00000000..4e023094 --- /dev/null +++ b/tests/model/test_annbatch_ooc.py @@ -0,0 +1,188 @@ +"""Out-of-core streaming over a real annbatch ``DatasetCollection``, incl. the ``chunk_size>1`` rule. + +``chunk_size>1`` reads contiguous slices, so every contiguous run of each category must be +``>= chunk_size`` (annbatch's run-length rule; a category may span several runs). We build collections +grouped via ``add_adatas(groupby=...)``, interleaved (short runs → error), and fragmented-but-valid +(multiple long runs per category → accepted), and check chunked streaming behaves accordingly. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest +import scipy.sparse as sp + +pytest.importorskip("annbatch") + +import anndata as ad +from annbatch import DatasetCollection + +import cellflow +from cellflow.data._annbatch import assert_source_chunkable +from dagloader import SamplerConfig + +_PREP = { + "sample_rep": "X", + "control_key": "control", + "perturbation_covariates": {"drug": ["drug"]}, + "split_covariates": ["cell_line"], +} + + +def _adata(*, interleaved: bool, n_per_combo=40, drugs=("control", "d1", "d2", "d3"), lines=("A", "B"), sparse=False): + rng = np.random.default_rng(0) + rows = [(cl, dr) for cl in lines for dr in drugs for _ in range(n_per_combo)] + if interleaved: + rng.shuffle(rows) # ungrouped input + obs = pd.DataFrame(rows, columns=["cell_line", "drug"]) + obs["control"] = obs["drug"] == "control" + obs.index = obs.index.astype(str) + x = rng.normal(size=(len(obs), 5)).astype("float32") + return ad.AnnData(X=sp.csr_matrix(x) if sparse else x, obs=obs) + + +def _collection(tmp_path, *, grouped: bool, sparse=False) -> DatasetCollection: + h5 = tmp_path / "a.h5ad" + _adata(interleaved=not grouped, sparse=sparse).write_h5ad(h5) + dc = DatasetCollection(str(tmp_path / "c.zarr"), mode="a") + if grouped: # sort by the grouping columns on add → contiguous category runs + dc.add_adatas([str(h5)], groupby=["cell_line", "drug"], shuffle=False) + else: # preserve the interleaved input order + dc.add_adatas([str(h5)], shuffle=False) + return dc + + +def _prepare_model_small(cf): + cf.prepare_model( + pooling="mean", condition_embedding_dim=8, time_encoder_dims=(8,), hidden_dims=(8,), decoder_dims=(8,) + ) + + +class TestOutOfCore: + def test_grouped_collection_chunked_trains(self, tmp_path): + cf = cellflow.model.CellFlowAnnbatch() + cf.prepare_data( + source=_collection(tmp_path, grouped=True), + sampler_config=SamplerConfig(batch_size=16, chunk_size=4, preload_nchunks=16), + **_PREP, + ) + assert cf._data_dim == 5 + _prepare_model_small(cf) + cf.train(num_iterations=2, valid_freq=100) + assert cf.solver is not None + + def test_ungrouped_collection_chunked_raises(self, tmp_path): + cf = cellflow.model.CellFlowAnnbatch() + with pytest.raises(ValueError, match="grouped|groupby"): + cf.prepare_data( + source=_collection(tmp_path, grouped=False), + sampler_config=SamplerConfig(batch_size=16, chunk_size=4, preload_nchunks=16), + **_PREP, + ) + + def test_ungrouped_collection_chunk1_ok(self, tmp_path): + # chunk_size=1 streams per-row → no grouping needed even for an interleaved collection + cf = cellflow.model.CellFlowAnnbatch() + cf.prepare_data( + source=_collection(tmp_path, grouped=False), + sampler_config=SamplerConfig(batch_size=16, chunk_size=1, preload_nchunks=16), + **_PREP, + ) + assert cf._dataloader.sample()["tgt_cell_data"].shape == (16, 5) + + def test_grouped_collection_split_chunked(self, tmp_path): + cf = cellflow.model.CellFlowAnnbatch() + cf.prepare_data( + source=_collection(tmp_path, grouped=True), + sampler_config=SamplerConfig(batch_size=16, chunk_size=4, preload_nchunks=16), + split_by=["drug"], + split_ratios={"train": 0.5, "val": 0.25, "test": 0.25}, + **_PREP, + ) + assert set(cf.split_eval_loaders) == {"val", "test"} # non-train splits read via DAGEvalLoader + + def test_in_memory_unsorted_source_is_auto_grouped(self): + # an in-memory (interleaved) AnnData source is grouped automatically → chunk_size>1 works + cf = cellflow.model.CellFlowAnnbatch() + cf.prepare_data( + source=_adata(interleaved=True), + sampler_config=SamplerConfig(batch_size=16, chunk_size=4, preload_nchunks=16), + **_PREP, + ) + assert cf._dataloader.sample()["src_cell_data"].shape == (16, 5) + + def test_fragmented_collection_chunked_ok(self, tmp_path): + # each category in TWO contiguous runs (fragmented) — valid as long as every run >= chunk_size. + # This is the case annbatch accepts but a "one run per class" check would wrongly reject. + block = 10 + rows = [(cl, dr) for _ in range(2) for cl in ("A", "B") for dr in ("control", "d1", "d2") for _ in range(block)] + obs = pd.DataFrame(rows, columns=["cell_line", "drug"]) + obs["control"] = obs["drug"] == "control" + obs.index = obs.index.astype(str) + x = np.random.default_rng(0).normal(size=(len(obs), 5)).astype("float32") + (tmp_path / "f.h5ad").parent.mkdir(exist_ok=True) + ad.AnnData(X=x, obs=obs).write_h5ad(tmp_path / "f.h5ad") + dc = DatasetCollection(str(tmp_path / "fc.zarr"), mode="a").add_adatas( + [str(tmp_path / "f.h5ad")], shuffle=False + ) + + from dagloader._io import leaf_codes, obs_columns + + codes, _ = leaf_codes(obs_columns(dc, ["cell_line", "drug"]), ["cell_line", "drug"]) + n_runs = 1 + int((np.diff(codes) != 0).sum()) + assert n_runs > 6, f"expected a fragmented collection (>6 runs for 6 categories), got {n_runs}" + + cf = cellflow.model.CellFlowAnnbatch() + cf.prepare_data( # chunk_size=4 <= run length 10 → must be accepted + source=dc, sampler_config=SamplerConfig(batch_size=8, chunk_size=4, preload_nchunks=8), **_PREP + ) + assert cf._dataloader.sample()["tgt_cell_data"].shape == (8, 5) + + def test_sparse_collection_prepares(self, tmp_path): + # a SPARSE-X collection stores X as a zarr *group* (no .shape); key_backings must wrap it so + # annbatch's add_datasets accepts it — regression for the "'Group' has no attribute 'shape'" error. + # Construction (which runs add_datasets) is the assertion here; reading sparse batches needs cupy. + cf = cellflow.model.CellFlowAnnbatch() + cf.prepare_data( + source=_collection(tmp_path, grouped=True, sparse=True), + sampler_config=SamplerConfig(batch_size=16, chunk_size=1, preload_nchunks=16), + **_PREP, + ) + assert cf._data_dim == 5 # DAGLoader built (add_datasets accepted the sparse-group backing) + + def test_control_in_memory_materializes(self, tmp_path): + # control_in_memory tells dagloader to materialize the ctrl node into RAM (Node.in_memory); + # prepare_data building the loader runs materialize_node over the (sparse) collection. + cf = cellflow.model.CellFlowAnnbatch() + cf.prepare_data( + source=_collection(tmp_path, grouped=True, sparse=True), + sampler_config=SamplerConfig(batch_size=16, chunk_size=1, preload_nchunks=16), + control_in_memory=True, + **_PREP, + ) + assert cf._scheme.nodes["ctrl"].in_memory is True # cellflow flagged it + assert isinstance(cf._dataloader._loader._nodes["ctrl"], ad.AnnData) # dagloader materialized it + + +class TestChunkableCheck: + """`assert_source_chunkable` matches annbatch's rule: runs >= chunk_size, fragmentation allowed.""" + + @staticmethod + def _frag(block): # each (cell_line, drug) category appears in 2 runs of `block` + rows = [(cl, dr) for _ in range(2) for cl in ("A", "B") for dr in ("d1", "d2") for _ in range(block)] + obs = pd.DataFrame(rows, columns=["cell_line", "drug"]) + obs.index = obs.index.astype(str) + return ad.AnnData(X=np.zeros((len(obs), 2), dtype="float32"), obs=obs) + + def test_accepts_fragmented_runs(self): + # fragmented but every run (10) >= chunk_size (4) → no raise (annbatch accepts this too) + assert_source_chunkable(self._frag(10), ("cell_line", "drug"), 4) + + def test_rejects_short_run(self): + # runs of 3 < chunk_size 4 → raise, with the groupby hint + with pytest.raises(ValueError, match="run of only 3|groupby|chunk_size"): + assert_source_chunkable(self._frag(3), ("cell_line", "drug"), 4) + + def test_chunk_size_one_always_ok(self): + assert_source_chunkable(self._frag(1), ("cell_line", "drug"), 1) # runs of 1 >= 1 diff --git a/tests/model/test_annbatch_path.py b/tests/model/test_annbatch_path.py new file mode 100644 index 00000000..37b67a2c --- /dev/null +++ b/tests/model/test_annbatch_path.py @@ -0,0 +1,45 @@ +"""Tests for the in-memory :class:`~cellflow.model.CellFlow` constructor / data prep. + +Covers the ``adata``-optional constructor (deprecation of the constructor argument) and passing +``adata`` to :meth:`prepare_data`. The streaming path now lives on +:class:`~cellflow.model.CellFlowAnnbatch`. +""" + +import anndata as ad +import pytest + +import cellflow + +PERT_COVARS = {"drug": ["drug1"]} +PERT_COVAR_REPS = {"drug": "drug"} + + +class TestAnnbatchPathScaffolding: + def test_constructor_adata_emits_futurewarning(self, adata_perturbation: ad.AnnData): + with pytest.warns(FutureWarning, match="prepare_data"): + cellflow.model.CellFlow(adata_perturbation) + + def test_constructor_without_adata_is_silent(self): + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("error", FutureWarning) + cf = cellflow.model.CellFlow() + assert cf.adata is None + + def test_prepare_data_with_adata_kwarg(self, adata_perturbation: ad.AnnData): + cf = cellflow.model.CellFlow() + cf.prepare_data( + sample_rep="X", + control_key="control", + perturbation_covariates=PERT_COVARS, + perturbation_covariate_reps=PERT_COVAR_REPS, + adata=adata_perturbation, + ) + assert cf.train_data is not None + assert cf.adata is adata_perturbation + + def test_prepare_data_without_any_adata_raises(self): + cf = cellflow.model.CellFlow() + with pytest.raises(ValueError, match="No `adata` provided"): + cf.prepare_data(sample_rep="X", control_key="control", perturbation_covariates=PERT_COVARS) diff --git a/tests/model/test_annbatch_save.py b/tests/model/test_annbatch_save.py new file mode 100644 index 00000000..6bab940e --- /dev/null +++ b/tests/model/test_annbatch_save.py @@ -0,0 +1,93 @@ +"""save/load and get_condition_embedding for annbatch-path models. + +`save` pickles the whole model incl. the streaming loaders. The loaders drop their live annbatch +iterators on pickle (generators aren't picklable) but keep the RNG/schedule state, so a reloaded model +resumes the same reproducible stream. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("annbatch") + +import anndata as ad + +import cellflow +from dagloader import SamplerConfig + +_PREP = { + "sample_rep": "X", + "control_key": "control", + "perturbation_covariates": {"drug": ["drug"]}, + "split_covariates": ["cell_line"], +} +_CFG = SamplerConfig(batch_size=8, chunk_size=1, preload_nchunks=8) + + +def _adata(): + rng = np.random.default_rng(0) + rows = [(cl, dr) for cl in ("A", "B") for dr in ("control", "d1", "d2") for _ in range(16)] + obs = pd.DataFrame(rows, columns=["cell_line", "drug"]) + obs["control"] = obs["drug"] == "control" + obs.index = obs.index.astype(str) + return ad.AnnData(X=rng.normal(size=(len(obs), 5)).astype("float32"), obs=obs) + + +def _prepared(seed=7): + cf = cellflow.model.CellFlowAnnbatch() + cf.prepare_data(source=_adata(), sampler_config=_CFG, seed=seed, **_PREP) + return cf + + +def _stream(cf, n): + return [cf._dataloader.sample()["tgt_cell_data"].copy() for _ in range(n)] + + +class TestAnnbatchSave: + def test_save_load_after_prepare(self, tmp_path): + cf = _prepared() + cf.save(str(tmp_path), file_prefix="p", overwrite=True) + loaded = cellflow.model.CellFlowAnnbatch.load(str(tmp_path / "p_CellFlowAnnbatch.pkl")) + assert loaded._scheme is not None + assert loaded._dataloader.sample()["tgt_cell_data"].shape == (8, 5) + + def test_save_load_mid_stream_resumes_deterministically(self, tmp_path): + # advance both models to the same point, save, load — the resumed streams must match + a, b = _prepared(), _prepared() + _stream(a, 2) + _stream(b, 2) + a.save(str(tmp_path), file_prefix="a", overwrite=True) + b.save(str(tmp_path), file_prefix="b", overwrite=True) + la = cellflow.model.CellFlowAnnbatch.load(str(tmp_path / "a_CellFlowAnnbatch.pkl")) + lb = cellflow.model.CellFlowAnnbatch.load(str(tmp_path / "b_CellFlowAnnbatch.pkl")) + ra, rb = _stream(la, 3), _stream(lb, 3) + assert all(np.array_equal(x, y) for x, y in zip(ra, rb, strict=True)) + # and the resumed state is preserved (not reset to the seed → differs from a fresh model) + fresh = _stream(_prepared(), 3) + assert not all(np.array_equal(x, y) for x, y in zip(ra, fresh, strict=True)) + + def test_save_load_after_training(self, tmp_path): + cf = _prepared() + cf.prepare_model( + pooling="mean", condition_embedding_dim=8, time_encoder_dims=(8,), hidden_dims=(8,), decoder_dims=(8,) + ) + cf.train(num_iterations=2, valid_freq=100) + cf.save(str(tmp_path), file_prefix="t", overwrite=True) + loaded = cellflow.model.CellFlowAnnbatch.load(str(tmp_path / "t_CellFlowAnnbatch.pkl")) + assert loaded.solver is not None and loaded._dataloader is not None + assert loaded._dataloader.sample()["tgt_cell_data"].shape == (8, 5) # loader still usable + + def test_get_condition_embedding_without_adata_warns(self, tmp_path): + cf = _prepared() + cf.prepare_model( + pooling="mean", condition_embedding_dim=8, time_encoder_dims=(8,), hidden_dims=(8,), decoder_dims=(8,) + ) + cf.train(num_iterations=2, valid_freq=100) + cov = _adata().obs.drop_duplicates(subset=["cell_line", "drug"]) # carries the control column too + with pytest.warns(UserWarning, match="streaming path"): + df_mean, df_var = cf.get_condition_embedding(cov) # default key_added, no adata → warns, no crash + assert isinstance(df_mean, pd.DataFrame) and isinstance(df_var, pd.DataFrame) + assert len(df_mean) == len(cov) # one embedding row per provided condition diff --git a/tests/model/test_annbatch_split.py b/tests/model/test_annbatch_split.py new file mode 100644 index 00000000..79000ecb --- /dev/null +++ b/tests/model/test_annbatch_split.py @@ -0,0 +1,159 @@ +"""Tests for the annbatch path split wiring on :class:`~cellflow.model.CellFlow`. + +The Scheme + condition + loaders are built from an in-memory ``AnnData`` used as the ``source`` (the +``dagloader`` is container-agnostic), so no ``DatasetCollection`` is needed. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("annbatch") # dagloader (and thus the annbatch path) needs annbatch + +import anndata as ad + +import cellflow +from dagloader import SamplerConfig, perturbation_scheme + +_CFG = SamplerConfig(batch_size=8, chunk_size=1, preload_nchunks=8) + + +def _toy_adata(n_per_combo: int = 8, drugs=("control", "d1", "d2", "d3", "d4", "d5"), lines=("A", "B")): + rng = np.random.default_rng(0) + rows = [(cl, dr) for cl in lines for dr in drugs for _ in range(n_per_combo)] + obs = pd.DataFrame(rows, columns=["cell_line", "drug"]) + obs["control"] = obs["drug"] == "control" + obs.index = obs.index.astype(str) + x = rng.normal(size=(len(obs), 5)).astype("float32") + return ad.AnnData(X=x, obs=obs) + + +def _toy_scheme(): + return perturbation_scheme( + _toy_adata(n_per_combo=2), + context=["cell_line"], + perturbation=["drug"], + control_values={"drug": "control"}, + key="X", + ) + + +class TestSplitAnnbatchData: + """`split_annbatch_data` operates on an already-built Scheme (injected here).""" + + def test_without_scheme_raises(self): + cf = cellflow.model.CellFlowAnnbatch() + with pytest.raises(ValueError, match="No annbatch `Scheme`"): + cf.split_annbatch_data(split_by=["drug"]) + + def test_split_with_injected_scheme(self): + cf = cellflow.model.CellFlowAnnbatch() + cf._scheme = _toy_scheme() + df = cf.split_annbatch_data(split_by=["drug"], random_state=0) + assert list(df.columns) == ["cell_line", "drug", "split"] + assert set(df["split"]) == {"train", "val", "test"} + assert set(cf._split_schemes) == {"train", "val", "test"} + for sch in cf._split_schemes.values(): # controls carried into every split + assert sch.nodes["ctrl"].weights == cf._scheme.nodes["ctrl"].weights + + +class TestPrepareAnnbatchData: + """`prepare_data` builds Scheme + condition + loaders from an AnnData source.""" + + def _prepare(self, cf, **kwargs): + return cf.prepare_data( + source=_toy_adata(), + sample_rep="X", + control_key="control", + perturbation_covariates={"drug": ["drug"]}, + split_covariates=["cell_line"], + sampler_config=_CFG, + **kwargs, + ) + + def test_requires_sampler_config(self): + cf = cellflow.model.CellFlowAnnbatch() + with pytest.raises(ValueError, match="sampler_config` is required"): + cf.prepare_data( + source=_toy_adata(), + sample_rep="X", + control_key="control", + perturbation_covariates={"drug": ["drug"]}, + ) + + def test_builds_scheme_condition_and_loader_without_split(self): + cf = cellflow.model.CellFlowAnnbatch() + self._prepare(cf) + assert cf._scheme is not None and cf._scheme.root == "pert" + assert cf._dataloader is not None # DAGTrainAdapter wired for train() + assert cf.split_eval_loaders == {} # no split → no eval loaders + # condition embeddings assembled (drug is categorical → one-hot), data dim from X + assert set(cf._condition_data) == {"drug"} + assert cf._data_dim == 5 + assert cf._split_schemes is None + + def test_split_builds_per_split_loaders(self): + cf = cellflow.model.CellFlowAnnbatch() + assert self._prepare(cf, split_by=["drug"], split_random_state=0) is None # prepare_* returns None + assert set(cf._split_assignment["split"]) == {"train", "val", "test"} + assert set(cf.split_eval_loaders) == {"val", "test"} # eval loaders for the non-train splits + assert set(cf._annbatch_sampler_configs) == {"train", "val", "test"} + assert cf._dataloader is not None + + def test_per_split_sampler_config(self): + cf = cellflow.model.CellFlowAnnbatch() + train = SamplerConfig(batch_size=8, chunk_size=1, preload_nchunks=8) + small = SamplerConfig(batch_size=4, chunk_size=1, preload_nchunks=4) + cf.prepare_data( + source=_toy_adata(), + sample_rep="X", + control_key="control", + perturbation_covariates={"drug": ["drug"]}, + split_covariates=["cell_line"], + split_by=["drug"], + sampler_config={"train": train, "val": small, "test": small}, + ) + assert cf._annbatch_sampler_configs["train"].batch_size == 8 + assert cf._annbatch_sampler_configs["val"].batch_size == 4 + + def test_per_split_missing_split_raises(self): + cf = cellflow.model.CellFlowAnnbatch() + with pytest.raises(ValueError, match="missing config"): + cf.prepare_data( + source=_toy_adata(), + sample_rep="X", + control_key="control", + perturbation_covariates={"drug": ["drug"]}, + split_covariates=["cell_line"], + split_by=["drug"], + sampler_config={"train": _CFG}, + ) + + def test_streamed_batch_shapes(self): + cf = cellflow.model.CellFlowAnnbatch() + self._prepare(cf) + batch = cf._dataloader.sample() + assert batch["src_cell_data"].shape == (8, 5) + assert batch["tgt_cell_data"].shape == (8, 5) + assert batch["condition"]["drug"].shape[0] == 1 # one condition per batch, leading axis + + def test_sparse_source_batches_densified(self): + # dagloader streams sparse for a sparse source; the model-boundary adapter densifies. + import scipy.sparse as sp + + adata = _toy_adata() + adata.X = sp.csr_matrix(adata.X) + cf = cellflow.model.CellFlowAnnbatch() + cf.prepare_data( + source=adata, + sample_rep="X", + control_key="control", + perturbation_covariates={"drug": ["drug"]}, + split_covariates=["cell_line"], + sampler_config=_CFG, + ) + batch = cf._dataloader.sample() + assert not sp.issparse(batch["src_cell_data"]) and not sp.issparse(batch["tgt_cell_data"]) + assert batch["src_cell_data"].shape == (8, 5) and batch["tgt_cell_data"].shape == (8, 5) diff --git a/tests/model/test_annbatch_train.py b/tests/model/test_annbatch_train.py new file mode 100644 index 00000000..3b67e3a7 --- /dev/null +++ b/tests/model/test_annbatch_train.py @@ -0,0 +1,72 @@ +"""End-to-end training over the annbatch/dagloader streaming path (in-memory AnnData as source).""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("annbatch") + +import anndata as ad + +import cellflow +from dagloader import SamplerConfig + + +def _toy_adata(n_per_combo: int = 20, drugs=("control", "d1", "d2", "d3"), lines=("A", "B")): + rng = np.random.default_rng(0) + rows = [(cl, dr) for cl in lines for dr in drugs for _ in range(n_per_combo)] + obs = pd.DataFrame(rows, columns=["cell_line", "drug"]) + obs["control"] = obs["drug"] == "control" + obs.index = obs.index.astype(str) + x = rng.normal(size=(len(obs), 5)).astype("float32") + return ad.AnnData(X=x, obs=obs) + + +def _prepare_model_small(cf): + cf.prepare_model( + pooling="mean", + condition_embedding_dim=8, + time_encoder_dims=(8,), + hidden_dims=(8,), + decoder_dims=(8,), + ) + + +class TestAnnbatchTraining: + def test_train_runs_without_split(self): + cf = cellflow.model.CellFlowAnnbatch() + cf.prepare_data( + source=_toy_adata(), + sample_rep="X", + control_key="control", + perturbation_covariates={"drug": ["drug"]}, + split_covariates=["cell_line"], + sampler_config=SamplerConfig(batch_size=16, chunk_size=1, preload_nchunks=16), + ) + _prepare_model_small(cf) + cf.train(num_iterations=2, valid_freq=100) + assert cf.solver is not None + assert cf.dataloader is not None + + def test_train_runs_on_train_split(self): + cf = cellflow.model.CellFlowAnnbatch() + cf.prepare_data( + source=_toy_adata(), + sample_rep="X", + control_key="control", + perturbation_covariates={"drug": ["drug"]}, + split_covariates=["cell_line"], + sampler_config=SamplerConfig(batch_size=16, chunk_size=1, preload_nchunks=16), + split_by=["drug"], + split_ratios={"train": 0.5, "val": 0.25, "test": 0.25}, + split_random_state=0, + ) + _prepare_model_small(cf) + cf.train(num_iterations=2, valid_freq=100) + assert cf.solver is not None + # val/test are read via DAGEvalLoader (control-rooted eval), not streamed as train-style loaders + assert set(cf.split_eval_loaders) == {"val", "test"} + val_out = next(cf.split_eval_loaders["val"].iter_conditions()) + assert "target" in val_out and "source" in val_out diff --git a/tests/model/test_annbatch_validation.py b/tests/model/test_annbatch_validation.py new file mode 100644 index 00000000..945e3836 --- /dev/null +++ b/tests/model/test_annbatch_validation.py @@ -0,0 +1,104 @@ +"""Validation in the streaming path: each condition's full cell set is read via ``DAGEvalLoader``. + +Validation preserves the legacy per-condition contract (``{source, condition, target}`` dicts, one entry +per condition) through :class:`~cellflow.data._dataloader.DAGEvalAdapter`, but reads cells via +annbatch slice reads instead of boolean-masking a materialized matrix. The regression that matters: cells +are read at the real ``sample_rep`` (e.g. an obsm key), not the internal ``"X"`` encoder-factory +placeholder. The ``val``/``test`` splits from ``split_by`` are auto-wired as evaluation sources. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("annbatch") + +import anndata as ad + +import cellflow +from cellflow.data._dataloader import DAGEvalAdapter +from dagloader import SamplerConfig + +_CFG = SamplerConfig(batch_size=8, chunk_size=1, preload_nchunks=8) + + +def _adata(*, rep_dim=None, x_dim=5, n_per_combo=8, drugs=("control", "d1", "d2", "d3"), lines=("A", "B"), seed=0): + rng = np.random.default_rng(seed) + rows = [(cl, dr) for cl in lines for dr in drugs for _ in range(n_per_combo)] + obs = pd.DataFrame(rows, columns=["cell_line", "drug"]) + obs["control"] = obs["drug"] == "control" + obs.index = obs.index.astype(str) + adata = ad.AnnData(X=rng.normal(size=(len(obs), x_dim)).astype("float32"), obs=obs) + if rep_dim is not None: + adata.obsm["X_pca"] = rng.normal(size=(len(obs), rep_dim)).astype("float32") + return adata + + +def _prepare(cf, source, *, sample_rep="X", **kwargs): + cf.prepare_data( + source=source, + sample_rep=sample_rep, + control_key="control", + perturbation_covariates={"drug": ["drug"]}, + split_covariates=["cell_line"], + sampler_config=_CFG, + **kwargs, + ) + + +class TestAnnbatchValidation: + def test_prepare_validation_without_setup_raises(self): + cf = cellflow.model.CellFlowAnnbatch() + with pytest.raises(ValueError, match="prepare_data"): + cf.prepare_validation_data(source=_adata(), name="val") + + def test_validation_sampler_reads_matched_source_and_target(self): + cf = cellflow.model.CellFlowAnnbatch() + _prepare(cf, _adata()) # sample_rep="X" + cf.prepare_validation_data(source=_adata(seed=1), name="val") + vs = cf.validation_data["val"] + assert isinstance(vs, DAGEvalAdapter) + batch = vs.sample("on_train_end") + assert set(batch) == {"source", "condition", "target"} + # control-rooted: one batch per control population (2 cell lines), keyed by the drawn (line, drug) + assert len(batch["target"]) == 2 + for k in batch["target"]: + assert k[1] != "control" # target is a perturbed condition + assert np.asarray(batch["target"][k]).shape[1] == 5 # X dim + assert np.asarray(batch["source"][k]).shape[1] == 5 # matched controls, same rep + + def test_validation_reads_obsm_sample_rep_not_x(self): + # training streams X_pca (dim 3); validation must read X_pca, NOT X (dim 5) + cf = cellflow.model.CellFlowAnnbatch() + _prepare(cf, _adata(rep_dim=3), sample_rep="X_pca") + assert cf._data_dim == 3 + cf.prepare_validation_data(source=_adata(rep_dim=3, seed=2), name="val") + batch = cf.validation_data["val"].sample("on_train_end") + for k in batch["target"]: + assert np.asarray(batch["target"][k]).shape[1] == 3 + assert np.asarray(batch["source"][k]).shape[1] == 3 + + def test_val_test_splits_auto_wired_as_eval_sources(self): + cf = cellflow.model.CellFlowAnnbatch() + _prepare( + cf, + _adata(), + split_by=["drug"], + split_ratios={"train": 0.5, "val": 0.25, "test": 0.25}, + split_random_state=0, + ) + # non-train splits become DAGEvalLoaders; "val" also feeds training-time validation + assert set(cf.split_eval_loaders) == {"val", "test"} + assert isinstance(cf.validation_data.get("val"), DAGEvalAdapter) + + def test_train_with_validation_runs(self): + cf = cellflow.model.CellFlowAnnbatch() + _prepare(cf, _adata()) + cf.prepare_validation_data(source=_adata(seed=3), name="val") + cf.prepare_model( + pooling="mean", condition_embedding_dim=8, time_encoder_dims=(8,), hidden_dims=(8,), decoder_dims=(8,) + ) + cf.train(num_iterations=2, valid_freq=2) # triggers a DAGEvalLoader validation pass + assert cf.solver is not None diff --git a/tests/solver/test_solver.py b/tests/solver/test_solver.py index 312336ed..5b3d75d2 100644 --- a/tests/solver/test_solver.py +++ b/tests/solver/test_solver.py @@ -204,6 +204,31 @@ def _make_otfm(condition_dropout_prob=0.0, guidance=None, condition_null="zero_e ) +def _make_genot(condition_dropout_prob=0.0, guidance=None, condition_null="zero_embedding"): + """Build a small GENOT solver for guidance tests.""" + vf = cellflow.networks.GENOTConditionalVelocityField( + output_dim=5, + max_combination_length=2, + condition_embedding_dim=12, + hidden_dims=(32, 32), + decoder_dims=(32, 32), + genot_source_dims=(32, 32), + condition_dropout_prob=condition_dropout_prob, + condition_null=condition_null, + ) + return _genot.GENOT( + vf=vf, + data_match_fn=match_linear, + probability_path=ConstantNoiseFlow(0.0), + optimizer=optax.adam(1e-3), + source_dim=5, + target_dim=5, + conditions={"drug": np.random.rand(2, 1, 3)}, + rng=vf_rng, + guidance=guidance, + ) + + class TestGuidance: def test_predict_guidance_none_matches_conditional_solve(self): """With ``guidance=None`` predict reproduces the pre-change conditional-only solve. @@ -261,7 +286,7 @@ def test_classifier_free_guidance_wraps_velocity(self): {"params": params}, t, x, condition, encoder_noise, train=False, force_uncond=True )[0] - guided_vf = solver.guidance.wrap(base_vf, solver.vf_state_inference) + guided_vf = solver.guidance.wrap(base_vf) v_guided = guided_vf(t, x, args) expected = v_null + w * (v_cond - v_null) @@ -282,7 +307,7 @@ def test_classifier_free_guidance_scale_one_is_conditional(self): base_vf = solver._base_velocity() v_cond = base_vf(t, x, args) - v_guided = solver.guidance.wrap(base_vf, solver.vf_state_inference)(t, x, args) + v_guided = solver.guidance.wrap(base_vf)(t, x, args) assert np.allclose(np.asarray(v_guided), np.asarray(v_cond), atol=1e-6) @@ -303,6 +328,28 @@ def test_from_ode_weight_parameterization(self): with pytest.raises(ValueError, match="cfg_ode_weight must be non-negative"): ClassifierFreeGuidance.from_ode_weight(-1.0) + def test_genot_classifier_free_guidance(self): + """GENOT supports CFG (shared, x_0-aware): guidance changes the prediction, ``scale=1.0`` + is the plain conditional solve, and it is gated on ``cfg_enabled``.""" + x = np.ones((4, 5), dtype=np.float32) + cond = {"drug": np.ones((1, 2, 3), dtype=np.float32)} + + g = _make_genot(condition_dropout_prob=0.5) + assert g.cfg_enabled + p_cond = g.predict(x, cond, guidance_scale=1.0, max_steps=10, throw=False) + p_guided = g.predict(x, cond, guidance_scale=2.0, max_steps=10, throw=False) + p_default = g.predict(x, cond, max_steps=10, throw=False) + assert np.all(np.isfinite(p_guided)) + assert not np.allclose(p_cond, p_guided) # guidance actually changes the field + assert np.allclose(p_cond, p_default) # scale=1.0 == no guidance + + # Without condition dropout, v_null is undefined -> guidance ignored (with warning). + g0 = _make_genot(condition_dropout_prob=0.0) + assert not g0.cfg_enabled + with pytest.warns(UserWarning, match="guidance_scale"): + p_ignored = g0.predict(x, cond, guidance_scale=3.0, max_steps=10, throw=False) + assert np.allclose(p_ignored, g0.predict(x, cond, max_steps=10, throw=False)) + @pytest.mark.parametrize("pooling", ["mean", "attention_token", "attention_seed"]) def test_mask_value_null_matches_mask_filled_condition(self, pooling): """With ``condition_null='mask_value'``, ``force_uncond`` equals evaluating the vf on a mask-filled condition.""" @@ -381,7 +428,7 @@ def test_theislab_parity_guidance_formula(self): {"params": params}, t, x, condition, encoder_noise, train=False, force_uncond=True )[0] - v_guided = solver.guidance.wrap(base_vf, solver.vf_state_inference)(t, x, args) + v_guided = solver.guidance.wrap(base_vf)(t, x, args) expected = (1.0 + w) * v_cond - w * v_null assert np.allclose(np.asarray(v_guided), np.asarray(expected), atol=1e-6) @@ -466,9 +513,7 @@ def test_per_call_guidance_scale_matches_construction_time(self): # Same vf_rng in _make_otfm -> identical init params -> the two paths are comparable. per_call = _make_otfm(condition_dropout_prob=0.5).predict(x, condition, guidance_scale=w) - construction = _make_otfm( - condition_dropout_prob=0.5, guidance=ClassifierFreeGuidance(w) - ).predict(x, condition) + construction = _make_otfm(condition_dropout_prob=0.5, guidance=ClassifierFreeGuidance(w)).predict(x, condition) assert np.allclose(per_call, construction, atol=1e-5) # And it is actually guiding: differs from the plain conditional (scale 1.0).