notebook/band-theory/FFT_and_planewaves.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8852f106",
   "metadata": {},
   "source": [
    "# **Fourier Transforms and Plane-Wave Expansions**\n",
    "\n",
    "**Authors:** Guoyuan Liu, Dou Du, Taylor James Baird and Giovanni Pizzi\n",
    "\n",
    "<i class=\"fa fa-home fa-2x\"></i><a href=\"../index.ipynb\" style=\"font-size: 20px\"> Go back to index</a>\n",
    "\n",
    "**Source code:** https://github.com/osscar-org/quantum-mechanics/blob/master/notebook/band-theory/FFT_and_planewaves.ipynb\n",
    "\n",
    "This notebook shows interactively how discrete Fourier series can represent a function with a limited amount of plane-wave components. A common way to represent a wavefunction when solving the Kohn-Sham equations is via its expansion in plane waves.\n",
    "This notebook focuses on a simple example (much simpler than a complete DFT calculation) in order to help the reader focus on the essential aspects of such a representation.\n",
    "\n",
    "<hr style=\"height:1px;border:none;color:#cccccc;background-color:#cccccc;\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f96673a6",
   "metadata": {},
   "source": [
    "## **Goals**\n",
    "\n",
    "* Understand how a plane-wave basis is directly related to a Fourier series.\n",
    "* Learn how to decompose a function using a FFT algorithm.\n",
    "* Examine how a function is reconstructed from a finite (possibly not complete) set of plane waves.\n",
    "* Understand the impact of the basis-set size on the convergence of the integral of the reconstructed function."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a158676c",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## **Background theory** \n",
    "\n",
    "[More on the background theory.](./theory/theory_fft.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70d65f1e",
   "metadata": {},
   "source": [
    "## **Tasks and exercises**\n",
    "\n",
    "<ol>\n",
    "<li>Prove that plane waves form an orthogonal basis set.\n",
    "<details>\n",
    "    <summary style=\"color: red\">Solution</summary>\n",
    "    \n",
    "We have to prove that $\\langle w_N^j, w_N^k \\rangle =\\langle w_N^k, w_N^j \\rangle= 0 $ for integer $j \\neq k$. We can simply carry out the inner product       \n",
    "$$ \\langle w_N^j, w_N^k \\rangle = \\langle w_N^k, w_N^j \\rangle = \\int_{-\\pi}^{\\pi} e^{ijx} e^{-ikx}dx = \\int_{-\\pi}^{\\pi} e^{i(j-k)x} dx = \\frac 1 {i(j-k)} [e^{i(j-k)x}]_{-\\pi}^{\\pi} = \\begin{cases} 0 & \\text{if j $\\neq$ k} \\\\ 2\\pi & \\text{if j = k}\\end{cases}$$\n",
    "</details></li>\n",
    "\n",
    "<li>How does the number of plane waves affect the approximation of the target function? Will a function with more \"oscillations\" require more components to be accurately represented?\n",
    "<details>\n",
    "    <summary style=\"color: red\">Solution</summary>\n",
    "\n",
    "Move the slider to try different numbers of Fourier components. Observe if the FFT interpolation approximates well the original function and if the integral of the square modulus is close to the convergence value. You can also change the objective function by the drop-down menu. Generally, more sampling yields more accurate representation. \n",
    "For functions with more oscillations (higher frequency components), more Fourier components are needed to reach the same level of accuracy.\n",
    "</details></li>\n",
    "\n",
    "<li>How can we reduce the number of plane waves needed in a DFT calculation, without sacrificing the accuracy of the representation?\n",
    "<details>\n",
    "    <summary style=\"color: red\">Solution</summary>\n",
    "\n",
    "Wavefunctions have the strongest oscillations near nucleus, and a very large number of plane waves is needed to accurately represent this region. Fortunately, core electrons are less relevant in chemical bonding, so we can simplify the problem and obtain a much smoother (pseudo)wavefunction by excluding the core electrons. To learn more about this approach, please check our <a href=\"./pseudopotential.ipynb\">notebook on pseudopotentials</a>. In general, the combination of pseudopotentials and a plane-wave expansion enables fast and accurate calculation of materials and their properties.\n",
    "</details></li>\n",
    "\n",
    "<li>In a DFT calculation, how can we control the number of plane waves used in the basis set?\n",
    "<details>\n",
    "    <summary style=\"color: red\">Solution</summary>\n",
    "  \n",
    "The kinetic energy of a plane wave of momentum $\\mathbf G$ is given by $\\frac {\\hbar^2}{2m} \\lvert \\mathbf G \\rvert^2$. By setting a cutoff energy, we can limit the size of the plane-wave basis set. The value of the cutoff depends on the system under investigation and the pseudopotential used, and convergence tests are normally required. To have a suggestion of a converged cutoff value based on the choice of pseudopotentials, you can check the <a href=\"https://www.materialscloud.org/discover/sssp/table/precision\">standard solid-state pseudopotentials (SSSP) library</a> on Materials Cloud.\n",
    "</details></li>\n",
    "</ol>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41c15519-11cd-4658-b957-1147e287edbf",
   "metadata": {},
   "source": [
    "<hr style=\"height:1px;border:none;color:#cccccc;background-color:#cccccc;\" />\n",
    "\n",
    "## Interactive visualization\n",
    "(be patient, it might take a few seconds to load)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a34187d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as ipw\n",
    "import numpy as np\n",
    "from matplotlib.ticker import MaxNLocator\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.fft as fft\n",
    "%matplotlib widget\n",
    "plt.rcParams['figure.autolayout'] = 'True' # turn on tight layout globally"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e85c3b8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# target functions\n",
    "def periodic_f(x):\n",
    "    # smooth\n",
    "    return np.exp(-((x-1)/0.15)**2) + 0.5 * np.exp(-((x-1.2)/0.1)**2) + 0.8 * np.exp(-((x-0.8)/0.1)**2)\n",
    "\n",
    "def periodic_f2(x):\n",
    "    # less smooth\n",
    "    return np.exp(-((x-1)/0.05)**2) -0.5*np.exp(-((x-1)/0.15)**2) + 0.5 * np.exp(-((x-1.2)/0.1)**2) + 0.8 * np.exp(-((x-0.8)/0.1)**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c463b6d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot x range\n",
    "x = np.linspace(0, 2, 201, endpoint=False)\n",
    "x_range = 2\n",
    "\n",
    "# widgets\n",
    "N_slider = ipw.IntSlider(description=r\"$N_{\\text{fft}}$\", min=6, max=40, value=6, step=1, continuous_update=False, layout={'margin':'0px 15px 0px 15px'})\n",
    "func_dropdown=ipw.Dropdown(description=\"Function\", options=[(\"Smooth\", \"periodic_f\"), (\"Less smooth\", \"periodic_f2\")], layout={ 'margin':'0px 15px 0px 15px'})\n",
    "reset_button = ipw.Button(description='Show all', icon='redo',  style={'description_width': 'initial'}, layout={'width':'220px', 'margin':'0px 20px 0px 60px'})\n",
    "hl_label = ipw.Label(value='(click on a FFT component to select it)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48d2b850",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_resampled(N_fft, x_range=2., function=periodic_f):\n",
    "    \"\"\" Compute FFT series with given number of sampling and target functions. \"\"\"\n",
    "    # Pick an even number to have zero\n",
    "    x_fft = np.linspace(0, x_range, N_fft+1, endpoint=False)# remove last point as it's the same as the first one by PBC\n",
    "    y_fft = function(x_fft)\n",
    "    # Fourier resampling\n",
    "    renormalization = len(x)/(len(y_fft))\n",
    "    y_resamp = fft.irfft(fft.rfft(y_fft), len(x)) * renormalization\n",
    "    \n",
    "    return x_fft, y_fft, y_resamp\n",
    "\n",
    "def get_integral_resampled(N_fft, x_range=2., function=periodic_f):\n",
    "    \"\"\" Compute the integral of the square modulus of the function. \"\"\" \n",
    "    x_fft, y_fft, _ = compute_resampled(N_fft, x_range, function=function)\n",
    "    return (y_fft**2).sum() * (x_fft[1] - x_fft[0])\n",
    "\n",
    "def plot_reconstruct(y_fft):\n",
    "    \"\"\" Plot Fourier expansions \"\"\"\n",
    "    ax2.clear()\n",
    "\n",
    "    coeffs = fft.rfft(y_fft)\n",
    "    N_rfft = 0 # number of fft expansions\n",
    "    for coeff, freq_int in list(zip(coeffs, range(len(coeffs)))):\n",
    "        freq = 2 * np.pi * freq_int / x_range\n",
    "        norm = 1 / (len(y_fft)) * 2\n",
    "        if freq_int == 0:\n",
    "            # The zero-frequency does not have a factor 2 because it's not a cosine\n",
    "            # summing the two complex conjugates, but just a constant\n",
    "            norm /= 2\n",
    "        this_frequency_contrib = ( coeff.real * np.cos(freq * x) - coeff.imag * np.sin(freq * x) ) * norm\n",
    "\n",
    "        ax2.plot(x, this_frequency_contrib + N_rfft) # plot components with vertical shift for visibility\n",
    "        # ax2.plot(x, this_frequency_contrib) # no shift\n",
    "        N_rfft += 1\n",
    "\n",
    "    ax2.axes.yaxis.set_ticks([]) # remove y ticks\n",
    "    ax2.set_title('Expansion Components')\n",
    "\n",
    "CONVERGE_SMOOTH = get_integral_resampled(N_fft=200, function=periodic_f)\n",
    "CONVERGE_ROUGH = get_integral_resampled(N_fft=200, function=periodic_f2)\n",
    "def plot_integral(func_name, func):\n",
    "    \"\"\" plot sum of the square modulus (integral) \"\"\"\n",
    "    ax3.clear()\n",
    "    converged_integral = CONVERGE_SMOOTH if func_name == \"periodic_f\" else CONVERGE_ROUGH\n",
    "    ax3.axhline(converged_integral, color='tab:red')\n",
    "    integrals = []\n",
    "    for N in range(6, 41):\n",
    "        integrals.append((N, get_integral_resampled(N, function=func)))\n",
    "    integrals_x, integrals_y = np.array(integrals).T\n",
    "    \n",
    "    ax3.plot(integrals_x, integrals_y, 'o--', alpha=0.8)\n",
    "    ax3.plot(integrals_x[N_fft-6], integrals_y[N_fft-6],'ro', markersize=11, label='current sampling')\n",
    "    ax3.set_xlabel('number of components')\n",
    "    ax3.set_ylabel(\"Integral of square modulus\")\n",
    "    ax3.set_title(\"Convergence of FFT\")\n",
    "    ax3.set_xlim(6,40)\n",
    "    ax3.xaxis.set_major_locator(MaxNLocator(integer=True))\n",
    "    ax3.legend(loc='best')\n",
    "\n",
    "def plot_sampling(func, x_fft, y_fft, y_resamp):\n",
    "    ax1.clear()\n",
    "    ax1.set_title('FFT interpolation')\n",
    "\n",
    "    x_fft, y_fft, y_resamp = compute_resampled(N_slider.value, function=func)\n",
    "    ax1.plot(x, func(x), 'k-', label='target')\n",
    "    ax1.plot(x_fft, y_fft, 'o', label='sampling')\n",
    "    ax1.fill_between(x, y_resamp, 0,ec='red', fc='yellow', label='FFT')\n",
    "    ax1.legend(loc='best')\n",
    "    ax1.set_ylim(-0.35,1.25)\n",
    "\n",
    "def on_plot_click(event):\n",
    "    \"\"\"handle mouse click event on expansion component plot\"\"\"\n",
    "    # line = event.artist\n",
    "    # xdata = line.get_xdata()\n",
    "    # ydata = line.get_ydata()\n",
    "    if event.inaxes != ax2:\n",
    "        return\n",
    "    for i in range(len(ax2.lines)):\n",
    "        ax2.lines[i].set_alpha(0.1)\n",
    "        ax2.lines[i].set_linewidth(1.1)\n",
    "\n",
    "    # get the id of the line2D object which is vertically closest to the mouse clicking position \n",
    "    id_line = min(enumerate(ax2.lines), key= lambda line: abs(np.mean(line[1].get_ydata())-event.ydata))[0]\n",
    "    ax2.lines[id_line].set_alpha(1)\n",
    "    ax2.lines[id_line].set_linewidth(2.0)\n",
    "\n",
    "    plot_sampling(func, x_fft, y_fft, y_resamp)\n",
    "    ax1.fill_between(ax2.lines[id_line].get_xdata(), ax2.lines[id_line].get_ydata()-id_line, 0, ec='tab:blue', fc='tab:green', alpha=0.5,label='component')\n",
    "    ax1.legend()\n",
    "\n",
    "def plot_update(change):\n",
    "    # get current widget value\n",
    "    global N_fft, x_fft, y_fft, y_resamp, func\n",
    "    N_fft = N_slider.value\n",
    "    func = globals()[func_dropdown.value] # get the function by function name\n",
    "    x_fft, y_fft, y_resamp = compute_resampled(N_fft, function=func)\n",
    "    # update sampling plot\n",
    "    plot_sampling(func, x_fft, y_fft, y_resamp)\n",
    "\n",
    "    # update reconstruct plot\n",
    "    plot_reconstruct(y_fft)\n",
    "\n",
    "    # udpate square modulus plot\n",
    "    plot_integral(func_dropdown.value, func)\n",
    "\n",
    "\n",
    "N_slider.observe(plot_update, names='value', type='change')\n",
    "func_dropdown.observe(plot_update, names='value', type='change')\n",
    "reset_button.on_click(plot_update)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "300944ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define layout by gridspec\n",
    "fig = plt.figure(constrained_layout=True, figsize=(7, 6))\n",
    "gs = fig.add_gridspec(3,4)\n",
    "ax1 = fig.add_subplot(gs[0:2,0:2])\n",
    "ax2 = fig.add_subplot(gs[0:2,2:4])\n",
    "ax3 = fig.add_subplot(gs[-1,:])\n",
    "\n",
    "# interactive plot 2 for line picking\n",
    "cid = fig.canvas.mpl_connect('button_press_event', on_plot_click)\n",
    "\n",
    "# show plots\n",
    "plot_update(None)\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f7c5461",
   "metadata": {},
   "outputs": [],
   "source": [
    "# display widgets\n",
    "display(ipw.HBox([N_slider, func_dropdown]))\n",
    "display(ipw.HBox([reset_button, hl_label]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17052811",
   "metadata": {},
   "source": [
    "## Legend\n",
    "\n",
    "The target function, sampling points and the reconstructed function are shown in the top left plot. The real part (cosine functions) and the constant term of the discrete Fourier series are shown in the top right panel.\n",
    "    \n",
    "Note that the components are shifted vertically for clarity. The integral of the square of the functions reconstructed from truncated Fourier series with different numbers of plane waves $N_{\\text{fft}}$ is shown in the bottom panel, where the current choice of sampling is indicated with a red dot. The converged value is also shown with a red horizontal line, obtained with a large number (200) of FFT components. \n",
    "\n",
    "The number of FFT components $N_{\\text{fft}}$ can be set by the slider. Two target functions can be chosen from the drop-down menu.\n",
    "\n",
    "By clicking one of the expansions in the top-right panel, the contribution of that component will be shown in the top left panel. Click the reset button to display all expansion coefficients again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0e704a2-b7ea-4551-85a1-726ef41954d0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "25dfdc17468cbd6567831e63e4a52a24e09010ff54215efa01b4e7aa765a5908"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  },
  "voila": {
   "authors": "Guoyuan Liu, Dou Du and Giovanni Pizzi"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}