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..cfe7f87d 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,20 @@ 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", + # `binded` — the declarative index-free loader, extracted from cellflow's former vendored + # `src/dagloader` package. Exposes `Loader` / `EvalLoader` / `Scheme` etc. Pinned by git URL until + # a PyPI release; direct URL so pip/uv/tox all resolve the branch. (Itself depends on the annbatch + # fork above — same URL, so it dedupes.) + "binded @ git+https://github.com/theislab/binded.git@feat/loader", + # 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 +103,7 @@ optional-dependencies.pp = [ "rdkit", ] optional-dependencies.test = [ + "cellflow-tools[annbatch]", "cellflow-tools[embedding]", "cellflow-tools[external]", "cellflow-tools[pp]", @@ -105,12 +120,19 @@ urls.Home-page = "https://github.com/theislab/cellflow" urls.Source = "https://github.com/theislab/cellflow" [tool.hatch.build.targets.wheel] +# `dagloader` was extracted into the standalone `binded` package (import path `binded`); cellflow now +# depends on it via the `annbatch` extra (git URL). Only `src/cellflow` ships in the wheel. packages = [ 'src/cellflow' ] [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" ] @@ -164,6 +186,7 @@ lint.pydocstyle.convention = "numpy" [tool.pytest.ini_options] testpaths = [ "tests" ] +pythonpath = [ "tests" ] # import shared test helpers (e.g. scheme_helpers) by bare module name xfail_strict = true addopts = [ "--import-mode=importlib", # allow using test files with same name @@ -172,6 +195,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..3105a862 --- /dev/null +++ b/src/cellflow/data/_annbatch.py @@ -0,0 +1,271 @@ +"""Build the annbatch/binded streaming training path from a CellFlow covariate spec. + +Turns the ``prepare_data`` covariate arguments into a :class:`binded.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 ``binded``. +""" + +from __future__ import annotations + +import os +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._logging import logger +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 binded import Container, Scheme + +Leaf = tuple[object, ...] # a scheme leaf: one value per grouping column + +__all__ = [ + "AnnbatchTraining", + "build_annbatch_training", + "sample_rep_to_key", +] + + +def sample_rep_to_key(sample_rep: str) -> str: + """CellFlow ``sample_rep`` → binded 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( + data: Container | str | os.PathLike | Sequence[str | os.PathLike | ad.AnnData], + *, + 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 = True, + min_cells_per_condition: int = 0, + chunk_size: int = 1, +) -> AnnbatchTraining: + """Assemble the :class:`binded.Scheme` + ``condition_fn`` for the streaming path (obs only). + + ``data`` is an out-of-core :class:`annbatch.DatasetCollection`, an in-memory ``AnnData``, an adata + zarr path, or a list of adata zarr paths (paths are resolved via :func:`~binded._io.open_source`). + ``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 binded to materialize the control (child) node into RAM (sets + :attr:`~binded.Node.in_memory`; binded owns the read via :func:`~binded._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). + + Two perturbed-only weight filters shape the root (target) node — controls are never filtered (with + ``control_in_memory`` the control node is materialized+sorted in RAM; otherwise its run-length is + annbatch's own concern, not ours): + + ``min_cells_per_condition`` zero-weights any perturbed condition with fewer than this many *total* + cells — a scientific filter on untrainable tiny conditions. Default ``0`` drops nothing. + + ``chunk_size`` (the streamed ``SamplerConfig.chunk_size``) drives the run-length filter. With + ``chunk_size > 1`` annbatch reads contiguous ``chunk_size``-long slices, so every run of a positive-weight + class must be ``>= chunk_size``. Any perturbed condition whose *smallest* contiguous run is shorter is + zero-weighted here (and thereby excluded from every split), so the rest stream chunked without annbatch + raising. This is the per-run guard a *total* filter can't provide — a big condition with a rare + sub-``chunk_size`` sliver in one plate is dropped. Dropped counts are logged. Default ``1`` filters + nothing; with both filters inactive (``min_cells_per_condition=0`` and ``chunk_size=1``) the root weights + are ``uniform`` — byte-identical to before. + """ + from binded import Bind, Node, Scheme, uniform + from binded._io import key_backings, obs_columns, open_source + + 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) + + # Accept a zarr path / list of adata zarr paths too: resolve to a Container (reads only `key` + the + # grouping obs). Path-backed data is out-of-core, so — like a DatasetCollection — it is never + # reordered here; only a user-supplied in-memory AnnData is stable-sorted below. + from_path = isinstance(data, str | os.PathLike | list | tuple) + if from_path: + data = open_source(data, keys=[key], cols=[*cols, control_key]) + + obs = obs_columns(data, [*cols, control_key]) + + # In-memory data: stable-sort by the grouping columns so `chunk_size > 1` reads contiguous slices + # (cheap, and cell order is irrelevant). Out-of-core data isn't reordered (expensive zarr re-sort) — + # it must be built grouped; the run-length filter below drops short-run perturbed conditions and + # annbatch validates the rest (and the controls) when it builds its samplers. + if isinstance(data, ad.AnnData) and not from_path: + order = obs[list(cols)].reset_index(drop=True).sort_values(list(cols), kind="stable").index.to_numpy() + data = data[order].copy() + obs = obs_columns(data, [*cols, control_key]) + + # The encoder and the scheme leaves depend only on the UNIQUE (grouping-cols, control) combinations — + # a few ×10^4 rows — not on the ~10^8 cells. So deduplicate ONCE here and drive the whole encoder + # (shell + DataManager + build_condition_data + enumerate_perturbations) and the pert/ctrl leaf lists + # off that tiny frame; feeding the full obs made every step O(n_cells) (a ~10-min prepare on Tahoe). + # Cast string grouping cols to `category` first so this single full-obs dedup hashes small integer + # codes, not raw strings. Parity-safe: `enumerate_perturbations` casts string cols the same way, so + # the leaf order is unchanged; only *object* (string) cols are cast, leaving numeric/bool covariates + # numeric (casting them would flip DataManager's numeric-vs-categorical detection) and preserving + # already-categorical cols' category order. + to_categorical = {c: "category" for c in cols if obs[c].dtype == object} + if to_categorical: + obs = obs.astype(to_categorical) + uniq = obs[[*cols, control_key]].drop_duplicates().reset_index(drop=True) + + # DataManager as a covariate-encoder factory: reads only obs + uns, so a cell-free, deduplicated shell + # suffices — it reads the unique category values (not per-cell counts), so the encoder is identical. + # `sample_rep` is stored for validation's `_get_cell_data` (verification is type-only, so it's safe here). + shell = ad.AnnData(obs=uniq.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( + uniq, + 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( + uniq, + 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). + # Built off the deduplicated frame — identical set of leaves to the full obs (order is irrelevant: + # `uniform` builds a dict and the loader resolves weights per string-sorted leaf). + ctrl_flag = uniq[control_key].to_numpy().astype(bool) + pert = [tuple(r) for r in uniq.loc[~ctrl_flag, list(cols)].drop_duplicates().to_numpy()] + ctrl = [tuple(r) for r in uniq.loc[ctrl_flag, list(cols)].drop_duplicates().to_numpy()] + + # Root (perturbed) leaf weights: `uniform(pert)` minus two perturbed-only filters (controls keep + # `uniform(ctrl)` — see the docstring). Both derive from ONE pass over the full `obs` (physical order) + # via `leaf_codes`: per-leaf total cells (bincount) and per-leaf smallest contiguous run (run-length + # min). Keyed by string-tuple so lookups survive `.to_numpy()` dtype quirks. When both filters are + # inactive (min_cells_per_condition=0 and chunk_size<=1) we skip the pass and use `uniform(pert)` — + # byte-identical to before. + if min_cells_per_condition > 0 or chunk_size > 1: + from binded._io import leaf_codes + + codes, leaves = leaf_codes(obs, list(cols)) + total = np.bincount(codes, minlength=len(leaves)) + run_starts = np.concatenate([[0], np.flatnonzero(np.diff(codes) != 0) + 1]) + run_len = np.diff(np.concatenate([run_starts, [len(codes)]])) + min_run = np.full(len(leaves), len(codes) + 1, dtype=np.int64) + np.minimum.at(min_run, codes[run_starts], run_len) # smallest run per leaf + stat = {tuple(map(str, lf)): (int(total[i]), int(min_run[i])) for i, lf in enumerate(leaves)} + + def _keep(leaf: Leaf) -> bool: # perturbed-only: total-cells filter AND per-run (chunk) filter + n_total, shortest = stat.get(tuple(map(str, leaf)), (0, 0)) + return n_total >= min_cells_per_condition and (chunk_size <= 1 or shortest >= chunk_size) + + pert_weights = {leaf: (1.0 if _keep(leaf) else 0.0) for leaf in pert} + n_kept = sum(w > 0 for w in pert_weights.values()) + if pert and n_kept == 0: + largest = max((stat.get(tuple(map(str, lf)), (0, 0))[0] for lf in pert), default=0) + longest = max((stat.get(tuple(map(str, lf)), (0, 0))[1] for lf in pert), default=0) + raise ValueError( + f"dropped every perturbed condition: none has >= min_cells_per_condition=" + f"{min_cells_per_condition} total cells (largest {largest}) and a contiguous run >= " + f"chunk_size={chunk_size} (longest run {longest}). Lower the thresholds, use chunk_size=1, or " + f"group the data so each condition forms runs >= chunk_size (e.g. `add_adatas(groupby=...)`)." + ) + if n_kept < len(pert): + dropped = sum(stat.get(tuple(map(str, lf)), (0, 0))[0] for lf, w in pert_weights.items() if w == 0) + allc = sum(stat.get(tuple(map(str, lf)), (0, 0))[0] for lf in pert) + logger.info( + f"annbatch streaming: dropped {len(pert) - n_kept}/{len(pert)} perturbed conditions " + f"({dropped:,}/{allc:,} cells = {dropped / max(allc, 1) * 100:.1f}%) — below " + f"min_cells_per_condition={min_cells_per_condition} or with a contiguous run < " + f"chunk_size={chunk_size}. Controls are unaffected." + ) + else: + pert_weights = uniform(pert) + + # `control_in_memory` just tells binded to materialize the control (child) node into RAM — the + # perturbed target keeps streaming out of core. binded owns the materialization (Node.in_memory); + # the bind still matches control↔target by the context columns. + scheme = Scheme( + sources={"data": data}, + nodes={ + "pert": Node("data", cols, key, pert_weights), + "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(data, 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..f9c52527 --- /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/binded 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..4dbf4352 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 binded import EvalLoader, Loader +__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. + ``binded`` 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 ``binded.Loader`` 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: "Loader"): + 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 ``Loader`` 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 ``binded.EvalLoader`` 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 ``EvalLoader`` 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:`binded.EvalLoader` 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: "EvalLoader", + *, + 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..c020a475 --- /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/``binded`` 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..a62c6050 --- /dev/null +++ b/src/cellflow/model/_cellflow_annbatch.py @@ -0,0 +1,457 @@ +"""Streaming (annbatch/binded) CellFlow: the default path. + +Trains, validates and predicts over an out-of-core :class:`~annbatch.DatasetCollection` or an +in-memory ``AnnData`` alike. +""" + +from __future__ import annotations + +import os +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 binded import EvalLoader, SamplerConfig, Scheme # optional deps — typing only + + # accepted `data` inputs: an in-memory ``AnnData`` / out-of-core ``DatasetCollection``, an adata zarr + # path, or a list of adata zarr paths (paths resolved via ``binded``'s ``open_source``). + DataInput = ad.AnnData | DatasetCollection | str | os.PathLike | Sequence[str | os.PathLike | ad.AnnData] + +__all__ = ["CellFlowAnnbatch"] + + +class CellFlowAnnbatch(BaseCellFlow): + """CellFlow over the annbatch/binded 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:`~binded.EvalLoader` (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 EvalLoader + 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, EvalLoader] = {} + self._config: Mapping[str, Any] | None = None # portable config captured by `from_config` + + @classmethod + def from_config(cls, config: Mapping[str, Any]) -> CellFlowAnnbatch: + """Config-first, data-free constructor — the single entry point. + + Selects the solver and captures the spec, but touches **no data**. The whole + ``prepare_data → prepare_model → train`` sequence collapses to + ``from_config → load_data(cells) → train_from_config``: load the config once, then attach the + annbatch cells after everything is set. + + Parameters + ---------- + config + A plain mapping (YAML/JSON-friendly, so the same config can drive a torch backend too — + only the constructed objects differ) with sections: + + - ``model`` — ``solver`` + :meth:`prepare_model` hyperparameters + (``condition_embedding_dim``, ``pooling``, ``hidden_dims``, ``decoder_dims``, + ``time_encoder_dims``, …; any further ``prepare_model`` kwargs under ``model.extra``). + - ``sampler`` — ``batch_size``, ``chunk_size``, ``prefetch_factor`` (→ ``preload_nchunks``), + ``min_cells_per_condition``, ``control_in_memory``. + - ``data`` — the covariate/rep **spec** only (``sample_rep``, ``control_key``, + ``perturbation_covariates``, ``split_covariates``, ``split_by``, ``split_ratios``, …). + The cells themselves are passed to :meth:`load_data`, not named here. + - ``trainer`` — ``num_iterations``, ``valid_freq`` (+ any :meth:`train` kwargs under + ``trainer.extra``). + - top-level ``seed``. + """ + model_cfg = dict(config.get("model", {})) + self = cls(solver=model_cfg.get("solver", "otfm")) + # shallow-copy nested blocks so later mutation of the caller's dict can't change setup + self._config = {k: (dict(v) if isinstance(v, Mapping) else v) for k, v in config.items()} + return self + + def load_data(self, data: DataInput) -> CellFlowAnnbatch: + """Attach the annbatch cells and finish setup (data + model) from the stored config. + + Run once after :meth:`from_config`. Builds the streaming ``Scheme`` from ``data``'s obs + (:meth:`prepare_data`) then the model (:meth:`prepare_model`), using the ``data`` / ``sampler`` + / ``model`` blocks captured by :meth:`from_config`. ``data`` is the cells only (an ``AnnData`` / + ``DatasetCollection`` / adata zarr path(s)); every other setting came from the config. + """ + if self._config is None: + raise ValueError("Call `from_config(...)` before `load_data(...)`.") + from binded import SamplerConfig + + cfg = self._config + d = dict(cfg.get("data", {})) + s = dict(cfg.get("sampler", {})) + seed = cfg.get("seed", 0) + + if "sample_rep" not in d: + raise ValueError("config['data']['sample_rep'] is required (the streamed representation, e.g. 'X_pca').") + batch_size = s.get("batch_size", 1024) + chunk_size = s.get("chunk_size", 1) + if chunk_size <= 0 or batch_size % chunk_size != 0: + raise ValueError(f"sampler.batch_size ({batch_size}) must be a positive multiple of chunk_size ({chunk_size}).") + preload = (batch_size // chunk_size) * s.get("prefetch_factor", 2) + + prep_kwargs: dict[str, Any] = { + "sample_rep": d["sample_rep"], + "control_key": d.get("control_key", "is_control"), + "perturbation_covariates": {k: list(v) for k, v in d.get("perturbation_covariates", {}).items()}, + "perturbation_covariate_reps": d.get("perturbation_covariate_reps") or None, + "sample_covariates": d.get("sample_covariates") or None, + "split_covariates": d.get("split_covariates") or None, + "sampler_config": SamplerConfig(batch_size=batch_size, chunk_size=chunk_size, preload_nchunks=preload), + "min_cells_per_condition": s.get("min_cells_per_condition", 0), + "split_by": d.get("split_by") or None, + "split_ratios": d.get("split_ratios"), + "seed": seed, + } + if s.get("control_in_memory") is not None: + prep_kwargs["control_in_memory"] = s["control_in_memory"] + self.prepare_data(data, **prep_kwargs) + + m = dict(cfg.get("model", {})) + m.pop("solver", None) # consumed at construction + m.pop("seed", None) # seed is top-level + extra = dict(m.pop("extra", {}) or {}) # pop BEFORE spreading m, else the "extra" key leaks into merged + merged = {**m, **extra} # extra may add/override prepare_model kwargs + for k in ("hidden_dims", "decoder_dims", "time_encoder_dims"): # flax needs hashable (tuple) static dims + if merged.get(k) is not None: + merged[k] = tuple(merged[k]) + merged.setdefault("seed", seed) + self.prepare_model(**merged) + return self + + def train_from_config(self, callbacks: Sequence[Any] = ()) -> Any: + """Train using the stored config's ``trainer`` / ``sampler`` blocks (run after :meth:`load_data`). + + ``callbacks`` (metrics + loggers) are supplied by the caller: they are app-level (e.g. cf-train + builds ``Metrics`` from ``trainer.metrics`` and a ``WandbLogger``) and intentionally not encoded + in the portable config. + """ + if self._config is None or self._dm is None: + raise ValueError("Call `from_config(...)` then `load_data(...)` before `train_from_config(...)`.") + t = dict(self._config.get("trainer", {})) + s = dict(self._config.get("sampler", {})) + return self.train( + num_iterations=t.get("num_iterations", 20000), + batch_size=s.get("batch_size", 1024), + valid_freq=t.get("valid_freq", 1000), + callbacks=list(callbacks), + **(t.get("extra", {}) or {}), + ) + + def prepare_data( + self, + data: DataInput, + 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 = True, + min_cells_per_condition: int = 0, + 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/binded 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 + ---------- + data + The cells to stream: an out-of-core :class:`annbatch.DatasetCollection`, an in-memory + ``AnnData`` (the ``binded`` is container-agnostic), an adata zarr path, or a list of adata + zarr paths. 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:`binded.SamplerConfig` for all splits, or a ``{split_name: SamplerConfig}`` mapping + covering every split (see :func:`binded.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 ``binded`` per-node RNG streams. + min_cells_per_condition + Drop (zero-weight) any perturbed condition with fewer than this many *total* cells — both a + scientific filter on untrainable tiny conditions and the lever that unblocks ``chunk_size > 1``: + a dropped (zero-weight) condition is exempt from annbatch's run-length rule, so its short run no + longer blocks chunked reads. Applied to the training source and to validation sources built from + the same spec. The default ``0`` drops nothing (behavior unchanged). Note this counts *total* + cells per condition, so it only unblocks ``chunk_size > 1`` when the kept conditions' per-plate + *runs* are also ``>= chunk_size``. + 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:`~binded.EvalLoader` (see :attr:`split_eval_loaders`), not streamed. + """ + from cellflow.data._annbatch import build_annbatch_training + from binded import Loader, resolve_split_configs + + # `sampler_config` is required; its (max) chunk_size drives the perturbed run-length filter in + # `build_annbatch_training` — short-run perturbed conditions are dropped so chunk_size>1 can stream. + if sampler_config is None: + raise ValueError("`sampler_config` is required: give a SamplerConfig or a {split: SamplerConfig} mapping.") + chunk_size = ( + sampler_config.chunk_size + if hasattr(sampler_config, "chunk_size") + else max(c.chunk_size for c in sampler_config.values()) + ) + + # Build the Scheme + condition_fn + condition embeddings from the covariate spec (obs only). + built = build_annbatch_training( + data, + 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, + min_cells_per_condition=min_cells_per_condition, + chunk_size=chunk_size, + ) + 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 EvalLoaders. + 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, + "min_cells_per_condition": min_cells_per_condition, + } + + # 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} + 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 + + # No separate cellflow-side run-length pre-check. The `chunk_size > 1` rule is handled upstream: + # `build_annbatch_training` zero-weights short-run perturbed conditions (and sub-threshold + # `min_cells_per_condition` ones), an in-memory control node samples at chunk_size=1 (see Loader), + # and annbatch validates whatever remains — the streamed perturbed layout and the controls — when it + # builds each node's sampler below. + + # Only the "train" split is streamed (feeds `train()`); val/test are read via EvalLoader + # below, so we don't build unused per-split streaming loaders. + self._dataloader = DAGTrainAdapter( + Loader(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 + # `EvalLoader` 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 binded import EvalLoader + + 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] = EvalLoader(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:`binded.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 binded 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 + # `EvalLoader`s (`split_eval_loaders`); here we only produce the schemes + table. + return split_assignment(self._split_schemes) + + def prepare_validation_data( + self, + data: DataInput, + 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:`~binded.EvalLoader` (no boolean masking). + + The ``data`` (an :class:`~anndata.AnnData`, an :class:`~annbatch.DatasetCollection`, an adata zarr + path, or a list of adata zarr paths) 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 + ---------- + data + The held-out validation cells (out-of-core, in-memory, or an adata zarr path / list of paths). + 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 binded import EvalLoader + + built = build_annbatch_training(data, **self._prep_kwargs) + eval_loader = EvalLoader(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, EvalLoader]: + """Per-split :class:`~binded.EvalLoader` 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..0225d4b7 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,31 @@ 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. + + Returns 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 +339,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 +440,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 +464,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 +687,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 +720,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 +746,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 +756,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 +784,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..6d4590a9 100644 --- a/src/cellflow/solvers/_otfm.py +++ b/src/cellflow/solvers/_otfm.py @@ -1,5 +1,6 @@ import warnings from collections.abc import Callable +from functools import partial from typing import Any, Protocol, runtime_checkable import diffrax @@ -28,12 +29,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 +84,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 @@ -160,7 +160,11 @@ def __init__( self.vf_step_fn = self._get_vf_step_fn() def _get_vf_step_fn(self) -> Callable: # type: ignore[type-arg] - @jax.jit + # Donate ``vf_state`` (params + optimizer moments + MultiSteps accumulator — the + # largest live buffers) so XLA updates them in place instead of allocating a fresh + # copy every step. Safe: the caller immediately rebinds ``self.vf_state`` to the + # returned state and never reuses the donated input. + @partial(jax.jit, donate_argnums=(1,)) def vf_step_fn( rng: jax.Array, vf_state: train_state.TrainState, @@ -282,9 +286,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 +355,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 +536,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..6e30f576 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 @@ -124,13 +124,30 @@ def train( sampler = dataloader if isinstance(dataloader, OOCTrainSampler): dataloader.set_sampler(num_iterations=num_iterations) + + # Keep per-step losses on-device and materialize them in bulk. Calling ``float(loss)`` + # every iteration forces a device->host sync that serializes host-side batch sampling + # behind GPU compute (GPU bubbles / underutilization). Buffering + a periodic flush + # removes ~all per-step syncs while bounding the async look-ahead so device memory + # cannot run away, and keeps ``training_logs["loss"]`` as plain floats (unchanged API). + loss_buf: list[Any] = [] + flush_every = min(valid_freq, 50) if valid_freq > 0 else 50 + + def _flush_losses() -> None: + if loss_buf: + self.training_logs["loss"].extend(np.asarray(jax.device_get(loss_buf)).ravel().tolist()) + loss_buf.clear() + for it in pbar: rng_jax, rng_step_fn = jax.random.split(rng_jax, 2) batch = sampler.sample(rng_np) loss = self.solver.step_fn(rng_step_fn, batch) - self.training_logs["loss"].append(float(loss)) + loss_buf.append(loss) # on-device; no host sync on the hot path + if len(loss_buf) >= flush_every: + _flush_losses() if ((it - 1) % valid_freq == 0) and (it > 1): + _flush_losses() # Get predictions from validation data valid_source_data, valid_true_data, valid_pred_data = self._validation_step( valid_loaders, mode="on_log_iteration" @@ -146,6 +163,7 @@ def train( postfix_dict["loss"] = round(mean_loss, 3) pbar.set_postfix(postfix_dict) + _flush_losses() if num_iterations > 0: valid_source_data, valid_true_data, valid_pred_data = self._validation_step( valid_loaders, mode="on_train_end" 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/data/test_cfsampler.py b/tests/data/test_cfsampler.py index 0ae5405a..ae114b4f 100644 --- a/tests/data/test_cfsampler.py +++ b/tests/data/test_cfsampler.py @@ -1,8 +1,8 @@ import numpy as np import pytest -from cellflow.data._dataloader import OOCTrainSampler, PredictionSampler, TrainSampler from cellflow.data._datamanager import DataManager +from cellflow.data._legacy import OOCTrainSampler, PredictionSampler, TrainSampler class TestTrainSampler: @@ -86,8 +86,8 @@ 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._datamanager import DataManager + from cellflow.data._legacy import ValidationSampler control_key = "control" sample_covariates = ["cell_type"] 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_min_cells.py b/tests/model/test_annbatch_min_cells.py new file mode 100644 index 00000000..ff11de81 --- /dev/null +++ b/tests/model/test_annbatch_min_cells.py @@ -0,0 +1,121 @@ +"""Tests for the perturbed-only weight filters on the annbatch streaming path. + +``min_cells_per_condition`` zero-weights perturbed conditions with too few *total* cells (a scientific +filter on untrainable tiny conditions). Independently, ``chunk_size > 1`` auto-drops perturbed conditions +whose smallest contiguous run is shorter than ``chunk_size`` (annbatch's run-length rule; controls are +exempt). Built from an in-memory ``AnnData`` source — which the build sorts, so there each condition is a +single run (total == smallest run); the per-run-vs-per-total distinction on a fragmented out-of-core source +is covered in ``test_annbatch_ooc.py``. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("annbatch") + +import anndata as ad + +import cellflow +from binded import SamplerConfig + +# a couple of tiny (sub-threshold) conditions alongside normal ones; the filter drops only the tiny ones. +# four "big" drugs so a default 60/20/20 split over the kept conditions still fills every split. +_COUNTS = {"control": 6, "big1": 4, "big2": 4, "big3": 4, "big4": 4, "tiny1": 1, "tiny2": 2} +_BIG = {"big1", "big2", "big3", "big4"} +_THRESHOLD = 3 # keeps control/big* (>= 3 cells), drops tiny1 (1) and tiny2 (2) + +_CFG_CHUNK1 = SamplerConfig(batch_size=4, chunk_size=1, preload_nchunks=4) +_CFG_CHUNK2 = SamplerConfig(batch_size=4, chunk_size=2, preload_nchunks=2) + + +def _toy_adata(counts: dict[str, int] = _COUNTS, line: str = "A") -> ad.AnnData: + """One cell line, `counts[drug]` cells per drug; rows shuffled so the build must sort them.""" + rng = np.random.default_rng(0) + rows = [(line, drug) for drug, n in counts.items() for _ in range(n)] + rows = [rows[i] for i in rng.permutation(len(rows))] # shuffled → build sorts into contiguous runs + 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(cf, *, sampler_config, **kwargs): + return cf.prepare_data( + data=_toy_adata(), + sample_rep="X", + control_key="control", + perturbation_covariates={"drug": ["drug"]}, + split_covariates=["cell_line"], + sampler_config=sampler_config, + **kwargs, + ) + + +def _positive_drugs(cf) -> set[str]: + """The drugs of the perturbed root node's positive-weight leaves (cols = (cell_line, drug)).""" + return {leaf[1] for leaf, w in cf._scheme.nodes["pert"].weights.items() if w > 0} + + +class TestMinCellsPerCondition: + def test_filter_drops_tiny_conditions_from_positive_leaves(self): + # (a) with the filter, the tiny conditions are absent from the root's positive-weight leaves. + cf = cellflow.model.CellFlowAnnbatch() + _prepare(cf, sampler_config=_CFG_CHUNK1, min_cells_per_condition=_THRESHOLD) + assert _positive_drugs(cf) == _BIG # tiny1/tiny2 dropped + # dropped leaves are retained as zero-weight (exemption), not deleted, so the source is unchanged. + all_drugs = {leaf[1] for leaf in cf._scheme.nodes["pert"].weights} + assert {"tiny1", "tiny2"} <= all_drugs + + def test_default_leaves_scheme_unchanged(self): + # (c) default (0) drops nothing → every perturbed leaf keeps weight 1.0 (== uniform). + cf = cellflow.model.CellFlowAnnbatch() + _prepare(cf, sampler_config=_CFG_CHUNK1) # min_cells_per_condition defaults to 0 + weights = cf._scheme.nodes["pert"].weights + assert _positive_drugs(cf) == _BIG | {"tiny1", "tiny2"} + assert set(weights.values()) == {1.0} # uniform: nothing zero-weighted + + def test_default_matches_explicit_zero(self): + # default and an explicit 0 produce identical root weights (byte-identical default path). + cf0, cf_explicit = cellflow.model.CellFlowAnnbatch(), cellflow.model.CellFlowAnnbatch() + _prepare(cf0, sampler_config=_CFG_CHUNK1) + _prepare(cf_explicit, sampler_config=_CFG_CHUNK1, min_cells_per_condition=0) + assert cf0._scheme.nodes["pert"].weights == cf_explicit._scheme.nodes["pert"].weights + + def test_chunk_gt_1_auto_drops_short_runs(self): + # chunk_size>1 auto-drops perturbed conditions whose smallest run < chunk_size — no filter needed. + # Sorted in-memory ⇒ one run == total, so tiny1 (1 cell, run 1 < 2) drops; tiny2 (2, run 2) is kept. + cf = cellflow.model.CellFlowAnnbatch() + _prepare(cf, sampler_config=_CFG_CHUNK2) # no min_cells_per_condition + assert _positive_drugs(cf) == _BIG | {"tiny2"} # only tiny1 (run 1 < 2) dropped + assert cf._dataloader.sample()["tgt_cell_data"].shape == (4, 5) + + def test_chunk_gt_1_ok_with_filter(self): + # (b) filtering the tiny (short-run) conditions unblocks chunk_size > 1 — no raise, and the loader + # streams: annbatch's ClassSampler exempts the zero-weight tiny leaves at iteration time too. + cf = cellflow.model.CellFlowAnnbatch() + _prepare(cf, sampler_config=_CFG_CHUNK2, min_cells_per_condition=_THRESHOLD) + assert _positive_drugs(cf) == _BIG + batch = cf._dataloader.sample() # end-to-end chunk_size=2 read + assert batch["tgt_cell_data"].shape == (4, 5) + + def test_filter_flows_through_split(self): + # zero-weighted conditions never reach the split universe (split_scheme splits positive weights). + cf = cellflow.model.CellFlowAnnbatch() + _prepare( + cf, + sampler_config=_CFG_CHUNK1, + min_cells_per_condition=_THRESHOLD, + split_by=["drug"], + split_random_state=0, + ) + assert set(cf._split_assignment["drug"]) == _BIG # no tiny drugs in any split + + def test_threshold_dropping_everything_raises(self): + # a threshold above every condition's count zero-weights the whole root → a clear error. + cf = cellflow.model.CellFlowAnnbatch() + with pytest.raises(ValueError, match="dropped every perturbed condition"): + _prepare(cf, sampler_config=_CFG_CHUNK1, min_cells_per_condition=1000) diff --git a/tests/model/test_annbatch_ooc.py b/tests/model/test_annbatch_ooc.py new file mode 100644 index 00000000..7cbc4905 --- /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 binded 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( + data=_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): + # annbatch enforces the run-length rule itself when building the sampler (no cellflow pre-check). + cf = cellflow.model.CellFlowAnnbatch() + with pytest.raises(ValueError, match="chunk_size|[Rr]e-chunk|run"): + cf.prepare_data( + data=_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( + data=_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( + data=_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 EvalLoader + + 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( + data=_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 binded._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 + data=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( + data=_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 # Loader built (add_datasets accepted the sparse-group backing) + + def test_control_in_memory_materializes(self, tmp_path): + # control_in_memory tells binded 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( + data=_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) # binded materialized it + + def test_chunk_drops_short_run_perturbed_only(self, tmp_path): + # A perturbed condition with a big TOTAL but a sub-chunk sliver run is dropped (per-RUN, not + # per-total); controls are never filtered. d2 has runs [10, 2] (total 12) → its run of 2 < chunk 4 + # drops it; d1 has runs [10, 10] → kept; control [10, 10] → kept and streamed. + def blk(dr, n): + return [("A", dr)] * n + + rows = blk("control", 10) + blk("d1", 10) + blk("d2", 10) + blk("control", 10) + blk("d1", 10) + blk("d2", 2) + 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") + ad.AnnData(X=x, obs=obs).write_h5ad(tmp_path / "s.h5ad") + dc = DatasetCollection(str(tmp_path / "sc.zarr"), mode="a").add_adatas( + [str(tmp_path / "s.h5ad")], shuffle=False + ) + + cf = cellflow.model.CellFlowAnnbatch() + cf.prepare_data(data=dc, sampler_config=SamplerConfig(batch_size=8, chunk_size=4, preload_nchunks=8), **_PREP) + pos = {leaf[1] for leaf, w in cf._scheme.nodes["pert"].weights.items() if w > 0} + assert pos == {"d1"} # d2 dropped for its run of 2 < chunk 4, despite total 12 + assert cf._dataloader.sample()["tgt_cell_data"].shape == (8, 5) diff --git a/tests/model/test_annbatch_ooc_memory.py b/tests/model/test_annbatch_ooc_memory.py new file mode 100644 index 00000000..36817962 --- /dev/null +++ b/tests/model/test_annbatch_ooc_memory.py @@ -0,0 +1,135 @@ +"""The core cluster invariant: the perturbed target streams OUT-OF-CORE while controls live in RAM. + +On Tahoe-scale data the perturbed population is the ~10^8-cell bulk that must never be materialized; +the matched controls are the small, re-drawn-every-batch population that belongs in memory. These tests +pin that split down explicitly — node identity, whether each node's backings are on-disk vs in-RAM, the +``control_in_memory`` toggle, and that the OOC/in-RAM split still matches control↔perturbed correctly — +so a regression that silently pulls the perturbed cells into RAM (an OOM on the cluster) is caught here. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("annbatch") + +import anndata as ad +from annbatch import DatasetCollection + +import cellflow +from binded import SamplerConfig +from binded._io import key_backings +from binded._loader import _is_backed + +_PREP = { + "sample_rep": "X", + "control_key": "control", + "perturbation_covariates": {"drug": ["drug"]}, + "split_covariates": ["cell_line"], +} +_LINE_CODE = {"A": 0.0, "B": 1.0} + + +def _adata(n_per_combo=30, drugs=("control", "d1", "d2"), lines=("A", "B")): + """X encodes each cell's identity: col0 = is_control (1/0), col1 = cell-line code; rest random. + + That lets a streamed batch be checked back to its provenance (perturbed vs control, which line) + without any obs — the loader streams X only. + """ + 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") + x[:, 0] = (obs["drug"] == "control").to_numpy().astype("float32") + x[:, 1] = obs["cell_line"].map(_LINE_CODE).to_numpy().astype("float32") + return ad.AnnData(X=x, obs=obs) + + +def _collection(tmp_path) -> DatasetCollection: + """A grouped out-of-core collection (contiguous category runs, as the cluster data is built).""" + h5 = tmp_path / "a.h5ad" + _adata().write_h5ad(h5) + dc = DatasetCollection(str(tmp_path / "c.zarr"), mode="a") + dc.add_adatas([str(h5)], groupby=["cell_line", "drug"], shuffle=False) + return dc + + +def _prepare(tmp_path, *, control_in_memory=True, **extra): + cf = cellflow.model.CellFlowAnnbatch() + cf.prepare_data( + data=_collection(tmp_path), + sampler_config=SamplerConfig(batch_size=8, chunk_size=1, preload_nchunks=8), + control_in_memory=control_in_memory, + **_PREP, + **extra, + ) + return cf + + +class TestOutOfCoreMemorySplit: + def test_pert_out_of_core_ctrl_in_memory(self, tmp_path): + # The default: perturbed target streams from disk, controls are materialized into RAM. + dc = _collection(tmp_path) + cf = cellflow.model.CellFlowAnnbatch() + cf.prepare_data( + data=dc, + sampler_config=SamplerConfig(batch_size=8, chunk_size=1, preload_nchunks=8), + control_in_memory=True, + **_PREP, + ) + ldr = cf._dataloader._loader + pert, ctrl = ldr._nodes["pert"], ldr._nodes["ctrl"] + + # pert node: still the SAME out-of-core collection object (no copy / no materialization), and + # every backing it streams is an on-disk backing — the whole point at 10^8 cells. + assert cf._scheme.nodes["pert"].in_memory is False + assert pert is dc, "perturbed source was copied/materialized — must stream the original collection" + assert isinstance(pert, DatasetCollection) + assert all(_is_backed(b) for b in key_backings(pert, "X")) + + # ctrl node: materialized into an in-memory AnnData, holding ONLY the control cells (the small + # re-drawn population), with in-RAM (not backed) arrays. + assert cf._scheme.nodes["ctrl"].in_memory is True + assert isinstance(ctrl, ad.AnnData) + assert not any(_is_backed(b) for b in key_backings(ctrl, "X")) + assert ctrl.n_obs == 60, "in-memory ctrl must hold exactly the controls (A/control 30 + B/control 30)" + # every materialized cell really is a control (col0 == 1) — no perturbed cell leaked into RAM. + assert np.all(np.asarray(ctrl.X)[:, 0] == 1.0) + + def test_control_in_memory_false_streams_both(self, tmp_path): + # Toggle off ⇒ controls stream out-of-core too (both nodes on-disk); nothing materialized. + cf = _prepare(tmp_path, control_in_memory=False) + ldr = cf._dataloader._loader + assert cf._scheme.nodes["ctrl"].in_memory is False + assert isinstance(ldr._nodes["ctrl"], DatasetCollection) + assert all(_is_backed(b) for b in key_backings(ldr._nodes["ctrl"], "X")) + assert all(_is_backed(b) for b in key_backings(ldr._nodes["pert"], "X")) + + def test_streamed_target_perturbed_inmem_source_control_matched(self, tmp_path): + # Correctness of the split: across many batches the OOC target is always perturbed, the in-RAM + # source is always a control, and both are the SAME cell line (the bind matches on cell_line). + cf = _prepare(tmp_path) + for _ in range(20): + b = cf._dataloader.sample() + tgt, src = np.asarray(b["tgt_cell_data"]), np.asarray(b["src_cell_data"]) + assert np.all(tgt[:, 0] == 0.0), "streamed target contained a control cell" + assert np.all(src[:, 0] == 1.0), "in-memory source contained a perturbed cell" + # class-coherent batch: one cell line per batch, and target/source share it (matched control). + assert len(np.unique(tgt[:, 1])) == 1 and len(np.unique(src[:, 1])) == 1 + assert tgt[0, 1] == src[0, 1], "control drawn from a different cell line than the target" + + def test_train_out_of_core_pert_in_memory_ctrl(self, tmp_path): + # End-to-end: the realistic cluster config trains, and the memory split still holds afterwards. + cf = _prepare(tmp_path) + 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) + assert cf.solver is not None + ldr = cf._dataloader._loader + assert isinstance(ldr._nodes["pert"], DatasetCollection) # target never pulled into RAM + assert isinstance(ldr._nodes["ctrl"], ad.AnnData) # controls stayed in RAM 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_prepare_parity.py b/tests/model/test_annbatch_prepare_parity.py new file mode 100644 index 00000000..c76d2179 --- /dev/null +++ b/tests/model/test_annbatch_prepare_parity.py @@ -0,0 +1,165 @@ +"""Parity guard for the optimized ``build_annbatch_training`` (streaming-path prepare). + +The prepare step now deduplicates the grouping columns ONCE and drives the whole covariate encoder off +that tiny frame (instead of the full ~10^8-row obs). This test pins the outputs to the canonical +in-memory :class:`~cellflow.data._datamanager.DataManager` path — the ground truth the streaming path +has always had to match — plus a naive full-obs computation of the scheme's leaves. Any divergence in +``condition_data`` / scheme leaves / ``data_dim`` / ``condition_fn`` wiring is a regression. +""" + +from __future__ import annotations + +import anndata as ad +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("annbatch") + +from cellflow.data._annbatch import build_annbatch_training +from cellflow.data._condition import _key_layout +from cellflow.data._datamanager import DataManager + +# (perturbation_covariates, perturbation_covariate_reps, split_covariates, sample_covariates, sample_reps) +SPECS = [ + ({"drug": ["drug1"]}, {"drug": "drug"}, ["cell_type"], [], {}), # embedding primary + split + ({"drug": ["drug1"]}, {}, ["cell_type"], [], {}), # one-hot primary + split + ({"drug": ["drug1", "drug2"]}, {"drug": "drug"}, ["cell_type"], [], {}), # combination length 2 + ({"drug": ["drug1"]}, {"drug": "drug"}, [], ["cell_type"], {"cell_type": "cell_type"}), # sample-covar path + ({"drug": ["drug1"], "dosage": ["dosage_a"]}, {"drug": "drug"}, ["cell_type"], [], {}), # numeric linked covar + ({"drug": ["drug1"]}, {}, [], [], {}), # bare one-hot, no split/sample +] + + +def _grouping_cols(pert_covars, split_covars, samp_covars): + """The build's deduped/ordered grouping columns (mirrors ``build_annbatch_training``).""" + pert_cols = tuple(c for grp in pert_covars.values() for c in grp) + return tuple(dict.fromkeys((*tuple(split_covars), *pert_cols, *tuple(samp_covars)))) + + +def _build(adata, pert_covars, pert_reps, split_covars, samp_covars, samp_reps): + return build_annbatch_training( + data=adata, + sample_rep="X", + control_key="control", + perturbation_covariates=pert_covars, + perturbation_covariate_reps=pert_reps or None, + split_covariates=split_covars, + sample_covariates=samp_covars, + sample_covariate_reps=samp_reps or None, + rep_dict=adata.uns, + ) + + +def _reference_dm(adata, pert_covars, pert_reps, split_covars, samp_covars, samp_reps): + dm = DataManager( + adata, + sample_rep="X", + control_key="control", + perturbation_covariates=pert_covars, + perturbation_covariate_reps=pert_reps or None, + split_covariates=split_covars, + sample_covariates=samp_covars, + sample_covariate_reps=samp_reps or None, + ) + return dm, dm._get_condition_data(adata=adata) + + +@pytest.mark.parametrize(("pert_covars", "pert_reps", "split_covars", "samp_covars", "samp_reps"), SPECS) +class TestBuildAnnbatchParity: + def test_condition_data_matches_in_memory( + self, adata_perturbation: ad.AnnData, pert_covars, pert_reps, split_covars, samp_covars, samp_reps + ): + built = _build(adata_perturbation, pert_covars, pert_reps, split_covars, samp_covars, samp_reps) + dm, ref = _reference_dm(adata_perturbation, pert_covars, pert_reps, split_covars, samp_covars, samp_reps) + + assert set(built.condition_data) == set(ref.condition_data) + for key in ref.condition_data: + np.testing.assert_array_equal(built.condition_data[key], ref.condition_data[key]) + assert built.max_combination_length == dm.max_combination_length + assert built.data_dim == adata_perturbation.n_vars + + def test_scheme_leaves_match_naive_full_obs( + self, adata_perturbation: ad.AnnData, pert_covars, pert_reps, split_covars, samp_covars, samp_reps + ): + built = _build(adata_perturbation, pert_covars, pert_reps, split_covars, samp_covars, samp_reps) + cols = _grouping_cols(pert_covars, split_covars, samp_covars) + obs = adata_perturbation.obs + ctrl = obs["control"].to_numpy().astype(bool) + + def _leaves(mask): # naive: dedup the FULL obs (the pre-optimization computation) + return {tuple(map(str, r)) for r in obs.loc[mask, list(cols)].drop_duplicates().to_numpy()} + + pert_got = {tuple(map(str, k)) for k in built.scheme.nodes["pert"].weights} + ctrl_got = {tuple(map(str, k)) for k in built.scheme.nodes["ctrl"].weights} + assert pert_got == _leaves(~ctrl) + assert ctrl_got == _leaves(ctrl) + + def test_condition_fn_maps_each_leaf_to_its_condition( + self, adata_perturbation: ad.AnnData, pert_covars, pert_reps, split_covars, samp_covars, samp_reps + ): + built = _build(adata_perturbation, pert_covars, pert_reps, split_covars, samp_covars, samp_reps) + dm, ref = _reference_dm(adata_perturbation, pert_covars, pert_reps, split_covars, samp_covars, samp_reps) + cols = _grouping_cols(pert_covars, split_covars, samp_covars) + # canonical tuple layout + reprojection cols->tuple_keys (same as DataManager / build) + _, tuple_keys = _key_layout(dm._perturb_covar_keys, list(split_covars), list(samp_covars)) + reorder = [cols.index(c) for c in tuple_keys] + cov_to_idx = {tuple(map(str, v)): k for k, v in ref.perturbation_idx_to_covariates.items()} + + for leaf in built.scheme.nodes["pert"].weights: # every perturbed leaf + idx = cov_to_idx[tuple(str(leaf[i]) for i in reorder)] + emitted = built.condition_fn(leaf) + assert set(emitted) == set(ref.condition_data) + for group in ref.condition_data: + np.testing.assert_array_equal(emitted[group], ref.condition_data[group][[idx]]) + + +def _tahoe_shaped_adata(n_lines=5, n_drugs=6, per_combo=8, seed=0) -> ad.AnnData: + """Tahoe-like obs: cell_line[category] / drug[object-str] / is_control[bool] (exercises the cast).""" + rng = np.random.default_rng(seed) + lines = [f"CL{i}" for i in range(n_lines)] + drugs = ["control"] + [f"drug{i}" for i in range(n_drugs - 1)] + rows = [(cl, dr) for cl in lines for dr in drugs for _ in range(per_combo)] + rows = [rows[i] for i in rng.permutation(len(rows))] # shuffled → build must sort in-memory + obs = pd.DataFrame(rows, columns=["cell_line", "drug"]) + obs["cell_line"] = obs["cell_line"].astype("category") + obs["drug"] = obs["drug"].astype(object) # object/str — NOT pre-categorical + obs["control"] = (obs["drug"] == "control").to_numpy() + obs.index = obs.index.astype(str) + X = rng.normal(size=(len(obs), 4)).astype("float32") + return ad.AnnData(X=X, obs=obs) + + +@pytest.mark.parametrize("pert_reps", [{"drug": "drug_emb"}, {}]) +def test_tahoe_shaped_object_drug_parity(pert_reps): + adata = _tahoe_shaped_adata() + if pert_reps: + drugs = list(pd.unique(adata.obs["drug"])) + adata.uns["drug_emb"] = {d: np.random.default_rng(1).normal(size=6).astype("float32") for d in drugs} + built = build_annbatch_training( + data=adata, + sample_rep="X", + control_key="control", + perturbation_covariates={"drug": ["drug"]}, + perturbation_covariate_reps=pert_reps or None, + split_covariates=["cell_line"], + rep_dict=adata.uns, + ) + dm = DataManager( + adata, + sample_rep="X", + control_key="control", + perturbation_covariates={"drug": ["drug"]}, + perturbation_covariate_reps=pert_reps or None, + split_covariates=["cell_line"], + ) + ref = dm._get_condition_data(adata=adata) + assert set(built.condition_data) == set(ref.condition_data) + for key in ref.condition_data: + np.testing.assert_array_equal(built.condition_data[key], ref.condition_data[key]) + # scheme leaves == naive full-obs dedup + obs = adata.obs + ctrl = obs["control"].to_numpy().astype(bool) + pert_got = {tuple(map(str, k)) for k in built.scheme.nodes["pert"].weights} + pert_ref = {tuple(map(str, r)) for r in obs.loc[~ctrl, ["cell_line", "drug"]].drop_duplicates().to_numpy()} + assert pert_got == pert_ref diff --git a/tests/model/test_annbatch_save.py b/tests/model/test_annbatch_save.py new file mode 100644 index 00000000..690e8999 --- /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 binded 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(data=_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..2c6c87af --- /dev/null +++ b/tests/model/test_annbatch_split.py @@ -0,0 +1,160 @@ +"""Tests for the annbatch path split wiring on :class:`~cellflow.model.CellFlow`. + +The Scheme + condition + loaders are built from an in-memory ``AnnData`` passed as ``data`` (the +``binded`` 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") # binded (and thus the annbatch path) needs annbatch + +import anndata as ad +from scheme_helpers import perturbation_scheme + +import cellflow +from binded import SamplerConfig + +_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( + data=_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( + data=_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( + data=_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( + data=_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): + # binded 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( + data=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..b83871d1 --- /dev/null +++ b/tests/model/test_annbatch_train.py @@ -0,0 +1,72 @@ +"""End-to-end training over the annbatch/binded 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 binded 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( + data=_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( + data=_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 EvalLoader (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..9ee713e3 --- /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 ``EvalLoader``. + +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 binded 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( + data=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(data=_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(data=_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(data=_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 EvalLoaders; "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(data=_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 EvalLoader validation pass + assert cf.solver is not None diff --git a/tests/scheme_helpers.py b/tests/scheme_helpers.py new file mode 100644 index 00000000..c154dfa8 --- /dev/null +++ b/tests/scheme_helpers.py @@ -0,0 +1,57 @@ +"""Test-only Scheme factory (moved out of ``binded``'s public API). + +``perturbation_scheme`` builds the cellflow-shaped two-node :class:`~binded.Scheme` (control → +perturbed, matched on context) from an obs table. Production cellflow assembles this scheme internally +(:func:`cellflow.data._annbatch.build_annbatch_training`); only the split tests need the standalone +factory, so it lives here rather than shipping as library surface. Importable as ``scheme_helpers`` via +the ``pythonpath = ["tests"]`` pytest setting. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence + +from binded import Bind, Node, Scheme, uniform +from binded._io import obs_columns +from binded._schema import Container + + +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/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).