Allow custom label_key for scib-metrics in scvi.autotune

[catheriz]

Is your feature request related to a problem? Please describe.
When using scvi.autotune.run_autotune with scIB metrics for multi-modal integration, I cannot specify which .obs columns should be used as labels for autotune with scib-metrics. The autotune framework only uses the model's internally registered labels fields from the data manager. setup_mudata does not support labels_key, so label-dependent scib metrics cannot access biologically meaningful labels like tissue or cell type that exist in my data but were not registered during model setup. This makes it impossible to tune hyperparameters using tissue-aware integration quality metrics, even though those labels are available in the dataset.

Describe the solution you'd like
I would like scvi.autotune.run_autotune and AutotuneExperiment to accept optional scib_label_key parameters or allow me to use label in the setup process

[ori-kron-wis]

Hey @catheriz, you are correct with your observation.
However, the best way to deal with this is to implement the label_key into the multi-modal model, running a setup_mudata. We already have the TOTALANVI, and we are working on MULTIANVI (MULTIVI + class label support).

Trying to forcefully fix and have the tuner accept even a different label (say, the above-mentioned solution is done) in a mudata settings is problematic as the tuner gets its inputs from the outputs of a trained model, which are already shuffled and the size of it also depends if a validation was done and so on, so this information has to come from whithin the model, which is either way something we need to do, not only for scib optimization.

Having said that, you can run the tuner without shuffling the data and no validation, with:

"train_size": 1.0,
"shuffle_set_split": False

as train_params in the search_space
Also, force shuffle=False in train_dataloader in the relevant DataSplitter and inject the values of the new labels you want to optimize by through the API from run_autone to the use of the Benchmarker. I will not implement it here, but it's rather straightforward for you to do, and again, this is not the solution we want in the codebase.

Once again, this is not what MULTIVI is optimized to do, which is rather the integration of its modalities, and not sure if those scores will have a meaning.