can train Novae on a subset of vars

Author: quentinblampey

docs/tutorials/proteins.ipynb

@@ -195,6 +195,19 @@
     "novae.data.quantile_scaling(adatas)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## (Optional) Remove low quality proteins\n",
+    "\n",
+    "You can check your proteins intensity, e.g., using `sc.pl.spatial(adata, color=adata.var_names, vmax=\"p95\", spot_size=20)`\n",
+    "\n",
+    "If you have proteins whose intensity **highly depends on the Field-Of-View** (typically, when you see a grid-like pattern) then you should not provide these to Novae. Otherwise, since you'll have FOV-specific intensities, Novae will detect the FOVs as specific domains.\n",
+    "\n",
+    "> Practically, you can either (i) subset your `AnnData` object, or (ii) provide `var_names=<list-of-proteins-to-keep>` to `novae.Novae` when initializing the model (see next section). In the latter case, make sure to update `embedding_size` to be lower than the number of proteins you have chosen.\n"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

novae/model.py

@@ -78,7 +78,7 @@ def __init__(
             sensitivity_noise_std: Standard deviation for the multiplicative for for the noise augmentation.
             dropout_rate: Dropout rate for the genes during augmentation.
             scgpt_model_dir: Path to a directory containing a scGPT checkpoint, i.e. a `vocab.json` and a `best_model.pt` file.
-            var_names: Only used when loading a pretrained model. Do not use it yourself.
+            var_names: Used when loading a pretrained model. Can also be used to specify the names of the variables to train on, e.g. to not consider low quality proteins whose intensity highly depends on the FOV.
         """
         super().__init__()
         ### Initialize cell embedder and prepare adata(s) object(s)

novae/utils/_validate.py

@@ -28,12 +28,13 @@ def prepare_adatas(
 
     Args:
         adata: An `AnnData` object, or a list of `AnnData` objects. Optional if the model was initialized with `adata`.
-        var_names: Only used when loading a pretrained model. Do not use it yourself.
+        var_names: Only used when loading a pretrained model, or to select a subset of vars to use.
 
     Returns:
         A list of `AnnData` objects ready to be used by the model. If only one `adata` object is provided, it will be wrapped in a list.
     """
     assert adata is not None or var_names is not None, "One of `adata` and `var_names` must not be None"
+    var_names = lower_var_names(var_names) if var_names is not None else None
 
     if adata is None:
         return None, var_names

tests/test_model.py

@@ -5,6 +5,7 @@
 import pandas as pd
 import pytest
 import torch
+from anndata import AnnData
 
 import novae
 from novae._constants import Keys
@@ -268,3 +269,16 @@ def test_init_prototypes():
     prototypes = model.swav_head.prototypes.data.clone()
     model.init_prototypes(adata)
     assert (model.swav_head.prototypes.data != prototypes).all()
+
+
+def test_var_name_subset():
+    adata = AnnData(np.random.rand(10, 30))
+    adata.var_names = [f"GENE{i}" for i in range(30)]
+    adata.obsm["spatial"] = np.random.randn(10, 2)
+
+    novae.spatial_neighbors(adata)
+
+    selected_genes = ["gene1"] + [f"GENE{i}" for i in range(5, 25)]  # check case insensitivity
+    model = novae.Novae(adata, var_names=selected_genes, embedding_size=10)
+
+    assert model.hparams.var_names == [x.lower() for x in selected_genes]