update expression that checks if stage in tracks.rec_stages

src/km3pipe/utils/h5extractf.py

@@ -14,6 +14,7 @@
     -h --help                   Show this screen.
 
 """
+
 import time
 from uuid import uuid4
 import numpy as np
@@ -305,7 +306,7 @@ def _yield_tracks(tracks, fields, without_full_reco=False):
     n_columns = max(km3io.definitions.fitparameters.values()) + 1
 
     for name, (stage, func) in track_fmap.items():
-        if stage is None or stage in tracks.rec_stages:
+        if (stage is None) or (ak.any(stage == tracks.rec_stages)):
             if func is None:
                 sel_tracks = tracks
             else: