Line Code
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144
##########################################################################
# Classes which provide the geometry information.
##

from builtins import range
import itertools
import os

import ROOT
ROOT.PyConfig.IgnoreCommandLineOptions = True
ROOT.gROOT.SetBatch()

import Alignment.MillePedeAlignmentAlgorithm.mpsvalidate.geometrydata as mpsv_geometrydata


class Alignables:
    """ Creates a list of the aligned strucutres. Get the fields out of the
    TrackerTree.root file.
    """

    def __init__(self, config):
        # list of Structure objects, contains structures which were aligned
        self.structures = []
        self.config = config

    def get_subdetid(self, objid):
        return mpsv_geometrydata.data[objid].subdetid

    def get_discriminator(self, objid):
        return mpsv_geometrydata.data[objid].discriminator

    def get_ndiscriminator(self, objid):
        subdetid = self.get_subdetid(objid)
        discriminator = self.get_discriminator(objid)
        ndiscriminator = {key: [] for key in discriminator}
        # open TrackerTree.root file
        treeFile = ROOT.TFile(os.path.join(self.config.jobDataPath,
                                           ".TrackerTree.root"))
        tree = treeFile.Get("TrackerTreeGenerator/TrackerTree/TrackerTree")

        for entry in tree:
            # check if entry is part of the structure
            if (entry.SubdetId == subdetid):
                for structure in discriminator:
                    ndiscriminator[structure].append(getattr(entry, structure))
        for structure in discriminator:
            ndiscriminator[structure] = [x for x in ndiscriminator[structure] if x != 0]

        return [len(set(ndiscriminator[structure]))
                for structure in discriminator]

    def create_list(self, MillePedeUser):
        # loop over output TTree
        for entry in MillePedeUser:
            # check which structures were aligned
            if (entry.ObjId != 1 and 999999 not in map(abs, entry.Par)):
                # check if structure is not yet in the list
                if not any(x.name == str(entry.Name) for x in self.structures):
                    # create new structure object
                    name = str(entry.Name)
                    subdetid = self.get_subdetid(entry.ObjId)
                    discriminator = self.get_discriminator(entry.ObjId)
                    ndiscriminator = self.get_ndiscriminator(entry.ObjId)
                    # create structure
                    self.structures.append(
                        Structure(name, subdetid, discriminator, ndiscriminator))
                    # add detids which belong to this structure
                    self.structures[-1].detids = self.get_detids(subdetid)

    def create_children_list(self):
        for struct in self.structures:
            # loop over discriminators -> create patterns
            # pattern {"half": 2, "side": 2, "layer": 6, ...}
            ranges = struct.ndiscriminator
            pranges = [list(range(1, x+1)) for x in ranges]
            # loop over all possible combinations of the values of the
            # discriminators
            for number in itertools.product(*pranges):
                # create pattern dict
                pattern = dict(zip(map(lambda x: x.lower(), struct.discriminator), number))
                # name out of pattern
                name = " ".join("{0} {1}".format(key.lower(), value)
                                for (key, value) in pattern.items())
                # get detids of child
                detids = self.get_detids(struct.subdetid, pattern)
                # create child and add it to parent
                child = Structure(name, struct.subdetid, detids=detids)
                struct.children.append(child)


    def get_detids(self, subdetid, pattern={}):
        # list of all detids in the structure
        detids = []
        # open TrackerTree.root file
        treeFile = ROOT.TFile(os.path.join(self.config.jobDataPath,
                                           ".TrackerTree.root"))
        tree = treeFile.Get("TrackerTreeGenerator/TrackerTree/TrackerTree")

        for entry in tree:
            # check if entry is part of the structure
            if (entry.SubdetId == subdetid):
                # to create a child also check the pattern
                structure_found = False
                for structure in ("Half", "Side", "Layer", "Rod", "Ring",
                                  "Petal", "Blade", "Panel", "OuterInner",
                                  "Module"):
                    if structure.lower() in pattern:
                        if getattr(entry, structure) != pattern[structure.lower()]:
                            structure_found = True
                            break
                if structure_found: continue

                detids.append(entry.RawId)
        return detids


class Structure:
    """ A object represents a physical strucutre
    """

    def __init__(self, name, subdetid, discriminator=[], ndiscriminator=[], detids=[]):
        # name of the structure
        self.name = name
        # fields to identify the DetIds which belong to the structure
        self.subdetid = subdetid
        # fields which allow to discriminate the parts of the structure
        self.discriminator = discriminator
        # number per discriminator
        self.ndiscriminator = ndiscriminator
        # all DetIds which belong to this structure
        self.detids = detids
        # fieldss of all parts of the structure
        self.children = []

    def get_name(self):
        return self.name

    def get_children(self):
        return self.children

    def contains_detid(self, detid):
        if detid in self.detids:
            return True
        return False