https://github.com/theislab/invae
Invariant Representation learning
https://github.com/theislab/invae
Last synced: 10 months ago
JSON representation
Invariant Representation learning
- Host: GitHub
- URL: https://github.com/theislab/invae
- Owner: theislab
- License: bsd-3-clause
- Created: 2023-07-12T18:33:54.000Z (about 3 years ago)
- Default Branch: main
- Last Pushed: 2024-06-26T12:40:54.000Z (about 2 years ago)
- Last Synced: 2025-09-06T14:19:41.747Z (11 months ago)
- Language: Jupyter Notebook
- Size: 28.5 MB
- Stars: 17
- Watchers: 3
- Forks: 3
- Open Issues: 1
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# inVAE
inVAE is a conditionally invariant variational autoencoder that identifies both spurious (distractors) and invariant features.
It leverages domain variability to learn conditionally invariant representations. We show that inVAE captures biological variations in single-cell datasets obtained from diverse conditions and labs.
inVAE incorporates biological covariates and mechanisms such as disease states, to learn an invariant data representation. This improves cell classification accuracy significantly.
## Installation
1. PyPI only
```pip install invae```
2. Development Version (latest version on github)
```git clone https://github.com/theislab/inVAE.git```
```cd inVAE```
```pip install .```
## Example
[Integration of Human Lung Cell Atlas using both healthy and disease samples](https://github.com/theislab/inVAE/blob/master/notebooks/inVAE_LungAtlas.ipynb)
## Usage
1. Load the data:
```adata = sc.read(path/to/data)```
2. Optional - Split the data into train, val, test (in supervised case for training classifier as well)
3. Initialize the model, either Factorized or Non-Factorized:
```
from inVAE import FinVAE, NFinVAE`
inv_covar_keys = {
'cont': [],
'cat': ['cell_type', 'disease'] #set to the keys in the adata
}
spur_covar_keys = {
'cont': [],
'cat': ['batch'] #set to the keys in the adata
}
model = FinVAE(
adata = adata_train,
layer = 'counts', # The layer where the raw counts are stored in adata (None for adata.X: default)
inv_covar_keys = inv_covar_keys,
spur_covar_keys = spur_covar_keys,
latent_dim_inv = 20,
latent_dim_spur = 5,
device = 'cpu',
decoder_dist = 'nb'
)
```
Set `inject_covar_in_latent= True` if you wish to add the spurious conditions directly to the latent (instead of learning the spurious latents). This gives you the most compatible version to SCVI.
For non-factorized model, use:
```
model = NFinVAE(
adata = adata_train,
layer = 'counts', # The layer where the raw counts are stored in adata (None for adata.X: default)
inv_covar_keys = inv_covar_keys,
spur_covar_keys = spur_covar_keys,
latent_dim_inv = 20,
latent_dim_spur = 5,
device = 'cpu',
decoder_dist = 'nb'
)
```
4. Train the generative model:
```model.train(n_epochs=500, lr_train=0.001, weight_decay=0.0001)```
5. Get the latent representation: In the case that covariates that were used are missing the encoder gets zeros as inputs for that sample and covariate
```
# This works for an arbitrary adata object not only for the training data
# Other options for the latent type are: full or spurious
latent = model.get_latent_representation(adata, latent_type='invariant')
```
6. Optional - Train the classifer (for cell types): if adata_val is not given or does not have labels the classifier is just trained on the adata object the generative model was trained on (here: adata_train)
```
model.train_classifier(
adata_val,
batch_key = 'batch',
label_key = 'cell_type',
)
```
7. Optional - Predict cell types:
```
# Other possible dataset_types: train or val
# train corresponds to the adata_train object above
# val to the adata used in the train_classifier function
# test is for a new unseen object
pred_test = model.predict(adata_test, dataset_type='test')
```
8. Optional - Infer latent representation via trained classifier:
```
# As key one can use 'val' or 'test' depending which key was used in the predict function above
# E.g. for invariant latent representation
# Otherwise do not subset for the full representation or subset to the last dimensions for the spurious one
latent_samples_inv = model.saved_latent['val'][:, :model.latent_dim_inv]
```
9. Optional - Saving and loading model:
```
model.save('./checkpoints/path.pt')
model.load('./checkpoints/path.pt')
```
Newest version now supports loading model parameters and weights in one:
```
# Same syntax for saving but now saves model params too
model.save('./checkpoints/path.pt')
# New loading function (old function can be used to load older checkpoints)
FinVAE.load_model('./checkpoints/path.pt', adata_train, device)
# or for NFinVAE
NFinVAE.load_model('./checkpoints/path.pt', adata_train, device)
```
## Dependencies
* scanpy==1.9.3
* torch==2.0.1
* tensorboard==2.13.0
* anndata==0.8.0
## Citation
[H. Aliee, F. Kapl, S. Hediyeh-Zadeh, F. J. Theis, Conditionally Invariant Representation Learning for Disentangling Cellular Heterogeneity, 2023](https://arxiv.org/abs/2307.00558)