[WIP] mrVI training pipeline

tools/models/scvi/mrvi-config.yaml

@@ -0,0 +1,19 @@
+census:
+  organism: 
+    "homo_sapiens"
+  obs_query: # Use if you want to train on a subset of the model
+    null
+  obs_query_model: # Required when loading data for model training. Do not change.
+    'is_primary_data == True and nnz >= 300'
+hvg:
+  top_n_hvg: 
+    5000
+  hvg_batch: 
+    [suspension_type, assay]
+anndata:
+  batch_key:
+    [dataset_id, assay, suspension_type, donor_id]
+  model_filename:
+    anndata_model.h5ad
+model:    
+  filename: "mrvi.model"

tools/models/scvi/mrvi-generate-embedding.py

@@ -0,0 +1,69 @@
+import scvi_v2
+
+print(scvi_v2.__file__)
+
+
+import torch
+
+torch.manual_seed(0)
+
+import anndata as ad
+import numpy as np
+import yaml
+
+file = "mrvi-config.yaml"
+
+if __name__ == "__main__":
+    with open(file) as f:
+        config = yaml.safe_load(f)
+
+    adata_config = config["anndata"]
+    filename = adata_config.get("model_filename")
+
+    census_config = config["census"]
+    experiment_name = census_config.get("organism")
+
+    scvi_dataset = ad.read_h5ad(filename)
+
+    scvi_dataset.obs["nuisance"] = (
+        # scvi_dataset.obs['dataset_id'].astype(str) + '_' +
+        scvi_dataset.obs["assay"].astype(str)
+        + "_"
+        + scvi_dataset.obs["suspension_type"].astype(str)
+    )
+    scvi_dataset.obs["sample"] = (
+        scvi_dataset.obs["dataset_id"].astype(str) + "_" + scvi_dataset.obs["donor_id"].astype(str)
+    )
+
+    # hv = pd.read_pickle("hv_genes.pkl")
+    # hv_idx = hv[hv].index
+
+    # census = cellxgene_census.open_soma(census_version="2023-12-15")
+
+    # obs_query = None # not for now
+
+    # query = census["census_data"][experiment_name].axis_query(
+    #     measurement_name="RNA",
+    #     obs_query=obs_query,
+    #     var_query=soma.AxisQuery(coords=(list(hv_idx),)),
+    # )
+
+    # idx = query.obs(column_names=["soma_joinid"]).concat().to_pandas().index.to_numpy()
+
+    # adata = query.to_anndata(X_name="raw")
+
+    model_config = config.get("model")
+    model_filename = model_config.get("filename")
+
+    # May or may not be necessary
+    # scvi_v2.MrVI.setup_anndata(scvi_dataset, sample_key="sample", batch_key="nuisance", labels_key="cell_type")
+
+    mrvi_model = scvi_v2.MrVI.load("mrvi.model", adata=scvi_dataset)
+
+    latent = mrvi_model.get_latent_representation(give_z=False)
+
+    # with open("mrvi-latent-idx.npy", "wb") as f:
+    #     np.save(f, idx)
+
+    with open("mrvi-latent.npy", "wb") as f:
+        np.save(f, latent)

tools/models/scvi/mrvi.py

@@ -0,0 +1,83 @@
+import scvi_v2
+
+print(scvi_v2.__file__)
+
+
+import torch
+
+torch.manual_seed(0)
+
+import anndata as ad
+import flax.linen as nn
+import yaml
+from lightning.pytorch.loggers import TensorBoardLogger
+
+file = "mrvi-config.yaml"
+
+if __name__ == "__main__":
+    with open(file) as f:
+        config = yaml.safe_load(f)
+
+    adata_config = config["anndata"]
+    filename = adata_config.get("model_filename")
+
+    scvi_dataset = ad.read_h5ad(filename)
+
+    train_kwargs = {
+        "early_stopping": True,
+    }
+
+    plan_kwargs = {"lr": 1e-3, "n_epochs_kl_warmup": 20}
+
+    model_kwargs = {
+        "n_latent": 100,
+        "n_latent_u": 20,
+        "qz_nn_flavor": "attention",
+        "px_nn_flavor": "attention",
+        "qz_kwargs": {"use_map": False, "stop_gradients": False, "stop_gradients_mlp": True, "dropout_rate": 0.03},
+        "px_kwargs": {
+            "stop_gradients": False,
+            "stop_gradients_mlp": True,
+            "h_activation": nn.softmax,
+            "dropout_rate": 0.03,
+            "low_dim_batch": True,
+        },
+        "learn_z_u_prior_scale": False,
+        "z_u_prior": False,
+        "u_prior_mixture": True,
+        "u_prior_mixture_k": 100,
+    }
+
+    scvi_dataset.obs["nuisance"] = (
+        # scvi_dataset.obs['dataset_id'].astype(str) + '_' +
+        scvi_dataset.obs["assay"].astype(str)
+        + "_"
+        + scvi_dataset.obs["suspension_type"].astype(str)
+    )
+    scvi_dataset.obs["sample"] = (
+        scvi_dataset.obs["dataset_id"].astype(str) + "_" + scvi_dataset.obs["donor_id"].astype(str)
+    )
+
+    model_config = config.get("model")
+    n_hidden = model_config.get("n_hidden")
+    n_latent = model_config.get("n_latent")
+    n_layers = model_config.get("n_layers")
+    dropout_rate = model_config.get("dropout_rate")
+    output_filename = model_config.get("filename")
+
+    scvi_v2.MrVI.setup_anndata(scvi_dataset, sample_key="sample", batch_key="nuisance", labels_key="cell_type")
+    mrvi_model = scvi_v2.MrVI(scvi_dataset, **model_kwargs)
+
+    logger = TensorBoardLogger("mrvi_tb_logs", name="mrvi_50_epochs")
+
+    mrvi_model.train(max_epochs=50, batch_size=4096, plan_kwargs=plan_kwargs, **train_kwargs)
+
+    mrvi_model.save(output_filename)
+
+    # # Get z representation
+    # adata.obsm["X_mrvi_z"] = mrvi_model.get_latent_representation(give_z=True)
+    # # Get u representation
+    # adata.obsm["X_mrvi_u"] = mrvi_model.get_latent_representation(give_z=False)
+    # sc.pp.neighbors(adata, use_rep="X_mrvi_u", key_added="neighbors_mrvi", method='rapids', n_neighbors=30)
+    # sc.tl.umap(adata, neighbors_key="neighbors_mrvi", method='rapids')
+    # sc.pl.umap(adata, color=['dataset_id', 'cell_subclass', 'suspension_type', 'sex'], ncols=1, frameon=False, wspace=0.4, title='mrVI (Census)')