Accessing a friend TTree using pyROOT

I have a pyROOT script where I use TChain::AddFriend to combine two TTrees:

from ROOT import TFile, gDirectory 

myfile = TFile("e.root")
mychain = gDirectory.Get("elec_gen")

friendFile = TFile("mu.root")
friendChain = gDirectory.Get("muon_gen")
mychain.AddFriend("muon_gen")

elec_gen_evtNum = mychain.evtNum
muon_gen_evtNum = mychain.muon_gen.evtNum

When I run this I get:

$ python friendTest.py 
Traceback (most recent call last):
  File "friendTest.py", line 11, in <module>
    muon_gen_evtNum = mychain.muon_gen.evtNum
AttributeError: 'TTree' object has no attribute 'muon_gen'

With the last line (muon_gen_evtNum = mychain.muon_gen.evtNum) commented-out it runs fine. So it appears I am not accessing the leaves of the friend tree (muon_gen) correctly. How do I access them?

I also tried combining the TTrees using

mychain.AddFriend("muon_gen","mu.root")

but this produces the same resut.

Hi,
sorry I am not an expert so I can’t give you a straight answer, but why are you expecting mychain to have a muon_gen field? Have you seen in an example (if so, can you point me to it?)

Hi,

See ROOT User’s Guide section entitled TChain::AddFriend.

Edit: to clarify, e.root contain the TTree elec_gen, mu.root contains the TTree muon_gen. When I add muon_gen as a Friend of elec_gen I believe should be able to access the muon_gen branches through mychain.

I don’t think what you are trying to do is supported. ROOT does not add attributes to python variables afaik.
The guide mentions the fact that inside a TTree::Draw string you can refer to the friend branches as <chainname>.<branchname>.<varname>.

Also note that the user’s guide is written with C++ in mind (and C++ example snippets), not python.

Does mychain.evtNum return the correct value?

Yes, mychain.evtNum returns the correct value. So it’s impossible to use AddFriend in pyROOT? Is there some equivalent functionality available in pyROOT?

I am really surprised mychain has a evtNum field. Are you sure you are not setting it yourself (or some library you are using)? When I construct a TChain in pyROOT it does not have that field:


from ROOT import TChain
mychain = TChain("file.root")
mychain
<ROOT.TChain object ("file.root") at 0x5bc4b10>
mychain.evtNum

AttributeErrorTraceback (most recent call last)
<ipython-input-11-76077e36a2d1> in <module>()
----> 1 mychain.evtNum

AttributeError: 'TChain' object has no attribute 'evtNum'

As per your real question: you can definitely call AddFriend, you can find the valid signatures here. This is very minimal but works (provided you have a tree “tree” in file “file.root” and a tree “tree2” in a file “file2.root”, and “tree2” has a branch named “d”):

from ROOT import TChain, TCanvas
c = TCanvas("c","c")

chain = TChain("tree")
chain.Add("file.root")
friend_chain = TChain("tree2")
friend_chain.Add("file2.root")
chain.AddFriend(friend_chain)
chain.Draw("d") # d is a branch contained in file2.root
c.Draw()

evtNum is the name of a branch in my elec_gen TTree (and my muon_gen TTree). If your file.root does not have a branch named evtNum then you will get an error. Maybe evtNum was a slightly confusing branch for me to choose for this question!

My problem is both of my files “e.root” and “mu.root” contain identically-named branches, e.g. both contain a branch called “PDG” (I will avoid using evtNum to avoid confusion!). Your example works when I just want to use Draw, but I do not know how to specify a tree when I want to assign to a variable:


from ROOT import TChain, TCanvas
c = TCanvas("c","c")

chain = TChain("elec_gen")
chain.Add("e.root")
friend_chain = TChain("muon_gen")
friend_chain.Add("mu.root")
chain.AddFriend(friend_chain)
chain.Draw("PDG") # PDG is a branch contained in e.root and in mu.root
c.Draw() # Draws PDG from elec_gen
chain.Draw("muon_gen.PDG")  # Draws PDG from muon_gen
chain.Draw("elec_gen.PDG") # Draws PDG from elec_gen

# So everything up to here works.
# Now I want to assign PDG to a variable:

elec_PDG = chain.PDG # assigns elec_gen.PDG to elec_PDG, I presume
muon_PDG = chain.muon_gen.PDG # I want assign muon_gen.PDG to muon_PDG
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: 'TChain' object has no attribute 'muon_gen'

muon_PDG = muon_gen.PDG # Try different syntax
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
NameError: name 'muon_gen' is not defined

So I still have the same problem from my original question.

Maybe @dpiparo can help :slight_smile:

Hi,

I am confused: if the trees have the same branches, why do you want to create a friend? I think you just want to have a TChain to process them all together, right? If yes, I propose a solution below. If no, I need more context.
Given that your trees have different names, there is an intermediate step to perform here which will have the same result.
So, I create two files with two trees inside with different names and identical branches, just one branch called evt_num which contains an integer running from 0 to 9*. This is to mock your dataset.
Then to process it in Python:

import ROOT

efile = ROOT.TFile("e.root")
elec_gen = efile.elec_gen

mufile = ROOT.TFile("mu.root")
mu_gen = mufile.mu_gen

for tree in (mu_gen, elec_gen):
   for event in tree:
      print event.evt_num

Cheers,
D

void create(){

   ROOT::Experimental::TDataFrame d(10);
   int i=0;
   d.Define("evt_num",[&i](){return i++;}).Snapshot("elec_gen","e.root");

   i=0;
   ROOT::Experimental::TDataFrame d2(10);
   d2.Define("evt_num",[&i](){return i++;}).Snapshot("mu_gen","mu.root");
}

Hi @dpiparo, thanks for coming to help.

The context is I have physics events split up between different files. So e.root and mu.root contain the same events, but different information from those events (e.root contains information about electrons, mu.root about muons, etc). So event == 10 in elec_gen is the same event as event == 10 in muon_gen, the trees just contain different information about that event.

So I want to loop over these trees simultaneously, so that I can access and manipulate all the information from a single event at once. AddFriend seems like it is designed for exactly this (maybe there is a better way?), but I am having difficulty specifying which tree I want to use data from. I am using some python code that is similar to a C++ TSelector:

from ROOT import TFile, TChain

mychain = TChain("elec_gen")
mychain.Add("e.root")
entries = mychain.GetEntriesFast()

friend_chain = TChain("muon_gen")
friend_chain.Add("mu.root")
mychain.AddFriend(friend_chain)

# Loop over entries (i.e. events)
for jentry in xrange( entries ):
   # Get the next tree in the chain and verify.
   ientry = mychain.LoadTree( jentry )
   if ientry < 0:
      break
                                                                                                                                                                                                 
   # Copy next entry into memory and verify.
   nb = mychain.GetEntry( jentry )
   if nb <= 0:
      continue

   # Now I want to access and manipulate the data in my files...
   PDG = mychain.PDG # this will default to elec_gen.PDG
   elec_PDG = mychain.elec_gen.PDG # try to access specific tree, does not work
   muon_PDG = mychain.muon_gen.PDG # try to access specific tree, does not work

   elec_cands = mychain.elec_gen.nCands # try to access specific tree, does not work
   muon_cands = mychain.muon_gen.nCands # try to access specific tree, does not work

   # ...so that I can do analysis, e.g.
   for muon_cand in range (0, muon_cands):
      if muon_PDG[muon_cand] == 13:
           # do something

When I run this I get errors like:

$ python friendTest.py
Traceback (most recent call last):
  File "friendTest.py", line 26, in <module>
    elec_PDG = mychain.elec_gen.PDG
AttributeError: 'TChain' object has no attribute 'elec_gen'

Hi,

now I understand.You are right, Friend trees are your friends in this case. But are the branches of these trees called the same?

Cheers,
D

Hi,

Yes, the branches of the trees have the same names.

Hi,

I am not sure this mix is possible in PyROOT. I need to look into this further.
As a workaround, would transforming the branches names recreating the datasets from the ones you have an option? How big are they?

Cheers,
D

Hi,

For now I am working with a small amount of data (only a few MB) just to test my code etc. But eventually it will be much bigger and there will be more file types (not just electron and muon), so changing all the branch names would not be ideal.

I have been trying some other options and I realised I can just load matching entries from the elec_gen and muon_gen trees simultaneously, without using AddFriend at all:

from ROOT import TFile, TChain

# Load trees as two separate trees. Note I do not use AddFriend.

elec_chain = TChain("elec_gen")
elec_chain.Add("e.root")

muon_chain = TChain("muon_gen")
muon_chain.Add("mu.root")

entries = elec_chain.GetEntriesFast()

# Loop over entries

for jentry in xrange( entries ):
   # Get the next tree in the chain and verify.
   ientry = elec_chain.LoadTree( jentry )
   if ientry < 0:
      break

   # Copy next entry into memory and verify.
   nb = elec_chain.GetEntry( jentry )
   if nb <= 0:
      continue

   # Load matching entry from muon_gen tree
   muon_entry = muon_chain.GetEntry(jentry)

   # Now I can access both trees:
   elec_PDG = elec_chain.PDG
   muon_PDG = muon_chain.PDG

   elec_cands = elec_chain.nCands
   muon_cands = muon_chain.nCands

   # etc

I have done a quick bit of testing and this seems to work. Is there any problem with doing this? Does AddFriend do anything that my method above does not do?

Hi,

good that you found a solution!
I think the problem is linked to the way data is structured and we cannot do much at ROOT level.

Cheers,
D

It’s a bit surprising there is no way to use Friend Trees in pyROOT, but maybe it is just not necessary since we can achieve the same functionality using my solution.

Anyway, thanks for your help @dpiparo and @eguiraud!

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.