scripts/imaging_calorimeter/utils.py

0001 '''
0002     A utility script to help the analysis of imaging calorimeter data
0003
0004     Author: Chao Peng (ANL)
0005     Date: 04/30/2021
0006 '''
0007
0008 import os
0009 import ROOT
0010 import numpy as np
0011 import pandas as pd
0012 import matplotlib
0013 import DDG4
0014 from ROOT import gROOT, gInterpreter
0015
0016
0017 # helper function to truncate color map (for a better view from the rainbow colormap)
0018 def truncate_colormap(cmap, minval=0.0, maxval=1.0, n=100):
0019     new_cmap = matplotlib.colors.LinearSegmentedColormap.from_list(
0020         'trunc({n},{a:.2f},{b:.2f})'.format(n=cmap.name, a=minval, b=maxval),
0021         cmap(np.linspace(minval, maxval, n)))
0022     return new_cmap
0023
0024
0025 # load root macros, input is an argument string
0026 def load_root_macros(arg_macros):
0027     for path in arg_macros.split(','):
0028         path = path.strip()
0029         if os.path.exists(path):
0030             gROOT.Macro(path)
0031         else:
0032             print('\"{}\" does not exist, skip loading it.'.format(path))
0033
0034
0035 # read mc particles from root file
0036 def get_mcp_data(path, evnums=None, branch='mcparticles2'):
0037     f = ROOT.TFile(path)
0038     events = f.events
0039     if evnums is None:
0040         evnums = np.arange(events.GetEntries())
0041     elif isinstance(evnums, int):
0042         evnums = [evnums]
0043
0044     dbuf = np.zeros(shape=(5000*len(evnums), 6))
0045     idb = 0
0046     for iev in evnums:
0047         if iev >= events.GetEntries():
0048             print('Error: event {:d} is out of range (0 - {:d})'.format(iev, events.GetEntries() - 1))
0049             continue
0050
0051         events.GetEntry(iev)
0052         # extract full mc particle data
0053         for part in getattr(events, branch):
0054             dbuf[idb] = (iev, part.psx, part.psy, part.psz, part.pdgID, part.status)
0055             idb += 1
0056     return pd.DataFrame(data=dbuf[:idb], columns=['event', 'px', 'py', 'pz', 'pid', 'status'])
0057
0058
0059 # read mc particles from root file
0060 def get_mcp_simple(path, evnums=None, branch='mcparticles2'):
0061     f = ROOT.TFile(path)
0062     events = f.events
0063     if evnums is None:
0064         evnums = np.arange(events.GetEntries())
0065     elif isinstance(evnums, int):
0066         evnums = [evnums]
0067
0068     dbuf = np.zeros(shape=(len(evnums), 6))
0069     idb = 0
0070     for iev in evnums:
0071         if iev >= events.GetEntries():
0072             print('Error: event {:d} is out of range (0 - {:d})'.format(iev, events.GetEntries() - 1))
0073             continue
0074
0075         events.GetEntry(iev)
0076         # extract full mc particle data
0077         part = getattr(events, branch)[2]
0078         dbuf[idb] = (iev, part.psx, part.psy, part.psz, part.pdgID, part.status)
0079         idb += 1
0080     return pd.DataFrame(data=dbuf[:idb], columns=['event', 'px', 'py', 'pz', 'pid', 'status'])
0081
0082
0083 # read hits data from root file
0084 def get_hits_data(path, evnums=None, branch='RecoEcalBarrelHits'):
0085     f = ROOT.TFile(path)
0086     events = f.events
0087     if evnums is None:
0088         evnums = np.arange(events.GetEntries())
0089     elif isinstance(evnums, int):
0090         evnums = [evnums]
0091
0092     dbuf = np.zeros(shape=(2000*len(evnums), 7))
0093     idb = 0
0094     for iev in evnums:
0095         if iev >= events.GetEntries():
0096             print('Error: event {:d} is out of range (0 - {:d})'.format(iev, events.GetEntries() - 1))
0097             continue
0098
0099         events.GetEntry(iev)
0100         for hit in getattr(events, branch):
0101             dbuf[idb] = (iev, hit.clusterID, hit.layerID, hit.position.x, hit.position.y, hit.position.z, hit.edep)
0102             idb += 1
0103
0104     return pd.DataFrame(data=dbuf[:idb], columns=['event', 'cluster', 'layer', 'x', 'y', 'z', 'edep'])
0105
0106
0107 # read layers data from root file
0108 def get_layers_data(path, evnums=None, branch="EcalBarrelClustersLayers"):
0109     f = ROOT.TFile(path)
0110     events = f.events
0111     if evnums is None:
0112         evnums = np.arange(events.GetEntries())
0113     elif isinstance(evnums, int):
0114         evnums = [evnums]
0115
0116     dbuf = np.zeros(shape=(2000*len(evnums), 7))
0117     idb = 0
0118     for iev in evnums:
0119         if iev >= events.GetEntries():
0120             print('Error: event {:d} is out of range (0 - {:d})'.format(iev, events.GetEntries() - 1))
0121             continue
0122
0123         events.GetEntry(iev)
0124         for layer in getattr(events, branch):
0125             dbuf[idb] = (iev, layer.clusterID, layer.layerID,
0126                          layer.position.x, layer.position.y, layer.position.z, layer.edep)
0127             idb += 1
0128
0129     return pd.DataFrame(data=dbuf[:idb], columns=['event', 'cluster', 'layer', 'x', 'y', 'z', 'edep'])
0130
0131
0132 # read clusters data from root file
0133 def get_clusters_data(path, evnums=None, branch='EcalBarrelClustersReco'):
0134     f = ROOT.TFile(path)
0135     events = f.events
0136     if evnums is None:
0137         evnums = np.arange(events.GetEntries())
0138     elif isinstance(evnums, int):
0139         evnums = [evnums]
0140
0141     dbuf = np.zeros(shape=(20*len(evnums), 6))
0142     idb = 0
0143     for iev in evnums:
0144         if iev >= events.GetEntries():
0145             print('Error: event {:d} is out of range (0 - {:d})'.format(iev, events.GetEntries() - 1))
0146             continue
0147
0148         events.GetEntry(iev)
0149         for k, cl in enumerate(getattr(events, branch)):
0150             dbuf[idb] = (iev, k, cl.nhits, cl.edep, cl.cl_theta, cl.cl_phi)
0151             idb += 1
0152
0153     return pd.DataFrame(data=dbuf[:idb], columns=['event', 'cluster', 'nhits', 'edep', 'cl_theta', 'cl_phi'])
0154
0155
0156 def compact_constants(path, names):
0157     if not os.path.exists(path):
0158         print('Cannot find compact file \"{}\".'.format(path))
0159         return []
0160     kernel = DDG4.Kernel()
0161     description = kernel.detectorDescription()
0162     kernel.loadGeometry("file:{}".format(path))
0163     try:
0164         vals = [description.constantAsDouble(n) for n in names]
0165     except:
0166         print('Fail to extract values from {}, return empty.'.format(names))
0167         vals = []
0168     kernel.terminate()
0169     return vals
0170
0171
0172