Integration with `illico`

[ilan-gold]

What kind of feature would you like to request?

Additional function parameters / changed functionality / changed defaults?

Please describe your wishes

See https://github.com/remydubois/illico/tree/main

The speed claims are awesome so I have a few questions to investigate, but would like to move forward here:

1. Are the results the same? This is not clear immediately from the docs but there do appear to be tests
2. Is what is implemented currently dask-ifiable like RSC or would it require refactoring or even "forking" into our codebase?
3. Are the memory requirements the same? (Presumably but again not documented)

Aside from this, I think a general review of the codebase might make sense. Adding new external packages has brought hesitancy so I think making sure we'd be comfortable bringing this into the scanpy codebase should illico ever become unmaintained makes sense (we don't want a repeat of external). My gut feeling is "yes" given the high degree of overlap with our existing codebase and the matching results (apparently, from the tests), but I'd like to make sure.

A really nice thing is that the dependencies are almost identical! Wonderful!

[remydubois]

I'm the maintainer of illico so I can help answer a few points. I think the main point of attention is Dask. See below for details

Are the results the same ?

Globally, yes. This test targets a precision of 1.e-6 between scanpy and illico. I also added a note in the documentation about it. In details:

1. illico's p-values & z-scores (hence gene ordering in the resulting recordarrays) will match rank_genes_groups as long as arguments do match (i.e: use_continuity=False is passed to illico, as scanpy does not apply continuity correction).
2. Fold-change: if the correct arguments are passed to illico (is_log1p=True and exp_post_agg=True), fold change returned will be close to scanpy but not equal as illico does not add a 1.e-9 to the average expression vectors. This can be easily fixed.

Are the memory requirements the same?

That's a good point. As a matter of fact, I have benchmarks available but never made the efforts to format them / push them. Memory requirements should be similar or lower, and controllable through batch_size and n_threads (both scaling linearly RAM footprint). In a "standard" setup (8 threads, batch_size~256 on the first 2k genes of K562), this is what a comparison between scanpy and illico would look like (heap memory as measured by memray)

Note: you can run more or less this benchmark with tox -e memory-bench.

Is what is implemented currently dask-ifiable

I am not totally sure as what it means to be dask-ifiable, but here are some ideas:

1. To give you some context, illico is multithreaded over data chunks: each thread takes as input a chunk index, eventually loads in RAM (if not already) the corresponding data, then processes it. This allows illico to support lazy h5-backed AnnData by loading data chunk by chunk.
2. A naive way to make illico compatible with Dask would be to add a DaskArray DataHandler (like this one) that would call .compute() on each chunk of the DaskArray (the same way h5 DataHandlers like this one just call data[:, lb:ub]). The actual rank-sum test computation would not be handled by Dask but remain under Numba/Rust, Dask would only be used to load the needed chunk in RAM. I doubt this would be the most efficient combination of Dask (multiprocessed data loading) + illico (multithreaded data processing), but at least it would run on AnnData read as Dask arrays, and should remain much faster than current implementation (albeit caveats like NB2 detailed below).
3. A much smarter way to make illico dask friendly would be that the kernel itself is a dask.delayed function (in other words: option 2 above is joblib.Parallel(kernel(chk.compute()) for chk in chunks), while option 3 would be dask.Parallel(kernel(chk).compute() for chk in chunks). In this latter scenario illico would be dask-multiprocessed and no longer joblib-multithreaded. I would let dask experts shine some light here.

NB1: Regarding the naive Dask option, what would happen under the hood is that each thread triggers a .compute(). This appears to be GIL-blocking, so there is no "CPU usage explosion" because threads don't perfectly run concurrently. Tests I ran with this branch indeed shows a much poorer scaling of illico when ran on Dask arrays (for k562 stored as dense array, all other params fixed, such as 4 cores in the Dask cluster):

Dask cores	illico threads	runtime	Speedup vs 1 thread
4	1	4min30s	1x
4	2	2min45s	1.5x (expected 2x)
4	4	2min	2.2x (expected 4x)
4	8	1min30s	3x (expected 8x)

NB2: In terms of performance, it is very hard to benchmark illico when lazy loading is involved. It highly depends on the underlying data format (dense, CSC or CSR), IO perf, and how the chunking is performed. Dask abstracts all of that so users can easily end up performing utterly underoptimized operations (like loading a vertical chunk in RAM, when the dask array is chunked horizontally).

NB3: Note that illico voluntarily rejects some data-format-X-test combos: lazy-loading a CSR dataset to perform OVR test (reference="rest" or None) can not work, because OVR test requires each column to be entirely in RAM (to sort it) but the CSR format can not load a full column without loading the entire .indices in RAM. At this point, just load everything in RAM once and for all. Once again, this guardrail can be lifted but I would advice against.

[flying-sheep]

I assume while we’re speeding up ranked_genes_groups, we can also just remove our performance based estimation bias: #864

@ilan-gold also check out #61 while you’re at it, it’s anndata-relevant

[ilan-gold]

Fold-change: if the correct arguments are passed to illico (is_log1p=True and exp_post_agg=True), fold change returned will be close to scanpy but not equal as illico does not add a 1.e-9 to the average expression vectors. This can be easily fixed.

So it looks like this LFC business that Phil mentioned (i.e., #864) is configurable in illico which would be top notch stuff. I assume that was when you meant by the above @remydubois ?

The memory usage looks great! Great call to highlight peak memory, it seems!

A much smarter way to make illico dask friendly would be that the kernel itself is a dask.delayed function (in other words: option 2 above is joblib.Parallel(kernel(chk.compute()) for chk in chunks), while option 3 would be dask.Parallel(kernel(chk).compute() for chk in chunks). In this latter scenario illico would be dask-multiprocessed and no longer joblib-multithreaded. I would let dask experts shine some light here.

Right this has roughly been our approach. IIUC the chunks come from contiguous sections of the dask array. map_blocks has generally worked for us in other places and is quite flexible: https://docs.dask.org/en/stable/generated/dask.array.map_blocks.html. Our dask handling is downstream of https://anndata.readthedocs.io/en/latest/generated/anndata.experimental.read_elem_lazy.html which enables chunked dask access i.e., (chunk_size, n_genes) CSR blocks, which looks a bit different than your branch.

So we would wrap the in-memory kernels themselves in dask. We might need control over kernel compilation to handle parlalel/not-parallel but that can be a follow-up. I need to do a bit more digging.