Source code for quscope.quantum_ctem.hamiltonian

"""
Quantum Hamiltonian for Conventional TEM

This module implements the complete quantum mechanical Hamiltonian for electron
wave propagation in Conventional Transmission Electron Microscopy, including:

1. Free propagation
2. Sample interaction (Weak Phase Object Approximation)
3. Lens aberrations (up to 5th order)
4. Evolution operator decomposition
5. Mapping to quantum circuits

Theoretical Framework:
    The electron wave function evolution in TEM can be described by the
    time-independent Schrödinger equation with a position-dependent potential.

    Total Hamiltonian:
        H_total = H_0 + H_sample + H_lens

    where:
        H_0: Free particle kinetic energy
        H_sample: Sample-electron interaction
        H_lens: Lens aberrations in momentum space

    The evolution operator is:
        U = exp(-iH_lens·t/ℏ) · exp(-iH_sample·t/ℏ) · exp(-iH_0·t/ℏ)

    This operator can be efficiently implemented as a quantum circuit.

References:
    - Messiah, A. (1961). Quantum Mechanics. North-Holland.
    - Kirkland, E. J. (2010). Advanced Computing in Electron Microscopy.
    - Nielsen & Chuang (2010). Quantum Computation and Quantum Information.

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

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

import numpy as np
import scipy.constants as const
from scipy.fft import fft2, fftfreq, ifft2


[docs] @dataclass class HamiltonianParameters: """ Parameters for the quantum Hamiltonian. Attributes: acceleration_voltage: Electron acceleration voltage (V) wavelength: Relativistic electron wavelength (Angstrom) grid_size_x: Number of grid points in x grid_size_y: Number of grid points in y pixel_size: Real-space pixel size (Angstrom) interaction_constant: σ = (2π·m_e·e·λ)/(h²) for WPOA """ acceleration_voltage: float wavelength: float grid_size_x: int grid_size_y: int pixel_size: float interaction_constant: float = None
[docs] def __post_init__(self): """Calculate interaction constant if not provided.""" if self.interaction_constant is None: # σ = interaction constant for WPOA: χ = σ·V # From Kirkland "Advanced Computing in Electron Microscopy": # σ [rad/(V·Å)] = me·e/(2π·ℏ²) × λ # Simplified: σ ≈ 0.01 × λ[Å] / (acceleration_voltage[kV]) # # For 200 kV: σ ≈ 3.2e-4 rad/(V·Å) # For 300 kV: σ ≈ 2.1e-4 rad/(V·Å) m_e = const.electron_mass # kg e = const.elementary_charge # C h = const.Planck # J·s hbar = h / (2 * np.pi) lambda_m = self.wavelength * 1e-10 # Convert Å to m # Correct formula: σ = (me·e)/(2π·ℏ²) × λ # Units: [kg·C]/[J·s]² × [m] = [kg·C·m]/[kg²·m⁴/s²·s²] = [C]/[kg·m³/s²] = 1/(V·m) self.interaction_constant = (m_e * e * lambda_m) / (2 * np.pi * hbar**2) # Convert from 1/(V·m) to 1/(V·Å) self.interaction_constant *= 1e-10
[docs] class FreeParticleHamiltonian: """ Free particle Hamiltonian H₀ = p²/(2m) = ℏ²k²/(2m). In the paraxial approximation (small scattering angles), this describes the kinetic energy of the electron beam. For high-energy electrons (80-300 kV), relativistic corrections are small but can be included in the wavelength calculation. """ def __init__(self, params: HamiltonianParameters): """ Initialize free particle Hamiltonian. Args: params: Hamiltonian parameters """ self.params = params self.kx, self.ky = self._generate_k_grid() self.k_squared = self.kx**2 + self.ky**2 def _generate_k_grid(self) -> Tuple[np.ndarray, np.ndarray]: """ Generate momentum space grid. Returns: kx, ky: 2D arrays of k-space coordinates (1/Angstrom) """ # Frequency grids freq_x = fftfreq(self.params.grid_size_x, d=self.params.pixel_size) freq_y = fftfreq(self.params.grid_size_y, d=self.params.pixel_size) # Convert to k-space (k = 2π·freq) kx, ky = np.meshgrid(2 * np.pi * freq_x, 2 * np.pi * freq_y, indexing="ij") return kx, ky
[docs] def energy(self) -> np.ndarray: """ Calculate kinetic energy E = ℏ²k²/(2m). In practice, we work with dimensionless units where ℏ=m=1, so E = k²/2. Returns: Energy at each k-point (2D array) """ return 0.5 * self.k_squared
[docs] def propagator(self, distance: float) -> np.ndarray: """ Free space propagation operator exp(-iH₀·z/ℏ). For small angles (paraxial approximation): exp(-iH₀·z/ℏ) ≈ exp(-iπλz·k²) Args: distance: Propagation distance (Angstrom) Returns: Propagator in momentum space (2D complex array) """ # Phase factor: -π·λ·z·k² phase = -np.pi * self.params.wavelength * distance * self.k_squared return np.exp(1j * phase)
[docs] def apply(self, psi: np.ndarray, distance: float) -> np.ndarray: """ Apply free space propagation to wave function. Args: psi: Wave function in real space (2D complex array) distance: Propagation distance (Angstrom) Returns: Propagated wave function in real space """ # Transform to momentum space psi_k = fft2(psi) # Apply propagator psi_k_propagated = psi_k * self.propagator(distance) # Transform back to real space return ifft2(psi_k_propagated)
[docs] class SampleHamiltonian: """ Sample interaction Hamiltonian under Weak Phase Object Approximation (WPOA). H_sample = V(x,y) where V is the projected atomic potential. The transmission function is: t(x,y) = exp(iσV(x,y)) where σ = (2π·m_e·e·λ)/(h²) is the interaction constant. WPOA is valid when: σV << 1 (weak phase modulation) Sample thickness << mean free path """ def __init__(self, params: HamiltonianParameters): """ Initialize sample Hamiltonian. Args: params: Hamiltonian parameters """ self.params = params self.potential = None
[docs] def set_potential(self, V: np.ndarray): """ Set the projected potential V(x,y). Args: V: Projected potential in V·Angstrom (2D array) """ assert V.shape == ( self.params.grid_size_x, self.params.grid_size_y, ), f"Potential shape {V.shape} doesn't match grid size" self.potential = V
[docs] def transmission_function(self) -> np.ndarray: """ Calculate transmission function t(x,y) = exp(iσV(x,y)). Returns: Transmission function (2D complex array) """ if self.potential is None: # No sample → pure transmission return np.ones( (self.params.grid_size_x, self.params.grid_size_y), dtype=complex ) # Apply phase shift: exp(iσV) phase = self.params.interaction_constant * self.potential return np.exp(1j * phase)
[docs] def apply(self, psi: np.ndarray) -> np.ndarray: """ Apply sample interaction to wave function. This is a simple multiplication in real space: ψ_exit(x,y) = ψ_incident(x,y) · t(x,y) Args: psi: Incident wave function (2D complex array) Returns: Exit wave function (2D complex array) """ return psi * self.transmission_function()
[docs] class LensHamiltonian: """ Lens aberration Hamiltonian in momentum space. The lens applies a k-dependent phase shift: H_lens → exp(iχ(k)) where χ(k) is the wave aberration function including all aberration coefficients up to 5th order. This operator is diagonal in momentum space, making it efficient to apply. """ def __init__(self, params: HamiltonianParameters, aberrations: Dict[str, float]): """ Initialize lens Hamiltonian. Args: params: Hamiltonian parameters aberrations: Dictionary of aberration coefficients Keys: 'defocus', 'c12a', 'c12b', 'c3', 'c5', etc. """ self.params = params self.aberrations = aberrations # Generate k-space grid self.kx, self.ky = self._generate_k_grid() self.k = np.sqrt(self.kx**2 + self.ky**2) self.theta = np.arctan2(self.ky, self.kx) # Pre-calculate aberration function self.chi = self.calculate_aberration_function() def _generate_k_grid(self) -> Tuple[np.ndarray, np.ndarray]: """Generate momentum space grid.""" freq_x = fftfreq(self.params.grid_size_x, d=self.params.pixel_size) freq_y = fftfreq(self.params.grid_size_y, d=self.params.pixel_size) kx, ky = np.meshgrid(2 * np.pi * freq_x, 2 * np.pi * freq_y, indexing="ij") return kx, ky
[docs] def calculate_aberration_function(self) -> np.ndarray: """ Calculate complete wave aberration function χ(k). Standard TEM convention (Kirkland, Spence & Zuo): χ(k) = π λ Δf k² + 0.5 π λ³ Cs k⁴ + ... where: λ = wavelength (Angstrom) Δf = defocus (Angstrom, positive = underfocus) k = spatial frequency (1/Angstrom) Cs = spherical aberration (Angstrom) Returns: χ(kx, ky) in radians (2D array) """ lam = self.params.wavelength k = self.k # Already in 1/Angstrom chi = np.zeros_like(k) # Defocus term: χ₁ = π λ Δf k² if "defocus" in self.aberrations: defocus = self.aberrations["defocus"] # Angstrom chi += np.pi * lam * defocus * k**2 # 2-fold astigmatism if "c12a" in self.aberrations and "c12b" in self.aberrations: C12 = self.aberrations["c12a"] phi12 = self.aberrations["c12b"] chi += np.pi * lam * C12 * k**2 * np.cos(2 * (self.theta - phi12)) # Spherical aberration Cs: χ₃ = 0.5 π λ³ Cs k⁴ if "cs" in self.aberrations: Cs = self.aberrations["cs"] * 1e7 # Convert mm to Angstrom chi += 0.5 * np.pi * (lam**3) * Cs * k**4 # Alternative key for Cs if "c3" in self.aberrations: C3 = self.aberrations["c3"] * 1e7 # Convert mm to Angstrom chi += 0.5 * np.pi * (lam**3) * C3 * k**4 # Higher order aberrations (keeping dimensionless form with 2π/λ factor later) # These use lam_k = λk for dimensional consistency lam_k = lam * k # 2nd order: Axial coma and 3-fold astigmatism if "c21a" in self.aberrations and "c21b" in self.aberrations: C21 = self.aberrations["c21a"] phi21 = self.aberrations["c21b"] chi += ( (2 * np.pi / lam) * (1 / 3) * C21 * lam_k**3 * np.cos(self.theta - phi21) ) if "c23a" in self.aberrations and "c23b" in self.aberrations: C23 = self.aberrations["c23a"] phi23 = self.aberrations["c23b"] chi += ( (2 * np.pi / lam) * (1 / 3) * C23 * lam_k**3 * np.cos(3 * (self.theta - phi23)) ) # 3rd order: 4-fold astigmatism if "c32a" in self.aberrations and "c32b" in self.aberrations: C32 = self.aberrations["c32a"] * 1e7 phi32 = self.aberrations["c32b"] chi += ( (2 * np.pi / lam) * 0.25 * C32 * lam_k**4 * np.cos(2 * (self.theta - phi32)) ) if "c34a" in self.aberrations and "c34b" in self.aberrations: C34 = self.aberrations["c34a"] * 1e7 phi34 = self.aberrations["c34b"] chi += ( (2 * np.pi / lam) * 0.25 * C34 * lam_k**4 * np.cos(4 * (self.theta - phi34)) ) # 4th order aberrations if "c41a" in self.aberrations and "c41b" in self.aberrations: C41 = self.aberrations["c41a"] * 1e7 phi41 = self.aberrations["c41b"] chi += (2 * np.pi / lam) * 0.2 * C41 * lam_k**5 * np.cos(self.theta - phi41) if "c43a" in self.aberrations and "c43b" in self.aberrations: C43 = self.aberrations["c43a"] * 1e7 phi43 = self.aberrations["c43b"] chi += ( (2 * np.pi / lam) * 0.2 * C43 * lam_k**5 * np.cos(3 * (self.theta - phi43)) ) if "c45a" in self.aberrations and "c45b" in self.aberrations: C45 = self.aberrations["c45a"] * 1e7 phi45 = self.aberrations["c45b"] chi += ( (2 * np.pi / lam) * 0.2 * C45 * lam_k**5 * np.cos(5 * (self.theta - phi45)) ) # 5th order: 5th order spherical aberration if "c5" in self.aberrations: C5 = self.aberrations["c5"] * 1e7 # Convert mm to Angstrom chi += (2 * np.pi / lam) * (1 / 6) * C5 * lam_k**6 if "c52a" in self.aberrations and "c52b" in self.aberrations: C52 = self.aberrations["c52a"] * 1e7 phi52 = self.aberrations["c52b"] chi += ( (2 * np.pi / lam) * (1 / 6) * C52 * lam_k**6 * np.cos(2 * (self.theta - phi52)) ) if "c54a" in self.aberrations and "c54b" in self.aberrations: C54 = self.aberrations["c54a"] * 1e7 phi54 = self.aberrations["c54b"] chi += ( (2 * np.pi / lam) * (1 / 6) * C54 * lam_k**6 * np.cos(4 * (self.theta - phi54)) ) if "c56a" in self.aberrations and "c56b" in self.aberrations: C56 = self.aberrations["c56a"] * 1e7 phi56 = self.aberrations["c56b"] chi += ( (2 * np.pi / lam) * (1 / 6) * C56 * lam_k**6 * np.cos(6 * (self.theta - phi56)) ) return chi
[docs] def transfer_function(self) -> np.ndarray: """ Calculate lens transfer function exp(iχ(k)). Returns: Transfer function in momentum space (2D complex array) """ return np.exp(1j * self.chi)
[docs] def ctf(self) -> np.ndarray: """ Calculate Contrast Transfer Function sin(χ(k)). For phase contrast imaging with weak phase objects. Returns: CTF (2D array) """ return np.sin(self.chi)
[docs] def apply(self, psi: np.ndarray) -> np.ndarray: """ Apply lens aberrations to wave function. This operates in momentum space: ψ_image(k) = exp(iχ(k)) · ψ_exit(k) Args: psi: Exit wave function in real space (2D complex array) Returns: Image wave function in real space (2D complex array) """ # Transform to momentum space psi_k = fft2(psi) # Apply transfer function psi_k_aberrated = psi_k * self.transfer_function() # Transform back to real space return ifft2(psi_k_aberrated)
[docs] class TEMHamiltonian: """ Complete TEM Hamiltonian combining all contributions. Total evolution: U = U_lens · U_sample · U_propagation This represents the complete electron wave propagation through: 1. Free space to sample 2. Sample interaction 3. Free space to image plane 4. Lens aberrations """ def __init__(self, params: HamiltonianParameters, aberrations: Dict[str, float]): """ Initialize complete TEM Hamiltonian. Args: params: Hamiltonian parameters aberrations: Aberration coefficients """ self.params = params self.aberrations = aberrations # Initialize component Hamiltonians self.H_free = FreeParticleHamiltonian(params) self.H_sample = SampleHamiltonian(params) self.H_lens = LensHamiltonian(params, aberrations)
[docs] def set_sample_potential(self, V: np.ndarray): """ Set sample potential. Args: V: Projected potential (V·Angstrom) """ self.H_sample.set_potential(V)
[docs] def propagate( self, psi_incident: np.ndarray, propagation_distance: float = 0.0 ) -> np.ndarray: """ Complete TEM wave propagation. Args: psi_incident: Incident wave function propagation_distance: Free space propagation before sample Returns: Final image wave function """ # Step 1: Free space propagation to sample (if needed) if propagation_distance > 0: psi = self.H_free.apply(psi_incident, propagation_distance) else: psi = psi_incident.copy() # Step 2: Sample interaction psi_exit = self.H_sample.apply(psi) # Step 3: Lens aberrations (in image plane) psi_image = self.H_lens.apply(psi_exit) return psi_image
[docs] def get_ctf(self) -> np.ndarray: """Get Contrast Transfer Function.""" return self.H_lens.ctf()
[docs] def get_aberration_function(self) -> np.ndarray: """Get wave aberration function χ(k).""" return self.H_lens.chi
[docs] def calculate_intensity(self, psi: np.ndarray) -> np.ndarray: """ Calculate image intensity I = |ψ|². Args: psi: Image wave function Returns: Intensity image """ return np.abs(psi) ** 2
[docs] def relativistic_wavelength(voltage: float) -> float: """ Calculate relativistic electron wavelength. λ = h / √(2·m₀·e·V·(1 + e·V/(2·m₀·c²))) Args: voltage: Acceleration voltage (V) Returns: Wavelength in Angstroms """ m0 = const.electron_mass # kg e = const.elementary_charge # C h = const.Planck # J·s c = const.speed_of_light # m/s # Relativistic correction factor gamma = 1 + e * voltage / (m0 * c**2) # Wavelength in meters lambda_m = h / np.sqrt(2 * m0 * e * voltage * gamma) # Convert to Angstroms return lambda_m * 1e10
# Example usage and validation if __name__ == "__main__": print("=" * 70) print("TEM Hamiltonian Module - Validation") print("=" * 70) print() # Setup parameters params = HamiltonianParameters( acceleration_voltage=200e3, wavelength=relativistic_wavelength(200e3), grid_size_x=4, grid_size_y=4, pixel_size=0.1, ) print(f"Acceleration voltage: {params.acceleration_voltage/1e3:.0f} kV") print(f"Wavelength: {params.wavelength:.4f} Å") print(f"Interaction constant σ: {params.interaction_constant:.4e} 1/(V·Å²)") print() # Define aberrations aberrations = {"defocus": 500.0, "c3": 1.3, "c5": 10.0} # Angstrom # mm # mm # Create Hamiltonian H = TEMHamiltonian(params, aberrations) print("Hamiltonian components initialized:") print(f" ✓ Free particle propagation") print(f" ✓ Sample interaction (WPOA)") print(f" ✓ Lens aberrations (5th order)") print() # Create test wave function (plane wave) psi_0 = np.ones((4, 4), dtype=complex) / 4.0 print("Test propagation:") psi_final = H.propagate(psi_0) print(f" Input norm: {np.linalg.norm(psi_0):.6f}") print(f" Output norm: {np.linalg.norm(psi_final):.6f}") print( f" Norm preserved: {np.allclose(np.linalg.norm(psi_0), np.linalg.norm(psi_final))}" ) print() print("✓ Hamiltonian module validated")