ptyrax.field

ptyrax.field#

Functions

defocus_phase_generator(old_field, ...)

Generate a quadratic defocus phase profile for a field.

initialize_new_probe(old_probe, ...[, key])

Create a new probe by reinitializing each mode with a given function.

multiply_field_phase(old_probe, ...)

Multiply the field by a phase-only factor.

plot_field(field[, show])

Plot a coherent field showing amplitude and phase.

remove_field_phase(old_probe)

Remove the phase of a field, keeping only the amplitude.

replace_defocus(old_field, defocus_amount[, ...])

Replace the field phase with a defocus along the specified axis or axes.

replace_defocus_x(old_field, defocus_amount)

Replace the field phase with a defocus along the x-axis only.

replace_defocus_xy(old_field, defocus_amount)

Replace the field phase with independent defocus along both axes.

replace_defocus_y(old_field, defocus_amount)

Replace the field phase with a defocus along the y-axis only.

replace_field_phase(old_probe, ...)

Replace the phase of a field with a new phase profile.

set_probe_wavelength(old_probe, wavelength)

Set the wavelength(s) of a probe field.

transpose_probe(old_probe)

Swap the spatial axes of a probe field.

Classes

CoherentField(data, wavelength, sampling[, ...])

Representation of a coherent field (probe or object) with sampling and wavelength metadata.
class ptyrax.field.CoherentField(data, wavelength, sampling, coordinate_system=<factory>, propagation_direction=<factory>, spatial_dims=(-2, -3), vector_dim=-1)[source]#

Bases: Module

Representation of a coherent field (probe or object) with sampling and wavelength metadata.

Parameters:
__call__()[source]#

Return the underlying field data, resolving any parametrization.

If the data is wrapped in an ArrayParametrization, this evaluates it to produce the raw array. Otherwise returns the data directly.

Returns:

The complex-valued field array.

Return type:

Float[Array, ‘* m n d’]

__plot__(show=True, gs=None, fig=None, **kwargs)[source]#

Visualize the field as amplitude and phase images.

Uses the first vector component (index 0 along the last axis) and shows amplitude/phase in the same plot using a bivariate colormap via plot().

Parameters:

  • show (bool) – Whether to call plt.show() after plotting.

  • gs (SubplotSpec) – Optional matplotlib SubplotSpec to draw into.

  • fig (Figure) – Optional matplotlib Figure to use.

  • **kwargs – Additional keyword arguments forwarded to plot().

Returns:

A tuple of (Figure, Axes, SubplotSpec).

Return type:

tuple[Figure, Axes, SubplotSpec]

property amplitude: Float[Array, '* m n d']#

Absolute value (amplitude) of the field, FFT-shifted if applicable.

coordinate_system: CoordinateSystem#
data: Float[Array, '* m n d'] | ArrayParametrization#
property intensity: Float[Array, '* m n d']#

\(|E|^2\).

The result is real-valued.

Type:

Squared magnitude (intensity) of the field

k_z(n_medium=1 + 0j)[source]#

Compute the z-component of the wave vector in the field’s internal coordinate system.

Parameters:

n_medium (complex)

Return type:

tuple[Float[Array, ‘m n’], Bool[Array, ‘m n’]]

multiply_fourier(factor)[source]#

Multiply the field by a Fourier-space factor.

The factor is applied to the first vector component in Fourier space, and the result is transformed back to real space.

Parameters:

factor (Complex[Array, 'm n']) – Complex array of shape (m, n) representing the Fourier space factor to apply.

Returns:

A new CoherentField with the modified data.

Return type:

CoherentField

multiply_real(factor)[source]#

Multiply the field by a real-valued factor in real space.

Parameters:

factor (Float[Array, 'm n'])

Return type:

CoherentField

propagate_fraunhofer(distance, inverse=False, fftshift=True)[source]#

Propagate the field to the far field using the Fraunhofer approximation.

Computes the Fourier transform (or inverse) of the first vector component and updates the sampling grid to far-field coordinates.

Parameters:

  • distance (jaxtyping.Float) – Propagation distance used to compute the far-field pixel size.

  • inverse (bool) – If True, use the inverse FFT (propagate back to near field).

  • fftshift (bool) – Whether to apply fftshift to the result.

Returns:

A new CoherentField with propagated data and far-field sampling.

Return type:

CoherentField

propagate_tilted_nearfield(displacement, n_medium=1 + 0j)[source]#

Propagate the field in the near-field regime along a tilted direction.

Performs angular-spectrum propagation by applying a phase ramp in Fourier space corresponding to a 3-D displacement vector. Evanescent components (those outside the Ewald sphere) are set to zero.

Parameters:

  • displacement (Float[Array, '3']) – 3-D displacement vector [dx, dy, dz] in real-space units. Only the z-component contributes to the propagation phase; the transverse components are currently ignored.

  • n_medium (complex)

Returns:

A new CoherentField with propagated data.

Return type:

CoherentField

propagation_direction: Array#
sampling: SamplingGrid#
property shape: tuple[TypeAliasForwardRef('jaxtyping.Integer'), ...]#

Shape of the underlying data array.

spatial_dims: tuple = (-2, -3)#
property to_fft_shifted: CoherentField#

Return an FFT-shifted version of this field.

Applies jnp.fft.fftshift along the spatial dimensions so that the zero-frequency component is centered. If the field is already shifted, returns self unchanged.

Returns:

A new CoherentField with shifted data and updated sampling metadata.

property to_fft_unshifted: CoherentField#

Return a non-FFT-shifted version of this field.

Applies jnp.fft.ifftshift along the spatial dimensions to restore the standard array ordering. If the field is already unshifted, returns self unchanged.

Returns:

A new CoherentField with unshifted data and updated sampling metadata.

vector_dim: int = -1#
wavelength: Float[Array, '']#
ptyrax.field.defocus_phase_generator(old_field, defocus_amount)[source]#

Generate a quadratic defocus phase profile for a field.

Uses gaussian() with a negative radius to produce a pure defocus (quadratic phase) without amplitude modulation.

Parameters:

  • old_field (CoherentField) – The field whose sampling grid defines the coordinate space.

  • defocus_amount (tuple[TypeAliasForwardRef('jaxtyping.Float'), TypeAliasForwardRef('jaxtyping.Float')]) – Tuple of (defocus_x, defocus_y) in meters.

Returns:

A complex-valued array of shape (n, m, 1) representing the defocus phase factor.

Return type:

Float[Array, ‘n m 1’]

ptyrax.field.initialize_new_probe(old_probe, new_initializer_functions, *, key=Array([0, 0], dtype=uint32))[source]#

Create a new probe by reinitializing each mode with a given function.

Each mode of old_probe is reinitialized using the corresponding initializer function. If a single initializer is provided it is broadcast to all modes.

Parameters:

  • old_probe (CoherentField) – The existing probe whose shape and metadata are preserved.

  • new_initializer_functions (tuple[Callable[[tuple[TypeAliasForwardRef('jaxtyping.Integer'), TypeAliasForwardRef('jaxtyping.Integer')]], CoherentField]]) – A tuple of callables, each accepting a shape tuple and a key keyword argument and returning an array for one mode. A single-element tuple is broadcast to all modes.

  • key (Key[Array, ''] | UInt32[Array, '2']) – JAX PRNG key used for random initialization.

Returns:

A new CoherentField with reinitialized data.

Raises:

ValueError – If the number of initializer functions does not match the number of probe modes (and is not exactly one).

Return type:

CoherentField

ptyrax.field.multiply_field_phase(old_probe, phase_or_generator)[source]#

Multiply the field by a phase-only factor.

Unlike replace_field_phase(), this adds phase on top of the existing field phase (and amplitude) by multiplying the field with a unit-amplitude complex factor.

Parameters:

  • old_probe (CoherentField) – The coherent field to modify.

  • phase_or_generator (Complex[Array, 'n m'] | Callable[[CoherentField], Complex[Array, 'n m']]) – Phase factor to multiply, provided as a complex array or as a callable that accepts a CoherentField and returns a complex array.

Returns:

A new CoherentField with the additional phase applied.

Return type:

CoherentField

ptyrax.field.plot_field(field, show=True, **kwargs)[source]#

Plot a coherent field showing amplitude and phase.

Extracts the first mode and first vector component, applies FFT-shift if needed, and delegates to plot().

Parameters:

  • field (CoherentField) – The coherent field to visualize.

  • show (bool) – Whether to call plt.show() after plotting.

  • **kwargs – Additional keyword arguments forwarded to plot().

Returns:

A tuple of (Figure, SubplotSpec, list of AxesImage).

Return type:

tuple[Figure, SubplotSpec, list[AxesImage]]

ptyrax.field.remove_field_phase(old_probe)[source]#

Remove the phase of a field, keeping only the amplitude.

Parameters:

old_probe (CoherentField) – The coherent field whose phase is to be removed.

Returns:

A new CoherentField with real-valued (amplitude-only) data.

Return type:

CoherentField

ptyrax.field.replace_defocus(old_field, defocus_amount, axis='xy')[source]#

Replace the field phase with a defocus along the specified axis or axes.

Parameters:

  • old_field (CoherentField) – The coherent field to modify.

  • defocus_amount (float | tuple[float, float]) – Defocus strength in meters. A single float for single-axis defocus, or a tuple (defocus_x, defocus_y) for axis="xy".

  • axis (str) – Which axis to defocus: "x", "y", or "xy".

Returns:

A new CoherentField with the specified defocus phase.

Raises:

ValueError – If axis is not one of "x", "y", or "xy".

Return type:

CoherentField

ptyrax.field.replace_defocus_x(old_field, defocus_amount)[source]#

Replace the field phase with a defocus along the x-axis only.

Parameters:
Return type:

CoherentField

ptyrax.field.replace_defocus_xy(old_field, defocus_amount)[source]#

Replace the field phase with independent defocus along both axes.

Parameters:

  • old_field (CoherentField)

  • defocus_amount (tuple[TypeAliasForwardRef('jaxtyping.Float'), TypeAliasForwardRef('jaxtyping.Float')])

Return type:

CoherentField

ptyrax.field.replace_defocus_y(old_field, defocus_amount)[source]#

Replace the field phase with a defocus along the y-axis only.

Parameters:
Return type:

CoherentField

ptyrax.field.replace_field_phase(old_probe, phase_or_generator)[source]#

Replace the phase of a field with a new phase profile.

The amplitude of the original field is preserved while its phase is replaced by the phase of the provided array, callable, or field.

Parameters:
Returns:

A new CoherentField with the original amplitude and the new phase.

Return type:

CoherentField

ptyrax.field.set_probe_wavelength(old_probe, wavelength)[source]#

Set the wavelength(s) of a probe field.

If a single float is given it is broadcast to all modes. An array of wavelengths must match the number of probe modes.

Parameters:

  • old_probe (CoherentField) – The probe field to update.

  • wavelength (float | list | tuple | np.ndarray | jnp.ndarray) – Wavelength value(s) in meters. Either a single float (applied to all modes) or a sequence matching the mode count.

Returns:

A new CoherentField with updated wavelength.

Raises:

ValueError – If the number of provided wavelengths does not match the number of probe modes.

Return type:

CoherentField

ptyrax.field.transpose_probe(old_probe)[source]#

Swap the spatial axes of a probe field.

Transposes the two spatial dimensions and interpolates onto a new SamplingGrid with swapped shape and pixel size.

Parameters:

old_probe (CoherentField) – The probe field to transpose.

Returns:

A new CoherentField with transposed spatial dimensions.

Return type:

CoherentField

Contents