Bootcamp 2 - Plots & LinearAlgebra.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a883d7b2-34b7-49e3-9325-abff04fe7353",
   "metadata": {},
   "source": [
    "<h1>Scientific coding bootcamp notebook 2: Plotting and linear algebra.</h1>\n",
    "Julia is designed from head to toe for scientific computing.  Much of its power comes from the state of the art packages that are part of the Julia ecosystem.  We'll acquaint ourselves with a few of those packages in this notebook. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1e718fe-705d-45a9-a48b-d6c7cbbc6b55",
   "metadata": {},
   "source": [
    "<h2>Plotting</h2>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd01776b-7c53-449f-afcc-3e087cb3b26f",
   "metadata": {},
   "outputs": [],
   "source": [
    "using Plots, LinearAlgebra"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c04ad27-3fe3-47bb-b493-9e49c898f3fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Basic plotting\n",
    "x = 0:0.1:10      # This is called a range.  \n",
    "y = x .^ 2\n",
    "plot(x,y,label=\"y=x^2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad53985a-dbdf-44f4-8e38-9374d9848342",
   "metadata": {},
   "outputs": [],
   "source": [
    "# There are two ways to plot multiple lines.  The first is to simply pass a vector of things to plot:\n",
    "plot(x,[y,x.^3]; labels = [\"y=x^2\" \"y=x^3\"], title=\"This is a title\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5274d9e-7779-42a3-80fe-72dab298e579",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The second is to use the function \"plot!\", which updates a plot:\n",
    "p = plot(x,y,label=\"y=x^2\")\n",
    "plot!(p,x,x.^4,label=\"y=x^4\",ylim=(0,100))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ef375a8-5e26-484f-add6-e629617edd69",
   "metadata": {},
   "source": [
    "</hline>\n",
    "<h1>Linear algebra</h1>\n",
    "<h3>SVD and pinv</h3>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0a42059-c0ec-4bd9-8c10-76f7d72d13fb",
   "metadata": {},
   "source": [
    "We now come to linear algebra.  The power of linear algebra is that it lets you reason geometrically about algebra, and it lets you reason algebraically about geometry.  We begin with the very useful singular value decomposition (SVD) and the related pseudo-inverse.\n",
    "\n",
    "SVD breaks an arbitrary matrix into a composition of a rotation, a stretching (diagonal matrix), and another rotation.  More explicitly, SVD of an arbitrary matrix M is given by U*S*V', where U and V are rotation matrices, S is diagonal, and ' denotes the conjugate-transpose.  The singular values tell you a lot about the matrix, such as whether it is invertible, and they are usually sorted from largest to smallest.  (Remarkably, this works even if M is not a square matrix.  Question: What do you think it means for a non-square matrix to be \"diagonal\"?)\n",
    "\n",
    "The pseudoinverse is the \"closest to an inverse you can get\".  If a matrix M is invertable, the pseudoinverse pinv(M) is just the regular inverse M^{-1}.  If M is not invertable but it is diagonal, then the pinv(M) is computed by inverting each non-zero diagonal element and leaving every other element zero.  In the general case, if the SVD of M is U*S*V', the pseudoinverse of M is V*pinv(S)*transpose(U). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9650e599-e608-4ea4-9a6f-74a6ebc2f7f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Try to interpret geometrically what this matrix does.\n",
    "M = [0 -2 0 ; 2 0 0 ; 0 0 0]\n",
    "F = svd(M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3072e05-c0c6-4970-9977-1eb9618e670d",
   "metadata": {},
   "outputs": [],
   "source": [
    "pinv(M)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7eac0b6-b9ea-4013-b3db-7e9c1a104a6b",
   "metadata": {},
   "source": [
    "Below are some problems to get you familiar with pinv and SVD. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c8a04e3-36e8-4731-91e0-4d01c2160297",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the pseudoinverse of the pseudoinverse of M, and interpret it. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e00a8da-c8d8-4f5f-942e-528a9efe699f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract the diagonal factor of svd(M), and turn it into a diagonal matrix. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e712788-1963-45f4-9e72-d32903fd5337",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find a matrix for which both rotations in the SVD are the identity matrix. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8462836c-8531-4675-80c4-320862eca0b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find a non-diagonal matrix for which the SVD diagonal factor has exactly one non-zero entry. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "886ece71-557c-4e73-a032-36bcac88269c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the determinant of a matrix using the diagonal factor of its SVD. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33a129d0-d524-47cd-b009-cda8e401c25d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the SVD of a non-square matrix.  What are the shapes of the factors?  Why does this make sense?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3d5dc28-d5e1-49e4-b82d-1688855862be",
   "metadata": {},
   "source": [
    "<h3>Least squares</h3>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "075b98e8-0c3f-41c9-aa14-0f432e60735a",
   "metadata": {},
   "source": [
    "A very useful application of pinv is for computing linear least squares fits.  Suppose we have data with x-coordinates xs and y coordinates ys.  A least squares fit is a slope A and intercept b such that (ys - A*xs - b) is pretty small.  This can be rearranged into \n",
    "\n",
    "ys ≈ [xs ones(length(xs))] * [A, b], \n",
    "\n",
    "or \n",
    "\n",
    "[A,b] = pinv([xs ones(length(xs))]) * ys."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0a18fc2-2f1a-48dd-9f20-666397b1e692",
   "metadata": {},
   "outputs": [],
   "source": [
    "xs = sort(rand(100))\n",
    "ys = sort(rand(100))\n",
    "A,b = pinv([xs ones(length(xs))]) * ys\n",
    "println(A)\n",
    "println(b)\n",
    "scatter(xs,ys,label=\"Data\")\n",
    "plot!(xs,A*xs .+ b,label=\"Fit\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5edc01a0-03dc-43cc-8fbd-aa8f8b640c4c",
   "metadata": {},
   "source": [
    "There's a convenient shorthand for the pseudoinverse, using the operator \\. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71a7f593-32c7-4112-ba44-e0a0c335dfe6",
   "metadata": {},
   "outputs": [],
   "source": [
    "A,b = [xs ones(length(xs))] \\ ys"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63b39b27-847f-43ef-91e4-19a71b16ca89",
   "metadata": {},
   "source": [
    "<h3>Function spaces</h3>\n",
    "\n",
    "Arguably the most powerful application of linear algebra is to vector spaces of functions.  To see how this works, we will start with spaces of polynomials.  Consider the set $P_N$ of all polynomials up to some fixed degree $N$.  We will take $N=4$ for concreteness in examples to follow.  An example of an element of $P_4$ is $p(x) = 12 - 4x/3 + \\pi^2 x^4$. \n",
    "\n",
    "**Question:** What is a basis for $P_4$? What is the dimension of $P_4$?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74df6394-ecf4-42fa-a1bd-4e875e1891e5",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "A downside of the \"obvious\" basis for $P_4$ is that it's not easy to figure out how to express a given function in terms of the basis polynomials.  For instance, if I give you a list of numbers like \n",
    "```\n",
    "fs=[0.0625, 0.0256, 0.0081, 0.0016, 0.0001, 0.0, 0.0001, 0.0016, 0.0081, 0.0256, 0.0625]\n",
    "```\n",
    "and tell you that these are the values of some polynomial sampled at the points `0.0:0.1:1.0`, it's probably not immediately clear how to recover the polynomial.  The key to doing so is to define an inner product on this vector space. \n",
    "\n",
    "An inner product is an abstract version of a dot product.  A typical example of an inner product of two real functions $f(x)$ and $g(x)$ is \n",
    "$$\n",
    "\\langle f | g \\rangle = \\int_0^1 f(x) g(x) dx.\n",
    "$$\n",
    "Given a basis $B$ and an arbitrary function $f\\in P_N$, we can always write $f(x)=\\sum_{p\\in B} c_p p(x)$ for some coefficients $c_p$.  We can use an inner product and some linear algebra to figure out what the coefficients are.  \n",
    "$$\n",
    "\\langle q | f\\rangle = \\sum_{p\\in B} \\langle q | p\\rangle c_p = \\Bigl[ \\langle q | p_0 \\rangle, q | p_1 \\rangle, \\langle q | p_2 \\rangle, \\langle q | p_3 \\rangle, \\langle q | p_4 \\rangle \\Bigr] \\begin{bmatrix}\n",
    "           c_0 \\\\ c_1 \\\\ c_2 \\\\ c_3 \\\\ c_4\n",
    "         \\end{bmatrix}\n",
    "$$\n",
    "Stacking together the equation for each $q \\in B$ gives a matrix equation\n",
    "$$\n",
    "\\begin{bmatrix}\n",
    "           \\langle p_0 | f\\rangle \\\\ \\langle p_1 | f\\rangle \\\\ \\langle p_2 | f\\rangle \\\\ \\langle p_3 | f\\rangle \\\\ \\langle p_4 | f\\rangle \n",
    "\\end{bmatrix} = \n",
    "\\begin{bmatrix}\n",
    "\\langle p_0 | p_0 \\rangle & \\langle p_0 | p_1 \\rangle & \\langle p_0 | p_2 \\rangle & \\langle p_0 | p_3 \\rangle & \\langle p_0 | p_4 \\rangle \\\\\n",
    "\\langle p_1 | p_0 \\rangle & \\langle p_1 | p_1 \\rangle & \\langle p_1 | p_2 \\rangle & \\langle p_1 | p_3 \\rangle & \\langle p_1 | p_4 \\rangle \\\\\n",
    "\\langle p_2 | p_0 \\rangle & \\langle p_2 | p_1 \\rangle & \\langle p_2 | p_2 \\rangle & \\langle p_2 | p_3 \\rangle & \\langle p_2 | p_4 \\rangle \\\\\n",
    "\\langle p_3 | p_0 \\rangle & \\langle p_3 | p_1 \\rangle & \\langle p_3 | p_2 \\rangle & \\langle p_3 | p_3 \\rangle & \\langle p_3 | p_4 \\rangle \\\\\n",
    "\\langle p_4 | p_0 \\rangle & \\langle p_4 | p_1 \\rangle & \\langle p_4 | p_2 \\rangle & \\langle p_4 | p_3 \\rangle & \\langle p_4 | p_4 \\rangle\n",
    "\\end{bmatrix}\n",
    "\\begin{bmatrix}\n",
    "           c_0 \\\\ c_1 \\\\ c_2 \\\\ c_3 \\\\ c_4 \n",
    "\\end{bmatrix}\n",
    "$$\n",
    "By matrix inversion we can get the coefficients in terms of the inner products $\\langle q | f\\rangle$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2bff1cb-6625-40bb-884b-f6af74709134",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Complete the following code to determine the polynomial coefficients associated to the data points fs. \n",
    "fs = [0.0625, 0.0256, 0.0081, 0.0016, 0.0001, 0.0, 0.0001, 0.0016, 0.0081, 0.0256, 0.0625]\n",
    "xs = 0.0:0.1:1.0\n",
    "\n",
    "function innerProduct01(f::Vector{Float64},g::Vector{Float64})\n",
    "    if length(f) != length(g) \n",
    "        error(\"Inputs have different lengths.\")\n",
    "    end\n",
    "    N = length(f) \n",
    "    dx = 1.0/(N-1)\n",
    "    # Your code here.  Compute the integral of fg from 0 to 1.  You may use any integration rule you like.  \n",
    "    # You can find examples of integration rules in Numerical Recipes 4.1.3. \n",
    "end\n",
    "\n",
    "innerProds = [innerProduct01(fs,(x-> x^k).(xs)) for k=0:4]\n",
    "coeffs = inv([innerProduct01((x-> x^j).(xs),(x-> x^k).(xs)) for j=0:4, k=0:4]) * innerProds  # Your code here. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "020e131e-4ec3-439c-ac9b-bc3bbbc02d1e",
   "metadata": {},
   "source": [
    "<h4>Polynomial fitting.</h4>\n",
    "\n",
    "You can use the same method as above to fit a fourth order polynomial to data which is *not* polynomial.  E.g. noisy data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8167ab82-33c7-41a3-bc15-84c75c191f1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "noisyfs = [0.0625, 0.0256, 0.0081, 0.0016, 0.0001, 0.0, 0.0001, 0.0016, 0.0081, 0.0256, 0.0625] + (rand(11) .- 0.5) * 0.02\n",
    "noisyInnerProds = [innerProduct01(noisyfs,(x-> x^k).(xs)) for k=0:4]\n",
    "noisyCoeffs = inv([innerProduct01((x-> x^j).(xs),(x-> x^k).(xs)) for j=0:4, k=0:4]) * noisyInnerProds  # Your code here. \n",
    "println(noisyCoeffs)\n",
    "# Your code here.  Make a scatter plot of xs vs. noisyfs. \n",
    "plot!((x-> sum(noisyCoeffs[j+1]*x^j for j=0:4)), xlim=(0,1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "516bce54-5442-410a-bd43-26df60c4175a",
   "metadata": {},
   "source": [
    "In practice, when you do curve fitting of this sort you should usually use a dedicated fitting package. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24850e0f-8880-4df2-938e-8c4c9170a053",
   "metadata": {},
   "source": [
    "<h3>Gaussian beam propagation</h3>\n",
    "\n",
    "We will apply the above ideas to a space of two-dimensional functions which model laser beam propagation. We begin by defining the Hermite-Gaussian modes, which are the most basic shapes that laser beams can assume.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7fc6a2a-4dc6-4437-b091-3a0dfe1e5665",
   "metadata": {},
   "outputs": [],
   "source": [
    "using SpecialPolynomials: Hermite, basis\n",
    "# Define a 64x64 computational grid\n",
    "N = 64\n",
    "const grid = (-sqrt(N)/2:sqrt(1/N):sqrt(N)/2-sqrt(1/N), -sqrt(N)/2:sqrt(1/N):sqrt(N)/2-sqrt(1/N))\n",
    "\n",
    "function HGMode(grid, n::Int, m::Int; waist::Float64=1.0, z::Float64=0.0, lam::Float64=1.0)\n",
    "    return HG1D(grid[1],n;waist=waist,z=z,lam=lam) .* transpose(HG1D(grid[2],m;waist=waist,z=z,lam=lam))\n",
    "end\n",
    "\n",
    "function HG1D(grid, n::Int; waist::Float64=1.0, z::Float64=0.0, lam::Float64=1.0)\n",
    "    zr = pi * waist^2 / lam\n",
    "    q = z + im*zr\n",
    "    w = waist*sqrt(1+(z/zr)^2)\n",
    "    prefactor = sqrt(sqrt(2/pi)/(2^n * factorial(n) * waist)) * sqrt(1/(1-im*z/zr)) * sqrt(-conj(q)/q)^n\n",
    "    polyfactor = basis(Hermite,n).(grid * sqrt(2) ./ w)\n",
    "    expfactor = exp.(-2*pi*im * grid.^2 / (2*lam*q))\n",
    "    return prefactor .* polyfactor .* expfactor\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51a352d9-b54c-422d-8624-7afd668f8204",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Why do we need to take abs here? \n",
    "heatmap(abs.(HGMode(grid,10,0)).^2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4aa0355a-fcc4-409a-984e-427f5fbabda0",
   "metadata": {},
   "source": [
    "Hermite-Gaussian modes $H_{n,m}(x,y)$ form a basis for 2D functions.  This means that any 2D function can be written as a sum of the form\n",
    "$$\n",
    "\\sum_{n,m} c_{n,m} H_{n,m}(x,y).\n",
    "$$\n",
    "How do we determine the coefficients $c_{n,m}$?  We could do the same matrix inversion as we did above, but HG modes have the special property of *orthonormality* which greatly simplifies the analysis.  The integral of a product of two HG modes (one of them complex conjugated) is zero if the modes are different and 1 if they are the same,\n",
    "$$\n",
    "\\int_{\\mathbb{R}^2} H_{n,m}(x,y) H^{*}_{p,q}(x,y) \\; dx \\; dy = \\delta_{np} \\delta_{mq}.\n",
    "$$\n",
    "\n",
    "If some function $f(x,y)$ equals $\\sum_{n,m} c_{n,m} H_{n,m}(x,y)$ for some choice of coefficients $c_{n,m}$, by multiplying both expressions against $H^{*}_{p,q}$ and integrating, we find\n",
    "$$\n",
    "c_{p,q} = \\int_{\\mathbb{R}^2} f(x,y) H^{*}_{p,q}(x,y) \\; dx \\; dy.\n",
    "$$\n",
    "You will implement this *mode decomposition* below for a function (represented as a matrix) sampled on the same grid as above. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b13d47c2-64ec-419d-b1a8-5cc922033c8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Finish the code below to implement the overlap integral in the mode decomposition coefficient formula. \n",
    "# Numerical Recipes 4.1.3 has relevant integration formulas, though you can probably devise your own.\n",
    "function HGOverlap(grid,f::Matrix{T},n::Int,m::Int; waist::Float64=1.0, z::Float64=0.0, lam::Float64=1.0) where {T<:Number}\n",
    "    mode = HGMode(grid, n, m; waist=waist, z=z, lam=lam)\n",
    "    integrand = f .* conj.(mode)\n",
    "    # Your code here. \n",
    "end\n",
    "\n",
    "LG01 = [exp(-x^2-y^2) * (1-2x^2-2y^2) for x in grid[1], y in grid[2]] * sqrt(2/pi)\n",
    "coeffs = [HGOverlap(grid,LG01,n,m) for n=0:5,m=0:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b27eca9-5748-4a3a-97fe-d9a061e24961",
   "metadata": {},
   "outputs": [],
   "source": [
    "# If you coded the above correctly, the sum of squared absolute values of the coefficients should be one.  \n",
    "sum(abs.(coeffs).^2) |> println\n",
    "heatmap(abs.(coeffs).^2)        # This visually shows where the coefficients are large. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e83abca-b3c7-48e6-aca7-fadb0070b578",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Examine the form of the profile LG01. Then use the coefficients you found above and the HGMode function evaluated at z=4 \n",
    "# to determine what a laser with initial profile LG01 will look like after propagating distance 4. \n",
    "heatmap(abs.(LG01).^2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d22c3d0-8f2c-4aa9-a712-a9d4156320b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "heatmap( #= Your code here.  Determine the form of the beam at z=4. =# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d6835e5-5a9b-4dac-98a5-160c16bc9844",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.10.3",
   "language": "julia",
   "name": "julia-1.10"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.10.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}