Batch correction for joint modality umap

[GirayEryilmaz]

Hi Muon Team!

I have been using muon to analyze some cite-seq data. I wanted to make a joint umap of the two modalities. Which I can do with the pipeline I found in the citeseq data integration tutorial.

batch_key = 'batch'
sc.external.pp.bbknn(mdata['rna'], batch_key=batch_key)
sc.external.pp.bbknn(mdata['adt'], batch_key=batch_key)

mu.pp.neighbors(mdata, key_added='wnn')
mu.tl.umap(mdata, neighbors_key='wnn', random_state=10)

However what I get suffers from batch effects. I had expected the umap to be batch corrected since muon.pp.neighbors works on batch corrected neighborhoods of individual modalities but that does not seem to be the case. Am I missing something?

Notes:

1. For clustering I use muon's leiden which produces better results than scanpy's leiden in terms of batch effects.
2. The umaps of individual modalities are batch corrected successfully.

Thanks!