Figure5_6.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "053309c2-cb6d-4623-83b7-7b0d947998e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import pytorch_lightning as pl\n",
    "import torchvision.transforms as tf\n",
    "from tqdm import tqdm\n",
    "from predict import *\n",
    "from HGGEP import *\n",
    "from dataset import ViT_HER2ST, ViT_SKIN\n",
    "from scipy.stats import pearsonr,spearmanr\n",
    "from torch.utils.data import DataLoader\n",
    "from pytorch_lightning.loggers import TensorBoardLogger\n",
    "from copy import deepcopy as dcp\n",
    "from collections import defaultdict as dfd\n",
    "from sklearn.metrics import adjusted_rand_score as ari_score\n",
    "from sklearn.metrics.cluster import normalized_mutual_info_score as nmi_score"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b132bfb4-bd29-44fa-bbca-3790b0848fd4",
   "metadata": {},
   "source": [
    "# Data Loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "6281d8f2-e5f7-4c26-88bf-5256197b1418",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "name=[*[f'A{i}' for i in range(2,7)],*[f'B{i}' for i in range(1,7)],\n",
    "      *[f'C{i}' for i in range(1,7)],*[f'D{i}' for i in range(1,7)],\n",
    "      *[f'E{i}' for i in range(1,4)],*[f'F{i}' for i in range(1,4)],*[f'G{i}' for i in range(1,4)]]\n",
    "patients = ['P2', 'P5', 'P9', 'P10']\n",
    "reps = ['rep1', 'rep2', 'rep3']\n",
    "skinname = []\n",
    "for i in patients:\n",
    "    for j in reps:\n",
    "        skinname.append(i+'_ST_'+j)\n",
    "device='cuda'\n",
    "tag='5-7-2-8-4-16-32'\n",
    "k,p,d1,d2,d3,h,c=map(lambda x:int(x),tag.split('-'))\n",
    "dropout=0.2\n",
    "random.seed(12000)\n",
    "np.random.seed(12000)\n",
    "torch.manual_seed(12000)\n",
    "torch.cuda.manual_seed(12000)\n",
    "torch.cuda.manual_seed_all(12000)  \n",
    "torch.backends.cudnn.benchmark = False\n",
    "torch.backends.cudnn.deterministic = True"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "45afd1c5-2495-4240-b3b6-92b415507582",
   "metadata": {},
   "source": [
    "# HGGEP Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "0f19d5d8-6510-493b-9287-ba7215155dc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "fold=12\n",
    "data='her2st'\n",
    "prune='Grid' if data=='her2st' else 'NA'\n",
    "genes=171 if data=='cscc' else 785"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "76cfcfb9-3169-46da-9288-2e2be7942d84",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['F1']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:02<00:00,  2.82s/it]\n"
     ]
    }
   ],
   "source": [
    "testset = pk_load(fold,'test',dataset=data,flatten=False,adj=True,ori=True,prune=prune)\n",
    "test_loader = DataLoader(testset, batch_size=1, num_workers=0, shuffle=False)\n",
    "label=testset.label[testset.names[0]]\n",
    "genes=785\n",
    "model=HGGEP(\n",
    "    depth1=d1, depth2=d2,depth3=d3,n_genes=genes, \n",
    "    kernel_size=k, patch_size=p,\n",
    "    heads=h, channel=c, dropout=0.2,\n",
    "    zinb=0.25, nb=False,\n",
    "    bake=5, lamb=0.5,\n",
    ")\n",
    "model.load_state_dict(torch.load(f'./model/HGGEP/{fold}-Hist2ST.ckpt'))\n",
    "pred, gt = test(model, test_loader,'cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4266cce3-6bc7-4f5b-81d5-d03d3ae089c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pearson Correlation: 0.1122036904738274\n"
     ]
    }
   ],
   "source": [
    "R=get_R(pred,gt)[0]\n",
    "print('Pearson Correlation:',np.nanmean(R))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4e7b1dfa",
   "metadata": {},
   "source": [
    "# Predict Top Gene: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "314fdbb4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:02<00:00,  2.01s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AnnData object with n_obs × n_vars = 691 × 785\n",
      "    obs: 'kmeans'\n",
      "    var: 'mean', 'std'\n",
      "    uns: 'pca', 'tsne'\n",
      "    obsm: 'spatial', 'X_pca', 'X_tsne'\n",
      "    varm: 'PCs' AnnData object with n_obs × n_vars = 691 × 785\n",
      "    var: 'mean', 'std'\n",
      "    obsm: 'spatial'\n"
     ]
    }
   ],
   "source": [
    "def comp_tsne_km(adata,k=10):\n",
    "    sc.pp.pca(adata)\n",
    "    sc.tl.tsne(adata)\n",
    "    kmeans = KMeans(n_clusters=k, init=\"k-means++\", random_state=0).fit(adata.obsm['X_pca'])\n",
    "    adata.obs['kmeans'] = kmeans.labels_.astype(str)\n",
    "    return adata\n",
    "\n",
    "model.load_state_dict(torch.load(f'./model/HGGEP/{fold}-Hist2ST.ckpt'))\n",
    "adata_pred, adata_truth = test(model, test_loader)\n",
    "adata_pred = comp_tsne_km(adata_pred,4)\n",
    "\n",
    "g = list(np.load('data/her_hvg_cut_1000.npy',allow_pickle=True))\n",
    "adata_pred.var_names = g\n",
    "adata_truth.var_names = g\n",
    "# print(len(g), g[:10]) # 785 gene\n",
    "sc.pp.scale(adata_pred)\n",
    "sc.pp.scale(adata_truth)\n",
    "print(adata_pred, adata_truth)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "77650b7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata_true = pd.DataFrame(data=adata_truth.X, index=adata_truth.obs_names, columns=adata_truth.var_names)\n",
    "adata_pre = pd.DataFrame(data=adata_pred.X, index=adata_pred.obs_names, columns=adata_pred.var_names)\n",
    "# list(adata_pre.columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f39e0b8c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>PCC</th>\n",
       "      <th>PV</th>\n",
       "      <th>-log10_PV</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>IGKC</th>\n",
       "      <td>0.710980</td>\n",
       "      <td>1.873444e-107</td>\n",
       "      <td>106.727359</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>IGLC2</th>\n",
       "      <td>0.426681</td>\n",
       "      <td>6.066996e-32</td>\n",
       "      <td>31.217026</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>FASN</th>\n",
       "      <td>0.402307</td>\n",
       "      <td>2.878096e-28</td>\n",
       "      <td>27.540895</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>SCD</th>\n",
       "      <td>0.387770</td>\n",
       "      <td>3.244230e-26</td>\n",
       "      <td>25.488888</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>IGHG3</th>\n",
       "      <td>0.384035</td>\n",
       "      <td>1.051743e-25</td>\n",
       "      <td>24.978090</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ANGPTL4</th>\n",
       "      <td>-0.035125</td>\n",
       "      <td>3.565669e-01</td>\n",
       "      <td>0.447859</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>CHST11</th>\n",
       "      <td>-0.046866</td>\n",
       "      <td>2.185400e-01</td>\n",
       "      <td>0.660469</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>TCEA3</th>\n",
       "      <td>-0.050173</td>\n",
       "      <td>1.877281e-01</td>\n",
       "      <td>0.726471</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>SCX</th>\n",
       "      <td>-0.062263</td>\n",
       "      <td>1.019823e-01</td>\n",
       "      <td>0.991475</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ENPP1</th>\n",
       "      <td>-0.079220</td>\n",
       "      <td>3.734730e-02</td>\n",
       "      <td>1.427741</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>785 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "              PCC             PV   -log10_PV\n",
       "IGKC     0.710980  1.873444e-107  106.727359\n",
       "IGLC2    0.426681   6.066996e-32   31.217026\n",
       "FASN     0.402307   2.878096e-28   27.540895\n",
       "SCD      0.387770   3.244230e-26   25.488888\n",
       "IGHG3    0.384035   1.051743e-25   24.978090\n",
       "...           ...            ...         ...\n",
       "ANGPTL4 -0.035125   3.565669e-01    0.447859\n",
       "CHST11  -0.046866   2.185400e-01    0.660469\n",
       "TCEA3   -0.050173   1.877281e-01    0.726471\n",
       "SCX     -0.062263   1.019823e-01    0.991475\n",
       "ENPP1   -0.079220   3.734730e-02    1.427741\n",
       "\n",
       "[785 rows x 3 columns]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_PCC(adata_pre, adata_true):\n",
    "    gene_names = adata_pre.columns\n",
    "\n",
    "    result_df = pd.DataFrame(index=['PCC', 'PV', '-log10_PV'], columns=gene_names)\n",
    "\n",
    "    for gene in gene_names:\n",
    "        corr_coef, p_value = pearsonr(adata_pre[gene], adata_true[gene])\n",
    "        result_df[gene] = [corr_coef, p_value, -np.log10(p_value)]\n",
    "\n",
    "    result_df = result_df.T\n",
    "\n",
    "    result_df_sorted = result_df.sort_values(by='PCC', ascending=False)\n",
    "\n",
    "    return result_df_sorted\n",
    "\n",
    "get_PCC(adata_pre, adata_true)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "040b31e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GNAS 0.34341222025864626\n",
      "UBA52 0.13513815851744454\n",
      "MUCL1 0.09126696780269167\n"
     ]
    }
   ],
   "source": [
    "# get the value\n",
    "from scipy.stats import pearsonr\n",
    "\n",
    "genes = ['GNAS', 'UBA52', 'MUCL1']\n",
    "for gene in genes:\n",
    "    gene_true = adata_true[gene]\n",
    "    gene_pre = adata_pre[gene]\n",
    "    # FASN_true.shape\n",
    "\n",
    "    pearson_corrcoef, p_value = pearsonr(gene_true, gene_pre)\n",
    "    # print('Pearson Correlation:',np.nanmean(R))\n",
    "    print(gene, pearson_corrcoef)\n",
    "\n",
    "# MYL12B\n",
    "# 8 0.6207262018691467\n",
    "# GNAS\n",
    "# 7 0.5906530987577501\n",
    "# FASN\n",
    "# 14 0.510873127279465\n",
    "# SCD\n",
    "# 10 0.6103439096566053"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "00bc3b1d",
   "metadata": {},
   "source": [
    "Predict gene expression:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11f1208b",
   "metadata": {},
   "outputs": [],
   "source": [
    "Gene_show = 'GNAS'  # ['GNAS', 'UBA52', 'MUCL1']\n",
    "\n",
    "model.load_state_dict(torch.load(f'./model/HGGEP/{fold}-Hist2ST.ckpt'))\n",
    "# device = torch.device('cpu')\n",
    "# dataset = ViT_HER2ST(train=False,sr=False,fold=fold)\n",
    "# test_loader = DataLoader(dataset, batch_size=1, num_workers=4)\n",
    "adata_pred, adata_truth = test(model, test_loader)\n",
    "adata_pred = comp_tsne_km(adata_pred,4)\n",
    "\n",
    "g = list(np.load('data/her_hvg_cut_1000.npy',allow_pickle=True))\n",
    "adata_pred.var_names = g\n",
    "adata_truth.var_names = g\n",
    "\n",
    "# Create the spatial plot\n",
    "sc.pl.spatial(adata_pred, img=None, color=Gene_show, spot_size=150, color_map='magma', show=False)\n",
    "\n",
    "# Save the figure with a pure white background and high resolution (e.g., 300 dpi)\n",
    "# plt.savefig('./figures/spatial_plot.png', dpi=600, facecolor='white')\n",
    "plt.savefig('./figures/5_spatial_'+str(Gene_show)+'.pdf')\n",
    "\n",
    "# If you want to display the plot as well\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7c7ab55",
   "metadata": {},
   "outputs": [],
   "source": [
    "# sc.pl.spatial(pred, img=None, color='kmeans', spot_size=112)\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "img_path = './data/her2st/data/ST-imgs/B/B1/BT10737_C1_HE.jpg'\n",
    "img_data = plt.imread(img_path)\n",
    "# print(pred.shape, pred[0][:10])\n",
    "sc.pl.spatial(pred, img=img_data, color='kmeans', spot_size=112)\n",
    "sc.pl.spatial(gt, img=img_data, color='kmeans', spot_size=112)\n",
    "\n",
    "sc.pl.spatial(adata_pred, img=None, color='FASN', spot_size=112, color_map='magma')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f613201-90c3-4cbf-be06-008b63155c29",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.spatial(pred, img=None, color='kmeans', spot_size=112)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9ebecb72",
   "metadata": {},
   "source": [
    "# Figure 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5075f01d",
   "metadata": {},
   "outputs": [],
   "source": [
    "Gene_show = 'GNAS'  # ['GNAS', 'UBA52', 'MUCL1']\n",
    "\n",
    "img_path = './data/her2st/data/ST-imgs/B/B3/BT10737_D1_HE.jpg'\n",
    "img_data = plt.imread(img_path)\n",
    "\n",
    "model.load_state_dict(torch.load(f'./model/her2st/{fold}-Hist2ST.ckpt'))\n",
    "# device = torch.device('cpu')\n",
    "# dataset = ViT_HER2ST(train=False,sr=False,fold=fold)\n",
    "# test_loader = DataLoader(dataset, batch_size=1, num_workers=4)\n",
    "adata_pred, adata_truth = test(model, test_loader)\n",
    "adata_pred = comp_tsne_km(adata_pred,4)\n",
    "\n",
    "g = list(np.load('data/her_hvg_cut_1000.npy',allow_pickle=True))\n",
    "adata_pred.var_names = g\n",
    "adata_truth.var_names = g\n",
    "\n",
    "# Create the spatial plot\n",
    "sc.pl.spatial(adata_pred, img=None, color=Gene_show, spot_size=150, color_map='magma', show=False)\n",
    "\n",
    "# Save the figure with a pure white background and high resolution (e.g., 300 dpi)\n",
    "plt.savefig('./figures/spatial_plot.pdf')\n",
    "\n",
    "# If you want to display the plot as well\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4073ea73",
   "metadata": {},
   "source": [
    "# Figure 5 : HGGEP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "edf905ab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['A2']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.50s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['A3']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.56s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['A4']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.22s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['A5']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.12s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['A6']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.98s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['B1']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.49s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['B2']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00,  1.06it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['B3']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.87s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['B4']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.25s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['B5']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.20s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['B6']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00,  1.01it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['C1']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00,  1.35it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['C2']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.31s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['C3']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00,  1.51it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['C4']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00,  1.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['C5']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00,  1.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['C6']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00,  1.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['D1']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:02<00:00,  2.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['D2']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.07s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['D3']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.32s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['D4']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.48s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['D5']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.78s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['D6']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.84s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['E1']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:02<00:00,  2.31s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['E2']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:02<00:00,  2.19s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['E3']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.87s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['F1']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:03<00:00,  3.05s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['F2']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:03<00:00,  3.61s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['F3']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:02<00:00,  2.47s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['G1']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:02<00:00,  2.69s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['G2']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.78s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "te_names: ['G3']\n",
      "Loading imgs...\n",
      "Loading metadata...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:01<00:00,  1.60s/it]\n"
     ]
    }
   ],
   "source": [
    "folds=[]\n",
    "\n",
    "from HGGEP import *\n",
    "for i in range(32):\n",
    "    folds.append(i)\n",
    "# print(folds)\n",
    "data='her2st'\n",
    "prune='Grid' if data=='her2st' else 'NA'\n",
    "genes=171 if data=='cscc' else 785\n",
    "gene_names = list(np.load('data/her_hvg_cut_1000.npy',allow_pickle=True))\n",
    "\n",
    "def comp_tsne_km(adata,k=10):\n",
    "    sc.pp.pca(adata)\n",
    "    sc.tl.tsne(adata)\n",
    "    kmeans = KMeans(n_clusters=k, init=\"k-means++\", random_state=0).fit(adata.obsm['X_pca'])\n",
    "    adata.obs['kmeans'] = kmeans.labels_.astype(str)\n",
    "    return adata\n",
    "\n",
    "def get_PCC(adata_pre, adata_true):\n",
    "    gene_names = adata_pre.columns\n",
    "    result_df = pd.DataFrame(index=['-log10_PV', 'PCC', 'PV',], columns=gene_names)\n",
    "\n",
    "    for gene in gene_names:\n",
    "        corr_coef, p_value = pearsonr(adata_pre[gene], adata_true[gene])\n",
    "        result_df[gene] = [-np.log10(p_value), corr_coef, p_value]\n",
    "\n",
    "    result_df = result_df.T\n",
    "\n",
    "    result_df_sorted = result_df.sort_values(by='-log10_PV', ascending=False)\n",
    "\n",
    "    return result_df_sorted\n",
    "\n",
    "HGGEP_results_all = []\n",
    "for fold in folds:\n",
    "    testset = pk_load(fold,'test',dataset=data,flatten=False,adj=True,ori=True,prune=prune)\n",
    "    test_loader = DataLoader(testset, batch_size=1, num_workers=0, shuffle=False)\n",
    "    label=testset.label[testset.names[0]]\n",
    "    model=HGGEP(\n",
    "        depth1=d1, depth2=d2,depth3=d3,n_genes=genes, \n",
    "        kernel_size=k, patch_size=p,\n",
    "        heads=h, channel=c, dropout=0.2,\n",
    "        zinb=0.25, nb=False,\n",
    "        bake=5, lamb=0.5,\n",
    "    )\n",
    "    model.load_state_dict(torch.load(f'./model/her2st/{fold}-Hist2ST.ckpt'))\n",
    "    adata_pred, adata_truth = test(model, test_loader)\n",
    "\n",
    "    adata_pred = comp_tsne_km(adata_pred,4)\n",
    "    sc.pp.scale(adata_pred)\n",
    "    sc.pp.scale(adata_truth)\n",
    "    adata_pred.var_names = gene_names\n",
    "    adata_truth.var_names = gene_names\n",
    "    adata_true = pd.DataFrame(data=adata_truth.X, index=adata_truth.obs_names, columns=adata_truth.var_names)\n",
    "    adata_pre = pd.DataFrame(data=adata_pred.X, index=adata_pred.obs_names, columns=adata_pred.var_names)\n",
    "    res = get_PCC(adata_pre, adata_true)\n",
    "    # print(res)\n",
    "    HGGEP_results_all.append(res)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "557de6ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MYL12B\n",
      "9 0.6516888665834257\n",
      "GNAS\n",
      "7 0.6371508417594052\n",
      "FASN\n",
      "16 0.5637526702753419\n",
      "SCD\n",
      "10 0.6495225608741068\n"
     ]
    }
   ],
   "source": [
    "from scipy.stats import pearsonr\n",
    "\n",
    "genes = ['MYL12B', 'GNAS', 'FASN', 'SCD\n",
    "']\n",
    "for gene in genes:\n",
    "    print(gene)\n",
    "    PCC_best, fold_best = 0, 0\n",
    "    for fold in folds:\n",
    "        pearson_corrcoef = HGGEP_results_all[fold].loc[gene]\n",
    "        PCC_value = pearson_corrcoef.loc['PCC']\n",
    "        if PCC_value > PCC_best:\n",
    "            PCC_best = PCC_value\n",
    "            fold_best = fold\n",
    "    print(fold_best, PCC_best)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "44cee80f",
   "metadata": {},
   "source": [
    "# ARI figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "845f8fc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import cv2\n",
    "from HGGEP import *\n",
    "from predict import *\n",
    "\n",
    "fold = 30\n",
    "resolution = 0.1\n",
    "\n",
    "# img_path = '/home/zy/libo/Hist2ST/data/her2st/data/ST-imgs/G/G2/HE_BT23903_C2a.jpg'\n",
    "# img_data = plt.imread(img_path)\n",
    "\n",
    "# gray_img = cv2.cvtColor(img_data, cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "# save_path = './figures/ARI_image_gray_'+str(fold)+'.png'\n",
    "# cv2.imwrite(save_path, gray_img)\n",
    "\n",
    "def comp_tsne_leiden(adata, resolution=1.0):\n",
    "    sc.pp.pca(adata)\n",
    "    sc.tl.tsne(adata)\n",
    "    \n",
    "    # Compute neighborhood graph\n",
    "    sc.pp.neighbors(adata)\n",
    "    \n",
    "    # Use Scanpy's Leiden clustering\n",
    "    sc.tl.leiden(adata, key_added='leiden', resolution=resolution)\n",
    "    \n",
    "    adata.obs['leiden'] = adata.obs['leiden'].astype(str)\n",
    "    \n",
    "    return adata\n",
    "\n",
    "def comp_tsne_km(adata,k=10):\n",
    "    sc.pp.pca(adata)\n",
    "    sc.tl.tsne(adata)\n",
    "    kmeans = KMeans(n_clusters=k, init=\"k-means++\", random_state=0).fit(adata.obsm['X_pca'])\n",
    "    adata.obs['kmeans'] = kmeans.labels_.astype(str)\n",
    "    return adata\n",
    "\n",
    "def get_PCC(adata_pre, adata_true):\n",
    "    gene_names = adata_pre.columns\n",
    "    result_df = pd.DataFrame(index=['-log10_PV', 'PCC', 'PV',], columns=gene_names)\n",
    "\n",
    "    for gene in gene_names:\n",
    "        corr_coef, p_value = pearsonr(adata_pre[gene], adata_true[gene])\n",
    "        result_df[gene] = [-np.log10(p_value), corr_coef, p_value]\n",
    "\n",
    "    result_df = result_df.T\n",
    "\n",
    "    result_df_sorted = result_df.sort_values(by='-log10_PV', ascending=False)\n",
    "\n",
    "    return result_df_sorted\n",
    "\n",
    "gene_names = list(np.load('data/her_hvg_cut_1000.npy',allow_pickle=True))\n",
    "testset = pk_load(fold,'test',dataset='her2st',flatten=False,adj=True,ori=True,prune=prune)\n",
    "test_loader = DataLoader(testset, batch_size=1, num_workers=0, shuffle=False)\n",
    "label=testset.label[testset.names[0]]\n",
    "model=HGGEP(\n",
    "    depth1=d1, depth2=d2,depth3=d3,n_genes=genes, \n",
    "    kernel_size=k, patch_size=p,\n",
    "    heads=h, channel=c, dropout=0.2,\n",
    "    zinb=0.25, nb=False,\n",
    "    bake=5, lamb=0.5,\n",
    ")\n",
    "ARI_best, model_fold = 0, 0\n",
    "for i in range(32):\n",
    "    model.load_state_dict(torch.load(f'./model/her2st/'+str(i)+'-HGGEP.ckpt'))\n",
    "\n",
    "    adata_pred, adata_truth = test(model, test_loader)\n",
    "\n",
    "    clus,ARI=cluster(adata_pred,label,method='leiden', resolution=resolution)\n",
    "    print(i, 'ARI:',ARI)\n",
    "    if ARI>ARI_best:\n",
    "        model_fold = i\n",
    "        ARI_best = ARI\n",
    "print('ARI_best:',ARI_best, model_fold)\n",
    "\n",
    "# model_fold = 17\n",
    "model.load_state_dict(torch.load(f'./model/her2st/'+str(model_fold)+'-HGGEP.ckpt'))\n",
    "\n",
    "adata_pred, adata_truth = test(model, test_loader)\n",
    "\n",
    "clus,ARI=cluster(adata_pred,label,method='leiden', resolution=resolution)  \n",
    "# 5: kmeans : ARI: 0.174, leiden : ARI: 0.112\n",
    "# 17: kmeans : ARI: 0.271, leiden : ARI: 0.121\n",
    "print('ARI:',ARI)\n",
    "\n",
    "# adata_pred = comp_tsne_km(adata_pred,3)\n",
    "adata_pred = comp_tsne_leiden(adata_pred, resolution=resolution)\n",
    "\n",
    "g = list(np.load('data/her_hvg_cut_1000.npy',allow_pickle=True))\n",
    "adata_pred.var_names = g\n",
    "adata_truth.var_names = g\n",
    "\n",
    "plt.gcf().set_facecolor('white')\n",
    "\n",
    "sc.pl.spatial(adata_pred, img=None, color='leiden', spot_size=112, show=False)\n",
    "# sc.pl.spatial(adata_pred, img=None, color='kmeans', spot_size=112, show=False) \n",
    "\n",
    "plt.savefig('./figures/Fig6_ARI_'+str(fold)+'.pdf')\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gtn",
   "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.13"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  },
  "vscode": {
   "interpreter": {
    "hash": "f5bc5193494f8e31bb35f143e2e9d2a9395cc76e28c0276d02df8e52647f48fb"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}