PSF Subtraction#

Overview about all models used in 4S. The classical way to use 4S is based on the Class explained in FourS. All other models are used internally and should usually not be used directly.

FourS#

class fours.models.psf_subtraction.FourS(science_cube, adi_angles, psf_template, noise_model_lambda, psf_fwhm=None, right_reason_mask_factor=1.5, rotation_grid_subsample=1, device=0, work_dir=None, verbose=True)#

Bases: object

This is the main class that should be used to perform PSF subtraction with the FourS algorithm. It combines the noise model, the normalization model, and the rotation model. For an example of how to use this class, see the Examples.

__init__(science_cube, adi_angles, psf_template, noise_model_lambda, psf_fwhm=None, right_reason_mask_factor=1.5, rotation_grid_subsample=1, device=0, work_dir=None, verbose=True)#

Initializes the FourS class for PSF subtraction.

Parameters

  • science_cube (ndarray) – A 3D numpy array representing the science data cube. Shape: (n_frames, image_size, image_size).

  • adi_angles (ndarray) – A 1D array of parallactic angles in radians.

  • psf_template (ndarray) – A 2D numpy array representing the PSF template.

  • noise_model_lambda (float) – Regularization parameter for controlling the L2 penalty applied to the weights of the noise model. This is the most important hyperparameter for the noise model. Start with a very large value (10^8) and decrease it by a factor of 10. Large values correspond to strong regularization (weak noise model, small risk of overfitting).

  • psf_fwhm (Optional[float]) – Full width at half maximum (FWHM) of the PSF. If not provided, it is estimated using the psf_template.

  • right_reason_mask_factor (float) – Size of the right reason mask in units of the FWHM of the PSF.

  • rotation_grid_subsample (int) – Sub-sampling factor for the rotation grid in the ADI process. Can be used if the GPU memory is not sufficient for the full resolution.

  • device (Union[int, str]) – GPU device identifier for computation. Use 0 for the first GPU, “cuda:1” for the second GPU, or “cpu” for CPU computation.

  • work_dir (Optional[str]) – Directory path to save models, residuals, and logs. If given, a tensorboard log is created. See documentation of tensorboard for an example how to use it.

  • verbose (bool) – If True, prints progress information during the computation.

classmethod create_from_checkpoint(noise_model_file, normalization_model_file, s4_work_dir, science_cube, adi_angles, psf_template, device, verbose=True)#

Create a FourS object from a saved checkpoint. This can be used to restore a trained FourS model for further processing. It is possible to restore the noise and continue the training with a stronger regularization.

Parameters

  • noise_model_file (str) – Path to the saved noise model checkpoint.

  • normalization_model_file (Optional[Union[str, Path]]) – Path to the saved normalization model checkpoint, or None if no normalization is required.

  • s4_work_dir (Optional[Union[str, Path]]) – Directory for saving/loading models and residuals. See __init__ for more information.

  • science_cube (ndarray) – The 3D numpy array representing the science data cube (frames x size x size).

  • adi_angles (ndarray) – Array of parallactic angles in radians.

  • psf_template (ndarray) – A 2D numpy array representing the PSF template.

  • device (Union[int, str]) – Device identifier for computation (e.g., “cuda:0”, “cuda:1”, or “cpu”).

  • verbose (bool) – If True, progress information will be printed.

Return type

FourS

Returns

A loaded FourS instance with the restored noise and normalization models.

save_models(file_name_noise_model, file_name_normalization_model)#

Saves the noise model and normalization model to the specified file paths.

Parameters

  • file_name_noise_model (str) – File name of the noise model to save.

  • file_name_normalization_model (str) – File name of the normalization model to save.

Return type

None

restore_models(file_noise_model=None, file_normalization_model=None)#

Restores the noise and normalization models from checkpoint files.

Parameters

  • file_noise_model (Optional[str]) – Path to the saved noise model file. If None, the noise model is not restored.

  • file_normalization_model (Optional[str]) – Path to the saved normalization model file. If None, the normalization model is not restored.

Returns

This method does not return any value. It updates the

noise_model and/or normalization_model in place.

Return type

None

fit_noise_model(num_epochs, training_name='', logging_interval=1, optimizer=None, optimizer_kwargs=None)#

Fits the noise model using the specified optimizer and training parameters. Call this function before calculating the residuals using the function compute_residuals.

Parameters

  • num_epochs (int) – The total number of epochs to train the noise model. The training process should converge within the number of epochs specified. Please check the tensorboard logs for the convergence of the training process. It is possible to resume training by calling this function again. You can even increase the regularization parameter to improve the noise model.

  • training_name (str) – A string to name the current training session. This string will be used for TensorBoard logs and residual folders.

  • logging_interval (int) – Number of epochs between logging residuals and losses to TensorBoard and saving residual FITS files.

  • optimizer (Optional[Optimizer]) – An optional optimizer for training. If not provided, a default LBFGS optimizer is used.

  • optimizer_kwargs (Optional[Dict]) – Optional arguments to initialize the optimizer. These include parameters such as “max_iter” and “history_size” for LBFGS. If the GPU memory is limited consider reducing the “history_size” parameter.

Returns

This method performs in-place updates to the noise model parameters and logs the training progress or residual results.

Return type

None

compute_residuals()#

Computes the final residuals for the science data by applying the trained noise and rotation models to the normalized input data.

Returns

  1. The mean residual image created by averaging the residuals in the aligned frame.

  2. The median residual image computed as the pixel-wise median of the aligned residual sequence.

Return type

A tuple containing two 2D arrays

Linear Noise Model#

class fours.models.noise.FourSNoise(data_image_size, psf_template, lambda_reg, cut_radius_psf, right_reason_mask_radius, convolve=True)#

Bases: Module

This class implements the noise model for the FourS algorithm. The noise model consists of a linear model which is masked by a right reason mask. Further, it is regularized by convolving the weights with a PSF template and a simple L2 regularization term (lambda_reg).

__init__(data_image_size, psf_template, lambda_reg, cut_radius_psf, right_reason_mask_radius, convolve=True)#

Initializes the FourSNoise noise model with the given parameters.

Parameters

  • data_image_size (int) – The size of the image data. Only square data supported. (pixels)

  • psf_template (ndarray) – A numpy array representing the point spread function (PSF), which is used to convolve the weights.

  • lambda_reg (float) – Regularization parameter for controlling the L2 penalty applied to the weights. This is the most important hyperparameter for the noise model. Start with a very large value and decrease it by a factor of 10. Large values correspond to strong regularization (weak noise model, small risk of overfitting).

  • cut_radius_psf (float) – The radius used to crop the PSF template.

  • right_reason_mask_radius (float) – Radius to construct the right reason mask (pixel). Should be around 1.5 times the FWHM of the PSF.

  • convolve (bool) – If True, enables convolution of weights with the PSF template; This is the default behavior.

save(file_path)#

Saves the noise model to a file. :type file_path: str :param file_path: The path to the file where the noise model should be saved.

Return type

None

classmethod load(file_path)#

Factory method to load a noise model from a file.

Parameters

file_path (str) – The path to the file where the noise model is saved.

Return type

FourSNoise

Returns

A FourSNoise object.

property betas#

Returns the betas after applying the right reason mask and the PSF template.

compute_betas()#

Called within the forward pass to compute the betas.

forward(science_norm_flatten)#

Performs a forward pass to predict the noise estimate for the given science data.

Parameters

science_norm_flatten (Tensor) – Input tensor of shape (time, x*y), where time is the temporal dimension and x*y is the flattened, already normalized image data.

Return type

Tensor

Returns

Predicted noise estimates with a shape of (time, x*y).

Image Rotation#

This class makes it possible to back-propagate through the ADI rotation using Spatial Transformer Networks. It might be useful for some applications.

class fours.models.rotation.FieldRotationModel(all_angles, input_size, subsample=1, inverse=False, register_grid=False)#

Bases: Module

This class implements a field rotation model which rotates the input data by a sequence of angles. This model is needed to back-propagate through the rotation as needed for the 4S loss function.

__init__(all_angles, input_size, subsample=1, inverse=False, register_grid=False)#

Initializes the FieldRotationModel with the given parameters.

Parameters

  • all_angles (ndarray) – Predefined rotation angles used to create the affine grid. This avoids expensive grid recalculation and improves computational efficiency.

  • input_size (int) – The size of the input frames, assumed to have equal height and width.

  • subsample (int) – Reduces the number of angles by sampling every nth angle based on this value. Defaults to 1 for no subsampling. This is useful for speeding up the forward pass and reduce GPU memory consumption.

  • inverse (bool) – If True, applies the inverse rotation using negative angles. Defaults to False.

  • register_grid (bool) – Specifies whether to precompute and register the affine grid as a buffer. If it is set as a buffer, the full grid will be moved to the GPU during the forward pass. Defaults to False.

forward(frame_stack, parang_idx=None, new_angles=None, output_dimensions=None)#

Rotates the input data by the specified angles using the predefined affine grid. The rotation is performed either using parang_idx or new_angles, which are mutually exclusive.

Parameters

  • frame_stack (Tensor) – Batch of input frames to be rotated, of shape (batch_size, channels, height, width).

  • parang_idx (Optional[Tensor]) – Indices to select precomputed affine grids for the rotation. Required if new_angles is not specified.

  • new_angles (Optional[Tensor]) – Custom angles in degrees to build new affine grids for rotation. Expensive operation, used for test data or finer control. Cannot be used with parang_idx.

  • output_dimensions – Target size for up/downscaling the rotated frames, specified as a tuple (height, width). Defaults to the model’s input size.

Return type

Tensor

Returns

Rotated tensor of shape (batch_size, channels, target_height, target_width), with frames rotated by given or indexed angles.

Normalisation#

class fours.models.normalization.FourSFrameNormalization(image_size, normalization_type='normal')#

Bases: Module

This class is used to normalize the science data.

property image_size#
save(file_path)#

Saves the normalization parameters and metadata to the given file path. The state dictionary includes the type of normalization, the image size, and precomputed statistics.

Parameters

file_path (str) – The path to save the parameters and metadata.

Return type

None

classmethod load(file_path)#

Factory method to load a normalization object from a file.

Parameters

file_path – Path to the file storing the normalization parameters and metadata.

Return type

FourSFrameNormalization

Returns

FourSFrameNormalization instance, initialized with loaded parameters.

prepare_normalization(science_data)#

Computes and sets the normalization parameters (mean and standard deviation or robust metrics) based on the provided science data.

Parameters

science_data (Tensor) – A 3D tensor containing science data frames used for normalization. Dimensions are expected as [frames, height, width].

Return type

None

Returns

None. The normalization parameters are updated and stored within the instance.

normalize_data(science_data, re_center=True)#

Normalizes the provided science data using prepared parameters.

Parameters

  • science_data (Tensor) – Input 3D tensor containing science data, with dimensions [frames, height, width]. Expected to match the dimensions used in prepare_normalization.

  • re_center (bool) – Optional; if True, the mean is subtracted from the data before normalization. Defaults to True.

Return type

Tensor

Returns

A tensor of normalized science data, where each value is scaled by the computed standard deviation or IQR. Any NaNs are replaced with 0.

forward(science_data)#

Performs forward pass by normalizing the provided science data tensor.

Parameters

science_data (Tensor) – A 3D tensor containing input science data with dimensions [frames, height, width]. Data should match the format used in prepare_normalization.

Return type

Tensor

Returns

A normalized tensor where the values are scaled by the computed normalization parameters and NaNs are replaced with 0.