add field id decoder

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']])