{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# scNym Cell Type Prediction Tutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scanpy.api as sc\n",
    "import os\n",
    "import os.path as osp\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.metrics import confusion_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scnym\n",
    "from scnym.predict import Predicter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download Example Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "mkdir: cannot create directory ‘../downloads’: File exists\n",
      "2019-07-18 12:46:35 URL:https://storage.googleapis.com/calico-website-mca-storage/20190604_tissue_specific_weights_h0256_l1.tar.gz [211326981/211326981] -> \"../downloads/tissue_specific_weights.tar.gz\" [1]\n",
      "2019-07-18 12:46:36 URL:https://storage.googleapis.com/calico-website-mca-storage/20190604_scnym_training_gene_names.txt [128251/128251] -> \"../downloads/training_gene_names.txt\" [1]\n",
      "2019-07-18 12:46:36 URL:https://storage.googleapis.com/calico-website-mca-storage/20190604_training_metadata.csv [5063191/5063191] -> \"../downloads/training_metadata.csv\" [1]\n",
      "2019-07-18 12:46:36 URL:https://storage.googleapis.com/calico-website-mca-storage/microwell_mouse_cell_atlas_lung.h5ad [20795497/20795497] -> \"../downloads/mca_lung_reannotated.h5ad\" [1]\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "mkdir ../downloads\n",
    "# wget -nv https://storage.googleapis.com/calico-website-mca-storage/20190604_tissue_independent_h0256_l1.pkl \\\n",
    "#     -O ../downloads/model_weights.pkl\n",
    "wget -nv https://storage.googleapis.com/calico-website-mca-storage/20190604_tissue_specific_weights_h0256_l1.tar.gz \\\n",
    "    -O ../downloads/tissue_specific_weights.tar.gz\n",
    "wget -nv https://storage.googleapis.com/calico-website-mca-storage/20190604_scnym_training_gene_names.txt \\\n",
    "    -O ../downloads/training_gene_names.txt\n",
    "wget -nv https://storage.googleapis.com/calico-website-mca-storage/20190604_training_metadata.csv \\\n",
    "    -O ../downloads/training_metadata.csv\n",
    "wget -nv https://storage.googleapis.com/calico-website-mca-storage/microwell_mouse_cell_atlas_lung.h5ad \\\n",
    "    -O ../downloads/mca_lung_reannotated.h5ad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tissue_specific_weights/\n",
      "tissue_specific_weights/Bladder.pkl\n",
      "tissue_specific_weights/Liver.pkl\n",
      "tissue_specific_weights/Thymus.pkl\n",
      "tissue_specific_weights/Kidney.pkl\n",
      "tissue_specific_weights/Trachea.pkl\n",
      "tissue_specific_weights/Heart_and_Aorta.pkl\n",
      "tissue_specific_weights/Mammary_Gland.pkl\n",
      "tissue_specific_weights/Spleen.pkl\n",
      "tissue_specific_weights/Tongue.pkl\n",
      "tissue_specific_weights/Marrow.pkl\n",
      "tissue_specific_weights/Limb_Muscle.pkl\n",
      "tissue_specific_weights/Lung.pkl\n",
      "/home/jacob/src/scnym/notebooks\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "cd ../downloads\n",
    "tar xvf tissue_specific_weights.tar.gz\n",
    "cd -"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Thymus.pkl',\n",
       " 'Kidney.pkl',\n",
       " 'Limb_Muscle.pkl',\n",
       " 'Liver.pkl',\n",
       " 'Spleen.pkl',\n",
       " 'Tongue.pkl',\n",
       " 'Trachea.pkl',\n",
       " 'Marrow.pkl',\n",
       " 'Bladder.pkl',\n",
       " 'Mammary_Gland.pkl',\n",
       " 'Heart_and_Aorta.pkl',\n",
       " 'Lung.pkl']"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.listdir('../downloads/tissue_specific_weights/')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Example Data\n",
    "\n",
    "We'll use lung data from the [\"Mouse Cell Atlas\"](https://www.ncbi.nlm.nih.gov/pubmed/29474909) that we have manually annotated with cell ontology class labels to demonstrate cell type prediction. Here, we use `scNym` models trained on the *Tabula Muris* 10X data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1500 cells and 25133 genes are in the data set.\n"
     ]
    }
   ],
   "source": [
    "adata = sc.read_h5ad(\n",
    "    '../downloads/mca_lung_reannotated.h5ad')\n",
    "print('%d cells and %d genes are in the data set.'%adata.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 57360 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sc.pl.tsne(adata, color='MCellType', palette='tab20') # MCellType are our manual cell type annotations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predict Cell Types using an scNym Model\n",
    "\n",
    "Here we use lung specific weights, available in our tissue specific weights download linked in the README."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "18099 training gene names.\n",
      "13 cell types in the training annotation.\n"
     ]
    }
   ],
   "source": [
    "model_weights = '../downloads/tissue_specific_weights/Lung.pkl'\n",
    "training_gene_names = np.loadtxt('../downloads/training_gene_names.txt', dtype='str')\n",
    "training_metadata = pd.read_csv('../downloads/training_metadata.csv')\n",
    "\n",
    "print('%d training gene names.'%len(training_gene_names))\n",
    "print('%d cell types in the training annotation.'%len((training_metadata\n",
    "                                                       .loc[training_metadata.tissue=='Lung', 'cell_ontology_class']\n",
    "                                                       .unique())))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare input matrix\n",
    "\n",
    "We first need to generate an expression matrix from our sample data with genes in the same order we used for training. \n",
    "\n",
    "scNym has a utility to make this easy. We find that scNym models tend to work well even when only a few thousand genes overlap between data sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Building matrix: 18099it [00:21, 856.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 16117 common genes.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from scnym.utils import build_classification_matrix\n",
    "\n",
    "X = build_classification_matrix(X=adata.X.todense(),\n",
    "                                model_genes=training_gene_names,\n",
    "                                sample_genes=np.array(adata.var_names))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['B cell', 'T cell', 'alveolar macrophage', 'ciliated columnar cell of tracheobronchial tree', 'classical monocyte', 'leukocyte', 'lung endothelial cell', 'mast cell', 'myeloid cell', 'natural killer cell', 'non-classical monocyte', 'stromal cell', 'type II pneumocyte']\n"
     ]
    }
   ],
   "source": [
    "cell_types = sorted(training_metadata.loc[training_metadata.tissue=='Lung','cell_ontology_class'].unique().tolist())\n",
    "print(cell_types)\n",
    "P = Predicter(model_weights=model_weights,\n",
    "              n_genes=len(training_gene_names),\n",
    "              n_cell_types=len(cell_types),\n",
    "              labels=cell_types,\n",
    "              n_hidden=256,\n",
    "              n_layers=1,\n",
    "              residual=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Finding cell types: 100%|██████████| 1500/1500 [00:04<00:00, 333.86it/s]\n"
     ]
    }
   ],
   "source": [
    "predictions, class_names, probabilities = P.predict(X=X, output='prob')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# our annotations don't include monocyte subclasses\n",
    "# so we group them together here\n",
    "class_names = np.array(class_names)\n",
    "class_names[class_names=='non-classical monocyte']='monocyte'\n",
    "class_names[class_names=='classical monocyte']='monocyte'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "... storing 'pred_type' as categorical\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 57360 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 62424 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "adata.obs['pred_type'] = class_names\n",
    "sc.pl.tsne(adata, color='MCellType', palette='tab20')\n",
    "sc.pl.tsne(adata, color='pred_type', palette='tab20')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that after a first pass, our model is mostly correct in annotating each cluster.\n",
    "However, some individual cells within each cluster appear to have the wrong label.\n",
    "\n",
    "To fix this, we use local neighborhood information to correct some cell type labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scnym.utils import knn_smooth_pred_class\n",
    "\n",
    "smooth_pred = knn_smooth_pred_class(X=adata.obsm['X_pca'],  # embedding for finding nearest neighbors\n",
    "                                    pred_class=np.array(adata.obs['pred_type']), # predicted classes\n",
    "                                    grouping=np.array(adata.obs['louvain']), # cluster partition\n",
    "                                    k=5, # k-nearest neighbors to consider\n",
    "                                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "... storing 'smooth_pred_type' as categorical\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 57360 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 60776 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "adata.obs['smooth_pred_type'] = smooth_pred\n",
    "sc.pl.tsne(adata, color='MCellType', palette='tab20')\n",
    "sc.pl.tsne(adata, color='smooth_pred_type', palette='tab20')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After *k*-NN smoothing, we find that the cell type predictions look mostly correct. The errors we do see are usually confusion between similar cell types in the training data (e.g. myeloid cells and B cells).\n",
    "\n",
    "Digging further, we find that some of the cells we originally annotated as myeloid might actually be B cells. See the expression pattern of B cell marker [*Cd72*](https://en.wikipedia.org/wiki/CD72) below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jacob/bin/envs/cbpytorch/lib/python3.6/site-packages/anndata/core/anndata.py:299: FutureWarning: In anndata v0.7+, arrays contained within an AnnData object will maintain their dimensionality. For example, prior to v0.7 `adata[0, 0].X` returned a scalar and `adata[0, :]` returned a 1d array, post v0.7 they will return two dimensional arrays. If you would like to get a one dimensional array from your AnnData object, consider using the `adata.obs_vector`, `adata.var_vector` methods or accessing the array directly.\n",
      "  warn_flatten()\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 50364 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "sc.pl.tsne(adata, color='Cd72')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "81.133333% accuracy before smoothing.\n",
      "84.733333% accuracy after smoothing.\n"
     ]
    }
   ],
   "source": [
    "y_true = np.array(adata.obs['MCellType'])\n",
    "y_pred = np.array(adata.obs['pred_type'])\n",
    "\n",
    "y_pred_s = np.array(adata.obs['smooth_pred_type'])\n",
    "\n",
    "acc = np.sum(y_true == y_pred)/len(y_true)\n",
    "acc_s = np.sum(y_true == y_pred_s)/len(y_true)\n",
    "print('{:f}% accuracy before smoothing.'.format(acc*100))\n",
    "print('{:f}% accuracy after smoothing.'.format(acc_s*100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Considering the train and test data are from different experimental platforms on different continents, this seems like a reasonable degree of accuracy for a first pass.\n",
    "\n",
    "Manual curation can be leveraged to rectify the remaining deficits."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cbpytorch",
   "language": "python",
   "name": "cbpytorch"
  },
  "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
