janssen.prop

Propagation methods for optical wavefronts.

Extended Summary

Various propagation algorithms for simulating optical field propagation through different media and materials. Includes lens-based propagation, material-based propagation, free-space propagation, and high-NA vector focusing methods.

Routine Listings

angular_spectrum_prop()

Angular spectrum propagation method (no paraxial approximation).

aplanatic_apodization()

Apply sqrt(cos(theta)) apodization for aplanatic lens systems.

compute_focal_volume()

Compute 3D focal volume at multiple z planes.

correct_propagator()

Automatically selects the most appropriate propagation method.

debye_wolf_focus()

Compute focal field using Debye-Wolf formulation.

digital_zoom()

Digital zoom transformation for optical fields.

fraunhofer_prop()

Fraunhofer (far-field) propagation.

fraunhofer_prop_scaled()

Fraunhofer propagation with output at specified pixel size.

fresnel_prop()

Fresnel (near-field) propagation.

high_na_focus()

Compute focal field using Richards-Wolf vector diffraction integrals.

lens_propagation()

Propagate optical wavefront through a lens.

multislice_propagation()

Propagate optical wavefront through a 3D material.

optical_path_length()

Compute the optical path length through a material.

optical_zoom()

Optical zoom transformation.

scalar_focus_for_comparison()

Compute scalar focal field for comparison with vector result.

total_transmit()

Compute the total transmission through a material.

Notes

All propagation functions are JAX-compatible and support automatic differentiation. The module is designed to be extensible for new propagation methods.

janssen.prop.angular_spectrum_prop(incoming: OpticalWavefront, z_move: int | float | complex | Num[Array, ''], refractive_index: int | float | complex | Num[Array, ''] | None = 1.0) OpticalWavefront[source]

Propagate a complex field using the angular spectrum method.

Parameters:
  • incoming (OpticalWavefront) –

    PyTree with the following parameters:

    fieldComplex[Array, “ hh ww”]

    Input complex field

    wavelengthFloat[Array, “ “]

    Wavelength of light in meters

    dxFloat[Array, “ “]

    Grid spacing in meters

    z_positionFloat[Array, “ “]

    Wave front position in meters

  • z_move (numeric) – Propagation distance in meters This is in free space.

  • refractive_index (Optional[ScalarNumeric], optional) – Index of refraction of the medium. Default is 1.0 (vacuum).

Returns:

propagated – Propagated wave front

Return type:

OpticalWavefront

Notes

The angular spectrum method is an exact solution to the Helmholtz equation for propagation in homogeneous media. It decomposes the field into plane waves, propagates each component, and reconstructs the field.

The transfer function is H(fx,fy) = exp(i*k*z*sqrt(1 - (lambda*fx)^2 - (lambda*fy)^2)) where k = 2*pi/lambda is the wavenumber. For spatial frequencies where (lambda*fx)^2 + (lambda*fy)^2 > 1, the waves become evanescent and are set to zero to prevent numerical instability.

This method makes no paraxial approximation and is valid for all propagation distances, though numerical accuracy may degrade for very large distances due to sampling limitations.

Algorithm:

  1. Compute spatial frequency grids fx, fy using FFT frequencies

  2. Build transfer function H = exp(i*k*z*sqrt(1 - lambda^2*(fx^2+fy^2)))

  3. Create evanescent mask where fx^2 + fy^2 <= 1/lambda^2

  4. FFT input field, multiply by masked transfer function, inverse FFT

  5. Return propagated wavefront with updated z_position

janssen.prop.correct_propagator(incoming: OpticalWavefront, z_move: int | float | complex | Num[Array, ''], refractive_index: int | float | complex | Num[Array, ''] | None = 1.0) OpticalWavefront[source]

Automatically select and apply the most appropriate propagator.

This function selects the optimal propagation method based on the Fresnel number and sampling criteria. It uses: - Angular spectrum for very short distances or high spatial frequencies - Fresnel propagation for intermediate distances - Fraunhofer propagation for far-field distances

Parameters:
  • incoming (OpticalWavefront) –

    PyTree with the following parameters: field : Complex[Array, “ hh ww”]

    Input complex field

    wavelengthFloat[Array, “ “]

    Wavelength of light in meters

    dxFloat[Array, “ “]

    Grid spacing in meters

    z_positionFloat[Array, “ “]

    Wave front position in meters

  • z_move (numeric) – Propagation distance in meters (in free space)

  • refractive_index (Optional[ScalarNumeric], optional) – Index of refraction of the medium. Default is 1.0 (vacuum)

Returns:

propagated – Propagated wave front using the most appropriate method

Return type:

OpticalWavefront

Notes

Implementation: 1. Get field dimensions (ny, nx) 2. Calculate field intensity distribution 3. Create coordinate arrays centered at field center 4. Calculate RMS width in both x and y directions 5. Use larger RMS width times 2 as characteristic aperture size 6. Account for refractive index in path length calculation 7. Calculate Fresnel number: F = a²/(λz) 8. Check angular spectrum validity criterion: dx < 0.5 * z * λ / L

where L is the field size

  1. Use nested jax.lax.cond to select propagator: - If F > 1.0 AND angular spectrum valid: use angular spectrum - Else if F > 0.1: use Fresnel propagation - Else: use Fraunhofer propagation (far-field)

Selection criteria: - Angular spectrum: F > 1 and sampling valid (most accurate, no

paraxial approximation)

  • Fresnel: 0.1 < F ≤ 1 (near to intermediate field)

  • Fraunhofer: F < 0.1 (far-field)

The angular spectrum method is preferred when applicable as it makes no paraxial approximations.

janssen.prop.digital_zoom(wavefront: OpticalWavefront, zoom_factor: int | float | complex | Num[Array, '']) OpticalWavefront[source]

Zoom an optical wavefront by a specified factor.

Key is this returns the same sized array as the original wavefront.

Parameters:
  • wavefront (OpticalWavefront) – Incoming optical wavefront.

  • zoom_factor (numeric) – Zoom factor (greater than 1 to zoom in, less than 1 to zoom out).

Returns:

zoomed_wavefront – Zoomed optical wavefront of the same spatial dimensions.

Return type:

OpticalWavefront

Notes

Algorithm:

For zoom in (zoom_factor >= 1.0): - Calculate the crop fraction (1 / zoom_factor) to determine the

central region to extract

  • Create interpolation coordinates for the zoomed region centered

    on the image

  • Use scipy.ndimage.map_coordinates with bilinear interpolation

    to sample the field

  • Return the zoomed field with adjusted pixel size (dx /

zoom_factor)

For zoom out (zoom_factor < 1.0): - Calculate the shrink fraction (zoom_factor) to determine the

final image size

  • Create a coordinate mapping from the full image to the shrunken

region - Use scipy.ndimage.map_coordinates to interpolate the original field - Apply a mask to zero out regions outside the shrunken

area (padding effect)

  • Return the zoomed field with adjusted pixel size (dx /

zoom_factor)

janssen.prop.fraunhofer_prop(incoming: OpticalWavefront, z_move: int | float | complex | Num[Array, ''], refractive_index: int | float | complex | Num[Array, ''] = 1.0) OpticalWavefront[source]

Propagate a complex field using the Fraunhofer approximation.

Parameters:
  • incoming (OpticalWavefront) –

    PyTree with the following parameters:

    fieldComplex[Array, “ hh ww”]

    Input complex field

    wavelengthFloat[Array, “ “]

    Wavelength of light in meters

    dxFloat[Array, “ “]

    Grid spacing in meters

    z_positionFloat[Array, “ “]

    Wave front position in meters

  • z_move (numeric) – Propagation distance in meters. This is in free space.

  • refractive_index (numeric, optional) – Index of refraction of the medium. Default is 1.0 (vacuum).

Returns:

propagated – Propagated wave front. Note that the output pixel size (dx) changes according to Fraunhofer scaling: dx_out = wavelength * z / (N * dx_in)

Return type:

OpticalWavefront

Notes

The Fraunhofer approximation represents far-field diffraction where the diffraction pattern is proportional to the Fourier transform of the aperture function. This is the limiting case of Fresnel diffraction when the propagation distance is very large.

The full Fraunhofer diffraction integral gives: U(x’,y’) = exp(i*k*z)/(i*lambda*z) * exp(i*k*(x’^2+y’^2)/(2*z)) *

FT{U(x,y)} * dx^2

where FT denotes the Fourier transform and the output coordinates are related to spatial frequencies by x’ = lambda*z*fx, y’ = lambda*z*fy.

The quadratic phase term exp(i*k*(x’^2+y’^2)/(2*z)) represents the spherical wavefront curvature in the output plane and is essential for coherent imaging and phase-sensitive applications.

The output pixel size changes according to the Fraunhofer scaling relation: dx_out = lambda * z / (N * dx_in), where N is the array size.

The Fraunhofer approximation is valid when the Fresnel number F = a^2/(λz) is small (typically F < 1), where a is the characteristic aperture size. For large Fresnel numbers, use fresnel_prop or angular_spectrum_prop.

Algorithm

  1. Compute centered FFT of input field using ifftshift/fft2/fftshift

  2. Compute output pixel size dx_out = lambda*z/(N*dx_in)

  3. Create output coordinate grid and compute quadratic phase

  4. Apply global phase factor exp(i*k*z)

  5. Apply amplitude scaling 1/(i*lambda*z) and area element dx^2

  6. Multiply by quadratic phase term

  7. Return propagated wavefront with new dx and updated z_position

janssen.prop.fraunhofer_prop_scaled(incoming: OpticalWavefront, z_move: int | float | complex | Num[Array, ''], output_dx: float | Float[Array, ''], refractive_index: int | float | complex | Num[Array, ''] = 1.0) OpticalWavefront[source]

Propagate using Fraunhofer with output at specified pixel size.

Performs Fraunhofer propagation with the output sampled at the desired pixel size. The output array has the same shape as the input.

Parameters:
  • incoming (OpticalWavefront) – Input optical wavefront.

  • z_move (numeric) – Propagation distance in meters.

  • output_dx (float) – Desired output pixel size in meters.

  • refractive_index (numeric, optional) – Index of refraction. Default is 1.0 (vacuum).

Returns:

propagated – Propagated wavefront with specified output pixel size. Output array shape equals input array shape.

Return type:

OpticalWavefront

Notes

The standard Fraunhofer relation is:

U(x’) = C * FT{U(x)} evaluated at fx = x’/(λz)

where C includes phase and amplitude factors.

For output pixel n (centered), we want x’_n = (n - N/2) * output_dx. This corresponds to spatial frequency fx_n = x’_n/(λz).

We achieve this by using a chirp-z transform (CZT) approach: instead of FFT which samples at fx = n/(N*dx_in), we use interpolation in Fourier space to sample at the desired frequencies.

The key insight is that the DFT samples at frequencies:

fx_fft[n] = (n - N/2) / (N * dx_in)

And we want to sample at:

fx_out[n] = (n - N/2) * output_dx / (λz)

The ratio is: fx_out / fx_fft = output_dx * N * dx_in / (λz)

This is equivalent to scaling the output coordinates, which we implement by interpolating the FFT result.

janssen.prop.fresnel_prop(incoming: OpticalWavefront, z_move: int | float | complex | Num[Array, ''], refractive_index: int | float | complex | Num[Array, ''] | None = 1.0) OpticalWavefront[source]

Propagate a complex field using the Fresnel approximation.

Parameters:
  • incoming (OpticalWavefront) –

    PyTree with the following parameters: field : Complex[Array, “ hh ww”]

    Input complex field

    wavelengthFloat[Array, “ “]

    Wavelength of light in meters

    dxFloat[Array, “ “]

    Grid spacing in meters

    z_positionFloat[Array, “ “]

    Wave front position in meters

  • z_move (numeric) – Propagation distance in meters This is in free space.

  • refractive_index (Optional[ScalarNumeric], optional) – Index of refraction of the medium. Default is 1.0 (vacuum).

Returns:

propagated – Propagated wave front

Return type:

OpticalWavefront

Notes

The Fresnel approximation is a paraxial approximation to scalar diffraction theory. It assumes that the propagation angle is small, which allows simplification of the angular spectrum transfer function using a Taylor expansion: sqrt(1 - lambda^2*(fx^2+fy^2)) ≈ 1 - lambda^2*(fx^2+fy^2)/2.

This leads to the Fresnel transfer function: H(fx,fy) = exp(-i*pi*lambda*z*(fx^2 + fy^2))

The output field is multiplied by a global phase factor exp(i*k*z) representing the on-axis propagation phase.

The Fresnel approximation is valid when the Fresnel number F = a^2/(λz) is large (typically F > 1), where a is the characteristic aperture size. For small Fresnel numbers, use fraunhofer_prop instead.

Algorithm:

  1. Compute spatial frequency grids fx, fy using FFT frequencies

  2. Build Fresnel transfer function H = exp(-i*pi*lambda*z*(fx^2+fy^2))

  3. FFT input field, multiply by transfer function, inverse FFT

  4. Multiply result by global phase exp(i*k*z)

  5. Return propagated wavefront with updated z_position

janssen.prop.lens_propagation(incoming: OpticalWavefront, lens: LensParams) OpticalWavefront[source]

Propagate an optical wavefront through a lens.

The lens is modeled as a thin lens with a given focal length and diameter.

Parameters:
  • incoming (OpticalWavefront) – The incoming optical wavefront

  • lens (LensParams) – The lens parameters including focal length and diameter

Returns:

outgoing – The propagated optical wavefront after passing through the lens

Return type:

OpticalWavefront

Notes

Algorithm:

  • Create a meshgrid of coordinates based on the incoming wavefront’s

    shape and pixel size.

  • Calculate the phase profile and transmission function of the lens.

  • Apply the phase screen to the incoming wavefront’s field.

  • Return the new optical wavefront with the updated field,

wavelength,

and pixel size.

janssen.prop.optical_zoom(wavefront: OpticalWavefront, zoom_factor: int | float | complex | Num[Array, '']) OpticalWavefront[source]

Modify the calibration of an optical wavefront without changing field.

Parameters:
  • wavefront (OpticalWavefront) – Incoming optical wavefront.

  • zoom_factor (numeric) – Zoom factor (greater than 1 to zoom in, less than 1 to zoom out).

Returns:

zoomed_wavefront – Zoomed optical wavefront of the same spatial dimensions.

Return type:

OpticalWavefront

janssen.prop.multislice_propagation(incoming: OpticalWavefront, material: SlicedMaterialFunction) OpticalWavefront[source]

Propagate optical wavefront through 3D material using multi-slice BPM.

Uses the split-step beam propagation method to simulate light propagation through a 3D material with spatially-varying complex refractive index. The material is divided into thin slices, and propagation alternates between material interaction and free-space propagation.

Parameters:
  • incoming (OpticalWavefront) – Input optical wavefront at the entrance of the material. If incoming.dx differs from material.dx, the wavefront will be automatically resampled to the smaller dx value for better resolution.

  • material (SlicedMaterialFunction) – 3D material with complex refractive index at each voxel. material.material[i, j, k] = n(i,j,k) + i*κ(i,j,k) where n is refractive index and κ is extinction coefficient.

Returns:

outgoing – Optical wavefront at the exit of the material, after propagating through all slices. The z_position is updated to reflect the total propagation distance.

Return type:

OpticalWavefront

Notes

Implementation Details:

Spatial Sampling Handling: If incoming.dx differs from material.dx, the incoming wavefront is automatically resampled using optical_zoom to match the smaller dx (better resolution). Tolerance for dx matching is 1e-10.

Algorithm: 1. Resample incoming wavefront to target dx if needed 2. Calculate wavenumber: k = 2π/λ 3. Initialize propagating field and z position 4. For each slice in material:

  1. Extract complex refractive index: n + iκ

  2. Compute phase shift: φ = k × (n - 1) × tz

  3. Compute absorption coefficient: α = 4πκ/λ

  4. Compute absorption: A = exp(-α × tz)

  5. Apply material interaction: field *= A × exp(iφ)

  6. Create wavefront at current z position

  7. Propagate to next slice using correct_propagator

  8. Update z position: z += tz

  1. Return final wavefront with updated field and z position

Propagation Method: The propagation method between slices is automatically selected by correct_propagator based on the Fresnel number and propagation distance. No manual method selection is required.

Multi-slice Approximation Validity: - Slice thickness << characteristic propagation distances - Paraxial approximation holds (small angles) - Refractive index varies slowly within each slice

Material Interaction Formulas: - Phase shift: φ = (2π/λ) × (n - 1) × tz - Absorption coefficient: α = 4πκ/λ - Absorption: A = exp(-α × tz) - Combined transmission: T = A × exp(iφ)

Warning

  • For highly absorbing materials (large κ), numerical precision may be reduced

  • Very thick slices (tz >> λ) may violate the thin-slice approximation

Examples

Propagate through a uniform glass slab:

>>> import jax.numpy as jnp
>>> from janssen.types import make_optical_wavefront,
>>>                           make_sliced_material_function
>>> from janssen.prop import multislice_propagation
>>>
>>> # Create input wavefront
>>> field = jnp.ones((128, 128), dtype=jnp.complex128)
>>> wavefront = make_optical_wavefront(
...     field=field, wavelength=550e-9, dx=1e-6, z_position=0.0
... )
>>>
>>> # Create glass material (n=1.5, no absorption)
>>> material_array = jnp.ones((128, 128, 20)) * (1.5 + 0.0j)
>>> material = make_sliced_material_function(
...     material=material_array, dx=1e-6, tz=5e-6
... )
>>>
>>> # Propagate through material
>>> output = multislice_propagation(wavefront, material)
>>> print(f"Total propagation: {output.z_position:.2e} m")

See also

correct_propagator

Automatic selection of propagation method

optical_zoom

Resampling function

janssen.prop.optical_path_length(material: SlicedMaterialFunction, x_idx: int | Int[Array, ''] | None = -1, y_idx: int | Int[Array, ''] | None = -1) Float[Array, 'H W'][source]

Compute optical path length through material.

Calculates the optical path length (OPL) along rays through the material. OPL accounts for both the physical distance and the refractive index.

Parameters:
  • material (SlicedMaterialFunction) – 3D material with complex refractive index.

  • x_idx (int, optional) – X-index for specific ray. If -1 (default), computes for all x. When specified, the result is tiled to maintain (H, W) shape.

  • y_idx (int, optional) – Y-index for specific ray. If -1 (default), computes for all y. When specified, the result is tiled to maintain (H, W) shape.

Returns:

opl – Optical path length in meters as a 2D array (H, W). When indices are specified, the result is broadcast/tiled: - Both indices: scalar value tiled to (H, W) - x_idx only: 1D result tiled along y to (H, W) - y_idx only: 1D result tiled along x to (H, W) - Neither: full 2D projection (H, W)

Users can extract specific values with indexing, e.g., opl[0, 0] for scalar, opl[0, :] for x-line, opl[:, 0] for y-line.

Return type:

Float[Array, " H W"]

Notes

Implementation: 1. Extract real part of material (refractive index n) 2. Get material dimensions (height, width) 3. Use nested jax.lax.cond to select computation:

  • All branches return same shape (H, W) for JIT compatibility

  • Both indices: broadcast scalar to (H, W)

  • Only x_idx: tile 1D result to (H, W)

  • Only y_idx: tile 1D result to (H, W)

  • Neither - compute 2D projection directly.

  1. Extract appropriate slice/element from 2D result

  2. Sum refractive indices along z and multiply by slice thickness

Uses jax.lax.cond with uniform output shapes, then extracts the relevant data. This ensures JIT compilation compatibility while supporting dynamic return shapes.

Formula: OPL = Σ_z n(x, y, z) × tz

This is the phase delay in distance units. To convert to phase: φ = (2π/λ) × OPL

Examples

>>> # Compute OPL for entire material
>>> opl_map = optical_path_length(material)
>>> # Compute OPL along center ray
>>> center_opl = optical_path_length(material, 64, 64)
janssen.prop.total_transmit(material: SlicedMaterialFunction, wavelength: float | Float[Array, '']) Float[Array, 'H W'][source]

Compute intensity transmission through material.

Calculates the total intensity transmission (squared amplitude) for light passing through the material, accounting for absorption at each slice.

Parameters:
  • material (SlicedMaterialFunction) – 3D material with complex refractive index.

  • wavelength (float) – Wavelength of light in meters.

Returns:

transmission – Intensity transmission map (0 to 1). Values < 1 indicate absorption.

Return type:

Float[Array, " H W"]

Notes

Implementation: 1. Convert wavelength to JAX array 2. Extract extinction coefficient κ (imaginary part of material) 3. Compute absorption coefficient: α = 4πκ/λ 4. Compute amplitude transmission per slice: exp(-α × tz) 5. Compute total amplitude transmission (product over all slices) 6. Compute intensity transmission: |amplitude|²

Formula: For each position (x, y):

T(x,y) = Π_z exp(-2 × α(x,y,z) × tz) where α = 4πκ/λ

The factor of 2 in the exponent comes from intensity = |amplitude|²

Examples

>>> transmission = total_transmit(material, 550e-9)
>>> absorption_percent = (1 - transmission) * 100
>>> print(f"Max absorption: {jnp.max(absorption_percent):.1f}%")
janssen.prop.aplanatic_apodization(cos_theta: Float[Array, 'ny nx']) Float[Array, 'ny nx'][source]

Apply aplanatic lens apodization factor.

For an aplanatic (sine-condition satisfying) lens, the amplitude apodization is √cos(θ) where θ is the convergence angle.

Parameters:

cos_theta (Float[Array, " ny nx"]) – Cosine of convergence angle at each pupil point.

Returns:

apodization – Apodization factor √cos(θ).

Return type:

Float[Array, " ny nx"]

Notes

The √cos(θ) factor arises from energy conservation when mapping a uniform pupil plane wave to converging spherical wavefronts. This is the standard apodization for microscope objectives.

janssen.prop.compute_focal_volume(pupil_field: OpticalWavefront, na: float | Float[Array, ''], focal_length: float | Float[Array, ''], z_positions: Float[Array, 'nz'], refractive_index: float | Float[Array, ''] = 1.0) tuple[Complex[Array, 'nz ny nx 3'], Float[Array, '']][source]

Compute 3D focal volume at multiple z planes.

Uses vmap to efficiently compute the focal field at multiple axial positions.

Parameters:
  • pupil_field (OpticalWavefront) – Input polarized field in pupil plane.

  • na (float) – Numerical aperture.

  • focal_length (float) – Focal length in meters.

  • z_positions (Float[Array, " nz"]) – Array of axial positions relative to focus.

  • refractive_index (float, optional) – Refractive index of focal medium. Default is 1.0.

Return type:

tuple[Complex[Array, 'nz ny nx 3'], Float[Array, '']]

Returns:

  • focal_volume (Complex[Array, " nz ny nx 3"]) – 3D vector field volume with [Ex, Ey, Ez] at each z.

  • dx_focal (Float[Array, " "]) – Transverse pixel size at focal plane.

Examples

>>> z_range = jnp.linspace(-2e-6, 2e-6, 41)  # ±2 μm
>>> volume, dx = compute_focal_volume(
...     pupil, na=0.9, focal_length=3e-3, z_positions=z_range
... )
>>> # volume has shape (41, ny, nx, 3)
janssen.prop.debye_wolf_focus(pupil_field: OpticalWavefront, na: float | Float[Array, ''], focal_length: float | Float[Array, ''], z_focus: float | Float[Array, ''] = 0.0, refractive_index: float | Float[Array, ''] = 1.0) VectorWavefront3D[source]

Compute focal field using Debye-Wolf formulation.

This is an alias for high_na_focus using the alternative naming convention from the Debye approximation literature.

Parameters:
  • pupil_field (OpticalWavefront) – Input polarized field in pupil plane.

  • na (float) – Numerical aperture.

  • focal_length (float) – Focal length in meters.

  • z_focus (float, optional) – Axial position relative to focus. Default is 0.0.

  • refractive_index (float, optional) – Refractive index of focal medium. Default is 1.0.

Returns:

focal_field – Vector field at focal plane.

Return type:

VectorWavefront3D

See also

high_na_focus

Main implementation with additional options.

janssen.prop.high_na_focus(pupil_field: OpticalWavefront, na: float | Float[Array, ''], focal_length: float | Float[Array, ''], z_focus: float | Float[Array, ''] = 0.0, output_dx: float | Float[Array, ''] | None = None, output_grid_size: tuple[int, int] | Int[Array, '2'] | None = None, refractive_index: float | Float[Array, ''] = 1.0, include_aplanatic_factor: bool = True) VectorWavefront3D[source]

Compute focal field using Richards-Wolf vector diffraction integrals.

This is the main entry point for high-NA vector focusing simulations. It takes a polarized pupil field and computes the full 3D vector field (Ex, Ey, Ez) at the focal plane.

Parameters:
  • pupil_field (OpticalWavefront) – Input field in the pupil plane. Must be polarized with shape (H, W, 2) containing [Ex, Ey] Jones components.

  • na (float) – Numerical aperture of the focusing lens.

  • focal_length (float) – Focal length of the lens in meters.

  • z_focus (float, optional) – Axial position relative to geometric focus in meters. z_focus = 0 gives the focal plane. Default is 0.0.

  • output_dx (float, optional) – Output pixel size in focal plane. If None, computed from diffraction limit: dx ≈ λ/(2*NA) / 4.

  • output_grid_size (Tuple[int, int], optional) – Output grid size. If None, uses same size as input.

  • refractive_index (float, optional) – Refractive index of focal medium. Default is 1.0 (air).

  • include_aplanatic_factor (bool, optional) – Whether to include √cos(θ) apodization. Default is True.

Returns:

focal_field – Vector field at focal plane with shape (H, W, 3) containing [Ex, Ey, Ez] components.

Return type:

VectorWavefront3D

Notes

The algorithm: 1. Create pupil coordinates and map to convergence angles 2. Apply aplanatic apodization √cos(θ) 3. Compute polarization rotation matrix P(θ, φ) 4. Apply rotation to get (Ex’, Ey’, Ez’) in focal coordinates 5. Apply defocus phase exp(ikz cos(θ)) 6. Use 2D FFT to evaluate the focusing integral 7. Package result as VectorWavefront3D

For the focal plane (z=0), the integral reduces to a 2D Fourier transform relationship, enabling efficient FFT-based computation.

Examples

>>> from janssen.models import radially_polarized_beam
>>> from janssen.prop import high_na_focus
>>>
>>> # Create radially polarized beam in pupil
>>> pupil = radially_polarized_beam(
...     wavelength=633e-9,
...     dx=10e-6,
...     grid_size=(256, 256),
...     beam_radius=1e-3,
... )
>>>
>>> # Focus with high-NA lens
>>> focal = high_na_focus(
...     pupil_field=pupil,
...     na=0.9,
...     focal_length=3e-3,
... )
>>>
>>> print(f"Ez peak: {jnp.max(jnp.abs(focal.ez)**2):.3e}")
janssen.prop.scalar_focus_for_comparison(pupil_field: OpticalWavefront, focal_length: float | Float[Array, ''], z_focus: float | Float[Array, ''] = 0.0) OpticalWavefront[source]

Compute scalar focal field for comparison with vector result.

This function ignores polarization and computes a simple Fourier transform focal field, demonstrating what scalar theory predicts. Useful for comparing with vector results to highlight the differences.

Parameters:
  • pupil_field (OpticalWavefront) – Input field (polarization ignored, uses total amplitude).

  • focal_length (float) – Focal length in meters.

  • z_focus (float, optional) – Axial position. Default is 0.0.

Returns:

scalar_focal – Scalar focal field (Ez effects not modeled).

Return type:

OpticalWavefront

Notes

The scalar approximation: - Ignores polarization rotation (no Ez generation) - Predicts identical PSF for all input polarizations - Fails at high NA where vector effects are significant