{ "cells": [ { "cell_type": "markdown", "id": "b877f551", "metadata": {}, "source": [ "(ch:test-spa)=\n", "# Practical: Single-Particle Analysis\n", "\n", "## The reconstruction problem\n", "\n", "In electron microscopy we want to deduce the 3D structure of a structurally homogeneous 'particle' — for example a protein — from cryo-EM images. Each image shows the particle in a different, unknown orientation. To make the problem tractable we work in 2D: we treat the unknown structure as a 2D image $A$ and each observation as a rotated, noisy version:\n", "\n", "$$\n", "X_i = R^{\\theta_i} A + \\sigma G_i\n", "$$ (eq:spa-obs)\n", "\n", "where $R^{\\theta_i}$ is rotation by the unknown angle $\\theta_i$ and $G_i$ is i.i.d. Gaussian noise. If we can estimate $\\theta_i$ for each image, we can recover $A$ by rotating each image back and averaging.\n", "\n", "We use the letter **Q** as our 2D 'particle' because it is immediately recognisable and its asymmetry makes orientation unambiguous. In real SPA the structure is of course unknown.\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", "### Helper functions\n", "\n", "The cell below defines all functions used throughout this chapter. Run it first." ] }, { "cell_type": "code", "execution_count": null, "id": "fd95f88a", "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 import ndimage\n", "from scipy.ndimage import zoom as nd_zoom, rotate as nd_rotate, gaussian_filter\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", "# ── image helpers ─────────────────────────────────────────────────────────────\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 make_circle(size=64, radius_frac=0.35):\n", " \"\"\"Filled circle — a featureless starting reference.\"\"\"\n", " y, x = np.mgrid[0:size, 0:size].astype(float)\n", " return ((x - size/2)**2 + (y - size/2)**2 < (size*radius_frac)**2).astype(float)\n", "\n", "def low_pass_filter(im, sigma=3.0):\n", " \"\"\"Gaussian low-pass filter (approximates tanh LP filter).\"\"\"\n", " return gaussian_filter(im, sigma=sigma)\n", "\n", "def show(im, ax=None, title='', vmin=None, vmax=None, cmap='gray'):\n", " if ax is None:\n", " _, ax = plt.subplots()\n", " ax.imshow(im, cmap=cmap, origin='lower',\n", " vmin=vmin or im.min(), vmax=vmax or im.max())\n", " ax.axis('off')\n", " if title:\n", " ax.set_title(title, fontsize=9)\n", "\n", "def fig2img(fig):\n", " buf = io.BytesIO(); fig.savefig(buf, format='png', dpi=100, bbox_inches='tight')\n", " buf.seek(0); plt.close(fig)\n", " return Image(data=buf.read())\n", "\n", "# ── simulation ────────────────────────────────────────────────────────────────\n", "\n", "def simulate_image(image, angle, sigma=1, rng=None):\n", " \"\"\"Rotate image and add Gaussian noise.\"\"\"\n", " rng = rng or np.random.default_rng()\n", " return nd_rotate(image, angle, reshape=False) + sigma * rng.standard_normal(image.shape)\n", "\n", "def simulate_images(image, sigma, N, seed=42):\n", " \"\"\"Generate N noisy randomly-rotated images.\"\"\"\n", " rng = np.random.default_rng(seed)\n", " angles = rng.uniform(0, 360, N)\n", " imgs = np.array([simulate_image(image, a, sigma, rng) for a in angles])\n", " return angles, imgs\n", "\n", "# ── alignment ─────────────────────────────────────────────────────────────────\n", "\n", "def rotate_reference(reference, delta_angle=5):\n", " \"\"\"Rotate reference through all candidate angles in [0, 360).\"\"\"\n", " angles = np.arange(0, 360, delta_angle)\n", " stack = np.array([nd_rotate(reference, a, reshape=False) for a in angles])\n", " return angles, stack\n", "\n", "def correlation(image, reference_stack):\n", " \"\"\"Cross-correlation score for each reference in the stack.\"\"\"\n", " return (image[None] * reference_stack).mean(axis=(1, 2))\n", "\n", "def align_images(images_stack, reference_stack, cand_angles):\n", " \"\"\"Return estimated angle for each image.\"\"\"\n", " estimated = []\n", " for img in images_stack:\n", " scores = correlation(img, reference_stack)\n", " best_idx = int(scores.argmax())\n", " estimated.append(cand_angles[best_idx])\n", " return np.array(estimated)\n", "\n", "# ── reconstruction ─────────────────────────────────────────────────────────────\n", "\n", "def reconstruct(images_stack, estimated_angles):\n", " \"\"\"Average back-rotated images.\"\"\"\n", " recon = np.zeros_like(images_stack[0])\n", " for angle, img in zip(estimated_angles, images_stack):\n", " recon += nd_rotate(img, -angle, reshape=False)\n", " return recon / len(images_stack)\n", "\n", "print(\"Helper functions loaded.\")" ] }, { "cell_type": "markdown", "id": "6123bc40", "metadata": {}, "source": [ "---\n", "\n", "## Part 0 – The test image" ] }, { "cell_type": "code", "execution_count": null, "id": "646d838d", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "# Pre-compute the letter Q once (used in all subsequent cells)\n", "Q = make_letter('Q', size=64)\n", "\n", "fig, ax = plt.subplots(figsize=(3, 3))\n", "show(Q, ax=ax, title=\"Test image: Q (64×64)\")\n", "display(fig2img(fig))" ] }, { "cell_type": "markdown", "id": "af56c188", "metadata": {}, "source": [ "---\n", "\n", "## Simulating cryo-EM images\n", "\n", "The interactive below shows a single simulated particle image. The image formation model is:\n", "\n", "$$\n", "X = R^\\theta A + \\sigma G\n", "$$\n", "\n", "Drag the **angle** slider to rotate the particle and the **σ** slider to add noise." ] }, { "cell_type": "code", "execution_count": null, "id": "8ba095fb", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "_sl = Layout(width=\"380px\")\n", "_sty = {\"description_width\": \"120px\"}\n", "angle_sl = IntSlider( value=30, min=0, max=359, step=1,\n", " description=\"Angle (°)\", style=_sty, layout=_sl)\n", "sigma_sl0 = FloatSlider(value=1.0, min=0.0, max=8.0, step=0.5,\n", " description=\"Noise σ\", style=_sty, layout=_sl)\n", "\n", "out_sim = Output()\n", "def _update_sim(_=None):\n", " rng = np.random.default_rng(7)\n", " img = simulate_image(Q, angle_sl.value, sigma_sl0.value, rng)\n", " fig, axes = plt.subplots(1, 3, figsize=(10, 3.5))\n", " show(Q, axes[0], title=\"True structure A\")\n", " show(nd_rotate(Q, angle_sl.value, reshape=False), axes[1],\n", " title=f\"Rotated: {angle_sl.value}°\")\n", " v = max(abs(img.min()), abs(img.max()))\n", " show(img, axes[2], vmin=-v, vmax=v,\n", " title=f\"Observed: θ={angle_sl.value}°, σ={sigma_sl0.value:.1f}\")\n", " with out_sim:\n", " clear_output(wait=True); display(fig2img(fig))\n", "\n", "for w in [angle_sl, sigma_sl0]:\n", " w.observe(_update_sim, names='value')\n", "display(VBox([angle_sl, sigma_sl0, out_sim]))\n", "_update_sim()" ] }, { "cell_type": "markdown", "id": "4a791177", "metadata": {}, "source": [ "---\n", "\n", "## Part 1 – Choosing an initial reference\n", "\n", "The alignment algorithm needs a starting reference to compare images against. Three strategies are common:\n", "\n", "| Reference | Advantage | Risk |\n", "|-----------|-----------|------|\n", "| Low-pass filtered true structure | Fast convergence | Only works if structure is already known (circular reasoning) |\n", "| Featureless circle | No reference bias | Slow convergence; poor angular discrimination |\n", "| Random noise / another structure | Quick to obtain | Risk of converging to wrong answer (model bias) |\n", "\n", "In practice, a featureless circle or a low-pass-filtered version of a prior reconstruction is used. The interactive below lets you choose the starting reference and see how much angular discrimination it provides for a single noisy image." ] }, { "cell_type": "code", "execution_count": null, "id": "b99e4d29", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "_sl2 = Layout(width=\"420px\")\n", "ref_dd = Dropdown(options=['Q (oracle)', 'Blurred Q (σ=5)', 'Circle',\n", " 'P (wrong)', 'O (wrong)'],\n", " value='Blurred Q (σ=5)', description=\"Reference:\", style=_sty)\n", "sigma_sl1 = FloatSlider(value=2.0, min=0.0, max=8.0, step=0.5,\n", " description=\"Noise σ\", style=_sty, layout=_sl2)\n", "\n", "_REFS = {\n", " 'Q (oracle)' : Q,\n", " 'Blurred Q (σ=5)' : low_pass_filter(Q, sigma=5),\n", " 'Circle' : make_circle(64),\n", " 'P (wrong)' : make_letter('P', 64),\n", " 'O (wrong)' : make_letter('O', 64),\n", "}\n", "# Pre-compute candidate rotations at 5° for each reference\n", "_CAND5 = np.arange(0, 360, 5)\n", "_STACKS5 = {name: np.array([nd_rotate(ref, a, reshape=False) for a in _CAND5])\n", " for name, ref in _REFS.items()}\n", "\n", "out_ref = Output()\n", "def _update_ref(_=None):\n", " rng = np.random.default_rng(3)\n", " true_ang = 130\n", " img = simulate_image(Q, true_ang, sigma_sl1.value, rng)\n", " ref_name = ref_dd.value\n", " ref = _REFS[ref_name]; stack = _STACKS5[ref_name]\n", " scores = correlation(img, stack)\n", " best_ang = _CAND5[scores.argmax()]\n", "\n", " fig, axes = plt.subplots(1, 4, figsize=(14, 3.5))\n", " show(ref, axes[0], title=f\"Reference\\n({ref_name})\")\n", " v = max(abs(img.min()), abs(img.max()))\n", " show(img, axes[1], vmin=-v, vmax=v,\n", " title=f\"Noisy image (σ={sigma_sl1.value:.1f})\\ntrue angle: {true_ang}°\")\n", " axes[2].plot(_CAND5, scores, color='steelblue', linewidth=1.4)\n", " axes[2].axvline(true_ang, color='red', linewidth=1.4, ls='--', label=f'True {true_ang}°')\n", " axes[2].axvline(best_ang, color='orange', linewidth=1.4, ls=':', label=f'Best {best_ang}°')\n", " axes[2].set_xlabel(\"Candidate angle (°)\"); axes[2].set_ylabel(\"Correlation\")\n", " axes[2].set_title(\"Alignment score curve\"); axes[2].legend(fontsize=8); axes[2].grid(alpha=0.3)\n", " aligned = nd_rotate(img, -best_ang, reshape=False)\n", " show(aligned, axes[3], title=f\"Image rotated back {best_ang}°\")\n", " with out_ref:\n", " clear_output(wait=True); display(fig2img(fig))\n", "\n", "for w in [ref_dd, sigma_sl1]:\n", " w.observe(_update_ref, names='value')\n", "display(VBox([ref_dd, sigma_sl1, out_ref]))\n", "_update_ref()" ] }, { "cell_type": "markdown", "id": "9d8b3eef", "metadata": {}, "source": [ "---\n", "\n", "## Part 2 – Alignment algorithm\n", "\n", "For a set of $N$ images $\\{X_i\\}$, we align each to the reference by searching over discrete candidate angles $\\Delta\\theta$ apart. For each image we rotate the reference to all candidate angles and pick the best-matching one:\n", "\n", "$$\n", "\\hat\\theta_i = \\arg\\max_{\\theta_j} \\bigl(X_i \\cdot R^{\\theta_j} A_\\text{ref}\\bigr)\n", "$$\n", "\n", "After estimating all angles, the scatter plot of true vs estimated angles reveals how well alignment worked. A tight diagonal indicates good alignment; spread indicates confusion.\n", "\n", "```{admonition} Task\n", ":class: note\n", "**Part 2, step 1:** Implement the alignment loop in the code cell below. For each image, compute the correlation against every rotated reference and record the best angle. Run the cell to see the scatter plot of true vs estimated angles.\n", "```" ] }, { "cell_type": "code", "execution_count": null, "id": "017ef043", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "# Parameters\n", "N_images = 200\n", "sigma_ex = 2.0\n", "dangle = 5 # candidate angle step in degrees\n", "\n", "rng_ex = np.random.default_rng(42)\n", "true_angles_ex, images_ex = simulate_images(Q, sigma_ex, N_images, seed=42)\n", "reference_ex = low_pass_filter(Q, sigma=5)\n", "cand_angles_ex, ref_stack_ex = rotate_reference(reference_ex, delta_angle=dangle)\n", "\n", "# Alignment (vectorised)\n", "estimated_angles_ex = align_images(images_ex, ref_stack_ex, cand_angles_ex)\n", "\n", "# Scatter: true vs estimated\n", "fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n", "axes[0].scatter(true_angles_ex, estimated_angles_ex, s=2, c='steelblue', alpha=0.6)\n", "axes[0].plot([0, 360], [0, 360], 'r--', linewidth=1)\n", "axes[0].set_xlabel(\"True angle (°)\"); axes[0].set_ylabel(\"Estimated angle (°)\")\n", "axes[0].set_title(f\"Angle estimation (N={N_images}, σ={sigma_ex})\")\n", "axes[0].set_xlim(0, 360); axes[0].set_ylim(0, 360)\n", "\n", "err = np.abs(((estimated_angles_ex - true_angles_ex + 180) % 360) - 180)\n", "axes[1].hist(err, bins=36, color='steelblue', edgecolor='k', linewidth=0.5)\n", "axes[1].set_xlabel(\"Angular error (°)\"); axes[1].set_ylabel(\"Count\")\n", "axes[1].set_title(f\"Median error: {np.median(err):.1f}°\")\n", "\n", "plt.tight_layout()\n", "display(fig2img(fig))" ] }, { "cell_type": "markdown", "id": "0df885b3", "metadata": {}, "source": [ "---\n", "\n", "## Part 3 – Reconstruction\n", "\n", "Given estimated angles, rotating each image back and averaging gives the reconstruction:\n", "\n", "$$\n", "A \\leftarrow \\frac{1}{N} \\sum_{i=1}^N (R^{\\hat\\theta_i})^{-1} X_i\n", "$$\n", "\n", "```{admonition} Task\n", ":class: note\n", "**Part 3:** Complete the reconstruction: for each image, rotate it back by $-\\hat\\theta_i$ and sum all rotated images. Divide by $N$ to get the average.\n", "```" ] }, { "cell_type": "code", "execution_count": null, "id": "17c16307", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "recon_ex = reconstruct(images_ex, estimated_angles_ex)\n", "\n", "fig, axes = plt.subplots(1, 3, figsize=(10, 3.5))\n", "show(Q, axes[0], title=\"True structure A\")\n", "show(images_ex[0], axes[1],\n", " vmin=images_ex[0].min(), vmax=images_ex[0].max(),\n", " title=f\"Single image (σ={sigma_ex})\")\n", "show(recon_ex, axes[2], title=f\"Reconstruction (N={N_images})\")\n", "plt.tight_layout()\n", "display(fig2img(fig))" ] }, { "cell_type": "markdown", "id": "889c6813", "metadata": {}, "source": [ "---\n", "\n", "## Interactive reconstruction\n", "\n", "The widget below runs the complete pipeline — from noisy images to reconstruction — with adjustable parameters. It also shows the scatter of estimated vs true angles and the correlation between reconstruction and true image as a quality metric." ] }, { "cell_type": "code", "execution_count": null, "id": "86485673", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "_sl3 = Layout(width=\"440px\")\n", "n_sl_r = IntSlider( value=100, min=10, max=500, step=10,\n", " description=\"# images\", style=_sty, layout=_sl3)\n", "sigma_sl_r= FloatSlider(value=2.0, min=0.0, max=10.0, step=0.5,\n", " description=\"Noise σ\", style=_sty, layout=_sl3)\n", "blur_sl_r = FloatSlider(value=5.0, min=0.0, max=15.0, step=1.0,\n", " description=\"LP blur σ\", style=_sty, layout=_sl3)\n", "da_sl_r = IntSlider( value=5, min=1, max=20, step=1,\n", " description=\"Δangle (°)\", style=_sty, layout=_sl3)\n", "\n", "out_recon = Output()\n", "def _update_recon(_=None):\n", " sigma = sigma_sl_r.value; N = n_sl_r.value\n", " blur = blur_sl_r.value; da = da_sl_r.value\n", " true_angs, imgs = simulate_images(Q, sigma, N, seed=42)\n", " ref_init = low_pass_filter(Q, sigma=blur) if blur > 0 else Q.copy()\n", " c_angs, r_stack = rotate_reference(ref_init, delta_angle=da)\n", " est_angs = align_images(imgs, r_stack, c_angs)\n", " recon = reconstruct(imgs, est_angs)\n", " # quality: correlation with true image\n", " qc = float((recon * Q).mean() /\n", " (np.sqrt((recon**2).mean() * (Q**2).mean()) + 1e-10))\n", " err = np.abs(((est_angs - true_angs + 180) % 360) - 180)\n", "\n", " fig, axes = plt.subplots(1, 4, figsize=(16, 3.8))\n", " show(ref_init, axes[0], title=\"Initial reference\")\n", " show(recon, axes[1], title=f\"Reconstruction (CC={qc:.2f})\")\n", " axes[2].scatter(true_angs, est_angs, s=1, c='steelblue', alpha=0.5)\n", " axes[2].plot([0,360],[0,360],'r--',lw=1)\n", " axes[2].set_xlabel(\"True angle (°)\"); axes[2].set_ylabel(\"Estimated (°)\")\n", " axes[2].set_title(f\"Angle accuracy (med. err={np.median(err):.0f}°)\")\n", " axes[3].hist(err, bins=36, color='steelblue', ec='k', lw=0.4)\n", " axes[3].set_xlabel(\"Angle error (°)\"); axes[3].set_ylabel(\"Count\")\n", " axes[3].set_title(\"Error histogram\")\n", " plt.tight_layout()\n", " with out_recon:\n", " clear_output(wait=True); display(fig2img(fig))\n", "\n", "for w in [n_sl_r, sigma_sl_r, blur_sl_r, da_sl_r]:\n", " w.observe(_update_recon, names='value')\n", "display(VBox([n_sl_r, sigma_sl_r, blur_sl_r, da_sl_r, out_recon]))\n", "_update_recon()" ] }, { "cell_type": "markdown", "id": "91da190e", "metadata": {}, "source": [ "---\n", "\n", "## Bonus 1 – Adding translation\n", "\n", "So far we assumed particles are perfectly centred. In reality, particle positions within the extracted box vary. The model extends to:\n", "\n", "$$\n", "X_i = T^{t_i}(R^{\\theta_i} A) + \\sigma G_i\n", "$$\n", "\n", "where $T^{t_i}$ is a translation by vector $\\mathbf{t}_i = (dx_i, dy_i)$. We can handle integer-pixel translations by rolling the image array. The search grid is now $N_\\theta \\times N_{dx} \\times N_{dy}$, so runtime increases by a factor of $(2 t_\\text{max}+1)^2$." ] }, { "cell_type": "code", "execution_count": null, "id": "9d7952ba", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "def translate_image(im, dx, dy):\n", " \"\"\"Translate by (dx, dy) pixels using np.roll.\"\"\"\n", " return np.roll(np.roll(im, dy, axis=0), dx, axis=1)\n", "\n", "def simulate_images_with_translation(image, sigma, N, max_t=10, seed=42):\n", " \"\"\"Images with random rotation AND translation.\"\"\"\n", " rng = np.random.default_rng(seed)\n", " angles = rng.uniform(0, 360, N)\n", " dxs = rng.integers(-max_t, max_t+1, N)\n", " dys = rng.integers(-max_t, max_t+1, N)\n", " imgs = []\n", " for ang, dx, dy in zip(angles, dxs, dys):\n", " rot = nd_rotate(image, ang, reshape=False)\n", " imgs.append(translate_image(rot, dx, dy) + sigma * rng.standard_normal(image.shape))\n", " return np.array(imgs), angles, dxs, dys\n", "\n", "# Demo: showcase translated images\n", "imgs_t, ang_t, dx_t, dy_t = simulate_images_with_translation(Q, sigma=1.0, N=6, max_t=15, seed=7)\n", "fig, axes = plt.subplots(1, 7, figsize=(16, 2.8))\n", "show(Q, axes[0], title=\"True A\")\n", "for i in range(6):\n", " axes[i+1].imshow(imgs_t[i], cmap='gray', origin='lower')\n", " axes[i+1].set_title(f\"θ={ang_t[i]:.0f}°\\ndx={dx_t[i]},dy={dy_t[i]}\", fontsize=7)\n", " axes[i+1].axis('off')\n", "plt.suptitle(\"Simulated images with random rotation AND translation\", fontsize=9)\n", "plt.tight_layout()\n", "display(fig2img(fig))" ] }, { "cell_type": "code", "execution_count": null, "id": "f8fe2c4f", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "def build_rotation_translation_stack(reference, delta_angle=10, max_t=8, t_step=4):\n", " \"\"\"\n", " Build a reference stack covering all combinations of rotation and translation.\n", " Returns (stack, angle_list, dx_list, dy_list).\n", " \"\"\"\n", " angles = np.arange(0, 360, delta_angle)\n", " ts = np.arange(-max_t, max_t+1, t_step)\n", " stack, ang_list, dx_list, dy_list = [], [], [], []\n", " for a in angles:\n", " rot = nd_rotate(reference, a, reshape=False)\n", " for dx in ts:\n", " for dy in ts:\n", " stack.append(translate_image(rot, dx, dy))\n", " ang_list.append(a); dx_list.append(dx); dy_list.append(dy)\n", " return np.array(stack), np.array(ang_list), np.array(dx_list), np.array(dy_list)\n", "\n", "def align_with_translation(images, ref_stack, ang_list, dx_list, dy_list):\n", " \"\"\"Vectorised correlation alignment over rotation × translation grid.\"\"\"\n", " n, h, w = images.shape\n", " R = ref_stack.shape[0]\n", " A = images.reshape(n, h*w).astype(float)\n", " B = ref_stack.reshape(R, h*w).astype(float)\n", " cc = (A @ B.T) / (h * w)\n", " best = cc.argmax(1)\n", " return ang_list[best], dx_list[best], dy_list[best]\n", "\n", "def reconstruct_from_rotation_translation(images, est_angles, est_dx, est_dy):\n", " recon = np.zeros_like(images[0])\n", " for img, ang, dx, dy in zip(images, est_angles, est_dx, est_dy):\n", " undone = translate_image(img, -dx, -dy)\n", " recon += nd_rotate(undone, -ang, reshape=False)\n", " return recon / len(images)\n", "\n", "# Run\n", "N_t = 80; sigma_t = 1.0; max_t = 10\n", "imgs_t, true_ang_t, true_dx, true_dy = simulate_images_with_translation(\n", " Q, sigma_t, N_t, max_t=max_t, seed=42)\n", "ref_t = low_pass_filter(Q, sigma=5)\n", "stack_t, ang_l, dx_l, dy_l = build_rotation_translation_stack(ref_t, delta_angle=10, max_t=max_t, t_step=4)\n", "est_ang_t, est_dx_t, est_dy_t = align_with_translation(imgs_t, stack_t, ang_l, dx_l, dy_l)\n", "recon_t = reconstruct_from_rotation_translation(imgs_t, est_ang_t, est_dx_t, est_dy_t)\n", "\n", "fig, axes = plt.subplots(1, 3, figsize=(10, 3.5))\n", "show(Q, axes[0], title=\"True A\")\n", "show(imgs_t[0], axes[1], vmin=imgs_t[0].min(), vmax=imgs_t[0].max(),\n", " title=f\"Sample image (σ={sigma_t})\")\n", "show(recon_t, axes[2], title=f\"Reconstruction with translation\\n(N={N_t})\")\n", "plt.tight_layout()\n", "display(fig2img(fig))" ] }, { "cell_type": "markdown", "id": "9d7b0f08", "metadata": {}, "source": [ "---\n", "\n", "## Bonus 2 – Bayesian maximum-likelihood alignment\n", "\n", "Hard assignment to the single best-fitting angle ignores alignment uncertainty. The **maximum-likelihood** approach instead computes a probability for each angle:\n", "\n", "$$\n", "P(X_i \\mid \\theta_j, A) \\propto \\exp\\!\\left[-\\frac{\\|X_i - R^{\\theta_j}A\\|^2}{2\\sigma^2}\\right]\n", "$$\n", "\n", "The **soft** reconstruction then uses all angles weighted by their posterior:\n", "\n", "$$\n", "w_{ij} = \\frac{P(X_i \\mid \\theta_j, A)}{\\sum_{j'} P(X_i \\mid \\theta_{j'}, A)}, \\qquad\n", "A \\leftarrow \\sum_i \\sum_j w_{ij} (R^{\\theta_j})^{-1} X_i\n", "$$\n", "\n", "At high SNR the weights collapse onto the single best angle (recovering hard assignment); at low SNR they spread across many angles, effectively averaging out alignment noise." ] }, { "cell_type": "code", "execution_count": null, "id": "750dd71e", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "def log_likelihood(image, reference, sigma):\n", " \"\"\"Log-likelihood of image given reference under Gaussian noise.\"\"\"\n", " return -0.5 * np.sum((image - reference)**2) / sigma**2\n", "\n", "def ml_reconstruct(images, ref_stack, cand_angles, sigma=2.0):\n", " \"\"\"Soft ML reconstruction: weighted sum over all candidate angles.\"\"\"\n", " recon = np.zeros_like(images[0])\n", " for img in images:\n", " log_liks = np.array([log_likelihood(img, ref, sigma) for ref in ref_stack])\n", " log_liks -= log_liks.max() # numerical stability\n", " weights = np.exp(log_liks)\n", " weights /= weights.sum()\n", " for w, ang, ref in zip(weights, cand_angles, ref_stack):\n", " recon += w * nd_rotate(img, -ang, reshape=False)\n", " return recon / len(images)\n", "\n", "# Compare hard vs soft at different noise levels\n", "rng_ml = np.random.default_rng(42)\n", "N_ml = 60\n", "ref_ml = low_pass_filter(Q, sigma=5)\n", "c_ml, s_ml = rotate_reference(ref_ml, delta_angle=10)\n", "\n", "fig, axes = plt.subplots(2, 5, figsize=(16, 7))\n", "for col, sig in enumerate([0.5, 1.5, 3.0, 5.0, 8.0]):\n", " true_angs_ml, imgs_ml = simulate_images(Q, sig, N_ml, seed=42)\n", " # Hard assignment\n", " est_hard = align_images(imgs_ml, s_ml, c_ml)\n", " recon_hard = reconstruct(imgs_ml, est_hard)\n", " # Soft (ML)\n", " recon_soft = ml_reconstruct(imgs_ml, s_ml, c_ml, sigma=sig)\n", " axes[0, col].imshow(recon_hard, cmap='gray', origin='lower')\n", " axes[0, col].axis('off'); axes[0, col].set_title(f\"Hard (σ={sig})\", fontsize=8)\n", " axes[1, col].imshow(recon_soft, cmap='gray', origin='lower')\n", " axes[1, col].axis('off'); axes[1, col].set_title(f\"ML soft (σ={sig})\", fontsize=8)\n", "\n", "axes[0, 0].set_ylabel(\"Hard assignment\", fontsize=9)\n", "axes[1, 0].set_ylabel(\"ML soft assignment\", fontsize=9)\n", "plt.suptitle(\"Hard vs ML reconstruction at increasing noise\", fontsize=10)\n", "plt.tight_layout()\n", "display(fig2img(fig))" ] }, { "cell_type": "markdown", "id": "9039ccb0", "metadata": {}, "source": [ "---\n", "\n", "## Bonus 3 – Iterative refinement\n", "\n", "The previous examples used the initial reference for all alignment. A better strategy is to use the reconstruction from one round as the reference for the next:\n", "\n", "1. Start from initial reference $A^{(0)}$ (e.g. blurred Q or a circle)\n", "2. Align images to $A^{(k)}$ → estimate $\\hat\\theta_i^{(k)}$\n", "3. Reconstruct $A^{(k+1)}$ from aligned images\n", "4. Repeat from step 2\n", "\n", "The quality of the reconstruction typically improves each iteration. The widget below lets you choose the starting reference and the number of iterations." ] }, { "cell_type": "code", "execution_count": null, "id": "0e5dad44", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "iter_sl = IntSlider( value=3, min=1, max=8, step=1,\n", " description=\"Iterations\", style=_sty, layout=_sl3)\n", "blur_sl2 = FloatSlider(value=5.0, min=0.0, max=15.0, step=1.0,\n", " description=\"Init. blur σ\", style=_sty, layout=_sl3)\n", "sigma_sl2= FloatSlider(value=2.0, min=0.0, max=10.0, step=0.5,\n", " description=\"Noise σ\", style=_sty, layout=_sl3)\n", "ref_dd2 = Dropdown(options=['Q (oracle)','Blurred Q','Circle','P (wrong)','O (wrong)'],\n", " value='Blurred Q', description=\"Start ref:\", style=_sty)\n", "\n", "out_iter = Output()\n", "def _update_iter(_=None):\n", " sig = sigma_sl2.value; iters = iter_sl.value; blur = blur_sl2.value\n", " name = ref_dd2.value\n", " true_angs_i, imgs_i = simulate_images(Q, sig, 100, seed=42)\n", " ref0 = (_REFS.get(name + ' (oracle)', _REFS.get(name, None))\n", " or low_pass_filter(Q, sigma=blur))\n", " if name == 'Blurred Q':\n", " ref0 = low_pass_filter(Q, sigma=blur)\n", " elif name == 'Q (oracle)':\n", " ref0 = Q.copy()\n", " elif name == 'Circle':\n", " ref0 = make_circle(64)\n", " elif name == 'P (wrong)':\n", " ref0 = make_letter('P', 64)\n", " elif name == 'O (wrong)':\n", " ref0 = make_letter('O', 64)\n", " ref = ref0.copy()\n", " history = [ref0.copy()]\n", " ccs = []\n", " for _ in range(iters):\n", " c_angs, r_stack = rotate_reference(ref, delta_angle=5)\n", " est = align_images(imgs_i, r_stack, c_angs)\n", " ref = reconstruct(imgs_i, est)\n", " history.append(ref.copy())\n", " ccs.append(float((ref * Q).mean() /\n", " (np.sqrt((ref**2).mean() * (Q**2).mean()) + 1e-10)))\n", "\n", " ncols = min(len(history), 6)\n", " fig, axes = plt.subplots(1, ncols + 1, figsize=(3*(ncols+1), 3.5))\n", " show(Q, axes[0], title=\"True A\")\n", " for col, (im, lab) in enumerate(\n", " zip(history[:ncols], ['Init'] + [f'Iter {i+1}' for i in range(ncols-1)]), 1):\n", " axes[col].imshow(im, cmap='gray', origin='lower')\n", " axes[col].axis('off')\n", " axes[col].set_title(lab if col == 1 else f\"{lab}\\nCC={ccs[col-2]:.2f}\", fontsize=8)\n", " plt.tight_layout()\n", " with out_iter:\n", " clear_output(wait=True); display(fig2img(fig))\n", "\n", "for w in [iter_sl, blur_sl2, sigma_sl2, ref_dd2]:\n", " w.observe(_update_iter, names='value')\n", "display(VBox([ref_dd2, blur_sl2, sigma_sl2, iter_sl, out_iter]))\n", "_update_iter()" ] }, { "cell_type": "markdown", "id": "58fe7377", "metadata": {}, "source": [ "---\n", "\n", "## Bonus 4 – 2D Classification\n", "\n", "Real datasets contain multiple particle types (different proteins, conformations, or viewing directions). **2D classification** assigns each image to one of $K$ classes simultaneously with the alignment:\n", "\n", "1. Maintain $K$ references $\\{A_1, \\ldots, A_K\\}$\n", "2. For each image, find the best (class, angle) pair by correlation\n", "3. Reconstruct each class from the images assigned to it\n", "4. Update references and iterate\n", "\n", "The interactive below mixes $K$ letters (Q, P, R) and attempts to recover them. The classification accuracy depends on the noise level and the number of images per class." ] }, { "cell_type": "code", "execution_count": null, "id": "c825b75e", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "K_sl = IntSlider( value=2, min=2, max=3, step=1,\n", " description=\"# classes K\", style=_sty, layout=_sl3)\n", "nper_sl = IntSlider( value=40, min=10, max=100, step=5,\n", " description=\"Images/class\", style=_sty, layout=_sl3)\n", "sigma_sl3= FloatSlider(value=1.5, min=0.0, max=6.0, step=0.5,\n", " description=\"Noise σ\", style=_sty, layout=_sl3)\n", "\n", "_CLS_LETTERS = ['Q', 'P', 'R']\n", "_CLS_IMGS = {lt: make_letter(lt, 64) for lt in _CLS_LETTERS}\n", "_CLS_CAND = np.arange(0, 360, 10)\n", "\n", "def _precompute_cls_stack(letter):\n", " return np.array([nd_rotate(_CLS_IMGS[letter], a, reshape=False) for a in _CLS_CAND])\n", "\n", "_CLS_STACKS = {lt: _precompute_cls_stack(lt) for lt in _CLS_LETTERS}\n", "\n", "out_cls4 = Output()\n", "def _update_cls4(_=None):\n", " K = K_sl.value; n_per = nper_sl.value; sig = sigma_sl3.value\n", " letters = _CLS_LETTERS[:K]\n", " rng = np.random.default_rng(42)\n", "\n", " # Generate mixed dataset\n", " all_imgs, true_labels = [], []\n", " for k, lt in enumerate(letters):\n", " idx = rng.integers(0, len(_CLS_CAND), n_per)\n", " noisy = _CLS_STACKS[lt][idx] + rng.standard_normal((n_per, 64, 64)) * sig\n", " all_imgs.append(noisy); true_labels.extend([k]*n_per)\n", " all_imgs = np.concatenate(all_imgs)\n", " true_labels= np.array(true_labels)\n", " shuf = rng.permutation(len(all_imgs))\n", " all_imgs = all_imgs[shuf]; true_labels = true_labels[shuf]\n", "\n", " # Build combined reference stack\n", " ref_stack_all = np.concatenate([_CLS_STACKS[lt] for lt in letters])\n", " n_cand_per = len(_CLS_CAND)\n", " n_all = all_imgs.shape[0]\n", " A = all_imgs.reshape(n_all, 64*64).astype(float)\n", " B = ref_stack_all.reshape(len(ref_stack_all), 64*64).astype(float)\n", " cc = (A @ B.T) / (64*64)\n", " best_global = cc.argmax(1)\n", " pred_class = best_global // n_cand_per\n", " best_ang_idx = best_global % n_cand_per\n", "\n", " acc = (pred_class == true_labels).mean() * 100\n", " class_recons = []\n", " for k in range(K):\n", " mask = (pred_class == k)\n", " recon_k = np.zeros((64, 64))\n", " count = 0\n", " for i in np.where(mask)[0]:\n", " recon_k += nd_rotate(all_imgs[i], -_CLS_CAND[best_ang_idx[i]], reshape=False)\n", " count += 1\n", " class_recons.append(recon_k / max(count, 1))\n", "\n", " naive_avg = all_imgs.mean(0)\n", " fig, axes = plt.subplots(1, K+3, figsize=(3.5*(K+3), 3.5))\n", " show(all_imgs[0], axes[0], vmin=all_imgs[0].min(), vmax=all_imgs[0].max(),\n", " title=\"Example image\\n(mixed)\")\n", " show(naive_avg, axes[1], title=\"Naive average\\n(blurry)\")\n", " for k, (recon_k, lt) in enumerate(zip(class_recons, letters)):\n", " n_in_cls = (pred_class==k).sum()\n", " axes[k+2].imshow(recon_k, cmap='gray', origin='lower')\n", " axes[k+2].axis('off')\n", " axes[k+2].set_title(f\"Class {k+1}: {lt}\\n({n_in_cls} imgs assigned)\", fontsize=8)\n", " for k, lt in enumerate(letters):\n", " axes[K+2].imshow(_CLS_IMGS[lt], cmap='gray', origin='lower',\n", " extent=[k, k+0.9, 0, 0.9], aspect='auto')\n", " axes[K+2].set_xlim(-0.1, K); axes[K+2].set_ylim(-0.05, 1)\n", " axes[K+2].axis('off'); axes[K+2].set_title(f\"True classes\\nAcc={acc:.0f}%\", fontsize=8)\n", " plt.suptitle(f\"2D Classification: K={K} classes, σ={sig}, {n_per} imgs/class\", fontsize=9)\n", " plt.tight_layout()\n", " with out_cls4:\n", " clear_output(wait=True); display(fig2img(fig))\n", "\n", "for w in [K_sl, nper_sl, sigma_sl3]:\n", " w.observe(_update_cls4, names='value')\n", "display(VBox([K_sl, nper_sl, sigma_sl3, out_cls4]))\n", "_update_cls4()" ] }, { "cell_type": "markdown", "id": "a2c0bc9b", "metadata": {}, "source": [ "---\n", "\n", "## Summary\n", "\n", "| Topic | Key concept |\n", "|-------|-------------|\n", "| Observation model | $X_i = R^{\\theta_i}A + \\sigma G_i$ |\n", "| Reference selection | Circular, blurred prior, or oracle; avoid reference bias |\n", "| Alignment | Maximise correlation $\\langle X_i, R^{\\theta_j}A\\rangle$ over discrete candidates |\n", "| Reconstruction | Average back-rotated images |\n", "| Model bias | Wrong reference → biased reconstruction; correct with iterative refinement |\n", "| ML alignment | Soft weights via $\\exp(-\\|X_i - R^\\theta A\\|^2/2\\sigma^2)$; better at low SNR |\n", "| Iterative refinement | Use reconstruction as new reference; converges toward true structure |\n", "| Classification | $K$-class joint alignment+assignment; separates heterogeneous datasets |" ] } ], "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, 40, 138, 144, 152, 166, 193, 209, 259, 278, 308, 325, 337, 345, 389, 403, 434, 485, 506, 548, 563, 618, 633, 714 ] }, "nbformat": 4, "nbformat_minor": 5 }