runLayerRecalibration.py

Go to the documentation of this file.

#!/usr/bin/env python
 
# Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
 
""" To compute corrected layer energies """
 
from glob import glob
import os.path
import os
from array import array
from itertools import izip
import logging
logging.basicConfig(level=logging.DEBUG)
 
import ROOT
ROOT.PyConfig.IgnoreCommandLineOptions = True
 
try:
    ROOT.gROOT.Macro("$ROOTCOREDIR/scripts/load_packages.C")
except RuntimeError:
    print("""you have to compile egammaLayerRecalibTool first:
rc find_packages
rc compile
""")
 
 
def LayerRecalibratedGenerator(input_tree, branches):
    ""
    tool = ROOT.egammaLayerRecalibTool("2012_alt_with_layer2")
 
    def _rec(eta, e0, e1, e2, e3):
        std_calibration_input = ROOT.StdCalibrationInputs(0, eta, e0, e1, e2, e3, 0)
        tool.scale_inputs(std_calibration_input)
        return std_calibration_input.E0raw, std_calibration_input.E1raw, std_calibration_input.E2raw, std_calibration_input.E3raw
    fcn = _rec
    entries = input_tree.GetEntries()
    for i in range(entries):
        input_tree.GetEntry(i)
        # check if branch exists
        try:
            inputs = list(map(input_tree.__getattr__, branches))
        except AttributeError:
            raise AttributeError("cannot find branches %s" % branches)
 
        try:
            yield fcn(*inputs)  # scalar
        except TypeError:
            result = [fcn(*ii) for ii in izip(*inputs)]  # vectors
            yield list(zip(*result))
 
if __name__ == "__main__":
    from optparse import OptionParser
 
    def foo_callback(option, opt, value, parser):
        setattr(parser.values, option.dest, value.split(','))
 
    parser = OptionParser()
    parser.add_option("-i", "--input", help="input file or directory")
    parser.add_option("--nevents", type=int, help="number of events")
    parser.add_option("-t", "--tree", help="input tree name")
    parser.add_option("-p", "--particle", type="choice", help="particle type (electron|photon)", choices=("electron","photon"))
    parser.add_option("-b", "--branches-to-copy", type='string',
                      action='callback',
                      callback=foo_callback,
                      help=r"comma separated list of branches to copy from the input tree, you can use \*, do not use spaces")
 
    (options, args) = parser.parse_args()
 
    input_filenames = None
    if os.path.isdir(options.input):
        input_filenames = glob(os.path.join(options.input, "*.root*"))
    else:
        input_filenames = [options.input]
    if not input_filenames:
        print("ERROR: no files found")
 
    if not options.particle:
        raise ValueError("you need to specify a value of --particle")
 
    logging.info("found %s files", len(input_filenames))
 
    try:
        ROOT.egammaLayerRecalibTool
    except Exception:
        raise AttributeError('You have to load egammaLayerRecalibTool package')
 
    prefix = 'el_' if options.particle == 'electron' else 'ph_'
 
    for filename in input_filenames:
        print("patching file %s" % filename)
 
        f = ROOT.TFile.Open(filename, "update")
        tree = f.Get(options.tree)
        if not tree:
            raise IOError("cannot find %s" % options.tree)
 
 
        branches = [prefix + 'cl_eta'] + [prefix + 'rawcl_Es%d' % i for i in range(4)]
        if not all(tree.GetBranch(i) for i in branches):
            logging.info("skipping %s*", prefix)
            continue
 
        branches_to_copy = options.branches_to_copy or []
        logging.info("copying these branches: %s", ",".join(branches_to_copy) if branches_to_copy else "None")
 
        new_branch_names = [prefix + 'rawcl_Es%d' % i for i in range(4)]
        if set(new_branch_names).intersection(branches_to_copy):
            logging.warning("raw_clE* variables in the copy list -> going to rename. PROBABILY YOU DON'T WANT TO DO IT!")
            new_branch_names = [prefix + 'rawcl_Es%d_new' % i for i in range(4)]
 
        tree.SetBranchStatus("*", 0)
        for branch in branches_to_copy:
            tree.SetBranchStatus(branch, 1)
        newtree = tree.CloneTree(0)
        newtree.SetName("layerRecalibratedTree")
        newtree.SetTitle("layerRecalibratedTree")
 
        # Disable all branches but the needed ones
        tree.SetBranchStatus('*', 0)
        for br in set(branches + branches_to_copy):
            logging.info("activating branch %s", br)
            tree.SetBranchStatus(br, 1)
 
        # Create new branch, the same type as the old one
        isVector = 'vector' in tree.GetLeaf(prefix + 'rawcl_Es0').GetTypeName()
        new_storages = []
        if isVector:
            for new_branch_name in new_branch_names:
                new_storage = ROOT.vector(float)()
                new_storages.append(new_storage)
                newtree.Branch(new_branch_name, new_storage)
        else:
            for new_branch_name in new_branch_names:
                new_storage = array('f', [0.])
                new_storages.append(new_storage)
                newtree.Branch(new_branch_name, new_storage, new_branch_name + '/F')
 
        for i, new_value in enumerate(LayerRecalibratedGenerator(tree, branches)):
            if i % 1000 == 0:
                print(i)
            if options.nevents and i >= options.nevents:
                break
 
            if isVector:
                for s, d in zip(new_storages, new_value):
                    s.clear()
                    map(s.push_back, d)
            else:
                for s, v in zip(new_storages, new_value):
                    s[0] = v
            newtree.Fill()
        print("loop ended after %d events" % i)
 
        logging.debug("all events done")
        tree.SetBranchStatus('*', 1)
        logging.debug("adding friend")
        tree.AddFriend(newtree)
        logging.info("updating old tree")
        tree.GetCurrentFile().Write("", ROOT.TObject.kOverwrite)
        logging.info("writing new tree")
        if f.Get("layerRecalibratedTree"):
            newtree.GetCurrentFile().Write("", ROOT.TObject.kOverwrite)
        else:
            newtree.GetCurrentFile().Write()
        logging.debug("closing file")
 
        # this is because ROOT is stupid
        tree.SetDirectory(0)
        newtree.SetDirectory(0)
 
        f.Close()
        logging.debug("closed")
 
        logging.info("creating check plots")
        chain = ROOT.TChain(options.tree)
        for filename in input_filenames:
            chain.Add(filename)
 
        for old_branch_name, new_branch_name in zip([prefix + 'rawcl_Es%d' % i for i in range(4)], new_branch_names):
            canvas = ROOT.TCanvas()
            chain.Draw("layerRecalibratedTree.%s/%s:%scl_eta>>(100, -2.5, 2.5, 100, 0.9, 1.1)" % (new_branch_name, old_branch_name, prefix),
                       "abs(%scl_eta) < 2.5 && %s != 0" % (prefix, old_branch_name), "", options.nevents or 1000000000)
            canvas.SaveAs("canvas_check_%s.png" % new_branch_name)