{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# HW5 Programming - Modeling Energy Usage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Table of Contents\n",
    "\n",
    "- [**Introduction**](#intro)\n",
    "- [**Q1** - Weighted Linear Regression](#q1)\n",
    "- [**Q2** - Iterative Training Methods](#q2)\n",
    "    - [**Q2a** - Helper Functions for Mini-Batch Stochastic Gradient Descent](#q2a)\n",
    "    - [**Q2b** - Implementing Mini-Batch Stochastic Gradient Descent for Weighted Linear Regression](#q2b)\n",
    "    - [**Q2c** - Programming Written Plotting](#q2c) (Writeup)\n",
    "    - [**Q2d** - Gradient Descent Plot Analysis](#q2d) (Writeup)\n",
    "    - [**Q2e & Q2f** - Big-O Complexity Questions](#q2e) (Writeup)\n",
    "- [**Q3** - Nonlinear Modeling](#q3)\n",
    "    - [**Q3a** - Polynomial Features](#q3a)\n",
    "    - [**Q3b** - K-Fold Cross Validation](#q3b)\n",
    "    - [**Q3c** - Programming Written Plot](#q3c) (Writeup)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup <a class=\"anchor\" name=\"setup\"></a>\n",
    "\n",
    "**Note:** running the second cell in this section with the `curl` and `unzip` command might require you to replace some files if you have ran it before, so if that is the case, then click where the cursor is blinking and enter the command to replace all.\n",
    "\n",
    "You'll need to run these cells, but you don't have to worry about their contents. You can look through them if you'd like of course."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install otter-grader if needed\n",
    "\n",
    "import importlib\n",
    "\n",
    "if importlib.util.find_spec(\"otter\") is None:\n",
    "    !pip install otter-grader\n",
    "    \n",
    "if importlib.util.find_spec(\"sklearn\") is None:\n",
    "    !pip install scikit-learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy additional files if needed\n",
    "# Note: csv files can be loaded from the internet while np.load doesn't support reading from the internet\n",
    "\n",
    "import os\n",
    "\n",
    "if not os.path.isdir(\"tests\") and \"data.npz\" not in os.listdir():\n",
    "    !curl https://www.cs.cmu.edu/~07280/assignments/hw5_additional_files.zip --output hw5_additional_files.zip\n",
    "    !unzip hw5_additional_files.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.neighbors import KNeighborsRegressor\n",
    "import re\n",
    "\n",
    "import otter\n",
    "\n",
    "# Bump up the default font size for matplotlib\n",
    "plt.rcParams.update({'font.size': 18})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Revert to older numpy repr strings to make otter-grader happy\n",
    "if repr(np.float64()).startswith('np.float'):\n",
    "     np.set_printoptions(legacy='1.25')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grader = otter.Notebook()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tests\n",
    "\n",
    "The tests that the Otter Grader uses are located in `test_case_code.py`. You can access this file by clicking the file icon on the left and double clicking on the python file (if you are using Google Colab). This file will contain the code for the test cases that the Otter Grader uses. You can copy and paste the code into code blocks and run it to debug your code. For the expected outputs, look at the corresponding Otter Grader tests. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction <a class=\"anchor\" name=\"intro\"></a>\n",
    "\n",
    "You are a recently hired machine learning engineer at Pat's Eclectic Energy Enterprise ⚡. Currently, the company is having trouble determining how much power to produce for the day. If the enterprise produces too much energy, then it goes to waste and serves as extra company cost, and if they do not produce enough energy, then the company receives bad reviews and customers are unhappy.\n",
    "\n",
    "As a result, Pat wanted to determine what is the best way to predict for energy consumption. After consulting his engineers, he realized that the temperature of the day could be a relatively strong indicator for energy consumption. Thus, since May till August, the company has been collecting data on the temperature of the day using a thermometer and the consumed energy for that day using expensive power meters. \n",
    "\n",
    "As the machine learning engineer for the company, Pat would like you to develop a model for the consumption of energy with respect to the temperature feature that they collected data on. A visualization of the data can be seen below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plotTempDataset(t):\n",
    "    '''plotting the weighted linear regression model t=(\"summer\"|\"winter\"|\"full\")'''\n",
    "    data = np.load(\"data.npz\")\n",
    "    X, y, R = data[f\"{t}_X\"], data[f\"{t}_y\"], data[f\"{t}_R\"]\n",
    "    y /= np.max(y)\n",
    "\n",
    "    plt.figure(figsize=(15,10))\n",
    "    plt.scatter(X, y, marker=\"P\")\n",
    "    plt.title(f\"Temperature vs. Usage Data points\")\n",
    "    plt.xlabel(\"Temperature\")\n",
    "    plt.ylabel(\"Usage\")\n",
    "    plt.show()\n",
    "    \n",
    "plotTempDataset(\"summer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Due to the linear relationship in the image, you believe that fitting a linear regression model would fare well. **Unfortunately**, Pat mentions that there is a problem with the dataset. When the day was extremely hot (over 80 degrees Fahrenheit), the power meters that measured the demand for data started malfunctioning. As a result, they produced measurements that are far more noisy than what is typically expected from the sensor.\n",
    "\n",
    "That is an issue for linear regression. Linear regression assumes that your errors/sample noise is uniform across all datapoints, but instead, we have different variances for data points associated with higher temperatures. As a result, from your studies and superb understanding of machine learning theory, you decide that performing weighted linear regression would be a better fit. That way, you can model the change in variance of the dataset that you are trying to predict on. The dataset comes with sample weights where each sample has an associated non-negative real-value number representing how well the data model's the true relationship between temperature and usage. (a higher variance gives the data point a lower weight). Below is a colored visualization of the data showing the data points where the sensor had low variance and the data points where the sensor had high variance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plotTempDatasetWeighted(t):\n",
    "    '''plotting the weighted linear regression model t=(\"summer\"|\"winter\"|\"full\")'''\n",
    "    data = np.load(\"data.npz\")\n",
    "    X, y, R = data[f\"{t}_X\"], data[f\"{t}_y\"], data[f\"{t}_R\"]\n",
    "    \n",
    "    idxs = X >= 80\n",
    "    X_low = X[~idxs]\n",
    "    X_high = X[idxs]\n",
    "    \n",
    "    y_low = y[~idxs]\n",
    "    y_high = y[idxs]\n",
    "    \n",
    "    plt.figure(figsize=(15,10))\n",
    "    plt.scatter(X_low, y_low, marker=\"P\", c=\"blue\", label=\"Low Variance (Weight=1)\")\n",
    "    plt.scatter(X_high, y_high, marker=\"P\", c=\"orange\", label=\"High Variance (Weight=0.25)\")\n",
    "    plt.title(f\"Temperature vs. Usage Data points\")\n",
    "    plt.xlabel(\"Temperature\")\n",
    "    plt.ylabel(\"Usage\")\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "    \n",
    "plotTempDatasetWeighted(\"summer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a result, you decide that you'll attempt to model the energy consumption as a function of a temperature feature using a weighted linear regression model in multiple ways. You'll gain a better understanding of how linear regression models are trained, and by trying different optimization techniques, you'll also develop intuition for various optimization algorithms such as the variants of gradient descent and understand the tradeoffs that come with them."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unfortunately, you don't have access to any libraries for training weighted linear regression models, and as a result, you'll have to implement the model on your own. Below, you've written an abstract class `Model` for defining the operations that are used with a machine learning model, such as `__init__` (or initializing) it with a specific hyperparameter configuration, `fit`-ing (or training) it given a dataset, `predict`-ing it on a new dataset, or `score`-ing it on a new set of samples with corresponding features. (Note: this is a similar interface that the `sklearn` Python machine learning package implements their machine learning models in)\n",
    "\n",
    "Using this abstract class, you'll implement weighted linear regression in multiple ways that will help you gain a better understanding of linear regression models as well as various optimization techniques.\n",
    "\n",
    "**Terminology Alert:** fitting a model to a dataset and training a model on a dataset are the same thing!\n",
    "\n",
    "**Note:** You don't have to implement anything in the `Model` class!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: this is an abstract class that is only meant for you to get an idea for a machine learning model does\n",
    "#       you do not need to implement/modify any functions here!\n",
    "\n",
    "class Model:\n",
    "    def __init__(self, *args):\n",
    "        \"\"\"initialize the machine learning model -- pass arguments that are hyperparameters for training\"\"\"\n",
    "        pass\n",
    "\n",
    "    def fit(X, y, R=None):\n",
    "        \"\"\" train the model on the dataset of samples X and their corresponding labels y\n",
    "            if R given, then samples have weights\n",
    "            Note: fitting is the same thing as training!\n",
    "        \"\"\"\n",
    "        pass\n",
    "\n",
    "    def predict(X):\n",
    "        \"\"\"predict on a new dataset of samples X\"\"\"\n",
    "        pass\n",
    "\n",
    "    def score(X, y):\n",
    "        \"\"\" evaluate how well the model performs on a dataset of samples X with labels y\n",
    "            compare the predicted values of the model with y and evaluate some loss or accuracy metric\n",
    "        \"\"\"\n",
    "        pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "---\n",
    "# Question 1 - Weighted Linear Regression <a class=\"anchor\" name=\"q1\"></a>\n",
    "\n",
    "Recall that you are using weighted linear regression as your model of choice. The reason for this is that we would like to weight the values of the the more-noisy datapoints with less importance (smaller $r_i$). Conversely, we would like to weight the values of less-noisy datapoints with more importance (larger $r_i$) because they have a better indication of what we would like to model.\n",
    "\n",
    "As a result, given a dataset of points (i.e. temperature values) in the form of the design matrix, $X$, which is (N,2) (because one a column of bias values and a column of temperatures), their corresponding targets (i.e. usage values) $\\mathbf{y}$ which is (N,1), and weights for each of the samples $R$ which is (N,N) ($R$ is a __diagonal__ matrix that contains weights for the *i*-th sample at $r_{ii}$), we can write down the weighted least squares loss of the weighted linear regression model as follows.\n",
    "\n",
    "$$J(\\mathbf{w})  = \\frac{1}{N}(\\mathbf{y}-X\\mathbf{w})^{\\top}R(\\mathbf{y}-X\\mathbf{w})$$\n",
    "\n",
    "where $\\mathbf{w}$ denotes the parameters for the linear regression model. You will be using your derivation of the closed-form solution and the gradient of this loss with respect to the weights to implement weighted least squares linear regression in the following iterative and non-iterative methods.\n",
    "\n",
    "1. Closed-Form Solution\n",
    "2. Gradient Descent\n",
    "3. Stochastic Gradient Descent\n",
    "4. Mini-Batch Stochastic Gradient Descent\n",
    "\n",
    "**Note:** for $R$, we are going to give you an `(N,)` element `np.ndarray` which you can then convert to an (N,N) matrix. We recommend looking into `np.diag` as a useful function to do so.\n",
    "\n",
    "**Additional Note:** We will represent the raw dataset of features as `Xraw` while we will represent the design matrix created from the raw features as `X`. Furthermore, for all matrix multiplications throughout this assignment, **do not use** `np.matmul`. Instead, **use** `@`.\n",
    "\n",
    "First, we'll get started with training weighted linear regression through a closed-form solution using the theory for the problem that you solved in the written portion of this assignment. First, in order to help with training linear regression models, implement the helper function below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "otter_answer_cell"
    ]
   },
   "outputs": [],
   "source": [
    "def makeDesignMatrix(Xraw):\n",
    "    \"\"\" Creates the design matrix from matrix Xraw by concatenating a \n",
    "        column of 1's to the left of the array\n",
    "        Used by linear regression models to create a bias feature\n",
    "        \n",
    "        >>> a\n",
    "        array([[ 4.,  7.],\n",
    "               [ 3.,  7.],\n",
    "               [12.,  7.]])\n",
    "        >>> makeDesignMatrix(a)\n",
    "        array([[ 1.,  4.,  7.],\n",
    "               [ 1.,  3.,  7.],\n",
    "               [ 1., 12.,  7.]])\n",
    "               \n",
    "        Input:\n",
    "        Xraw: np.ndarray of shape (N,M)\n",
    "        \n",
    "        Returns:\n",
    "        X: design matrix made from Xraw\n",
    "    \"\"\"\n",
    "    ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "# make sure the outputs are what you expect!\n",
    "Xraw = np.array([[1,2],[3,4],[5,6]])\n",
    "X = makeDesignMatrix(Xraw)\n",
    "X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "Xraw = np.array([[0,0,0,2],[2,3,5,9],[0,3,0,8],[0,3,0,7]])\n",
    "X = makeDesignMatrix(Xraw)\n",
    "X"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "Now, with having implemented the matrix to create our design matrix which adds a 1's feature for creating a bias in our linear regression model, we will now train our weighted linear regression model with the closed form solution."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "### Implement `.fit` and `.predict` of Class `WeightedLinReg_ClosedForm`\n",
    "\n",
    "**Note:** if you are running your code and the test cases are not passing even if you believe that you are passing the test cases, try resetting your notebook and clicking `Run All`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "otter_answer_cell"
    ]
   },
   "outputs": [],
   "source": [
    "class WeightedLinReg_ClosedForm(Model):\n",
    "    # closed form solution to weighted linear regression\n",
    "    def __init__(self):\n",
    "        # Do not need to implement anything here\n",
    "        pass\n",
    "\n",
    "    def fit(self, Xraw, y, R):\n",
    "        \"\"\" Calculates the closed-form solution parameters for weighted linear regression\n",
    "            Note: be aware of the shapes of your input when you are doing NumPy operations!\n",
    "        \n",
    "            Input:\n",
    "            Xraw: np.ndarray of shape (N, M) representing N data points each consisting of M features\n",
    "            y: np.ndarray of shape (N,1) representing the targets for the each of the N data points\n",
    "                - data point X[i,:] has feature y[i]\n",
    "                - y is a column vector\n",
    "            R: np.ndarray of shape (N,) representing the weights for each of the samples\n",
    "            \n",
    "            No output -- but after this function call, self.params must store the most up-to-date model parameters\n",
    "        \"\"\"\n",
    "        self.params = None\n",
    "        ...\n",
    "\n",
    "    def predict(self, Xraw):\n",
    "        \"\"\" Calculates the predicted target values of the datapoints in Xraw and returns as a column vector\n",
    "\n",
    "            Note: make sure to convert Xraw to the design matrix before making your predictions\n",
    "            Note: predict can only be used after .fit is called, use self.params to perform prediction\n",
    "            Note: weights are not necessary at prediction time -- only at training\n",
    "            \n",
    "            Input:\n",
    "            Xraw: np.ndarray of shape (N, M) representing N data points each consisting of M features\n",
    "            \n",
    "            Output:\n",
    "            y_hat: np.ndarray of shape (N,1) representing the predicted target values for the N data points in X\n",
    "                - y_hat[i] is the predicted target value for data point X[i,:]\n",
    "            \n",
    "        \"\"\"\n",
    "        ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "Feel free to write your own code to test and visualize your `WeightedLinReg_ClosedForm` model implementation in as many cells as you would like below!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "Xraw = np.array([[1,1],[2,2],[3,3.5]])\n",
    "y = 3 + Xraw @ np.array([1,1]).reshape((-1,1))\n",
    "R = np.array([1,1,1])\n",
    "\n",
    "lr = WeightedLinReg_ClosedForm()\n",
    "lr.fit(Xraw, y, R)\n",
    "np.round(lr.params, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "Xraw = np.array([[1,1],[2,2],[3,3.5]])\n",
    "y = 3 + Xraw @ np.array([1,1]).reshape((-1,1))\n",
    "R = np.array([1,1,1])\n",
    "\n",
    "lr = WeightedLinReg_ClosedForm()\n",
    "lr.fit(Xraw, y, R)\n",
    "predictions = lr.predict(np.array([[1,2],[3,3],[4,5]]))\n",
    "np.round(predictions, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "grader.check(\"Q1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing Weighted Linear Regression\n",
    "Run the `plotWLR` function to visualize your model on the dataset that you were provided by the company and ensure that the weighted linear regression model appropriately fit to the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plotWLR():\n",
    "    '''plotting the weighted linear regression model'''\n",
    "    data = np.load(\"data.npz\")\n",
    "    Xraw, y, R = data[\"summer_X\"], data[\"summer_y\"], data[\"summer_R\"]\n",
    "    R = np.diag(R)\n",
    "\n",
    "    lr = WeightedLinReg_ClosedForm()\n",
    "    lr.fit(Xraw, y, R)\n",
    "    \n",
    "    plt.figure(figsize=(15,10))\n",
    "    plt.scatter(Xraw, y, marker=\"P\")\n",
    "    XSample = np.linspace(np.min(Xraw), np.max(Xraw), 400).reshape((-1, 1))\n",
    "    plt.plot(XSample, lr.predict(XSample), linestyle=\"solid\", linewidth=4, color=\"lightgreen\")\n",
    "    np.set_printoptions(precision=4)\n",
    "    plt.title(f\"Model Type: Closed Form WeightedLinReg b={lr.params[0]} w={lr.params[1]}\")\n",
    "    plt.xlabel(\"Temperature\")\n",
    "    plt.ylabel(\"Usage\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plotWLR()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Submit your code to Gradescope early and often\n",
    "\n",
    "There is no limit on the number of submissions to Gradescope, so as you complete parts of the assignment it is a really good idea to save your notebook and upload it to Gradescope.\n",
    "\n",
    "Not all of the tests are included in the local autograder. Some of the tests are \"hidden\" and only run in the server autograder on Gradescope.\n",
    "\n",
    "Before continuing with the rest of the assignment, go ahead and save your notebook (or click File->Download->Download .ipynb) and then upload your hw5.ipynb file to Gradescope under assignment hw5 (programming)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Question 2 - Iterative Training Methods <a class=\"anchor\" name=\"q2\"></a>\n",
    "- [**Q2a** - Helper Functions for Mini-Batch Stochastic Gradient Descent](#q2a)\n",
    "- [**Q2b** - Implementing Mini-Batch Stochastic Gradient Descent for Weighted Linear Regression](#q2b)\n",
    "- [**Q2c** - Programming Written Plotting](#q2c) (Writeup)\n",
    "- [**Q2d** - Gradient Descent Plot Analysis](#q2d) (Writeup)\n",
    "- [**Q2e & Q2f** - Big-O Complexity Questions](#q2e) (Writeup)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that you've implemented the weighted linear regression model, you are interested in finding different training methods for weighted linear regression, specifically through different variants of gradient descent.\n",
    "\n",
    "There are 3 main different types of gradient descent methods that you will be implementing:\n",
    "\n",
    "1. Gradient Descent (GD) - the entire dataset is used to estimate the gradient before each update\n",
    "2. Stochastic Gradient Descent (SGD) - a single sample is used to estimate the gradient before each update\n",
    "3. Mini-Batch Stochastic Gradient Descent (MBSGD) - a batch of samples is used to estimate the gradient before each update based on a batch size\n",
    "\n",
    "**However,** instead of implementing the learning algorithms for all three gradient descent methods, notice that there is a relationship between MBSGD and the other two gradient descent methods, GD and SGD, with respect to the batch size. So, if you can train a model with MBSGD, you can already train it with GD and SGD!\n",
    "\n",
    "With this understanding, you decide to implement the mini-batch stochastic gradient descent algorithm, but first, there are some helper functions necessary to do so."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Question 2a - Helper Functions for Mini-Batch Stochastic Gradient Descent <a class=\"anchor\" name=\"q2a\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "`shuffleDataset` will be given to you. When shuffling a dataset use this function, do not use any other random function as that will cause testing difficulties. Note: `shuffleDataset` doesn't actually appropriately shuffle the dataset for you by creating a random permutation of data. It simply cycles the array by one element to simulate shuffling. In practice, we would actually run the code that is in the comments below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "def shuffleDataset(X, y, R):\n",
    "        \"\"\" Shuffles the dataset into a random permutation\n",
    "            \n",
    "            Input:\n",
    "            X: np.ndarray of shape (N,M) that is the design matrix used for training the model\n",
    "            y: np.ndarray of shape (N,1) that contains the target values for each sample in the design matrix\n",
    "            R: np.ndarray of shape (N,) that is an array containing the weights of each sample\n",
    "            \n",
    "            Output:\n",
    "            XShuffled: np.ndarray of shape (N,M) shuffled version of X\n",
    "            yShuffled: np.ndarray of shape (N,1) shuffled version of y\n",
    "            RShuffled: np.ndarray of shape (N,) shuffled version of R\n",
    "        \"\"\"\n",
    "        idxs = np.arange(X.shape[0])\n",
    "        np.random.seed(280) # NOTE: required for test cases, but in practice this line should be removed\n",
    "        np.random.shuffle(idxs)\n",
    "        XShuffled, yShuffled, RShuffled = X[idxs, :], y[idxs, :], R[idxs]\n",
    "\n",
    "        return X[idxs, :], y[idxs, :], R[idxs]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "### Implement `initParams`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "Any gradient descent method needs a starting point at which to compute the gradient. Define the function `initParams` which initializes the parameters that will be learned via gradient descent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "otter_answer_cell"
    ]
   },
   "outputs": [],
   "source": [
    "def initParams(numParameters):\n",
    "    \"\"\" Initializes the parameters for a model to be trained with gradient descent.\n",
    "        Iniitialize all parameter values to 1\n",
    "        \n",
    "        When calling this function, make sure to consider the bias\n",
    "        \n",
    "        Input:\n",
    "        numParameters: the number of parameters that the linear model will have\n",
    "        \n",
    "        Output:\n",
    "        params: an (numParameters,1) column vector of ones\n",
    "    \"\"\"\n",
    "    ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "### Implement `computeBatches`\n",
    "With mini-batch stochastic gradient descent, at the start of each epoch, the dataset is shuffled and then partitioned into groups of batches. Implement this functionality in the function below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "otter_answer_cell"
    ]
   },
   "outputs": [],
   "source": [
    "def computeBatches(X, y, R, batchSize):\n",
    "    \"\"\" Places the dataset provided into a list of batches of the dataset according to batch size\n",
    "\n",
    "        Input:\n",
    "        X: np.ndarray of shape (N,M) that is the design matrix used for training the model\n",
    "        y: np.ndarray of shape (N,1) that contains the target values for each sample in the design matrix\n",
    "        R: np.ndarray of shape (N,) that is an array containing the weights of each sample\n",
    "        batchSize: (int) the size of the batches\n",
    "\n",
    "        NOTE: if the size of the dataset is not divisible by the batchSize, leave the last batch \n",
    "        to be the size of the remainder of the dataset (i.e X.shape[0] % batchSize)\n",
    "\n",
    "        Output:\n",
    "        batches: list of 3-tuples of np.ndarrays Xsub, ysub, and Rsub\n",
    "            - batches[i] = (Xsub, ysub, Rsub) for samples from batchSize * i to batchSize * (i + 1)\n",
    "            - the size of each batch is either batchSize or X.shape[0] % batchSize (for the edge case)\n",
    "    \"\"\"\n",
    "    batches = []\n",
    "    ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "### Implement `computeGradient`\n",
    "Finally, regardless of the method of gradient descent, the gradient of the loss of the weighted dataset w.r.t. the parameters is always the same. Using your theory, implement a method to compute the gradient of the weighted loss function with respect to the parameters for weighted linear regression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "otter_answer_cell"
    ]
   },
   "outputs": [],
   "source": [
    "def computeGradient(p, X, y, R):\n",
    "    \"\"\" Computes the gradient of the loss w.r.t. the parameters p for a linear model.\n",
    "        The loss is the weighted mean squared error function\n",
    "        \n",
    "        \n",
    "        Inputs:\n",
    "        p: np.ndarray of shape (M,1) that contains the parameters for some linear regression model\n",
    "        X: np.ndarray of shape (N,M) that is the design matrix used for training the model\n",
    "        y: np.ndarray of shape (N,1) that contains the target values for each sample in the design matrix\n",
    "        R: np.ndarray of shape (N,) that is an array containing the weights of each sample\n",
    "        \n",
    "        Outputs:\n",
    "        g: np.ndarray of shape (M,1) that contains the gradients calculated from the data\n",
    "            - g[i] corresponds to the gradient of the parameter p[i]\n",
    "            - Note: the weighted mean squared error function is the average of the weighted squared error for each datapoint\n",
    "    \"\"\"\n",
    "    ...\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "otter_answer_cell"
    ]
   },
   "outputs": [],
   "source": [
    "X = np.array([[1,1,-1],[1,2,-2],[1,3,-3],[1,4,-4],[1,5,-5],[1,6,-6]])\n",
    "y = np.array([1,2,3,4,5,6]).reshape((-1,1))\n",
    "R = np.array([0.25,0.25,1,0.75,0,0])\n",
    "batchSize = 2\n",
    "\n",
    "batches = computeBatches(X, y, R, batchSize)\n",
    "batches[0][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "params = initParams(3)\n",
    "params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "X = np.array([[1,1,-1],[1,2,-2],[1,3,-3],[1,4,-4],[1,5,-5],[1,6,-6]])\n",
    "y = np.array([1,2,3,4,5,6]).reshape((-1,1))\n",
    "R = np.array([0.25,0.25,1,0.75,0,0])\n",
    "batchSize = 2\n",
    "\n",
    "batches = computeBatches(X, y, R, batchSize)\n",
    "batches[1][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "p = np.array([-1,1]).reshape((-1,1))\n",
    "X = np.array([[1,-1],[1,2]])\n",
    "y = np.array([-3,3]).reshape((-1,1))\n",
    "R = np.array([1,0.5])\n",
    "\n",
    "grad = computeGradient(p, X, y, R)\n",
    "grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "grader.check(\"Q2a\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Question 2b - Implementing Mini-Batch Stochastic Gradient Descent for Weighted Linear Regression <a class=\"anchor\" name=\"q2b\"></a>\n",
    "Finally, with all of the helper methods in place, we can now use mini-batch stochastic gradient descent to find the optimal parameters for weighted linear regression. Now, you can implement the `WeightedLinReg_MBSGD` class.\n",
    "\n",
    "For the purposes of grading and nice plotting, below is the definition and implementation of the `TrainingLogger` class that defines a way to store the parameters and training error over multiple iterations of gradient descent and view the training process. Overall, the details of this class are unimportant, but when implementing mini-batch stochastic gradient descent, **define an object of this class at the start of training, log the initialized parameters and error, and then, for every update to the parameters in `WeightedLinReg_MBSGD.fit(...)`, log their values after the update**.\n",
    "\n",
    "Note: you will need to calculate the weighted mean-squared error to pass into the `error` parameter of `TrainingLogger.log(params, error)`. The weighted mean-square error is defined below\n",
    "\n",
    "$$MSE = \\frac{1}{N} \\sum_{i=1}^N R_{i,i} \\left(\\mathbf{y}^{(i)} - \\hat{\\mathbf{y}}^{(i)}\\right)^2$$\n",
    "\n",
    "You will implement this loss function in the helper function `getWeightedMSE(...)`\n",
    "\n",
    "Note: do not modify the `TrainingLogger` class!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrainingLogger:\n",
    "    def __init__(self, Xraw, y, R):\n",
    "        \"\"\" Initializes the training logger with the specific weighted dataset (X, y, R)\n",
    "            Every call to .log(params) will evaluate the given params against this dataset\n",
    "            using the MSE metric\n",
    "            \n",
    "            Inputs:\n",
    "            Xraw: np.ndarray of shape (N, M) representing N data points each consisting of M features\n",
    "            y: np.ndarray of shape (N,1) representing the targets for the each of the N data points\n",
    "            R: np.ndarray matrix of shape (N,) that is an array containing the weights of each sample\n",
    "        \"\"\"\n",
    "        self.data = (Xraw, y, np.diag(R))\n",
    "        self.allIterParams = []\n",
    "        self.allIterErrors = []\n",
    "\n",
    "    def log(self, params, error):\n",
    "        \"\"\" Log the parameter values and calculate their error w.r.t. the params\n",
    "        \n",
    "            Inputs:\n",
    "            params: np.ndarray of shape (M+1,1) that contains the parameters for some linear regression model\n",
    "                - note that the +1 is due to the bias parameter\n",
    "            error: (float) represents the weighted mean-squared error of the model on the training dataset\n",
    "        \"\"\"\n",
    "        self.allIterParams.append(np.zeros(params.shape) + params)\n",
    "        self.allIterErrors.append(error)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "### Implement `getWeightedMSE(...)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "otter_answer_cell"
    ]
   },
   "outputs": [],
   "source": [
    "def getWeightedMSE(Xraw, y, R, params):\n",
    "    \"\"\" Calculates the weighted mean squared error of the given dataset on the given set of parameters\n",
    "        \n",
    "        Input:\n",
    "        Xraw: np.ndarray of shape (N, M) representing N data points (i.e. samples) each consisting of M features\n",
    "        y: np.ndarray of shape (N,1) representing the true target values for the N data points in Xraw\n",
    "        R: np.ndarray of shape (N,) representing the weights for each of the N data points in Xraw\n",
    "        params: np.ndarray of shape (M+1,1) that are the current parameters of the model to evaluate\n",
    "        \n",
    "        Output:\n",
    "        error: float equal to the weighted mean-squared error of the dataset on the parameters\n",
    "    \"\"\"\n",
    "    ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "### Implement `.fit(...)`, and `.predict(...)` of `WeightedLinReg_MBSGD`\n",
    "Implement the following functions listed above that are missing their implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "otter_answer_cell"
    ]
   },
   "outputs": [],
   "source": [
    "class WeightedLinReg_MBSGD(Model):\n",
    "\n",
    "    def __init__(self, lr=0.02, epochs=500, batchSize=256):\n",
    "        \"\"\" Initializes a model that can be trained\n",
    "            with the mini-batch gradient descent algorithm using the .fit(...) method\n",
    "            \n",
    "            gradient descent algorithms can be:\n",
    "            1. regular gradient descent (GD)\n",
    "            2. stochastic gradient descent (SGD)\n",
    "            3. mini-batch stochastic gradient descent (MBSGD)\n",
    "            (note: MBSGD is a generalization of GD and SGD, how?)\n",
    "        \n",
    "            Input:\n",
    "            lr: (float) the learning rate to use during gradient descent\n",
    "            epochs: (int) the number of iterations to perform gradient descent\n",
    "                - note: the number of epochs is the number of times you iterate\n",
    "                        over the entire dataset\n",
    "            batchSize: (int) the size of each batch in training\n",
    "        \"\"\"\n",
    "        # Note: do not need to modify __init__ method\n",
    "        self.lr = lr\n",
    "        self.epochs=epochs\n",
    "        self.batchSize = batchSize\n",
    "\n",
    "    def fit(self, Xraw, y, R):\n",
    "        \"\"\" Performs mini-batch stochastic gradient descent to train the model\n",
    "            for self.epochs iterations and with a learning rate of self.lr and a batch size of self.batchSize\n",
    "            \n",
    "            1) This function initializes parameters, constructs the design matrix, and performs multiple epochs.\n",
    "            Each epoch will shuffle the data, batch it, and then iterate through the batches, and update the model \n",
    "            parameters using that batch using the gradient of the loss w.r.t. the parameters.\n",
    "\n",
    "            2) Additionally, this function will log the value of the parameters at the start of training and \n",
    "            once after every single update to the parameter values, and store the log in self.logger.\n",
    "\n",
    "            Note: Make sure to shuffle and then batch the original dataset that you were given. Do not reshuffle a\n",
    "            version of the dataset that has been already shuffled.\n",
    "            \n",
    "            Input:\n",
    "            Xraw: np.ndarray of shape (N,M) that is the dataset of N data points each with M features\n",
    "            y: np.ndarray of shape (N,1) that contains the target values for each sample in the design matrix\n",
    "            R: np.ndarray of shape (N,) that is an array containing the weights of each sample \n",
    "    \n",
    "            No output -- self.params and self.logger must be defined after this function has completed\n",
    "                - self.params must equal the fitted parameters after performing mini-batch stochastic gradient descent\n",
    "                - self.logger must contain the training logger that has logged 1 + numBatches * self.epochs times\n",
    "                    - once after the parameters were initialize\n",
    "                    - numBatches * self.epochs times after the parameters were updated\n",
    "        \"\"\"\n",
    "        self.params, self.logger = None, None\n",
    "        ...\n",
    "        \n",
    "    def predict(self, Xraw):\n",
    "        \"\"\" Calculates the predicted target values of the datapoints in Xraw. Should only be called after\n",
    "            .fit is called, use self.params to perform prediction\n",
    "        \n",
    "            Additional Note: sample weights are not necessary at prediction time -- only at training\n",
    "            \n",
    "            Input:\n",
    "            Xraw: np.ndarray of shape (N, M) representing N data points each consisting of M features\n",
    "            \n",
    "            Output:\n",
    "            y_hat: np.ndarray of shape (N,1) representing the predicted target values for the N data points in Xraw\n",
    "        \"\"\"\n",
    "        ...\n",
    "\n",
    "    def getTrainingLog(self):\n",
    "        \"\"\" Returns the training logger that contains the information accumulated during training\n",
    "            .log has been called on it for 1 + self.epochs * number_parameter_updates\n",
    "                - once when the parameters were first initialized\n",
    "                - once after every single update to the parameters\n",
    "            Only callable after training\n",
    "        \"\"\"\n",
    "        return self.logger"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "### Implement `trainMultipleModels()`\n",
    "Now that you have implemented the `WeightedLinReg_MGSD` class, we will define a function to train 3 weighted linear regression models on the same dataset but each using a different variant of gradient descent.\n",
    "1. For each variant, use a learning rate of 0.02\n",
    "2. For each variant, train for 500 epochs  \n",
    "3. The first should use gradient descent, the second should use stochastic gradient descent, and the third should use mini-batch stochastic gradient descent with a batch size of 256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "otter_answer_cell"
    ]
   },
   "outputs": [],
   "source": [
    "def trainMultipleModels(Xraw, y, R):\n",
    "    \"\"\" Trains multiple weighted linear regression models using various gradient descent methods\n",
    "        Given the training dataset as input, trains 3 models of type WeightedLinReg_MBSGD\n",
    "        \n",
    "        All models are trained with 500 epochs and a learning rate of 0.02.\n",
    "        However, the first model is trained via gradient descent. The second model is trained via \n",
    "        stochastic gradient descent. The third model is trained via mini-batch stochastic gradient descent\n",
    "        with a batch size of 256.\n",
    "        \n",
    "        Input:\n",
    "        Xraw: np.ndarray of shape (N,M) that is the dataset of N data points each with M features\n",
    "        y: np.ndarray of shape (N,1) that contains the target values for each sample in the design matrix\n",
    "        R: np.ndarray of shape (N,) that is an array containing the weights of each sample \n",
    "        \n",
    "        Output:\n",
    "        gd: WeightedLinReg_MBSGD trained using gradient descent\n",
    "        sgd: WeightedLinReg_MBSGD trained using stochastic gradient descent\n",
    "        mbsgd: WeightedLinReg_MBSGD trained using mini-batch stochastic gradient descent with batch size 256\n",
    "    \"\"\"\n",
    "    \n",
    "    ...\n",
    "    return gd, sgd, mbsgd\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "Xraw = np.array([[-1,1],[1,-2,],[3,9],[-4,5],[1,1]])\n",
    "params = np.array([3,-1,1]).reshape(-1,1)\n",
    "y = params[2,:] + Xraw @ params[:2,:]\n",
    "R = np.array([0.5,1,1,1,0.25])\n",
    "\n",
    "err = getWeightedMSE(Xraw, y, R, np.array([-3,1,1]).reshape(-1,1))\n",
    "err"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "Xraw = np.array([[1,1],[2,2],[3,3.5]])\n",
    "y = 3 + Xraw @ np.array([1,1]).reshape((-1,1))\n",
    "R = np.array([1,1,1])\n",
    "\n",
    "lr = WeightedLinReg_MBSGD(lr=0.5, epochs=1, batchSize=1)\n",
    "lr.fit(Xraw, y, R)\n",
    "predictions = lr.predict(np.array([[1,2],[3,3],[4,5]]))\n",
    "np.round(predictions, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "Xraw = np.array([[1,1],[2,2],[3,3],[2,2.5]])\n",
    "y = np.array([-1,-2,-3,-2.3]).reshape((-1,1))\n",
    "R = np.array([1,1,1,0.5])\n",
    "\n",
    "lr = WeightedLinReg_MBSGD(lr=0.001, epochs=3, batchSize=2)\n",
    "lr.fit(Xraw, y, R)\n",
    "np.round(lr.params, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "grader.check(\"Q2b\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are just some helper functions to plot the data and the models that you've implemented. Feel free to look into the data, but make sure that you do not modify these functions as they will be used for the written portion of the programming assignment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loadGDData():\n",
    "    data = np.load(\"data.npz\")\n",
    "    Xraw, y, R = data[\"summer_X\"], data[\"summer_y\"], data[\"summer_R\"]\n",
    "    Xraw /= np.max(Xraw) # we scale the dataset to prevent the gradients from exploding\n",
    "    y /= np.max(y)\n",
    "    R = np.diag(R)\n",
    "    \n",
    "    return Xraw, y, R\n",
    "\n",
    "def plotGDModels(X, y, R, gd, sgd, mbsgd):\n",
    "    '''plotting the GD, SGD, MB-SGD models'''\n",
    "\n",
    "    # plot the fitted functions\n",
    "    models = [gd, sgd, mbsgd]\n",
    "    fig, axes = plt.subplots(3, 1, figsize=(10,15), sharex=True, constrained_layout=True)\n",
    "    X_sample = np.array([np.min(X), np.max(X)]).reshape((-1, 1))\n",
    "    colors = [\"orange\", \"green\", \"purple\"]\n",
    "    labels = [\"gd\", \"sgd\", \"mbsgd\"]\n",
    "\n",
    "    for i in range(len(models)):\n",
    "        axes[i].scatter(X, y, marker=\"P\")\n",
    "        axes[i].plot(X_sample, models[i].predict(X_sample), linestyle=\"solid\", linewidth=4, color=colors[i])\n",
    "        axes[i].set_title(f\"Model Type:{labels[i]}, b={models[i].params[0]}, w={models[i].params[1]}\")\n",
    "        axes[i].set_xlabel(\"Temperature: scaled to [0,1]\")\n",
    "        axes[i].set_ylabel(\"Usage: scaled to [0,1]\")\n",
    "\n",
    "    plt.show()\n",
    "    \n",
    "def plotGDLosses(X, y, R, gd, sgd, mbsgd):\n",
    "    '''plot the error rates between the 3 iterative models'''\n",
    "\n",
    "    # plot the error rates of the models\n",
    "    models = [gd, sgd, mbsgd]\n",
    "    colors = [\"orange\", \"green\", \"purple\"]\n",
    "    labels = [\"gd\", \"sgd\", \"mbsgd\"]\n",
    "\n",
    "    epochs = 500\n",
    "    epochsToShow = 4\n",
    "    \n",
    "    plt.figure(figsize=(10,10))\n",
    "\n",
    "    for m, c, l in zip(models, colors, labels):\n",
    "        logger = m.getTrainingLog()\n",
    "        log = logger.allIterErrors\n",
    "        samplesPerEpoch = (len(log)-1) // epochs\n",
    "        nPoints = epochsToShow * samplesPerEpoch\n",
    "        x = np.linspace(0, epochsToShow, nPoints+1)\n",
    "        plt.plot(x, log[:nPoints+1], color=c, label=l)\n",
    "    \n",
    "    plt.xlabel(\"Epoch #\")\n",
    "    plt.ylabel(\"Loss (Weighted MSE)\")\n",
    "    plt.title(\"Loss Over Epochs for Iterative Training Methods\")\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q2c - Programming Written Question <a class=\"anchor\" name=\"q2c\"></a>\n",
    "\n",
    "Run the two plots below, and paste them into the writeup of the programming section of the homework for question **Q2c**. Afterwards, answer the following analysis questions in **Q2d through Q2f**.\n",
    "\n",
    "Note: you may have gotten different parameter values for the closed form solution of weighted linear regression in comparison to the various gradient descent methods. This is perfectly fine because you will notice that the parameters for gradient descent will converge to the closed form solution if they trained for more and more epochs (around 5000 epochs). If you would like to test this out, change the number of epochs that `WeightedLinReg_GD` and `WeightedLinReg_MB_SGD` are trained for to 5000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# note: this function will take some time due to a larger number of epochs used for training \n",
    "#       and the logging that is happening on every parameter update (it's especially time-consuming for SGD)\n",
    "Xraw, y, R = loadGDData()\n",
    "gd, sgd, mbsgd = trainMultipleModels(Xraw, y, R)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plotGDModels(Xraw, y, R, gd, sgd, mbsgd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plotGDLosses(X, y, R, gd, sgd, mbsgd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q2d - Gradient Descent Analysis <a class=\"anchor\" name=\"q2d\"></a>\n",
    "Analyze the gradient descent plots and answer question **Q2d** in the writeup of this assignment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q2e and Q2f - Big-O Analysis <a class=\"anchor\" name=\"q2e\"></a>\n",
    "Analyze the big-O complexity of closed form weighted linear regression and compare it to the complexity of a single iteration of gradient descent in **Q2e** and **Q2f** of the writeup of this assignment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Submit your code to Gradescope early and often\n",
    "\n",
    "There is no limit on the number of submissions to Gradescope, so as you complete parts of the assignment it is a really good idea to save your notebook and upload it to Gradescope.\n",
    "\n",
    "Not all of the tests are included in the local autograder. Some of the tests are \"hidden\" and only run in the server autograder on Gradescope.\n",
    "\n",
    "Before continuing with the rest of the assignment, go ahead and save your notebook (or click File->Download->Download .ipynb) and then upload your hw5.ipynb file to Gradescope under assignment hw5 (programming)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Question 3 - Nonlinear Modeling <a class=\"anchor\" name=\"q3\"></a>\n",
    "\n",
    "- [**Q3a** - Polynomial Features](#q3a)\n",
    "- [**Q3b** - K-Fold Cross Validation](#q3b)\n",
    "- [**Q3c** - Programming Written Plot](#q3c) (Writeup)\n",
    "\n",
    "With your success modeling the summer dataset, the company decided to collect more data throughout the year to get a better understanding of consumer demand depending on the temperature throughout the day. Earlier you were working with data collected from May to August. Now the company gathered data throughout the rest of the year, September to April. Now with data from the full year you can now create a model that is made from temperatures recorded throughout the whole year. Below, you can find the updated dataset, distinguished by the data collected from the summmer (May to August) and the winter (September to April)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plotDatasetSeason():\n",
    "    '''plotting the weighted linear regression model t=(\"summer\"|\"winter\"|\"full\")'''\n",
    "    data = np.load(\"data.npz\")\n",
    "    \n",
    "    X_summer, y_summer = data[\"summer_X\"], data[\"summer_y\"]\n",
    "    X_winter, y_winter = data[\"winter_X\"], data[\"winter_y\"]\n",
    "    \n",
    "    plt.figure(figsize=(15,10))\n",
    "    plt.scatter(X_summer, y_summer, marker=\"P\", c=\"sandybrown\", label=\"Summer\")\n",
    "    plt.scatter(X_winter, y_winter, marker=\"P\", c=\"lightskyblue\", label=\"Winter\")\n",
    "    plt.title(f\"Temperature vs. Usage Data points\")\n",
    "    plt.xlabel(\"Temperature\")\n",
    "    plt.ylabel(\"Usage\")\n",
    "    plt.legend()\n",
    "    \n",
    "plotDatasetSeason()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Oh no!** We were using linear models to model our dataset; however, after collecting data for the whole year, our dataset appears to have a non-linear relationship. Using a vanilla linear regression on its own may not model this dataset that well. As a result, you will take a first look at feature engineering, where you will engineer polynomial features from your dataset, and use that for linear regression. For example, typically in linear regression, you are modelling the following function (assume 1 feature and 1 target)\n",
    "\n",
    "$$\\mathbf{y}=\\phi_0 + \\phi_1\\mathbf{x}$$\n",
    "\n",
    "where $\\phi_0$ and $\\phi_1$ are your parameters. However, if you were to engineer features for up to the third degree polynomial, you would now be modelling this equation.\n",
    "\n",
    "$$\\mathbf{y} = \\phi_0 + \\phi_1\\mathbf{x} + \\phi_2\\mathbf{x}^2 + \\phi_3\\mathbf{x}^3$$\n",
    "\n",
    "Note, while the data that you are using represents polynomial features, the actual function is still linear w.r.t to the polynomial features. You are performing a linear combination of the polynomial features. As a result, by creating polynomial features, you can still use linear regression to model this dataset!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Question 3a - Polynomial Features <a class=\"anchor\" name=\"q3a\"></a>\n",
    "\n",
    "Before you can start implementing the weighted linear regression model with polynomial features, we first need to engineer polynomial features from our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "otter_answer_cell"
    ]
   },
   "outputs": [],
   "source": [
    "def makePolynomialDesignMatrix(Xraw, degree):\n",
    "    \"\"\" Creates polynomial features from Xraw by concatenating a \n",
    "        column of 1's to the left of the array and concatenating higher degree\n",
    "        columns of the original columns to the right\n",
    "        \n",
    "        >>> a\n",
    "        array([[ 4,  7],\n",
    "               [ 3,  8],\n",
    "               [12,  9]])\n",
    "        >>> makePolynomialDesignMatrix(a, 3)\n",
    "        array([[1,   4,    7,   16,   49,   64,  343],\n",
    "               [1,   3,    8,    9,   64,   27,  512],\n",
    "               [1,  12,    9,  144,   81, 1728,  729]])\n",
    "               \n",
    "        Input:\n",
    "        Xraw: np.ndarray of shape (N,M)\n",
    "        degree: int such that degree >= 1\n",
    "        Returns:\n",
    "        polynomial_matrix: matrix with polynomial features up to degree (N, 1+M*degree)\n",
    "    \"\"\"\n",
    "    ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "A = np.array([[4,7],[3,8],[12,9]])\n",
    "X = makePolynomialDesignMatrix(A, 3).astype(np.int64)\n",
    "X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "A = np.array([[1,1,1],[2,2,2],[3,3,3],[4,4,4]])\n",
    "X = makePolynomialDesignMatrix(A, 2)\n",
    "X"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "### Implement `.fit`, `.predict`, and `.score` for `PolyWeightedLinReg_ClosedForm`\n",
    "Implement the closed form solution to the weighted linear regression classifier below; however, this time use polynomial features for the given input degree `self.degree` in the class. Furthermore, implement the function for prediction as well.\n",
    "\n",
    "Finally, because you would like to compare this model with other models when doing model selection, we will have you compute the score function so that you can calculate how well the model is performing on an input dataset. However, instead of calculating the usual mean-squared error, we are going to have you calculate a new metric that is useful for evaluating regression models called the **coefficient of determination**, $R^2$, which is defined below.\n",
    "\n",
    "$$R^2 = 1 - \\frac{RSS}{TSS}$$\n",
    "\n",
    "where the residual sum of squares is $RSS = \\sum_{i=1}^N(y^{(i)} - f(x^{(i)}))^2$ and total sum of squares is $TSS =  \\sum_{i=1}^N(y^{(i)} - \\bar{y})^2$ where $\\bar{y} = \\frac{1}{N} \\sum_{i=1}^Ny^{(i)}$. Implement this metric in the `.score` method for `PolyWeightedLinReg_ClosedForm`. The coefficient of determination is a metric that is maximized at value 1, and a higher value represents a model with a better fit to the dataset. We are using the coefficient of determination because it is a commonly used metric for evaluating model performance on various regression models and is commonly used throughout some machine learning packages, in addition to mean-squared error.\n",
    "\n",
    "Note: You won't have to implement the training logger for this algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "otter_answer_cell"
    ]
   },
   "outputs": [],
   "source": [
    "class PolyWeightedLinReg_ClosedForm(Model):\n",
    "    def __init__(self, degree=1):\n",
    "        # Note: do not need to modify __init__ method\n",
    "        self.degree = degree\n",
    "    \n",
    "    def fit(self, Xraw, y, R):\n",
    "        \"\"\" Calculates the closed-form solution parameters for weighted linear regression\n",
    "            with polynomial features.\n",
    "            \n",
    "            Note: be aware of the shapes of your input when you are doing NumPy operations!\n",
    "        \n",
    "            Input:\n",
    "            Xraw: np.ndarray of shape (N, M) representing N data points each consisting of M features\n",
    "            y: np.ndarray of shape (N,1) representing the targets for the each of the N data points\n",
    "                - data point X[i,:] has target y[i]\n",
    "                - y is a column vector\n",
    "            R: np.ndarray of shape (N,) representing the weights for each of the samples\n",
    "            \n",
    "            No output, but self.params is saved with the learned parameters for the model\n",
    "        \"\"\"\n",
    "        ...\n",
    "\n",
    "    def predict(self, Xraw):\n",
    "        \"\"\" Makes predictions for the new input data Xraw\n",
    "        \n",
    "            Input:\n",
    "            Xraw: np.ndarray of shape (N, M) representing N data points each consisting of M features\n",
    "            \n",
    "            Output:\n",
    "            y_hat: np.ndarray of shape (N,1) representing the predicted target values for the each of the datapoints\n",
    "        \"\"\"\n",
    "        ...\n",
    "\n",
    "    def score(self, Xraw, y):\n",
    "        ''' Calculates the coefficient of determination by predicting on X and evaluating the metric on y\n",
    "            \n",
    "            Input:\n",
    "            X: np.ndarray of shape (N, M) representing N data points each consisting of M features\n",
    "            y: np.ndarray of shape (N, 1) representing the predicted target valeus for each datapoints\n",
    "        '''\n",
    "        ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "Xraw = np.array([1,2,3,4,5,6,7,8]).reshape((-1,1))\n",
    "y = 8 + 7 * Xraw + 3 * Xraw ** 2 + -3 * Xraw ** 3\n",
    "R = np.ones((Xraw.shape[0]))\n",
    "\n",
    "lr = PolyWeightedLinReg_ClosedForm(degree=3)\n",
    "lr.fit(Xraw, y, R)\n",
    "score = lr.score(Xraw, y)\n",
    "score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "Xraw = np.array([[1,1],[2,2],[0,3]])\n",
    "y = 3 + Xraw @ np.array([1,1]).reshape((-1,1))\n",
    "R = np.array([1,1,1])\n",
    "\n",
    "lr = PolyWeightedLinReg_ClosedForm(degree=3)\n",
    "lr.fit(Xraw, y, R)\n",
    "predictions = lr.predict(np.array([[1,2],[3,3],[4,5]]))\n",
    "np.round(predictions, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "grader.check(\"Q3a\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plotModelTemp(model, t):\n",
    "    '''plot the dataset and the model's predictions on it'''\n",
    "    data = np.load(\"data.npz\")\n",
    "    Xraw, y, R = data[f\"{t}_X\"], data[f\"{t}_y\"], data[f\"{t}_R\"]\n",
    "    \n",
    "    model.fit(Xraw, y, np.diag(R))\n",
    "    \n",
    "    Xmin = np.min(Xraw)\n",
    "    Xmax = np.max(Xraw)\n",
    "    Xmodel = np.linspace(Xmin, Xmax, 400).reshape((-1,1))\n",
    "    ymodel = model.predict(Xmodel)\n",
    "    \n",
    "    X_summer, y_summer = data[\"summer_X\"], data[\"summer_y\"]\n",
    "    X_winter, y_winter = data[\"winter_X\"], data[\"winter_y\"]\n",
    "    \n",
    "    print(f\"Score: {model.score(Xraw, y)}\")\n",
    "    \n",
    "    plt.figure(figsize=(12,7))\n",
    "    plt.scatter(X_summer, y_summer, marker=\"P\", c=\"sandybrown\", label=\"Summer\")\n",
    "    plt.scatter(X_winter, y_winter, marker=\"P\", c=\"lightskyblue\", label=\"Winter\")\n",
    "    plt.plot(Xmodel, ymodel, linewidth=4, label=\"model predictions\")\n",
    "    plt.title(f\"Temperature vs. Usage Plot with Polynomial Model of Degree {model.degree}\")\n",
    "    plt.xlabel(\"Temperature\")\n",
    "    plt.ylabel(\"Usage\")\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plotModelTemp(PolyWeightedLinReg_ClosedForm(degree=2), \"full\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plotModelTemp(PolyWeightedLinReg_ClosedForm(degree=5), \"full\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plotModelTemp(PolyWeightedLinReg_ClosedForm(degree=11), \"full\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Question 3b - K-Fold Cross Validation <a class=\"anchor\" name=\"q3b\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "Now that you've created a non-linear model and learned how to fit it to the dataset that you are working with, you feel that you are ready to start training multiple models and choosing the model that you believe will perform the best for the enterprise. However, you would like some way of comparing the generalizability of the models. Recalling different model selection strategies, you decide to perform **k-fold cross validation** to perform model selection, but you will have to implement it from scratch.\n",
    "\n",
    "There are 3 major parts to performing k-fold cross validation\n",
    "1. splitting the dataset into k folds\n",
    "2. computing the training dataset and validation dataset from k folds after choosing a specific fold for validation\n",
    "3. training, evaluating model across all folds and averaging the metric.\n",
    "\n",
    "We will have you compute the second two parts of this process for performing k-fold cross validation. Below is the implementation of the first part of the k-fold cross validation process. Do not modify it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "def createKFolds(Xraw, y, R, k):\n",
    "    \"\"\" Splits/partitions the given dataset into K-fold cross validation datasets usable for model selection.\n",
    "        Produces k-folds of roughly even sizes of associated data points.\n",
    "        \n",
    "        Input:\n",
    "        Xraw: np.ndarray of shape (N, M) representing N data points each consisting of M features\n",
    "        y: np.ndarray of shape (N,1) representing the targets for the each of the N data points\n",
    "            - data point X[i,:] has target y[i]\n",
    "            - y is a column vector\n",
    "        R: np.ndarray of shape (N,) representing the weights for each of the samples\n",
    "        k: (int) the number of folds to produce in the dataset (k <= Xraw.shape[0])\n",
    "        \n",
    "        Output:\n",
    "        XrawFolds: list of length k of subset of dataset of Xraw\n",
    "        yFolds: list of length k of subset of dataset of y\n",
    "        RFolds: list of length k of subset of dataset of R\n",
    "        \n",
    "    \"\"\"\n",
    "    \n",
    "    # initialize storage for data\n",
    "    XrawFolds = []\n",
    "    yFolds = []\n",
    "    RFolds = []\n",
    "    \n",
    "    # shuffle the data\n",
    "    N, M = Xraw.shape\n",
    "    \n",
    "    # construct the folds\n",
    "    idxs = np.rint(np.linspace(0, N, k+1)).astype(np.int64)\n",
    "    for i in range(idxs.shape[0] - 1):\n",
    "        s = idxs[i]\n",
    "        e = idxs[i+1]\n",
    "        XrawFolds.append(Xraw[s:e,:])\n",
    "        yFolds.append(y[s:e,:])\n",
    "        RFolds.append(R[s:e])\n",
    "    return XrawFolds, yFolds, RFolds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "# example to see how the k-fold works (feel free to modify this cell block or add any others)\n",
    "N, M = 100, 3\n",
    "k = 6\n",
    "Xraw = np.random.randn(N,M)\n",
    "y = np.random.randn(N,1)\n",
    "r = np.random.randn(N)\n",
    "XrawFolds, yFolds, RFolds = createKFolds(Xraw, y, r, k)\n",
    "\n",
    "# print the shapes of the folds\n",
    "print(\"ith column of shapes represents the shape of the ith fold\")\n",
    "print(\"X-folds:\", end=\" \"); [print(XrawFolds[i].shape, end=\" \") for i in range(len(XrawFolds))]; print()\n",
    "print(\"y-folds:\", end=\" \"); [print(yFolds[i].shape, end=\" \") for i in range(len(yFolds))]; print()\n",
    "print(\"r-folds:\", end=\" \"); [print(RFolds[i].shape, end=\" \") for i in range(len(RFolds))]; print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "otter_answer_cell"
    ]
   },
   "outputs": [],
   "source": [
    "def getTrainValFromFold(XrawFolds, yFolds, RFolds, index):\n",
    "    \"\"\" Constructs a training and validation dataset from a k-fold cross validation dataset \n",
    "        where index is the fold that will be the validation set while the rest are training data.\n",
    "        \n",
    "        Input:\n",
    "        XrawFolds: list of length k of np.ndarray subsets of the dataset Xraw\n",
    "        yFolds: list of length k of np.ndarray subsets of target values y\n",
    "        RFolds: list of length k of np.ndarray subsets of sample weights R\n",
    "        index: (int) from 0 to k-1 representing the fold that will be the validation dataset\n",
    "        \n",
    "        Output:\n",
    "        trainXraw: training features composed of vertically stacking the below folds\n",
    "            - XFolds[0], XFolds[1], ... XFolds[index-1], XFolds[index+1], ..., XFolds[k-1]\n",
    "        trainy: training targets composed by vertically stacking the below folds\n",
    "            - yFolds[0], yFolds[1], ... yFolds[index-1], yFolds[index+1], ..., yFolds[k-1]\n",
    "        trainR: sample weights stored in a 1-dimensional vector composed concatenating the below folds\n",
    "            - RFolds[0], RFolds[1], ... RFolds[index-1], RFolds[index+1], ..., RFolds[k-1]\n",
    "        validX: validation features composed of XFolds[index]\n",
    "        validy: validation targets composed of yFolds[index]\n",
    "    \"\"\"\n",
    "    ...\n",
    "    return trainXraw, trainy, trainR, validX, validy\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "N, M = 20, 3\n",
    "k = 5\n",
    "idx = 4\n",
    "Xraw = np.arange(N*M).reshape((N,M))\n",
    "y = np.arange(N).reshape((N,1))\n",
    "r = np.ones(N)\n",
    "XFolds, yFolds, RFolds = createKFolds(Xraw, y, r, k)\n",
    "trainX, trainy, trainR, validX, validy = getTrainValFromFold(XFolds, yFolds, RFolds, idx)\n",
    "\n",
    "assert(trainX.shape[0] == 16)\n",
    "assert(trainy.shape[0] == 16)\n",
    "assert(trainR.shape[0] == 16)\n",
    "assert(validX.shape[0] == 4)\n",
    "assert(validy.shape[0] == 4)\n",
    "\n",
    "assert(np.all(validy == yFolds[idx]))\n",
    "trainX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "For the final part of K-Fold Cross Validation, you will need to compute the k-fold cross validation accuracy of a given model that implements `.fit` and `.score`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "otter_answer_cell"
    ]
   },
   "outputs": [],
   "source": [
    "def kFoldCrossValidation(model, XrawFolds, yFolds, RFolds):\n",
    "    \"\"\" Computes the k-fold cross validation score of the given model.\n",
    "        This validation score is calculated as follows: for each fold, calculate the training and validation datasets\n",
    "        and average the score across the folds.\n",
    "        \n",
    "        Input:\n",
    "        model: (Model) implements the `.fit(X, y, R)` method and `.score(X, y)` method that computes\n",
    "               the coefficient of determination\n",
    "        XrawFolds: list of length k of np.ndarray subsets of dataset Xraw\n",
    "        yFolds: list of length k of np.ndarray subsets of target values y\n",
    "        RFolds: list of length k of np.ndarray subsets of sample weights R\n",
    "        \n",
    "        Output:\n",
    "        cvScore: the K-Fold Cross Validation score for the input model\n",
    "    \"\"\"\n",
    "    ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "k = 3\n",
    "Xraw = np.arange(10).reshape((-1,1))\n",
    "y = 3 + 5 * Xraw + 3 * np.array([0.2,-0.4,0.4,0.2,0.4,-0.2,0.2,-0.2,0.4,0.2]).reshape((10,1))\n",
    "r = np.ones(10)\n",
    "XrawFolds, yFolds, RFolds = createKFolds(Xraw, y, r, k)\n",
    "\n",
    "model = PolyWeightedLinReg_ClosedForm(degree=3)\n",
    "score = kFoldCrossValidation(model, XrawFolds, yFolds, RFolds)\n",
    "np.round(score, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "grader.check(\"Q3b\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Q3c - Training and Selecting Models <a class=\"anchor\" name=\"q3b\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "Now that you've completed implementing the k-fold cross validation process, you will train two different types of models, each with 5 different hyperparameters, leading to 10 different models (which we will provide you). You will evaluate the models using k-fold cross validation, and then, you will determine the best model of the 10 from the their cross-validation scores. Note: all of the models will follow the same interface, so they can have the same functions called on them.\n",
    "\n",
    "The first type of model that you will train is the `PolyWeightedLinReg_ClosedForm` model for `degree=[1,3,5,7,9]`.\n",
    "The second type of model that you will train is the `KNNRegressor` model for `n_neighbors=[1,3,9,15,25]` (class defined below). We will provide these models for you to compare using cross-validation.\n",
    "\n",
    "These will be the ten models that you will be training. Note that the `KNNRegressor` just wraps an `sklearn.neighbors.KNeighborsRegressor` model, allowing it to take weighted inputs for fitting but not really using them. This allows the `KNNRegressor` and the `PolyWeightedLinReg_ClosedForm` to have the same interface making them easy to use together in coding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "otter_answer_cell"
    ]
   },
   "outputs": [],
   "source": [
    "def selectModel(models, Xraw, y, R, kFolds=5):\n",
    "    \"\"\" Given a list of models to train (models) and a dataset (Xraw, y, R) and the number of folds (kFolds),\n",
    "        train and evaluate the models using K-Fold Cross Validation and return their scores and the corresponding best\n",
    "        model.\n",
    "        \n",
    "        Input:\n",
    "        models: a list of objects of type Model that are usable for performing K-Fold Cross Validation\n",
    "        Xraw: np.ndarray of shape (N,M) that is the dataset to train the model with N data points of M features\n",
    "        y: np.ndarray of shape (N,1) that contains the target values for each sample in the design matrix\n",
    "        R: np.ndarray of shape (N,) that is an array containing the weights of each sample\n",
    "        kFolds: (int) number of folds to perform k-fold cross validation for\n",
    "        \n",
    "        Output:\n",
    "        bestIdx: the index of the model in models with the highest cross-validation score\n",
    "        cvScores: list of k-fold cross validation scores for input models\n",
    "            - models[i] has k-fold cross validation score cvScore[i]\n",
    "    \"\"\"\n",
    "    ...\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "otter_answer_cell"
    ]
   },
   "outputs": [],
   "source": [
    "# DO NOT EDIT this code block\n",
    "from sklearn.neighbors import KNeighborsRegressor\n",
    "class KNNRegressor(Model):\n",
    "    \n",
    "    def __init__(self, n_neighbors=5):\n",
    "        self.knn = KNeighborsRegressor(n_neighbors=n_neighbors)\n",
    "        \n",
    "    def fit(self, Xraw, y, R):\n",
    "        self.knn.fit(Xraw, y)\n",
    "        \n",
    "    def predict(self, Xraw):\n",
    "        return self.knn.predict(Xraw).reshape((-1,1))\n",
    "        \n",
    "    def score(self, Xraw, y):\n",
    "        return self.knn.score(Xraw, y)\n",
    "\n",
    "# DO NOT EDIT initModels()\n",
    "def initModels():\n",
    "        \"\"\" Initializes models that is meant for training\n",
    "        \"\"\"\n",
    "        polyModels = [PolyWeightedLinReg_ClosedForm(degree=i) for i in [1,3,5,7,9]]\n",
    "        knnModels = [KNNRegressor(n_neighbors=i) for i in [1,3,9,15,25]]\n",
    "        return polyModels + knnModels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "kFolds = 2\n",
    "Xraw = np.arange(50).reshape((-1,1))\n",
    "y = 3 + Xraw @ np.array([2]).reshape((-1,1))\n",
    "R = np.ones(50)\n",
    "\n",
    "models = initModels()\n",
    "bestIdx, scores = selectModel(models, Xraw, y, R, kFolds=kFolds)\n",
    "assert(bestIdx == np.argmax(scores))\n",
    "np.round(np.array(scores).astype(np.float64).reshape((-1,1)), 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "grader.check(\"Q3c\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Performing K-Fold Cross Validation with Different Models on Energy Usage Dataset\n",
    "Run the following commands below to plot the models described on the energy usage dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plotModelTemp(model, title=\"Temperature vs. Usage Plot with Polynomial Model\"):\n",
    "    '''plot the dataset and the model's predictions on it'''\n",
    "    data = np.load(\"data.npz\")\n",
    "    Xraw, y, R = data[f\"full_X\"], data[f\"full_y\"], data[f\"full_R\"]\n",
    "    \n",
    "    Xmin = np.min(Xraw)\n",
    "    Xmax = np.max(Xraw)\n",
    "    Xmodel = np.linspace(Xmin, Xmax, 400).reshape((-1,1))\n",
    "    ymodel = model.predict(Xmodel)\n",
    "    \n",
    "    X_summer, y_summer = data[\"summer_X\"], data[\"summer_y\"]\n",
    "    X_winter, y_winter = data[\"winter_X\"], data[\"winter_y\"]\n",
    "    \n",
    "    plt.figure(figsize=(12,7))\n",
    "    plt.scatter(X_summer, y_summer, marker=\"P\", c=\"sandybrown\", label=\"Summer\")\n",
    "    plt.scatter(X_winter, y_winter, marker=\"P\", c=\"lightskyblue\", label=\"Winter\")\n",
    "    plt.plot(Xmodel, ymodel, linewidth=4, label=\"model predictions\")\n",
    "    plt.title(title)\n",
    "    plt.xlabel(\"Temperature\")\n",
    "    plt.ylabel(\"Usage\")\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the dataset and define the number of folds\n",
    "data = np.load(\"data.npz\")\n",
    "Xraw, y, R = data[f\"full_X\"], data[f\"full_y\"], np.diag(data[f\"full_R\"])\n",
    "kFolds = 2\n",
    "\n",
    "# init the models and train and evaluate using K-Fold Cross Validation\n",
    "models = initModels()\n",
    "bestIdx, scores = selectModel(models, Xraw, y, R, kFolds=kFolds)\n",
    "\n",
    "# plot all models with their dataset\n",
    "for i, m in enumerate(models):\n",
    "    if i < 5:\n",
    "        modelString = f\"PolyWLR(degree={m.degree})\"\n",
    "    else: \n",
    "        n_neighbors = m.knn.get_params()[\"n_neighbors\"]\n",
    "        modelString = f\"KNN(n_neighbors={n_neighbors})\"\n",
    "    title = modelString + f\" --- k-fold score: {round(scores[i],2)}\"\n",
    "    \n",
    "    # fit model fully to training dataset and then plot\n",
    "    m.fit(Xraw, y, R)\n",
    "    plotModelTemp(m, title=title)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q3c - Programming Written Plot: plotting the best model <a class=\"anchor\" name=\"q3c\"></a>\n",
    "Below, we have provided code to plot the best model (according to the k-fold cross validation score). Please provide it in **Q3c** in the writeup of the programming portion of the assignment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# of the models above, plot the best model based on k-fold cross validation score\n",
    "m = models[bestIdx]\n",
    "if bestIdx < 5:\n",
    "    modelString = f\"PolyWLR(degree={m.degree})\"\n",
    "else: \n",
    "    n_neighbors = m.knn.get_params()[\"n_neighbors\"]\n",
    "    modelString = f\"KNN(n_neighbors={n_neighbors})\"\n",
    "title = modelString + f\" --- k-fold score: {round(scores[i],2)}\"\n",
    "\n",
    "# fit model fully to training dataset and then plot\n",
    "m.fit(Xraw, y, R)\n",
    "plotModelTemp(m, title=title)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Submit your code to Gradescope early and often\n",
    "\n",
    "There is no limit on the number of submissions to Gradescope, so as you complete parts of the assignment it is a really good idea to save your notebook and upload it to Gradescope.\n",
    "\n",
    "Not all of the tests are included in the local autograder. Some of the tests are \"hidden\" and only run in the server autograder on Gradescope.\n",
    "\n",
    "Before continuing with the rest of the assignment, go ahead and save your notebook (or click File->Download->Download .ipynb) and then upload your hw5.ipynb file to Gradescope under assignment hw5 (programming)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "280-testing",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
