ptyrax.models.interaction

ptyrax.models.interaction#

Functions

to_single_slice_approximation(multislice[, ...])

Convert a multi-slice model to a single-slice approximation.

Classes

FresnelReflection(coordinates, sampling[, ...])

Thin-object Fresnel reflection interaction model.

InteractionModel(coordinates, sampling[, ...])

Abstract base class for interaction models in ptychographic reconstruction.

MultiSlice(coordinates, sampling, ...[, ...])

Multi-slice interaction model for thick samples.
class ptyrax.models.interaction.FresnelReflection(coordinates, sampling, forward_sampling=None, initializer=<function uniform>, parametrization_type=None, regularization_functions=(), normalize=True)[source]#

Bases: InteractionModel

Thin-object Fresnel reflection interaction model.

Models the interaction of the illumination beam with a thin planar sample using a multiplicative approximation. The exit field is computed as:

$$ psi_{text{exit}}(mathbf{r}) = O(mathbf{r}) cdot P(mathbf{r}) $$

where $O$ is the complex reflection coefficient of the sample and $P$ is the probe illumination. The reflected field propagation direction is computed by reflecting the incident direction about the sample surface normal.

Variables:

  • coordinates – Sample coordinate system with position and rotation for each scan point, wrapped as an IndexSliceParameter.

  • surface_normal – Unit vector normal to the sample surface in the sample’s local coordinate frame.

  • reflection_coefficient – Complex-valued 2D array representing the sample’s spatially-resolved reflection coefficient.

  • sampling – Pixel grid defining the sample discretization.

  • forward_sampling – Optional separate grid for the forward model computation. If None, uses the illumination field’s sampling.

  • regularization_functions – Tuple of callables that compute scalar regularization penalties from the reflection coefficient.

  • normalize – Whether to apply normalization by the square root of the number of pixels.

Parameters:
__call__(illumination_field, normalize=True)[source]#

Compute the thin-object multiplicative interaction.

Interpolates the sample reflection coefficient onto the illumination grid, multiplies it element-wise with the probe, and constructs the reflected output field with updated propagation direction.

Parameters:

  • illumination_field (CoherentField) – Incoming coherent illumination field.

  • normalize (bool) – Whether to apply pixel-count normalization to the interpolated object. If the original array is normalized, this ensures the output remains approximally normalized after interpolation.

Returns:

The reflected coherent field after sample interaction.

Return type:

CoherentField

__plot__(*args, **kwargs)[source]#

Plot the reflection coefficient as a complex-valued image.

Parameters:

  • *args – Positional arguments forwarded to plot().

  • **kwargs – Keyword arguments forwarded to plot().

Returns:

A tuple of (figure, subplot_spec, plotted_array) as produced by the plotting utility.

Return type:

tuple[Figure, SubplotSpec, Array | list[Array]]

coordinates: IndexSliceParameter[CoordinateSystem]#
forward_sampling: SamplingGrid#
normalize: bool = True#
reflection_coefficient: Shaped[Array, '* m n']#
regularization_functions: tuple[Callable[[Complex[Array, '* m n']], float]]#
sampling: SamplingGrid#
property shape: tuple[TypeAliasForwardRef('jaxtyping.Integer'), TypeAliasForwardRef('jaxtyping.Integer')]#

Spatial shape (m, n) of the sample reflection coefficient array.

surface_normal: Float[Array, '3']#
class ptyrax.models.interaction.InteractionModel(coordinates, sampling, initializer=None)[source]#

Bases: Module

Abstract base class for interaction models in ptychographic reconstruction.

An interaction model describes how an incoming coherent illumination field interacts with a sample to produce an outgoing field. Subclasses implement specific physical interaction mechanisms such as thin-object Fresnel reflection or multi-slice propagation.

All interaction models are JAX-compatible Equinox modules that can be differentiated through and used in JIT-compiled pipelines.

Variables:

coordinates – The coordinate system defining the sample position and orientation in the global reference frame.

Parameters:
abstractmethod __call__(input_field)[source]#

Compute the interaction of an illumination field with the sample.

Parameters:

input_field (CoherentField) – The incoming coherent illumination field incident on the sample.

Returns:

The outgoing coherent field after interaction with the sample, including updated propagation direction and coordinate system.

Return type:

CoherentField

coordinates: CoordinateSystem#
class ptyrax.models.interaction.MultiSlice(coordinates, sampling, slice_distances, initializer=<function uniform>, parametrization_type=<class 'ptyrax.parametrizations.DirectArrayParametrization'>, regularization_functions=(), inner_interactions=None, **kwargs)[source]#

Bases: InteractionModel

Multi-slice interaction model for thick samples.

Divides the sample into multiple axial slices, each modeled as an independent FresnelReflection interaction. The illumination is near-field propagated to each slice depth, interacted, and propagated back. The total exit field is the coherent sum of contributions from all slices:

\[\psi_{\text{exit}} = \sum_{j=1}^{N} \mathcal{P}(z_j)\left[ O_j \cdot \mathcal{P}(z_j)[P] \right]\]

where $mathcal{P}(z)$ denotes near-field propagation by distance $z$ and $O_j$ is the reflection coefficient of the $j$-th slice.

Variables:

  • coordinates – Sample coordinate system shared across all slices.

  • inner_interactions – Vmapped stack of per-slice interaction models.

  • slice_distances – Array of axial distances from the reference plane to each slice.

  • separable_in_z – If True, uses a single shared interaction model for all slices (only varying the propagation distance).

  • inverted_bottom – If True, inverts the ordering for the bottom half of the slices.

Parameters:
__call__(illumination_field, normalize=True)[source]#

Compute the multi-slice interaction as a coherent sum over slices.

Propagates the illumination to each slice depth, computes the per-slice interaction, propagates back, and sums all contributions coherently.

Parameters:

  • illumination_field (CoherentField) – Incoming coherent illumination field.

  • normalize (bool) – Whether to apply normalization in the per-slice interactions.

Returns:

The total reflected coherent field (coherent sum of all slices).

Return type:

CoherentField

coordinates: IndexSliceParameter[CoordinateSystem]#
forward_fields(illumination_field)[source]#

Compute per-slice reflected fields using independent interactions.

Each slice uses its own interaction model. This is the forward pass when separable_in_z=False.

Parameters:

illumination_field (CoherentField) – Incoming coherent illumination field.

Returns:

A vmapped stack of reflected fields, one per slice.

Return type:

CoherentField

classmethod from_interactions(interactions, distances, **kwargs)[source]#

Construct a MultiSlice from pre-built per-slice interaction models.

Parameters:

  • interactions (tuple[InteractionModel]) – Tuple of interaction models, one per slice.

  • distances (Float[Array, 'n_slices']) – Array of axial slice distances.

  • **kwargs – Additional keyword arguments forwarded to the constructor (e.g., separable_in_z, inverted_bottom).

Returns:

A new MultiSlice instance wrapping the provided interactions.

Return type:

Type[MultiSlice]

inner_interactions: FresnelReflection#
inverted_bottom: bool = False#
property n_slices: jaxtyping.Integer#

Number of axial slices in the model.

separable_forward_fields(illumination_field)[source]#

Compute per-slice reflected fields using a single shared interaction.

Uses only the first slice’s interaction model for all depths, varying only the propagation distance. This is the forward pass when separable_in_z=True.

Parameters:

illumination_field (CoherentField) – Incoming coherent illumination field.

Returns:

A vmapped stack of reflected fields, one per slice.

Return type:

CoherentField

separable_in_z: bool = False#
static single_slice_forward(illumination_field, z_displacement, interaction)[source]#

Forward model for a single slice: propagate, interact, propagate back.

Parameters:

  • illumination_field (CoherentField) – Incoming coherent illumination field.

  • z_displacement (jaxtyping.Float) – Axial distance from the reference plane to this slice (in meters).

  • interaction (InteractionModel) – The interaction model for this slice.

Returns:

The back-reflected field after round-trip propagation and interaction with the slice.

Return type:

CoherentField

slice_distances: Float[Array, 'n_slices']#
ptyrax.models.interaction.to_single_slice_approximation(multislice, tilt_angle=0.0, wavelength=1.0, interaction_type=<class 'ptyrax.models.interaction.FresnelReflection'>)[source]#

Convert a multi-slice model to a single-slice approximation.

Collapses the multi-slice reflection coefficients into a single effective reflection coefficient by coherently summing the per-slice contributions with depth-dependent phase factors:

\[O_{\text{eff}}(\mathbf{r}) = \sum_j O_j(\mathbf{r}) \exp\!\left(i \frac{4\pi}{\lambda} z_j \cos\theta\right)\]

where $z_j$ is the slice distance and $theta$ is the tilt angle.

Parameters:

  • multislice (MultiSlice) – The multi-slice interaction model to collapse.

  • tilt_angle (jaxtyping.Float) – Beam tilt angle in degrees relative to the surface normal. Affects the depth-dependent phase factor.

  • wavelength (jaxtyping.Float) – Illumination wavelength (in meters).

  • interaction_type (type[InteractionModel]) – Class of the output single-slice interaction model.

Returns:

A single-slice interaction model with the coherently summed reflection coefficient.

Raises:

TypeError – If multislice is not a MultiSlice instance.

Return type:

FresnelReflection

Contents