{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2780405c-b820-4c85-809f-ec98e9f5c8da",
   "metadata": {},
   "source": [
    "# Integrating the four S1 slices from MISAR-seq MB dataset\n",
    "\n",
    "Before proceeding, ensure that INSTINCT is installed.  \n",
    "You can follow the instruction provided on [Github](https://github.com/yyLIU12138/INSTINCT) for installation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa35ca95-d271-4c3b-8697-e9cab3648349",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import csv\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import anndata as ad\n",
    "import networkx as nx\n",
    "\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.mixture import GaussianMixture\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "import INSTINCT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a965acf6-5ba2-4c01-b5fc-721fa0702dbb",
   "metadata": {},
   "source": [
    "### Load the raw data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3eb7707d-f8ff-41c4-ae7d-ea1e74caf549",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = '../../data/spMOdata/EpiTran_MouseBrain_Jiang2023/preprocessed/'\n",
    "slice_name_list = ['E11_0-S1', 'E13_5-S1', 'E15_5-S1', 'E18_5-S1']\n",
    "slice_index_list = list(range(len(slice_name_list)))\n",
    "\n",
    "save_dir = '../../results/MouseBrain_Jiang2023/'\n",
    "if not os.path.exists(save_dir):\n",
    "    os.makedirs(save_dir)\n",
    "\n",
    "cas_dict = {}\n",
    "for sample in slice_name_list:\n",
    "    sample_data = ad.read_h5ad(data_dir + sample + '_atac.h5ad')\n",
    "\n",
    "    if 'insertion' in sample_data.obsm:\n",
    "        del sample_data.obsm['insertion']\n",
    "\n",
    "    cas_dict[sample] = sample_data\n",
    "cas_list = [cas_dict[sample] for sample in slice_name_list]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e97c69d-e9f9-4f75-b361-e41baf245d35",
   "metadata": {},
   "source": [
    "### Merge the peaks\n",
    "Initially, nearly no common peaks are shared among the four slices.  \n",
    "Therefore, we nead to merge the four peak sets into a common merged peak set, such that the four slices have indentical peak sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a014cdc8-4f99-4ac9-80f9-fb0b25abf52e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# merge peaks\n",
    "cas_list = INSTINCT.peak_sets_alignment(cas_list)\n",
    "\n",
    "# save the merged data\n",
    "for idx, adata in enumerate(cas_list):\n",
    "    adata.write_h5ad(f'{data_dir}merged_{slice_name_list[idx]}_atac.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e050b0b-73a2-4fbf-b8d5-0e0836e4fd91",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the merged data\n",
    "cas_list = [ad.read_h5ad(data_dir + 'merged_' + sample + '_atac.h5ad') for sample in slice_name_list]\n",
    "\n",
    "# rename the spots\n",
    "for j in range(len(cas_list)):\n",
    "    cas_list[j].obs_names = [x + '_' + slice_name_list[j] for x in cas_list[j].obs_names]\n",
    "\n",
    "# concatenation\n",
    "adata_concat = ad.concat(cas_list, label=\"slice_name\", keys=slice_name_list)\n",
    "# adata_concat.obs_names_make_unique()\n",
    "print(adata_concat.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "beeb83e9-9d78-46fd-bb51-d4873b66b086",
   "metadata": {},
   "source": [
    "### Data preprocessing\n",
    "The concatenated data matrix is read count matrix, so we set use_fragment_count=True to transform it into fragment count matrix.  \n",
    "\n",
    "Please not that data samples in cas_list are only filtered but not transformed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03cdf564-56b9-4a81-a19e-694af5e1f4cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# preprocess CAS data\n",
    "print('Start preprocessing')\n",
    "INSTINCT.preprocess_CAS(cas_list, adata_concat, use_fragment_count=True, min_cells_rate=0.03)\n",
    "print(adata_concat.shape)\n",
    "print('Done!')\n",
    "\n",
    "# save the data\n",
    "adata_concat.write_h5ad(save_dir + f\"preprocessed_concat_atac.h5ad\")\n",
    "for i in range(len(slice_name_list)):\n",
    "    cas_list[i].write_h5ad(save_dir + f\"filtered_merged_{slice_name_list[i]}_atac.h5ad\")\n",
    "\n",
    "# load the data\n",
    "cas_list = [ad.read_h5ad(save_dir + f\"filtered_merged_{sample}_atac.h5ad\") for sample in slice_name_list]\n",
    "origin_concat = ad.concat(cas_list, label=\"slice_name\", keys=slice_name_list)\n",
    "adata_concat = ad.read_h5ad(save_dir + f\"preprocessed_concat_atac.h5ad\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ec89d9c-6173-447d-93e8-da9e26426662",
   "metadata": {},
   "source": [
    "### Perform PCA\n",
    "We performed PCA and the first 100 pcs are used as input.  \n",
    "The PCA results should be added to .obsm['X_pca'] of the concatenated data, or you can name it by yourself and set the 'input_mat_key' of INSTINCT_Model to the corresponding name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8476e23c-4ff9-4be3-8e0b-f2a2de7cc5f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'Applying PCA to reduce the feature dimension to 100 ...')\n",
    "pca = PCA(n_components=100, random_state=1234)\n",
    "input_matrix = pca.fit_transform(adata_concat.X.toarray())\n",
    "np.save(save_dir + 'input_matrix_atac.npy', input_matrix)\n",
    "print('Done !')\n",
    "\n",
    "input_matrix = np.load(save_dir + 'input_matrix_atac.npy')\n",
    "adata_concat.obsm['X_pca'] = input_matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3df90da-72fd-42ed-8aa6-a83fd1e7bb62",
   "metadata": {},
   "source": [
    "### Create neighbor graph\n",
    "The neighbor graph for each slice is added to .uns['graph_list'] of the concatenated data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c1f466e-c7fb-42d1-aecc-38d843e092f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate the spatial graph\n",
    "INSTINCT.create_neighbor_graph(cas_list, adata_concat)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3977477a-e6f0-4edc-aef5-8ef16485e787",
   "metadata": {},
   "source": [
    "### Run the model\n",
    "The latent representation matrices containing representations for spots are included to .obsm['INSTINCT_latent'] of corresponding slices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b35e5a34-cf2d-4056-a374-cc004775025b",
   "metadata": {},
   "outputs": [],
   "source": [
    "INSTINCT_model = INSTINCT.INSTINCT_Model(cas_list, adata_concat, device=device)\n",
    "\n",
    "INSTINCT_model.train(report_loss=True, report_interval=100)\n",
    "\n",
    "INSTINCT_model.eval(cas_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "156ca191-a276-41ab-8b51-1d12f67dbbad",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = ad.concat(cas_list, label=\"slice_name\", keys=slice_name_list)\n",
    "\n",
    "with open(save_dir + 'INSTINCT_embed.csv', 'w', newline='') as file:\n",
    "    writer = csv.writer(file)\n",
    "    writer.writerows(result.obsm['INSTINCT_latent'])\n",
    "\n",
    "with open(save_dir + 'INSTINCT_noise_embed.csv', 'w', newline='') as file:\n",
    "    writer = csv.writer(file)\n",
    "    writer.writerows(result.obsm['INSTINCT_latent_noise'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "596f2a1b-9a77-4690-a266-d566353a8a1c",
   "metadata": {},
   "source": [
    "### Clustering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcdb2515-122d-4bdb-a401-16ae6e6aad7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def match_cluster_labels(true_labels, est_labels):\n",
    "    true_labels_arr = np.array(list(true_labels))\n",
    "    est_labels_arr = np.array(list(est_labels))\n",
    "\n",
    "    org_cat = list(np.sort(list(pd.unique(true_labels))))\n",
    "    est_cat = list(np.sort(list(pd.unique(est_labels))))\n",
    "\n",
    "    B = nx.Graph()\n",
    "    B.add_nodes_from([i + 1 for i in range(len(org_cat))], bipartite=0)\n",
    "    B.add_nodes_from([-j - 1 for j in range(len(est_cat))], bipartite=1)\n",
    "\n",
    "    for i in range(len(org_cat)):\n",
    "        for j in range(len(est_cat)):\n",
    "            weight = np.sum((true_labels_arr == org_cat[i]) * (est_labels_arr == est_cat[j]))\n",
    "            B.add_edge(i + 1, -j - 1, weight=-weight)\n",
    "\n",
    "    match = nx.algorithms.bipartite.matching.minimum_weight_full_matching(B)\n",
    "\n",
    "    if len(org_cat) >= len(est_cat):\n",
    "        return np.array([match[-est_cat.index(c) - 1] - 1 for c in est_labels_arr])\n",
    "    else:\n",
    "        unmatched = [c for c in est_cat if not (-est_cat.index(c) - 1) in match.keys()]\n",
    "        l = []\n",
    "        for c in est_labels_arr:\n",
    "            if (-est_cat.index(c) - 1) in match:\n",
    "                l.append(match[-est_cat.index(c) - 1] - 1)\n",
    "            else:\n",
    "                l.append(len(org_cat) + unmatched.index(c))\n",
    "        return np.array(l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "196ceb3e-134d-4b97-861c-b25d8ba5a89b",
   "metadata": {},
   "outputs": [],
   "source": [
    "gm = GaussianMixture(n_components=16, covariance_type='tied', random_state=1234)\n",
    "y = gm.fit_predict(result.obsm['INSTINCT_latent'], y=None)\n",
    "result.obs[\"gm_clusters\"] = pd.Series(y, index=result.obs.index, dtype='category')\n",
    "result.obs['matched_clusters'] = pd.Series(match_cluster_labels(result.obs['Annotation_for_Combined'],\n",
    "                                                                result.obs[\"gm_clusters\"]),\n",
    "                                           index=result.obs.index, dtype='category')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f376a87-3f99-4c85-8c5d-f6c15aec3924",
   "metadata": {},
   "source": [
    "### Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57deeab5-7175-46a8-9b94-e0fa99bda7d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.metrics.cluster import adjusted_rand_score\n",
    "from sklearn.metrics.cluster import normalized_mutual_info_score\n",
    "from sklearn.metrics.cluster import fowlkes_mallows_score\n",
    "from sklearn.metrics.cluster import homogeneity_score\n",
    "from sklearn.metrics.cluster import adjusted_mutual_info_score\n",
    "from sklearn.metrics.cluster import completeness_score\n",
    "import sklearn\n",
    "import sklearn.neighbors\n",
    "import pandas as pd\n",
    "import networkx as nx\n",
    "\n",
    "import scib\n",
    "import scanpy as sc\n",
    "\n",
    "\n",
    "def cluster_metrics(target, pred):\n",
    "    target = np.array(target)\n",
    "    pred = np.array(pred)\n",
    "    \n",
    "    ari = adjusted_rand_score(target, pred)\n",
    "    ami = adjusted_mutual_info_score(target, pred)\n",
    "    nmi = normalized_mutual_info_score(target, pred)\n",
    "    fmi = fowlkes_mallows_score(target, pred)\n",
    "    comp = completeness_score(target, pred)\n",
    "    homo = homogeneity_score(target, pred)\n",
    "    print('ARI: %.3f, AMI: %.3f, NMI: %.3f, FMI: %.3f, Comp: %.3f, Homo: %.3f' % (ari, ami, nmi, fmi, comp, homo))\n",
    "    \n",
    "    return ari, ami, nmi, fmi, comp, homo\n",
    "\n",
    "\n",
    "def mean_average_precision(x: np.ndarray, y: np.ndarray, k: int=30, **kwargs) -> float:\n",
    "    r\"\"\"\n",
    "    Mean average precision\n",
    "    Parameters\n",
    "    ----------\n",
    "    x\n",
    "        Coordinates\n",
    "    y\n",
    "        Cell_type/Layer labels\n",
    "    k\n",
    "        k neighbors\n",
    "    **kwargs\n",
    "        Additional keyword arguments are passed to\n",
    "        :class:`sklearn.neighbors.NearestNeighbors`\n",
    "    Returns\n",
    "    -------\n",
    "    map\n",
    "        Mean average precision\n",
    "    \"\"\"\n",
    "    \n",
    "    def _average_precision(match: np.ndarray) -> float:\n",
    "        if np.any(match):\n",
    "            cummean = np.cumsum(match) / (np.arange(match.size) + 1)\n",
    "            return cummean[match].mean().item()\n",
    "        return 0.0\n",
    "    \n",
    "    y = np.array(y)\n",
    "    knn = sklearn.neighbors.NearestNeighbors(n_neighbors=min(y.shape[0], k + 1), **kwargs).fit(x)\n",
    "    nni = knn.kneighbors(x, return_distance=False)\n",
    "    match = np.equal(y[nni[:, 1:]], np.expand_dims(y, 1))\n",
    "    \n",
    "    return np.apply_along_axis(_average_precision, 1, match).mean().item()\n",
    "\n",
    "\n",
    "def bio_conservation_metrics(adata, use_rep, label_key, batch_key, k_map=30, threshold=1):\n",
    "    if label_key not in adata.obs or batch_key not in adata.obs or use_rep not in adata.obsm:\n",
    "        print(\"KeyError\")\n",
    "        return None\n",
    "\n",
    "    adata.obs[label_key] = adata.obs[label_key].astype(str).astype(\"category\")\n",
    "    adata.obs[batch_key] = adata.obs[batch_key].astype(str).astype(\"category\")\n",
    "    # sc.pp.neighbors(adata, use_rep=use_rep, random_state=1234)\n",
    "\n",
    "    MAP = mean_average_precision(adata.obsm[use_rep].copy(), adata.obs[label_key], k=k_map)\n",
    "    cell_type_ASW = scib.me.silhouette(adata, label_key=label_key, embed=use_rep)\n",
    "    isolated_asw = scib.me.isolated_labels_asw(adata, batch_key=batch_key, label_key=label_key, embed=use_rep,\n",
    "                                               iso_threshold=threshold)\n",
    "    isolated_f1 = scib.me.isolated_labels_f1(adata, batch_key=batch_key, label_key=label_key, embed=use_rep,\n",
    "                                             iso_threshold=threshold)\n",
    "\n",
    "    print('mAP: %.3f, Cell type ASW: %.3f, Isolated label ASW: %.3f, Isolated label F1: %.3f' %\n",
    "          (MAP, cell_type_ASW, isolated_asw, isolated_f1))\n",
    "\n",
    "    return MAP, cell_type_ASW, isolated_asw, isolated_f1\n",
    "\n",
    "\n",
    "def batch_correction_metrics(adata, origin_concat, use_rep, label_key, batch_key):\n",
    "    if label_key not in adata.obs or batch_key not in adata.obs or use_rep not in adata.obsm:\n",
    "        print(\"KeyError\")\n",
    "        return None\n",
    "\n",
    "    adata.obs[label_key] = adata.obs[label_key].astype(str).astype(\"category\")\n",
    "    adata.obs[batch_key] = adata.obs[batch_key].astype(str).astype(\"category\")\n",
    "    origin_concat.X = origin_concat.X.astype(float)\n",
    "    sc.pp.neighbors(adata, use_rep=use_rep, random_state=1234)\n",
    "\n",
    "    # g_iLISI = scib.me.ilisi_graph(adata, batch_key=batch_key, type_=\"embed\", use_rep=use_rep)\n",
    "    batch_ASW = scib.me.silhouette_batch(adata, batch_key=batch_key, label_key=label_key, embed=use_rep, verbose=False)\n",
    "    batch_PCR = scib.me.pcr_comparison(origin_concat, adata, covariate=batch_key, embed=use_rep)\n",
    "    kBET = scib.me.kBET(adata, batch_key=batch_key, label_key=label_key, type_='embed', embed=use_rep)\n",
    "    g_conn = scib.me.graph_connectivity(adata, label_key=label_key)\n",
    "    print('Batch ASW: %.3f, Batch PCR: %.3f, kBET: %.3f, Graph connectivity: %.3f' %\n",
    "          (batch_ASW, batch_PCR, kBET, g_conn))\n",
    "    return batch_ASW, batch_PCR, kBET, g_conn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd920483-a4db-4830-b019-431876789be2",
   "metadata": {},
   "outputs": [],
   "source": [
    "ari, ami, nmi, fmi, comp, homo = cluster_metrics(result.obs['Annotation_for_Combined'],\n",
    "                                                 result.obs['matched_clusters'].tolist())\n",
    "map, c_asw, i_asw, i_f1 = bio_conservation_metrics(result, use_rep='INSTINCT_latent',\n",
    "                                                   label_key='Annotation_for_Combined', batch_key='slice_name')\n",
    "b_asw, b_pcr, kbet, g_conn = batch_correction_metrics(result, origin_concat, use_rep='INSTINCT_latent',\n",
    "                                                      label_key='Annotation_for_Combined',\n",
    "                                                      batch_key='slice_name')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebb9f129-2faf-4a7c-a26a-af592adddf34",
   "metadata": {},
   "source": [
    "### Spatial domain identification and UMAP visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e1218ab-f27c-40b7-9f21-6007cc0fb817",
   "metadata": {},
   "outputs": [],
   "source": [
    "cls_list = ['Primary_brain_1', 'Primary_brain_2', 'Midbrain', 'Diencephalon_and_hindbrain', 'Basal_plate_of_hindbrain',\n",
    "            'Subpallium_1', 'Subpallium_2', 'Cartilage_1', 'Cartilage_2', 'Cartilage_3', 'Cartilage_4',\n",
    "            'Mesenchyme', 'Muscle', 'Thalamus', 'DPallm', 'DPallv']\n",
    "\n",
    "colors_for_clusters = ['red', 'tomato', 'chocolate', 'orange', 'goldenrod',\n",
    "                       'b', 'royalblue', 'g', 'limegreen', 'lime', 'springgreen',\n",
    "                       'deepskyblue', 'pink', 'fuchsia', 'yellowgreen', 'olivedrab']\n",
    "\n",
    "order_for_clusters = [11, 12, 9, 7, 0, 13, 14, 1, 2, 3, 4, 8, 10, 15, 5, 6]\n",
    "\n",
    "cluster_to_color_map = {cluster: color for cluster, color in zip(cls_list, colors_for_clusters)}\n",
    "order_to_cluster_map = {order: cluster for order, cluster in zip(order_for_clusters, cls_list)}\n",
    "\n",
    "from umap.umap_ import UMAP\n",
    "reducer = UMAP(n_neighbors=30, n_components=2, metric=\"correlation\", n_epochs=None, learning_rate=1.0,\n",
    "               min_dist=0.3, spread=1.0, set_op_mix_ratio=1.0, local_connectivity=1, repulsion_strength=1,\n",
    "               negative_sample_rate=5, a=None, b=None, random_state=1234, metric_kwds=None,\n",
    "               angular_rp_forest=False, verbose=False)\n",
    "\n",
    "gm = GaussianMixture(n_components=len(cls_list), covariance_type='tied', random_state=1234)\n",
    "y = gm.fit_predict(result.obsm['INSTINCT_latent'], y=None)\n",
    "result.obs[\"gm_clusters\"] = pd.Series(y, index=result.obs.index, dtype='category')\n",
    "result.obs['matched_clusters'] = pd.Series(match_cluster_labels(\n",
    "    result.obs['Annotation_for_Combined'], result.obs[\"gm_clusters\"]),\n",
    "    index=result.obs.index, dtype='category')\n",
    "my_clusters = np.sort(list(set(result.obs['matched_clusters'])))\n",
    "matched_colors = [cluster_to_color_map[order_to_cluster_map[order]] for order in my_clusters]\n",
    "matched_to_color_map = {matched: color for matched, color in zip(my_clusters, matched_colors)}\n",
    "\n",
    "spots_count = [0]\n",
    "n = 0\n",
    "for sample in cas_list:\n",
    "    num = sample.shape[0]\n",
    "    n += num\n",
    "    spots_count.append(n)\n",
    "\n",
    "for i in range(len(cas_list)):\n",
    "    cas_list[i].obs['matched_clusters'] = result.obs['matched_clusters'][spots_count[i]:spots_count[i+1]]\n",
    "\n",
    "sp_embedding = reducer.fit_transform(result.obsm['INSTINCT_latent'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d55a9cb-7439-4d3e-8178-1fd6c114455c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from matplotlib.lines import Line2D\n",
    "import matplotlib as mpl\n",
    "mpl.rcParams['pdf.fonttype'] = 42\n",
    "mpl.rcParams['ps.fonttype'] = 42\n",
    "\n",
    "\n",
    "def plot_mousebrain(cas_list, adata_concat, ground_truth_key, matched_clusters_key, model,\n",
    "                    cluster_to_color_map, matched_to_color_map, cluster_orders, slice_name_list, cls_list,\n",
    "                    sp_embedding, save_root=None, frame_color=None, save=False, plot=False):\n",
    "\n",
    "    fig, axs = plt.subplots(2, 4, figsize=(15, 7))\n",
    "    fig.suptitle(f'{model} Clustering Results', fontsize=16)\n",
    "    for i in range(len(cas_list)):\n",
    "        real_colors = list(cas_list[i].obs[ground_truth_key].astype('str').map(cluster_to_color_map))\n",
    "        axs[0, i].scatter(cas_list[i].obsm['spatial'][:, 0], cas_list[i].obsm['spatial'][:, 1], linewidth=0.5, s=30,\n",
    "                          marker=\".\", color=real_colors, alpha=0.9)\n",
    "        axs[0, i].set_title(f'{slice_name_list[i]} (Ground Truth)', size=12)\n",
    "        axs[0, i].invert_yaxis()\n",
    "        axs[0, i].axis('off')\n",
    "\n",
    "        cluster_colors = list(cas_list[i].obs[matched_clusters_key].map(matched_to_color_map))\n",
    "        axs[1, i].scatter(cas_list[i].obsm['spatial'][:, 0], cas_list[i].obsm['spatial'][:, 1], linewidth=0.5, s=30,\n",
    "                          marker=\".\", color=cluster_colors, alpha=0.9)\n",
    "        axs[1, i].set_title(f'{slice_name_list[i]} (Cluster Results)', size=12)\n",
    "        axs[1, i].invert_yaxis()\n",
    "        axs[1, i].axis('off')\n",
    "\n",
    "    legend_handles_1 = [\n",
    "        Line2D([0], [0], marker='o', color='w', markersize=8, markerfacecolor=cluster_to_color_map[cluster],\n",
    "               label=cluster) for cluster in cls_list\n",
    "    ]\n",
    "    axs[0, 3].legend(\n",
    "        handles=legend_handles_1,\n",
    "        fontsize=8, title='Spot-types', title_fontsize=10, bbox_to_anchor=(1, 1.15))\n",
    "    legend_handles_2 = [\n",
    "        Line2D([0], [0], marker='o', color='w', markersize=8, markerfacecolor=matched_to_color_map[order],\n",
    "               label=f'{i}') for i, order in enumerate(cluster_orders)\n",
    "    ]\n",
    "    axs[1, 3].legend(\n",
    "        handles=legend_handles_2,\n",
    "        fontsize=8, title='Clusters', title_fontsize=10, bbox_to_anchor=(1, 1.1))\n",
    "    plt.gcf().subplots_adjust(left=0.05, top=None, bottom=None, right=0.85)\n",
    "    if save:\n",
    "        save_path = save_root + f'/{model}_clustering_results.pdf'\n",
    "        plt.savefig(save_path, dpi=500)\n",
    "\n",
    "    n_spots = adata_concat.shape[0]\n",
    "    size = 10000 / n_spots\n",
    "    order = np.arange(n_spots)\n",
    "    colors_for_slices = [[0.2298057, 0.29871797, 0.75368315],\n",
    "                         [0.70567316, 0.01555616, 0.15023281],\n",
    "                         [0.2298057, 0.70567316, 0.15023281],\n",
    "                         [0.5830223, 0.59200322, 0.12993134]]\n",
    "    slice_cmap = {slice_name_list[i]: colors_for_slices[i] for i in range(len(slice_name_list))}\n",
    "    colors = list(adata_concat.obs['slice_name'].astype('str').map(slice_cmap))\n",
    "    plt.figure(figsize=(5, 5))\n",
    "    if frame_color:\n",
    "        plt.rc('axes', edgecolor=frame_color, linewidth=2)\n",
    "    plt.scatter(sp_embedding[order, 0], sp_embedding[order, 1], s=size, c=colors)\n",
    "    plt.tick_params(axis='both', bottom=False, top=False, left=False, right=False,\n",
    "                    labelleft=False, labelbottom=False, grid_alpha=0)\n",
    "    # legend_handles = [\n",
    "    #     Line2D([0], [0], marker='o', color='w', markersize=8, markerfacecolor=slice_cmap[slice_name_list[i]], label=slice_name_list[i])\n",
    "    #     for i in range(len(slice_name_list))\n",
    "    # ]\n",
    "    # plt.legend(handles=legend_handles, fontsize=8, title='Slices', title_fontsize=10,\n",
    "    #            loc='upper left')\n",
    "    plt.title(f'Slices ({model})', fontsize=16)\n",
    "    if save:\n",
    "        save_path = save_root + f\"/{model}_slices_umap.pdf\"\n",
    "        plt.savefig(save_path)\n",
    "\n",
    "    colors = list(adata_concat.obs[ground_truth_key].astype('str').map(cluster_to_color_map))\n",
    "    plt.figure(figsize=(5, 5))\n",
    "    if frame_color:\n",
    "        plt.rc('axes', edgecolor=frame_color, linewidth=2)\n",
    "    plt.scatter(sp_embedding[order, 0], sp_embedding[order, 1], s=size, c=colors)\n",
    "    plt.tick_params(axis='both', bottom=False, top=False, left=False, right=False,\n",
    "                    labelleft=False, labelbottom=False, grid_alpha=0)\n",
    "    plt.title(f'Annotated Spot-types ({model})', fontsize=16)\n",
    "    if save:\n",
    "        save_path = save_root + f\"/{model}_annotated_clusters_umap.pdf\"\n",
    "        plt.savefig(save_path)\n",
    "\n",
    "    colors = list(adata_concat.obs[matched_clusters_key].map(matched_to_color_map))\n",
    "    plt.figure(figsize=(5, 5))\n",
    "    if frame_color:\n",
    "        plt.rc('axes', edgecolor=frame_color, linewidth=2)\n",
    "    plt.scatter(sp_embedding[order, 0], sp_embedding[order, 1], s=size, c=colors)\n",
    "    plt.tick_params(axis='both', bottom=False, top=False, left=False, right=False,\n",
    "                    labelleft=False, labelbottom=False, grid_alpha=0)\n",
    "    plt.title(f'Identified Clusters ({model})', fontsize=16)\n",
    "    if save:\n",
    "        save_path = save_root + f\"/{model}_identified_clusters_umap.pdf\"\n",
    "        plt.savefig(save_path)\n",
    "\n",
    "    if plot:\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0737792a-8fd7-40e9-bf10-9008af367c09",
   "metadata": {},
   "outputs": [],
   "source": [
    "save = False\n",
    "plot = True\n",
    "plot_mousebrain(cas_list, result, 'Annotation_for_Combined', 'matched_clusters', 'INSTINCT', cluster_to_color_map,\n",
    "                matched_to_color_map, my_clusters, slice_name_list, cls_list, sp_embedding, save_root=save_dir,\n",
    "                frame_color='darkviolet', save=save, plot=plot)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}