BunchGroupSet.py

Go to the documentation of this file.

# Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
 
from itertools import groupby
 
from AthenaCommon.Logging import logging
from AthenaConfiguration.Enums import BeamType
 
from .Limits import Limits
from .L1MenuFlags import L1MenuFlags
 
 
log = logging.getLogger(__name__)
 
def createDefaultBunchGroupSetMC():
    """
    create BunchGroupSet for simulation
 
    This BGS is independent from the menu, and contains only BG0 (BCRVeto) and BG1 (Paired)
    """
    bgs = BunchGroupSet("MC")
    bgs.addBunchGroup( BunchGroupSet.BunchGroup(name = 'BCRVeto', internalNumber = 0).addRange(0,3539).addRange(3561,3563).normalize() )\
       .addBunchGroup( BunchGroupSet.BunchGroup(name = 'Paired', internalNumber = 1).addTrain(0,3564).normalize() )
    return bgs
 
 
def createDefaultBunchGroupSet(flags):
    """
    sets default bunchgroups for all menus, needed for simulation.
    """
    if hasattr(L1MenuFlags, "BunchGroupNames"): # if flag has been set
        # if menu defines bunchgroup names, then we generate a bunchgroup set from that
        name = L1MenuFlags.MenuSetup().partition('_')[0]
        bgs = BunchGroupSet(name)
        bgs.addBunchGroup( BunchGroupSet.BunchGroup(name = 'BCRVeto', internalNumber = 0).addRange(0,3539).addRange(3561,3563).normalize() )\
           .addBunchGroup( BunchGroupSet.BunchGroup(name = 'Paired', internalNumber = 1).addTrain(0,3445).addTrain(3536,4).addTrain(3561,3).normalize() )
        if flags.Beam.Type is BeamType.Cosmics:
            bgs.addBunchGroup( BunchGroupSet.BunchGroup(name = 'EMPTY', internalNumber = 3).addTrain(0,3564).normalize() )
 
        bunchgroupnames = L1MenuFlags.BunchGroupNames()[:Limits.NumBunchgroups]
        bunchgroupnames += ['NotUsed'] * (Limits.NumBunchgroups - len(bunchgroupnames))
        for i,bgname in enumerate(bunchgroupnames):
            bgs.bunchGroups[i].name = bgname
    else:
        bgs = createDefaultBunchGroupSetMC()
 
    return bgs
 
 
class BunchGroupSet(object):
 
    class BunchGroup(object):
        nameUndefBG = "NotUsed"
        def __init__(self, internalNumber, name = nameUndefBG, partition=0, bunches=None):
            # check if internalNumber is within limits
            if internalNumber<0 or internalNumber >= Limits.NumBunchgroups:
                raise RuntimeError('Cannot create bunchgroup with internal number %i (must be between 0 and %i)' % (internalNumber, Limits.NumBunchgroups-1))
            self.name           = name # name of the bunchgroup
            self.internalNumber = internalNumber # number between 0 and 15
            self.menuPartition  = partition # CTP partition that the bunchgroup belongs to
            self.bunches        = bunches if bunches else [] # list of bcids
            self.normalized     = [] # list of trains (filled in normalized())
        def __str__(self):
            return "bunchgroup %s(%i) with %i bunches" % (self.name, self.internalNumber, len(self))
        def __len__(self):
            return len(self.bunches)
        def addBunch(self, bcid):
            self.bunches += [ bcid ]
            return self
        def addRange(self, firstbcid, lastbcid):
            self.bunches += list(range(firstbcid, lastbcid+1))
            return self
        def addTrain(self, firstbcid, length):
            self.bunches += list(range(firstbcid, firstbcid+length))
            return self
        def normalize(self):
            """ turn list of bunches into trains """
            self.normalized = []
            self.bunches = sorted(list(set(self.bunches))) # uniq and sort
            if any(map(lambda bcid : bcid<0 or bcid >= 3564, self.bunches)):
                raise RuntimeError("Found bcid outside range 0..3563 in bunchgroup %s" % self.name)
            for k,g in groupby(enumerate(self.bunches), lambda x : x[1]-x[0]):
                train = list(g)
                self.normalized += [ (train[0][1], len(train)) ]
            return self
        def json(self):
            confObj = {
                "name":  self.name,
                "id":    self.internalNumber,
                "info":  f"{len(self)} bunches, {len(self.normalized)} groups",
                "bcids": [ {"first": first, "length": length} for first, length in self.normalized ]
            }
            return confObj
        
 
    def __init__(self, name='', menuPartition=0):
        self.name          = name
        self.menuPartition = menuPartition
        self.bunchGroups   = [ BunchGroupSet.BunchGroup(i) for i in range(0,Limits.NumBunchgroups) ] # initialize all BunchGroups
 
    def __len__(self):
        return len(self.bunchGroups)
 
 
    @classmethod
    def partitioning(cls):
        first = L1MenuFlags.BunchGroupPartitioning()
        last = first[1:] + [ Limits.NumBunchgroups ]
        partitioning = dict( zip([1,2,3],zip(first,last)) )
        return partitioning
 
 
    def resize(self, newsize):
        if type(self.bunchGroups) is not list:
            self.bunchGroups = []
        self.bunchGroups += newsize * [None]
        self.bunchGroups = self.bunchGroups[:newsize]
 
 
    def addBunchGroup(self, bunchGroup):
        # check if one already exists with this number
        doesExist = self.bunchGroups[bunchGroup.internalNumber].name != BunchGroupSet.BunchGroup.nameUndefBG
        if doesExist:
            raise RuntimeError("Adding bunchgroup with internal number %i, which already exists" % bunchGroup.internalNumber)
 
        partition=0
        if hasattr(L1MenuFlags, "BunchGroupPartitioning"):
            for lowestBG in L1MenuFlags.BunchGroupPartitioning():
                if bunchGroup.internalNumber >= lowestBG:
                    partition += 1
        bunchGroup.partition = partition
        self.bunchGroups[bunchGroup.internalNumber] = bunchGroup
        return self
 
 
    def json(self):
        confObj = {f"BGRP{bg.internalNumber}": bg.json() for bg in self.bunchGroups}
        return confObj
 
    def writeJSON(self, outputFile, destdir="./", pretty=True):
        outputFile = destdir.rstrip('/') + '/' + outputFile
        confObj = {
            "name": self.name,
            "filetype": "bunchgroupset",
            "bunchGroups": self.json()
        }
        with open( outputFile, mode="wt" ) as fh:
            import json
            json.dump(confObj, fh, indent = 4 if pretty else None, separators=(',', ': '))
        log.info("Wrote %s", outputFile)
        return outputFile