Skip to content
Snippets Groups Projects
Commit 2fc5789d authored by Lukas Hennig's avatar Lukas Hennig
Browse files

improved selection of first entry from record

parent eefed662
No related branches found
Tags v0.1.2 v1.2.0
1 merge request!88Improve ak.record indexing
Pipeline #56189 passed
...@@ -263,8 +263,7 @@ def best_track(tracks, startend=None, minmax=None, stages=None): ...@@ -263,8 +263,7 @@ def best_track(tracks, startend=None, minmax=None, stages=None):
out = apply_mask(tracks, m3) out = apply_mask(tracks, m3)
if original_ndim == 1: if original_ndim == 1:
if isinstance(out, ak.Record): if isinstance(out, ak.Record):
first_and_only_element_mask = ak.Array([True]) return select_first_entry_from_record(out)
return apply_mask(out, first_and_only_element_mask)
return out[0] return out[0]
return out[:, 0] return out[:, 0]
...@@ -280,6 +279,18 @@ def apply_mask(record, mask): ...@@ -280,6 +279,18 @@ def apply_mask(record, mask):
return record[mask] return record[mask]
def select_first_entry_from_record(record):
assert isinstance(
record, ak.Record
), f"Input has to be ak.Record, but is {type(record)}"
first_entry_data = {}
for field in record.fields:
value = record[field]
first_entry_data[field] = value[0]
return ak.Record(first_entry_data)
def mask(arr, sequence=None, startend=None, minmax=None, atleast=None): def mask(arr, sequence=None, startend=None, minmax=None, atleast=None):
"""Return a boolean mask which mask each nested sub-array for a condition. """Return a boolean mask which mask each nested sub-array for a condition.
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment