Solo: scvi.external.SOLO.from_scvi_model() errors if scvi.data.setup_anndata() used a layer

[sjfleming]

If you train an scvi batch-effect correction model, and then use that trained VAE as an input for Solo, then it seems that an AssertionError occurs upon scvi.external.SOLO.from_scvi_model() if the original scvi.data.setup_anndata() specified a layer that was not None.

The reason appears to be that when Solo creates simulated doublets here
https://github.com/YosefLab/scvi-tools/blob/334c03822b34d2e1dd9b87ecff6bfb0f4803f05e/scvi/external/solo/_model.py#L232-L236
it does not copy the layer information into the doublets_ad object. Then when scVI calls _validate_anndata(adata) from get_latent_representation(), the below error appears.

To reproduce

import scvi

# assume adata is a loaded AnnData object with a layer called "cellbender", and adata.obs['scvi_batch']

scvi.data.setup_anndata(adata, layer='cellbender', batch_key='scvi_batch')
vae = scvi.model.SCVI(adata)
vae.train()

print('running SOLO doublet detection algorithm on the scVI model')
for batch_label in adata.obs['scvi_batch'].unique():
    solo_batch = scvi.external.SOLO.from_scvi_model(vae, restrict_to_batch=batch_label)  # error occurs here
    solo_batch.train()
    solo_df = solo_batch.predict()
    adata.uns['solo_' + batch_label] = solo_df.copy()

The error

running SOLO doublet detection algorithm on the scVI model
�[34mINFO    �[0m Creating doublets, preparing SOLO model.                               
/home/user/miniconda/lib/python3.7/site-packages/scvi/data/_anndata.py:1232: UserWarning: Categorical encoding for batch is similar but not equal between the anndata used to train and the anndata passed in. Will attempt transfer. Expected categories: [0 1]. Received categories: [0].
 
  warnings.warn(warning_msg.format(category, mapping1, mapping2))
Traceback (most recent call last):
  File "<stdin>", line 81, in <module>
  File "/home/user/miniconda/lib/python3.7/site-packages/scvi/external/solo/_model.py", line 180, in from_scvi_model
    doublet_latent_rep = scvi_model.get_latent_representation(doublet_adata)
  File "/home/user/miniconda/lib/python3.7/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/user/miniconda/lib/python3.7/site-packages/scvi/model/base/_vaemixin.py", line 158, in get_latent_representation
    adata = self._validate_anndata(adata)
  File "/home/user/miniconda/lib/python3.7/site-packages/scvi/model/base/_base_model.py", line 153, in _validate_anndata
    transfer_anndata_setup(self.scvi_setup_dict_, adata)
  File "/home/user/miniconda/lib/python3.7/site-packages/scvi/data/_anndata.py", line 435, in transfer_anndata_setup
    x_loc, x_key = _setup_x(adata_target, layer)
  File "/home/user/miniconda/lib/python3.7/site-packages/scvi/data/_anndata.py", line 798, in _setup_x
    ), "{} is not a valid key in adata.layers".format(layer)
AssertionError: cellbender is not a valid key in adata.layers

But "cellbender" is a valid key in adata.layers.

AnnData object with n_obs × n_vars = 94169 × 36601
    obs: 'sample', 'exon_fraction', 'mito_fraction', 'batch', 'scrublet', 'entropy', 'individual', 'tissue', 'ribo_fraction', 'mito_ribo_fraction', 'n_gene', 'n_umi', 'outlier', 'scvi_batch'
    var: 'gene_ids', 'feature_types', 'genome'
    layers: 'cellbender', 'cellranger'

It's just not a valid key in the simulated doublets_ad created by Solo.

Versions:

0.11.0

[adamgayoso]

Thanks for raising this. We will address this soon.

[adamgayoso]

Btw, we have this method to make the logger output nicer when writing to a file.

can be accessed like scvi.settings.reset_logging_handler()

[sjfleming]

Thanks!!