JP-3669: Updating the C Extension to do CHARGELOSS Read Noise Recalcu…

…lations (#275)

.gitignore

@@ -4,6 +4,9 @@ __pycache__/
 *$py.class
 *~
 
+# Temp files
+*.*.swp
+
 # C extensions
 *.so
 

CHANGES.rst

@@ -23,6 +23,12 @@ General
 
 - Add TweakReg submodule. [#267]
 
+ramp_fitting
+~~~~~~~~~~~~
+
+- Move the CHARGELOSS read noise variance recalculation from the JWST step
+  code to the C extension to simplify the code and improve performance.[#275]
+
 Changes to API
 --------------
 

changes/275.general.rst

@@ -0,0 +1,2 @@
+[ramp_fitting] Moving the read noise recalculation due to CHARGELOSS flagging from
+the JWST ramp fit step code into the STCAL ramp fit C-extension.

src/stcal/ramp_fitting/ols_fit.py

@@ -920,6 +920,7 @@ def discard_miri_groups(ramp_data):
     data = ramp_data.data
     err = ramp_data.err
     groupdq = ramp_data.groupdq
+    orig_gdq = ramp_data.orig_gdq
 
     n_int, ngroups, nrows, ncols = data.shape
 
@@ -949,6 +950,8 @@ def discard_miri_groups(ramp_data):
     if num_bad_slices > 0:
         data = data[:, num_bad_slices:, :, :]
         err = err[:, num_bad_slices:, :, :]
+        if orig_gdq is not None:
+            orig_gdq = orig_gdq[:, num_bad_slices:, :, :]
 
     log.info("Number of leading groups that are flagged as DO_NOT_USE: %s", num_bad_slices)
 
@@ -968,6 +971,8 @@ def discard_miri_groups(ramp_data):
         data = data[:, :-1, :, :]
         err = err[:, :-1, :, :]
         groupdq = groupdq[:, :-1, :, :]
+        if orig_gdq is not None:
+            orig_gdq = orig_gdq[:, :-1, :, :]
 
         log.info("MIRI dataset has all pixels in the final group flagged as DO_NOT_USE.")
 
@@ -981,6 +986,8 @@ def discard_miri_groups(ramp_data):
     ramp_data.data = data
     ramp_data.err = err
     ramp_data.groupdq = groupdq
+    if orig_gdq is not None:
+        ramp_data.orig_gdq = orig_gdq
 
     return True
 

src/stcal/ramp_fitting/ramp_fit.py

@@ -30,7 +30,7 @@
 BUFSIZE = 1024 * 300000  # 300Mb cache size for data section
 
 
-def create_ramp_fit_class(model, dqflags=None, suppress_one_group=False):
+def create_ramp_fit_class(model, algorithm, dqflags=None, suppress_one_group=False):
     """
     Create an internal ramp fit class from a data model.
 
@@ -58,11 +58,24 @@ def create_ramp_fit_class(model, dqflags=None, suppress_one_group=False):
     else:
         dark_current_array = model.average_dark_current
 
+    orig_gdq = None
+    if algorithm.upper() == "OLS_C":
+        wh_chargeloss = np.where(np.bitwise_and(model.groupdq.astype(np.uint32), dqflags['CHARGELOSS']))
+        if len(wh_chargeloss[0]) > 0:
+            orig_gdq = model.groupdq.copy()
+        del wh_chargeloss
+
     if isinstance(model.data, u.Quantity):
         ramp_data.set_arrays(model.data.value, model.err.value, model.groupdq,
                              model.pixeldq, dark_current_array)
     else:
-        ramp_data.set_arrays(model.data, model.err, model.groupdq, model.pixeldq, dark_current_array)
+        ramp_data.set_arrays(
+            model.data,
+            model.err,
+            model.groupdq,
+            model.pixeldq,
+            dark_current_array,
+            orig_gdq)
 
     # Attribute may not be supported by all pipelines.  Default is NoneType.
     drop_frames1 = model.meta.exposure.drop_frames1 if hasattr(model, "drop_frames1") else None
@@ -78,6 +91,7 @@ def create_ramp_fit_class(model, dqflags=None, suppress_one_group=False):
     if "zero_frame" in model.meta.exposure and model.meta.exposure.zero_frame:
         ramp_data.zeroframe = model.zeroframe
 
+    ramp_data.algorithm = algorithm
     ramp_data.set_dqflags(dqflags)
     ramp_data.start_row = 0
     ramp_data.num_rows = ramp_data.data.shape[2]
@@ -170,7 +184,7 @@ def ramp_fit(
     # Create an instance of the internal ramp class, using only values needed
     # for ramp fitting from the to remove further ramp fitting dependence on
     # data models.
-    ramp_data = create_ramp_fit_class(model, dqflags, suppress_one_group)
+    ramp_data = create_ramp_fit_class(model, algorithm, dqflags, suppress_one_group)
 
     if algorithm.upper() == "OLS_C":
         ramp_data.run_c_code = True

src/stcal/ramp_fitting/ramp_fit_class.py

@@ -10,6 +10,10 @@ def __init__(self):
         self.pixeldq = None
         self.average_dark_current = None
 
+        # Needed for CHARGELOSS recomputation
+        self.orig_gdq = None
+        self.algorithm = None
+
         # Meta information
         self.instrument_name = None
 
@@ -25,6 +29,7 @@ def __init__(self):
         self.flags_saturated = None
         self.flags_no_gain_val = None
         self.flags_unreliable_slope = None
+        self.flags_chargeloss = None
 
         # ZEROFRAME
         self.zframe_mat = None
@@ -41,13 +46,15 @@ def __init__(self):
 
         # C code debugging switch.
         self.run_c_code = False
+        self.run_chargeloss = True
+        # self.run_chargeloss = False
 
         self.one_groups_locs = None  # One good group locations.
         self.one_groups_time = None  # Time to use for one good group ramps.
 
         self.current_integ = -1
 
-    def set_arrays(self, data, err, groupdq, pixeldq, average_dark_current):
+    def set_arrays(self, data, err, groupdq, pixeldq, average_dark_current, orig_gdq=None):
         """
         Set the arrays needed for ramp fitting.
 
@@ -72,6 +79,11 @@ def set_arrays(self, data, err, groupdq, pixeldq, average_dark_current):
         average_dark_current : ndarray (float32)
             2-D array containing the average dark current. It has
             dimensions (nrows, ncols)
+
+        orig_gdq : ndarray
+            4-D array containing a copy of the original group DQ array.  Since
+            the group DQ array can be modified during ramp fitting, this keeps
+            around the original group DQ flags passed to ramp fitting.
         """
         # Get arrays from the data model
         self.data = data
@@ -80,6 +92,8 @@ def set_arrays(self, data, err, groupdq, pixeldq, average_dark_current):
         self.pixeldq = pixeldq
         self.average_dark_current = average_dark_current
 
+        self.orig_gdq = orig_gdq
+
     def set_meta(self, name, frame_time, group_time, groupgap, nframes, drop_frames1=None):
         """
         Set the metainformation needed for ramp fitting.
@@ -131,6 +145,8 @@ def set_dqflags(self, dqflags):
         self.flags_saturated = dqflags["SATURATED"]
         self.flags_no_gain_val = dqflags["NO_GAIN_VALUE"]
         self.flags_unreliable_slope = dqflags["UNRELIABLE_SLOPE"]
+        if self.algorithm is not None and self.algorithm.upper() == "OLS_C":
+            self.flags_chargeloss = dqflags["CHARGELOSS"]
 
     def dbg_print_types(self):
         # Arrays from the data model
@@ -200,6 +216,16 @@ def dbg_print_pixel_info(self, row, col):
         # print(f"    err :\n{self.err[:, :, row, col]}")
         # print(f"    pixeldq :\n{self.pixeldq[row, col]}")
 
+    def dbg_print_info(self):
+        print(" ")
+        nints, ngroups, nrows, ncols = self.data.shape
+        for row in range(nrows):
+            for col in range(ncols):
+                print("=" * 80)
+                print(f"**** Pixel ({row}, {col}) ****")
+                self.dbg_print_pixel_info(row, col)
+        print("=" * 80)
+
     def dbg_write_ramp_data_pix_pre(self, fname, row, col, fd):
         fd.write("def create_ramp_data_pixel():\n")
         indent = INDENT