feat: add differential abundance for scvi-tools models

[lordy5]

Still in progress

I still need to add tests, etc.

[codecov[bot]]

Codecov Report

:white_check_mark: All modified and coverable lines are covered by tests. :white_check_mark: Project coverage is 74.75%. Comparing base (e4ae4d2) to head (39b89fc).

:exclamation: There is a different number of reports uploaded between BASE (e4ae4d2) and HEAD (39b89fc). Click for more details.

HEAD has 45 uploads less than BASE

Flag	BASE (e4ae4d2)	HEAD (39b89fc)
	48	3

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #3618       +/-   ##
===========================================
- Coverage   86.51%   74.75%   -11.76%     
===========================================
  Files         225      226        +1     
  Lines       21631    21671       +40     
===========================================
- Hits        18713    16200     -2513     
- Misses       2918     5471     +2553     

Files with missing lines	Coverage Δ
src/scvi/model/base/_da_pytorch.py	100.00% <100.00%> (ø)

... and 62 files with indirect coverage changes

:rocket: New features to boost your workflow:

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

[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.