Merge pull request #1845 from pypeit/flexure_diag

Move flexure diagnostic code to class methods

pypeit/core/flexure.py

@@ -37,7 +37,7 @@
 from pypeit.datamodel import DataContainer
 from pypeit.images.detector_container import DetectorContainer
 from pypeit.images.mosaic import Mosaic
-from pypeit import specobj, specobjs, spec2dobj
+from pypeit import specobj, specobjs
 from pypeit import wavemodel
 
 from IPython import embed
@@ -1395,87 +1395,6 @@ def sky_em_residuals(wave:np.ndarray, flux:np.ndarray,
     return dwave[m], diff[m], diff_err[m], los[m], los_err[m]
 
 
-def flexure_diagnostic(file, file_type='spec2d', flexure_type='spec', chk_version=False):
-    """
-    Print the spectral or spatial flexure of a spec2d or spec1d file.
-
-    Args:
-        file (:obj:`str`, `Path`_):
-            Filename of the spec2d or spec1d file to check
-        file_type (:obj:`str`, optional):
-            Type of the file to check. Options are 'spec2d' or 'spec1d'. Default is 'spec2d'.
-        flexure_type (:obj:`str`, optional):
-            Type of flexure to check. Options are 'spec' or 'spat'. Default is 'spec'.
-        chk_version (:obj:`bool`, optional):
-            If True, raise an error if the datamodel version or type check failed.
-            If False, throw a warning only. Default is False.
-
-    Returns:
-        :obj:`astropy.table.Table`, :obj:`float`:  Type depends on the keyword
-        arguments:
-
-            - If file_type is 'spec2d' and flexure_type is 'spec', return a
-              table with the spectral flexure
-
-            - If file_type is 'spec2d' and flexure_type is 'spat', return the
-              spatial flexure
-
-            - If file_type is 'spec1d', return a table with the spectral flexure
-
-    """
-
-    # value to return
-    return_flex = None
-
-    if file_type == 'spec2d':
-        # load the spec2d file
-        allspec2D = spec2dobj.AllSpec2DObj.from_fits(file, chk_version=chk_version)
-        # Loop on Detectors
-        for det in allspec2D.detectors:
-            print('')
-            print('=' * 50 + f'{det:^7}' + '=' * 51)
-            # get and print the spectral flexure
-            if flexure_type == 'spec':
-                spec_flex = allspec2D[det].sci_spec_flexure
-                spec_flex.rename_column('sci_spec_flexure', 'global_spec_shift')
-                if np.all(spec_flex['global_spec_shift'] != None):
-                    spec_flex['global_spec_shift'].format = '0.3f'
-                # print the table
-                spec_flex.pprint_all()
-                # return the table
-                return_flex = spec_flex
-            # get and print the spatial flexure
-            if flexure_type == 'spat':
-                spat_flex = allspec2D[det].sci_spat_flexure
-                # print the value
-                print(f'Spat shift: {spat_flex}')
-                # return the value
-                return_flex = spat_flex
-    elif file_type == 'spec1d':
-        # no spat flexure in spec1d file
-        if flexure_type == 'spat':
-            msgs.error("Spat flexure not available in the spec1d file, try with a spec2d file")
-        # load the spec1d file
-        sobjs = specobjs.SpecObjs.from_fitsfile(file, chk_version=chk_version)
-        spec_flex = Table()
-        spec_flex['NAME'] = sobjs.NAME
-        spec_flex['global_spec_shift'] = sobjs.FLEX_SHIFT_GLOBAL
-        if np.all(spec_flex['global_spec_shift'] != None):
-            spec_flex['global_spec_shift'].format = '0.3f'
-        spec_flex['local_spec_shift'] = sobjs.FLEX_SHIFT_LOCAL
-        if np.all(spec_flex['local_spec_shift'] != None):
-            spec_flex['local_spec_shift'].format = '0.3f'
-        spec_flex['total_spec_shift'] = sobjs.FLEX_SHIFT_TOTAL
-        if np.all(spec_flex['total_spec_shift'] != None):
-            spec_flex['total_spec_shift'].format = '0.3f'
-        # print the table
-        spec_flex.pprint_all()
-        # return the table
-        return_flex = spec_flex
-
-    return return_flex
-
-
 # TODO -- Consider separating the methods from the DataContainer as per calibrations
 class MultiSlitFlexure(DataContainer):
     """

pypeit/scripts/chk_flexure.py

@@ -15,10 +15,13 @@ def get_parser(cls, width=None):
         parser = super().get_parser(description='Print QA on flexure to the screen',
                                     width=width)
 
-        parser.add_argument('input_file', type=str, nargs='+', help='One or more PypeIt spec2d or spec1d file')
+        parser.add_argument('input_file', type=str, nargs='+',
+                            help='One or more PypeIt spec2d or spec1d file')
         inp = parser.add_mutually_exclusive_group(required=True)
-        inp.add_argument('--spec', default=False, action='store_true', help='Check the spectral flexure')
-        inp.add_argument('--spat', default=False, action='store_true', help='Check the spatial flexure')
+        inp.add_argument('--spec', default=False, action='store_true',
+                         help='Check the spectral flexure')
+        inp.add_argument('--spat', default=False, action='store_true',
+                         help='Check the spatial flexure')
         parser.add_argument('--try_old', default=False, action='store_true',
                             help='Attempt to load old datamodel versions.  A crash may ensue..')
         return parser
@@ -29,9 +32,11 @@ def main(args):
         from IPython import embed
         from astropy.io import fits
         from pypeit import msgs
-        from pypeit.core import flexure
+        from pypeit import specobjs
+        from pypeit import spec2dobj
 
         chk_version = not args.try_old
+        flexure_type = 'spat' if args.spat else 'spec'
 
         # Loop over the input files
         for in_file in args.input_file:
@@ -41,23 +46,21 @@ def main(args):
             # What kind of file are we??
             hdul = fits.open(in_file)
             head0 = hdul[0].header
-            file_type = None
+
             if 'PYP_CLS' in head0.keys() and head0['PYP_CLS'].strip() == 'AllSpec2DObj':
-                file_type = 'spec2d'
+                # load the spec2d file
+                allspec2D = spec2dobj.AllSpec2DObj.from_fits(in_file, chk_version=chk_version)
+                allspec2D.flexure_diagnostics(flexure_type=flexure_type)
             elif 'DMODCLS' in head0.keys() and head0['DMODCLS'].strip() == 'SpecObjs':
-                file_type = 'spec1d'
+                if flexure_type == 'spat':
+                    msgs.error("Spat flexure not available in the spec1d file, try with a "
+                               "spec2d file")
+                # load the spec1d file
+                sobjs = specobjs.SpecObjs.from_fitsfile(in_file, chk_version=chk_version)
+                sobjs.flexure_diagnostics()
             else:
                 msgs.error("Bad file type input!")
 
-            # Check the flexure
-            flexure.flexure_diagnostic(in_file, file_type=file_type, flexure_type='spat' if args.spat else 'spec',
-                                       chk_version=chk_version)
-
             #  space between files for clarity
             print('')
 
-
-
-
-
-

pypeit/spec2dobj.py

@@ -728,3 +728,45 @@ def __repr__(self):
         txt += ') >'
         return txt
 
+    def flexure_diagnostics(self, flexure_type='spat'):
+        """
+        Print and return the spectral or spatial flexure of a spec2d file.
+
+        Args:
+            flexure_type (:obj:`str`, optional):
+                Type of flexure to check. Options are 'spec' or 'spat'. Default
+                is 'spec'.
+
+        Returns:
+            :obj:`dict`: Dictionary with the flexure values for each detector. If
+            flexure_type is 'spec', the spectral flexure is stored in an astropy table.
+            If flexure_type is 'spat', the spatial flexure is stored in a float.
+
+        """
+        if flexure_type not in ['spat', 'spec']:
+            msgs.error(f'flexure_type must be spat or spec, not {flexure_type}')
+        return_flex = {}
+        # Loop on Detectors
+        for det in self.detectors:
+            print('')
+            print('=' * 50 + f'{det:^7}' + '=' * 51)
+            # get and print the spectral flexure
+            if flexure_type == 'spec':
+                spec_flex = self[det].sci_spec_flexure
+                spec_flex.rename_column('sci_spec_flexure', 'global_spec_shift')
+                if np.all(spec_flex['global_spec_shift'] != None):
+                    spec_flex['global_spec_shift'].format = '0.3f'
+                # print the table
+                spec_flex.pprint_all()
+                # return the table
+                return_flex[det] = spec_flex
+            # get and print the spatial flexure
+            if flexure_type == 'spat':
+                spat_flex = self[det].sci_spat_flexure
+                # print the value
+                print(f'Spat shift: {spat_flex}')
+                # return the value
+                return_flex[det] = spat_flex
+
+        return return_flex
+

pypeit/specobjs.py

@@ -1027,6 +1027,29 @@ def get_extraction_groups(self, model_full_slit=False) -> List[List[int]]:
 
         return groups
 
+    def flexure_diagnostics(self):
+        """
+        Print and return the spectral flexure of a spec1d file.
+
+        Returns:
+            :obj:`astropy.table.Table`: Table with the spectral flexure.
+        """    
+        spec_flex = Table()
+        spec_flex['NAME'] = self.NAME
+        spec_flex['global_spec_shift'] = self.FLEX_SHIFT_GLOBAL
+        if np.all(spec_flex['global_spec_shift'] != None):
+            spec_flex['global_spec_shift'].format = '0.3f'
+        spec_flex['local_spec_shift'] = self.FLEX_SHIFT_LOCAL
+        if np.all(spec_flex['local_spec_shift'] != None):
+            spec_flex['local_spec_shift'].format = '0.3f'
+        spec_flex['total_spec_shift'] = self.FLEX_SHIFT_TOTAL
+        if np.all(spec_flex['total_spec_shift'] != None):
+            spec_flex['total_spec_shift'].format = '0.3f'
+        # print the table
+        spec_flex.pprint_all()
+        # return the table
+        return spec_flex
+
 
 #TODO Should this be a classmethod on specobjs??
 def get_std_trace(detname, std_outfile, chk_version=True):