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