wf_psf.psf_models.tf_modules

TensorFlow-Based PSF Modeling.

A module containing TensorFlow implementations for modeling monochromatic PSFs using Zernike polynomials and Fourier optics.

Author:: Tobias Liaudat <tobiasliaudat@gmail.com>

Classes

`TFBuildPhase`(args, *kwargs)	Build a complex phase map from an Optical Path Difference (OPD) map.
`TFFftDiffract`(args, *kwargs)	Diffract the wavefront into a monochromatic PSF.
`TFMonochromaticPSF`(args, *kwargs)	Calculate a monochromatic Point Spread Function (PSF) from an OPD map.
`TFZernikeMonochromaticPSF`(args, *kwargs)	Build a monochromatic Point Spread Function (PSF) from Zernike coefficients.
`TFZernikeOPD`(args, *kwargs)	Convert Zernike coefficients into an Optical Path Difference (OPD).

class wf_psf.psf_models.tf_modules.TFBuildPhase(*args: Any, **kwargs: Any)[source]

Bases: Module

Build a complex phase map from an Optical Path Difference (OPD) map.

This class takes an OPD map and converts it into a complex phase map. It applies necessary obscurations (such as apertures or masks) and zero-padding to match the required size for diffraction simulations. The resulting phase map is essential for further optical modeling, such as diffraction simulations or other optical system analysis.

phase_N

The desired size of the phase map (e.g., pixel count for height and width).

Type:: int

lambda_obs

The observed wavelength used for phase calculations, typically in meters.

Type:: float

obscurations

A tensor representing the obscurations (e.g., apertures or masks) to be applied to the phase.

Type:: tf.Tensor

Methods

`__call__`(opd)	Convert an OPD map to a padded and obscured phase map.
`apply_obscurations`(phase)	Apply obscurations to the phase map.
`opd_to_phase`(opd)	Convert an OPD map to a complex phase map.
`zero_padding_diffraction`(no_pad_phase)	Pad the phase map with zeros based on the required size.

apply_obscurations(phase: tensorflow.Tensor) → tensorflow.Tensor[source]

Apply obscurations to the phase map.

This method multiplies the phase map element-wise with the obscurations tensor. The obscurations tensor can represent apertures or masks that block or modify portions of the phase map.

Parameters:: phase (tf.Tensor) – The phase map to which obscurations will be applied. Expected shape is [batch_size, height, width].
Returns:: The phase map after applying the obscurations.
Return type:: tf.Tensor

opd_to_phase(opd: tensorflow.Tensor) → tensorflow.Tensor[source]

Convert an OPD map to a complex phase map.

This method takes an optical path difference (OPD) map and converts it into a complex phase map using the formula: phase = exp(i * (2 * pi / lambda_obs) * opd).

Parameters:: opd (tf.Tensor) – The optical path difference map. Expected shape is [batch_size, height, width].
Returns:: The complex phase map resulting from the OPD.
Return type:: tf.Tensor

zero_padding_diffraction(no_pad_phase)[source]

Pad the phase map with zeros based on the required size.

This method adds zero-padding to the input phase map to match the required size for diffraction calculations. The padding is computed based on the phase_N attribute and the input phase map size.

Parameters:: no_pad_phase (tf.Tensor) – The phase map that needs to be padded. Expected shape is [batch_size, height, width].
Returns:: padded_phase – The padded phase map with shape [batch_size, phase_N, phase_N].
Return type:: tf.Tensor

class wf_psf.psf_models.tf_modules.TFFftDiffract(*args: Any, **kwargs: Any)[source]

Bases: Module

Diffract the wavefront into a monochromatic PSF.

output_dim

Dimension of the output square postage stamp

Type:: int

output_Q

Downsampling factor. Must be integer.

Type:: int

Methods

`__call__`(input_phase)	Calculate the normalized Point Spread Function (PSF) from a phase array.
`normalize_psf`(psf)	Normalize the Point Spread Function (PSF).
`tf_crop_img`(image, output_crop_dim)	Crop images using TensorFlow methods.

normalize_psf(psf)[source]

Normalize the Point Spread Function (PSF).

This function normalizes a given Point Spread Function (PSF) by summing over the spatial dimensions and dividing the PSF by the resulting sum. The PSF is expected to have at least 3 dimensions, with the first dimension representing the batch size and the remaining two dimensions representing the spatial dimensions (height and width).

Parameters:: psf (tf.Tensor) – A tensor representing the Point Spread Function (PSF) with shape [batch, height, width]. The PSF is expected to be a 3D tensor, where the first dimension corresponds to the batch size, and the other two dimensions represent the spatial dimensions of the PSF.
Returns:: The normalized PSF with the same shape as the input, [batch, height, width], where each PSF has been normalized by the sum of the PSF over the spatial dimensions.
Return type:: tf.Tensor

tf_crop_img(image, output_crop_dim)[source]

Crop images using TensorFlow methods.

This method handles a batch of 2D images and crops them to the specified dimension. The images are expected to have the shape [batch, width, height], and the method uses TensorFlow’s crop_to_bounding_box to crop each image in the batch.

Parameters:

image (tf.Tensor) – A batch of 2D images with shape [batch, height, width]. The images are expected to be 3D tensors where the second and third dimensions represent the height and width.
output_crop_dim (int) – The dimension of the square crop. The image will be cropped to this dimension.

Returns:

The cropped images with shape [batch, output_crop_dim, output_crop_dim].

Return type:

tf.Tensor

class wf_psf.psf_models.tf_modules.TFMonochromaticPSF(*args: Any, **kwargs: Any)[source]

Bases: Module

Calculate a monochromatic Point Spread Function (PSF) from an OPD map.

This class computes the monochromatic Point Spread Function (PSF) by first converting the Optical Path Difference (OPD) map into a phase map. Then, it applies diffraction using Fast Fourier Transform (FFT) techniques to simulate the PSF, which is essential in optical system simulations.

output_Q

The output quality factor used for diffraction simulations.

Type:: int

tf_build_phase

A module that builds the phase map from the OPD map, applying necessary zero-padding and obscurations.

Type:: TFBuildPhase

tf_fft_diffract

A module that performs the diffraction simulation using FFT.

Type:: TFFftDiffract

Parameters:

phase_N (int) – The size of the phase map (e.g., pixel count for the height and width).
lambda_obs (float) – The observed wavelength used for phase calculations.
obscurations (tf.Tensor) – A tensor representing the obscurations (e.g., apertures or masks) to be applied to the phase.
output_Q (int) – The output quality factor used for diffraction simulations.
output_dim (int, optional) – The output dimension for the PSF, by default 64.
name (str, optional) – The name for the TensorFlow module, by default None.

Methods

__call__(opd)

Compute the PSF from an OPD map.

class wf_psf.psf_models.tf_modules.TFZernikeMonochromaticPSF(*args: Any, **kwargs: Any)[source]

Bases: Module

Build a monochromatic Point Spread Function (PSF) from Zernike coefficients.

This class computes the monochromatic PSF by following the Zernike model. It involves multiple stages, including the calculation of the OPD (Optical Path Difference), the phase from the OPD, and diffraction via FFT-based operations. The Zernike coefficients are used to generate the PSF.

Parameters:

phase_N (int) – The size of the phase grid, typically a square matrix dimension.
lambda_obs (float) – The wavelength of the observed light.
obscurations (tf.Tensor) – A tensor representing the obscurations in the system, which will be applied to the phase.
zernike_maps (tf.Tensor) – A tensor containing the Zernike maps, with the shape (num_coeffs, x_dim, y_dim), where num_coeffs is the number of Zernike coefficients and x_dim, y_dim are the dimensions of the Zernike maps.
output_dim (int, optional, default=64) – The output dimension of the PSF, i.e., the size of the resulting image.
name (str, optional) – The name of the module. Default is None.

tf_build_opd_zernike

A module used to generate the OPD from the Zernike coefficients.

Type:: TFZernikeOPD

tf_build_phase

A module used to compute the phase from the OPD.

Type:: TFBuildPhase

tf_fft_diffract

A module that performs the diffraction calculation using FFT-based methods.

Type:: TFFftDiffract

Methods

__call__(z_coeffs)

Compute the monochromatic PSF from Zernike coefficients.

class wf_psf.psf_models.tf_modules.TFZernikeOPD(*args: Any, **kwargs: Any)[source]

Bases: Module

Convert Zernike coefficients into an Optical Path Difference (OPD).

This class performs the weighted sum of Zernike coefficients and Zernike maps to compute the OPD. The Zernike maps and the corresponding Zernike coefficients are required to perform the calculation.

Parameters:

zernike_maps (tf.Tensor) – A tensor containing the Zernike maps. The shape should be (num_coeffs, x_dim, y_dim), where num_coeffs is the number of Zernike coefficients and x_dim, y_dim are the dimensions of each map.
name (str, optional) – The name of the module. Default is None.

Returns:

A tensor representing the OPD, with shape (num_star, x_dim, y_dim), where num_star corresponds to the number of stars and x_dim, y_dim are the dimensions of the OPD map.

Return type:

tf.Tensor

Methods

__call__(z_coeffs)

Compute the OPD from Zernike coefficients and maps.