{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import optimizers\n",
    "import datamanip\n",
    "import plotting\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.model_selection import train_test_split\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "FIG_DIR = os.path.abspath(\n",
    "    os.path.join(os.path.dirname(plotting.__file__), \"..\", \"figures\")\n",
    ")\n",
    "print(FIG_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(-1, 1, 100_000)\n",
    "y = datamanip.noise_data(datamanip.runge_function(x), 1.0)\n",
    "x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=plotting.get_figsize(0.5))\n",
    "ax.scatter(x_train, y_train, label=\"Train Data\", s=1)\n",
    "ax.scatter(x_test, y_test, label=\"Test Data\", s=1)\n",
    "ax.set_xlabel(\"$x$\")\n",
    "ax.set_ylabel(\"$y$\")\n",
    "ax.legend()\n",
    "fig.savefig(os.path.join(FIG_DIR, \"data_scatter.png\"), dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "beta_OLS_list = []\n",
    "mse_list = []\n",
    "train_mse_list = []\n",
    "r2_list = []\n",
    "train_r2_list = []\n",
    "\n",
    "polynomial_degrees = np.arange(1, 21)\n",
    "\n",
    "for polynomial_degree in polynomial_degrees:\n",
    "    X_train = datamanip.polynomial_features(x_train, polynomial_degree, False)\n",
    "    X_test = datamanip.polynomial_features(x_test, polynomial_degree, False)\n",
    "    X_train_scaled, X_test_scaled = datamanip.scale_data(X_train, X_test)\n",
    "    y_train_scaled, y_test_scaled = datamanip.scale_data(y_train, y_test)\n",
    "\n",
    "    beta = optimizers.OLS_parameters(X_train_scaled, y_train_scaled)\n",
    "\n",
    "    y_pred = X_test_scaled @ beta\n",
    "    mse, r2 = datamanip.evaluate_model(y_test_scaled, y_pred)\n",
    "    y_train_pred = X_train_scaled @ beta\n",
    "    train_mse, train_r2 = datamanip.evaluate_model(y_train_scaled, y_train_pred)\n",
    "\n",
    "    train_mse_list.append(train_mse)\n",
    "    train_r2_list.append(train_r2)\n",
    "    beta_OLS_list.append(beta)\n",
    "    mse_list.append(mse)\n",
    "    r2_list.append(r2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plotting.mse_r2_plot(\n",
    "    polynomial_degrees, train_mse_list, mse_list, train_r2_list, r2_list\n",
    ")\n",
    "fig.savefig(os.path.join(FIG_DIR, \"ols_mse_r2.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plotting.parameter_plot(polynomial_degrees, beta_OLS_list)\n",
    "fig.savefig(os.path.join(FIG_DIR, \"ols_parameter_plot.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "beta_ridge_list = []\n",
    "mse_ridge_list = []\n",
    "train_mse_ridge_list = []\n",
    "r2_ridge_list = []\n",
    "train_r2_ridge_list = []\n",
    "\n",
    "lambda_values = [0.5]\n",
    "\n",
    "for polynomial_degree in polynomial_degrees:\n",
    "    X_train = datamanip.polynomial_features(x_train, polynomial_degree, False)\n",
    "    X_test = datamanip.polynomial_features(x_test, polynomial_degree, False)\n",
    "    X_train_scaled, X_test_scaled = datamanip.scale_data(X_train, X_test)\n",
    "    y_train_scaled, y_test_scaled = datamanip.scale_data(y_train, y_test)\n",
    "\n",
    "    for lambda_ in lambda_values:\n",
    "        beta = optimizers.Ridge_parameters(X_train_scaled, y_train_scaled, lambda_)\n",
    "\n",
    "        y_pred = X_test_scaled @ beta\n",
    "        mse, r2 = datamanip.evaluate_model(y_test_scaled, y_pred)\n",
    "        y_train_pred = X_train_scaled @ beta\n",
    "        train_mse, train_r2 = datamanip.evaluate_model(y_train_scaled, y_train_pred)\n",
    "\n",
    "        train_mse_ridge_list.append(train_mse)\n",
    "        train_r2_ridge_list.append(train_r2)\n",
    "        beta_ridge_list.append(beta)\n",
    "        mse_ridge_list.append(mse)\n",
    "        r2_ridge_list.append(r2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plotting.mse_r2_plot(\n",
    "    polynomial_degrees,\n",
    "    train_mse_ridge_list,\n",
    "    mse_ridge_list,\n",
    "    train_r2_ridge_list,\n",
    "    r2_ridge_list,\n",
    ")\n",
    "fig.savefig(os.path.join(FIG_DIR, \"ridge_mse_r2.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "beta_ridge_list = []\n",
    "mse_ridge_list = []\n",
    "train_mse_ridge_list = []\n",
    "r2_ridge_list = []\n",
    "train_r2_ridge_list = []\n",
    "\n",
    "lambda_values = np.logspace(-5, 3, 20)\n",
    "polynomial_degrees = np.array([10])\n",
    "\n",
    "for polynomial_degree in polynomial_degrees:\n",
    "    X_train = datamanip.polynomial_features(x_train, polynomial_degree, False)\n",
    "    X_test = datamanip.polynomial_features(x_test, polynomial_degree, False)\n",
    "    X_train_scaled, X_test_scaled = datamanip.scale_data(X_train, X_test)\n",
    "    y_train_scaled, y_test_scaled = datamanip.scale_data(y_train, y_test)\n",
    "\n",
    "    for lambda_ in lambda_values:\n",
    "        beta = optimizers.Ridge_parameters(X_train_scaled, y_train_scaled, lambda_)\n",
    "\n",
    "        y_pred = X_test_scaled @ beta\n",
    "        mse, r2 = datamanip.evaluate_model(y_test_scaled, y_pred)\n",
    "        y_train_pred = X_train_scaled @ beta\n",
    "        train_mse, train_r2 = datamanip.evaluate_model(y_train_scaled, y_train_pred)\n",
    "\n",
    "        train_mse_ridge_list.append(train_mse)\n",
    "        train_r2_ridge_list.append(train_r2)\n",
    "        beta_ridge_list.append(beta)\n",
    "        mse_ridge_list.append(mse)\n",
    "        r2_ridge_list.append(r2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plotting.mse_r2_plot(\n",
    "    lambda_values,\n",
    "    train_mse_ridge_list,\n",
    "    mse_ridge_list,\n",
    "    train_r2_ridge_list,\n",
    "    r2_ridge_list,\n",
    "    labels={\"xlabel\": \"$\\\\lambda$\"},\n",
    ")\n",
    "for ax in axs:\n",
    "    ax.set_xscale(\"log\")\n",
    "fig.savefig(os.path.join(FIG_DIR, \"ridge_mse_r2_lambda.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plotting.parameter_plot(\n",
    "    lambda_values, beta_ridge_list, labels={\"xlabel\": \"$\\\\lambda$\"}\n",
    ")\n",
    "\n",
    "\n",
    "# Set xtick labels to scientific notation\n",
    "def format_number(num: float) -> str:\n",
    "    if num == 0:\n",
    "        return \"$0$\"\n",
    "    exponent = int(np.floor(np.log10(abs(num))))\n",
    "    coefficient = num / 10**exponent\n",
    "    return f\"${coefficient:.0f}\\\\cdot10^{{{exponent}}}$\"\n",
    "\n",
    "\n",
    "ax.set_xticklabels(\n",
    "    [f\"${format_number(tick)}$\" for tick in lambda_values[::3]], rotation=45\n",
    ")\n",
    "fig.savefig(os.path.join(FIG_DIR, \"ridge_parameter_plot.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = datamanip.polynomial_features(x_train, 10, False)\n",
    "X_test = datamanip.polynomial_features(x_test, 10, False)\n",
    "X_tr, X_te = datamanip.scale_data(X_train, X_test)\n",
    "y_tr, y_te = datamanip.scale_data(y_train, y_test)\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=plotting.get_figsize(0.5))\n",
    "num_iters = 1_000\n",
    "\n",
    "for learning_rate in np.logspace(-4, 0, 5):\n",
    "    OLSGrad_Desc = optimizers.OLSGradientDescent(\n",
    "        learning_rate=learning_rate, num_iterations=num_iters\n",
    "    )\n",
    "    OLSGrad_Desc.fit(X_tr, y_tr)\n",
    "    cost_history = OLSGrad_Desc.cost_history\n",
    "    RidgeGrad_Desc = optimizers.RidgeGradientDescent(\n",
    "        learning_rate=learning_rate, num_iterations=num_iters, lam=0.1\n",
    "    )\n",
    "    RidgeGrad_Desc.fit(X_tr, y_tr)\n",
    "    cost_history_ridge = RidgeGrad_Desc.cost_history\n",
    "    ax1.plot(cost_history, label=f\"$\\\\eta = 10^{{{int(np.log10(learning_rate))}}}$\")\n",
    "    ax2.plot(\n",
    "        cost_history_ridge, label=f\"$\\\\eta = 10^{{{int(np.log10(learning_rate))}}}$\"\n",
    "    )\n",
    "ax1.set_xlabel(\"Iteration\")\n",
    "ax1.set_ylabel(\"OLS Cost\")\n",
    "ax1.legend()\n",
    "ax1.set_ylim(bottom=0.465, top=0.505)\n",
    "ax2.set_ylim(bottom=0.4775, top=0.505)\n",
    "ax2.set_xlabel(\"Iteration\")\n",
    "ax2.set_ylabel(\"Ridge Cost\")\n",
    "ax2.legend()\n",
    "fig.tight_layout()\n",
    "fig.savefig(os.path.join(FIG_DIR, \"gradient_descent_convergence.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_cost_history(optimizer, X, y, **kwargs):\n",
    "    opt = optimizer(**kwargs)\n",
    "    opt.fit(X, y)\n",
    "    return opt.cost_history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=plotting.get_figsize(0.8))\n",
    "\n",
    "X_train = datamanip.polynomial_features(x_train, 10, False)\n",
    "X_test = datamanip.polynomial_features(x_test, 10, False)\n",
    "X_tr, X_te = datamanip.scale_data(X_train, X_test)\n",
    "y_tr, y_te = datamanip.scale_data(y_train, y_test)\n",
    "\n",
    "num_iters = 250\n",
    "learning_rate_ols = 0.1\n",
    "learning_rate_ridge = 0.01\n",
    "lam = 0.1\n",
    "\n",
    "optimizers_ols = [\n",
    "    (optimizers.OLSGradientDescent, {\"learning_rate\": learning_rate_ols}),\n",
    "    (optimizers.OLSMomentum, {\"learning_rate\": learning_rate_ols, \"delta\": 0.9}),\n",
    "    (optimizers.OLSAdaGrad, {\"learning_rate\": learning_rate_ols}),\n",
    "    (optimizers.OLSRMSProp, {\"learning_rate\": learning_rate_ols, \"gamma\": 0.9}),\n",
    "    (\n",
    "        optimizers.OLSAdam,\n",
    "        {\"learning_rate\": learning_rate_ols, \"beta1\": 0.9, \"beta2\": 0.999},\n",
    "    ),\n",
    "]\n",
    "\n",
    "optimizers_ridge = [\n",
    "    (\n",
    "        optimizers.RidgeGradientDescent,\n",
    "        {\"learning_rate\": learning_rate_ridge, \"lam\": lam},\n",
    "    ),\n",
    "    (\n",
    "        optimizers.RidgeMomentum,\n",
    "        {\"learning_rate\": learning_rate_ridge, \"delta\": 0.9, \"lam\": lam},\n",
    "    ),\n",
    "    (optimizers.RidgeAdaGrad, {\"learning_rate\": learning_rate_ridge, \"lam\": lam}),\n",
    "    (\n",
    "        optimizers.RidgeRMSProp,\n",
    "        {\"learning_rate\": learning_rate_ridge, \"gamma\": 0.9, \"lam\": lam},\n",
    "    ),\n",
    "    (\n",
    "        optimizers.RidgeAdam,\n",
    "        {\n",
    "            \"learning_rate\": learning_rate_ridge,\n",
    "            \"beta1\": 0.9,\n",
    "            \"beta2\": 0.999,\n",
    "            \"lam\": lam,\n",
    "        },\n",
    "    ),\n",
    "]\n",
    "\n",
    "\n",
    "def plot_optimizers(ax, optimizer_list, X, y, num_iters, ylabel=\"Cost\"):\n",
    "    for opt_class, params in optimizer_list:\n",
    "        history = get_cost_history(opt_class, X, y, num_iterations=num_iters, **params)\n",
    "        ax.plot(history, label=f\"{opt_class.__name__}\")\n",
    "    ax.set_xlabel(\"Iteration\")\n",
    "    ax.set_ylabel(ylabel)\n",
    "    ax.legend()\n",
    "\n",
    "\n",
    "plot_optimizers(ax1, optimizers_ols, X_tr, y_tr, num_iters, ylabel=\"Cost (OLS)\")\n",
    "ax1.set_ylim(bottom=0.465, top=0.505)\n",
    "\n",
    "plot_optimizers(ax2, optimizers_ridge, X_tr, y_tr, num_iters, ylabel=\"Cost (Ridge)\")\n",
    "ax2.set_ylim(bottom=0.4775, top=0.505)\n",
    "fig.tight_layout()\n",
    "fig.savefig(os.path.join(FIG_DIR, \"optimizer_comparison.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=plotting.get_figsize(0.5))\n",
    "\n",
    "ols_adam = optimizers.OLSAdam(learning_rate=0.01, beta1=0.9, beta2=0.999)\n",
    "ridge_adam = optimizers.RidgeAdam(learning_rate=0.01, beta1=0.9, beta2=0.999, lam=0.1)\n",
    "lasso_adam = optimizers.LASSOAdam(learning_rate=0.01, beta1=0.9, beta2=0.999, lam=0.1)\n",
    "ols_gd = optimizers.OLSGradientDescent(learning_rate=0.1)\n",
    "ridge_gd = optimizers.RidgeGradientDescent(learning_rate=0.01, lam=0.1)\n",
    "lasso_gd = optimizers.LASSOGradientDescent(learning_rate=0.01, lam=0.1)\n",
    "\n",
    "for i, (optimizer, label) in enumerate(\n",
    "    zip(\n",
    "        [ols_adam, ridge_adam, lasso_adam, ols_gd, ridge_gd, lasso_gd],\n",
    "        [\"OLS Adam\", \"Ridge Adam\", \"LASSO Adam\", \"OLS GD\", \"Ridge GD\", \"LASSO GD\"],\n",
    "    )\n",
    "):\n",
    "    optimizer.fit(X_tr, y_tr)\n",
    "    cost_history = optimizer.cost_history\n",
    "    ls = \"-\" if i < 3 else \"--\"\n",
    "    c = \"C\" + str(i % 3)\n",
    "    label = label.split()[0] if i < 3 else \"\"\n",
    "    ax.plot(cost_history, label=label, linestyle=ls, color=c)\n",
    "ax.set_xlabel(\"Iteration\")\n",
    "ax.set_ylabel(\"Cost\")\n",
    "# Add linestyles manually to legend\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "handles += [\n",
    "    plt.Line2D([0], [0], color=\"k\", linestyle=\"-\"),\n",
    "    plt.Line2D([0], [0], color=\"k\", linestyle=\"--\"),\n",
    "]\n",
    "labels += [\"Adam Opt.\", \"Grad. Desc.\"]\n",
    "ax.legend(handles, labels)\n",
    "fig.tight_layout()\n",
    "fig.savefig(os.path.join(FIG_DIR, \"cost_function_comparison.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "project1",
   "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.13.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
