Support instantiation with 'adata=None' + LightningDataModule for training

[xiaotong743]

Description of feature

Hi Amir,

Thanks for the amazing work!

We’re training scVI on large datasets that require approximate shuffling, which isn’t supported by AnnData. In scVI, there is an EXPERIMENTAL feature where user can instantiate without adata and pass the datamodule during training:

model = scvi.model.SCVI(adata=None)
model.train(datamodule=datamodule)

It would be great to have the same support in DRVI. Thanks again.

[moinfar]

Hi,

Thanks!

Thanks for the suggestion. This technically depends on Issue #22, as the datamodule was recently added to scvi-tools (DRVI depends on an older version).
I will inform you and ask for your help with testing it when it is done.

[moinfar]

@xiaotong743
I have created #33 to solve this problem.
Can you please check the update-scvi-version branch? In theory, it should be possible to do the same in this branch.

[xiaotong743]

@moinfar
Right now, it still returns:
'DRVI.init() missing 1 required positional argument: 'adata''.
It looks like the 'BaseModelClass' in scVI explicitly allows only GIMVI and SCVI to be instantiated without adata or a registry by a condition, could be the reason:

elif (self.__class__.__name__ == "GIMVI") or (self.__class__.__name__ == "SCVI"):
            # note some models do accept empty registry/adata (e.g: gimvi)
            pass

Meanwhile, it also triggers this error in DRVI:

if isinstance(adata, AnnData):
            pass
        elif MerlinData is not None and isinstance(adata, MerlinData):
            self._data_splitter_cls = MerlinDataSplitter
        else:
            raise ValueError(
                "Only AnnData and MerlinData are supported. "
                "If you have passed an instalce of MerlinData and still get this error, "
                "make sure merlin is installed as a dependency."
            )

[moinfar]

@xiaotong743
Do you have any registry? What are the covariate information?

[moinfar]

In addition to my previous question, I took a look at the scvi-tools experimental changes. The new code makes the extension of scVI much more complicated and, in some cases, impossible. I prefer to defer any further alignment in this direction as I am sure some changes will be reverted.

I have two suggestion:

1. For now, I made passing adata=None but passing registry possible (please pull again).
   I think this way, you should be able to train the model using your custom loader. Please check here. You should hopefully be able to get a proper registry dictionary using a dummy anndata and initialize the model with that.

2. Alternatively, I think maybe it was better if scvi-tools didn't make that much complication, and a user could initialize the model using an empty anndata with same structure of obs and var:

n_cells = 100_000_000
n_genes = 4000

# Create a sparse CSR matrix with all zeros
X = sp.csr_matrix((n_cells, n_genes), dtype=np.float32)

# Create the AnnData object
adata = ad.AnnData(X)
adata.obs = ...
adata.var = ...

Then the user theoretically should be able to pass a compatible datamodule to model.train(...).