model_user_guide.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TZEUaYSjC6CZ"
      },
      "source": [
        "# Constructing a high-level model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "5uAmHirHtUwd"
      },
      "outputs": [],
      "source": [
        "!pip install --quiet scvi-colab\n",
        "from scvi_colab import install\n",
        "install()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tcMSn-GSxAeC",
        "outputId": "831e06d8-1d75-4eeb-a095-74da9b8bf3d3"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Global seed set to 0\n"
          ]
        }
      ],
      "source": [
        "import numpy as np\n",
        "import scvi\n",
        "from scvi import REGISTRY_KEYS\n",
        "from scvi.data import AnnDataManager\n",
        "from scvi.data.fields import (\n",
        "    LayerField,\n",
        "    CategoricalObsField,\n",
        "    NumericalObsField,\n",
        "    CategoricalJointObsField,\n",
        "    NumericalJointObsField,\n",
        ")\n",
        "import torch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J7a98f8MiykA"
      },
      "source": [
        "At this point we have covered\n",
        "\n",
        "1. Data registration via `setup_anndata` and dataloaders via `AnnDataLoader`\n",
        "2. Building a probabilistic model by subclassing `BaseModuleClass`\n",
        "\n",
        "In this tutorial, we will cover the highest-level classes in `scvi-tools`: the model classes. The main purpose of these classes (e.g., `scvi.model.SCVI`) is to wrap the actions of module instantiation, training, and subsequent posterior queries of our module into a convenient interface. These model classes are the fundamental objects driving scientific analysis of data with `scvi-tools`. Out of convention, we will refer to these objects as \"models\" and the lower-level objects presented in the previous tutorial as \"modules\".\n",
        "\n",
        "## A simple model class\n",
        "Here we will walkthrough an example of building the `scvi.model.SCVI` class. We will progressively add functionality to the class.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SUwavqvn0BBk"
      },
      "source": [
        "\n",
        "### Sketch of `BaseModelClass`\n",
        "\n",
        "Let us start by providing a high level overview of `BaseModelClass` that we will inherit. Note that this is pseudocode to provide intuition. We see that `BaseModelClass` contains some unverisally applicable methods, and some private methods (conventionally starting with `_` in Python) that will become useful after training the model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3_rabB4JoGG3"
      },
      "source": [
        "```python\n",
        "class MyModel(UnsupervisedTrainingMixin, BaseModelClass)\n",
        "\n",
        "    def __init__(self, adata):\n",
        "        # sets some basic attributes like is_trained_\n",
        "        # record the setup_dict registered in the adata\n",
        "        self.adata = adata\n",
        "        self.scvi_setup_dict_ = adata.uns[\"_scvi\"]\n",
        "        self.summary_stats = self.scvi_setup_dict_[\"summary_stats\"]\n",
        "\n",
        "    def _validate_anndata(self, adata):\n",
        "        # check that anndata is equivalent by comparing\n",
        "        # to the initial setup_dict\n",
        "\n",
        "    def _make_dataloader(adata):\n",
        "        # return a dataloader to iterate over adata\n",
        "        \n",
        "    def train(...):\n",
        "        # Universal train method provided by UnsupservisedTrainingMixin\n",
        "        # BaseModelClass does not come with train\n",
        "        # In general train methods are straightforward to compose manually\n",
        "\n",
        "    def save(...):\n",
        "        # universal save method\n",
        "        # saves modules, anndata setup dict, and attributes ending with _\n",
        "\n",
        "    def load(...):\n",
        "        # universal load method\n",
        "\n",
        "    @classmethod\n",
        "    def setup_anndata(cls, adata, ...)\n",
        "        # anndata registration method\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4vvYzZa_ukC4"
      },
      "source": [
        "### Baseline version of `SCVI` class\n",
        "\n",
        "Let's now create the simplest possible version of the `SCVI` class. We inherit the `BaseModelClass`, and write our `__init__` method.\n",
        "\n",
        "We take care to do the following:\n",
        "\n",
        "1. Set the `module` attribute to be equal to our `VAE` module, which here is the torch-level version of scVI.\n",
        "2. Add a `_model_summary_string` attr, which will be used as a representation for the model.\n",
        "3. Run `self.init_params_ = self._get_init_params(locals())`, which stores the arguments used to initialize the model, facilitating saving/loading of the model later.\n",
        "\n",
        "To initialize the `VAE`, we can use the information in `self.summary_stats`, which is information that was stored in the anndata object at `setup_anndata()` time. In this example, we have only exposed `n_latent` to users through `SCVI`. In practice, we try to expose only the most relevant parameters, as all other parameters can be accessed by passing `model_kwargs`.\n",
        "\n",
        "Finally, we write the `setup_anndata` class method, which is used to register the appropriate matrices inside the anndata that we will use to load data into the model. This method uses the `AnnDataManager` class to orchestrate the data registration process. More details about the `AnnDataManager` can be learned in our data handling tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "Q0zQzJD4jNoV"
      },
      "outputs": [],
      "source": [
        "from typing import Optional\n",
        "\n",
        "from anndata import AnnData\n",
        "from scvi.module import VAE\n",
        "from scvi.model.base import BaseModelClass, UnsupervisedTrainingMixin\n",
        "\n",
        "class SCVI(UnsupervisedTrainingMixin, BaseModelClass):\n",
        "    \"\"\"\n",
        "    single-cell Variational Inference [Lopez18]_.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(\n",
        "        self,\n",
        "        adata: AnnData,\n",
        "        n_latent: int = 10,\n",
        "        **model_kwargs,\n",
        "    ):\n",
        "        super(SCVI, self).__init__(adata)\n",
        "\n",
        "        self.module = VAE(\n",
        "            n_input=self.summary_stats[\"n_vars\"],\n",
        "            n_batch=self.summary_stats[\"n_batch\"],\n",
        "            n_latent=n_latent,\n",
        "            **model_kwargs,\n",
        "        )\n",
        "        self._model_summary_string = (\n",
        "            \"SCVI Model with the following params: \\nn_latent: {}\"\n",
        "        ).format(\n",
        "            n_latent,\n",
        "        )\n",
        "        self.init_params_ = self._get_init_params(locals())\n",
        "\n",
        "    @classmethod\n",
        "    def setup_anndata(\n",
        "        cls,\n",
        "        adata: AnnData,\n",
        "        batch_key: Optional[str] = None,\n",
        "        layer: Optional[str] = None,\n",
        "        **kwargs,\n",
        "    ) -> Optional[AnnData]:\n",
        "        setup_method_args = cls._get_setup_method_args(**locals())\n",
        "        anndata_fields = [\n",
        "            LayerField(REGISTRY_KEYS.X_KEY, layer, is_count_data=True),\n",
        "            CategoricalObsField(REGISTRY_KEYS.BATCH_KEY, batch_key),\n",
        "            # Dummy fields required for VAE class.\n",
        "            CategoricalObsField(REGISTRY_KEYS.LABELS_KEY, None),\n",
        "            NumericalObsField(\n",
        "                REGISTRY_KEYS.SIZE_FACTOR_KEY, None, required=False\n",
        "            ),\n",
        "            CategoricalJointObsField(\n",
        "                REGISTRY_KEYS.CAT_COVS_KEY, None\n",
        "            ),\n",
        "            NumericalJointObsField(\n",
        "                REGISTRY_KEYS.CONT_COVS_KEY, None\n",
        "            ),\n",
        "        ]\n",
        "        adata_manager = AnnDataManager(\n",
        "            fields=anndata_fields, setup_method_args=setup_method_args\n",
        "        )\n",
        "        adata_manager.register_fields(adata, **kwargs)\n",
        "        cls.register_manager(adata_manager)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-Qw-uCsuwhgO"
      },
      "source": [
        "Now we explore what we can and cannot do with this model. Let's get some data and initialize a `SCVI` instance. Of note, for testing purposes we like to use `scvi.data.synthetic_iid()` which returns a simple, small anndata object that was already run through `setup_anndata()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Tw3DQF9IwftC",
        "outputId": "a5485318-6e8f-413f-cbde-43b0aab7fcba"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "AnnData object with n_obs × n_vars = 400 × 100\n",
              "    obs: 'batch', 'labels'\n",
              "    uns: 'protein_names'\n",
              "    obsm: 'protein_expression'"
            ]
          },
          "execution_count": 4,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "adata = scvi.data.synthetic_iid()\n",
        "adata"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Above we saw in the `setup_anndata()` implementation that we ended the function with `cls.register_manager(adata_manager)`. This function stores the newly created `AnnDataManager` instance in a class-specific dictionary called `_setup_adata_manager_store`. Specifically, this maps from UUIDs (specific to each `AnnData` object; stored on `adata.uns[\"_scvi_uuid\"]`) to `AnnDataManager` instances instantiated by that class's `setup_anndata()` function.\n",
        "\n",
        "On model initialization, the model instance retrieves the `AnnDataManager` object specific to the passed in `adata`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 230
        },
        "id": "cAg0lsNpuCYG",
        "outputId": "553265a8-f2e9-4495-8fcb-513a8e724d7a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "adata UUID (assigned by setup_anndata): ae1459ee-966f-4f25-9ca0-48ae640be873\n",
            "AnnDataManager: <scvi.data._manager.AnnDataManager object at 0x28e73b610>\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">SCVI Model with the following params: \n",
              "n_latent: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">10</span>\n",
              "Training status: Not Trained\n",
              "</pre>\n"
            ],
            "text/plain": [
              "SCVI Model with the following params: \n",
              "n_latent: \u001b[1;36m10\u001b[0m\n",
              "Training status: Not Trained\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/plain": []
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "SCVI.setup_anndata(adata, batch_key=\"batch\")\n",
        "print(f\"adata UUID (assigned by setup_anndata): {adata.uns['_scvi_uuid']}\")\n",
        "print(f\"AnnDataManager: {SCVI._setup_adata_manager_store[adata.uns['_scvi_uuid']]}\")\n",
        "model = SCVI(adata)\n",
        "model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### More `AnnDataManager` Details"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The `AnnDataManager` class stores state on data registered with scvi-tools. Since each manager is specific to a single `AnnData`, each model instance has an `AnnDataManager` instance for every `AnnData` object it has interfaced with. In addition to `setup_anndata()`, new `AnnDataManager` objects are created via the `_validate_anndata()` method when called on new `AnnData` objects (not the `AnnData` the model instance was initialized with). `_validate_anndata()` should be called in any method that references data on the `AnnData` object, via the scvi-tools data handling strategy (e.g. `get_latent_representation()`). Any instance-specific `AnnDataManager` objects are stored in a separate class-specific manager store called `_per_instance_manager_store`, which maps model instance UUIDs (assigned on initialization) and `AnnData` UUIDs to `AnnDataManager` instances. This avoids the issue of incorrect `AnnDataManager` retrieval when working with two model instances working over the same `AnnData` object."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "model instance UUID: 7fdfe5d2-601a-492a-b9fa-9048a23c954b\n",
            "adata UUID: ae1459ee-966f-4f25-9ca0-48ae640be873\n",
            "AnnDataManager for adata: <scvi.data._manager.AnnDataManager object at 0x28e73b610>\n"
          ]
        }
      ],
      "source": [
        "print(f\"model instance UUID: {model.id}\")\n",
        "print(f\"adata UUID: {adata.uns['_scvi_uuid']}\")\n",
        "print(f\"AnnDataManager for adata: {SCVI._per_instance_manager_store[model.id][adata.uns['_scvi_uuid']]}\") # { model instance UUID: { adata UUID: AnnDataManager } }"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[34mINFO    \u001b[0m Input AnnData not setup with scvi-tools. attempting to transfer AnnData setup       \n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "AnnData object with n_obs × n_vars = 400 × 100\n",
              "    obs: 'batch', 'labels', '_scvi_batch', '_scvi_labels'\n",
              "    uns: 'protein_names', '_scvi_uuid', '_scvi_manager_uuid'\n",
              "    obsm: 'protein_expression'"
            ]
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "adata2 = scvi.data.synthetic_iid()\n",
        "model._validate_anndata(adata2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "adata2 UUID: ae1459ee-966f-4f25-9ca0-48ae640be873\n",
            "Model instance specific manager store: {'ae1459ee-966f-4f25-9ca0-48ae640be873': <scvi.data._manager.AnnDataManager object at 0x28e73b610>, '42b64083-6b32-449e-9da3-5ccaca06cf04': <scvi.data._manager.AnnDataManager object at 0x28e73b9d0>}\n"
          ]
        }
      ],
      "source": [
        "print(f\"adata2 UUID: {adata.uns['_scvi_uuid']}\")\n",
        "print(f\"Model instance specific manager store: {SCVI._per_instance_manager_store[model.id]}\") # { model instance UUID: { adata UUID: AnnDataManager } }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Additionally, the data registration process can modify or add data on the `AnnData` object directly. As a result, if calls between two models are interleaved, it is possible that we refer to fields created by another model instance's data registration incorrectly. In order to avoid this, `_validate_anndata()` additionally checks the `AnnData` object for an `AnnDataManager`-specific UUID stored in `adata.uns['_scvi_manager_uuid']`. If this UUID is inconsistent with the `AnnDataManager` fetched from the manager store, this means the data registration must be replayed on the `AnnData` object before referencing any data on the `AnnData`. This is automatically done in `_validate_anndata()`.\n",
        "\n",
        "As a result, we can interleave method calls on two model instances without worrying about this clobbering issue."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Manager UUID: 038cead6-6cc4-4fcb-93c2-845fb66eec7c\n",
            "Last setup with manager UUID: 038cead6-6cc4-4fcb-93c2-845fb66eec7c\n",
            "Encoded batch obs field: 0      0\n",
            "1      0\n",
            "2      0\n",
            "3      0\n",
            "4      0\n",
            "      ..\n",
            "395    0\n",
            "396    0\n",
            "397    0\n",
            "398    0\n",
            "399    0\n",
            "Name: _scvi_batch, Length: 400, dtype: int8\n"
          ]
        }
      ],
      "source": [
        "SCVI.setup_anndata(adata, batch_key=None) # No batch correction.\n",
        "model2 = SCVI(adata)\n",
        "print(f\"Manager UUID: {model2.adata_manager.id}\")\n",
        "print(f\"Last setup with manager UUID: {adata.uns['_scvi_manager_uuid']}\")\n",
        "print(f\"Encoded batch obs field: {adata.obs['_scvi_batch']}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Manager UUID: 0b11c46e-e008-4993-9f70-b8e521e9c05f\n",
            "Last setup with manager UUID: 0b11c46e-e008-4993-9f70-b8e521e9c05f\n",
            "Encoded batch obs field: 0      0\n",
            "1      0\n",
            "2      0\n",
            "3      0\n",
            "4      0\n",
            "      ..\n",
            "395    1\n",
            "396    1\n",
            "397    1\n",
            "398    1\n",
            "399    1\n",
            "Name: _scvi_batch, Length: 400, dtype: int8\n"
          ]
        }
      ],
      "source": [
        "model._validate_anndata(adata) # Replays registration on adata\n",
        "print(f\"Manager UUID: {model.adata_manager.id}\")\n",
        "print(f\"Last setup with manager UUID: {adata.uns['_scvi_manager_uuid']}\")\n",
        "print(f\"Encoded batch obs field: {adata.obs['_scvi_batch']}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iG1AU4cI0Fe1"
      },
      "source": [
        "### The `train` method"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "A model can be trained simply by calling the `train` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vOau9VkQxyIp",
        "outputId": "10865d7f-345b-46cb-9d15-0d15b3aafdb6"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "GPU available: False, used: False\n",
            "TPU available: False, using: 0 TPU cores\n",
            "IPU available: False, using: 0 IPUs\n",
            "/Users/jhong/miniforge3/envs/scvi-env/lib/python3.9/site-packages/pytorch_lightning/trainer/data_loading.py:432: UserWarning: The number of training samples (3) is smaller than the logging interval Trainer(log_every_n_steps=10). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.\n",
            "  rank_zero_warn(\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 20/20: 100%|██████████| 20/20 [00:00<00:00, 45.34it/s, loss=306, v_num=1]\n"
          ]
        }
      ],
      "source": [
        "model.train(max_epochs=20)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EDFKIKMdx-TO"
      },
      "source": [
        "We were able to train this model, as this method is inherited in the class. Let us now take a look at psedocode of the `train` method of `UnsupervisedTrainingMixin`. The function of each of these objects is described in the API reference.\n",
        "\n",
        "```python\n",
        "    def train(\n",
        "        self,\n",
        "        max_epochs: Optional[int] = 100,\n",
        "        use_gpu: Optional[bool] = None,\n",
        "        train_size: float = 0.9,\n",
        "        **kwargs,\n",
        "    ):\n",
        "        \"\"\"\n",
        "        Train the model.\n",
        "        \"\"\"\n",
        "        # object to make train/test/val dataloaders\n",
        "        data_splitter = DataSplitter(\n",
        "            self.adata,\n",
        "            train_size=train_size,\n",
        "            validation_size=validation_size,\n",
        "            batch_size=batch_size,\n",
        "            use_gpu=use_gpu,\n",
        "        )\n",
        "        # defines optimizers, training step, val step, logged metrics\n",
        "        training_plan = TrainingPlan(\n",
        "            self.module, len(data_splitter.train_idx),\n",
        "        )\n",
        "        # creates Trainer, pre and post training procedures (Trainer.fit())\n",
        "        runner = TrainRunner(\n",
        "            self,\n",
        "            training_plan=training_plan,\n",
        "            data_splitter=data_splitter,\n",
        "            max_epochs=max_epochs,\n",
        "            use_gpu=use_gpu,\n",
        "            **kwargs,\n",
        "        )\n",
        "        return runner()\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9fT3-7MQ0S6G"
      },
      "source": [
        "We notice two new things:\n",
        "\n",
        "1. A training plan (`training_plan`)\n",
        "2. A train runner (`runner`)\n",
        "\n",
        "The `TrainRunner` is a lightweight wrapper of the PyTorch lightning's [`Trainer`](https://pytorch-lightning.readthedocs.io/en/stable/trainer.html#trainer-class-api), which is a completely black-box method once a `TrainingPlan` is defined. So what does the `TrainingPlan` do?\n",
        "\n",
        "1. Configures optimizers (e.g., Adam), learning rate schedulers.\n",
        "2. Defines the training step, which runs a minibatch of data through the model and records the loss.\n",
        "3. Defines the validation step, same as training step, but for validation data.\n",
        "4. Records relevant metrics, such as the ELBO.\n",
        "\n",
        "In `scvi-tools` we have `scvi.lightning.TrainingPlan`, which should cover many use cases, from VAEs and VI, to MLE and MAP estimation. Developers may find that they need a custom `TrainingPlan` for e.g,. multiple optimizers and complex training scheme. These can be written and used by the model class.\n",
        "\n",
        "Developers may also overwrite this train method to add custom functionality like Early Stopping (see TOTALVI's train method). In most cases the higher-level train method can call `super().train()`, which would be the `BaseModelClass` train method."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bEqgLc8c2-Gf"
      },
      "source": [
        "### Save and load"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-3aR7N5O3BKR"
      },
      "source": [
        "We can also save and load this model object, as it follows the expected structure."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5kjJoYZ42_7h",
        "outputId": "c39de0f8-4e60-4974-d307-5e33676a9e37"
      },
      "outputs": [],
      "source": [
        "model.save(\"saved_model/\", save_anndata=True)\n",
        "model = SCVI.load(\"saved_model/\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZAGP3_R23TnF"
      },
      "source": [
        "## Writing methods to query the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QsTrpi2GCuRW"
      },
      "source": [
        "So we have a model that wraps a module that has been trained. How can we get information out of the module and present in cleanly to our users? Let's implement a simple example: getting the latent representation out of the VAE.\n",
        "\n",
        "This method has the following structure:\n",
        "\n",
        "1. Validate the user-supplied data\n",
        "2. Create a data loader\n",
        "3. Iterate over the data loader and feed into the VAE, getting the tensor of interest out of the VAE."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "PehIGudfDXRr"
      },
      "outputs": [],
      "source": [
        "from typing import Optional, Sequence\n",
        "\n",
        "@torch.no_grad()\n",
        "def get_latent_representation(\n",
        "    self,\n",
        "    adata: Optional[AnnData] = None,\n",
        "    indices: Optional[Sequence[int]] = None,\n",
        "    batch_size: Optional[int] = None,\n",
        ") -> np.ndarray:\n",
        "    r\"\"\"\n",
        "    Return the latent representation for each cell.\n",
        "\n",
        "    Parameters\n",
        "    ----------\n",
        "    adata\n",
        "        AnnData object with equivalent structure to initial AnnData. If `None`, defaults to the\n",
        "        AnnData object used to initialize the model.\n",
        "    indices\n",
        "        Indices of cells in adata to use. If `None`, all cells are used.\n",
        "    batch_size\n",
        "        Minibatch size for data loading into model. Defaults to `scvi.settings.batch_size`.\n",
        "\n",
        "    Returns\n",
        "    -------\n",
        "    latent_representation : np.ndarray\n",
        "        Low-dimensional representation for each cell\n",
        "    \"\"\"\n",
        "    if self.is_trained_ is False:\n",
        "        raise RuntimeError(\"Please train the model first.\")\n",
        "\n",
        "    adata = self._validate_anndata(adata)\n",
        "    dataloader = self._make_dataloader(adata=adata, indices=indices, batch_size=batch_size)\n",
        "    latent = []\n",
        "    for tensors in dataloader:\n",
        "        inference_inputs = self.module._get_inference_input(tensors)\n",
        "        outputs = self.module.inference(**inference_inputs)\n",
        "        qz_m = outputs[\"qz_m\"]\n",
        "\n",
        "        latent += [qz_m.cpu()]\n",
        "    return torch.cat(latent).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5hhvcl6zGGIV"
      },
      "source": [
        "<div class=\"alert alert-info\">\n",
        "\n",
        "Note\n",
        "\n",
        "Validating the anndata is critical to the user experience. If `None` is passed it just returns the anndata used to initialize the model, but if a different object is passed, it checks that this new object is equivalent in structure to the anndata passed to the model. We took great care in engineering this function so as to allow passing anndata objects with potentially missing categories (e.g., model was trained on batches `[\"A\", \"B\", \"C\"]`, but the passed anndata only has `[\"B\", \"C\"]`). These sorts of checks will ensure that your module will see data that it expects, and the user will get the results they expect without advanced data manipulations.\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h4a-4Ssmu6qk"
      },
      "source": [
        "As a convention, we like to keep the module code as bare as possible and leave all posterior manipulation of module tensors to the model class methods. However, it would have been possible to write a `get_z` method in the module, and just have the model class that method."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wmsozWgrDObo"
      },
      "source": [
        "## Mixing in pre-coded features\n",
        "\n",
        "We have a number of Mixin classes that can add functionality to your model through inheritance. Here we demonstrate the [`VAEMixin`](https://www.scvi-tools.org/en/stable/api/reference/scvi.model.base.VAEMixin.html#scvi.model.base.VAEMixin) class.\n",
        "\n",
        "Let's try to get the latent representation from the object we already created."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KntvOQekxq0K",
        "outputId": "f47f524e-9b42-413e-8a4f-b06067980539"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "This function does not exist\n"
          ]
        }
      ],
      "source": [
        "try:\n",
        "    model.get_latent_representation()\n",
        "except AttributeError:\n",
        "    print(\"This function does not exist\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FPf4xHI_L-oI"
      },
      "source": [
        "This method becomes avaialble once the `VAEMixin` is inherited. Here's an overview of the mixin methods, which are coded generally enough that they should be broadly useful to those building VAEs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QD3HWo2lLKSC"
      },
      "source": [
        "```python\n",
        "class VAEMixin:\n",
        "    @torch.no_grad()\n",
        "    def get_elbo(\n",
        "        self,\n",
        "        adata: Optional[AnnData] = None,\n",
        "        indices: Optional[Sequence[int]] = None,\n",
        "        batch_size: Optional[int] = None,\n",
        "    ) -> float:\n",
        "        pass\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def get_marginal_ll(\n",
        "        self,\n",
        "        adata: Optional[AnnData] = None,\n",
        "        indices: Optional[Sequence[int]] = None,\n",
        "        n_mc_samples: int = 1000,\n",
        "        batch_size: Optional[int] = None,\n",
        "    ) -> float:\n",
        "        pass\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def get_reconstruction_error(\n",
        "        self,\n",
        "        adata: Optional[AnnData] = None,\n",
        "        indices: Optional[Sequence[int]] = None,\n",
        "        batch_size: Optional[int] = None,\n",
        "    ) -> Union[float, Dict[str, float]]:\n",
        "        pass\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def get_latent_representation(\n",
        "        self,\n",
        "        adata: Optional[AnnData] = None,\n",
        "        indices: Optional[Sequence[int]] = None,\n",
        "        give_mean: bool = True,\n",
        "        mc_samples: int = 5000,\n",
        "        batch_size: Optional[int] = None,\n",
        "    ) -> np.ndarray:\n",
        "        pass\n",
        "\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nA2ES0nANQN8"
      },
      "source": [
        "Let's now inherit the mixin into our SCVI class."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "Lu4Rb9nw2217"
      },
      "outputs": [],
      "source": [
        "from scvi.model.base import VAEMixin, UnsupervisedTrainingMixin\n",
        "\n",
        "class SCVI(VAEMixin, UnsupervisedTrainingMixin, BaseModelClass):\n",
        "    \"\"\"\n",
        "    single-cell Variational Inference [Lopez18]_.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(\n",
        "        self,\n",
        "        adata: AnnData,\n",
        "        n_latent: int = 10,\n",
        "        **model_kwargs,\n",
        "    ):\n",
        "        super(SCVI, self).__init__(adata)\n",
        "\n",
        "        self.module = VAE(\n",
        "            n_input=self.summary_stats[\"n_vars\"],\n",
        "            n_batch=self.summary_stats[\"n_batch\"],\n",
        "            n_latent=n_latent,\n",
        "            **model_kwargs,\n",
        "        )\n",
        "        self._model_summary_string = (\n",
        "            \"SCVI Model with the following params: \\nn_latent: {}\"\n",
        "        ).format(\n",
        "            n_latent,\n",
        "        )\n",
        "        self.init_params_ = self._get_init_params(locals())\n",
        "\n",
        "\n",
        "    @classmethod\n",
        "    def setup_anndata(\n",
        "        cls,\n",
        "        adata: AnnData,\n",
        "        batch_key: Optional[str] = None,\n",
        "        layer: Optional[str] = None,\n",
        "        **kwargs,\n",
        "    ) -> Optional[AnnData]:\n",
        "        setup_method_args = cls._get_setup_method_args(**locals())\n",
        "        anndata_fields = [\n",
        "            LayerField(REGISTRY_KEYS.X_KEY, layer, is_count_data=True),\n",
        "            CategoricalObsField(REGISTRY_KEYS.BATCH_KEY, batch_key),\n",
        "            # Dummy fields required for VAE class.\n",
        "            CategoricalObsField(REGISTRY_KEYS.LABELS_KEY, None),\n",
        "            NumericalObsField(\n",
        "                REGISTRY_KEYS.SIZE_FACTOR_KEY, None, required=False\n",
        "            ),\n",
        "            CategoricalJointObsField(\n",
        "                REGISTRY_KEYS.CAT_COVS_KEY, None\n",
        "            ),\n",
        "            NumericalJointObsField(\n",
        "                REGISTRY_KEYS.CONT_COVS_KEY, None\n",
        "            ),\n",
        "        ]\n",
        "        adata_manager = AnnDataManager(\n",
        "            fields=anndata_fields, setup_method_args=setup_method_args\n",
        "        )\n",
        "        adata_manager.register_fields(adata, **kwargs)\n",
        "        cls.register_manager(adata_manager)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8erM-DwwNV9V",
        "outputId": "7ca877a6-2767-4873-df30-e422135e02d1"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "GPU available: False, used: False\n",
            "TPU available: False, using: 0 TPU cores\n",
            "IPU available: False, using: 0 IPUs\n",
            "/Users/jhong/miniforge3/envs/scvi-env/lib/python3.9/site-packages/pytorch_lightning/trainer/data_loading.py:432: UserWarning: The number of training samples (3) is smaller than the logging interval Trainer(log_every_n_steps=10). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.\n",
            "  rank_zero_warn(\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 10/10: 100%|██████████| 10/10 [00:00<00:00, 36.06it/s, loss=318, v_num=1]\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "array([[-0.7370485 ,  0.28592116,  0.24874236, ...,  0.7185316 ,\n",
              "        -0.6551055 ,  1.2049382 ],\n",
              "       [-1.1715283 ,  1.1206872 ,  0.05082523, ...,  0.05453417,\n",
              "        -0.8783005 ,  0.718166  ],\n",
              "       [-0.64500284, -0.5552139 ,  0.05247398, ...,  0.84965473,\n",
              "        -0.75963885,  1.4641552 ],\n",
              "       ...,\n",
              "       [-0.43021405, -0.01531297,  0.7627437 , ...,  0.4630339 ,\n",
              "        -0.02551817,  1.3263805 ],\n",
              "       [-1.1203899 ,  0.6630224 , -0.43606877, ...,  1.2753956 ,\n",
              "        -0.5158905 ,  1.1000206 ],\n",
              "       [-0.6726756 ,  0.73773426, -0.12030821, ...,  0.4862482 ,\n",
              "        -0.4421261 ,  1.7749897 ]], dtype=float32)"
            ]
          },
          "execution_count": 16,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "SCVI.setup_anndata(adata, batch_key=\"batch\")\n",
        "model = SCVI(adata)\n",
        "model.train(10)\n",
        "model.get_latent_representation()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qWvfa7q9NggC"
      },
      "source": [
        "## Summary"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "59_79HYQNjv_"
      },
      "source": [
        "We learned the structure of the high-level model classes in scvi-tools, and learned how a simple version of `SCVI` is implemented.\n",
        "\n",
        "Questions? Comments? Keep the discussion going on our [forum](https://discourse.scvi-tools.org/)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "model_user_guide.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}