{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1bba0128",
   "metadata": {},
   "source": [
    "# Supervised regression\n",
    "\n",
    "The goal is to build and optimize neural networks with PyTorch with the objective of fitting a known feature in the training dataset.\n",
    "\n",
    "The fit is achieved by minimizing the mean-squared-error between the prediction and the target features. As noticed during the presentation, by minimizing the mean-squared-error, we assume that the underlying error distribution is Gaussian. One could use the mean absolute error instead when assuming the distribution is Laplacian.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "23feddde",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple\n",
    "\n",
    "# import standard PyTorch modules\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import axes3d\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb1286f0",
   "metadata": {},
   "source": [
    "We start by generating some fake dataset, which is simple enough that we can visualize the results easily. For this reason, the dataset will contain only one independent variable.\n",
    "\n",
    "The simulated example data will be $f(x) = (3 + \\kappa) x^2 + \\epsilon$, where $\\epsilon \\sim \\mathcal{N}(\\mu=0, \\sigma=10)$ and $\\kappa \\sim \\mathcal{N}(\\mu=0, \\sigma=0.03)$.\n",
    "\n",
    "In this case we do know the true model, so it is interesting to take some time to pinpoint the role of $\\kappa$ and $\\epsilon$. These variables add fluctuation to the results. $\\epsilon$ adds Gaussian noise in a way that is completely independent from $x$ and cannot be traced down to a particular functional dependence. $\\kappa$ changes a specific parameter of the model, in this case the coefficient 3, by around 1%.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5d457cd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_data(N: int) -> np.ndarray:\n",
    "    x = 2*np.random.randn(N, 1)\n",
    "    epsilon = 10*np.random.randn(N, 1)\n",
    "    kappa = 0.03*np.random.randn(N, 1)\n",
    "    z = (3 + kappa)*x**2 + epsilon\n",
    "    return np.concatenate((x, z), axis=1).astype(np.float32)\n",
    "\n",
    "train_data = generate_data(N=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48433f6f",
   "metadata": {},
   "source": [
    "PyTorch allows you to create a class that outputs a single data entry and use that to feed input to your neural network. We will use the following class to feed the data to the neural network. This looks useless if all your data fits in a Numpy array, but notice that if you have a lot of data and cannot load it all in memory, this allows you to read data on demand, as you need it and only the needed samples are stored at a single time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "30205402",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, data: np.ndarray):\n",
    "        self.data = data\n",
    "    def __len__(self) -> int:\n",
    "        \"\"\"How many samples do I have?\"\"\"\n",
    "        return len(self.data)\n",
    "    def __getitem__(self, idx):\n",
    "        # give me item with index idx\n",
    "        return {\"data\": self.data[idx, 0:1], \"target\": self.data[idx, 1:]}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cc0b0774",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1000\n"
     ]
    }
   ],
   "source": [
    "my_dataset = MyDataset(train_data)\n",
    "print(len(my_dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6dccfac6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'data': array([0.24012195], dtype=float32), 'target': array([14.341419], dtype=float32)}\n"
     ]
    }
   ],
   "source": [
    "print(my_dataset[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "527089bd",
   "metadata": {},
   "source": [
    "Plot some of the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "067b8105",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 30032 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "ax.scatter(train_data[:, 0], train_data[:, 1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e517975c",
   "metadata": {},
   "source": [
    "And now let us define the neural network. In PyTorch, neural networks always extend `nn.Module`. They define their sub-parts in their constructor, which are convolutional layers and fully connected linear layers in this case, and the method `forward` is expected to receive an input image and output the network target.\n",
    "\n",
    "The network parameters are the weights of the `Linear` layers, which are conveniently hidden here, but can be accessed if you try to access their `weights` elements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d908ef86",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Network(nn.Module):\n",
    "    \"\"\"\n",
    "        This is our parametrized function.\n",
    "        It stores all the parametrized weights theta inside the model object.\n",
    "        For such a simple example data, it was not necessary to have such a complex model:\n",
    "        this was only done here to show the interface provided by PyTorch.\n",
    "        The forward function receives the x values and outputs an estimate of the target.\n",
    "        \n",
    "        The nn.Sequential object allows one to apply each step in the given list of parameter\n",
    "        in a sequential way. An alternative would be to create each of these layers manually\n",
    "        and apply them one after the other in the forward method.\n",
    "    \"\"\"\n",
    "    def __init__(self, input_dimension: int=1, output_dimension: int=1):\n",
    "        \"\"\"\n",
    "        Constructor. Here we initialize the weights.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "\n",
    "        hidden_layer = 100\n",
    "        self.model = nn.Sequential(\n",
    "                                   nn.Linear(input_dimension, hidden_layer),\n",
    "                                   nn.ReLU(),\n",
    "                                   nn.Linear(hidden_layer, hidden_layer),\n",
    "                                   nn.ReLU(),\n",
    "                                   nn.Linear(hidden_layer, output_dimension)\n",
    "                                    )\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        This function is called when one does my_network(x) and it represents the action\n",
    "        of our parametrized function in the input.\n",
    "        \"\"\"\n",
    "        return self.model(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c5620dc",
   "metadata": {},
   "source": [
    "Let us create one instance of this network. We also create an instance of PyTorch's `DataLoader`, which has the task of taking a given number of data elements and outputing it in a single object. This \"mini-batch\" of data is used during training, so that we do not need to load the entire data in memory during the optimization procedure.\n",
    "\n",
    "We also create an instance of the Adam optimizer, which is used to tune the parameters of the network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "988e1979",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/daniloefl/miniconda3/envs/ml/lib/python3.12/site-packages/torch/cuda/__init__.py:619: UserWarning: Can't initialize NVML\n",
      "  warnings.warn(\"Can't initialize NVML\")\n"
     ]
    }
   ],
   "source": [
    "network = Network()\n",
    "B = 10\n",
    "loader = torch.utils.data.DataLoader(my_dataset, batch_size=B)\n",
    "optimizer = torch.optim.Adam(network.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ee54520",
   "metadata": {},
   "source": [
    "Now we actually repeatedly try to optimize the network parameters. Each time we go through all the data we have, we go through one \"epoch\". For each epoch, we take several \"mini-batches\" of data (given by the `DataLoader` in `loader`) and use it to make one training step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d15d655d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0/100: average loss 292.65962\n",
      "Epoch 1/100: average loss 168.60480\n",
      "Epoch 2/100: average loss 144.85957\n",
      "Epoch 3/100: average loss 131.24882\n",
      "Epoch 4/100: average loss 125.56896\n",
      "Epoch 5/100: average loss 122.62807\n",
      "Epoch 6/100: average loss 120.44151\n",
      "Epoch 7/100: average loss 118.44099\n",
      "Epoch 8/100: average loss 116.76986\n",
      "Epoch 9/100: average loss 115.28613\n",
      "Epoch 10/100: average loss 114.02854\n",
      "Epoch 11/100: average loss 112.92324\n",
      "Epoch 12/100: average loss 112.03872\n",
      "Epoch 13/100: average loss 111.31309\n",
      "Epoch 14/100: average loss 110.73915\n",
      "Epoch 15/100: average loss 110.27166\n",
      "Epoch 16/100: average loss 109.91651\n",
      "Epoch 17/100: average loss 109.58956\n",
      "Epoch 18/100: average loss 109.31038\n",
      "Epoch 19/100: average loss 109.14096\n",
      "Epoch 20/100: average loss 108.96208\n",
      "Epoch 21/100: average loss 108.82673\n",
      "Epoch 22/100: average loss 108.67057\n",
      "Epoch 23/100: average loss 108.56558\n",
      "Epoch 24/100: average loss 108.45631\n",
      "Epoch 25/100: average loss 108.37027\n",
      "Epoch 26/100: average loss 108.29927\n",
      "Epoch 27/100: average loss 108.26275\n",
      "Epoch 28/100: average loss 108.19757\n",
      "Epoch 29/100: average loss 108.11492\n",
      "Epoch 30/100: average loss 108.07109\n",
      "Epoch 31/100: average loss 108.04241\n",
      "Epoch 32/100: average loss 108.01575\n",
      "Epoch 33/100: average loss 107.99630\n",
      "Epoch 34/100: average loss 107.98242\n",
      "Epoch 35/100: average loss 107.95160\n",
      "Epoch 36/100: average loss 107.93830\n",
      "Epoch 37/100: average loss 107.89870\n",
      "Epoch 38/100: average loss 107.88430\n",
      "Epoch 39/100: average loss 107.87700\n",
      "Epoch 40/100: average loss 107.84878\n",
      "Epoch 41/100: average loss 107.83601\n",
      "Epoch 42/100: average loss 107.80727\n",
      "Epoch 43/100: average loss 107.80601\n",
      "Epoch 44/100: average loss 107.79060\n",
      "Epoch 45/100: average loss 107.77407\n",
      "Epoch 46/100: average loss 107.76293\n",
      "Epoch 47/100: average loss 107.71370\n",
      "Epoch 48/100: average loss 107.70025\n",
      "Epoch 49/100: average loss 107.67249\n",
      "Epoch 50/100: average loss 107.63180\n",
      "Epoch 51/100: average loss 107.60612\n",
      "Epoch 52/100: average loss 107.56095\n",
      "Epoch 53/100: average loss 107.57543\n",
      "Epoch 54/100: average loss 107.55797\n",
      "Epoch 55/100: average loss 107.54764\n",
      "Epoch 56/100: average loss 107.52183\n",
      "Epoch 57/100: average loss 107.49846\n",
      "Epoch 58/100: average loss 107.51441\n",
      "Epoch 59/100: average loss 107.48613\n",
      "Epoch 60/100: average loss 107.48426\n",
      "Epoch 61/100: average loss 107.45074\n",
      "Epoch 62/100: average loss 107.46100\n",
      "Epoch 63/100: average loss 107.44333\n",
      "Epoch 64/100: average loss 107.44383\n",
      "Epoch 65/100: average loss 107.40632\n",
      "Epoch 66/100: average loss 107.40891\n",
      "Epoch 67/100: average loss 107.40680\n",
      "Epoch 68/100: average loss 107.38369\n",
      "Epoch 69/100: average loss 107.37979\n",
      "Epoch 70/100: average loss 107.36220\n",
      "Epoch 71/100: average loss 107.37177\n",
      "Epoch 72/100: average loss 107.35137\n",
      "Epoch 73/100: average loss 107.34140\n",
      "Epoch 74/100: average loss 107.33566\n",
      "Epoch 75/100: average loss 107.33218\n",
      "Epoch 76/100: average loss 107.31788\n",
      "Epoch 77/100: average loss 107.31176\n",
      "Epoch 78/100: average loss 107.29611\n",
      "Epoch 79/100: average loss 107.29593\n",
      "Epoch 80/100: average loss 107.29700\n",
      "Epoch 81/100: average loss 107.28151\n",
      "Epoch 82/100: average loss 107.26731\n",
      "Epoch 83/100: average loss 107.26950\n",
      "Epoch 84/100: average loss 107.25340\n",
      "Epoch 85/100: average loss 107.25051\n",
      "Epoch 86/100: average loss 107.24163\n",
      "Epoch 87/100: average loss 107.24350\n",
      "Epoch 88/100: average loss 107.23004\n",
      "Epoch 89/100: average loss 107.23470\n",
      "Epoch 90/100: average loss 107.21753\n",
      "Epoch 91/100: average loss 107.22043\n",
      "Epoch 92/100: average loss 107.20536\n",
      "Epoch 93/100: average loss 107.20931\n",
      "Epoch 94/100: average loss 107.18696\n",
      "Epoch 95/100: average loss 107.20320\n",
      "Epoch 96/100: average loss 107.18551\n",
      "Epoch 97/100: average loss 107.19133\n",
      "Epoch 98/100: average loss 107.18022\n",
      "Epoch 99/100: average loss 107.16567\n"
     ]
    }
   ],
   "source": [
    "epochs = 100\n",
    "# for each epoch\n",
    "for epoch in range(epochs):\n",
    "    losses = list()\n",
    "    # for each mini-batch given by the loader:\n",
    "    for batch in loader:\n",
    "        # get the input in the mini-batch\n",
    "        # this has size (B, C)\n",
    "        # where B is the mini-batch size\n",
    "        # C is the number of features (1 in this case)\n",
    "        features = batch[\"data\"]\n",
    "        # get the targets in the mini-batch (there shall be B of them)\n",
    "        target = batch[\"target\"]\n",
    "        # get the output of the neural network:\n",
    "        prediction = network(features)\n",
    "        \n",
    "        # calculate the loss function being minimized\n",
    "        # in this case, it is the mean-squared error between the prediction and the target values\n",
    "        loss = F.mse_loss(prediction, target)\n",
    "        # exactly equivalent to:\n",
    "        #loss = ((prediction - target)**2).mean()\n",
    "\n",
    "        # clean the optimizer temporary gradient storage\n",
    "        optimizer.zero_grad()\n",
    "        # calculate the gradient of the loss function as a function of the gradients\n",
    "        loss.backward()\n",
    "        # ask the Adam optimizer to change the parameters in the direction of - gradient\n",
    "        # Adam scales the gradient by a constant which is adaptively tuned\n",
    "        # take a look at the Adam paper for more details: https://arxiv.org/abs/1412.6980\n",
    "        optimizer.step()\n",
    "        losses.append(loss.detach().cpu().item())\n",
    "    avg_loss = np.mean(np.array(losses))\n",
    "    print(f\"Epoch {epoch}/{epochs}: average loss {avg_loss:.5f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4980bf4",
   "metadata": {},
   "source": [
    "Let us check what the network says about some new data it has never seen before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "09646d29",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data = generate_data(N=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e315b5dc",
   "metadata": {},
   "source": [
    "And now we can plot again the new images, now showing what the network tells us about it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7a06a4c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "predicted = network(torch.from_numpy(test_data[:,0:1])).detach().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "bab0ce43",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 47296 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "ax.scatter(test_data[:, 0], test_data[:, 1], label=\"Test data\")\n",
    "ax.scatter(test_data[:, 0], predicted, label=\"Predicted\")\n",
    "ax.set(xlabel=\"$x$\", ylabel=\"$f(x)$\")\n",
    "#ax.set_yscale('log')\n",
    "plt.legend(frameon=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce33b86f-2de6-481f-a663-7ccb3949946c",
   "metadata": {},
   "source": [
    "## What is the impact of the loss function?\n",
    "\n",
    "What if instead of the mean squared error, you used the mean absolute error (L1 loss)? How would the result change? What does that mean for the assumptions of the model?\n",
    "\n",
    "Try replacing `F.mse_loss` above with `F.l1_loss` or `F.huber_loss`. Look for the documentation on those losses to understand what they do."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98ed597c-ef43-456a-a651-6660d4f7d62b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "27e5ca02-6df6-4bc9-a72a-6ff38f5d6aeb",
   "metadata": {},
   "source": [
    "## Regularization\n",
    "\n",
    "You may try to apply regularization to the fit by setting the `weight_decay` argument of the optimizer. For example:\n",
    "\n",
    "```\n",
    "optimizer = torch.optim.Adam(network.parameters(), lr=1e-3, weight_decay=1e-3)\n",
    "```\n",
    "\n",
    "This adds an extra term in the loss function proportional to the sum of square of weights, similarly to how it was discussed in the lecture.\n",
    "\n",
    "If one wants to add other types of regularization terms to the loss function, one may also do that explicitly, by changing the way the loss function is calculated.\n",
    "For example, to add an L1 regularization (absolute value of the parameters, which corresponds to the Laplace prior of parameters), one may add the following lines after calculating the loss function:\n",
    "\n",
    "```\n",
    "l1_lambda = 0.001\n",
    "l1_norm = sum(torch.linalg.norm(p, 1) for p in model.parameters())\n",
    "loss = loss + l1_lambda * l1_norm\n",
    "```\n",
    "\n",
    "Try those out!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d90a6c7-60e1-4622-acee-181d41e46b32",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
