Draft: Feature/real data metrics conference original

config/inference_conf.yaml

@@ -0,0 +1,37 @@
+
+inference:
+  # Inference batch size
+  batch_size: 16
+  # Cycle to use for inference. Can be: 1, 2, ...
+  cycle: 2
+
+  # Paths to the configuration files and trained model directory
+  configs:
+    # Path to the directory containing the trained model
+    trained_model_path: /path/to/trained/model/
+
+    # Subdirectory name of the trained model, e.g. psf_model
+    model_subdir: model
+
+    # Relative Path to the training configuration file used to train the model
+    trained_model_config_path: config/training_config.yaml
+
+    # Path to the data config file (this could contain prior information)
+    data_config_path: 
+
+  # The following parameters will overwrite the `model_params` in the training config file.
+  model_params:
+    # Num of wavelength bins to reconstruct polychromatic objects.
+    n_bins_lda: 8
+
+    # Downsampling rate to match the oversampled model to the specified telescope's sampling.
+    output_Q: 1
+
+    # Dimension of the pixel PSF postage stamp
+    output_dim: 64
+
+    # Flag to perform centroid error correction
+    correct_centroids: False
+
+    # Flag to perform CCD misalignment error correction
+    add_ccd_misalignments: True

config/metrics_config.yaml

@@ -12,6 +12,9 @@ metrics:
   # 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 chi2 metric.
+  eval_chi2_metric: True
+
   # Evaluate the monchromatic RMSE metric.
   eval_mono_metric: True
 

data/generation/README.md

@@ -8,6 +8,16 @@ To run the script:
 python data_generation_script.py -c data_generation_params_v.0.1.0.yml
 ```
 
+> ⚠️ Warning 
+>
+> There are some differences with the original dataset used for the WaveDiff paper (Liaudat et al. 2023) even if we use it as a reference:
+> - The assingment of SEDs for each star will not match that one of the original dataset. Although the same templates are used.
+> - The assigned noise level (SNR) for each star will not match the original dataset. Although, the same distribution will be used.
+> 
+> Nevertheless, the `C_poly` will match and the `positions` will match. Therefore, results from the new datasets will not be exactly the same as in the original datasets.
+
+
+### Dataset description
 
 **Dataset 0.x.x:**
 
@@ -56,3 +66,5 @@ python data_generation_script.py -c data_generation_params_v.0.1.0.yml
 - v3.2.1/2/3/4 with dummy (unitary) masks
 - v3.3.1/2/3/4 with realistic masks
 
+
+

data/generation/data_generation_params_v.0.1.0.yml

@@ -106,7 +106,7 @@ dataset_features:
   # Centroid shift options
   add_intrapixel_shifts: False
   # In pixels (should be abs(limits)<0.5)
-  intrapixel_shift_range: [-0.5, 0.5]
+  intrapixel_shift_range: [-0.3, 0.3]
 
   # CCD misalignment options
   add_ccd_misalignments: False

data/generation/data_generation_params_v.0.2.0.yml

@@ -106,7 +106,7 @@ dataset_features:
   # Centroid shift options
   add_intrapixel_shifts: False
   # In pixels (should be abs(limits)<0.5)
-  intrapixel_shift_range: [-0.5, 0.5]
+  intrapixel_shift_range: [-0.3, 0.3]
 
   # CCD misalignment options
   add_ccd_misalignments: False

data/generation/data_generation_params_v.0.3.0.yml

@@ -106,7 +106,7 @@ dataset_features:
   # Centroid shift options
   add_intrapixel_shifts: False
   # In pixels (should be abs(limits)<0.5)
-  intrapixel_shift_range: [-0.5, 0.5]
+  intrapixel_shift_range: [-0.3, 0.3]
 
   # CCD misalignment options
   add_ccd_misalignments: False

data/generation/data_generation_params_v.1.1.0.yml

@@ -106,7 +106,7 @@ dataset_features:
   # Centroid shift options
   add_intrapixel_shifts: True
   # In pixels (should be abs(limits)<0.5)
-  intrapixel_shift_range: [-0.5, 0.5]
+  intrapixel_shift_range: [-0.3, 0.3]
 
   # CCD misalignment options
   add_ccd_misalignments: True

data/generation/data_generation_params_v.1.2.0.yml

@@ -106,7 +106,7 @@ dataset_features:
   # Centroid shift options
   add_intrapixel_shifts: True
   # In pixels (should be abs(limits)<0.5)
-  intrapixel_shift_range: [-0.5, 0.5]
+  intrapixel_shift_range: [-0.3, 0.3]
 
   # CCD misalignment options
   add_ccd_misalignments: True

data/generation/data_generation_params_v.1.3.0.yml

@@ -106,7 +106,7 @@ dataset_features:
   # Centroid shift options
   add_intrapixel_shifts: True
   # In pixels (should be abs(limits)<0.5)
-  intrapixel_shift_range: [-0.5, 0.5]
+  intrapixel_shift_range: [-0.3, 0.3]
 
   # CCD misalignment options
   add_ccd_misalignments: True

data/generation/data_generation_params_v.2.1.0.yml

@@ -106,7 +106,7 @@ dataset_features:
   # Centroid shift options
   add_intrapixel_shifts: True
   # In pixels (should be abs(limits)<0.5)
-  intrapixel_shift_range: [-0.5, 0.5]
+  intrapixel_shift_range: [-0.3, 0.3]
 
   # CCD misalignment options
   add_ccd_misalignments: True

data/generation/data_generation_params_v.2.2.0.yml

@@ -106,7 +106,7 @@ dataset_features:
   # Centroid shift options
   add_intrapixel_shifts: True
   # In pixels (should be abs(limits)<0.5)
-  intrapixel_shift_range: [-0.5, 0.5]
+  intrapixel_shift_range: [-0.3, 0.3]
 
   # CCD misalignment options
   add_ccd_misalignments: True

data/generation/data_generation_params_v.2.3.0.yml

@@ -106,7 +106,7 @@ dataset_features:
   # Centroid shift options
   add_intrapixel_shifts: True
   # In pixels (should be abs(limits)<0.5)
-  intrapixel_shift_range: [-0.5, 0.5]
+  intrapixel_shift_range: [-0.3, 0.3]
 
   # CCD misalignment options
   add_ccd_misalignments: True

data/generation/data_generation_params_v.3.1.0.yml

@@ -106,7 +106,7 @@ dataset_features:
   # Centroid shift options
   add_intrapixel_shifts: True
   # In pixels (should be abs(limits)<0.5)
-  intrapixel_shift_range: [-0.5, 0.5]
+  intrapixel_shift_range: [-0.3, 0.3]
 
   # CCD misalignment options
   add_ccd_misalignments: True

data/generation/data_generation_params_v.3.2.0.yml

@@ -106,7 +106,7 @@ dataset_features:
   # Centroid shift options
   add_intrapixel_shifts: True
   # In pixels (should be abs(limits)<0.5)
-  intrapixel_shift_range: [-0.5, 0.5]
+  intrapixel_shift_range: [-0.3, 0.3]
 
   # CCD misalignment options
   add_ccd_misalignments: True

data/generation/data_generation_params_v.3.3.0.yml

@@ -106,7 +106,7 @@ dataset_features:
   # Centroid shift options
   add_intrapixel_shifts: True
   # In pixels (should be abs(limits)<0.5)
-  intrapixel_shift_range: [-0.5, 0.5]
+  intrapixel_shift_range: [-0.3, 0.3]
 
   # CCD misalignment options
   add_ccd_misalignments: True

data/generation/data_generation_script.py

@@ -25,9 +25,9 @@
 )
 from wf_psf.utils.read_config import read_conf, RecursiveNamespace
 from wf_psf.sims.psf_simulator import PSFSimulator
-from wf_psf.utils.preprocessing import shift_x_y_to_zk1_2_wavediff
+from wf_psf.data.data_zernike_utils import shift_x_y_to_zk1_2_wavediff
 from wf_psf.sims.spatial_varying_psf import SpatialVaryingPSF, ZernikeHelper
-from wf_psf.utils.ccd_misalignments import CCDMisalignmentCalculator
+from wf_psf.instrument.ccd_misalignments import CCDMisalignmentCalculator
 
 
 # Pre-defined colormap
@@ -561,6 +561,7 @@ def main(args):
 
     # ------------ #
     # Centroid shifts
+    no_shift_test_zks = None
 
     if add_intrapixel_shifts:
 
@@ -599,6 +600,7 @@ def main(args):
 
         # Add the centroid shifts to the Zernike coefficients
         train_zks += train_delta_centroid_shifts
+        no_shift_test_zks = np.copy(test_zks)
         test_zks += test_delta_centroid_shifts
 
     # ------------ #
@@ -752,12 +754,24 @@ def main(args):
             sim_PSF_toolkit.generate_poly_PSF(test_SED_list[it], n_bins=n_bins)
         )
 
+    # Generate test polychromatic PSFs without shifts
+    if no_shift_test_zks is not None:
+        test_poly_psf_noshift_list = []
+        print("Generate test PSFs at observation resolution without shifts")
+        for it in tqdm(range(no_shift_test_zks.shape[0])):
+            sim_PSF_toolkit.set_z_coeffs(no_shift_test_zks[it, :])
+            test_poly_psf_noshift_list.append(
+                sim_PSF_toolkit.generate_poly_PSF(test_SED_list[it], n_bins=n_bins)
+            )
+
     # Generate numpy arrays from the lists
     train_poly_psf_np = np.array(train_poly_psf_list)
     train_SED_np = np.array(train_SED_list)
 
     test_poly_psf_np = np.array(test_poly_psf_list)
     test_SED_np = np.array(test_SED_list)
+    if no_shift_test_zks is not None:
+        test_poly_psf_noshift_np = np.array(test_poly_psf_noshift_list)
 
     # Generate the noisy train stars
     # Copy the training stars
@@ -776,11 +790,28 @@ def main(args):
         axis=0,
     )
 
+    # Also add noise to the test stars
+    noisy_test_poly_psf_np = np.copy(test_poly_psf_np)
+    # Generate a dataset with a SNR varying randomly within the desired range
+    rand_SNR = (
+        np.random.rand(noisy_test_poly_psf_np.shape[0]) * (SNR_range[1] - SNR_range[0])
+    ) + SNR_range[0]
+    # Add Gaussian noise to the observations
+    noisy_test_poly_psf_np = np.stack(
+        [
+            add_noise(_im, desired_SNR=_SNR)
+            for _im, _SNR in zip(noisy_test_poly_psf_np, rand_SNR)
+        ],
+        axis=0,
+    )
+
     # ------------ #
     # Generate masks
 
     if add_masks:
 
+        masked_noisy_test_poly_psf_np = np.copy(noisy_test_poly_psf_np)
+
         if mask_type == "random":
             # Generate random train masks
             train_masks = generate_n_mask(
@@ -797,6 +828,12 @@ def main(args):
                 noisy_train_poly_psf_np.dtype
             )
 
+            # Apply the random masks to the test stars
+            masked_noisy_test_poly_psf_np = (
+                masked_noisy_test_poly_psf_np
+                * test_masks.astype(noisy_test_poly_psf_np.dtype)
+            )
+
             # Turn masks to SHE convention. 1 (True) means to mask and 0 (False) means to keep
             train_masks = ~train_masks
             test_masks = ~test_masks
@@ -1093,8 +1130,21 @@ def main(args):
             SR_sim_PSF_toolkit.generate_poly_PSF(test_SED_list[it_j], n_bins=n_bins)
         )
 
+    # Generate the test super resolved (SR) polychromatic PSFs without shifts
+    if no_shift_test_zks is not None:
+        SR_test_poly_psf_noshift_list = []
+
+        print("Generate testing SR PSFs no shifts")
+        for it_j in tqdm(range(n_test_stars)):
+            SR_sim_PSF_toolkit.set_z_coeffs(no_shift_test_zks[it_j, :])
+            SR_test_poly_psf_noshift_list.append(
+                SR_sim_PSF_toolkit.generate_poly_PSF(test_SED_list[it_j], n_bins=n_bins)
+            )
+
     # Generate numpy arrays from the lists
     SR_test_poly_psf_np = np.array(SR_test_poly_psf_list)
+    if no_shift_test_zks is not None:
+        SR_test_poly_psf_noshift_np = np.array(SR_test_poly_psf_noshift_list)
 
     # ------------ #
     # Save test datasets
@@ -1131,17 +1181,23 @@ def main(args):
     test_psf_dataset = {
         "stars": test_poly_psf_np,
         "SR_stars": SR_test_poly_psf_np,
+        "noisy_stars": noisy_test_poly_psf_np,
         "positions": test_positions,
         "SEDs": test_SED_np,
         "zernike_GT": test_zks,
     }
 
     if add_masks:
         test_psf_dataset["masks"] = test_masks
+        test_psf_dataset["masked_noisy_stars"] = masked_noisy_test_poly_psf_np
 
     if add_ccd_misalignments:
         test_psf_dataset["zernike_ccd_misalignments"] = test_delta_Z3_arr
 
+    if no_shift_test_zks is not None:
+        test_psf_dataset["stars_noshift"] = test_poly_psf_noshift_np
+        test_psf_dataset["SR_stars_noshift"] = SR_test_poly_psf_noshift_np
+
     if add_intrapixel_shifts:
         test_psf_dataset["zernike_centroid_shifts"] = test_delta_centroid_shifts
         test_psf_dataset["pix_centroid_shifts"] = np.stack(

src/wf_psf/__init__.py

@@ -1,7 +1,7 @@
 import importlib
 
 # Dynamically import modules to trigger side effects when wf_psf is imported
-importlib.import_module('wf_psf.psf_models.psf_models')
-importlib.import_module('wf_psf.psf_models.psf_model_semiparametric')
-importlib.import_module('wf_psf.psf_models.psf_model_physical_polychromatic')
-importlib.import_module('wf_psf.psf_models.tf_psf_field') 
+importlib.import_module("wf_psf.psf_models.psf_models")
+importlib.import_module("wf_psf.psf_models.models.psf_model_semiparametric")
+importlib.import_module("wf_psf.psf_models.models.psf_model_physical_polychromatic")
+importlib.import_module("wf_psf.psf_models.tf_modules.tf_psf_field")