python/mpsvalidate/geometry.py

0001 ##########################################################################
0002 # Classes which provide the geometry information.
0003 ##
0004
0005 from builtins import range
0006 import itertools
0007 import os
0008
0009 import ROOT
0010 ROOT.PyConfig.IgnoreCommandLineOptions = True
0011 ROOT.gROOT.SetBatch()
0012
0013 import Alignment.MillePedeAlignmentAlgorithm.mpsvalidate.geometrydata as mpsv_geometrydata
0014
0015
0016 class Alignables:
0017     """ Creates a list of the aligned strucutres. Get the fields out of the
0018     TrackerTree.root file.
0019     """
0020
0021     def __init__(self, config):
0022         # list of Structure objects, contains structures which were aligned
0023         self.structures = []
0024         self.config = config
0025
0026     def get_subdetid(self, objid):
0027         return mpsv_geometrydata.data[objid].subdetid
0028
0029     def get_discriminator(self, objid):
0030         return mpsv_geometrydata.data[objid].discriminator
0031
0032     def get_ndiscriminator(self, objid):
0033         subdetid = self.get_subdetid(objid)
0034         discriminator = self.get_discriminator(objid)
0035         ndiscriminator = {key: [] for key in discriminator}
0036         # open TrackerTree.root file
0037         treeFile = ROOT.TFile(os.path.join(self.config.jobDataPath,
0038                                            ".TrackerTree.root"))
0039         tree = treeFile.Get("TrackerTreeGenerator/TrackerTree/TrackerTree")
0040
0041         for entry in tree:
0042             # check if entry is part of the structure
0043             if (entry.SubdetId == subdetid):
0044                 for structure in discriminator:
0045                     ndiscriminator[structure].append(getattr(entry, structure))
0046         for structure in discriminator:
0047             ndiscriminator[structure] = [x for x in ndiscriminator[structure] if x != 0]
0048
0049         return [len(set(ndiscriminator[structure]))
0050                 for structure in discriminator]
0051
0052     def create_list(self, MillePedeUser):
0053         # loop over output TTree
0054         for entry in MillePedeUser:
0055             # check which structures were aligned
0056             if (entry.ObjId != 1 and 999999 not in map(abs, entry.Par)):
0057                 # check if structure is not yet in the list
0058                 if not any(x.name == str(entry.Name) for x in self.structures):
0059                     # create new structure object
0060                     name = str(entry.Name)
0061                     subdetid = self.get_subdetid(entry.ObjId)
0062                     discriminator = self.get_discriminator(entry.ObjId)
0063                     ndiscriminator = self.get_ndiscriminator(entry.ObjId)
0064                     # create structure
0065                     self.structures.append(
0066                         Structure(name, subdetid, discriminator, ndiscriminator))
0067                     # add detids which belong to this structure
0068                     self.structures[-1].detids = self.get_detids(subdetid)
0069
0070     def create_children_list(self):
0071         for struct in self.structures:
0072             # loop over discriminators -> create patterns
0073             # pattern {"half": 2, "side": 2, "layer": 6, ...}
0074             ranges = struct.ndiscriminator
0075             pranges = [list(range(1, x+1)) for x in ranges]
0076             # loop over all possible combinations of the values of the
0077             # discriminators
0078             for number in itertools.product(*pranges):
0079                 # create pattern dict
0080                 pattern = dict(zip(map(lambda x: x.lower(), struct.discriminator), number))
0081                 # name out of pattern
0082                 name = " ".join("{0} {1}".format(key.lower(), value)
0083                                 for (key, value) in pattern.items())
0084                 # get detids of child
0085                 detids = self.get_detids(struct.subdetid, pattern)
0086                 # create child and add it to parent
0087                 child = Structure(name, struct.subdetid, detids=detids)
0088                 struct.children.append(child)
0089
0090
0091     def get_detids(self, subdetid, pattern={}):
0092         # list of all detids in the structure
0093         detids = []
0094         # open TrackerTree.root file
0095         treeFile = ROOT.TFile(os.path.join(self.config.jobDataPath,
0096                                            ".TrackerTree.root"))
0097         tree = treeFile.Get("TrackerTreeGenerator/TrackerTree/TrackerTree")
0098
0099         for entry in tree:
0100             # check if entry is part of the structure
0101             if (entry.SubdetId == subdetid):
0102                 # to create a child also check the pattern
0103                 structure_found = False
0104                 for structure in ("Half", "Side", "Layer", "Rod", "Ring",
0105                                   "Petal", "Blade", "Panel", "OuterInner",
0106                                   "Module"):
0107                     if structure.lower() in pattern:
0108                         if getattr(entry, structure) != pattern[structure.lower()]:
0109                             structure_found = True
0110                             break
0111                 if structure_found: continue
0112
0113                 detids.append(entry.RawId)
0114         return detids
0115
0116
0117 class Structure:
0118     """ A object represents a physical strucutre
0119     """
0120
0121     def __init__(self, name, subdetid, discriminator=[], ndiscriminator=[], detids=[]):
0122         # name of the structure
0123         self.name = name
0124         # fields to identify the DetIds which belong to the structure
0125         self.subdetid = subdetid
0126         # fields which allow to discriminate the parts of the structure
0127         self.discriminator = discriminator
0128         # number per discriminator
0129         self.ndiscriminator = ndiscriminator
0130         # all DetIds which belong to this structure
0131         self.detids = detids
0132         # fieldss of all parts of the structure
0133         self.children = []
0134
0135     def get_name(self):
0136         return self.name
0137
0138     def get_children(self):
0139         return self.children
0140
0141     def contains_detid(self, detid):
0142         if detid in self.detids:
0143             return True
0144         return False