"""
Quantum-Enhanced Weak Phase Object Simulations
======================================
This module implements a hybrid quantum-classical framework for simulating
CTEM and STEM images of thin specimen based on weak phase object approximation.
This leverages QFTs and inverse QFTs, replacing various FFTs and iFFTs to reduce
computational overhead.
"""
import matplotlib.pyplot as plt
import numpy as np
from quscope.simulations.quantum_utils import TEMQFT
from quscope.utils.constants import PhysicalConstants
from quscope.utils.kirkland import KirklandPotential
[docs]
class ThinCTEM:
"""
Quantum CTEM simulation for thin specimens using weak phase object approximation.
This class implements:
- Weak phase object approximation for thin specimens
- QFT replacing classical FFT
- Support for abritrary atomic structures
- For CTEM at the moment
"""
def __init__(
self,
image_size=50.0,
n_qubits=8,
beam_energy=200e3,
kirkland_params_file="kirkland.json",
):
"""
Initialize thin specimen simulator.
Parameters:
-----------
image_size : float
Lateral size of the image in Angstroms.
n_qubits : int
Number of qubits per dimension (n_qubits=8 gives 256x256 pixels)
beam_energy : float
Electron beam energy in eV.
"""
self.image_size = image_size
self.n_qubits = n_qubits
self.pixels = 2**n_qubits
self.beam_energy = beam_energy
# Calculate beam parameters
self.wavelength = PhysicalConstants.calculate_wavelength(beam_energy)
self.sigma = PhysicalConstants.calculate_sigma(beam_energy)
# Load Kirkland parameters
self.params = KirklandPotential(kirkland_params_file)
# Set up QFTs
self.qfts = TEMQFT(n_qubits)
# Create coordinate grids (real space)
self.dx = self.image_size / self.pixels
x = (np.arange(self.pixels) - self.pixels / 2 + 0.5) * self.dx
self.x = x
self.y = x
self.X, self.Y = np.meshgrid(x, x, indexing="xy")
[docs]
def calculate_transmission_function(self, atom_positions, atom_z_values):
"""
Calculate transmission function for weak phase object.
Parameters:
-----------
atom_positions : list
List of (x,y) positions in Angstroms.
atom_z_values : list
List of atomic numbers corresponding to positions
Returns:
--------
transmission : np.ndarray
Complex transmission function.
"""
# Get projected potential
V_total = np.zeros((self.pixels, self.pixels))
print("\nAtomic potential peaks:")
for (x_atom, y_atom), Z in zip(atom_positions, atom_z_values):
V_atom = self.params.kirkland_potential_2d(
self.X, self.Y, x_atom, y_atom, Z
)
V_total += V_atom
idx_x = np.argmin(np.abs(self.x - x_atom))
idx_y = np.argmin(np.abs(self.y - y_atom))
V_peak = V_atom[idx_y, idx_x]
element = self.params.get_element_symbol(Z)
print(f"{element} (Z={Z}): V_peak = {V_peak:.2f} eV")
phase = self.sigma * V_total
print(f"\nPhase range: [{np.min(phase):.4f}, {np.max(phase):.4f}] radians")
transmission = np.exp(1j * phase)
return transmission, V_total
[docs]
def objective_lens_transfer_function(self, kx, ky, defocus, Cs, alpha_max=None):
"""
Apply objective lens transfer function in reciprocal space.
Parameters:
-----------
kx, ky : float
Spatial components in the x and y directions in space.
defocus : float
Defocus in Angstroms.
Cs : float
Spherical aberration coefficient in Angstroms.
alpha_max : float, optional
Objective aperture semi-angle in mrad.
Returns:
--------
H : np.ndarray
Transfer function.
"""
k2 = kx**2 + ky**2
k = np.sqrt(k2)
chi = (
np.pi
* self.wavelength
* k2
* (0.5 * Cs * self.wavelength**2 * k2 - defocus)
)
H = np.exp(-1j * chi)
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, atom_z_values, defocus=700, Cs=1.3e7, alpha_max=None
):
"""
Simulate CTEM image using quantum algorithms.
Parameters:
-----------
atom_positions : list
List of (x,y) positions in Angstroms.
atom_z_values : list
List of atomic numbers.
defocus : float
Defocus in Angstroms.
Cs : float
Spherical aberration in Angstroms.
alpha_max : float, optional
Objective aperture in mrad.
Returns:
--------
results : Dict
Dictionary containing:
- 'intensity': Final image intensity.
- 'transmission': Complex transmission function.
- 'psi': Exit wave function.
- 'potential': Projected potential.
"""
if alpha_max is not None:
alpha_max = alpha_max * 1e-3 # mrad to rad
# Get transmission function
transmission, potential = self.calculate_transmission_function(
atom_positions, atom_z_values
)
print("\nApplying QFT...")
# QFT
psi_k = self.qfts.qft_2d(transmission)
psi_k = np.fft.fftshift(psi_k) # Classical postprocessing
# Get 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")
# Apply objective lens transfer function
H = self.objective_lens_transfer_function(KX, KY, defocus, Cs, alpha_max)
psi_k *= H
print("Apply Inverse QFT...")
# iQFT
psi_k_shifted = np.fft.ifftshift(psi_k)
psi = self.qfts.iqft_2d(psi_k_shifted)
# Calculate intensity
intensity = np.abs(psi) ** 2
return {
"transmission": transmission,
"intensity": intensity,
"potential": potential,
"psi": psi,
"atom_positions": atom_positions,
"atom_z_values": atom_z_values,
}
[docs]
def plot_transmission_function(self, atom_positions, atom_z_values):
"""Plot transmission function line scan (reproduces Figure 5.11)"""
transmission, _ = self.calculate_transmission_function(
atom_positions, atom_z_values
)
center_idx = self.pixels // 2
line_real = np.real(transmission[center_idx, :])
line_imag = np.imag(transmission[center_idx, :])
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 10))
# Real part
ax1.plot(self.x, line_real, "b-", linewidth=1.5, label="Quantum")
ax1.set_xlim(-25, 25)
ax1.set_ylim(0, 1.2)
ax1.set_ylabel("Real Part")
ax1.grid(True, alpha=0.3)
ax1.legend()
elements = [self.params.get_element_symbol(Z) for Z in atom_z_values]
for (x, y), elem in zip(atom_positions, elements):
ax1.text(x, 1.15, elem, ha="center", va="bottom", fontsize=10)
ax1.axvline(x, color="gray", linestyle=":", alpha=0.5)
ax1.set_title(
"Line scan of the complex transmission function\n"
+ f"for {len(atom_positions)} atoms ({self.beam_energy/1e3:.0f} keV)",
fontsize=12,
)
# Imaginary part
ax2.plot(self.x, line_imag, "b-", linewidth=1.5, label="Quantum")
ax2.set_xlim(-25, 25)
ax2.set_ylim(0, 1.0)
ax2.set_xlabel("position x (in Ang)")
ax2.set_ylabel("Imag Part")
ax2.grid(True, alpha=0.3)
ax2.legend()
for (x, y), elem in zip(atom_positions, elements):
ax2.axvline(x, color="gray", linestyle=":", alpha=0.5)
plt.tight_layout()
plt.show()
return fig
[docs]
def plot_phase_contrast_image(self, results, title_suffix=""):
"""Plot phase contrast image and line scan (reproduces Figure 5.12)"""
intensity = results["intensity"]
atom_positions = results["atom_positions"]
atom_z_values = results["atom_z_values"]
elements = [self.params.get_element_symbol(Z) for Z in atom_z_values]
# 2D image
fig, ax = plt.subplots(1, 1, figsize=(8, 8))
extent = [
-self.image_size / 2,
self.image_size / 2,
-self.image_size / 2,
self.image_size / 2,
]
print(f"\nIntensity range: [{np.min(intensity):.3f}, {np.max(intensity):.3f}]")
im = ax.imshow(
intensity,
extent=extent,
cmap="gray",
origin="lower",
interpolation="bilinear",
vmin=0.7,
vmax=1.1,
)
for x, y in atom_positions:
circle = plt.Circle((x, y), 2.0, fill=False, edgecolor="red", linewidth=1.5)
ax.add_patch(circle)
ax.set_xlabel("x (Å)")
ax.set_ylabel("y (Å)")
ax.set_title(
f"Coherent bright field phase contrast image{title_suffix}", fontsize=12
)
ax.set_xlim(-25, 25)
ax.set_ylim(-25, 25)
plt.colorbar(im, ax=ax, label="Intensity")
plt.tight_layout()
plt.show()
# Line scan through atoms
fig2, ax2 = plt.subplots(1, 1, figsize=(8, 5))
center_idx = self.pixels // 2
line_intensity = intensity[center_idx, :]
ax2.plot(self.x, line_intensity, "b-", linewidth=1.5, label="Quantum")
ax2.set_xlabel("position x (in Ang)")
ax2.set_ylabel("Image Intensity")
ax2.set_title("Line scan through the center of the atoms", fontsize=12)
ax2.grid(True, alpha=0.3)
ax2.set_xlim(-25, 25)
ax2.set_ylim(0.5, 1.1)
ax2.legend()
for (x_pos, y_pos), elem in zip(atom_positions, elements):
ax2.axvline(x_pos, color="gray", linestyle=":", alpha=0.5)
ax2.text(x_pos, 1.08, elem, ha="center", va="bottom", fontsize=10)
plt.tight_layout()
plt.show()
return fig, fig2
# Kirkland example structure
[docs]
def create_five_atoms_example():
"""Create the classic 5-atom example from Kirkland Figure 5.11"""
elements = ["C", "Si", "Cu", "Au", "U"]
z_values = [6, 14, 29, 79, 92]
positions = []
for i in range(5):
x_pos = (i - 2) * 10.0
positions.append([x_pos, 0.0])
return positions, z_values, elements