allow attachWavelengthSolution() to progress without a distortion model (since there may be no lines in some GNIRS XD orders)

Author: chris-simpson

geminidr/core/primitives_spect.py

@@ -715,15 +715,15 @@ def attachWavelengthSolution(self, adinputs=None, **params):
                 # the warning at the end
                 try:
                     if 'distortion_corrected' not in wcs.available_frames:
-                        distortion_models = []
-                        break
+                        m_distcorr = None  # because this is used in later code
+                        distortion_models.append(None)
                 except AttributeError:
                     distortion_models = []
                     break
-
-                m_distcorr = wcs.get_transform(wcs.input_frame,
-                                               'distortion_corrected')
-                distortion_models.append(m_distcorr)
+                else:
+                    m_distcorr = wcs.get_transform(wcs.input_frame,
+                                                   'distortion_corrected')
+                    distortion_models.append(m_distcorr)
 
                 try:
                     wave_model = am.get_named_submodel(wcs.forward_transform,
@@ -737,9 +737,7 @@ def attachWavelengthSolution(self, adinputs=None, **params):
                                         if isinstance(frame, cf.SpectralFrame)])
 
             if not distortion_models:
-                log.warning("Could not find a 'distortion_corrected' frame "
-                            f"in arc {arc.filename} extension {ext.id} - "
-                            "continuing")
+                log.warning("Problem with one or more WCS objects - continuing")
                 if 'sq' in self.mode:
                     fail = True
                 continue
@@ -767,7 +765,8 @@ def attachWavelengthSolution(self, adinputs=None, **params):
                     geotable = import_module('.geometry_conf', self.inst_lookups)
                     transform.add_mosaic_wcs(ad, geotable)
                     for ext in ad:
-                        if 'mosaic' in ext.wcs.available_frames:
+                        if ('mosaic' in ext.wcs.available_frames and
+                                m_distcorr is not None):
                             ext.wcs.insert_frame('mosaic', m_distcorr,
                                                  cf.Frame2D(name='distortion_corrected'))
 
@@ -859,10 +858,12 @@ def attachWavelengthSolution(self, adinputs=None, **params):
                         # No mosaicking, so we can just do a shift
                         m_shift = (models.Shift((ad_detsec.x1 - arc_detsec.x1) / xbin) &
                                    models.Shift((ad_detsec.y1 - arc_detsec.y1) / ybin))
-                        m_distcorr = m_shift | m_distcorr
+                        m_distcorr = (m_shift if m_distcorr is None else
+                                      m_shift | m_distcorr)
 
-                    ad[0].wcs.insert_frame(ad[0].wcs.input_frame, m_distcorr,
-                                           cf.Frame2D(name='distortion_corrected'))
+                    if m_distcorr is not None:
+                        ad[0].wcs.insert_frame(ad[0].wcs.input_frame, m_distcorr,
+                                               cf.Frame2D(name='distortion_corrected'))
 
                 if wave_model is None:
                     log.warning(f"{arc.filename} has no wavelength solution")
@@ -887,11 +888,15 @@ def attachWavelengthSolution(self, adinputs=None, **params):
                     # applying the distortion correction
                     shifts = [c1 - c2 for c1, c2 in zip(ext.detector_section(),
                                                         ext_arc.detector_section())]
-                    dist_model = (models.Shift(shifts[0] / xbin) &
-                                  models.Shift(shifts[1] / ybin)) | dist_model
-
-                    ext.wcs.insert_frame(ad[0].wcs.input_frame, dist_model,
-                                         cf.Frame2D(name='distortion_corrected'))
+                    if shifts.count(0) < len(shifts):
+                        shift_model = (models.Shift(shifts[0] / xbin) &
+                                      models.Shift(shifts[1] / ybin))
+                        dist_model =  (shift_model if dist_model is None else
+                                       shift_model | dist_model)
+
+                    if dist_model is not None:
+                        ext.wcs.insert_frame(ad[0].wcs.input_frame, dist_model,
+                                             cf.Frame2D(name='distortion_corrected'))
 
                     if wave_model is None:
                         log.warning(f"{arc.filename} extension {ext.id} has "