{
 "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 pyoptim.optimizers as optimizers\n",
    "import pyoptim.datamanip as datamanip\n",
    "import pyoptim.plotting as 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(\n",
    "    x, y, test_size=0.2, random_state=datamanip.get_RNG().integers(0, 1e6)\n",
    ")\n",
    "\n",
    "fig, ax = plotting.scatter_dataset(x_train, x_test, y_train, y_test)\n",
    "fig.set_layout_engine(\"compressed\")\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.set_layout_engine(\"compressed\")\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.set_layout_engine(\"compressed\")\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.set_layout_engine(\"compressed\")\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.set_layout_engine(\"compressed\")\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",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig(os.path.join(FIG_DIR, \"ridge_parameter_plot.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "lambda_values = np.logspace(-5, 3, 60)\n",
    "polynomial_degrees = np.arange(3, 20, dtype=int)\n",
    "\n",
    "test_mse_ridge_list = np.zeros((len(polynomial_degrees), len(lambda_values)))\n",
    "\n",
    "for i, polynomial_degree in enumerate(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 j, lambda_ in enumerate(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",
    "        test_mse_ridge_list[i, j] = mse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=plotting.get_figsize(0.5))\n",
    "c = ax.pcolormesh(\n",
    "    lambda_values,\n",
    "    polynomial_degrees,\n",
    "    test_mse_ridge_list,\n",
    "    shading=\"auto\",\n",
    "    cmap=\"inferno\",\n",
    ")\n",
    "fig.colorbar(c, ax=ax, label=\"Test MSE\")\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_xlabel(\"$\\\\lambda$\")\n",
    "ax.set_ylabel(\"Polynomial Degree\")\n",
    "fig.savefig(os.path.join(FIG_DIR, \"ridge_mse_heatmap.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "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(2, 1, figsize=plotting.get_figsize(1.0))\n",
    "num_iters = 1_000\n",
    "\n",
    "learning_rates = np.logspace(-4, 0, 5)\n",
    "ols_optimizers = [\n",
    "    (\n",
    "        optimizers.OLSGradientDescent,\n",
    "        {\"learning_rate\": learning_rate, \"num_iterations\": num_iters},\n",
    "    )\n",
    "    for learning_rate in learning_rates\n",
    "]\n",
    "ridge_optimizers = [\n",
    "    (\n",
    "        optimizers.RidgeGradientDescent,\n",
    "        {\"learning_rate\": learning_rate, \"num_iterations\": num_iters, \"lam\": 0.1},\n",
    "    )\n",
    "    for learning_rate in learning_rates\n",
    "]\n",
    "labels = [\n",
    "    f\"$\\\\eta = 10^{{{int(np.log10(learning_rate))}}}$\"\n",
    "    for learning_rate in learning_rates\n",
    "]\n",
    "\n",
    "plotting.plot_optimizers(ax1, ols_optimizers, X_tr, y_tr, \"OLS Cost\", labels=labels)\n",
    "plotting.plot_optimizers(ax2, ridge_optimizers, X_tr, y_tr, \"Ridge Cost\", labels=labels)\n",
    "ax1.set_ylim(bottom=0.465, top=0.505)\n",
    "ax2.set_ylim(bottom=0.4775, top=0.505)\n",
    "fig.set_layout_engine(\"compressed\")\n",
    "fig.savefig(os.path.join(FIG_DIR, \"gradient_descent_convergence.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=plotting.get_figsize(1.0))\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",
    "general_kwargs = {\n",
    "    \"num_iterations\": num_iters,\n",
    "}\n",
    "optimizer_kwargs = [\n",
    "    {},\n",
    "    {\"delta\": 0.9},\n",
    "    {},\n",
    "    {\"gamma\": 0.995},\n",
    "    {\"beta1\": 0.9, \"beta2\": 0.999},\n",
    "]\n",
    "ols_kwargs = {\"learning_rate\": learning_rate_ols, **general_kwargs}\n",
    "ridge_kwargs = {\n",
    "    \"learning_rate\": learning_rate_ridge,\n",
    "    \"lam\": lam,\n",
    "    **general_kwargs,\n",
    "}\n",
    "optimizers_ols = [\n",
    "    (opt, {**ols_kwargs, **opt_kwargs})\n",
    "    for opt, opt_kwargs in zip(optimizers.OLS_GD_OPTIMIZERS, optimizer_kwargs)\n",
    "]\n",
    "optimizers_ridge = [\n",
    "    (opt, {**ridge_kwargs, **opt_kwargs})\n",
    "    for opt, opt_kwargs in zip(optimizers.RIDGE_GD_OPTIMIZERS, optimizer_kwargs)\n",
    "]\n",
    "\n",
    "\n",
    "plotting.plot_optimizers(ax1, optimizers_ols, X_tr, y_tr, ylabel=\"Cost (OLS)\")\n",
    "plotting.plot_optimizers(ax2, optimizers_ridge, X_tr, y_tr, ylabel=\"Cost (Ridge)\")\n",
    "\n",
    "ax1.set_ylim(bottom=0.465, top=0.505)\n",
    "ax2.set_ylim(bottom=0.4775, top=0.505)\n",
    "fig.set_layout_engine(\"compressed\")\n",
    "fig.savefig(os.path.join(FIG_DIR, \"optimizer_comparison.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=plotting.get_figsize(0.6))\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.set_layout_engine(\"compressed\")\n",
    "fig.savefig(os.path.join(FIG_DIR, \"cost_function_comparison.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_size = 1_000_000\n",
    "num_epochs = 1000\n",
    "batches_per_epoch = 100\n",
    "batch_size = 512\n",
    "\n",
    "ols_gd = optimizers.OLSGradientDescent(learning_rate=0.1, num_iterations=num_epochs)\n",
    "ols_sgd = optimizers.OLSStochasticGradientDescent(\n",
    "    learning_rate=0.1,\n",
    "    num_iterations=num_epochs * batches_per_epoch,\n",
    "    batch_size=batch_size,\n",
    "    batches_per_epoch=batches_per_epoch,\n",
    ")\n",
    "ridge_gd = optimizers.RidgeGradientDescent(\n",
    "    learning_rate=0.01, lam=0.1, num_iterations=num_epochs\n",
    ")\n",
    "ridge_sgd = optimizers.RidgeStochasticGradientDescent(\n",
    "    learning_rate=0.01,\n",
    "    lam=0.1,\n",
    "    num_iterations=num_epochs * batches_per_epoch,\n",
    "    batch_size=batch_size,\n",
    "    batches_per_epoch=batches_per_epoch,\n",
    ")\n",
    "\n",
    "x = np.linspace(-1, 1, X_size)\n",
    "y = datamanip.noise_data(datamanip.runge_function(x), 0.1)  # LOWER NOISE FOR SGD\n",
    "x_train, x_test, y_train, y_test = train_test_split(\n",
    "    x, y, test_size=0.2, random_state=datamanip.get_RNG().integers(0, 1e6)\n",
    ")\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",
    "optimizer_list = [ols_sgd, ols_gd, ridge_sgd, ridge_gd]\n",
    "labels = [\n",
    "    \"Stochastic GD\",\n",
    "    \"Gradient Descent\",\n",
    "    \"Ridge Stochastic GD\",\n",
    "    \"Ridge Gradient Descent\",\n",
    "]\n",
    "fig, ax = plotting.optimization_performance_evaluation(\n",
    "    optimizer_list, labels, X_tr, y_tr\n",
    ")\n",
    "fig.set_layout_engine(\"compressed\")\n",
    "fig.savefig(os.path.join(FIG_DIR, \"optimization_performance.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_size = 1_000_000\n",
    "num_epochs = 500\n",
    "batches_per_epoch = 100\n",
    "batch_size = 128\n",
    "\n",
    "learning_rate_ols = 0.1\n",
    "learning_rate_ridge = 0.001\n",
    "learning_rate_lasso = 0.001\n",
    "lam_ridge = 0.1\n",
    "lam_lasso = 0.1\n",
    "\n",
    "x = np.linspace(-1, 1, X_size)\n",
    "y = datamanip.noise_data(datamanip.runge_function(x), 0.1)  # LOWER NOISE FOR SGD\n",
    "x_train, x_test, y_train, y_test = train_test_split(\n",
    "    x, y, test_size=0.2, random_state=datamanip.get_RNG().integers(0, 1e6)\n",
    ")\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",
    "general_kwargs = {\n",
    "    \"num_iterations\": num_epochs * batches_per_epoch,\n",
    "    \"batch_size\": batch_size,\n",
    "    \"batches_per_epoch\": batches_per_epoch,\n",
    "}\n",
    "optimizer_kwargs = [\n",
    "    {},\n",
    "    {\"delta\": 0.9},\n",
    "    {},\n",
    "    {\"gamma\": 0.9},\n",
    "    {\"beta1\": 0.9, \"beta2\": 0.999},\n",
    "]\n",
    "ols_kwargs = {\"learning_rate\": learning_rate_ols, **general_kwargs}\n",
    "ridge_kwargs = {\n",
    "    \"learning_rate\": learning_rate_ridge,\n",
    "    \"lam\": lam_ridge,\n",
    "    **general_kwargs,\n",
    "}\n",
    "lasso_kwargs = {\n",
    "    \"learning_rate\": learning_rate_lasso,\n",
    "    \"lam\": lam_lasso,\n",
    "    **general_kwargs,\n",
    "}\n",
    "optimizers_ols = [\n",
    "    (opt, {**kwargs, **ols_kwargs})\n",
    "    for opt, kwargs in zip(\n",
    "        optimizers.OLS_SGD_OPTIMIZERS,\n",
    "        optimizer_kwargs,\n",
    "    )\n",
    "]\n",
    "optimizers_ridge = [\n",
    "    (opt, {**kwargs, **ridge_kwargs})\n",
    "    for opt, kwargs in zip(\n",
    "        optimizers.RIDGE_SGD_OPTIMIZERS,\n",
    "        optimizer_kwargs,\n",
    "    )\n",
    "]\n",
    "optimizers_lasso = [\n",
    "    (opt, {**kwargs, **lasso_kwargs})\n",
    "    for opt, kwargs in zip(\n",
    "        optimizers.LASSO_SGD_OPTIMIZERS,\n",
    "        optimizer_kwargs,\n",
    "    )\n",
    "]\n",
    "\n",
    "labels = [\"SGD\", \"Mom. SGD\", \"AdaGrad\", \"RMSProp\", \"Adam\"]\n",
    "\n",
    "fig, axs = plt.subplots(\n",
    "    3, 1, figsize=plotting.get_figsize(1.0, full_width=False), sharex=True\n",
    ")\n",
    "plotting.plot_optimizers(\n",
    "    axs[0], optimizers_ols, X_tr, y_tr, ylabel=\"OLS Cost\", labels=labels\n",
    ")\n",
    "plotting.plot_optimizers(\n",
    "    axs[1], optimizers_ridge, X_tr, y_tr, ylabel=\"Ridge Cost\", labels=labels\n",
    ")\n",
    "plotting.plot_optimizers(\n",
    "    axs[2], optimizers_lasso, X_tr, y_tr, ylabel=\"LASSO Cost\", labels=labels\n",
    ")\n",
    "for ax in axs:\n",
    "    ax.set_xlabel(None)\n",
    "\n",
    "axs[2].set_xlabel(\"Epoch\")\n",
    "\n",
    "# axs[0].set_title(\"OLS\")\n",
    "# axs[1].set_title(\"Ridge\")\n",
    "# axs[2].set_title(\"LASSO\")\n",
    "fig.set_layout_engine(\"compressed\")\n",
    "fig.savefig(os.path.join(FIG_DIR, \"stochastic_gradient_descent_convergence.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(-1, 1, 100)\n",
    "y = datamanip.runge_function(x)\n",
    "x_train, x_test, y_train, y_test_noise_free = train_test_split(\n",
    "    x, y, test_size=0.2, random_state=datamanip.get_RNG().integers(0, 1e6)\n",
    ")\n",
    "y_train = datamanip.noise_data(y_train, 1.0)\n",
    "y_test = datamanip.noise_data(y_test_noise_free, 1.0)\n",
    "\n",
    "polynomial_degrees = np.arange(1, 15)\n",
    "n_bootstraps = 5000\n",
    "mses = np.zeros((len(polynomial_degrees), n_bootstraps))\n",
    "biases = np.zeros((len(polynomial_degrees), n_bootstraps))\n",
    "variances = np.zeros((len(polynomial_degrees), n_bootstraps))\n",
    "\n",
    "for i, polynomial_degree in enumerate(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_nf = datamanip.scale_data(y_train, y_test_noise_free)\n",
    "    y_train_scaled, y_test_scaled = datamanip.scale_data(y_train, y_test)\n",
    "\n",
    "    for b, (X_, y_) in enumerate(\n",
    "        datamanip.bootstrap_resample(X_train_scaled, y_train_scaled, n_bootstraps)\n",
    "    ):\n",
    "        beta = optimizers.Ridge_parameters(X_, y_, lam=1e-10)  # approx OLS but stable\n",
    "        y_pred = X_test_scaled @ beta\n",
    "        mse, _ = datamanip.evaluate_model(y_test_scaled, y_pred)\n",
    "        mses[i, b] = mse\n",
    "        biases[i, b] = np.mean((y_test_scaled_nf - np.mean(y_pred)) ** 2)\n",
    "        variances[i, b] = np.var(y_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=plotting.get_figsize(0.5))\n",
    "mse_mean = np.mean(mses, axis=1)\n",
    "bias_mean = np.mean(biases, axis=1)\n",
    "var_mean = np.mean(variances, axis=1)\n",
    "\n",
    "ax.plot(polynomial_degrees, bias_mean, label=\"Bias$^2$\", color=\"C2\", ls=\"-\")\n",
    "ax.fill_between(\n",
    "    polynomial_degrees,\n",
    "    np.zeros_like(bias_mean),\n",
    "    bias_mean,\n",
    "    color=\"C2\",\n",
    "    alpha=0.3,\n",
    ")\n",
    "ax.plot(\n",
    "    polynomial_degrees,\n",
    "    var_mean + bias_mean,\n",
    "    label=\"Variance + Bias$^2$\",\n",
    "    color=\"C1\",\n",
    "    ls=\"-\",\n",
    ")\n",
    "ax.fill_between(\n",
    "    polynomial_degrees,\n",
    "    bias_mean,\n",
    "    bias_mean + var_mean,\n",
    "    color=\"C1\",\n",
    "    alpha=0.3,\n",
    ")\n",
    "ax.fill_between(\n",
    "    polynomial_degrees,\n",
    "    bias_mean + var_mean,\n",
    "    mse_mean,\n",
    "    color=\"C0\",\n",
    "    alpha=0.3,\n",
    "    label=\"Irreducible Error\",\n",
    "    hatch=\"//\",\n",
    "    edgecolor=\"white\",\n",
    ")\n",
    "ax.plot(polynomial_degrees, mse_mean, label=\"MSE\", color=\"C0\", ls=\"-\")\n",
    "\n",
    "\n",
    "ax.set_xlabel(\"Polynomial Degree\")\n",
    "ax.set_ylabel(\"Error\")\n",
    "ax.legend()\n",
    "\n",
    "fig.set_layout_engine(\"compressed\")\n",
    "fig.savefig(os.path.join(FIG_DIR, \"bias_variance_tradeoff.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "k_folds = 5\n",
    "k_fold_mses = np.zeros((len(polynomial_degrees), k_folds))\n",
    "splits = datamanip.k_fold_split(x, datamanip.noise_data(y), k_folds)\n",
    "\n",
    "for i, polynomial_degree in enumerate(polynomial_degrees):\n",
    "    for k, (x_tr, y_tr, x_val, y_val) in enumerate(splits):\n",
    "        X_train = datamanip.polynomial_features(x_tr, polynomial_degree, False)\n",
    "        X_val = datamanip.polynomial_features(x_val, polynomial_degree, False)\n",
    "        X_train_scaled, X_val_scaled = datamanip.scale_data(X_train, X_val)\n",
    "        y_train_scaled, y_val_scaled = datamanip.scale_data(y_tr, y_val)\n",
    "\n",
    "        beta = optimizers.Ridge_parameters(X_train_scaled, y_train_scaled, lam=1e-10)\n",
    "        y_pred = X_val_scaled @ beta\n",
    "        mse, _ = datamanip.evaluate_model(y_val_scaled, y_pred)\n",
    "        k_fold_mses[i, k] = mse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "SINGULAR_PLOT = True\n",
    "if not SINGULAR_PLOT:\n",
    "    fig, (ax, ax2) = plt.subplots(\n",
    "        1, 2, figsize=plotting.get_figsize(0.35, True), sharey=True\n",
    "    )\n",
    "else:\n",
    "    fig, ax = plt.subplots(1, 1, figsize=plotting.get_figsize(0.5, False))\n",
    "    ax2 = ax\n",
    "mse_mean = np.mean(mses, axis=1)\n",
    "bias_mean = np.mean(biases, axis=1)\n",
    "var_mean = np.mean(variances, axis=1)\n",
    "\n",
    "ax.plot(polynomial_degrees, bias_mean, label=\"Bias$^2$\", color=\"C2\", ls=\"-\")\n",
    "ax.fill_between(\n",
    "    polynomial_degrees,\n",
    "    np.zeros_like(bias_mean),\n",
    "    bias_mean,\n",
    "    color=\"C2\",\n",
    "    alpha=0.3,\n",
    ")\n",
    "ax.plot(\n",
    "    polynomial_degrees,\n",
    "    var_mean + bias_mean,\n",
    "    label=\"Variance + Bias$^2$\",\n",
    "    color=\"C1\",\n",
    "    ls=\"-\",\n",
    ")\n",
    "ax.fill_between(\n",
    "    polynomial_degrees,\n",
    "    bias_mean,\n",
    "    bias_mean + var_mean,\n",
    "    color=\"C1\",\n",
    "    alpha=0.3,\n",
    ")\n",
    "ax.fill_between(\n",
    "    polynomial_degrees,\n",
    "    bias_mean + var_mean,\n",
    "    mse_mean,\n",
    "    color=\"C0\",\n",
    "    alpha=0.3,\n",
    "    label=\"Irreducible Error\",\n",
    "    hatch=\"//\",\n",
    "    edgecolor=\"white\",\n",
    ")\n",
    "ax.plot(polynomial_degrees, mse_mean, label=\"MSE\", color=\"C0\", ls=\"-\")\n",
    "\n",
    "\n",
    "ax.set_xlabel(\"Polynomial Degree\")\n",
    "ax.set_ylabel(\"Error\")\n",
    "ax.legend()\n",
    "\n",
    "fig.tight_layout()\n",
    "\n",
    "if not SINGULAR_PLOT:\n",
    "    ax2.plot(polynomial_degrees, mse_mean, label=\"Bootstrapping MSE\", color=\"C1\")\n",
    "ax2.plot(\n",
    "    polynomial_degrees,\n",
    "    np.mean(k_fold_mses, axis=1),\n",
    "    label=f\"{k_folds}-Fold Crossvalidation MSE\",\n",
    "    color=\"black\",\n",
    "    ls=\"--\",\n",
    ")\n",
    "\n",
    "\n",
    "ax2.set_xlabel(\"Polynomial Degree\")\n",
    "# ax2.set_ylabel(\"Mean Squared Error\")\n",
    "ax2.legend()\n",
    "fig.set_layout_engine(\"compressed\")\n",
    "fig.savefig(os.path.join(FIG_DIR, \"bias_variance_tradeoff_combined_plot.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "polynomial_degrees = np.arange(1, 25)\n",
    "k_folds = 5\n",
    "\n",
    "k_fold_mses_ols = np.zeros((len(polynomial_degrees), k_folds))\n",
    "k_fold_mses_ridge = np.zeros_like(k_fold_mses_ols)\n",
    "k_fold_mses_lasso = np.zeros_like(k_fold_mses_ols)\n",
    "\n",
    "splits = datamanip.k_fold_split(x, datamanip.noise_data(y), k_folds)\n",
    "\n",
    "for i, polynomial_degree in enumerate(polynomial_degrees):\n",
    "    for k, (x_tr, y_tr, x_val, y_val) in enumerate(splits):\n",
    "        X_train = datamanip.polynomial_features(x_tr, polynomial_degree, False)\n",
    "        X_val = datamanip.polynomial_features(x_val, polynomial_degree, False)\n",
    "        X_train_scaled, X_val_scaled = datamanip.scale_data(X_train, X_val)\n",
    "        y_train_scaled, y_val_scaled = datamanip.scale_data(y_tr, y_val)\n",
    "\n",
    "        for optim, results_array in zip(\n",
    "            [\n",
    "                optimizers.OLSGradientDescent,\n",
    "                optimizers.RidgeGradientDescent,\n",
    "                optimizers.LASSOGradientDescent,\n",
    "            ],\n",
    "            [k_fold_mses_ols, k_fold_mses_ridge, k_fold_mses_lasso],\n",
    "        ):\n",
    "            Optimizer = optim(num_iterations=1000, learning_rate=0.1, lam=0.1)\n",
    "            theta = Optimizer.fit(X_train_scaled, y_train_scaled)\n",
    "            y_pred = X_val_scaled @ theta\n",
    "            mse, _ = datamanip.evaluate_model(y_val_scaled, y_pred)\n",
    "            results_array[i, k] = mse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "FILL_BETWEEN = False\n",
    "fig, ax = plt.subplots(figsize=plotting.get_figsize(0.5))\n",
    "\n",
    "for i, (mse_array, label) in enumerate(\n",
    "    zip(\n",
    "        [k_fold_mses_ols, k_fold_mses_ridge, k_fold_mses_lasso],\n",
    "        [\"OLS\", \"Ridge Regression\", \"Lasso Regression\"],\n",
    "    )\n",
    "):\n",
    "    c = f\"C{i}\"\n",
    "    ax.plot(polynomial_degrees, np.mean(mse_array, axis=1), label=label, c=c)\n",
    "    if FILL_BETWEEN:\n",
    "        ax.fill_between(\n",
    "            polynomial_degrees,\n",
    "            np.mean(mse_array, axis=1) - np.std(mse_array, axis=1),\n",
    "            np.mean(mse_array, axis=1) + np.std(mse_array, axis=1),\n",
    "            color=c,\n",
    "            alpha=0.3,\n",
    "        )\n",
    "\n",
    "ax.set_xlabel(\"Polynomial Degree\")\n",
    "ax.set_ylabel(f\"MSE ({k_folds}-fold)\")\n",
    "ax.legend()\n",
    "fig.set_layout_engine(\"compressed\")\n",
    "fig.savefig(os.path.join(FIG_DIR, \"kfold_mse_comparison_per_cost_function.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "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
}
