Add CLR normalization for ADT data

Author: Marius1311

src/methods/cellmapper_scvi/config.vsh.yaml

@@ -58,6 +58,10 @@ arguments:
     choices: ["clr", "log_cp10k"]
     default: "clr"
     description: Normalization method for ADT data, clr = centered log ratio. 
+  - name: "--plot_umap"
+    type: boolean
+    default: false
+    description: Whether to plot the UMAP embedding of the latent space (for diagnoscic purposes)
 resources:
   - type: python_script
     path: script.py

src/methods/cellmapper_scvi/script.py

@@ -7,14 +7,15 @@
 # Note: this section is auto-generated by viash at runtime. To edit it, make changes
 # in config.vsh.yaml and then run `viash config inject config.vsh.yaml`.
 par = {
-    'input_train_mod1': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/swap/train_mod1.h5ad',
-    'input_train_mod2': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/swap/train_mod2.h5ad',
-    'input_test_mod1': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/swap/test_mod1.h5ad',
+    'input_train_mod1': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/train_mod1.h5ad',
+    'input_train_mod2': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/train_mod2.h5ad',
+    'input_test_mod1': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/test_mod1.h5ad',
     'output': 'output.h5ad',
     'n_neighbors': 30, 
     'kernel_method': 'hnoca',
-    'use_hvg': True,
-    'adt_normalization': 'clr',
+    'use_hvg': False,
+    'adt_normalization': 'clr',  # Normalization method for ADT data
+    'plot_umap': True,
 
 }
 meta = {
@@ -43,7 +44,9 @@
 
 # Compute a latent representation using an appropriate model based on the modality
 print("Get latent representation", flush=True)
-adata = get_representation(adata=adata, modality=mod1, use_hvg=par['use_hvg'], adt_normalization=par['adt_normalization'])
+adata = get_representation(
+    adata=adata, modality=mod1, use_hvg=par['use_hvg'], adt_normalization=par['adt_normalization'], plot_umap=par['plot_umap']
+    )
 
 # Place the representation back into individual objects
 input_train_mod1.obsm["X_scvi"] = adata[adata.obs["split"] == "train"].obsm["X_scvi"].copy()

src/methods/cellmapper_scvi/utils.py

@@ -3,10 +3,16 @@
 import scvi 
 from scipy.sparse import issparse, csr_matrix, csc_matrix
 import muon
+import scanpy as sc
 
 
 def get_representation(
-        adata: ad.AnnData, modality: Literal["GEX", "ADT", "ATAC"], use_hvg: bool = True, adt_normalization: Literal["clr", "log_cp10k"] = "clr") -> ad.AnnData:
+        adata: ad.AnnData, 
+        modality: Literal["GEX", "ADT", "ATAC"], 
+        use_hvg: bool = True, 
+        adt_normalization: Literal["clr", "log_cp10k"] = "clr",
+        plot_umap: bool = False,
+    ) -> ad.AnnData:
     """
     Get a joint latent space representation of the data based on the modality.
     
@@ -29,6 +35,9 @@ def get_representation(
         Normalization method for ADT data. Options are:
          - "clr" (centered log-ratio transformation)
          - "log_cp10k" (normalization to 10k counts per cell and logarithm transformation)
+    plot_umap
+        Purely for diagnostic purposes, to see whether the data integration looks ok, this optionally computes 
+        a UMAP in shared latent space and stores a plot. 
 
     Returns
     -------
@@ -46,8 +55,9 @@ def get_representation(
     # Setup the AnnData object for scVI
     if modality == "GEX":
         layer = "counts"
-        scvi.model.SCVI.setup_anndata(adata, batch_key="batch", layer=layer)
-        model = scvi.model.SCVI(adata, gene_likelihood="nb", n_layers=2, n_latent=30)
+        scvi.model.SCVI.setup_anndata(adata, layer=layer, categorical_covariate_keys=["split", "batch"])
+        model = scvi.model.SCVI(adata)
+
     elif modality == "ADT":
         print(f"Normalizing the ADT data using method '{adt_normalization}'")
         if adt_normalization == "clr":
@@ -60,11 +70,11 @@ def get_representation(
             raise ValueError(f"Unknown ADT normalization method: {adt_normalization}")
 
         layer = "adt_normalized"
-        scvi.model.SCVI.setup_anndata(adata, batch_key="batch", layer=layer)
+        scvi.model.SCVI.setup_anndata(adata, layer=layer, categorical_covariate_keys=["split", "batch"])
         model = scvi.model.SCVI(adata, gene_likelihood="normal", n_layers=1, n_latent=10)
     elif modality == "ATAC":
         layer = "counts"
-        scvi.model.PEAKVI.setup_anndata(adata, batch_key="batch", layer=layer)
+        scvi.model.PEAKVI.setup_anndata(adata, layer=layer, categorical_covariate_keys=["split", "batch"])
         model = scvi.model.PEAKVI(adata)
     else:
         raise ValueError(f"Unknown modality: {modality}")
@@ -80,4 +90,11 @@ def get_representation(
     # Get the latent representation
     adata.obsm["X_scvi"] = model.get_latent_representation()
 
+    if plot_umap:
+        sc.pp.neighbors(adata, use_rep="X_scvi")
+        sc.tl.umap(adata)
+
+        plot_name = f"_{modality}_{adt_normalization}_use_hvg_{use_hvg}.png"
+        sc.pl.embedding(adata, basis="umap", color=["batch", "split"], show=False, save=plot_name)
+
     return adata