Full_notebook_with_k_estimation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dbe935ed-2854-44d5-9b60-262def371e01",
   "metadata": {},
   "source": [
    "# Clustering performance when $k$ must be estimated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d113ba7-db27-403f-bab8-f832bccd3654",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score\n",
    "from sklearn.cluster import SpectralClustering\n",
    "from sklearn.metrics import silhouette_score\n",
    "from sklearn.cluster import KMeans\n",
    "from tqdm import tqdm\n",
    "import json\n",
    "\n",
    "from models.baseline_model import ArticleBaselineModel\n",
    "from models.tabularncd_model import TabularNCDModel\n",
    "from models.NCD_Spectral_Clustering import *\n",
    "from models.NCD_Kmeans import k_means_pp\n",
    "from models.PBN_model import PBNModel\n",
    "from src.dataset_utils import *\n",
    "from src.utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "839e05a9-1c4e-427a-b9c0-1da7103d5ebb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from plotly.subplots import make_subplots\n",
    "import plotly.graph_objects as go\n",
    "import plotly.io as pio\n",
    "\n",
    "pio.renderers.default = 'iframe'  # So we can view the plots in the jupyter notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bfdd2d2-c85e-4596-a8e2-7febfa7dfc7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e68d990-0057-47b2-b706-334d6a794162",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = setup_device(use_cuda=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a43d8da-0bba-493e-ac1f-916d3853e2f6",
   "metadata": {},
   "source": [
    "## Table of contents:\n",
    "* [1. Define some functions](#1)\n",
    "* [2. $k$ estimation in the original feature space](#2)\n",
    "* [2. Clustering models](#2)\n",
    "    * [2.1 $k$-means](#2.1)\n",
    "    * [2.2 NCD $k$-means](#2.2)\n",
    "    * [2.3 Spectral Clustering](#2.3)\n",
    "    * [2.4 NCD Spectral Clustering](#2.4)\n",
    "* [3. NCD models](#3)\n",
    "    * [3.1 Baseline](#3.1)\n",
    "    * [3.2 PBN](#3.2)\n",
    "    * [3.3 TabularNCD](#32.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "417650bb-d9cd-4e39-9c1d-af7ed3868fd7",
   "metadata": {},
   "source": [
    "### Import the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce1c28fb-a1c8-4817-8a05-bea61fe0170c",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_names = ['HumanActivityRecognition', 'LetterRecognition', 'Pendigits', 'USCensus1990', 'multiple_feature', 'optdigits', 'cnae_9']\n",
    "\n",
    "dataset_name = 'HumanActivityRecognition'\n",
    "\n",
    "x_train, y_train, x_test, y_test, unknown_class_value, y_train_save, y_test_save = import_dataset_with_name(dataset_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d8eb61d-2a57-4e39-943c-58bf4af7f00d",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train = torch.tensor(x_train, dtype=torch.float, device=device)\n",
    "x_test = torch.tensor(x_test, dtype=torch.float, device=device)\n",
    "\n",
    "# For plots\n",
    "y_train_unknown_save = y_train_save[y_train == unknown_class_value]\n",
    "y_test_unknown_save = y_test_save[y_test == unknown_class_value]\n",
    "\n",
    "# For evaluation\n",
    "y_train_unknown_save_codes = np.array(pd.Series(y_train_unknown_save).astype('category').cat.codes)\n",
    "y_test_unknown_save_codes = np.array(pd.Series(y_test_unknown_save).astype('category').cat.codes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e46bb877-836d-4c11-bdeb-b6ef369f49b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_classes_distribution(y_train_save[y_train!=unknown_class_value], y_train_save[y_train==unknown_class_value], y_test_save, dataset_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fc5f5d1-e843-4929-940d-e3881c532e40",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train_known = x_train[y_train != unknown_class_value]\n",
    "x_train_unknown = x_train[y_train == unknown_class_value]\n",
    "y_train_known = y_train[y_train != unknown_class_value]\n",
    "\n",
    "x_test_known = x_test[y_test != unknown_class_value]\n",
    "x_test_unknown = x_test[y_test == unknown_class_value]\n",
    "y_test_known = y_test[y_test != unknown_class_value]\n",
    "\n",
    "# We need the targets to be in {0, ..., C^l} exactly\n",
    "classifier_mapper, classifier_ind = np.unique(y_train_known, return_inverse=True)\n",
    "classifier_mapping_dict = dict(zip(y_train_known, classifier_ind))\n",
    "\n",
    "y_train_known = np.array(list(map(classifier_mapping_dict.get, y_train_known)))\n",
    "y_test_known = np.array(list(map(classifier_mapping_dict.get, y_test_known)))\n",
    "\n",
    "gt_k = len(np.unique(y_test_unknown_save_codes))  # Only used to print errors"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0dde57a0-4835-41b8-956b-ab6e6546c98d",
   "metadata": {},
   "source": [
    "# 1) Define some functions <a class=\"anchor\" id=\"1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fffba994-1dd4-42ff-941a-781824ed137d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_k_means_predictions(x, k_range, avg=5):\n",
    "    \"\"\"\n",
    "    return: scores : shape=(len(range), avg)\n",
    "    \"\"\"\n",
    "    predictions = []\n",
    "    for k_clust in tqdm(k_range):\n",
    "        tmp_predictions = []\n",
    "        for _ in np.arange(avg):\n",
    "            km = KMeans(n_clusters=k_clust, init='k-means++', n_init=10)\n",
    "            y_pred = km.fit_predict(x.cpu().numpy())\n",
    "            tmp_predictions.append(y_pred)\n",
    "        predictions.append(tmp_predictions)\n",
    "    return np.array(predictions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da12b5c1-0f59-4dd1-8057-e8b853d15f24",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_silhouette_curve(x, preds):\n",
    "    scores = []\n",
    "    for i in range(preds.shape[0]):\n",
    "        tmp_scores = []\n",
    "        for j in range(preds.shape[1]):\n",
    "            tmp_scores.append(silhouette_score(x, preds[i, j]))\n",
    "        scores.append(tmp_scores)\n",
    "    return np.array(scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edcf119f-7dac-4904-9e37-2e4594121a4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def estimate_k_with_silhouette(k_min, k_max, x):\n",
    "    k_range = np.arange(k_min, k_max + 1, 1)\n",
    "\n",
    "    predictions = get_k_means_predictions(x, k_range=k_range, avg=3)\n",
    "\n",
    "    silhouette_scores = get_silhouette_curve(x.cpu().numpy(), predictions)\n",
    "    \n",
    "    sil_k_estimation = k_range[np.argmax(silhouette_scores.mean(axis=1))]\n",
    "\n",
    "    return sil_k_estimation, silhouette_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cf188df-3eef-40e6-a82e-51961c6e79b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_k_means_perf_curve(metric, range, n_classes, metric_name=\"Silhouette\"):\n",
    "    fig = make_subplots(rows=1, cols=1, horizontal_spacing=0.05, vertical_spacing=0.1)\n",
    "    \n",
    "    fig.add_trace(go.Scatter(x=range, y=metric.mean(axis=1),\n",
    "                             error_y=dict(type='data', array=metric.std(axis=1), visible=True),\n",
    "                             mode=\"lines+markers\", name=metric_name), row=1, col=1, secondary_y=False)\n",
    "    \n",
    "    best_perf = range[np.argmax(metric.mean(axis=1))]\n",
    "    pos = \"top right\" if best_perf < n_classes else \"top left\"\n",
    "    fig.add_vline(x=n_classes, annotation_text=f\"Ground truth: {n_classes}\", annotation_position=pos, line_color='red', annotation_font_color='red', secondary_y=False)\n",
    "    pos = \"top left\" if best_perf < n_classes else \"top right\"\n",
    "    fig.add_vline(x=best_perf, annotation_text=f\"Estimation: {best_perf}\", annotation_position=pos, line_color='green', annotation_font_color='green', secondary_y=False)\n",
    "    \n",
    "    fig.update_xaxes(title=\"k clusters\")\n",
    "    fig.update_yaxes(title=metric_name)\n",
    "    fig.update_layout(height=400, width=700, showlegend=False)  # , title=dataset_name\n",
    "    fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cee2829-35f8-470c-8bea-669c495e851f",
   "metadata": {},
   "source": [
    "# 2) $k$ estimation in the original feature space <a class=\"anchor\" id=\"2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8862c402-3884-40fb-8a81-5e21ae78b68b",
   "metadata": {},
   "outputs": [],
   "source": [
    "k_est, silhouette_scores = estimate_k_with_silhouette(k_min=2, k_max=30, x=x_train_unknown)\n",
    "\n",
    "plot_k_means_perf_curve(silhouette_scores, np.arange(2, 30, 1), gt_k)\n",
    "\n",
    "print(f\"Estimated k={k_est}, ground truth={gt_k} (difference={np.abs(gt_k - k_est)})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4530f162-292d-4f1f-a20a-7e2c75e66ee3",
   "metadata": {},
   "source": [
    "# 3) Clustering models <a class=\"anchor\" id=\"3\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6205dea-5827-42da-b520-ec7a9e51aa10",
   "metadata": {},
   "source": [
    "## 3.1 $k$-means <a class=\"anchor\" id=\"3.1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccca8964-ad56-4b46-a395-febdf832170c",
   "metadata": {},
   "outputs": [],
   "source": [
    "accs_1, nmis_1, aris_1 = [], [], []\n",
    "\n",
    "for i in tqdm(range(10)):\n",
    "    km = KMeans(n_clusters=k_est, init='random', n_init=10).fit(x_train_unknown.cpu().numpy())\n",
    "    \n",
    "    y_test_unknown_pred = km.predict(x_test_unknown.cpu().numpy())\n",
    "    accs_1.append(hungarian_accuracy(y_test_unknown_save_codes, y_test_unknown_pred))\n",
    "    nmis_1.append(normalized_mutual_info_score(y_test_unknown_save_codes, y_test_unknown_pred))\n",
    "    aris_1.append(adjusted_rand_score(y_test_unknown_save_codes, y_test_unknown_pred))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e13043e8-c37a-4d09-ba4d-8d64eb67ff8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"TEST: ACC={:.1f}±{:.1f} | NMI={:.1f}±{:.1f} | ARI={:.1f}±{:.1f}\".format(np.mean(accs_1)*100, np.std(accs_1)*100, np.mean(nmis_1)*100, np.std(nmis_1)*100, np.mean(aris_1)*100, np.std(aris_1)*100))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "249a45de-50bf-41f8-8517-c5d33d3fb9b8",
   "metadata": {},
   "source": [
    "## 3.2 NCD $k$-means <a class=\"anchor\" id=\"3.2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80204bf3-e6a0-4956-998e-87087fa9d949",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For this method, we define the centroids of the known classes with ground truth as initial centroids\n",
    "known_classes_centroids = torch.stack([x_train_known[y_train_known==c].mean(axis=0) for c in np.unique(y_train_known)])\n",
    "\n",
    "centroid_to_class_dict = dict(enumerate(np.unique(y_train[y_train != unknown_class_value])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc9319c9-b38e-4d45-a910-a8936c3f7b2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "accs_2, nmis_2, aris_2 = [], [], []\n",
    "\n",
    "for i in tqdm(range(10)):\n",
    "    kmpp = k_means_pp(pre_centroids=known_classes_centroids, k_new_centroids=k_est)\n",
    "    kmpp.fit(x_train_unknown, tolerance=1e-4, n_iterations=300)\n",
    "    \n",
    "    y_test_unknown_pred = kmpp.predict_unknown_data(x_test_unknown).cpu().numpy()\n",
    "    accs_2.append(hungarian_accuracy(y_test_unknown_save_codes, y_test_unknown_pred))\n",
    "    nmis_2.append(normalized_mutual_info_score(y_test_unknown_save_codes, y_test_unknown_pred))\n",
    "    aris_2.append(adjusted_rand_score(y_test_unknown_save_codes, y_test_unknown_pred))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f0986e2-f79c-4846-ac3f-1a8f8d681095",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"TEST: ACC={:.1f}±{:.1f} | NMI={:.1f}±{:.1f} | ARI={:.1f}±{:.1f}\".format(np.mean(accs_2)*100, np.std(accs_2)*100, np.mean(nmis_2)*100, np.std(nmis_2)*100, np.mean(aris_2)*100, np.std(aris_2)*100))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abbeef85-2941-4d46-865b-f08ca24d380c",
   "metadata": {},
   "source": [
    "## 3.3 Spectral Clustering <a class=\"anchor\" id=\"3.3\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d631280f-7ee3-45e5-948c-6d37d7b5ac26",
   "metadata": {},
   "source": [
    "#### /!\\ We feed our custom adjacency matrix to sklearn's Spectral Clustering /!\\\n",
    "Instead of a sigma, we use a custom parameter, min_dist.\n",
    "See the paper for more details"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13e0bc3c-c5d1-46e8-92b6-fe8ff739d7d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma = estimate_sigma(batch_outer_difference_normed(x_train_unknown), 0.001)\n",
    "print(\"Using min_dist = 0.001, estimated a sigma =\", sigma)\n",
    "\n",
    "test_outer_diff = batch_outer_difference_normed(x_test_unknown)\n",
    "test_A = get_adjacency_matrix(test_outer_diff, sigma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e1a3639-0b66-48a1-a6f7-3ec46655ae37",
   "metadata": {},
   "outputs": [],
   "source": [
    "accs_3, nmis_3, aris_3 = [], [], []\n",
    "for i in tqdm(range(10)):\n",
    "    sc = SpectralClustering(n_clusters=k_est, affinity='precomputed', n_jobs=-1).fit(test_A.cpu().numpy())\n",
    "    test_pred = sc.labels_\n",
    "    \n",
    "    accs_3.append(hungarian_accuracy(test_pred, y_test_unknown_save_codes))\n",
    "    nmis_3.append(normalized_mutual_info_score(test_pred, y_test_unknown_save_codes))\n",
    "    aris_3.append(adjusted_rand_score(test_pred, y_test_unknown_save_codes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e635c920-13b0-4ec8-87da-f96c84f00e90",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"TEST: ACC={:.1f}±{:.1f} | NMI={:.1f}±{:.1f} | ARI={:.1f}±{:.1f}\".format(np.mean(accs_3)*100, np.std(accs_3)*100, np.mean(nmis_3)*100, np.std(nmis_3)*100, np.mean(aris_3)*100, np.std(aris_3)*100))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e052a22f-3b38-463c-a3c3-ab4aaf6303ba",
   "metadata": {},
   "source": [
    "## 3.4 NCD Spectral Clustering <a class=\"anchor\" id=\"3.4\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3fd76c8-fc99-467e-8248-4eb3a880b9e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We load the hyperparameters that were optimized through grid search\n",
    "d = json.load(open(\"hyperparameters.json\"))\n",
    "ncdsc_params = d[\"NCD SC\"][\"Estimated k\"][dataset_name]\n",
    "print(\"Using hyperparameters:\", ncdsc_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2decd4d-4b17-45a0-a5eb-b808a41cabd5",
   "metadata": {},
   "source": [
    "#### /!\\ In this case, we estimate the number of clusters in the spectral embedding before the $k$-means /!\\"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6da2045e-7266-45bf-99dc-741038e60e21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# (1) Get the spectral embedding for all the data (since the hyperparameters are adapted for all the data, not only the novel data)\n",
    "x_test_full = torch.concat([x_test_known, x_test_unknown], axis=0)\n",
    "full_spectral_embedding = get_spectral_embedding(x_test_full, ncdsc_params['n_components'], ncdsc_params['min_dist'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b3dcfb8-f69f-4f2a-a5c9-13d5a44e484d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# (2) Estimate the number of clusters from the UNKNOWN DATA ONLY\n",
    "unknown_spectral_embedding = full_spectral_embedding[len(x_test_known):]\n",
    "ncd_sc_k_est, ncd_sc_silhouette_scores = estimate_k_with_silhouette(k_min=2, k_max=30, x=unknown_spectral_embedding)\n",
    "\n",
    "plot_k_means_perf_curve(ncd_sc_silhouette_scores, np.arange(2, 30, 1), gt_k)\n",
    "\n",
    "print(f\"Estimated k={ncd_sc_k_est}, ground truth={gt_k} (difference={np.abs(gt_k - ncd_sc_k_est)})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b3831ae-a0ac-4767-8951-002a6485294c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# (3) With the estimations, run k-means on spectral embedding of novel data only\n",
    "accs_4, nmis_4, aris_4 = [], [], []\n",
    "for i in tqdm(range(5)):\n",
    "    kmpp = k_means_pp(pre_centroids=None, k_new_centroids=len(np.unique(y_train_known)) + ncd_sc_k_est)\n",
    "    kmpp.fit(full_spectral_embedding, tolerance=1e-10, n_iterations=1000, n_init=10)\n",
    "    sil_y_test_full_pred = kmpp.predict_unknown_data(full_spectral_embedding).cpu().numpy()\n",
    "    sil_y_test_unknown_pred = sil_y_test_full_pred[len(x_test_known):]\n",
    "    \n",
    "    accs_4.append(hungarian_accuracy(y_test_unknown_save_codes, sil_y_test_unknown_pred))\n",
    "    nmis_4.append(normalized_mutual_info_score(y_test_unknown_save_codes, sil_y_test_unknown_pred))\n",
    "    aris_4.append(adjusted_rand_score(y_test_unknown_save_codes, sil_y_test_unknown_pred))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2cb2f7b-6593-4432-8af5-265eb081021e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"TEST: ACC={:.1f}±{:.1f} | NMI={:.1f}±{:.1f} | ARI={:.1f}±{:.1f}\".format(np.mean(accs_4)*100, np.std(accs_4)*100, np.mean(nmis_4)*100, np.std(nmis_4)*100, np.mean(aris_4)*100, np.std(aris_4)*100))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "848933f6-710f-456d-8274-548cdbf62786",
   "metadata": {},
   "source": [
    "# 4) NCD models <a class=\"anchor\" id=\"4\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ceb00946-e2f6-4650-9f0a-ab599db49ee4",
   "metadata": {},
   "source": [
    "## 4.1 Baseline <a class=\"anchor\" id=\"4.1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc218415-13f2-4b83-91f8-37c1cf7ccd6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_config = {\n",
    "    'input_size': x_train.shape[1],\n",
    "    'hidden_layers_dims': [math.floor(3*x_train.shape[1]/4), math.floor(2*x_train.shape[1]/4)],\n",
    "    'activation_fct': 'relu',  # relu or sigmoid or tanh or None\n",
    "    'use_batchnorm': True,  # True or False\n",
    "    'use_norm': 'l2',  # None or 'l1' or 'l2'\n",
    "    \n",
    "    'n_classes': len(np.unique(y_train_known)),\n",
    "    # 'n_clusters': ,  /!\\ we will estimate this one in this notebook\n",
    "    \n",
    "    'clustering_model': 'kmeans',  # kmeans or ncd_kmeans or spectral_clustering or ncd_spectral_clustering\n",
    "    'clustering_runs': 3,  # To compute the average accuracy of the clustering\n",
    "    \n",
    "    'batch_size': 512,\n",
    "    'epochs': 200,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "477d39fb-12df-4ddb-87c0-aedcb8251131",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We load the hyperparameters that were optimized through grid search\n",
    "d = json.load(open(\"hyperparameters.json\"))\n",
    "config = d[\"Baseline\"][\"Estimated k\"][dataset_name]\n",
    "print(\"Using hyperparameters:\", config)\n",
    "\n",
    "b_config = base_config.copy()\n",
    "b_config.update(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a60a835-c695-45eb-a18a-097212cee0f7",
   "metadata": {},
   "source": [
    "#### /!\\ In this case, we estimate the number of clusters in the latent space before the $k$-means /!\\"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23eb95ba-a3fb-4e32-8d08-75c3c95c2649",
   "metadata": {},
   "outputs": [],
   "source": [
    "accs_5, nmis_5, aris_5 = [], [], []\n",
    "\n",
    "for i in range(10):\n",
    "    model = ArticleBaselineModel(b_config).to(device)\n",
    "    \n",
    "    # (1) Train the latent space\n",
    "    losses_dict = model.train_on_known_classes(x_train_known=x_train_known, y_train_known=y_train_known,\n",
    "                                               x_test_unknown=None, y_test_unknown=None, x_test_known=None, y_test_known=None,\n",
    "                                               batch_size=b_config['batch_size'], lr=b_config['lr'], epochs=b_config['epochs'],\n",
    "                                               n_clusters=None, clustering_runs=None, evaluate=False, disable_tqdm=True)\n",
    "    model.eval()\n",
    "\n",
    "    # (2) Estimate the number of clusters from the UNKNOWN DATA ONLY in the latent space\n",
    "    with torch.no_grad():\n",
    "        z_unknown = model.encoder_forward(x_train[y_train == unknown_class_value])\n",
    "    \n",
    "    baseline_k_est, baseline_silhouette_scores = estimate_k_with_silhouette(k_min=2, k_max=30, x=z_unknown)\n",
    "    \n",
    "    print(f\"Estimated k={baseline_k_est}, ground truth={gt_k} (difference={np.abs(gt_k - baseline_k_est)})\")\n",
    "\n",
    "    # (3) In the latent space, run k-means with the estimated number of clusters\n",
    "    accs, nmis, aris = [], [], []\n",
    "    for i in range(5):\n",
    "        y_test_unknown_pred = model.predict_new_data(baseline_k_est, x_test_unknown, x_test_known, y_test_known)\n",
    "        accs.append(hungarian_accuracy(y_test_unknown_save_codes, y_test_unknown_pred))\n",
    "        nmis.append(normalized_mutual_info_score(y_test_unknown_save_codes, y_test_unknown_pred))\n",
    "        aris.append(adjusted_rand_score(y_test_unknown_save_codes, y_test_unknown_pred))\n",
    "    \n",
    "    accs_5.append(np.mean(accs))\n",
    "    nmis_5.append(np.mean(nmis))\n",
    "    aris_5.append(np.mean(aris))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "341acd53-4e46-47a7-91f1-c954aa1f390c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"TEST: ACC={:.1f}±{:.1f} | NMI={:.1f}±{:.1f} | ARI={:.1f}±{:.1f}\".format(np.mean(accs_5)*100, np.std(accs_5)*100, np.mean(nmis_5)*100, np.std(nmis_5)*100, np.mean(aris_5)*100, np.std(aris_5)*100))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da0a3a87-9717-4fc1-95d4-92c7301419fe",
   "metadata": {},
   "source": [
    "## 4.2 PBN <a class=\"anchor\" id=\"4.2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f3df44e-1e79-4f52-ae5c-04e4baf182c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_config = {\n",
    "    'input_size': x_train.shape[1],\n",
    "    'hidden_layers_dims': [math.floor(3*x_train.shape[1]/4), math.floor(2*x_train.shape[1]/4)],\n",
    "    'activation_fct': 'relu',  # relu or sigmoid or tanh or None\n",
    "    'use_batchnorm': True,  # True or False\n",
    "    'use_norm': 'l2',  # None or 'l1' or 'l2'\n",
    "    \n",
    "    'n_classes': len(np.unique(y_train_known)),\n",
    "    # 'n_clusters': ?,  /!\\ we will estimate this one in this notebook\n",
    "    \n",
    "    'clustering_model': 'kmeans',  # kmeans or ncd_kmeans or spectral_clustering or ncd_spectral_clustering\n",
    "    'clustering_runs': 1,  # To compute the average accuracy of the clustering\n",
    "    \n",
    "    'batch_size': 512,\n",
    "    'epochs': 200,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "355b4099-b07b-4f20-9b78-3389288f7644",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We load the hyperparameters that were optimized through grid search\n",
    "d = json.load(open(\"hyperparameters.json\"))\n",
    "config = d[\"PBN\"][\"Estimated k\"][dataset_name]\n",
    "print(\"Using hyperparameters:\", config)\n",
    "\n",
    "pbn_config = base_config.copy()\n",
    "pbn_config.update(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45840822-b5ea-43c7-8814-9c01ca70feda",
   "metadata": {},
   "outputs": [],
   "source": [
    "accs_6, nmis_6, aris_6 = [], [], []\n",
    "\n",
    "for i in range(10):\n",
    "    model = PBNModel(pbn_config).to(device)\n",
    "\n",
    "    # (1) Train the latent space\n",
    "    losses_dict = model.train_on_known_classes(x_train=x_train, y_train=y_train, unknown_class_value=unknown_class_value, x_test_unknown=x_test_unknown, y_test_unknown=y_test_unknown_save_codes, x_test_known=x_test_known, y_test_known=y_test_known,\n",
    "                                               batch_size=pbn_config['batch_size'], lr=pbn_config['lr'], epochs=pbn_config['epochs'], clustering_runs=pbn_config['clustering_runs'], w=pbn_config['w'],\n",
    "                                               evaluate=False, disable_tqdm=True)\n",
    "    model.eval()\n",
    "\n",
    "    # (2) Estimate the number of clusters from the UNKNOWN DATA ONLY in the latent space\n",
    "    with torch.no_grad():\n",
    "        z_unknown = model.encoder_forward(x_train[y_train == unknown_class_value])\n",
    "    \n",
    "    baseline_k_est, baseline_silhouette_scores = estimate_k_with_silhouette(k_min=2, k_max=30, x=z_unknown)\n",
    "    \n",
    "    print(f\"Estimated k={baseline_k_est}, ground truth={gt_k} (difference={np.abs(gt_k - baseline_k_est)})\")\n",
    "\n",
    "    # (3) In the latent space, run k-means with the estimated number of clusters\n",
    "    accs, nmis, aris = [], [], []\n",
    "    for i in range(5):\n",
    "        y_test_unknown_pred = model.predict_new_data(baseline_k_est, x_test_unknown, x_test_known, y_test_known)\n",
    "        accs.append(hungarian_accuracy(y_test_unknown_save_codes, y_test_unknown_pred))\n",
    "        nmis.append(normalized_mutual_info_score(y_test_unknown_save_codes, y_test_unknown_pred))\n",
    "        aris.append(adjusted_rand_score(y_test_unknown_save_codes, y_test_unknown_pred))\n",
    "    \n",
    "    accs_6.append(np.mean(accs))\n",
    "    nmis_6.append(np.mean(nmis))\n",
    "    aris_6.append(np.mean(aris))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c6ab58c-901c-4646-9df3-bbae1a12f468",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"TEST: ACC={:.1f}±{:.1f} | NMI={:.1f}±{:.1f} | ARI={:.1f}±{:.1f}\".format(np.mean(accs_6)*100, np.std(accs_6)*100, np.mean(nmis_6)*100, np.std(nmis_6)*100, np.mean(aris_6)*100, np.std(aris_6)*100))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3205bde-7444-47cd-a28b-5456fad553d5",
   "metadata": {},
   "source": [
    "## 4.3 TabularNCD <a class=\"anchor\" id=\"4.3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0149071d-0ed4-4aa9-a2cb-6a9642423a75",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_config = {\n",
    "    'hidden_layers_dims': [math.floor(3*x_train.shape[1]/4), math.floor(2*x_train.shape[1]/4)],\n",
    "    'input_size': x_train.shape[1],\n",
    "    'n_known_classes': len(np.unique(y_train)),  # Takes into account the unknown class\n",
    "    'activation_fct': 'relu',\n",
    "    'use_batchnorm': True,\n",
    "    'batch_size': 512,\n",
    "    'epochs': 200,\n",
    "    'M': 2000,\n",
    "    \n",
    "    'n_unknown_classes': k_est,  # /!\\ We use the estimation here /!\\\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5eb11f9f-1ad8-4c97-a460-ccfa08cb433d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We load the hyperparameters that were optimized through grid search\n",
    "d = json.load(open(\"hyperparameters.json\"))\n",
    "config = d[\"TabularNCD\"][\"Estimated k\"][dataset_name]\n",
    "print(\"Using hyperparameters:\", config)\n",
    "\n",
    "tncd_config = base_config.copy()\n",
    "tncd_config.update(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75f0f283-e9eb-450e-aace-b6ccb504d75d",
   "metadata": {},
   "outputs": [],
   "source": [
    "accs_7, nmis_7, aris_7 = [], [], []\n",
    "\n",
    "for i in tqdm(range(10)):\n",
    "    model = TabularNCDModel(tncd_config).to(device)\n",
    "    losses_dict = model.joint_training(config=tncd_config,\n",
    "                                       x_train=x_train, y_train=y_train,\n",
    "                                       x_test_known=x_test_known, y_test_known=y_test_known,\n",
    "                                       x_test_unknown=x_test_unknown, y_test_unknown=y_test_unknown_save_codes,\n",
    "                                       y_train_unknown=y_train_unknown_save_codes,\n",
    "                                       unknown_class_value=unknown_class_value,\n",
    "                                       disable_tqdm=True)\n",
    "    \n",
    "    accs_7.append(losses_dict[\"test cluster ACC\"][-1])\n",
    "    nmis_7.append(losses_dict[\"test cluster NMI\"][-1])\n",
    "    aris_7.append(losses_dict[\"test cluster ARI\"][-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72473a84-7c78-4498-8f5d-80e8297e5374",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"TEST: ACC={:.1f}±{:.1f} | NMI={:.1f}±{:.1f} | ARI={:.1f}±{:.1f}\".format(np.mean(accs_7)*100, np.std(accs_7)*100, np.mean(nmis_7)*100, np.std(nmis_7)*100, np.mean(aris_7)*100, np.std(aris_7)*100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb6dc496-fc02-4eff-926d-e5499ae14c62",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15eda257-43c8-4ff9-98af-7edbd647764d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venvexploration",
   "language": "python",
   "name": "venvexploration"
  },
  "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}