{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6131203a",
   "metadata": {},
   "source": [
    "# E190BO Homework 1 — Gradient Descent\n",
    "\n",
    "### Fall 2026\n",
    "\n",
    "Submit **this notebook**, run top to bottom with all output visible, to Canvas.\n",
    "If you did any derivations on paper, scan them and attach them as a separate PDF.\n",
    "\n",
    "In your one-on-one check-in meeting, we will ask you questions about your notebook to\n",
    "assess your level of understanding.  Your goal is to understand this notebook deeply.\n",
    "\n",
    "### How this notebook works\n",
    "\n",
    "Almost all of the code is written for you. Everywhere you see\n",
    "\n",
    "```\n",
    "# YOUR CODE HERE (1 line): ...\n",
    "raise NotImplementedError(\"Replace this line with your code.\")\n",
    "```\n",
    "\n",
    "you delete that `raise` line and write your own line (or two, or three — it always says how\n",
    "many) in its place. Every one of those lines is a piece of *math* you derived, translated\n",
    "into numpy. The plotting, the printing, the bookkeeping, and the loops are all provided,\n",
    "because none of that is what this assignment is about.\n",
    "\n",
    "**Everything must be implemented from scratch in numpy.** Do not call\n",
    "`np.polyfit`, `np.linalg.lstsq`, `sklearn`, or any other library that solves the\n",
    "regression for you.\n",
    "\n",
    "### AI policy for this assignment\n",
    "\n",
    "The **course tutor bot** is the permitted AI resource. It will ask you questions\n",
    "and give you hints, but it will not hand you derivations or working code — that is\n",
    "deliberate, and it is the point. Other AI assistants are not permitted on this\n",
    "assignment.\n",
    "\n",
    "Fill in the disclosure cell below. An honest \"I was stuck on the Problem 3 gradient\n",
    "for two hours and the tutor bot walked me through the product rule\" costs you\n",
    "nothing, and it is exactly the sort of thing worth talking about at our check-in."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "870a3238",
   "metadata": {},
   "source": [
    "### Disclosure\n",
    "\n",
    "*Roughly how long did this assignment take you?*\n",
    "\n",
    "*Where did you get stuck, and what did you use to get unstuck (tutor bot, classmates,\n",
    "office hours, notes)?*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cb798f0",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Run this cell first. It loads the exam-score data and standardizes every column\n",
    "(subtract the mean, divide by the standard deviation) so that each column has mean 0\n",
    "and standard deviation 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79624770",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from pathlib import Path\n",
    "\n",
    "# Reads testScores.csv from next to this notebook if you have it, and otherwise\n",
    "# straight from the course site — so the notebook runs on its own.\n",
    "CSV = \"testScores.csv\"\n",
    "if not Path(CSV).exists():\n",
    "    CSV = \"https://course.aiprototypes.org/resources/files/eng190bo-fa26-lab-1-testscores.csv\"\n",
    "\n",
    "df = pd.read_csv(CSV).dropna(axis=1, how=\"all\")\n",
    "z = (df - df.mean()) / df.std(ddof=0)        # standardize every column\n",
    "print(df.shape, \"->\", list(z.columns))\n",
    "z.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f1c398e",
   "metadata": {},
   "source": [
    "### Provided: the gradient checker\n",
    "\n",
    "You will use this tool in all three problems, so it is worth reading the docstring once now.\n",
    "\n",
    "The idea: you derive a gradient by hand and code it up. `grad_check` then estimates the *same*\n",
    "gradient a completely different way — by nudging one parameter at a time and watching how much\n",
    "the loss changes. If your derivation and your code are both correct, the two answers agree to\n",
    "about ten decimal places. If they disagree, something is wrong, and you want to know that\n",
    "*before* you spend an hour debugging a fit that will not converge.\n",
    "\n",
    "Notice the convention used everywhere in this notebook: **every loss function and every\n",
    "gradient function takes the data first and the parameters last.** So `grad_check` can call\n",
    "your functions as `loss_fn(X, y, w)` and `grad_fn(X, y, w)` without knowing anything else\n",
    "about them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95b48a5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- PROVIDED --- do not modify\n",
    "\n",
    "def grad_check(loss_fn, grad_fn, w, X, y, eps=1e-5):\n",
    "    \"\"\"Compare your analytical gradient against a numerical estimate of the same thing.\n",
    "\n",
    "    The definition of a partial derivative says that, for a small step eps,\n",
    "\n",
    "        de/dw_j   is approximately   ( e(w + eps in slot j) - e(w - eps in slot j) ) / (2*eps)\n",
    "\n",
    "    This function computes that estimate for every slot j and compares it to whatever your\n",
    "    gradient function returns.\n",
    "\n",
    "    Inputs\n",
    "    ------\n",
    "    loss_fn : your loss function.  It gets called as loss_fn(X, y, w).\n",
    "    grad_fn : your gradient function.  It gets called as grad_fn(X, y, w).\n",
    "    w       : the parameter values to check at.  Use RANDOM values, not all zeros -- a\n",
    "              wrong gradient can accidentally look correct at a special point like zero.\n",
    "    X, y    : the data.  These are passed straight through to your two functions.\n",
    "    eps     : the size of the small nudge.  The default is fine; leave it alone.\n",
    "\n",
    "    Output\n",
    "    ------\n",
    "    rel_err : one number.  Below 1e-7 counts as a pass.\n",
    "    \"\"\"\n",
    "    w = np.array(w, dtype=float)\n",
    "    analytic = np.array(grad_fn(X, y, w), dtype=float)\n",
    "    numeric = np.zeros(w.size)\n",
    "\n",
    "    for j in range(w.size):\n",
    "        w_plus = w.copy()\n",
    "        w_minus = w.copy()\n",
    "        w_plus[j] = w_plus[j] + eps\n",
    "        w_minus[j] = w_minus[j] - eps\n",
    "        numeric[j] = (loss_fn(X, y, w_plus) - loss_fn(X, y, w_minus)) / (2 * eps)\n",
    "\n",
    "    size = np.linalg.norm(analytic) + np.linalg.norm(numeric)\n",
    "    if size == 0:\n",
    "        rel_err = 0.0\n",
    "    else:\n",
    "        rel_err = np.linalg.norm(analytic - numeric) / size\n",
    "\n",
    "    print(\"  your gradient:\", np.array2string(analytic, precision=6))\n",
    "    print(\"  numerical    :\", np.array2string(numeric, precision=6))\n",
    "    if rel_err < 1e-7:\n",
    "        print(f\"  relative error = {rel_err:.3e}   -> PASS\")\n",
    "    else:\n",
    "        print(f\"  relative error = {rel_err:.3e}   -> FAIL\")\n",
    "    return rel_err"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8ab93b5",
   "metadata": {},
   "source": [
    "---\n",
    "# Problem 1 — Vectorizing the in-class activity\n",
    "\n",
    "In class you fit $\\hat{y}_i = a x_i + b$ with an explicit `agrad` and `bgrad`. That works\n",
    "for two parameters. It does not scale.\n",
    "\n",
    "Here you will rewrite the same fit in **matrix form**. Collect the data into a design matrix\n",
    "$X$ of shape $(N, m)$ where each **row** is one data sample, and let the **last column be all\n",
    "ones** so that the last weight is the intercept:\n",
    "\n",
    "$$\\hat{\\bar y} = X \\bar w, \\qquad e = \\frac{1}{N}(\\bar y - X\\bar w)^T(\\bar y - X\\bar w)$$\n",
    "\n",
    "For this 1-D problem, $m = 2$ and $\\bar w = [a, b]$.\n",
    "\n",
    "Write `compute_loss` and `compute_grad` for a **general** $m$ — you will reuse them\n",
    "unchanged in Problem 2, so do not hard-code $m = 2$.\n",
    "\n",
    "**(a)** Build the design matrix `X1`.\n",
    "**(b)** Implement `compute_loss` and `compute_grad`.\n",
    "**(c)** Verify `compute_grad` with `grad_check`.\n",
    "**(d)** Implement the gradient descent step, and confirm it reproduces your in-class answer."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe801f72",
   "metadata": {},
   "source": [
    "### 1(a) Build the design matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0f9b497",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The data for Problem 1: use EXAM1 to predict FINAL.\n",
    "x = z[\"EXAM1\"].values        # shape (25,)  -- the input\n",
    "y = z[\"FINAL\"].values        # shape (25,)  -- the target\n",
    "N = len(y)                   # 25 data samples\n",
    "\n",
    "# YOUR CODE HERE (1 line): build the design matrix X1, which should have shape (N, 2).\n",
    "# Column 0 is x.  Column 1 is all ones -- that is the intercept column.\n",
    "# Two functions that help: np.column_stack([a, b]) glues 1-D arrays together as columns,\n",
    "# and np.ones(N) makes an array of N ones.\n",
    "raise NotImplementedError(\"Replace this line with your code.\")\n",
    "\n",
    "print(\"X1 shape:\", X1.shape, \" (should be (25, 2))\")\n",
    "print(X1[:3])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd40e97d",
   "metadata": {},
   "source": [
    "### 1(b) The loss and its gradient\n",
    "\n",
    "Derive $\\partial e / \\partial \\bar w$ in matrix form first (it is on the L2 slides), then\n",
    "implement it below. In both functions the residual vector\n",
    "\n",
    "$$\\bar r = \\bar y - X\\bar w$$\n",
    "\n",
    "is computed for you — it is the one quantity both formulas are built out of."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c4e6bc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_loss(X, y, w):\n",
    "    \"\"\"Mean squared error of the linear model yhat = X @ w.\n",
    "\n",
    "    Inputs\n",
    "    ------\n",
    "    X : the design matrix, shape (N, m).  Each ROW is one data sample, and the LAST\n",
    "        column is the column of ones.\n",
    "    y : the targets, shape (N,).\n",
    "    w : the weights, shape (m,).  w[-1] is the intercept.\n",
    "\n",
    "    Output\n",
    "    ------\n",
    "    A single number: the mean squared error.\n",
    "    \"\"\"\n",
    "    residuals = y - X @ w                # shape (N,).  X @ w is the prediction yhat.\n",
    "\n",
    "    # YOUR CODE HERE (1 line): return the MEAN of the squared residuals.\n",
    "    # np.mean(a) averages an array, and a ** 2 squares every entry of it.\n",
    "    raise NotImplementedError(\"Replace this line with your code.\")\n",
    "\n",
    "\n",
    "def compute_grad(X, y, w):\n",
    "    \"\"\"Gradient of compute_loss with respect to w.\n",
    "\n",
    "    Inputs\n",
    "    ------\n",
    "    X, y, w : exactly the same as in compute_loss above.\n",
    "\n",
    "    Output\n",
    "    ------\n",
    "    An array of shape (m,) -- one partial derivative per weight.  It must have the same\n",
    "    shape as w, NOT a single number.\n",
    "    \"\"\"\n",
    "    N = len(y)\n",
    "    residuals = y - X @ w                # shape (N,), same as above\n",
    "\n",
    "    # YOUR CODE HERE (1 line): return the gradient you derived.\n",
    "    # In numpy, X.T is the transpose of X, and A @ b is matrix-vector multiplication.\n",
    "    # No for loop over the N data samples -- the matrix multiply IS the sum over samples.\n",
    "    raise NotImplementedError(\"Replace this line with your code.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ee26ec4",
   "metadata": {},
   "source": [
    "Some checks. These should all pass before you go any further."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7c99243",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- PROVIDED --- do not modify\n",
    "\n",
    "assert np.isclose(compute_loss(X1, y, np.array([0.0, 0.0])), np.mean(y ** 2))\n",
    "assert np.isclose(compute_loss(X1, y, np.array([1.0, 0.0])), np.mean((y - x) ** 2))\n",
    "assert compute_grad(X1, y, np.array([0.5, 0.1])).shape == (2,)\n",
    "np.testing.assert_allclose(compute_grad(X1, y, np.array([0.0, 0.0])),\n",
    "                           [-2 * np.mean(x * y), -2 * np.mean(y)], atol=1e-12)\n",
    "print(\"assertions passed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33ae0c3d",
   "metadata": {},
   "source": [
    "### 1(c) Gradient checking\n",
    "\n",
    "Nothing to write here — just run it. Check at a **random** $\\bar w$, not at $\\bar w = 0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb967a83",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- PROVIDED --- do not modify\n",
    "\n",
    "rng = np.random.default_rng(0)\n",
    "w_test = rng.normal(size=2)\n",
    "\n",
    "err = grad_check(compute_loss, compute_grad, w_test, X1, y)\n",
    "assert err < 1e-7, \"gradient check failed -- fix compute_grad before continuing\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30b8b084",
   "metadata": {},
   "source": [
    "### 1(d) Gradient descent\n",
    "\n",
    "One line to write: the update step itself. Everything around it — starting the weights,\n",
    "recording the loss at every iteration, returning the history for plotting — is done for you.\n",
    "\n",
    "You will call this same function in Problem 2 and again in Problem 3. It never needs to\n",
    "change, because it does not know or care *which* loss and gradient it is being handed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbf0e911",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gradient_descent(loss_fn, grad_fn, X, y, w0, gamma, numIterations):\n",
    "    \"\"\"Batch gradient descent.\n",
    "\n",
    "    Inputs\n",
    "    ------\n",
    "    loss_fn : a loss function, called as loss_fn(X, y, w)\n",
    "    grad_fn : a gradient function, called as grad_fn(X, y, w)\n",
    "    X, y    : the data, passed straight through to loss_fn and grad_fn\n",
    "    w0      : starting values for the parameters\n",
    "    gamma   : the learning rate (the step size)\n",
    "    numIterations : how many steps to take\n",
    "\n",
    "    Outputs\n",
    "    -------\n",
    "    w         : the parameters after the final step\n",
    "    loss_hist : the loss after every step, length numIterations + 1\n",
    "    \"\"\"\n",
    "    w = np.array(w0, dtype=float)\n",
    "    loss_hist = [loss_fn(X, y, w)]\n",
    "\n",
    "    for i in range(numIterations):\n",
    "        # YOUR CODE HERE (1 line): take one step downhill.\n",
    "        # The new w is the old w, minus gamma times the gradient at the current w.\n",
    "        # The gradient at the current w is grad_fn(X, y, w).\n",
    "        raise NotImplementedError(\"Replace this line with your code.\")\n",
    "\n",
    "        loss_hist.append(loss_fn(X, y, w))\n",
    "\n",
    "    return w, np.array(loss_hist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf659ec3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- PROVIDED --- do not modify\n",
    "\n",
    "w1, hist1 = gradient_descent(compute_loss, compute_grad, X1, y,\n",
    "                             w0=np.zeros(2), gamma=0.1, numIterations=500)\n",
    "print(f\"a = {w1[0]:.6f},  b = {w1[1]:.6f}\")\n",
    "\n",
    "# In class you should have found a = 0.946071, b = 0.000000\n",
    "np.testing.assert_allclose(w1, [0.9460708318102029, 0.0], atol=1e-8)\n",
    "print(\"matches the in-class scalar result to 1e-8\")\n",
    "\n",
    "plt.figure(figsize=(7, 4))\n",
    "plt.semilogy(hist1)\n",
    "plt.xlabel(\"iteration\"); plt.ylabel(\"MSE\")\n",
    "plt.title(\"1(d) convergence\"); plt.grid(alpha=.3); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "442b78b2",
   "metadata": {},
   "source": [
    "---\n",
    "# Problem 2 — $m$-dimensional linear regression\n",
    "\n",
    "Now use **all three** standardized midterm scores to predict the standardized final:\n",
    "\n",
    "$$\\hat y_i = w_1 x_{i,1} + w_2 x_{i,2} + w_3 x_{i,3} + b$$\n",
    "\n",
    "Four weights instead of two. Your Problem 1 code should work unchanged — that was the whole\n",
    "point of writing it for a general $m$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a88ac4af",
   "metadata": {},
   "source": [
    "## 2(a) Ordinary least squares\n",
    "\n",
    "One line to write: the design matrix, exactly like 1(a) but with three input columns instead\n",
    "of one. Then run the two provided cells.\n",
    "\n",
    "Afterwards, answer in the markdown cell below: **which midterm carries the most weight, and\n",
    "why is it meaningful to compare these coefficients directly here when usually it is not?**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4930e5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# YOUR CODE HERE (1 line): build the design matrix X, of shape (N, 4).\n",
    "# Columns 0, 1, 2 are EXAM1, EXAM2, EXAM3.  Column 3 is all ones (the intercept).\n",
    "# The three input columns are z[\"EXAM1\"].values, z[\"EXAM2\"].values, z[\"EXAM3\"].values.\n",
    "# Same tools as 1(a): np.column_stack([...]) and np.ones(N).\n",
    "raise NotImplementedError(\"Replace this line with your code.\")\n",
    "\n",
    "print(\"X shape:\", X.shape, \" (should be (25, 4))\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "506c8413",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- PROVIDED --- gradient check, then fit, then plot.  Do not modify.\n",
    "\n",
    "rng = np.random.default_rng(1)\n",
    "err = grad_check(compute_loss, compute_grad, rng.normal(size=4), X, y)\n",
    "assert err < 1e-7\n",
    "\n",
    "w2, hist2 = gradient_descent(compute_loss, compute_grad, X, y,\n",
    "                             w0=np.zeros(4), gamma=0.1, numIterations=2000)\n",
    "print(\"\\nweights [EXAM1, EXAM2, EXAM3, intercept] =\", np.round(w2, 6))\n",
    "print(\"final MSE =\", round(hist2[-1], 6))\n",
    "\n",
    "plt.figure(figsize=(7, 4))\n",
    "plt.semilogy(hist2)\n",
    "plt.xlabel(\"iteration\"); plt.ylabel(\"MSE\")\n",
    "plt.title(\"2(a) OLS convergence\"); plt.grid(alpha=.3); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ee00ed3",
   "metadata": {},
   "source": [
    "*Your answer here.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "198bd747",
   "metadata": {},
   "source": [
    "## 2(b) Ridge regression\n",
    "\n",
    "Ridge regression is a linear model with a slightly different loss function. It is the same MSE,\n",
    "plus a penalty on the size of the weights:\n",
    "\n",
    "$$e_{\\text{ridge}}(\\bar w) = \\frac{1}{N}\\sum_{i=1}^{N}(y_i - \\hat y_i)^2 + \\lambda \\sum_{j=1}^{m-1} w_j^2$$\n",
    "\n",
    "**Note the sum runs to $m-1$: the intercept $b = w_m$ is *not* penalized.** Shrinking the\n",
    "intercept would mean penalizing the model for predicting a nonzero average, which is not\n",
    "something we want to discourage.\n",
    "\n",
    "**(i)** Derive $\\partial e_{\\text{ridge}} / \\partial \\bar w$ by hand. It is your Problem 2(a)\n",
    "gradient plus one extra term. Be careful about the intercept component.\n",
    "**(ii)** Implement it and pass gradient checking.\n",
    "**(iii)** Fit with $\\lambda = 1$ and report the weights.\n",
    "**(iv)** Now fit with $\\lambda = 100$ and report those weights. Compare them to your\n",
    "$\\lambda = 1$ weights: what changed, and what does that tell you about what $\\lambda$ does?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5454d6b0",
   "metadata": {},
   "source": [
    "**(i)** *Your derivation here.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46430521",
   "metadata": {},
   "source": [
    "### 2(b)(ii) Implement the ridge loss and gradient\n",
    "\n",
    "Two lines to write. Both functions reuse the ones you already wrote, so neither needs to\n",
    "recompute anything about the data.\n",
    "\n",
    "The one idea to get right is the intercept. In numpy, `w[:-1]` means \"all of w except the\n",
    "last entry\" — that is exactly the set of weights the penalty applies to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd52bc9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_ridge_loss(X, y, w, lam):\n",
    "    \"\"\"The ridge objective: mean squared error, plus lam times the sum of the squared\n",
    "    NON-intercept weights.\n",
    "\n",
    "    Inputs\n",
    "    ------\n",
    "    X, y : the data, exactly as in compute_loss\n",
    "    w    : the weights, shape (m,).  w[-1] is the intercept.\n",
    "    lam  : lambda, the penalty strength.  A single number.\n",
    "\n",
    "    Output\n",
    "    ------\n",
    "    A single number.\n",
    "    \"\"\"\n",
    "    mse = compute_loss(X, y, w)          # the first term -- already written in Problem 1\n",
    "    penalized = w[:-1]                   # every weight EXCEPT the intercept\n",
    "\n",
    "    # YOUR CODE HERE (1 line): return mse plus lam times the sum of the squared\n",
    "    # entries of \"penalized\".  np.sum(a) adds up an array.\n",
    "    raise NotImplementedError(\"Replace this line with your code.\")\n",
    "\n",
    "\n",
    "def compute_ridge_grad(X, y, w, lam):\n",
    "    \"\"\"Gradient of compute_ridge_loss with respect to w.\n",
    "\n",
    "    Inputs\n",
    "    ------\n",
    "    X, y, w, lam : exactly as in compute_ridge_loss above.\n",
    "\n",
    "    Output\n",
    "    ------\n",
    "    An array of shape (m,), same shape as w.\n",
    "    \"\"\"\n",
    "    g = compute_grad(X, y, w).copy()     # the MSE part of the gradient, from Problem 1\n",
    "\n",
    "    # YOUR CODE HERE (1 line): add the penalty term to the first m-1 entries of g, and\n",
    "    # leave the last entry (the intercept) alone.  Write it as\n",
    "    #     g[:-1] = g[:-1] + <the extra term you derived for w[:-1]>\n",
    "    raise NotImplementedError(\"Replace this line with your code.\")\n",
    "\n",
    "    return g"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30a6405a",
   "metadata": {},
   "source": [
    "Your two ridge functions take a fourth argument, `lam`, but `grad_check` and\n",
    "`gradient_descent` expect functions of exactly `(X, y, w)`. The fix is a pair of tiny\n",
    "wrapper functions that pin `lam` to one value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05f5c224",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- PROVIDED --- do not modify\n",
    "\n",
    "def ridge_loss_lam1(X, y, w):\n",
    "    \"\"\"compute_ridge_loss with lambda pinned to 1.\"\"\"\n",
    "    return compute_ridge_loss(X, y, w, 1.0)\n",
    "\n",
    "\n",
    "def ridge_grad_lam1(X, y, w):\n",
    "    \"\"\"compute_ridge_grad with lambda pinned to 1.\"\"\"\n",
    "    return compute_ridge_grad(X, y, w, 1.0)\n",
    "\n",
    "\n",
    "rng = np.random.default_rng(2)\n",
    "err = grad_check(ridge_loss_lam1, ridge_grad_lam1, rng.normal(size=4), X, y)\n",
    "assert err < 1e-7"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9545305",
   "metadata": {},
   "source": [
    "### 2(b)(iii) Fit with $\\lambda = 1$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d7fec7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- PROVIDED --- do not modify\n",
    "\n",
    "w_ridge1, hist_r1 = gradient_descent(ridge_loss_lam1, ridge_grad_lam1, X, y,\n",
    "                                     w0=np.zeros(4), gamma=0.1, numIterations=2000)\n",
    "print(\"lambda = 1 weights:\", np.round(w_ridge1, 6))\n",
    "print(\"OLS weights       :\", np.round(w2, 6))\n",
    "\n",
    "plt.figure(figsize=(7, 4))\n",
    "plt.semilogy(hist_r1)\n",
    "plt.xlabel(\"iteration\"); plt.ylabel(\"ridge objective\")\n",
    "plt.title(\"2(b) ridge convergence, lambda = 1\"); plt.grid(alpha=.3); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ecb2bae",
   "metadata": {},
   "source": [
    "### 2(b)(iv) Fit with $\\lambda = 100$\n",
    "\n",
    "Same thing with a much larger penalty. Run the cell, then compare the three rows of weights\n",
    "it prints and answer the question in the markdown cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "772843ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- PROVIDED --- do not modify\n",
    "\n",
    "def ridge_loss_lam100(X, y, w):\n",
    "    \"\"\"compute_ridge_loss with lambda pinned to 100.\"\"\"\n",
    "    return compute_ridge_loss(X, y, w, 100.0)\n",
    "\n",
    "\n",
    "def ridge_grad_lam100(X, y, w):\n",
    "    \"\"\"compute_ridge_grad with lambda pinned to 100.\"\"\"\n",
    "    return compute_ridge_grad(X, y, w, 100.0)\n",
    "\n",
    "\n",
    "w_ridge100, hist_r100 = gradient_descent(ridge_loss_lam100, ridge_grad_lam100, X, y,\n",
    "                                         w0=np.zeros(4), gamma=0.004, numIterations=20000)\n",
    "print(\"lambda = 100 weights:\", np.round(w_ridge100, 6))\n",
    "print(\"lambda = 1   weights:\", np.round(w_ridge1, 6))\n",
    "print(\"OLS          weights:\", np.round(w2, 6))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b082fdd",
   "metadata": {},
   "source": [
    "*Your answer here.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d69f1309",
   "metadata": {},
   "source": [
    "---\n",
    "# Problem 3 — Parameter estimation for a nonlinear model\n",
    "\n",
    "Everything so far has been linear in the parameters, which makes the gradient easy and the\n",
    "loss surface a bowl with exactly one minimum. Real models are usually neither.\n",
    "\n",
    "You are given noisy measurements of the response of an underdamped second-order system:\n",
    "\n",
    "$$y = A\\,e^{-bx}\\sin(\\omega x + \\varphi)$$\n",
    "\n",
    "with four unknown parameters: amplitude $A$, decay rate $b$, angular frequency $\\omega$, and\n",
    "phase $\\varphi$. Estimate all four by gradient descent.\n",
    "\n",
    "The structure is the same as Problems 1 and 2, and you will reuse the *same*\n",
    "`gradient_descent` function. The only genuinely new thing is that the model is nonlinear in\n",
    "its parameters, so the four partial derivatives take real work to derive.\n",
    "\n",
    "Throughout Problem 3 the parameters live in one array `p = [A, b, omega, phi]`, playing\n",
    "exactly the role `w` played before.\n",
    "\n",
    "### What to turn in\n",
    "\n",
    "**(a)** Derive $\\partial e/\\partial A$, $\\partial e/\\partial b$, $\\partial e/\\partial \\omega$,\n",
    "$\\partial e/\\partial \\varphi$ by hand, for the MSE loss. You will need the chain rule and,\n",
    "for two of them, the product rule.\n",
    "\n",
    "**(b)** Implement the model and the gradient, and pass gradient checking.\n",
    "\n",
    "**(c)** Fit the model. Plot the loss vs iteration for three different learning rates that all\n",
    "converge.\n",
    "\n",
    "**(d)** Plot the data with your fitted curve overlaid.\n",
    "\n",
    "### Notes — read these, they will save you time\n",
    "\n",
    "- $\\varphi$ is only defined modulo $2\\pi$, and flipping the sign of $A$ while shifting\n",
    "  $\\varphi$ by $\\pi$ gives the *identical* curve. If your parameters look different from a\n",
    "  classmate's but your curves lie on top of each other, you are both right.\n",
    "- A single $\\gamma$ shared across four parameters of quite different scales will converge\n",
    "  slowly. That is expected here."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5b59326",
   "metadata": {},
   "source": [
    "### Provided: the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e6d1ab6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- PROVIDED --- do not modify\n",
    "\n",
    "def generate_data(seed=190):\n",
    "    rng = np.random.default_rng(seed)\n",
    "    A_true, b_true, w_true, p_true = 3.0, 0.6, 4.0, 0.5\n",
    "    xs = np.sort(rng.uniform(0, 4, size=40))\n",
    "    ys = A_true * np.exp(-b_true * xs) * np.sin(w_true * xs + p_true)\n",
    "    ys = ys + rng.normal(0, 0.10, size=xs.shape)\n",
    "    return xs, ys, np.array([A_true, b_true, w_true, p_true])\n",
    "\n",
    "xs, ys, p_true = generate_data()      # xs and ys both have shape (40,)\n",
    "\n",
    "plt.figure(figsize=(7, 4))\n",
    "plt.scatter(xs, ys, s=18)\n",
    "plt.xlabel(\"x\"); plt.ylabel(\"y\")\n",
    "plt.title(\"Problem 3 data\"); plt.grid(alpha=.3); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cee4488",
   "metadata": {},
   "source": [
    "### 3(a) Derivations\n",
    "\n",
    "*Your derivations here (or on paper, attached as a PDF).*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8387ad4",
   "metadata": {},
   "source": [
    "### 3(b) The model, the loss, and the gradient\n",
    "\n",
    "Five short lines to write in total, spread over three functions.\n",
    "\n",
    "The third function is the heart of the assignment. Notice that its structure is identical to\n",
    "`compute_grad` back in Problem 1 — residuals times a derivative, summed over the data, scaled\n",
    "by $-2/N$. The only difference is that here there are four different derivatives instead of\n",
    "one matrix multiply. The three repeated quantities $E$, $S$, $C$ from your 3(a) derivation are\n",
    "computed for you, and the $A$ derivative is worked out as an example of the pattern."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a257121",
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward(x, p):\n",
    "    \"\"\"Evaluate the model at the points in x.\n",
    "\n",
    "    Inputs\n",
    "    ------\n",
    "    x : an array of x values, any shape.\n",
    "    p : the four parameters, in the order [A, b, omega, phi].\n",
    "\n",
    "    Output\n",
    "    ------\n",
    "    An array the same shape as x: the predicted y at each x.\n",
    "    \"\"\"\n",
    "    A, b, omega, phi = p                 # unpack the four parameters into named variables\n",
    "\n",
    "    # YOUR CODE HERE (1 line): return the model outputs corresponding to the inputs in x\n",
    "    # Use np.exp(...) and np.sin(...).\n",
    "    raise NotImplementedError(\"Replace this line with your code.\")\n",
    "\n",
    "\n",
    "def model_loss(x, y, p):\n",
    "    \"\"\"Mean squared error between the measurements y and the model output at x.\n",
    "\n",
    "    Inputs\n",
    "    ------\n",
    "    x : the measured x values, shape (N,)\n",
    "    y : the measured y values, shape (N,)\n",
    "    p : the four parameters [A, b, omega, phi]\n",
    "\n",
    "    Output\n",
    "    ------\n",
    "    A single number.\n",
    "    \"\"\"\n",
    "    residuals = y - forward(x, p)        # shape (N,) -- same idea as in Problem 1\n",
    "\n",
    "    # YOUR CODE HERE (1 line): return the mean of the squared residuals.\n",
    "    raise NotImplementedError(\"Replace this line with your code.\")\n",
    "\n",
    "\n",
    "def model_grad(x, y, p):\n",
    "    \"\"\"Gradient of model_loss with respect to p.\n",
    "\n",
    "    Inputs\n",
    "    ------\n",
    "    x, y, p : exactly as in model_loss above.\n",
    "\n",
    "    Output\n",
    "    ------\n",
    "    An array of shape (4,): the partial derivatives with respect to A, b, omega, phi,\n",
    "    in that order.\n",
    "    \"\"\"\n",
    "    N = len(y)\n",
    "    A, b, omega, phi = p\n",
    "\n",
    "    E = np.exp(-b * x)                   # shape (N,)\n",
    "    S = np.sin(omega * x + phi)          # shape (N,)\n",
    "    C = np.cos(omega * x + phi)          # shape (N,)\n",
    "    residuals = y - A * E * S            # shape (N,) -- this is y - forward(x, p)\n",
    "\n",
    "    # Each of these is d(yhat)/d(parameter), evaluated at every data point, shape (N,).\n",
    "    # The first one is done for you as an example; write the other three from your 3(a)\n",
    "    # derivations.  Note that x, E, S, and C are all arrays, so ordinary * multiplies\n",
    "    # them entry by entry -- there is no loop over the data anywhere here.\n",
    "    dyhat_dA = E * S\n",
    "\n",
    "    # YOUR CODE HERE (3 lines): set dyhat_db, dyhat_domega, and dyhat_dphi.\n",
    "    raise NotImplementedError(\"Replace this line with your code.\")\n",
    "\n",
    "    # Combine: each entry is -2/N times the sum over data points of residual * derivative.\n",
    "    # The @ operator between two 1-D arrays does exactly that sum of products.\n",
    "    return -2.0 / N * np.array([residuals @ dyhat_dA,\n",
    "                                residuals @ dyhat_db,\n",
    "                                residuals @ dyhat_domega,\n",
    "                                residuals @ dyhat_dphi])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a9e20c8",
   "metadata": {},
   "source": [
    "Gradient check, at three random starting points. Do not move on until it passes — a\n",
    "wrong gradient will cost you far more time in 3(c) than it costs you to find it here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dc3c395",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- PROVIDED --- do not modify\n",
    "\n",
    "rng = np.random.default_rng(7)\n",
    "for trial in range(3):\n",
    "    p_random = np.array([rng.uniform(0.5, 5.0), rng.uniform(0.0, 2.0),\n",
    "                         rng.uniform(0.5, 10.0), rng.uniform(0.0, 2 * np.pi)])\n",
    "    print(f\"trial {trial}, p = {np.round(p_random, 3)}\")\n",
    "    err = grad_check(model_loss, model_grad, p_random, xs, ys)\n",
    "    assert err < 1e-7\n",
    "\n",
    "print(\"\\ngradient check passes at every test point\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d05f0131",
   "metadata": {},
   "source": [
    "### 3(c) Fit the model with three different learning rates\n",
    "\n",
    "One line to write: the call to your `gradient_descent` from Problem 1. It works here without\n",
    "any changes, because `model_loss` and `model_grad` follow the same `(data, data, parameters)`\n",
    "convention that `compute_loss` and `compute_grad` did.\n",
    "\n",
    "The starting point is provided. It is not arbitrary: the scatter plot shows about 2.5 cycles\n",
    "over $x \\in [0, 4]$, so the period is roughly $1.6$ and $\\omega = 2\\pi/1.6 \\approx 3.9$. Make\n",
    "sure you see where that came from."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0b29bd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_start = np.array([3.0, 0.5, 3.9, 0.0])       # [A, b, omega, phi]\n",
    "\n",
    "plt.figure(figsize=(7, 4))\n",
    "for gamma in [0.002, 0.01, 0.05]:\n",
    "    # YOUR CODE HERE (1 line): run gradient descent from p_start with this gamma for\n",
    "    # 20000 iterations, storing the results in p_fit and hist, like this:\n",
    "    #     p_fit, hist = gradient_descent(<loss fn>, <grad fn>, <x data>, <y data>,\n",
    "    #                                    <start>, <learning rate>, <num iterations>)\n",
    "    raise NotImplementedError(\"Replace this line with your code.\")\n",
    "\n",
    "    print(f\"gamma = {gamma:<7} ->  p = {np.round(p_fit, 4)},  final MSE = {hist[-1]:.4f}\")\n",
    "    plt.semilogy(hist, label=f\"gamma = {gamma}\")\n",
    "\n",
    "plt.xlabel(\"iteration\"); plt.ylabel(\"MSE\"); plt.legend(); plt.grid(alpha=.3)\n",
    "plt.title(\"3(c) three converging learning rates\"); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4519753",
   "metadata": {},
   "source": [
    "### 3(d) Plot the fit on top of the data\n",
    "\n",
    "One line to write: evaluate your fitted model on the smooth grid `x_grid` so it can be\n",
    "plotted as a curve. The numbers printed above the plot let you compare your four estimates\n",
    "against the values the data was actually generated from."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e11ffdca",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_best, hist_best = gradient_descent(model_loss, model_grad, xs, ys, p_start, 0.05, 20000)\n",
    "print(\"your estimate:\", np.round(p_best, 4))\n",
    "print(\"true values  :\", np.round(p_true, 4))\n",
    "\n",
    "x_grid = np.linspace(0, 4, 400)        # 400 evenly spaced x values, for a smooth curve\n",
    "\n",
    "# YOUR CODE HERE (1 line): evaluate your fitted model on x_grid, storing it in y_curve.\n",
    "# Your forward function accepts a whole array of x values at once.\n",
    "raise NotImplementedError(\"Replace this line with your code.\")\n",
    "\n",
    "plt.figure(figsize=(7, 4))\n",
    "plt.scatter(xs, ys, s=18, label=\"data\")\n",
    "plt.plot(x_grid, y_curve, \"r-\", lw=2, label=\"fitted\")\n",
    "plt.xlabel(\"x\"); plt.ylabel(\"y\"); plt.legend(); plt.grid(alpha=.3)\n",
    "plt.title(\"3(d) fitted damped sinusoid\"); plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
