diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ab36d348a560c71fc683ffc63c6e47b5ad500a67..10611ad5446d8a2ceed111ef9bc87a94f5b6737a 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -1,6 +1,11 @@ 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 --------- diff --git a/km3io/tools.py b/km3io/tools.py index 56855789ee8e7307f9e712afcb8028a9d16c5871..538c9660d739c6d8408278bba0b2dabe32cfca5a 100644 --- a/km3io/tools.py +++ b/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) diff --git a/tests/test_tools.py b/tests/test_tools.py index 6610a4b28264973af36ecdf09f8f757ec9c450f0..753db2b9d900a7827bd0dea694c5f5fca1bfba7c 100644 --- a/tests/test_tools.py +++ b/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), + )