wf_psf.psf_models.models.psf_model_physical_polychromatic

PSF Model Physical Semi-Parametric Polychromatic.

A module which defines the classes and methods to manage the parameters of the psf physical polychromatic model.

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

Classes

`PhysicalPolychromaticFieldFactory`()	Factory class for the TensorFlow Physical Polychromatic PSF Field Model.
`TFPhysicalPolychromaticField`(args, *kwargs)	TensorFlow Physical Polychromatic PSF Field class.

class wf_psf.psf_models.models.psf_model_physical_polychromatic.PhysicalPolychromaticFieldFactory[source]

Bases: PSFModelBaseFactory

Factory class for the TensorFlow Physical Polychromatic PSF Field Model.

This factory class is responsible for instantiating instances of the TensorFlow Physical Polychromatic PSF Field Model. It is registered with the PSF model factory registry.

Parameters:: ids (tuple) – A tuple containing identifiers for the factory class.

get_model_instance(model_params, training_params, data=None, coeff_mat=None)[source]: Instantiates an instance of the TensorFlow Physical Polychromatic Field class with the provided parameters.

Methods

get_model_instance(model_params, ...[, ...])

Create an instance of the TensorFlow Physical Polychromatic Field model.

get_model_instance(model_params, training_params, data, coeff_mat=None)[source]

Create an instance of the TensorFlow Physical Polychromatic Field model.

This method instantiates a TFPhysicalPolychromaticField object with the given model and training parameters, and data containing prior information like Zernike coefficients. Optionally, a coefficient matrix can be provided.

Parameters:

model_params (Recursive Namespace) – A Recursive Namespace object containing parameters for this PSF model class.
training_params (Recursive Namespace) – A Recursive Namespace object containing training hyperparameters for this PSF model class.
data (DataConfigHandler) – A DataConfigHandler object that provides access to training and test datasets, as well as prior knowledge like Zernike coefficients.
coeff_mat (Tensor or None, optional) – Coefficient matrix defining the parametric PSF field model.

Returns:

An instance of the TensorFlow Physical Polychromatic Field model.

Return type:

TFPhysicalPolychromaticField

ids = ('physical_poly',)

class wf_psf.psf_models.models.psf_model_physical_polychromatic.TFPhysicalPolychromaticField(*args, **kwargs)[source]

Bases: Model

TensorFlow Physical Polychromatic PSF Field class.

This class represents a polychromatic PSF field model with a physical layer. It incorporates parametric and non-parametric modeling approaches to accurately reconstruct the point spread function (PSF) across multiple wavelengths.

The model provides functionalities for: - Initializing model parameters and defining the physical PSF layer. - Performing forward passes and computing wavefront transformations. - Handling Zernike parameterization and coefficient matrices. - Evaluating model performance and saving optimization history.

See individual method docstrings for more details.

Attributes:

activity_regularizer: Optional regularizer function for the output of this layer.
compute_dtype: The dtype of the layer’s computations.
distribute_reduction_method: The method employed to reduce per-replica values during training.
distribute_strategy: The tf.distribute.Strategy this model was created under.
dtype: The dtype of the layer weights.
dtype_policy: The dtype policy associated with this layer.
dynamic: Whether the layer is dynamic (eager-only); set in the constructor.
inbound_nodes: Return Functional API nodes upstream of this layer.
input: Retrieves the input tensor(s) of a layer.
input_mask: Retrieves the input mask tensor(s) of a layer.
input_shape: Retrieves the input shape(s) of a layer.
input_spec: InputSpec instance(s) describing the input format for this layer.
layers
losses: List of losses added using the add_loss() API.
metrics: Returns the model’s metrics added using compile(), add_metric() APIs.
metrics_names: Returns the model’s display labels for all outputs.
n_zks_total: Get the total number of Zernike coefficients.
name: Name of the layer (string), set in the constructor.
name_scope: Returns a tf.name_scope instance for this class.
non_trainable_variables: Sequence of non-trainable variables owned by this module and its submodules.
non_trainable_weights: List of all non-trainable weights tracked by this layer.
obscurations
opd_dim
outbound_nodes: Return Functional API nodes downstream of this layer.
output: Retrieves the output tensor(s) of a layer.
output_mask: Retrieves the output mask tensor(s) of a layer.
output_shape: Retrieves the output shape(s) of a layer.
run_eagerly: Settable attribute indicating whether the model should run eagerly.
save_nonparam_history: Check if the model should save the optimization history for non-parametric features.
save_param_history: Check if the model should save the optimization history for parametric features.
state_updates: Deprecated, do NOT use!
stateful
submodules: Sequence of all sub-modules.
supports_masking: Whether this layer supports computing a mask using compute_mask.
tf_np_poly_opd: Lazy loading of the non-parametric polynomial variations OPD layer.
tf_physical_layer: Lazy loading of the physical layer of the PSF model.
tf_poly_Z_field: Lazy loading of the polynomial Zernike field layer.
tf_zernike_OPD: Lazy loading of the Zernike Optical Path Difference (OPD) layer.
trainable
trainable_variables: Sequence of trainable variables owned by this module and its submodules.
trainable_weights: List of all trainable weights tracked by this layer.
updates
variable_dtype: Alias of Layer.dtype, the dtype of the weights.
variables: Returns the list of all layer variables/weights.
weights: Returns the list of all layer variables/weights.
zernike_maps: Get Zernike maps.
zks_total_contribution

Methods

`add_loss`(losses, **kwargs)	Add loss tensor(s), potentially dependent on layer inputs.
`add_metric`(value[, name])	Adds metric tensor to the layer.
`add_update`(updates)	Add update op(s), potentially dependent on layer inputs.
`add_variable`(args, *kwargs)	Deprecated, do NOT use! Alias for add_weight.
`add_weight`([name, shape, dtype, ...])	Adds a new variable to the layer.
`assign_coeff_matrix`(coeff_mat)	Assign a coefficient matrix to the parametric PSF field model.
`build`(input_shape)	Builds the model based on input shapes received.
`call`(inputs[, training])	Define the PSF (Point Spread Function) field forward model.
`compile`([optimizer, loss, metrics, ...])	Configures the model for training.
`compute_loss`([x, y, y_pred, sample_weight])	Compute the total loss, validate it, and return it.
`compute_mask`(inputs[, mask])	Computes an output mask tensor.
`compute_metrics`(x, y, y_pred, sample_weight)	Update metric states and collect all metrics to be returned.
`compute_output_shape`(input_shape)	Computes the output shape of the layer.
`compute_output_signature`(input_signature)	Compute the output tensor signature of the layer based on the inputs.
`compute_zernikes`(input_positions)	Compute Zernike coefficients at a batch of positions.
`count_params`()	Count the total number of scalars composing the weights.
`evaluate`([x, y, batch_size, verbose, ...])	Returns the loss value & metrics values for the model in test mode.
`evaluate_generator`(generator[, steps, ...])	Evaluates the model on a data generator.
`finalize_state`()	Finalizes the layers state after updating layer weights.
`fit`([x, y, batch_size, epochs, verbose, ...])	Trains the model for a fixed number of epochs (iterations on a dataset).
`fit_generator`(generator[, steps_per_epoch, ...])	Fits the model on data yielded batch-by-batch by a Python generator.
`from_config`(config[, custom_objects])	Creates a layer from its config.
`get_coeff_matrix`()	Get coefficient matrix.
`get_config`()	Returns the config of the Model.
`get_input_at`(node_index)	Retrieves the input tensor(s) of a layer at a given node.
`get_input_mask_at`(node_index)	Retrieves the input mask tensor(s) of a layer at a given node.
`get_input_shape_at`(node_index)	Retrieves the input shape(s) of a layer at a given node.
`get_layer`([name, index])	Retrieves a layer based on either its name (unique) or index.
`get_metrics_result`()	Returns the model's metrics values as a dict.
`get_output_at`(node_index)	Retrieves the output tensor(s) of a layer at a given node.
`get_output_mask_at`(node_index)	Retrieves the output mask tensor(s) of a layer at a given node.
`get_output_shape_at`(node_index)	Retrieves the output shape(s) of a layer at a given node.
`get_weight_paths`()	Retrieve all the variables and their paths for the model.
`get_weights`()	Retrieves the weights of the model.
`load_weights`(filepath[, by_name, ...])	Loads all layer weights, either from a TensorFlow or an HDF5 weight file.
`make_predict_function`([force])	Creates a function that executes one step of inference.
`make_test_function`([force])	Creates a function that executes one step of evaluation.
`make_train_function`([force])	Creates a function that executes one step of training.
`pad_zernikes`(zk_param, zk_prior)	Pad the Zernike coefficients to match the maximum length.
`predict`(x[, batch_size, verbose, steps, ...])	Generates output predictions for the input samples.
`predict_generator`(generator[, steps, ...])	Generates predictions for the input samples from a data generator.
`predict_mono_psfs`(input_positions, ...)	Predict a set of monochromatic Point Spread Functions (PSFs) at desired positions.
`predict_on_batch`(x)	Returns predictions for a single batch of samples.
`predict_opd`(input_positions)	Predict the OPD at some positions.
`predict_step`(data[, evaluate_step])	Predict (inference) step.
`predict_zernikes`(input_positions)	Predict Zernike coefficients at a batch of positions.
`reset_metrics`()	Resets the state of all the metrics in the model.
`save`(filepath[, overwrite, ...])	Saves the model to Tensorflow SavedModel or a single HDF5 file.
`save_spec`([dynamic_batch])	Returns the tf.TensorSpec of call inputs as a tuple (args, kwargs).
`save_weights`(filepath[, overwrite, ...])	Saves all layer weights.
`set_nonzero_nonparam`()	Set the non-parametric part to non-zero values.
`set_output_Q`(output_Q[, output_dim])	Set the output sampling rate (output_Q) for PSF generation.
`set_trainable_layers`([param_bool, nonparam_bool])	Set the layers to be trainable.
`set_weights`(weights)	Sets the weights of the layer, from NumPy arrays.
`set_zero_nonparam`()	Set the non-parametric part of the OPD (Optical Path Difference) to zero.
`summary`([line_length, positions, print_fn, ...])	Prints a string summary of the network.
`test_on_batch`(x[, y, sample_weight, ...])	Test the model on a single batch of samples.
`test_step`(data)	The logic for one evaluation step.
`to_json`(**kwargs)	Returns a JSON string containing the network configuration.
`to_yaml`(**kwargs)	Returns a yaml string containing the network configuration.
`train_on_batch`(x[, y, sample_weight, ...])	Runs a single gradient update on a single batch of data.
`train_step`(data)	The logic for one training step.
`with_name_scope`(method)	Decorator to automatically enter the module name scope.

__call__
get_obs_pos
reset_states

assign_coeff_matrix(coeff_mat: Tensor | None) → None[source]

Assign a coefficient matrix to the parametric PSF field model.

This method updates the coefficient matrix used by the parametric PSF field model, allowing for customization or modification of the model’s parameters. If coeff_mat is None, the model will revert to using its default coefficient matrix.

Parameters:: coeff_mat (Optional[tf.Tensor]) – A TensorFlow tensor representing the coefficient matrix for the PSF field model. If None, the model will use the default coefficient matrix.
Return type:: None

call(inputs, training=True)[source]

Define the PSF (Point Spread Function) field forward model.

This method defines the forward model of the PSF field, which involves several steps: 1. Transforming input positions into Zernike coefficients. 2. Converting Zernike coefficients into Optical Path Difference (OPD) maps. 3. Combining OPD maps with Spectral Energy Distribution (SED) information to generate

polychromatic PSFs.

Parameters:

inputs (list) –
List containing input data required for PSF computation. It should contain two elements: - input_positions: Tensor [batch_dim, 2]

Positions at which to compute the PSFs.
- packed_SEDs: Tensor [batch_dim, …]
  Packed Spectral Energy Distributions (SEDs) for the corresponding positions.
training (bool, optional) – Indicates whether the model is being trained or used for inference. Defaults to True.

Returns:

poly_psfs – Polychromatic PSFs generated by the forward model.

Return type:

Tensor

Notes

The input_positions tensor should have a shape of [batch_dim, 2], where each row represents the x and y coordinates of a position.
The packed_SEDs tensor should have a shape of [batch_dim, …], containing the SED information for each position.
During training, this method computes the Zernike coefficients from the input positions and calculates the corresponding OPD maps. Additionally, it adds an L2 loss term based on the parametric OPD maps.
During inference, this method generates predictions using precomputed OPD maps or by propagating through the forward model.

Examples

# Usage during training inputs = [input_positions, packed_SEDs] poly_psfs = psf_model(inputs)

# Usage during inference inputs = [input_positions, packed_SEDs] poly_psfs = psf_model(inputs, training=False)

compute_zernikes(input_positions)[source]

Compute Zernike coefficients at a batch of positions.

This method computes the Zernike coefficients for a batch of input positions using both the parametric model and the physical layer.

Parameters:: input_positions (Tensor [batch_dim, 2]) – Positions for which to compute the Zernike coefficients.
Returns:: zernike_coefficients – Computed Zernike coefficients for the input positions.
Return type:: Tensor [batch, n_zks_total, 1, 1]

Notes

This method combines the predictions from both the parametric model and the physical layer to obtain the final Zernike coefficients.

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

get_obs_pos()[source]

property n_zks_total: Get the total number of Zernike coefficients.

property obscurations

property opd_dim

pad_zernikes(zk_param, zk_prior)[source]

Pad the Zernike coefficients to match the maximum length.

Pad the input Zernike coefficient tensors to match the length of the maximum number of Zernike coefficients among the parametric and prior parts.

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].

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].

predict_mono_psfs(input_positions, lambda_obs, phase_N)[source]

Predict a set of monochromatic Point Spread Functions (PSFs) at desired positions.

This method calculates monochromatic PSFs based on the provided input positions, observed wavelength, and required wavefront dimension.

Parameters:

input_positions (Tensor [batch_dim, 2]) – Positions at which to compute the PSFs.
lambda_obs (float) – Observed wavelength in micrometers (um).
phase_N (int) – Required wavefront dimension. This should be calculated using a SimPSFToolkit instance. Example: ` simPSF_np = wf.SimPSFToolkit(...) phase_N = simPSF_np.feasible_N(lambda_obs) `

Returns:

mono_psf_batch – Batch of monochromatic PSFs.

Return type:

Tensor

predict_opd(input_positions)[source]

Predict the OPD at some positions.

Parameters:: input_positions (Tensor [batch_dim, 2]) – Positions to predict the OPD.
Returns:: opd_maps – OPD at requested positions.
Return type:: Tensor [batch, opd_dim, opd_dim]

predict_step(data, evaluate_step=False)[source]

Predict (inference) step.

A method to enable a special type of interpolation (different from training) for the physical layer.

Parameters:

data (NOT SURE)
evaluate_step (bool) – Boolean flag to evaluate step

Returns:

Instance of TFBatchPolychromaticPSF class containing computed polychromatic PSFs.

Return type:

poly_psfs TFBatchPolychromaticPSF

predict_zernikes(input_positions)[source]

Predict Zernike coefficients at a batch of positions.

This method predicts the Zernike coefficients for a batch of input positions using both the parametric model and the physical layer. During training, the prediction from the physical layer is typically not used.

Parameters:: input_positions (Tensor [batch_dim, 2]) – Positions for which to predict the Zernike coefficients.
Returns:: zernike_coeffs – Predicted Zernike coefficients for the input positions.
Return type:: Tensor [batch, n_zks_total, 1, 1]

Notes

At training time, the prediction from the physical layer may not be utilized, as the model might be trained to rely solely on the parametric part.

property save_nonparam_history: bool: Check if the model should save the optimization history for non-parametric features.

property save_param_history: bool: Check if the model should save the optimization history for parametric features.

set_nonzero_nonparam()[source]

Set the non-parametric part to non-zero values.

This method sets the non-parametric component of the Optical Path Difference (OPD) to non-zero values, allowing it to contribute to the overall PSF (Point Spread Function).

set_output_Q(output_Q: float, output_dim: int | None = None) → None[source]

Set the output sampling rate (output_Q) for PSF generation.

This method updates the output_Q parameter, which defines the resampling factor for generating PSFs at different resolutions relative to the telescope’s native sampling. It also allows optionally updating output_dim, which sets the output resolution of the PSF model.

If output_dim is provided, the PSF model’s output resolution is updated. The method then reinitializes the batch polychromatic PSF generator to reflect the updated parameters.

Parameters:

output_Q (float) – The resampling factor that determines the output PSF resolution relative to the telescope’s native sampling.
output_dim (Optional[int], default=None) – The new output dimension for the PSF model. If None, the output dimension remains unchanged.

Return type:

None

set_trainable_layers(param_bool=True, nonparam_bool=True)[source]

Set the layers to be trainable.

A method to set layers to be trainable.

Parameters:

param_bool (bool) – Boolean flag for parametric model layers
nonparam_bool (bool) – Boolean flag for non-parametric model layers

set_zero_nonparam()[source]

Set the non-parametric part of the OPD (Optical Path Difference) to zero.

This method sets the non-parametric component of the Optical Path Difference (OPD) to zero, effectively removing its contribution from the overall PSF (Point Spread Function).

property tf_np_poly_opd: Lazy loading of the non-parametric polynomial variations OPD layer.

property tf_physical_layer: Lazy loading of the physical layer of the PSF model.

property tf_poly_Z_field: Lazy loading of the polynomial Zernike field layer.

property tf_zernike_OPD: Lazy loading of the Zernike Optical Path Difference (OPD) layer.

property zernike_maps: Get Zernike maps.

property zks_total_contribution