wf_psf.psf_models.tf_layers

TensorFlow layers for PSF modelling.

This module contains TensorFlow layers to model PSF variations across the field of view.

Author:: Tobias Liaudat <tobias.liaudat@cea.fr>

Classes

`TFBatchMonochromaticPSF`(args, *kwargs)	Calculate a monochromatic PSF from a batch of OPDs.
`TFBatchPolychromaticPSF`(args, *kwargs)	Calculate a polychromatic PSF from an OPD and stored SED values.
`TFNonParametricGraphOPD`(args, *kwargs)	Non-parametric OPD generation with only graph-cosntraint variations.
`TFNonParametricMCCDOPDv2`(args, *kwargs)	Non-parametric OPD generation with hybrid-MCCD variations.
`TFNonParametricPolynomialVariationsOPD`(...)	Non-parametric OPD generation with polynomial variations.
`TFPhysicalLayer`(args, *kwargs)	The Zernike physical layer.
`TFPolynomialZernikeField`(args, *kwargs)	Calculate the zernike coefficients for a given position.
`TFZernikeOPD`(args, *kwargs)	Calculate the OPD from Zernike maps and coefficients.

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

Bases: Layer

Calculate a monochromatic PSF from a batch of OPDs.

The calculation of the phase_N variable is done with the PSFSimulator class but outside the TF class.

Parameters:

obscurations (Tensor [opd_dim, opd_dim]) – Obscurations to apply to the wavefront.
psf_batch (Tensor [batch_size, output_dim, output_dim]) – Tensor containing the psfs that will be updated each time a calculation is required. Can be started with zeros.
output_Q (int) – Output oversampling value.
output_dim (int) – Output PSF stamp dimension.

Methods

`__call__`(args, *kwargs)	Call self as a function.
`calculate_monochromatic_PSF`(current_opd)	Calculate monochromatic PSF from OPD info.
`call`(opd_batch)	Calculate the batch poly PSFs.
`init_mono_PSF`()	Initialise or restart the PSF generator.
`set_lambda_phaseN`([phase_N, lambda_obs])	Set the lambda value for monochromatic PSFs and the phaseN.
`set_output_params`(output_Q, output_dim)	Set output patams, Q and dimension.

calculate_monochromatic_PSF(current_opd)[source]: Calculate monochromatic PSF from OPD info.

call(opd_batch)[source]: Calculate the batch poly PSFs.

init_mono_PSF()[source]: Initialise or restart the PSF generator.

set_lambda_phaseN(phase_N=914, lambda_obs=0.7)[source]: Set the lambda value for monochromatic PSFs and the phaseN.

set_output_params(output_Q, output_dim)[source]: Set output patams, Q and dimension.

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

Bases: Layer

Calculate a polychromatic PSF from an OPD and stored SED values.

The calculation of the packed values with the respective SED is done with the PSFSimulator class but outside the TF class.

Parameters:

obscurations (Tensor [opd_dim, opd_dim]) – Obscurations to apply to the wavefront.
packed_SED_data (Tensor [batch_size, 3, n_bins_lda]) –
Comes from: tf.convert_to_tensor(list(list(Tensor,Tensor,Tensor))) Where each inner list consist of a packed_elem:

packed_elems: Tuple of tensors Contains three 1D tensors with the parameters needed for the calculation of one monochromatic PSF.

packed_elems[0]: phase_N packed_elems[1]: lambda_obs packed_elems[2]: SED_norm_val

The SED data is constant in a FoV.
psf_batch (Tensor [batch_size, output_dim, output_dim]) – Tensor containing the psfs that will be updated each time a calculation is required. REMOVED!

Methods

`__call__`(args, *kwargs)	Call self as a function.
`calculate_monochromatic_PSF`(packed_elems)	Calculate monochromatic PSF from packed elements.
`calculate_polychromatic_PSF`(packed_elems)	Calculate a polychromatic PSF.
`call`(inputs)	Calculate the batch polychromatic PSFs.

calculate_monochromatic_PSF(packed_elems)[source]

Calculate monochromatic PSF from packed elements.

packed_elems[0]: phase_N packed_elems[1]: lambda_obs packed_elems[2]: SED_norm_val

calculate_polychromatic_PSF(packed_elems)[source]: Calculate a polychromatic PSF.

call(inputs)[source]: Calculate the batch polychromatic PSFs.

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

Bases: Layer

Non-parametric OPD generation with only graph-cosntraint variations.

Parameters:

obs_pos (tensor(n_stars, 2)) – Observed positions of the n_stars in the dataset. The indexing of the positions has to correspond to the indexing in the spatial_dic.
spatial_dic (tensor(n_stars, n_dic_elems)) – Dictionary containing the spatial-constraint dictionary. n_stars corresponds to the total number of stars in the dataset. n_dic_elems corresponds to the number of elements of the dictionary, not to be confounded with n_comp, the total number of non-parametric features of the wavefront-PSF.
x_lims ([int, int]) – Limits of the x axis.
y_lims ([int, int]) – Limits of the y axis.
graph_comps (int) – Number of wavefront-PSF features correspondign to the graph constraint.
d_max (int) – Max degree of polynomial determining the FoV variations. The number of wavefront-PSF features of the polynomial part is computed (d_max+1)*(d_max+2)/2.
opd_dim (int) – Dimension of the OPD maps. Same as pupil diameter.

Methods

`__call__`(args, *kwargs)	Call self as a function.
`call`(positions)	Calculate the OPD maps for the given positions.
`init_vars`()	Initialize trainable variables.
`predict`(positions)	Prediction step.
`set_alpha_identity`()	Set alpha matrix to the identity.
`set_alpha_zero`()	Set alpha matrix to zero.

call(positions)[source]

Calculate the OPD maps for the given positions.

Calculating: batch(spatial_dict) x alpha x S

Parameters:: positions (Tensor(batch, 2)) – First element is x-axis, second is y-axis.
Returns:: opd_maps
Return type:: Tensor(batch, opd_dim, opd_dim)

init_vars()[source]

Initialize trainable variables.

Basic initialization. Random uniform for S and identity for alpha.

predict(positions)[source]: Prediction step.

set_alpha_identity()[source]: Set alpha matrix to the identity.

set_alpha_zero()[source]: Set alpha matrix to zero.

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

Bases: Layer

Non-parametric OPD generation with hybrid-MCCD variations.

Parameters:

obs_pos (tensor(n_stars, 2)) – Observed positions of the n_stars in the dataset. The indexing of the positions has to correspond to the indexing in the spatial_dic.
spatial_dic (tensor(n_stars, n_dic_elems)) – Dictionary containing the spatial-constraint dictionary. n_stars corresponds to the total number of stars in the dataset. n_dic_elems corresponds to the number of elements of the dictionary, not to be confounded with n_comp, the total number of non-parametric features of the wavefront-PSF.
x_lims ([int, int]) – Limits of the x axis.
y_lims ([int, int]) – Limits of the y axis.
graph_comps (int) – Number of wavefront-PSF features correspondign to the graph constraint.
d_max (int) – Max degree of polynomial determining the FoV variations. The number of wavefront-PSF features of the polynomial part is computed (d_max+1)*(d_max+2)/2.
opd_dim (int) – Dimension of the OPD maps. Same as pupil diameter.

Methods

`__call__`(args, *kwargs)	Call self as a function.
`call`(positions)	Calculate the OPD maps for the given positions.
`init_vars`()	Initialize trainable variables.
`predict`(positions)	Prediction step.
`set_alpha_identity`()	Set alpha matrix to the identity.
`set_alpha_zero`()	Set alpha matrix to zero.

call(positions)[source]

Calculate the OPD maps for the given positions.

Calculating: batch(spatial_dict) x alpha x S

Parameters:: positions (Tensor(batch, 2)) – First element is x-axis, second is y-axis.
Returns:: opd_maps
Return type:: Tensor(batch, opd_dim, opd_dim)

init_vars()[source]

Initialize trainable variables.

Basic initialization. Random uniform for S and identity for alpha.

predict(positions)[source]: Prediction step.

set_alpha_identity()[source]: Set alpha matrix to the identity.

set_alpha_zero()[source]: Set alpha matrix to zero.

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

Bases: Layer

Non-parametric OPD generation with polynomial variations.

Parameters:

x_lims ([int, int]) – Limits of the x axis.
y_lims ([int, int]) – Limits of the y axis.
random_seed (int) – Random seed initialization for Tensor Flow
d_max (int) – Max degree of polynomial determining the FoV variations.
opd_dim (int) – Dimension of the OPD maps. Same as pupil diameter.

Methods

`__call__`(args, *kwargs)	Call self as a function.
`assign_S_mat`(S_mat)	Assign DD features matrix.
`call`(positions)	Calculate the OPD maps for the given positions.
`init_vars`()	Initialize trainable variables.
`set_alpha_identity`()	Set alpha matrix to the identity.
`set_alpha_zero`()	Set alpha matrix to zero.

assign_S_mat(S_mat)[source]: Assign DD features matrix.

call(positions)[source]

Calculate the OPD maps for the given positions.

Calculating: Pi(pos) x alpha x S

Parameters:: positions (Tensor(batch, 2)) – First element is x-axis, second is y-axis.
Returns:: opd_maps
Return type:: Tensor(batch, opd_dim, opd_dim)

init_vars()[source]

Initialize trainable variables.

Basic initialization. Random uniform for S and identity for alpha.

set_alpha_identity()[source]: Set alpha matrix to the identity.

set_alpha_zero()[source]: Set alpha matrix to zero.

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

Bases: Layer

The Zernike physical layer.

This layer gives the Zernike contribution of the physical layer. It is fixed and not trainable. It can interpolate the Zernike coefficients at the input positions using different interpolation schemes.

Parameters:

obs_pos (Tensor (n_stars, 2)) – Observed positions of the n_stars in the dataset. The indexing of the positions has to correspond to the indexing in the zks_prior.
n_zernikes (int) – Number of Zernike polynomials
zks_prior (Tensor (n_stars, n_zernikes)) – Zernike coefficients for each position
interpolation_type (str) – Type of interpolation to be used. Options are: None, ‘all’, ‘top_K’, ‘independent_Zk’. Default is None.
interpolation_args (dict) – Interpolation hyper-parameters. The order of the RBF interpolation, and the K elements in the top_K interpolation.

Methods

`__call__`(args, *kwargs)	Call self as a function.
`call`(positions)	Calculate the prior Zernike coefficients for a batch of positions.
`interpolate_all`(positions)	Interpolate using all the input elements.
`interpolate_independent_Zk`(positions)	Interpolate each Zernike independently.
`interpolate_top_K`(positions)	Interpolate using only the K closest elements.

call(positions)[source]

Calculate the prior Zernike coefficients for a batch of positions.

This method calculates the Zernike coefficients for a batch of input positions based on the pre-computed Zernike coefficients for observed positions.

Parameters:: positions (tf.Tensor) – Tensor of shape (batch_size, 2) representing the positions. The first element represents the x-axis, and the second element represents the y-axis.
Returns:: zernike_coeffs – Tensor of shape (batch_size, n_zernikes, 1, 1) containing the prior Zernike coefficients corresponding to the input positions.
Return type:: tf.Tensor

Notes

The method retrieves the Zernike coefficients for each input position from the pre-computed Zernike coefficients stored for observed positions. It matches each input position with the closest observed position and retrieves the corresponding Zernike coefficients.

Before calling this method, ensure that the position polynomial matrix and the corresponding Zernike coefficients have been precomputed and set for the layer.

Raises:: ValueError – If the shape of the input positions tensor is not compatible.

interpolate_all(positions)[source]

Interpolate using all the input elements.

The TensorFlow Addons function tfa.image.interpolate_spline is used to perform the RBF interpolation of the Zernike coefficients.

Parameters:: positions (tf.Tensor) – Tensor of shape (batch_size, 2) representing the positions. The first element represents the x-axis, and the second element represents the y-axis.
Returns:: interp_zks – Tensor of shape (batch_size, n_zernikes, 1, 1) containing the interpolated Zernike coefficients corresponding to the input positions.
Return type:: tf.Tensor

interpolate_independent_Zk(positions)[source]

Interpolate each Zernike independently.

The class wf.utils.IndependentZernikeInterpolation allows to interpolate each order of the Zernike polynomials independently using all the points avaialble to build the interpolant.

Parameters:: positions (tf.Tensor) – Tensor of shape (batch_size, 2) representing the positions. The first element represents the x-axis, and the second element represents the y-axis.
Returns:: interp_zks – Tensor of shape (batch_size, n_zernikes, 1, 1) containing the interpolated Zernike coefficients corresponding to the input positions.
Return type:: tf.Tensor

interpolate_top_K(positions)[source]

Interpolate using only the K closest elements.

The class wf.utils.ZernikeInterpolation allows to interpolate the Zernike coefficients using only the K closest points to build the interpolant.

Parameters:: positions (tf.Tensor) – Tensor of shape (batch_size, 2) representing the positions. The first element represents the x-axis, and the second element represents the y-axis.
Returns:: interp_zks – Tensor of shape (batch_size, n_zernikes, 1, 1) containing the interpolated Zernike coefficients corresponding to the input positions.
Return type:: tf.Tensor

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

Bases: Layer

Calculate the zernike coefficients for a given position.

This module implements a polynomial model of Zernike coefficient variation.

Parameters:

n_zernikes (int) – Number of Zernike polynomials to consider
d_max (int) – Max degree of polynomial determining the FoV variations.

Methods

`__call__`(args, *kwargs)	Call self as a function.
`assign_coeff_matrix`(coeff_mat)	Assign coefficient matrix.
`call`(positions)	Calculate the zernike coefficients for a given position.
`get_coeff_matrix`()	Get coefficient matrix.
`get_poly_coefficients_shape`()	Return the shape of the coefficient matrix.
`init_coeff_matrix`()	Initialize coefficient matrix.

assign_coeff_matrix(coeff_mat)[source]: Assign coefficient matrix.

call(positions)[source]

Calculate the zernike coefficients for a given position.

The position polynomial matrix and the coefficients should be set before calling this function.

Parameters:: positions (Tensor(batch, 2)) – First element is x-axis, second is y-axis.
Returns:: zernikes_coeffs
Return type:: Tensor(batch, n_zernikes, 1, 1)

get_coeff_matrix()[source]: Get coefficient matrix.

get_poly_coefficients_shape()[source]: Return the shape of the coefficient matrix.

init_coeff_matrix()[source]: Initialize coefficient matrix.

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

Bases: Layer

Calculate the OPD from Zernike maps and coefficients.

This class generates OPD maps from Zernike coefficients and Zernike maps. Both Zernike maps and Zernike coefficients must be provided to the class.

Parameters:

zernike_maps (Tensor (Num_coeffs, x_dim, y_dim))
z_coeffs (Tensor (batch_size, n_zernikes, 1, 1))

Returns:

opd

Return type:

Tensor (batch_size, x_dim, y_dim)

Methods

`__call__`(args, *kwargs)	Call self as a function.
`call`(z_coeffs)	Perform the weighted sum of Zernikes coeffs and maps.

call(z_coeffs)[source]

Perform the weighted sum of Zernikes coeffs and maps.

Returns:: opd
Return type:: Tensor (batch_size, x_dim, y_dim)