Skip to content
Snippets Groups Projects
Commit 4d80f759 authored by Tamas Gal's avatar Tamas Gal :speech_balloon:
Browse files

Merge branch 'improve_ak.Record_indexing' into 'master'

Improve ak.record indexing

See merge request !88
parents eefed662 2fc5789d
No related branches found
Tags v1.2.3
1 merge request!88Improve ak.record indexing
Pipeline #62151 passed
......@@ -263,8 +263,7 @@ def best_track(tracks, startend=None, minmax=None, stages=None):
out = apply_mask(tracks, m3)
if original_ndim == 1:
if isinstance(out, ak.Record):
first_and_only_element_mask = ak.Array([True])
return apply_mask(out, first_and_only_element_mask)
return select_first_entry_from_record(out)
return out[0]
return out[:, 0]
......@@ -280,6 +279,18 @@ def apply_mask(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):
"""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