plotEmbedding_jointCNNgae_Starmap.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3cf5bca1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('/home/xinyiz/pamrats')\n",
    "\n",
    "import time\n",
    "import os\n",
    "\n",
    "import scanpy\n",
    "import numpy as np\n",
    "import scipy.sparse as sp\n",
    "\n",
    "import torch\n",
    "from torch import optim\n",
    "\n",
    "import image.loadImage as loadImage\n",
    "import gae.gae.optimizer as optimizer\n",
    "import image.modelsCNN as modelsCNN\n",
    "\n",
    "import pickle\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import umap\n",
    "import pandas as pd\n",
    "from sklearn.preprocessing import scale\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "from sklearn.cluster import DBSCAN,MiniBatchKMeans,AgglomerativeClustering\n",
    "from sklearn import metrics\n",
    "\n",
    "import anndata as ad\n",
    "import gc\n",
    "import matplotlib.font_manager as fm\n",
    "from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5fac6a59",
   "metadata": {},
   "outputs": [],
   "source": [
    "#load pretrained GAE\n",
    "\n",
    "datadir='/home/xinyiz/2021-01-13-mAD-test-dataset'\n",
    "sampleidx={'disease13':'AD_mouse9494','control13':'AD_mouse9498','disease8':'AD_mouse9723','control8':'AD_mouse9735'}\n",
    "\n",
    "\n",
    "scaleddata=scanpy.read_h5ad(datadir+'/2020-12-27-starmap-mAD-raw.h5ad')\n",
    "\n",
    "cellCoord={}\n",
    "for s in sampleidx.keys():\n",
    "    sampleidx_s=sampleidx[s] \n",
    "    cellCoord[s]=((scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx_s,['y','x']].to_numpy())/0.3).astype(int)\n",
    "scaleddata=None\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c7d2b898",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\" \n",
    "ifplot=True\n",
    "ifcluster=True\n",
    "\n",
    "inverseAct=None\n",
    "# inverseAct=None\n",
    "plottype='umap'\n",
    "pca=PCA()\n",
    "minCells=15 #min number of cells for analysis\n",
    "# clustermethod=['kmeanbatch']\n",
    "# clustermethod=['leiden','agglomerative','kmeanbatch']\n",
    "clustermethod=['leiden']\n",
    "#umap/leiden clustering parameters\n",
    "n_neighbors=10\n",
    "min_dist=0.25\n",
    "n_pcs=40 #for clustering\n",
    "# resolution=[0.5,0.8,1,1.5]\n",
    "resolution=[0.05,0.1,0.2,0.3,0.5,0.8,1,1.5]\n",
    "plotepoch=36\n",
    "savenameAdd=''\n",
    "#DBscan\n",
    "epslist= [6,8,10]\n",
    "min_sampleslist=[15,30,45] \n",
    "#agglomerative\n",
    "nclusterlist=[2,3,4,5,8,10,15]\n",
    "aggMetric=['euclidean']\n",
    "\n",
    "\n",
    "combineCelltype={'glia':['Astro','Micro', 'OPC', 'Oligo'],'CA':['CA1', 'CA2', 'CA3']}\n",
    "\n",
    "use_cuda=True\n",
    "fastmode=False #Validate during training pass\n",
    "seed=3\n",
    "kernel_size=4\n",
    "stride=2\n",
    "padding=1\n",
    "\n",
    "# fc_dim1=6000\n",
    "hidden1=64 #Number of channels in hidden layer 1\n",
    "hidden2=128 \n",
    "hidden3=256\n",
    "hidden4=256\n",
    "hidden5=96\n",
    "fc_dim1=96*25*25\n",
    "fc_dim2=5000\n",
    "# hidden1=64 #Number of channels in hidden layer 1\n",
    "# hidden2=128 \n",
    "# hidden3=256\n",
    "# hidden4=512\n",
    "# hidden5=512\n",
    "# fc_dim1=512*25*25\n",
    "# fc_dim2=1024\n",
    "# # fc_dim3=128\n",
    "# # fc_dim4=128\n",
    "# # gcn_dim1=2600\n",
    "# adv_hidden=128\n",
    "\n",
    "model_str='cnn_vae'\n",
    "targetBatch=None\n",
    "diamThresh_mul=800\n",
    "minThresh_mul=12\n",
    "overlap=int(diamThresh_mul*0.5)\n",
    "name='cd13_thresh25_02'\n",
    "logsavepath='/mnt/external_ssd/xinyi/log/train_jointGAEcnn_starmap/'+name\n",
    "modelsavepath='/mnt/external_ssd/xinyi/models/train_jointGAEcnn_starmap/'+name\n",
    "plotsavepath='/mnt/external_ssd/xinyi/plots/train_jointGAEcnn_starmap/'+name\n",
    "# modelsavepath='/mnt/external_ssd/xinyi/models/train_cnn_starmap/'+name\n",
    "# plotsavepath='/mnt/external_ssd/xinyi/plots/train_cnn_starmap/'+name\n",
    "\n",
    "#Load data\n",
    "plot_samples={'disease13':'AD_mouse9494','control13':'AD_mouse9498','disease8':'AD_mouse9723','control8':'AD_mouse9735'}\n",
    "datadir='/home/xinyiz/2021-01-13-mAD-test-dataset'    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "63332f3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set cuda and seed\n",
    "np.random.seed(seed)\n",
    "if use_cuda and (not torch.cuda.is_available()):\n",
    "    print('cuda not available')\n",
    "    use_cuda=False\n",
    "torch.manual_seed(seed)\n",
    "if use_cuda:\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.backends.cudnn.enabled = True\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b790fd69",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create model\n",
    "if model_str=='cnn_vae':\n",
    "    model = modelsCNN.CNN_VAE(kernel_size, stride, padding, 1, hidden1, hidden2, hidden3, hidden4, hidden5, fc_dim1,fc_dim2)\n",
    "if model_str=='cnn_vae_alexnet':\n",
    "    model = modelsCNN.CNN_VAE_alexnet(fc_dim1)      \n",
    "if use_cuda:\n",
    "    model.cuda()    \n",
    "model.load_state_dict(torch.load(os.path.join(modelsavepath,str(plotepoch)+'.pt')))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9cd59972",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(seed)\n",
    "def inverseLeakyRelu(v,slope=0.01):\n",
    "    vnegidx=(v<0)\n",
    "    v[vnegidx]=1/slope*v[vnegidx]\n",
    "    return v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "07e34f59",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(22210, 22344)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "findfont: Font family ['arial'] not found. Falling back to DejaVu Sans.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(22210, 22344)\n",
      "(22210, 22344)\n",
      "(22210, 22344)\n",
      "(22210, 22344)\n",
      "(22210, 22344)\n",
      "(22210, 22344)\n",
      "(22210, 22344)\n",
      "(22355, 18953)\n",
      "(22355, 18953)\n",
      "(22355, 18953)\n",
      "(22355, 18953)\n",
      "(22355, 18953)\n",
      "(22355, 18953)\n",
      "(22355, 18953)\n",
      "(22355, 18953)\n",
      "(22355, 18953)\n",
      "(22294, 19552)\n",
      "(22294, 19552)\n",
      "(22294, 19552)\n",
      "(22294, 19552)\n",
      "(22294, 19552)\n",
      "(22294, 19552)\n",
      "(22294, 19552)\n",
      "(22294, 19552)\n",
      "(22452, 19616)\n",
      "(22452, 19616)\n",
      "(22452, 19616)\n",
      "(22452, 19616)\n",
      "(22452, 19616)\n",
      "(22452, 19616)\n",
      "(22452, 19616)\n",
      "(22452, 19616)\n"
     ]
    }
   ],
   "source": [
    "#recon examples by clusters\n",
    "gaeClusterPath='/mnt/external_ssd/xinyi/plots/train_gae_starmap/allk20XA_02_dca_over/combinedlogminmax_beforeAct/cluster/leiden_nn10mdist025n_pcs40res0.1epoch9990'\n",
    "with open(gaeClusterPath, 'rb') as input:\n",
    "    gaeclusterlabels = pickle.load(input)\n",
    "gaeclusterlabels=gaeclusterlabels.astype(int)\n",
    "\n",
    "sampleNames=None\n",
    "for s in sampleidx.keys():\n",
    "    sampleidx_s=sampleidx[s] \n",
    "    if sampleNames is None:\n",
    "        sampleNames=np.repeat(s,cellCoord[s].shape[0])\n",
    "    else:\n",
    "        sampleNames=np.concatenate((sampleNames,np.repeat(s,cellCoord[s].shape[0])),axis=None)\n",
    "\n",
    "\n",
    "np.random.seed(3)\n",
    "\n",
    "examplesavepath=os.path.join(plotsavepath,'reconExamples')\n",
    "if not os.path.exists(examplesavepath):\n",
    "    os.mkdir(examplesavepath)\n",
    "model.eval()\n",
    "for s in plot_samples.keys():\n",
    "    training_samples_t=s\n",
    "    for c in np.unique(gaeclusterlabels):\n",
    "#         if c in [0,1]:\n",
    "#             continue\n",
    "#     for c in [0,1]:\n",
    "        cidx=np.arange(cellCoord[training_samples_t].shape[0])[gaeclusterlabels[sampleNames==s]==c]\n",
    "        if cidx.size==0:\n",
    "            continue\n",
    "        cidx=np.random.choice(cidx,8,replace=False)\n",
    "        imgInputnp=loadImage.load_cellCentroid(cellCoord[training_samples_t][cidx],sampleidx[training_samples_t],datadir,diamThresh_mul,ifFlip=False,seed=3,imagename='pi_sum.tif',minmaxscale=True,nchannels=1)\n",
    "        for i in range(imgInputnp.shape[0]):\n",
    "            imgInput=imgInputnp[[i]]\n",
    "            if use_cuda:\n",
    "                imgInput=torch.tensor(imgInput).cuda().float()\n",
    "            recon,z, mu, logvar = model(imgInput)\n",
    "            recon=recon.cpu().detach().numpy()\n",
    "            reconmin=np.min(recon)\n",
    "            reconmax=np.max(recon)\n",
    "            recon=(recon-reconmin)/(reconmax-reconmin)\n",
    "            imgInput=imgInput.cpu().detach().numpy()\n",
    "#             plt.imsave(os.path.join(examplesavepath,s+'cluster'+str(c)+'_'+str(i)+'epoch'+str(plotepoch)+'_recon.jpg'),np.stack((recon.reshape((diamThresh_mul,diamThresh_mul)),np.zeros((diamThresh_mul,diamThresh_mul)),np.zeros((diamThresh_mul,diamThresh_mul))),axis=2))\n",
    "#             plt.imsave(os.path.join(examplesavepath,s+'cluster'+str(c)+'_'+str(i)+'epoch'+str(plotepoch)+'_input.jpg'),np.stack((imgInput.reshape((diamThresh_mul,diamThresh_mul)),np.zeros((diamThresh_mul,diamThresh_mul)),np.zeros((diamThresh_mul,diamThresh_mul))),axis=2))\n",
    "            \n",
    "            fontprops = fm.FontProperties(size=2,family='arial')\n",
    "            fig, ax = plt.subplots(figsize=(1, 1), dpi=800)\n",
    "            ax.imshow(np.stack((recon.reshape((diamThresh_mul,diamThresh_mul)),np.zeros((diamThresh_mul,diamThresh_mul)),np.zeros((diamThresh_mul,diamThresh_mul))),axis=2))\n",
    "            scalebar = AnchoredSizeBar(ax.transData,\n",
    "                           106, u'10 \\u03bcm', 'lower left', \n",
    "                           pad=0.1,\n",
    "                           color='white',\n",
    "                           frameon=False,\n",
    "                           size_vertical=0.3,\n",
    "                           fontproperties=fontprops)\n",
    "\n",
    "            ax.add_artist(scalebar)\n",
    "            ax.set_yticks([])\n",
    "            ax.set_xticks([])\n",
    "            plt.savefig(os.path.join(examplesavepath,s+'cluster'+str(c)+'_'+str(i)+'epoch'+str(plotepoch)+'_recon_scalebar.jpg'),dpi=800)\n",
    "            plt.close()\n",
    "            \n",
    "            fig, ax = plt.subplots(figsize=(1, 1), dpi=800)\n",
    "            ax.imshow(np.stack((imgInput.reshape((diamThresh_mul,diamThresh_mul)),np.zeros((diamThresh_mul,diamThresh_mul)),np.zeros((diamThresh_mul,diamThresh_mul))),axis=2))\n",
    "            scalebar = AnchoredSizeBar(ax.transData,\n",
    "                           106, u'10 \\u03bcm', 'lower left', \n",
    "                           pad=0.1,\n",
    "                           color='white',\n",
    "                           frameon=False,\n",
    "                           size_vertical=0.3,\n",
    "                           fontproperties=fontprops)\n",
    "\n",
    "            ax.add_artist(scalebar)\n",
    "            ax.set_yticks([])\n",
    "            ax.set_xticks([])\n",
    "            plt.savefig(os.path.join(examplesavepath,s+'cluster'+str(c)+'_'+str(i)+'epoch'+str(plotepoch)+'_input_scalebar.jpg'),dpi=800)\n",
    "            plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39e2aba7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#recon examples\n",
    "examplesavepath=os.path.join(plotsavepath,'reconExamples')\n",
    "if not os.path.exists(examplesavepath):\n",
    "    os.mkdir(examplesavepath)\n",
    "model.eval()\n",
    "for s in plot_samples.keys():\n",
    "    training_samples_t=s\n",
    "    imgInputnp=loadImage.load_cellCentroid(cellCoord[training_samples_t][100:105],sampleidx[training_samples_t],datadir,diamThresh_mul,ifFlip=False,seed=3,imagename='pi_sum.tif',minmaxscale=True,nchannels=1)\n",
    "    for i in range(imgInputnp.shape[0]):\n",
    "        imgInput=imgInputnp[[i]]\n",
    "        if use_cuda:\n",
    "            imgInput=torch.tensor(imgInput).cuda().float()\n",
    "        recon,z, mu, logvar = model(imgInput)\n",
    "        recon=recon.cpu().detach().numpy()\n",
    "        reconmin=np.min(recon)\n",
    "        reconmax=np.max(recon)\n",
    "        recon=(recon-reconmin)/(reconmax-reconmin)\n",
    "        imgInput=imgInput.cpu().detach().numpy()\n",
    "        plt.imsave(os.path.join(examplesavepath,s+str(i)+'epoch'+str(plotepoch)+'_recon.jpg'),np.stack((recon.reshape((diamThresh_mul,diamThresh_mul)),np.zeros((diamThresh_mul,diamThresh_mul)),np.zeros((diamThresh_mul,diamThresh_mul))),axis=2))\n",
    "        plt.imsave(os.path.join(examplesavepath,s+str(i)+'epoch'+str(plotepoch)+'_input.jpg'),np.stack((imgInput.reshape((diamThresh_mul,diamThresh_mul)),np.zeros((diamThresh_mul,diamThresh_mul)),np.zeros((diamThresh_mul,diamThresh_mul))),axis=2))"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}