JP-3777: Select first and last groups in ramp fitting

changes/9095.ramp_fitting.rst

@@ -0,0 +1,4 @@
+Add code to enable users to select firstgroup and lastgroup parameters when performing
+ramp fitting. It works by setting the DO_NOT_USE bit in the GROUPDQ extension for groups
+outside the selected range.  Added a unit test and updated the docs.
+

docs/jwst/ramp_fitting/arguments.rst

@@ -16,6 +16,11 @@ The ramp fitting step has the following optional arguments that can be set by th
 * ``--int_name``: A string that can be used to override the default name
   for the per-integration product.
 
+* ``--firstgroup``, ``--lastgroup``: A pair of integers that can be used to set the first and last groups
+  to be used in the ramp fitting procedure.  These values are zero-indexed.  If the first group is set to <0,
+  or None, or the last group is set to > (#groups - 1) or None, or the first group is set to > last group,
+  the settings are ignored and the full ramp is fit.  Default values for both parameters are None.
+
 * ``--suppress_one_group``: A boolean to suppress computations for saturated ramps
   with only one good (unsaturated) sample.  The default is set to True to suppress these computations,
   which will compute all values for the ramp the same as if the entire ramp were

docs/jwst/ramp_fitting/description.rst

@@ -40,6 +40,10 @@ segment of length greater than one.
 Once all segments have been determined, slopes and variances are determined for
 each one.
 
+The range of groups to be fitted can be limited using the ``--firstgroup`` and
+``--lastgroup`` parameters.  This works by setting the DO_NOT_USE DQ bit in the GROUPDQ
+attribute for all groups outside the range selected.
+
 Pixels are processed simultaneously in blocks using the array-based functionality of numpy.
 The size of the block depends on the image size and the number of groups per
 integration.

jwst/ramp_fitting/ramp_fit_step.py

@@ -12,6 +12,7 @@
 from stdatamodels.jwst import datamodels
 from stdatamodels.jwst.datamodels import dqflags
 
+
 from ..stpipe import Step
 
 from ..lib import reffile_utils
@@ -23,10 +24,6 @@
 log.setLevel(logging.DEBUG)
 
 
-
-__all__ = ["RampFitStep"]
-
-
 def get_reference_file_subarrays(model, readnoise_model, gain_model, nframes):
     """
     Get readnoise array for calculation of variance of noiseless ramps, and
@@ -391,6 +388,8 @@
         save_opt = boolean(default=False) # Save optional output
         opt_name = string(default='')
         suppress_one_group = boolean(default=True)  # Suppress saturated ramps with good 0th group
+        firstgroup = integer(default=None)   # Ignore groups before this one (zero indexed)
+        lastgroup = integer(default=None)   # Ignore groups after this one (zero indexed)
         maximum_cores = string(default='1')
         # cores for multiprocessing. Can be an integer, 'half', 'quarter', or 'all'
     """
@@ -414,7 +413,7 @@
         # Open the input data model
         with datamodels.RampModel(step_input) as input_model:
 
-            # Cork on a copy
+            # Work on a copy
             result = input_model.copy()
 
             max_cores = self.maximum_cores
@@ -456,6 +455,15 @@
 
             int_times = result.int_times
 
+            # Set the DO_NOT_USE bit in the groupdq values for groups before firstgroup
+            # and groups after lastgroup
+            firstgroup = self.firstgroup
+            lastgroup = self.lastgroup
+            groupdqflags = dqflags.group
+
+            if firstgroup is not None or lastgroup is not None:
+                set_groupdq(firstgroup, lastgroup, ngroups, result.groupdq, groupdqflags)
+
             # Before the ramp_fit() call, copy the input model ("_W" for weighting)
             # for later reconstruction of the fitting array tuples.
             input_model_W = result.copy()
@@ -550,3 +558,45 @@
         del result
 
         return out_model, int_model
+
+def set_groupdq(firstgroup, lastgroup, ngroups, groupdq, groupdqflags):
+    """
+    Set the groupdq flags based on the values of firstgroup, lastgroup
+
+    Parameters:
+    -----------
+
+    firstgroup : int or None
+        The first group to be used in the ramp fitting
+
+    lastgroup : int or None
+        The last group to be used in the ramp fitting
+
+    ngroups : int
+        The number of groups in the ramp
+
+    groupdq : ndarray
+        The groupdq array to be modified in place
+
+    groupdqflags : dict
+        The groupdq flags dict
+
+    """
+    if firstgroup is None:
+        firstgroup = 0
+    if lastgroup is None:
+        lastgroup = ngroups - 1
+    if firstgroup < 0:
+        log.warning("first group < 0, reset to 0")
+        firstgroup = 0
+    if lastgroup >= ngroups:
+        log.warning(f"Last group number >= #groups ({ngroups}), reset to {ngroups-1}")
+    if firstgroup >= lastgroup:
+        log.warning(f"firstgroup ({firstgroup}) cannot be >= lastgroup ({lastgroup})")
+        log.warning("Group selectors ignored")
+        firstgroup = 0
+        lastgroup = ngroups - 1
+    if firstgroup > 0:
+        groupdq[:,:firstgroup] = np.bitwise_or(groupdq[:,:firstgroup], groupdqflags['DO_NOT_USE'])
+    if lastgroup < (ngroups - 1):
+        groupdq[:,(lastgroup+1):] = np.bitwise_or(groupdq[:,(lastgroup+1):], groupdqflags['DO_NOT_USE'])

jwst/ramp_fitting/tests/test_ramp_fit_step.py

@@ -1,9 +1,21 @@
 import pytest
+import logging
 import numpy as np
 
 from stdatamodels.jwst.datamodels import dqflags, RampModel, GainModel, ReadnoiseModel
 
-from jwst.ramp_fitting.ramp_fit_step import RampFitStep
+from jwst.ramp_fitting.ramp_fit_step import RampFitStep, set_groupdq
+from jwst.tests.helpers import LogWatcher
+
+@pytest.fixture
+def log_watcher(monkeypatch):
+    # Set a log watcher to check for a log message at any level
+    # in RampFitStep
+    watcher = LogWatcher('')
+    logger = logging.getLogger('stpipe')
+    for level in ['debug', 'info', 'warning', 'error']:
+        monkeypatch.setattr(logger, level, watcher)
+    return watcher
 
 DELIM = "-" * 80
 
@@ -258,6 +270,35 @@ def test_int_times2(generate_miri_reffiles, setup_inputs):
 
     assert len(cube_model.int_times) == nints
 
+def test_set_groups(generate_miri_reffiles, setup_inputs):
+    # Test results when using the firstgroup and lastgroup options
+    ngroups = 20
+    rampmodel, gdq, rnmodel, pixdq, err, gain = setup_inputs(ngroups=ngroups)
+    # Set up data array as group# squared
+    # This has the property that the slope=(firstgroup+lastgroup)
+    squares = np.array([k*k for k in range(ngroups)],dtype=np.float32)
+    rampmodel.data[0,:] = squares[:, np.newaxis, np.newaxis]
+    firstgroup = 3
+    lastgroup = 11
+    slopes, cubemodel = RampFitStep.call(rampmodel, override_gain=gain, override_readnoise=rnmodel,
+                                         firstgroup=firstgroup, lastgroup=lastgroup)
+    np.testing.assert_allclose(slopes.data, firstgroup+lastgroup, rtol=1e-7)
+
+def test_warnings(log_watcher):
+    # Test user warnings
+    ngroups = 20
+    groupdqflags = dqflags.group
+    groupdq = np.zeros((1, ngroups, 1024, 1024), dtype=np.uint16)
+    firstgroup = -10
+    lastgroup = None
+    log_watcher.message = "first group < 0, reset to 0"
+    set_groupdq(firstgroup, lastgroup, ngroups, groupdq, groupdqflags)
+    log_watcher.assert_seen()
+    firstgroup = 3
+    lastgroup = ngroups
+    log_watcher.message = f"Last group number >= #groups ({ngroups}), reset to {ngroups-1}"
+    set_groupdq(firstgroup, lastgroup, ngroups, groupdq, groupdqflags)
+    log_watcher.assert_seen()
 
 def one_group_suppressed(nints, suppress, setup_inputs):
     """