From 2337326392a4b9e847e961833defe1f070e08932 Mon Sep 17 00:00:00 2001 From: ggmarshall Date: Wed, 9 Oct 2024 13:44:27 +0200 Subject: [PATCH 1/3] make more generic to handle new tiers --- src/pygama/evt/utils.py | 91 ++++++++++++++++++++++------------------- 1 file changed, 48 insertions(+), 43 deletions(-) diff --git a/src/pygama/evt/utils.py b/src/pygama/evt/utils.py index 4f8391353..fa559f405 100644 --- a/src/pygama/evt/utils.py +++ b/src/pygama/evt/utils.py @@ -16,16 +16,16 @@ H5DataLoc = namedtuple( "H5DataLoc", ("file", "group", "table_fmt"), defaults=3 * (None,) ) - -DataInfo = namedtuple( - "DataInfo", ("raw", "tcm", "dsp", "hit", "evt"), defaults=5 * (None,) -) +DataInfo = namedtuple("DataInfo", ("raw"), defaults=1 * (None,)) TCMData = namedtuple("TCMData", ("id", "idx", "cumulative_length")) def make_files_config(data: dict): - if not isinstance(data, DataInfo): + if not isinstance(data, tuple): + DataInfo = namedtuple( + "DataInfo", tuple(data.keys()), defaults=len(data.keys()) * (None,) + ) return DataInfo( *[ H5DataLoc(*data[tier]) if tier in data else H5DataLoc() @@ -72,7 +72,7 @@ def find_parameters( idx_ch, field_list, ) -> dict: - """Finds and returns parameters from `hit` and `dsp` tiers. + """Finds and returns parameters from non `tcm`, `evt` tiers. Parameters ---------- @@ -83,43 +83,38 @@ def find_parameters( idx_ch index array of entries to be read from datainfo. field_list - list of tuples ``(tier, field)`` to be found in the `hit/dsp` tiers. + list of tuples ``(tier, field)`` to be found in non `tcm`, `evt` tiers. """ f = make_files_config(datainfo) - # find fields in either dsp, hit - dsp_flds = [e[1] for e in field_list if e[0] == f.dsp.group] - hit_flds = [e[1] for e in field_list if e[0] == f.hit.group] + final_dict = {} - hit_dict, dsp_dict = {}, {} + for name, tier in f._asdict().items(): + if name not in ["tcm", "evt"] and tier.file is not None: # skip other tables + keys = [ + k.split("/")[-1] + for k in lh5.ls(tier.file, f"{ch.replace('/', '')}/{tier.group}/") + ] + flds = [e[1] for e in field_list if e[0] == name and e[1] in keys] - if len(hit_flds) > 0: - hit_ak = lh5.read_as( - f"{ch.replace('/', '')}/{f.hit.group}/", - f.hit.file, - field_mask=hit_flds, - idx=idx_ch, - library="ak", - ) + if len(flds) > 0: + tier_ak = lh5.read_as( + f"{ch.replace('/', '')}/{tier.group}/", + tier.file, + field_mask=flds, + idx=idx_ch, + library="ak", + ) - hit_dict = dict( - zip([f"{f.hit.group}_" + e for e in ak.fields(hit_ak)], ak.unzip(hit_ak)) - ) + tier_dict = dict( + zip( + [f"{name}_" + e for e in ak.fields(tier_ak)], + ak.unzip(tier_ak), + ) + ) + final_dict = final_dict | tier_dict - if len(dsp_flds) > 0: - dsp_ak = lh5.read_as( - f"{ch.replace('/', '')}/{f.dsp.group}/", - f.dsp.file, - field_mask=dsp_flds, - idx=idx_ch, - library="ak", - ) - - dsp_dict = dict( - zip([f"{f.dsp.group}_" + e for e in ak.fields(dsp_ak)], ak.unzip(dsp_ak)) - ) - - return hit_dict | dsp_dict + return final_dict def get_data_at_channel( @@ -178,10 +173,14 @@ def get_data_at_channel( # evaluate expression # move tier+dots in expression to underscores (e.g. evt.foo -> evt_foo) + + new_expr = expr + for name in f._asdict(): + if name not in ["tcm", "raw"]: + new_expr = new_expr.replace(f"{name}.", f"{name}_") + res = eval( - expr.replace(f"{f.dsp.group}.", f"{f.dsp.group}_") - .replace(f"{f.hit.group}.", f"{f.hit.group}_") - .replace(f"{f.evt.group}.", ""), + new_expr, var, ) @@ -231,17 +230,23 @@ def get_mask_from_query( # get sub evt based query condition if needed if isinstance(query, str): - query_lst = re.findall(r"(hit|dsp).([a-zA-Z_$][\w$]*)", query) + query_lst = re.findall( + rf"({'|'.join(f._asdict().keys())}).([a-zA-Z_$][\w$]*)", query + ) query_var = find_parameters( datainfo=datainfo, ch=ch, idx_ch=idx_ch, field_list=query_lst, ) + + new_query = query + for name in f._asdict(): + if name not in ["tcm", "evt"]: + new_query = new_query.replace(f"{name}.", f"{name}_") + limarr = eval( - query.replace(f"{f.dsp.group}.", f"{f.dsp.group}_").replace( - f"{f.hit.group}.", f"{f.hit.group}_" - ), + new_query, query_var, ) From 8e823c0e42c22c252bcb819a3a0ff616fcc159d6 Mon Sep 17 00:00:00 2001 From: ggmarshall Date: Wed, 9 Oct 2024 14:39:23 +0200 Subject: [PATCH 2/3] fix for tests --- src/pygama/evt/utils.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/pygama/evt/utils.py b/src/pygama/evt/utils.py index fa559f405..3644e141d 100644 --- a/src/pygama/evt/utils.py +++ b/src/pygama/evt/utils.py @@ -16,13 +16,19 @@ H5DataLoc = namedtuple( "H5DataLoc", ("file", "group", "table_fmt"), defaults=3 * (None,) ) -DataInfo = namedtuple("DataInfo", ("raw"), defaults=1 * (None,)) +DataInfo = namedtuple("DataInfo", ("raw", "tcm", "evt"), defaults=3 * (None,)) TCMData = namedtuple("TCMData", ("id", "idx", "cumulative_length")) def make_files_config(data: dict): if not isinstance(data, tuple): + if "raw" not in data: + data["raw"] = (None,) + if "tcm" not in data: + data["tcm"] = (None,) + if "evt" not in data: + data["evt"] = (None,) DataInfo = namedtuple( "DataInfo", tuple(data.keys()), defaults=len(data.keys()) * (None,) ) From 2d143cd327f0e7fa6925b53f2485386ecc740c0b Mon Sep 17 00:00:00 2001 From: ggmarshall Date: Wed, 9 Oct 2024 14:52:17 +0200 Subject: [PATCH 3/3] evt replaced by empty string --- src/pygama/evt/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/pygama/evt/utils.py b/src/pygama/evt/utils.py index 3644e141d..4aedc1438 100644 --- a/src/pygama/evt/utils.py +++ b/src/pygama/evt/utils.py @@ -182,7 +182,9 @@ def get_data_at_channel( new_expr = expr for name in f._asdict(): - if name not in ["tcm", "raw"]: + if name == "evt": + new_expr = new_expr.replace(f"{name}.", "") + elif name not in ["tcm", "raw"]: new_expr = new_expr.replace(f"{name}.", f"{name}_") res = eval(