Feat/annbatch loader#297
Open
selmanozleyen wants to merge 30 commits into
Open
Conversation
CellFlow: - Make `adata` optional; passing it to the constructor is deprecated (FutureWarning) in favour of `prepare_data(adata=...)`. `adata is None` selects the annbatch path. - Add `prepare_annbatch_data` stub (signature + docstring + mode guard, no body yet). Shared condition logic (single source of truth in cellflow.data._condition): - `get_max_combination_length` extracted; DataManager delegates to it. - `enumerate_perturbations`: pandas target-combo enumeration (no dask/masks), parity-tested to match DataManager._get_condition_data exactly. dagloader: - SamplerConfig uses `kw_only=True` so the required `preload_nchunks` (no hidden default) imports cleanly; README examples pass it explicitly. Tests: tests/data/test_condition.py, tests/model/test_annbatch_path.py. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- cellflow.data._condition: add `build_condition_data`, which enumerates perturbations and reuses DataManager._get_embeddings over a plain serial loop to assemble condition_data (idx-aligned embeddings) with no dask/masks. Add shared `_key_layout` helper; refactor `enumerate_perturbations` onto it. - Parity-tested: build_condition_data output equals DataManager._get_condition_data(...).condition_data exactly (incl. the one-hot-encoder path, sample covariates, and combinations). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Split a Scheme's target combinations into named train/val/test splits — a
weights-only transform (weight = the selection), holding out whole condition
combinations (OOD split), mirroring CellFlow2's combination-level splitter.
Controls / bound children are carried through unchanged.
- dagloader._split: `split_scheme` (Scheme -> {split: Scheme}) and
`split_assignment` (inspection table); exported from the package.
- CellFlow.split_annbatch_data: class-level split step; also wired into
`prepare_annbatch_data` via split_* params. Building the Scheme from an
annbatch source remains a TODO (NotImplementedError on the normal path).
- Tests: pure split_scheme/split_assignment (no DatasetCollection, no loader)
and the CellFlow split step (via an injected Scheme).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- prepare_annbatch_data takes `sampler_config`: a single SamplerConfig applied
to every split, or a {split_name: SamplerConfig} mapping that must cover all
splits. Resolved via new `dagloader.resolve_split_configs` into
{split: SamplerConfig} (no-split case → a single "train"). Replaces the old
scalar batch_size/chunk_size params.
- SamplerConfig.chunk_size is now required (no hidden default of 1), matching
preload_nchunks; README examples and tests updated to pass it explicitly.
Building the Scheme/loader from an annbatch source remains a TODO; the split +
config-resolution steps are exercised via an injected Scheme (no collection,
no loader).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… train)
prepare_annbatch_data now builds the dagloader Scheme (perturbed root, matched-
control child) + a condition_fn from the covariate spec, one DAGLoader per split,
and wires the "train" split to train(); val/test loaders are kept on
_annbatch_loaders. Training runs end-to-end over the streamed batches.
- cellflow/data/_annbatch.py: build_annbatch_training(...) — reuses DataManager
(as a cell-free encoder factory over obs+uns, never X) and build_condition_data,
so condition embeddings are byte-identical to the in-memory path (parity-checked).
New `rep_dict` arg supplies the uns-equivalent embeddings (no adata in this path).
- cellflow/data/_dataloader.py: DAGLoaderTrainSampler — key-rename adapter mapping
the DAGLoader {target,source,condition} stream to the {tgt,src,condition}_cell_data
.sample(rng) contract, so both paths reach the solver identically.
- _cellflow: prepare_model / train branch to source condition data + dataloader from
the annbatch path; in-memory behavior unchanged.
- dagloader/_loader.py: fix DAGLoader.__next__ calling a non-existent _node_reps
(should be _nodes_next) — it crashed on every bound child, so no source batch
could be produced. One-line bug fix; no design change.
- Tests: build/split/per-split-config + end-to-end training over an in-memory
AnnData source (no DatasetCollection needed).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- In-memory AnnData sources are grouped (stable-sorted by the grouping columns) in build_annbatch_training, so chunk_size>1 streams contiguous slices out of the box (cheap one-time reorder; cell order is irrelevant to sampling). - Out-of-core DatasetCollections are NOT reordered (a physical zarr re-sort is expensive): document the assumption and validate it from obs. When any split uses chunk_size>1, assert_source_grouped checks each category is one contiguous run >= chunk_size and raises a clear error pointing at add_adatas(groupby=...) otherwise. - Tests over real DatasetCollections: grouped (add_adatas groupby) trains with chunk_size=4; interleaved raises the clear error; chunk_size=1 works ungrouped; split+chunked; and an in-memory interleaved source is auto-grouped. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The previous check required each category to be a single contiguous run, which is stricter than annbatch: its ClassSampler only requires every contiguous run of a (non-excluded) class to be >= chunk_size — a class MAY span multiple runs. So a "fragmented but valid" source (each category in several long runs) was wrongly rejected. Verified empirically: annbatch streams a 2-runs-of-10 layout at chunk_size=4, and raises only when a run is genuinely shorter than chunk_size. - Rename assert_source_grouped -> assert_source_chunkable; drop the one-run-per-class clause, keep only "every run >= chunk_size" (annbatch's actual rule), with a clear error pointing at add_adatas(groupby=...). Docs updated to match. - Tests: a fragmented DatasetCollection (>6 runs for 6 categories) now trains at chunk_size=4; unit tests assert the check accepts fragmented long runs, rejects a genuine short run, and is a no-op at chunk_size=1. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Ruff-clean the vendored dagloader package (all pre-existing, unrelated to the annbatch work): - _loader.py: drop the unused `densify` import (F401), sort imports (I001), `dict.fromkeys` for the shared-config map (C420), `zip(..., strict=True)` where the two sequences are equal-length by construction (B905), and split two summary/description docstrings (D205/D209). - _io.py: capitalize a docstring first word (D403). - __init__.py / _schema.py: raw-string docstrings that carry RST `\s` escapes (D301). No behavior change; dagloader imports and the full annbatch/dagloader test suite pass unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
dagloader now streams whatever the source stores (sparse or dense) and stays representation-agnostic — the unused `densify` helper is removed from `dagloader/_io.py`. Densification happens where the solver needs it: the `DAGLoaderTrainSampler` adapter materializes a sparse cell batch to dense (dense batches pass through untouched), so the in-memory and streaming paths reach the solver identically. Test: a sparse (csr) in-memory source streams dense `src`/`tgt` batches. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Validation is inherently in-memory (metrics need materialized cells), so it takes a held-out AnnData in the streaming path too. Two fixes make that correct: - build_annbatch_training now gives the encoder-factory DataManager the REAL `sample_rep` instead of a hardcoded "X". `_verify_sample_rep` is type-only (it does not check obsm existence) and cells are only read later against the actual validation adata, so the real value is safe on the cell-free shell. Previously validation would have read `adata.X` regardless of `sample_rep` — wrong for obsm reps. - prepare_validation_data's guard no longer requires the in-memory `train_data`; it accepts a model prepared via prepare_data OR prepare_annbatch_data. Docstring notes the streaming-path requirements (adata carries `sample_rep` + covariate reps in .uns) and that this is unrelated to the `val` split loader. Tests: validation prepares in the annbatch path; reads obsm `sample_rep` (dim 3) not X (dim 5); training runs with a validation pass; and still errors before any setup. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
`save` cloudpickles the whole model, which for the annbatch path includes the streaming loaders. Two things blocked that once iteration had started (i.e. after training — exactly when you save): - DAGLoader held live annbatch iterators (generators aren't picklable). Add DAGLoader.__getstate__/__setstate__ that drop `_iters` while keeping the per-node RNG streams, samplers, drawn schedules and configs — so a reloaded loader resumes the SAME reproducible sequence (a fresh pass drawn from the restored RNG state). - A suspended pass leaves the sampler's `_rng` as the transient `_ScheduleRng` wrapper, whose `__getattr__` recursed during unpickling. Give `_ScheduleRng` a `__reduce__` that pickles it AS the underlying real Generator (it is rebuilt each pass anyway) and guard `__getattr__` against delegating dunders. Also fixes `get_condition_embedding` in the annbatch path: it stored the embedding under `self.adata.uns`, but there is no `adata` in the streaming path (was an AttributeError by default). Now it warns and returns the embeddings when `adata` is absent. Tests: DAGLoader pickles mid-stream and two restores resume identically; CellFlow save/load after prepare, mid-stream (deterministic resume, state preserved), and after training; get_condition_embedding warns instead of crashing without adata. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Trim the verbose docstrings and a few long comments added for the annbatch path: drop repeated rationale, "mirrors X"/"byte-for-byte"/"parity-tested" asides, and parentheticals that restate the summary. Kept the actual contract — numpydoc Parameters/Returns, the chunk_size run-length rule, the sample_rep/validation note, and the "controls carried through" / "RNG preserved on pickle" facts. Docstrings/comments only — no behavior change; full annbatch + dagloader suite passes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Replace `Any` in the annbatch code with real types and make the additions mypy-clean: - _annbatch.py: `source: Container` (AnnData | DatasetCollection), a `Leaf` alias (`tuple[object, ...]`) for scheme leaves, and `rep_dict: Mapping[str, Mapping[str, ArrayLike]]`. - _dataloader.py: `DAGLoaderTrainSampler.__init__(loader: DAGLoader)`, `sample(rng: np.random.Generator | None) -> dict[str, np.ndarray | dict[str, np.ndarray]]`, `_densify(x: ArrayLike)` (getattr avoids a spurious no-`todense` error). - _cellflow.py: type the annbatch attributes (`_scheme: Scheme | None`, `_split_schemes`, `_annbatch_sampler_configs`, `_annbatch_loaders`), and use `Mapping[str, object]` for force-training values. `prepare_model` now narrows the condition-data / max-combination-length source with an if/elif/else (no `int | None`). - dagloader: parametrize `DAGLoader.__getstate__/__setstate__` dicts and annotate `_ScheduleRng.__reduce__`. mypy on the changed files: 0 new errors and 5 pre-existing ones fixed (down 22 -> 17; the rest are pre-existing debt in unrelated methods). No behavior change; annbatch suite + in-memory prepare_model smoke pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
uv.lock was added by an early branch commit but is not tracked on main; drop it from the tree and gitignore it so it stays out of the PR. Regenerable with `uv lock`. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #297 +/- ##
======================================
Coverage 0.00% 0.00%
======================================
Files 40 49 +9
Lines 2858 3464 +606
======================================
- Misses 2858 3464 +606
🚀 New features to boost your workflow:
|
for more information, see https://pre-commit.ci
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
So I am keeping this as a draft because I'd like to share the state of the annbatch adaptation. Currently there is a dagloader which is planned to be shared with sc-flow-tools. That code itself would need to change but it's a glue code to not wait on annbatch changes needed for this
TODO: need to check if these merges has broken something or not