TreeToH5.py

# standard library imports
from __future__ import absolute_import, division, print_function

# standard numerical library imports
import numpy as np
from numpy import stack, vstack, hstack, hsplit
import matplotlib.pyplot as plt

import math

# uproot, to parse the trees
import uproot

branches = [
    'fatjet_pt',
    'fatjet_eta',
    'fatjet_eta_detector',
    'fatjet_phi',
    'fatjet_CaloMass',
    'fatjet_TrackAssistedMassCalibrated',
    'fatjet_Split12',
    'fatjet_Split23',
    'fatjet_ECF1',
    'fatjet_ECF2',
    'fatjet_ECF3',
    'fatjet_C2',
    'fatjet_D2',
    'fatjet_Tau1_wta',
    'fatjet_Tau2_wta',
    'fatjet_Tau3_wta',
    'fatjet_Tau21_wta',
    'fatjet_Tau32_wta',
    'fatjet_KtDR',
    'fatjet_Qw',
    'fatjet_FoxWolfram20',
    'fatjet_Angularity',
    'fatjet_Aplanarity',
    'fatjet_Charge',
    'fatjet_dRmatched_maxEParton_flavor',
    'fatjet_ghost_assc_flavor',
    'fatjet_Lund1_Pt1',
    'fatjet_Lund1_Pt2',
    'fatjet_Lund1_Z',
    'fatjet_Lund1_DeltaR',
    'fatjet_Lund2_Pt1',
    'fatjet_Lund2_Pt2',
    'fatjet_Lund2_Z',
    'fatjet_Lund2_DeltaR'
    ]
input_dict = dict((branch,np.array([0],dtype=np.float64)) for branch in branches)
#for k, v in input_dict.items():
#    print(k, v)

idx = 0
for arrays in uproot.iterate("user.mleblanc.17582546._000013.tree.root", "SubstructureJetTree/nominal", branches, reportpath=True, reportentries=True):
    print(arrays)

    for k,v in arrays[3].items():
        print(k,v)
        input_dict[k] = np.append(input_dict[k],v)

    idx+=1

# Get rid of the first entries, since they're placeholders anyway
input_dict = {k: input_dict[k][1:] for k in input_dict}

print(input_dict)