Extensibility API for niche calling

Description of feature

Hi @selmanozleyen, @shashkat and @Intron7,

in https://github.com/scverse/squidpy/pull/1245 a new, modular framework for niche calling was implemented. While I already provided implementation-specific feedback in https://github.com/scverse/squidpy/issues/1277, here I want to propose an "extensibility API" that allows to build custom niche calling algorithms based on the existing building blocks, similar to what's already implemented for the neighborhood algorithms.

That means defining an interface an "embedder" and a "clusterer" and a calculate_niche_custom function as a high level API (analogous to spatial_neighbors_from_builder. Compared to the current implementation, I would increase the level of abstraction. An embedder or clusterer doesn't need to know anything about the AnnData object or its structure, it just gets array(s) as input and produces array(s) as output. That makes them easier to test and reason about and increases the compatibility with the sklearn ecosystem (see below).

Clusterer

A clusterer takes an embedding and generates labels. Instead of rolling our own interface, I think we should piggyback on scikit-learn. This is nice, because (1) it's a well-known API and (2) we can just use off-the-shelf scikit-learn algorithms without any custom code!

Anything that implements the ClusterMixin interfaces qualifies:

• _GMMClusterer goes away entirely, just use sklearn.mixture.GaussianMixture.
• Want to use kmeans for clustering for its speed? Use sklearn.cluster.KMeans!
• Want to try BayesianGMM with cellcharter? Use sklearn.mixture.BayesianGaussianMixture!
• or implement your own (in our case, leiden)

  class LeidenClusterer(ClusterMixin, BaseEstimator):
      def __init__(self): ... # set parameters
      def fit(self, X, y=None): 
           tmp_ad = ad.AnnData(X=X)
           sc.pp.neighbors(tmp_ad)
           sc.tl.leiden(tmp_ad)
           return tmp_od.obs["leiden"]
  

Embedder

This is essentially a Transformer in the scikit-learn world, but it's a bit trickier, because in addition to a feature space, our Embedders typically also require a spatial neighborhood graph. Relying on the scikit-learn API directly doesn't cover us. So I'd suggest to go with something that is still sklearn-like:

class BaseSpatialTransformer(BaseEstimator):
    def fit_transform(self, X, y=None, *, S=None, **fit_params):
        if y is None:
            return self.fit(X, S=S, **fit_params).transform(X, S=S)
        return self.fit(X, y, S=S, **fit_params).transform(X, S=S)
    
    def fit(self, X, y=None, *, S=None): ...
    def transform(self, X, *, S=None): ...

Additionally, one could look into the experimental metadata routing which pass the spatial matrix along when using scikit-learn features such as Pipeline() or GridSearchCV(), but I don't think it's that relevant at this point.

Piecing it together

Moving out (1) extraction of the matrics from AnnData and (2) hyperparameter search from the Embedders and Clusterers means that the outer functions need to do a bit more. I still think it's better, as it makes the individual components more reusable. I think the solution could be that instead of a clusterer and an embedder class, "recipe" callback function is passed to the calculate_niche_custom function that takes an AnnData slice and returns one or multiple .obs columns. Alternatively, there could be another base class that defines an abstract function for this.

For calculate_niche_neighborhood this would look somewhat like this

def calculate_niche_neighborhood(adata, groups, *, mask, library_key, resolutions, spatial_connectivities_key, n_neighbors, scale, distance, abs_nhood):
    embedder = NeighborhoodProfileEmbedder(scale=scale, distance=distance, abs_nhood=abs_nhood)
    clusterer = LeidenClusterer(n_neighbors=n_neighbors, ...)
    
    def _recipe(adata_slice):
         embedding = embedder.fit_transform(X = adata.obs[groups], S = adata.obsp[spatial_connectivities_key])
         obs = {}
         for res in resolutions:
             obs[f"leiden_{res}"] = clusterer.fit_predict(embedding)
        return obs
   
    run_niche_pipeline(adata, maks=mask, library_key=library_key, recipe=_recipe)

Implementing a "custom" niche calling algorithm, e.g. combining neighborhood embedding with GaussianMixture could be easily realized that way.

Let me know what you think!

Start by reading the modular niche-calling framework in PR #1245, the implementation feedback in issue #1277, and Squidpy's neighborhood extensibility documentation. Map the proposed ClusterMixin-based clusterer, spatial transformer, and custom recipe API onto the existing niche-calling implementation. Done means agreeing on the abstraction and delivering a reusable custom pipeline with compatible built-in behavior.