plotEmbedding_10xADFFPE_extBatchCorrection.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('/home/xinyiz/pamrats')\n",
    "sys.path.append('/home/xinyi/anaconda3/envs/pytorch3/lib/python3.10/site-packages/')\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 gae.gae.optimizer as optimizer\n",
    "import gae.gae.model\n",
    "import gae.gae.preprocessing as preprocessing\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",
    "\n",
    "import json\n",
    "import matplotlib.image as mpimg\n",
    "from skimage import io\n",
    "from sklearn.neighbors import NearestNeighbors\n",
    "from sklearn.metrics.pairwise import euclidean_distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed=3\n",
    "\n",
    "np.random.seed(seed)\n",
    "def plotembeddingbyCT(ctlist,savename,excludelist,embedding,savepath,plotname,plotdimx=0,plotdimy=1,savenameAdd='',img=None,ncolors=None):\n",
    "    \n",
    "    celltypes=np.unique(ctlist)\n",
    "    if ncolors is None:\n",
    "        colortest=sns.color_palette(\"husl\", celltypes.size)\n",
    "    else:\n",
    "        colortest=sns.color_palette(\"husl\", ncolors)\n",
    "    fig, ax = plt.subplots(dpi=400)\n",
    "    if not img is None:\n",
    "        plt.imshow(img)\n",
    "    for ct in celltypes:\n",
    "        if ct in excludelist:\n",
    "            continue\n",
    "        idx=(ctlist==ct)\n",
    "        if not img is None:\n",
    "            ax.scatter(\n",
    "                embedding[idx, plotdimy],\n",
    "                embedding[idx, plotdimx],\n",
    "                color=colortest[int(ct)],label=ct,s=1.5,alpha=0.5\n",
    "                )\n",
    "        else:\n",
    "            ax.scatter(\n",
    "                embedding[idx, plotdimx],\n",
    "                embedding[idx, plotdimy],\n",
    "                color=colortest[int(ct)],label=ct,s=2.5,alpha=1\n",
    "                )\n",
    "\n",
    "    plt.gca().set_aspect('equal', 'datalim')\n",
    "    fig.set_figheight(5)\n",
    "    fig.set_figwidth(5)\n",
    "    box = ax.get_position()\n",
    "    ax.set_position([box.x0, box.y0 + box.height * 0.1,\n",
    "                     box.width, box.height * 0.9])\n",
    "    ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.05),\n",
    "              fancybox=True,ncol=5, shadow=True,prop={'size': 6})\n",
    "    plt.title(plotname+' embedding', fontsize=12)\n",
    "    plt.savefig(os.path.join(savepath,savename+savenameAdd+'.jpg'))\n",
    "    plt.close('all')\n",
    "    \n",
    "    gc.collect()\n",
    "    \n",
    "def plotembeddingbyCT_str(ctlist,savename,excludelist,embedding,savepath,plotname,plotdimx=0,plotdimy=1,savenameAdd=''):\n",
    "    celltypes=np.unique(ctlist)\n",
    "    celltypes_dict={}\n",
    "    idx=0\n",
    "    for ct in celltypes:\n",
    "        celltypes_dict[ct]=idx\n",
    "        idx+=1\n",
    "        \n",
    "    colortest=sns.color_palette(\"husl\", celltypes.size)\n",
    "    fig, ax = plt.subplots(dpi=400)\n",
    "    for ct in celltypes:\n",
    "        if ct in excludelist:\n",
    "            continue\n",
    "        idx=(ctlist==ct)\n",
    "        ax.scatter(\n",
    "            embedding[idx, plotdimx],\n",
    "            embedding[idx, plotdimy],\n",
    "            color=colortest[celltypes_dict[ct]],label=ct,s=1.5,alpha=0.5\n",
    "            )\n",
    "\n",
    "    plt.gca().set_aspect('equal', 'datalim')\n",
    "    fig.set_figheight(5)\n",
    "    fig.set_figwidth(5)\n",
    "    box = ax.get_position()\n",
    "    ax.set_position([box.x0, box.y0 + box.height * 0.1,\n",
    "                     box.width, box.height * 0.9])\n",
    "    # Put a legend below current axis\n",
    "    ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.05),\n",
    "              fancybox=True,ncol=2, shadow=True,prop={'size': 6})\n",
    "    plt.title(plotname+' embedding', fontsize=24)\n",
    "    plt.savefig(os.path.join(savepath,savename+savenameAdd+'.jpg'))\n",
    "    plt.close('all')\n",
    "    \n",
    "    gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(seed)\n",
    "def plotembeddingbyCT_contrast(ctlist,savename,excludelist,embedding,savepath,plotname,plotdimx=0,plotdimy=1,savenameAdd='',maxplot=None): \n",
    "    celltypes=np.unique(ctlist)\n",
    "    celltypes_dict={}\n",
    "    idx=0\n",
    "    for ct in celltypes:\n",
    "        celltypes_dict[ct]=idx\n",
    "        idx+=1\n",
    "\n",
    "    colortest=sns.color_palette(\"tab10\")\n",
    "    if not os.path.exists(os.path.join(savepath)):\n",
    "        os.makedirs(savepath)\n",
    "\n",
    "    for ct in celltypes:\n",
    "        if maxplot and int(ct)>maxplot:\n",
    "            continue\n",
    "        fig, ax = plt.subplots()\n",
    "        if ct == 'Unassigned':\n",
    "            continue\n",
    "\n",
    "        idx=(ctlist!=ct)\n",
    "        ax.scatter(\n",
    "            embedding[idx, plotdimx],\n",
    "            embedding[idx, plotdimy],\n",
    "            color=colortest[1],label='others',s=5,alpha=1\n",
    "            )\n",
    "\n",
    "        idx=(ctlist==ct)\n",
    "        ax.scatter(\n",
    "            embedding[idx, plotdimx],\n",
    "            embedding[idx, plotdimy],\n",
    "            color=colortest[0],label=ct,s=6,alpha=1\n",
    "            )\n",
    "\n",
    "        plt.gca().set_aspect('equal', 'datalim')\n",
    "        fig.set_figheight(10)\n",
    "        fig.set_figwidth(10)\n",
    "        ax.legend()\n",
    "        plt.title(plotname+' embedding', fontsize=24)\n",
    "        plt.gcf().savefig(os.path.join(savepath,savename+'_'+str(ct)+savenameAdd+'.jpg'))\n",
    "        plt.close('all')\n",
    "        gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plotCTcomp(labels,ctlist,savepath,savenamecluster,addname=''):\n",
    "    res=np.zeros((np.unique(labels).size,np.unique(ctlist).size))\n",
    "    for li in range(res.shape[0]):\n",
    "        l=np.unique(labels)[li]\n",
    "        nl=np.sum(labels==l)\n",
    "        ctlist_l=ctlist[labels==l]\n",
    "        for ci in range(res.shape[1]):\n",
    "            c=np.unique(ctlist)[ci]\n",
    "            res[li,ci]=np.sum(ctlist_l==c)\n",
    "#             res[li,ci]=np.sum(ctlist_l==c)/nl\n",
    "    if not byCT:\n",
    "        addname+=''\n",
    "        for li in range(res.shape[0]):\n",
    "            l=np.unique(labels)[li]\n",
    "            nl=np.sum(labels==l)\n",
    "            res[li]=res[li]/nl\n",
    "    else:\n",
    "        addname+='_normbyCT'\n",
    "        for ci in range(res.shape[1]):\n",
    "            c=np.unique(ctlist)[ci]\n",
    "            nc=np.sum(ctlist==c)\n",
    "            res[:,ci]=res[:,ci]/nc\n",
    "    \n",
    "    fig, ax = plt.subplots(figsize=(10, 10))\n",
    "    im = ax.imshow(res,cmap='binary')\n",
    "    ax.set_yticks(np.arange(np.unique(labels).size))\n",
    "    ax.set_yticklabels(np.unique(labels))\n",
    "    ax.set_xticks(np.arange(np.unique(ctlist).size))\n",
    "    ax.set_xticklabels(np.unique(ctlist))\n",
    "    plt.setp(ax.get_xticklabels(), rotation=45, ha=\"right\",rotation_mode=\"anchor\")\n",
    "    fig.tight_layout()\n",
    "    plt.savefig(os.path.join(savepath,savenamecluster+'_ctComposition'+addname+'.jpg'))\n",
    "    plt.close()\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/xinyi/.local/lib/python3.8/site-packages/anndata/_core/anndata.py:1830: UserWarning: Variable names are not unique. To make them unique, call `.var_names_make_unique`.\n",
      "  utils.warn_names_duplicates(\"var\")\n"
     ]
    }
   ],
   "source": [
    "datadir='/home/xinyi/staci_validation/10xVisiumADFFPE/'\n",
    "tissuepospath='VisiumFFPE_Mouse_Brain_Alzheimers_AppNote_aggr_tissue_positions_list.csv'\n",
    "tissuepos=pd.read_csv(os.path.join(datadir,tissuepospath),header=None)\n",
    "\n",
    "tissuepos.index=tissuepos.iloc[:,0]\n",
    "scalefactor=0.150015\n",
    "libraryID=pd.read_csv(os.path.join(datadir,'VisiumFFPE_Mouse_Brain_Alzheimers_AppNote_aggregation.csv'))\n",
    "features=scanpy.read_10x_h5(os.path.join(datadir,'VisiumFFPE_Mouse_Brain_Alzheimers_AppNote_filtered_feature_bc_matrix.h5'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "features.obs['x']=tissuepos.loc[features.obs.index,4]*scalefactor\n",
    "features.obs['y']=tissuepos.loc[features.obs.index,5]*scalefactor\n",
    "features.obs['barcodes']=features.obs.index\n",
    "sampleidx=features.obs['barcodes'].apply(lambda x: x.split('-')[1])\n",
    "samplenameList=np.zeros(sampleidx.size).astype(str)\n",
    "for s in np.unique(sampleidx):\n",
    "    samplenameList[sampleidx==s]=libraryID['library_id'][int(s)-1]\n",
    "features.obs['samplename']=samplenameList\n",
    "features.var_names_make_unique()\n",
    "features=features[:,np.array(np.sum(features.X,axis=0)>3).flatten()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_samples=np.unique(samplenameList)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "ifplot=True\n",
    "ifcluster=True\n",
    "\n",
    "inverseAct='leakyRelu'\n",
    "# inverseAct=None\n",
    "plottype='umap'\n",
    "pca=PCA()\n",
    "npc=50 #for pca var ration\n",
    "npc_plot=10 #for pairwise pc plots\n",
    "\n",
    "minCells=15 #min number of cells for analysis\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.2,0.3,0.4,0.5,0.6,0.8]\n",
    "plotepoch=9310\n",
    "savenameAdd=''\n",
    "\n",
    "\n",
    "use_cuda=True\n",
    "fastmode=False #Validate during training pass\n",
    "seed=3\n",
    "useSavedMaskedEdges=False\n",
    "maskedgeName='knn6_connectivity'\n",
    "nneighbors=6\n",
    "hidden1=32 #Number of units in hidden layer 1\n",
    "hidden2=32 #Number of units in hidden layer 2\n",
    "fc_dim1=32\n",
    "\n",
    "dropout=0.01\n",
    "model_str='gcn_vae_xa_e2_d1_dca_sharded'\n",
    "adj_decodeName=None #gala or None\n",
    "plot_sample_X=['mnn']\n",
    "plotRecon='' #'meanRecon'\n",
    "standardizeX=False\n",
    "name='10xAD_01_dca_over_mnn' \n",
    "logsavepath='/data/xinyi/log/train_gae_visium_validation/'+name\n",
    "modelsavepath='/data/xinyi/models/train_gae_visium_validation/'+name\n",
    "plotsavepath='/data/xinyi/plots/train_gae_visium_validation/'+name\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getA_knn(samplename,k,a_mode,savepath=None):\n",
    "    sobj_coord_np=features.obs.loc[features.obs.index[features.obs['samplename']==samplename],['x','y']].to_numpy()\n",
    "    nbrs = NearestNeighbors(n_neighbors=k, algorithm='ball_tree').fit(sobj_coord_np)\n",
    "    a=nbrs.kneighbors_graph(sobj_coord_np,mode=a_mode)\n",
    "    if a_mode=='connectivity':\n",
    "        a=a-sp.identity(sobj_coord_np.shape[0],format='csr')\n",
    "    if a_mode=='distance':\n",
    "        a[a!=0]=1/a[a!=0]\n",
    "    if savepath !=None:\n",
    "        sp.save_npz(savepath,a)\n",
    "    return a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/xinyi/.local/lib/python3.8/site-packages/scanpy/preprocessing/_highly_variable_genes.py:478: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
      "  hvg = hvg.append(missing_hvg, ignore_index=True)\n",
      "/home/xinyi/.local/lib/python3.8/site-packages/scanpy/preprocessing/_highly_variable_genes.py:478: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
      "  hvg = hvg.append(missing_hvg, ignore_index=True)\n",
      "/home/xinyi/.local/lib/python3.8/site-packages/scanpy/preprocessing/_highly_variable_genes.py:478: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
      "  hvg = hvg.append(missing_hvg, ignore_index=True)\n",
      "/home/xinyi/.local/lib/python3.8/site-packages/scanpy/preprocessing/_highly_variable_genes.py:478: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
      "  hvg = hvg.append(missing_hvg, ignore_index=True)\n",
      "/home/xinyi/.local/lib/python3.8/site-packages/scanpy/preprocessing/_highly_variable_genes.py:478: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
      "  hvg = hvg.append(missing_hvg, ignore_index=True)\n",
      "/home/xinyi/.local/lib/python3.8/site-packages/scanpy/preprocessing/_highly_variable_genes.py:478: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
      "  hvg = hvg.append(missing_hvg, ignore_index=True)\n",
      "/home/xinyi/.local/lib/python3.8/site-packages/scanpy/preprocessing/_highly_variable_genes.py:478: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
      "  hvg = hvg.append(missing_hvg, ignore_index=True)\n",
      "/home/xinyi/.local/lib/python3.8/site-packages/scanpy/preprocessing/_highly_variable_genes.py:478: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
      "  hvg = hvg.append(missing_hvg, ignore_index=True)\n",
      "/home/xinyi/.local/lib/python3.8/site-packages/scanpy/preprocessing/_highly_variable_genes.py:478: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
      "  hvg = hvg.append(missing_hvg, ignore_index=True)\n",
      "/home/xinyi/.local/lib/python3.8/site-packages/scanpy/preprocessing/_highly_variable_genes.py:478: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
      "  hvg = hvg.append(missing_hvg, ignore_index=True)\n",
      "/home/xinyi/.local/lib/python3.8/site-packages/scanpy/preprocessing/_highly_variable_genes.py:478: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
      "  hvg = hvg.append(missing_hvg, ignore_index=True)\n",
      "/home/xinyi/.local/lib/python3.8/site-packages/scanpy/preprocessing/_highly_variable_genes.py:478: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
      "  hvg = hvg.append(missing_hvg, ignore_index=True)\n",
      "/home/xinyi/anaconda3/envs/pytorch3/lib/python3.10/site-packages/mnnpy/utils.py:30: NumbaPerformanceWarning: \u001b[1m\u001b[1m\u001b[1mnp.dot() is faster on contiguous arrays, called on (array(float32, 1d, A), array(float32, 1d, A))\u001b[0m\u001b[0m\u001b[0m\n",
      "  dist[i, j] = np.dot(m[i], n[j])\n",
      "/home/xinyi/anaconda3/envs/pytorch3/lib/python3.10/site-packages/mnnpy/utils.py:205: NumbaPerformanceWarning: \u001b[1m\u001b[1m\u001b[1mnp.dot() is faster on contiguous arrays, called on (array(float32, 1d, C), array(float32, 1d, A))\u001b[0m\u001b[0m\u001b[0m\n",
      "  scale = np.dot(working, grad)\n",
      "/home/xinyi/anaconda3/envs/pytorch3/lib/python3.10/site-packages/mnnpy/utils.py:215: NumbaPerformanceWarning: \u001b[1m\u001b[1m\u001b[1mnp.dot() is faster on contiguous arrays, called on (array(float32, 1d, C), array(float32, 1d, A))\u001b[0m\u001b[0m\u001b[0m\n",
      "  curproj = np.dot(grad, curcell)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cython module _utils not initialized. Fallback to python.\n",
      "Performing cosine normalization...\n",
      "Starting MNN correct iteration. Reference batch: 0\n",
      "Step 1 of 11: processing batch 1\n",
      "  Looking for MNNs...\n",
      "  Computing correction vectors...\n",
      "  Adjusting variance...\n",
      "  Applying correction...\n",
      "Step 2 of 11: processing batch 2\n",
      "  Looking for MNNs...\n",
      "  Computing correction vectors...\n",
      "  Adjusting variance...\n",
      "  Applying correction...\n",
      "Step 3 of 11: processing batch 3\n",
      "  Looking for MNNs...\n",
      "  Computing correction vectors...\n",
      "  Adjusting variance...\n",
      "  Applying correction...\n",
      "Step 4 of 11: processing batch 4\n",
      "  Looking for MNNs...\n",
      "  Computing correction vectors...\n",
      "  Adjusting variance...\n",
      "  Applying correction...\n",
      "Step 5 of 11: processing batch 5\n",
      "  Looking for MNNs...\n",
      "  Computing correction vectors...\n",
      "  Adjusting variance...\n",
      "  Applying correction...\n",
      "Step 6 of 11: processing batch 6\n",
      "  Looking for MNNs...\n",
      "  Computing correction vectors...\n",
      "  Adjusting variance...\n",
      "  Applying correction...\n",
      "Step 7 of 11: processing batch 7\n",
      "  Looking for MNNs...\n",
      "  Computing correction vectors...\n",
      "  Adjusting variance...\n",
      "  Applying correction...\n",
      "Step 8 of 11: processing batch 8\n",
      "  Looking for MNNs...\n",
      "  Computing correction vectors...\n",
      "  Adjusting variance...\n",
      "  Applying correction...\n",
      "Step 9 of 11: processing batch 9\n",
      "  Looking for MNNs...\n",
      "  Computing correction vectors...\n",
      "  Adjusting variance...\n",
      "  Applying correction...\n",
      "Step 10 of 11: processing batch 10\n",
      "  Looking for MNNs...\n",
      "  Computing correction vectors...\n",
      "  Adjusting variance...\n",
      "  Applying correction...\n",
      "Step 11 of 11: processing batch 11\n",
      "  Looking for MNNs...\n",
      "  Computing correction vectors...\n",
      "  Adjusting variance...\n",
      "  Applying correction...\n",
      "MNN correction complete. Gathering output...\n",
      "Packing AnnData object...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/xinyi/.local/lib/python3.8/site-packages/scipy/sparse/_index.py:137: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil_matrix is more efficient.\n",
      "  self._set_arrayXarray_sparse(i, j, x)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done.\n",
      "Transgenic_17p9_rep1\n",
      "Transgenic_17p9_rep2\n",
      "Transgenic_2p5_rep1\n",
      "Transgenic_2p5_rep2\n",
      "Transgenic_5p7_rep1\n",
      "Transgenic_5p7_rep2\n",
      "Wildtype_13p4_rep1\n",
      "Wildtype_13p4_rep2\n",
      "Wildtype_2p5_rep1\n",
      "Wildtype_2p5_rep2\n",
      "Wildtype_5p7_rep1\n",
      "Wildtype_5p7_rep2\n",
      "17186\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n"
     ]
    }
   ],
   "source": [
    "savedirlist={}\n",
    "featureslist={}\n",
    "features_raw_list={}\n",
    "adj_list={}\n",
    "coordlist={}\n",
    "commonGenes=[]\n",
    "    \n",
    "    \n",
    "if plot_sample_X[0]=='logminmax':\n",
    "    for samplename in np.unique(features.obs['samplename']):\n",
    "        print(samplename)\n",
    "    \n",
    "        featurelog_train=np.log2(features.X[features.obs['samplename']==samplename].toarray()+1/2)\n",
    "        scaler = MinMaxScaler()\n",
    "        featurelog_train_minmax=np.transpose(scaler.fit_transform(np.transpose(featurelog_train)))\n",
    "        featureslist[samplename+'X_'+plot_sample_X[0]]=torch.tensor(featurelog_train_minmax)\n",
    "elif plot_sample_X[0]=='mnn':\n",
    "    feature_corrected=features.copy()\n",
    "    scanpy.pp.normalize_total(feature_corrected, target_sum=1e4)\n",
    "    scanpy.pp.log1p(feature_corrected)\n",
    "    scanpy.pp.highly_variable_genes(feature_corrected,batch_key='samplename')\n",
    "    featureslist_correct={}\n",
    "    for b in np.unique(feature_corrected.obs['samplename']):\n",
    "        featureslist_correct[b]=feature_corrected[feature_corrected.obs['samplename']==b]\n",
    "    featureslist_correct=scanpy.external.pp.mnn_correct(featureslist_correct['Transgenic_17p9_rep1'],\n",
    "                                            featureslist_correct['Transgenic_17p9_rep2'],\n",
    "                                            featureslist_correct['Transgenic_2p5_rep1'],\n",
    "                                            featureslist_correct['Transgenic_2p5_rep2'],\n",
    "                                            featureslist_correct['Transgenic_5p7_rep1'],\n",
    "                                            featureslist_correct['Transgenic_5p7_rep2'],\n",
    "                                            featureslist_correct['Wildtype_13p4_rep1'],\n",
    "                                            featureslist_correct['Wildtype_13p4_rep2'],\n",
    "                                            featureslist_correct['Wildtype_2p5_rep1'],\n",
    "                                            featureslist_correct['Wildtype_2p5_rep2'],\n",
    "                                            featureslist_correct['Wildtype_5p7_rep1'],\n",
    "                                            featureslist_correct['Wildtype_5p7_rep2'],batch_key='samplename',var_subset=feature_corrected.var.index[feature_corrected.var.highly_variable],do_concatenate=False)\n",
    "    for b in range(len(featureslist_correct[0])):\n",
    "        samplename=featureslist_correct[0][b].obs['samplename'][0]\n",
    "        featureslist[samplename+'X_'+plot_sample_X[0]]=torch.tensor(featureslist_correct[0][b].X.toarray())\n",
    "elif plot_sample_X[0]=='combat':\n",
    "    feature_corrected=features.copy()\n",
    "    scanpy.pp.normalize_total(feature_corrected, target_sum=1e4)\n",
    "    scanpy.pp.log1p(feature_corrected)\n",
    "    scanpy.pp.scale(feature_corrected)\n",
    "    scanpy.pp.combat(feature_corrected,key='samplename')\n",
    "    for b in np.unique(features.obs['samplename']):\n",
    "        featureslist[b+'X_'+plot_sample_X[0]]=torch.tensor(feature_corrected[feature_corrected.obs['samplename']==b])\n",
    "for samplename in np.unique(features.obs['samplename']):\n",
    "    print(samplename)    \n",
    "    adj_list[samplename]=getA_knn(samplename,nneighbors+1,'connectivity')\n",
    "    features_raw_list[samplename+'X_'+'raw']=torch.tensor(features.X[features.obs['samplename']==samplename].toarray())\n",
    "    \n",
    "    \n",
    "num_features=features.shape[1]\n",
    "print(num_features)\n",
    "adjnormlist={}\n",
    "pos_weightlist={}\n",
    "normlist={}\n",
    "for ai in adj_list.keys():\n",
    "    adjnormlist[ai]=preprocessing.preprocess_graph(adj_list[ai])\n",
    "    \n",
    "    pos_weightlist[ai] = torch.tensor(float(adj_list[ai].shape[0] * adj_list[ai].shape[0] - adj_list[ai].sum()) / adj_list[ai].sum()) #using full unmasked adj\n",
    "    normlist[ai] = adj_list[ai].shape[0] * adj_list[ai].shape[0] / float((adj_list[ai].shape[0] * adj_list[ai].shape[0] - adj_list[ai].sum()) * 2)\n",
    "    \n",
    "    adj_label=adj_list[ai] + sp.eye(adj_list[ai].shape[0])\n",
    "    adj_list[ai]=torch.tensor(adj_label.toarray())\n",
    "    \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load model\n",
    "num_nodes,num_features = list(featureslist.values())[0].shape\n",
    "if model_str=='gcn_vae_xa':\n",
    "    model  = gae.gae.model.GCNModelVAE_XA(num_features, hidden1, hidden2,fc_dim1,fc_dim2,fc_dim3,fc_dim4, dropout)\n",
    "elif model_str=='fc1':\n",
    "    model  = gae.gae.model.FCVAE1(num_features, hidden1,dropout)\n",
    "elif model_str == 'gcn_vae_xa_e2_d1':\n",
    "    model  = gae.gae.model.GCNModelVAE_XA_e2_d1(num_features, hidden1,hidden2, dropout)\n",
    "elif model_str == 'gcn_vae_gcnX_inprA':\n",
    "    model = gae.gae.model.GCNModelVAE_gcnX_inprA(num_features, hidden1, hidden2,gcn_dim1, dropout)\n",
    "elif model_str=='fc1_dca':\n",
    "    model = gae.gae.model.FCVAE1_DCA(num_features, hidden1,fc_dim1, dropout)\n",
    "elif model_str=='fc1_dca_sharded':\n",
    "    model = gae.gae.model.FCVAE1_DCA_sharded(num_features, hidden1,fc_dim1, dropout)\n",
    "elif model_str=='gcn_vae_xa_e2_d1_dca':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e2_d1_DCA(num_features, hidden1,hidden2,fc_dim1, dropout)\n",
    "elif model_str=='gcn_vae_xa_e2_d1_dca_sharded':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e2_d1_DCA_sharded(num_features, hidden1,hidden2,fc_dim1, dropout)\n",
    "elif model_str=='gcn_vae_xa_e2_d1_dcaFork':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e2_d1_DCAfork(num_features, hidden1,hidden2,fc_dim1, dropout)\n",
    "elif model_str=='gcn_vae_xa_e2_d1_dcaElemPi':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e2_d1_DCAelemPi(num_features, hidden1,hidden2,fc_dim1, dropout,shareGenePi)\n",
    "elif model_str=='gcn_vae_xa_e2_d1_dcaConstantDisp':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e2_d1_DCA_constantDisp(num_features, hidden1,hidden2,fc_dim1, dropout,shareGenePi)\n",
    "else:\n",
    "    print('model not found')\n",
    "model.load_state_dict(torch.load(os.path.join(modelsavepath,str(plotepoch)+'.pt')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "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": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(seed)\n",
    "def clusterLeiden_single(inArray,n_neighbors,n_pcs,min_dist,resolution,randseed=seed):\n",
    "    n_pcs=np.min([inArray.shape[0]-1,inArray.shape[1]-1,n_pcs])\n",
    "    adata=ad.AnnData(inArray)\n",
    "    scanpy.tl.pca(adata, svd_solver='arpack')\n",
    "    scanpy.pp.neighbors(adata, n_neighbors=n_neighbors, n_pcs=n_pcs)\n",
    "    scanpy.tl.umap(adata,min_dist=min_dist,random_state=randseed)\n",
    "    scanpy.tl.leiden(adata,resolution=resolution,random_state=randseed)\n",
    "    return adata.obs['leiden'].to_numpy()\n",
    "\n",
    "def clusterLeiden(inArray,n_neighbors,n_pcs,min_dist,resolution,sobj_coord_np,randseed=seed):\n",
    "    for r in resolution:\n",
    "        clusterRes=clusterLeiden_single(inArray,n_neighbors,n_pcs,min_dist,r,randseed=seed)\n",
    "#         print(clusterRes.shape)\n",
    "        savenamecluster='leiden_nn'+str(n_neighbors)+'mdist0'+str(int(min_dist*100))+'n_pcs'+str(n_pcs)+'res'+str(r)+'epoch'+str(plotepoch)\n",
    "        with open(os.path.join(clustersavedir,savenamecluster), 'wb') as output:\n",
    "            pickle.dump(clusterRes, output, pickle.HIGHEST_PROTOCOL)\n",
    "        plotembeddingbyCT_contrast(clusterRes,'leiden_location',[],sobj_coord_np,os.path.join(savedir,'contrast'),'location'+' of '+s,savenameAdd=savenamecluster)\n",
    "\n",
    "def clusterLeiden_allsample(embedding,savedir,clustersavedir,inArray,n_neighbors,n_pcs,min_dist,resolution,sobj_coord_np,samplenameList,randseed=seed):\n",
    "    for r in resolution:\n",
    "        clusterRes=clusterLeiden_single(inArray,n_neighbors,n_pcs,min_dist,r,randseed=seed)\n",
    "        savenamecluster='leiden_nn'+str(n_neighbors)+'mdist0'+str(int(min_dist*100))+'n_pcs'+str(n_pcs)+'res'+str(r)+'epoch'+str(plotepoch)\n",
    "        with open(os.path.join(clustersavedir,savenamecluster), 'wb') as output:\n",
    "            pickle.dump(clusterRes, output, pickle.HIGHEST_PROTOCOL)\n",
    "        plotembeddingbyCT(clusterRes,'leiden',[],embedding,savedir,plottype+' of all samples',savenameAdd=savenamecluster)\n",
    "        plotembeddingbyCT_contrast(clusterRes,'leiden',[],embedding,os.path.join(savedir,'contrast'),plottype+' of all samples',savenameAdd=savenamecluster,maxplot=50)\n",
    "\n",
    "        with open(os.path.join(clustersavedir,savenamecluster), 'rb') as output:\n",
    "            clusterRes=pickle.load(output)\n",
    "        for s in plot_samples:\n",
    "            sidx=(samplenameList==s)\n",
    "            img=None\n",
    "#             img=mpimg.imread(os.path.join(datadir,'spatial',s,'tissue_hires_image.png'))\n",
    "            plotembeddingbyCT(clusterRes[sidx],'leiden_location'+s,[],sobj_coord_np[sidx],savedir,'location'+' of '+s,savenameAdd=savenamecluster+'_noImg',img=img,ncolors=np.unique(clusterRes).size)\n",
    "            plotembeddingbyCT_contrast(clusterRes[sidx],'leiden_location'+s,[],sobj_coord_np[sidx],os.path.join(savedir,'contrast'),'location'+' of '+s,savenameAdd=savenamecluster,maxplot=50)\n",
    "\n",
    "         "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n",
      "0\n"
     ]
    }
   ],
   "source": [
    "#compute embeddings\n",
    "mulist={}\n",
    "for s in np.unique(samplenameList):\n",
    "    adj=adj_list[s]\n",
    "    adj_norm = preprocessing.preprocess_graph(adj)\n",
    "    adj_decode=None\n",
    "    if adj_decodeName == 'gala':\n",
    "        adj_decode=preprocessing.preprocess_graph_sharp(adj)\n",
    "    for xcorr in plot_sample_X:\n",
    "        samplename=s+'X_'+xcorr\n",
    "        features_s=featureslist[samplename]\n",
    "        \n",
    "        model.eval()\n",
    "        if adj_decodeName==None:\n",
    "            adj_recon,mu,logvar,z, features_recon = model(features_s, adj_norm)\n",
    "        else:\n",
    "            adj_recon,mu,logvar,z, features_recon = model(features_s, adj_norm,adj_decode)\n",
    "        if inverseAct=='leakyRelu':\n",
    "            muplot=inverseLeakyRelu(mu.cpu().detach().numpy())\n",
    "        else:\n",
    "            muplot=mu.cpu().detach().numpy()\n",
    "        if plotRecon:\n",
    "            if plotRecon=='meanRecon':\n",
    "                mulist[samplename]=features_recon[3].cpu().detach().numpy()\n",
    "        else:\n",
    "            mulist[samplename]=muplot\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# combine all latents to one plot \n",
    "np.random.seed(seed)\n",
    "for xcorr in plot_sample_X:\n",
    "    latents=None\n",
    "    samplenameList_plot=None\n",
    "    sobj_coord_np=None\n",
    "    \n",
    "    for s in np.unique(samplenameList):\n",
    "        samplename=s+'X_'+xcorr\n",
    "        muplot=np.copy(mulist[samplename])\n",
    "            \n",
    "        if latents is None:\n",
    "            latents=muplot\n",
    "            sobj_coord_np=features.obs.loc[:,('x','y')].to_numpy()[samplenameList==s]\n",
    "            samplenameList_plot=np.repeat(s,muplot.shape[0])\n",
    "        else:\n",
    "            latents=np.vstack((latents,muplot))\n",
    "            sobj_coord_np=np.concatenate((sobj_coord_np,features.obs.loc[:,('x','y')].to_numpy()[samplenameList==s]),axis=0)\n",
    "            samplenameList_plot=np.concatenate((samplenameList_plot,np.repeat(s,muplot.shape[0])),axis=None)\n",
    "\n",
    "    sampledir=os.path.join(plotsavepath,'combined'+xcorr)\n",
    "    if inverseAct:\n",
    "        sampledir+='_beforeAct'\n",
    "    savedir=os.path.join(sampledir,'embedding_'+plottype)\n",
    "    clustersavedir=os.path.join(sampledir,'cluster')\n",
    "    if not os.path.exists(sampledir):\n",
    "        os.mkdir(sampledir)\n",
    "    if not os.path.exists(savedir):\n",
    "        os.mkdir(savedir)\n",
    "    if not os.path.exists(clustersavedir):\n",
    "        os.mkdir(clustersavedir)\n",
    "    \n",
    "    if plottype=='umap':\n",
    "        npc_plot=2\n",
    "        reducer = umap.UMAP(n_neighbors=n_neighbors,min_dist=min_dist,random_state=seed)\n",
    "        embedding = reducer.fit_transform(latents)\n",
    "        savenameAdd='_nn'+str(n_neighbors)+'mdist0'+str(int(min_dist*100))+'epoch'+str(plotepoch)\n",
    "    elif plottype=='pca':\n",
    "        pca.fit(latents)\n",
    "        fig, ax = plt.subplots(dpi=400)\n",
    "        fig.set_figheight(2.5)\n",
    "        fig.set_figwidth(10)\n",
    "        plt.bar(np.arange(npc),pca.explained_variance_ratio_[:npc])\n",
    "        plt.savefig(os.path.join(savedir,'varRatio_'+str(npc)+'_epoch'+str(plotepoch)+'.jpg'))\n",
    "        plt.close()\n",
    "        fig, ax = plt.subplots(dpi=400)\n",
    "        fig.set_figheight(2.5)\n",
    "        fig.set_figwidth(10)\n",
    "        plt.bar(np.arange(npc),pca.explained_variance_[:npc])\n",
    "        plt.savefig(os.path.join(savedir,'var_'+str(npc)+'_epoch'+str(plotepoch)+'.jpg'))\n",
    "        plt.close()\n",
    "        embedding=pca.transform(latents)\n",
    "#         embedding=pca.fit_transform(latents)\n",
    "        savenameAdd='_epoch'+str(plotepoch)\n",
    "    if ifplot:\n",
    "        for dim1 in range(npc_plot-1):\n",
    "            for dim2 in range(dim1+1,npc_plot):\n",
    "                plotembeddingbyCT_str(samplenameList,'sample',[],embedding,savedir,plottype+'of all samples',plotdimx=dim1,plotdimy=dim2,savenameAdd=savenameAdd+'_pc'+str(dim1)+'pc'+str(dim2))\n",
    "    \n",
    "    if embedding.shape[0]<minCells:\n",
    "        continue\n",
    "    if ifcluster:\n",
    "        if 'leiden' in clustermethod:\n",
    "            clusterLeiden_allsample(embedding,savedir,clustersavedir,latents,n_neighbors,n_pcs,min_dist,resolution,sobj_coord_np,samplenameList,randseed=seed)\n",
    "            assert np.sum(muplot-np.copy(mulist[s+'X_'+xcorr]))==0\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}