test_UMAP_varient.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TreeVAE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Table Of Contents\n",
    "1. [Data Loading](#section_1)\n",
    "2. [Generations](#section_2)\n",
    "3. [Reconstructions](#section_3)\n",
    "4. [Tree and Representation Analysis](#section_4)\n",
    "5. [CelebA Attributes](#section_5)\n",
    "\n",
    "This is the notebook for analyzing and visualizing the trees learnt by TreeVAE. \n",
    "\n",
    "Trees can be learnt by running main.py and stored by setting the option save_model to True in the config file.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a name=\"section_1\"></a> 1. Data Loading"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Always execute this section first. This section loads the data and model and computes the NMI to ensure that the model was loaded correctly. Make sure to set the path in the second cell to the specific model that you want to analyze."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "#import torch \n",
    "import numpy as np\n",
    "from scipy.io import mmread\n",
    "import scanpy as sc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "matrix_dp = mmread('/home/junyi/code/treevae/data/Variant/bpdcn712_assemble_trimmed_aligned_mt_addtag_cellSNP0_WaldVariant_paperCell/passed_dp.mtx')\n",
    "matrix_ad = mmread('/home/junyi/code/treevae/data/Variant/bpdcn712_assemble_trimmed_aligned_mt_addtag_cellSNP0_WaldVariant_paperCell/passed_ad.mtx')\n",
    "# Read the cell label\n",
    "#matrix_dp.shape\n",
    "df_cell_label = pd.read_csv('/home/junyi/code/treevae/data/Variant/bpdcn712_assemble_trimmed_aligned_mt_addtag_cellSNP0_WaldVariant_paperCell/cell_label.csv')\n",
    "var_name = pd.read_csv('/home/junyi/code/treevae/data/Variant/bpdcn712_assemble_trimmed_aligned_mt_addtag_cellSNP0_WaldVariant_paperCell/passed_variant_names.txt',header=None)\n",
    "sample_name = pd.read_csv('/home/junyi/code/treevae/data/Variant/bpdcn712_assemble_trimmed_aligned_mt_addtag_cellSNP0_WaldVariant_paperCell/passed_sample_names.txt',header=None)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_39702/1113072881.py:4: RuntimeWarning: Mean of empty slice\n",
      "  AF_mean = np.nanmean(matrix_ad_dp, 0)\n"
     ]
    }
   ],
   "source": [
    "matrix_ad_dp = matrix_ad/matrix_dp\n",
    "matrix_ad_dp = matrix_ad_dp.T\n",
    "\n",
    "AF_mean = np.nanmean(matrix_ad_dp, 0)\n",
    "nan_indices = np.isnan(AF_mean)\n",
    "AF_mean[nan_indices] = 0\n",
    "\n",
    "\n",
    "matrix_ad_dp[np.isnan(matrix_ad_dp)] = np.outer(np.ones(matrix_ad_dp.shape[0]), AF_mean)[np.isnan(matrix_ad_dp)]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2237, 16569)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "matrix_ad_dp.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata = sc.AnnData(np.asarray(matrix_ad_dp))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "ar_name =var_name\n",
    "df_obs = df_cell_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_obs.set_index('barcode', inplace=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "ar_name.set_index(0, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1G&gt;A</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2A&gt;C</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3T&gt;A</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4C&gt;A</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5A&gt;C</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16565C&gt;A</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16566G&gt;A</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16567A&gt;C</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16568T&gt;A</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16569G&gt;A</th>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>16569 rows × 0 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "Empty DataFrame\n",
       "Columns: []\n",
       "Index: [1G>A, 2A>C, 3T>A, 4C>A, 5A>C, 6C>A, 7A>C, 8G>A, 9G>A, 10T>A, 11C>A, 12T>A, 13A>T, 14T>A, 15C>A, 16A>C, 17C>A, 18C>A, 19C>A, 20T>G, 21A>T, 22T>A, 23T>G, 24A>C, 25A>C, 26C>A, 27C>A, 28A>C, 29C>A, 30T>A, 31C>A, 32A>C, 33C>G, 34G>C, 35G>A, 36G>A, 37A>C, 38G>T, 39C>G, 40T>A, 41C>A, 42T>C, 43C>A, 44C>A, 45A>C, 46T>G, 47G>A, 48C>G, 49A>C, 50T>A, 51T>G, 52T>G, 53G>T, 54G>T, 55T>G, 56A>T, 57T>A, 58T>G, 59T>G, 60T>A, 61C>G, 62G>A, 63T>A, 64C>A, 65T>G, 66G>T, 67G>A, 68G>A, 69G>A, 70G>T, 71G>A, 72T>G, 73G>A, 74T>A, 75G>C, 76C>G, 77A>C, 78C>A, 79G>A, 80C>A, 81G>T, 82A>C, 83T>G, 84A>C, 85G>T, 86C>A, 87A>C, 88T>G, 89T>A, 90G>A, 91C>A, 92G>A, 93A>C, 94G>A, 95A>T, 96C>A, 97G>A, 98C>G, 99T>G, 100G>C, ...]\n",
       "\n",
       "[16569 rows x 0 columns]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ar_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.var.index = ar_name.index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.obs = df_obs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "View of AnnData object with n_obs × n_vars = 2237 × 2\n",
       "    obs: 'true_label'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata[:,[\"2593G>A\",\"683G>A\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array((adata.X>0).sum(axis=0)).ravel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#adata.var[\"n_cells\"] = np.array((adata.X>0).sum(axis=0)).ravel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#adata.var[\"VAF_0.5_counts\"] = np.array((adata.X>=0.5).sum(axis=0)).ravel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#gene_mask = adata.var[\"VAF_0.5_counts\"] > 10\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#adata.var[\"VAF_0.5_counts\"].max()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#filtered_adata = adata[:, gene_mask]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#np.array(((filtered_adata.X)>0.5).sum(axis=1)> filtered_adata.shape[1]*0.01).ravel() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#filtered_adata.obs['cells_0.01_v'] = np.array(((filtered_adata.X)>0.5).sum(axis=1)> filtered_adata.shape[1]*0.05).ravel() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#filtered_adata.obs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.layers['counts'] = adata.X.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#sc.pp.normalize_total(adata)\n",
    "# # # Logarithmize the data\n",
    "# sc.pp.log1p(adata)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.tl.pca(adata)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.pca_variance_ratio(adata, n_pcs=50, log=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.pca(adata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#sc.pp.highly_variable_genes(adata, n_top_genes=2000)\n",
    "sc.pp.neighbors(adata)\n",
    "sc.tl.umap(adata)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.umap(\n",
    "    adata,\n",
    "    color=\"true_label\",\n",
    "    # Setting a smaller point size to get prevent overlap\n",
    "    size=60,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.umap(\n",
    "    adata,\n",
    "    color=\"6185A>G\",\n",
    "    # Setting a smaller point size to get prevent overlap\n",
    "    size=60,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.obsm[\"X_treevae_n0\"] = node_embeddings[0][\"z_sample\"]\n",
    "adata.obsm[\"X_treevae_n1\"] = node_embeddings[1][\"z_sample\"]\n",
    "adata.obsm[\"X_treevae_n2\"] = node_embeddings[2][\"z_sample\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.obs[\"node_lv1\"] = (((node_embeddings[1][\"prob\"]<0.5)*1)+1).astype(str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pp.neighbors(adata,use_rep=\"X_treevae_n0\",n_neighbors=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.tl.umap(adata)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.settings.set_figure_params(dpi=100, facecolor=\"white\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.umap(\n",
    "    adata,\n",
    "    color=\"node_lv1\",\n",
    "    # Setting a smaller point size to get prevent overlap\n",
    "    size=500,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.umap(\n",
    "    adata,\n",
    "    color=\"true_label\",\n",
    "    # Setting a smaller point size to get prevent overlap\n",
    "    size=20,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.var"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.tl.rank_genes_groups(adata, 'node_lv1', method='wilcoxon',key_added='rank_genes_node_lv1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_degs=sc.get.rank_genes_groups_df(adata,'1',key=\"rank_genes_node_lv1\",pval_cutoff=0.05)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_degs.sort_values(\"logfoldchanges\",ascending=False).to_csv(\"results/bpdcn712_Wald_degs.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_degs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "df_degs.head() shows the degree distribution of the tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_dgs= df_degs.sort_values(\"logfoldchanges\",ascending=False)\n",
    "df_dgs.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_dgs.tail(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.heatmap(adata,  list(df_dgs.head(20).names.values) + list(df_dgs.tail(20).names.values), groupby='node_lv1',swap_axes=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute Train NMI\n",
    "prob_leaves = predict(gen_train_eval, model, device,'prob_leaves')\n",
    "y = np.squeeze(np.argmax(prob_leaves, axis=-1))\n",
    "print('Train NMI:',normalized_mutual_info_score(y, np.squeeze(y_train)))\n",
    "\n",
    "tot_counts = []\n",
    "print(\"                                  Leaf\", np.arange(10))\n",
    "for i in np.unique(y_test):\n",
    "    list_y_hat, counts = np.unique(y[np.squeeze(y_train)==i], return_counts=True)\n",
    "    for j in range(n_d):\n",
    "        if j not in list_y_hat:\n",
    "            list_y_hat = np.insert(list_y_hat, j, j)\n",
    "            counts = np.insert(counts, j, 0)\n",
    "    tot_counts.append(counts)\n",
    "    print(f\"Class {i:<10}\", list_y_hat, counts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Compute Test NMI\n",
    "prob_leaves = predict(gen_test, model, device,'prob_leaves')\n",
    "y = np.squeeze(np.argmax(prob_leaves, axis=-1))\n",
    "print('Test NMI:', normalized_mutual_info_score(y, np.squeeze(y_test)))\n",
    "\n",
    "tot_counts = []\n",
    "print(\"                                  Leaf\", np.arange(10))\n",
    "for i in np.unique(y_test):\n",
    "    list_y_hat, counts = np.unique(y[np.squeeze(y_test)==i], return_counts=True)\n",
    "    for j in range(n_d):\n",
    "        if j not in list_y_hat:\n",
    "            list_y_hat = np.insert(list_y_hat, j, j)\n",
    "            counts = np.insert(counts, j, 0)\n",
    "    tot_counts.append(counts)\n",
    "    print(f\"Class {i:<10}\", list_y_hat, counts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a name=\"section_2\"></a> 2. Generations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This section is concerned with unconditionally generating new samples as opposed to reconstructing existing data points."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Clusterwise generations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, given one unconditional random sampling from the root, we visualize the generations for each leaf. That is, each row corresponds to one sample and each column corresponds to one leaf. Above each generation, we provide the probability of falling into the respective leaf for this sample. \n",
    "\n",
    "This way of visualization can provide insights on the characteristics each leaf is associated with. Observe that the generations differ across the leaves, as each leaf decodes the sample in the style of the cluster that it learnt. It is likely that cluster-differences are observed more strongly than in the reconstructions' section, as here, we have no guiding information from the bottom-up."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "n_imgs = 15\n",
    "with torch.no_grad():\n",
    "    reconstructions, p_c_z = model.generate_images(n_imgs, device)\n",
    "reconstructions = move_to(reconstructions, 'cpu')\n",
    "for i in range(n_imgs):\n",
    "    fig, axs = plt.subplots(1, n_d, figsize=(15, 15))\n",
    "    for c in range(n_d):\n",
    "        axs[c].imshow(display_image(reconstructions[c][i]), cmap=plt.get_cmap('gray'))\n",
    "        axs[c].set_title(f\"L{c}: \" + f\"p=%.2f\" % torch.round(p_c_z[i][c],))\n",
    "        axs[c].axis('off')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate new images according to cluster assignment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this subsection, given a leaf, we store the first 100 generations, for which this leaf is their most likely cluster assignment. This allows us to gain insights on the cluster and characteristics that each leaf learnt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here, we store generations for each leaf simultaneously until\n",
    "# every leaf has n_imgs associated generations, or we iterated through max_iter batches.\n",
    "n_imgs = configs['training']['batch_size']\n",
    "max_iter = 200\n",
    "\n",
    "with torch.no_grad():\n",
    "    reconstructions, p_c_z = model.generate_images(n_imgs, device)\n",
    "reconstructions = move_to(reconstructions, 'cpu')\n",
    "clusterwise_reconst = [torch.zeros_like(reconstructions[0][0:2]) for i in range(len(reconstructions))]\n",
    "n_iter=0\n",
    "while min([clusterwise_reconst[leaf_ind].shape[0] for leaf_ind in range(len(reconstructions))]) < n_imgs+2 and n_iter < max_iter:\n",
    "    for i in range(n_imgs):\n",
    "        leaf_ind = torch.argmax(p_c_z[i])\n",
    "        if clusterwise_reconst[leaf_ind].shape[0] < n_imgs+2:\n",
    "            clusterwise_reconst[leaf_ind] = torch.vstack([clusterwise_reconst[leaf_ind], reconstructions[leaf_ind][i].unsqueeze(0)])\n",
    "    with torch.no_grad():\n",
    "        reconstructions, p_c_z = model.generate_images(n_imgs, device)\n",
    "    reconstructions = move_to(reconstructions, 'cpu')\n",
    "    n_iter += 1\n",
    "    if n_iter %10 == 0:\n",
    "        print(n_iter)\n",
    "for i in range(len(reconstructions)):\n",
    "    clusterwise_reconst[i] = clusterwise_reconst[i][2:,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(clusterwise_reconst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# For each leaf, we visualize n_grid x n_grid generations, \n",
    "# which have highest probability of being assigned to this cluster\n",
    "n_leaves = len(clusterwise_reconst)\n",
    "n_grid = min(5,int((clusterwise_reconst[leaf_ind].shape[0])**.5))\n",
    "\n",
    "k=0\n",
    "for l in range(n_leaves):\n",
    "        fig, axs = plt.subplots(n_grid, n_grid, figsize=(4,4))\n",
    "        i=0\n",
    "        for a in range(n_grid):\n",
    "            for b in range(n_grid):\n",
    "                try:\n",
    "                        axs[a,b].set_axis_off()\n",
    "                        axs[a,b].imshow(display_image(clusterwise_reconst[k][i]), cmap=plt.get_cmap('gray'))\n",
    "                        i+=1\n",
    "                except Exception as e:\n",
    "                    print(e)\n",
    "                    pass\n",
    "        fig.suptitle(f\"Leaf {k} samples\",fontsize=25)\n",
    "        fig.tight_layout()\n",
    "        fig.subplots_adjust(top=0.87)\n",
    "        k+=1\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# For closer inspection, one can select a specific leaf by leaf_ind and investigate more generations.\n",
    "leaf_ind = 0\n",
    "    \n",
    "n_grid = int((clusterwise_reconst[leaf_ind].shape[0])**.5)\n",
    "fig, axs = plt.subplots(n_grid, n_grid, figsize=(15,15))\n",
    "\n",
    "i=0\n",
    "for a in range(n_grid):\n",
    "    for b in range(n_grid):\n",
    "        axs[a,b].set_axis_off()\n",
    "        axs[a,b].imshow(display_image(clusterwise_reconst[leaf_ind][i]), cmap=plt.get_cmap('gray'))\n",
    "        i+=1\n",
    "fig.suptitle(f\"Leaf {leaf_ind} samples\",fontsize=25)\n",
    "fig.tight_layout()\n",
    "fig.subplots_adjust(top=0.95)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a name=\"section_3\"></a> 3. Reconstructions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This section is concerned with computing reconstructions of input samples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Clusterwise reconstructions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, given one input image, we visualize the reconstructions for each leaf. That is, each row corresponds to one input image and each column corresponds to one leaf. Above each reconstruction, we provide the probability of falling into the respective leaf for this sample. \n",
    "\n",
    "This way of visualization can provide insights on the characteristics each leaf is associated with. Observe that the reconstructions differ across the leaves, as each leaf reconstructs the image in the style of the cluster that it learnt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Training Set\n",
    "gen_train_eval_iter = iter(gen_train_eval)\n",
    "inputs, labels = next(gen_train_eval_iter)\n",
    "\n",
    "\n",
    "inputs_gpu, labels_gpu = inputs.to(device), labels.to(device)\n",
    "with torch.no_grad():\n",
    "    reconstructions_gpu, node_leaves_gpu = model.compute_reconstruction(inputs_gpu)\n",
    "reconstructions = move_to(reconstructions_gpu, 'cpu')\n",
    "node_leaves = move_to(node_leaves_gpu, 'cpu')\n",
    "\n",
    "\n",
    "for i in range(10):\n",
    "    print(\"Class:\", labels[i].item())\n",
    "    fig, axs = plt.subplots(1, n_d+1, figsize=(15, 15))\n",
    "    axs[n_d].imshow(display_image(inputs[i]), cmap=plt.get_cmap('gray'))\n",
    "    axs[n_d].set_title(\"Original\")\n",
    "    axs[n_d].axis('off')\n",
    "    for c in range(n_d):\n",
    "        axs[c].imshow(display_image(reconstructions[c][i]), cmap=plt.get_cmap('gray'))\n",
    "        axs[c].set_title(f\"L{c}: \" + f\"p=%.2f\" % torch.round(node_leaves[c]['prob'][i]))\n",
    "        axs[c].axis('off')\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Test Set\n",
    "gen_test_iter = iter(gen_test)\n",
    "inputs, labels = next(gen_test_iter)\n",
    "inputs_gpu, labels_gpu = inputs.to(device), labels.to(device)\n",
    "with torch.no_grad():\n",
    "    reconstructions_gpu, node_leaves_gpu = model.compute_reconstruction(inputs_gpu)\n",
    "reconstructions = move_to(reconstructions_gpu, 'cpu')\n",
    "node_leaves = move_to(node_leaves_gpu, 'cpu')\n",
    "\n",
    "\n",
    "for i in range(10):\n",
    "    print(\"Class:\", labels[i].item())\n",
    "    fig, axs = plt.subplots(1, n_d+1, figsize=(15, 15))\n",
    "    axs[n_d].imshow(display_image(inputs[i]), cmap=plt.get_cmap('gray'))\n",
    "    axs[n_d].set_title(\"Original\")\n",
    "    axs[n_d].axis('off')\n",
    "    for c in range(n_d):\n",
    "        axs[c].imshow(display_image(reconstructions[c][i]), cmap=plt.get_cmap('gray'))\n",
    "        axs[c].set_title(f\"L{c}: \" + f\"p=%.2f\" % torch.round(node_leaves[c]['prob'][i]))\n",
    "        axs[c].axis('off')\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Group reconstructions according to cluster assignment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this subsection, given a leaf, we store the reconstructions of the first 100 samples, for which this leaf is their most likely cluster assignment. This allows us to visualize for each leaf, which samples fall into it, in order to gain insights on the cluster that each leaf learnt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test Set\n",
    "# Here, we store samples for each leaf simultaneously by iterating through the training set until\n",
    "# every leaf has n_imgs associated samples, or we iterated through max_iter batches.\n",
    "max_iter = 100\n",
    "n_imgs = configs['training']['batch_size']\n",
    "\n",
    "n_iter=0\n",
    "gen_test_iter = iter(gen_test)\n",
    "inputs, labels = next(gen_test_iter)\n",
    "inputs_gpu, labels_gpu = inputs.to(device), labels.to(device)\n",
    "with torch.no_grad():\n",
    "    reconstructions_gpu, node_leaves_gpu = model.compute_reconstruction(inputs_gpu)\n",
    "reconstructions = move_to(reconstructions_gpu, 'cpu')\n",
    "node_leaves = move_to(node_leaves_gpu, 'cpu')\n",
    "p_c_z = torch.stack([node_leaves[i]['prob'] for i in range(len(node_leaves))],1)\n",
    "clusterwise_reconst = [torch.zeros_like(reconstructions[0][0:2]) for i in range(len(reconstructions))]\n",
    "while min([clusterwise_reconst[leaf_ind].shape[0] for leaf_ind in range(len(reconstructions))]) < n_imgs+2 and n_iter < max_iter:\n",
    "    n_iter += 1\n",
    "    if n_iter %10 == 0:\n",
    "        print(n_iter)\n",
    "    for i in range(n_imgs):\n",
    "        leaf_ind = p_c_z[i].numpy().argmax()\n",
    "        if clusterwise_reconst[leaf_ind].shape[0] < n_imgs+2:\n",
    "            clusterwise_reconst[leaf_ind] = torch.vstack([clusterwise_reconst[leaf_ind], reconstructions[leaf_ind][i].unsqueeze(0)])\n",
    "    inputs, labels = next(gen_test_iter)\n",
    "    inputs_gpu, labels_gpu = inputs.to(device), labels.to(device)\n",
    "    with torch.no_grad():\n",
    "        reconstructions_gpu, node_leaves_gpu = model.compute_reconstruction(inputs_gpu)\n",
    "    reconstructions = move_to(reconstructions_gpu, 'cpu')\n",
    "    node_leaves = move_to(node_leaves_gpu, 'cpu')\n",
    "    p_c_z = torch.stack([node_leaves[i]['prob'] for i in range(len(node_leaves))],1)\n",
    "for i in range(len(reconstructions)):\n",
    "    clusterwise_reconst[i] = clusterwise_reconst[i][2:,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# For each leaf, we visualize n_grid x n_grid reconstructions of samples, \n",
    "# which have highest probability of being assigned to this cluster\n",
    "n_leaves = len(clusterwise_reconst)\n",
    "n_grid = min(5,int((clusterwise_reconst[leaf_ind].shape[0])**.5))\n",
    "\n",
    "k=0\n",
    "for l in range(n_leaves):\n",
    "        fig, axs = plt.subplots(n_grid, n_grid, figsize=(4,4))\n",
    "        i=0\n",
    "        for a in range(n_grid):\n",
    "            for b in range(n_grid):\n",
    "                axs[a,b].set_axis_off()\n",
    "                axs[a,b].imshow(display_image(clusterwise_reconst[k][i]), cmap=plt.get_cmap('gray'))\n",
    "                i+=1\n",
    "        fig.suptitle(f\"Leaf {k} samples\",fontsize=25)\n",
    "        fig.tight_layout()\n",
    "        fig.subplots_adjust(top=0.87)\n",
    "        k+=1\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# For closer inspection, one can select a specific leaf by leaf_ind and investigate more reconstructions.\n",
    "leaf_ind = 0\n",
    "    \n",
    "n_grid = int((clusterwise_reconst[leaf_ind].shape[0])**.5)\n",
    "fig, axs = plt.subplots(n_grid, n_grid, figsize=(15,15))\n",
    "\n",
    "i=0\n",
    "for a in range(n_grid):\n",
    "    for b in range(n_grid):\n",
    "        axs[a,b].set_axis_off()\n",
    "        axs[a,b].imshow(display_image(clusterwise_reconst[leaf_ind][i]), cmap=plt.get_cmap('gray'))\n",
    "        i+=1\n",
    "fig.suptitle(f\"Leaf {leaf_ind} samples\",fontsize=25)\n",
    "fig.tight_layout()\n",
    "fig.subplots_adjust(top=0.95)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a name=\"section_4\"></a> 4. Tree and Representation Analysis "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section we explore the structure of the learnt tree as well as the representations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tree embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below, we visualize the learnt embeddings by performing PCA on each node. Set use_pca to False if you want to directly see the first two dimensions without dimensionality reduction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Do you want to look at pca embeddings or learnt representations\n",
    "use_pca = True\n",
    "\n",
    "# pick data loader\n",
    "data_loader = gen_test\n",
    "\n",
    "\n",
    "# each entry in node_embeddings is a dictionary with keys 'prob' and 'z_sample' for each leaf\n",
    "nb_nodes = len(data_tree)\n",
    "node_embeddings = [{'prob': [], 'z_sample': []} for _ in range(nb_nodes)]\n",
    "label_list = []\n",
    "\n",
    "# iterate over test data points\n",
    "for inputs, labels in tqdm(data_loader):\n",
    "    inputs_gpu, labels_gpu = inputs.to(device), labels.to(device)\n",
    "\n",
    "    label_list.append(labels)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        node_info = get_node_embeddings(model, inputs_gpu)\n",
    "    node_info = move_to(node_info, 'cpu')\n",
    "\n",
    "    # for each node, append the probability and z_sample to the list\n",
    "\n",
    "    k = 0 # need this variable to skip \"no digits\" nodes\n",
    "    for i in range(nb_nodes):\n",
    "        j = i - k\n",
    "        if data_tree[i][1] == 'no digits':\n",
    "            k += 1\n",
    "            continue\n",
    "\n",
    "        node_embeddings[i]['prob'].append(node_info[j]['prob'].numpy())\n",
    "        node_embeddings[i]['z_sample'].append(node_info[j]['z_sample'].numpy())\n",
    "\n",
    "# flatten the lists\n",
    "k = 0\n",
    "for i in range(nb_nodes):\n",
    "    if data_tree[i][1] == 'no digits':\n",
    "        node_embeddings[i]['prob'] = []\n",
    "        node_embeddings[i]['z_sample'] = []\n",
    "        continue\n",
    "    \n",
    "    node_embeddings[i]['prob'] = np.concatenate(node_embeddings[i]['prob'])\n",
    "    node_embeddings[i]['z_sample'] = np.concatenate(node_embeddings[i]['z_sample'])\n",
    "\n",
    "label_list = np.concatenate(label_list)\n",
    "\n",
    "# Draw the tree graph with scatter plots as nodes and arrows for edges\n",
    "draw_tree_with_scatter_plots(data_tree, node_embeddings, label_list, pca = use_pca)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Leaf embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below, we visualize the learnt leaf embeddings after performing PCA. This allows for a closer inspection of the leaf embeddings, which are also visualized in the tree above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get leaf embeddings for each test data point\n",
    "gen_test_iter = iter(gen_test)\n",
    "inputs, labels = next(gen_test_iter)\n",
    "inputs_gpu, labels_gpu = inputs.to(device), labels.to(device)\n",
    "with torch.no_grad():\n",
    "    reconstructions_gpu, node_leaves_gpu = model.compute_reconstruction(inputs_gpu)\n",
    "reconstructions = move_to(reconstructions_gpu, 'cpu')\n",
    "node_leaves = move_to(node_leaves_gpu, 'cpu')\n",
    "\n",
    "# each entry in node_leaves is a dictionary with keys 'prob' and 'z_sample' for each leaf\n",
    "node_leaves = [{'prob': [], 'z_sample': []} for _ in range(n_d)]\n",
    "label_list = []\n",
    "\n",
    "for inputs, labels in tqdm(gen_test):\n",
    "    inputs_gpu, labels_gpu = inputs.to(device), labels.to(device)\n",
    "\n",
    "    label_list.append(labels)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        _, node_leaves_gpu = model.compute_reconstruction(inputs_gpu)\n",
    "        node_leaves_cpu = move_to(node_leaves_gpu, 'cpu')\n",
    "        \n",
    "    # for each leaf, append the probability and z_sample to the list\n",
    "    for i in range(n_d):\n",
    "        node_leaves[i]['prob'].append(node_leaves_cpu[i]['prob'].numpy())\n",
    "        node_leaves[i]['z_sample'].append(node_leaves_cpu[i]['z_sample'].numpy())\n",
    "\n",
    "# flatten the lists\n",
    "for i in range(n_d):\n",
    "    node_leaves[i]['prob'] = np.concatenate(node_leaves[i]['prob'])\n",
    "    node_leaves[i]['z_sample'] = np.concatenate(node_leaves[i]['z_sample'])\n",
    "\n",
    "label_list = np.concatenate(label_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize z_sample for each leaf, do PCA and plot in 2D\n",
    "from sklearn.decomposition import PCA\n",
    "\n",
    "# PCA on node_leaves['z_sample']\n",
    "colors = label_list\n",
    "plt.figure(figsize=(20, 10))\n",
    "\n",
    "for i in range(n_d):\n",
    "    z_sample = node_leaves[i]['z_sample']\n",
    "    weights = node_leaves[i]['prob']\n",
    "\n",
    "    pca = PCA(n_components=2)\n",
    "    z_sample_pca = pca.fit_transform(z_sample)\n",
    "\n",
    "    plt.subplot(2, -(-len(node_leaves)//2), i+1)\n",
    "    plt.scatter(z_sample_pca[:, 0], z_sample_pca[:, 1], c=colors, cmap='tab10', alpha=weights)\n",
    "    plt.title(f\"Leaf {i}\")\n",
    "    plt.colorbar()\n",
    "    plt.xlabel(\"PC1\")\n",
    "    plt.ylabel(\"PC2\")\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a name=\"section_5\"></a> 5. CelebA attributes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This section is designated for analyzing the learnt splits and clusters of datasets without ground truth cluster labels, but various attributes. It is designed with a focus on CelebA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert configs['data']['data_name'] == 'celeba'\n",
    "import pandas as pd\n",
    "data_dir = './data/celeba/'\n",
    "attr = pd.read_csv(data_dir+'/list_attr_celeba.txt', sep=\"\\s+\", skiprows=1)\n",
    "y_test = attr[182637:]\n",
    "y_train = attr[:162770]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate cluster-matching attributes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Preprocessing step where we store for every node, the indeces of the test samples, whose most likely path went through said node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change to leafwise view of samples\n",
    "prob_leaves = predict(gen_test, model, device,'prob_leaves')\n",
    "y = np.squeeze(np.argmax(prob_leaves, axis=-1))\n",
    "sample_ind = []\n",
    "for i in range(len(np.unique(y))):\n",
    "    sample_ind.append([])\n",
    "for i in np.unique(y):\n",
    "    sample_ind[i] = np.where(y==i)[0]\n",
    "    \n",
    "# Fill all internal nodes and create datatree with corresponding samples\n",
    "data_tree_ids = []\n",
    "for i in range(len(data_tree)):\n",
    "    data_tree_ids.append([i,[]])\n",
    "for listnode in reversed(data_tree_ids):\n",
    "    i = listnode[0]\n",
    "    if data_tree[i][3] == 1:\n",
    "        # If leaf, just copy samples from above\n",
    "        data_tree_ids[i][1] = sample_ind[i-(len(data_tree_ids)-len(sample_ind))]\n",
    "    else:\n",
    "        # If internal node, take samples from children\n",
    "        children = []\n",
    "        for j in range(len(data_tree)):\n",
    "            if data_tree[j][2] == i:\n",
    "                children.append(j)\n",
    "        assert len(children)==2\n",
    "        data_tree_ids[i][1] = np.sort(np.concatenate((data_tree_ids[children[0]][1],data_tree_ids[children[1]][1])))\n",
    "        \n",
    "        \n",
    "        \n",
    "# Final ID-tree, where for each node, we store which test sample went through it\n",
    "data_tree_ids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each split, we additionally store the five attribute that correlate most highly with the the split. This gives an intuition on what attributes the split is based on, i.e. which characteristics the split differentiates between.\n",
    "\n",
    "Note that for CelebA, the \"ground truth\" attributes are in our opinion not the most descriptive ones regarding overall image&cluster impression and focus sometimes on details, on which we don't pick up on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Highest correlated features per split\n",
    "data_tree_new = data_tree.copy()\n",
    "for i in range(len(data_tree_ids)):\n",
    "    in_leaf = False\n",
    "    node_ind = data_tree_ids[i][1]\n",
    "    # Samples in node before split\n",
    "    node_samples = y_test.iloc[node_ind]\n",
    "    # Split of samples\n",
    "    node_split = np.zeros(len(y_test))\n",
    "    children = []\n",
    "    for j in range(len(data_tree)):\n",
    "        if data_tree[j][2] == i:\n",
    "            children.append(j)\n",
    "        if children == []:\n",
    "            in_leaf = True\n",
    "        else:\n",
    "            in_leaf = False\n",
    "    if not in_leaf: \n",
    "        child_left = children[0]\n",
    "        node_split[data_tree_ids[child_left][1]] = 1\n",
    "        node_split = node_split[node_ind]\n",
    "        # Store corr coefficients\n",
    "        corr = np.corrcoef(np.concatenate((np.array(node_samples),np.expand_dims(node_split,1)),1).T)[len(y_test.columns),0:len(y_test.columns)]\n",
    "        data_tree_ids[i].append(corr)\n",
    "        \n",
    "        # Store 5 strongest correlations\n",
    "        ind = np.abs(corr).argsort()[-5:][::-1]\n",
    "        features = y_test.columns[ind].tolist()\n",
    "        for k in range(len(ind)):\n",
    "            if corr[ind[k]] < 0:\n",
    "                features[k] = 'not ' + features[k]\n",
    "            features[k] = features[k] + ' ({})'.format(round(corr[ind[k]], 2))\n",
    "        data_tree_ids[i].append(features)\n",
    "        \n",
    "data_tree_ids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a summary, for each attribute, we print the split that has the highest correlation with it. This gives an intuition on what internal node differentiates the most according to a given attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Attributewise node with highest correlation (i.e. internal node that was splitting attribute the most)\n",
    "attr_maxnode = y_test.columns.tolist()\n",
    "for i in range(len(y_test.columns)):\n",
    "    attrcorr = []\n",
    "    for node in range(len(data_tree_ids)):\n",
    "        if len(data_tree_ids[node])==len(data_tree_ids[0]):\n",
    "            attrcorr.append(data_tree_ids[node][2][i])\n",
    "    attrcorr = np.array(attrcorr)\n",
    "    if len(np.argwhere(np.isnan(attrcorr)).squeeze(1))>0:\n",
    "        attrcorr[np.argwhere(np.isnan(attrcorr)).squeeze(1)] = 0\n",
    "    ind = np.argmax(np.abs(attrcorr))\n",
    "    attr_maxnode[i] = attr_maxnode[i] + \": \" + f'{ind}' ' ({})'.format(round(attrcorr[ind], 2))\n",
    "attr_maxnode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation of clusters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can analyze the clustering quality according to certain attributes. To do this, in the second cell, pick the indeces of the attributes, whose intersections you want to determine as ground truth clusterings. Then, the NMI is calculated for treating the selected attributes' intersections as ground-truth clusters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_test.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pick labels here\n",
    "label_ind = [2,20,39]\n",
    "print([attr_maxnode[i] for i in label_ind])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if len(label_ind)==2:\n",
    "    label_dict = {\n",
    "        (-1, -1): 0,\n",
    "        (-1, 1): 1,\n",
    "        (1, -1): 2,\n",
    "        (1, 1): 3\n",
    "    }\n",
    "else:\n",
    "    label_dict = {\n",
    "        (-1, -1, -1): 0,\n",
    "        (-1, -1, 1): 1,\n",
    "        (-1, 1, -1): 2,\n",
    "        (-1, 1, 1): 3,\n",
    "        (1, -1, -1): 4,\n",
    "        (1, -1, 1): 5,\n",
    "        (1, 1, -1): 6,\n",
    "        (1, 1, 1): 7\n",
    "    }\n",
    "selected_classes = np.array(y_test.iloc[:, label_ind])\n",
    "selected_classes = [tuple([x for x in a]) for a in selected_classes]\n",
    "label_true = [label_dict[sample_labels] for sample_labels in selected_classes]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('NMI:')\n",
    "normalized_mutual_info_score(y, label_true)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create attribute-wise percentage table for leaves"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This subsection presents the frequency of the attributes for each leaf. The numbers indicate the percentage of samples assigned to a given leaf, that contain a certain attribute. For example: 67% of all people assigned to leaf 3 are blonde."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "leaf_attr = []\n",
    "n_leaves = len(np.unique(y))\n",
    "for i in range(1,1+n_leaves):\n",
    "    data_tree_ids[-i].append((y_test.iloc[data_tree_ids[-i][1]] == 1).mean())\n",
    "    leaf_attr.append((y_test.iloc[data_tree_ids[-i][1]] == 1).mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "leaf_attr_table = pd.DataFrame(np.stack(leaf_attr)[::-1])\n",
    "leaf_attr_table.columns = y_test.columns\n",
    "\n",
    "pd.set_option('display.max_columns', None)\n",
    "leaf_attr_table.round(3)*100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, one can create new attributes by combining previous attributes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_vars =[]\n",
    "for i in range(1,1+n_leaves):\n",
    "    temp = y_test.iloc[data_tree_ids[-i][1]] == 1\n",
    "    temp['Hair_Loss'] = np.clip(temp['Bald'] + temp['Receding_Hairline'],0,1)\n",
    "    temp['Dark_Hair'] = np.clip(temp['Brown_Hair'] + temp['Black_Hair'],0,1)\n",
    "    temp['Happy'] = np.clip(temp['Smiling'] + temp['Mouth_Slightly_Open'],0,1)\n",
    "    temp['Light_Hair'] = np.clip(temp['Blond_Hair'] + temp['Gray_Hair'],0,1)\n",
    "    temp['Beard'] = np.clip(temp['5_o_Clock_Shadow'] + 1-temp['No_Beard'],0,1)\n",
    "\n",
    "    new_vars.append([temp['Hair_Loss'].mean(),temp['Dark_Hair'].mean(),temp['Happy'].mean(),temp['Light_Hair'].mean(),temp['Beard'].mean()])\n",
    "    \n",
    "new_vars_table = pd.DataFrame(np.stack(new_vars)[::-1])\n",
    "new_vars_table.columns = ['Hair_Loss','Dark_Hair','Happy','Light_Hair','Beard']\n",
    "new_vars_table.round(3)*100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "RNA-FM",
   "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}