gaia_models.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9ead029",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "import pickle\n",
    "import seaborn as sns\n",
    "\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, r2_score, confusion_matrix\n",
    "from sklearn.metrics import mean_absolute_error as mae\n",
    "from sklearn.metrics import mean_squared_error as mse\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import StandardScaler, MinMaxScaler\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import Sequential, Input\n",
    "from tensorflow.keras.callbacks import EarlyStopping\n",
    "from tensorflow.keras.layers import Conv1D, MaxPooling1D, Flatten, Dense\n",
    "from tensorflow.keras.models import load_model\n",
    "from tensorflow.python.keras.layers.kernelized import RandomFourierFeatures\n",
    "\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c3a5262-b0e2-435c-be98-23030a828a50",
   "metadata": {},
   "source": [
    "# Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbf1e093",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = \"data/classification/dr3\"\n",
    "def load_data(data_dir):\n",
    "    list_5d = []\n",
    "    list_6d = []\n",
    "    for file in sorted(os.listdir(data_dir)):\n",
    "        if \"apg\" in file:\n",
    "            continue\n",
    "        df = pd.read_csv(os.path.join(data_dir, file), index_col=0)\n",
    "        if \"5d\" in file:\n",
    "            list_5d.append(df)\n",
    "        else:\n",
    "            list_6d.append(df)\n",
    "            \n",
    "    df_5d = pd.concat(list_5d, axis=0, ignore_index=True)\n",
    "    df_6d = pd.concat(list_6d, axis=0, ignore_index=True)\n",
    "    return df_5d, df_6d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d099aabc",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_5d, df_6d = load_data(data_dir)\n",
    "\n",
    "df_6d[\"class_6d\"] = df_6d[\"class\"]\n",
    "df_6d[\"gm_p_high_6d\"] = df_6d[\"gm_p_high\"]\n",
    "\n",
    "df = pd.merge(df_5d, df_6d[[\"Host\", \"class_6d\", \"gm_p_high_6d\"]], on=\"Host\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fbfe818",
   "metadata": {},
   "outputs": [],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67f7dcbf",
   "metadata": {},
   "source": [
    "# EDA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7284143b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def label_groups(row):\n",
    "    if row[\"class\"] == 0 and row[\"class_6d\"] == 2:\n",
    "        return 0\n",
    "    elif row[\"class\"] == 2 and row[\"class_6d\"] == 0:\n",
    "        return 1\n",
    "    elif row[\"class\"] == 0 and row[\"class_6d\"] == 1:\n",
    "        return 2\n",
    "    elif row[\"class\"] == 1 and row[\"class_6d\"] == 0:\n",
    "        return 3\n",
    "    elif row[\"class\"] == 2 and row[\"class_6d\"] == 1:\n",
    "        return 4\n",
    "    elif row[\"class\"] == 1 and row[\"class_6d\"] == 2:\n",
    "        return 5\n",
    "    elif row[\"class\"] == 0 and row[\"class_6d\"] == 0:\n",
    "        return 6\n",
    "    elif row[\"class\"] == 1 and row[\"class_6d\"] == 1:\n",
    "        return 7\n",
    "    elif row[\"class\"] == 2 and row[\"class_6d\"] == 2:\n",
    "        return 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1364af2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pie_groups(row):\n",
    "    if row[\"class_group\"] <= 1:\n",
    "        return 0\n",
    "    elif 5 >= row[\"class_group\"] >= 2:\n",
    "        return 1\n",
    "    else:\n",
    "        return 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9e95469",
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"class_group\"] = df.apply(lambda row: label_groups(row), axis=1)\n",
    "groups = [\"5d low density - 6d high density\", \"6d low density - 5d high density\",\n",
    "          \"5d low density - 6d ambigous\", \"5d ambigous - 6d low density\",\n",
    "          \"5d high density - 6d ambigous\", \"5d ambigous - 6d high density\",\n",
    "          \"5d low density - 6d low density\", \"5d ambigous - 6d ambigous\",\n",
    "          \"5d high density - 6d high density\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae680f78",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pie, ax = plt.subplots(figsize=[10,6], facecolor=\"w\")\n",
    "ax = sns.countplot(y=df[\"class_group\"])\n",
    "ax.set_yticklabels(labels=groups)\n",
    "plt.ylabel(\"\")\n",
    "plt.xlabel(\"Count\")\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"report_images/density_class_prediction_breakdown.png\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e614c22e",
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"pie_group\"] = df.apply(lambda row: pie_groups(row), axis=1)\n",
    "outs = df[df[\"pie_group\"] == 0][\"Host\"]\n",
    "pie_groups = [\"Matching classes\", \"Missmatched as ambigous\", \"Mismatched density classes\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "434a0478",
   "metadata": {},
   "outputs": [],
   "source": [
    "pie, ax = plt.subplots(figsize=[9,6], facecolor=\"w\")\n",
    "plt.pie(x=df[\"pie_group\"].value_counts(), autopct=\"%.1f%%\", labels=pie_groups, pctdistance=0.5)\n",
    "plt.title(\"\", fontsize=14)\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"report_images/density_class_prediction_comparison.png\", bbox_inches = \"tight\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56ecb6f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(25, 15), facecolor=\"white\")\n",
    "mask = np.triu(np.ones_like(df.drop([\"Host\", \"class\"], axis=1).corr(), dtype=bool))\n",
    "cmap = [\"#f781bf\", \"#a65628\", \"#ffff33\", \"#ff7f00\", \"#984ea3\", \"#4daf4a\", \"#377eb8\",\"#e41a1c\"]\n",
    "sns.heatmap(df.drop([\"Host\", \"class\"], axis=1).corr(), mask=mask, vmin=-1, vmax=1, cmap=cmap,square=True, linewidths=.5, annot=False)\n",
    "plt.savefig(\"report_images/corr_graph.png\", bbox_inches=\"tight\", pad_inches=0.05);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f606f82-59f0-46d7-96fa-b1a34444603b",
   "metadata": {},
   "source": [
    "# Helping functions for model training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90b8d341-dfce-46f9-99b9-5bde7a8b897d",
   "metadata": {},
   "source": [
    "### Encoding data for neural networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "008fd47e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode_for_nn(y_train, y_val=None):\n",
    "    depth = tf.constant(3)\n",
    "    y_train = tf.one_hot(indices=y_train, depth=depth)\n",
    "    if y_val is None:\n",
    "        return y_train\n",
    "    y_val = tf.one_hot(indices=y_val, depth=depth)\n",
    "    \n",
    "    return y_train, y_val#, y_test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe8fa151-5356-459a-89a2-4d044d3f9ab9",
   "metadata": {},
   "source": [
    "### Scaling data for convolutional neural network input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af9a1c6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cnn_scale(x_train, x_val, x_test):\n",
    "    sc = StandardScaler()\n",
    "    x_train = sc.fit_transform(x_train)\n",
    "    x_test = sc.transform(x_test)\n",
    "    x_val = sc.transform(x_val)\n",
    "    x_train = np.reshape(x_train, (x_train.shape[0], x_train.shape[1], 1))\n",
    "    x_test = np.reshape(x_test, (x_test.shape[0], x_test.shape[1], 1))\n",
    "    x_val = np.reshape(x_val, (x_val.shape[0], x_val.shape[1], 1))\n",
    "    \n",
    "    return x_train, x_val, x_test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e73976a-656f-4b13-a64a-413ee9cf3b7c",
   "metadata": {},
   "source": [
    "### Create train/test sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9d9ff45-2097-41ae-a964-d8190dc07fff",
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_data(train_df, regression=False):\n",
    "    \n",
    "    if regression:\n",
    "        output = train_df[\"gm_p_high_6d\"]\n",
    "    else:\n",
    "        output = train_df[\"class_6d\"]\n",
    "    \n",
    "    x_train = train_df.iloc[20000:, :]\n",
    "    y_train = output[20000:]\n",
    "    x_test = train_df.iloc[:20000, :]\n",
    "    y_test = output[:20000]\n",
    "\n",
    "    hosts = x_test[\"Host\"]\n",
    "    x_train = x_train[[\"target_density\", \"densities_max\", \"densities_min\", \"densities_mean\", \"gm_cov_high\", \"gm_mean_high\"]]\n",
    "    x_test = x_test[[\"target_density\", \"densities_max\", \"densities_min\", \"densities_mean\", \"gm_cov_high\", \"gm_mean_high\"]]\n",
    "\n",
    "    return x_train, x_test, y_train, y_test, hosts"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6526247e-76fa-4ab7-be1a-b7a1e7581db8",
   "metadata": {},
   "source": [
    "### Return classification performance metrics of the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b766074-6a0a-40f9-9c45-ecdbdc9f24d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_scores(model, model_name, features, params, x_test, y_test, hosts, train_time, columns):\n",
    "\n",
    "    start = time.perf_counter()\n",
    "    if \"Random Forest\" in model_name:\n",
    "        predictions = model.predict(x_test)\n",
    "    else:\n",
    "        predictions = np.argmax(model.predict(x_test), axis=-1)\n",
    "    end = time.perf_counter()\n",
    "    pred_time = end-start\n",
    "\n",
    "    test_accuracy = accuracy_score(y_test, predictions)\n",
    "    test_precision = precision_score(y_test, predictions, average=\"macro\")\n",
    "    test_recall = recall_score(y_test, predictions, average=\"macro\")\n",
    "    test_f1 = f1_score(y_test, predictions, average=\"macro\")\n",
    "    \n",
    "    df = pd.DataFrame([[model_name, features, params, test_accuracy, test_precision, test_recall, test_f1, train_time, pred_time]], columns=columns)\n",
    "    return df, predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b76feb1e-7b0f-40f2-af7a-f413111250e3",
   "metadata": {},
   "source": [
    "### Return regression performance metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9729710-1a05-43c5-8df1-a2816def6c94",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_reg_scores(model, model_name, features, params, x_test, y_test, hosts, train_time, columns):\n",
    "\n",
    "    start = time.perf_counter()\n",
    "    \n",
    "    predictions = model.predict(x_test)\n",
    " \n",
    "    end = time.perf_counter()\n",
    "    pred_time = end-start\n",
    "\n",
    "    mse_score = mse(y_test, predictions)\n",
    "    r2 = r2_score(y_test, predictions)\n",
    "    \n",
    "    df = pd.DataFrame([[model_name, features, params, mse_score, r2, train_time, pred_time]], columns=columns)\n",
    "    return df, predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74ec06b3-8dec-458a-b23b-c96491544c0b",
   "metadata": {},
   "source": [
    "### Random Forest model setups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c211f07-9340-4a4b-ad82-aed4ca62f868",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Classification\n",
    "def rfc(x_train, y_train):\n",
    "    model = RandomForestClassifier()\n",
    "    old = model.get_params()\n",
    "    model.set_params(n_estimators=100, max_depth=25, class_weight=\"balanced\", n_jobs=-1)\n",
    "    new = model.get_params()\n",
    "    \n",
    "    start = time.perf_counter()\n",
    "    model.fit(x_train, y_train)\n",
    "    end = time.perf_counter()\n",
    "    t = end-start\n",
    "    \n",
    "    params = {k: new[k] for k in new if k in old and new[k] != old[k]}\n",
    "    \n",
    "    return model, params, t, \"Random Forest\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1db66ecf-30b2-4d1b-a086-5ed625fcc72f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Regression\n",
    "def rfr(x_train, y_train):\n",
    "    model = RandomForestRegressor()\n",
    "    old = model.get_params()\n",
    "    model.set_params(n_jobs=-1)\n",
    "    new = model.get_params()\n",
    "    \n",
    "    start = time.perf_counter()\n",
    "    model.fit(x_train, y_train)\n",
    "    end = time.perf_counter()\n",
    "    t = end-start\n",
    "    \n",
    "    params = {k: new[k] for k in new if k in old and new[k] != old[k]}\n",
    "    \n",
    "    return model, params, t, \"Random Forest\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26b912c4-bf68-40d9-b95c-4893b795f857",
   "metadata": {},
   "source": [
    "### Convolutional neural network architectures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f89f4b9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Classification\n",
    "def cnn_1d(x_train, y_train, x_val, y_val, params):\n",
    "    params = []\n",
    "    opt = tf.keras.optimizers.Adam(learning_rate=0.0001)\n",
    "    stopper = EarlyStopping(monitor=\"val_loss\", patience=5)\n",
    "    n_timesteps, n_features, n_outputs = x_train.shape[1], x_train.shape[2], 3\n",
    "    activation = \"relu\"\n",
    "    epochs = 100\n",
    "    \n",
    "    model = Sequential()\n",
    "    model.add(Input(shape=(n_timesteps, n_features)))\n",
    "    params.append(f\"Input: ({n_timesteps}, {n_features})\")\n",
    "    model.add(Conv1D(filters=256, kernel_size=4, activation=activation))\n",
    "    params.append(f\"Conv1D: filters=256, kernel_size=4, activation={activation}\")\n",
    "\n",
    "    model.add(MaxPooling1D(pool_size=3))\n",
    "    params.append(f\"MaxPooling1D: pool_size=3\")\n",
    "    model.add(Flatten())\n",
    "    params.append(f\"Flatten\")\n",
    "    model.add(Dense(units=128, activation=activation))\n",
    "    params.append(f\"Dense: units=128, activation={activation}\")\n",
    "    model.add(Dense(units=64, activation=activation))\n",
    "    params.append(f\"Dense: units=64, activation={activation}\")\n",
    "    model.add(Dense(units=3, activation=\"softmax\"))\n",
    "    params.append(f\"Dense: units=3, activation=softmax\")\n",
    "    model.compile(optimizer=opt, loss=\"categorical_crossentropy\", metrics=[\"categorical_accuracy\"])\n",
    "    params.append(f\"Compile: Optimizer=Adam(learning_rate=0.0001), loss=categorical_crossentropy, metrics=categorical_accuracy\")\n",
    "    params.append(f\"Args: Epochs={epochs}, Callbacks=EarlyStopping(monitor=val_loss, patience=5)\")\n",
    "    \n",
    "    start = time.perf_counter()\n",
    "    history = model.fit(x_train, y_train, validation_data=(x_val, y_val), epochs=epochs, callbacks=stopper, verbose=0)\n",
    "    end = time.perf_counter()\n",
    "    t = end-start\n",
    "    \n",
    "    return model, history, params, t, \"1dCNN\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b174dea-bea6-4071-9311-fa3790b0ba95",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Regression\n",
    "def cnn_1d_reg(x_train, y_train, x_val, y_val, params):\n",
    "    tf.keras.callbacks.History()\n",
    "    opt = tf.keras.optimizers.Adam(learning_rate=0.0001)\n",
    "    stopper = EarlyStopping(monitor=\"val_loss\", patience=5)\n",
    "    n_timesteps, n_features, n_outputs = x_train.shape[1], x_train.shape[2], 1\n",
    "    activation = \"relu\"\n",
    "    epochs = 100\n",
    "    \n",
    "    model = Sequential()\n",
    "    model.add(Input(shape=(n_timesteps, n_features)))\n",
    "    params.append(f\"Input: ({n_timesteps}, {n_features})\")\n",
    "    model.add(Conv1D(filters=256, kernel_size=3, activation=activation))\n",
    "    params.append(f\"Conv1D: filters=256, kernel_size=3, activation={activation}\")\n",
    "    model.add(Conv1D(filters=32, kernel_size=3, activation=activation))\n",
    "    params.append(f\"Conv1D: filters=32, kernel_size=3, activation={activation}\")\n",
    "    model.add(MaxPooling1D(pool_size=2))\n",
    "    params.append(f\"MaxPooling1D: pool_size=2\")\n",
    "    model.add(Flatten())\n",
    "    params.append(f\"Flatten\")\n",
    "    model.add(Dense(units=64, activation=activation))\n",
    "    params.append(f\"Dense: units=64, activation={activation}\")\n",
    "    model.add(Dense(units=32, activation=activation))\n",
    "    params.append(f\"Dense: units=32, activation={activation}\")\n",
    "\n",
    "    model.add(Dense(units=n_outputs, activation=\"sigmoid\"))\n",
    "    params.append(f\"Dense: units={n_outputs}, activation=sigmoid\")\n",
    "    model.compile(optimizer=opt, loss=\"mse\")\n",
    "    params.append(f\"Compile: Optimizer=Adam(learning_rate=0.0001), loss=mse\")\n",
    "    params.append(f\"Args: Epochs={epochs}, Callbacks=EarlyStopping(monitor=val_loss, patience=5)\")\n",
    "    \n",
    "    start = time.perf_counter()\n",
    "    history = model.fit(x_train, y_train, validation_data=(x_val, y_val), epochs=epochs, callbacks=stopper, verbose=0)\n",
    "    end = time.perf_counter()\n",
    "    t = end-start\n",
    "    \n",
    "    return model, history, params, t, \"1dCNN\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d74e1c10-4ffa-4180-ac44-6a700d45ef51",
   "metadata": {},
   "source": [
    "### Classification neural network with RFF layer architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbf315a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Classification\n",
    "def ann(x_train, y_train):\n",
    "    params = []\n",
    "    opt = tf.keras.optimizers.Adam(learning_rate=0.0001)\n",
    "\n",
    "    stopper = EarlyStopping(monitor=\"val_loss\", patience=5)\n",
    "    n_features = x_train.shape[1]\n",
    "    activation = \"relu\"\n",
    "    epochs = 100\n",
    "    \n",
    "    model = tf.keras.Sequential()\n",
    "    model.add(Input(shape=(n_features, )))\n",
    "    params.append(f\"Input: ({n_features}, )\")\n",
    "    model.add(RandomFourierFeatures(output_dim=16384, scale=3., kernel_initializer=\"laplacian\"))\n",
    "    params.append(f\"RandomFourierFeatures: output_dim=16384, scale=3., kernel_initializer=laplacian\")\n",
    "    \n",
    "    model.add(Dense(units=64, activation=activation))\n",
    "    params.append(f\"Dense: units=128\")\n",
    "    model.add(Dense(units=16, activation=activation))\n",
    "    params.append(f\"Dense: units=64\")\n",
    "    \n",
    "    model.add(Dense(units=3, activation=\"softmax\"))\n",
    "    params.append(f\"Dense: units=3, activation=softmax\")\n",
    "    \n",
    "    model.compile(optimizer=opt, loss=\"categorical_crossentropy\", metrics=[\"categorical_accuracy\"])\n",
    "    \n",
    "    params.append(f\"Compile: Optimizer=Adam(learning_rate=0.0001), loss=categorical_crossentropy, metrics=categorical_accuracy\")\n",
    "    params.append(f\"Args: Epochs={epochs}, Callbacks=EarlyStopping(monitor=val_loss, patience=5)\")\n",
    "    \n",
    "    start = time.perf_counter()\n",
    "    history = model.fit(x_train, y_train, validation_split=0.1, epochs=epochs, callbacks=stopper, verbose=0)\n",
    "    end = time.perf_counter()\n",
    "    t = end-start\n",
    "    \n",
    "    params.append(f\"Validation split: 0.1\")\n",
    "    \n",
    "    return model, history, params, t, \"ANN\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a709ea12-6b65-471f-b6e5-3e1a43db00ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Regression\n",
    "def ann_reg(x_train, y_train):\n",
    "    tf.keras.callbacks.History()\n",
    "    params = []\n",
    "    opt = tf.keras.optimizers.Adam(learning_rate=0.0001)\n",
    "\n",
    "    stopper = EarlyStopping(monitor=\"val_loss\", patience=5)\n",
    "    n_features = x_train.shape[1]\n",
    "    activation = \"relu\"\n",
    "    epochs = 100\n",
    "    \n",
    "    model = tf.keras.Sequential()\n",
    "    model.add(Input(shape=(n_features, )))\n",
    "    params.append(f\"Input: ({n_features}, )\")\n",
    "    model.add(RandomFourierFeatures(output_dim=16384, scale=2., kernel_initializer=\"laplacian\"))\n",
    "    params.append(f\"RandomFourierFeatures: output_dim=16384, scale=2., kernel_initializer=laplacian\")\n",
    "    \n",
    "    model.add(Dense(units=64, activation=activation))\n",
    "    params.append(f\"Dense: units=64\")\n",
    "    model.add(Dense(units=16, activation=activation))\n",
    "    params.append(f\"Dense: units=16\")\n",
    "    \n",
    "    model.add(Dense(units=1, activation=\"sigmoid\"))\n",
    "    params.append(f\"Dense: units=1, activation=sigmoid\")\n",
    "    model.compile(optimizer=opt, loss=\"mse\")\n",
    "    params.append(f\"Compile: Optimizer=Adam(learning_rate=0.0001), loss=mse\")\n",
    "    params.append(f\"Args: Epochs={epochs}, Callbacks=EarlyStopping(monitor=val_loss, patience=5)\")\n",
    "    \n",
    "    start = time.perf_counter()\n",
    "    history = model.fit(x_train, y_train, validation_split=0.1, epochs=epochs, callbacks=stopper, verbose=0)\n",
    "    end = time.perf_counter()\n",
    "    t = end-start\n",
    "    \n",
    "    params.append(f\"Validation split: 0.1\")\n",
    "    \n",
    "    return model, history, params, t, \"ANN\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9a53f2c",
   "metadata": {},
   "source": [
    "# Classification models training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d3fcd1b-bce3-4873-8452-81908ef87db0",
   "metadata": {},
   "source": [
    "### Each training session is done 10 times and the mean of all results is saved"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a6fd552",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "columns = [\"Classifier\", \"Features\", \"Parameters\", \"Accuracy\", \"Precision\", \"Recall\", \"F1-score\", \"Train Duration\", \"Predict Duration\"]\n",
    "set_size = df_5d.shape[0]\n",
    "results_dir = \"results/classification\"\n",
    "results_file = \"results\" + \"_\" + str(set_size) + \".csv\"\n",
    "iters = 1\n",
    "results = pd.DataFrame(columns=columns)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbd18e50-0540-4e98-a92f-93b5e2e96096",
   "metadata": {},
   "source": [
    "### Random Forest training pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "635d4282-203a-461c-aa22-136de56663f7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train_rf(results):\n",
    "    scores = pd.DataFrame(columns=columns)\n",
    "    best_score = 0\n",
    "    best_pred = None\n",
    "    best_model = None\n",
    "\n",
    "    j = 0\n",
    "    for file in os.listdir(\"saved_models\"):\n",
    "        if \"RF_model\" in file:\n",
    "            j += 1\n",
    "\n",
    "    for i in range(iters):\n",
    "        x_train, x_test, y_train, y_test, hosts = split_data(df)\n",
    "        features = x_test.columns.values\n",
    "        model, params, train_time, model_name = rfc(x_train, y_train)\n",
    "        model_name = model_name + f\" {str(j)}\"\n",
    "        score, predictions = get_scores(model, model_name, features, params, x_test, y_test, hosts, train_time, columns)\n",
    "        scores = scores.append(score, ignore_index=True)\n",
    "        if score[\"F1-score\"][0] > best_score:\n",
    "            best_score = score[\"F1-score\"][0]\n",
    "            best_pred = predictions\n",
    "            best_model = model\n",
    "\n",
    "    cm = confusion_matrix(y_test, best_pred)\n",
    "\n",
    "    fig = plt.figure(figsize=(7, 7), facecolor=\"w\")\n",
    "    ticklabels = [\"Underdensity\", \"Ambigous\", \"Overdensity\"]\n",
    "    sns.heatmap(cm, yticklabels=ticklabels, xticklabels=ticklabels, cmap=\"viridis\", center=0, annot=True, fmt=\"g\", cbar=False, annot_kws={\"size\": 15})\n",
    "    plt.savefig(f\"report_images/RF_matrix_{j}.png\")\n",
    "    plt.close()\n",
    "    pickle.dump(best_model, open(f\"saved_models/RF_model_{j}.sav\", \"wb\"))\n",
    "\n",
    "    results = results.append([pd.concat([scores.iloc[0, :3], scores.mean()], axis=0)])\n",
    "    \n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68779548-3202-4345-a0ca-bd4fcad21d6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train_rf(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b3f699c-b7ee-4b45-baab-f621d1c78a21",
   "metadata": {},
   "source": [
    "### 1dCNN training pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dddccd8f-d94d-454e-8de1-d1d47a83643f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train_1dcnn(results):\n",
    "    scores = pd.DataFrame(columns=columns)\n",
    "    best_score = 0\n",
    "    best_pred = None\n",
    "    best_model = None\n",
    "    best_history = None\n",
    "    \n",
    "    j = 0\n",
    "    for file in os.listdir(\"saved_models/1DCNN\"):\n",
    "        if \"1DCNN_model\" in file:\n",
    "            j += 1\n",
    "    \n",
    "    for i in range(iters):\n",
    "        params = []\n",
    "        x_train, x_test, y_train, y_test, hosts = split_data(df)\n",
    "        features = x_test.columns.values\n",
    "        x_train, x_val, y_train, y_val = train_test_split(x_train, y_train, test_size=0.2)\n",
    "        params.append(\"Validation split: 0.2\")\n",
    "        y_train, y_val = encode_for_nn(y_train, y_val)\n",
    "        x_train, x_val, x_test = cnn_scale(x_train, x_val, x_test)\n",
    "        model, history, params, train_time, model_name = cnn_1d(x_train, y_train, x_val, y_val, params)\n",
    "        model_name = model_name + f\" {str(j)}\"\n",
    "        score, predictions = get_scores(model, model_name, features, params, x_test, y_test, hosts, train_time, columns)\n",
    "        scores = scores.append(score, ignore_index=True)\n",
    "        if score[\"F1-score\"][0] > best_score:\n",
    "            best_score = score[\"F1-score\"][0]\n",
    "            best_pred = predictions\n",
    "            best_model = model\n",
    "            best_history = history\n",
    "\n",
    "    cm = confusion_matrix(y_test, best_pred)\n",
    "    fig = plt.figure(figsize=(7, 5), facecolor=\"w\")\n",
    "    ticklabels = [\"Underdensity\", \"Ambigous\", \"Overdensity\"]\n",
    "    sns.heatmap(cm, yticklabels=ticklabels, xticklabels=ticklabels, cmap=\"viridis\", center=0, annot=True, fmt=\"g\", cbar=False, annot_kws={\"size\": 15})\n",
    "    plt.savefig(f\"report_images/1DCNN/1DCNN_matrix_{j}.png\")\n",
    "    #plt.show()\n",
    "    plt.close()\n",
    "    \n",
    "    fig = plt.figure(figsize=(7, 5), facecolor=\"w\")\n",
    "    plt.plot(best_history.history[\"categorical_accuracy\"])\n",
    "    plt.plot(best_history.history[\"val_categorical_accuracy\"])\n",
    "    plt.title(\"model accuracy\")\n",
    "    plt.ylabel(\"accuracy\")\n",
    "    plt.xlabel(\"epoch\")\n",
    "    plt.ylim(0.89, 0.935)\n",
    "    plt.legend([\"training\", \"validation\"], loc=\"upper left\")\n",
    "    plt.savefig(f\"report_images/1DCNN/1DCNN_accuracy_{j}.png\")\n",
    "    #plt.show()\n",
    "    plt.close()\n",
    "    \n",
    "    fig = plt.figure(figsize=(7, 5), facecolor=\"w\")\n",
    "    plt.plot(best_history.history[\"loss\"])\n",
    "    plt.plot(best_history.history[\"val_loss\"])\n",
    "    plt.title(\"model loss\")\n",
    "    plt.ylabel(\"loss\")\n",
    "    plt.xlabel(\"epoch\")\n",
    "    plt.ylim(0.16, 0.26)\n",
    "    plt.legend([\"training\", \"validation\"], loc=\"upper left\")\n",
    "    plt.savefig(f\"report_images/1DCNN/1DCNN_loss_{j}.png\")\n",
    "    plt.close()\n",
    "    model.save(f\"saved_models/1DCNN/1DCNN_model_{j}\")\n",
    "\n",
    "    results = results.append([pd.concat([scores.iloc[0, :3], scores.mean()], axis=0)])\n",
    "    \n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dcdafdc-1698-4d01-850c-d74bd8f4211f",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train_1dcnn(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2122b742-67ba-43b9-9987-0eb3d62daf77",
   "metadata": {},
   "source": [
    "### ANN training pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30c4b393-611e-41d5-97fa-c3b902cde483",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train_ann(results):\n",
    "    scores = pd.DataFrame(columns=columns)\n",
    "    best_score = 0\n",
    "    best_pred = None\n",
    "    best_model = None\n",
    "    best_history = None\n",
    "    \n",
    "    j = 0\n",
    "    for file in os.listdir(\"saved_models/ANN\"):\n",
    "        if \"ANN_model\" in file:\n",
    "            j += 1\n",
    "    \n",
    "    for i in range(iters):\n",
    "        x_train, x_test, y_train, y_test, hosts = split_data(df)\n",
    "        features = x_test.columns.values\n",
    "        y_train= encode_for_nn(y_train)\n",
    "        model, history, params, train_time, model_name = ann(x_train, y_train)\n",
    "        model_name = model_name + f\" {str(j)}\"\n",
    "        score, predictions = get_scores(model, model_name, features, params, x_test, y_test, hosts, train_time, columns)\n",
    "        scores = scores.append(score, ignore_index=True)\n",
    "        if score[\"F1-score\"][0] > best_score:\n",
    "            best_score = score[\"F1-score\"][0]\n",
    "            best_pred = predictions\n",
    "            best_model = model\n",
    "            best_history = history\n",
    "\n",
    "    cm = confusion_matrix(y_test, best_pred)\n",
    "    fig = plt.figure(figsize=(7, 5), facecolor=\"w\")\n",
    "    ticklabels = [\"Underdensity\", \"Ambigous\", \"Overdensity\"]\n",
    "    sns.heatmap(cm, yticklabels=ticklabels, xticklabels=ticklabels, cmap=\"viridis\", center=0, annot=True, fmt=\"g\", cbar=False, annot_kws={\"size\": 15})\n",
    "    plt.savefig(f\"report_images/ANN/ANN_matrix_{j}.png\")\n",
    "    plt.close()\n",
    "    \n",
    "    fig = plt.figure(figsize=(7, 5), facecolor=\"w\")\n",
    "    plt.plot(best_history.history[\"categorical_accuracy\"])\n",
    "    plt.plot(best_history.history[\"val_categorical_accuracy\"])\n",
    "    plt.title(\"model accuracy\")\n",
    "    plt.ylabel(\"accuracy\")\n",
    "    plt.xlabel(\"epoch\")\n",
    "    plt.ylim(0.89, 0.935)\n",
    "    plt.legend([\"training\", \"validation\"], loc=\"upper left\")\n",
    "    plt.savefig(f\"report_images/ANN/ANN_accuracy_{j}.png\")\n",
    "    plt.close()\n",
    "    \n",
    "    fig = plt.figure(figsize=(7, 5), facecolor=\"w\")\n",
    "    plt.plot(best_history.history[\"loss\"])\n",
    "    plt.plot(best_history.history[\"val_loss\"])\n",
    "    plt.title(\"model loss\")\n",
    "    plt.ylabel(\"loss\")\n",
    "    plt.xlabel(\"epoch\")\n",
    "    plt.ylim(0.16, 0.26)\n",
    "    plt.legend([\"training\", \"validation\"], loc=\"upper left\")\n",
    "    plt.savefig(f\"report_images/ANN/ANN_loss_{j}.png\")\n",
    "    plt.close()\n",
    "    \n",
    "    model.save(f\"saved_models/ANN/ANN_model_{j}\")\n",
    "\n",
    "    results = results.append([pd.concat([scores.iloc[0, :3], scores.mean()], axis=0)])\n",
    "    \n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eeb003f0-82f6-49c5-9eb8-59e320d118da",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train_ann(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6312401-2c73-461a-a86e-f7e84d0dabf5",
   "metadata": {},
   "source": [
    "### Save classification results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecaef84e-801c-4eae-b731-52d584581c39",
   "metadata": {},
   "outputs": [],
   "source": [
    "if Path(os.path.join(results_dir, results_file)).is_file():\n",
    "    results.to_csv(os.path.join(results_dir, results_file), mode=\"a\", header=False, index=False)\n",
    "else:\n",
    "    results.to_csv(os.path.join(results_dir, results_file), index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c3b4353-21ae-4fd8-9a79-8a0bd7d5dabf",
   "metadata": {},
   "source": [
    "# Regression models training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05c4ecbe-06f9-489b-a01a-5b2616da2290",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "columns = [\"Regressor\", \"Features\", \"Parameters\", \"MSE\", \"R2 score\", \"Train Duration\", \"Predict Duration\"]\n",
    "set_size = df_5d.shape[0]\n",
    "results_dir = \"results/regression\"\n",
    "results_file = \"results\" + \"_\" + str(set_size) + \".csv\"\n",
    "iters = 1\n",
    "results = pd.DataFrame(columns=columns)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d8d4eef-134a-439d-be1c-66fff2c6614b",
   "metadata": {},
   "source": [
    "### Random Forest training pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33dcd79f-a99d-40c8-aa4b-6fd0023e0594",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train_rf(results):\n",
    "    scores = pd.DataFrame(columns=columns)\n",
    "    best_score = 0\n",
    "    best_pred = None\n",
    "    best_model = None\n",
    "\n",
    "    j = 0\n",
    "    for file in os.listdir(\"saved_models/RFR\"):\n",
    "        if \"RF_model\" in file:\n",
    "            j += 1\n",
    "    \n",
    "    for i in range(iters):\n",
    "        x_train, x_test, y_train, y_test, hosts = split_data(df, regression=True)\n",
    "        features = x_test.columns.values\n",
    "        model, params, train_time, model_name = rfr(x_train, y_train)\n",
    "        model_name = model_name + f\" {str(j)}\"\n",
    "        score, predictions = get_reg_scores(model, model_name, features, params, x_test, y_test, hosts, train_time, columns)\n",
    "        scores = scores.append(score, ignore_index=True)\n",
    "        if score[\"R2 score\"][0] > best_score:\n",
    "            best_score = score[\"R2 score\"][0]\n",
    "            best_pred = predictions\n",
    "            best_model = model\n",
    "\n",
    "    fig = plt.figure(figsize=(7, 5), facecolor=\"w\")\n",
    "    plt.scatter(y_test, predictions)\n",
    "    plt.xlabel(\"Prediction\")\n",
    "    plt.ylabel(\"Ground truth\")\n",
    "    plt.savefig(f\"report_images/RFR/RF_loss_{j}.png\")\n",
    "    plt.close()\n",
    "    \n",
    "    pickle.dump(best_model, open(f\"saved_models/RFR/RF_model_{j}.sav\", \"wb\"))\n",
    "\n",
    "    results = results.append([pd.concat([scores.iloc[0, :3], scores.mean()], axis=0)])\n",
    "    \n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5549fa55-bdf0-4f81-a74e-7dcac12a3a73",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train_rf(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3716636-4eeb-4a8c-a7eb-611c40cb65a2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train_1dcnn(results):\n",
    "    scores = pd.DataFrame(columns=columns)\n",
    "    best_score = 0\n",
    "    best_pred = None\n",
    "    best_model = None\n",
    "    best_history = None\n",
    "    \n",
    "    j = 0\n",
    "    for file in os.listdir(\"saved_models/1DCNN_reg\"):\n",
    "        if \"1DCNN_model\" in file:\n",
    "            j += 1\n",
    "    \n",
    "    for i in range(iters):\n",
    "        params = []\n",
    "        x_train, x_test, y_train, y_test, hosts = split_data(df, regression=True)\n",
    "        features = x_test.columns.values\n",
    "        x_train, x_val, y_train, y_val = train_test_split(x_train, y_train, test_size=0.2)\n",
    "        params.append(\"Validation split: 0.2\")\n",
    "        x_train, x_val, x_test = cnn_scale(x_train, x_val, x_test)\n",
    "        model, history, params, train_time, model_name = cnn_1d_reg(x_train, y_train, x_val, y_val, params)\n",
    "        model_name = model_name + f\" {str(j)}\"\n",
    "        score, predictions = get_reg_scores(model, model_name, features, params, x_test, y_test, hosts, train_time, columns)\n",
    "        scores = scores.append(score, ignore_index=True)\n",
    "        if score[\"R2 score\"][0] > best_score:\n",
    "            best_score = score[\"R2 score\"][0]\n",
    "            best_pred = predictions\n",
    "            best_model = model\n",
    "            best_history = history\n",
    "    \n",
    "    fig = plt.figure(figsize=(7, 5), facecolor=\"w\")\n",
    "    plt.scatter(y_test, predictions)\n",
    "    plt.xlabel(\"Prediction\")\n",
    "    plt.ylabel(\"Ground truth\")\n",
    "    plt.savefig(f\"report_images/1DCNN_reg/1DCNN_preds_{j}.png\")\n",
    "    plt.close()\n",
    "    \n",
    "    fig = plt.figure(figsize=(7, 5), facecolor=\"w\")\n",
    "    plt.plot(best_history.history[\"loss\"])\n",
    "    plt.plot(best_history.history[\"val_loss\"])\n",
    "    plt.title(\"model loss\")\n",
    "    plt.ylabel(\"loss\")\n",
    "    plt.xlabel(\"epoch\")\n",
    "    plt.legend([\"training\", \"validation\"], loc=\"upper left\")\n",
    "    plt.savefig(f\"report_images/1DCNN_reg/1DCNN_loss_{j}.png\")\n",
    "    plt.close()\n",
    "    model.save(f\"saved_models/1DCNN_reg/1DCNN_model_{j}\")\n",
    "\n",
    "    results = results.append([pd.concat([scores.iloc[0, :3], scores.mean()], axis=0)])\n",
    "    \n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58928e35-7891-447e-a712-14cf1845d4c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train_1dcnn(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ead8db8-9c01-4e9f-b3c0-59cabd5bd0fb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train_ann(results):\n",
    "    scores = pd.DataFrame(columns=columns)\n",
    "    best_score = 0\n",
    "    best_pred = None\n",
    "    best_model = None\n",
    "    best_history = None\n",
    "    \n",
    "    j = 0\n",
    "    for file in os.listdir(\"saved_models/ANN_reg\"):\n",
    "        if \"ANN_model\" in file:\n",
    "            j += 1\n",
    "    \n",
    "    for i in range(iters):\n",
    "        x_train, x_test, y_train, y_test, hosts = split_data(df, regression=True)\n",
    "        features = x_test.columns.values\n",
    "        model, history, params, train_time, model_name = ann_reg(x_train, y_train)\n",
    "        model_name = model_name + f\" {str(j)}\"\n",
    "        score, predictions = get_reg_scores(model, model_name, features, params, x_test, y_test, hosts, train_time, columns)\n",
    "        scores = scores.append(score, ignore_index=True)\n",
    "        if score[\"R2 score\"][0] > best_score:\n",
    "            best_score = score[\"R2 score\"][0]\n",
    "            best_pred = predictions\n",
    "            best_model = model\n",
    "            best_history = history\n",
    "\n",
    "    fig = plt.figure(figsize=(7, 5), facecolor=\"w\")\n",
    "    plt.scatter(y_test, predictions)\n",
    "    plt.xlabel(\"Prediction\")\n",
    "    plt.ylabel(\"Ground truth\")\n",
    "    plt.savefig(f\"report_images/ANN_reg/ANN_preds_{j}.png\")\n",
    "    plt.close()\n",
    "    \n",
    "    fig = plt.figure(figsize=(7, 5), facecolor=\"w\")\n",
    "    plt.plot(best_history.history[\"loss\"])\n",
    "    plt.plot(best_history.history[\"val_loss\"])\n",
    "    plt.title(\"model loss\")\n",
    "    plt.ylabel(\"loss\")\n",
    "    plt.xlabel(\"epoch\")\n",
    "    plt.legend([\"training\", \"validation\"], loc=\"upper left\")\n",
    "    plt.savefig(f\"report_images/ANN_reg/ANN_loss_{j}.png\")\n",
    "    plt.close()\n",
    "    model.save(f\"saved_models/ANN_reg/ANN_model_{j}\")\n",
    "    \n",
    "    results = results.append([pd.concat([scores.iloc[0, :3], scores.mean()], axis=0)])\n",
    "    \n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5be4e268-ba8f-44d2-aa45-fe4fb37d3ded",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train_ann(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d5cdfe2-9fbb-47ba-b5e9-dbba7621037c",
   "metadata": {},
   "source": [
    "### Save regression results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d8b23c7-6edc-4e3b-ba1a-655b8e0892fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "if Path(os.path.join(results_dir, results_file)).is_file():\n",
    "    results.to_csv(os.path.join(results_dir, results_file), mode=\"a\", header=False, index=False)\n",
    "else:\n",
    "    results.to_csv(os.path.join(results_dir, results_file), index=False)"
   ]
  }
 ],
 "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}