plotEmbedding_Starmap.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\" \n",
    "ifplot=True #if umap plots and clustering plots should be generated\n",
    "ifcluster=True #if clustering should be performed (This is to provide the option of adjusting plotting parameters without re-run clustering.)\n",
    "\n",
    "inverseAct='leakyRelu' #This specifies that the latent should be plotted before the leakyRelu activation function. Use None to plot the latent after leakyRelu.\n",
    "# inverseAct=None\n",
    "plottype='umap' #umap or pca\n",
    "pca=PCA()\n",
    "npc=50 #number of principal components to plot in the variance ratio plot\n",
    "npc_plot=10 #for pairwise pc plots\n",
    "\n",
    "#plotting and clustering parameters as defined in the respective functions\n",
    "minCells=15 #min number of cells for analysis\n",
    "clustermethod=['leiden'] #specify which clustering methods to use\n",
    "#umap/leiden clustering parameters\n",
    "n_neighbors=10\n",
    "min_dist=0.25\n",
    "n_pcs=40 #for clustering\n",
    "resolution=[0.05,0.1,0.2,0.3,0.5,0.8,1,1.5]\n",
    "plotepoch=9270\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",
    "#In the plots and clusters of each cell type, it's also possible to plot/cluster multiple related cell types together.\n",
    "combineCelltype={'glia':['Astro','Micro', 'OPC', 'Oligo'],'CA':['CA1', 'CA2', 'CA3']}\n",
    "\n",
    "#network parameters; should be kept the same as in the training script\n",
    "use_cuda=True\n",
    "fastmode=False #Validate during training pass\n",
    "seed=3\n",
    "useSavedMaskedEdges=False\n",
    "maskedgeName='knn20_connectivity'\n",
    "hidden1=5000 #Number of units in hidden layer 1\n",
    "hidden2=5000 #Number of units in hidden layer 2\n",
    "# hidden3=16\n",
    "fc_dim1=5000\n",
    "# fc_dim2=2112\n",
    "# fc_dim3=2112\n",
    "# fc_dim4=2112\n",
    "\n",
    "dropout=0.01\n",
    "model_str='gcn_vae_xa_e2_d1_dca'\n",
    "adj_decodeName=None #gala or None\n",
    "plot_samples={'disease13':'AD_mouse9494','control13':'AD_mouse9498','disease8':'AD_mouse9723','control8':'AD_mouse9735'}\n",
    "plot_sample_X=['logminmax']\n",
    "plotRecon='' #'meanRecon'\n",
    "standardizeX=False\n",
    "name='allk20XA_05_dca_over_lossXreconOnly_wKL'\n",
    "modelsavepath='/mnt/external_ssd/xinyi/models/train_gae_starmap/'+name\n",
    "plotsavepath='/mnt/external_ssd/xinyi/plots/train_gae_starmap/'+name\n",
    "datadir='/home/xinyiz/2021-01-13-mAD-test-dataset'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Load data\n",
    "savedir=os.path.join('/home/xinyiz/starmap')\n",
    "adj_dir=os.path.join(savedir,'a')\n",
    "\n",
    "featureslist={}\n",
    "if plot_sample_X[0] in ['corrected','scaled']:\n",
    "    scaleddata=scanpy.read_h5ad(datadir+'/2020-12-27-starmap-mAD-scaled.h5ad')\n",
    "    \n",
    "    for s in plot_samples.keys():\n",
    "        featureslist[s+'X_'+'corrected']=torch.tensor(scaleddata.layers['corrected'][scaleddata.obs['sample']==plot_samples[s]])\n",
    "        featureslist[s+'X_'+'scaled']=torch.tensor(scaleddata.layers['scaled'][scaleddata.obs['sample']==plot_samples[s]])\n",
    "    \n",
    "else:\n",
    "    scaleddata=scanpy.read_h5ad(datadir+'/2020-12-27-starmap-mAD-raw.h5ad')\n",
    "    \n",
    "    for s in plot_samples.keys():\n",
    "        scaleddata_train=scaleddata.X[scaleddata.obs['sample']==plot_samples[s]]\n",
    "\n",
    "        if plot_sample_X[0]=='logminmax':\n",
    "            featurelog_train=np.log2(scaleddata_train+1/2)\n",
    "            scaler = MinMaxScaler()\n",
    "            featurelog_train_minmax=np.transpose(scaler.fit_transform(np.transpose(featurelog_train)))\n",
    "            featureslist[s+'X_'+plot_sample_X[0]]=torch.tensor(featurelog_train_minmax)\n",
    "\n",
    "if protein: ##adjust for scaled/corrected\n",
    "    proteinsavepath=os.path.join('/mnt/xinyi/','starmap','protein')\n",
    "    for s in plot_samples.keys():\n",
    "        pmtx=sp.load_npz(os.path.join(proteinsavepath,plot_samples[s]+'_'+protein+'.npz'))\n",
    "        pmtx=preprocessing.sparse_mx_to_torch_sparse_tensor(pmtx)\n",
    "        pmtx=pmtx.to_dense()\n",
    "        scalefactor=torch.sum(featureslist[s+'X_'+plot_sample_X[0]])/torch.sum(pmtx)*proteinWeights\n",
    "        featureslist[s+'X_'+plot_sample_X[0]]=torch.cat((featureslist[s+'X_'+plot_sample_X[0]],pmtx*scalefactor),dim=1)\n",
    "\n",
    "adj_list={}\n",
    "adj_list['disease13']=sp.load_npz(os.path.join(adj_dir,maskedgeName+'_AD_mouse9494.npz'))\n",
    "adj_list['control13']=sp.load_npz(os.path.join(adj_dir,maskedgeName+'_AD_mouse9498.npz'))\n",
    "adj_list['disease8']=sp.load_npz(os.path.join(adj_dir,maskedgeName+'_AD_mouse9723.npz'))\n",
    "adj_list['control8']=sp.load_npz(os.path.join(adj_dir,maskedgeName+'_AD_mouse9735.npz'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 5,
     "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=='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",
    "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": 6,
   "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",
    "    colortest=sns.color_palette(\"husl\", 13)\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'))\n",
    "#     plt.show()\n",
    "    \n",
    "#     fig.clf()\n",
    "    plt.close('all')\n",
    "    \n",
    "    gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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": 8,
   "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": 9,
   "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",
    "        for s in plot_samples.keys():\n",
    "            sidx=(samplenameList==s)\n",
    "            plotembeddingbyCT(clusterRes[sidx],'leiden_location'+s,[],sobj_coord_np[sidx],savedir,'location'+' of '+s,savenameAdd=savenamecluster)\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": 10,
   "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": 11,
   "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",
    "\n",
    "            for s in plot_samples.keys():\n",
    "                sidx=(samplenameList==s)\n",
    "                plotembeddingbyCT(clusterRes[sidx],'agg_location'+s,[],sobj_coord_np[sidx],savedir,'location'+' of '+s,savenameAdd=savenamecluster)\n",
    "                plotembeddingbyCT_contrast(clusterRes[sidx],'agg_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": 12,
   "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,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "0\n",
      "0\n",
      "0\n"
     ]
    }
   ],
   "source": [
    "#compute embeddings\n",
    "mulist={}\n",
    "for s in plot_samples.keys():\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=featureslist[samplename]\n",
    "        if standardizeX:\n",
    "            features=torch.tensor(scale(features,axis=0, with_mean=True, with_std=True, copy=True))\n",
    "        if use_cuda:\n",
    "            model.cuda()\n",
    "            features = features.cuda().float()\n",
    "            adj_norm=adj_norm.cuda()\n",
    "            if adj_decodeName:\n",
    "                adj_decode=adj_decode.cuda()\n",
    "        \n",
    "        model.eval()\n",
    "        if adj_decodeName==None:\n",
    "            adj_recon,mu,logvar,z, features_recon = model(features, adj_norm)\n",
    "        else:\n",
    "            adj_recon,mu,logvar,z, features_recon = model(features, 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": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# combine all latents to one plot\n",
    "ifplot=False\n",
    "np.random.seed(seed)\n",
    "for xcorr in plot_sample_X:\n",
    "    latents=None\n",
    "    celltype_broad=None\n",
    "    celltype_sub=None\n",
    "    region=None\n",
    "    samplenameList=None\n",
    "    sobj_coord_np=None\n",
    "    \n",
    "    for s in plot_samples.keys():\n",
    "        sampleidx=plot_samples[s]        \n",
    "        samplename=s+'X_'+xcorr\n",
    "        muplot=np.copy(mulist[samplename])\n",
    "            \n",
    "        if latents is None:\n",
    "            latents=muplot\n",
    "            celltype_broad=scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx,'top_level']\n",
    "            celltype_sub=scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx,'cell_type_label']\n",
    "            region=scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx,'region']\n",
    "            sobj_coord_np=scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx,['x','y']].to_numpy()\n",
    "            samplenameList=np.repeat(s,muplot.shape[0])\n",
    "        else:\n",
    "            latents=np.vstack((latents,muplot))\n",
    "            celltype_broad=np.concatenate((celltype_broad,scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx,'top_level']),axis=None)\n",
    "            celltype_sub=np.concatenate((celltype_sub,scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx,'cell_type_label']),axis=None)\n",
    "            region=np.concatenate((region,scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx,'region']),axis=None)\n",
    "            sobj_coord_np=np.concatenate((sobj_coord_np,scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx,['x','y']].to_numpy()),axis=0)\n",
    "            samplenameList=np.concatenate((samplenameList,np.repeat(s,muplot.shape[0])),axis=None)\n",
    "       \n",
    "    origCT=np.unique(celltype_broad)\n",
    "    celltypeplot=np.concatenate((origCT,list(combineCelltype.keys())),axis=None)\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(samplenameList,'sample',[],embedding,savedir,plottype+'of all samples',plotdimx=dim1,plotdimy=dim2,savenameAdd=savenameAdd+'_pc'+str(dim1)+'pc'+str(dim2))\n",
    "                plotembeddingbyCT(celltype_broad,'celltype_broad',[],embedding,savedir,plottype+'all samples',plotdimx=dim1,plotdimy=dim2,savenameAdd=savenameAdd+'_pc'+str(dim1)+'pc'+str(dim2))\n",
    "    #             plotembeddingbyCT(celltype_sub,'celltype_sub',[],embedding,savedir,plottype+'all samples',savenameAdd=savenameAdd)\n",
    "                plotembeddingbyCT(region,'region',[],embedding,savedir,plottype+'all samples',plotdimx=dim1,plotdimy=dim2,savenameAdd=savenameAdd+'_pc'+str(dim1)+'pc'+str(dim2))\n",
    "\n",
    "    #             plotembeddingbyCT_contrast(celltype_sub,'celltype_sub',[],embedding,os.path.join(savedir,'contrast'),plottype+'all samples',savenameAdd=savenameAdd)    \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",
    "        if 'dbscan' in clustermethod:\n",
    "            clusterDBscan_allsample(embedding,savedir,clustersavedir,latents,epslist,min_sampleslist,n_pcs,sobj_coord_np,samplenameList)\n",
    "            assert np.sum(muplot-np.copy(mulist[s+'X_'+xcorr]))==0\n",
    "        if 'agglomerative' in clustermethod:\n",
    "            clusterAgg_allsample(embedding,savedir,clustersavedir,latents,nclusterlist,aggMetric,n_pcs,sobj_coord_np,samplenameList)\n",
    "            assert np.sum(muplot-np.copy(mulist[s+'X_'+xcorr]))==0\n",
    "        if 'kmeanbatch' in clustermethod:\n",
    "            clusterMinibatchKmean_allsample(embedding,savedir,clustersavedir,latents,nclusterlist,n_pcs,sobj_coord_np,samplenameList)\n",
    "            assert np.sum(muplot-np.copy(mulist[s+'X_'+xcorr]))==0\n",
    "#     #by region\n",
    "#     for reg in np.unique(region):\n",
    "#         savedir=os.path.join(sampledir,'embedding_'+plottype+'_'+reg)\n",
    "#         clustersavedir=os.path.join(sampledir,'cluster'+'_'+reg)\n",
    "#         if not os.path.exists(savedir):\n",
    "#             os.mkdir(savedir)\n",
    "#         if not os.path.exists(clustersavedir):\n",
    "#             os.mkdir(clustersavedir)\n",
    "\n",
    "#         reg_idx=region==reg\n",
    "\n",
    "# #         if plottype=='umap':\n",
    "# #             reducer = umap.UMAP(n_neighbors=n_neighbors,min_dist=min_dist,random_state=seed)\n",
    "# #             embedding = reducer.fit_transform(latents[reg_idx])\n",
    "# #             savenameAdd='_nn'+str(n_neighbors)+'mdist0'+str(int(min_dist*100))+'epoch'+str(plotepoch)\n",
    "# #         elif plottype=='pca':\n",
    "# #             embedding=pca.fit_transform(latents[reg_idx])\n",
    "# #             savenameAdd='_epoch'+str(plotepoch)\n",
    "# #         if ifplot:\n",
    "# #             plotembeddingbyCT(samplenameList[reg_idx],'sample',[],embedding,savedir,plottype+' of '+'all samples'+' '+reg,savenameAdd=savenameAdd)\n",
    "# #             plotembeddingbyCT(celltype_broad[reg_idx],'celltype_broad',[],embedding,savedir,plottype+' of '+'all samples'+' '+reg,savenameAdd=savenameAdd)\n",
    "# #             plotembeddingbyCT(celltype_sub[reg_idx],'celltype_sub',[],embedding,savedir,plottype+' of '+'all samples'+' '+reg,savenameAdd=savenameAdd)\n",
    "# #     #             plotembeddingbyCT(region,'region',[],embedding[reg_idx],savedir,'UMAP of '+s)\n",
    "\n",
    "# #             plotembeddingbyCT_contrast(celltype_sub[reg_idx],'celltype_sub',[],embedding,os.path.join(savedir,'contrast'),plottype+' of '+'all samples'+' '+reg,savenameAdd=savenameAdd)\n",
    "        \n",
    "# #         if embedding.shape[0]<minCells:\n",
    "# #             continue\n",
    "# #         if ifcluster:\n",
    "# #             if 'leiden' in clustermethod:\n",
    "# #                 clusterLeiden_allsample(embedding,savedir,clustersavedir,latents[reg_idx],n_neighbors,n_pcs,min_dist,resolution,sobj_coord_np[reg_idx],samplenameList[reg_idx],randseed=seed)\n",
    "# #                 assert np.sum(muplot-np.copy(mulist[s+'X_'+xcorr]))==0\n",
    "# #             if 'dbscan' in clustermethod:\n",
    "# #                 clusterDBscan_allsample(embedding,savedir,clustersavedir,latents[reg_idx],epslist,min_sampleslist,n_pcs,sobj_coord_np[reg_idx],samplenameList[reg_idx])\n",
    "# #                 assert np.sum(muplot-np.copy(mulist[s+'X_'+xcorr]))==0\n",
    "# #             if 'agglomerative' in clustermethod:\n",
    "# #                 clusterAgg_allsample(embedding,savedir,clustersavedir,latents[reg_idx],nclusterlist,aggMetric,n_pcs,sobj_coord_np[reg_idx],samplenameList[reg_idx])\n",
    "# #                 assert np.sum(muplot-np.copy(mulist[s+'X_'+xcorr]))==0\n",
    "# #             if 'kmeanbatch' in clustermethod:\n",
    "# #                 clusterMinibatchKmean_allsample(embedding,savedir,clustersavedir,latents[reg_idx],nclusterlist,n_pcs,sobj_coord_np[reg_idx],samplenameList[reg_idx])\n",
    "# #                 assert np.sum(muplot-np.copy(mulist[s+'X_'+xcorr]))==0\n",
    "#         #by region and celltype\n",
    "#         for ct in celltypeplot:\n",
    "# #             if not ((reg=='Cortex' and ct in ['Ex']) or (reg=='Hippocampus' and ct in ['CA1','DG','Micro','CA'])):\n",
    "#             if not ct =='Ex':\n",
    "# #             if not (ct in ['Micro'] or (reg=='Hippocampus' and ct in ['CA1'])):\n",
    "#                 continue\n",
    "#             print(reg+ct)\n",
    "#             savedir=os.path.join(sampledir,'embedding_'+plottype+'_'+reg+ct)\n",
    "#             clustersavedir=os.path.join(sampledir,'cluster'+'_'+reg+ct)\n",
    "#             if not os.path.exists(savedir):\n",
    "#                 os.mkdir(savedir)\n",
    "#             if not os.path.exists(clustersavedir):\n",
    "#                 os.mkdir(clustersavedir)\n",
    "\n",
    "            \n",
    "#             if ct in origCT:\n",
    "#                 ct_idx=celltype_broad==ct\n",
    "#             else:\n",
    "#                 ct_idx=False\n",
    "#                 for i in combineCelltype[ct]:\n",
    "#                     ct_idx=np.logical_or(ct_idx,celltype_broad==i)\n",
    "#             ct_idx=np.logical_and(reg_idx,ct_idx)      \n",
    "            \n",
    "#             if np.sum(ct_idx)<3:\n",
    "#                 continue\n",
    "#             if plottype=='umap':\n",
    "#                 reducer = umap.UMAP(n_neighbors=n_neighbors,min_dist=min_dist,random_state=seed)\n",
    "#                 embedding = reducer.fit_transform(latents[ct_idx])\n",
    "#                 savenameAdd='_nn'+str(n_neighbors)+'mdist0'+str(int(min_dist*100))+'epoch'+str(plotepoch)\n",
    "#             elif plottype=='pca':\n",
    "#                 pca.fit(latents[ct_idx])\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)+'.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[ct_idx])\n",
    "# #                 embedding=pca.fit_transform(latents[ct_idx])\n",
    "#                 savenameAdd='_epoch'+str(plotepoch)\n",
    "                \n",
    "#             if ifplot:\n",
    "#                 for dim1 in range(npc_plot-1):\n",
    "#                     for dim2 in range(dim1+1,npc_plot):\n",
    "#                         plotembeddingbyCT(samplenameList[ct_idx],'sample',[],embedding,savedir,plottype+' of '+reg+' all samples'+' '+ct,plotdimx=dim1,plotdimy=dim2,savenameAdd=savenameAdd+'_pc'+str(dim1)+'pc'+str(dim2))\n",
    "#                 #         plotembeddingbyCT(celltype_broad[reg_idx],'celltype_broad',[],embedding,savedir,plottype+' of '+'all samples'+' '+reg)\n",
    "#                         plotembeddingbyCT(celltype_sub[ct_idx],'celltype_sub',[],embedding,savedir,plottype+' of '+reg+' all samples'+' '+ct,plotdimx=dim1,plotdimy=dim2,savenameAdd=savenameAdd+'_pc'+str(dim1)+'pc'+str(dim2))\n",
    "#             #             plotembeddingbyCT(region[ct_idx],'region',[],embedding,savedir,plottype+' of '+reg+' all samples'+' '+ct)\n",
    "\n",
    "# #                         plotembeddingbyCT_contrast(celltype_sub[ct_idx],'celltype_sub',[],embedding,os.path.join(savedir,'contrast'),plottype+' of '+reg+' all samples'+' '+ct,savenameAdd=savenameAdd)\n",
    "        \n",
    "#             if embedding.shape[0]<minCells:\n",
    "#                 continue\n",
    "#             if ifcluster:\n",
    "#                 if 'leiden' in clustermethod:\n",
    "#                     clusterLeiden_allsample(embedding,savedir,clustersavedir,latents[ct_idx],n_neighbors,n_pcs,min_dist,resolution,sobj_coord_np[ct_idx],samplenameList[ct_idx],randseed=seed)\n",
    "#                     assert np.sum(muplot-np.copy(mulist[s+'X_'+xcorr]))==0\n",
    "#                 if 'dbscan' in clustermethod:\n",
    "#                     clusterDBscan_allsample(embedding,savedir,clustersavedir,latents[ct_idx],epslist,min_sampleslist,n_pcs,sobj_coord_np[ct_idx],samplenameList[ct_idx])\n",
    "#                     assert np.sum(muplot-np.copy(mulist[s+'X_'+xcorr]))==0\n",
    "#                 if 'agglomerative' in clustermethod:\n",
    "#                     clusterAgg_allsample(embedding,savedir,clustersavedir,latents[ct_idx],nclusterlist,aggMetric,n_pcs,sobj_coord_np[ct_idx],samplenameList[ct_idx])\n",
    "#                     assert np.sum(muplot-np.copy(mulist[s+'X_'+xcorr]))==0\n",
    "#                 if 'kmeanbatch' in clustermethod:\n",
    "#                     clusterMinibatchKmean_allsample(embedding,savedir,clustersavedir,latents[ct_idx],nclusterlist,n_pcs,sobj_coord_np[ct_idx],samplenameList[ct_idx])\n",
    "#                     assert np.sum(muplot-np.copy(mulist[s+'X_'+xcorr]))==0\n",
    "# #     #by celltype\n",
    "# #     for ct in celltypeplot:\n",
    "# #         print(ct)\n",
    "# #         savedir=os.path.join(sampledir,'embedding_'+plottype+'_'+ct)\n",
    "# #         clustersavedir=os.path.join(sampledir,'cluster'+'_'+ct)\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 ct in origCT:\n",
    "# #             ct_idx=celltype_broad==ct\n",
    "# #         else:\n",
    "# #             ct_idx=False\n",
    "# #             for i in combineCelltype[ct]:\n",
    "# #                 ct_idx=np.logical_or(ct_idx,celltype_broad==i)\n",
    "\n",
    "# #         if plottype=='umap':\n",
    "# #             reducer = umap.UMAP(n_neighbors=n_neighbors,min_dist=min_dist,random_state=seed)\n",
    "# #             embedding = reducer.fit_transform(latents[ct_idx])\n",
    "# #             savenameAdd='_nn'+str(n_neighbors)+'mdist0'+str(int(min_dist*100))+'epoch'+str(plotepoch)\n",
    "# #         elif plottype=='pca':\n",
    "# #             embedding=pca.fit_transform(latents[ct_idx])\n",
    "# #             savenameAdd='_epoch'+str(plotepoch)\n",
    "# #         if ifplot:\n",
    "# #             plotembeddingbyCT(samplenameList[ct_idx],'sample',[],embedding,savedir,plottype+' of '+'all samples'+' '+ct,savenameAdd=savenameAdd)\n",
    "# #     #         plotembeddingbyCT(celltype_broad[reg_idx],'celltype_broad',[],embedding,savedir,plottype+' of '+'all samples'+' '+reg)\n",
    "# #             plotembeddingbyCT(celltype_sub[ct_idx],'celltype_sub',[],embedding,savedir,plottype+' of '+'all samples'+' '+ct,savenameAdd=savenameAdd)\n",
    "# #             plotembeddingbyCT(region[ct_idx],'region',[],embedding,savedir,plottype+' of '+'all samples'+' '+ct,savenameAdd=savenameAdd)\n",
    "\n",
    "# #             plotembeddingbyCT_contrast(celltype_sub[ct_idx],'celltype_sub',[],embedding,os.path.join(savedir,'contrast'),plottype+' of '+'all samples'+' '+ct,savenameAdd=savenameAdd)\n",
    "        \n",
    "# #         if embedding.shape[0]<minCells:\n",
    "# #             continue\n",
    "# #         if ifcluster:\n",
    "# #             if 'leiden' in clustermethod:\n",
    "# #                 clusterLeiden_allsample(embedding,savedir,clustersavedir,latents[ct_idx],n_neighbors,n_pcs,min_dist,resolution,sobj_coord_np[ct_idx],samplenameList[ct_idx],randseed=seed)\n",
    "# #                 assert np.sum(muplot-np.copy(mulist[s+'X_'+xcorr]))==0\n",
    "# #             if 'dbscan' in clustermethod:\n",
    "# #                 clusterDBscan_allsample(embedding,savedir,clustersavedir,latents[ct_idx],epslist,min_sampleslist,n_pcs,sobj_coord_np[ct_idx],samplenameList[ct_idx])\n",
    "# #                 assert np.sum(muplot-np.copy(mulist[s+'X_'+xcorr]))==0\n",
    "# #             if 'agglomerative' in clustermethod:\n",
    "# #                 clusterAgg_allsample(embedding,savedir,clustersavedir,latents[ct_idx],nclusterlist,aggMetric,n_pcs,sobj_coord_np[ct_idx],samplenameList[ct_idx])\n",
    "# #                 assert np.sum(muplot-np.copy(mulist[s+'X_'+xcorr]))==0\n",
    "# #             if 'kmeanbatch' in clustermethod:\n",
    "# #                 clusterMinibatchKmean_allsample(embedding,savedir,clustersavedir,latents[ct_idx],nclusterlist,n_pcs,sobj_coord_np[ct_idx],samplenameList[ct_idx])\n",
    "# #                 assert np.sum(muplot-np.copy(mulist[s+'X_'+xcorr]))==0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# combine all latents to one plot -- subcluster \n",
    "np.random.seed(seed)\n",
    "\n",
    "ifplot=True\n",
    "ifcluster=True\n",
    "\n",
    "#leiden\n",
    "subResolution=[0.05,0.1]\n",
    "\n",
    "plottype='umap'\n",
    "pca=PCA()\n",
    "minCells_sub=10 #min number of cells for analysis\n",
    "# clustermethod=['leiden','agglomerative','kmeanbatch']\n",
    "clustermethod=['leiden']\n",
    "#umap/leiden clustering parameters\n",
    "n_neighbors_sub=5\n",
    "min_dist_sub=0.25\n",
    "n_pcs_sub=10 #for clustering\n",
    "# resolution_sub=[0.12,0.15,0.17]\n",
    "resolution_sub=[0.05,0.075,0.1,0.15,0.2,0.25,0.3]\n",
    "#DBscan\n",
    "epslist_sub= [6,8,10]\n",
    "min_sampleslist_sub=[15,30,45] \n",
    "#agglomerative\n",
    "nclusterlist_sub=[2,4,8,12]\n",
    "aggMetric_sub=['euclidean']\n",
    "\n",
    "xcorr=plot_sample_X[0]\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",
    "\n",
    "def subcluster(clustermethod,labels,savepath,addname):\n",
    "    latents=None\n",
    "    celltype_broad=None\n",
    "    celltype_sub=None\n",
    "    region=None\n",
    "    samplenameList=None\n",
    "    sobj_coord_np=None\n",
    "    for s in plot_samples.keys():\n",
    "        sampleidx=plot_samples[s]        \n",
    "        samplename=s+'X_'+xcorr\n",
    "        muplot=np.copy(mulist[samplename])\n",
    "\n",
    "        if latents is None:\n",
    "            latents=muplot\n",
    "            celltype_broad=scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx,'top_level']\n",
    "            celltype_sub=scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx,'cell_type_label']\n",
    "            region=scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx,'region']\n",
    "            sobj_coord_np=scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx,['x','y']].to_numpy()\n",
    "            samplenameList=np.repeat(s,muplot.shape[0])\n",
    "        else:\n",
    "            latents=np.vstack((latents,muplot))\n",
    "            celltype_broad=np.concatenate((celltype_broad,scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx,'top_level']),axis=None)\n",
    "            celltype_sub=np.concatenate((celltype_sub,scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx,'cell_type_label']),axis=None)\n",
    "            region=np.concatenate((region,scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx,'region']),axis=None)\n",
    "            sobj_coord_np=np.concatenate((sobj_coord_np,scaleddata.obs.loc[scaleddata.obs['sample']==sampleidx,['x','y']].to_numpy()),axis=0)\n",
    "            samplenameList=np.concatenate((samplenameList,np.repeat(s,muplot.shape[0])),axis=None)\n",
    "\n",
    "    sampledir=os.path.join(savepath,'combined')\n",
    "    savedir_all=os.path.join(sampledir,'embedding_'+plottype)\n",
    "    clustersavedir_all=os.path.join(sampledir,'cluster')\n",
    "    if not os.path.exists(savepath):\n",
    "        os.mkdir(savepath)\n",
    "    if not os.path.exists(sampledir):\n",
    "        os.mkdir(sampledir)\n",
    "    if not os.path.exists(savedir_all):\n",
    "        os.mkdir(savedir_all)\n",
    "    if not os.path.exists(clustersavedir_all):\n",
    "        os.mkdir(clustersavedir_all)\n",
    "    \n",
    "    for l in np.unique(labels):\n",
    "        clusteridx=(labels==l)\n",
    "#         origCT=np.unique(celltype_broad)\n",
    "#         celltypeplot=np.concatenate((origCT,list(combineCelltype.keys())),axis=None)\n",
    "        savedir=os.path.join(savedir_all,l)\n",
    "        clustersavedir=os.path.join(clustersavedir_all,l)\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",
    "            reducer = umap.UMAP(n_neighbors=n_neighbors,min_dist=min_dist,random_state=seed)\n",
    "            embedding = reducer.fit_transform(latents[clusteridx])\n",
    "            savenameAdd='_nn'+str(n_neighbors)+'mdist0'+str(int(min_dist*100))+'epoch'+str(plotepoch)\n",
    "        elif plottype=='pca':\n",
    "            embedding=pca.fit_transform(latents[clusteridx])\n",
    "            savenameAdd='_epoch'+str(plotepoch)\n",
    "        if ifplot:\n",
    "            plotembeddingbyCT(samplenameList[clusteridx],'sample',[],embedding,savedir,plottype+'of all samples',savenameAdd=savenameAdd)\n",
    "            plotembeddingbyCT(celltype_broad[clusteridx],'celltype_broad',[],embedding,savedir,plottype+'all samples',savenameAdd=savenameAdd)\n",
    "            plotembeddingbyCT(celltype_sub[clusteridx],'celltype_sub',[],embedding,savedir,plottype+'all samples',savenameAdd=savenameAdd)\n",
    "            plotembeddingbyCT(region[clusteridx],'region',[],embedding,savedir,plottype+'all samples',savenameAdd=savenameAdd)\n",
    "\n",
    "            plotembeddingbyCT_contrast(celltype_sub[clusteridx],'celltype_sub',[],embedding,os.path.join(savedir,'contrast'),plottype+'all samples',savenameAdd=savenameAdd)    \n",
    "\n",
    "        if embedding.shape[0]<minCells_sub:\n",
    "            continue\n",
    "        if ifcluster:\n",
    "            if clustermethod=='leiden':\n",
    "                clusterLeiden_allsample(embedding,savedir,clustersavedir,latents[clusteridx],n_neighbors_sub,n_pcs_sub,min_dist_sub,resolution_sub,sobj_coord_np[clusteridx],samplenameList[clusteridx],randseed=seed)\n",
    "            elif clustermethod=='dbscan':\n",
    "                clusterDBscan_allsample(embedding,savedir,clustersavedir,latents[clusteridx],epslist_sub,min_sampleslist_sub,n_pcs_sub,sobj_coord_np[clusteridx],samplenameList[clusteridx])\n",
    "            elif clustermethod=='agglomerative':\n",
    "                clusterAgg_allsample(embedding,savedir,clustersavedir,latents[clusteridx],nclusterlist_sub,aggMetric_sub,n_pcs_sub,sobj_coord_np[clusteridx],samplenameList[clusteridx])\n",
    "            elif clustermethod=='kmeanbatch':\n",
    "                clusterMinibatchKmean_allsample(embedding,savedir,clustersavedir,latents[clusteridx],nclusterlist_sub,n_pcs_sub,sobj_coord_np[clusteridx],samplenameList[clusteridx])\n",
    "\n",
    "def subclusterLeiden(n_neighbors,n_pcs,min_dist,resolution,addname=''):\n",
    "    for r in resolution:\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",
    "        readpath=os.path.join(clustersavedir,savenamecluster)\n",
    "        if not os.path.exists(readpath):\n",
    "            print('DNE: '+readpath)\n",
    "            continue\n",
    "        with open(readpath, 'rb') as input:\n",
    "            labels = pickle.load(input)\n",
    "        labels=np.array(labels)\n",
    "        if np.unique(labels).shape[0]==1:\n",
    "            continue\n",
    "        \n",
    "        savepath=os.path.join(clustersavedir,savenamecluster+'_subcluster')\n",
    "        subcluster('leiden',labels,savepath,addname)\n",
    "\n",
    "def subclusterDBscan(epsL,min_samplesL,n_pcs,addname=''):\n",
    "    for eps in epsL:\n",
    "        for min_samples in min_samplesL:\n",
    "            savenamecluster='dbscan_eps'+str(eps)+'msamples'+str(min_samples)+'n_pcs'+str(n_pcs)+'epoch'+str(plotepoch)\n",
    "            readpath=os.path.join(clustersavedir,savenamecluster)\n",
    "            if not os.path.exists(readpath):\n",
    "                print('DNE: '+readpath)\n",
    "                continue\n",
    "            with open(readpath, 'rb') as input:\n",
    "                labels = pickle.load(input)\n",
    "            labels=np.array(labels)\n",
    "            if np.unique(labels).shape[0]==1:\n",
    "                continue\n",
    "            \n",
    "            savepath=os.path.join(clustersavedir,savenamecluster+'_subcluster')\n",
    "            subcluster('dbscan',labels,savepath,addname)\n",
    "\n",
    "def subclusterAgg(nclusterL,aggmetricL,n_pcs,addname=''):\n",
    "    for ncluster in nclusterL:\n",
    "        for aggmetric in aggmetricL:\n",
    "            savenamecluster='agg_ncluster'+str(ncluster)+aggmetric+'n_pcs'+str(n_pcs)+'epoch'+str(plotepoch)\n",
    "            readpath=os.path.join(clustersavedir,savenamecluster)\n",
    "            if not os.path.exists(readpath):\n",
    "                print('DNE: '+readpath)\n",
    "                continue\n",
    "            with open(readpath, 'rb') as input:\n",
    "                labels = pickle.load(input)\n",
    "            labels=np.array(labels)\n",
    "            if np.unique(labels).shape[0]==1:\n",
    "                continue\n",
    "            savepath=os.path.join(clustersavedir,savenamecluster+'_subcluster')\n",
    "            subcluster('agglomerative',labels,savepath,addname)\n",
    "            \n",
    "def subclusterMinibatchKmean(nclusterL,n_pcs,addname=''):\n",
    "    for ncluster in nclusterL:\n",
    "        savenamecluster='minibatchkmean_ncluster'+str(ncluster)+'n_pcs'+str(n_pcs)+'epoch'+str(plotepoch)\n",
    "        readpath=os.path.join(clustersavedir,savenamecluster)\n",
    "        if not os.path.exists(readpath):\n",
    "            print('DNE: '+readpath)\n",
    "            continue\n",
    "        with open(readpath, 'rb') as input:\n",
    "            labels = pickle.load(input)\n",
    "        labels=np.array(labels)\n",
    "        if np.unique(labels).shape[0]==1:\n",
    "            continue\n",
    "            \n",
    "        savepath=os.path.join(clustersavedir,savenamecluster+'_subcluster')\n",
    "        subcluster('kmeanbatch',labels,savepath,addname)\n",
    "        \n",
    "subclusterLeiden(n_neighbors,n_pcs,min_dist,subResolution,addname='')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "savedir"
   ]
  }
 ],
 "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
}