ptyrax.models.propagation

ptyrax.models.propagation#

Functions

detector_sphere_coordinates(...)

Compute unit vectors on the Ewald sphere for each detector pixel.

nearfield_propagation_coefficient_fourier(...)

propagate_tilted(field, ...[, ...])

Propagate a coherent field to a tilted detector plane.

source_fourier_occupancy_from_tilt_interpolation(...)

Compute which source Fourier pixels are actually observed by the tilted detector.

source_fourier_support_from_tilt_interpolation(...)

Compute the filled Fourier support mask from tilted detector coverage.

tilt_interpolation_indices(field, ...)

Compute interpolation indices for mapping a tilted detector onto the source Fourier grid.

tilt_target_to_incident(sampling, ...)

Transform detector coordinates into the incident field's Fourier frame.

Classes

FarfieldPropagator()

Far-field (Fraunhofer) propagator for ptychographic reconstruction.

Propagator()

Abstract propagator interface mapping input fields to a target geometry.
class ptyrax.models.propagation.FarfieldPropagator[source]#

Bases: Propagator

Far-field (Fraunhofer) propagator for ptychographic reconstruction.

Implements propagation from the sample exit plane to a detector in the far-field regime using Fraunhofer diffraction. The detected field is the Fourier transform of the exit field, mapped onto the detector through the Ewald sphere geometry.

This propagator supports arbitrarily tilted detectors (e.g. for reflection-geometry ptychography) by delegating to propagate_tilted().

Example

>>> propagator = FarfieldPropagator()
>>> detected_field, mask = propagator(exit_field, detector_coords, detector_grid)
__call__(input_field, output_coordinates, output_grid)[source]#

Propagate an input field to the far-field detector plane.

Applies Fraunhofer propagation with tilt correction to map the exit field onto the detector grid.

Parameters:

  • input_field (CoherentField) – The coherent exit field at the sample plane.

  • output_coordinates (CoordinateSystem) – Coordinate system of the detector.

  • output_grid (SamplingGrid) – Sampling grid of the detector pixels.

Returns:

  • The propagated field at the detector plane.

  • A boolean mask indicating valid detector pixels.

Return type:

A tuple of

input_coordinates(output_coordinates, output_grid, wavelength)[source]#

Compute the sample-plane coordinates from detector geometry.

Derives the input (sample-side) coordinate system and sampling grid from the detector configuration using the far-field reciprocal relationship. The input grid pixel size is determined by the wavelength and propagation distance.

Parameters:

  • output_coordinates (CoordinateSystem) – Coordinate system of the detector.

  • output_grid (SamplingGrid) – Sampling grid of the detector.

  • wavelength (jaxtyping.Float) – The illumination wavelength in meters.

Returns:

  • The coordinate system at the sample plane (same rotation as detector, zero translation).

  • The sampling grid at the sample plane, computed via the far-field inverse relationship.

Return type:

A tuple of

class ptyrax.models.propagation.Propagator[source]#

Bases: Module

Abstract propagator interface mapping input fields to a target geometry.

A propagator models how the exit field from a sample propagates to the detector plane. Subclasses implement specific propagation physics (e.g. Fraunhofer far-field diffraction, Fresnel near-field propagation).

Example

>>> propagator = FarfieldPropagator()
>>> output_field, valid_mask = propagator(input_field, detector_coords, detector_grid)
abstractmethod __call__(input_field, output_coordinates, output_grid)[source]#

Propagate an input field to the output (detector) geometry.

Parameters:

  • input_field (CoherentField) – The coherent field at the sample exit plane.

  • output_coordinates (CoordinateSystem) – The coordinate system of the detector, including its position and orientation relative to the global frame.

  • output_grid (SamplingGrid) – The sampling grid defining detector pixel positions.

Returns:

  • The propagated coherent field at the detector plane.

  • A boolean mask indicating which detector pixels receive valid signal (pixels outside the Ewald sphere coverage are marked False).

Return type:

A tuple of

abstractmethod static input_coordinates(output_coordinates, output_grid, wavelength)[source]#

Compute the input (sample-side) coordinates from detector geometry.

Given the detector coordinate system and sampling grid, determines the corresponding coordinate system and sampling grid at the sample exit plane. This is used to initialize the probe and sample grids consistently with the detector geometry.

Parameters:

  • output_coordinates (CoordinateSystem) – The coordinate system of the detector.

  • output_grid (SamplingGrid) – The sampling grid of the detector.

  • wavelength (jaxtyping.Float) – The illumination wavelength in meters.

Returns:

  • The coordinate system at the sample exit plane.

  • The sampling grid at the sample exit plane, derived from the far-field reciprocal relationship.

Return type:

A tuple of

ptyrax.models.propagation.detector_sphere_coordinates(target_coordinate_system, target_sampling)[source]#

Compute unit vectors on the Ewald sphere for each detector pixel.

Projects each detector pixel position into 3D global coordinates using the detector’s coordinate system (position and orientation), then normalizes to obtain unit vectors pointing from the sample to each pixel. These unit vectors lie on the Ewald sphere (when multiplied by the wavenumber) and represent the outgoing wavevector directions \(\hat{\mathbf{k}}_\text{out}\).

Parameters:

  • target_coordinate_system (CoordinateSystem) – Coordinate system of the detector, defining its translation (distance from sample) and rotation (orientation).

  • target_sampling (SamplingGrid) – Sampling grid defining the physical positions of detector pixels.

Returns:

An array of shape (n, 3) containing the unit vectors for each of the n detector pixels in global coordinates.

Return type:

Float[Array, ‘n 3’]

ptyrax.models.propagation.nearfield_propagation_coefficient_fourier(k_z, z_distance, valid=None)[source]#
Parameters:

  • k_z (Float[Array, 'm n'])

  • z_distance (jaxtyping.Float)

  • valid (Bool[Array, 'm n'])

Return type:

Complex[Array, ‘m n’]

ptyrax.models.propagation.propagate_tilted(field, target_coordinate_system, target_sampling, interpolation_mode='real_imaginary', propagator_type='fraunhofer', apply_jacobian=False, skip_interpolate=False)[source]#

Propagate a coherent field to a tilted detector plane.

Performs far-field or near-field propagation followed by interpolation onto the tilted detector grid. The detector may be arbitrarily oriented relative to the sample frame, as specified by target_coordinate_system.

The propagation involves:

  1. Computing the Fourier transform of the input field (Fraunhofer or Fresnel).

  2. Mapping the tilted detector pixels onto the Ewald sphere in the source Fourier frame.

  3. Interpolating the propagated field at the mapped positions.

Parameters:

  • field (CoherentField) – The coherent exit field at the sample plane.

  • target_coordinate_system (CoordinateSystem) – Coordinate system of the detector, defining its position and orientation.

  • target_sampling (SamplingGrid) – Sampling grid of the detector pixels.

  • interpolation_mode (Literal['real_imaginary', 'amplitude_phase']) – How to interpolate complex values. "real_imaginary" interpolates real and imaginary parts independently. "amplitude_phase" interpolates amplitude directly and uses the phase of the real/imaginary interpolation.

  • propagator_type (Literal['fresnel', 'fraunhofer', 'fresnel_nograd']) – The propagation model to use. "fraunhofer" for far-field, "fresnel" for near-field with quadratic phase, "fresnel_nograd" for Fresnel with the quadratic phase factor excluded from gradient computation.

  • apply_jacobian (bool) – If True, applies the Jacobian correction factor from the Ewald sphere projection to preserve energy density.

  • skip_interpolate (bool)

Returns:

  • The propagated CoherentField at the detector plane.

  • A boolean mask of shape target_sampling.shape where True indicates valid detector pixels covered by the source Fourier space.

Return type:

A tuple of

ptyrax.models.propagation.source_fourier_occupancy_from_tilt_interpolation(field, target_coordinate_system, target_sampling)[source]#

Compute which source Fourier pixels are actually observed by the tilted detector.

Determines the set of source Fourier-space pixels that are “hit” by at least one detector pixel when mapped through the Ewald sphere geometry. This is useful for understanding which spatial frequencies of the sample are accessible given the detector placement.

Parameters:

  • field (CoherentField) – The coherent field defining the source Fourier grid.

  • target_coordinate_system (CoordinateSystem) – Coordinate system of the tilted detector.

  • target_sampling (SamplingGrid) – Sampling grid of the detector.

Returns:

  • source_occupancy: A boolean array of shape field.sampling.shape where True indicates that the corresponding Fourier pixel is observed by at least one detector pixel.

  • detector_mask: A boolean array of shape target_sampling.shape indicating valid detector pixels.

  • field_frame_target_indices: The interpolation indices array of shape (2, d) mapping detector pixels to source Fourier indices.

Return type:

A tuple of

ptyrax.models.propagation.source_fourier_support_from_tilt_interpolation(field, target_coordinate_system, target_sampling)[source]#

Compute the filled Fourier support mask from tilted detector coverage.

Extends source_fourier_occupancy_from_tilt_interpolation() by constructing a convex polygon from the detector boundary pixels and filling the interior to produce a continuous support mask. This accounts for the fact that not every interior Fourier pixel may be directly hit by a detector pixel due to discrete sampling, but is still within the accessible frequency range.

Parameters:

  • field (CoherentField) – The coherent field defining the source Fourier grid.

  • target_coordinate_system (CoordinateSystem) – Coordinate system of the tilted detector.

  • target_sampling (SamplingGrid) – Sampling grid of the detector.

Returns:

  • source_support: A boolean array of shape field.sampling.shape representing the filled Fourier support region (union of polygon interior and direct occupancy).

  • source_occupancy: A boolean array of directly occupied Fourier pixels (same as from source_fourier_occupancy_from_tilt_interpolation()).

  • detector_mask: A boolean array of valid detector pixels.

  • field_frame_target_indices: The interpolation indices array of shape (2, d).

Return type:

A tuple of

ptyrax.models.propagation.tilt_interpolation_indices(field, target_coordinate_system, target_sampling)[source]#

Compute interpolation indices for mapping a tilted detector onto the source Fourier grid.

Maps detector pixel positions through the Ewald sphere geometry to determine which source Fourier-space pixels they correspond to. This enables interpolation of the propagated field onto an arbitrarily oriented detector in the far-field.

Parameters:

  • field (CoherentField) – The coherent field whose Fourier grid defines the source frame.

  • target_coordinate_system (CoordinateSystem) – Coordinate system of the tilted detector.

  • target_sampling (SamplingGrid) – Sampling grid of the detector pixels.

Returns:

  • Interpolation indices of shape (2, d) where d is the total number of detector pixels. Each column gives the fractional (x, y) index into the source Fourier grid.

  • A boolean mask of shape (d,) indicating which detector pixels fall within the source Fourier grid bounds.

  • The Jacobian correction factor of shape (d,) for energy density preservation under the coordinate transformation.

Return type:

A tuple of

ptyrax.models.propagation.tilt_target_to_incident(sampling, coordinate_system, target_sampling, target_coordinate_system, illumination_direction, wavlen)[source]#

Transform detector coordinates into the incident field’s Fourier frame.

Computes the scattering vector for each detector pixel by projecting detector positions onto the Ewald sphere and subtracting the illumination direction. The resulting scattering vectors are then rotated into the source field’s local coordinate frame.

The scattering vector is defined as:

\[\mathbf{q} = \hat{\mathbf{k}}_\text{out} - \hat{\mathbf{k}}_\text{in}\]

where \(\hat{\mathbf{k}}_\text{out}\) is the unit vector from sample to detector pixel (Ewald sphere projection) and \(\hat{\mathbf{k}}_\text{in}\) is the unit illumination direction.

Parameters:

  • sampling (SamplingGrid) – The sampling grid of the source field.

  • coordinate_system (CoordinateSystem) – The coordinate system of the source field.

  • target_sampling (SamplingGrid) – The sampling grid of the detector.

  • target_coordinate_system (CoordinateSystem) – The coordinate system of the detector.

  • illumination_direction (Float[Array, '3']) – The illumination wave vector direction as a 3D vector (not necessarily unit length).

  • wavlen (jaxtyping.Float) – The illumination wavelength in meters.

Returns:

  • The far-field sampling grid corresponding to the source field.

  • The (x, y) coordinates of each detector pixel in the source field’s Fourier frame, shape (2, d).

  • The Jacobian correction factor (z-component of Ewald sphere unit vectors) of shape (d,).

Return type:

A tuple of

Contents