Merge pull request #597 from rosannadeckert/dev-build

Changed `evt/spms` to allow for event building with zero-suppression data

Author: gipert

src/pygama/evt/build_evt.py

@@ -414,7 +414,12 @@ def evaluate_expression(
         #     pygama.evt.modules.spms.my_func([...], arg1=val, arg2=val)
 
         # get arguments list passed to the function (outermost parentheses)
-        args_str = re.search(r"\((.*)\)$", expr.strip()).group(1)
+        result = re.search(r"\((.*)\)$", expr.strip(), re.DOTALL)
+        if result is None:
+            msg = f"could not parse the function arguments in '{expr}'"
+            raise RuntimeError(msg)
+
+        args_str = result.group(1)
 
         # handle tier scoping: evt.<>
         args_str = args_str.replace("evt.", "")

src/pygama/evt/modules/spms.py

@@ -88,11 +88,30 @@ def gather_pulse_data(
         lgdo_obj = lh5.read(
             f"/{channel}/{tierinfo.group}/{column}", tierinfo.file, idx=idx
         )
-        data = lgdo_obj.view_as(library="ak")
+        data = lgdo_obj.view_as(library="np")
 
         # remove nans (this happens when SiPM data is stored as ArrayOfEqualSizedArrays)
         data = ak.drop_none(ak.nan_to_none(data))
 
+        # number of channels per event
+        evt_length = np.diff(np.insert(tcm.cumulative_length, 0, 0))
+
+        # construct global event ID's
+        glob_ids = np.repeat(
+            np.arange(0, tcm.cumulative_length.shape[0], 1), evt_length
+        )
+        glob_ids_ch = glob_ids[
+            tcm.id == table_id
+        ]  # global ID's where channel had a trigger
+
+        # count number of hits per channel for global events with trigger in channel, else 0
+        glob_ids_cts = np.zeros(tcm.cumulative_length.shape[0], dtype=int)
+        glob_ids_cts[glob_ids_ch] = ak.count(data, axis=1)
+
+        # insert empty row [] for global events with no trigger in channel
+        # unflatten to the number of hits in channel otherwise
+        data = ak.unflatten(ak.flatten(data), glob_ids_cts)
+
         # increase the dimensionality by one (events)
         data = ak.unflatten(data, np.full(data.layout.length, 1, dtype="uint8"))
 
@@ -161,17 +180,18 @@ def gather_tcm_data(
     for field in ("id", "idx"):
         tcm_vov[field] = types.VectorOfVectors(
             flattened_data=tcm._asdict()[field], cumulative_length=tcm.cumulative_length
-        ).view_as("ak")
+        )
 
     # list user wanted table names
     table_ids = [
         utils.get_tcm_id_by_pattern(datainfo.hit.table_fmt, id) for id in table_names
     ]
     # find them in tcm.id (we'll filter the rest out)
-    locs = np.isin(tcm_vov["id"], table_ids)
+    tcm_id_padded = tcm_vov["id"].to_aoesa().view_as("np")
+    locs = np.isin(tcm_id_padded, table_ids)
 
     # select tcm field requested by the user
-    data = tcm_vov[tcm_field]
+    data = tcm_vov[tcm_field].view_as("ak")
 
     # apply mask
     # NOTE: need to cast to irregular axes, otherwise the masking result is
@@ -201,6 +221,7 @@ def gather_tcm_data(
             raise ValueError(msg)
 
         # convert the 3D mask to a 2D mask (can be used to filter table_ids)
+        pulse_mask = pulse_mask[ak.num(pulse_mask, axis=2) > 0]
         ch_mask = ak.sum(pulse_mask, axis=-1) > 0
 
         # apply the mask

src/pygama/hit/build_hit.py

@@ -109,8 +109,8 @@ def build_hit(
             # sanitize config
             hit_config = utils.load_dict(hit_config)
 
+        lh5_tables_config = {}
         if lh5_tables is None:
-            lh5_tables_config = {}
             if "dsp" in ls(infile):
                 log.debug("found candidate table /dsp")
                 lh5_tables_config["dsp"] = hit_config