feat: add differential abundance for scvi-tools models

[lordy5]

@ori-kron-wis @canergen I've finished fixing some issues with the code and have now added tests. Let me know if there's anything else I should do.

[codecov]

Codecov Report

❌ Patch coverage is 95.34884% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 74.92%. Comparing base (7ce63e0) to head (210ddf0).

Files with missing lines	Patch %	Lines
src/scvi/model/base/_base_model.py	33.33%	2 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (7ce63e0) and HEAD (210ddf0). Click for more details.

HEAD has 118 uploads less than BASE

Flag	BASE (7ce63e0)	HEAD (210ddf0)
	121	3

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3618      +/-   ##
==========================================
- Coverage   84.70%   74.92%   -9.78%     
==========================================
  Files         225      225              
  Lines       21637    21680      +43     
==========================================
- Hits        18327    16244    -2083     
- Misses       3310     5436    +2126     

Files with missing lines	Coverage Δ
src/scvi/model/base/_vaemixin.py	92.24% <100.00%> (-2.14%)	⬇️
src/scvi/model/base/_base_model.py	76.02% <33.33%> (-5.04%)	⬇️

... and 44 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[ori-kron-wis]

Thanks @lordy5, I have a few questions/clarifications, mainly engineering wise.
The general idea is to do as much code reuse, add more tests and have the support for the generalization of this function to all models in a proper way.

So far, we had the DA support for CytoVi and MrVI only.

1. Is this implementation here should be general enough for all models?
2. If this is the case, I would expect this PR to have more tests, for other models as well. I would at least add tests for SCANVI, TOTALVI (with mudata input) and a spatial model (RESOLVI, SCVIVA). This is something users are going to try out and lets make sure it works.
3. If theres no support for specific models I would add a warning of "DA is currently not supported" - This is something we did for get_normalized_expression function.
4. However, (3) was achieved because get_normalized_expression and for that manner, differential_expression, are implemented in _vaemixin.py. So model that does not inherit this mixin will not have it and the message of "not implemented" will appear. In your implementation all is implemented directly in base path in da_pytoch, but like _de_core, we might be missing the wrapper to the DA in _vaemixin.py (like we do for DE) - so this is how I would recommend it to be implemented, otherwise it is has less control and will likely break for users who try out anything.
5. So, seems CytoVi and mrVI added their own things on top of your DA, its fine and we have many cases that a model expand on a base/mixin function, question is whether your implementation can support those models implementations with minimum additions, if so, much code can be removed (we want to reuse as much as possible). In such case need to update those models as well and validate for correctness in their tutorials.
6. (5) is also relvant as @canergen responded to a user that we are going to have the memory saving version of DA in MRVI soon (#3624 (comment)) however as I wrote right now, the implementation here has no connection to the MRVI implementation unless we will do something about it and we should (update mrvi's DA or have this one supports it).
7. Add changlog

[canergen]

Some comments

1. Is this implementation here should be general enough for all models?
   Yes, we should check with @florianingelfinger for cytoVI though.
2. If this is the case, I would expect this PR to have more tests, for other models as well. I would at least add tests for SCANVI, TOTALVI (with mudata input) and a spatial model (RESOLVI, SCVIVA). This is something users are going to try out and lets make sure it works.
   Sounds good. Let’s not blast up test time though.
3. If theres no support for specific models I would add a warning of "DA is currently not supported" - This is something we did for get_normalized_expression function.
   It should be for all models except Stereoscope and CondSCVI/DestVI.
4. However, (3) was achieved because get_normalized_expression and for that manner, differential_expression, are implemented in _rnamixin.py. So model that does not inherit this mixin will not have it and the message of "not implemented" will appear. In your implementation all is implemented directly in base path in da_pytoch, but like _de_core, we might be missing the wrapper to the DA in _rnamixin.py (like we do for DE) - so this is how I would recommend it to be implemented, otherwise it is has less control and will likely break for users who try out anything.
   It shouldn’t be in rnamixin but yes we should also print an error.
5. So, seems CytoVi and mrVI added their own things on top of your DA, its fine and we have many cases that a model expand on a base/mixin function, question is whether your implementation can support those models implementations with minimum additions, if so, much code can be removed (we want to reuse as much as possible). In such case need to update those models as well and validate for correctness in their tutorials.
   The MrVI changes don’t need to be reflected. For Jax it requires its own function but I wouldn’t write a wrapper to have same API in Jax.

[lordy5]

Hi Ori and Can,
Thanks for the feedback and thoughts. I was also thinking about many of these things.

@ori-kron-wis

1. Yes however I haven't tested with other models yet as you said, and I guess we're not sure about cytovi yet.
2. Sounds, good, I will go ahead and add tests for other models. I was planning to, but wasn't sure which models it would be relevant for. I think I know which models to test with now after Can's feedback above.
3. Makes sense.
4. Yes, I definitely agree giving an error like (3) is a good way to handle unsupported models. However, I don't fully understand which way you're suggesting to do this. Is it that I should keep the base DA functionality in da_pytorch, and then make a wrapper in one of the mixins? Then for models that don't inherit from these mixins, there will be an error when trying to call DA? If that's the case, where would make the most sense to put this DA wrapper?
5. I'll take a look at the CytoVI implementation but I'm pretty sure I can use my implementation to remove duplicate code. But just to clarify, @canergen, were you saying above that the DA for MrVI should be standalone?
6. Ok will make sure to add that.

[florianingelfinger]

I think for CytoVI I would keep the DA as is (otherwise there may be some adaptation needed). I would also expose the aggregation function as for many cases the logsumexp favors high DA scores for outlier samples compared to sth like the logmedian.