aminnj / pdroot Goto Github PK

import pandas as pd
import numpy as np
import pdroot # powers up pandas

df = pd.DataFrame(dict(foo=np.random.random(100), bar=np.random.random(100)))

# write out dataframe to a ROOT file
df.to_root("test.root")

# read ROOT files and optionally specify certain columns and/or a range of rows
df = pd.read_root("test.root", columns=["foo"], entry_start=0, entry_stop=50)

# expression string and a query/selection string
df.draw("mass+0.1", "0.1 < foo < 0.2", bins="200,0,10")

# 2D with "x:y"
df.draw("mass:foo+1", "0.1 < foo < 0.2")

df = pd.read_root("nano.root", columns=["/Electron_(pt|eta|phi|mass)$/", "MET_pt"])
df.head()

>>> df["Electron_pt"].ak() 

<Array [[], [121], ... [179], [19.7, 14.4]] type='4273 * var * float64'>

>>> df[["Electron_pt","Electron_eta"]].ak()

<Array [{Electron_pt: [], ... -1.82]}] type='4273 * {"Electron_pt": option[var * fl...'>

from pdroot import to_pandas
dfc["Electron_good"] = to_pandas(dfc["Electron_pt"].ak() > 40)

# supports reduction operators:
#    min/max/sum/mean/length/len/argmin/margmax
#    (ROOT's Length$ -> length, etc)
df.draw("length(Jet_pt)")
df.draw("sum(Jet_pt>10)", "MET_pt>40", bins="5,-0.5,4.5")
df.draw("max(abs(Jet_eta))")
df.draw("Jet_eta[argmax(Jet_pt)]")

# combine event-level and object-level selection
df.draw("Jet_pt", "abs(Jet_eta) > 1.0 and MET_pt > 10")

# 2D
df.draw("Jet_pt:Jet_eta", "MET_pt > 40.")

# indexing; imagine you are operating row-by-row, so Jet_pt[0], not Jet_pt[:,0]
df.draw("Jet_pt[0]:Jet_eta[0]", "MET_pt > 10")

# combine reduction operators with fancy indexing
df.draw("sum(Jet_pt[abs(Jet_eta)<2.0])", bins="100,0,100")

# use the underlying array before a histogram is created
df["ht"] = df.draw("sum(Jet_pt[Jet_pt>40])", to_array=True)
df["ht"] = df.adraw("sum(Jet_pt[Jet_pt>40])") # think "*a*rray draw"

>>> pdroot.parse.to_ak_expr("sum(Jet_pt[:2])") # sum of first/leading two jet pTs

'ak.sum(ak.pad_none(Jet_pt, 3, clip=True)[:, :2], axis=-1)'

df = pdroot.ChunkDataFrame(filename="nano.root", entry_start=0, entry_stop=100e3)

# keeps track of original indexing, so if subsetted,
# any newly read arrays will be modified to match
df = df[df["MET_pt"]>40]

pt = df["Jet_pt"].ak()
df["ht"] = ak.sum(pt[pt>40])
# or
df["ht"] = df.adraw("sum(Jet_pt*(Jet_pt>40))")

df.head()