ptyrad.optics.probe

ptyrad.optics.probe#

Numpy-based electron / x-ray probe generation functions

Functions

hermite_like(fundam, M, N)

Generates orthogonal Hermite-like probe modes from a fundamental mode.

make_aberration_surface_krivanek_cartesian(...)

Calculates the aberration phase surface using recursive Cartesian polynomials.

make_aberration_surface_krivanek_complex(...)

Calculates the aberration phase surface chi(k) using Krivanek Complex form.

make_aberration_surface_krivanek_polar(...)

Calculates the aberration phase surface chi(k) using Krivanek Polar form.

make_fzp_probe(beam_kev, Npix, dx, Ls, Rn, ...)

Generates a Fresnel zone plate probe with internal Fresnel propagation for x-ray ptychography simulations.

make_mixed_probe(probe, pmodes, pmode_init_pows)

Make a mixed state probe from a single state probe

make_stem_probe(kv, conv_angle, Npix, dx, ...)

Simulates a STEM probe in real space using the specified methods for chi(k) calculations.

orthogonalize_modes_vec_np(modes[, sort])

Orthogonalize probe modes via Gram matrix eigendecomposition (NumPy version).

sort_by_mode_int_np(modes)

Sorts a set of modes in descending order based on their total intensity.
ptyrad.optics.probe.make_aberration_surface_krivanek_polar(aberrations, kX, kY, wavelength)[source]#

Calculates the aberration phase surface chi(k) using Krivanek Polar form.

Implements the standard polar expansion as defined in Kirkland Eqn. 2.22.

Parameters:

  • aberrations (Dict[Tuple[int, int], Dict[str, float]]) – A dictionary mapping order (n, m) to polar coefficients. Format: {(n, m): {‘mag’: float, ‘phi’: float}} ‘mag’: Coefficient magnitude (e.g., C_s) in Angstroms. ‘phi’: Azimuthal angle in degrees.

  • kX (ndarray) – Spatial frequency coordinate X (1/Angstrom).

  • kY (ndarray) – Spatial frequency coordinate Y (1/Angstrom).

  • wavelength (float) – Electron wavelength in Angstroms.

Returns:

The aberration phase surface in radians.

Return type:

np.ndarray

ptyrad.optics.probe.make_aberration_surface_krivanek_complex(aberrations, kX, kY, wavelength)[source]#

Calculates the aberration phase surface chi(k) using Krivanek Complex form.

Implements the complex power series expansion (Kirkland Eqn. 2.19/2.20). This form utilizes the complex coordinate omega = alpha_x + i*alpha_y. Note that we swapped the exponents of omega and conj(omega) so the angle convention is consistent with Cartesian and Polar form.

Parameters:

  • aberrations (Dict[Tuple[int, int], complex]) – A dictionary mapping order (n, m) to complex coefficients. Format: {(n, m): complex_value}

  • kX (ndarray) – Spatial frequency coordinate X (1/Angstrom).

  • kY (ndarray) – Spatial frequency coordinate Y (1/Angstrom).

  • wavelength (float) – Electron wavelength in Angstroms.

Returns:

The aberration phase surface in radians (real-valued).

Return type:

np.ndarray

ptyrad.optics.probe.make_aberration_surface_krivanek_cartesian(aberrations, kX, kY, wavelength)[source]#

Calculates the aberration phase surface using recursive Cartesian polynomials.

This method is significantly faster than polar or complex forms for large arrays as it avoids expensive trigonometric operations by using recursive multiplication.

Mathematical Derivation:#

  1. Start with the standard Polar form (Kirkland Eq 2.22): $chi(alpha, phi) = frac{2pi}{lambda} frac{1}{n+1} C_{nm} alpha^{n+1} cos[m(phi - phi_{nm})]$

  2. Expand the cosine term $cos(mphi - mphi_{nm})$ and ignore the prefactor and summation for now: $chi propto alpha^{n+1} [ cos(mphi)cos(mphi_{nm}) + sin(mphi)sin(mphi_{nm}) ]$

  3. Define Cartesian coefficients $C_{nma}, C_{nmb}$ to substitute $C_{nm}$ and $phi_{nm}$: $C_{nma} = C_{nm} cos(mphi_{nm})$ $C_{nmb} = C_{nm} sin(mphi_{nm})$

  4. Split the radial term $alpha^{n+1}$ into $alpha^{n+1-m} cdot alpha^m$ to isolate angular parts: $chi propto alpha^{n+1-m} [ C_{nma} (alpha^m cos mphi) + C_{nmb} (alpha^m sin mphi) ]$

  5. Define Cartesian Angular Polynomials $X_m, Y_m$ using complex variable $Z = alpha_x + ialpha_y$: $Z^m = (alpha e^{iphi})^m = alpha^m (cos mphi + i sin mphi)$

    Therefore: $X_m = text{Re}(Z^m) = alpha^m cos(mphi)$ $Y_m = text{Im}(Z^m) = alpha^m sin(mphi)$

  6. Final Calculation: $X_m, Y_m$ are pre-calculated using the recurrence $Z_{m+1} = Z_m cdot Z$. $chi = frac{2pi}{lambda} sum frac{1}{n+1} (alpha^2)^{frac{n+1-m}{2}} [ C_{nma}X_m + C_{nmb}Y_m ]$

param aberrations:

A dictionary mapping order (n, m) to Cartesian coefficients. Format: {(n, m): {‘a’: float, ‘b’: float}} ‘a’: Cnma, cosine-like coefficient (Real part). ‘b’: Cnmb, sine-like coefficient (Imaginary part).

param kX:

Spatial frequency coordinate X (1/Angstrom).

param kY:

Spatial frequency coordinate Y (1/Angstrom).

param wavelength:

Electron wavelength in Angstroms.

returns:

The aberration phase surface in radians.

rtype:

np.ndarray

Parameters:

  • aberrations (Dict[Tuple[int, int], Dict[str, float]])

  • kX (ndarray)

  • kY (ndarray)

  • wavelength (float)

Return type:

ndarray

ptyrad.optics.probe.make_stem_probe(kv, conv_angle, Npix, dx, aberrations, method='cartesian')[source]#

Simulates a STEM probe in real space using the specified methods for chi(k) calculations. The three methods (polar, cartesian, complex) give identical result within numerical precision, while ‘cartesian’ is chosen as the default as it’s the fastest (though they’re all just few ms).

Constructs the probe by defining the aperture and aberrations in Fourier space, applying the phase shift, and performing an inverse FFT to obtain the real-space complex wave function.

Parameters:

  • kv (float) – Acceleration voltage in kilovolts (kV).

  • conv_angle (float) – Convergence semi-angle in milliradians (mrad).

  • Npix (int) – Number of pixels for the square simulation grid.

  • dx (float) – Real-space pixel size in Angstroms.

  • aberrations (dict | Aberrations) – An Aberrations instance, or dictionary of aberration coefficients. The dictionary can be in Haider (e.g., {‘C1’: 10}), or Krivanek (e.g., {‘C12’: 10, ‘phi12’: 30}) notation in polar / cartesian / complex form. Mix-match and aliases like ‘defocus’, ‘Cs’ are supported.

  • method (Literal['polar', 'cartesian', 'complex']) – The computation approach for chi(k) calculation. Options: - ‘polar’: Standard Krivanek polar form (C_nm * alpha^(n+1) * cos[m(phi-phi_nm)]). - ‘cartesian’: Recursive Cartesian polynomials (C_nma * X[m] + C_nmb * Y[m]). - ‘complex’: Analytic complex power series (C_nm * w*^(n+1-s) * w^s).

Returns:

A 2D complex array representing the probe wave function in real space, normalized such that the total intensity sums to 1.

Return type:

np.ndarray

ptyrad.optics.probe.make_fzp_probe(beam_kev, Npix, dx, Ls, Rn, dRn, D_FZP, D_H)[source]#

Generates a Fresnel zone plate probe with internal Fresnel propagation for x-ray ptychography simulations.

Parameters:

  • beam_kev (float) – Energy of the x-ray photon.

  • Npix (int) – Number of pixels.

  • dx (float) – Pixel size (in meters) in the sample plane.

  • Ls (float) – Distance (in meters) from the focal plane to the sample.

  • Rn (float) – Radius of outermost zone (in meters).

  • dRn (float) – Width of outermost zone (in meters).

  • D_FZP (float) – Diameter of pinhole.

  • D_H (float) – Diameter of the central beamstop (in meters).

Returns:

Calculated probe field in the sample plane.

Return type:

ndarray

ptyrad.optics.probe.make_mixed_probe(probe, pmodes, pmode_init_pows)[source]#

Make a mixed state probe from a single state probe

ptyrad.optics.probe.hermite_like(fundam, M, N)[source]#

Generates orthogonal Hermite-like probe modes from a fundamental mode.

This function takes a base probe (the fundamental mode) and multiplies it by Hermitian functions up to a maximum $x$-order $M$ and $y$-order $N$ to compute higher-order modes. The resulting modes are then iteratively orthonormalized against all previously generated modes.

Note

This is a Python implementation ported from ptycho/+core/hermite_like.m in PtychoShelves, with the following modifications:

  • Array indexing is converted from MATLAB (1-based) to Python (0-based).

  • The X and Y spatial meshgrids are generated internally rather than passed as arguments.

  • The output tensor H has the shape (pmode, Ny, Nx) to be consistent with PtyRAD conventions.

  • The function always outputs $(M+1)(N+1)$ modes, which may be slightly more than a user’s target pmode count (requiring subsequent truncation).

Parameters:

  • fundam (numpy.ndarray) – The base fundamental probe function, typically a 2D complex array of shape (Ny, Nx).

  • M (int or float) – The maximum $x$-order of the Hermite basis.

  • N (int or float) – The maximum $y$-order of the Hermite basis.

Returns:

A 3D complex array of the generated orthonormalized modes with shape ((M+1)*(N+1), Ny, Nx).

Return type:

numpy.ndarray

ptyrad.optics.probe.sort_by_mode_int_np(modes)[source]#

Sorts a set of modes in descending order based on their total intensity.

The intensity of each mode is calculated as the sum of its squared amplitude across all spatial dimensions. This is commonly used to ensure the dominant probe or object modes are positioned at the lowest indices.

Parameters:

modes (numpy.ndarray) – An N-dimensional array of modes, where the first dimension represents the mode index (e.g., (pmode, Ny, Nx) for 2D modes or (omode, Nz, Ny, Nx) for 3D modes).

Returns:

The input array sorted in descending order of total intensity along the first dimension.

Return type:

numpy.ndarray

ptyrad.optics.probe.orthogonalize_modes_vec_np(modes, sort=False)[source]#

Orthogonalize probe modes via Gram matrix eigendecomposition (NumPy version).

Mirrors the logic of orthogonalize_modes_vec (PyTorch) but operates on NumPy arrays, making it suitable for probe initialization where PyTorch is not yet involved.

Uses np.linalg.eigh (Hermitian eigensolver) for numerical stability. eigh dispatches to dsyev/cheev on all LAPACK backends (Accelerate on macOS, MKL on Windows/Linux), unlike eig which used the numerically weaker cgeev. The Gram matrix is computed and decomposed in complex128 to guard against float32/64 precision loss.

Falls back silently to the original modes if the result is invalid (poor orthogonality or norm change), matching the behaviour of the Torch version.

Parameters:

  • modes (np.ndarray) – Input modes of shape (Nmode, Ny, Nx), complex.

  • sort (bool, optional) – Whether to sort modes by their intensity (descending), by default False.

Returns:

Orthogonalized modes of the same shape and dtype as input, or the original modes unchanged if the validity check fails.

Return type:

np.ndarray

Contents