diff --git a/benchmarks/imaging_ecal/scripts/get_layerids.py b/benchmarks/imaging_ecal/scripts/get_layerids.py index 724572897bf860b63759f3442f487285f75300db..1264f532d54dbb439dd349b5700c69a248432379 100644 --- a/benchmarks/imaging_ecal/scripts/get_layerids.py +++ b/benchmarks/imaging_ecal/scripts/get_layerids.py @@ -2,13 +2,63 @@ A simple analysis script to extract some basic info of Monte-Carlo hits ''' import os -import DDG4 import ROOT import pandas as pd import numpy as np import argparse from matplotlib import pyplot as plt import matplotlib.ticker as ticker +from lxml import etree as ET + + +class AthenaDecoder: + def __init__(self, compact, readout): + self.readouts = self.getReadouts(compact) + self.changeReadout(readout) + + def changeReadout(self, readout): + self.fieldsmap = self.decomposeIDs(self.readouts[readout]) + + def get(self, idvals, field): + start, width = self.fieldsmap[field] + return np.bitwise_and(np.right_shift(idvals, start), (1 << width) - 1) + + @staticmethod + def getReadouts(path): + res = dict() + AthenaDecoder.__getReadoutsRecur(path, res) + return res + + @staticmethod + def __getReadoutsRecur(path, res): + if not os.path.exists(path): + print('Xml file {} not exist! Ignored it.'.format(path)) + return + lccdd = ET.parse(path).getroot() + readouts = lccdd.find('readouts') + if readouts is not None: + for readout in readouts.getchildren(): + ids = readout.find('id') + if ids is not None: + res[readout.attrib['name']] = ids.text + for child in lccdd.getchildren(): + if child.tag == 'include': + root_dir = os.path.dirname(os.path.realpath(path)) + AthenaDecoder.__getReadoutsRecur(os.path.join(root_dir, child.attrib['ref']), res) + + @staticmethod + def decomposeIDs(id_str): + res = dict() + curr_bit = 0 + for field_bits in id_str.split(','): + elements = field_bits.split(':') + field_name = elements[0] + bit_width = abs(int(elements[-1])) + if len(elements) == 3: + curr_bit = int(elements[1]) + res[field_name] = (curr_bit, bit_width) + curr_bit += bit_width + return res # read from RDataFrame and flatten a given collection, return pandas dataframe @@ -50,23 +100,30 @@ if __name__ == '__main__': help='Readout name for the hits collection') args = parser.parse_args() + # decoder + decoder = AthenaDecoder(args.compact, args.readout) + # get hits rdf_rec = ROOT.RDataFrame('events', args.rec_file) df = flatten_collection(rdf_rec, args.coll) df.rename(columns={c: c.replace(args.coll + '.', '') for c in df.columns}, inplace=True) # initialize dd4hep detector - kernel = DDG4.Kernel() - description = kernel.detectorDescription() - kernel.loadGeometry("file:{}".format(args.compact)) - - decoder = description.readout(args.readout).idSpec().decoder() - lindex = decoder.index('layer') - get_layer_id = np.vectorize(lambda cid: decoder.get(cid, lindex)) + # import DDG4 + # kernel = DDG4.Kernel() + # description = kernel.detectorDescription() + # kernel.loadGeometry("file:{}".format(args.compact)) + # decoder = description.readout(args.readout).idSpec().decoder() + # lindex = decoder.index('layer') + # get_layer_id = np.vectorize(lambda cid: decoder.get(cid, lindex)) - df.loc[:, 'layerID'] = get_layer_id(df['cellID'].astype(int).values) - print(df[['cellID', 'layerID', 'position.x', 'position.y', 'position.z', 'energy']]) + # df.loc[:, 'layerID'] = get_layer_id(df['cellID'].astype(int).values) + # print(df[['cellID', 'layerID', 'position.x', 'position.y', 'position.z', 'energy']]) # always terminate dd4hep kernel - kernel.terminate() + # kernel.terminate() + + # faster way to get layerids + df.loc[:, 'layerID'] = decoder.get(df['cellID'].values, 'layer') + print(df[['cellID', 'layerID', 'position.x', 'position.y', 'position.z', 'energy']])