Merge pull request #142 from CosmoStat/case_study_psf_decontamination

WaveDiff: New PSF Model Features, Refactoring, Rotation Obscuration Mask, and PEP8 Compliance

This commit merges the 'case_study_psf_decontamination' feature/refactor branch into `develop`, introducing 
significant new features, architectural changes, bug fixes, and CI/tooling improvements.

Highlights:

New Features
---------------
- Added physical Zernike prior layer to PSF model (Closes #123)
- Added obscuration mask rotation support to PSF models
- Included module for simulating spatially-varying PSFs with validation tests (Closes #116, #133)
- Integrated centroid estimation correction (Closes #113)
- Corrected CCD z-axis misalignment handling
- Added phase retrieval module with random seed support (Closes #91)
- Added new data generation script to support structured model validation

Bug Fixes
-----------
- Fixed random seed bug in data-driven PSF reinitialization (Closes #132)

Refactors and Internal Architecture
-------------------------------------
- Replaced `TrainingDataHandler` and `TestDataHandler` with unified `DataHandler`
- Refactored `train_utils.py` and `train.py` using the Strategy Pattern (Closes #89, #152)
- Introduced helper functions and modular callbacks
- Added custom loss/metric functions for masked datasets (e.g. for centroid and sample weights)
- Renamed functions, classes, and variables to comply with PEP8 (Closes #139)
- Removed deprecated code and unused imports

Testing and Validation
-------------------------
- Added `test_training/` with unit tests for `train.py` and `train_utils.py`

CI / Tooling Updates
------------------------
- Replaced `black` with `ruff` for code formatting (Closes #119, #120, #145)
- Unpinned `pytest` (Closes #125)
- Increased TensorFlow version to >=2.11.0 (Closes #90)

⚠️ Notes
--------
- This commit is large and includes many interdependent changes. Follow-up work includes:
  - Further refactoring of the data generation script
  - Cleanup of temporary or duplicate logic introduced during the transition
  - Finalisation of naming schemes and docstrings across modules
  - Preparation of a release-ready, slimmed-down develop branch

Co-authored-by many developers

Author: jeipollack

.github/workflows/ci.yml

@@ -28,15 +28,7 @@ jobs:
 
       - name: Install dependencies
         run: python -m pip install ".[test]"
-
+      
       - name: Test with pytest
         run: python -m pytest
 
-      # Add Black formatter
-      - name: Install Black formatter
-        run: python -m pip install black
-      
-      - name: Check code formatting with Black
-        run: black . --check --diff
-
-

config/data_config.yaml

@@ -1,8 +1,9 @@
 # Training and test data sets for training and/or metrics evaluation
 data:
   training:
-    # Specify directory path to data; Default setting is /path/to/repo/data
+    # Specify directory path to training dataset
     data_dir: data/coherent_euclid_dataset/
+    # Provide name of training dataset
     file: train_Euclid_res_200_TrainStars_id_001.npy
     # if training data set file does not exist, generate a new one by setting values below
     stars: null
@@ -26,7 +27,9 @@ data:
       euclid_obsc: true
       n_stars: 200
   test:
+    # Specify directory path to training dataset
     data_dir: data/coherent_euclid_dataset/
+    # Provide name of test dataset
     file: test_Euclid_res_id_001.npy
     # If test data set file not provided produce a new one
     stars: null

config/metrics_config.yaml

@@ -1,46 +1,85 @@
 metrics:
    # Specify the type of model weights to load by entering "psf_model" to load weights of final psf model or "checkpoint" to load weights from a checkpoint callback.
   model_save_path: <enter psf_model or checkpoint>
+  
   # Choose the training cycle for which to evaluate the psf_model. Can be: 1, 2, ...
   saved_training_cycle: 2
+  
   # Metrics-only run: Specify model_params for a pre-trained model else leave blank if running training + metrics
   # Specify path to Parent Directory of Trained Model 
   trained_model_path: </path/to/parent/directory/of/trained/model>
+  
   # Name of the Trained Model Config file stored in config sub-directory in the trained_model_path parent directory
   trained_model_config: <enter name of trained model config file>
-  #Evaluate the monchromatic RMSE metric.
-  eval_mono_metric_rmse: True
-  #Evaluate the OPD RMSE metric.
-  eval_opd_metric_rmse: True
-  #Evaluate the super-resolution and the shape RMSE metrics for the train dataset.
-  eval_train_shape_sr_metric_rmse: True
+  
+  # Evaluate the monchromatic RMSE metric.
+  eval_mono_metric: True
+  
+  # Evaluate the OPD RMSE metric.
+  eval_opd_metric: True
+  
+  # Evaluate the super-resolution and the shape RMSE metrics for the train dataset.
+  eval_train_shape_results_dict: False
+
+  # Evaluate the super-resolution and the shape RMSE metrics for the test dataset.
+  eval_test_shape_results_dict: False
+  
   # Name of Plotting Config file - Enter name of yaml file to run plot metrics else if empty run metrics evaluation only
   plotting_config: <enter name of plotting_config .yaml file or leave empty>
+  
   ground_truth_model:
     model_params:
-      #Model used as ground truth for the evaluation. Options are: 'poly' for polychromatic and 'physical' [not available].
-      model_name: poly 
+      # PSF model used as ground truth for the evaluation. Options are: 'ground_truth_poly' for polychromatic and 'ground_truth_physical_poly' for polychromatic model with the physical layer extension PSF models.
+      model_name: <ground_truth_poly or ground_truth_physical_poly>
 
       # Evaluation parameters
-      #Number of bins used for the ground truth model poly PSF generation
+      # Number of bins used for the ground truth model poly PSF generation
       n_bins_lda: 20
 
-      #Downsampling rate to match the oversampled model to the specified telescope's sampling.
+      # Downsampling rate to match the oversampled model to the specified telescope's sampling.
       output_Q: 3
 
-      #Oversampling rate used for the OPD/WFE PSF model.
+      # Oversampling rate used for the OPD/WFE PSF model.
       oversampling_rate: 3 
 
-      #Dimension of the pixel PSF postage stamp
+      # Dimension of the pixel PSF postage stamp
       output_dim: 32
 
-      #Dimension of the OPD/Wavefront space."
+      # Dimension of the OPD/Wavefront space."
       pupil_diameter: 256
 
-      #Boolean to define if we use sample weights based on the noise standard deviation estimation
+      # Top-hat filter to avoid the aliasing effect in the obscuration mask
+      LP_filter_length: 2
+
+      # Boolean to define if we use sample weights based on the noise standard deviation estimation
       use_sample_weights: True 
 
-      #Interpolation type for the physical poly model. Options are: 'none', 'all', 'top_K', 'independent_Zk'."
+      # Flag to use Zernike prior
+      use_prior: False
+
+      # Correct centroids
+      correct_centroids: False
+
+      # Sigma of the window function used to compute the centroids. Default Euclid value is 2.5
+      sigma_centroid_window: 2.5
+
+      # Default reference_shifts value for observations at Euclid conditions, i.e., pixel sampling and telescope parameters.
+      reference_shifts: [-1/3, -1/3]
+
+      # Rotation angle (in degrees) for the obscuration mask.
+      # Must be a multiple of 90 (0, 90, 180, 270, etc.).
+      # Rotation is counterclockwise.
+      obscuration_rotation_angle: 0
+
+      # Consider CCD missalignments
+      add_ccd_misalignments: False
+
+      # CCD missalignments input file path
+      # This should be refactored. It might be better to directly look for the `tiles.npy` in
+      # the `data/assets/` directory in the repository
+      ccd_misalignments_input_path:
+      
+      # Interpolation type for the physical poly model. Options are: 'none', 'all', 'top_K', 'independent_Zk'."
       interpolation_type: None
 
       # SED intepolation points per bin
@@ -55,68 +94,68 @@ metrics:
       # Standard deviation of the multiplicative SED Gaussian noise.
       sed_sigma: 0
 
-      #Limits of the PSF field coordinates for the x axis.
+      # Limits of the PSF field coordinates for the x axis.
       x_lims: [0.0, 1.0e+3]
 
-      #Limits of the PSF field coordinates for the y axis.
+      # Limits of the PSF field coordinates for the y axis.
       y_lims: [0.0, 1.0e+3]
 
-      # Hyperparameters for Parametric model 
+      # Hyperparameters for the Parametric model 
       param_hparams:
         # Random seed for Tensor Flow Initialization
         random_seed: 3877572
 
         # Parameter for the l2 loss function for the Optical path differences (OPD)/WFE
         l2_param: 0.
 
-        #Zernike polynomial modes to use on the parametric part.
+        # Zernike polynomial modes to use on the parametric part.
         n_zernikes: 45
 
-        #Max polynomial degree of the parametric part.
+        # Max polynomial degree of the parametric part.
         d_max: 2  
 
-        #Flag to save optimisation history for parametric model
+        # Flag to save optimisation history for parametric model
         save_optim_history_param: true
 
       # Hyperparameters for non-parametric model
       nonparam_hparams:
-        #Max polynomial degree of the non-parametric part. 
+        # Max polynomial degree of the non-parametric part. 
         d_max_nonparam: 5 
 
         # Number of graph features
         num_graph_features: 10
 
-        #L1 regularisation parameter for the non-parametric part."
+        # L1 regularisation parameter for the non-parametric part."
         l1_rate: 1.0e-8
 
-        #Flag to enable Projected learning for DD_features to be used with `poly` or `semiparametric` model.
+        # Flag to enable Projected learning for DD_features to be used with `poly` or `semiparametric` model.
         project_dd_features: False
 
-        #Flag to reset DD_features to be used with `poly` or `semiparametric` model
+        # Flag to reset DD_features to be used with `poly` or `semiparametric` model
         reset_dd_features: False
 
-        #Flag to save optimisation history for non-parametric model
+        # Flag to save optimisation history for non-parametric model
         save_optim_history_nonparam: True 
 
   metrics_hparams:
     # Batch size to use for the evaluation.
     batch_size: 16
 
-    #Save RMS error for each super resolved PSF in the test dataset in addition to the mean across the FOV."
-    #Flag to get Super-Resolution pixel PSF RMSE for each individual test star. 
-    #If `True`, the relative pixel RMSE of each star is added to ther saving dictionary.
+    # Save RMS error for each super resolved PSF in the test dataset in addition to the mean across the FOV."
+    # Flag to get Super-Resolution pixel PSF RMSE for each individual test star. 
+    # If `True`, the relative pixel RMSE of each star is added to ther saving dictionary.
     opt_stars_rel_pix_rmse: False
 
     ## Specific parameters
     # Parameter for the l2 loss of the OPD.
     l2_param: 0.
 
     ## Define the resolution at which you'd like to measure the shape of the PSFs
-    #Downsampling rate from the high-resolution pixel modelling space.
+    # Downsampling rate from the high-resolution pixel modelling space.
     # Recommended value: 1
     output_Q: 1
 
-    #Dimension of the pixel PSF postage stamp; it should be big enough so that most of the signal is contained inside the postage stamp. 
+    # Dimension of the pixel PSF postage stamp; it should be big enough so that most of the signal is contained inside the postage stamp. 
     # It also depends on the Q values used.
     # Recommended value: 64 or higher
     output_dim: 64