{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "616f8110",
   "metadata": {},
   "source": [
    "(ch:test-fourier)=\n",
    "# Practical: Fourier Optics and Image Processing\n",
    "\n",
    "## Introduction\n",
    "\n",
    "Digital images are two-dimensional intensity functions $f(x,y)$. The information in an image is encoded in a pattern of variation — bright and dark regions, sharp edges, smooth gradients. Fourier analysis gives us a way to decompose any such pattern into a sum of simple sinusoidal waves, each characterised by a **frequency**, an **amplitude**, and a **phase**. This decomposition is not merely mathematical convenience: in electron microscopy it is the physical basis of image formation (through the contrast transfer function) and of reconstruction algorithms (through the Fourier slice theorem).\n",
    "\n",
    "This practical builds up Fourier intuition from the ground up:\n",
    "1. What is a wave?\n",
    "2. Fourier series — decomposing arbitrary 1D signals\n",
    "3. Frequency spectra — looking at signals through the lens of their frequency content\n",
    "4. 2D Fourier analysis — extending to images\n",
    "5. The convolution theorem — efficient filtering\n",
    "6. The contrast transfer function — how the microscope shapes the signal\n",
    "\n",
    "```{admonition} Interactive elements\n",
    ":class: tip\n",
    "Click **Live Code** in the top toolbar to activate the kernel, then expand each **Show code** toggle and click ▶ to run the cell.\n",
    "```\n",
    "\n",
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d20a25e6",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import io, warnings\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "matplotlib.use('agg')\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.ndimage import rotate as nd_rotate, zoom as nd_zoom, gaussian_filter\n",
    "from scipy.signal import convolve2d\n",
    "from ipywidgets import (IntSlider, FloatSlider, Dropdown, VBox, HBox,\n",
    "                        Layout, Output)\n",
    "from IPython.display import display, Image, clear_output\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "_sl  = Layout(width=\"420px\")\n",
    "_sty = {\"description_width\": \"150px\"}\n",
    "\n",
    "def fig2img(fig, dpi=100):\n",
    "    buf = io.BytesIO()\n",
    "    fig.savefig(buf, format='png', dpi=dpi, bbox_inches='tight')\n",
    "    buf.seek(0); plt.close(fig)\n",
    "    return Image(data=buf.read())\n",
    "\n",
    "def make_letter(text, size=64):\n",
    "    \"\"\"Render a letter as a 2D float array via matplotlib Agg canvas.\"\"\"\n",
    "    fig, ax = plt.subplots(figsize=(2, 2))\n",
    "    ax.axis('off')\n",
    "    fig.text(0.23, 0.26, text, fontsize=100)\n",
    "    fig.canvas.draw()\n",
    "    buf = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)\n",
    "    w, h = fig.canvas.get_width_height()\n",
    "    data = buf.reshape(h, w, 4)[:, :, :3].mean(2)\n",
    "    plt.close(fig)\n",
    "    data = (data - data.min()) / (data.max() - data.min() + 1e-10)\n",
    "    data = (~data.astype(bool))[::-1].astype(float)\n",
    "    return nd_zoom(data, size / data.shape[0])\n",
    "\n",
    "def power_spectrum_2d(im):\n",
    "    return np.log(np.abs(np.fft.fftshift(np.fft.fft2(im)))**2 + 1)\n",
    "\n",
    "print(\"Setup complete.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6a26174",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Part 1 – What is a wave?\n",
    "\n",
    "A sinusoidal wave has three fundamental parameters:\n",
    "\n",
    "$$\n",
    "f(t) = A \\cdot \\cos(2\\pi f \\cdot t - \\varphi)\n",
    "$$\n",
    "\n",
    "| Parameter | Symbol | Effect |\n",
    "|-----------|--------|--------|\n",
    "| Amplitude | $A$    | Height of the wave |\n",
    "| Frequency | $f$    | Number of cycles per unit time |\n",
    "| Phase     | $\\varphi$ | Horizontal shift |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07b8e7df",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "amp_sl   = FloatSlider(value=2.0, min=0.0, max=4.0,   step=0.5,\n",
    "                        description=\"Amplitude A\",  style=_sty, layout=_sl)\n",
    "freq_sl  = IntSlider(  value=4,   min=1,   max=12,   step=1,\n",
    "                        description=\"Frequency f\",   style=_sty, layout=_sl)\n",
    "phase_sl = FloatSlider(value=0.0, min=-np.pi, max=np.pi, step=0.25,\n",
    "                        description=\"Phase φ (rad)\", style=_sty, layout=_sl)\n",
    "out_wave = Output()\n",
    "\n",
    "def _update_wave(_=None):\n",
    "    t = np.linspace(0, 1, 1000)\n",
    "    f = amp_sl.value * np.cos(2*np.pi*freq_sl.value*t - phase_sl.value)\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(10, 3.5))\n",
    "    ax.plot(t, f, color='steelblue', linewidth=2)\n",
    "    ax.axhline(0, color='k', linewidth=0.8, linestyle='--')\n",
    "    ax.set_xlim(0, 1); ax.set_ylim(-4.5, 4.5)\n",
    "    ax.set_xlabel(\"t\", fontsize=11); ax.set_ylabel(\"f(t)\", fontsize=11)\n",
    "    ax.set_title(f\"f(t) = {amp_sl.value:.1f}·cos(2π·{freq_sl.value}·t − {phase_sl.value:.2f})\",\n",
    "                 fontsize=11)\n",
    "    ax.grid(alpha=0.3)\n",
    "    with out_wave:\n",
    "        clear_output(wait=True); display(fig2img(fig))\n",
    "\n",
    "for w in [amp_sl, freq_sl, phase_sl]:\n",
    "    w.observe(_update_wave, names='value')\n",
    "display(VBox([amp_sl, freq_sl, phase_sl, out_wave]))\n",
    "_update_wave()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6d2339b",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Part 2 – Fourier series\n",
    "\n",
    "Any periodic signal can be expressed as a sum of sinusoidal waves (**Fourier series**):\n",
    "\n",
    "$$\n",
    "b(x) = A_0 + \\sum_{k=1}^{F} A_k \\cdot \\cos\\!\\left(\\frac{2\\pi k x}{P} - \\varphi_k\\right)\n",
    "$$\n",
    "\n",
    "where $F$ is the maximum frequency, $P$ the period, and $A_k$, $\\varphi_k$ the amplitude and phase at frequency $k$.\n",
    "\n",
    "The amplitudes and phases are computed from the signal via the **Fourier coefficients**:\n",
    "\n",
    "$$\n",
    "a_k = \\frac{2}{P}\\sum_{i=0}^{P} b(i)\\cos\\!\\frac{2\\pi k i}{P}, \\qquad\n",
    "b_k = \\frac{2}{P}\\sum_{i=0}^{P} b(i)\\sin\\!\\frac{2\\pi k i}{P}, \\qquad\n",
    "A_k = \\sqrt{a_k^2 + b_k^2}\n",
    "$$\n",
    "\n",
    "### The box function\n",
    "\n",
    "A classic example is the box (rectangular) function:\n",
    "\n",
    "$$\n",
    "b(x) = \\begin{cases} 1 & -a/2 < x < a/2 \\\\ 0 & \\text{elsewhere} \\end{cases}\n",
    "$$\n",
    "\n",
    "Its Fourier series contains only odd harmonics, and the coefficients decay as $1/k$. As we add more terms, the reconstruction improves but Gibbs ringing remains at the edges."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a04c112",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def getbox(width=200, N=1000):\n",
    "    \"\"\"Box function centred at N//2 with given pixel width.\"\"\"\n",
    "    sig = np.zeros(N)\n",
    "    half = width // 2\n",
    "    sig[N//2 - half : N//2 + half + 1] = 1\n",
    "    return sig\n",
    "\n",
    "def fourier_series_reconstruct(signal, tot_freq):\n",
    "    \"\"\"Reconstruct signal from its first tot_freq Fourier components.\"\"\"\n",
    "    N    = len(signal)\n",
    "    x    = np.arange(N)\n",
    "    recon = np.zeros(N)\n",
    "    for k in range(tot_freq):\n",
    "        ak = 2/N * np.sum(signal * np.cos(2*np.pi*k*x/N))\n",
    "        bk = 2/N * np.sum(signal * np.sin(2*np.pi*k*x/N))\n",
    "        Ak = np.sqrt(ak**2 + bk**2)\n",
    "        pk = np.arctan2(bk, ak)\n",
    "        recon += Ak * np.cos(2*np.pi*k*x/N - pk)\n",
    "    recon -= recon.mean() - signal.mean()\n",
    "    return recon\n",
    "\n",
    "width_sl  = IntSlider(value=200, min=50, max=500, step=50,\n",
    "                       description=\"Box width (px)\", style=_sty, layout=_sl)\n",
    "nwave_sl  = IntSlider(value=15,  min=1,  max=80,  step=1,\n",
    "                       description=\"# waves F\",      style=_sty, layout=_sl)\n",
    "out_box = Output()\n",
    "\n",
    "def _update_box(_=None):\n",
    "    sig   = getbox(width_sl.value)\n",
    "    recon = fourier_series_reconstruct(sig, nwave_sl.value)\n",
    "    xax   = np.linspace(-500, 500, len(sig))\n",
    "    fig, axes = plt.subplots(1, 2, figsize=(13, 3.8))\n",
    "    axes[0].plot(xax, sig,   color='gray',     lw=1.2, label='Box function')\n",
    "    axes[0].plot(xax, recon, color='steelblue', lw=1.6, label=f'Reconstruction (F={nwave_sl.value})')\n",
    "    axes[0].set_xlabel(\"x (pixels)\"); axes[0].set_ylabel(\"Amplitude\")\n",
    "    axes[0].set_title(f\"Fourier series reconstruction: {nwave_sl.value} waves\")\n",
    "    axes[0].set_ylim(-0.4, 1.5); axes[0].legend(fontsize=9); axes[0].grid(alpha=0.3)\n",
    "\n",
    "    # Amplitude spectrum\n",
    "    N    = len(sig)\n",
    "    ks   = np.arange(1, nwave_sl.value+1)\n",
    "    x    = np.arange(N)\n",
    "    amps = []\n",
    "    for k in ks:\n",
    "        ak = 2/N * np.sum(sig * np.cos(2*np.pi*k*x/N))\n",
    "        bk = 2/N * np.sum(sig * np.sin(2*np.pi*k*x/N))\n",
    "        amps.append(np.sqrt(ak**2 + bk**2))\n",
    "    axes[1].bar(ks, amps, color='steelblue', edgecolor='k', linewidth=0.4)\n",
    "    axes[1].set_xlabel(\"Frequency k\"); axes[1].set_ylabel(\"Amplitude $A_k$\")\n",
    "    axes[1].set_title(\"Amplitude spectrum of the box function\")\n",
    "    axes[1].grid(alpha=0.3, axis='y')\n",
    "    plt.tight_layout()\n",
    "    with out_box:\n",
    "        clear_output(wait=True); display(fig2img(fig))\n",
    "\n",
    "for w in [width_sl, nwave_sl]:\n",
    "    w.observe(_update_box, names='value')\n",
    "display(VBox([width_sl, nwave_sl, out_box]))\n",
    "_update_box()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b83c8b80",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Part 3 – Frequency spectrum (1D FFT)\n",
    "\n",
    "Instead of computing Fourier coefficients term by term, the **discrete Fourier transform (DFT)** computes all of them at once:\n",
    "\n",
    "$$\n",
    "F(k) = \\frac{1}{P} \\sum_{m=0}^{P-1} b(m) \\cdot e^{-i 2\\pi k m / P}\n",
    "$$\n",
    "\n",
    "The fast Fourier transform (FFT) computes this in $O(N \\log N)$ rather than $O(N^2)$. The **frequency spectrum** (plot of $|F(k)|$ vs $k$) reveals which frequencies are present in the signal.\n",
    "\n",
    "### Nyquist sampling theorem\n",
    "\n",
    "If a signal contains no frequencies higher than $W$ Hz, it is fully determined by samples taken every $1/(2W)$ seconds:\n",
    "\n",
    "$$\n",
    "f_\\text{max} \\leq \\frac{1}{2 \\Delta t} \\quad \\Leftrightarrow \\quad \\Delta t \\leq \\frac{1}{2 f_\\text{max}}\n",
    "$$\n",
    "\n",
    "The **Nyquist frequency** $f_N = 1/(2\\Delta t)$ is the highest frequency measurable at a given sampling rate. In EM images, the pixel size $\\Delta x$ sets $f_N = 1/(2\\Delta x)$, and only features larger than $2\\Delta x$ can be resolved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a90ee24",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "a1_sl = FloatSlider(value=1.0, min=0.0, max=5.0, step=0.5, description=\"Amp₁\", style=_sty, layout=_sl)\n",
    "f1_sl = IntSlider(  value=10,  min=1,   max=100, step=1,   description=\"Freq₁\", style=_sty, layout=_sl)\n",
    "a2_sl = FloatSlider(value=1.0, min=0.0, max=5.0, step=0.5, description=\"Amp₂\", style=_sty, layout=_sl)\n",
    "f2_sl = IntSlider(  value=20,  min=1,   max=100, step=1,   description=\"Freq₂\", style=_sty, layout=_sl)\n",
    "a3_sl = FloatSlider(value=1.0, min=0.0, max=5.0, step=0.5, description=\"Amp₃\", style=_sty, layout=_sl)\n",
    "f3_sl = IntSlider(  value=40,  min=1,   max=100, step=1,   description=\"Freq₃\", style=_sty, layout=_sl)\n",
    "dc_sl = FloatSlider(value=5.0, min=-5.0, max=10.0, step=0.5, description=\"DC offset\", style=_sty, layout=_sl)\n",
    "out_fft1 = Output()\n",
    "\n",
    "def _update_fft1(_=None):\n",
    "    tot_time = 1.0; N = 1000\n",
    "    t = np.linspace(0, tot_time, N)\n",
    "    sig = (dc_sl.value\n",
    "           + a1_sl.value * np.cos(2*np.pi*f1_sl.value*t)\n",
    "           + a2_sl.value * np.cos(2*np.pi*f2_sl.value*t)\n",
    "           + a3_sl.value * np.cos(2*np.pi*f3_sl.value*t))\n",
    "    dft   = np.fft.fft(sig)\n",
    "    freqs = np.fft.fftfreq(N, tot_time/N)\n",
    "    amp   = np.abs(dft) / N\n",
    "\n",
    "    fig, axes = plt.subplots(1, 2, figsize=(13, 3.8))\n",
    "    axes[0].plot(t, sig, color='steelblue', lw=1.4)\n",
    "    axes[0].set_xlabel(\"Time (s)\"); axes[0].set_ylabel(\"Amplitude\")\n",
    "    axes[0].set_title(\"Signal\"); axes[0].grid(alpha=0.3)\n",
    "    flim = min(100, max(f1_sl.value, f2_sl.value, f3_sl.value) + 20)\n",
    "    axes[1].bar(freqs, amp, width=freqs[1]-freqs[0] if len(freqs) > 1 else 1,\n",
    "                color='steelblue', edgecolor='none')\n",
    "    axes[1].set_xlabel(\"Frequency (Hz)\"); axes[1].set_ylabel(\"Amplitude\")\n",
    "    axes[1].set_xlim(-flim, flim)\n",
    "    axes[1].set_title(\"Frequency spectrum (FFT)\")\n",
    "    axes[1].grid(alpha=0.3)\n",
    "    plt.tight_layout()\n",
    "    with out_fft1:\n",
    "        clear_output(wait=True); display(fig2img(fig))\n",
    "\n",
    "for w in [a1_sl, f1_sl, a2_sl, f2_sl, a3_sl, f3_sl, dc_sl]:\n",
    "    w.observe(_update_fft1, names='value')\n",
    "display(VBox([HBox([a1_sl, f1_sl]), HBox([a2_sl, f2_sl]),\n",
    "              HBox([a3_sl, f3_sl]), dc_sl, out_fft1]))\n",
    "_update_fft1()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c3d6d9f",
   "metadata": {},
   "source": [
    "```{admonition} Question 1\n",
    ":class: seealso\n",
    "Given a detector that records 15 frames per 100 s. What is the maximum frequency component that can be recovered? What does this imply for the pixel size required to resolve a feature of width 2 Å in a cryo-EM image?\n",
    "```\n",
    "\n",
    "---\n",
    "\n",
    "## Part 4 – 2D Fourier analysis\n",
    "\n",
    "Images are 2D signals. A 2D sinusoidal wave has the form:\n",
    "\n",
    "$$\n",
    "f(x,y) = A \\cdot \\sin(2\\pi f_x x + 2\\pi f_y y + \\varphi)\n",
    "$$\n",
    "\n",
    "with separate spatial frequencies $f_x$ (cycles per pixel along $x$) and $f_y$ (cycles per pixel along $y$). The 2D Fourier transform of an image gives the amplitude and phase at every spatial frequency $(f_x, f_y)$.\n",
    "\n",
    "### Properties\n",
    "\n",
    "**Translation property**: Shifting an image in real space does not change the amplitude spectrum — only the phase spectrum.\n",
    "\n",
    "**Rotation property**: Rotating an image rotates its Fourier transform by the same angle.\n",
    "\n",
    "The interactive below shows 2D sine waves at different spatial frequencies, and the effects of translation and rotation on the Fourier transform of a real image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ea8541e",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "fx_sl  = FloatSlider(value=3.0, min=0.0, max=10.0, step=0.5,\n",
    "                      description=\"fx (cycles/im)\", style=_sty, layout=_sl)\n",
    "fy_sl  = FloatSlider(value=1.0, min=0.0, max=10.0, step=0.5,\n",
    "                      description=\"fy (cycles/im)\", style=_sty, layout=_sl)\n",
    "ph_sl2 = FloatSlider(value=0.0, min=-np.pi, max=np.pi, step=0.25,\n",
    "                      description=\"Phase φ\",         style=_sty, layout=_sl)\n",
    "out_2dw = Output()\n",
    "\n",
    "def _update_2dw(_=None):\n",
    "    N  = 128\n",
    "    x  = np.linspace(0, 1, N)\n",
    "    X, Y = np.meshgrid(x, x)\n",
    "    wave = np.sin(2*np.pi*fx_sl.value*X + 2*np.pi*fy_sl.value*Y + ph_sl2.value)\n",
    "    ps   = power_spectrum_2d(wave)\n",
    "\n",
    "    fig, axes = plt.subplots(1, 2, figsize=(10, 4.5))\n",
    "    axes[0].imshow(wave, cmap='RdBu', origin='lower', vmin=-1, vmax=1)\n",
    "    axes[0].set_title(f\"2D wave: fx={fx_sl.value}, fy={fy_sl.value}, φ={ph_sl2.value:.2f}\")\n",
    "    axes[0].axis('off')\n",
    "    axes[1].imshow(ps, cmap='inferno', origin='lower')\n",
    "    axes[1].set_title(\"Power spectrum (log|FT|²)\")\n",
    "    axes[1].axis('off')\n",
    "    plt.tight_layout()\n",
    "    with out_2dw:\n",
    "        clear_output(wait=True); display(fig2img(fig))\n",
    "\n",
    "for w in [fx_sl, fy_sl, ph_sl2]:\n",
    "    w.observe(_update_2dw, names='value')\n",
    "display(VBox([fx_sl, fy_sl, ph_sl2, out_2dw]))\n",
    "_update_2dw()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85a6acf0",
   "metadata": {},
   "source": [
    "### 2D FT of real images\n",
    "\n",
    "Different image features produce characteristic Fourier signatures. A ring in real space produces a ring in Fourier space; a line produces a perpendicular line through the origin."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae863448",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import io, numpy as np, matplotlib\n",
    "matplotlib.use('agg')\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.ndimage import gaussian_filter\n",
    "from IPython.display import display, Image\n",
    "\n",
    "def make_ring(n=128, r=0.3, w=0.04):\n",
    "    y, x = np.mgrid[0:n, 0:n].astype(float) / n - 0.5\n",
    "    R = np.sqrt(x**2 + y**2)\n",
    "    return ((R > r-w/2) & (R < r+w/2)).astype(float)\n",
    "\n",
    "def make_hline(n=128, pos=0.0, width=0.03):\n",
    "    y = np.linspace(-0.5, 0.5, n)\n",
    "    return (np.abs(y - pos) < width).astype(float)[:, None] * np.ones((n, n))\n",
    "\n",
    "def make_dot(n=128, r=0.06):\n",
    "    y, x = np.mgrid[0:n, 0:n].astype(float) / n - 0.5\n",
    "    return (np.sqrt(x**2 + y**2) < r).astype(float)\n",
    "\n",
    "n = 128\n",
    "examples = [\n",
    "    (make_letter('Q', n), \"Letter Q\"),\n",
    "    (make_ring(n), \"Ring\"),\n",
    "    (make_hline(n), \"Horizontal line\"),\n",
    "    (make_dot(n), \"Small dot\"),\n",
    "    (np.random.default_rng(0).standard_normal((n, n)), \"Gaussian noise\"),\n",
    "    (gaussian_filter(make_letter('Q', n), sigma=5), \"Low-pass Q\"),\n",
    "]\n",
    "\n",
    "fig, axes = plt.subplots(2, len(examples), figsize=(3.5*len(examples), 7))\n",
    "for col, (im, title) in enumerate(examples):\n",
    "    ps = power_spectrum_2d(im)\n",
    "    axes[0, col].imshow(im,  cmap='gray',    origin='lower'); axes[0, col].axis('off')\n",
    "    axes[0, col].set_title(title, fontsize=8)\n",
    "    axes[1, col].imshow(ps,  cmap='inferno', origin='lower'); axes[1, col].axis('off')\n",
    "\n",
    "axes[0, 0].set_ylabel(\"Real space\",   fontsize=9)\n",
    "axes[1, 0].set_ylabel(\"Fourier space (log power)\", fontsize=9)\n",
    "plt.tight_layout()\n",
    "_buf = io.BytesIO(); fig.savefig(_buf, format='png', bbox_inches='tight', dpi=96); _buf.seek(0)\n",
    "display(Image(_buf.read()))\n",
    "plt.close('all')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e80bf7dd",
   "metadata": {},
   "source": [
    "### Translation and rotation properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "330dc8b3",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "prop_dd  = Dropdown(options=['Translation (+dx, +dy)', 'Rotation (30°)', 'Both'],\n",
    "                    value='Translation (+dx, +dy)', description=\"Property:\", style=_sty)\n",
    "out_prop = Output()\n",
    "\n",
    "_Q128 = make_letter('Q', 128)\n",
    "_Q128_ft = np.fft.fftshift(np.fft.fft2(_Q128))\n",
    "\n",
    "def _update_prop(_=None):\n",
    "    choice = prop_dd.value\n",
    "    if 'Translation' in choice:\n",
    "        Q2 = np.roll(np.roll(_Q128, 20, axis=0), 30, axis=1)\n",
    "    elif 'Rotation' in choice:\n",
    "        Q2 = nd_rotate(_Q128, 30, reshape=False)\n",
    "    else:\n",
    "        Q2 = nd_rotate(np.roll(np.roll(_Q128, 20, axis=0), 30, axis=1), 30, reshape=False)\n",
    "    F2  = np.fft.fftshift(np.fft.fft2(Q2))\n",
    "    amp1 = np.log(np.abs(_Q128_ft) + 1)\n",
    "    amp2 = np.log(np.abs(F2)       + 1)\n",
    "    pha1 = np.angle(_Q128_ft)\n",
    "    pha2 = np.angle(F2)\n",
    "\n",
    "    fig, axes = plt.subplots(2, 4, figsize=(16, 8))\n",
    "    for row, (im, ft, lab) in enumerate([(_Q128, _Q128_ft, 'Original'), (Q2, F2, choice)]):\n",
    "        axes[row, 0].imshow(im,   cmap='gray',     origin='lower'); axes[row, 0].set_title(f\"{lab}\\nReal space\")\n",
    "        axes[row, 1].imshow(np.log(np.abs(ft)+1), cmap='inferno', origin='lower')\n",
    "        axes[row, 1].set_title(\"FT amplitude\"); axes[row, 2].set_title(\"FT phase\")\n",
    "        axes[row, 2].imshow(np.angle(ft), cmap='hsv', origin='lower', vmin=-np.pi, vmax=np.pi)\n",
    "        for ax in axes[row, :3]: ax.axis('off')\n",
    "    # Difference in amplitude\n",
    "    axes[0, 3].imshow(np.abs(amp1 - amp2), cmap='hot', origin='lower')\n",
    "    axes[0, 3].set_title(\"|Δ amplitude|\"); axes[0, 3].axis('off')\n",
    "    axes[1, 3].imshow(np.abs(pha1 - pha2) % np.pi, cmap='hot', origin='lower')\n",
    "    axes[1, 3].set_title(\"|Δ phase| mod π\"); axes[1, 3].axis('off')\n",
    "    plt.suptitle(f\"Fourier properties: {choice}\", fontsize=10)\n",
    "    plt.tight_layout()\n",
    "    with out_prop:\n",
    "        clear_output(wait=True); display(fig2img(fig))\n",
    "\n",
    "prop_dd.observe(_update_prop, names='value')\n",
    "display(VBox([prop_dd, out_prop]))\n",
    "_update_prop()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "305fc753",
   "metadata": {},
   "source": [
    "```{admonition} Question 2\n",
    ":class: seealso\n",
    "Play with the translation and rotation of the image and observe the amplitude spectrum. What changes and what stays the same in each case? Can you explain why?\n",
    "```\n",
    "\n",
    "---\n",
    "\n",
    "## Part 5 – The convolution theorem\n",
    "\n",
    "**Convolution** describes how one function modifies another by sweeping over all positions:\n",
    "\n",
    "$$\n",
    "(f * g)(y) = \\int_{-\\infty}^{\\infty} f(x)\\,g(y-x)\\,dx\n",
    "$$\n",
    "\n",
    "In image processing, convolving an image with a **kernel** (a small matrix) applies a local operation: smoothing, sharpening, or edge detection. The **convolution theorem** states:\n",
    "\n",
    "$$\n",
    "\\mathcal{F}\\{f * g\\} = \\mathcal{F}\\{f\\} \\cdot \\mathcal{F}\\{g\\}\n",
    "$$\n",
    "\n",
    "Convolution in real space equals *multiplication* in Fourier space. This makes filtering in Fourier space much faster for large kernels than direct spatial convolution.\n",
    "\n",
    "### Common kernels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c0cf558",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import io, numpy as np, matplotlib\n",
    "matplotlib.use('agg')\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.signal import convolve2d\n",
    "from IPython.display import display, Image\n",
    "\n",
    "def sobel_x(n=3):\n",
    "    return np.array([[-1,0,1],[-2,0,2],[-1,0,1]], float)\n",
    "\n",
    "def sobel_y(n=3):\n",
    "    return np.array([[1,2,1],[0,0,0],[-1,-2,-1]], float)\n",
    "\n",
    "def laplacian():\n",
    "    return np.array([[0,1,0],[1,-4,1],[0,1,0]], float)\n",
    "\n",
    "def gaussian_kernel(size=9, sigma=2.0):\n",
    "    ax = np.arange(-(size//2), size//2+1)\n",
    "    k  = np.exp(-ax**2/(2*sigma**2))\n",
    "    k  = np.outer(k, k); return k / k.sum()\n",
    "\n",
    "def box_blur(size=5):\n",
    "    return np.ones((size,size)) / size**2\n",
    "\n",
    "def sharpen():\n",
    "    return np.array([[0,-1,0],[-1,5,-1],[0,-1,0]], float)\n",
    "\n",
    "Q64 = make_letter('Q', 64)\n",
    "\n",
    "kernels = {\n",
    "    'Gaussian blur (σ=2)': gaussian_kernel(9, 2.0),\n",
    "    'Box blur (5×5)':       box_blur(5),\n",
    "    'Laplacian':            laplacian(),\n",
    "    'Sobel X':              sobel_x(),\n",
    "    'Sobel Y':              sobel_y(),\n",
    "    'Sharpen':              sharpen(),\n",
    "}\n",
    "\n",
    "fig, axes = plt.subplots(3, len(kernels), figsize=(3.5*len(kernels), 9))\n",
    "for col, (name, kernel) in enumerate(kernels.items()):\n",
    "    filtered = convolve2d(Q64, kernel, mode='same', boundary='wrap')\n",
    "    kshow    = kernel.copy()\n",
    "    axes[0, col].imshow(Q64,      cmap='gray',    origin='lower'); axes[0, col].axis('off')\n",
    "    axes[0, col].set_title(name,  fontsize=7)\n",
    "    n_k      = kshow.shape[0]\n",
    "    axes[1, col].imshow(kshow,    cmap='RdBu',    origin='lower',\n",
    "                        vmin=-abs(kshow).max(), vmax=abs(kshow).max())\n",
    "    axes[1, col].set_title(f\"Kernel ({n_k}×{n_k})\", fontsize=7); axes[1, col].axis('off')\n",
    "    axes[2, col].imshow(filtered, cmap='gray',    origin='lower'); axes[2, col].axis('off')\n",
    "    axes[2, col].set_title(\"Result\", fontsize=7)\n",
    "\n",
    "axes[0, 0].set_ylabel(\"Input\", fontsize=9)\n",
    "axes[1, 0].set_ylabel(\"Kernel\", fontsize=9)\n",
    "axes[2, 0].set_ylabel(\"Filtered output\", fontsize=9)\n",
    "plt.tight_layout()\n",
    "_buf = io.BytesIO(); fig.savefig(_buf, format='png', bbox_inches='tight', dpi=96); _buf.seek(0)\n",
    "display(Image(_buf.read()))\n",
    "plt.close('all')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02950465",
   "metadata": {},
   "source": [
    "### Interactive kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a10b98f",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "ker_dd  = Dropdown(options=['Gaussian blur','Box blur','Laplacian','Sobel X','Sobel Y','Sharpen'],\n",
    "                   value='Gaussian blur', description=\"Kernel:\", style=_sty)\n",
    "ksz_sl  = IntSlider(value=9, min=3, max=21, step=2,\n",
    "                     description=\"Kernel size\", style=_sty, layout=_sl)\n",
    "out_ker = Output()\n",
    "\n",
    "def _update_ker(_=None):\n",
    "    name = ker_dd.value; ksz = ksz_sl.value\n",
    "    if   name == 'Gaussian blur': kernel = gaussian_kernel(ksz, ksz/6)\n",
    "    elif name == 'Box blur':      kernel = box_blur(ksz)\n",
    "    elif name == 'Laplacian':     kernel = laplacian()\n",
    "    elif name == 'Sobel X':       kernel = sobel_x()\n",
    "    elif name == 'Sobel Y':       kernel = sobel_y()\n",
    "    else:                         kernel = sharpen()\n",
    "\n",
    "    filtered  = convolve2d(Q64, kernel, mode='same', boundary='wrap')\n",
    "    ps_input  = power_spectrum_2d(Q64)\n",
    "    ps_kernel = power_spectrum_2d(np.pad(kernel, (32-kernel.shape[0]//2,\n",
    "                                                  32-kernel.shape[0]//2+1)))\n",
    "    ps_out    = power_spectrum_2d(filtered)\n",
    "\n",
    "    fig, axes = plt.subplots(2, 3, figsize=(13, 8))\n",
    "    axes[0,0].imshow(Q64,       cmap='gray',    origin='lower'); axes[0,0].axis('off')\n",
    "    axes[0,0].set_title(\"Input image\")\n",
    "    axes[0,1].imshow(kernel, cmap='RdBu', origin='lower',\n",
    "                     vmin=-abs(kernel).max(), vmax=abs(kernel).max())\n",
    "    axes[0,1].set_title(f\"Kernel: {name}\\n({kernel.shape[0]}×{kernel.shape[1]})\")\n",
    "    axes[0,1].axis('off')\n",
    "    axes[0,2].imshow(filtered,  cmap='gray',    origin='lower'); axes[0,2].axis('off')\n",
    "    axes[0,2].set_title(\"Filtered output (real space)\")\n",
    "    axes[1,0].imshow(ps_input,  cmap='inferno', origin='lower'); axes[1,0].axis('off')\n",
    "    axes[1,0].set_title(\"FT of input\")\n",
    "    axes[1,1].imshow(ps_kernel, cmap='inferno', origin='lower'); axes[1,1].axis('off')\n",
    "    axes[1,1].set_title(\"FT of kernel (transfer function)\")\n",
    "    axes[1,2].imshow(ps_out,    cmap='inferno', origin='lower'); axes[1,2].axis('off')\n",
    "    axes[1,2].set_title(\"FT of output = FT(in) × FT(kernel)\")\n",
    "    plt.suptitle(\"Convolution theorem: multiplication in Fourier space\", fontsize=10)\n",
    "    plt.tight_layout()\n",
    "    with out_ker:\n",
    "        clear_output(wait=True); display(fig2img(fig))\n",
    "\n",
    "for w in [ker_dd, ksz_sl]:\n",
    "    w.observe(_update_ker, names='value')\n",
    "display(VBox([ker_dd, ksz_sl, out_ker]))\n",
    "_update_ker()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a620de0a",
   "metadata": {},
   "source": [
    "```{admonition} Question 3\n",
    ":class: seealso\n",
    "What does the FT of a Gaussian blur kernel look like? What does this tell you about which spatial frequencies are attenuated? How does this compare to the Laplacian kernel?\n",
    "```\n",
    "\n",
    "---\n",
    "\n",
    "## Part 6 – Low-pass and high-pass filtering in Fourier space\n",
    "\n",
    "We can design filters by directly modifying the Fourier transform: set Fourier coefficients outside (or inside) a radius to zero before inverse-transforming. This is equivalent to convolving with the Fourier transform of the mask function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1da7b8e",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "ftype_dd = Dropdown(options=['Low-pass', 'High-pass', 'Band-pass'],\n",
    "                    value='Low-pass', description=\"Filter type:\", style=_sty)\n",
    "fcut_sl  = FloatSlider(value=0.2, min=0.02, max=0.5, step=0.02,\n",
    "                        description=\"Cutoff f₁\",    style=_sty, layout=_sl)\n",
    "fcut2_sl = FloatSlider(value=0.4, min=0.02, max=0.5, step=0.02,\n",
    "                        description=\"Cutoff f₂ (band)\", style=_sty, layout=_sl)\n",
    "out_lp   = Output()\n",
    "\n",
    "def freq_mask(n, f1, f2=None, ftype='Low-pass'):\n",
    "    \"\"\"Build a circular frequency mask.\"\"\"\n",
    "    y, x  = np.mgrid[0:n, 0:n].astype(float)\n",
    "    y -= n//2; x -= n//2\n",
    "    R = np.sqrt(x**2 + y**2) / n\n",
    "    if ftype == 'Low-pass':\n",
    "        return (R <= f1).astype(float)\n",
    "    elif ftype == 'High-pass':\n",
    "        return (R >= f1).astype(float)\n",
    "    else:\n",
    "        return ((R >= f1) & (R <= (f2 or f1))).astype(float)\n",
    "\n",
    "def _update_lp(_=None):\n",
    "    n      = Q64.shape[0]\n",
    "    ftype  = ftype_dd.value\n",
    "    f1     = fcut_sl.value\n",
    "    f2     = fcut2_sl.value\n",
    "    mask   = freq_mask(n, f1, f2, ftype)\n",
    "    F      = np.fft.fftshift(np.fft.fft2(Q64))\n",
    "    F_filt = F * mask\n",
    "    filtered = np.real(np.fft.ifft2(np.fft.ifftshift(F_filt)))\n",
    "\n",
    "    fig, axes = plt.subplots(1, 4, figsize=(16, 4.2))\n",
    "    axes[0].imshow(Q64,      cmap='gray',    origin='lower'); axes[0].axis('off')\n",
    "    axes[0].set_title(\"Original\")\n",
    "    axes[1].imshow(mask,     cmap='Blues',   origin='lower'); axes[1].axis('off')\n",
    "    axes[1].set_title(f\"Fourier mask\\n({ftype}, f₁={f1:.2f})\")\n",
    "    axes[2].imshow(np.log(np.abs(F_filt)+1), cmap='inferno', origin='lower')\n",
    "    axes[2].axis('off'); axes[2].set_title(\"Filtered FT\")\n",
    "    axes[3].imshow(filtered, cmap='gray',    origin='lower'); axes[3].axis('off')\n",
    "    axes[3].set_title(\"Filtered image\")\n",
    "    plt.tight_layout()\n",
    "    with out_lp:\n",
    "        clear_output(wait=True); display(fig2img(fig))\n",
    "\n",
    "for w in [ftype_dd, fcut_sl, fcut2_sl]:\n",
    "    w.observe(_update_lp, names='value')\n",
    "display(VBox([ftype_dd, fcut_sl, fcut2_sl, out_lp]))\n",
    "_update_lp()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cd0a128",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Part 7 – The contrast transfer function (CTF)\n",
    "\n",
    "In phase-contrast electron microscopy, the image intensity is related to the projected potential of the specimen through the **contrast transfer function (CTF)**. The CTF modulates the Fourier amplitudes as a function of spatial frequency $k$:\n",
    "\n",
    "$$\n",
    "\\text{CTF}(k) = -\\sin\\!\\left[\\Delta\\varphi + \\frac{-\\pi}{2}C_s\\lambda^3 k^4 + \\pi\\lambda\\Delta_f k^2\\right]\n",
    "$$\n",
    "\n",
    "where:\n",
    "\n",
    "| Symbol | Quantity |\n",
    "|--------|----------|\n",
    "| $k$ | spatial frequency (Å⁻¹) |\n",
    "| $C_s$ | spherical aberration coefficient |\n",
    "| $\\lambda$ | relativistic electron wavelength |\n",
    "| $\\Delta_f$ | defocus (positive = underfocus) |\n",
    "| $\\Delta\\varphi$ | additional phase shift (e.g. from phase plate) |\n",
    "\n",
    "The electron wavelength depends on the accelerating voltage $V$ (in eV):\n",
    "\n",
    "$$\n",
    "\\lambda = \\frac{12.264}{\\sqrt{V(1 + V \\cdot 0.98 \\times 10^{-6})}} \\; \\text{Å}\n",
    "$$\n",
    "\n",
    "Defocusing causes **phase reversals** at specific spatial frequencies (the CTF zeros). Frequencies near these zeros are not faithfully transferred to the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc4761f2",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def relativistic_lambda(voltage_eV):\n",
    "    \"\"\"Relativistic electron wavelength in Angstrom.\"\"\"\n",
    "    return 12.264 / np.sqrt(voltage_eV * (1 + voltage_eV * 0.98e-6))\n",
    "\n",
    "def ctf_1d(defocus_um, voltage_kV=300, Cs_mm=2.7, delta_phi=0, B_factor=0,\n",
    "           n_pts=256, max_freq=0.5):\n",
    "    \"\"\"Compute 1D CTF curve.\"\"\"\n",
    "    k       = np.linspace(0, max_freq, n_pts)\n",
    "    V_eV    = voltage_kV * 1e3\n",
    "    lam     = relativistic_lambda(V_eV)\n",
    "    df_A    = defocus_um * 1e4        # µm → Å\n",
    "    Cs_A    = Cs_mm * 1e7             # mm → Å\n",
    "    gamma   = (-np.pi/2)*Cs_A*lam**3*k**4 + np.pi*lam*df_A*k**2\n",
    "    ctf     = -np.sin(delta_phi + gamma)\n",
    "    if B_factor > 0:\n",
    "        ctf *= np.exp(-B_factor * k**2)\n",
    "    return k, ctf\n",
    "\n",
    "df_sl   = FloatSlider(value=2.0,  min=0.5,  max=10.0, step=0.5,\n",
    "                       description=\"Defocus (µm)\",   style=_sty, layout=_sl)\n",
    "volt_sl = FloatSlider(value=300,  min=100,  max=300,  step=100,\n",
    "                       description=\"Voltage (kV)\",   style=_sty, layout=_sl)\n",
    "cs_sl   = FloatSlider(value=2.7,  min=0.0,  max=5.0,  step=0.1,\n",
    "                       description=\"Cs (mm)\",         style=_sty, layout=_sl)\n",
    "phi_sl  = FloatSlider(value=0.0,  min=-np.pi/2, max=np.pi/2, step=0.1,\n",
    "                       description=\"Phase shift Δφ\", style=_sty, layout=_sl)\n",
    "bf_sl   = FloatSlider(value=0.0,  min=0.0,  max=200.0, step=10.0,\n",
    "                       description=\"B-factor (Ų)\",  style=_sty, layout=_sl)\n",
    "out_ctf = Output()\n",
    "\n",
    "def _update_ctf(_=None):\n",
    "    k, ctf_curve = ctf_1d(df_sl.value, volt_sl.value, cs_sl.value,\n",
    "                           phi_sl.value, bf_sl.value)\n",
    "    lam = relativistic_lambda(volt_sl.value * 1e3)\n",
    "\n",
    "    # Apply CTF to Q image\n",
    "    n     = Q64.shape[0]\n",
    "    apix  = 1.0\n",
    "    ky    = np.fft.fftfreq(n, apix)\n",
    "    kx    = np.fft.rfftfreq(n, apix)\n",
    "    KX, KY = np.meshgrid(kx, ky)\n",
    "    K2D   = np.sqrt(KX**2 + KY**2) * 0.5 / (n * apix)  # scale to ~0–0.5 Å⁻¹ range\n",
    "    df_A  = df_sl.value * 1e4\n",
    "    Cs_A  = cs_sl.value * 1e7\n",
    "    gamma2 = (-np.pi/2)*Cs_A*lam**3*K2D**4 + np.pi*lam*df_A*K2D**2\n",
    "    ctf2d = -np.sin(phi_sl.value + gamma2)\n",
    "    if bf_sl.value > 0:\n",
    "        ctf2d *= np.exp(-bf_sl.value * K2D**2)\n",
    "\n",
    "    Fq    = np.fft.rfftn(Q64)\n",
    "    Q_ctf = np.fft.irfftn(Fq * ctf2d, Q64.shape)\n",
    "    Q_noisy = Q_ctf + np.random.default_rng(1).standard_normal(Q64.shape)*0.3\n",
    "\n",
    "    fig, axes = plt.subplots(1, 4, figsize=(16, 4.5))\n",
    "    axes[0].plot(k, ctf_curve, color='steelblue', lw=1.8)\n",
    "    axes[0].axhline(0, color='k', lw=0.8, ls='--')\n",
    "    axes[0].set_xlabel(\"Spatial frequency (Å⁻¹)\"); axes[0].set_ylabel(\"CTF\")\n",
    "    axes[0].set_title(f\"CTF (Δf={df_sl.value} µm, V={volt_sl.value:.0f} kV)\")\n",
    "    axes[0].set_ylim(-1.2, 1.2); axes[0].grid(alpha=0.3)\n",
    "\n",
    "    ps_ctf2d = np.log(np.abs(np.fft.fftshift(np.fft.fft2(Q_ctf)))**2 + 1)\n",
    "    axes[1].imshow(Q64,     cmap='gray', origin='lower'); axes[1].axis('off')\n",
    "    axes[1].set_title(\"True image\")\n",
    "    axes[2].imshow(Q_noisy, cmap='gray', origin='lower'); axes[2].axis('off')\n",
    "    axes[2].set_title(\"CTF-modulated image\")\n",
    "    axes[3].imshow(ps_ctf2d, cmap='inferno', origin='lower'); axes[3].axis('off')\n",
    "    axes[3].set_title(\"Power spectrum (Thon rings)\")\n",
    "    plt.tight_layout()\n",
    "    with out_ctf:\n",
    "        clear_output(wait=True); display(fig2img(fig))\n",
    "\n",
    "for w in [df_sl, volt_sl, cs_sl, phi_sl, bf_sl]:\n",
    "    w.observe(_update_ctf, names='value')\n",
    "display(VBox([df_sl, volt_sl, cs_sl, phi_sl, bf_sl, out_ctf]))\n",
    "_update_ctf()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e785bbb3",
   "metadata": {},
   "source": [
    "### CTF correction\n",
    "\n",
    "To recover the true image from a CTF-modulated observation, we need to correct for the CTF. Three standard methods:\n",
    "\n",
    "**Method 1 – Phase flipping**: Multiply Fourier amplitudes by the sign of the CTF. Brings all amplitudes to positive, but does not correct their magnitudes.\n",
    "\n",
    "**Method 2 – Full CTF correction with threshold**: Divide by the CTF, but ignore frequencies where $|\\text{CTF}| < \\epsilon$ (near zeros) to avoid division instability.\n",
    "\n",
    "**Method 3 – Wiener filter**: Divide by CTF with regularisation:\n",
    "$$\n",
    "\\hat{F}(k) = \\frac{F_\\text{obs}(k) \\cdot \\text{CTF}(k)}{\\text{CTF}(k)^2 + \\text{SNR}^{-1}}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87cfb657",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "corr_dd = Dropdown(options=['Phase flip', 'Full CTF (threshold)', 'Wiener filter'],\n",
    "                   value='Wiener filter', description=\"Correction:\", style=_sty)\n",
    "snr_sl  = FloatSlider(value=1.0,   min=0.01, max=10.0, step=0.1,\n",
    "                       description=\"SNR (Wiener)\", style=_sty, layout=_sl)\n",
    "thr_sl  = FloatSlider(value=0.05,  min=0.01, max=0.5,  step=0.01,\n",
    "                       description=\"Threshold ε\",  style=_sty, layout=_sl)\n",
    "df2_sl  = FloatSlider(value=2.0,   min=0.5,  max=10.0, step=0.5,\n",
    "                       description=\"Defocus (µm)\", style=_sty, layout=_sl)\n",
    "out_ctfc = Output()\n",
    "\n",
    "def _update_ctfc(_=None):\n",
    "    lam   = relativistic_lambda(300e3)\n",
    "    n     = Q64.shape[0]; apix = 1.0\n",
    "    ky    = np.fft.fftfreq(n, apix); kx = np.fft.rfftfreq(n, apix)\n",
    "    KX, KY = np.meshgrid(kx, ky)\n",
    "    K2D   = np.sqrt(KX**2 + KY**2) * 0.5 / (n * apix)\n",
    "    df_A  = df2_sl.value * 1e4; Cs_A = 2.7e7\n",
    "    gamma2 = (-np.pi/2)*Cs_A*lam**3*K2D**4 + np.pi*lam*df_A*K2D**2\n",
    "    ctf2d = -np.sin(gamma2)\n",
    "\n",
    "    rng   = np.random.default_rng(1)\n",
    "    Fq    = np.fft.rfftn(Q64)\n",
    "    F_obs = Fq * ctf2d + rng.standard_normal(Fq.shape)*0.3\n",
    "\n",
    "    method = corr_dd.value\n",
    "    if method == 'Phase flip':\n",
    "        F_corr = F_obs * np.sign(ctf2d)\n",
    "    elif method == 'Full CTF (threshold)':\n",
    "        eps    = thr_sl.value\n",
    "        F_corr = np.where(np.abs(ctf2d) >= eps, F_obs / ctf2d, 0.0)\n",
    "    else:\n",
    "        snr    = snr_sl.value\n",
    "        F_corr = F_obs * ctf2d / (ctf2d**2 + 1.0/snr)\n",
    "\n",
    "    Q_corr  = np.fft.irfftn(F_corr, Q64.shape)\n",
    "    Q_ctf_i = np.fft.irfftn(F_obs,  Q64.shape)\n",
    "\n",
    "    fig, axes = plt.subplots(1, 3, figsize=(12, 4))\n",
    "    axes[0].imshow(Q64,     cmap='gray', origin='lower'); axes[0].axis('off')\n",
    "    axes[0].set_title(\"True image\")\n",
    "    axes[1].imshow(Q_ctf_i, cmap='gray', origin='lower'); axes[1].axis('off')\n",
    "    axes[1].set_title(f\"CTF-modulated (Δf={df2_sl.value} µm)\")\n",
    "    axes[2].imshow(Q_corr,  cmap='gray', origin='lower'); axes[2].axis('off')\n",
    "    axes[2].set_title(f\"Corrected: {method}\")\n",
    "    plt.tight_layout()\n",
    "    with out_ctfc:\n",
    "        clear_output(wait=True); display(fig2img(fig))\n",
    "\n",
    "for w in [corr_dd, snr_sl, thr_sl, df2_sl]:\n",
    "    w.observe(_update_ctfc, names='value')\n",
    "display(VBox([df2_sl, corr_dd, thr_sl, snr_sl, out_ctfc]))\n",
    "_update_ctfc()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "814a1e25",
   "metadata": {},
   "source": [
    "```{admonition} Question 4\n",
    ":class: seealso\n",
    "Compare the three CTF correction methods. What are the trade-offs?\n",
    "- Phase flipping: what information is still lost?\n",
    "- Full CTF correction with threshold: what happens to frequencies near the CTF zeros?\n",
    "- Wiener filter: what is the optimal SNR for your simulated data?\n",
    "```\n",
    "\n",
    "---\n",
    "\n",
    "## Summary\n",
    "\n",
    "| Topic | Key concept |\n",
    "|-------|-------------|\n",
    "| Sinusoidal waves | $f(t) = A\\cos(2\\pi f t - \\varphi)$; three parameters: $A$, $f$, $\\varphi$ |\n",
    "| Fourier series | Any periodic signal = sum of sinusoids; coefficients via $a_k$, $b_k$ integrals |\n",
    "| FFT | Fast algorithm ($O(N\\log N)$) for discrete FT; frequency spectrum = $|F(k)|$ |\n",
    "| Nyquist | Sampling at $\\Delta t$ can recover frequencies up to $1/(2\\Delta t)$ |\n",
    "| 2D FT | Generalises 1D; amplitude spectrum = $|F(f_x,f_y)|$; rotation and translation properties |\n",
    "| Convolution theorem | $\\mathcal{F}\\{f*g\\} = \\mathcal{F}\\{f\\}\\cdot\\mathcal{F}\\{g\\}$; filter in Fourier space |\n",
    "| CTF | $-\\sin(\\gamma(k))$; phase reversals at zeros; correct by phase flip, division, or Wiener filter |"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "text_representation": {
    "extension": ".md",
    "format_name": "myst",
    "format_version": 0.13,
    "jupytext_version": "1.10.3"
   }
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "mystnb": {
   "execution_timeout": 180
  },
  "source_map": [
   14,
   38,
   79,
   97,
   126,
   158,
   219,
   243,
   285,
   312,
   344,
   350,
   395,
   399,
   442,
   469,
   529,
   533,
   580,
   593,
   642,
   672,
   749,
   764,
   818
  ]
 },
 "nbformat": 4,
 "nbformat_minor": 5
}