"""
Weak Phase Object Approximation (WPOA) CTEM Simulator
Implements classical CTEM image simulation using the weak phase object approximation
from Earl J. Kirkland's "Advanced Computing in Electron Microscopy", 2nd Edition, Chapter 5.
The WPOA assumes the transmission function is:
t(x,y) = exp(iσV(x,y))
where:
- σ = interaction parameter (depends on beam energy)
- V(x,y) = projected atomic potential (from KirklandPotential)
The simulation pipeline:
1. Calculate projected potential V(x,y) for all atoms
2. Calculate transmission function t(x,y) = exp(iσV)
3. Fourier transform to reciprocal space
4. Apply objective lens transfer function (CTF with aberrations)
5. Inverse Fourier transform back to real space
6. Calculate intensity I = |ψ|²
Reference:
Kirkland, E. J. (2010). Advanced Computing in Electron Microscopy (2nd ed.).
Springer. Chapter 5, Figures 5.11-5.12
"""
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import numpy as np
from .kirkland_potential import KirklandPotential
[docs]
class WPOASimulator:
"""
Classical CTEM simulator using Weak Phase Object Approximation.
This simulator implements the WPOA method for conventional transmission
electron microscopy (CTEM) image simulation. It handles:
- Relativistic wavelength calculation
- Interaction parameter σ calculation
- Transmission function with weak phase object approximation
- Objective lens transfer function with spherical aberration
- Coherent bright field image formation
Attributes:
image_size: Physical size of image in Angstroms
pixels: Number of pixels (assumed square image)
beam_energy: Electron beam energy in eV
wavelength: Relativistic electron wavelength in Angstroms
sigma: Interaction parameter in rad/eV
dx: Pixel size in Angstroms
X, Y: 2D coordinate grids in Angstroms
potential_calculator: KirklandPotential instance for atomic potentials
Example:
>>> from quscope.ctem import WPOASimulator
>>>
>>> # Initialize simulator for 200 keV electrons
>>> sim = WPOASimulator(
... image_size=50.0, # 50 Angstroms
... pixels=512,
... beam_energy=200e3 # 200 keV
... )
>>>
>>> # Define atom positions: (x, y, Z)
>>> atoms = [
... (-20, 0, 6), # C at -20 Å
... (-10, 0, 14), # Si at -10 Å
... (0, 0, 29), # Cu at 0 Å
... (10, 0, 79), # Au at 10 Å
... (20, 0, 92), # U at 20 Å
... ]
>>>
>>> # Simulate CTEM image
>>> results = sim.simulate_image(
... atom_positions=atoms,
... defocus=700.0, # 700 Å underfocus
... Cs=1.3e7, # 1.3 mm spherical aberration
... alpha_max=10.37 # 10.37 mrad aperture
... )
>>>
>>> # Access results
>>> intensity = results['intensity']
>>> transmission = results['transmission']
>>> potential = results['potential']
"""
# Physical constants
M0C2 = 511.0e3 # Electron rest mass energy in eV
HC = 12.2639 # Constant for wavelength calculation (from Kirkland)
def __init__(
self,
image_size: float = 50.0,
pixels: int = 512,
beam_energy: float = 200e3,
params_file: Optional[Path] = None,
):
"""
Initialize WPOA simulator.
Args:
image_size: Physical size of image in Angstroms
pixels: Number of pixels (square image assumed)
beam_energy: Electron beam energy in eV (e.g., 200e3 for 200 keV)
params_file: Optional path to Kirkland parameters file
Notes:
- Coordinate system: origin at image center
- Pixel grid: centered at pixel midpoints
- Fourier space: fftshift convention for zero-frequency at center
"""
self.image_size = image_size
self.pixels = pixels
self.beam_energy = beam_energy
# Calculate relativistic wavelength and interaction parameter
self.wavelength = self._calculate_wavelength()
self.sigma = self._calculate_sigma()
# Create coordinate grids (pixel-centered)
self.dx = self.image_size / self.pixels
x = (np.arange(self.pixels) - self.pixels / 2 + 0.5) * self.dx
y = (np.arange(self.pixels) - self.pixels / 2 + 0.5) * self.dx
self.x = x
self.y = y
self.X, self.Y = np.meshgrid(x, y, indexing="xy")
# Initialize potential calculator
self.potential_calculator = KirklandPotential(params_file=params_file)
def _calculate_wavelength(self) -> float:
"""
Calculate relativistic electron wavelength.
Uses Kirkland Equation 5.2:
λ = 12.2639 / √(V + 0.97845×10⁻⁶·V²)
where:
- V = beam energy in eV (NOT keV)
- Constant is 12.2639 (not 12.398)
Returns:
Wavelength in Angstroms
Example:
>>> sim = WPOASimulator(beam_energy=200e3)
>>> sim.wavelength
0.02508... # ~0.025 Å for 200 keV
"""
V = self.beam_energy # in eV
# Use proper constant from Kirkland
wavelength = 12.2639 / np.sqrt(V + 0.97845e-6 * V**2)
return wavelength
def _calculate_sigma(self) -> float:
"""
Calculate interaction parameter σ.
Uses Kirkland Equation 5.6:
σ = (2πmeλ/h²) · (m₀c² + eV)/(2m₀c² + eV)
In practical units (Angstroms and eV):
σ ≈ C · γ / (λ · V)
where C ≈ 60 for λ in Angstroms, V in eV.
Returns:
Interaction parameter in rad/eV
Notes:
- For 200 keV: σ ≈ 0.00729 rad/eV
- For 1000 eV potential: phase ≈ 7.29 radians
"""
V = self.beam_energy # in eV
# Relativistic factor
gamma = (self.M0C2 + V) / (2 * self.M0C2 + V)
# Interaction parameter
# Using empirically determined constant (from Kirkland figures)
# Note: This constant was calibrated with wavelength formula
V_keV = V / 1000.0
sigma = 0.00335 * gamma / (self.wavelength * V_keV)
return sigma
[docs]
def calculate_transmission_function(
self, atom_positions: List[Tuple[float, float, int]]
) -> Tuple[np.ndarray, np.ndarray]:
"""
Calculate transmission function for given atom positions.
Steps:
1. Calculate total projected potential V(x,y) from all atoms
2. Calculate phase shift: φ(x,y) = σ·V(x,y)
3. Calculate transmission: t(x,y) = exp(i·φ)
Args:
atom_positions: List of (x, y, Z) tuples where:
- x, y: atom coordinates in Angstroms
- Z: atomic number
Returns:
transmission: Complex transmission function, shape (pixels, pixels)
potential: Total projected potential in eV, shape (pixels, pixels)
Example:
>>> atoms = [(0, 0, 6), (5, 0, 14)] # C and Si
>>> t, V = sim.calculate_transmission_function(atoms)
>>> np.abs(t) # Magnitude (always 1 in WPOA)
array([[1., 1., ...], ...])
>>> np.angle(t) # Phase shift
array([[0.12, 0.15, ...], ...])
"""
# Calculate total potential from all atoms
V_total = self.potential_calculator.calculate_multiple_atoms(
self.X, self.Y, atom_positions
)
# Calculate phase: φ = σ·V
phase = self.sigma * V_total
# Transmission function: t = exp(i·φ)
transmission = np.exp(1j * phase)
return transmission, V_total
[docs]
def objective_lens_transfer_function(
self,
kx: np.ndarray,
ky: np.ndarray,
defocus: float,
Cs: float,
alpha_max: Optional[float] = None,
) -> np.ndarray:
"""
Calculate objective lens transfer function (CTF).
Implements Kirkland Equation 5.27:
H(k) = exp(-i·χ(k))
where the phase aberration is:
χ(k) = π·λ·k²·(½·Cs·λ²·k² - Δf)
Args:
kx: X spatial frequencies in 1/Angstrom, shape (M, N)
ky: Y spatial frequencies in 1/Angstrom, shape (M, N)
defocus: Defocus value Δf in Angstroms (positive = underfocus)
Cs: Spherical aberration coefficient in Angstroms
alpha_max: Optional aperture semi-angle in radians
Returns:
Transfer function H(k), complex array shape (M, N)
Notes:
- Positive defocus = underfocus (weaker lens)
- Negative defocus = overfocus (stronger lens)
- Scherzer defocus: Δf = -1.2·(Cs·λ)^(1/2)
- If alpha_max provided, applies hard aperture cutoff
Example:
>>> kx = np.fft.fftfreq(512, d=0.1)
>>> ky = np.fft.fftfreq(512, d=0.1)
>>> KX, KY = np.meshgrid(kx, ky)
>>> H = sim.objective_lens_transfer_function(
... KX, KY, defocus=700, Cs=1.3e7, alpha_max=0.01037
... )
"""
# Spatial frequency magnitude
k2 = kx**2 + ky**2
k = np.sqrt(k2)
# Phase aberration χ(k) - Kirkland Eq. 5.27
chi = (
np.pi
* self.wavelength
* k2
* (0.5 * Cs * self.wavelength**2 * k2 - defocus)
)
# Transfer function
H = np.exp(-1j * chi)
# Apply objective aperture if specified
if alpha_max is not None:
k_max = alpha_max / self.wavelength
aperture = k <= k_max
H *= aperture
return H
[docs]
def simulate_image(
self,
atom_positions: List[Tuple[float, float, int]],
defocus: float = 700.0,
Cs: float = 1.3e7,
alpha_max: Optional[float] = None,
return_wavefunction: bool = False,
) -> Dict[str, np.ndarray]:
"""
Simulate CTEM image using weak phase object approximation.
Full simulation pipeline:
1. Calculate transmission function t(x,y)
2. Fourier transform: ψ(k) = FFT[t(x,y)]
3. Apply lens CTF: ψ'(k) = ψ(k)·H(k)
4. Inverse transform: ψ(x,y) = IFFT[ψ'(k)]
5. Calculate intensity: I(x,y) = |ψ(x,y)|²
Args:
atom_positions: List of (x, y, Z) atom coordinates and atomic numbers
defocus: Defocus in Angstroms (positive = underfocus)
Cs: Spherical aberration in Angstroms (e.g., 1.3e7 = 1.3 mm)
alpha_max: Aperture semi-angle in milliradians (converted internally)
return_wavefunction: If True, include complex wavefunction in results
Returns:
Dictionary containing:
- 'intensity': Image intensity I = |ψ|², shape (pixels, pixels)
- 'transmission': Complex transmission function
- 'potential': Projected atomic potential in eV
- 'positions': Atom positions used
- 'psi': Complex wavefunction (if return_wavefunction=True)
Example:
>>> # Reproduce Kirkland Figure 5.12
>>> atoms = [(-20,0,6), (-10,0,14), (0,0,29), (10,0,79), (20,0,92)]
>>> results = sim.simulate_image(
... atom_positions=atoms,
... defocus=700.0, # 700 Å underfocus
... Cs=1.3e7, # 1.3 mm
... alpha_max=10.37 # 10.37 mrad
... )
>>> intensity = results['intensity']
>>> # Intensity range should be ~0.72 to 1.03 (Kirkland value)
"""
# Convert aperture angle from milliradians to radians
if alpha_max is not None:
alpha_max = alpha_max * 1e-3
# Step 1: Calculate transmission function
transmission, potential = self.calculate_transmission_function(atom_positions)
# Step 2: Fourier transform (fftshift for zero-frequency at center)
psi_k = np.fft.fftshift(np.fft.fft2(transmission))
# Step 3: Get spatial frequency coordinates
kx = np.fft.fftshift(np.fft.fftfreq(self.pixels, d=self.dx))
ky = np.fft.fftshift(np.fft.fftfreq(self.pixels, d=self.dx))
KX, KY = np.meshgrid(kx, ky, indexing="xy")
# Step 4: Apply objective lens transfer function
H = self.objective_lens_transfer_function(KX, KY, defocus, Cs, alpha_max)
psi_k *= H
# Step 5: Inverse Fourier transform
psi = np.fft.ifft2(np.fft.ifftshift(psi_k))
# Step 6: Calculate intensity
intensity = np.abs(psi) ** 2
# Prepare results
results = {
"intensity": intensity,
"transmission": transmission,
"potential": potential,
"positions": atom_positions,
}
if return_wavefunction:
results["psi"] = psi
return results
[docs]
def get_simulation_parameters(self) -> Dict[str, float]:
"""
Get current simulation parameters.
Returns:
Dictionary with simulation parameters:
- beam_energy: in eV
- wavelength: in Angstroms
- sigma: interaction parameter in rad/eV
- pixel_size: in Angstroms
- image_size: in Angstroms
- pixels: number of pixels
Example:
>>> params = sim.get_simulation_parameters()
>>> print(f"Wavelength: {params['wavelength']:.5f} Å")
Wavelength: 0.02508 Å
"""
return {
"beam_energy": self.beam_energy,
"wavelength": self.wavelength,
"sigma": self.sigma,
"pixel_size": self.dx,
"image_size": self.image_size,
"pixels": self.pixels,
}