Merge branch '88-retrieving-trigger-bit-information-from-trigger_mask' into 'master'

Resolve "Retrieving trigger bit information from trigger_mask"

Closes #88

See merge request !61

CHANGELOG.rst

 Unreleased changes
 ------------------
 
+* Added ``km3io.tools.is_bit_set()`` along with some special methods to check
+  if a given ``trigger_mask`` (of an event or a hit) has a specific trigger
+  bit set, via ``km3io.tools.is_3dmuon``, ``km3io.tools.is_3dshower`` and
+  ``km3io.tools.is_mxshower``
+
 Version 0
 ---------
 

km3io/tools.py

@@ -491,3 +491,47 @@ def usr(objects, field):
     available_fields = objects.usr_names[0].tolist()
     idx = available_fields.index(field)
     return objects.usr[:, idx]
+
+
+def is_bit_set(value, bit_position):
+    """Returns true if a bit at the given position is 1.
+
+    value: int or array(int)
+      The value to check, can be a single value or an array of values.
+    bit_position: int
+      0 for the first position, 1 for the second etc.
+    """
+    return (np.array(value) & (1 << bit_position)).astype(bool)
+
+
+def is_3dshower(trigger_mask):
+    """Returns True if the trigger mask contains the 3D shower flag.
+
+    Parameters
+    ----------
+    trigger_mask : int or array(int)
+      A value or an array of the trigger_mask, either of an event, or a hit.
+    """
+    return is_bit_set(trigger_mask, ktrg.JTRIGGER3DSHOWER)
+
+
+def is_mxshower(trigger_mask):
+    """Returns True if the trigger mask contains the MX shower flag.
+
+    Parameters
+    ----------
+    trigger_mask : int or array(int)
+      A value or an array of the trigger_mask, either of an event, or a hit.
+    """
+    return is_bit_set(trigger_mask, ktrg.JTRIGGERMXSHOWER)
+
+
+def is_3dmuon(trigger_mask):
+    """Returns True if the trigger mask contains the 3D muon flag.
+
+    Parameters
+    ----------
+    trigger_mask : int or array(int)
+      A value or an array of the trigger_mask, either of an event, or a hit.
+    """
+    return is_bit_set(trigger_mask, ktrg.JTRIGGER3DMUON)

tests/test_tools.py

@@ -8,7 +8,6 @@ from pathlib import Path
 from km3net_testdata import data_path
 from km3io.definitions import fitparameters as kfit
 
-
 from km3io import OfflineReader
 from km3io.tools import (
     to_num,
@@ -28,6 +27,10 @@ from km3io.tools import (
     best_dusjshower,
     is_cc,
     usr,
+    is_bit_set,
+    is_3dshower,
+    is_mxshower,
+    is_3dmuon,
 )
 
 OFFLINE_FILE = OfflineReader(data_path("offline/km3net_offline.root"))
@@ -581,3 +584,61 @@ class TestUsr(unittest.TestCase):
             usr(OFFLINE_MC_TRACK_USR.mc_tracks[0], "bx").tolist(),
             atol=0.0001,
         )
+
+
+class TestIsBitSet(unittest.TestCase):
+    def test_is_bit_set_for_single_values(self):
+        value = 2  # 10
+        assert not is_bit_set(value, 0)
+        assert is_bit_set(value, 1)
+        value = 42  # 101010
+        assert not is_bit_set(value, 0)
+        assert is_bit_set(value, 1)
+        assert not is_bit_set(value, 2)
+        assert is_bit_set(value, 3)
+        assert not is_bit_set(value, 4)
+        assert is_bit_set(value, 5)
+
+    def test_is_bit_set_for_lists(self):
+        values = [2, 42, 4]
+        assert np.allclose([True, True, False], is_bit_set(values, 1))
+
+    def test_is_bit_set_for_numpy_arrays(self):
+        values = np.array([2, 42, 4])
+        assert np.allclose([True, True, False], is_bit_set(values, 1))
+
+    def test_is_bit_set_for_awkward_arrays(self):
+        values = ak.Array([2, 42, 4])
+        assert np.allclose([True, True, False], is_bit_set(values, 1))
+
+
+class TestTriggerMaskChecks(unittest.TestCase):
+    def test_is_3dshower(self):
+        assert np.allclose(
+            [True, True, True, True, True, False, False, True, True, True],
+            is_3dshower(OFFLINE_FILE.events.trigger_mask),
+        )
+        assert np.allclose(
+            [True, True, True, True, True, True, True, True, True, False],
+            is_3dshower(GENHEN_OFFLINE_FILE.events.trigger_mask),
+        )
+
+    def test_is_mxshower(self):
+        assert np.allclose(
+            [True, True, True, True, True, True, True, True, True, True],
+            is_mxshower(OFFLINE_FILE.events.trigger_mask),
+        )
+        assert np.allclose(
+            [False, False, False, False, False, False, False, False, False, False],
+            is_mxshower(GENHEN_OFFLINE_FILE.events.trigger_mask),
+        )
+
+    def test_is_3dmuon(self):
+        assert np.allclose(
+            [True, True, True, True, True, False, False, True, True, True],
+            is_3dmuon(OFFLINE_FILE.events.trigger_mask),
+        )
+        assert np.allclose(
+            [False, False, False, True, False, False, True, False, True, True],
+            is_3dmuon(GENHEN_OFFLINE_FILE.events.trigger_mask),
+        )