train_gae_visium_10xADFFPE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook follows the same setup as train_gae_starmap_multisamples.ipynb.\n",
    "The major difference is that different layers of the model is defined across multiple GPUs to utilize more memory for overparameterization. The user needs to specify the GPU ID in each \"*.cuda(id)\" command. Our current code assumes GPU 0-3 are available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import os\n",
    "\n",
    "import scanpy as sc\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.preprocessing import MinMaxScaler\n",
    "import json\n",
    "import matplotlib.image as mpimg\n",
    "import networkx as nx\n",
    "\n",
    "from sklearn.neighbors import NearestNeighbors\n",
    "from sklearn.metrics.pairwise import euclidean_distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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=sc.read_10x_h5(os.path.join(datadir,'VisiumFFPE_Mouse_Brain_Alzheimers_AppNote_filtered_feature_bc_matrix.h5'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "features.obs['x']=tissuepos.loc[features.obs.index,4]\n",
    "features.obs['y']=tissuepos.loc[features.obs.index,5]\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Settings\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\" \n",
    "cudadevice0=0\n",
    "cudadevice1=1\n",
    "use_cuda=True\n",
    "fastmode=False #Validate during training pass\n",
    "seed=3\n",
    "useSavedMaskedEdges=True\n",
    "maskedgeName='knn6_connectivity'\n",
    "nneighbors=6\n",
    "epochs=13240\n",
    "saveFreq=490\n",
    "lr=0.01 #initial learning rate\n",
    "lr_adv=0.001\n",
    "weight_decay=0 #Weight for L2 loss on embedding matrix.\n",
    "\n",
    "hidden1=30000 #Number of units in hidden layer 1\n",
    "hidden2=30000 #Number of units in hidden layer 2\n",
    "fc_dim1=30000\n",
    "adv_hidden=128\n",
    "\n",
    "dropout=0.01\n",
    "testNodes=0.1 #fraction of total nodes for testing\n",
    "valNodes=0.05 #fraction of total nodes for validation\n",
    "XreconWeight=240\n",
    "# clfweight=20\n",
    "advWeight=2\n",
    "# randFeatureSubset=None\n",
    "model_str='gcn_vae_xa_e2_d1_dca_sharded'\n",
    "clf=None\n",
    "adv=None #'clf_fc1_control_eq' #'clf_fc1'\n",
    "protein=None #'scaled_binary'\n",
    "proteinWeights=0.05\n",
    "adj_decodeName=None #gala or None\n",
    "ridgeL=0.01\n",
    "shareGenePi=True\n",
    "\n",
    "training_sample_X='logminmax'\n",
    "switchFreq=10\n",
    "standardizeX=False\n",
    "name='10xAD' \n",
    "useA=True #set to True to incorporate adjacency reconstruction loss\n",
    "logsavepath='log/train_gae_visium_validation/'+name\n",
    "modelsavepath='models/train_gae_visium_validation/'+name\n",
    "plotsavepath='plots/train_gae_visium_validation/'+name\n",
    "\n",
    "if not os.path.exists(logsavepath):\n",
    "    os.mkdir(logsavepath)\n",
    "if not os.path.exists(modelsavepath):\n",
    "    os.mkdir(modelsavepath)\n",
    "if not os.path.exists(plotsavepath):\n",
    "    os.mkdir(plotsavepath)\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",
    "# Load data\n",
    "# if randFeatureSubset != None:\n",
    "#     idx=np.random.choice(features.shape[1],randFeatureSubset,replace=False)\n",
    "#     features=features[:,idx]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plotA(coord,adj):\n",
    "    g=nx.from_numpy_matrix(adj)\n",
    "    pos={}\n",
    "    for n in range(coord.shape[0]):\n",
    "        pos[n]=(coord[n][0],coord[n][1])\n",
    "    fig, ax = plt.subplots(dpi=400)\n",
    "    nx.draw(g,pos,node_size=1,width=0.1)\n",
    "    plt.savefig(os.path.join(plotsavepath,samplename+'_connectedA'))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "savedirlist={}\n",
    "featureslist={}\n",
    "features_raw_list={}\n",
    "adj_list={}\n",
    "coordlist={}\n",
    "commonGenes=[]\n",
    "    \n",
    "for samplename in np.unique(features.obs['samplename']):\n",
    "    print(samplename)\n",
    "    \n",
    "    adj_list[samplename]=getA_knn(samplename,nneighbors+1,'connectivity')\n",
    "    if training_sample_X=='logminmax':\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_'+training_sample_X]=torch.tensor(featurelog_train_minmax)\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for samplename in adj_list.keys():\n",
    "    print(samplename)\n",
    "    sobj_coord_np=features.obs.loc[features.obs.index[features.obs['samplename']==samplename],['x','y']].to_numpy()\n",
    "    plotA(sobj_coord_np,adj_list[samplename].numpy()-sp.eye(sobj_coord_np.shape[0]))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loop over all train/validation sets\n",
    "\n",
    "mse=torch.nn.MSELoss()\n",
    "# mse=torch.nn.MSELoss(reduction=None)\n",
    "# Create model\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",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "elif model_str == 'gcn_vae_gcnX_inprA':\n",
    "    model = gae.gae.model.GCNModelVAE_gcnX_inprA(num_features, hidden1, hidden2,gcn_dim1, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "    \n",
    "elif model_str == 'gcn_vae_gcnX_inprA_w':\n",
    "    model = gae.gae.model.GCNModelVAE_gcnX_inprA_w(num_features, hidden1, hidden2,gcn_dim1, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_a=optimizer.optimizer_CE    \n",
    "    \n",
    "elif model_str=='gcn_vae_xa_e3':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e3(num_features, hidden1, hidden2,hidden3,fc_dim1,fc_dim2,fc_dim3,fc_dim4, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "    \n",
    "elif model_str == 'gcn_vae_xa_e1':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e1(num_features, hidden1, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "\n",
    "elif model_str == 'gcn_vae_xa_e2_d1':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e2_d1(num_features, hidden1,hidden2, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "    \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",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_zinb\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "    \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",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_zinb\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "    \n",
    "elif model_str=='gcn_vae_xa_e2_d1_dca_fca':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e2_d1_DCA_fca(num_features, hidden1,hidden2,fc_dim1, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_zinb\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "\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",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_zinb\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "\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",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_zinb\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "\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",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_zinb\n",
    "    loss_a=optimizer.optimizer_CE    \n",
    "    \n",
    "elif model_str == 'gcn_vae_xa_e4_d1':\n",
    "    model = gae.gae.model.GCNModelVAE_XA_e4_d1(num_features, hidden1,hidden2,hidden3,hidden4, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_a=optimizer.optimizer_CE    \n",
    "    \n",
    "elif model_str=='fc':\n",
    "    model = gae.gae.model.FCVAE(num_features, hidden1, hidden2,hidden3,hidden4,hidden5,fc_dim1,fc_dim2,fc_dim3,fc_dim4, dropout)\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "#     loss_x=mse\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "elif model_str=='fcae':\n",
    "    model = gae.gae.model.FCAE(num_features, hidden1, hidden2,hidden3,hidden4,hidden5,fc_dim1,fc_dim2,fc_dim3,fc_dim4, dropout)\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "#     loss_x=mse\n",
    "\n",
    "elif model_str=='fcae1':\n",
    "    model = gae.gae.model.FCAE1(num_features, dropout,hidden1)\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "#     loss_x=mse\n",
    "elif model_str=='fcae2':\n",
    "    model = gae.gae.model.FCAE2(num_features, dropout,hidden1,hidden2)\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "\n",
    "elif model_str=='fc1':\n",
    "    model = gae.gae.model.FCVAE1(num_features, hidden1,dropout)\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "\n",
    "elif model_str=='fc1_fca':\n",
    "    model = gae.gae.model.FCVAE1_fca(num_features, hidden1,dropout)\n",
    "    loss_x=optimizer.optimizer_MSE\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "    \n",
    "elif model_str=='fc1_dca':\n",
    "    model = gae.gae.model.FCVAE1_DCA(num_features, hidden1,fc_dim1, dropout)\n",
    "    loss_x=optimizer.optimizer_zinb\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "    \n",
    "elif model_str=='fc1_dca_sharded':\n",
    "    model = gae.gae.model.FCVAE1_DCA_sharded(num_features, hidden1,fc_dim1, dropout)\n",
    "    loss_x=optimizer.optimizer_zinb\n",
    "    loss_kl=optimizer.optimizer_kl\n",
    "    loss_a=optimizer.optimizer_CE\n",
    "\n",
    "if clf=='clf_fc1':\n",
    "    modelClf=gae.gae.model.Clf_fc1(hidden2, dropout,clf_hidden,ct_unique.size)\n",
    "    loss_clf=torch.nn.CrossEntropyLoss()\n",
    "    \n",
    "if adv=='clf_fc1' or adv=='clf_fc1_eq' or adv=='clf_fc1_control' or adv=='clf_fc1_control_eq':\n",
    "    modelAdv=gae.gae.model.Clf_fc1(hidden2, dropout,adv_hidden,sampleLabellist_ae['control13'].size()[1])\n",
    "    loss_adv=optimizer.optimizer_CEclf\n",
    "    \n",
    "\n",
    "\n",
    "optimizerVAEXA = optim.SGD(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
    "if clf:\n",
    "    optimizerClf=optim.Adam(modelClf.parameters(), lr=lr, weight_decay=weight_decay)\n",
    "if adv:\n",
    "    optimizerAdv=optim.Adam(modelAdv.parameters(), lr=lr_adv, weight_decay=weight_decay)    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "    \n",
    "if not os.path.exists(os.path.join(plotsavepath,'trace')):\n",
    "    os.mkdir(os.path.join(plotsavepath,'trace'))\n",
    "    \n",
    "def train(epoch):\n",
    "    t = time.time()\n",
    "    model.train()\n",
    "    optimizerVAEXA.zero_grad(True)\n",
    "    \n",
    "    if adj_decodeName==None:\n",
    "        adj_recon,mu,logvar,z,features_recon = model(features, adj_norm)\n",
    "    #     features_recon, z, mu, logvar=model(features.float())\n",
    "    else:\n",
    "        adj_recon,mu,logvar,z,features_recon = model(features, adj_norm,adj_decode)\n",
    "        \n",
    "    adj_recon=adj_recon.cuda(0)\n",
    "    mu=mu.cuda(0)\n",
    "    logvar=logvar.cuda(0)\n",
    "    z=z.cuda(0)\n",
    "#     features_recon=features_recon.cuda(0)\n",
    "    \n",
    "    if clf:\n",
    "        modelClf.train()\n",
    "        optimizerClf.zero_grad()\n",
    "        clfOut=modelClf(z)\n",
    "    if adv and (training_samples_t in list(sampleLabellist_ae.keys())) and (training_samples_t != targetBatch):\n",
    "        modelAdv.eval()\n",
    "        advOut=modelAdv(z)\n",
    "    \n",
    "    loss_kl_train=loss_kl(mu, logvar, train_nodes_idx)\n",
    "    if 'dca' in model_str:\n",
    "        loss_x_train=loss_x(features_recon, features.cuda(0),train_nodes_idx,XreconWeight,ridgeL,features_raw)\n",
    "    else:\n",
    "        loss_x_train=loss_x(features_recon, features.cuda(0),train_nodes_idx,XreconWeight,mse)\n",
    "    loss_a_train=loss_a(adj_recon, adj_label, pos_weight, norm,train_nodes_idx)\n",
    "    loss=loss_kl_train+loss_x_train \n",
    "    if useA:\n",
    "        loss=loss+loss_a_train\n",
    "    loss.backward()\n",
    "    optimizerVAEXA.step()\n",
    "        \n",
    "    with torch.no_grad():\n",
    "        if not fastmode:\n",
    "            # Evaluate validation set performance separately,\n",
    "            # deactivates dropout during validation run & no variation in z.\n",
    "            model.eval()\n",
    "            if adj_decodeName==None:\n",
    "                adj_recon,mu,logvar,z, features_recon = model(features, adj_norm)\n",
    "        #         features_recon, z, mu, logvar=model(features.float())\n",
    "            else:\n",
    "                adj_recon,mu,logvar,z, features_recon = model(features, adj_norm,adj_decode)\n",
    "\n",
    "        adj_recon=adj_recon.cuda(0)\n",
    "        mu=mu.cuda(0)\n",
    "        logvar=logvar.cuda(0)\n",
    "        z=z.cuda(0)\n",
    "    #     features_recon=features_recon.cuda(0)\n",
    "\n",
    "        if clf:\n",
    "            modelClf.eval()\n",
    "            clfOut=modelClf(z)\n",
    "        if adv and (training_samples_t in list(sampleLabellist_ae.keys())) and (training_samples_t != targetBatch):\n",
    "            advOut=modelAdv(z)\n",
    "        if 'dca' in model_str:\n",
    "            loss_x_val=loss_x(features_recon, features.cuda(0),val_nodes_idx,XreconWeight,ridgeL,features_raw)\n",
    "        else:\n",
    "            loss_x_val=loss_x(features_recon, features,val_nodes_idx,XreconWeight,mse)\n",
    "        loss_a_val=loss_a(adj_recon, adj_label, pos_weight, norm,val_nodes_idx)\n",
    "        loss_val=loss_x_val\n",
    "        if useA:\n",
    "            loss_val=loss_val+loss_a_val\n",
    "    print(training_samples_t+' Epoch: {:04d}'.format(epoch),\n",
    "#           'loss_train: {:.4f}'.format(loss.item()),\n",
    "          'loss_kl_train: {:.4f}'.format(loss_kl_train.item()),\n",
    "          'loss_x_train: {:.4f}'.format(loss_x_train.item()),\n",
    "          'loss_a_train: {:.4f}'.format(loss_a_train.item()),\n",
    "#           'loss_clf_train: {:.4f}'.format(loss_clf_train.item()),\n",
    "#           'loss_val: {:.4f}'.format(loss_val.item()),\n",
    "          'loss_x_val: {:.4f}'.format(loss_x_val.item()),\n",
    "          'loss_a_val: {:.4f}'.format(loss_a_val.item()),\n",
    "#           'loss_clf_val: {:.4f}'.format(loss_clf_val.item()),\n",
    "          'time: {:.4f}s'.format(time.time() - t))\n",
    "    return time.time() - t,float(loss),float(loss_kl_train),float(loss_x_train),float(loss_a_train),float(loss_val),float(loss_x_val),float(loss_a_val)        \n",
    "\n",
    "    \n",
    "# print('cross-validation ',seti)\n",
    "train_loss_ep=[None]*epochs\n",
    "train_loss_kl_ep=[None]*epochs\n",
    "train_loss_x_ep=[None]*epochs\n",
    "train_loss_a_ep=[None]*epochs\n",
    "train_loss_clf_ep=[None]*epochs\n",
    "train_loss_adv_ep=[None]*epochs\n",
    "train_loss_advD_ep=[None]*epochs\n",
    "val_loss_ep=[None]*epochs\n",
    "val_loss_x_ep=[None]*epochs\n",
    "val_loss_a_ep=[None]*epochs\n",
    "val_loss_clf_ep=[None]*epochs\n",
    "val_loss_adv_ep=[None]*epochs\n",
    "val_loss_advD_ep=[None]*epochs\n",
    "\n",
    "timePerEpoch=[None]*epochs\n",
    "\n",
    "t_ep=time.time()\n",
    "\n",
    "def trace_handler(p):\n",
    "    output= p.key_averages().table(sort_by=\"self_cuda_time_total\", row_limit=10)\n",
    "    with open(os.path.join(plotsavepath,'trace', \"averages.txt\"), 'a') as f:\n",
    "        f.write(output)\n",
    "    print(output)\n",
    "    p.export_chrome_trace(os.path.join(plotsavepath,'trace',str(p.step_num) + \".json\"))\n",
    "\n",
    "with torch.profiler.profile(\n",
    "    activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],\n",
    "    profile_memory=True,\n",
    "    schedule=torch.profiler.schedule(\n",
    "        wait=1,\n",
    "        warmup=3,\n",
    "        active=6),\n",
    "    on_trace_ready=trace_handler\n",
    ") as p:        \n",
    "\n",
    "    for ep in range(epochs):\n",
    "    # for ep in range(10000,20000):\n",
    "        t=int(ep/switchFreq)%np.unique(samplenameList).size\n",
    "        training_samples_t=np.unique(samplenameList)[t]\n",
    "        print(training_samples_t)\n",
    "        \n",
    "        adj_norm=adjnormlist[training_samples_t]\n",
    "    #     adj_label=adj_list[training_samples_t]\n",
    "        features=featureslist[training_samples_t+'X_'+training_sample_X]\n",
    "    #     adj_norm=adjnormlist[training_samples_t].cuda().float()\n",
    "        adj_label=adj_list[training_samples_t].cuda(0).float()\n",
    "    #     features=featureslist[training_samples_t+'X_'+training_sample_X].cuda().float()\n",
    "        pos_weight=pos_weightlist[training_samples_t]\n",
    "        norm=normlist[training_samples_t]\n",
    "        if 'dca' in model_str:\n",
    "    #         features_raw=features_raw_list[training_samples_t+'X_raw']\n",
    "            features_raw=features_raw_list[training_samples_t+'X_raw'].cuda()\n",
    "        num_nodes,_ = features.shape\n",
    "\n",
    "        maskedgeres= preprocessing.mask_nodes_edges(features.shape[0],testNodeSize=testNodes,valNodeSize=valNodes)\n",
    "        train_nodes_idx,val_nodes_idx,test_nodes_idx = maskedgeres\n",
    "        if use_cuda:\n",
    "            train_nodes_idx=train_nodes_idx.cuda(0)\n",
    "            val_nodes_idx=val_nodes_idx.cuda(0)\n",
    "            test_nodes_idx=test_nodes_idx.cuda(0)\n",
    "            \n",
    "        with torch.profiler.record_function(\"train_and_validate\"):\n",
    "            timePerEpoch[ep],train_loss_ep[ep],train_loss_kl_ep[ep],train_loss_x_ep[ep],train_loss_a_ep[ep],val_loss_ep[ep],val_loss_x_ep[ep],val_loss_a_ep[ep]=train(ep)\n",
    "\n",
    "\n",
    "        if ep%saveFreq == 0 and ep!=0:\n",
    "            torch.save(model.state_dict(), os.path.join(modelsavepath,str(ep)+'.pt'))\n",
    "        p.step()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "    \n",
    "with open(os.path.join(logsavepath,'train_loss'), 'wb') as output:\n",
    "    pickle.dump(train_loss_ep, output, pickle.HIGHEST_PROTOCOL)\n",
    "with open(os.path.join(logsavepath,'train_loss_kl'), 'wb') as output:\n",
    "    pickle.dump(train_loss_kl_ep, output, pickle.HIGHEST_PROTOCOL)\n",
    "with open(os.path.join(logsavepath,'train_loss_x'), 'wb') as output:\n",
    "    pickle.dump(train_loss_x_ep, output, pickle.HIGHEST_PROTOCOL)\n",
    "with open(os.path.join(logsavepath,'train_loss_a'), 'wb') as output:\n",
    "    pickle.dump(train_loss_a_ep, output, pickle.HIGHEST_PROTOCOL)\n",
    "with open(os.path.join(logsavepath,'val_loss'), 'wb') as output:\n",
    "    pickle.dump(val_loss_ep, output, pickle.HIGHEST_PROTOCOL)\n",
    "with open(os.path.join(logsavepath,'val_loss_x'), 'wb') as output:\n",
    "    pickle.dump(val_loss_x_ep, output, pickle.HIGHEST_PROTOCOL)\n",
    "with open(os.path.join(logsavepath,'val_loss_a'), 'wb') as output:\n",
    "    pickle.dump(val_loss_a_ep, output, pickle.HIGHEST_PROTOCOL)\n",
    "with open(os.path.join(logsavepath,'time'), 'wb') as output:\n",
    "    pickle.dump(timePerEpoch, output, pickle.HIGHEST_PROTOCOL)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# plt.plot(np.arange(epochs),train_loss_ep)\n",
    "# plt.plot(np.arange(epochs),val_loss_ep)\n",
    "plt.plot(np.arange(epochs),train_loss_x_ep)\n",
    "plt.plot(np.arange(epochs),val_loss_x_ep)\n",
    "# plt.plot(np.arange(epochs),train_loss_a_ep)\n",
    "# plt.plot(np.arange(epochs),val_loss_a_ep)\n",
    "plt.plot(np.arange(epochs),train_loss_kl_ep)\n",
    "# plt.ylim((0,30))\n",
    "# plt.xlim((8800,9400))\n",
    "plt.legend(['training x recon loss','validation x recon loss','training kl loss'],loc='upper right')\n",
    "# plt.legend(['training loss','validation loss','training x recon loss','validation x recon loss','training a recon loss','validation a recon loss','training kl loss'],loc='upper right')\n",
    "plt.savefig(os.path.join(plotsavepath,'loss_seed3_xkl.jpg'))\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "testepoch=6500\n",
    "model.load_state_dict(torch.load(os.path.join(modelsavepath,str(testepoch)+'.pt')))\n",
    "model.eval()\n",
    "for s in training_samples:\n",
    "    print(s)\n",
    "    \n",
    "#     adj_norm=adjnormlist[s].cuda().float()\n",
    "    adj_label=adj_list[s].cuda(0).float()\n",
    "    features=featureslist[s+'X_'+training_sample_X].cuda().float()\n",
    "    adj_norm=adjnormlist[s]\n",
    "#     features=featureslist[s+'X_'+training_sample_X]\n",
    "    pos_weight=pos_weightlist[s]\n",
    "    norm=normlist[s]\n",
    "    \n",
    "    if 'dca' in model_str:\n",
    "        features_raw=features_raw_list[s+'X_raw'].cuda()\n",
    "    num_nodes,num_features = features.shape\n",
    "    maskedgeres= preprocessing.mask_nodes_edges(features.shape[0],testNodeSize=testNodes,valNodeSize=valNodes)\n",
    "    train_nodes_idx,val_nodes_idx,test_nodes_idx = maskedgeres\n",
    "    \n",
    "    if s in training_samples:\n",
    "        test_nodes_idx_s=test_nodes_idx\n",
    "    else:\n",
    "        test_nodes_idx_s=torch.tensor(np.arange(num_nodes))\n",
    "        \n",
    "\n",
    "    if adj_decodeName==None:\n",
    "        adj_recon,mu,logvar,z, features_recon = model(features, adj_norm)\n",
    "    else:\n",
    "        adj_decode=adj_decode.cuda()\n",
    "        adj_recon,mu,logvar,z, features_recon = model(features, adj_norm,adj_decode)\n",
    "    if adv and (s in list(sampleLabellist_ae.keys())):\n",
    "        sampleLabel_ae=sampleLabellist_ae[s].cuda().float()\n",
    "        modelAdv.eval()\n",
    "        advOut=modelAdv(z)\n",
    "    if 'dca' in model_str:\n",
    "        loss_x_test=loss_x(features_recon, features,test_nodes_idx_s,XreconWeight,ridgeL,features_raw)\n",
    "    else:\n",
    "        loss_x_test=loss_x(features_recon, features,test_nodes_idx_s,XreconWeight,mse)\n",
    "    loss_a_test=loss_a(adj_recon, adj_label, pos_weight, norm,test_nodes_idx_s)\n",
    "    loss_test = loss_x_test+loss_a_test\n",
    "    \n",
    "    if adv and (s in list(sampleLabellist_ae.keys())):\n",
    "        loss_adv_test=loss_adv(advOut,sampleLabel_ae,test_nodes_idx)\n",
    "        print('loss_adv_test: {:.4f}'.format(loss_adv_test.item()))\n",
    "        \n",
    "    print('test results',\n",
    "          'loss_test: {:.4f}'.format(loss_test.item()),\n",
    "          'loss_x_test: {:.4f}'.format(loss_x_test.item()),\n",
    "          'loss_a_test: {:.4f}'.format(loss_a_test.item()))\n",
    "#          'loss_adv_test: {:.4f}'.format(loss_adv_test.item()))\n",
    "    if protein:\n",
    "        test_nodes_idx_s_genes=torch.clone(test_nodes_idx_s)\n",
    "        test_nodes_idx_s_genes[2112:]=0\n",
    "        test_nodes_idx_s_proteins=torch.clone(test_nodes_idx_s)\n",
    "        test_nodes_idx_s_proteins[:2112]=0\n",
    "        if 'dca' in model_str:\n",
    "            loss_genes_test=loss_x(features_recon, features,test_nodes_idx_s_genes,XreconWeight,ridgeL,features_raw)\n",
    "            loss_proteins_test=loss_x(features_recon, features,test_nodes_idx_s_proteins,XreconWeight,ridgeL,features_raw)\n",
    "        print('loss_x_genes: {:.4f}'.format(loss_genes_test.item()),\n",
    "          'loss_x_proteins: {:.4f}'.format(loss_proteins_test.item()))\n",
    "\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": 4
}