{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "nbsphinx": "hidden"
   },
   "source": [
    "# Vitessce Widget Tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualization of single-cell RNA seq data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Import dependencies\n",
    "\n",
    "We need to import the classes and functions that we will be using from the corresponding packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from os.path import join, isfile, isdir\n",
    "from urllib.request import urlretrieve\n",
    "from anndata import read_h5ad\n",
    "import scanpy as sc\n",
    "\n",
    "from vitessce import (\n",
    "    VitessceConfig,\n",
    "    Component as cm,\n",
    "    CoordinationType as ct,\n",
    "    AnnDataWrapper,\n",
    ")\n",
    "from vitessce.data_utils import (\n",
    "    optimize_adata,\n",
    "    VAR_CHUNK_SIZE,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Download the data\n",
    "\n",
    "For this example, we need to download a dataset from the COVID-19 Cell Atlas https://www.covid19cellatlas.org/index.healthy.html#habib17."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata_filepath = join(\"data\", \"habib17.processed.h5ad\")\n",
    "if not isfile(adata_filepath):\n",
    "    os.makedirs(\"data\", exist_ok=True)\n",
    "    urlretrieve('https://covid19.cog.sanger.ac.uk/habib17.processed.h5ad', adata_filepath)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Load the data\n",
    "\n",
    "Note: this function may print a `FutureWarning`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata = read_h5ad(adata_filepath)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 3.1. Preprocess the Data For Visualization\n",
    "\n",
    "This dataset contains 25,587 genes.  We prepare to visualize the top 50 highly variable genes for the heatmap as ranked by dispersion norm, although one may use any boolean array filter for the heatmap."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_dispersion = adata.var[\"dispersions_norm\"][\n",
    "    sorted(\n",
    "        range(len(adata.var[\"dispersions_norm\"])),\n",
    "        key=lambda k: adata.var[\"dispersions_norm\"][k],\n",
    "    )[-51:][0]\n",
    "]\n",
    "adata.var[\"top_highly_variable\"] = (\n",
    "    adata.var[\"dispersions_norm\"] > top_dispersion\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.2 Save the Data to Zarr store\n",
    "\n",
    "We want to convert the original `h5ad` file to a [Zarr](https://zarr.readthedocs.io/en/stable/) store, which Vitessce is able to load. We can use the `optimize_adata` function to ensure that all arrays and dataframe columns that we intend to use in our visualization are in the optimal format to be loaded by Vitessce. This function will cast arrays to numerical data types that take up less space (as long as the values allow). Note: unused arrays and columns (i.e., not specified in any of the parameters to `optimize_adata`) will not be copied into the new AnnData object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "zarr_filepath = join(\"data\", \"habib17.processed.zarr\")\n",
    "if not isdir(zarr_filepath):\n",
    "    adata = optimize_adata(\n",
    "        adata,\n",
    "        obs_cols=[\"CellType\"],\n",
    "        obsm_keys=[\"X_umap\"],\n",
    "        optimize_X=True,\n",
    "        var_cols=[\"top_highly_variable\"],\n",
    "    )\n",
    "    adata.write_zarr(zarr_filepath, chunks=[adata.shape[0], VAR_CHUNK_SIZE])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Create the Vitessce widget configuration\n",
    "\n",
    "Vitessce needs to know which pieces of data we are interested in visualizing, the visualization types we would like to use, and how we want to coordinate (or link) the views."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1. Instantiate a `VitessceConfig` object\n",
    "\n",
    "Use the `VitessceConfig` constructor to create an instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vc = VitessceConfig(schema_version=\"1.0.15\", name='Habib et al', description='COVID-19 Healthy Donor Brain')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2. Add a dataset to the `VitessceConfig` instance\n",
    "\n",
    "In Vitessce, a dataset is a container for one file per data type. The `.add_dataset(name)` method on the `vc` instance sets up and returns a new dataset instance.\n",
    "\n",
    "Then, we can call the dataset's `.add_object(wrapper_object)` method to attach a \"data wrapper\" instance to our new dataset. For example, the `AnnDataWrapper` helps to configure AnnData Zarr stores for use in the Vitessce configuration.\n",
    "\n",
    "Dataset wrapper classes may require additional parameters to resolve ambiguities. For instance, `AnnData` objects may store multiple clusterings or cell type annotation columns in the `adata.obs` dataframe. We can use the parameter `obs_set_paths` to tell Vitessce that certain columns of the `obs` dataframe correspond to cell type annotations or cell clusterings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = vc.add_dataset(name='Brain').add_object(AnnDataWrapper(\n",
    "        adata_path=zarr_filepath,\n",
    "        obs_embedding_paths=[\"obsm/X_umap\"],\n",
    "        obs_embedding_names=[\"UMAP\"],\n",
    "        obs_set_paths=[\"obs/CellType\"],\n",
    "        obs_set_names=[\"Cell Type\"],\n",
    "        obs_feature_matrix_path=\"X\",\n",
    "        initial_feature_filter_path=\"var/top_highly_variable\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.3. Add visualizations to the `VitessceConfig` instance\n",
    "\n",
    "Now that we have added a dataset, we can configure visualizations. The `.add_view` method adds a view (i.e. visualization or controller component) to the configuration.\n",
    "\n",
    "The `Component` enum class (which we have imported as `cm` here) can be used to fill in the `component_type` parameter.\n",
    "\n",
    "For convenience, the `SCATTERPLOT` component type takes the extra `mapping` keyword argument, which specifies which embedding should be used for mapping cells to (x,y) points on the plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scatterplot = vc.add_view(cm.SCATTERPLOT, dataset=dataset, mapping=\"UMAP\")\n",
    "cell_sets = vc.add_view(cm.OBS_SETS, dataset=dataset)\n",
    "genes = vc.add_view(cm.FEATURE_LIST, dataset=dataset)\n",
    "heatmap = vc.add_view(cm.HEATMAP, dataset=dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.4. Define the visualization layout\n",
    "\n",
    "The `vc.layout(view_concat)` method allows us to specify how our views will be arranged in the layout grid in the widget. The `|` and `/` characters are magic syntax for `hconcat(v1, v2)` and `vconcat(v1, v2)`, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vc.layout((scatterplot | cell_sets) / (heatmap | genes));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Create the widget\n",
    "\n",
    "The `vc.widget()` method returns the configured widget instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vw = vc.widget()\n",
    "vw"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}