CWRU.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "901ae9e1-41e0-4e0d-b1ce-e09a5c2b39d9",
   "metadata": {},
   "source": [
    "# CWRU Experiments"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e87143e1-a8e5-49a0-89d0-f661041f42af",
   "metadata": {},
   "source": [
    "This notebook contains the basic code to reconstruct the results shown in the paper, in terms of feature selection, model selection and model evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79db7880-58d7-4eb1-a250-08a10c04a0d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import torch\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "from kan import KAN\n",
    "from kan.utils import ex_round\n",
    "\n",
    "from utils import get_data, feature_selection, model_selection, plot_heatmaps, plot_pareto, plot_cm\n",
    "\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import classification_report\n",
    "\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2b42755-18ac-4982-8c14-af0fefcd9909",
   "metadata": {},
   "source": [
    "The setup is based on the classification problem, but it can be trivially extended to fault detection or severity classification by wrangling the `df` dataframe. An example for fault detection is provided in the commented code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d209d6a4-91c2-4bc6-93b6-7690c6e81573",
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment_name = 'cwru_classification'\n",
    "exp_seed = 42 # Add here the random seed\n",
    "\n",
    "# Read the data\n",
    "df = pd.read_csv(os.path.join('data','cwru.csv')).drop(columns=['severity'])\n",
    "\n",
    "\n",
    "# The following is an example of trivial wrangling to study the fault detection problem instead of classification\n",
    "\"\"\"\n",
    "# Keep all rows with label 0\n",
    "zeros_df = df[df['label'] == 0]\n",
    "\n",
    "# Filter out label 0 and sample 30 rows for each label from 1 to 5\n",
    "filtered_df = df[df['label'] != 0]\n",
    "sampled_df = filtered_df.groupby('label', group_keys=False).apply(lambda x: x.sample(min(len(x), 30), random_state=exp_seed))\n",
    "\n",
    "# Concatenate the two dataframes (zeros_df and sampled_df)\n",
    "df = pd.concat([zeros_df, sampled_df]).reset_index(drop=True)\n",
    "\n",
    "# Cast labels to binary format\n",
    "df['label'] = df['label'].apply(lambda x: 1 if x != 0 else 0)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5870490-eea6-41ec-97c9-637b605f5798",
   "metadata": {},
   "source": [
    "## Feature Selection"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9861ec4e-bf33-4d4d-ad4d-c7071aafd6fc",
   "metadata": {},
   "source": [
    "We run a grid search over $\\lambda$ and $\\tau$ to find the combination that retains the smallest number of features, with the highest attained F1-Score."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21c014ef-3a48-4b18-b0b7-3783673e7dd5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Parameters\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "grid_size = 5\n",
    "grid_eps = 0.05\n",
    "k = 3\n",
    "epochs = 80\n",
    "use_scaler = True\n",
    "data_split = (70,15,15)\n",
    "optim = \"Adam\"\n",
    "\n",
    "thresholds = np.linspace(0.01, 0.1, 20)\n",
    "lambdas = np.linspace(0.001, 0.01, 20)\n",
    "\n",
    "featsdf = feature_selection(df=df, grid_size=grid_size, grid_eps=grid_eps, k=k,\n",
    "                            thresholds=thresholds, lambdas=lambdas, optim=optim, epochs=epochs, \n",
    "                            use_scaler=use_scaler, data_split=data_split, \n",
    "                            device=device, exp_seed=exp_seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa5e3649-51a1-4656-a2d4-f4a83408ca1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the results directory if it doesn't exist\n",
    "os.makedirs('results', exist_ok=True)\n",
    "\n",
    "# Save the dataframe locally as pickle file to avoid having to run everything again\n",
    "featsdf.to_pickle(os.path.join('results', f'{experiment_name}_featsdf.pkl'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae3592dc-b1b1-4dd9-abfa-c68aef43b441",
   "metadata": {},
   "source": [
    "The results of the feature selection process in terms of the grid search, as well as the generated Pareto set can be seen in the following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bea1b63-0f92-42d1-88c6-722dbbf3e3d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "indices = {'x': 'thresholds', 'y': 'lambdas', 'z0': 'f1_scores', 'z1': 'num_feats'}\n",
    "savepath = os.path.join('results', f'{experiment_name}_feat_heatmaps.pdf')\n",
    "titles = ['Interpolated F1-Score', 'Interpolated Number of Features']\n",
    "x_label, y_label = r'$\\tau$ Values', r'$\\lambda$ Values'\n",
    "cbar_labels = ['','']\n",
    "\n",
    "plot_heatmaps(featsdf, indices, savepath, interpolation='spline36', cmap='Spectral', titles=titles, x_label=x_label, y_label=y_label, cbar_labels=cbar_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba5ac606-2b65-41a2-84be-9e5baf5d746f",
   "metadata": {},
   "outputs": [],
   "source": [
    "cmap = plt.get_cmap('Spectral')\n",
    "\n",
    "plotcols = ['num_feats', 'f1_scores', 'pareto']\n",
    "savepath = os.path.join('results', f'{experiment_name}_feat_pareto.pdf')\n",
    "bg_col, pareto_col, nonpareto_col = '#f7fafa', cmap(0.25), cmap(1.0)\n",
    "labels=['Points outside of the Pareto set', 'Points within the Pareto set']\n",
    "title='F1-Score vs Number of Retained Features'\n",
    "x_label, y_label = 'Number of Features', 'F1-Score (%)'\n",
    "\n",
    "plot_pareto(featsdf, savepath, plotcols=plotcols, bg_col=bg_col, pareto_col=pareto_col, nonpareto_col=nonpareto_col,\n",
    "            labels=labels, title=title, x_label=x_label, y_label=y_label)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e30b4f9-3606-4228-843b-4ba1d2409710",
   "metadata": {},
   "source": [
    "Given the returned Pareto Set, we decide to work with the feature set that has the highest F1-Score, but contains at most 10 features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e43f0009-71b8-4fe3-86af-bd8a24e48c8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the index corresponding to the \"best\" point of the pareto set\n",
    "# Objective: Maximize F1-Score keeping up to 10 features\n",
    "\n",
    "# Get pareto set\n",
    "fpset = featsdf[featsdf['pareto']==True]\n",
    "\n",
    "# Select points with num_feats <= 10\n",
    "under_10_pset = fpset.loc[fpset['num_feats'] <= 10]\n",
    "\n",
    "# Get the index of the highest F1-Score among these points\n",
    "idx = under_10_pset.loc[under_10_pset['f1_scores'] == under_10_pset['f1_scores'].max()].index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a6c03ca-1c24-412d-a7da-06d6d4d282f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "fpset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33a597ad-3004-4f54-8f7a-811d4ef3ae01",
   "metadata": {},
   "source": [
    "This means that feature selection is complete and we are working with the following features:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "627c0ce9-8836-4d23-8f2a-29f1de421287",
   "metadata": {},
   "outputs": [],
   "source": [
    "kept_feats = featsdf.iloc[idx]['features'].values[0].tolist()\n",
    "lamb = featsdf.iloc[idx]['lambdas'].values[0]\n",
    "print(f\"Working with the {len(kept_feats)} following features:\\n{kept_feats} \\nand lambda = {lamb:.3f}.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01a50425-58df-4983-b5e7-4f0d5c9245be",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "de2ade9b-c853-4e7c-875a-2aedbe3d5980",
   "metadata": {},
   "source": [
    "## Model Selection"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c583ddb8-64c7-4854-8bad-757e4813a8e0",
   "metadata": {},
   "source": [
    "We run another grid search, this time for the model parameters $G$ and $g_e$, using only the retained features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a27ff4ad-a671-4485-8cd9-9f2846f05e6a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if use_scaler == True:\n",
    "    scaler = StandardScaler()\n",
    "else:\n",
    "    scaler = None\n",
    "\n",
    "dataset = get_data(df, scaler=scaler, data_split=data_split, final_eval=False, feat_idxs=kept_feats, device=device, exp_seed=exp_seed)\n",
    "\n",
    "k = 4\n",
    "alpha_param = 0.05\n",
    "beta_param = 1.5\n",
    "r2_threshold = 0.0\n",
    "\n",
    "optim = \"Adam\"\n",
    "epochs = 200\n",
    "grid_update_num = 10\n",
    "stop_grid_update_step = 150\n",
    "\n",
    "grid_sizes = [8, 10, 12, 15, 20, 30, 40, 50]\n",
    "grid_es = [0.0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1.0]\n",
    "\n",
    "modelsdf = model_selection(dataset, grid_sizes, grid_es, lamb=0.0, k=k, optim=optim, epochs=epochs, grid_update_num=grid_update_num,\n",
    "                           stop_grid_update_step=stop_grid_update_step, alpha=alpha_param, beta=beta_param,\n",
    "                           r2_threshold=r2_threshold, device=device, exp_seed=exp_seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b60564d5-d35b-4e80-a0d7-98dbff5a8f58",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Separate the grid search for adaptive grids and the one for non adaptive (g_e = 1.0)\n",
    "staticmodelsdf = modelsdf[modelsdf['grid_es'] == 1.0].reset_index(drop=True)\n",
    "modelsdf = modelsdf[modelsdf['grid_es'] != 1.0].reset_index(drop=True)\n",
    "\n",
    "# Save the dataframe locally as pickle file to avoid having to run everything again\n",
    "modelsdf.to_pickle(os.path.join('results', f'{experiment_name}_modelsdf.pkl'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "972f6d62-1d12-41c7-87c0-fd8b9bc7cc1c",
   "metadata": {},
   "source": [
    "As previously, the results of the model selection process in terms of the grid search, as well as the generated Pareto set can be seen in the following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c18a6564-4e6d-43bc-9096-8a03242ad59c",
   "metadata": {},
   "outputs": [],
   "source": [
    "indices = {'x': 'grid_es', 'y': 'grid_sizes', 'z0': 'f1_kan', 'z1': 'f1_sym'}\n",
    "savepath = os.path.join('results', f'{experiment_name}_model_heatmaps.pdf')\n",
    "titles = ['Interpolated KAN F1-Score', 'Interpolated Symbolic F1-Score']\n",
    "x_label, y_label = 'Grid Adaptability', 'Grid Size'\n",
    "cbar_labels = ['','']\n",
    "\n",
    "plot_heatmaps(modelsdf, indices, savepath, interpolation='spline36', cmap='Spectral', titles=titles, x_label=x_label, y_label=y_label, cbar_labels=cbar_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf18d987-003b-4623-a9d0-5c4b87e50941",
   "metadata": {},
   "outputs": [],
   "source": [
    "cmap = plt.get_cmap('Spectral')\n",
    "\n",
    "plotcols = ['f1_kan', 'f1_sym', 'pareto']\n",
    "savepath = os.path.join('results', f'{experiment_name}_model_pareto.pdf')\n",
    "bg_col, pareto_col, nonpareto_col = '#f7fafa', cmap(0.25), cmap(1.0)\n",
    "labels=['Points outside of the Pareto set', 'Points within the Pareto set']\n",
    "title='Achieved F1-Scores per run'\n",
    "x_label, y_label = 'KAN F1-Score (%)', 'Symbolic F1-Score (%)'\n",
    "\n",
    "\n",
    "plot_pareto(modelsdf, savepath, plotcols=plotcols, bg_col=bg_col, pareto_col=pareto_col, nonpareto_col=nonpareto_col,\n",
    "            labels=labels, title=title, x_label=x_label, y_label=y_label)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2634e592-accf-4f69-8d29-f4065bbe8527",
   "metadata": {},
   "source": [
    "Given the returned Pareto Set, we decide to work with the model that has the highest Symbolic F1-Score."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bf17a95-5c34-414c-af15-1b705a9bae29",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the index corresponding to the \"best\" point of the pareto set\n",
    "# Objective: Maximize Mean F1-Score\n",
    "\n",
    "# Get pareto set\n",
    "mpset = modelsdf[modelsdf['pareto']==True]\n",
    "\n",
    "# Calculate mean F1-Score\n",
    "mean_f1 = 0.5*(mpset['f1_kan'] + mpset['f1_sym'])\n",
    "\n",
    "# Select point with maximum Symbolic F1-Score\n",
    "idx = mean_f1[mean_f1 == mean_f1.max()].index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31dce9df-ee00-4909-a315-b39c69116eb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "mpset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b82c3ce6-5a77-4bd8-b179-54a3ab8dae5e",
   "metadata": {},
   "source": [
    "This suggests that model selection is complete and we are working with the following hyperparameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58c49b40-7cf3-4234-9323-b5ab9a02f5ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "G, g_e = modelsdf.iloc[idx]['grid_sizes'].values[0], modelsdf.iloc[idx]['grid_es'].values[0]\n",
    "print(f\"Working with G = {G} and g_e = {g_e}.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f912caf-7245-44a2-8e15-42957fab44d9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "d0bf8164-786b-423a-976f-6f2747eabaf1",
   "metadata": {},
   "source": [
    "## Final Evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e92be86b-603e-4562-80a0-47c3bba52c3b",
   "metadata": {},
   "source": [
    "To perform the final evaluation with the selected features and model hyperparameters, we concatenate the training and validation data and use them to train a final model instance. Then, we assess its performance on the test data, which have not been used anywhere so far."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba21f25f-bf29-428f-9bad-10fc72e09d76",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Concatenate validation and training data\n",
    "if use_scaler == True:\n",
    "    final_scaler = StandardScaler()\n",
    "else:\n",
    "    final_scaler = None\n",
    "    \n",
    "final_data = get_data(df, scaler=final_scaler, data_split=data_split, final_eval=True, feat_idxs=kept_feats, device=device, exp_seed=exp_seed)\n",
    "\n",
    "# Initialize a model instance\n",
    "kan_input = final_data['train_input'].shape[1]\n",
    "kan_output = final_data['train_label'].unique().shape[0]\n",
    "model = KAN(width=[kan_input, kan_output], grid=G, k=4, grid_eps=g_e, sparse_init=False, seed=exp_seed, auto_save=False, device=device)\n",
    "\n",
    "# Train Model\n",
    "results = model.fit(final_data, opt=optim, steps=epochs, lamb=0.0, update_grid=True, grid_update_num=grid_update_num,\n",
    "                    stop_grid_update_step=stop_grid_update_step, loss_fn=torch.nn.CrossEntropyLoss())\n",
    "\n",
    "print(f\"Trained final model with {len(kept_feats)} features, using G = {G} and g_e = {g_e}.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93b8b6c3-fcad-4b93-a51a-a2d9e25fcc05",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate on test data\n",
    "preds = torch.argmax(model.forward(final_data['test_input']).detach(), dim=1).cpu()\n",
    "truth = final_data['test_label'].cpu()\n",
    "\n",
    "# Get classification report\n",
    "print(classification_report(truth, preds))\n",
    "\n",
    "# Plot Confusion Matrix\n",
    "class_names = [\"N\", \"F\",]\n",
    "cm_title = 'Confusion Matrix for Fault Classification'\n",
    "x_label = 'Predicted Class'\n",
    "y_label = 'True Class'\n",
    "plot_cm(preds, truth, class_names, percs=False, cmap='Blues', title=cm_title, x_label=x_label, y_label=y_label)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62f8ef7c-b8cc-4c3a-a36a-e160cfe3aa1e",
   "metadata": {},
   "source": [
    "The final step is the \"symbolification\" of the trained KAN and its re-evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef0a2df5-352a-4592-9415-317c765e616b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Symbolify Model\n",
    "model.auto_symbolic(verbose=1, alpha=alpha_param, beta=beta_param, r2_threshold=r2_threshold)\n",
    "\n",
    "# Evaluate Symbolic Version\n",
    "preds_sym = torch.argmax(model.forward(final_data['test_input']).detach(), dim=1).cpu()\n",
    "\n",
    "# Get classification report\n",
    "print(classification_report(truth, preds_sym))\n",
    "\n",
    "# Plot Confusion Matrix\n",
    "plot_cm(preds_sym, truth, class_names, percs=False, cmap='Blues', title=cm_title, x_label=x_label, y_label=y_label)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "506409e9-c15d-4c67-b494-0b97f3a73d1e",
   "metadata": {},
   "source": [
    "Extracting the symbolic formulas:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8050dfd-b13f-4870-ad18-b32379c37a61",
   "metadata": {},
   "outputs": [],
   "source": [
    "symbf = model.symbolic_formula()[0]\n",
    "\n",
    "for i in range(len(symbf)):\n",
    "    sf = symbf[i]\n",
    "    print(ex_round(sf, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "642d5847-8494-4f78-b8d7-d21cabe0d255",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}