Source code for quscope.quantum_ctem.ctf_calculator

"""
Contrast Transfer Function (CTF) Calculator and Visualization

This module provides comprehensive CTF analysis for Conventional TEM including:
- 1D radial CTF plots
- 2D CTF visualization in momentum space
- Multi-voltage comparison
- Individual aberration contributions
- Scherzer defocus calculation
- Resolution limits

Designed for publication-quality figures suitable for both quantum computing
and electron microscopy audiences.

References:
    - Kirkland, E. J. (2010). Advanced Computing in Electron Microscopy.
    - Krivanek, O. L., et al. (2008). Ultramicroscopy 108(3), 179-195.
    - Spence, J. C. H. (2013). High-Resolution Electron Microscopy (4th ed.).

Author: QuScope Development Team
Date: January 2025
"""

from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

import matplotlib.pyplot as plt
import numpy as np
import scipy.constants as const
from matplotlib import cm

try:
    from .hamiltonian import (
        HamiltonianParameters,
        LensHamiltonian,
        relativistic_wavelength,
    )
except ImportError:
    from hamiltonian import (
        HamiltonianParameters,
        LensHamiltonian,
        relativistic_wavelength,
    )


[docs] @dataclass class CTFParameters: """ Parameters for CTF calculation. Attributes: voltage: Acceleration voltage (V) defocus: Defocus C₁ (Angstrom) cs: Spherical aberration C₃ (mm) c5: 5th order spherical aberration (mm) aperture: Objective aperture semi-angle (mrad) aberrations: Complete aberration dictionary """ voltage: float defocus: float cs: float = 0.0 c5: float = 0.0 aperture: float = 10.0 aberrations: Dict[str, float] = None
[docs] def __post_init__(self): """Initialize aberrations dictionary if not provided.""" if self.aberrations is None: self.aberrations = {"defocus": self.defocus, "c3": self.cs, "c5": self.c5}
[docs] class CTFCalculator: """ Calculate Contrast Transfer Function for TEM. The CTF describes how spatial frequencies are transferred from the sample to the image, including effects of defocus and aberrations. For phase contrast imaging (weak phase object): CTF(k) = A(k) · sin(χ(k)) - B(k) · cos(χ(k)) where χ(k) is the wave aberration function and A(k), B(k) are envelope functions describing partial coherence and damping effects. For simplicity, we often use: CTF(k) = sin(χ(k)) """ def __init__( self, params: CTFParameters, max_k: float = 10.0, n_points: int = 1000 ): """ Initialize CTF calculator. Args: params: CTF parameters max_k: Maximum spatial frequency (1/Angstrom) n_points: Number of points for radial CTF """ self.params = params self.wavelength = relativistic_wavelength(params.voltage) self.max_k = max_k self.n_points = n_points # Generate radial k-space self.k_radial = np.linspace(0, max_k, n_points)
[docs] def chi(self, k: np.ndarray, theta: np.ndarray = None) -> np.ndarray: """ Calculate wave aberration function χ(k). For axially symmetric aberrations (no astigmatism/coma): χ(k) = π·λ·k²·C₁ + π/2·(λk)⁴·C₃ + π/3·(λk)⁶·C₅ Args: k: Spatial frequency (1/Angstrom), scalar or array theta: Azimuthal angle (radians), for non-axial aberrations Returns: χ(k) in radians """ lam = self.wavelength lam_k = lam * k # Defocus term chi = np.pi * lam * k**2 * self.params.defocus # Spherical aberration C₃ if self.params.cs != 0: cs_angstrom = self.params.cs * 1e7 # mm to Angstrom chi += 0.5 * np.pi * cs_angstrom * lam_k**4 # 5th order spherical aberration C₅ if self.params.c5 != 0: c5_angstrom = self.params.c5 * 1e7 chi += (np.pi / 3) * c5_angstrom * lam_k**6 return chi
[docs] def ctf(self, k: np.ndarray, theta: np.ndarray = None) -> np.ndarray: """ Calculate CTF = sin(χ(k)). Args: k: Spatial frequency theta: Azimuthal angle (optional) Returns: CTF value """ return np.sin(self.chi(k, theta))
[docs] def calculate_scherzer_defocus(self) -> float: """ Calculate Scherzer defocus for optimal phase contrast. Scherzer defocus balances defocus and spherical aberration to maximize contrast transfer at medium spatial frequencies. For C₃-dominated systems: Δf_Scherzer = -1.2 · √(C₃·λ) Returns: Scherzer defocus (Angstrom), negative = overfocus """ if self.params.cs == 0: return 0.0 cs_angstrom = self.params.cs * 1e7 return -1.2 * np.sqrt(cs_angstrom * self.wavelength)
[docs] def calculate_point_resolution(self) -> float: """ Calculate point resolution (Scherzer limit). d_Scherzer = 0.66 · (C₃·λ³)^(1/4) This is the finest detail that can be resolved with optimal defocus in a Cs-uncorrected microscope. Returns: Point resolution (Angstrom) """ if self.params.cs == 0: # For Cs-corrected microscopes, limited by higher orders if self.params.c5 != 0: c5_angstrom = self.params.c5 * 1e7 return 0.8 * (c5_angstrom * self.wavelength**5) ** (1 / 6) return 0.5 * self.wavelength # Optimistic estimate cs_angstrom = self.params.cs * 1e7 return 0.66 * (cs_angstrom * self.wavelength**3) ** 0.25
[docs] def find_first_zero(self) -> float: """ Find first zero of CTF (point resolution crossover). Returns: k value of first CTF zero (1/Angstrom) """ ctf_values = self.ctf(self.k_radial) # Find sign changes sign_changes = np.diff(np.sign(ctf_values)) zeros = np.where(sign_changes != 0)[0] if len(zeros) > 0: return self.k_radial[zeros[0]] return self.max_k
[docs] def calculate_information_limit(self) -> float: """ Calculate information limit (highest usable spatial frequency). This is limited by damping envelopes due to: - Chromatic aberration - Partial spatial coherence - Partial temporal coherence - Instabilities Simplified estimate based on aberrations: k_max ≈ 1 / (acceptable phase error) Returns: Information limit (1/Angstrom) """ # Find where |χ| exceeds π (destructive interference) chi_values = np.abs(self.chi(self.k_radial)) exceeds_pi = np.where(chi_values > np.pi)[0] if len(exceeds_pi) > 0: return self.k_radial[exceeds_pi[0]] return self.max_k
[docs] class CTFVisualizer: """ Generate publication-quality CTF visualizations. """ def __init__(self, figsize: Tuple[int, int] = (12, 10)): """ Initialize visualizer. Args: figsize: Figure size (width, height) in inches """ self.figsize = figsize self.colors = plt.cm.tab10(np.linspace(0, 1, 10))
[docs] def plot_1d_ctf( self, calculators: Dict[str, CTFCalculator], ax: Optional[plt.Axes] = None, show_zeros: bool = True, show_envelope: bool = False, ) -> plt.Figure: """ Plot 1D radial CTF for multiple conditions. Args: calculators: Dictionary of {label: CTFCalculator} ax: Matplotlib axes (creates new if None) show_zeros: Mark CTF zeros show_envelope: Show envelope function Returns: Figure object """ if ax is None: fig, ax = plt.subplots(figsize=(10, 6)) else: fig = ax.get_figure() for idx, (label, calc) in enumerate(calculators.items()): k = calc.k_radial ctf = calc.ctf(k) # Plot CTF ax.plot(k, ctf, label=label, linewidth=2, color=self.colors[idx]) # Mark first zero (point resolution) if show_zeros: k_zero = calc.find_first_zero() ax.axvline( k_zero, color=self.colors[idx], linestyle="--", alpha=0.5, linewidth=1, ) ax.text(k_zero, -0.9, f"{1/k_zero:.2f} Å", fontsize=9, ha="center") # Formatting ax.set_xlabel("Spatial Frequency k (1/Å)", fontsize=14, fontweight="bold") ax.set_ylabel("CTF Amplitude", fontsize=14, fontweight="bold") ax.set_title("Contrast Transfer Function", fontsize=16, fontweight="bold") ax.grid(True, alpha=0.3) ax.legend(fontsize=11, framealpha=0.9) ax.set_xlim(0, max([c.max_k for c in calculators.values()])) ax.set_ylim(-1.1, 1.1) # Add resolution markers ax.axhline(0, color="black", linewidth=0.8, alpha=0.5) ax.fill_between([0, ax.get_xlim()[1]], -1, 1, alpha=0.05, color="gray") plt.tight_layout() return fig
[docs] def plot_2d_ctf( self, calc: CTFCalculator, n_points: int = 512, cmap: str = "RdBu_r" ) -> plt.Figure: """ Plot 2D CTF in momentum space. Args: calc: CTF calculator n_points: Number of points in each dimension cmap: Colormap name Returns: Figure object with 2D CTF visualization """ # Generate 2D k-space grid k_max = calc.max_k kx = np.linspace(-k_max, k_max, n_points) ky = np.linspace(-k_max, k_max, n_points) KX, KY = np.meshgrid(kx, ky) # Calculate k and theta K = np.sqrt(KX**2 + KY**2) Theta = np.arctan2(KY, KX) # Calculate 2D CTF CTF_2D = calc.ctf(K, Theta) # Apply aperture aperture_k = calc.params.aperture * 1e-3 / calc.wavelength # mrad to 1/Å aperture_mask = K <= aperture_k CTF_2D = CTF_2D * aperture_mask # Create figure fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6)) # Plot 2D CTF extent = [-k_max, k_max, -k_max, k_max] im1 = ax1.imshow( CTF_2D, extent=extent, cmap=cmap, vmin=-1, vmax=1, origin="lower", interpolation="bilinear", ) ax1.set_xlabel("kₓ (1/Å)", fontsize=13) ax1.set_ylabel("kᵧ (1/Å)", fontsize=13) ax1.set_title("2D Contrast Transfer Function", fontsize=14, fontweight="bold") # Add aperture circle circle = plt.Circle( (0, 0), aperture_k, fill=False, edgecolor="white", linewidth=2, linestyle="--", ) ax1.add_patch(circle) # Add resolution circles d_scherzer = calc.calculate_point_resolution() k_scherzer = 1 / d_scherzer circle_res = plt.Circle( (0, 0), k_scherzer, fill=False, edgecolor="yellow", linewidth=1.5, linestyle=":", ) ax1.add_patch(circle_res) ax1.text( 0, -k_scherzer * 1.15, f"Scherzer: {d_scherzer:.2f} Å", ha="center", va="top", color="yellow", fontsize=10, bbox=dict(boxstyle="round", facecolor="black", alpha=0.7), ) plt.colorbar(im1, ax=ax1, label="CTF Amplitude") # Plot radial average k_radial = calc.k_radial ctf_radial = calc.ctf(k_radial) ax2.plot(k_radial, ctf_radial, "b-", linewidth=2.5, label="CTF") ax2.axhline(0, color="black", linewidth=0.8, alpha=0.5) ax2.axvline( aperture_k, color="red", linestyle="--", linewidth=2, label=f"Aperture ({calc.params.aperture} mrad)", ) ax2.axvline( k_scherzer, color="orange", linestyle=":", linewidth=2, label=f"Point Resolution ({d_scherzer:.2f} Å)", ) ax2.set_xlabel("Spatial Frequency k (1/Å)", fontsize=13) ax2.set_ylabel("CTF Amplitude", fontsize=13) ax2.set_title("Radial Average", fontsize=14, fontweight="bold") ax2.grid(True, alpha=0.3) ax2.legend(fontsize=10) ax2.set_xlim(0, k_max) ax2.set_ylim(-1.1, 1.1) plt.tight_layout() return fig
[docs] def plot_multi_voltage_comparison( self, voltages: List[float], cs: float = 1.3, defocus: float = None ) -> plt.Figure: """ Compare CTF for different acceleration voltages. Args: voltages: List of voltages (V), e.g., [80e3, 120e3, 200e3, 300e3] cs: Spherical aberration (mm) defocus: Defocus (Angstrom), uses Scherzer if None Returns: Figure with multi-voltage comparison """ fig, axes = plt.subplots(2, 2, figsize=(14, 12)) axes = axes.flatten() for idx, voltage in enumerate(voltages): ax = axes[idx] # Calculate Scherzer defocus for this voltage if not specified if defocus is None: wavelength = relativistic_wavelength(voltage) cs_angstrom = cs * 1e7 df = -1.2 * np.sqrt(cs_angstrom * wavelength) else: df = defocus # Create CTF calculator params = CTFParameters(voltage=voltage, defocus=df, cs=cs) calc = CTFCalculator(params, max_k=5.0) # Plot CTF k = calc.k_radial ctf_vals = calc.ctf(k) ax.plot(k, ctf_vals, "b-", linewidth=2.5) # Mark point resolution d_res = calc.calculate_point_resolution() k_res = 1 / d_res ax.axvline(k_res, color="red", linestyle="--", linewidth=2, alpha=0.7) # Formatting ax.set_xlabel("Spatial Frequency k (1/Å)", fontsize=12) ax.set_ylabel("CTF", fontsize=12) ax.set_title( f"{voltage/1e3:.0f} kV (λ={calc.wavelength:.4f} Å, " f"d={d_res:.2f} Å)", fontsize=13, fontweight="bold", ) ax.grid(True, alpha=0.3) ax.axhline(0, color="black", linewidth=0.8, alpha=0.5) ax.set_xlim(0, 5.0) ax.set_ylim(-1.1, 1.1) # Add text annotation ax.text( 0.98, 0.97, f"Cs = {cs} mm\nΔf = {df:.0f} Å\nd_res = {d_res:.2f} Å", transform=ax.transAxes, fontsize=10, verticalalignment="top", horizontalalignment="right", bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.8), ) plt.suptitle( "CTF Comparison: Multi-Voltage Analysis", fontsize=16, fontweight="bold", y=0.995, ) plt.tight_layout() return fig
# Example usage if __name__ == "__main__": print("=" * 70) print("CTF Calculator - Validation and Demonstration") print("=" * 70) print() # Test different voltages voltages = [80e3, 120e3, 200e3, 300e3] print("Calculating wavelengths for standard voltages:") for V in voltages: lam = relativistic_wavelength(V) print(f" {V/1e3:>3.0f} kV: λ = {lam:.5f} Å") print() # Create CTF calculator for 200 kV params = CTFParameters(voltage=200e3, defocus=-500.0, cs=1.3, c5=10.0) calc = CTFCalculator(params) print(f"200 kV TEM with Cs = 1.3 mm:") print(f" Scherzer defocus: {calc.calculate_scherzer_defocus():.1f} Å") print(f" Point resolution: {calc.calculate_point_resolution():.3f} Å") print(f" Information limit: {1/calc.calculate_information_limit():.3f} Å") print() # Create visualizations print("Generating CTF visualizations...") # 1D CTF comparison calculators = {} for V in [80e3, 200e3, 300e3]: lam = relativistic_wavelength(V) cs_ang = 1.3 * 1e7 df = -1.2 * np.sqrt(cs_ang * lam) params = CTFParameters(voltage=V, defocus=df, cs=1.3) calculators[f"{V/1e3:.0f} kV"] = CTFCalculator(params, max_k=5.0) viz = CTFVisualizer() fig1 = viz.plot_1d_ctf(calculators) plt.savefig("ctf_1d_comparison.png", dpi=300, bbox_inches="tight") print(" ✓ Saved: ctf_1d_comparison.png") # 2D CTF params_2d = CTFParameters(voltage=200e3, defocus=-500.0, cs=1.3) calc_2d = CTFCalculator(params_2d, max_k=5.0) fig2 = viz.plot_2d_ctf(calc_2d, n_points=512) plt.savefig("ctf_2d_visualization.png", dpi=300, bbox_inches="tight") print(" ✓ Saved: ctf_2d_visualization.png") # Multi-voltage comparison fig3 = viz.plot_multi_voltage_comparison([80e3, 120e3, 200e3, 300e3], cs=1.3) plt.savefig("ctf_multi_voltage.png", dpi=300, bbox_inches="tight") print(" ✓ Saved: ctf_multi_voltage.png") print() print("✓ CTF module validated and figures generated") print() print("Publication-ready figures:") print(" • ctf_1d_comparison.png - 1D radial CTF") print(" • ctf_2d_visualization.png - 2D momentum space CTF") print(" • ctf_multi_voltage.png - Multi-voltage comparison")