plotEmbedding_Starmap_scrnaseq.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7c66eed6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\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",
    "\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0bbc15e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\" \n",
    "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','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.3]\n",
    "# resolution=[0.05,0.1,0.2,0.3,0.5,0.8,1,1.5]\n",
    "resolution=[0.05,0.1,0.2,0.3,0.5,0.8]\n",
    "plotepoch=279\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",
    "\n",
    "use_cuda=True\n",
    "fastmode=False #Validate during training pass\n",
    "seed=3\n",
    "useSavedMaskedEdges=False\n",
    "hidden1=30000 #Number of units in hidden layer 1\n",
    "hidden2=30000 #Number of units in hidden layer 2\n",
    "# hidden3=16\n",
    "fc_dim1=30000\n",
    "# fc_dim2=2112\n",
    "# fc_dim3=2112\n",
    "# fc_dim4=2112\n",
    "# gcn_dim1=2048\n",
    "\n",
    "protein=None #'scaled_binary'\n",
    "# proteinWeights=0.05\n",
    "dropout=0.01\n",
    "# randFeatureSubset=None\n",
    "model_str='fc1_dca_sharded'\n",
    "adj_decodeName=None #gala or None\n",
    "# plot_samples={'control13':'AD_mouse9498','control8':'AD_mouse9735'}\n",
    "plot_sample_X=['logminmax']\n",
    "plotRecon='' #'meanRecon'\n",
    "# plot_sample_X=['corrected','scaled']\n",
    "standardizeX=False\n",
    "tissue='Bone_Marrow'\n",
    "name='allk20XA_02_dca_over_FCXonly_scrnaseq_'+tissue\n",
    "modelsavepath='/data/xinyi/models/train_gae_scrnaseq/'+name\n",
    "plotsavepath='/data/xinyi/plots/train_gae_scrnaseq/'+name\n",
    "datadir='/data/xinyi/STACI/scrnaseq/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f4b00a02",
   "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": 8,
   "id": "d9a881e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "data=scanpy.read_h5ad(os.path.join(datadir,'Immune_ALL_human.h5ad'))\n",
    "data=data[:,np.array(np.sum(data.layers['counts'],axis=0)>3).flatten()]\n",
    "if tissue!= 'all':\n",
    "    data=data[data.obs['tissue']==tissue]\n",
    "    \n",
    "if plot_sample_X[0]=='logminmax':\n",
    "    data_train=np.log2(data.layers['counts']+1/2)\n",
    "    scaler = MinMaxScaler()\n",
    "    data_train=np.transpose(scaler.fit_transform(np.transpose(data_train)))\n",
    "    data_train=torch.tensor(data_train)\n",
    "    \n",
    "\n",
    "# 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",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a7ac2477",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load model\n",
    "num_nodes,num_features = data.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=='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_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",
    "elif model_str=='fc1_dca_sharded':\n",
    "    model = gae.gae.model.FCVAE1_DCA_sharded(num_features, hidden1,fc_dim1, dropout)\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": 10,
   "id": "fa8e0c2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(seed)\n",
    "# np.random.seed(4)\n",
    "def plotembeddingbyCT(ctlist,savename,excludelist,embedding,savepath,plotname,plotdimx=0,plotdimy=1,savenameAdd=''):\n",
    "    celltypes=np.unique(ctlist)\n",
    "#     print(celltypes)\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",
    "#     np.random.shuffle(colortest)\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",
    "#             color=colortest[int(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, shadow=True, ncol=5)\n",
    "#     ax.legend(ncol=3)\n",
    "    plt.title(plotname+' embedding', fontsize=12)\n",
    "    plt.savefig(os.path.join(savepath,savename+savenameAdd+'.jpg'),bbox_inches='tight')\n",
    "#     plt.show()\n",
    "    \n",
    "#     fig.clf()\n",
    "    plt.close('all')\n",
    "    \n",
    "    gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1e586e7b",
   "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=1,alpha=0.5\n",
    "            )\n",
    "\n",
    "        idx=(ctlist==ct)\n",
    "        ax.scatter(\n",
    "            embedding[idx, plotdimx],\n",
    "            embedding[idx, plotdimy],\n",
    "            color=colortest[0],label=ct,s=3,alpha=0.5\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.show()\n",
    "#         nplot+=1\n",
    "        \n",
    "    \n",
    "#         fig.clf()\n",
    "        plt.close('all')\n",
    "        gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ad4a9032",
   "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": 13,
   "id": "146b9360",
   "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(clusterRes,'leiden',[],embedding,savedir,plottype+' of '+s,savenameAdd=savenamecluster)\n",
    "        plotembeddingbyCT_contrast(clusterRes,'leiden',[],embedding,os.path.join(savedir,'contrast'),plottype+' of '+s,savenameAdd=savenamecluster)\n",
    "\n",
    "        plotembeddingbyCT(clusterRes,'leiden_location',[],sobj_coord_np,savedir,'location'+' of '+s,savenameAdd=savenamecluster)\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "972b2cc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(seed)\n",
    "def clusterDBscan_single(inArray,eps,min_samples,n_pcs):\n",
    "    n_pcs=np.min([inArray.shape[0]-1,inArray.shape[1]-1,n_pcs])\n",
    "    inArray=pca.fit_transform(inArray)\n",
    "    labels = DBSCAN(eps=eps, min_samples=min_samples).fit_predict(inArray[:,:n_pcs])\n",
    "#     db = DBSCAN(eps=eps, min_samples=min_samples).fit(inArray[:,:n_pcs])\n",
    "#     core_samples_mask = np.zeros_like(db.labels_, dtype=bool)\n",
    "#     core_samples_mask[db.core_sample_indices_] = True\n",
    "#     labels = db.labels_\n",
    "    return labels\n",
    "\n",
    "def clusterDBscan(inArray,epsL,min_samplesL,n_pcs,sobj_coord_np):\n",
    "    for eps in epsL:\n",
    "        for min_samples in min_samplesL:\n",
    "            clusterRes=clusterDBscan_single(inArray,eps,min_samples,n_pcs)\n",
    "    #         print(clusterRes.shape)\n",
    "            savenamecluster='dbscan_eps'+str(eps)+'msamples'+str(min_samples)+'n_pcs'+str(n_pcs)+'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,'dbscan',[],embedding,savedir,plottype+' of '+s,savenameAdd=savenamecluster)\n",
    "            plotembeddingbyCT_contrast(clusterRes,'dbscan',[],embedding,os.path.join(savedir,'contrast'),plottype+' of '+s,savenameAdd=savenamecluster)\n",
    "\n",
    "            plotembeddingbyCT(clusterRes,'dbscan_location',[],sobj_coord_np,savedir,'location'+' of '+s,savenameAdd=savenamecluster)\n",
    "            plotembeddingbyCT_contrast(clusterRes,'dbscan_location',[],sobj_coord_np,os.path.join(savedir,'contrast'),'location'+' of '+s,savenameAdd=savenamecluster)\n",
    "\n",
    "def clusterDBscan_allsample(embedding,savedir,clustersavedir,inArray,epsL,min_samplesL,n_pcs,sobj_coord_np,samplenameList):\n",
    "    for eps in epsL:\n",
    "        for min_samples in min_samplesL:\n",
    "            clusterRes=clusterDBscan_single(inArray,eps,min_samples,n_pcs)\n",
    "            savenamecluster='dbscan_eps'+str(eps)+'msample'+str(min_samples)+'n_pcs'+str(n_pcs)+'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,'dbscan',[],embedding,savedir,plottype+' of all samples',savenameAdd=savenamecluster)\n",
    "            plotembeddingbyCT_contrast(clusterRes,'dbscan',[],embedding,os.path.join(savedir,'contrast'),plottype+' of all samples',savenameAdd=savenamecluster,maxplot=50)\n",
    "\n",
    "            for s in plot_samples.keys():\n",
    "                sidx=(samplenameList==s)\n",
    "                plotembeddingbyCT(clusterRes[sidx],'dbscan_location'+s,[],sobj_coord_np[sidx],savedir,'location'+' of '+s,savenameAdd=savenamecluster)\n",
    "                plotembeddingbyCT_contrast(clusterRes[sidx],'dbscan_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": 15,
   "id": "319eb751",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(seed)\n",
    "def clusterAgg_single(inArray,ncluster,aggmetric,n_pcs):\n",
    "    n_pcs=np.min([inArray.shape[0]-1,inArray.shape[1]-1,n_pcs])\n",
    "    inArray=pca.fit_transform(inArray)\n",
    "    labels = AgglomerativeClustering(n_clusters=ncluster,affinity=aggmetric).fit_predict(inArray[:,:n_pcs])\n",
    "#     labels = agg.labels_\n",
    "    return labels\n",
    "\n",
    "def clusterAgg(inArray,nclusterL,aggmetricL,n_pcs,sobj_coord_np):\n",
    "    for ncluster in nclusterL:\n",
    "        for aggmetric in aggmetricL:\n",
    "            clusterRes=clusterAgg_single(inArray,ncluster,aggmetric,n_pcs)\n",
    "    #         print(clusterRes.shape)\n",
    "            savenamecluster='agg_ncluster'+str(ncluster)+aggmetric+'n_pcs'+str(n_pcs)+'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,'agg',[],embedding,savedir,plottype+' of '+s,savenameAdd=savenamecluster)\n",
    "            plotembeddingbyCT_contrast(clusterRes,'agg',[],embedding,os.path.join(savedir,'contrast'),plottype+' of '+s,savenameAdd=savenamecluster)\n",
    "\n",
    "            plotembeddingbyCT(clusterRes,'agg_location',[],sobj_coord_np,savedir,'location'+' of '+s,savenameAdd=savenamecluster)\n",
    "            plotembeddingbyCT_contrast(clusterRes,'agg_location',[],sobj_coord_np,os.path.join(savedir,'contrast'),'location'+' of '+s,savenameAdd=savenamecluster)\n",
    "\n",
    "def clusterAgg_allsample(embedding,savedir,clustersavedir,inArray,nclusterL,aggmetricL,n_pcs,sobj_coord_np,samplenameList):\n",
    "    for ncluster in nclusterL:\n",
    "        for aggmetric in aggmetricL:\n",
    "            clusterRes=clusterAgg_single(inArray,ncluster,aggmetric,n_pcs)\n",
    "            savenamecluster='agg_ncluster'+str(ncluster)+aggmetric+'n_pcs'+str(n_pcs)+'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,'agg',[],embedding,savedir,plottype+' of all samples',savenameAdd=savenamecluster)\n",
    "            plotembeddingbyCT_contrast(clusterRes,'agg',[],embedding,os.path.join(savedir,'contrast'),plottype+' of all samples',savenameAdd=savenamecluster,maxplot=50)\n",
    "         "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "e0a6f276",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(seed)\n",
    "def clusterMinibatchKmean_single(inArray,ncluster,n_pcs,batchsize=100):\n",
    "    n_pcs=np.min([inArray.shape[0]-1,inArray.shape[1]-1,n_pcs])\n",
    "    batchsize=int(np.min([(inArray.shape[0]-1)/3,(inArray.shape[1]-1)/3,batchsize]))\n",
    "    inArray=pca.fit_transform(inArray)\n",
    "    labels = MiniBatchKMeans(n_clusters=ncluster,random_state=seed,batch_size=batchsize).fit_predict(inArray[:,:n_pcs])\n",
    "    return labels\n",
    "\n",
    "def clusterMinibatchKmean(inArray,nclusterL,n_pcs,sobj_coord_np):\n",
    "    for ncluster in nclusterL:\n",
    "        clusterRes=clusterMinibatchKmean_single(inArray,ncluster,n_pcs)\n",
    "#         print(clusterRes.shape)\n",
    "        savenamecluster='minibatchkmean_ncluster'+str(ncluster)+'n_pcs'+str(n_pcs)+'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,'minibatchkmean',[],embedding,savedir,plottype+' of '+s,savenameAdd=savenamecluster)\n",
    "        plotembeddingbyCT_contrast(clusterRes,'minibatchkmean',[],embedding,os.path.join(savedir,'contrast'),plottype+' of '+s,savenameAdd=savenamecluster)\n",
    "\n",
    "        plotembeddingbyCT(clusterRes,'minibatchkmean_location',[],sobj_coord_np,savedir,'location'+' of '+s,savenameAdd=savenamecluster)\n",
    "        plotembeddingbyCT_contrast(clusterRes,'minibatchkmean_location',[],sobj_coord_np,os.path.join(savedir,'contrast'),'location'+' of '+s,savenameAdd=savenamecluster)\n",
    "\n",
    "def clusterMinibatchKmean_allsample(embedding,savedir,clustersavedir,inArray,nclusterL,n_pcs,sobj_coord_np,samplenameList):\n",
    "    for ncluster in nclusterL:\n",
    "        clusterRes=clusterMinibatchKmean_single(inArray,ncluster,n_pcs)\n",
    "        savenamecluster='minibatchkmean_ncluster'+str(ncluster)+'n_pcs'+str(n_pcs)+'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,'minibatchkmean',[],embedding,savedir,plottype+' of all samples',savenameAdd=savenamecluster)\n",
    "        plotembeddingbyCT_contrast(clusterRes,'minibatchkmean',[],embedding,os.path.join(savedir,'contrast'),plottype+' of all samples',savenameAdd=savenamecluster,maxplot=50)\n",
    "\n",
    "        for s in plot_samples.keys():\n",
    "            sidx=(samplenameList==s)\n",
    "            plotembeddingbyCT(clusterRes[sidx],'minibatchkmean_location'+s,[],sobj_coord_np[sidx],savedir,'location'+' of '+s,savenameAdd=savenamecluster)\n",
    "            plotembeddingbyCT_contrast(clusterRes[sidx],'minibatchkmean_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": 13,
   "id": "28a7bf05",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([9581, 12303])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8608580",
   "metadata": {},
   "outputs": [],
   "source": [
    "mulist=None\n",
    "batchsize=3000\n",
    "valIdx=np.arange(data_train.shape[0])\n",
    "with torch.no_grad():\n",
    "    model.eval()\n",
    "    nvalBatches=int(np.ceil(data_train.shape[0]/batchsize))\n",
    "    for i in range(nvalBatches):\n",
    "        valIdx_i=valIdx[i*batchsize:min((i+1)*batchsize,valIdx.shape[0])]\n",
    "\n",
    "        features=data_train[valIdx_i].cuda(0).float()\n",
    "        adj_recon,mu,logvar,z, features_recon = model(features,None)\n",
    "        if inverseAct=='leakyRelu':\n",
    "            muplot=inverseLeakyRelu(mu.cpu().detach().numpy())\n",
    "        else:\n",
    "            muplot=mu.cpu().detach().numpy()\n",
    "        if mulist is None:\n",
    "            mulist=muplot\n",
    "        else:\n",
    "            mulist=np.vstack((mulist,muplot))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0ff9ee42",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(modelsavepath,'latent'), 'wb') as output:\n",
    "    pickle.dump(mulist, output, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "27091041",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(modelsavepath,'latent'), 'rb') as output:\n",
    "    mulist=pickle.load(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "4ba90461",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/data/xinyi/models/train_gae_scrnaseq/allk20XA_02_dca_over_FCXonly_scrnaseq_Bone_Marrow'"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "modelsavepath"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abb961e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "npc_plot=2\n",
    "reducer = umap.UMAP(n_neighbors=n_neighbors,min_dist=min_dist,random_state=seed)\n",
    "embedding = reducer.fit_transform(data.layers['counts'])\n",
    "savenameAdd='_nn'+str(n_neighbors)+'mdist0'+str(int(min_dist*100))+'epoch'+str(plotepoch)\n",
    "plotembeddingbyCT(batch,'data_batch',[],embedding,savedir,plottype+'of all samples',plotdimx=dim1,plotdimy=dim2,savenameAdd=savenameAdd+'_pc'+str(dim1)+'pc'+str(dim2))\n",
    "plotembeddingbyCT(celltype,'data_celltype',[],embedding,savedir,plottype+'all samples',plotdimx=dim1,plotdimy=dim2,savenameAdd=savenameAdd+'_pc'+str(dim1)+'pc'+str(dim2))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "72f24ee1",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'muplot' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-18-ed0e9552c05d>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     56\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0;34m'leiden'\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mclustermethod\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     57\u001b[0m         \u001b[0mclusterLeiden_allsample\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0membedding\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0msavedir\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mclustersavedir\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mlatents\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mn_neighbors\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mn_pcs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mmin_dist\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mresolution\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mrandseed\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mseed\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 58\u001b[0;31m         \u001b[0;32massert\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmuplot\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcopy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmulist\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0ms\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;34m'X_'\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0mxcorr\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m==\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     59\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0;34m'dbscan'\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mclustermethod\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     60\u001b[0m         \u001b[0mclusterDBscan_allsample\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0membedding\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0msavedir\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mclustersavedir\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mlatents\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mepslist\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mmin_sampleslist\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mn_pcs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0msobj_coord_np\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0msamplenameList\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'muplot' is not defined"
     ]
    }
   ],
   "source": [
    "# combine all latents to one plot\n",
    "np.random.seed(seed)\n",
    "for xcorr in plot_sample_X:\n",
    "    latents=mulist\n",
    "    celltype=data.obs['final_annotation']\n",
    "    tissue=data.obs['tissue']\n",
    "    batch=data.obs['batch']\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(batch,'batch',[],embedding,savedir,plottype+'of all samples',plotdimx=dim1,plotdimy=dim2,savenameAdd=savenameAdd+'_pc'+str(dim1)+'pc'+str(dim2))\n",
    "            plotembeddingbyCT(celltype,'celltype',[],embedding,savedir,plottype+'all samples',plotdimx=dim1,plotdimy=dim2,savenameAdd=savenameAdd+'_pc'+str(dim1)+'pc'+str(dim2))\n",
    "#             plotembeddingbyCT(tissue,'tissue',[],embedding,savedir,plottype+'all samples',plotdimx=dim1,plotdimy=dim2,savenameAdd=savenameAdd+'_pc'+str(dim1)+'pc'+str(dim2))\n",
    "\n",
    "            plotembeddingbyCT_contrast(celltype,'celltype',[],embedding,os.path.join(savedir,'contrast'),plottype+'all samples',savenameAdd=savenameAdd)    \n",
    "            plotembeddingbyCT_contrast(batch,'batch',[],embedding,os.path.join(savedir,'contrast'),plottype+'all samples',savenameAdd=savenameAdd)    \n",
    "\n",
    "\n",
    "\n",
    "if ifcluster:\n",
    "    if 'leiden' in clustermethod:\n",
    "        clusterLeiden_allsample(embedding,savedir,clustersavedir,latents,n_neighbors,n_pcs,min_dist,resolution,None,batch,randseed=seed)\n",
    "    if 'dbscan' in clustermethod:\n",
    "        clusterDBscan_allsample(embedding,savedir,clustersavedir,latents,epslist,min_sampleslist,n_pcs,sobj_coord_np,samplenameList)\n",
    "    if 'agglomerative' in clustermethod:\n",
    "        clusterAgg_allsample(embedding,savedir,clustersavedir,latents,nclusterlist,aggMetric,n_pcs,sobj_coord_np,samplenameList)\n",
    "    if 'kmeanbatch' in clustermethod:\n",
    "        clusterMinibatchKmean_allsample(embedding,savedir,clustersavedir,latents,nclusterlist,n_pcs,sobj_coord_np,samplenameList)\n"
   ]
  }
 ],
 "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
}