Merge pull request #757 from kevin218/dev/ra/optimizer/jwst

Eureka! optimizer for Stages 1, 3, and 4.

Author: kevin218

demos/JWST/CalibratedStellarSpectra/run_eureka.py

@@ -5,6 +5,10 @@
 import eureka.S3_data_reduction.s3_reduce as s3
 import eureka.S4cal_StellarSpectra.s4cal_StellarSpec as s4cal
 
+# Set up some parameters to make plots look nicer.
+# You can set usetex=True if you have LaTeX installed
+eureka.lib.plots.set_rc(style='eureka', usetex=False, filetype='.png')
+
 # eventlabel = 'imaging_template'
 # eventlabel = 'miri_lrs_template'
 eventlabel = 'nirspec_fs_template'

demos/JWST/Optimizer/S1opt_template.ecf

@@ -0,0 +1,36 @@
+# Eureka! Control File for Optimizer Stage 1
+
+# Optimizer Documentation: https://eurekadocs.readthedocs.io/en/latest/ecf.html#stages-1-3-4-optimizer
+# Rename this file to match the event label given in run_eureka_optimization.py
+# and run_eureka.py (e.g., S1opt_nirspec_fs_template.ecf).
+
+# Fitness Scaling
+# Use these settings to override the default fitness scaling factors.
+scaling_MAD_spec    0.01            # Optimizer scaling for 2D MAD value
+scaling_MAD_white   1.0             # Optimizer scaling for WLC MAD value
+
+# List of parameters to optimize in Stage 1
+# Commenting out this line will use all single parameters listed below.
+# Single parameter options: jump_rejection_threshold, expand_mask, bg_deg, bg_method, p3thresh, window_len
+# Double parameter options: any combination of the above, joined by two underscores (e.g., expand_mask__p3thresh)
+params_to_optimize_s1  ['jump_rejection_threshold', 'expand_mask', 'p3thresh']
+
+# Range for parameter values to optimize
+# Only specify if overriding the default values in S1opt_meta.py
+# For single parameters, use format: sweep_<parameter_name> = range(min, max+1) or np.arange(min, max+step, step) if step != 1
+# For double parameters, use format: sweep_<parameter1>__<parameter2> = [range(min1, max1+1), range(min2, max2+1)]
+sweep_expand_mask  range(5,10)         # Example: optimize BG inner edge between 5 and 9 pixels
+
+# Diagnostics
+isplots_S1opt       1               # Generate none (0), few (1), some (3), or many (5) figures (Options: 1 - 5)
+delete_intermediate True           # If True, intermediate directories will be deleted after each optimization step
+delete_final        True           # If True, intermediate directories will be deleted upon completion of optimization run
+verbose             True           # If True, more details will be printed about steps in the optimization process
+hide_plots          True            # If True, plots will automatically be closed rather than popping up
+
+# Project directory
+topdir              /home/User/Data/JWST-Sim/optimizer/
+
+# Directories relative to topdir
+inputdir            Stage0
+outputdir           Stage1opt

demos/JWST/Optimizer/S3opt_template.ecf

@@ -0,0 +1,42 @@
+# Eureka! Control File for Optimizer Stage 3
+
+# Optimizer Documentation: https://eurekadocs.readthedocs.io/en/latest/ecf.html#stages-1-3-4-optimizer
+# Rename this file to match the event label given in run_eureka_optimization.py
+# and run_eureka.py (e.g., S3opt_nirspec_fs_template.ecf).
+
+# Fitness Scaling
+# Use these settings to override the default fitness scaling factors.
+scaling_MAD_spec    0.01            # Optimizer scaling for 2D MAD value
+scaling_MAD_white   1.0             # Optimizer scaling for WLC MAD value
+
+# List of single parameter options to optimize in Stage 3:
+#   dqmask, bg_deg, bg_thresh, bg_hw, bg_method, p3thresh, spec_hw, median_thresh, window_len, p7thresh
+# Double parameter options: any combination of the above, joined by two underscores (e.g., bg_hw__p3thresh)
+# Special cases: spec_hw__bg_hw (requires spec_hw < bg_hw)
+params_to_optimize_s3   ['spec_hw__bg_hw', 'dqmask', 'bg_deg', 'bg_thresh', 'bg_method', 'p3thresh', 'median_thresh__window_len', 'p7thresh']
+
+# List of single parameter options to optimize in Stage 4:
+#   mad_sigma, mad_box_width, sigma, box_width
+# Suggested double parameter options:
+#   mad_sigma__mad_box_width, sigma__box_width
+params_to_optimize_s4   ['mad_sigma__mad_box_width', 'sigma__box_width']
+
+# Range for parameter values to optimize
+# Only specify if overriding the default values in S3opt_meta.py
+# For single parameters, use format: sweep_<parameter_name> = range(min, max+1) or np.arange(min, max+step, step) if step != 1
+# For double parameters, use format: sweep_<parameter1>__<parameter2> = [range(min1, max1+1), range(min2, max2+1)]
+sweep_bg_thresh    np.arange(3,5.5,0.5)  # Example: optimize BG threshold between 3 and 5 sigma, in steps of 0.5
+
+# Diagnostics
+isplots_S3opt       1               # Generate none (0), few (1), some (3), or many (5) figures (Options: 1 - 5)
+delete_intermediate False           # If True, intermediate directories will be deleted after each optimization step
+delete_final        True            # If True, intermediate directories will be deleted upon completion of optimization run
+verbose             True            # If True, more details will be printed about steps in the optimization process
+hide_plots          True            # If True, plots will automatically be closed rather than popping up
+
+# Project directory
+topdir              /home/User/Data/JWST-Sim/optimizer/
+
+# Directories relative to topdir
+inputdir            Stage2
+outputdir           Stage3opt

demos/JWST/Optimizer/run_eureka_optimization.py

@@ -0,0 +1,46 @@
+import os
+import eureka.lib.plots
+import eureka.optimizer.S1opt_optimizer as s1opt
+import eureka.optimizer.S3opt_optimizer as s3opt
+
+# Set up some parameters to make plots look nicer.
+# You can set usetex=True if you have LaTeX installed
+eureka.lib.plots.set_rc(style='eureka', usetex=False, filetype='.png')
+
+"""
+Eureka! Optimization
+--------------------
+
+Description:
+This script is designed to optimize the Stage 1, 3, and 4 ECF parameters for
+JWST time-series observations.
+
+Inputs:
+- Loaded from a Eureka! control file, see S1opt_template.ecf/S3opt_template.ecf.
+- Parameters not specified for optimization will adopt the values listed in the
+    standard ECF (e.g., S1_<eventlabel>.ecf or S3_<eventlabel>.ecf).
+
+Outputs:
+- The metadata object
+- The fitness score after optimizing each parameter.
+- The best parameter values found during the optimization.
+"""
+
+eventlabel = 'nirspec_fs_template'
+ecf_path = '.'+os.sep
+
+if __name__ == "__main__":
+    # To skip one or more stages that were already run,
+    # just comment them out below
+
+    # Stage 1 optimization
+    s1opt_meta, history, best = s1opt.wrapper(eventlabel,
+                                              ecf_path=ecf_path,
+                                              initial_run=True,
+                                              final_run=True)
+
+    # Stages 3 and 4 optimization
+    s3opt_meta, history, best = s3opt.wrapper(eventlabel,
+                                              ecf_path=ecf_path,
+                                              initial_run=True,
+                                              final_run=True)

demos/JWST/run_eureka.py

@@ -7,6 +7,10 @@
 import eureka.S5_lightcurve_fitting.s5_fit as s5
 import eureka.S6_planet_spectra.s6_spectra as s6
 
+# Set up some parameters to make plots look nicer.
+# You can set usetex=True if you have LaTeX installed
+eureka.lib.plots.set_rc(style='eureka', usetex=False, filetype='.png')
+
 # eventlabel = 'imaging_template'
 # eventlabel = 'miri_lrs_template'
 eventlabel = 'nirspec_fs_template'

docs/source/ecf.rst

@@ -257,6 +257,89 @@ In short this weights each pixel, :math:`i`, within a slope following :math:`w_i
 ``custom``: As with default, except a custom SNR to :math:`P` lookup table can be defined through the ``default_ramp_fit_custom_snr_bounds`` and ``default_ramp_fit_custom_exponents`` (see example .ecf file).
 
 
+
+Stages 1/3/4 Optimizer
+----------------------
+This tool allows users to optimize the parameter values of Stages 1, 3, and 4 by performing a parametric sweep over a specified set of parameter values and evaluating the resulting spectra with a user-defined weighting of the fitness function.  The optimizer starts with an intial run of a stage using the default parameter values, performs a sweep over the specified parameter values, and then completes a final run with the best parameter values from the sweep.
+
+The fitness function is a weighted combination of the 2D MAD and the white light curve MAD.  The optimizer will generate a plot showing how the fitness function improves with each parameter sweep.  In the rare event that the final fitness value is worse than the initial value, users should consider trying smaller step sizes or a different set of parameter values that include the default parameter values.
+
+**Important Warning:** The optimizer is designed as a tool for parameter exploration and fine-tuning, not as a black-box solution. Users must critically investigate all outputs before and after optimization to ensure the final tunings are appropriate and scientifically valid. The fitness metric is inherently imperfect and may lead to undesirable outcomes in complex scenarios, such as observations affected by nearby binary stars or other astrophysical contaminants. Always review the spectra from different parameter values in the sweep to understand their impact, and remember that optimal parameter values can vary significantly between datasets. Thus, re-running the optimizer for each new dataset is recommended.
+
+scaling_MAD_spec
+''''''''''''''''
+Scaling factor applied to the pixel-level 2D Median Absolute Difference (MAD) value in the fitness function. Higher values prioritize spectral quality.
+
+scaling_MAD_white
+'''''''''''''''''
+Scaling factor applied to the white light curve MAD value in the fitness function. Higher values prioritize band-integrated quality.
+
+params_to_optimize_s1
+''''''''''''''''''''''
+List of parameters to optimize in Stage 1. Commenting out this line will use all single parameters.  Single parameter options: jump_rejection_threshold, expand_mask, bg_deg, bg_method, p3thresh, window_len. Double parameter options: any combination of the above, joined by two underscores (e.g., expand_mask__p3thresh). Example:
+
+.. code-block:: python
+
+   params_to_optimize = ['jump_rejection_threshold', 'expand_mask', 'p3thresh']
+
+params_to_optimize_s3
+'''''''''''''''''''''
+List of parameters to optimize in Stage 3. Commenting out this line will use all single parameters. Single parameter options: dqmask, bg_deg, bg_thresh, bg_hw, bg_method, p3thresh, spec_hw, median_thresh, window_len, p7thresh. Double parameter options: any combination of the above, joined by two underscores (e.g., bg_hw__p3thresh). Special cases: spec_hw__bg_hw (requires spec_hw < bg_hw). Recommended parameters:
+
+.. code-block:: python
+
+   params_to_optimize_s3 = ['spec_hw__bg_hw', 'dqmask', 'bg_deg', 'bg_thresh', 'bg_method', 'p3thresh', 'median_thresh__window_len', 'p7thresh']
+
+params_to_optimize_s4
+'''''''''''''''''''''
+List of parameters to optimize in Stage 4. Single parameter options: mad_sigma, mad_box_width, sigma, box_width. Double parameter options: mad_sigma__mad_box_width, sigma__box_width. Recommended parameters:
+
+.. code-block:: python
+
+    params_to_optimize_s4 = ['mad_sigma__mad_box_width', 'sigma__box_width']
+
+sweep_<parameter_name>
+''''''''''''''''''''''
+Set of parameter values to optimize. Only specify if overriding the default values in the meta Python files. For single parameters, use any of the following formats:
+
+.. code-block:: python
+
+   sweep_<parameter_name> = range(min, max+1)
+   sweep_<parameter_name> = np.arange(min, max+step, step)
+   sweep_<parameter_name> = [3, 5, 6, 7]
+   sweep_<parameter_name> = [True, False]
+
+sweep_<parameter1>__<parameter2>
+''''''''''''''''''''''''''''''''
+Set of parameter values to optimize. Only specify if overriding the default values in the meta Python files. For double parameters, use the following format:
+
+.. code-block:: python
+
+   sweep_<parameter1>__<parameter2> = [range(min1, max1+1), range(min2, max2+1)]
+
+isplots_S1opt and isplots_S3opt
+'''''''''''''''''''''''''''''''
+Sets how many plots should be saved during optimization. Can generate none (0), few (1), some (3), or many (5) figures (Options: 1 - 5).
+
+delete_final
+''''''''''''
+If True, intermediate directories will be deleted after the optimizer finishes. Outputs from the optimizer's final run remain in the usual output location. Set to ``False`` if you want to retain all trial outputs for
+inspection or debugging.
+
+delete_intermediate
+'''''''''''''''''''
+If True, intermediate directories will be deleted after each optimization step.  Outputs from the optimizer's final run remain in the usual output location. This may be useful for Stage 1 optimization since the outputs can be quite large. For Stage 3 optimization, this can usually be set to ``False``.
+
+verbose
+'''''''
+If True, more details will be printed about steps in the optimization process.
+
+hide_plots
+''''''''''
+If True, plots will automatically be closed rather than popping up on the screen.
+
+
+
 Stage 2
 -------
 
@@ -727,7 +810,7 @@ Optional. The path to a file that contains the time array you want to use instea
 
 
 Stage 4
---------
+-------
 
 .. include:: ../media/S4_template.ecf
    :literal:

environment.yml

@@ -88,6 +88,7 @@ dependencies = [
     "pandas",
     "photutils",
     "requests",
+    "pyyaml",
     "scipy>=1.12.0", # Lower limit needed for stcal.ramp_fitting.likely_fit
     "shapely",
     "tqdm",
@@ -100,7 +101,7 @@ dependencies = [
     { pip = "exotic-ld @ git+https://github.com/taylorbell57/ExoTiC-LD.git@0bb2ff8" }, # Pin needed for pkg_resources deprecation
     { pip = "fleck" },
     { pip = "george" }, # Needed for GP
-    { pip = "planet-harmonica @ git+https://github.com/DavoGrant/harmonica@main" }, # PyPI build currently only supports for python 3.8 -- 3.10
+    { pip = "planet-harmonica @ git+https://github.com/Exo-TiC/harmonica@main" }, # PyPI build currently only supports for python 3.8 -- 3.10
     { pip = "pastasoss>=1.2.3" }, # Lower limit needed for pkg_resources deprecation
     { pip = "setuptools_scm>=6.2" }, # Needed when importing from a source tree without generated version.py
     { pip = "stdatamodels" },
@@ -114,7 +115,7 @@ jwst = [
     { pip = "stcal>=1.15.2" }, # Lower limit needed for updated stcal.ramp_fitting.ramp_fit.ramp_fit arguments
 ]
 hst = [
-    { pip = "image_registration @ git+https://github.com/keflavich/image_registration@70e3880" }, # Need GitHub version to avoid np.float issue
+    { pip = "image_registration>=0.2.10" }, # Lower bound needed to avoid np.float issue
 ]
 docs = [
     { pip = "myst-parser>=0.18.0" }, # Lower bound for stable MyST support

pyproject.toml

@@ -65,7 +65,7 @@ def GLBS(input_model, log, meta):
         if meta.inst == 'miri':
             meta.isrotate = 0
 
-        if meta.isplots_S1 == 4:
+        if meta.isplots_S1 >= 4:
             # Plot all groups
             isplots_S1 = meta.isplots_S1
         # Otherwise, only show plots for the last good group

src/eureka/S1_detector_processing/group_level.py

@@ -36,6 +36,10 @@ def __init__(self, folder=None, file=None, eventlabel=None, **kwargs):
         # Set a default data file suffix
         self.suffix = getattr(self, 'suffix', 'uncal')
 
+        # Set optimization flag
+        self.firstSegOnly_S1 = getattr(self, 'firstSegOnly_S1', False)
+        self.firstSegOnly_S3 = getattr(self, 'firstSegOnly_S3', False)
+
     def set_defaults(self):
         '''Set Stage 1 specific defaults for generic instruments.
         '''
@@ -188,6 +192,7 @@ def set_defaults(self):
         self.nplots = getattr(self, 'nplots', 5)
         self.hide_plots = getattr(self, 'hide_plots', True)
         self.testing_S1 = getattr(self, 'testing_S1', False)
+        self.save_results = getattr(self, 'save_results', not self.testing_S1)
         self.verbose = getattr(self, 'verbose', True)
 
         # Project directory