WIP JP-3610: Memory usage with Detector1 pipeline

changes/8588.general.rst

@@ -0,0 +1 @@
+Created two functions to workaround the memory increase due to datamodel.open. The new functions in the basic_utils script are use_datamodels and copy_datamodel. The first function determines if it is necessary to open a datamodel and returns it or simply returns the datamodel given as input. The second function makes a copy of the datamodel depending on whether the step is being run as part of a pipeline or not. Both functions are now called in all steps of the Detector1 pipeline.

jwst/charge_migration/charge_migration.py

@@ -1,5 +1,6 @@
 #  Module for charge migration
 #
+import gc
 import logging
 import numpy as np
 
@@ -45,6 +46,7 @@ def charge_migration(output_model, signal_threshold):
     # Save the flags in the output GROUPDQ array
     output_model.groupdq = gdq_new
 
+    gc.collect()
     return output_model
 
 

jwst/charge_migration/charge_migration_step.py

@@ -1,10 +1,12 @@
 #! /usr/bin/env python
+import gc
 import logging
 
 from stdatamodels.jwst import datamodels
 from ..stpipe import Step
 
 from . import charge_migration
+from jwst.lib.basic_utils import use_datamodel, copy_datamodel
 
 log = logging.getLogger(__name__)
 log.setLevel(logging.DEBUG)
@@ -24,24 +26,24 @@
         skip = boolean(default=True)
     """
 
-    def process(self, step_input):
+    def process(self, input_model):
 
         # Open the input data model
-        with datamodels.RampModel(step_input) as input_model:
+        with use_datamodel(input_model, model_class=datamodels.RampModel) as input_model:
 
-            if (input_model.data.shape[1] < 3):  # skip step if only 1 or 2 groups/integration
-                log.info('Too few groups per integration; skipping charge_migration')
+            result, input_model = copy_datamodel(input_model, self.parent)
 
-                input_model.meta.cal_step.charge_migration = 'SKIPPED'
-                return input_model
+            if (result.data.shape[1] < 3):  # skip step if only 1 or 2 groups/integration
+                log.info('Too few groups per integration; skipping charge_migration')
 
-            # Work on a copy
-            result = input_model.copy()
+                result.meta.cal_step.charge_migration = 'SKIPPED'
+                return result
 
             # Retrieve the parameter value(s)
             signal_threshold = self.signal_threshold
 
             result = charge_migration.charge_migration(result, signal_threshold)
             result.meta.cal_step.charge_migration = 'COMPLETE'
 
+        gc.collect()
         return result

jwst/dark_current/dark_current_step.py

@@ -1,7 +1,9 @@
+import gc
 from stdatamodels.jwst import datamodels
 
 from ..stpipe import Step
 from stcal.dark_current import dark_sub
+from jwst.lib.basic_utils import use_datamodel, copy_datamodel
 import numpy as np
 
 
@@ -23,24 +25,24 @@
 
     reference_file_types = ['dark']
 
-    def process(self, step_input):
+    def process(self, input_model):
 
         # Open the input data model
-        with datamodels.RampModel(step_input) as input_model:
+        with use_datamodel(input_model, model_class=datamodels.RampModel) as input_model:
+
+            result, input_model = copy_datamodel(input_model, self.parent)
 
             # Get the name of the dark reference file to use
-            self.dark_name = self.get_reference_file(input_model, 'dark')
+            self.dark_name = self.get_reference_file(result, 'dark')
             self.log.info('Using DARK reference file %s', self.dark_name)
 
             # Check for a valid reference file
             if self.dark_name == 'N/A':
                 self.log.warning('No DARK reference file found')
                 self.log.warning('Dark current step will be skipped')
-                input_model.meta.cal_step.dark = 'SKIPPED'
-                return input_model
-
-            # Work on a copy
-            result = input_model.copy()
+                result.meta.cal_step.dark = 'SKIPPED'
+                gc.collect()
+                return result
 
             # Create name for the intermediate dark, if desired.
             dark_output = self.dark_output
@@ -77,7 +79,8 @@
             # Cleanup
             del dark_model
             del result
-
+
+        gc.collect()
         return out_ramp
 
     def set_average_dark_current(self, input_model, dark_model):

jwst/dq_init/dq_init_step.py

@@ -1,8 +1,10 @@
 #! /usr/bin/env python
+import gc
 from stdatamodels.jwst import datamodels
 
 from ..stpipe import Step
 from . import dq_initialization
+from jwst.lib.basic_utils import use_datamodel, copy_datamodel
 
 
 __all__ = ["DQInitStep"]
@@ -41,7 +43,7 @@
 
         # Try to open the input as a regular RampModel
         try:
-            input_model = datamodels.RampModel(step_input)
+            input_model = use_datamodel(step_input, model_class=datamodels.RampModel)
             # Check to see if it's Guider raw data
             if input_model.meta.exposure.type in dq_initialization.guider_list:
                 # Reopen as a GuiderRawModel
@@ -61,19 +63,18 @@
             self.log.error("Can't open input")
             raise
 
+        result, input_model = copy_datamodel(input_model, self.parent)
+
         # Retrieve the mask reference file name
-        self.mask_filename = self.get_reference_file(input_model, 'mask')
+        self.mask_filename = self.get_reference_file(result, 'mask')
         self.log.info('Using MASK reference file %s', self.mask_filename)
 
         # Check for a valid reference file
         if self.mask_filename == 'N/A':
             self.log.warning('No MASK reference file found')
             self.log.warning('DQ initialization step will be skipped')
-            input_model.meta.cal_step.dq_init = 'SKIPPED'
-            return input_model
-
-        # Work on a copy
-        result = input_model.copy()
+            result.meta.cal_step.dq_init = 'SKIPPED'
+            return result
 
         # Load the reference file
         mask_model = datamodels.MaskModel(self.mask_filename)
@@ -82,7 +83,8 @@
         result = dq_initialization.correct_model(result, mask_model)
 
         # Cleanup
+        del step_input
         del mask_model
-        del input_model
+        gc.collect()
 
         return result

jwst/dq_init/dq_initialization.py

@@ -1,4 +1,5 @@
 import logging
+import gc
 
 import numpy as np
 
@@ -81,6 +82,7 @@ def do_dqinit(output_model, mask_model):
 
     output_model.meta.cal_step.dq_init = 'COMPLETE'
 
+    gc.collect()
     return output_model
 
 

jwst/emicorr/emicorr.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import logging
+import gc
 from astropy.stats import sigma_clipped_stats as scs
 from stdatamodels.jwst import datamodels
 
@@ -176,6 +177,7 @@ def apply_emicorr(output_model, emicorr_model,
     readpatt = output_model.meta.exposure.readpatt
     xsize = output_model.meta.subarray.xsize   # SUBSIZE1 keyword
     xstart = output_model.meta.subarray.xstart   # SUBSTRT1 keyword
+
     # get the number of samples, 10us sample times per pixel (1 for fastmode, 9 for slowmode)
     nsamples = output_model.meta.exposure.nsamples
 
@@ -438,25 +440,24 @@ def apply_emicorr(output_model, emicorr_model,
         # This is the phase matched noise model to subtract from each pixel of the input image
         dd_noise = lut[(phaseall * period_in_pixels).astype(int)]
 
-        # Interleave (straight copy) into 4 amps
-        noise = np.zeros((nints, ngroups, ny, nx))   # same size as input data
-        noise_x = np.arange(nx4) * 4
-        for k in range(4):
-            noise[:, :, :, noise_x + k] = dd_noise
-
         # Safety catch; anywhere the noise value is not finite, set it to zero
-        noise[~np.isfinite(noise)] = 0.0
+        dd_noise[~np.isfinite(dd_noise)] = 0.0
 
         # Subtract EMI noise from the input data
         log.info('Subtracting EMI noise from data')
-        output_model.data = output_model.data - noise
+
+        # Interleave (straight copy) into 4 amps
+        noise_x = np.arange(nx4) * 4
+        for k in range(4):
+            output_model.data[..., noise_x + k] = output_model.data[..., noise_x + k] - dd_noise
 
         # clean up
         del data
         del dd_all
         del times_this_int
         del phaseall
-        del noise
+        del dd_noise
+        gc.collect()
 
     if save_intermediate_results and save_onthefly_reffile is not None:
         if 'FAST' in readpatt:
@@ -471,7 +472,7 @@ def apply_emicorr(output_model, emicorr_model,
         }
         freq_pa_dict['subarray_cases'] = on_the_fly_subarr_case
         mk_reffile(freq_pa_dict, save_onthefly_reffile)
-
+        
     return output_model
 
 

jwst/emicorr/emicorr_step.py

@@ -1,8 +1,10 @@
 #! /usr/bin/env python
 
+import gc
 from stdatamodels.jwst import datamodels
 from ..stpipe import Step
 from . import emicorr
+from jwst.lib.basic_utils import use_datamodel, copy_datamodel
 
 
 __all__ = ["EmiCorrStep"]
@@ -28,27 +30,26 @@
 
     reference_file_types = ['emicorr']
 
-    def process(self, step_input):
+    def process(self, input_model):
 
         # Open the input data model
-        with datamodels.open(step_input) as input_model:
+        with use_datamodel(input_model, model_class=datamodels.RampModel) as input_model:
+
+            result, input_model = copy_datamodel(input_model, self.parent)
 
             # Catch the cases to skip
-            instrument = input_model.meta.instrument.name
+            instrument = result.meta.instrument.name
             if instrument != 'MIRI':
                 self.log.warning('EMI correction not implemented for instrument: {}'.format(instrument))
-                input_model.meta.cal_step.emicorr = 'SKIPPED'
-                return input_model
+                result.meta.cal_step.emicorr = 'SKIPPED'
+                return result
 
-            readpatt = input_model.meta.exposure.readpatt
+            readpatt = result.meta.exposure.readpatt
             allowed_readpatts = ['FAST', 'FASTR1', 'SLOW', 'SLOWR1']
             if readpatt.upper() not in allowed_readpatts:
                 self.log.warning('EMI correction not implemented for read pattern: {}'.format(readpatt))
-                input_model.meta.cal_step.emicorr = 'SKIPPED'
-                return input_model
-
-            # Work on a copy
-            result = input_model.copy()
+                result.meta.cal_step.emicorr = 'SKIPPED'
+                return result
 
             # Setup parameters
             pars = {
@@ -103,4 +104,5 @@
             # Cleanup
             del emicorr_model
 
+        gc.collect()
         return result

jwst/firstframe/firstframe_step.py

@@ -1,6 +1,7 @@
+import gc
 from ..stpipe import Step
 from . import firstframe_sub
-from stdatamodels.jwst import datamodels
+from jwst.lib.basic_utils import use_datamodel, copy_datamodel
 
 
 __all__ = ["FirstFrameStep"]
@@ -18,23 +19,22 @@ class FirstFrameStep(Step):
     spec = """
     """
 
-    def process(self, step_input):
+    def process(self, input_model):
 
         # Open the input data model
-        with datamodels.open(step_input) as input_model:
+        with use_datamodel(input_model) as input_model:
+
+            result, input_model = copy_datamodel(input_model, self.parent)
 
             # check the data is MIRI data
-            detector = input_model.meta.instrument.detector.upper()
-            if detector[:3] != 'MIR':
+            detector = result.meta.instrument.detector.upper()
+            if detector[:3] == 'MIR':
+                # Do the firstframe correction subtraction
+                result = firstframe_sub.do_correction(result)
+            else:
                 self.log.warning('First Frame Correction is only for MIRI data')
                 self.log.warning('First frame step will be skipped')
-                input_model.meta.cal_step.firstframe = 'SKIPPED'
-                return input_model
-
-            # Cork on a copy
-            result = input_model.copy()
-
-            # Do the firstframe correction subtraction
-            result = firstframe_sub.do_correction(result)
+                result.meta.cal_step.firstframe = 'SKIPPED'
 
+        gc.collect()
         return result

jwst/firstframe/firstframe_sub.py

@@ -1,6 +1,7 @@
 #
 #  Module for the firstframe correction for MIRI science data sets
 #
+import gc
 import numpy as np
 import logging
 
@@ -45,4 +46,5 @@ def do_correction(output):
         log.warning("Step will be skipped")
         output.meta.cal_step.firstframe = 'SKIPPED'
 
+    gc.collect()
     return output

jwst/gain_scale/gain_scale.py

@@ -1,3 +1,4 @@
+import gc
 import numpy as np
 import logging
 
@@ -44,4 +45,5 @@ def do_correction(output_model, gain_factor):
     output_model.meta.exposure.gain_factor = gain_factor
     output_model.meta.cal_step.gain_scale = 'COMPLETE'
 
+    gc.collect()
     return output_model