ptyrax.models.illumination

ptyrax.models.illumination#

Functions

default_mode_initializer(mode_index, ...)

Create a default probe mode for multi-mode illumination models.

default_weight_initializer(shape)

Create default mode weights for an orthogonalized probe.

Classes

DirectIllumination(probe_data, *args[, ...])

Direct illumination model storing an explicit probe field.

IlluminationModel()

Abstract base class for illumination models in ptychographic reconstruction.

NamedLoss(tag, value)

A named loss value produced during regularization.

OrthogonalizedProbe(_, wavelength, ...[, ...])

Illumination model that represents a set of orthogonal probe modes with per-scan singular values.

PupilIllumination(*args, **kwargs)

Illumination model exposing the pupil representation of a DirectIllumination probe.
class ptyrax.models.illumination.DirectIllumination(probe_data, *args, n_scan=1, key=None, mode_data_or_initializer=None, parametrization_type=None, regularization_functions=(), **kwargs)[source]#

Bases: IlluminationModel

Direct illumination model storing an explicit probe field.

This is the simplest illumination model: it holds a single CoherentField as the probe and returns it unchanged on every call. Supports optional reparametrization of the underlying array and per-probe regularization functions.

Variables:

  • _probe – The stored coherent probe field.

  • regularization_functions – Tuple of callables applied to the probe during regularization.

Parameters:
__call__()[source]#

Return the stored probe field.

Returns:

The coherent probe field.

Return type:

CoherentField

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

Plot the probe field.

Delegates to CoherentField.__plot__().

Return type:

None

classmethod from_coherent_field(probe, **kwargs)[source]#

Construct a DirectIllumination from an existing CoherentField.

Parameters:

  • probe (CoherentField) – A fully constructed coherent field to use as the probe.

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

Returns:

A new DirectIllumination instance wrapping the given field.

Return type:

Self

property probe: CoherentField#

The coherent probe field stored by this illumination model.

regularization_functions: tuple[Callable[[DirectIllumination], TypeAliasForwardRef('jaxtyping.Float')], ...]#
class ptyrax.models.illumination.IlluminationModel[source]#

Bases: Module

Abstract base class for illumination models in ptychographic reconstruction.

An illumination model describes the probe beam that illuminates the sample. Subclasses must implement __call__() to produce the probe field and __regularize__() to compute any regularization losses on the probe parameters.

abstractmethod __call__()[source]#

Compute and return the probe illumination field.

Returns:

The coherent probe field used to illuminate the sample.

Return type:

CoherentField

class ptyrax.models.illumination.NamedLoss(tag, value)[source]#

Bases: Module

A named loss value produced during regularization.

Pairs a human-readable tag with a scalar loss value so that individual regularization contributions can be tracked and logged separately.

Variables:

  • tag – Identifier string for this loss contribution (e.g. "probe.smoothness.0").

  • value – Scalar loss value.

Parameters:

  • tag (str)

  • value (jaxtyping.Float)

tag: str#
value: float#
class ptyrax.models.illumination.OrthogonalizedProbe(_, wavelength, probe_grid, coordinates, n_scan, n_modes=1, mode_data_or_initializer=None, weight_initializer=<function default_weight_initializer>, regularization_functions=(), *, key)[source]#

Bases: IlluminationModel

Illumination model that represents a set of orthogonal probe modes with per-scan singular values.

The probe_modes are combined with singular_values to produce the effective probe for each scan index.

Parameters:

  • wavelength (float)

  • probe_grid (SamplingGrid)

  • coordinates (CoordinateSystem)

  • n_scan (jaxtyping.Integer)

  • n_modes (jaxtyping.Integer)

  • mode_data_or_initializer (bool)

  • weight_initializer (Callable[[tuple[int, ...]], np.ndarray])

  • regularization_functions (tuple[Callable[[CoherentField], TypeAliasForwardRef('jaxtyping.Float')], ...])

  • key (jax.random.PRNGKey)

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

Compute the effective probe by combining orthogonalized modes.

Blends the raw probe_modes with their QR-orthogonalized versions according to relaxed_orthogonalization, then contracts with the per-scan singular_values to produce the output probe field.

Returns:

A CoherentField representing the combined probe for the current scan index.

Return type:

CoherentField

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

Plot the orthogonalized probe modes weighted by their singular values.

Returns:

Tuple of (figure, subplot_spec, plotted_data) as expected by the plotting framework.

Return type:

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

n_modes: int = 1#
n_scan: int#
property orthogonal_probes: Float[Array, '* m n d']#

Orthogonalized probe modes obtained via QR decomposition.

Flattens the spatial dimensions of the stored probe modes, computes a QR factorization, and reshapes the Q factor back to the original spatial dimensions to yield an orthonormal set of probe modes.

Returns:

Array of orthogonalized probe mode data with the same shape as the raw probe modes.

probe_modes: PyTree[CoherentField]#
regularization_functions: tuple[Callable[[CoherentField], float], ...]#
relaxed_orthogonalization: float = 0.0#
singular_values: Array#
class ptyrax.models.illumination.PupilIllumination(*args, **kwargs)[source]#

Bases: DirectIllumination

Illumination model exposing the pupil representation of a DirectIllumination probe.

Computes the pupil via an inverse Fourier transform of the probe and provides methods to obtain propagated fields from the pupil representation.

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

Compute the probe field from the pupil via a forward Fourier transform.

Returns:

A CoherentField in the sample plane obtained by Fourier-transforming the stored pupil.

Return type:

CoherentField

__plot__(*args, **kwargs)#

Plot the probe field.

Delegates to CoherentField.__plot__().

Return type:

None

classmethod from_coherent_field(probe, **kwargs)#

Construct a DirectIllumination from an existing CoherentField.

Parameters:

  • probe (CoherentField) – A fully constructed coherent field to use as the probe.

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

Returns:

A new DirectIllumination instance wrapping the given field.

Return type:

Self

property probe: CoherentField#

Compute the probe field from the pupil, optionally with defocus propagation.

When propagation_distance is provided, a quadratic phase factor is applied to the pupil before Fourier-transforming, effectively propagating the probe by that distance.

Parameters:

propagation_distance – Optional scalar propagation distance. If None, no defocus is applied.

Returns:

A CoherentField representing the probe in the sample plane.

pupil: CoherentField#
regularization_functions: tuple[Callable[[DirectIllumination], TypeAliasForwardRef('jaxtyping.Float')], ...]#
ptyrax.models.illumination.default_mode_initializer(mode_index, sampling, max_index, wavelength, probe_grid, coordinates, key)[source]#

Create a default probe mode for multi-mode illumination models.

Generates an aperture-based probe mode whose radius scales linearly with mode_index and applies a random amplitude modulation combined with a quadratic phase ramp to break symmetry between modes.

Parameters:

  • mode_index (jaxtyping.Integer) – Zero-based index of the mode being initialized.

  • sampling (SamplingGrid) – Sampling grid describing the pixel spacing.

  • max_index (jaxtyping.Integer) – Total number of modes (used to scale the aperture radius).

  • wavelength (jaxtyping.Float) – Probe wavelength in consistent units.

  • probe_grid (SamplingGrid) – Sampling grid of the output probe field.

  • coordinates (CoordinateSystem) – Coordinate system for the output field.

  • key (PRNGKey) – JAX PRNG key for random amplitude generation.

Returns:

A CoherentField representing the initialized probe mode.

Return type:

CoherentField

ptyrax.models.illumination.default_weight_initializer(shape)[source]#

Create default mode weights for an orthogonalized probe.

Returns weights that decay as \(1 / (2k + 1)\) for uniform initialization across modes.

Parameters:

shape (tuple[TypeAliasForwardRef('jaxtyping.Integer'), ...]) – Shape of the weight array, typically (n_scan, n_modes).

Returns:

An array of mode weights with the given shape.

Return type:

Float[Array, ‘…’]

Contents