diff --git a/imap_processing/quality_flags.py b/imap_processing/quality_flags.py new file mode 100644 index 000000000..5a1c2c24f --- /dev/null +++ b/imap_processing/quality_flags.py @@ -0,0 +1,68 @@ +"""Bitwise flagging.""" + +from enum import IntFlag + + +class FlagNameMixin(IntFlag): + """Modifies flags for Python versions < 3.11.""" + + @property + def name(self) -> str: + """ + Override the default name property to handle combined flags. + + Returns + ------- + combined_name : str + The combined name of the individual flags. + """ + if self._name_ is not None: + return self._name_ + + members = [member for member in self.__class__ if member & self == member] + return "|".join(str(m).split(".", 1)[-1] for m in members if m != 0x0) + + +class CommonFlags(FlagNameMixin): + """Common quality flags.""" + + NONE = 0x0 + INF = 2**0 # bit 0, Infinite value + NEG = 2**1 # bit 1, Negative value + + +class ENAFlags(FlagNameMixin): + """Common ENA flags.""" + + BADSPIN = 2**2 # bit 2, Bad spin + + +class ImapUltraFlags(FlagNameMixin): + """IMAP Ultra flags.""" + + NONE = CommonFlags.NONE + INF = CommonFlags.INF # bit 0 + NEG = CommonFlags.NEG # bit 1 + BADSPIN = ENAFlags.BADSPIN # bit 2 + FLAG1 = 2**3 # bit 2 + + +class ImapLoFlags(FlagNameMixin): + """IMAP Lo flags.""" + + NONE = CommonFlags.NONE + INF = CommonFlags.INF # bit 0 + NEG = CommonFlags.NEG # bit 1 + BADSPIN = ENAFlags.BADSPIN # bit 2 + FLAG2 = 2**3 # bit 2 + + +class HitFlags( + FlagNameMixin, +): + """Hit flags.""" + + NONE = CommonFlags.NONE + INF = CommonFlags.INF # bit 0 + NEG = CommonFlags.NEG # bit 1 + FLAG3 = 2**2 # bit 2 diff --git a/imap_processing/tests/test_quality_flags.py b/imap_processing/tests/test_quality_flags.py new file mode 100644 index 000000000..cb16f1d66 --- /dev/null +++ b/imap_processing/tests/test_quality_flags.py @@ -0,0 +1,71 @@ +"""Test bitwise flagging.""" + +import numpy as np + +from imap_processing.quality_flags import HitFlags, ImapLoFlags, ImapUltraFlags + + +def test_quality_flags(): + """Test the bitwise operations.""" + + # Test individual flags + assert HitFlags.NONE == 0x0 + assert ImapUltraFlags.NONE == 0x0 + assert ImapLoFlags.NONE == 0x0 + assert HitFlags.INF == 2**0 + assert ImapUltraFlags.INF == 2**0 + assert ImapLoFlags.INF == 2**0 + assert HitFlags.NEG == 2**1 + assert ImapUltraFlags.NEG == 2**1 + assert ImapLoFlags.NEG == 2**1 + assert ImapUltraFlags.BADSPIN == 2**2 + assert ImapLoFlags.BADSPIN == 2**2 + assert ImapUltraFlags.FLAG1 == 2**3 + assert ImapLoFlags.FLAG2 == 2**3 + assert HitFlags.FLAG3 == 2**2 + + # Test combined flags for Ultra + flag = ( + ImapUltraFlags.INF + | ImapUltraFlags.NEG + | ImapUltraFlags.BADSPIN + | ImapUltraFlags.FLAG1 + ) + assert flag & ImapUltraFlags.INF + assert flag & ImapUltraFlags.BADSPIN + assert flag & ImapUltraFlags.FLAG1 + assert flag.name == "INF|NEG|BADSPIN|FLAG1" + assert flag.value == 15 + + # Test combined flags for Lo + flag = ImapLoFlags.INF | ImapLoFlags.NEG | ImapLoFlags.BADSPIN | ImapLoFlags.FLAG2 + assert flag & ImapLoFlags.INF + assert flag & ImapLoFlags.BADSPIN + assert flag & ImapLoFlags.FLAG2 + assert flag.name == "INF|NEG|BADSPIN|FLAG2" + assert flag.value == 15 + + # Test combined flags for HIT + flag = HitFlags.INF | HitFlags.NEG | HitFlags.FLAG3 + assert flag & HitFlags.INF + assert flag & HitFlags.FLAG3 + assert flag.name == "INF|NEG|FLAG3" + assert flag.value == 7 + + # Test use-case for Ultra + data = np.array([-6, np.inf, 2, 3]) + quality = np.array( + [ + ImapUltraFlags.INF | ImapUltraFlags.NEG, + ImapUltraFlags.INF, + ImapUltraFlags.NONE, + ImapUltraFlags.NONE, + ] + ) + # Select data without INF flags + non_inf_mask = (quality & ImapUltraFlags.INF.value) == 0 + np.array_equal(data[non_inf_mask], np.array([-6, 2, 3])) + + # Select data without NEG or INF flags + non_neg_mask = (quality & ImapUltraFlags.NEG.value) == 0 + np.array_equal(data[non_inf_mask & non_neg_mask], np.array([2, 3]))