Reading multiple trees using uproot.iterate()

[ligerlac]

Assume a dataset spread across many files and you want to read branch_a_1 and branch_a_2 from tree_a and branch_b from tree_b. The baseline approach would be:

for fn in glob.glob('files*.root'):
    with uproot.open(fn) as f:
        a_1 = f['tree_a/branch_a_1']
        a_2 = f['tree_a/branch_a_2']
        b = f['tree_b/branch_b']
        do_something ...

How can this be achieved with uproot.iterate()? I came up with this non-optimal work-around:

for tup in zip(
    uproot.iterate('files*.root:tree_a', ['branch_a_1', 'branch_a_2']),
    uproot.iterate('files*.root:tree_b', ['branch_b'])
):
    arr = ak.zip({f: a[f] for a in tup for f in ak.fields(a)})
    a_1 = arr['branch_a_1']
    a_2 = arr['branch_a_1']
    b = arr['branch_b']
    do_something ...

Is there a more elegant solution?

[jpivarski]

What you've described is a reasonable way to go about it. You probably want to set step_size in uproot.iterate to an integer, to ensure that the arrays can always be ak.zipped together (ensure that they broadcast, by having the same number of entries).

[ligerlac]

Thanks for the swift answer, @jpivarski. The thing with step_size is a good point. I assume the work-around will result in opening each file twice? At least that would explain the ~20% longer run time I observed compared to the baseline approach.

[jpivarski]

I missed the fact that it was the same files, but different TTrees. If you want to ensure that you open the files only once, use uproot.TTree.iterate instead of uproot.iterare. This involves more bookkeeping, as you have to uproot.open each file manually.