Integrating Multi-source Gene Expression (GEX)

This is a repository for the code for integrating multi-source gene expression (GEX) data using the scatlasvae.model.scAtlasVAE from the scAtlasVAE package.

1import scatlasvae
2
3adata = scatlasvae.read_h5ad("path/to/adata.h5ad")

The adata is a anndata.AnnData object with raw GEX count matrix stored in adata.X.

Training the VAE model using batch key

The following use sample_name as the batch key. The batch index is converted to 10-dimensional embedding for the decoder part of the model to remove the batch effect.

1vae_model = scatlasvae.model.scAtlasVAE(
2  adata=adata,
3  batch_key="sample_name",
4  batch_embedding='embedding',
5  device='cuda:0',
6  batch_hidden_dim=10,
7)
8loss_record = vae_model.fit()
9adata.obsm['X_gex'] = vae_model.get_latent_embedding()

The following use sample_name as the batch key. The batch index is converted to one-hot encoding for the decoder part of the model to remove the batch effect.

1vae_model = scatlasvae.model.scAtlasVAE(
2  adata=adata,
3  batch_key="sample_name",
4  batch_embedding='onehot',
5  device='cuda:0',
6)
7loss_record = vae_model.fit()
8adata.obsm['X_gex'] = vae_model.get_latent_embedding()

Training the VAE model using batch key and categorical covariates (e.g. study_name)

1vae_model = scatlasvae.model.scAtlasVAE(
2  adata=adata,
3  batch_key=["sample_name","study_name"],
4  batch_embedding='embedding',
5  device='cuda:0',
6  batch_hidden_dim=10
7)
8loss_record = vae_model.fit()
9adata.obsm['X_gex'] = vae_model.get_latent_embedding()

Training the VAE model using batch key and label key (e.g. cell_type)

 1vae_model = scatlasvae.model.scAtlasVAE(
 2  adata=adata,
 3  batch_key="sample_name",
 4  label_key='cell_type',
 5  batch_embedding='embedding',
 6  device='cuda:0',
 7  batch_hidden_dim=10,
 8)
 9loss_record = vae_model.fit()
10adata.obsm['X_gex'] = vae_model.get_latent_embedding()

Training the VAE model using multiple batch keys and mutiple label keys

When integrating datasets from different atlas with different batch keys and cell type annotations, the batch_key and label_key can be a list of keys, and value in label_key can be ‘undefined’ if the cell type annotation is not available in the dataset. These information will not be used in the model training.

After training the multi-batch and multi-label model, one can use :py:method:`scatlasvae.ut.cell_type_alignment` to visualize the alignment of cell types.

 1vae_model = scatlasvae.model.scAtlasVAE(
 2  adata=adata,
 3  batch_key=["sample_name", "study_name"],
 4  label_key=["cell_type_1","cell_type_2"],
 5  batch_embedding='embedding',
 6  device='cuda:0',
 7  batch_hidden_dim=10,
 8)
 9loss_record = vae_model.fit()
10adata.obsm['X_gex'] = vae_model.get_latent_embedding()
11
12predictions = vae_model.predict_labels(return_pandas=True)
13predictions.columns = list(map(lambda x: 'predicted_'+x, predictions.columns))
14adata.obs = adata.obs.join(predictions)
15
16predictions_logits = vae_model.predict_labels(return_pandas=False)
17adata.uns['predictions_logits'] = predictions_logits

1count, fig = scatlasvae.ut.cell_type_alignment(
2  adata,
3  obs_1='predicted_cell_type_1',
4  obs_2='predicted_cell_type_2',
5  return_fig=True
6)
7fig.show()

Saving the VAE model

The save_to_disk method saves the VAE model to the path.

1vae_model.save_to_disk(path)