allow IFUImageModel

Author: tapastro

jwst/pixel_replace/pixel_replace.py

@@ -55,9 +55,9 @@ def __init__(self, input_model, **pars):
         try:
             self.algorithm = self.algorithm_dict[self.pars['algorithm']]
 
-        except:
-            log.critical(f"Algorithm name provided does not match an"
-                         f"implemented algorithm!")
+        except KeyError:
+            log.critical(f"Algorithm name {self.pars['algorithm']} provided does "
+                         f"not match an implemented algorithm!")
             raise Exception
 
     def replace(self):
@@ -67,10 +67,10 @@ def replace(self):
         to each 2D spectrum in input.
         """
         # ImageModel inputs (MIR_LRS-FIXEDSLIT)
-        if isinstance(self.input, datamodels.ImageModel):
+        if isinstance(self.input, (datamodels.ImageModel, datamodels.IFUImageModel)):
             self.output = self.algorithm(self.input)
             n_replaced = np.sum(self.output.dq & self.REPLACED)
-            log.info(f"Input ImageModel had {n_replaced} pixels replaced.")
+            log.info(f"Input model had {n_replaced} pixels replaced.")
 
         # MultiSlitModel inputs (WFSS, NRS_FIXEDSLIT, ?)
         elif isinstance(self.input, datamodels.MultiSlitModel):
@@ -86,7 +86,7 @@ def replace(self):
 
         else:
             # This should never happen, as these would be caught in the step code.
-            log.critical(f"How did we get here?")
+            log.critical("How did we get here?")
             raise Exception
 
     def fit_profile(self, model):
@@ -124,8 +124,8 @@ def fit_profile(self, model):
         # Truncate array to region where good pixels exist
         good_pixels = np.where(~model.dq & 1)
         if np.any(0 in np.shape(good_pixels)):
-            log.warning(f"No good pixels in at least one dimension of "
-                        f"data array - skipping pixel replacement.")
+            log.warning("No good pixels in at least one dimension of "
+                        "data array - skipping pixel replacement.")
             return model
         x_range = [np.min(good_pixels[0]), np.max(good_pixels[0]) + 1]
         y_range = [np.min(good_pixels[1]), np.max(good_pixels[1]) + 1]

jwst/pixel_replace/pixel_replace_step.py

@@ -59,6 +59,8 @@ def process(self, input):
             # MIRI_LRS-FIXEDSLIT comes in ImageModel - any others?
             elif isinstance(input_model, datamodels.ImageModel):
                 self.log.debug(f'Input ImageModel is of exptype: {input_model.meta.exposure.type}')
+            elif isinstance(input_model, datamodels.IFUImageModel):
+                self.log.debug(f'Input is an IFUImageModel.')
                 '''
                 elif isinstance(input_model, datamodels.CubeModel):
                     # It's a 3-D multi-integration model