Update DRVI and align with the latest scvi-tools version (recreated branch)

[moinfar]

This is a substitution of unexpectedly merged #33

Steps:

• Relax all dependencies
• Align with the latest scvi-tools
• Automatic tests passed
• Check backward compatibility
• Test and make sure metrics match
• Test to see if additional features (new in scvi-tools) work with no problem.
• Put limits on dependencies with respect to the newest packages.
• Allow batch_key and cat_covs_keys to be None when provided in setup anndata.
• Update tutorials.

Goal: solve #22

[codecov]

Codecov Report

❌ Patch coverage is 79.01786% with 47 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.46%. Comparing base (3d5c996) to head (934ac00).
⚠️ Report is 24 commits behind head on main.

Files with missing lines	Patch %	Lines
src/drvi/scvi_tools_based/module/_constants.py	36.11%	23 Missing ⚠️
...c/drvi/scvi_tools_based/model/base/_archesmixin.py	74.07%	7 Missing ⚠️
src/drvi/scvi_tools_based/model/_drvi.py	87.23%	6 Missing ⚠️
...i/scvi_tools_based/model/base/_generative_mixin.py	55.55%	4 Missing ⚠️
src/drvi/nn_modules/noise_model.py	92.10%	3 Missing ⚠️
src/drvi/scvi_tools_based/module/_drvi.py	89.65%	3 Missing ⚠️
src/drvi/nn_modules/prior.py	75.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #53      +/-   ##
==========================================
- Coverage   81.05%   80.46%   -0.59%     
==========================================
  Files          41       41              
  Lines        2291     2309      +18     
==========================================
+ Hits         1857     1858       +1     
- Misses        434      451      +17     

Files with missing lines	Coverage Δ
src/drvi/nn_modules/embedding.py	82.94% <100.00%> (ø)
src/drvi/nn_modules/encodig.py	92.85% <100.00%> (+0.17%)	⬆️
src/drvi/nn_modules/freezable.py	100.00% <100.00%> (ø)
src/drvi/nn_modules/layer/factory.py	75.55% <100.00%> (+0.55%)	⬆️
src/drvi/nn_modules/layer/linear_layer.py	97.61% <100.00%> (+0.05%)	⬆️
src/drvi/nn_modules/layer/structures.py	66.66% <100.00%> (+4.16%)	⬆️
...drvi/scvi_tools_based/data/fields/_fixed_fields.py	95.83% <100.00%> (+0.18%)	⬆️
src/drvi/scvi_tools_based/nn/_base_components.py	83.79% <100.00%> (-0.05%)	⬇️
src/drvi/utils/metrics/_aggregation.py	100.00% <100.00%> (ø)
src/drvi/utils/metrics/_pairwise.py	90.76% <100.00%> (+0.14%)	⬆️
... and 12 more

🚀 New features to boost your workflow:

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

[moinfar]

Checked training and inference using dummy anndata and separate dataloader with code below:

# # Train with custom dataloader

# +
import scvi
import scanpy as sc
from scvi.data import AnnDataManager
from scvi.data.fields import LayerField, CategoricalObsField, NumericalObsField

fields = [
    LayerField(registry_key="X", layer="counts"),
    CategoricalObsField(registry_key="batch", attr_key="batch"),
]

adata_manager = AnnDataManager(fields=fields, setup_method_args={})
adata_manager.register_fields(adata)

# +
from scvi.dataloaders import AnnDataLoader

dataloader = AnnDataLoader(
    adata_manager, 
    batch_size=1024, 
    shuffle=True, 
    drop_last=True
)
# -

for batch in dataloader:
    # 'batch' is a dictionary containing tensors
    x = batch["X"]
    b = batch["batch"]
    
    print(f"Batch X shape: {x.shape}")
    print(x.max())
    break

# +
adata_tmp = adata.copy()
adata_tmp.X = adata_tmp.X * 0
adata_tmp.layers['counts'] = adata_tmp.layers['counts'] * 0

# Setup data
DRVI.setup_anndata(
    adata_tmp,
    layer="counts",
    batch_key="batch",
    is_count_data=True,
)

# construct the model
model = DRVI(
    adata_tmp,
    n_latent=128,
    encoder_dims=[128, 128],
    decoder_dims=[128, 128],
)
model
# -

n_epochs = 400
model.train(
    max_epochs=n_epochs,
    early_stopping=False,
    early_stopping_patience=20,
    datamodule=dataloader,
)



# +
dataloader = AnnDataLoader(
    adata_manager, 
    batch_size=1024, 
    shuffle=False, 
    drop_last=False,
)

latent = model.get_latent_representation(dataloader=dataloader)
latent.shape

# +
embed = ad.AnnData(latent, obs=adata.obs)
sc.pp.subsample(embed, fraction=1.0)  # Shuffling for better visualization of overlapping colors

sc.pp.neighbors(embed, n_neighbors=10, use_rep="X", n_pcs=embed.X.shape[1])
sc.tl.umap(embed, spread=1.0, min_dist=0.5, random_state=123)
sc.pp.pca(embed)
# -

sc.pl.umap(embed, color=["batch", "final_annotation"], ncols=1, frameon=False)

[moinfar]

Disentanglement metrics checked.

[moinfar]

Loading from the previous version is also checked and is OK.
Unfortunately, the pandas problem occurs, which is an open issue in scvi-tools.

So if one faces:

ModuleNotFoundError: No module named 'pandas.core.indexes.numeric'

They can workaround the problem by:

import sys
import pandas
sys.modules['pandas.core.indexes.numeric'] = pandas.core.indexes.base
pandas.core.indexes.base.Int64Index = pandas.core.indexes.base.Index

Other than that loading from older version is fine.