settylab / Mellon

Non-parametric density inference for single-cell analysis.
https://mellon.readthedocs.io
GNU General Public License v3.0
63 stars 3 forks source link

mellon.FunctionEstimator() exhausts 500GB of RAM with >100,000 cells #6

Closed katsturgess closed 1 year ago

katsturgess commented 1 year ago

Hi,

I am using Palantir to infer trajectories and gene expression patterns in a dataset of 198,000 haematopoietic cells. It runs beautifully for all trajectories apart from the most abundant trajectory which contains 116,000 cells. When attempting to run palantir.presults.compute_gene_trends() on 9000 genes for this lineage and 9000 genes, I receive the error: XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory allocating 108969792800 bytes. I am running this on our hpc with 500GB of RAM.

Looking at the Palantir api for the function palantir.presults.compute_gene_trends() I can see that the model fitting is being performed by mellon.FunctionEstimator(sigma=1, ls=10, n_landmarks=0, cov_func_curry= mellon.cov.Matern52); I have attempted to run just this command separately on a df of gene counts from the same trajectory and a series of pseudotime values as per the Palantir code, to reduce the number of duplicated dataframes utilising memory, but receive the same error message.

Maybe it's just too many cells or genes, but if there may be another solution I haven't tried I'd love to get it work as it was so successful for the rest of my dataset, and I'd like to apply more of your pipeline to my dataset.

Thanks for any advice.

katosh commented 1 year ago

Hi @katsturgess! You are right, for this amount of cells the underlying Gaussian Process will use exceedingly much memory. The solution we intend for this situation is to use Mellons sparse Gaussian Process implementation. You can use the sparse implementation by setting the number of landmarks to a positive integer. For the gene trend analysis a small number like 500 should produce sufficient resolution. This should work with

palantir.presults.compute_gene_trends(..., n_landmarks=500, ls=1)

or

mellon.FunctionEstimator(sigma=1, ls=1, n_landmarks=500, cov_func_curry=mellon.cov.Matern52)

Note that I also changed the length scale ls to active similar smoothness of the trends. We will probably also change the current default in the Palantir function to avoid this problem in the future. Thank you for your report!

katosh commented 1 year ago

The fix was merged to Palantir with https://github.com/dpeerlab/Palantir/pull/114. You can

pip install 'git+https://github.com/dpeerlab/Palantir'

to get the new version. Please reopen the issue if this did not solve your problem @katsturgess!

katsturgess commented 1 year ago

Thanks for the prompt reply - changing the number of landmarks fixed it using less than 1/3 of the hpc resources I had needed before. Looking forward to aplying more mellon analysis to my dataset - it looks super-interesting!

katosh commented 1 year ago

I am glad to hear this worked out. Please don't hesitate to reach out or share any other issues you encounter!