Merge pull request #426 from GreenBankObservatory/418-readwrite-of-fl…

…ags-in-sdfits-scan-and-spectrum

418 read/write of flags in GBTFITSLoad

.pre-commit-config.yaml

@@ -1,7 +1,7 @@
 ---
 # See https://pre-commit.com for more information
 default_language_version:
-  python: python3.10
+  python: python3
 
 # See https://pre-commit.com/hooks.html for more hooks
 repos:

docs/source/how-tos/index.rst

@@ -77,7 +77,7 @@ Practical step-by-step guides to help you achieve a specific goal. Most useful w
 
         :material-outlined:`compare_arrows;3em;green` **Align Spectra**
 
-        How to read and save data
+        How to align spectra
 
         .. button-link:: examples/align_spectra.html
             :color: primary
@@ -86,6 +86,21 @@ Practical step-by-step guides to help you achieve a specific goal. Most useful w
 
             Align Spectra
 
+    .. grid-item-card::
+        :shadow: md
+        :margin: 2 2 0 0
+
+        :material-outlined:`flag;3em;green` **Flagging**
+
+        How to flag data
+
+        .. button-link:: examples/flagging.html
+            :color: primary
+            :outline:
+            :click-parent:
+
+            Flagging
+
 .. toctree::
    :maxdepth: 4
    :hidden:
@@ -95,3 +110,4 @@ Practical step-by-step guides to help you achieve a specific goal. Most useful w
    examples/smoothing
    examples/dataIO
    examples/align_spectra
+   examples/flagging

src/dysh/fits/gbtfitsload.py

@@ -52,12 +52,14 @@ class GBTFITSLoad(SDFITSLoad, HistoricalBase):
     hdu : int or list
         Header Data Unit to select from input file. Default: all HDUs
 
+    skipflags: bool
+        If True, do not read any flag files associated with these data. Default:False
+
     """
 
     @log_call_to_history
-    def __init__(self, fileobj, source=None, hdu=None, **kwargs):
+    def __init__(self, fileobj, source=None, hdu=None, skipflags=False, **kwargs):
         kwargs_opts = {
-            "fix": False,  # fix non-standard header elements
             "index": True,  # only set to False for performance testing.
             "verbose": False,
         }
@@ -70,7 +72,7 @@ def __init__(self, fileobj, source=None, hdu=None, **kwargs):
         self.GBT = Observatory["GBT"]
         if path.is_file():
             logger.debug(f"Treating given path {path} as a file")
-            self._sdf.append(SDFITSLoad(fileobj, source, hdu, **kwargs_opts))
+            self._sdf.append(SDFITSLoad(path, source, hdu, **kwargs_opts))
         elif path.is_dir():
             logger.debug(f"Treating given path {path} as a directory")
             # Find all the FITS files in the directory and sort alphabetically
@@ -92,13 +94,14 @@ def __init__(self, fileobj, source=None, hdu=None, **kwargs):
                 self.add_comment(h.header.get("COMMENT", []))
         self._remove_duplicates()
         if kwargs_opts["index"]:
-            self._create_index_if_needed()
+            self._create_index_if_needed(skipflags)
             self._update_radesys()
         # We cannot use this to get mmHg as it will disable all default astropy units!
         # https://docs.astropy.org/en/stable/api/astropy.units.cds.enable.html#astropy.units.cds.enable
         # cds.enable()  # to get mmHg
 
-        if kwargs.get("verbose", None):
+        # ushow/udata depend on the index being present, so check that index is created.
+        if kwargs.get("verbose", None) and kwargs_opts["index"]:
             print("==GBTLoad %s" % fileobj)
             self.ushow("OBJECT", 0)
             self.ushow("SCAN", 0)
@@ -116,7 +119,10 @@ def __init__(self, fileobj, source=None, hdu=None, **kwargs):
         lsdf = len(self._sdf)
         if lsdf > 1:
             print(f"Loaded {lsdf} FITS files")
-        self.add_history(f"Project ID: {self.projectID}", add_time=True)
+        if kwargs_opts["index"]:
+            self.add_history(f"Project ID: {self.projectID}", add_time=True)
+        else:
+            print("Reminder: No index created; many functions won't work.")
 
     def __repr__(self):
         return str(self.files)
@@ -139,7 +145,6 @@ def projectID(self):
         -------
         str
             The project ID string
-
         """
         return uniq(self["PROJID"])[0]
 
@@ -779,7 +784,9 @@ def apply_flags(self):
             for i, ((fi, bi), g) in enumerate(dfs):
                 chan_mask = convert_array_to_mask(chan, self._sdf[fi].nchan(bi))
                 rows = g["ROW"].to_numpy()
-                self._sdf[fi]._flagmask[bi][rows] = chan_mask
+                logger.debug(f"Applying {chan} to {rows=}")
+                logger.debug(f"{np.where(chan_mask)}")
+                self._sdf[fi]._flagmask[bi][rows] |= chan_mask
 
     @log_call_to_history
     def clear_flags(self):
@@ -788,7 +795,19 @@ def clear_flags(self):
             sdf._init_flags()
         self._flag.clear()
 
-    def _create_index_if_needed(self):
+    def _create_index_if_needed(self, skipflags=False):
+        """
+        Parameters
+        ----------
+        skipflags : bool, optional
+            If True, do not read any flag files associated with these data.  The default is False.
+
+        Returns
+        -------
+        None.
+
+        """
+
         if self._selection is not None:
             return
         i = 0
@@ -809,6 +828,25 @@ def _create_index_if_needed(self):
         self._construct_procedure()
         self._construct_integration_number()
 
+        if skipflags:
+            return
+
+        # for directories with multiple FITS files and possibly multiple FLAG files
+        # we have to ensure the right flag file goes with the right FITS tile.
+        # The GBT convention is the same filename with '.flag' instead of '.fits'.
+        # We construct the flagfile and also pass in FITSINDEX column to ensure
+        # only the data associated with that file are flagged.
+        found_flags = False
+        for s in self._sdf:
+            p = Path(s.filename)
+            flagfile = p.with_suffix(".flag")
+            if flagfile.exists():
+                fi = uniq(s["FITSINDEX"])[0]
+                self.flags.read(flagfile, fitsindex=fi)
+                found_flags = True
+        if found_flags:
+            print("Flags were created from existing flag files. Use GBTFITSLoad.flags.show() to see them.")
+
     def _construct_procedure(self):
         """
         Construct the procedure string (PROC) from OBSMODE and add it to the index (i.e., a new SDFITS column).
@@ -2283,6 +2321,7 @@ def write(
         self,
         fileobj,
         multifile=True,
+        flags=True,
         verbose=False,
         output_verify="exception",
         overwrite=False,
@@ -2300,6 +2339,8 @@ def write(
         multifile: bool, optional
             If True, write to multiple files if and only if there are multiple SDFITS files in this GBTFITSLoad.
             Otherwise, write to a single SDFITS file.
+        flags: bool, optional
+            If True, write the applied flags to a `FLAGS` column in the binary table.
         verbose: bool, optional
             If True, print out some information about number of rows written per file
         output_verify : str
@@ -2349,6 +2390,12 @@ def write(
                     rows.sort()
                     lr = len(rows)
                     if lr > 0:
+                        if flags:  # update the flags before we select rows
+                            flagval = self._sdf[k]._flagmask[b].astype(np.uint8)
+                            dim1 = np.shape(flagval)[1]
+                            form = f"{dim1}B"
+                            c = fits.Column(name="FLAGS", format=form, array=flagval)
+                            self._sdf[k]._update_binary_table_column({"FLAGS": c})
                         ob = self._sdf[k]._bintable_from_rows(rows, b)
                         if len(ob.data) > 0:
                             outhdu.append(ob)
@@ -2373,7 +2420,7 @@ def write(
             if verbose:
                 print(f"Total of {total_rows_written} rows written to files.")
         else:
-            hdu = self._sdf[fi[0]]._hdu[fi[0]].copy()
+            hdu = self._sdf[fi[0]]._hdu[0].copy()
             outhdu = fits.HDUList(hdu)
             for k in fi:
                 df = select_from("FITSINDEX", k, _final)
@@ -2383,6 +2430,13 @@ def write(
                     rows.sort()
                     lr = len(rows)
                     if lr > 0:
+                        if flags:  # update the flags before we select rows
+                            flagval = self._sdf[k]._flagmask[b].astype(np.uint8)
+                            dim1 = np.shape(flagval)[1]
+                            form = f"{dim1}B"
+                            # tdim = f"({dim1}, 1, 1, 1)" # let fitsio do this
+                            c = fits.Column(name="FLAGS", format=form, array=flagval)
+                            self._sdf[k]._update_binary_table_column({"FLAGS": c})
                         ob = self._sdf[k]._bintable_from_rows(rows, b)
                         if len(ob.data) > 0:
                             outhdu.append(ob)
@@ -2471,7 +2525,7 @@ def __setitem__(self, items, values):
         col_exists = len(set(self.columns).intersection(iset)) > 0
         # col_in_selection =
         if col_exists:
-            warnings.warn("Changing an existing SDFITS column")
+            warnings.warn(f"Changing an existing SDFITS column {items}")
         # now deal with values as arrays
         is_array = False
         if isinstance(values, (Sequence, np.ndarray)) and not isinstance(values, str):