FwkIO/scripts/DQMIO2histo.py

0001 #!/bin/env python3
0002
0003 """
0004 Script converting DQM I/O format input file into folder structured ROOT file.
0005 Ouput files historgrams are easy browseable by ROOT. When there are more than 1 run
0006 in input file it creates a Run X named folder for each run.
0007 Thanks for Marco Rovere for giving example script/class needed to browse DQM I/O
0008 formatted input.
0009 """
0010 from __future__ import print_function
0011
0012 from builtins import range
0013 import ROOT as R
0014 import sys
0015 import re
0016 import os
0017 import argparse
0018
0019 class DQMIO:
0020     """
0021     Class responsible for browsing the content of a DQM file produced
0022     with the DQMIO I/O framework of CMSSW
0023     """
0024     types=["Ints","Floats","Strings",  ##defined DQMIO types
0025          "TH1Fs","TH1Ss","TH1Ds",
0026          "TH2Fs", "TH2Ss", "TH2Ds",
0027          "TH3Fs", "TProfiles","TProfile2Ds", "kNIndicies"]
0028
0029     def __init__(self, input_filename, output_filename):
0030         self._filename = input_filename
0031         self._canvas = None
0032         self.f = R.TFile(output_filename, "RECREATE")
0033         self.already_defined = {"TProfiles" : False, "TProfile2Ds" : False,
0034                 "TH2Fs" : False, "TH2Ds" : False}
0035
0036         if os.path.exists(self._filename): #we try to open input file if fail
0037             self._root_file = R.TFile.Open(self._filename) #-> close script
0038             if args.debug:
0039                 print("## DEBUG ##:")
0040                 print("    Input: %s\n    Output: %s" % (input_filename,
0041                     output_filename))
0042
0043         else:
0044             print("File %s does not exists" % self._filename)
0045             sys.exit(1)
0046
0047     def print_index(self):
0048         """
0049         Loop over the complete index and dump it on the screen.
0050         """
0051         indices = self._root_file.Get("Indices")
0052         if args.debug:
0053             print("## DEBUG ##:")
0054             print("Run,\tLumi,\tType,\t\tFirstIndex,\tLastIndex")
0055             for i in range(indices.GetEntries()):
0056                 indices.GetEntry(i)
0057                 print('{0:4d}\t{1:4d}\t{2:4d}({3:s})\t\t{4:4d}\t{5:4d}'.format(
0058                     indices.Run, indices.Lumi, indices.Type,
0059                     DQMIO.types[indices.Type], indices.FirstIndex, indices.LastIndex))
0060
0061         for i in range(indices.GetEntries()):
0062             indices.GetEntry(i)
0063             if indices.Type < len(DQMIO.types):
0064                 self.write_to_file(self.types[indices.Type],
0065                     [indices.FirstIndex,indices.LastIndex], str(indices.Run))
0066
0067             else:
0068                 print("Unknown histogram type. Type numer: %s" % (indices.Type))
0069         self.f.Close()
0070
0071     def write_to_file(self, hist_type, index_range, run):
0072         """
0073         Method looping over entries for specified histogram type and
0074         writing to FullName path to output ROOT File
0075         """
0076         print("Working on: %s indexes: %s..%s" % (hist_type ,index_range[0],
0077             index_range[1]))
0078         t_tree = self._root_file.Get(hist_type)
0079         __run_dir = "Run %s" % (run)
0080         ###we set Branch for the needed type
0081         if hist_type == "TProfiles":
0082             if not self.already_defined["TProfiles"]:
0083                 R.gROOT.ProcessLine("TProfile* _tprof;")
0084                 self.already_defined["TProfiles"] = True
0085             t_tree.SetBranchAddress("Value", R._tprof)
0086             t_tree.GetEntry(index_range[0])
0087         elif hist_type == "TProfile2Ds":
0088             if not self.already_defined["TProfile2Ds"]:
0089                 R.gROOT.ProcessLine("TProfile2D* _tprof2d;")
0090                 self.already_defined["TProfile2Ds"] = True
0091             t_tree.SetBranchAddress("Value", R._tprof2d)
0092             t_tree.GetEntry(index_range[0])
0093         elif hist_type == "TH2Fs":
0094             if not self.already_defined["TH2Fs"]:
0095                 R.gROOT.ProcessLine("TH2F* _th2f;")
0096                 self.already_defined["TH2Fs"] = True
0097             t_tree.SetBranchAddress("Value", R._th2f)
0098             t_tree.GetEntry(index_range[0])
0099         elif hist_type == "TH2Ds":
0100             if not self.already_defined["TH2Ds"]:
0101                 R.gROOT.ProcessLine("TH2D* _th2d;")
0102                 self.already_defined["TH2Ds"] = True
0103             t_tree.SetBranchAddress("Value", R._th2d)
0104             t_tree.GetEntry(index_range[0])
0105
0106         for i in range(0,t_tree.GetEntries()+1): ##iterate on entries for specified type
0107             if i >= index_range[0] and i <= index_range[1]: ##if entries as in range:
0108                 t_tree.GetEntry(i)
0109                 name = str(t_tree.FullName)
0110                # print "  %s:  %s" % (i, name)
0111                 file_path = name.split("/")[:-1]            ##  same run/lumi histograms
0112                 __directory = "%s/%s" % (os.path.join("DQMData", __run_dir),
0113                     "/".join(file_path))
0114                 directory_ret = self.f.GetDirectory(__directory)
0115                 if not directory_ret:
0116                     self.f.mkdir(os.path.join(__directory))
0117                 self.f.cd(os.path.join(__directory))
0118                 if hist_type == "Strings":
0119                     construct_str = '<%s>s=%s</%s>' % (name.split("/")[-1:][0],
0120                         t_tree.Value, name.split("/")[-1:][0])
0121                     tmp_str = R.TObjString(construct_str)
0122                     tmp_str.Write()
0123                 elif hist_type == "Ints":
0124                     construct_str = '<%s>i=%s</%s>' % (name.split("/")[-1:][0],
0125                         t_tree.Value, name.split("/")[-1:][0])
0126                     tmp_str = R.TObjString(construct_str)
0127                     tmp_str.Write()
0128                 elif hist_type == "Floats":
0129                     construct_str = '<%s>f=%s</%s>' % (name.split("/")[-1:][0],
0130                         t_tree.Value, name.split("/")[-1:][0])
0131                     tmp_str = R.TObjString(construct_str)
0132                     tmp_str.Write()
0133                 else:
0134                     if hist_type in ["TProfiles", "TProfile2Ds", "TH2Fs", "TH2Ds"]:
0135                         if hist_type == "TProfiles": #if type is specific we write it.
0136                             R._tprof.Write()
0137                         elif hist_type == "TProfile2Ds":
0138                             R._tprof2d.Write()
0139                         elif hist_type == "TH2Fs":
0140                             R._th2f.Write()
0141                         elif hist_type == "TH2Ds":
0142                             R._th2d.Write()
0143                     else: #else we wirte Leafs Value which is a histogram
0144                         t_tree.Value.Write()
0145
0146 if __name__ == '__main__':
0147     parser = argparse.ArgumentParser()
0148     parser.add_argument("-in", "--input", help = "Input DQMIO ROOT file")
0149     parser.add_argument("-o", "--output", help = "Output filename",
0150         default = "DQMIO_converter_output.root")
0151     parser.add_argument("--debug", help = "Debug mode to spam you console",
0152         action = "store_true")
0153
0154     args = parser.parse_args()
0155     __in_file = args.input
0156     __out_file = args.output
0157     dqmio = DQMIO(__in_file, __out_file)
0158     dqmio.print_index()