wf_psf.utils.utils

Utility functions for the PSF simulation and modeling.

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

Functions

`add_noise`(image, desired_SNR)	Add noise to an image to obtain a desired SNR.
`calc_poly_position_mat`(pos, x_lims, y_lims, ...)	Calculate a matrix with position polynomials.
`compute_unobscured_zernike_projection`(tf_z1, ...)	Compute a zernike projection for unobscured wavefronts (OPDs).
`convert_to_tf`(data, tf_dtype)	Convert a sequence of array-like objects to TensorFlow tensors with a specified dtype.
`decimate_im`(input_im, decim_f)	Decimate image.
`decompose_tf_obscured_opd_basis`(tf_opd, ...)	Decompose obscured OPD into a basis using an iterative algorithm.
`downsample_im`(input_im, output_dim)	Downsample image to (output_dim, output_dim).
`generalised_sigmoid`(x[, max_val, power_k])	Apply a generalized sigmoid function to the input.
`generate_SED_elems`(SED, psf_simulator[, n_bins])	Generate SED elements for PSF modeling.
`generate_SED_elems_in_tensorflow`(SED, ...[, ...])	Generate SED Elements in TensorFlow Units.
`generate_n_mask`(shape[, n_masks])	Generate n masks with random 2D cosine waves.
`generate_packed_elems`(SED, psf_simulator[, ...])	Generate packed SED elements as TensorFlow tensors.
`load_multi_cycle_params_click`(args)	Load multiple cycle training parameters.
`single_mask_generator`(shape)	Generate a single mask with random 2D cosine waves.
`zernike_generator`(n_zernikes, wfe_dim)	Generate Zernike maps.

Classes

`IndependentZernikeInterpolation`(tf_pos, tf_zks)	Interpolate each Zernike polynomial independently.
`NoiseEstimator`(img_dim, win_rad)	A class for estimating noise levels in an image.
`ZernikeInterpolation`(tf_pos, tf_zks[, k, order])	Interpolate Zernike coefficients using K-nearest RBF splines.

class wf_psf.utils.utils.IndependentZernikeInterpolation(tf_pos, tf_zks, order=2)[source]

Bases: object

Interpolate each Zernike polynomial independently.

The interpolation is done independently for each Zernike polynomial.

Parameters:

tf_pos (Tensor (n_sources, 2)) – Positions
tf_zks (Tensor (n_sources, n_zernikes)) – Zernike coefficients for each position
order (int) – Order of the RBF interpolation. Default is 2, corresponds to thin plate interp (r^2*log(r))

Methods

`interp_one_zk`(zk_prior)	Interpolate a single Zernike polynomial across target positions.
`interpolate_zks`(target_pos)	Vectorize to interpolate to each Zernike.

interp_one_zk(zk_prior)[source]

Interpolate a single Zernike polynomial across target positions.

Each Zernike coefficient in zk_prior is interpolated independently using a spline.

Parameters:: zk_prior (tf.Tensor of shape (n_sources,)) – Zernike coefficients for a single Zernike polynomial, defined at the source positions self.tf_pos.
Returns:: Interpolated Zernike coefficients at the target positions self.target_pos.
Return type:: tf.Tensor of shape (n_targets,)

Notes

This function uses tfa.image.interpolate_spline, which requires the input to have a batch dimension. The extra dimension is removed before returning the result.

interpolate_zks(target_pos)[source]

Vectorize to interpolate to each Zernike.

Each zernike is computed indepently from the others.

Parameters:: target_pos (Tensor (n_targets, 2)) – Positions to interpolate to.
Return type:: Tensor (n_targets, n_zernikes)

class wf_psf.utils.utils.NoiseEstimator(img_dim: tuple[int, int], win_rad: int)[source]

Bases: object

A class for estimating noise levels in an image.

Parameters:

img_dim (tuple of int) – The dimensions of the image as (height, width).
win_rad (int) – The radius of the exclusion window (in pixels).

Methods

`apply_mask`([mask])	Apply a given mask to the exclusion window.
`estimate_noise`(image[, mask])	Estimates the noise level of an image using the MAD estimator.
`sigma_mad`(x)	Robustly estimate the standard deviation using the Median Absolute Deviation (MAD).

apply_mask(mask: ndarray | None = None) → ndarray[source]

Apply a given mask to the exclusion window.

Parameters:: mask (np.ndarray, optional) – A boolean mask to apply to the exclusion window. If None, the exclusion window is returned without any modification.
Returns:: The resulting boolean array after applying the mask to the exclusion window.
Return type:: np.ndarray

estimate_noise(image: ndarray, mask: ndarray | None = None) → float[source]

Estimates the noise level of an image using the MAD estimator.

Parameters:

image (np.ndarray) – The input image for noise estimation.
mask (np.ndarray, optional) – A boolean mask specifying which pixels to include in the noise estimation. If None, the default exclusion window is used. The mask should have the same shape as image.

Returns:

The estimated noise standard deviation (MAD of the image pixels within the window or mask).

Return type:

float

static sigma_mad(x)[source]

Robustly estimate the standard deviation using the Median Absolute Deviation (MAD).

Computes MAD = median(|x - median(x)|) and scales it by 1.4826 to make the estimator consistent with the standard deviation for a Gaussian distribution:

sigma ≈ 1.4826 * MAD

Parameters:: x (array-like) – Input data. The values are flattened before computation. NaNs are not specially handled and will propagate; remove or mask them prior to calling if needed.
Returns:: Robust estimate of the standard deviation of the input data.
Return type:: float

Notes

The MAD-based estimator is much less sensitive to outliers than the sample standard deviation, making it appropriate for noisy data with occasional large deviations.
The constant 1.4826 is the scaling factor for consistency with the standard deviation of a normal distribution.

class wf_psf.utils.utils.ZernikeInterpolation(tf_pos, tf_zks, k=50, order=2)[source]

Bases: object

Interpolate Zernike coefficients using K-nearest RBF splines.

This class provides utilities to interpolate Zernike-coefficient vectors defined at a set of source positions to arbitrary query positions using a local RBF spline fitted to the K nearest source samples.

The interpolation pipeline: - For a given query position, compute Euclidean distances to all source

positions and select the K nearest neighbors.

Use tfa.image.interpolate_spline (RBF / spline interpolation) on the selected neighbor positions and their Zernike coefficient vectors to compute the interpolated coefficients at the query location.

Parameters:

tf_pos (tf.Tensor, shape (n_sources, 2)) – Source/sample positions (x, y). Expected dtype float32 or convertible.
tf_zks (tf.Tensor, shape (n_sources, n_zernikes)) – Zernike coefficient vectors at the source positions.
k (int, default 50) – Number of nearest neighbors to use for the local interpolation. If larger than the number of sources, all sources are used.
order (int, default 2) – Spline order passed to tfa.image.interpolate_spline (e.g. 2 for thin plate style interpolation).

tf_pos, tf_zks, k, order: Stored copies of the constructor inputs.

Notes

This class relies on TensorFlow Addons’ interpolate_spline, which requires inputs to include a leading batch dimension; the implementation handles that automatically.
For best numerical stability and compatibility with TFA, use float32 tensors for inputs when possible.
Two main methods are provided:
- interpolate_zk(single_pos): interpolate a single position -> 1D vector.
- interpolate_zks(interp_positions): vectorized interpolation for many query positions (uses tf.map_fn under the hood).

Methods

`interpolate_zk`(single_pos)	Interpolate Zernike coefficients at a single query position using K-nearest neighbors.
`interpolate_zks`(interp_positions)	Interpolate Zernike coefficient vectors at multiple query positions.

interpolate_zk(single_pos)[source]

Interpolate Zernike coefficients at a single query position using K-nearest neighbors.

Finds the K nearest training positions to the query position and uses RBF spline interpolation to estimate Zernike coefficients at that location.

Parameters:: single_pos (tf.Tensor, shape (2,)) – Query position coordinates as (x, y).
Returns:: Interpolated Zernike coefficient vector at the query position.
Return type:: tf.Tensor, shape (n_zernikes,)

interpolate_zks(interp_positions)[source]

Interpolate Zernike coefficient vectors at multiple query positions.

Vectorized wrapper that applies self.interpolate_zk to each row of interp_positions using tf.map_fn.

Parameters:: interp_positions (tf.Tensor, shape (n_targets, 2)) – Query positions where Zernike coefficients should be interpolated. Each row is an (x, y) coordinate.
Returns:: Interpolated Zernike coefficient vectors for each query position. tf.map_fn may introduce an extra singleton dimension; this is removed by tf.squeeze before returning.
Return type:: tf.Tensor, shape (n_targets, n_zernikes), dtype=tf.float32

Notes

self.interpolate_zk expects a 1-D tensor of shape (2,) and returns a 1-D tensor of length n_zernikes.
This function uses tf.map_fn with fn_output_signature=tf.float32 and swap_memory=True for efficient batching.

wf_psf.utils.utils.add_noise(image, desired_SNR)[source]: Add noise to an image to obtain a desired SNR.

wf_psf.utils.utils.calc_poly_position_mat(pos, x_lims, y_lims, d_max)[source]

Calculate a matrix with position polynomials.

Scale positions to the square: [self.x_lims[0], self.x_lims[1]] x [self.y_lims[0], self.y_lims[1]] to the square [-1,1] x [-1,1]

wf_psf.utils.utils.compute_unobscured_zernike_projection(tf_z1, tf_z2, norm_factor=None)[source]

Compute a zernike projection for unobscured wavefronts (OPDs).

Compute internal product between zernikes and OPDs.

Defined such that Zernikes are orthonormal to each other.

First one should compute: norm_factor = unobscured_zernike_projection(tf_zernike,tf_zernike) for futur calls: unobscured_zernike_projection(OPD,tf_zernike_k, norm_factor)

If the OPD has obscurations, or is not an unobscured circular aperture, the Zernike polynomials are no longer orthonormal. Therefore, you should consider using the function tf_decompose_obscured_opd_basis that takes into account the obscurations in the projection.

wf_psf.utils.utils.convert_to_tf(data, tf_dtype)[source]

Convert a sequence of array-like objects to TensorFlow tensors with a specified dtype.

Parameters:

data (Iterable) – An iterable (e.g., list, tuple) of array-like objects (numpy arrays, Python lists/tuples, tf.Tensor, etc.) to be converted to TensorFlow tensors.
tf_dtype (tf.DType) – The TensorFlow dtype to cast each element to (for example tf.float32, tf.int32, etc.).

Returns:

A list where each element is the result of calling tf.convert_to_tensor on the corresponding item from data, cast to tf_dtype.

Return type:

list of tf.Tensor

Raises:

TypeError – If data is not an iterable. A TypeError may also be raised by tf.convert_to_tensor for individual elements that cannot be converted.

Notes

The function preserves the top-level sequence structure by returning a list regardless of the input sequence type.
Element-wise conversion uses TensorFlow’s conversion semantics; shape inference and broadcasting follow TensorFlow rules.

wf_psf.utils.utils.decimate_im(input_im, decim_f)[source]

Decimate image.

Decimated by a factor of decim_f. Based on the PIL library using the default interpolator. Default: PIL.Image.BICUBIC.

wf_psf.utils.utils.decompose_tf_obscured_opd_basis(tf_opd, tf_obscurations, tf_zk_basis, n_zernike, iters=20)[source]

Decompose obscured OPD into a basis using an iterative algorithm.

Tensorflow implementation.

Parameters:

tf_opd (tf.Tensor) – Input OPD that requires to be decomposed on tf_zk_basis. The tensor shape is (opd_dim, opd_dim).
tf_obscurations (tf.Tensor) – Tensor with the obscuration map. The tensor shape is (opd_dim, opd_dim).
tf_zk_basis (tf.Tensor) – Zernike polynomial maps. The tensor shape is (n_batch, opd_dim, opd_dim)
n_zernike (int) – Number of Zernike polynomials to project on.
iters (int) – Number of iterations of the algorithm.

Returns:

obsc_coeffs – Array of size n_zernike with projected Zernike coefficients

Return type:

np.ndarray

Raises:

ValueError – If n_zernike is bigger than tf_zk_basis.shape[0].

wf_psf.utils.utils.downsample_im(input_im, output_dim)[source]

Downsample image to (output_dim, output_dim).

Uses OpenCV INTER_AREA when available, otherwise falls back to scikit-image local mean downsampling.

Parameters:

input_im (np.ndarray) – Input 2D image to be downsampled.
output_dim (int) – Desired output dimension (both height and width).

Returns:

Downsampled 2D image of shape (output_dim, output_dim).

Return type:

np.ndarray

wf_psf.utils.utils.generalised_sigmoid(x, max_val=1, power_k=1)[source]

Apply a generalized sigmoid function to the input.

This function computes a smooth, S-shaped curve that generalizes the standard sigmoid function. It’s useful for scaling values while maintaining a bounded output.

Parameters:

x (array_like) – Input value(s) to which the generalized sigmoid is applied.
max_val (float, optional) – Maximum output value. Default is 1.
power_k (float, optional) – Power parameter controlling the steepness of the curve. Default is 1. Higher values create steeper transitions.

Returns:

Output value(s) scaled by the generalized sigmoid function, bounded between -max_val and max_val.

Return type:

array_like

Notes

When power_k=1, this reduces to a standard rational sigmoid function. The function is odd, meaning generalised_sigmoid(-x) = -generalised_sigmoid(x).

wf_psf.utils.utils.generate_SED_elems(SED, psf_simulator, n_bins=20)[source]

Generate SED elements for PSF modeling.

Computes feasible Zernike mode numbers, wavelength values, and normalized SED for a given spectral energy distribution (SED) sampled across specified wavelength bins. These elements are required for PSF simulation and modeling with the TensorFlow-based PSF classes.

Parameters:

SED (np.ndarray) – The unfiltered SED with shape (n_wavelengths, 2). The first column contains wavelength positions (in wavelength units), and the second column contains the corresponding SED flux values.
psf_simulator (PSFSimulator) – An instance of the PSFSimulator class initialized with the correct optical and instrumental parameters.
n_bins (int, optional) – Number of wavelength bins to sample the SED. Default is 20.

Returns:

feasible_Nnp.ndarray, shape (n_bins,)
Feasible Zernike mode numbers at each wavelength bin.
feasible_wvnp.ndarray, shape (n_bins,)
Sampled wavelength values across the SED.
SED_normnp.ndarray or float
Normalized SED values corresponding to feasible wavelengths.

Return type:

tuple of (np.ndarray, np.ndarray, np.ndarray or float)