[FEA] Dask Array Support for rsc.pp.scrublet: A Straightforward implementation

[MPebworthEpana]

Right now, rsc.pp.scrublet doesn't support Dask arrays, and there's a relatively straightforward path to implement one (at least, from what I know).

Background:

1. Scrublet only really needs to run within a sample, or batch. This is provided to the function as a 'batch_key'
2. These samples/batches are typically on the order of < 100k cells for batches, or < 10,000 for samples, meaning that they can fit within a typical GPU's memory.

Implementation concept:

1. Check the the anndata object has a Dask array. If so, require a batch_key be provided.
2. Rechunk the dask array by batch_key - one dask array for each batch_key
3. Run scrublet in memory on each GPU (.compute_chunk_sizes())
4. Save results in obs as normal.

[Intron7]

This won't work for large scale datasets, eg. for Lake Tahoe each batch is 10Million cells. This will lead to a OOM error

[MPebworthEpana]

Good news: Batch size is not the number to evaluate- it's sample size that matters (Using batch is just faster for small datasets).

TLDR: Doublets are only really an issue within each sample, due to the biochemistry of single cell sequencing. The Lake Tahoe dataset, with 10M per batch, has 1k samples = 10k cells per sample. OOM errors won't occur at that scale. By iterating over samples, RAPIDS single cell could massively boost the speed of scrublet analysis, especially if it worked with multi-GPUs. Even if methods improve with 20k or 30k cells per sample, this implementation will still work, without OOM errors.

Longer explanation:
Running scrublet on sample, not batch, is critical for the Lake Tahoe dataset from Parse Bio, due to the biochemistry of the assay. With Parse Bio's kits (and 10x's multiplex kits), you only need to run scrublet to remove doublets within a given sample. All doublets between samples will be removed during pre-processing steps, well before any RAPIDS analysis: those types of doublets are very easily identified and removed because they have barcodes from two samples mixed together. When a doublet has barcodes from the sample sample, you then need scrublet to guess whether two different kinds of cells were mixed together. Same thing for 10x cell hashing. Also same for non-multiplexed kits where the samples were not mixed.

The within-sample doublet rate with Parse Bio is roughly 3-5%. One batch within the Lake Tahoe dataset has 10M cells over 1,000 samples - or about 10k cells per sample. Which is well within the capabilities of one GPU's memory. And critically necessary, because with 10M cells, you will have roughly 300k to 500k doublets to removed.

Let me know if you have any other questions, or concerns. I think there's a pretty straightforward path here. With this implementation, any OOM error will be because the user provided the wrong column for sample_id.

References:
https://support.parsebiosciences.com/hc/en-us/articles/360053107311-What-is-the-expected-doublet-rate
https://genomebiology.biomedcentral.com/articles/10.1186/s13059-018-1603-1
"Cell hashing can also aid in detection of doublets, whether or not multiple samples are pooled. If the input sample or samples are partitioned and labeled with multiple unique HTOs, then doublets are likely to contain cells with distinct HTOs, which can be detected during HTO demultiplexing ([Stoeckius et al., 2018]"

• https://academic.oup.com/bioinformatics/article/38/10/2791/6565315?login=false

[Intron7]

@MPebworthEpana Your approach technically works, but it's not ideal. Rechunking and sorting the array seem to trigger persistence, which causes the entire dataset to be loaded into memory. On top of that, it appears to introduce additional memory leaks.

The best workaround I’ve found so far is to keep the data as a CPU-backed Dask array, then rechunk by batch within the function and load the relevant batch onto the device. However, even this results in host RAM usage exceeding 150 GB for a 2-million-cell dataset.

I also experimented with delayed and futures for subsetting and computing on each batch, but that approach also led to excessive memory consumption.

Overall, I’m not really satisfied with any of the current options—but I’ll keep exploring.

[MPebworthEpana]

Okay, thank you for trying. Is there no way to identify the optimal size chunks (with an overlap), and then pull in fresh from a zarr store with that batch size?

I.e. chunk 1 would include batches 1, 2, and part of 3. Chunk 2 would end up including batch 3, 4, and part of 5.
Scrublet would be run within each chunk and the results of 'part of 3' would just be thrown out. Same for the 'part of 5'.