refactor quality flags to use quality class like others

IMAP-Science-Operations-Center · Aug 30, 2024 · 3802533 · 3802533
1 parent 9c8c457
commit 3802533
Show file tree

Hide file tree

Showing 5 changed files with 69 additions and 87 deletions.
diff --git a/imap_processing/cdf/config/imap_swapi_variable_attrs.yaml b/imap_processing/cdf/config/imap_swapi_variable_attrs.yaml
@@ -11,31 +11,13 @@ energy:
   VALIDMIN: 0
   VAR_TYPE: support_data
 
-flags:
-  CATDESC: Data quality flags
-  FIELDNAM: Quality Flags
-  FILLVAL: -9223370000000000000
-  FORMAT: I2
-  LABLAXIS: Flags
-  SCALE_TYP: linear
-  UNITS: " "
-  VALIDMAX: 1
-  VALIDMIN: 0
-  VAR_TYPE: support_data
-
 # <=== LABL_PTR_i Attributes ===>
 energy_label:
   CATDESC: Energy step id in lookup table
   FIELDNAM: Energy Step
   FORMAT: A2
   VAR_TYPE: metadata
 
-flags_label:
-  CATDESC: Data quality flags
-  FIELDNAM: Quality Flags
-  FORMAT: A20
-  VAR_TYPE: metadata
-
 # <=== Data Variables ===>
 counts_default: &counts_default
   DEPEND_0: epoch
@@ -55,11 +37,9 @@ flags_default:
   FIELDNAM: Quality Flag
   LABLAXIS: Flags
   DEPEND_0: epoch
-  DEPEND_1: flags
-  DEPEND_2: energy
+  DEPEND_1: energy
   DISPLAY_TYPE: spectrogram
-  LABL_PTR_1: flags_label
-  LABL_PTR_2: energy_label
+  LABL_PTR_1: energy_label
   FILLVAL: 4294967295
   FORMAT: I1
   UNITS: ' '

diff --git a/imap_processing/quality_flags.py b/imap_processing/quality_flags.py
@@ -66,3 +66,26 @@ class HitFlags(
     INF = CommonFlags.INF  # bit 0
     NEG = CommonFlags.NEG  # bit 1
     FLAG3 = 2**2  # bit 2
+
+
+class SWAPIFlags(
+    FlagNameMixin,
+):
+    """SWAPI flags."""
+
+    NONE = CommonFlags.NONE
+    INF = CommonFlags.INF  # bit 0
+    NEG = CommonFlags.NEG  # bit 1
+    SWP_PCEM_COMP = 2**2  # bit 2
+    SWP_SCEM_COMP = 2**3  # bit 3
+    SWP_COIN_COMP = 2**4  # bit 4
+    OVR_T_ST = 2**5  # bit 5
+    UND_T_ST = 2**6  # bit 6
+    PCEM_CNT_ST = 2**7  # bit 7
+    SCEM_CNT_ST = 2**8  # bit 8
+    PCEM_V_ST = 2**9  # bit 9
+    PCEM_I_ST = 2**10  # bit 10
+    PCEM_INT_ST = 2**11  # bit 11
+    SCEM_V_ST = 2**12  # bit 12
+    SCEM_I_ST = 2**13  # bit 13
+    SCEM_INT_ST = 2**14  # bit 14
diff --git a/imap_processing/swapi/l1/swapi_l1.py b/imap_processing/swapi/l1/swapi_l1.py
@@ -8,6 +8,7 @@
 
 from imap_processing import imap_module_directory
 from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
+from imap_processing.quality_flags import SWAPIFlags
 from imap_processing.swapi.swapi_utils import SWAPIAPID, SWAPIMODE
 from imap_processing.utils import packet_file_to_datasets
 
@@ -475,26 +476,12 @@ def process_swapi_science(
     # ===================================================================
     # Quality flags
     # ===================================================================
-    flags_name = [
-        "swp_pcem_comp",
-        "swp_scem_comp",
-        "swp_coin_comp",
-        "ovr_t_st",
-        "und_t_st",
-        "pcem_cnt_st",
-        "scem_cnt_st",
-        "pcem_v_st",
-        "pcem_i_st",
-        "pcem_int_st",
-        "scem_v_st",
-        "scem_i_st",
-        "scem_int_st",
-    ]
-    quality_flags_data = []
+    quality_flags_data = np.zeros((total_full_sweeps, 72), np.uint16)
+
     # Add science data quality flags
-    quality_flags_data.append(pcem_compression_flags)
-    quality_flags_data.append(scem_compression_flags)
-    quality_flags_data.append(coin_compression_flags)
+    quality_flags_data[pcem_compression_flags == 1] |= SWAPIFlags.SWP_PCEM_COMP
+    quality_flags_data[scem_compression_flags == 1] |= SWAPIFlags.SWP_SCEM_COMP
+    quality_flags_data[coin_compression_flags == 1] |= SWAPIFlags.SWP_COIN_COMP
 
     # Add housekeeping-derived quality flags
     # --------------------------------------
@@ -506,47 +493,46 @@ def process_swapi_science(
     good_sweep_times = good_sweep_sci["epoch"].data
     good_sweep_hk_data = hk_dataset.sel({"epoch": good_sweep_times})
 
-    # Since there will be one SWAPI HK packet for each SWAPI SCI packets, as
-    # both are recorded at 1 Hz(1 second), we can use this information to set
-    # quality flag for each science packet's data. Each packet contains
-    # information of one sequence data. Each sequence's
-    # PCEM_CNT0, PCEM_CNT1, PCEM_CNT2, PCEM_CNT3, PCEM_CNT4,
-    # PCEM_CNT5 data is coming from the same sciene packet and therefore
-    # it's going to share the same HK quality flag. That's why HK quality
-    # flag is repeated 6 times for each sequence(aka packet).
-    for flag_name in flags_name[3:]:
-        quality_flags_data.append(
-            np.repeat(good_sweep_hk_data[flag_name].data, 6).reshape(-1, 72)
-        )
+    # Since there is one SWAPI HK packet for each SWAPI SCI packet,
+    # and both are recorded at 1 Hz (1 packet per second),
+    # we can leverage this to set the quality flags for each science
+    # packet's data. Each SWAPI science packet represents
+    # one sequence of data, where the sequence includes measurements
+    # like PCEM_CNT0, PCEM_CNT1, PCEM_CNT2, PCEM_CNT3,
+    # PCEM_CNT4, and PCEM_CNT5. Because all these measurements come
+    # from the same science packet, they should share
+    # the same HK quality flag. This is why the HK quality flag is
+    # repeated 6 times, once for each measurement within
+    # the sequence (each packet corresponds to one sequence).
+
+    hk_flags_name = [
+        "OVR_T_ST",
+        "UND_T_ST",
+        "PCEM_CNT_ST",
+        "SCEM_CNT_ST",
+        "PCEM_V_ST",
+        "PCEM_I_ST",
+        "PCEM_INT_ST",
+        "SCEM_V_ST",
+        "SCEM_I_ST",
+        "SCEM_INT_ST",
+    ]
 
-    # Before, individual quality flags data were in this shape:
-    #   (number_of_sweep, energy_step)
-    # Now, to group array of flags data into this shape:
-    #   (number_of_sweep, number_of_flags, energy_step)
-    # With this, each sweep has its quality flag.
-    quality_flags_data = np.stack(quality_flags_data, axis=1)
+    for flag_name in hk_flags_name:
+        current_flag = np.repeat(good_sweep_hk_data[flag_name.lower()].data, 6).reshape(
+            -1, 72
+        )
+        # Use getattr to dynamically access the flag in SWAPIFlags class
+        flag_to_set = getattr(SWAPIFlags, flag_name)
+        # set the quality flag for each data
+        quality_flags_data[current_flag == 1] |= flag_to_set
 
     swp_flags = xr.DataArray(
-        np.array(quality_flags_data, dtype=np.uint8),
-        dims=["epoch", "flags", "energy"],
+        quality_flags_data,
+        dims=["epoch", "energy"],
         attrs=cdf_manager.get_variable_attributes("flags_default"),
     )
 
-    # Quality flags of sweep data
-    flags = xr.DataArray(
-        np.arange(13),
-        name="flags",
-        dims=["flags"],
-        attrs=cdf_manager.get_variable_attributes("flags"),
-    )
-
-    flags_label = xr.DataArray(
-        flags_name,
-        name="flags_label",
-        dims=["flags_label"],
-        attrs=cdf_manager.get_variable_attributes("flags_label"),
-    )
-
     # ===================================================================
     # Step 3: Create xarray.Dataset
     # ===================================================================
@@ -589,8 +575,6 @@ def process_swapi_science(
             "epoch": epoch_time,
             "energy": energy,
             "energy_label": energy_label,
-            "flags": flags,
-            "flags_label": flags_label,
         },
         attrs=l1_global_attrs,
     )
@@ -620,8 +604,9 @@ def process_swapi_science(
     # Uncertainty is quantified for the PCEM, SCEM, and COIN counts.
     # The Poisson contribution is
     #   uncertainty = sqrt(count)
-
-    # TODO: fill in with actual formula once SWAPI provides it.
+    # TODO:
+    # Above uncertaintly formula will change in the future.
+    # Replace it with actual formula once SWAPI provides it.
     # Right now, we are using sqrt(count) as a placeholder
     dataset["swp_pcem_err_plus"] = xr.DataArray(
         np.sqrt(swp_pcem_counts),

diff --git a/imap_processing/swapi/l2/swapi_l2.py b/imap_processing/swapi/l2/swapi_l2.py
@@ -50,8 +50,6 @@ def swapi_l2(l1_dataset: xr.Dataset, data_version: str) -> xr.Dataset:
         "epoch",
         "energy",
         "energy_label",
-        "flags",
-        "flags_label",
         "swp_flags",
     ]
     l2_dataset = l1_dataset[l1_data_keys]

diff --git a/imap_processing/tests/swapi/test_swapi_l1.py b/imap_processing/tests/swapi/test_swapi_l1.py
@@ -79,7 +79,7 @@ def test_decompress_count():
     # compressed + no-overflow, compressed + overflow, no compression
     raw_values = np.array([[12, 0xFFFF, 12]])
     compression_flag = np.array([[1, 1, 0]])
-    expected = np.array([[12 * 16, -1, 12]])
+    expected = np.array([[12 * 16, np.iinfo(np.int32).max, 12]], dtype=np.int32)
     returned_value = decompress_count(raw_values, compression_flag)
     np.testing.assert_array_equal(returned_value, expected)
 
@@ -220,8 +220,6 @@ def test_process_swapi_science(decom_test_data):
         "epoch": 3,
         "energy": 72,
         "energy_label": 72,
-        "flags": 13,
-        "flags_label": 13,
     }
     # Test epoch data is correct
     expected_epoch_datetime = met_to_j2000ns([48, 60, 72])
@@ -321,8 +319,6 @@ def test_process_swapi_science(decom_test_data):
         "epoch": 2,
         "energy": 72,
         "energy_label": 72,
-        "flags": 13,
-        "flags_label": 13,
     }
 
     # Test CDF File