Writing TTree in ROOT file inside a loop

[a-akram]

Dear All,

I am trying to write a root file inside a for-loop in Python. I have a process() that returns a DataFrame per event with just two columns: hit_id, track_id. Now I want to write this DataFrame as a TTree to a root file. Note that the size of DataFrame changes in each event. This is what I have tried:

root_file = uproot.recreate(os.path.join(outputdir, "trackml.root"))

for filename in all_files[:max_evts]:
    df = process(filename, **vars(args))    # df >> {hit_id, track_id}
    root_file["TrackML"] = df

root_file.close()

Everything runs perfectly. But when I read the trackml.root then I see multiple instances of TrackML TTree as TrackML;1, TrackML;2,..., TrackML;100. What I was expecting is a single TrackML TTree that has events up to max_evts so I can loop over trackml.root access hit_id, track_id event by event later inside ROOT environment.

Am I missing something? Is everything all right?

[jpivarski]

I answered this on Gitter, but I should copy our conversation here so that it's useful to others in the future.

Jim:

Hi @a-akram! What's happening here is that Uproot follows ROOT's way of adding same-name objects to a directory: it gives each one a different "cycle number" to distinguish them. If root_file were a Python dict, each df would overwrite the previous one, but ROOT keeps them all, and the cycle number is like a version number.

But I think it's likely that you want to combine all of the dfs into one big TTree. In that case, use assignment (what you have above) to get the first one in, which defines the Tree, and all subsequent dfs can fill the TTree using uproot.WritableTree.extend.

In writing these scripts, I end up creating a first_time boolean to switch between assignment on the first time and extend on all other times through the loop.

Adeel:

Thanks for clarification. It not clear to me how to apply the extend() here. In above code example, is the statement root_file["TrackML"].extend(df) will work, ofcourse that will used from the second iteration?

(later)

Thanks it perfectly worked. I have another question. Is there a way to store TClonesArray?

Jim:

There is not a way to write the Python data as a specified ROOT/C++ type, such as TClonesArray. Sorry! (Each different ROOT/C++ type is a different serialization format, so it would essentially be a new project for each one. TClonesArray has a particularly odd serialization, as seen in the code that reads TClonesArray.)