{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "048e86d0",
   "metadata": {},
   "source": [
    "# 6.8300 Problem Set 2\n",
    "\n",
    "This homework will cover\n",
    "\n",
    "- How to steer a sine wave\n",
    "- How to derive the 2D Fourier Transform from first principles\n",
    "- How to properly compute image derivatives\n",
    "- How to orient convolutions in any direction with [steerable filters]( https://people.csail.mit.edu/billf/publications/Design_and_Use_of_Steerable_Filters.pdf )\n",
    "\n",
    "It will be graded mostly on an autograder with some manual review of images!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78d12886",
   "metadata": {},
   "source": [
    "## Part 0: Imports + Helpers\n",
    "- Please pip install any of the libraries if you are missing anything (having the most recent imports should work!).\n",
    "- Please feel free to use these implemented functions anywhere in the notebook. Note some of these functions may use functions that you need to properly implement to use correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87b3654c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jaxtyping import Float, Complex\n",
    "from torch import Tensor\n",
    "from typing import Tuple, Callable\n",
    "import matplotlib.animation as animation\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "from IPython.display import HTML\n",
    "import torch.nn.functional as F\n",
    "from skimage.color import rgb2gray\n",
    "from skimage.data import astronaut as _astronaut\n",
    "from skimage.data import hubble_deep_field\n",
    "from skimage.util import img_as_float\n",
    "from scipy.ndimage import gaussian_filter as gaussian_filter_scipy\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1a864f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_animation():\n",
    "    \"\"\"\n",
    "    Return an animation of a sine wave being steered \n",
    "    by changing the coefficients a and b.\n",
    "    \"\"\"\n",
    "    x = torch.linspace(0, 2 * np.pi, 200)\n",
    "    x_np = x.numpy()\n",
    "    fig, ax = plt.subplots(figsize=(6, 4))\n",
    "    (line_a,) = ax.plot(x_np, np.zeros_like(x_np), label=r\"$a\\cos(x)$\", color=\"C0\")\n",
    "    (line_b,) = ax.plot(x_np, np.zeros_like(x_np), label=r\"$b\\sin(x)$\", color=\"C1\")\n",
    "    (line_sum,) = ax.plot(\n",
    "        x_np, np.zeros_like(x_np), label=r\"$a\\cos(x)+b\\sin(x)$\", color=\"C2\", linewidth=2\n",
    "    )\n",
    "    ax.set_xlim(0, 2 * np.pi)\n",
    "    ax.set_ylim(-1.5, 1.5)\n",
    "    ax.legend(loc=\"upper right\")\n",
    "    angles = np.linspace(0, 2 * np.pi, 60)\n",
    "\n",
    "    def update(frame):\n",
    "        angle = angles[frame]\n",
    "        a, b = angle_to_coeffs(torch.tensor(angle))\n",
    "\n",
    "        y_a = a * torch.cos(x)\n",
    "        y_b = b * torch.sin(x)\n",
    "        y_sum = coeffs_to_sine(a, b, x)\n",
    "\n",
    "        line_a.set_ydata(y_a.numpy())\n",
    "        line_b.set_ydata(y_b.numpy())\n",
    "        line_sum.set_ydata(y_sum.numpy())\n",
    "\n",
    "        ax.set_title(f\"Sine Steering (a = {a:.2f}, b = {b:.2f})\")\n",
    "        return line_a, line_b, line_sum\n",
    "\n",
    "    ani = animation.FuncAnimation(\n",
    "        fig, update, frames=len(angles), interval=100, blit=True\n",
    "    )\n",
    "    return ani"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a251253b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gaussian_filter(\n",
    "    img: torch.Tensor,  # The input tensor\n",
    "    sigma: float,  # Standard deviation for the Gaussian kernel\n",
    "    order: int | list = 0,  # The order of the filter's derivative along each dim\n",
    "    mode: str = \"reflect\",  # Padding mode for `torch.nn.functional.pad`\n",
    "    truncate: float = 4.0,  # Number of standard deviations to sample the filter\n",
    ") -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    Convolves an image with a Gaussian kernel (or its derivatives).\n",
    "    \"\"\"\n",
    "\n",
    "    # Specify the dimensions of the convolution to use\n",
    "    ndim = img.ndim - 2\n",
    "    if isinstance(order, int):\n",
    "        order = [order] * ndim\n",
    "    else:\n",
    "        assert len(order) == ndim, \"Specify the Gaussian derivative order for each dim\"\n",
    "    convfn = getattr(F, f\"conv{ndim}d\")\n",
    "\n",
    "    # Convolve along the rows, columns, and depth (optional)\n",
    "    for dim, derivative_order in enumerate(order):\n",
    "        img = _conv(img, convfn, sigma, derivative_order, truncate, mode, dim)\n",
    "    return img\n",
    "\n",
    "\n",
    "def _conv(\n",
    "    img: torch.Tensor,\n",
    "    convfn: Callable,\n",
    "    sigma: float,\n",
    "    order: int,\n",
    "    truncate: float,\n",
    "    mode: str,\n",
    "    dim: int,\n",
    "):\n",
    "    # Make a 1D kernel and pad such that the image size remains the same\n",
    "    kernel = _gaussian_filter_1d(sigma, order, truncate, img.dtype, img.device)\n",
    "    padding = len(kernel) // 2\n",
    "\n",
    "    # Specify the padding dimensions\n",
    "    pad = [0] * 2 * (img.ndim - 2)\n",
    "    for idx in range(2 * dim, 2 * dim + 2):\n",
    "        pad[idx] = padding\n",
    "    pad = pad[::-1]\n",
    "    x = F.pad(img, pad, mode=mode)\n",
    "\n",
    "    # Specify the dimension along which to do the convolution\n",
    "    view = [1] * img.ndim\n",
    "    view[dim + 2] *= -1\n",
    "\n",
    "    return convfn(x, weight=kernel.view(*view))\n",
    "\n",
    "\n",
    "def astronaut(dtype: torch.dtype = torch.float32):\n",
    "    img = _astronaut()\n",
    "    img = img_as_float(img)\n",
    "    img = rgb2gray(img)\n",
    "    img = torch.from_numpy(img).to(dtype)\n",
    "    return img[None, None]\n",
    "\n",
    "\n",
    "def hubble(dtype: torch.dtype = torch.float32):\n",
    "    img = hubble_deep_field()[0:500, 0:500]\n",
    "    img = img_as_float(img)\n",
    "    img = rgb2gray(img)\n",
    "    img = torch.from_numpy(img).to(dtype)\n",
    "    return img[None, None]\n",
    "\n",
    "\n",
    "def imshow(*imgs):\n",
    "    imgs = [img.squeeze().detach().cpu() for img in imgs]\n",
    "    n_imgs = len(imgs)\n",
    "    fig, axs = plt.subplots(ncols=n_imgs, figsize=(n_imgs * 2 + 2, 3))\n",
    "    for ax, img in zip(axs, imgs):\n",
    "        ax.imshow(img, cmap=\"gray\")\n",
    "        ax.axis(\"off\")\n",
    "    plt.tight_layout()\n",
    "    return fig, axs\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2641cb5b",
   "metadata": {},
   "source": [
    "## Part 1: Steerable Bases\n",
    "> Steerable bases refer to a set of basis functions that can be linearly combined to produce a function in the same family, oriented (or phased) in any direction.\n",
    "\n",
    "In this part, we consider the simplest case: 1D sinusoids. We will see that a sine and cosine of the same frequency form a steerable basis for sinusoids. Any linear combination of a sine and a cosine (of the same frequency) is just another sine wave of that same frequency, with a different phase shift. In other words, by adjusting the coefficients in front of sin(x) and cos(x), we are “steering” the phase of the resulting sine wave without changing its frequency."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb96b191",
   "metadata": {},
   "source": [
    "### Part 1.1 Sine + Cosine = Sine\n",
    "\n",
    "**(a)** Let’s derive the central identity for sine steering.  \n",
    "Given the basis functions $\\sin(x)$ and $\\cos(x)$, show algebraically that the linear combination $a\\cos(x) + b\\sin(x)$ is a sinusoid of the same frequency.\n",
    "\n",
    "In other words, solve for $R$ and $\\theta$ such that\n",
    "\n",
    "$$\n",
    "y(x) = a\\cos(x) + b\\sin(x) = R\\sin(x + \\theta).\n",
    "$$\n",
    "Hint: $\\sin(a+b) = \\sin(a)\\cos(b) + \\cos(a)\\sin(b).$\n",
    "\n",
    "\n",
    "With these results, complete the following function:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc84ff55",
   "metadata": {},
   "outputs": [],
   "source": [
    "def identity(\n",
    "    a: Float[Tensor, \"...\"], b: Float[Tensor, \"...\"]\n",
    ") -> tuple[Float[Tensor, \"...\"], Float[Tensor, \"...\"]]:\n",
    "    \"\"\" \n",
    "    Compute and return the amplitude and phase shift for the problem.\n",
    "    \"\"\"\n",
    "    # STUDENT CODE HERE\n",
    "    \n",
    "    # END STUDENT CODE\n",
    "    return R, theta"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fd45d11",
   "metadata": {},
   "source": [
    "### Part 1.2 Steer a sine!\n",
    "Now that we can rewrite our functions in terms of a single sine, we can control its amplitude and phase using the coefficients a and b. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7482d3a1",
   "metadata": {},
   "source": [
    "(a) Complete the function coeffs_to_sine such that, given the coefficients a and b, it returns the sine from Part 1.1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45d6f0db",
   "metadata": {},
   "outputs": [],
   "source": [
    "def coeffs_to_sine(\n",
    "    a: Float[Tensor, \"...\"], b: Float[Tensor, \"...\"], x: Float[Tensor, \"...\"]\n",
    ") -> Float[Tensor, \"...\"]:\n",
    "    \"\"\"\n",
    "    Computes a linear combination of cosine and sine functions given coefficients a and b.\n",
    "\n",
    "    The function returns the value a*cos(x) + b*sin(x) for each element in x.\n",
    "\n",
    "    Args:\n",
    "        a: Coefficient for the cosine term.\n",
    "        b: Coefficient for the sine term.\n",
    "        x: Tensor of angles.\n",
    "\n",
    "    Returns:\n",
    "        A tensor of the same shape as x representing a*cos(x) + b*sin(x).\n",
    "    \"\"\"\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "302e8d4a",
   "metadata": {},
   "source": [
    "(b) Complete the function angle_to_coeffs such that, given an angle (in radians), it returns the a and b coefficients in the unit circle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bfb6134",
   "metadata": {},
   "outputs": [],
   "source": [
    "def angle_to_coeffs(\n",
    "    angle: Float[Tensor, \"...\"],\n",
    ") -> Tuple[Float[Tensor, \"...\"], Float[Tensor, \"...\"]]:\n",
    "    \"\"\"\n",
    "    Converts an angle to its corresponding cosine and sine coefficients.\n",
    "\n",
    "    This function computes the cosine and sine of the provided angle, effectively mapping the angle\n",
    "    to a point on the unit circle.\n",
    "\n",
    "    Args:\n",
    "        angle: Tensor of angles.\n",
    "\n",
    "    Returns:\n",
    "        A tuple (cos(angle), sin(angle)) of tensors, each with the same shape as the input.\n",
    "    \"\"\"\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return cos_angle, sin_angle"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c74d712f",
   "metadata": {},
   "source": [
    "(c) Run the cell below to show how the sine is steered by varying our coefficients! Confirm your code works by comparing the output to /outputs/1.2 (same for rest of images produced!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "526495f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "ani = return_animation()\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e3e5716",
   "metadata": {},
   "source": [
    "## Part 2: The Fourier Transform\n",
    "> How does the Fourier Transform arise naturally from our image processing toolkit?\n",
    "\n",
    "Note: Some of the graphics in this section are taken from this [blog post](https://medium.com/data-science/deriving-convolution-from-first-principles-4ff124888028).\n",
    "\n",
    "In Lecture 3, we discussed how images can be viewed as functions in an inner product space. As such, they can be represented as a linear combination of basis functions that span our image space. We saw two different bases:\n",
    "- The pixel basis, where each basis function is the image with a single pixel set to 1 and all others to 0.\n",
    "- The (real) Fourier basis, where each basis function is a sine or cosine wave of a certain frequency.\n",
    "\n",
    "The upshot of being able to represent images in these bases is that different linear operators lend themselves to being more interpretable / easier to compute in different bases. In fact, in Lecture 4 we saw that we can derive the Fourier Transform as precisely the change of basis in the inner product space where operators that preserve translational symmetry in the image space are diagonalized.\n",
    "\n",
    "An example of such an operator is the blurring convolution, as seen below:\n",
    "\n",
    "<img src=\"img/blur_shift.png\" width=\"50%\">\n",
    "\n",
    "So, intuitively, the Fourier basis is precisely the basis in which this blur (and the other shift-commutative operators) is easy to compute and interpret!\n",
    "\n",
    "If you have encountered the Fourier Transform before (e.g., in a traditional signals course), the above characterization might not be immediately obvious, so let's take a look at how we arrive at this."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43b85d1e",
   "metadata": {},
   "source": [
    "### Part 2.1: Torus Shifts\n",
    "\n",
    "When applying our image processing operations on 2D images, we assume circular boundaries. For instance, when we shift an image horizontally or vertically, we expect the pixels that are \"pushed outside the borders\" to loop back around to the other side of the image.\n",
    "\n",
    "Effectively, we are \"glueing\" the horizontal and vertical edges of the image together. Topologically, this corresponds to constructing a torus $T^2 = S^1 \\times S^1$, as seen below:\n",
    "\n",
    "![](img/torus_construction.gif)\n",
    "\n",
    "We also saw in Lecture how shifts applied as linear operators on a 1D signal with $n$ discrete entries take the form of a $n \\times n$ permutation matrix that looks like a shifted identity:\n",
    "\n",
    "<img src=\"img/1D_shift.png\" width=\"60%\">\n",
    "\n",
    "With these facts, implement the function shift_operator, which extends this idea to 2D images. This function takes an image shape $(h, w)$ and how many pixels it should be shifted horizontally/vertically, and returns the linear operator corresponding to the shift that operates on a flattened image with circular boundary conditions (that is, a $hw \\times hw$ matrix)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f54d995b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def shift_operator(img_shape: Tuple[int, int], shift_x: int, shift_y: int) -> Tensor:\n",
    "    \"\"\"\n",
    "    Constructs a 2D shift operator for an image with circular boundaries.\n",
    "\n",
    "    Args:\n",
    "        img_shape: Tuple[int, int]\n",
    "            The (height, width) dimensions of the image.\n",
    "        shift_x: int\n",
    "            The number of pixels to shift horizontally.\n",
    "        shift_y: int\n",
    "            The number of pixels to shift vertically.\n",
    "\n",
    "    Returns:\n",
    "        Tensor of shape (h*w, h*w)\n",
    "            A matrix that, when applied to a flattened image, shifts it by the specified amounts.\n",
    "    \"\"\"\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return shift_matrix\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a20769d4",
   "metadata": {},
   "source": [
    "See /outputs/2.1 for what the below function should look like (same for rest of notebook)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86bd66da",
   "metadata": {},
   "outputs": [],
   "source": [
    "img = torch.tensor(plt.imread(\"img/dog.png\"))[..., 0]\n",
    "\n",
    "S = shift_operator(img.shape, shift_x=15, shift_y=10)  # Experiment with these values!\n",
    "shifted_img = (S @ img.flatten()).reshape(img.shape)\n",
    "\n",
    "# Visualize the original and shifted images\n",
    "plt.subplot(121)\n",
    "plt.imshow(img, cmap=\"gray\")\n",
    "plt.title(\"Original Image\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "plt.subplot(122)\n",
    "plt.imshow(shifted_img, cmap=\"gray\")\n",
    "plt.title(\"Shifted Image\")\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "611bf9aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test your implementation against the explicit method for shifting a matrix\n",
    "shift_x = torch.randint(-100, 100, (1,)).item()\n",
    "shift_y = torch.randint(-100, 100, (1,)).item()\n",
    "\n",
    "torch.allclose(\n",
    "    (shift_operator(img.shape, shift_x, shift_y) @ img.flatten()),\n",
    "    torch.roll(img, [-shift_x, -shift_y], dims=(1, 0)).flatten(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c11c8071",
   "metadata": {},
   "source": [
    "### Part 2.2: Convolutions in the Torus\n",
    "\n",
    "In our quest for the Fourier transform, the next step is to examine linear operators that commute with shifts, namely convolutions. In the one-dimensional case, we saw how convolving a kernel with a discrete signal of size $n$ corresponds to appplying a linear operator that takes the form of a $n \\times n$ circulant matrix.\n",
    "\n",
    "<img src=\"img/circulant.png\" width=\"60%\">\n",
    "\n",
    "Implement the function matrix_from_convolution_kernel, which constructs the circulant matrix of size $n \\times n$ given a convolution kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af6e56f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def matrix_from_convolution_kernel(\n",
    "    kernel: Float[Tensor, \"*\"], n: int\n",
    ") -> Float[Tensor, \"n n\"]:\n",
    "    \"\"\"\n",
    "    Constructs a circulant matrix of size n x n from a 1D convolution kernel with periodic alignment.\n",
    "\n",
    "    Args:\n",
    "        kernel: Tensor\n",
    "            A 1D convolution kernel.\n",
    "        n: int\n",
    "            The desired size of the circulant matrix.\n",
    "\n",
    "    Returns:\n",
    "        Tensor of shape (n, n)\n",
    "            The circulant matrix representing the convolution with periodic boundary conditions.\n",
    "    \"\"\"\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return circulant_matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2edcdc20",
   "metadata": {},
   "source": [
    "Similar to what we did with the shift operator in Part 2.1, we want to \"lift\" the convolution operator to act on our two-dimensional flattened images living on the torus. It turns out that, similar to how we can separate a 2D shift into a vertical and horizontal shift, we can often separate a 2D convolution filter into two 1D filters (see [separable filters](https://en.wikipedia.org/wiki/Separable_filter)).\n",
    "\n",
    "Implement the function image_operator_from_sep_kernels, which returns the $hw \\times hw$ matrix corresponding to the linear operator of the 2D convolution given by its separable components applied to the flattened image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d633cd59",
   "metadata": {},
   "outputs": [],
   "source": [
    "def image_operator_from_sep_kernels(\n",
    "    img_shape: Tuple[int, int],\n",
    "    kernel_x: Float[Tensor, \"*\"],\n",
    "    kernel_y: Float[Tensor, \"*\"],\n",
    ") -> Float[Tensor, \"N N\"]:\n",
    "    \"\"\"\n",
    "    Constructs a 2D convolution operator for an image by combining separable 1D kernels.\n",
    "\n",
    "    Args:\n",
    "        img_shape: Tuple[int, int]\n",
    "            The (height, width) dimensions of the image.\n",
    "        kernel_x: Tensor\n",
    "            The 1D convolution kernel to be applied horizontally.\n",
    "        kernel_y: Tensor\n",
    "            The 1D convolution kernel to be applied vertically.\n",
    "\n",
    "    Returns:\n",
    "        Tensor of shape (h*w, h*w)\n",
    "            The 2D convolution operator acting on a flattened image.\n",
    "    \"\"\"\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return conv_operator\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0675fa4c",
   "metadata": {},
   "source": [
    "See /outputs/2.2 for what the below cell should look like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e508104a",
   "metadata": {},
   "outputs": [],
   "source": [
    "img = torch.tensor(plt.imread(\"img/dog.png\"))[..., 0]\n",
    "\n",
    "# kernel = torch.tensor([1.0, -2.0, 1.0])  # Laplacian\n",
    "kernel = torch.tensor([1, 1, 1]) / 3  # Blur\n",
    "M = image_operator_from_sep_kernels(img.shape, kernel, kernel)\n",
    "\n",
    "convolution_output = M @ img.flatten()\n",
    "\n",
    "# Visualize the original and shifted images\n",
    "plt.subplot(121)\n",
    "plt.imshow(img, cmap=\"gray\")\n",
    "plt.title(\"Original Image\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "plt.subplot(122)\n",
    "plt.imshow(convolution_output.reshape(img.shape), cmap=\"gray\")\n",
    "plt.title(\"Convolution Operator on Image\")\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a496221",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note that shifts and convolutions commute!\n",
    "torch.allclose(S @ M, M @ S)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63c57aa2",
   "metadata": {},
   "source": [
    "Note that both our 1D circulant matrices and 2D convolution operators are symmetric matrices!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e26907a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "kernel = torch.tensor([1, 1, 1]) / 3  # Blur\n",
    "\n",
    "circulant_check = matrix_from_convolution_kernel(kernel, 10)\n",
    "M_check = image_operator_from_sep_kernels((10, 10), kernel, kernel)\n",
    "\n",
    "plt.subplot(121)\n",
    "plt.imshow(circulant_check)\n",
    "plt.title(\"1D Circulant Matrix (Blur)\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "plt.subplot(122)\n",
    "plt.imshow(M_check)\n",
    "plt.title(\"2D Convolution Operator (Blur)\")\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56687a0b",
   "metadata": {},
   "source": [
    "### Part 2.3: Spectral Analysis\n",
    "\n",
    "Recall that many important operators in signal and image processing (such as convolution and blurring) can be understood by studying their eigenvalues and eigenvectors. In particular, the eigenvectors describe the “modes” of the system, while the eigenvalues indicate how strongly each mode is amplified or attenuated.\n",
    "\n",
    "Implement eigendecomposition, which computes the eigenvalues and eigenvectors of a given self-adjoint linear operator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a31f7287",
   "metadata": {},
   "outputs": [],
   "source": [
    "def eigendecomposition(\n",
    "    operator: Float[Tensor, \"N N\"], descending: bool = True\n",
    ") -> Tuple[Float[Tensor, \"N\"], Float[Tensor, \"N N\"]]:\n",
    "    \"\"\"\n",
    "    Computes the eigenvalues and eigenvectors of a self-adjoint (Hermitian) linear operator.\n",
    "\n",
    "    Args:\n",
    "        operator: Tensor of shape (N, N)\n",
    "            A self-adjoint linear operator.\n",
    "        descending: bool\n",
    "            If True, sort the eigenvalues and eigenvectors in descending order.\n",
    "\n",
    "    Returns:\n",
    "        A tuple (eigenvalues, eigenvectors) where:\n",
    "            eigenvalues: Tensor of shape (N,)\n",
    "            eigenvectors: Tensor of shape (N, N)\n",
    "    \"\"\"\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return eigenvalues, eigenvectors\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d931b0c",
   "metadata": {},
   "source": [
    "By inspecting the eigendecomposition of the convolution operator, we can recreate the figures from the Lecture!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1860c089",
   "metadata": {},
   "outputs": [],
   "source": [
    "# DECOMPOSITION CHECK\n",
    "# Note the kernel is smaller so that we can see the eigenvectors better\n",
    "M = image_operator_from_sep_kernels((10, 10), kernel, kernel)\n",
    "D, Q = eigendecomposition(M, descending=True)\n",
    "\n",
    "plt.figure(figsize=(12, 4))\n",
    "\n",
    "# Plot the eigenvalues using a scatter plot\n",
    "plt.subplot(131)\n",
    "plt.plot(D, \"o\", markersize=1)\n",
    "plt.title(\"Eigenvalues\")\n",
    "plt.xlabel(\"Eigenvector Index\")\n",
    "plt.ylabel(\"Eigenvalue\")\n",
    "\n",
    "# Plot all the eigenvectors\n",
    "plt.subplot(132)\n",
    "im = plt.imshow(Q, aspect=\"auto\", cmap=\"gray\")\n",
    "plt.title(\"Eigenvectors\")\n",
    "plt.xlabel(\"Eigenvector index\")\n",
    "plt.ylabel(\"Component index\")\n",
    "\n",
    "# Plot a few of the eigenvectors\n",
    "plt.subplot(133)\n",
    "plt.plot(D[1] * Q[:, 1], label=\"Eigenvector 2\")\n",
    "plt.plot(D[2] * Q[:, 2], linestyle=\"--\", label=\"Eigenvector 3\")\n",
    "plt.legend()\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc0e82d6",
   "metadata": {},
   "source": [
    "### Part 2.4: The Fourier Transform!\n",
    "\n",
    "With the eigendecomposition of our operator, we can change the basis of any flattened image or linear operator to be expressed in terms of the eigenvectors of our operator. For our 2D convolution, this is precisely the 2D Fourier transform!\n",
    "\n",
    "Implement fourier_transform, which projects a flattened image to an eigenbasis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55c3514a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fourier_transform(\n",
    "    img: Float[Tensor, \"N\"], basis: Float[Tensor, \"N N\"]\n",
    ") -> Float[Tensor, \"N\"]:\n",
    "    \"\"\"\n",
    "    Projects a flattened image onto the Fourier (eigen) basis.\n",
    "\n",
    "    Args:\n",
    "        img: Tensor of shape (N,)\n",
    "            A flattened image.\n",
    "        basis: Tensor of shape (N, N)\n",
    "            The Fourier eigenbasis.\n",
    "\n",
    "    Returns:\n",
    "        Tensor of shape (N,)\n",
    "            The image represented in the Fourier domain.\n",
    "    \"\"\"\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return img_fourier"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d425f7d7",
   "metadata": {},
   "source": [
    "Implement fourier_transform_operator, which applies a change of basis defined by an eigenbasis to a given linear operator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2cbf667",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fourier_transform_operator(\n",
    "    operator: Float[Tensor, \"N N\"], basis: Float[Tensor, \"N N\"]\n",
    ") -> Float[Tensor, \"N N\"]:\n",
    "    \"\"\"\n",
    "    Computes the representation of a linear operator in the Fourier (eigen) basis.\n",
    "\n",
    "    Args:\n",
    "        operator: Tensor of shape (N, N)\n",
    "            The original linear operator in pixel space.\n",
    "        basis: Tensor of shape (N, N)\n",
    "            The Fourier eigenbasis.\n",
    "\n",
    "    Returns:\n",
    "        Tensor of shape (N, N)\n",
    "            The operator represented in the Fourier basis.\n",
    "    \"\"\"\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return operator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4aec6679",
   "metadata": {},
   "source": [
    "Implement inv_fourier_transform, which returns a flattened image projected onto an eigenbasis back to pixel space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efe6020b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def inv_fourier_transform(\n",
    "    fourier_img: Float[Tensor, \"N\"], basis: Float[Tensor, \"N N\"]\n",
    ") -> Float[Tensor, \"N\"]:\n",
    "    \"\"\"\n",
    "    Reconstructs an image in pixel space from its Fourier coefficients using the provided eigenbasis.\n",
    "\n",
    "    Args:\n",
    "        fourier_img: Tensor of shape (N,)\n",
    "            The image in the Fourier domain.\n",
    "        basis: Tensor of shape (N, N)\n",
    "            The Fourier eigenbasis used in the forward transform.\n",
    "\n",
    "    Returns:\n",
    "        Tensor of shape (N,)\n",
    "            The reconstructed image in pixel space.\n",
    "    \"\"\"\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return img_reconstructed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d70b921",
   "metadata": {},
   "source": [
    "Run the cell below to notice how the function fourier_transform_operator creates operators that are (block)-diagonalized. \n",
    "\n",
    "This is because convolution in the spatial domain corresponds to element-wise multiplication in the frequency domain. Each block corresponds to a specific frequency range, and the operator acts independently on each frequency component. This structure is highly desirable because it makes computations much more efficient (similar to how FFT makes convolution faster), as operations on diagonal matrices have O(n) complexity! Additionally, this block-diagonal structure makes it easier to analyze the effect of the operator on different frequency components, which is useful for designing filters and understanding their behaviors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "072374f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For the convolution operator (here, M) we expect a diagonal matrix\n",
    "M = image_operator_from_sep_kernels(img.shape, kernel, kernel)\n",
    "_, Q = eigendecomposition(M)\n",
    "\n",
    "M_fourier = fourier_transform_operator(M, Q)\n",
    "diag_error = torch.norm(M_fourier - torch.diag(torch.diag(M_fourier)))\n",
    "print(\"Off-diagonal norm for convolution operator:\", diag_error)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "517b9792",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For the shift operator:\n",
    "S_fourier = fourier_transform_operator(S, Q)\n",
    "\n",
    "# For another convolution:\n",
    "kernel = torch.tensor([1.0, -2.0, 1.0])  # Laplacian\n",
    "# kernel = torch.tensor([1, 1, 1]) / 3 # Blur\n",
    "M1 = image_operator_from_sep_kernels(img.shape, kernel, kernel)\n",
    "M1_fourier = fourier_transform_operator(M1, Q)\n",
    "\n",
    "# For their composition:\n",
    "M2 = S @ M1\n",
    "M2_fourier = fourier_transform_operator(M2, Q)\n",
    "\n",
    "plt.figure(figsize=(12, 4))\n",
    "plt.subplot(131)\n",
    "plt.imshow(S_fourier[-100:, -100:], cmap=\"bwr\", vmin=-1, vmax=1)\n",
    "plt.title(\"FT Shift\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "plt.subplot(132)\n",
    "plt.imshow(M1_fourier[-100:, -100:], cmap=\"bwr\", vmin=-1, vmax=1)\n",
    "plt.title(\"FT Convolution\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "plt.subplot(133)\n",
    "plt.imshow(M2_fourier[-100:, -100:], cmap=\"bwr\", vmin=-1, vmax=1)\n",
    "plt.title(\"FT Composition\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c160352b",
   "metadata": {},
   "source": [
    "Further, notice how the bright regions in the Fourier-domain image are concentrated in a small set of frequency components. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73d812ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "fourier_img = fourier_transform(img.flatten(), Q)\n",
    "fourier_img = M_fourier @ S_fourier @ fourier_img  # Feel free to experiment here!\n",
    "pixel_img = inv_fourier_transform(fourier_img, Q)\n",
    "\n",
    "plt.subplot(121)\n",
    "plt.imshow(fourier_img.reshape(img.shape), cmap=\"gray\")\n",
    "plt.title(\"Fourier domain\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "plt.subplot(122)\n",
    "plt.imshow(pixel_img.reshape(img.shape), cmap=\"gray\")\n",
    "plt.title(\"Pixel domain\")\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60c031cc",
   "metadata": {},
   "source": [
    "## Part 3: Gaussian Filters\n",
    ">In this section, we will discuss the best way to compute [image derivatives](https://en.wikipedia.org/wiki/Image_derivative).\n",
    "\n",
    "To compute image derivatives, you may have previously seen various finite-differences approximations of the gradient (e.g., [Sobel filter](https://en.wikipedia.org/wiki/Sobel_operator)). However, [these filters have numerous problems](https://www.crisluengo.net/archives/22/). Chief amongst these issues is the preservation of high frequency noise, which makes the gradients output of these filters a very poor approximation of the image derivative. Here, you will derive and implement Gaussian filters, a highly accurate, exact (up to a smoothing), and efficient way to compute image derivatives.\n",
    "\n",
    "To motivate Gaussian filters, let $I(x,y)$ be a 2D image and let $G(x,y)$ be a 2D Gaussian with mean $\\mu = 0$ and isotropic variance $\\sigma^2$. From the properties of convolution, we have:\n",
    "\n",
    "$$\n",
    "\\frac{d}{dx}(I * G)\n",
    "= \\frac{dI}{dx} * G\n",
    "= \\frac{dG}{dx} * I\n",
    "= G^{(1)} * I.\n",
    "$$\n",
    "\n",
    "That is, the convolution of the gradient of our image with a Gaussian is equivalent to the convolution of the derivative of a Gaussian with our image!\n",
    "\n",
    "Why is this exciting? Because we can now compute the exact derivative of a blurred version of our image (Gaussian derivatives have closed-form formulas), instead of relying on finite differences. This also let's us easily calculate higher-order image derivatives with $G^{(n)} \\times I.$\n",
    "\n",
    "It's worth noting that you also have to smooth an image prior to applying a Sobel filter anyways (because they amplify high-frequency noise), so you really gain nothing by avoiding Gaussian filters!\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c47655ab",
   "metadata": {},
   "source": [
    "### Part 3.1: Implementing a 1D Gaussian filter\n",
    "\n",
    "In this part, you will implement a 1D Gaussian filter with $n$-th order derivatives in `_gaussian_kernel_1d`.\n",
    "\n",
    "Derivatives of Gaussian functions have nice closed-form formulas. However, since we sample finite-length Gaussian kernels for convolution, some numerical errors will creep in if we simply use the chain rule to obtain higher-order Gaussian filters.\n",
    "\n",
    "---\n",
    "\n",
    "#### Polynomial Response Conditions\n",
    "\n",
    "1. Derive the response of $G^{(n)}$ to $f(x) = 1$ for $n \\ge 1$.\n",
    "\n",
    "   $$\n",
    "   G^{(n)} * 1\n",
    "   $$\n",
    "\n",
    "   What does this tell you the **mean** of your Gaussian derivative kernel should be?\n",
    "\n",
    "2. Derive the response of $G^{(n)}$ to $f(x) = x^n$ for $n \\ge 1$.\n",
    "\n",
    "   $$\n",
    "   G^{(n)} * x^n\n",
    "   $$\n",
    "\n",
    "   What does this tell you about the **sum of your kernel when multiplied by $x^n$**?\n",
    "\n",
    "#### Implementation Details\n",
    "\n",
    "- Feel free to hard-code the second-order derivatives or derive and use a closed-form expression for the $n$-th derivative of a Gaussian (see: [hermite polynomials](https://en.wikipedia.org/wiki/Hermite_polynomials)).\n",
    "- Let the radius of the filter be $r = \\lceil \\sigma t \\rceil, $ where $t$ is the truncation hyperparameter. Note that the total length of the filter should therefore be $2r + 1.$\n",
    "- You must use the polynomial response conditions to guide your answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4be7400",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _gaussian_filter_1d(\n",
    "    sigma: float,  # Standard deviation of the Gaussian\n",
    "    order: int,  # Order of the derivative\n",
    "    truncate: float = 4.0,  # Truncate the filter at this many standard deviations\n",
    "    dtype: torch.dtype = torch.float32,  # Data type to run the computation in\n",
    "    device: torch.device = torch.device(\"cpu\"),  # Device to run the computation on\n",
    ") -> Float[Tensor, \" filter_size\"]:\n",
    "    \n",
    "    def hermite(x, n):\n",
    "        # STUDENT CODE HERE\n",
    "\n",
    "        # END STUDENT CODE\n",
    "        pass\n",
    "    \n",
    "    # STUDENT CODE HERE\n",
    "    radius = ... \n",
    "    # END STUDENT CODE\n",
    "\n",
    "    positions = torch.arange(-radius, radius + 1, device=device, dtype=dtype)\n",
    "    u = positions / sigma\n",
    "    k = math.sqrt(1 / (2 * torch.pi * sigma**2))\n",
    "    kernel = k * (-(u**2) / 2).exp()\n",
    "\n",
    "    if order == 0:\n",
    "        return kernel\n",
    "\n",
    "    # Compute the derivative of the Gaussian kernel\n",
    "    kernel *= (-1) ** order * hermite(u, order)\n",
    "\n",
    "    # STUDENT CODE HERE\n",
    "    kernel -= ...\n",
    "    kernel /= ...\n",
    "    # END STUDENT CODE\n",
    "\n",
    "    return kernel"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28e23470",
   "metadata": {},
   "source": [
    "### Part 3.2: Implementing a 2D Gaussian filter\n",
    "\n",
    "Use the separability of the 2D Gaussian filter to implement the following **derivatives**.\n",
    "\n",
    "Note: A 2D Gaussian can be written as the product of two 1D Gaussians: $ G(x,y) = G(x)\\,G(y). $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e84fc4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def zeroth_order(img: Float[Tensor, \"*H W\"], sigma: float) -> Float[Tensor, \"*H W\"]:\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return img_modified\n",
    "\n",
    "def first_order_x(img: Float[Tensor, \"*H W\"], sigma: float) -> Float[Tensor, \"*H W\"]:\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return img_dx\n",
    "\n",
    "\n",
    "def first_order_y(img: Float[Tensor, \"*H W\"], sigma: float) -> Float[Tensor, \"*H W\"]:\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return img_dy\n",
    "\n",
    "\n",
    "def first_order_xy(img: Float[Tensor, \"*H W\"], sigma: float) -> Float[Tensor, \"*H W\"]:\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return img_dxy\n",
    "\n",
    "def second_order_xx(img: Float[Tensor, \"*H W\"], sigma: float) -> Float[Tensor, \"*H W\"]:\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return img_dxx\n",
    "\n",
    "def second_order_yy(img: Float[Tensor, \"*H W\"], sigma: float) -> Float[Tensor, \"*H W\"]:\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return img_dyy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f3e85ac",
   "metadata": {},
   "source": [
    "Implementing a numerically stable second-order Gaussian filter can be complicated. Compare your implementation, which should return a constant zero-valued second-derivative in the y-direction, with a package that implements Gaussian filters incorrectly... `scipy`!\n",
    "\n",
    "`scipy`'s implementation has errors of magnitude $10^3$, while yours should be < $10^{-13}$.\n",
    "\n",
    "Why does this occur? This discrepancy occurs because computing higher-order derivatives of Gaussian filters is numerically sensitive, especially when the filter is truncated to a finite length. In theory, the second derivative of a linear function is exactly zero. Since the test image is constructed so that its second derivative in the $y$-direction is zero, an ideal Gaussian second-derivative filter should produce values very close to zero everywhere. However, in practice, Gaussian filters must be sampled and truncated. If derivatives are computed by repeatedly applying finite-difference or discrete differentiation rules, small approximation errors are introduced. These errors accumulate and are amplified when taking higher-order derivatives, leading to large numerical inaccuracies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd2d5f34",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct an image with a zero-valued second-derivative in the column-direction\n",
    "img = torch.zeros(1, 1, 31, 31, dtype=torch.float64)\n",
    "img[..., 7:24] += torch.arange(1, 18)\n",
    "img[..., 24:] += 17\n",
    "\n",
    "for sigma in [0.5]: # Modify as desired\n",
    "    # Apply your second-derivative filter\n",
    "    our_second_derivative = second_order_yy(img, sigma=sigma)\n",
    "\n",
    "    # Apply the second-derivative filter using scipy\n",
    "    scipy_second_derivative = gaussian_filter_scipy(\n",
    "        img.numpy().squeeze(), sigma=sigma, order=[2, 0]\n",
    "    )\n",
    "\n",
    "    # Plot the original image and the filtered images\n",
    "    plt.figure(figsize=(15, 3.75))\n",
    "    plt.subplot(131)\n",
    "    plt.imshow(img.squeeze(), cmap=\"gray\")\n",
    "    plt.title(\"Original Image\")\n",
    "    plt.colorbar()\n",
    "\n",
    "    plt.subplot(132)\n",
    "    plt.imshow(our_second_derivative.squeeze(), cmap=\"turbo\")\n",
    "    plt.title(\"Your Second Derivative\")\n",
    "    plt.colorbar()\n",
    "\n",
    "    plt.subplot(133)\n",
    "    plt.imshow(scipy_second_derivative, cmap=\"turbo\")\n",
    "    plt.title(\"Scipy Second Derivative\")\n",
    "    plt.colorbar()\n",
    "\n",
    "    plt.suptitle(f\"Sigma = {sigma}\")\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f91588be",
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma = 1\n",
    "fig, axs = imshow(\n",
    "    img := astronaut(),\n",
    "    zeroth_order(img, sigma),\n",
    "    first_order_x(img, sigma),\n",
    "    first_order_y(img, sigma),\n",
    "    second_order_xx(img, sigma),\n",
    "    second_order_yy(img, sigma),\n",
    ")\n",
    "axs[0].set_title(\"Original image\")\n",
    "axs[1].set_title(\"Smoothed\")\n",
    "axs[2].set_title(\"First x-derivative\")\n",
    "axs[3].set_title(\"First y-derivative\")\n",
    "axs[4].set_title(\"Second x-derivative\")\n",
    "axs[5].set_title(\"Second y-derivative\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "427ad27b",
   "metadata": {},
   "source": [
    "### Part 3.3: Implementing the Laplacian of Gaussian (LoG)\n",
    "The Laplacian operator is defined as\n",
    "\n",
    "$$\n",
    "\\Delta = \\frac{\\partial^2}{\\partial x^2} + \\frac{\\partial^2}{\\partial y^2}.\n",
    "$$\n",
    "\n",
    "When implemented using Gaussian filters, this operator is called the **Laplacian of Gaussian (LoG)**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29a49386",
   "metadata": {},
   "outputs": [],
   "source": [
    "def log(img: Float[Tensor, \"*H W\"], sigma: float) -> Float[Tensor, \"*H W\"]:\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return img_log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d624bae",
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma = 1\n",
    "fig, axs = imshow(\n",
    "    img := astronaut(),\n",
    "    log(img, sigma),\n",
    ")\n",
    "axs[0].set_title(\"Original image\")\n",
    "axs[1].set_title(\"LoG filtered image\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0a36dd8",
   "metadata": {},
   "source": [
    "### Part 3.4: Steerable Filters\n",
    "> In this part, we will discuss [steerable filters](https://people.csail.mit.edu/billf/publications/Design_and_Use_of_Steerable_Filters.pdf), a method of computing filter responses at any orientation.\n",
    "\n",
    "A steerable first-order Gaussian filter can be oriented at any angle using the formula\n",
    "\n",
    "$$\n",
    "G_\\theta^{(1)}\n",
    "=\n",
    "G_{0^\\circ}^{(1)} \\cos\\theta\n",
    "+\n",
    "G_{90^\\circ}^{(1)} \\sin\\theta.\n",
    "$$\n",
    "\n",
    "First, start by implementing a steerable first-order Gaussian filter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86bea514",
   "metadata": {},
   "outputs": [],
   "source": [
    "def oriented_filter(theta: float, sigma: float, **kwargs) -> Float[Tensor, \"N N\"]:\n",
    "    \"\"\"\n",
    "    Return an oriented first-order Gaussian filter\n",
    "    given an angle (in radians) and standard deviation.\n",
    "\n",
    "    Hint:\n",
    "    - Use `.gauss._gaussian_filter_1d`!\n",
    "\n",
    "    Implementation details:\n",
    "    - **kwargs are passed to `_gaussian_filter_1d`\n",
    "    \"\"\"\n",
    "    theta = torch.tensor(theta)\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return filter_2d"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4f128b9",
   "metadata": {},
   "source": [
    "A 2D function to convolve an image with a steerable filter:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9faa2d30",
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv(\n",
    "    img: Float[Tensor, \"B 1 H W\"],  # Input image\n",
    "    kernel: Float[Tensor, \"N N\"] | Complex[Tensor, \"N N\"],  # Convolutional kernel\n",
    "    mode: str = \"reflect\",  # Padding mode\n",
    ") -> Float[Tensor, \"B 1 H W\"]:\n",
    "    \"\"\"\n",
    "    Convolve an image with a 2D kernel (assume N < H and N < W).\n",
    "    \"\"\"\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return img_convolved"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76f85082",
   "metadata": {},
   "source": [
    "A function to filter an image with a oriented filter:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d06bc62b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def steer_the_filter(\n",
    "    img: Float[Tensor, \"B 1 H W\"], theta: float, sigma: float, **kwargs\n",
    ") -> Float[Tensor, \"B 1 H W\"]:\n",
    "    \"\"\"\n",
    "    Return the image convolved with a steered filter.\n",
    "    \"\"\"\n",
    "    theta = torch.tensor(theta)\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return img_steered"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bec93c7",
   "metadata": {},
   "source": [
    "And finally, instead of steering the function basis, start by constructing a filtered image basis and steer these images instead!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aab8a906",
   "metadata": {},
   "outputs": [],
   "source": [
    "def steer_the_images(\n",
    "    img: Float[Tensor, \"B 1 H W\"], theta: float, sigma: float, **kwargs\n",
    ") -> Float[Tensor, \"B 1 H W\"]:\n",
    "    \"\"\"\n",
    "    Return the steered image convolved with a filter.\n",
    "    \"\"\"\n",
    "    theta = torch.tensor(theta)\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return img_steered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b77dfc17",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify your 2D filters are properly oriented\n",
    "fig, axs = plt.subplots(2, 4, figsize=(8, 4))\n",
    "for idx, ax in enumerate(axs.flatten()):\n",
    "    ax.imshow(oriented_filter(idx * torch.pi / 4, 2.5), cmap=\"bwr\")\n",
    "    ax.axis(\"off\")\n",
    "    ax.set_title(f\"$\\\\theta = {idx} \\\\pi / 4$\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8d88813",
   "metadata": {},
   "source": [
    "Run the cell below. What do you notice about the outputs of steer_the_filter and steer_the_images?\n",
    "- What does this tell you the relationship between rotations in the image space and the space of oriented filters?\n",
    "- Does this tell you anything interesting about the relationship between the Fourier Transform and rotations?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25ee57f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Steer the filter\n",
    "steered_filter = []\n",
    "fig, axs = plt.subplots(2, 4, figsize=(8, 4))\n",
    "for idx, ax in enumerate(axs.flatten()):\n",
    "    theta = idx * torch.pi / 4\n",
    "    img = steer_the_filter(astronaut(), theta=theta, sigma=2.5)\n",
    "    ax.imshow(img.squeeze(), cmap=\"bwr\")  # , vmin=-.125, vmax=.125)\n",
    "    ax.axis(\"off\")\n",
    "    ax.set_title(f\"$\\\\theta = {idx} \\\\pi / 4$\")\n",
    "    steered_filter.append(img)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Steer the image\n",
    "steered_images = []\n",
    "fig, axs = plt.subplots(2, 4, figsize=(8, 4))\n",
    "for idx, ax in enumerate(axs.flatten()):\n",
    "    theta = idx * torch.pi / 4\n",
    "    out = steer_the_images(astronaut(), theta=theta, sigma=2.5)\n",
    "    ax.imshow(out.squeeze(), cmap=\"bwr\")  # |, vmin=-.125, vmax=.125)\n",
    "    ax.axis(\"off\")\n",
    "    ax.set_title(f\"$\\\\theta = {idx} \\\\pi / 4$\")\n",
    "    steered_images.append(out)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30b23da5",
   "metadata": {},
   "source": [
    "We should notice the outputs of steer_the_filter and steer_the_images are identical! This is because convolution is linear and distributes over addition.\n",
    "\n",
    "In steer_the_filter, we compute $I * \\bigl(G_0 \\cos\\theta + G_{90} \\sin\\theta\\bigr),$ while in 'steer_the_images', we compute $(I * G_0)\\cos\\theta + (I * G_{90})\\sin\\theta.$\n",
    "\n",
    "By the distributive property of convolution, these expressions are mathematically equivalent, so both methods produce the same result. This equivalence shows that rotations in image space correspond directly to rotations in filter space. Rotating a filter and then applying it to an image is equivalent to steering the responses of fixed basis filters. Thus, oriented filters form a representation that is consistent under rotations. Finally, this property is analogous to the Fourier transform, where spatial shifts become simple phase changes in the frequency domain. Similarly, steerable filters convert rotations in the spatial domain into linear combinations of basis responses, making rotational structure easier to analyze and compute.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3246fa41",
   "metadata": {},
   "source": [
    "### Part 3.5: Measuring Edge Orientation in an Image\n",
    "\n",
    "As a final application, we explore how to measure the orientation of edges in an image.\n",
    "\n",
    "We do this using a complex-valued filter that encodes the image gradient in a single convolution. The filter is constructed as $G = G_x + i G_y$, where $G_x$ and $G_y$ are Gaussian-smoothed derivatives in the $x$ and $y$ directions. After convolving this filter with the image, we obtain $I_x + i I_y$, which represents the gradient field of the image.\n",
    "\n",
    "The magnitude $|I_x + i I_y| = \\sqrt{I_x^2 + I_y^2}$ measures the edge strength, while the phase $\\arg(I_x + i I_y)$ gives the edge orientation.\n",
    "\n",
    "Finally, multiplying the phase by the normalized magnitude suppresses unreliable orientations in flat regions, ensuring that only strong edges contribute meaningful direction estimates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5da28f66",
   "metadata": {},
   "outputs": [],
   "source": [
    "def measure_orientation(img, sigma: float, **kwargs):\n",
    "    \"\"\"\n",
    "    Measure the local orientation of an image using a complex-valued filter. \n",
    "\n",
    "    Return a complex tensor where the magnitude encodes the strength of the orientation\n",
    "    and the phase encodes the angle of the orientation.\n",
    "    \"\"\"\n",
    "    def euler_filters(sigma: float, **kwargs):\n",
    "        smooth   = _gaussian_filter_1d(sigma=sigma, order=0, **kwargs)\n",
    "        grad     = _gaussian_filter_1d(sigma=sigma, order=1, **kwargs)\n",
    "        real_kernel = torch.outer(smooth, grad)\n",
    "        imag_kernel = torch.outer(grad, smooth)\n",
    "        return real_kernel, imag_kernel\n",
    "\n",
    "    def conv_real(img, kernel, mode=\"reflect\"):\n",
    "        N = kernel.shape[-1]\n",
    "        pad = N // 2\n",
    "        img_padded = F.pad(img, (pad, pad, pad, pad), mode=mode)\n",
    "        weight = kernel.unsqueeze(0).unsqueeze(0)\n",
    "        return F.conv2d(img_padded, weight)\n",
    "\n",
    "    # STUDENT CODE HERE\n",
    "\n",
    "    # END STUDENT CODE\n",
    "    return orientation_measure"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdf3b2f3",
   "metadata": {},
   "source": [
    "Confirm your image looks like the one in 3.5!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f25b2263",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Measure the orientation of edges in the image\n",
    "img = hubble()\n",
    "orientations = measure_orientation(img, sigma=1)\n",
    "\n",
    "# Plot the original image and the orientation of edges\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8.5, 4), constrained_layout=True)\n",
    "ax1.imshow(img.squeeze(), cmap=\"gray\")\n",
    "ax1.set_title(\"Original Image\")\n",
    "ax1.axis(\"off\")\n",
    "\n",
    "im = ax2.imshow(orientations.squeeze(), cmap=\"bwr\", vmin=-torch.pi, vmax=torch.pi)\n",
    "ax2.set_title(\"Edge Orientation\")\n",
    "ax2.axis(\"off\")\n",
    "\n",
    "divider = make_axes_locatable(ax2)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.2)\n",
    "plt.colorbar(\n",
    "    im,\n",
    "    cax=cax,\n",
    "    ticks=[-torch.pi, 0, torch.pi],\n",
    "    format=lambda x, _: \"-π\" if x < 0 else (\"π\" if x > 0 else \"0\"),\n",
    ")\n",
    "\n",
    "ax1.set_aspect(\"equal\")\n",
    "ax2.set_aspect(\"equal\")\n",
    "\n",
    "plt.savefig(\"img/hubble.png\", bbox_inches=\"tight\", pad_inches=0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e577814",
   "metadata": {},
   "source": [
    "## Summary\n",
    "Overall PSET2 connects signal processing, linear algebra, and image analysis to core ideas in computer vision. \n",
    "\n",
    "---\n",
    "\n",
    "### Part 1: Steerable Bases\n",
    "\n",
    "In this section, we learned how sine and cosine functions form a **steerable basis**.\n",
    "\n",
    "- Showed that `a cos(x) + b sin(x)` is equivalent to a single shifted sine wave.\n",
    "- Interpreted coefficients `(a, b)` as controlling amplitude and phase.\n",
    "- Used the unit circle to steer signals by changing phase.\n",
    "- Implemented functions to generate and visualize steered sinusoids.\n",
    "\n",
    "Takeaway: Simple basis functions can be combined to control orientation and phase, which motivates steerable filters in vision.\n",
    "\n",
    "---\n",
    "\n",
    "### Part 2: Fourier Transform\n",
    "\n",
    "This section explored how image operations can be represented.\n",
    "\n",
    "- Built shift operators with circular boundaries.\n",
    "- Represented convolution as circulant and block matrices.\n",
    "- Verified that shifts and convolutions commute.\n",
    "- Computed eigenvalues and eigenvectors of convolution operators.\n",
    "- Implemented the Fourier transform as a change of basis.\n",
    "\n",
    "Takeaway: The Fourier transform helps us understand and implement image filtering by separating images into independent frequency components.\n",
    "\n",
    "---\n",
    "\n",
    "### Part 3: Gaussian Filters\n",
    "\n",
    "In the final section, we focused on computing image derivatives and orientations.\n",
    "\n",
    "- Implemented Gaussian filters and their derivatives.\n",
    "- Built stable first- and second-order derivative operators.\n",
    "- Implemented the Laplacian of Gaussian (LoG).\n",
    "- Constructed steerable filters for arbitrary orientations.\n",
    "- Measured edge strength and direction using complex filters.\n",
    "\n",
    "Takeaway: We saw how image filtering can be analyzed using eigen-decomposition, and how this naturally leads to Fourier representations.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1705650f",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
