Experimental Feature: Function delegation to rapids_singlecell when gpu is available

[selmanozleyen]

hi @flying-sheep @ilan-gold @timtreis @Intron7 , for squidpy 2.0 we are planning to delegate to the gpu functionality when there is a device present by default. I am trying to come up with an API for this. I saw that the scanpy settings class was very complicated so I though maybe we need something simpler for our case.

This is what I came up with. I will come up with the co_occurence support after lunch but I wanted an initial feedback on this.

Example usage

from squidpy._utils import resolve_device_arg

def co_occurrence(
    adata: AnnData | SpatialData,
    device: Literal["cpu", "gpu"] | None = None,
) -> tuple[NDArrayA, NDArrayA] | None:

    effective_device = resolve_device_arg(device)

    if effective_device == "gpu":
        from rapids_singlecell.squidpy import co_occurrence as rsc_co_occurrence

        return rsc_co_occurrence(
            adata,
        )
     ...

Installation instructions

pip install -e '.[docs,test,dev,gpu-cuda12]' --extra-index-url https://pypi.nvidia.com

[codecov]

Codecov Report

❌ Patch coverage is 92.98246% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 66.62%. Comparing base (469a7d0) to head (50ece87).

Files with missing lines	Patch %	Lines
src/squidpy/_settings/_dispatch.py	93.87%	1 Missing and 2 partials ⚠️
src/squidpy/_settings/_settings.py	92.50%	2 Missing and 1 partial ⚠️
src/squidpy/gr/_ppatterns.py	77.77%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1093      +/-   ##
==========================================
+ Coverage   66.21%   66.62%   +0.40%     
==========================================
  Files          44       46       +2     
  Lines        7163     7274     +111     
  Branches     1217     1230      +13     
==========================================
+ Hits         4743     4846     +103     
- Misses       1943     1948       +5     
- Partials      477      480       +3     

Files with missing lines	Coverage Δ
src/squidpy/_docs.py	94.73% <100.00%> (+0.53%)	⬆️
src/squidpy/_utils.py	57.22% <100.00%> (+0.97%)	⬆️
src/squidpy/gr/_ligrec.py	77.51% <100.00%> (+0.22%)	⬆️
src/squidpy/gr/_ppatterns.py	77.82% <77.77%> (-0.10%)	⬇️
src/squidpy/_settings/_dispatch.py	93.87% <93.87%> (ø)
src/squidpy/_settings/_settings.py	92.50% <92.50%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ilan-gold]

I think this can generally be wrapped up in a decorator where we:

1. check effective_device
2. Subset the inbound arguments from squidpy to those of rapids' function (I imagine n_jobs doesn't apply, this can probably be done with python's inspect module)
3. run the rapids function
4. Optionally bring the data back to the CPU (I would guess people want this, or at least it should be behind a flag)

[flying-sheep]

Anndata’s settings are indeed much nicer than scanpy’s. They do need a few parts to work, but with all that in place they’re great!

1. the dynamic docstring generation
2. the type hints
3. this test that makes sure its type hints are in sync.

[selmanozleyen]

Ok knowing what you guys would want was important thanks. Currently I will go with the inspect module and a decorator even though I think we might not be able to support 1-1 dispatch in certain configurations in the future. Also will probably have some friction on docstrings and type-hints if we are going to find common arguments dynamically.But lets see.

he dynamic docstring generation
the type hints
this test that makes sure its type hints are in sync.

yeah these are also important to me so I will ensure this.

[selmanozleyen]

Before thinking more about settings. I'd like to clarify that we might even want to do this on a fork of rsc, this is the decision we made with @timtreis in order to speed up things. But I really want to make sure we don't diverge too much and it will be in a state where we can incorperate our changes to rsc back easily. I think this direction would be good for the health of both repos.

I have a plan for squidpy functions like spatial_neighbors and nhood_enrichment to also provide more granular functions for the power users. These two functions are special compared to the rest because they are quite big and general. Not sure if it can/should support full gpu support.

spatial_neighbours -> might need splitting up for power users as discussed in #1047, more clear for sure but also can support faster more specialized functions. This way RSC can also just support spatial_neighbours_knn for example and not worry about the rest.

nhood_enrichment -> people might want to extend this or add their own flavor as seen in https://github.com/scverse/squidpy/pull/1023/changes. I think we can do this without bloating the main function also, which was the reason it wasn't merged. This one is not related to RSC but I wanted to mention it to point out how this might be a good direction for the rsc-squidpy sync and also good for squidpy.

Another benefit: Power users are important and will keep this repo alive.

In general if it wasn't for spatial_neighbours I'd be fine with using inspectmodule even though I dislike meta-progamming very much. But what do you think about these two functions? I feel like because their names are quite general, the users would like to extend them or modify them. Should we consolidate their signatures? The rest of the functions in squidpy can match with rsc 1-1 except for some function independent args like n_jobs and stuff.

[selmanozleyen]

@flying-sheep

Anndata’s settings are indeed much nicer than scanpy’s. They do need a few parts to work, but with all that in place they’re great!
...

ok I understand a bit more the scanpy settings but I don't think I'd need dynamic typing just to store a global preferred device right?

[selmanozleyen]

I checked and it works in our clusters.
I have to note, spatial_autocorr is only equal between gpu and cpu when rtol=1e-3