wf_psf.data.data_zernike_utils

Utilities for Zernike Data Handling.

This module provides utilities for converting between physical optical misalignments (centroid shifts, defocus, CCD alignment effects) and Zernike coefficients used in wavefront modelling.

It supports assembling Zernike contributions from multiple sources, including centroid corrections, CCD misalignment terms, and optional priors, for use in PSF modelling pipelines.

This module does not handle data loading or model execution; it only provides deterministic transformations between physical parameters and Zernike representations.

Authors:: Tobias Liaudat <tobias.liaudat@cea.fr> and Jennifer Pollack <jennifer.pollack@cea.fr>

Functions

`assemble_zernike_contributions`(model_params)	Assemble Zernike contributions from prior, centroid correction, and CCD misalignment.
`combine_zernike_contributions`(contributions)	Combine multiple Zernike contribution arrays into a single array.
`compute_zernike_tip_tilt`(star_images[, ...])	Compute Zernike tip-tilt corrections for a batch of PSF images.
`defocus_to_zk4_wavediff`(dz[, ...])	Compute Zernike 4 value for a given defocus in WaveDiff conventions.
`defocus_to_zk4_zemax`(dz[, tel_focal_length, ...])	Compute Zernike 4 value for a given defocus in zemax conventions.
`pad_contribution_to_order`(contribution, ...)	Pad a Zernike contribution array to the max Zernike order.
`pad_tf_zernikes`(zk_param, zk_prior, n_zks_total)	Pad the Zernike coefficient tensors to match the specified total number of Zernikes.
`shift_x_y_to_zk1_2_wavediff`(dxy[, ...])	Compute Zernike 1(2) for a given shifts in x(y) in WaveDiff conventions.

Classes

ZernikeDataset(container)

Domain-specific canonical view over a DatasetContainer PSF modelling with Zernike input layers.

class wf_psf.data.data_zernike_utils.ZernikeDataset(container: DatasetContainer)[source]

Bases: object

Domain-specific canonical view over a DatasetContainer PSF modelling with Zernike input layers.

Provides access to Zernike-related inputs, priors, and datasets for centroid and misalignment corrections

container

Container storing the underlying datasets.

Type:: DatasetContainer

zernike_prior

True Zernike prior (e.g., from PDC). May be None.

Type:: Optional[np.ndarray]

sources

2D image stamps of sources for centroid correction. May be None.

Type:: Optional[np.ndarray]

masks

Masks corresponding to the sources. May be None.

Type:: Optional[np.ndarray]

positions

Positions for CCD misalignment corrections. May be None.

Type:: Optional[np.ndarray]

Attributes:

centroid_inputs: Return a dictionary suitable for centroid correction.
masks: Return masks corresponding to source images.
positions: Return positions for CCD misalignment correction.
sources: Return 2D source images (stamps).
zernike_prior: Return Zernike prior.

property centroid_inputs: dict | None

Return a dictionary suitable for centroid correction.

Includes ‘stamps’ (sources) and optionally ‘masks’ if available. Returns None if sources are not available.

container: DatasetContainer

property masks: ndarray | None: Return masks corresponding to source images.

property positions: ndarray | None: Return positions for CCD misalignment correction.

property sources: ndarray | None: Return 2D source images (stamps).

property zernike_prior: ndarray | None: Return Zernike prior.

wf_psf.data.data_zernike_utils.assemble_zernike_contributions(model_params, zernike_prior=None, centroid_dataset=None, positions=None, batch_size=16)[source]

Assemble Zernike contributions from prior, centroid correction, and CCD misalignment.

This function checks the model parameters to determine which contributions to include, computes each contribution as needed, and combines them into a single Zernike contribution tensor. It handles the logic for when certain contributions are not used or not available, ensuring that the final output is correctly shaped and contains the appropriate information based on the configuration.

Parameters:

model_params (RecursiveNamespace) – Parameters controlling which contributions to apply.
zernike_prior (Optional[np.ndarray or tf.Tensor]) – The precomputed Zernike prior. Can be either a NumPy array or a TensorFlow tensor. If a Tensor, will be converted to NumPy in eager mode.
centroid_dataset (Optional[object]) – Dataset used to compute centroid correction. Must have both training and validation sets.
positions (Optional[np.ndarray or tf.Tensor]) – Positions used for computing CCD misalignment. Must be available in inference mode.
batch_size (int) – Batch size for centroid correction.

Returns:

A tensor representing the full Zernike contribution map.

Return type:

tf.Tensor

wf_psf.data.data_zernike_utils.combine_zernike_contributions(contributions: list[ndarray]) → ndarray[source]

Combine multiple Zernike contribution arrays into a single array.

Each contribution is zero-padded along the second dimension (Zernike order) to match the maximum order across inputs, then summed element-wise.

Parameters:: contributions (list[np.ndarray]) – List of arrays of shape (n_samples, n_orders_i), where all arrays must share the same number of samples (first dimension) but may differ in Zernike order (second dimension).
Returns:: Array of shape (n_samples, max_order) containing the summed contributions after zero-padding
Return type:: np.ndarray
Raises:: ValueError – If the list is empty or contributions have inconsistent number of samples.

wf_psf.data.data_zernike_utils.compute_zernike_tip_tilt(star_images: ndarray, star_masks: ndarray | None = None, pixel_sampling: float = 1.2e-05, reference_shifts: list[float] = [-0.3333333333333333, -0.3333333333333333], sigma_init: float = 2.5, n_iter: int = 20) → ndarray[source]

Compute Zernike tip-tilt corrections for a batch of PSF images.

This function estimates the centroid shifts of multiple PSFs and computes the corresponding Zernike tip-tilt corrections to align them with a reference.

Parameters:

star_images (np.ndarray) – A batch of PSF images (3D array of shape (num_images, height, width)).
star_masks (np.ndarray, optional) – A batch of masks (same shape as star_postage_stamps). Each mask can have: - 0 to ignore the pixel. - 1 to fully consider the pixel. - Values in (0,1] as weights for partial consideration. Defaults to None.
pixel_sampling (float, optional) – The pixel size in meters. Defaults to 12e-6 m (12 microns).
reference_shifts (list[float], optional) – The target centroid shifts in pixels, specified as [dy, dx]. Defaults to [-1/3, -1/3] (nominal Euclid conditions).
sigma_init (float, optional) – Initial standard deviation for centroid estimation. Default is 2.5.
n_iter (int, optional) – Number of iterations for centroid refinement. Default is 20.

Returns:

An array of shape (num_images, 2), where: - Column 0 contains Zk1 (tip) values. - Column 1 contains Zk2 (tilt) values.

Return type:

np.ndarray

Notes

This function processes all images at once using vectorized operations.
The Zernike coefficients are computed in the WaveDiff convention.

wf_psf.data.data_zernike_utils.defocus_to_zk4_wavediff(dz, tel_focal_length=24.5, tel_diameter=1.2)[source]

Compute Zernike 4 value for a given defocus in WaveDiff conventions.

All inputs should be in [m].

The output zernike coefficient is in [um] units as expected by wavediff.

Parameters:

dz (float) – Shift in the z-axis, perpendicular to the focal plane. Units in [m].
tel_focal_length (float) – Telescope focal length in [m].
tel_diameter (float) – Telescope aperture diameter in [m].

wf_psf.data.data_zernike_utils.defocus_to_zk4_zemax(dz, tel_focal_length=24.5, tel_diameter=1.2)[source]

Compute Zernike 4 value for a given defocus in zemax conventions.

All inputs should be in [m].

Parameters:

dz (float) – Shift in the z-axis, perpendicular to the focal plane. Units in [m].
tel_focal_length (float) – Telescope focal length in [m].
tel_diameter (float) – Telescope aperture diameter in [m].

wf_psf.data.data_zernike_utils.pad_contribution_to_order(contribution: ndarray, max_order: int) → ndarray[source]

Pad a Zernike contribution array to the max Zernike order.

Parameters:

contribution (np.ndarray) – Array of shape (n_samples, n_orders) representing Zernike contributions.
max_order (int) – Target number of Zernike order; determines the number of columns after padding

Returns:

Padded array of shape (n_samples, max_order) with zeros appended to the right if max_order > current number of orders.

Return type:

np.ndarray

wf_psf.data.data_zernike_utils.pad_tf_zernikes(zk_param: Tensor, zk_prior: Tensor, n_zks_total: int)[source]

Pad the Zernike coefficient tensors to match the specified total number of Zernikes.

Parameters:

zk_param (tf.Tensor) – Zernike coefficients for the parametric part. Shape [batch, n_zks_param, 1, 1].
zk_prior (tf.Tensor) – Zernike coefficients for the prior part. Shape [batch, n_zks_prior, 1, 1].
n_zks_total (int) – Total number of Zernikes to pad to.

Returns:

padded_zk_param (tf.Tensor) – Padded Zernike coefficients for the parametric part. Shape [batch, n_zks_total, 1, 1].
padded_zk_prior (tf.Tensor) – Padded Zernike coefficients for the prior part. Shape [batch, n_zks_total, 1, 1].

wf_psf.data.data_zernike_utils.shift_x_y_to_zk1_2_wavediff(dxy, tel_focal_length=24.5, tel_diameter=1.2)[source]

Compute Zernike 1(2) for a given shifts in x(y) in WaveDiff conventions.

All inputs should be in [m]. A displacement of, for example, 0.5 pixels should be scaled with the corresponding pixel scale, e.g. 12[um], to get a displacement in [m], which would be dxy=0.5*12e-6.

The output zernike coefficient is in [um] units as expected by wavediff.

To apply match the centroid with a dx that has a corresponding zk1, the new PSF should be generated with -zk1.

The same applies to dy and zk2.

Parameters:

dxy (float) – Centroid shift in [m]. It can be on the x-axis or the y-axis.
tel_focal_length (float) – Telescope focal length in [m].
tel_diameter (float) – Telescope aperture diameter in [m].