dev/data_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "critical-vitamin",
   "metadata": {
    "id": "critical-vitamin"
   },
   "source": [
    "# Data handling in scvi-tools\n",
    "\n",
    "In this tutorial we will cover how data is handled in scvi-tools.\n",
    "\n",
    "Sections:\n",
    "\n",
    "1. Introduction to the `registry` comprised of `data_registry`, `state_registry`, and `summary_stats`.\n",
    "1. Explanation of `AnnDataField` classes and how they populate the `registry` via the `AnnDataManager`.\n",
    "1. Data loading with `AnnDataLoader()` outside of scvi-tools models.\n",
    "1. Writing a `setup_anndata()` function for an scvi-tools model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36e2bc80",
   "metadata": {},
   "source": [
    "Uncomment the following lines in Google Colab in order to install `scvi-tools`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "22c58cb8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:38.613882Z",
     "iopub.status.busy": "2023-12-07T00:51:38.613741Z",
     "iopub.status.idle": "2023-12-07T00:51:38.616003Z",
     "shell.execute_reply": "2023-12-07T00:51:38.615618Z"
    }
   },
   "outputs": [],
   "source": [
    "# !pip install --quiet scvi-colab\n",
    "# from scvi_colab import install\n",
    "\n",
    "# install()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "formed-cleaners",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:38.617461Z",
     "iopub.status.busy": "2023-12-07T00:51:38.617359Z",
     "iopub.status.idle": "2023-12-07T00:51:41.501175Z",
     "shell.execute_reply": "2023-12-07T00:51:41.500679Z"
    },
    "id": "formed-cleaners",
    "outputId": "2767b15d-c984-4fe7-e5ca-667d398f83cf"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scvi\n",
    "import torch.nn\n",
    "from scvi.data import AnnDataManager\n",
    "from scvi.data.fields import CategoricalObsField, LayerField\n",
    "from scvi.dataloaders import AnnDataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8d127d00",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.503211Z",
     "iopub.status.busy": "2023-12-07T00:51:41.502839Z",
     "iopub.status.idle": "2023-12-07T00:51:41.506537Z",
     "shell.execute_reply": "2023-12-07T00:51:41.506121Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Seed set to 0\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Last run with scvi-tools version: 1.1.0\n"
     ]
    }
   ],
   "source": [
    "scvi.settings.seed = 0\n",
    "print(\"Last run with scvi-tools version:\", scvi.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "physical-johnston",
   "metadata": {
    "id": "physical-johnston"
   },
   "source": [
    "## Recording AnnData state with object registration\n",
    "\n",
    "Scvi-tools knows what subset of AnnData to load into models during training/inference via a data registration process handled by `setup_anndata()`.\n",
    "\n",
    "This setup process is orchestrated by an `AnnDataManager` object which wraps the `AnnData` object and creates a corresponding `registry`.\n",
    "\n",
    "In this section we enumerate the fields in the `registry` object. The registry takes the form of a nested dictionary and is stored as an instance variable of an `AnnDataManager` object, `adata_manager.registry`.\n",
    "\n",
    "The top level of the registry contains the following keys:\n",
    "\n",
    "- `scvi_version` keeps track of the version of scvi-tools used to setup the AnnData Object.\n",
    "- `model_name` and `setup_args` keep track of the model and arguments used to run `setup_anndata()`. These fields are optional, since the `AnnDataManager` can also be created outside of a `setup_anndata()` function.\n",
    "- `field_registries` is dictionary which maps registry keys (e.g. `batch`, `labels`) to additional field-specific information.\n",
    "\n",
    "Within each field registry, there the following three keys:\n",
    "\n",
    "- `data_registry` contains the location of data to load. This is what is used by the DataLoaders to iterate over the AnnData.\n",
    "- `state_registry` contains any state (e.g. categorical mappings for batch) relevant to the field during `register_field()`.\n",
    "- `summary_stats` contains summary statistics relevant to the field."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2baae67b",
   "metadata": {
    "id": "2baae67b"
   },
   "source": [
    "Here we construct an `AnnDataManager` and create a `registry` via `register_fields`. In the next section, we will breakdown how the `registry` is populated as a function of the `AnnDataFields`. We can visualize the contents of the `registry` via the function `view_registry()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "entertaining-degree",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.508061Z",
     "iopub.status.busy": "2023-12-07T00:51:41.507952Z",
     "iopub.status.idle": "2023-12-07T00:51:41.661303Z",
     "shell.execute_reply": "2023-12-07T00:51:41.660863Z"
    },
    "id": "entertaining-degree",
    "outputId": "ed94367c-6a12-4f19-b70a-1c7d79bcbf73"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['scvi_version', 'model_name', 'setup_args', 'field_registries', '_scvi_uuid'])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<frozen abc>:119: FutureWarning: SparseDataset is deprecated and will be removed in late 2024. It has been replaced by the public classes CSRDataset and CSCDataset.\n",
      "\n",
      "For instance checks, use `isinstance(X, (anndata.experimental.CSRDataset, anndata.experimental.CSCDataset))` instead.\n",
      "\n",
      "For creation, use `anndata.experimental.sparse_dataset(X)` instead.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "adata = scvi.data.synthetic_iid()\n",
    "\n",
    "anndata_fields = [\n",
    "    LayerField(registry_key=\"x\", layer=None, is_count_data=True),\n",
    "    CategoricalObsField(registry_key=\"batch\", attr_key=\"batch\"),\n",
    "]\n",
    "adata_manager = AnnDataManager(fields=anndata_fields)\n",
    "adata_manager.register_fields(adata)\n",
    "print(\n",
    "    adata_manager.registry.keys()\n",
    ")  # There is additionally a _scvi_uuid key which is used to uniquely identify AnnData objects for subsequent retrieval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a39faeb8",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 385
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.662896Z",
     "iopub.status.busy": "2023-12-07T00:51:41.662776Z",
     "iopub.status.idle": "2023-12-07T00:51:41.674077Z",
     "shell.execute_reply": "2023-12-07T00:51:41.673739Z"
    },
    "id": "a39faeb8",
    "outputId": "6873c6e8-4ac3-4609-f062-ce43d954cd8b"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Anndata setup with scvi-tools version <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.1</span>.<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>.\n",
       "</pre>\n"
      ],
      "text/plain": [
       "Anndata setup with scvi-tools version \u001b[1;36m1.1\u001b[0m.\u001b[1;36m0\u001b[0m.\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">     Summary Statistics     </span>\n",
       "┏━━━━━━━━━━━━━━━━━━┳━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> Summary Stat Key </span>┃<span style=\"font-weight: bold\"> Value </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━╇━━━━━━━┩\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">     n_batch      </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">   2   </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">       n_x        </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">  100  </span>│\n",
       "└──────────────────┴───────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[3m     Summary Statistics     \u001b[0m\n",
       "┏━━━━━━━━━━━━━━━━━━┳━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mSummary Stat Key\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mValue\u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━╇━━━━━━━┩\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m    n_batch     \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m  2  \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m      n_x       \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m 100 \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "└──────────────────┴───────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">               Data Registry               </span>\n",
       "┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> Registry Key </span>┃<span style=\"font-weight: bold\">   scvi-tools Location    </span>┃\n",
       "┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">    batch     </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\"> adata.obs['_scvi_batch'] </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">      x       </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">         adata.X          </span>│\n",
       "└──────────────┴──────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[3m               Data Registry               \u001b[0m\n",
       "┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mRegistry Key\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m  scvi-tools Location   \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m   batch    \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128madata.obs['_scvi_batch']\u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m     x      \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m        adata.X         \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "└──────────────┴──────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">                  batch State Registry                   </span>\n",
       "┏━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">  Source Location   </span>┃<span style=\"font-weight: bold\"> Categories </span>┃<span style=\"font-weight: bold\"> scvi-tools Encoding </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\"> adata.obs['batch'] </span>│<span style=\"color: #008000; text-decoration-color: #008000\">  batch_0   </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">          0          </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">                    </span>│<span style=\"color: #008000; text-decoration-color: #008000\">  batch_1   </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">          1          </span>│\n",
       "└────────────────────┴────────────┴─────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[3m                  batch State Registry                   \u001b[0m\n",
       "┏━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m Source Location  \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mCategories\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mscvi-tools Encoding\u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33madata.obs['batch']\u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m batch_0  \u001b[0m\u001b[32m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m         0         \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m                  \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m batch_1  \u001b[0m\u001b[32m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m         1         \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "└────────────────────┴────────────┴─────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "adata_manager.view_registry()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2698e80d",
   "metadata": {
    "id": "2698e80d"
   },
   "source": [
    "The above summary incorporates all three of the components making up each field registry as mentioned before."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "delayed-spirituality",
   "metadata": {
    "id": "delayed-spirituality"
   },
   "source": [
    "### Data Registry\n",
    "\n",
    "First, lets turn our attention to the `data_registry`.\n",
    "\n",
    "This is used by the `AnnDataLoader` during training to\n",
    "\n",
    "- Access the correct \"slots\" of AnnData\n",
    "- Minibatch the data, while optionally densifying sparse-formatted data\n",
    "\n",
    "Each key of the `data_registry` is the name of tensor and is used to retreive the data from the dataloader output.\n",
    "\n",
    "- All the data registered via `register_fields()` have registry keys associated with them, as defined in the `AnnDataField` class.\n",
    "\n",
    "The value of each key in the `data_registry` is a dictionary with two keys: `attr_name` and `attr_key`.\n",
    "\n",
    "- `attr_name` is the attribute of `adata` to load data from eg. `obs`, `obsm`, `layers`.\n",
    "- `attr_key` is the key of the attribute to access the data.\n",
    "\n",
    "For example, based off the following data_registry, batch information is loaded from `adata.obs['_scvi_batch']` and will be accessible via `batch`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27f489f6",
   "metadata": {
    "id": "27f489f6"
   },
   "source": [
    "While the data registry dictionary is stored within the `registry`, the `AnnDataManager` provides a helper method, `adata_manager.data_registry`, which coalesces the full data registry across each of the fields. This helper method additionally wraps the dictionary in a custom `attrdict` class which allows dictionary access via dot notation (e.g. `data_registry.batch.attr_name`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "suspected-bennett",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.675591Z",
     "iopub.status.busy": "2023-12-07T00:51:41.675478Z",
     "iopub.status.idle": "2023-12-07T00:51:41.678692Z",
     "shell.execute_reply": "2023-12-07T00:51:41.678420Z"
    },
    "id": "suspected-bennett",
    "outputId": "c753f115-069f-458b-d965-1a3c36830053"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "batch:\n",
       "  attr_key: _scvi_batch\n",
       "  attr_name: obs\n",
       "x:\n",
       "  attr_key: null\n",
       "  attr_name: X"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_registry = adata_manager.data_registry\n",
    "data_registry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "rapid-monday",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.680138Z",
     "iopub.status.busy": "2023-12-07T00:51:41.679932Z",
     "iopub.status.idle": "2023-12-07T00:51:41.682061Z",
     "shell.execute_reply": "2023-12-07T00:51:41.681749Z"
    },
    "id": "rapid-monday",
    "outputId": "3ed2fabd-66be-4aeb-8298-9645ad2cc104"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "attr_key: _scvi_batch\n",
      "attr_name: obs\n",
      "\n",
      "_scvi_batch\n"
     ]
    }
   ],
   "source": [
    "print(data_registry[\"batch\"])\n",
    "print(data_registry.batch.attr_key)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b15123f7",
   "metadata": {
    "id": "b15123f7"
   },
   "source": [
    "### State Registries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "immediate-wonder",
   "metadata": {
    "id": "immediate-wonder"
   },
   "source": [
    "During the data registration process, we also keep track of additional information from the registration process, necessary for model initialization or downstream functionality. For example, for the batch field, scvi-tools keeps track of the location of the original data as well as the categorical to integer mappings.\n",
    "\n",
    "The batch state registry holds the following two keys:\n",
    "\n",
    "- `original_key` is the original key passed in by the user to load the data.\n",
    "- `categorical_mapping` is the categorical to integer mapping of the data. The index of the category is its corresponding integer representation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51aff770",
   "metadata": {
    "id": "51aff770"
   },
   "source": [
    "We can access a state registry via the function `AnnDataManager.get_state_registry()` which takes a registry key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "executed-vampire",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.683577Z",
     "iopub.status.busy": "2023-12-07T00:51:41.683385Z",
     "iopub.status.idle": "2023-12-07T00:51:41.685586Z",
     "shell.execute_reply": "2023-12-07T00:51:41.685316Z"
    },
    "id": "executed-vampire",
    "outputId": "ff495511-1ba9-4796-abc3-dd2f64d84be8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['categorical_mapping', 'original_key']\n",
      "Categorical mapping: ['batch_0' 'batch_1']\n",
      "Original key: batch\n"
     ]
    }
   ],
   "source": [
    "batch_state_registry = adata_manager.get_state_registry(\"batch\")\n",
    "print(batch_state_registry.keys())\n",
    "\n",
    "print(f\"Categorical mapping: {batch_state_registry.categorical_mapping}\")\n",
    "print(f\"Original key: {batch_state_registry.original_key}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efce512d",
   "metadata": {
    "id": "efce512d"
   },
   "source": [
    "### Summary Stats"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5784b4fe",
   "metadata": {
    "id": "5784b4fe"
   },
   "source": [
    "Lastly, we have the summary stats dictionary which is a dictionary meant to store summary statistics frequently used in models, to avoid redundancy and for summarization in `view_registry()`. Like the other two components, the `AnnDataManager` has a helper method in the form of the property `adata_manager.summary_stats`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1d95bffe",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.687124Z",
     "iopub.status.busy": "2023-12-07T00:51:41.686914Z",
     "iopub.status.idle": "2023-12-07T00:51:41.689556Z",
     "shell.execute_reply": "2023-12-07T00:51:41.689295Z"
    },
    "id": "1d95bffe",
    "outputId": "7ec08c87-c7e4-44de-b2f9-87408b852991"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "n_batch: 2\n",
       "n_x: 100"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata_manager.summary_stats"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1e8f0af",
   "metadata": {
    "id": "c1e8f0af"
   },
   "source": [
    "## AnnDataManager and AnnDataFields"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed6440c3",
   "metadata": {
    "id": "ed6440c3"
   },
   "source": [
    "Now that we have gone over the registered state of an `AnnDataManager`, we can go over how the underlying logic is organized.\n",
    "\n",
    "While the `AnnDataManager` provides the main interface to the data registration components, the logic specific to each field is encapsulated in `AnnDataField` classes (any child class of `BaseAnnDataField`).\n",
    "\n",
    "An `AnnDataField` class contains four main functions to be implemented:\n",
    "\n",
    "1. `register_field` sets up the relevant field on the AnnData object and returns the state registry for this field.\n",
    "1. `validate_field` is a function called before `register_field`. E.g. checks if the data field is present on the AnnData object.\n",
    "1. `transfer_field` is a function similar to `register_field`, but additionally takes a source `state_registry` which can modify the behavior of registration. E.g. for categorical fields we may want to maintain the source categories and append any additional categories on the target AnnData object for downstream transfer learning.\n",
    "1. `get_summary_stats` is a function that takes a `state_registry` and outputs the summary stat dictionary. Note, this means the summary statistics must be a function of what is stored in `state_registry`.\n",
    "\n",
    "Together, the set of `AnnDataField`s produces the `registry` detailed in part 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2bc8138d",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.691063Z",
     "iopub.status.busy": "2023-12-07T00:51:41.690862Z",
     "iopub.status.idle": "2023-12-07T00:51:41.692694Z",
     "shell.execute_reply": "2023-12-07T00:51:41.692427Z"
    },
    "id": "2bc8138d",
    "outputId": "5425c408-71c6-4a7b-c82d-8642bbd852d1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[<scvi.data.fields._layer_field.LayerField object at 0x7fdd032d9ed0>, <scvi.data.fields._dataframe_field.CategoricalObsField object at 0x7fde78add450>]\n"
     ]
    }
   ],
   "source": [
    "print(adata_manager.fields)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "39205485",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.694098Z",
     "iopub.status.busy": "2023-12-07T00:51:41.693908Z",
     "iopub.status.idle": "2023-12-07T00:51:41.710696Z",
     "shell.execute_reply": "2023-12-07T00:51:41.710424Z"
    },
    "id": "39205485",
    "outputId": "823816ec-bf19-45a9-8871-aea83c90317e"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before register_field:\n",
      "AnnData object with n_obs × n_vars = 400 × 100\n",
      "    obs: 'batch', 'labels'\n",
      "    uns: 'protein_names'\n",
      "    obsm: 'protein_expression', 'accessibility'\n",
      "\n",
      "After register_field:\n",
      "AnnData object with n_obs × n_vars = 400 × 100\n",
      "    obs: 'batch', 'labels', '_scvi_batch'\n",
      "    uns: 'protein_names'\n",
      "    obsm: 'protein_expression', 'accessibility'\n",
      "\n",
      "State registry: {'categorical_mapping': array(['batch_0', 'batch_1'], dtype=object), 'original_key': 'batch'}\n"
     ]
    }
   ],
   "source": [
    "adata2 = scvi.data.synthetic_iid()\n",
    "batch_field = adata_manager.fields[1]\n",
    "print(\"Before register_field:\")\n",
    "print(adata2)\n",
    "print()\n",
    "\n",
    "batch_state_registry = batch_field.register_field(adata2)\n",
    "print(\"After register_field:\")\n",
    "print(adata2)\n",
    "print()\n",
    "print(f\"State registry: {batch_state_registry}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3409c0d",
   "metadata": {
    "id": "b3409c0d"
   },
   "source": [
    "Notice how in this case the batch field's `register_field()` function adds an additional `.obs` field that stores an encoded version of the `adata.obs['batch']` column. The categorical mapping array order corresponds to the integer encoding of the category."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "Lt8fQpMfiM0s",
   "metadata": {
    "id": "Lt8fQpMfiM0s"
   },
   "source": [
    "```{important}\n",
    "Adding data inplace to an AnnData object should be done with care. Most fields do not add any data; however, for categorical data it is faster to encode the categories inplace instead of on-the-fly during data loading.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5e27699",
   "metadata": {},
   "source": [
    "`adata_manager.transfer_fields()` can be used to produce a new `AnnDataManager` for a target `AnnData` object that follows the same structure as the original `AnnData`. This can be useful in models that are trained with an `AnnData` object, then used to make predictions on new query data. Under the hood, this calls the `transfer_field()` function of each field in the `AnnDataManager`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c409de8",
   "metadata": {},
   "source": [
    "In the example below, we try to transfer the batch field onto a new `AnnData` object with an extra batch category. In this case, the `CategoricalObsField.transfer_field()` function is parameterized with a `extend_categories` kwarg which, when `True`, will extend the batch category set as necessary. If marked `extend_categories=False`, `transfer_field()` will raise an error, which may be desired behavior in cases where we want to ensure that query data does not contain any batch categories missing from the train data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9563a62d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.712263Z",
     "iopub.status.busy": "2023-12-07T00:51:41.712155Z",
     "iopub.status.idle": "2023-12-07T00:51:41.734816Z",
     "shell.execute_reply": "2023-12-07T00:51:41.734418Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Category batch_2 not found in source registry. Cannot transfer setup without `extend_categories = True`.\n"
     ]
    }
   ],
   "source": [
    "adata3 = scvi.data.synthetic_iid(n_batches=3)\n",
    "try:\n",
    "    adata_manager.transfer_fields(adata_target=adata3)\n",
    "except ValueError as e:\n",
    "    print(\n",
    "        e\n",
    "    )  # Errors due to additional batch_2 category absent from adata_manager.adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "55d4cea3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.736193Z",
     "iopub.status.busy": "2023-12-07T00:51:41.736086Z",
     "iopub.status.idle": "2023-12-07T00:51:41.744260Z",
     "shell.execute_reply": "2023-12-07T00:51:41.743998Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Anndata setup with scvi-tools version <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.1</span>.<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>.\n",
       "</pre>\n"
      ],
      "text/plain": [
       "Anndata setup with scvi-tools version \u001b[1;36m1.1\u001b[0m.\u001b[1;36m0\u001b[0m.\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">     Summary Statistics     </span>\n",
       "┏━━━━━━━━━━━━━━━━━━┳━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> Summary Stat Key </span>┃<span style=\"font-weight: bold\"> Value </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━╇━━━━━━━┩\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">     n_batch      </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">   3   </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">       n_x        </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">  100  </span>│\n",
       "└──────────────────┴───────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[3m     Summary Statistics     \u001b[0m\n",
       "┏━━━━━━━━━━━━━━━━━━┳━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mSummary Stat Key\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mValue\u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━╇━━━━━━━┩\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m    n_batch     \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m  3  \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m      n_x       \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m 100 \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "└──────────────────┴───────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">               Data Registry               </span>\n",
       "┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> Registry Key </span>┃<span style=\"font-weight: bold\">   scvi-tools Location    </span>┃\n",
       "┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">    batch     </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\"> adata.obs['_scvi_batch'] </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">      x       </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">         adata.X          </span>│\n",
       "└──────────────┴──────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[3m               Data Registry               \u001b[0m\n",
       "┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mRegistry Key\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m  scvi-tools Location   \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m   batch    \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128madata.obs['_scvi_batch']\u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m     x      \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m        adata.X         \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "└──────────────┴──────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">                  batch State Registry                   </span>\n",
       "┏━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">  Source Location   </span>┃<span style=\"font-weight: bold\"> Categories </span>┃<span style=\"font-weight: bold\"> scvi-tools Encoding </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\"> adata.obs['batch'] </span>│<span style=\"color: #008000; text-decoration-color: #008000\">  batch_0   </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">          0          </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">                    </span>│<span style=\"color: #008000; text-decoration-color: #008000\">  batch_1   </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">          1          </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">                    </span>│<span style=\"color: #008000; text-decoration-color: #008000\">  batch_2   </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">          2          </span>│\n",
       "└────────────────────┴────────────┴─────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[3m                  batch State Registry                   \u001b[0m\n",
       "┏━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m Source Location  \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mCategories\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mscvi-tools Encoding\u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33madata.obs['batch']\u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m batch_0  \u001b[0m\u001b[32m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m         0         \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m                  \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m batch_1  \u001b[0m\u001b[32m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m         1         \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m                  \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m batch_2  \u001b[0m\u001b[32m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m         2         \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "└────────────────────┴────────────┴─────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "adata3_manager = adata_manager.transfer_fields(\n",
    "    adata_target=adata3, extend_categories=True\n",
    ")\n",
    "adata3_manager.view_registry()  # batch_2 appended to the batch category mapping"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "timely-practitioner",
   "metadata": {
    "id": "timely-practitioner"
   },
   "source": [
    "## DataLoaders\n",
    "\n",
    "`AnnDataLoader` is the base dataloader for scvi-tools. In this section we show how the data registered is loaded by `AnnDataLoader`.\n",
    "\n",
    "Parameters of `AnnDataLoader`:\n",
    "\n",
    "- `adata_manager`: `AnnDataManager` object to load data from.\n",
    "- `shuffle`: if True will shuffle the data beforehand.\n",
    "- `indices`: can provide a subset of indices to load from (Useful when doing train/test splits).\n",
    "- `data_and_attributes`: a dictionary where the key corresponds to its key in the `data_registry` and the value is the numpy data type. By default, all data is passed to the model as `np.float32`.\n",
    "- `data_loader_kwargs`: additional arguments from `torch.utils.data.DataLoader`.\n",
    "\n",
    "First, we construct an `AnnDataLoader` and get the first batch. Then we will enumerate all the values in the batch. The variable **data_batch** contains the first batch of data. It is a dictionary whose values are the tensors registered in the previous section via `register_fields()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "incident-brighton",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.745738Z",
     "iopub.status.busy": "2023-12-07T00:51:41.745610Z",
     "iopub.status.idle": "2023-12-07T00:51:41.748736Z",
     "shell.execute_reply": "2023-12-07T00:51:41.748472Z"
    },
    "id": "incident-brighton"
   },
   "outputs": [],
   "source": [
    "# initialize an AnnDataLoader which will iterate over our anndata\n",
    "adl = AnnDataLoader(adata_manager, shuffle=False, batch_size=10)\n",
    "\n",
    "# get the first batch of data\n",
    "data_batch = next(iter(adl))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "governmental-emergency",
   "metadata": {
    "id": "governmental-emergency"
   },
   "source": [
    "Notice that the keys in **data_batch** are the same as the keys in the `data_registry`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "incident-headline",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.750192Z",
     "iopub.status.busy": "2023-12-07T00:51:41.750085Z",
     "iopub.status.idle": "2023-12-07T00:51:41.751957Z",
     "shell.execute_reply": "2023-12-07T00:51:41.751691Z"
    },
    "id": "incident-headline",
    "outputId": "6be0a5b0-57c3-4119-85ce-ad7143c8680a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data_batch_keys:\n",
      "dict_keys(['batch', 'x'])\n"
     ]
    }
   ],
   "source": [
    "print(\"data_batch_keys:\")\n",
    "print(data_batch.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "nonprofit-grass",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.753378Z",
     "iopub.status.busy": "2023-12-07T00:51:41.753174Z",
     "iopub.status.idle": "2023-12-07T00:51:41.755506Z",
     "shell.execute_reply": "2023-12-07T00:51:41.755245Z"
    },
    "id": "nonprofit-grass",
    "outputId": "627dc7fc-81f7-43b5-b673-20c3bd4d2c15"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['batch', 'x']"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata_manager.data_registry.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "metallic-equivalent",
   "metadata": {
    "id": "metallic-equivalent"
   },
   "source": [
    "If we look at the labels for the first batch from the data loader, it corresponds to the labels of the first 10 cells of our AnnData."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "adolescent-moment",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.756980Z",
     "iopub.status.busy": "2023-12-07T00:51:41.756790Z",
     "iopub.status.idle": "2023-12-07T00:51:41.759632Z",
     "shell.execute_reply": "2023-12-07T00:51:41.759375Z"
    },
    "id": "adolescent-moment",
    "outputId": "e9e231ea-8f88-416b-f2a3-e840a32f6913"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0    batch_0\n",
       "1    batch_0\n",
       "2    batch_0\n",
       "3    batch_0\n",
       "4    batch_0\n",
       "5    batch_0\n",
       "6    batch_0\n",
       "7    batch_0\n",
       "8    batch_0\n",
       "9    batch_0\n",
       "Name: batch, dtype: category\n",
       "Categories (2, object): ['batch_0', 'batch_1']"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata.obs[\"batch\"][:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "facial-price",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.761013Z",
     "iopub.status.busy": "2023-12-07T00:51:41.760840Z",
     "iopub.status.idle": "2023-12-07T00:51:41.763280Z",
     "shell.execute_reply": "2023-12-07T00:51:41.763021Z"
    },
    "id": "facial-price",
    "outputId": "f9a45ad8-7912-4211-ad35-2e95e9463ab0"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0],\n",
       "        [0],\n",
       "        [0],\n",
       "        [0],\n",
       "        [0],\n",
       "        [0],\n",
       "        [0],\n",
       "        [0],\n",
       "        [0],\n",
       "        [0]])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# CategoricalObsField.register_field() automatically encoded the categorical labels as integers\n",
    "data_batch[\"batch\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "verified-custom",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.764652Z",
     "iopub.status.busy": "2023-12-07T00:51:41.764470Z",
     "iopub.status.idle": "2023-12-07T00:51:41.766398Z",
     "shell.execute_reply": "2023-12-07T00:51:41.766127Z"
    },
    "id": "verified-custom",
    "outputId": "c9821fd6-afd3-4c84-abaa-6431d53f13fc"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10, 100])\n",
      "torch.Size([10, 1])\n"
     ]
    }
   ],
   "source": [
    "print(data_batch[\"x\"].shape)  # shape is batch_size x n_genes\n",
    "print(data_batch[\"batch\"].shape)  # shape is batch_size x 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "geological-collector",
   "metadata": {
    "id": "geological-collector"
   },
   "source": [
    "By default, all the data loaded in scvi-tools is `np.float32`. If you wish to load as a different datatype, you can pass in a dictionary where the key corresponds to a key in the data registry and the value is the datatype.\n",
    "\n",
    "In the following snippet, we load some continuous data as `np.float64` and integer data as `np.long32`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "architectural-embassy",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.767805Z",
     "iopub.status.busy": "2023-12-07T00:51:41.767629Z",
     "iopub.status.idle": "2023-12-07T00:51:41.770498Z",
     "shell.execute_reply": "2023-12-07T00:51:41.770226Z"
    },
    "id": "architectural-embassy",
    "outputId": "df8d51fc-71a3-4d89-b8b6-35e56334a31a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.float32\n",
      "torch.int64\n"
     ]
    }
   ],
   "source": [
    "adl = AnnDataLoader(adata_manager, shuffle=False, batch_size=10)\n",
    "data_batch = next(iter(adl))\n",
    "\n",
    "# by default data has the dtype np.float32\n",
    "print(data_batch[\"x\"].dtype)\n",
    "print(data_batch[\"batch\"].dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "hidden-sudan",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.771881Z",
     "iopub.status.busy": "2023-12-07T00:51:41.771703Z",
     "iopub.status.idle": "2023-12-07T00:51:41.773822Z",
     "shell.execute_reply": "2023-12-07T00:51:41.773567Z"
    },
    "id": "hidden-sudan",
    "outputId": "9e44c675-657f-48cf-f900-fbe6c4b35d03"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['batch', 'x'])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_batch.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "conventional-county",
   "metadata": {
    "id": "conventional-county"
   },
   "source": [
    "To specify the datatype of each key, we can use the `data_and_attributes` parameter of AnnDataLoader. Here we make make `X` an `np.long` and our `cat_covs` an `np.float64`, but keep everything else as `np.float32`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "interested-craft",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.775332Z",
     "iopub.status.busy": "2023-12-07T00:51:41.775135Z",
     "iopub.status.idle": "2023-12-07T00:51:41.777143Z",
     "shell.execute_reply": "2023-12-07T00:51:41.776877Z"
    },
    "id": "interested-craft",
    "outputId": "2063f5f1-f456-45ce-fe14-500ad9b1cf27"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data Registry keys: ['batch', 'x']\n"
     ]
    }
   ],
   "source": [
    "# the keys of data_and_attributes should correspond to keys in the data registry\n",
    "data_registry_keys = adata_manager.data_registry.keys()\n",
    "print(\"Data Registry keys:\", data_registry_keys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "powered-maldives",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.778533Z",
     "iopub.status.busy": "2023-12-07T00:51:41.778354Z",
     "iopub.status.idle": "2023-12-07T00:51:41.780485Z",
     "shell.execute_reply": "2023-12-07T00:51:41.780218Z"
    },
    "id": "powered-maldives",
    "outputId": "8100c10c-6aed-4460-a5a1-8898ba4df2fb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'batch': <class 'numpy.float32'>, 'x': <class 'numpy.int64'>}\n"
     ]
    }
   ],
   "source": [
    "data_and_attributes = {}\n",
    "for key in data_registry_keys:\n",
    "    if key == \"x\":\n",
    "        data_and_attributes[key] = np.int64\n",
    "    else:\n",
    "        data_and_attributes[key] = np.float32\n",
    "print(data_and_attributes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "moved-score",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.781895Z",
     "iopub.status.busy": "2023-12-07T00:51:41.781694Z",
     "iopub.status.idle": "2023-12-07T00:51:41.784548Z",
     "shell.execute_reply": "2023-12-07T00:51:41.784281Z"
    },
    "id": "moved-score",
    "outputId": "65250cfb-972d-40e9-b87b-523cbe911e9d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.int64\n",
      "torch.float32\n"
     ]
    }
   ],
   "source": [
    "adl = AnnDataLoader(\n",
    "    adata_manager, shuffle=False, batch_size=10, data_and_attributes=data_and_attributes\n",
    ")\n",
    "data_batch = next(tensors for tensors in adl)\n",
    "\n",
    "# by default data has the dtype np.float32\n",
    "print(data_batch[\"x\"].dtype)\n",
    "print(data_batch[\"batch\"].dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "thermal-signal",
   "metadata": {
    "id": "thermal-signal"
   },
   "source": [
    "Finally, if the `data_and_attributes` parameter is used, it will only load the keys of the passed in dictionary. For example, if the only key in the dictionary passed in to `data_and_attributes` is X, the data loader will only load X."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "relevant-tournament",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.786002Z",
     "iopub.status.busy": "2023-12-07T00:51:41.785791Z",
     "iopub.status.idle": "2023-12-07T00:51:41.788308Z",
     "shell.execute_reply": "2023-12-07T00:51:41.788039Z"
    },
    "id": "relevant-tournament",
    "outputId": "9e45aed1-a78e-4edd-a780-3dea3b14cd3c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['x'])\n"
     ]
    }
   ],
   "source": [
    "data_and_attributes = {\"x\": float}\n",
    "adl = AnnDataLoader(\n",
    "    adata_manager, shuffle=False, batch_size=10, data_and_attributes=data_and_attributes\n",
    ")\n",
    "data_batch = next(iter(adl))\n",
    "\n",
    "print(data_batch.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62cf3179",
   "metadata": {
    "id": "62cf3179"
   },
   "source": [
    "### Using the DataLoader"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9db47ad3",
   "metadata": {
    "id": "9db47ad3"
   },
   "source": [
    "Below we demonstrate a toy use case where we can take advantage of the `AnnDataLoader` to minibatch data from our `AnnData` object into a model. In this example, we train a simple linear regression model."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1zxE4hVEgwfF",
   "metadata": {
    "id": "1zxE4hVEgwfF"
   },
   "source": [
    "```{important}\n",
    "The DataLoader will not move data to a device (e.g., GPU). This is the responsibility of the user. Alternatively, frameworks like [PyTorch Lightning](https://pytorch-lightning.readthedocs.io/en/latest/) will do this autmoatically for users.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "c9d691b6",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:41.789831Z",
     "iopub.status.busy": "2023-12-07T00:51:41.789634Z",
     "iopub.status.idle": "2023-12-07T00:51:42.580996Z",
     "shell.execute_reply": "2023-12-07T00:51:42.580562Z"
    },
    "id": "c9d691b6",
    "outputId": "11bb0eab-5b60-4756-f66a-b8485a460b06"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[iteration 1] loss: 338.4957580566406\n",
      "[iteration 2] loss: 94.9266357421875\n",
      "[iteration 3] loss: 101.34793090820312\n",
      "[iteration 4] loss: 73.22296905517578\n",
      "[iteration 5] loss: 10.78959846496582\n"
     ]
    }
   ],
   "source": [
    "# Initialize synthetic_iid data and register with an AnnDataManager\n",
    "n_genes, n_labels = 10, 3\n",
    "adata = scvi.data.synthetic_iid(n_genes=n_genes, n_labels=n_labels)\n",
    "anndata_fields = [\n",
    "    LayerField(registry_key=\"x\", layer=None, is_count_data=True),\n",
    "    CategoricalObsField(registry_key=\"labels\", attr_key=\"labels\"),\n",
    "]\n",
    "adata_manager = AnnDataManager(fields=anndata_fields)\n",
    "adata_manager.register_fields(adata)\n",
    "\n",
    "# Regression model\n",
    "linear_reg_model = torch.nn.Linear(n_genes, 1)\n",
    "\n",
    "# Define loss and optimize\n",
    "loss_fn = torch.nn.MSELoss(reduction=\"sum\")\n",
    "optim = torch.optim.Adam(linear_reg_model.parameters(), lr=0.05)\n",
    "\n",
    "\n",
    "def train(x, labels):\n",
    "    labels = labels.float()\n",
    "    # run the model forward on the data\n",
    "    label_pred = linear_reg_model(x).squeeze(-1)\n",
    "    # calculate the mse loss\n",
    "    loss = loss_fn(label_pred, labels.squeeze())\n",
    "    # initialize gradients to zero\n",
    "    optim.zero_grad()\n",
    "    # backpropagate\n",
    "    loss.backward()\n",
    "    # take a gradient step\n",
    "    optim.step()\n",
    "    return loss\n",
    "\n",
    "\n",
    "# Create AnnDataLoader\n",
    "# drop a minibatch if it has 3 or fewer observations\n",
    "data_loader = AnnDataLoader(\n",
    "    adata_manager,\n",
    "    batch_size=128,\n",
    "    shuffle=True,\n",
    ")\n",
    "\n",
    "for i in range(5):\n",
    "    for data in data_loader:\n",
    "        loss = train(data[\"x\"], data[\"labels\"])\n",
    "    print(f\"[iteration {i + 1}] loss: {loss.item()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fewer-growth",
   "metadata": {
    "id": "fewer-growth"
   },
   "source": [
    "## `scvi-tools` Data Registration\n",
    "\n",
    "Scvi-tools models produce an `AnnDataManager` instance in the `setup_anndata()` function for the purpose of data registration.\n",
    "\n",
    "`setup_anndata()` is used to setup data fields specific to each model.\n",
    "\n",
    "Here we will go over the parameters of one instance of a `setup_anndata()` method, `scvi.model.SCVI.setup_anndata()`:\n",
    "\n",
    "- `adata` is the input `AnnData` object.\n",
    "- `layer` is the key in `adata.layers` to use for the input data matrix. By default, this is None and the input data matrix will be pulled from `adata.X`.\n",
    "- `batch_key` is the key in `adata.obs` for batch information. If this is None, will assume that all the data is the same batch.\n",
    "- `labels_key` is the key in `adata.obs` for label information. If this is None, will assume that all the data has the same label.\n",
    "- `size_factor_key` is the key in `adata.obs` that optionally stores size factors for computing the likelihood. If this is None, the library size is used to compute the size factor.\n",
    "- `categorical_covariate_keys` is a list of keys in `adata.obs` for categorical covariates.\n",
    "- `continuous_covariate_key` is a list of keys in `adata.obs` for continuous covariates.\n",
    "\n",
    "Under the hood:\n",
    "\n",
    "- For all categorical data (batch, labels, categorical covariates), scvi will automatically compute a mapping from values to integers. Eg. `['a','b','c','a']` will become `[0,1,2,0]`.\n",
    "- For data fields registered with `scvi.model.SCVI.setup_anndata()`, scvi will copy the data to a seperate field in the anndata.\n",
    "  - `batch_key` is copied to `scvi.obs['_scvi_batch']` with its integer encoding\n",
    "  - `labels_key` is copied to `scvi.obs['_scvi_labels']` with its integer encoding\n",
    "  - keys in `categorical_covariate_keys` are concatenated and saved as a pandas DataFrame and stored in `adata.obsm['_scvi_extra_categorical_covs']` with its integer encoding.\n",
    "  - keys in `continuous_covariate_keys` are concatenated and saved as a pandas DataFrame and stored in `adata.obsm['_scvi_extra_continuous_covs']`\n",
    "\n",
    "In the following code, we first format an example AnnData Object to setup for scvi-tools, then call `scvi.model.SCVI.setup_anndata()` to register all the tensors we want to load to the model during training.\n",
    "For our example AnnData Object, we build off the `synthetic_iid()` dataset, copy X to a layer, and add continuous and categorical covariates to the AnnData."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "south-enclosure",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:42.582609Z",
     "iopub.status.busy": "2023-12-07T00:51:42.582499Z",
     "iopub.status.idle": "2023-12-07T00:51:42.598815Z",
     "shell.execute_reply": "2023-12-07T00:51:42.598520Z"
    },
    "id": "south-enclosure",
    "outputId": "f6b1bf0e-f9dd-4492-b525-bf2a03f7a212"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AnnData object with n_obs × n_vars = 400 × 100\n",
      "    obs: 'batch', 'labels', 'my_categorical_covariate', 'my_continuous_covariate'\n",
      "    uns: 'protein_names'\n",
      "    obsm: 'protein_expression', 'accessibility'\n",
      "    layers: 'raw_counts'\n"
     ]
    }
   ],
   "source": [
    "adata = scvi.data.synthetic_iid()\n",
    "adata.layers[\"raw_counts\"] = adata.X.copy()\n",
    "adata.obs[\"my_categorical_covariate\"] = [\"A\"] * 200 + [\"B\"] * 200\n",
    "adata.obs[\"my_continuous_covariate\"] = np.random.randint(0, 100, 400)\n",
    "print(adata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "through-postcard",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:42.600206Z",
     "iopub.status.busy": "2023-12-07T00:51:42.600099Z",
     "iopub.status.idle": "2023-12-07T00:51:42.604847Z",
     "shell.execute_reply": "2023-12-07T00:51:42.604461Z"
    },
    "id": "through-postcard"
   },
   "outputs": [],
   "source": [
    "scvi.model.SCVI.setup_anndata(\n",
    "    adata,\n",
    "    batch_key=\"batch\",\n",
    "    labels_key=\"labels\",\n",
    "    layer=\"raw_counts\",\n",
    "    categorical_covariate_keys=[\"my_categorical_covariate\"],\n",
    "    continuous_covariate_keys=[\"my_continuous_covariate\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c53b21e",
   "metadata": {
    "id": "1c53b21e"
   },
   "source": [
    "Under the hood, this method creates an `AnnDataManager` instance and stores it in a model-specific manager store until a model is initialized with the same `AnnData` object. We can visualize the resulting `registry` via `view_anndata_setup()`. This calls `view_registry()` under the hood."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "e8fe3987",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 977
    },
    "execution": {
     "iopub.execute_input": "2023-12-07T00:51:42.606179Z",
     "iopub.status.busy": "2023-12-07T00:51:42.606076Z",
     "iopub.status.idle": "2023-12-07T00:51:42.624603Z",
     "shell.execute_reply": "2023-12-07T00:51:42.624347Z"
    },
    "id": "e8fe3987",
    "outputId": "9bae34f2-6ef7-4548-f865-51966da91e1d"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Anndata setup with scvi-tools version <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.1</span>.<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>.\n",
       "</pre>\n"
      ],
      "text/plain": [
       "Anndata setup with scvi-tools version \u001b[1;36m1.1\u001b[0m.\u001b[1;36m0\u001b[0m.\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Setup via `SCVI.setup_anndata` with arguments:\n",
       "</pre>\n"
      ],
      "text/plain": [
       "Setup via `SCVI.setup_anndata` with arguments:\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">{</span>\n",
       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'layer'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'raw_counts'</span>,\n",
       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'batch_key'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'batch'</span>,\n",
       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'labels_key'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'labels'</span>,\n",
       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'size_factor_key'</span>: <span style=\"color: #800080; text-decoration-color: #800080; font-style: italic\">None</span>,\n",
       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'categorical_covariate_keys'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'my_categorical_covariate'</span><span style=\"font-weight: bold\">]</span>,\n",
       "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #008000; text-decoration-color: #008000\">'continuous_covariate_keys'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'my_continuous_covariate'</span><span style=\"font-weight: bold\">]</span>\n",
       "<span style=\"font-weight: bold\">}</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m{\u001b[0m\n",
       "\u001b[2;32m│   \u001b[0m\u001b[32m'layer'\u001b[0m: \u001b[32m'raw_counts'\u001b[0m,\n",
       "\u001b[2;32m│   \u001b[0m\u001b[32m'batch_key'\u001b[0m: \u001b[32m'batch'\u001b[0m,\n",
       "\u001b[2;32m│   \u001b[0m\u001b[32m'labels_key'\u001b[0m: \u001b[32m'labels'\u001b[0m,\n",
       "\u001b[2;32m│   \u001b[0m\u001b[32m'size_factor_key'\u001b[0m: \u001b[3;35mNone\u001b[0m,\n",
       "\u001b[2;32m│   \u001b[0m\u001b[32m'categorical_covariate_keys'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'my_categorical_covariate'\u001b[0m\u001b[1m]\u001b[0m,\n",
       "\u001b[2;32m│   \u001b[0m\u001b[32m'continuous_covariate_keys'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'my_continuous_covariate'\u001b[0m\u001b[1m]\u001b[0m\n",
       "\u001b[1m}\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">         Summary Statistics         </span>\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">     Summary Stat Key     </span>┃<span style=\"font-weight: bold\"> Value </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━┩\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">         n_batch          </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">   2   </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">         n_cells          </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">  400  </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\"> n_extra_categorical_covs </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">   1   </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\"> n_extra_continuous_covs  </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">   1   </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">         n_labels         </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">   3   </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">          n_vars          </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">  100  </span>│\n",
       "└──────────────────────────┴───────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[3m         Summary Statistics         \u001b[0m\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m    Summary Stat Key    \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mValue\u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━┩\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m        n_batch         \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m  2  \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m        n_cells         \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m 400 \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33mn_extra_categorical_covs\u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m  1  \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33mn_extra_continuous_covs \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m  1  \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m        n_labels        \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m  3  \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m         n_vars         \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m 100 \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "└──────────────────────────┴───────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">                             Data Registry                             </span>\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">      Registry Key      </span>┃<span style=\"font-weight: bold\">            scvi-tools Location             </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">           X            </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">         adata.layers['raw_counts']         </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">         batch          </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">          adata.obs['_scvi_batch']          </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\"> extra_categorical_covs </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\"> adata.obsm['_scvi_extra_categorical_covs'] </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\"> extra_continuous_covs  </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\"> adata.obsm['_scvi_extra_continuous_covs']  </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">         labels         </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">         adata.obs['_scvi_labels']          </span>│\n",
       "└────────────────────────┴────────────────────────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[3m                             Data Registry                             \u001b[0m\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m     Registry Key     \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m           scvi-tools Location            \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m          X           \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m        adata.layers['raw_counts']        \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m        batch         \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m         adata.obs['_scvi_batch']         \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33mextra_categorical_covs\u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128madata.obsm['_scvi_extra_categorical_covs']\u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33mextra_continuous_covs \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128madata.obsm['_scvi_extra_continuous_covs'] \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m        labels        \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m        adata.obs['_scvi_labels']         \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "└────────────────────────┴────────────────────────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">                  batch State Registry                   </span>\n",
       "┏━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">  Source Location   </span>┃<span style=\"font-weight: bold\"> Categories </span>┃<span style=\"font-weight: bold\"> scvi-tools Encoding </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\"> adata.obs['batch'] </span>│<span style=\"color: #008000; text-decoration-color: #008000\">  batch_0   </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">          0          </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">                    </span>│<span style=\"color: #008000; text-decoration-color: #008000\">  batch_1   </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">          1          </span>│\n",
       "└────────────────────┴────────────┴─────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[3m                  batch State Registry                   \u001b[0m\n",
       "┏━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m Source Location  \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mCategories\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mscvi-tools Encoding\u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33madata.obs['batch']\u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m batch_0  \u001b[0m\u001b[32m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m         0         \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m                  \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m batch_1  \u001b[0m\u001b[32m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m         1         \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "└────────────────────┴────────────┴─────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">                  labels State Registry                   </span>\n",
       "┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">   Source Location   </span>┃<span style=\"font-weight: bold\"> Categories </span>┃<span style=\"font-weight: bold\"> scvi-tools Encoding </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\"> adata.obs['labels'] </span>│<span style=\"color: #008000; text-decoration-color: #008000\">  label_0   </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">          0          </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">                     </span>│<span style=\"color: #008000; text-decoration-color: #008000\">  label_1   </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">          1          </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">                     </span>│<span style=\"color: #008000; text-decoration-color: #008000\">  label_2   </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">          2          </span>│\n",
       "└─────────────────────┴────────────┴─────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[3m                  labels State Registry                   \u001b[0m\n",
       "┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m  Source Location  \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mCategories\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mscvi-tools Encoding\u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33madata.obs['labels']\u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m label_0  \u001b[0m\u001b[32m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m         0         \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m                   \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m label_1  \u001b[0m\u001b[32m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m         1         \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m                   \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m label_2  \u001b[0m\u001b[32m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m         2         \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "└─────────────────────┴────────────┴─────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">                   extra_categorical_covs State Registry                    </span>\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">            Source Location            </span>┃<span style=\"font-weight: bold\"> Categories </span>┃<span style=\"font-weight: bold\"> scvi-tools Encoding </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\"> adata.obs['my_categorical_covariate'] </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     A      </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">          0          </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">                                       </span>│<span style=\"color: #008000; text-decoration-color: #008000\">     B      </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">          1          </span>│\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\">                                       </span>│<span style=\"color: #008000; text-decoration-color: #008000\">            </span>│<span style=\"color: #af00d7; text-decoration-color: #af00d7\">                     </span>│\n",
       "└───────────────────────────────────────┴────────────┴─────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[3m                   extra_categorical_covs State Registry                    \u001b[0m\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m           Source Location           \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mCategories\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mscvi-tools Encoding\u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33madata.obs['my_categorical_covariate']\u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    A     \u001b[0m\u001b[32m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m         0         \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m                                     \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m    B     \u001b[0m\u001b[32m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m         1         \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33m                                     \u001b[0m\u001b[38;5;33m \u001b[0m│\u001b[32m \u001b[0m\u001b[32m          \u001b[0m\u001b[32m \u001b[0m│\u001b[38;5;128m \u001b[0m\u001b[38;5;128m                   \u001b[0m\u001b[38;5;128m \u001b[0m│\n",
       "└───────────────────────────────────────┴────────────┴─────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-style: italic\">  extra_continuous_covs State Registry  </span>\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">           Source Location            </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #0087ff; text-decoration-color: #0087ff\"> adata.obs['my_continuous_covariate'] </span>│\n",
       "└──────────────────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[3m  extra_continuous_covs State Registry  \u001b[0m\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m          Source Location           \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[38;5;33m \u001b[0m\u001b[38;5;33madata.obs['my_continuous_covariate']\u001b[0m\u001b[38;5;33m \u001b[0m│\n",
       "└──────────────────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model = scvi.model.SCVI(adata)\n",
    "model.view_anndata_setup()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35a7ca64",
   "metadata": {
    "id": "35a7ca64"
   },
   "source": [
    "Each model defines a set of appropriate `AnnDataField`s and orchestrates calls to these functions and stores the resulting `registry` as an instance variable. As mentioned before, the `AnnDataManager` is constructed during `setup_anndata()` and retrieved during model initialization.\n",
    "\n",
    "Here we have an abbreviated version of a `setup_anndata()` implementation for a model that only takes a `layer` kwarg and a `batch_key`:\n",
    "\n",
    "```python\n",
    "@classmethod\n",
    "def setup_anndata(\n",
    "    cls,\n",
    "    adata: AnnData,\n",
    "    layer: Optional[str] = None,\n",
    "    batch_key: Optional[str] = None,\n",
    "    **kwargs,  # Used when loading a model with a new AnnData object.\n",
    "):\n",
    "    setup_method_args = cls._get_setup_method_args(\n",
    "        **locals()\n",
    "    )  # Used for saving/loading purposes.\n",
    "    anndata_fields = [\n",
    "        LayerField(REGISTRY_KEYS.X_KEY, layer, is_count_data=True),\n",
    "        CategoricalObsField(REGISTRY_KEYS.BATCH_KEY, batch_key),\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(\n",
    "        adata_manager\n",
    "    )  # Stores the AnnDataManager in a class-specific manager store.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f45dbc64",
   "metadata": {
    "id": "f45dbc64"
   },
   "source": [
    "The `setup_anndata()` function itself is quite simple since any complexity in preprocessing is contained within the `AnnDataField` functions. By factorizing the preprocessing steps into each subclass, model developers can easily extend and reuse logic across models and fields."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "data_tutorial.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "scvi-tools-dev",
   "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.11.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "b5142939ddaa1787bd1bfcf4c0ad4d35be0fa2237c553f986d37efcb39f03f79"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}