SortedCollectionCreator.py

Go to the documentation of this file.

# Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
 
__author__ = "Marcin Nowak"
__doc__ = """
Create a sorted collection of Event references from a set of Athnea files containing EventInfoTags
(or from other compatible APR Event collections)
"""
 
from CollectionUtilities.GetCollectionType import getCollectionType
 
 
class SortedCollectionCreator:
   """Creates a sorted (APR) collection of event references from input files"""
   
   def __init__(self, name="sortEvents"):
      from AthenaCommon import Logging
      Logging.log.name = name
      self.name = name
      self.info = Logging.log.info
      self.debug = Logging.log.debug
      self.verbose = Logging.log.verbose
      self.collDescription = None
      self.collSvc = None
 
   def loadRoot(self):
      """import ROOT, create CollectionSvc"""
      import ROOT
      self.pool = ROOT.pool
      self.collSvc = self.pool.CollectionService()
 
   def readCollectionDescription(self, collection):
      """read Collection Description and remember it"""
      desc =  collection.description()
      self.debug("Reading Collection Description from {}".format(desc.name()))
      # read attributes' names and types from the description and remember them
      for an in range(0, desc.numberOfAttributeColumns()):
         attr = desc.attributeColumn(an)
         name = attr.name()
         self.attrNames.append(name)
         self.attrTypes[name] = attr.type()
      self.tokenName = desc.eventReferenceColumnName()
      # make a local copy of the description
      self.collDescription = self.pool.CollectionDescription( desc )
      # prevent overwriting input by accident
      self.collDescription.setName("")
 
 
   def readInputCollections(self, inputCollections):
      """read all input collections into memore"""
      self.collDescription = None
      self.allRows = []
      self.attrNames = []
      self.attrTypes = {}
      self.tokenName = None
      for inFileName in inputCollections:
         self.debug("Opening {}".format(inFileName))
         iColl = self.collSvc.open( "PFN:"+inFileName, getCollectionType(inFileName))
         self.debug("{} opened".format(inFileName))
         if self.collDescription is None:
            self.readCollectionDescription(iColl)
         self.info("Reading Events from {}".format(inFileName))
 
         query = iColl.newQuery()
         query.selectAll()
         cursor = query.execute()
      
         while cursor.next():
            row = cursor.currentRow()
            # put the token first in the attribute list, for convenience
            t = [ row.tokenList()[self.tokenName].toString() ]
            for nam in self.attrNames:
               t.append( row.attributeList()[nam].data[ self.attrTypes[nam] ]() )
            self.allRows.append(t)   
       
         iColl.close()
 
      for t in self.allRows:
         self.verbose( (t[1:], t[0]) )
      self.verbose('='*80)
      self.info("Finished reading input collections, total events read: {}".format(len(self.allRows)) )
 
   def sortEvents(self, sortAttrName, sortReverse=False):
      """sort the events based on an attribute name"""
      self.info("Sorting on attribute {}, sort order {}".format(sortAttrName, ("Descending" if sortReverse else "Ascending") ))
      # add 1 to offsets because the Ref is first
      attrPos = self.attrNames.index(sortAttrName)+1
      self.allRows.sort( key=lambda t: t[attrPos], reverse=sortReverse )
 
      for t in self.allRows:
         self.verbose( (t[1:], t[0]) )
      self.verbose('='*80)
 
   def writeCollection(self, outputCollection, outputCollectionType):
      """write sorted collection into a Collection file"""
      self.info("Writing Event collection {}".format(outputCollection))
      self.collDescription.setName(outputCollection)
      self.collDescription.setType(outputCollectionType)
      # create the output collection (file)
      dstColl = self.collSvc.create(self.collDescription)
      row = self.pool.CollectionRowBuffer()
      dstColl.initNewRow( row )
      for t in self.allRows:
         row.tokenList()[0].fromString( t[0] )
         for idx,nam in enumerate(self.attrNames):
            type = self.attrTypes[nam]
            row.attributeList()[nam].setValue[type]( t[idx+1] )
         dstColl.insertRow( row )
      
      dstColl.commit()
      dstColl.close()
       
   def execute(self, inputCollections, outputCollection="PFN:collection.root", sortAttribute="LumiBlockN",
               sortOrder="Ascending", outputCollectionType="RootCollection"):
      sort_opts = ("Ascending", "Descending")
      self.info("Executing SortedCollectionCreator, inputs={}, output='{}' ({}), sort by: {}, order: {}"
                .format(inputCollections, outputCollection, outputCollectionType, sortAttribute, sortOrder))
      if isinstance(inputCollections, str):
         inputs = [inputCollections]
      else:
         inputs = inputCollections
      if sortOrder.lower() not in [opt.lower() for opt in sort_opts]:
          raise Exception(self.name + ": Accepted sortOrder values are: " + str(sort_opts))
      sortReverse = ( sortOrder.lower()[0] == "d" )
      self.loadRoot()
      self.readInputCollections(inputs)
      self.sortEvents(sortAttribute, sortReverse)
      self.writeCollection(outputCollection, outputCollectionType)